scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

feat: save mcp call log via responses

Browse source
This commit is contained in:
Yuta Saito 2026-01-21 13:54:24 +09:00
parent 6267f1689b
commit ae4d92ad50
4 changed files with 196 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/responses/main.py
View file

@ -301,6 +301,8 @@ async def aresponses_api_with_mcp(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers_from_request,
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
)
if tool_results:
@ -349,6 +351,7 @@ async def aresponses_api_with_mcp(
tool_server_map=tool_server_map,
base_iterator=final_response,
mcp_events=tool_execution_events,
user_api_key_auth=user_api_key_auth,
)
# Add custom output elements to the final response (for non-streaming)

2
litellm/responses/mcp/chat_completions_handler.py
View file

@ -142,6 +142,8 @@ async def acompletion_with_mcp(
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_call_id=kwargs.get("litellm_call_id"),
litellm_trace_id=kwargs.get("litellm_trace_id"),
)
if not tool_results:

190
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -1,3 +1,5 @@
import traceback
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
@ -11,14 +13,18 @@ from typing import (
)
from litellm._logging import verbose_logger
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
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
from litellm.types.utils import CallTypes, Choices, ModelResponse, StandardLoggingMCPToolCall
from litellm.utils import Rules, function_setup
if TYPE_CHECKING:
from mcp.types import Tool as MCPTool
from litellm.proxy.utils import ProxyLogging
else:
MCPTool = Any
@ -470,6 +476,8 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
litellm_call_id: Optional[str] = None,
litellm_trace_id: Optional[str] = None,
) -> List[Dict[str, Any]]:
"""Execute tool calls and return results."""
from fastapi import HTTPException
@ -478,10 +486,16 @@ class LiteLLM_Proxy_MCP_Handler:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy.proxy_server import proxy_logging_obj
from litellm._uuid import uuid
tool_results = []
tool_call_id: Optional[str] = None
rules_obj = Rules()
for tool_call in tool_calls:
logging_request_data: Dict[str, Any] = {}
tool_name: str = ""
try:
(
tool_name,
@ -514,6 +528,101 @@ class LiteLLM_Proxy_MCP_Handler:
):
sanitized_tool_name = unprefixed_name
start_time = datetime.now()
logging_input = [
{
"role": "tool",
"content": {
"tool_name": sanitized_tool_name,
"arguments": parsed_arguments,
},
}
]
tool_logging_call_id = litellm_call_id or str(uuid.uuid4())
logging_request_data: Dict[str, Any] = {
"model": f"MCP: {tool_name}",
"metadata": {
"tool_call_id": tool_call_id,
"tool_name": sanitized_tool_name,
"server_name": server_name,
},
"input": logging_input,
"call_type": CallTypes.call_mcp_tool.value,
"litellm_call_id": tool_logging_call_id,
}
if litellm_trace_id:
logging_request_data["litellm_trace_id"] = litellm_trace_id
user_identifier = None
if user_api_key_auth is not None:
user_api_key = getattr(user_api_key_auth, "api_key", None)
if user_api_key:
logging_request_data["metadata"]["user_api_key"] = user_api_key
user_identifier = getattr(user_api_key_auth, "end_user_id", None) or getattr(
user_api_key_auth, "user_id", None
)
if user_identifier:
logging_request_data["user"] = user_identifier
litellm_logging_obj: Optional[LiteLLMLoggingObj] = None
try:
litellm_logging_obj, _ = function_setup(
original_function="call_mcp_tool",
rules_obj=rules_obj,
start_time=start_time,
**logging_request_data,
)
except Exception as logging_error:
verbose_logger.debug(
"Failed to initialize logging for MCP tool call %s: %s",
tool_name,
logging_error,
)
litellm_logging_obj = None
logging_request_data["litellm_logging_obj"] = litellm_logging_obj
logging_request_data["arguments"] = parsed_arguments
if litellm_logging_obj:
try:
litellm_logging_obj.pre_call(
input=logging_input,
api_key="",
)
except Exception:
verbose_logger.exception(
"Failed to run pre_call for MCP tool logging"
)
standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = {
"name": sanitized_tool_name,
"arguments": parsed_arguments,
"namespaced_tool_name": tool_name,
}
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(
tool_name
)
if mcp_server:
mcp_info = mcp_server.mcp_info or {}
standard_logging_mcp_tool_call["mcp_server_name"] = (
mcp_info.get("server_name")
or getattr(mcp_server, "server_name", None)
or server_name
)
logo_url = mcp_info.get("logo_url")
if logo_url:
standard_logging_mcp_tool_call["mcp_server_logo_url"] = logo_url
cost_info = mcp_info.get("mcp_server_cost_info")
if cost_info:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = cost_info
if litellm_logging_obj:
litellm_logging_obj.model_call_details[
"mcp_tool_call_metadata"
] = standard_logging_mcp_tool_call
litellm_logging_obj.model = f"MCP: {tool_name}"
litellm_logging_obj.call_type = CallTypes.call_mcp_tool.value
result = await global_mcp_server_manager.call_tool(
server_name=server_name,
name=sanitized_tool_name,
@ -526,6 +635,26 @@ class LiteLLM_Proxy_MCP_Handler:
proxy_logging_obj=proxy_logging_obj,
)
if litellm_logging_obj:
try:
litellm_logging_obj.post_call(original_response=result)
end_time = datetime.now()
await litellm_logging_obj.async_post_mcp_tool_call_hook(
kwargs=litellm_logging_obj.model_call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
)
await litellm_logging_obj.async_success_handler(
result=result,
start_time=start_time,
end_time=end_time,
)
except Exception:
verbose_logger.exception(
"Failed to log MCP tool call success for %s", tool_name
)
# Format result for inclusion in response
result_text = LiteLLM_Proxy_MCP_Handler._parse_mcp_result(result)
tool_results.append(
@ -537,6 +666,12 @@ class LiteLLM_Proxy_MCP_Handler:
)
except BlockedPiiEntityError as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
request_data=logging_request_data,
error=e,
)
verbose_logger.error(
f"BlockedPiiEntityError in MCP tool call: {str(e)}"
)
@ -549,6 +684,12 @@ class LiteLLM_Proxy_MCP_Handler:
}
)
except GuardrailRaisedException as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
request_data=logging_request_data,
error=e,
)
verbose_logger.error(
f"GuardrailRaisedException in MCP tool call: {str(e)}"
)
@ -561,12 +702,28 @@ class LiteLLM_Proxy_MCP_Handler:
}
)
except HTTPException as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
request_data=logging_request_data,
error=e,
)
verbose_logger.error(f"HTTPException in MCP tool call: {str(e)}")
error_message = f"Tool call failed: {str(e.detail) if hasattr(e, 'detail') else 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 Exception as e:
await LiteLLM_Proxy_MCP_Handler._log_mcp_tool_failure(
proxy_logging_obj=proxy_logging_obj,
user_api_key_auth=user_api_key_auth,
request_data=logging_request_data,
error=e,
)
verbose_logger.exception(f"Error executing MCP tool call: {e}")
tool_results.append(
{
@ -718,6 +875,33 @@ class LiteLLM_Proxy_MCP_Handler:
**call_params,
)
@staticmethod
async def _log_mcp_tool_failure(
*,
proxy_logging_obj: Optional["ProxyLogging"],
user_api_key_auth: Any,
request_data: Dict[str, Any],
error: Exception,
) -> None:
"""Log MCP tool failures via proxy logging hooks."""
if proxy_logging_obj is None or user_api_key_auth is None:
return
try:
traceback_str = traceback.format_exc(
limit=MAXIMUM_TRACEBACK_LINES_TO_LOG
)
await proxy_logging_obj.post_call_failure_hook(
request_data=request_data,
original_exception=error,
user_api_key_dict=user_api_key_auth,
route="/responses/mcp/call_tool",
traceback_str=traceback_str,
)
except Exception:
verbose_logger.exception("Failed to log MCP tool call failure")
@staticmethod
def _create_mcp_streaming_response(
input: Union[str, Any],
@ -758,7 +942,7 @@ class LiteLLM_Proxy_MCP_Handler:
mcp_events=mcp_discovery_events, # Pre-generated MCP discovery events
tool_server_map=tool_server_map,
mcp_tools_with_litellm_proxy=mcp_tools_with_litellm_proxy,
user_api_key_auth=kwargs.get("user_api_key_auth"),
user_api_key_auth=kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth"),
original_request_params=request_params,
)

4
litellm/responses/mcp/mcp_streaming_iterator.py
View file

@ -298,6 +298,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
self.custom_llm_provider = self.original_request_params.get(
"custom_llm_provider", None
)
self.litellm_call_id = self.original_request_params.get("litellm_call_id")
self.litellm_trace_id = self.original_request_params.get("litellm_trace_id")
self._extract_mcp_headers_from_params()
@ -568,6 +570,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
mcp_server_auth_headers=self.mcp_server_auth_headers,
oauth2_headers=self.oauth2_headers,
raw_headers=self.raw_headers,
litellm_call_id=self.litellm_call_id,
litellm_trace_id=self.litellm_trace_id,
)
# Create completion events and output_item.done events for tool execution