""" Conversation thread management system for AgentPress. This module provides comprehensive conversation management, including: - Thread creation and persistence - Message handling with support for text and images - Tool registration and execution - LLM interaction with streaming support - Error handling and cleanup """ import json import uuid from typing import List, Dict, Any, Optional, Type, Union, AsyncGenerator, Tuple, Callable, Literal from services.llm import make_llm_api_call from agentpress.tool import Tool, ToolResult from agentpress.tool_registry import ToolRegistry from agentpress.response_processor import ( ResponseProcessor, ProcessorConfig ) from services.supabase import DBConnection from utils.logger import logger # Type alias for tool choice ToolChoice = Literal["auto", "required", "none"] class ThreadManager: """Manages conversation threads with LLM models and tool execution. Provides comprehensive conversation management, handling message threading, tool registration, and LLM interactions with support for both standard and XML-based tool execution patterns. """ def __init__(self): """Initialize ThreadManager.""" self.db = DBConnection() self.tool_registry = ToolRegistry() self.response_processor = ResponseProcessor( tool_registry=self.tool_registry, add_message_callback=self.add_message ) def add_tool(self, tool_class: Type[Tool], function_names: Optional[List[str]] = None, **kwargs): """Add a tool to the ThreadManager.""" self.tool_registry.register_tool(tool_class, function_names, **kwargs) async def add_message( self, thread_id: str, type: str, content: Union[Dict[str, Any], List[Any], str], is_llm_message: bool = False, metadata: Optional[Dict[str, Any]] = None ): """Add a message to the thread in the database. Args: thread_id: The ID of the thread to add the message to. type: The type of the message (e.g., 'text', 'image_url', 'tool_call'). content: The content of the message. Can be a dictionary, list, or string. It will be stored as JSONB in the database. is_llm_message: Flag indicating if the message originated from the LLM. Defaults to False (user message). metadata: Optional dictionary for additional message metadata. Defaults to None, stored as an empty JSONB object if None. """ logger.debug(f"Adding message of type '{type}' to thread {thread_id}") client = await self.db.client # Prepare data for insertion data_to_insert = { 'thread_id': thread_id, 'type': type, 'content': json.dumps(content) if isinstance(content, (dict, list)) else content, 'is_llm_message': is_llm_message, 'metadata': json.dumps(metadata or {}), # Ensure metadata is always a JSON object } try: result = await client.table('messages').insert(data_to_insert).execute() logger.info(f"Successfully added message to thread {thread_id}") except Exception as e: logger.error(f"Failed to add message to thread {thread_id}: {str(e)}", exc_info=True) raise async def get_messages(self, thread_id: str) -> List[Dict[str, Any]]: """Get all messages for a thread. Args: thread_id: The ID of the thread to get messages for. Returns: List of message objects. """ logger.debug(f"Getting messages for thread {thread_id}") client = await self.db.client try: result = await client.rpc('get_llm_formatted_messages', {'p_thread_id': thread_id}).execute() # Parse the returned data which might be stringified JSON if not result.data: return [] # Return properly parsed JSON objects messages = [] for item in result.data: if isinstance(item, str): try: parsed_item = json.loads(item) messages.append(parsed_item) except json.JSONDecodeError: logger.error(f"Failed to parse message: {item}") else: messages.append(item) return messages except Exception as e: logger.error(f"Failed to get messages for thread {thread_id}: {str(e)}", exc_info=True) return [] async def run_thread( self, thread_id: str, system_prompt: Dict[str, Any], stream: bool = True, temporary_message: Optional[Dict[str, Any]] = None, llm_model: str = "gpt-4o", llm_temperature: float = 0, llm_max_tokens: Optional[int] = None, processor_config: Optional[ProcessorConfig] = None, tool_choice: ToolChoice = "auto", ) -> Union[Dict[str, Any], AsyncGenerator]: """Run a conversation thread with LLM integration and tool execution. Args: thread_id: The ID of the thread to run system_prompt: System message to set the assistant's behavior stream: Use streaming API for the LLM response temporary_message: Optional temporary user message for this run only llm_model: The name of the LLM model to use llm_temperature: Temperature parameter for response randomness (0-1) llm_max_tokens: Maximum tokens in the LLM response processor_config: Configuration for the response processor tool_choice: Tool choice preference ("auto", "required", "none") Returns: An async generator yielding response chunks or error dict """ logger.info(f"Starting thread execution for thread {thread_id}") logger.debug(f"Parameters: model={llm_model}, temperature={llm_temperature}, max_tokens={llm_max_tokens}") try: # 1. Get messages from thread for LLM call messages = await self.get_messages(thread_id) # 2. Prepare messages for LLM call + add temporary message if it exists prepared_messages = [system_prompt] # Find the last user message index last_user_index = -1 for i, msg in enumerate(messages): if msg.get('role') == 'user': last_user_index = i # Insert temporary message before the last user message if it exists if temporary_message and last_user_index >= 0: prepared_messages.extend(messages[:last_user_index]) prepared_messages.append(temporary_message) prepared_messages.extend(messages[last_user_index:]) logger.debug("Added temporary message before the last user message") else: # If no user message or no temporary message, just add all messages prepared_messages.extend(messages) if temporary_message: prepared_messages.append(temporary_message) logger.debug("Added temporary message to the end of prepared messages") # 3. Create or use processor config if processor_config is None: processor_config = ProcessorConfig() logger.debug(f"Processor config: XML={processor_config.xml_tool_calling}, Native={processor_config.native_tool_calling}, " f"Execute tools={processor_config.execute_tools}, Strategy={processor_config.tool_execution_strategy}") # Check if native_tool_calling is enabled and throw an error if it is if processor_config.native_tool_calling: error_message = "Native tool calling is not supported in this version" logger.error(error_message) return { "status": "error", "message": error_message } # 4. Prepare tools for LLM call openapi_tool_schemas = None if processor_config.native_tool_calling: openapi_tool_schemas = self.tool_registry.get_openapi_schemas() logger.debug(f"Retrieved {len(openapi_tool_schemas) if openapi_tool_schemas else 0} OpenAPI tool schemas") # 5. Track this agent run in the database run_id = str(uuid.uuid4()) client = await self.db.client run_data = { 'id': run_id, 'thread_id': thread_id, 'status': 'running', 'started_at': 'now()', } await client.table('agent_runs').insert(run_data).execute() logger.debug(f"Created agent run record with ID: {run_id}") # 6. Make LLM API call logger.info("Making LLM API call") try: llm_response = await make_llm_api_call( prepared_messages, llm_model, temperature=llm_temperature, max_tokens=llm_max_tokens, tools=openapi_tool_schemas, tool_choice=tool_choice if processor_config.native_tool_calling else None, stream=stream ) logger.debug("Successfully received LLM API response") except Exception as e: # Update agent_run status to error await client.table('agent_runs').update({ 'status': 'error', 'error': str(e), 'completed_at': 'now()' }).eq('id', run_id).execute() logger.error(f"Failed to make LLM API call: {str(e)}", exc_info=True) raise # 7. Process LLM response using the ResponseProcessor if stream: logger.info("Processing streaming response") response_generator = self.response_processor.process_streaming_response( llm_response=llm_response, thread_id=thread_id, config=processor_config ) # Wrap the generator to update the agent_run when complete async def wrapped_generator(): responses = [] try: async for chunk in response_generator: responses.append(chunk) yield chunk # Update agent_run to completed when done await client.table('agent_runs').update({ 'status': 'completed', 'responses': json.dumps(responses), 'completed_at': 'now()' }).eq('id', run_id).execute() logger.debug(f"Updated agent run {run_id} to completed status") except Exception as e: # Update agent_run to error await client.table('agent_runs').update({ 'status': 'error', 'error': str(e), 'completed_at': 'now()' }).eq('id', run_id).execute() logger.error(f"Error in streaming response: {str(e)}", exc_info=True) raise return wrapped_generator() else: logger.info("Processing non-streaming response") try: response = await self.response_processor.process_non_streaming_response( llm_response=llm_response, thread_id=thread_id, config=processor_config ) # Update agent_run to completed await client.table('agent_runs').update({ 'status': 'completed', 'responses': json.dumps([response]), 'completed_at': 'now()' }).eq('id', run_id).execute() logger.debug(f"Updated agent run {run_id} to completed status") return response except Exception as e: # Update agent_run to error await client.table('agent_runs').update({ 'status': 'error', 'error': str(e), 'completed_at': 'now()' }).eq('id', run_id).execute() logger.error(f"Error in non-streaming response: {str(e)}", exc_info=True) raise except Exception as e: logger.error(f"Error in run_thread: {str(e)}", exc_info=True) return { "status": "error", "message": str(e) }