From 18240662db00fd8cda89bf8aa968440a94d39c56 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 22 Jan 2026 15:18:24 +0530 Subject: [PATCH 1/4] Add custom vertex ai mapping to the output --- litellm/constants.py | 2 +- .../google_genai/adapters/transformation.py | 3 +- litellm/litellm_core_utils/core_helpers.py | 6 +- .../vertex_and_google_ai_studio_gemini.py | 4 +- ...odel_prices_and_context_window_backup.json | 188 +++++++++++++++++- litellm/types/llms/openai.py | 2 +- litellm/types/llms/vertex_ai.py | 2 + litellm/types/utils.py | 3 +- ...test_vertex_and_google_ai_studio_gemini.py | 34 +++- 9 files changed, 231 insertions(+), 13 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 5da0cf1489f..d346bc46f51 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1067,7 +1067,7 @@ known_tokenizer_config = { } -OPENAI_FINISH_REASONS = ["stop", "length", "function_call", "content_filter", "null"] +OPENAI_FINISH_REASONS = ["stop", "length", "function_call", "content_filter", "null", "finish_reason_unspecified", "malformed_function_call"] HUMANLOOP_PROMPT_CACHE_TTL_SECONDS = int( os.getenv("HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", 60) ) # 1 minute diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 58a52666d38..0a296012210 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -2,7 +2,6 @@ import json from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast from litellm import verbose_logger - from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema from litellm.types.llms.openai import ( AllMessageValues, @@ -771,6 +770,8 @@ class GoogleGenAIAdapter: "content_filter": "SAFETY", "tool_calls": "STOP", "function_call": "STOP", + "finish_reason_unspecified": "FINISH_REASON_UNSPECIFIED", + "malformed_function_call": "MALFORMED_FUNCTION_CALL", } return mapping.get(finish_reason, "STOP") diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index dadb36f3fd7..9cb0a00d9fc 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -79,9 +79,11 @@ def map_finish_reason( elif finish_reason == "eos_token" or finish_reason == "stop_sequence": return "stop" elif ( - finish_reason == "FINISH_REASON_UNSPECIFIED" or finish_reason == "STOP" + finish_reason == "FINISH_REASON_UNSPECIFIED" ): # vertex ai - got from running `print(dir(response_obj.candidates[0].finish_reason))`: ['FINISH_REASON_UNSPECIFIED', 'MAX_TOKENS', 'OTHER', 'RECITATION', 'SAFETY', 'STOP',] - return "stop" + return "finish_reason_unspecified" + elif finish_reason == "MALFORMED_FUNCTION_CALL": + return "malformed_function_call" elif finish_reason == "SAFETY" or finish_reason == "RECITATION": # vertex ai return "content_filter" elif finish_reason == "STOP": # vertex ai diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 2d2e07e74db..b78ac8f9e98 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1199,7 +1199,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): and what it means """ return { - "FINISH_REASON_UNSPECIFIED": "stop", # openai doesn't have a way of representing this + "FINISH_REASON_UNSPECIFIED": "finish_reason_unspecified", "STOP": "stop", "MAX_TOKENS": "length", "SAFETY": "content_filter", @@ -1209,7 +1209,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "BLOCKLIST": "content_filter", "PROHIBITED_CONTENT": "content_filter", "SPII": "content_filter", - "MALFORMED_FUNCTION_CALL": "stop", # openai doesn't have a way of representing this + "MALFORMED_FUNCTION_CALL": "malformed_function_call", # openai doesn't have a way of representing this "IMAGE_SAFETY": "content_filter", } diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a74b80e7373..ab034d9f51b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -16863,14 +16863,14 @@ "supports_vision": true }, "gpt-4o-audio-preview": { - "input_cost_per_audio_token": 0.0001, + "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_audio_token": 0.0002, + "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 1e-05, "supports_audio_input": true, "supports_audio_output": true, @@ -16880,14 +16880,14 @@ "supports_tool_choice": true }, "gpt-4o-audio-preview-2024-10-01": { - "input_cost_per_audio_token": 0.0001, + "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", "max_input_tokens": 128000, "max_output_tokens": 16384, "max_tokens": 16384, "mode": "chat", - "output_cost_per_audio_token": 0.0002, + "output_cost_per_audio_token": 8e-05, "output_cost_per_token": 1e-05, "supports_audio_input": true, "supports_audio_output": true, @@ -16930,6 +16930,186 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "gpt-audio": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-2025-08-28": { + "input_cost_per_audio_token": 3.2e-05, + "input_cost_per_token": 2.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 6.4e-05, + "output_cost_per_token": 1e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini-2025-10-06": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, + "gpt-audio-mini-2025-12-15": { + "input_cost_per_audio_token": 1e-05, + "input_cost_per_token": 6e-07, + "litellm_provider": "openai", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_audio_token": 2e-05, + "output_cost_per_token": 2.4e-06, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses", + "/v1/realtime", + "/v1/batch" + ], + "supported_modalities": [ + "text", + "audio" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": false, + "supports_reasoning": false, + "supports_response_schema": false, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": false + }, "gpt-4o-mini": { "cache_read_input_token_cost": 7.5e-08, "cache_read_input_token_cost_priority": 1.25e-07, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 3f9842de7da..694cdcabc92 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2002,7 +2002,7 @@ class OpenAIBatchResult(TypedDict, total=False): OpenAIChatCompletionFinishReason = Literal[ - "stop", "content_filter", "function_call", "tool_calls", "length" + "stop", "content_filter", "function_call", "tool_calls", "length", "finish_reason_unspecified", "malformed_function_call" # last 2 are vertex ai specific ] diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 0a4a2d0f14c..049a5010c79 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -396,6 +396,8 @@ class Candidates(TypedDict, total=False): "BLOCKLIST", "PROHIBITED_CONTENT", "SPII", + "MALFORMED_FUNCTION_CALL", + "IMAGE_SAFETY", ] safetyRatings: List[SafetyRatings] citationMetadata: CitationMetadata diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cac2fe85541..e840589af6b 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -46,6 +46,7 @@ from .llms.openai import ( FineTuningJob, ImageURLListItem, OpenAIChatCompletionChunk, + OpenAIChatCompletionFinishReason, OpenAIFileObject, OpenAIRealtimeStreamList, ResponsesAPIResponse, @@ -1254,7 +1255,7 @@ class Delta(SafeAttributeModel, OpenAIObject): class Choices(SafeAttributeModel, OpenAIObject): - finish_reason: str + finish_reason: OpenAIChatCompletionFinishReason index: int message: Message logprobs: Optional[Union[ChoiceLogprobs, Any]] = None diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 5be080b53fa..ac099a0168c 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -10,11 +10,11 @@ from pydantic import BaseModel import litellm from litellm import ModelResponse, completion +from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.types.llms.vertex_ai import UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import CustomStreamWrapper @@ -606,6 +606,38 @@ def test_check_finish_reason(): ) +def test_finish_reason_unspecified_and_malformed_function_call(): + """ + Test that FINISH_REASON_UNSPECIFIED and MALFORMED_FUNCTION_CALL + return their lowercase values instead of being mapped to 'stop' + since we don't have good mappings for these. + """ + finish_reason_mappings = VertexGeminiConfig.get_finish_reason_mapping() + + # Test FINISH_REASON_UNSPECIFIED returns lowercase version + assert finish_reason_mappings["FINISH_REASON_UNSPECIFIED"] == "finish_reason_unspecified" + assert ( + VertexGeminiConfig._check_finish_reason( + chat_completion_message=None, finish_reason="FINISH_REASON_UNSPECIFIED" + ) + == "finish_reason_unspecified" + ) + + # Test MALFORMED_FUNCTION_CALL returns lowercase version + assert finish_reason_mappings["MALFORMED_FUNCTION_CALL"] == "malformed_function_call" + assert ( + VertexGeminiConfig._check_finish_reason( + chat_completion_message=None, finish_reason="MALFORMED_FUNCTION_CALL" + ) + == "malformed_function_call" + ) + + # Ensure these values are in the OpenAI finish reasons constant + from litellm import OPENAI_FINISH_REASONS + assert "finish_reason_unspecified" in OPENAI_FINISH_REASONS + assert "malformed_function_call" in OPENAI_FINISH_REASONS + + def test_vertex_ai_usage_metadata_response_token_count(): """For Gemini Live API""" from litellm.types.utils import PromptTokensDetailsWrapper From 8a622c51f5f916310baf837d101e3e9903f5b1a8 Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 22 Jan 2026 15:32:04 +0900 Subject: [PATCH 2/4] feat: Add MCP tools response to chat completions --- .../litellm_core_utils/streaming_handler.py | 44 ++++ .../responses/mcp/chat_completions_handler.py | 98 ++++++++- tests/mcp_tests/test_mcp_chat_completions.py | 206 ++++++++++++++++++ .../mcp/test_chat_completions_handler.py | 152 +++++++++++++ 4 files changed, 497 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3093a37c26a..3304759f749 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1571,6 +1571,46 @@ class CustomStreamWrapper: ) return chunk + def _add_mcp_metadata_to_final_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + """ + Add MCP metadata from _hidden_params to the final chunk's delta.provider_specific_fields. + + This method checks if MCP metadata is stored in _hidden_params and adds it to + the chunk's delta.provider_specific_fields, similar to how RAG adds search results. + """ + try: + # Check if MCP metadata should be added to final chunk + if not hasattr(self, "_hidden_params") or not self._hidden_params: + return chunk + + mcp_metadata = self._hidden_params.get("mcp_metadata") + if not mcp_metadata: + return chunk + + # Add MCP metadata to delta.provider_specific_fields + if hasattr(chunk, "choices") and chunk.choices: + for choice in chunk.choices: + if isinstance(choice, StreamingChoices) and hasattr(choice, "delta") and choice.delta: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(choice.delta, "provider_specific_fields", None) or {} + ) + + # Add MCP metadata + if isinstance(mcp_metadata, dict): + provider_fields.update(mcp_metadata) + + # Set the provider_specific_fields + setattr(choice.delta, "provider_specific_fields", provider_fields) + + except Exception as e: + from litellm._logging import verbose_logger + verbose_logger.exception( + f"Error adding MCP metadata to final chunk: {str(e)}" + ) + + return chunk + def cache_streaming_response(self, processed_chunk, cache_hit: bool): """ Caches the streaming response @@ -1712,6 +1752,8 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) response._hidden_params["usage"] = usage + # Add MCP metadata to final chunk if present + response = self._add_mcp_metadata_to_final_chunk(response) # RETURN RESULT return response @@ -1884,6 +1926,8 @@ class CustomStreamWrapper: processed_chunk ) ) + # Add MCP metadata to final chunk if present (after hooks) + processed_chunk = self._add_mcp_metadata_to_final_chunk(processed_chunk) return processed_chunk raise StopAsyncIteration diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 26853b30596..0b0004d1548 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -15,6 +15,69 @@ from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper +def _add_mcp_metadata_to_response( + response: Union[ModelResponse, CustomStreamWrapper], + openai_tools: Optional[List], + tool_calls: Optional[List] = None, + tool_results: Optional[List] = None, +) -> None: + """ + Add MCP metadata to response's provider_specific_fields. + + This function adds MCP-related information to the response so that + clients can access which tools were available, which were called, and + what results were returned. + + For ModelResponse: adds to choices[].message.provider_specific_fields + For CustomStreamWrapper: stores in _hidden_params and automatically adds to + final chunk's delta.provider_specific_fields via CustomStreamWrapper._add_mcp_metadata_to_final_chunk() + """ + if isinstance(response, CustomStreamWrapper): + # For streaming, store MCP metadata in _hidden_params + # CustomStreamWrapper._add_mcp_metadata_to_final_chunk() will automatically + # add it to the final chunk's delta.provider_specific_fields + if not hasattr(response, "_hidden_params"): + response._hidden_params = {} + + mcp_metadata = {} + if openai_tools: + mcp_metadata["mcp_list_tools"] = openai_tools + if tool_calls: + mcp_metadata["mcp_tool_calls"] = tool_calls + if tool_results: + mcp_metadata["mcp_call_results"] = tool_results + + if mcp_metadata: + response._hidden_params["mcp_metadata"] = mcp_metadata + return + + if not isinstance(response, ModelResponse): + return + + if not hasattr(response, "choices") or not response.choices: + return + + # Add MCP metadata to all choices' messages + for choice in response.choices: + message = getattr(choice, "message", None) + if message is not None: + # Get existing provider_specific_fields or create new dict + provider_fields = ( + getattr(message, "provider_specific_fields", None) or {} + ) + + # Add MCP metadata + if openai_tools: + provider_fields["mcp_list_tools"] = openai_tools + if tool_calls: + provider_fields["mcp_tool_calls"] = tool_calls + if tool_results: + provider_fields["mcp_call_results"] = tool_results + + # Set the provider_specific_fields + setattr(message, "provider_specific_fields", provider_fields) + + async def acompletion_with_mcp( model: str, messages: List, @@ -103,7 +166,13 @@ async def acompletion_with_mcp( # If not auto-executing, just make the call with transformed tools if not should_auto_execute: - return await litellm_acompletion(**base_call_args) + response = await litellm_acompletion(**base_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + ) + return response # For auto-execute: disable streaming for initial call stream = kwargs.get("stream", False) @@ -130,7 +199,17 @@ async def acompletion_with_mcp( if stream: retry_args = dict(base_call_args) retry_args["stream"] = stream - return await litellm_acompletion(**retry_args) + response = await litellm_acompletion(**retry_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + ) + return response + _add_mcp_metadata_to_response( + response=initial_response, + openai_tools=openai_tools, + ) return initial_response # Execute tool calls @@ -147,6 +226,11 @@ async def acompletion_with_mcp( ) if not tool_results: + _add_mcp_metadata_to_response( + response=initial_response, + openai_tools=openai_tools, + tool_calls=tool_calls, + ) return initial_response # Create follow-up messages with tool results @@ -161,4 +245,12 @@ async def acompletion_with_mcp( follow_up_call_args["messages"] = follow_up_messages follow_up_call_args["stream"] = stream - return await litellm_acompletion(**follow_up_call_args) + response = await litellm_acompletion(**follow_up_call_args) + if isinstance(response, (ModelResponse, CustomStreamWrapper)): + _add_mcp_metadata_to_response( + response=response, + openai_tools=openai_tools, + tool_calls=tool_calls, + tool_results=tool_results, + ) + return response diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py index 973301abfb2..8857f016df8 100644 --- a/tests/mcp_tests/test_mcp_chat_completions.py +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -312,3 +312,209 @@ async def test_completion_mcp_with_streaming_no_timeout_error(monkeypatch): # Verify acompletion was called (should be called by acompletion_with_mcp) assert len(acompletion_calls) >= 1, "acompletion should be called" + + +@pytest.mark.asyncio +async def test_mcp_metadata_in_streaming_final_chunk(monkeypatch): + """ + Test that MCP metadata is added to the final streaming chunk's + delta.provider_specific_fields when using MCP tools with streaming. + """ + from types import SimpleNamespace + from unittest.mock import patch + + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.litellm_core_utils.litellm_logging import Logging + + dummy_tool = SimpleNamespace( + name="local_search", + description="search", + inputSchema={"type": "object", "properties": {}}, + ) + + async def fake_process(user_api_key_auth, mcp_tools_with_litellm_proxy): + return [dummy_tool], {"local_search": "local"} + + async def fake_execute(**kwargs): + tool_calls = kwargs.get("tool_calls") or [] + call_entry = tool_calls[0] + call_id = call_entry.get("id") or call_entry.get("call_id") or "call" + return [ + { + "tool_call_id": call_id, + "result": "executed", + "name": call_entry.get("name", "local_search"), + } + ] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + fake_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + fake_execute, + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda secret_fields, tools: (None, None, None, None)), + ) + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None): + return ModelResponseStream( + id="test-stream", + model="gpt-4o-mini", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + chunks = [ + create_chunk("Hello"), + create_chunk(" world"), + create_chunk("!", finish_reason="stop"), # Final chunk + ] + + # Create a proper CustomStreamWrapper with logging_obj + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="gpt-4o-mini", + logging_obj=logging_obj, + ) + self.chunks = chunks + self._index = 0 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + if self._index == len(self.chunks): + self.sent_last_chunk = True + # Call the method that adds MCP metadata to final chunk + chunk = self._add_mcp_metadata_to_final_chunk(chunk) + return chunk + raise StopIteration + + # Track calls to acompletion + acompletion_calls = [] + + async def mock_acompletion(**kwargs): + acompletion_calls.append(kwargs) + # First call (non-streaming for tool extraction) + if not kwargs.get("stream", False): + return ModelResponse( + id="test-1", + model="gpt-4o-mini", + choices=[{ + "message": { + "role": "assistant", + "tool_calls": [{ + "id": "call-1", + "type": "function", + "function": { + "name": "local_search", + "arguments": "{}" + } + }] + }, + "finish_reason": "tool_calls" + }], + created=0, + object="chat.completion", + ) + # Second call (streaming follow-up) + return MockStreamingResponse() + + with patch("litellm.acompletion", side_effect=mock_acompletion): + response = litellm.completion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "server_label": "local", + "require_approval": "never", + } + ], + stream=True, + mock_response="Final answer", + mock_tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + ) + + import asyncio + assert asyncio.iscoroutine(response) + result = await response + + assert isinstance(result, CustomStreamWrapper) + + # Verify _hidden_params contains mcp_metadata + assert hasattr(result, "_hidden_params") + assert "mcp_metadata" in result._hidden_params + mcp_metadata = result._hidden_params["mcp_metadata"] + assert "mcp_list_tools" in mcp_metadata + assert "mcp_tool_calls" in mcp_metadata + assert "mcp_call_results" in mcp_metadata + + # Consume the stream and check final chunk + all_chunks = list(result) + assert len(all_chunks) > 0 + + # Find the final chunk (with finish_reason) + final_chunk = None + for chunk in all_chunks: + if hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "finish_reason") and choice.finish_reason: + final_chunk = chunk + break + + # If no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # Verify MCP metadata is in the final chunk's delta.provider_specific_fields + if hasattr(final_chunk, "choices") and final_chunk.choices: + choice = final_chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + provider_fields = getattr(choice.delta, "provider_specific_fields", None) + assert provider_fields is not None, "Final chunk should have provider_specific_fields" + assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools" + assert "mcp_tool_calls" in provider_fields, "Should have mcp_tool_calls" + assert "mcp_call_results" in provider_fields, "Should have mcp_call_results" diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py index 03a749a8083..bc1f4fb72b5 100644 --- a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -177,3 +177,155 @@ async def test_acompletion_with_mcp_auto_exec_performs_follow_up(monkeypatch): assert first_call["stream"] is False assert second_call["messages"] == ["follow-up"] assert second_call["stream"] is True + + +@pytest.mark.asyncio +async def test_acompletion_with_mcp_adds_metadata_to_streaming(monkeypatch): + """ + Test that acompletion_with_mcp adds MCP metadata to CustomStreamWrapper + and it appears in the final chunk's delta.provider_specific_fields. + """ + from litellm.utils import CustomStreamWrapper + from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta + from litellm.litellm_core_utils.litellm_logging import Logging + + tools = [{"type": "mcp", "server_url": "litellm_proxy/mcp/local"}] + openai_tools = [{"type": "function", "function": {"name": "local_search"}}] + tool_calls = [{"id": "call-1", "type": "function", "function": {"name": "local_search"}}] + tool_results = [{"tool_call_id": "call-1", "result": "executed"}] + + # Create mock streaming chunks + def create_chunk(content, finish_reason=None): + return ModelResponseStream( + id="test-stream", + model="test-model", + created=1234567890, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=content, + role="assistant", + ), + finish_reason=finish_reason, + ) + ], + ) + + chunks = [ + create_chunk("Hello"), + create_chunk(" world", finish_reason="stop"), # Final chunk + ] + + # Create a proper CustomStreamWrapper + from unittest.mock import MagicMock + logging_obj = MagicMock() + logging_obj.model_call_details = {} + + class MockStreamingResponse(CustomStreamWrapper): + def __init__(self): + super().__init__( + completion_stream=None, + model="test-model", + logging_obj=logging_obj, + ) + self.chunks = chunks + self._index = 0 + self.sent_last_chunk = False + + def __iter__(self): + return self + + def __next__(self): + if self._index < len(self.chunks): + chunk = self.chunks[self._index] + self._index += 1 + if self._index == len(self.chunks): + self.sent_last_chunk = True + # Call the method that adds MCP metadata to final chunk + chunk = self._add_mcp_metadata_to_final_chunk(chunk) + return chunk + raise StopIteration + + mock_acompletion = AsyncMock(return_value=MockStreamingResponse()) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_use_litellm_mcp_gateway", + staticmethod(lambda tools: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_parse_mcp_tools", + staticmethod(lambda tools: (tools, [])), + ) + async def mock_process(**_): + return (tools, {"local_search": "local"}) + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_process_mcp_tools_without_openai_transform", + mock_process, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_transform_mcp_tools_to_openai", + staticmethod(lambda *_, **__: openai_tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: False), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + with patch("litellm.acompletion", mock_acompletion): + result = await acompletion_with_mcp( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + tools=tools, + stream=True, + ) + + # Verify result is CustomStreamWrapper + assert isinstance(result, CustomStreamWrapper) + + # Verify _hidden_params contains mcp_metadata + assert hasattr(result, "_hidden_params") + assert "mcp_metadata" in result._hidden_params + mcp_metadata = result._hidden_params["mcp_metadata"] + assert "mcp_list_tools" in mcp_metadata + assert mcp_metadata["mcp_list_tools"] == openai_tools + + # Consume the stream and check final chunk + all_chunks = list(result) + assert len(all_chunks) > 0 + + # Find the final chunk (with finish_reason) + final_chunk = None + for chunk in all_chunks: + if hasattr(chunk, "choices") and chunk.choices: + choice = chunk.choices[0] + if hasattr(choice, "finish_reason") and choice.finish_reason: + final_chunk = chunk + break + + # If no chunk with finish_reason, use the last chunk + if final_chunk is None and all_chunks: + final_chunk = all_chunks[-1] + + assert final_chunk is not None, "Should have a final chunk" + + # Verify MCP metadata is in the final chunk's delta.provider_specific_fields + if hasattr(final_chunk, "choices") and final_chunk.choices: + choice = final_chunk.choices[0] + if hasattr(choice, "delta") and choice.delta: + provider_fields = getattr(choice.delta, "provider_specific_fields", None) + assert provider_fields is not None, "Final chunk should have provider_specific_fields" + assert "mcp_list_tools" in provider_fields, "Should have mcp_list_tools" + assert provider_fields["mcp_list_tools"] == openai_tools From 1657af81ebab679ebbb70d0ec4a94ea8814e413e Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Thu, 22 Jan 2026 16:11:56 +0900 Subject: [PATCH 3/4] feat: display mcp output on the play ground --- .../components/playground/chat_ui/ChatUI.tsx | 5 +- .../playground/llm_calls/chat_completion.tsx | 68 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index a3bdaccd805..1db4d3a5857 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -902,6 +902,7 @@ const ChatUI: React.FC = ({ customProxyBaseUrl || undefined, mcpServers, mcpServerToolRestrictions, + handleMCPEvent, ); } else if (endpointType === EndpointType.IMAGE) { // For image generation @@ -1664,7 +1665,7 @@ const ChatUI: React.FC = ({ {message.role === "assistant" && index === chatHistory.length - 1 && mcpEvents.length > 0 && - endpointType === EndpointType.RESPONSES && ( + (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && (
@@ -1797,7 +1798,7 @@ const ChatUI: React.FC = ({ {/* Show MCP events during loading if no assistant message exists yet */} {isLoading && mcpEvents.length > 0 && - endpointType === EndpointType.RESPONSES && + (endpointType === EndpointType.RESPONSES || endpointType === EndpointType.CHAT) && chatHistory.length > 0 && chatHistory[chatHistory.length - 1].role === "user" && (
diff --git a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx index 24112ca1666..17cddc3ab16 100644 --- a/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx +++ b/ui/litellm-dashboard/src/components/playground/llm_calls/chat_completion.tsx @@ -4,6 +4,7 @@ import { TokenUsage } from "../chat_ui/ResponseMetrics"; import { VectorStoreSearchResponse } from "../chat_ui/types"; import { getProxyBaseUrl } from "@/components/networking"; import { MCPServer } from "../../mcp_tools/types"; +import { MCPEvent } from "../chat_ui/MCPEventsDisplay"; export async function makeOpenAIChatCompletionRequest( chatHistory: { role: string; content: string | any[] }[], @@ -27,6 +28,7 @@ export async function makeOpenAIChatCompletionRequest( customBaseUrl?: string, mcpServers?: MCPServer[], mcpServerToolRestrictions?: Record, + onMCPEvent?: (event: MCPEvent) => void, ) { // base url should be the current base_url const isLocal = process.env.NODE_ENV === "development"; @@ -56,6 +58,13 @@ export async function makeOpenAIChatCompletionRequest( // For collecting complete response text let fullResponseContent = ""; let fullReasoningContent = ""; + + // Track MCP metadata from final chunk + let mcpMetadata: { + mcp_list_tools?: any[]; + mcp_tool_calls?: any[]; + mcp_call_results?: any[]; + } | null = null; // Build tools array const tools: any[] = []; @@ -158,6 +167,19 @@ export async function makeOpenAIChatCompletionRequest( onSearchResults(delta.provider_specific_fields.search_results); } + // Check for MCP metadata in provider_specific_fields (typically in final chunk) + if (delta && delta.provider_specific_fields) { + const providerFields = delta.provider_specific_fields; + if (providerFields.mcp_list_tools || providerFields.mcp_tool_calls || providerFields.mcp_call_results) { + mcpMetadata = { + mcp_list_tools: providerFields.mcp_list_tools, + mcp_tool_calls: providerFields.mcp_tool_calls, + mcp_call_results: providerFields.mcp_call_results, + }; + console.log("MCP metadata found in chunk:", mcpMetadata); + } + } + // Check for usage data using type assertion const chunkWithUsage = chunk as any; if (chunkWithUsage.usage && onUsageData) { @@ -182,6 +204,52 @@ export async function makeOpenAIChatCompletionRequest( } } + // Process MCP metadata from final chunk and convert to MCPEvent format + if (mcpMetadata && onMCPEvent) { + // Convert mcp_list_tools to MCPEvent + if (mcpMetadata.mcp_list_tools && mcpMetadata.mcp_list_tools.length > 0) { + const toolsEvent: MCPEvent = { + type: "response.output_item.done", + item: { + type: "mcp_list_tools", + tools: mcpMetadata.mcp_list_tools.map((tool: any) => ({ + name: tool.function?.name || tool.name || "", + description: tool.function?.description || tool.description || "", + input_schema: tool.function?.parameters || tool.input_schema || {}, + })), + }, + timestamp: Date.now(), + }; + onMCPEvent(toolsEvent); + } + + // Convert mcp_tool_calls and mcp_call_results to MCPEvent[] + if (mcpMetadata.mcp_tool_calls && mcpMetadata.mcp_tool_calls.length > 0) { + mcpMetadata.mcp_tool_calls.forEach((toolCall: any, index: number) => { + const functionName = toolCall.function?.name || toolCall.name || ""; + const functionArgs = toolCall.function?.arguments || toolCall.arguments || "{}"; + + // Find corresponding result + const result = mcpMetadata.mcp_call_results?.find( + (r: any) => r.tool_call_id === toolCall.id || r.tool_call_id === toolCall.call_id + ) || mcpMetadata.mcp_call_results?.[index]; + + const callEvent: MCPEvent = { + type: "response.output_item.done", + item: { + type: "mcp_call", + name: functionName, + arguments: typeof functionArgs === "string" ? functionArgs : JSON.stringify(functionArgs), + output: result?.result ? (typeof result.result === "string" ? result.result : JSON.stringify(result.result)) : undefined, + }, + item_id: toolCall.id || toolCall.call_id, + timestamp: Date.now(), + }; + onMCPEvent(callEvent); + }); + } + } + const endTime = Date.now(); const totalLatency = endTime - startTime; if (onTotalLatency) { From 991fee056f89060a91ce884a252e4886b149fe63 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 22 Jan 2026 19:22:29 +0530 Subject: [PATCH 4/4] Fix batch tests --- .../anthropic_passthrough_logging_handler.py | 6 +++--- .../vertex_passthrough_logging_handler.py | 6 +++--- .../test_vertex_ai_batch_passthrough.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 4e1112329ee..e70d6cb7fca 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -396,7 +396,7 @@ class AnthropicPassthroughLoggingHandler: # Add batch-specific metadata to indicate this is a pending batch job litellm_model_response.choices = [Choices( - finish_reason="batch_pending", + finish_reason="stop", index=0, message={ "role": "assistant", @@ -438,7 +438,7 @@ class AnthropicPassthroughLoggingHandler: # Add error-specific metadata litellm_model_response.choices = [Choices( - finish_reason="batch_error", + finish_reason="stop", index=0, message={ "role": "assistant", @@ -472,7 +472,7 @@ class AnthropicPassthroughLoggingHandler: # Add error-specific metadata litellm_model_response.choices = [Choices( - finish_reason="batch_error", + finish_reason="stop", index=0, message={ "role": "assistant", diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 0962fafe3f6..3d5c529a3bb 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -619,7 +619,7 @@ class VertexPassthroughLoggingHandler: # Add batch-specific metadata to indicate this is a pending batch job litellm_model_response.choices = [Choices( - finish_reason="batch_pending", + finish_reason="stop", index=0, message={ "role": "assistant", @@ -661,7 +661,7 @@ class VertexPassthroughLoggingHandler: # Add error-specific metadata litellm_model_response.choices = [Choices( - finish_reason="batch_error", + finish_reason="stop", index=0, message={ "role": "assistant", @@ -695,7 +695,7 @@ class VertexPassthroughLoggingHandler: # Add error-specific metadata litellm_model_response.choices = [Choices( - finish_reason="batch_error", + finish_reason="stop", index=0, message={ "role": "assistant", diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 68c6bf98cb2..66c063d47d8 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -127,7 +127,7 @@ class TestVertexAIBatchPassthroughHandler: assert result is not None assert "result" in result assert "kwargs" in result - assert result["result"].choices[0].finish_reason == "batch_error" + assert result["result"].choices[0].finish_reason == "stop" assert result["kwargs"]["batch_job_state"] == "JOB_STATE_FAILED" def test_get_actual_model_id_from_router_with_router(self):