import logging from typing import Dict, Any, AsyncGenerator, Callable from agentpress.tool_parser import ToolParser from agentpress.tool_executor import ToolExecutor import asyncio class LLMResponseProcessor: """ Handles LLM response processing and tool execution management. This class manages both streaming and non-streaming responses from Language Models, coordinating tool execution timing and order. It maintains message state and handles the execution of tool calls either immediately or after collecting a complete response, with support for both parallel and sequential execution patterns. Attributes: thread_id (str): Identifier for the conversation thread tool_executor (ToolExecutor): Executor for handling tool calls tool_parser (ToolParser): Parser for processing LLM responses available_functions (Dict): Registry of available tool functions add_message (Callable): Callback to add messages to the thread update_message (Callable): Callback to update existing messages parallel_tool_execution (bool): Whether to execute tools in parallel or sequentially tool_calls_buffer (Dict): Buffer for storing incomplete tool calls during streaming processed_tool_calls (set): Set of already processed tool call IDs current_message (Dict): Current message being processed in streaming mode content_buffer (str): Buffer for accumulating content during streaming tool_calls_accumulated (list): List of tool calls accumulated during streaming message_added (bool): Flag to indicate if a message has been added to the thread """ def __init__( self, thread_id: str, tool_executor: ToolExecutor, tool_parser: ToolParser, available_functions: Dict, add_message_callback: Callable, update_message_callback: Callable ): self.thread_id = thread_id self.tool_executor = tool_executor self.tool_parser = tool_parser self.available_functions = available_functions self.add_message = add_message_callback self.update_message = update_message_callback # State tracking for streaming responses self.tool_calls_buffer = {} self.processed_tool_calls = set() self.content_buffer = "" self.tool_calls_accumulated = [] self.message_added = False async def process_stream( self, response_stream: AsyncGenerator, execute_tools: bool = True, immediate_execution: bool = True ) -> AsyncGenerator: """ Process streaming LLM response and handle tool execution. Yields chunks immediately as they arrive, while handling tool execution and message management in the background. """ pending_tool_calls = [] async def handle_message_management(chunk): # Accumulate content if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content: self.content_buffer += chunk.choices[0].delta.content # Parse and accumulate tool calls parsed_message, is_complete = await self.tool_parser.parse_stream( chunk, self.tool_calls_buffer ) if parsed_message and 'tool_calls' in parsed_message: self.tool_calls_accumulated = parsed_message['tool_calls'] # Handle message management and tool execution if chunk.choices[0].finish_reason or (self.content_buffer and self.tool_calls_accumulated): message = { "role": "assistant", "content": self.content_buffer } if self.tool_calls_accumulated: message["tool_calls"] = self.tool_calls_accumulated if not self.message_added: await self.add_message(self.thread_id, message) self.message_added = True else: await self.update_message(self.thread_id, message) # Handle tool execution if execute_tools and self.tool_calls_accumulated: new_tool_calls = [ tool_call for tool_call in self.tool_calls_accumulated if tool_call['id'] not in self.processed_tool_calls ] if new_tool_calls: if immediate_execution: results = await self.tool_executor.execute_tool_calls( tool_calls=new_tool_calls, available_functions=self.available_functions, thread_id=self.thread_id, executed_tool_calls=self.processed_tool_calls ) for result in results: await self.add_message(self.thread_id, result) self.processed_tool_calls.add(result['tool_call_id']) else: pending_tool_calls.extend(new_tool_calls) # Handle end of stream if chunk.choices[0].finish_reason: if not immediate_execution and pending_tool_calls: results = await self.tool_executor.execute_tool_calls( tool_calls=pending_tool_calls, available_functions=self.available_functions, thread_id=self.thread_id, executed_tool_calls=self.processed_tool_calls ) for result in results: await self.add_message(self.thread_id, result) self.processed_tool_calls.add(result['tool_call_id']) pending_tool_calls.clear() async for chunk in response_stream: # Start background task for message management and tool execution asyncio.create_task(handle_message_management(chunk)) # Immediately yield the chunk yield chunk async def process_response( self, response: Any, execute_tools: bool = True ) -> None: """ Process complete LLM response and execute tools. Handles non-streaming responses, parsing the complete response and executing any tool calls according to the configured execution strategy. Args: response: Complete response from the LLM execute_tools: Whether to execute identified tool calls """ try: assistant_message = await self.tool_parser.parse_response(response) await self.add_message(self.thread_id, assistant_message) if execute_tools and 'tool_calls' in assistant_message and assistant_message['tool_calls']: results = await self.tool_executor.execute_tool_calls( tool_calls=assistant_message['tool_calls'], available_functions=self.available_functions, thread_id=self.thread_id, executed_tool_calls=self.processed_tool_calls ) for result in results: await self.add_message(self.thread_id, result) logging.info(f"Tool execution result: {result}") except Exception as e: logging.error(f"Error processing response: {e}") response_content = response.choices[0].message.get('content', '') await self.add_message(self.thread_id, { "role": "assistant", "content": response_content or "" })