From 6451f10f2c9f685fde5cd2b463bb2bd2e79bb3ab Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:34:01 -0700 Subject: [PATCH] fix: fix ruff check --- .../mcp_server/mcp_server_manager.py | 140 ++++++++++-------- .../mcp_server/rest_endpoints.py | 6 +- .../proxy/_experimental/mcp_server/server.py | 2 +- 3 files changed, 79 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bc69314bcc..9fb39e7498 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -561,6 +561,77 @@ class MCPServerManager: ) return prefixed_tools + async def pre_call_tool_check( + self, + name: str, + arguments: Dict[str, Any], + server_name_from_prefix: str, + user_api_key_auth: Optional[UserAPIKeyAuth], + proxy_logging_obj: ProxyLogging, + ): + pre_hook_kwargs = { + "name": name, + "arguments": arguments, + "server_name": server_name_from_prefix, + "user_api_key_auth": user_api_key_auth, + "user_api_key_user_id": ( + getattr(user_api_key_auth, "user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_team_id": ( + getattr(user_api_key_auth, "team_id", None) + if user_api_key_auth + else None + ), + "user_api_key_end_user_id": ( + getattr(user_api_key_auth, "end_user_id", None) + if user_api_key_auth + else None + ), + "user_api_key_hash": ( + getattr(user_api_key_auth, "api_key_hash", None) + if user_api_key_auth + else None + ), + } + + # Create MCP request object for processing + mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs( + pre_hook_kwargs + ) + + # Convert to LLM format for existing guardrail compatibility + synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( + mcp_request_obj, pre_hook_kwargs + ) + + try: + # Use standard pre_call_hook with call_type="mcp_call" + modified_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_auth, # type: ignore + data=synthetic_llm_data, + call_type="mcp_call", # type: ignore + ) + if modified_data: + # Convert response back to MCP format and apply modifications + modified_kwargs = ( + proxy_logging_obj._convert_mcp_hook_response_to_kwargs( + modified_data, pre_hook_kwargs + ) + ) + if modified_kwargs.get("arguments") != arguments: + arguments = modified_kwargs["arguments"] + + except ( + BlockedPiiEntityError, + GuardrailRaisedException, + HTTPException, + ) as e: + # Re-raise guardrail exceptions to properly fail the MCP call + verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}") + raise e + async def call_tool( self, name: str, @@ -614,71 +685,14 @@ class MCPServerManager: # Using standard pre_call_hook with call_type="mcp_call" ######################################################### if proxy_logging_obj: - pre_hook_kwargs = { - "name": name, - "arguments": arguments, - "server_name": server_name_from_prefix, - "user_api_key_auth": user_api_key_auth, - "user_api_key_user_id": ( - getattr(user_api_key_auth, "user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_team_id": ( - getattr(user_api_key_auth, "team_id", None) - if user_api_key_auth - else None - ), - "user_api_key_end_user_id": ( - getattr(user_api_key_auth, "end_user_id", None) - if user_api_key_auth - else None - ), - "user_api_key_hash": ( - getattr(user_api_key_auth, "api_key_hash", None) - if user_api_key_auth - else None - ), - } - - # Create MCP request object for processing - mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs( - pre_hook_kwargs + await self.pre_call_tool_check( + name=original_tool_name, + arguments=arguments, + server_name_from_prefix=server_name_from_prefix, + user_api_key_auth=user_api_key_auth, + proxy_logging_obj=proxy_logging_obj, ) - # Convert to LLM format for existing guardrail compatibility - synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format( - mcp_request_obj, pre_hook_kwargs - ) - - try: - # Use standard pre_call_hook with call_type="mcp_call" - modified_data = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_auth, # type: ignore - data=synthetic_llm_data, - call_type="mcp_call", # type: ignore - ) - if modified_data: - # Convert response back to MCP format and apply modifications - modified_kwargs = ( - proxy_logging_obj._convert_mcp_hook_response_to_kwargs( - modified_data, pre_hook_kwargs - ) - ) - if modified_kwargs.get("arguments") != arguments: - arguments = modified_kwargs["arguments"] - - except ( - BlockedPiiEntityError, - GuardrailRaisedException, - HTTPException, - ) as e: - # Re-raise guardrail exceptions to properly fail the MCP call - verbose_logger.error( - f"Guardrail blocked MCP tool call pre call: {str(e)}" - ) - raise e - # Get server-specific auth header if available server_auth_header = None if mcp_server_auth_headers and mcp_server.alias: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d2f16154f3..ef770a9d43 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,11 +1,7 @@ import importlib from typing import Dict, List, Optional -from urllib.parse import urlencode, urlparse, urlunparse -import httpx -from fastapi import APIRouter, Depends, HTTPException, Query, Request -from fastapi.params import Form -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi import APIRouter, Depends, Query, Request from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7f16179390..38b0cfd99a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -364,7 +364,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: Optional[Dict[str, str]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: - f""" + """ Helper method to fetch tools from MCP servers based on server filtering criteria. Args: