Merge pull request #19552 from BerriAI/litellm_feat_completions_mcp_output

[feat] Add MCP tools response to chat completions

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

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

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"

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