diff --git a/litellm/responses/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 7d53452c1c0..8801e561915 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -1,19 +1,10 @@ -from litellm._uuid import uuid -from typing import ( - TYPE_CHECKING, - Any, - Dict, - List, - Optional, - Union, - cast, -) +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast from litellm._logging import verbose_logger -from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, -) +from litellm._uuid import uuid +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( + BaseLiteLLMOpenAIResponseObject, MCPCallArgumentsDeltaEvent, MCPCallArgumentsDoneEvent, MCPCallCompletedEvent, @@ -38,22 +29,24 @@ async def create_mcp_list_tools_events( mcp_tools_with_litellm_proxy: List[ToolParam], user_api_key_auth: Any, base_item_id: str, - pre_processed_mcp_tools: List[Any] + pre_processed_mcp_tools: List[Any], ) -> List[ResponsesAPIStreamingResponse]: """Create MCP discovery events using pre-processed tools from the parent""" - + events: List[ResponsesAPIStreamingResponse] = [] - + try: # Extract MCP server names mcp_servers = [] for tool in mcp_tools_with_litellm_proxy: if isinstance(tool, dict) and "server_url" in tool: server_url = tool.get("server_url") - if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"): + if isinstance(server_url, str) and server_url.startswith( + "litellm_proxy/mcp/" + ): server_name = server_url.split("/")[-1] mcp_servers.append(server_name) - + # Emit list tools in progress event in_progress_event = MCPListToolsInProgressEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS, @@ -62,21 +55,21 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(in_progress_event) - + # Use the pre-processed MCP tools that were already fetched, filtered, and deduplicated by the parent filtered_mcp_tools = pre_processed_mcp_tools - + # Convert tools to dict format for the event mcp_tools_dict = [] for tool in filtered_mcp_tools: - if hasattr(tool, 'model_dump') and callable(getattr(tool, 'model_dump')): + if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")): # Type cast to help mypy understand this is safe after hasattr check mcp_tools_dict.append(cast(Any, tool).model_dump()) - elif hasattr(tool, '__dict__'): + elif hasattr(tool, "__dict__"): mcp_tools_dict.append(tool.__dict__) else: - mcp_tools_dict.append({"name": getattr(tool, 'name', str(tool))}) - + mcp_tools_dict.append({"name": getattr(tool, "name", str(tool))}) + # Emit list tools completed event completed_event = MCPListToolsCompletedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED, @@ -85,7 +78,7 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(completed_event) - + # Add output_item.done event with the actual tools list (matching OpenAI format) from litellm.types.llms.openai import OutputItemDoneEvent @@ -95,45 +88,50 @@ async def create_mcp_list_tools_events( first_tool = mcp_tools_with_litellm_proxy[0] if isinstance(first_tool, dict): server_label_value = first_tool.get("server_label", "") - server_label = str(server_label_value) if server_label_value is not None else "" - + server_label = ( + str(server_label_value) if server_label_value is not None else "" + ) + # Format tools for OpenAI output_item.done format formatted_tools = [] for tool in filtered_mcp_tools: tool_dict = { - "name": getattr(tool, 'name', 'unknown'), - "description": getattr(tool, 'description', ''), + "name": getattr(tool, "name", "unknown"), + "description": getattr(tool, "description", ""), "annotations": {"read_only": False}, } - + # Add input_schema if available - if hasattr(tool, 'inputSchema'): - tool_dict["input_schema"] = getattr(tool, 'inputSchema') - elif hasattr(tool, 'input_schema'): - tool_dict["input_schema"] = getattr(tool, 'input_schema') - + if hasattr(tool, "inputSchema"): + tool_dict["input_schema"] = getattr(tool, "inputSchema") + elif hasattr(tool, "input_schema"): + tool_dict["input_schema"] = getattr(tool, "input_schema") + formatted_tools.append(tool_dict) - + # Create the output_item.done event with MCP tools list output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": base_item_id, - "type": "mcp_list_tools", - "server_label": server_label, - "tools": formatted_tools - } + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": base_item_id, + "type": "mcp_list_tools", + "server_label": server_label, + "tools": formatted_tools, + } + ), ) events.append(output_item_done_event) - + verbose_logger.debug(f"Created {len(events)} MCP discovery events") - + except Exception as e: verbose_logger.error(f"Error creating MCP list tools events: {e}") import traceback + traceback.print_exc() - + # Emit failed event on error failed_event = MCPListToolsFailedEvent( type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED, @@ -142,37 +140,39 @@ async def create_mcp_list_tools_events( item_id=base_item_id, ) events.append(failed_event) - + # Still emit output_item.done event even on failure (with empty tools list) from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": base_item_id, - "type": "mcp_list_tools", - "server_label": "", - "tools": [] - } + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": base_item_id, + "type": "mcp_list_tools", + "server_label": "", + "tools": [], + } + ), ) events.append(output_item_done_event) - + return events def create_mcp_call_events( - tool_name: str, - tool_call_id: str, + tool_name: str, + tool_call_id: str, arguments: str, result: Optional[str] = None, base_item_id: Optional[str] = None, - sequence_start: int = 1 + sequence_start: int = 1, ) -> List[ResponsesAPIStreamingResponse]: """Create MCP call events following OpenAI's specification""" events: List[ResponsesAPIStreamingResponse] = [] item_id = base_item_id or f"mcp_{uuid.uuid4().hex[:8]}" - + # MCP call in progress event in_progress_event = MCPCallInProgressEvent( type=ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS, @@ -181,7 +181,7 @@ def create_mcp_call_events( item_id=item_id, ) events.append(in_progress_event) - + # MCP call arguments delta event (streaming the arguments) arguments_delta_event = MCPCallArgumentsDeltaEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA, @@ -191,7 +191,7 @@ def create_mcp_call_events( sequence_number=sequence_start + 1, ) events.append(arguments_delta_event) - + # MCP call arguments done event arguments_done_event = MCPCallArgumentsDoneEvent( type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE, @@ -201,7 +201,7 @@ def create_mcp_call_events( sequence_number=sequence_start + 2, ) events.append(arguments_done_event) - + # MCP call completed event (or failed if result indicates failure) if result is not None: completed_event = MCPCallCompletedEvent( @@ -211,23 +211,25 @@ def create_mcp_call_events( output_index=0, ) events.append(completed_event) - + # Add output_item.done event with the tool call result from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": item_id, - "type": "mcp_call", - "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", - "arguments": arguments, - "error": None, - "name": tool_name, - "output": result, - "server_label": "litellm" - }, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": item_id, + "type": "mcp_call", + "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", + "arguments": arguments, + "error": None, + "name": tool_name, + "output": result, + "server_label": "litellm", + } + ), ) events.append(output_item_done_event) else: @@ -238,7 +240,7 @@ def create_mcp_call_events( output_index=0, ) events.append(failed_event) - + return events @@ -250,51 +252,60 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): 3. Handles tool execution and follow-up calls for auto-execute tools 4. Emits tool execution events in the stream """ - + def __init__( self, base_iterator: Any, # Can be None - will be created internally mcp_events: List[ResponsesAPIStreamingResponse], mcp_tools_with_litellm_proxy: Optional[List[Any]] = None, user_api_key_auth: Any = None, - original_request_params: Optional[Dict[str, Any]] = None + original_request_params: Optional[Dict[str, Any]] = None, ): # MCP setup self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or [] self.user_api_key_auth = user_api_key_auth self.original_request_params = original_request_params or {} self.should_auto_execute = self._should_auto_execute_tools() - + # Streaming state management self.phase = "mcp_discovery" # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished self.finished = False - + # Event queues and generation flags - self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = mcp_events # Pre-generated MCP discovery events + self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = ( + mcp_events # Pre-generated MCP discovery events + ) self.tool_execution_events: List[ResponsesAPIStreamingResponse] = [] self.mcp_discovery_generated = True # Events are already generated - self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility - + self.mcp_events = ( + mcp_events # Store the initial MCP events for backward compatibility + ) + # Iterator references - self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = base_iterator # Will be created when needed + self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = ( + base_iterator # Will be created when needed + ) self.follow_up_iterator: Optional[Any] = None - + # Response collection for tool execution self.collected_response: Optional[ResponsesAPIResponse] = None - + # Set up model metadata (will be updated when we get the real iterator) - self.model = self.original_request_params.get('model', 'unknown') + self.model = self.original_request_params.get("model", "unknown") self.litellm_metadata = {} - self.custom_llm_provider = self.original_request_params.get('custom_llm_provider', None) - + self.custom_llm_provider = self.original_request_params.get( + "custom_llm_provider", None + ) + # Mark as async iterator self.is_async = True - + def _should_auto_execute_tools(self) -> bool: """Check if tools should be auto-executed""" from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) + return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( self.mcp_tools_with_litellm_proxy ) @@ -306,45 +317,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): """ Phase-based streaming: 1. mcp_discovery - Emit MCP discovery events - 2. initial_response - Stream the first LLM response + 2. initial_response - Stream the first LLM response 3. tool_execution - Emit tool execution events 4. follow_up_response - Stream the follow-up response 5. finished - End iteration """ - + # Phase 1: MCP Discovery Events if self.phase == "mcp_discovery": # Generate MCP discovery events if not already done # MCP discovery events are already generated and available - + # Emit MCP discovery events if self.mcp_discovery_events: return self.mcp_discovery_events.pop(0) - + # All MCP discovery events emitted, move to next phase - verbose_logger.debug("MCP discovery phase complete, transitioning to initial_response") + verbose_logger.debug( + "MCP discovery phase complete, transitioning to initial_response" + ) self.phase = "initial_response" await self._create_initial_response_iterator() # Fall through to process the initial response immediately - + # Phase 2: Initial Response Stream if self.phase == "initial_response": if self.base_iterator: # Check if base_iterator is actually iterable - if hasattr(self.base_iterator, '__anext__'): + if hasattr(self.base_iterator, "__anext__"): try: chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined] - + # If auto-execution is enabled, check for completed responses - if self.should_auto_execute and self._is_response_completed(chunk): + if self.should_auto_execute and self._is_response_completed( + chunk + ): # Collect the response for tool execution - response_obj = getattr(chunk, 'response', None) + response_obj = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj # Move to tool execution phase after emitting this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() - + return chunk except StopAsyncIteration: # Initial response ended, move to next phase @@ -357,24 +372,26 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): else: # base_iterator is not async iterable (likely a ResponsesAPIResponse) # Collect it for tool execution if needed - if self.should_auto_execute and isinstance(self.base_iterator, ResponsesAPIResponse): + if self.should_auto_execute and isinstance( + self.base_iterator, ResponsesAPIResponse + ): self.collected_response = self.base_iterator self.phase = "tool_execution" await self._generate_tool_execution_events() else: self.phase = "finished" raise StopAsyncIteration - + # Phase 3: Tool Execution Events if self.phase == "tool_execution": # Emit any queued tool execution events if self.tool_execution_events: return self.tool_execution_events.pop(0) - + # Move to follow-up response phase self.phase = "follow_up_response" await self._create_follow_up_iterator() - + # Phase 4: Follow-up Response Stream if self.phase == "follow_up_response": if self.follow_up_iterator: @@ -386,20 +403,22 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): else: self.phase = "finished" raise StopAsyncIteration - + # Phase 5: Finished if self.phase == "finished": raise StopAsyncIteration - + # Should not reach here raise StopAsyncIteration - + def _is_response_completed(self, chunk: ResponsesAPIStreamingResponse) -> bool: """Check if this chunk indicates the response is completed""" from litellm.types.llms.openai import ResponsesAPIStreamEvents - return getattr(chunk, 'type', None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED - - + + return ( + getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + async def _create_initial_response_iterator(self) -> None: """Create the initial response iterator by making the first LLM call""" try: @@ -408,38 +427,45 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # Make the initial response API call - but avoid the MCP wrapper params = self.original_request_params.copy() - params['stream'] = True # Ensure streaming - + params["stream"] = True # Ensure streaming + # Use the pre-fetched all_tools from original_request_params (no re-processing needed) params_for_llm = {} for key, value in params.items(): - params_for_llm[key] = value # Copy all params as-is since tools are already processed - - tools_count = len(params_for_llm.get('tools', [])) + params_for_llm[key] = ( + value # Copy all params as-is since tools are already processed + ) + + tools_count = len(params_for_llm.get("tools", [])) verbose_logger.debug(f"Making LLM call with {tools_count} tools") response = await aresponses(**params_for_llm) - + # Set the base iterator - if hasattr(response, '__aiter__') or hasattr(response, '__iter__'): + if hasattr(response, "__aiter__") or hasattr(response, "__iter__"): self.base_iterator = response # Copy metadata from the real iterator - self.model = getattr(response, 'model', self.model) - self.litellm_metadata = getattr(response, 'litellm_metadata', {}) - self.custom_llm_provider = getattr(response, 'custom_llm_provider', self.custom_llm_provider) - verbose_logger.debug(f"Created base iterator: {type(self.base_iterator)}") + self.model = getattr(response, "model", self.model) + self.litellm_metadata = getattr(response, "litellm_metadata", {}) + self.custom_llm_provider = getattr( + response, "custom_llm_provider", self.custom_llm_provider + ) + verbose_logger.debug( + f"Created base iterator: {type(self.base_iterator)}" + ) else: # Non-streaming response - this shouldn't happen but handle it verbose_logger.warning(f"Got non-streaming response: {type(response)}") self.base_iterator = None self.phase = "finished" - + except Exception as e: verbose_logger.error(f"Error creating initial response iterator: {e}") import traceback + traceback.print_exc() self.base_iterator = None self.phase = "finished" - + async def _generate_tool_execution_events(self) -> None: """Generate tool execution events and execute tools""" if not self.collected_response: @@ -447,7 +473,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - + try: # Extract tool calls from the response if self.collected_response is not None: @@ -456,9 +482,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): tool_calls = [] if not tool_calls: return - + for tool_call in tool_calls: - tool_name, tool_arguments, tool_call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + tool_name, tool_arguments, tool_call_id = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + ) if tool_name and tool_call_id: # Create MCP call events for this tool execution call_events = create_mcp_call_events( @@ -467,34 +495,35 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): arguments=tool_arguments or "{}", # JSON string with arguments result=None, # Will be set after execution base_item_id=f"mcp_{uuid.uuid4().hex[:8]}", - sequence_start=len(self.tool_execution_events) + 1 + sequence_start=len(self.tool_execution_events) + 1, ) # Add the in_progress and arguments events (not the completed event yet) self.tool_execution_events.extend(call_events[:-1]) - + # Execute the tools tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( - tool_calls=tool_calls, - user_api_key_auth=self.user_api_key_auth + tool_calls=tool_calls, user_api_key_auth=self.user_api_key_auth ) - + # Create completion events and output_item.done events for tool execution for tool_result in tool_results: tool_call_id = tool_result.get("tool_call_id", "unknown") result_text = tool_result.get("result", "") - + # Find matching tool name and arguments tool_name = "unknown" tool_arguments = "{}" for tool_call in tool_calls: - name, args, call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + name, args, call_id = ( + LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call) + ) if call_id == tool_call_id: tool_name = name or "unknown" tool_arguments = args or "{}" break - + item_id = f"mcp_{uuid.uuid4().hex[:8]}" - + # Create the completion event completed_event = MCPCallCompletedEvent( type=ResponsesAPIStreamEvents.MCP_CALL_COMPLETED, @@ -503,79 +532,84 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): output_index=0, ) self.tool_execution_events.append(completed_event) - + # Create output_item.done event with the tool call result from litellm.types.llms.openai import OutputItemDoneEvent - + output_item_done_event = OutputItemDoneEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, output_index=0, - item={ - "id": item_id, - "type": "mcp_call", - "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", - "arguments": tool_arguments, - "error": None, - "name": tool_name, - "output": result_text, - "server_label": "litellm" # or extract from tool config - }, + item=BaseLiteLLMOpenAIResponseObject( + **{ + "id": item_id, + "type": "mcp_call", + "approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}", + "arguments": tool_arguments, + "error": None, + "name": tool_name, + "output": result_text, + "server_label": "litellm", # or extract from tool config + } + ), ) self.tool_execution_events.append(output_item_done_event) - + # Store tool results for follow-up call self.tool_results = tool_results - + except Exception as e: verbose_logger.error(f"Error in tool execution: {e}") import traceback + traceback.print_exc() self.tool_results = [] - + async def _create_follow_up_iterator(self) -> None: """Create the follow-up response iterator with tool results""" - if not self.collected_response or not hasattr(self, 'tool_results'): + if not self.collected_response or not hasattr(self, "tool_results"): return - + from litellm.responses.main import aresponses from litellm.responses.mcp.litellm_proxy_mcp_handler import ( LiteLLM_Proxy_MCP_Handler, ) - + try: # Create follow-up input if self.collected_response is not None: follow_up_input = LiteLLM_Proxy_MCP_Handler._create_follow_up_input( response=self.collected_response, # type: ignore[arg-type] tool_results=self.tool_results, - original_input=self.original_request_params.get('input') + original_input=self.original_request_params.get("input"), ) - + # Make follow-up call with streaming follow_up_params = self.original_request_params.copy() - follow_up_params.update({ - 'input': follow_up_input, - 'previous_response_id': self.collected_response.id, # type: ignore[attr-defined] - 'stream': True - }) + follow_up_params.update( + { + "input": follow_up_input, + "previous_response_id": self.collected_response.id, # type: ignore[attr-defined] + "stream": True, + } + ) else: return # Remove tool_choice to avoid forcing more tool calls - follow_up_params.pop('tool_choice', None) - + follow_up_params.pop("tool_choice", None) + follow_up_response = await aresponses(**follow_up_params) - + # Set up the follow-up iterator - if hasattr(follow_up_response, '__aiter__'): + if hasattr(follow_up_response, "__aiter__"): self.follow_up_iterator = follow_up_response - + except Exception as e: verbose_logger.error(f"Error creating follow-up iterator: {e}") import traceback + traceback.print_exc() self.follow_up_iterator = None - def __iter__(self): return self @@ -583,11 +617,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): # First, emit any queued MCP events if self.mcp_events: # type: ignore[attr-defined] return self.mcp_events.pop(0) # type: ignore[attr-defined] - + # Then delegate to the base iterator if not self.is_async: try: - if self.base_iterator and hasattr(self.base_iterator, '__next__'): + if self.base_iterator and hasattr(self.base_iterator, "__next__"): return next(cast(Any, self.base_iterator)) # type: ignore[arg-type] else: raise StopIteration