From 69efae55cf90f0ade941d3b9542ab67e1d164d6f Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 12:35:39 -0700 Subject: [PATCH] fix: fix linting error --- .../mcp_server/cost_calculator.py | 35 ++++++++++++++----- litellm/proxy/auth/oauth2_proxy_hook.py | 4 +-- 2 files changed, 28 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/cost_calculator.py b/litellm/proxy/_experimental/mcp_server/cost_calculator.py index eea10924a11..b8fdba23d92 100644 --- a/litellm/proxy/_experimental/mcp_server/cost_calculator.py +++ b/litellm/proxy/_experimental/mcp_server/cost_calculator.py @@ -1,6 +1,7 @@ """ Cost calculator for MCP tools. """ + from typing import TYPE_CHECKING, Any, Optional, cast from litellm.types.mcp import MCPServerCostInfo @@ -13,11 +14,12 @@ if TYPE_CHECKING: else: LitellmLoggingObject = Any + class MCPCostCalculator: @staticmethod def calculate_mcp_tool_call_cost( litellm_logging_obj: Optional[LitellmLoggingObject], - ) -> float: + ) -> float: """ Calculate the cost of an MCP tool call. @@ -25,28 +27,43 @@ class MCPCostCalculator: """ if litellm_logging_obj is None: return 0.0 - + ######################################################### # Get the response cost from logging object model_call_details # This is set when a user modifies the response in a post_mcp_tool_call_hook ######################################################### - response_cost = litellm_logging_obj.model_call_details.get("response_cost", None) + response_cost = litellm_logging_obj.model_call_details.get( + "response_cost", None + ) if response_cost is not None: return response_cost - + ######################################################### # Unpack the mcp_tool_call_metadata ######################################################### - mcp_tool_call_metadata: StandardLoggingMCPToolCall = cast(StandardLoggingMCPToolCall, litellm_logging_obj.model_call_details.get("mcp_tool_call_metadata", {})) or {} - mcp_server_cost_info: MCPServerCostInfo = mcp_tool_call_metadata.get("mcp_server_cost_info", {}) or {} + mcp_tool_call_metadata: StandardLoggingMCPToolCall = ( + cast( + StandardLoggingMCPToolCall, + litellm_logging_obj.model_call_details.get( + "mcp_tool_call_metadata", {} + ), + ) + or {} + ) + mcp_server_cost_info: MCPServerCostInfo = ( + mcp_tool_call_metadata.get("mcp_server_cost_info") or MCPServerCostInfo() + ) ######################################################### # User defined cost per query ######################################################### - default_cost_per_query = mcp_server_cost_info.get("default_cost_per_query", None) - tool_name_to_cost_per_query: dict = mcp_server_cost_info.get("tool_name_to_cost_per_query", {}) or {} + default_cost_per_query = mcp_server_cost_info.get( + "default_cost_per_query", None + ) + tool_name_to_cost_per_query: dict = ( + mcp_server_cost_info.get("tool_name_to_cost_per_query", {}) or {} + ) tool_name = mcp_tool_call_metadata.get("name", "") - ######################################################### # 1. If tool_name is in tool_name_to_cost_per_query, use the cost per query # 2. If tool_name is not in tool_name_to_cost_per_query, use the default cost per query diff --git a/litellm/proxy/auth/oauth2_proxy_hook.py b/litellm/proxy/auth/oauth2_proxy_hook.py index a1db5d842c4..7e517092b8a 100644 --- a/litellm/proxy/auth/oauth2_proxy_hook.py +++ b/litellm/proxy/auth/oauth2_proxy_hook.py @@ -14,8 +14,8 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth: verbose_proxy_logger.debug("Handling oauth2 proxy request") # Define the OAuth2 config mappings - oauth2_config_mappings: Dict[str, str] = general_settings.get( - "oauth2_config_mappings", None + oauth2_config_mappings: Dict[str, str] = ( + general_settings.get("oauth2_config_mappings") or {} ) verbose_proxy_logger.debug(f"Oauth2 config mappings: {oauth2_config_mappings}")