diff --git a/docs/my-website/docs/completion/input.md b/docs/my-website/docs/completion/input.md index bdbd0b0492..7df4f77017 100644 --- a/docs/my-website/docs/completion/input.md +++ b/docs/my-website/docs/completion/input.md @@ -174,11 +174,11 @@ def completion( - `seed`: *integer or null (optional)* - This feature is in Beta. If specified, our system will make a best effort to sample deterministically, such that repeated requests with the same seed and parameters should return the same result. Determinism is not guaranteed, and you should refer to the `system_fingerprint` response parameter to monitor changes in the backend. -- `tools`: *array (optional)* - A list of tools the model may call. Currently, only functions are supported as a tool. Use this to provide a list of functions the model may generate JSON inputs for. +- `tools`: *array (optional)* - A list of tools the model may call. Use this to provide a list of functions the model may generate JSON inputs for. - - `type`: *string* - The type of the tool. Currently, only function is supported. + - `type`: *string* - The type of the tool. You can set this to `"function"` or `"mcp"` (matching the `/responses` schema) to call LiteLLM-registered MCP servers directly from `/chat/completions`. - - `function`: *object* - Required. + - `function`: *object* - Required for function tools. - `tool_choice`: *string or object (optional)* - Controls which (if any) function is called by the model. none means the model will not call a function and instead generates a message. auto means the model can pick between generating a message or calling a function. Specifying a particular function via `{"type": "function", "function": {"name": "my_function"}}` forces the model to call that function. @@ -247,4 +247,3 @@ def completion( - `eos_token`: *string (optional)* - Initial string applied at the end of a sequence - `hf_model_name`: *string (optional)* - [Sagemaker Only] The corresponding huggingface name of the model, used to pull the right chat template for the model. - diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index a9f7e24913..f9c9cbb456 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -1137,6 +1137,37 @@ curl --location '/v1/responses' \ }' ``` +## Use MCP tools with `/chat/completions` + +:::tip Works with all providers +This flow is **provider-agnostic**: the same MCP tool definition works for _every_ LLM backend behind LiteLLM (OpenAI, Azure OpenAI, Anthropic, Amazon Bedrock, Vertex, self-hosted deployments, etc.). +::: + +LiteLLM Proxy also supports MCP-aware tooling on the classic `/v1/chat/completions` endpoint. Provide the MCP tool definition directly in the `tools` array and LiteLLM will fetch and transform the MCP server's tools into OpenAI-compatible function calls. When `require_approval` is set to `"never"`, the proxy automatically executes the returned tool calls and feeds the results back into the model before returning the assistant response. + +```bash title="Chat Completions with MCP Tools" showLineNumbers +curl --location '/v1/chat/completions' \ +--header 'Content-Type: application/json' \ +--header "Authorization: Bearer $LITELLM_API_KEY" \ +--data '{ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "Summarize the latest open PR."} + ], + "tools": [ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/github", + "server_label": "github_mcp", + "require_approval": "never" + } + ] +}' +``` + +If you omit `require_approval` or set it to any value other than `"never"`, the MCP tool calls are returned to the client so that you can review and execute them manually, matching the upstream OpenAI behavior. + + ## LiteLLM Proxy - Walk through MCP Gateway LiteLLM exposes an MCP Gateway for admins to add all their MCP servers to LiteLLM. The key benefits of using LiteLLM Proxy with MCP are: diff --git a/litellm/main.py b/litellm/main.py index 20089b4c23..b08ffd16e3 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -69,6 +69,7 @@ from litellm.constants import ( ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, @@ -299,7 +300,6 @@ MOCK_RESPONSE_TYPE = Union[str, Exception, dict, ModelResponse, ModelResponseStr class LiteLLM: - def __init__( self, *, @@ -1091,6 +1091,22 @@ def completion( # type: ignore # noqa: PLR0915 tools = validate_and_fix_openai_tools(tools=tools) # validate tool_choice tool_choice = validate_chat_completion_tool_choice(tool_choice=tool_choice) + + skip_mcp_handler = kwargs.pop("_skip_mcp_handler", False) + if not skip_mcp_handler and tools: + from litellm.responses.mcp.chat_completions_handler import ( + handle_chat_completion_with_mcp, + ) + + mcp_handler_context = locals().copy() + completion_callable = globals().get("acompletion") + mcp_result = run_async_function( + handle_chat_completion_with_mcp, + mcp_handler_context, + completion_callable, + ) + if mcp_result is not None: + return mcp_result ######### unpacking kwargs ##################### args = locals() api_base = kwargs.get("api_base", None) @@ -1181,7 +1197,6 @@ def completion( # type: ignore # noqa: PLR0915 prompt_id=prompt_id, non_default_params=non_default_params ) ): - ( model, messages, @@ -2130,7 +2145,7 @@ def completion( # type: ignore # noqa: PLR0915 config = litellm.GenAIHubOrchestrationConfig.get_config() for k, v in config.items(): if ( - k not in optional_params + k not in optional_params ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in optional_params[k] = v @@ -2300,7 +2315,6 @@ def completion( # type: ignore # noqa: PLR0915 try: if use_base_llm_http_handler: - response = base_llm_http_handler.completion( model=model, messages=messages, @@ -3413,9 +3427,9 @@ def completion( # type: ignore # noqa: PLR0915 "aws_region_name" not in optional_params or optional_params["aws_region_name"] is None ): - optional_params["aws_region_name"] = ( - aws_bedrock_client.meta.region_name - ) + optional_params[ + "aws_region_name" + ] = aws_bedrock_client.meta.region_name bedrock_route = BedrockModelInfo.get_bedrock_route(model) if bedrock_route == "converse": @@ -3633,7 +3647,6 @@ def completion( # type: ignore # noqa: PLR0915 if api_key is not None and "Authorization" not in headers: headers["Authorization"] = f"Bearer {api_key}" - response = base_llm_http_handler.completion( model=model, stream=stream, @@ -3769,7 +3782,6 @@ def completion( # type: ignore # noqa: PLR0915 ) raise e elif custom_llm_provider == "gradient_ai": - api_base = litellm.api_base or api_base response = base_llm_http_handler.completion( model=model, @@ -4420,7 +4432,7 @@ def embedding( # noqa: PLR0915 litellm_params=litellm_params_dict, ) elif custom_llm_provider == "github_copilot": - api_key = (api_key or litellm.api_key) + api_key = api_key or litellm.api_key response = base_llm_http_handler.embedding( model=model, input=input, @@ -5585,9 +5597,9 @@ def adapter_completion( new_kwargs = translation_obj.translate_completion_input_params(kwargs=kwargs) response: Union[ModelResponse, CustomStreamWrapper] = completion(**new_kwargs) # type: ignore - translated_response: Optional[Union[BaseModel, AdapterCompletionStreamWrapper]] = ( - None - ) + translated_response: Optional[ + Union[BaseModel, AdapterCompletionStreamWrapper] + ] = None if isinstance(response, ModelResponse): translated_response = translation_obj.translate_completion_output_params( response=response @@ -6292,9 +6304,9 @@ def speech( # noqa: PLR0915 ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY ] = query_params - litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = ( - voice_id - ) + litellm_params_dict[ + ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY + ] = voice_id if api_base is not None: litellm_params_dict["api_base"] = api_base @@ -6344,16 +6356,16 @@ def speech( # noqa: PLR0915 text_to_speech_provider_config = VertexAITextToSpeechConfig() # Cast to specific Vertex AI config type to access dispatch method - vertex_config = cast( - VertexAITextToSpeechConfig, text_to_speech_provider_config - ) + vertex_config = cast(VertexAITextToSpeechConfig, text_to_speech_provider_config) # Store Vertex AI specific params in litellm_params_dict - litellm_params_dict.update({ - "vertex_project": generic_optional_params.vertex_project, - "vertex_location": generic_optional_params.vertex_location, - "vertex_credentials": generic_optional_params.vertex_credentials, - }) + litellm_params_dict.update( + { + "vertex_project": generic_optional_params.vertex_project, + "vertex_location": generic_optional_params.vertex_location, + "vertex_credentials": generic_optional_params.vertex_credentials, + } + ) response = vertex_config.dispatch_text_to_speech( model=model, @@ -6724,9 +6736,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(content_chunks) > 0: - response["choices"][0]["message"]["content"] = ( - processor.get_combined_content(content_chunks) - ) + response["choices"][0]["message"][ + "content" + ] = processor.get_combined_content(content_chunks) thinking_blocks = [ chunk @@ -6737,9 +6749,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(thinking_blocks) > 0: - response["choices"][0]["message"]["thinking_blocks"] = ( - processor.get_combined_thinking_content(thinking_blocks) - ) + response["choices"][0]["message"][ + "thinking_blocks" + ] = processor.get_combined_thinking_content(thinking_blocks) reasoning_chunks = [ chunk @@ -6750,9 +6762,9 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(reasoning_chunks) > 0: - response["choices"][0]["message"]["reasoning_content"] = ( - processor.get_combined_reasoning_content(reasoning_chunks) - ) + response["choices"][0]["message"][ + "reasoning_content" + ] = processor.get_combined_reasoning_content(reasoning_chunks) annotation_chunks = [ chunk diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py new file mode 100644 index 0000000000..1957e5fa92 --- /dev/null +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -0,0 +1,199 @@ +"""Helpers for handling MCP-aware `/chat/completions` requests.""" + +from typing import ( + Any, + Awaitable, + Callable, + Dict, + Iterable, + Optional, + Union, + cast, +) + +from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, +) +from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.llms.openai import ToolParam +from litellm.types.utils import ModelResponse +from litellm.utils import CustomStreamWrapper + +CompletionCallable = Callable[..., Awaitable[Union[ModelResponse, CustomStreamWrapper]]] + +_CHAT_COMPLETION_CALL_ARG_KEYS = [ + "model", + "messages", + "functions", + "function_call", + "timeout", + "temperature", + "top_p", + "n", + "stream", + "stream_options", + "stop", + "max_tokens", + "max_completion_tokens", + "modalities", + "prediction", + "audio", + "presence_penalty", + "frequency_penalty", + "logit_bias", + "user", + "response_format", + "seed", + "tools", + "tool_choice", + "parallel_tool_calls", + "logprobs", + "top_logprobs", + "deployment_id", + "reasoning_effort", + "verbosity", + "safety_identifier", + "service_tier", + "base_url", + "api_version", + "api_key", + "model_list", + "extra_headers", + "thinking", + "web_search_options", + "shared_session", +] + + +def _build_call_args_from_context(call_context: Dict[str, Any]) -> Dict[str, Any]: + """Build kwargs for `acompletion` from the `completion` call context.""" + + call_args = { + key: call_context.get(key) + for key in _CHAT_COMPLETION_CALL_ARG_KEYS + if key in call_context + } + additional_kwargs = dict(call_context.get("kwargs") or {}) + call_args.update(additional_kwargs) + return call_args + + +async def _call_acompletion_internal( + completion_callable: CompletionCallable, **call_args: Any +) -> Union[ModelResponse, CustomStreamWrapper]: + """Invoke `acompletion` while skipping MCP interception to avoid recursion.""" + + safe_args = dict(call_args) + safe_args["_skip_mcp_handler"] = True + safe_args.pop("acompletion", None) + return await completion_callable(**safe_args) + + +async def handle_chat_completion_with_mcp( + call_context: Dict[str, Any], + completion_callable: CompletionCallable, +) -> Optional[Union[ModelResponse, CustomStreamWrapper]]: + """Handle MCP-enabled tool execution for chat completion requests.""" + + call_args = _build_call_args_from_context(call_context) + + tools = call_args.get("tools") + if not tools: + return None + + tools_for_mcp = cast(Optional[Iterable[ToolParam]], tools) + + if not LiteLLM_Proxy_MCP_Handler._should_use_litellm_mcp_gateway( + tools=tools_for_mcp + ): + return None + + mcp_tools, _ = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) + if not mcp_tools: + return None + + base_call_args = dict(call_args) + + user_api_key_auth = call_args.get("user_api_key_auth") or ( + (call_args.get("metadata", {}) or {}).get("user_api_key_auth") + ) + ( + deduplicated_mcp_tools, + tool_server_map, + ) = await LiteLLM_Proxy_MCP_Handler._process_mcp_tools_without_openai_transform( + user_api_key_auth=user_api_key_auth, + mcp_tools_with_litellm_proxy=mcp_tools, + ) + + openai_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( + deduplicated_mcp_tools, + target_format="chat", + ) + + base_call_args["tools"] = openai_tools or None + + should_auto_execute = LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools( + mcp_tools_with_litellm_proxy=mcp_tools + ) + + ( + mcp_auth_header, + mcp_server_auth_headers, + oauth2_headers, + raw_headers, + ) = ResponsesAPIRequestUtils.extract_mcp_headers_from_request( + secret_fields=base_call_args.get("secret_fields"), + tools=tools, + ) + + if not should_auto_execute: + return await _call_acompletion_internal(completion_callable, **base_call_args) + + mock_tool_calls = base_call_args.pop("mock_tool_calls", None) + + initial_call_args = dict(base_call_args) + initial_call_args["stream"] = False + if mock_tool_calls is not None: + initial_call_args["mock_tool_calls"] = mock_tool_calls + + initial_response = await _call_acompletion_internal( + completion_callable, **initial_call_args + ) + if not isinstance(initial_response, ModelResponse): + return initial_response + + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( + response=initial_response + ) + + if not tool_calls: + if base_call_args.get("stream"): + retry_args = dict(base_call_args) + retry_args["stream"] = call_args.get("stream") + return await _call_acompletion_internal(completion_callable, **retry_args) + return initial_response + + tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls( + tool_server_map=tool_server_map, + tool_calls=tool_calls, + user_api_key_auth=user_api_key_auth, + mcp_auth_header=mcp_auth_header, + mcp_server_auth_headers=mcp_server_auth_headers, + oauth2_headers=oauth2_headers, + raw_headers=raw_headers, + ) + + if not tool_results: + return initial_response + + follow_up_messages = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages=call_args.get("messages", []), + response=initial_response, + tool_results=tool_results, + ) + + follow_up_call_args = dict(base_call_args) + follow_up_call_args["messages"] = follow_up_messages + follow_up_call_args["stream"] = call_args.get("stream") + + return await _call_acompletion_internal(completion_callable, **follow_up_call_args) diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index a41d6f4f5a..4dda665f70 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,10 +1,21 @@ -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterable, + List, + Optional, + Tuple, + Union, + Literal, +) from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.utils import split_server_prefix_from_name from litellm.responses.main import aresponses from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ResponsesAPIResponse, ToolParam +from litellm.types.utils import Choices, ModelResponse if TYPE_CHECKING: from mcp.types import Tool as MCPTool @@ -163,7 +174,7 @@ class LiteLLM_Proxy_MCP_Handler: if len(allowed_mcp_servers) == 1: tool_server_map[tool_name] = allowed_mcp_servers[0] else: - tool_server_map[tool_name], _ = split_server_prefix_from_name( + _, tool_server_map[tool_name] = split_server_prefix_from_name( tool_name ) @@ -274,15 +285,23 @@ class LiteLLM_Proxy_MCP_Handler: return deduplicated_mcp_tools, tool_server_map @staticmethod - def _transform_mcp_tools_to_openai(mcp_tools: List[Any]) -> List[Any]: + def _transform_mcp_tools_to_openai( + mcp_tools: List[Any], + target_format: Literal["responses", "chat"] = "responses", + ) -> List[Any]: """Transform MCP tools to OpenAI-compatible format.""" from litellm.experimental_mcp_client.tools import ( transform_mcp_tool_to_openai_responses_api_tool, + transform_mcp_tool_to_openai_tool, ) - openai_tools = [] + openai_tools: List[Any] = [] for mcp_tool in mcp_tools: - openai_tool = transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) + openai_tool: Any + if target_format == "chat": + openai_tool = transform_mcp_tool_to_openai_tool(mcp_tool) + else: + openai_tool = transform_mcp_tool_to_openai_responses_api_tool(mcp_tool) openai_tools.append(openai_tool) return openai_tools @@ -325,22 +344,59 @@ class LiteLLM_Proxy_MCP_Handler: return tool_calls + @staticmethod + def _extract_tool_calls_from_chat_response(response: ModelResponse) -> List[Any]: + """Extract tool calls from a chat completion response.""" + tool_calls: List[Any] = [] + + try: + for choice in response.choices: + message = getattr(choice, "message", None) + if message is None: + continue + tool_call_entries = getattr(message, "tool_calls", None) + if tool_call_entries: + for tool_call in tool_call_entries: + if hasattr(tool_call, "model_dump"): + tool_calls.append(tool_call.model_dump()) + else: + tool_calls.append(tool_call) + except Exception: + verbose_logger.exception( + "Failed to extract tool calls from chat completion response" + ) + + return tool_calls + @staticmethod def _extract_tool_call_details( tool_call, ) -> Tuple[Optional[str], Optional[str], Optional[str]]: """Extract tool name, arguments, and call_id from a tool call.""" if isinstance(tool_call, dict): - tool_name = tool_call.get("name") - tool_arguments = tool_call.get("arguments") tool_call_id = tool_call.get("call_id") or tool_call.get("id") + + # OpenAI chat completions wrap tool info under a `function` block + function_block = tool_call.get("function") + if isinstance(function_block, dict): + tool_name = function_block.get("name") + tool_arguments = function_block.get("arguments") + else: + tool_name = tool_call.get("name") + tool_arguments = tool_call.get("arguments") else: - tool_name = getattr(tool_call, "name", None) - tool_arguments = getattr(tool_call, "arguments", None) tool_call_id = getattr(tool_call, "call_id", None) or getattr( tool_call, "id", None ) + function_obj = getattr(tool_call, "function", None) + if function_obj is not None: + tool_name = getattr(function_obj, "name", None) + tool_arguments = getattr(function_obj, "arguments", None) + else: + tool_name = getattr(tool_call, "name", None) + tool_arguments = getattr(tool_call, "arguments", None) + return tool_name, tool_arguments, tool_call_id @staticmethod @@ -399,8 +455,8 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod async def _execute_tool_calls( - tool_server_map: dict[str, str], - tool_calls: List[Any], + tool_server_map: dict[str, str], + tool_calls: List[Any], user_api_key_auth: Any, mcp_auth_header: Optional[str] = None, mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None, @@ -453,7 +509,11 @@ class LiteLLM_Proxy_MCP_Handler: # Format result for inclusion in response result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result) tool_results.append( - {"tool_call_id": tool_call_id, "result": result_text} + { + "tool_call_id": tool_call_id, + "result": result_text, + "name": tool_name, + } ) except BlockedPiiEntityError as e: @@ -462,7 +522,11 @@ class LiteLLM_Proxy_MCP_Handler: ) error_message = f"Tool call blocked: PII entity '{getattr(e, 'entity_type', 'unknown')}' detected by guardrail '{getattr(e, 'guardrail_name', 'unknown')}'. {str(e)}" tool_results.append( - {"tool_call_id": tool_call_id, "result": error_message} + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } ) except GuardrailRaisedException as e: verbose_logger.error( @@ -470,7 +534,11 @@ class LiteLLM_Proxy_MCP_Handler: ) error_message = f"Tool call blocked: Guardrail '{getattr(e, 'guardrail_name', 'unknown')}' violation. {str(e)}" tool_results.append( - {"tool_call_id": tool_call_id, "result": error_message} + { + "tool_call_id": tool_call_id, + "result": error_message, + "name": tool_name, + } ) except HTTPException as e: verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}") @@ -484,11 +552,55 @@ class LiteLLM_Proxy_MCP_Handler: { "tool_call_id": tool_call_id, "result": f"Error executing tool: {str(e)}", + "name": tool_name, } ) return tool_results + @staticmethod + def _create_follow_up_messages_for_chat( + original_messages: List[Any], + response: ModelResponse, + tool_results: List[Dict[str, Any]], + ) -> List[Any]: + """Create follow-up chat messages that include tool execution results.""" + from copy import deepcopy + + from litellm.utils import convert_list_message_to_dict + + follow_up_messages: List[Any] = convert_list_message_to_dict( + deepcopy(original_messages) + ) + + if not follow_up_messages: + follow_up_messages = [] + + message_to_append: Optional[dict] = None + try: + first_choice = response.choices[0] + if isinstance(first_choice, Choices) and getattr( + first_choice, "message", None + ): + message_to_append = first_choice.message.model_dump(exclude_none=True) + except Exception: + verbose_logger.exception("Failed to convert assistant message for MCP flow") + + if message_to_append: + follow_up_messages.append(message_to_append) + + for tool_result in tool_results: + follow_up_messages.append( + { + "role": "tool", + "tool_call_id": tool_result.get("tool_call_id"), + "name": tool_result.get("name"), + "content": tool_result.get("result", ""), + } + ) + + return follow_up_messages + @staticmethod def _create_follow_up_input( response: ResponsesAPIResponse, diff --git a/tests/mcp_tests/test_mcp_chat_completions.py b/tests/mcp_tests/test_mcp_chat_completions.py new file mode 100644 index 0000000000..ae13b6ca6e --- /dev/null +++ b/tests/mcp_tests/test_mcp_chat_completions.py @@ -0,0 +1,143 @@ +import pytest + +import litellm +from litellm.types.utils import ModelResponse + + +@pytest.mark.asyncio +async def test_acompletion_mcp_auto_exec(monkeypatch): + from types import SimpleNamespace + + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + + 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): + fake_execute.called = True # type: ignore[attr-defined] + tool_calls = kwargs.get("tool_calls") or [] + assert tool_calls, "tool calls should be present during auto execution" + 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"), + } + ] + + fake_execute.called = False # type: ignore[attr-defined] + + 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)), + ) + + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "server_label": "local", + "require_approval": "never", + } + ], + mock_response="Final answer", + mock_tool_calls=[ + { + "id": "call-1", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + ) + + assert isinstance(response, ModelResponse) + assert response.choices[0].message.content == "Final answer" + assert fake_execute.called is True # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_acompletion_mcp_respects_manual_approval(monkeypatch): + from types import SimpleNamespace + + from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, + ) + from litellm.responses.utils import ResponsesAPIRequestUtils + + 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): + pytest.fail("auto execution should not run when approval is required") + + 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)), + ) + + response = await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + tools=[ + { + "type": "mcp", + "server_url": "litellm_proxy/mcp/local", + "server_label": "local", + "require_approval": "manual", + } + ], + mock_response="Pending tool", + mock_tool_calls=[ + { + "id": "call-2", + "type": "function", + "function": {"name": "local_search", "arguments": "{}"}, + } + ], + ) + + assert isinstance(response, ModelResponse) + tool_calls = response.choices[0].message.tool_calls + assert tool_calls is not None and len(tool_calls) == 1 diff --git a/tests/test_litellm/responses/mcp/test_chat_completions_handler.py b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py new file mode 100644 index 0000000000..96e7c39aee --- /dev/null +++ b/tests/test_litellm/responses/mcp/test_chat_completions_handler.py @@ -0,0 +1,167 @@ +import pytest +from unittest.mock import AsyncMock + +from litellm.types.utils import ModelResponse + +from litellm.responses.mcp.chat_completions_handler import ( + handle_chat_completion_with_mcp, +) +from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, +) +from litellm.responses.utils import ResponsesAPIRequestUtils + + +@pytest.mark.asyncio +async def test_handle_chat_completion_returns_none_without_tools(): + completion_callable = AsyncMock() + + result = await handle_chat_completion_with_mcp({}, completion_callable) + + assert result is None + completion_callable.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_handle_chat_completion_without_auto_execution_calls_model(monkeypatch): + tools = [{"type": "function", "function": {"name": "tool"}}] + completion_callable = AsyncMock(return_value="ok") + + 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 ([], {}) + + 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-tool"]), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: False), + ) + captured_secret_fields = {} + + def mock_extract(**kwargs): + captured_secret_fields["value"] = kwargs.get("secret_fields") + return (None, None, None, None) + + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(mock_extract), + ) + + call_context = { + "tools": tools, + "messages": [], + "kwargs": {"secret_fields": {"api_key": "value"}}, + } + result = await handle_chat_completion_with_mcp(call_context, completion_callable) + + assert result == "ok" + completion_callable.assert_awaited_once() + kwargs = completion_callable.await_args.kwargs + assert kwargs.get("_skip_mcp_handler") is True + assert kwargs.get("tools") == ["openai-tool"] + assert captured_secret_fields["value"] == {"api_key": "value"} + + +@pytest.mark.asyncio +async def test_handle_chat_completion_auto_exec_performs_follow_up(monkeypatch): + tools = [{"type": "function", "function": {"name": "tool"}}] + initial_response = ModelResponse( + id="1", + model="test", + choices=[], + created=0, + object="chat.completion", + ) + follow_up_response = ModelResponse( + id="2", + model="test", + choices=[], + created=0, + object="chat.completion", + ) + completion_callable = AsyncMock( + side_effect=[initial_response, follow_up_response] + ) + + 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, {"tool": "server"})), + ) + async def mock_process(**_): + return (tools, {"tool": "server"}) + + 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 *_, **__: tools), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_should_auto_execute_tools", + staticmethod(lambda **_: True), + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_extract_tool_calls_from_chat_response", + staticmethod(lambda **_: ["call"]), + ) + async def mock_execute(**_): + return ["result"] + + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_execute_tool_calls", + mock_execute, + ) + monkeypatch.setattr( + LiteLLM_Proxy_MCP_Handler, + "_create_follow_up_messages_for_chat", + staticmethod(lambda **_: ["follow-up"]), + ) + monkeypatch.setattr( + ResponsesAPIRequestUtils, + "extract_mcp_headers_from_request", + staticmethod(lambda **_: (None, None, None, None)), + ) + + call_context = {"tools": tools, "messages": ["msg"], "stream": True} + result = await handle_chat_completion_with_mcp(call_context, completion_callable) + + assert result is follow_up_response + assert completion_callable.await_count == 2 + first_call = completion_callable.await_args_list[0].kwargs + second_call = completion_callable.await_args_list[1].kwargs + assert first_call["stream"] is False + assert second_call["messages"] == ["follow-up"] + assert second_call["stream"] is True diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py new file mode 100644 index 0000000000..9d4e0aeded --- /dev/null +++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py @@ -0,0 +1,144 @@ +import pytest + +from litellm.responses.mcp.litellm_proxy_mcp_handler import ( + LiteLLM_Proxy_MCP_Handler, +) +from litellm.types.utils import ModelResponse + + +def test_deduplicate_mcp_tools_single_allowed_server(): + tools = [{"name": "search"}, {"name": "search"}] # duplicate on purpose + + deduped, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + tools, + ["everything"], + ) + + assert len(deduped) == 1 + assert server_map == {"search": "everything"} + + +@pytest.mark.parametrize( + "tool_name,expected_server", + [ + ("alpha-tool", "alpha"), + ("beta-another_tool", "beta"), + ], +) +def test_deduplicate_mcp_tools_prefixed_names(tool_name, expected_server): + tools = [{"name": tool_name}] + + _, server_map = LiteLLM_Proxy_MCP_Handler._deduplicate_mcp_tools( + tools, + ["alpha", "beta"], + ) + + assert server_map[tool_name] == expected_server + + +def test_extract_tool_calls_from_chat_response_handles_tool_calls(): + response = ModelResponse( + id="resp-1", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-123", + "type": "function", + "function": {"name": "foo", "arguments": "{}"}, + } + ], + }, + "finish_reason": "tool_calls", + } + ], + model="gpt", + created=0, + object="chat.completion", + ) + + tool_calls = LiteLLM_Proxy_MCP_Handler._extract_tool_calls_from_chat_response( + response + ) + + assert len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "foo" + + +def test_create_follow_up_messages_for_chat_appends_tool_results(): + original_messages = [{"role": "user", "content": "hi"}] + response = ModelResponse( + id="resp-2", + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call-abc", + "type": "function", + "function": {"name": "foo", "arguments": "{}"}, + } + ], + }, + } + ], + model="gpt", + created=0, + object="chat.completion", + ) + tool_results = [ + { + "tool_call_id": "call-abc", + "name": "foo", + "result": "done", + } + ] + + follow_up = LiteLLM_Proxy_MCP_Handler._create_follow_up_messages_for_chat( + original_messages, + response, + tool_results, + ) + + assert follow_up[0]["role"] == "user" + assert follow_up[-1]["role"] == "tool" + assert follow_up[-1]["name"] == "foo" + assert follow_up[-1]["content"] == "done" + + +def test_transform_mcp_tools_to_openai_uses_chat_format(monkeypatch): + captured = {} + + def fake_transform_chat(tool): + captured.setdefault("chat", []).append(tool) + return {"chat": True} + + def fake_transform_responses(tool): + captured.setdefault("responses", []).append(tool) + return {"responses": True} + + monkeypatch.setattr( + "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool", + fake_transform_chat, + ) + monkeypatch.setattr( + "litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_responses_api_tool", + fake_transform_responses, + ) + + chat_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai( + ["tool"], target_format="chat" + ) + resp_tools = LiteLLM_Proxy_MCP_Handler._transform_mcp_tools_to_openai(["tool"]) + + assert chat_tools == [{"chat": True}] + assert resp_tools == [{"responses": True}] + assert captured["chat"] == ["tool"] + assert captured["responses"] == ["tool"]