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 78

fix: fix linting error

Browse source
This commit is contained in:
Krrish Dholakia 2025-09-27 12:35:39 -07:00
parent b0e6b8f0ee
commit 69efae55cf
2 changed files with 28 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/proxy/_experimental/mcp_server/cost_calculator.py
View file

@ -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

4
litellm/proxy/auth/oauth2_proxy_hook.py
View file

@ -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}")