tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #17747 from BerriAI/litellm_feat_mcp-chat-completions

Browse Source 
feat: add support for using MCPs on /chat/completions
This commit is contained in:
YutaSaito
2025-12-13 05:05:21 +09:00
committed by GitHub
parent b651012fdc 8cccf3e9f0
commit 8899b63fa4
8 changed files with 858 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/completion/input.md
+3 -4
View File
@@ -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.
docs/my-website/docs/mcp.md
+31
View File
@@ -1137,6 +1137,37 @@ curl --location '<your-litellm-proxy-base-url>/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 '<your-litellm-proxy-base-url>/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:
litellm/main.py
+45 -33
View File
@@ -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
litellm/responses/mcp/chat_completions_handler.py
+199
View File
@@ -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)
litellm/responses/mcp/litellm_proxy_mcp_handler.py
+126 -14
View File
@@ -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,
tests/mcp_tests/test_mcp_chat_completions.py
+143
View File
@@ -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
tests/test_litellm/responses/mcp/test_chat_completions_handler.py
+167
View File
@@ -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
tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
+144
View File
@@ -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"]