From 5cad0dd94b090c2fb309cedad0b407b178d311ce Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 7 Jul 2025 21:08:10 -0700 Subject: [PATCH] [Feat] Add MCP Cost Tracking (#12385) * fix MCP_TOOL_NAME_PREFIX * test_mcp_cost_tracking * init MCPCostCalculator * add call_mcp_tool * calculate_mcp_tool_call_cost * add MCPServerCostInfo * add mcp_server_cost_info * add mcp_server_cost_info * pass through litellm_logging_obj * logged_standard_logging_payload * fix logging MCP tool call * test_mcp_cost_tracking_per_tool * fix NewMCPServerRequest * fix add_update_server * add MCP info to schema.prisma * fix create_mcp_server * working custom cost per call * fix mcp server cost * fix MCPCostCalculator * TestMCPCostCalculator --- .../litellm_proxy_extras/schema.prisma | 1 + litellm/cost_calculator.py | 14 +- litellm/litellm_core_utils/litellm_logging.py | 47 ++-- .../mcp_server/cost_calculator.py | 50 ++++ litellm/proxy/_experimental/mcp_server/db.py | 8 + .../mcp_server/mcp_server_manager.py | 9 +- .../proxy/_experimental/mcp_server/server.py | 22 +- litellm/proxy/_types.py | 5 +- litellm/proxy/schema.prisma | 1 + litellm/types/mcp.py | 15 +- .../types/mcp_server/mcp_server_manager.py | 2 + litellm/types/utils.py | 13 +- schema.prisma | 1 + tests/mcp_tests/test_mcp_logging.py | 245 ++++++++++++++++++ .../mcp_server/test_mcp_cost_calculator.py | 85 ++++++ .../mcp_tools/create_mcp_server.tsx | 32 ++- .../mcp_tools/mcp_server_cost_config.tsx | 83 ++++++ .../components/mcp_tools/mcp_server_edit.tsx | 126 ++++++--- .../src/components/mcp_tools/types.tsx | 8 + 19 files changed, 696 insertions(+), 71 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/cost_calculator.py create mode 100644 tests/mcp_tests/test_mcp_logging.py create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_cost_calculator.py create mode 100644 ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_config.tsx diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 9b0fbbaa8f2..32c50e7ffeb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -175,6 +175,7 @@ model LiteLLM_MCPServerTable { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + mcp_info Json? @default("{}") } // Generate Tokens for Proxy diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index d65411e0f98..d0a5cac14e4 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -95,7 +95,11 @@ from litellm.utils import ( ) if TYPE_CHECKING: - pass + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LitellmLoggingObject, + ) +else: + LitellmLoggingObject = Any def _cost_per_token_custom_pricing_helper( @@ -594,6 +598,7 @@ def completion_cost( # noqa: PLR0915 standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, litellm_model_name: Optional[str] = None, router_model_id: Optional[str] = None, + litellm_logging_obj: Optional[LitellmLoggingObject] = None, ) -> float: """ Calculate the cost of a given completion call fot GPT-3.5-turbo, llama2, any litellm supported llm. @@ -837,6 +842,11 @@ def completion_cost( # noqa: PLR0915 custom_llm_provider=custom_llm_provider, litellm_model_name=model, ) + elif call_type == CallTypes.call_mcp_tool.value: + from litellm.proxy._experimental.mcp_server.cost_calculator import ( + MCPCostCalculator, + ) + return MCPCostCalculator.calculate_mcp_tool_call_cost(litellm_logging_obj=litellm_logging_obj) # Calculate cost based on prompt_tokens, completion_tokens if ( "togethercomputer" in model @@ -999,6 +1009,7 @@ def response_cost_calculator( standard_built_in_tools_params: Optional[StandardBuiltInToolsParams] = None, litellm_model_name: Optional[str] = None, router_model_id: Optional[str] = None, + litellm_logging_obj: Optional[LitellmLoggingObject] = None, ) -> float: """ Returns @@ -1031,6 +1042,7 @@ def response_cost_calculator( standard_built_in_tools_params=standard_built_in_tools_params, litellm_model_name=litellm_model_name, router_model_id=router_model_id, + litellm_logging_obj=litellm_logging_obj, ) return response_cost except Exception as e: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 42d91ab6d42..0b5fea85176 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1140,6 +1140,7 @@ class Logging(LiteLLMLoggingBaseClass): "prompt": prompt, "standard_built_in_tools_params": self.standard_built_in_tools_params, "router_model_id": router_model_id, + "litellm_logging_obj": self, } except Exception as e: # error creating kwargs for cost calculation debug_info = StandardLoggingModelCostFailureDebugInformation( @@ -1345,22 +1346,7 @@ class Logging(LiteLLMLoggingBaseClass): and result is not None and self.stream is not True ): - if ( - isinstance(logging_result, ModelResponse) - or isinstance(logging_result, ModelResponseStream) - or isinstance(logging_result, EmbeddingResponse) - or isinstance(logging_result, ImageResponse) - or isinstance(logging_result, TranscriptionResponse) - or isinstance(logging_result, TextCompletionResponse) - or isinstance(logging_result, HttpxBinaryResponseContent) # tts - or isinstance(logging_result, RerankResponse) - or isinstance(logging_result, FineTuningJob) - or isinstance(logging_result, LiteLLMBatch) - or isinstance(logging_result, ResponsesAPIResponse) - or isinstance(logging_result, OpenAIFileObject) - or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) - or isinstance(logging_result, OpenAIModerationResponse) - ): + if self._is_recognized_call_type_for_logging(logging_result=logging_result): ## HIDDEN PARAMS ## hidden_params = getattr(logging_result, "_hidden_params", {}) if hidden_params: @@ -1456,6 +1442,35 @@ class Logging(LiteLLMLoggingBaseClass): return start_time, end_time, result except Exception as e: raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}") + + def _is_recognized_call_type_for_logging( + self, + logging_result: Any, + ): + """ + Returns True if the call type is recognized for logging (eg. ModelResponse, ModelResponseStream, etc.) + """ + if ( + isinstance(logging_result, ModelResponse) + or isinstance(logging_result, ModelResponseStream) + or isinstance(logging_result, EmbeddingResponse) + or isinstance(logging_result, ImageResponse) + or isinstance(logging_result, TranscriptionResponse) + or isinstance(logging_result, TextCompletionResponse) + or isinstance(logging_result, HttpxBinaryResponseContent) # tts + or isinstance(logging_result, RerankResponse) + or isinstance(logging_result, FineTuningJob) + or isinstance(logging_result, LiteLLMBatch) + or isinstance(logging_result, ResponsesAPIResponse) + or isinstance(logging_result, OpenAIFileObject) + or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject) + or isinstance(logging_result, OpenAIModerationResponse) + or ( + self.call_type == CallTypes.call_mcp_tool.value + ) + ): + return True + return False def _flush_passthrough_collected_chunks_helper( self, diff --git a/litellm/proxy/_experimental/mcp_server/cost_calculator.py b/litellm/proxy/_experimental/mcp_server/cost_calculator.py new file mode 100644 index 00000000000..da47682be6b --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/cost_calculator.py @@ -0,0 +1,50 @@ +""" +Cost calculator for MCP tools. +""" +from typing import TYPE_CHECKING, Any, Optional, cast + +from litellm.types.mcp import MCPServerCostInfo +from litellm.types.utils import StandardLoggingMCPToolCall + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LitellmLoggingObject, + ) +else: + LitellmLoggingObject = Any + +class MCPCostCalculator: + @staticmethod + def calculate_mcp_tool_call_cost(litellm_logging_obj: Optional[LitellmLoggingObject]) -> float: + """ + Calculate the cost of an MCP tool call. + + Default is 0.0, unless user specifies a custom cost per request for MCP tools. + """ + if litellm_logging_obj is None: + return 0.0 + + ######################################################### + # 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 {} + ######################################################### + # 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 {} + 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 + # 3. Default to 0.0 if no cost per query is found + ######################################################### + cost_per_query: float = 0.0 + if tool_name in tool_name_to_cost_per_query: + cost_per_query = tool_name_to_cost_per_query[tool_name] + elif default_cost_per_query is not None: + cost_per_query = default_cost_per_query + return cost_per_query diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 605b1b6792d..6bb2c5a1985 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -215,14 +215,22 @@ async def create_mcp_server( """ Create a new mcp server record in the db """ + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps if data.server_id is None: data.server_id = str(uuid.uuid4()) + + # json dumps mcp_info + mcp_info: Optional[str] = None + if data.mcp_info is not None: + mcp_info = safe_dumps(data.mcp_info) + del data.mcp_info mcp_server_record = await prisma_client.db.litellm_mcpservertable.create( data={ **data.model_dump(), "created_by": touched_by, "updated_by": touched_by, + "mcp_info": mcp_info, } ) return mcp_server_record diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 6cafaeec3c4..932146820ee 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -20,6 +20,12 @@ from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) +from litellm.proxy._experimental.mcp_server.utils import ( + add_server_prefix_to_tool_name, + get_server_name_prefix_tool_mcp, + is_tool_name_prefixed, + normalize_server_name, +) from litellm.proxy._types import ( LiteLLM_MCPServerTable, MCPAuthType, @@ -30,7 +36,6 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer -from litellm.proxy._experimental.mcp_server.utils import add_server_prefix_to_tool_name, normalize_server_name, get_server_name_prefix_tool_mcp, is_tool_name_prefixed class MCPServerManager: @@ -121,6 +126,7 @@ class MCPServerManager: def add_update_server(self, mcp_server: LiteLLM_MCPServerTable): if mcp_server.server_id not in self.get_registry(): + _mcp_info: MCPInfo = mcp_server.mcp_info or {} new_server = MCPServer( server_id=mcp_server.server_id, name=mcp_server.alias or mcp_server.server_id, @@ -131,6 +137,7 @@ class MCPServerManager: mcp_info=MCPInfo( server_name=mcp_server.alias or mcp_server.server_id, description=mcp_server.description, + mcp_server_cost_info=_mcp_info.get("mcp_server_cost_info", None), ), ) self.registry[mcp_server.server_id] = new_server diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 930d099f83c..7362ec31dbe 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -23,7 +23,7 @@ from litellm.proxy._experimental.mcp_server.utils import ( normalize_server_name, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer from litellm.types.utils import StandardLoggingMCPToolCall from litellm.utils import client @@ -310,7 +310,7 @@ if MCP_AVAILABLE: ) # Remove prefix from tool name for logging and processing - original_tool_name, server_name_from_prefix = get_server_name_prefix_tool_mcp( + original_tool_name, _ = get_server_name_prefix_tool_mcp( name) standard_logging_mcp_tool_call: StandardLoggingMCPToolCall = ( @@ -326,15 +326,20 @@ if MCP_AVAILABLE: litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = ( standard_logging_mcp_tool_call ) - litellm_logging_obj.model_call_details["model"] = ( - f"{MCP_TOOL_NAME_PREFIX}: {standard_logging_mcp_tool_call.get('name') or ''}" - ) + model_name = f"MCP: {MCP_TOOL_NAME_PREFIX}: {standard_logging_mcp_tool_call.get('name') or ''}" + litellm_logging_obj.model = model_name + litellm_logging_obj.model_call_details["model"] = model_name litellm_logging_obj.model_call_details["custom_llm_provider"] = ( standard_logging_mcp_tool_call.get("mcp_server_name") ) - + ######################################################### + # Managed MCP Server Tool # Try managed server tool first (pass the full prefixed name) - if name in global_mcp_server_manager.tool_name_to_mcp_server_name_mapping: + # Primary and recommended way to use MCP servers + ######################################################### + mcp_server: Optional[MCPServer] = global_mcp_server_manager._get_mcp_server_from_tool_name(name) + if mcp_server: + standard_logging_mcp_tool_call["mcp_server_cost_info"] = (mcp_server.mcp_info or {}).get("mcp_server_cost_info") return await _handle_managed_mcp_tool( name=name, # Pass the full name (potentially prefixed) arguments=arguments, @@ -343,6 +348,9 @@ if MCP_AVAILABLE: ) # Fall back to local tool registry (use original name) + ######################################################### + # Deprecated: Local MCP Server Tool + ######################################################### return await _handle_local_mcp_tool(original_tool_name, arguments) def _get_standard_logging_mcp_tool_call( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index aa2f3b7a3a8..0cee7907c8f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -24,6 +24,7 @@ from litellm.types.mcp import ( MCPTransport, MCPTransportType, ) +from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem from litellm.types.utils import ( @@ -845,6 +846,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): spec_version: MCPSpecVersionType = MCPSpecVersion.mar_2025 auth_type: Optional[MCPAuthType] = None url: str + mcp_info: Optional[MCPInfo] = None class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): @@ -855,6 +857,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): spec_version: MCPSpecVersionType = MCPSpecVersion.mar_2025 auth_type: Optional[MCPAuthType] = None url: str + mcp_info: Optional[MCPInfo] = None class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): @@ -871,7 +874,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): created_by: Optional[str] = None updated_at: Optional[datetime] = None updated_by: Optional[str] = None - + mcp_info: Optional[MCPInfo] = None class NewUserRequestTeam(LiteLLMPydanticObjectBase): team_id: str diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 9b0fbbaa8f2..32c50e7ffeb 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -175,6 +175,7 @@ model LiteLLM_MCPServerTable { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + mcp_info Json? @default("{}") } // Generate Tokens for Proxy diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 988b0ef8e8a..5f8876946b8 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -1,5 +1,5 @@ import enum -from typing import Literal, Optional +from typing import Dict, Literal, Optional from pydantic import BaseModel, ConfigDict from typing_extensions import TypedDict @@ -27,3 +27,16 @@ MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025] MCPAuthType = Optional[ Literal[MCPAuth.none, MCPAuth.api_key, MCPAuth.bearer_token, MCPAuth.basic] ] + + + +class MCPServerCostInfo(TypedDict, total=False): + default_cost_per_query: Optional[float] + """ + Default cost per query for the MCP server tool call + """ + + tool_name_to_cost_per_query: Optional[Dict[str, float]] + """ + Granular, set a custom cost for each tool in the MCP server + """ \ No newline at end of file diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 2a7c5e97366..1eeab925e2e 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -4,12 +4,14 @@ from pydantic import BaseModel, ConfigDict from typing_extensions import TypedDict from litellm.proxy._types import MCPAuthType, MCPSpecVersionType, MCPTransportType +from litellm.types.mcp import MCPServerCostInfo class MCPInfo(TypedDict, total=False): server_name: str description: Optional[str] logo_url: Optional[str] + mcp_server_cost_info: Optional[MCPServerCostInfo] class MCPServer(BaseModel): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 4ff678d329c..515eca1e060 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -37,6 +37,7 @@ from litellm.types.llms.base import ( BaseLiteLLMOpenAIResponseObject, LiteLLMPydanticObjectBase, ) +from litellm.types.mcp import MCPServerCostInfo from ..litellm_core_utils.core_helpers import map_finish_reason from .guardrails import GuardrailEventHooks @@ -273,6 +274,11 @@ class CallTypes(Enum): generate_content_stream = "generate_content_stream" agenerate_content_stream = "agenerate_content_stream" + ######################################################### + # MCP Call Types + ######################################################### + call_mcp_tool = "call_mcp_tool" + CallTypesLiteral = Literal[ "embedding", @@ -1815,7 +1821,7 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_request_route: Optional[str] -class StandardLoggingMCPToolCall(TypedDict, total=False): +class StandardLoggingMCPToolCall(TypedDict, total=False): name: str """ Name of the tool to call @@ -1841,6 +1847,11 @@ class StandardLoggingMCPToolCall(TypedDict, total=False): (this is to render the logo on the logs page on litellm ui) """ + mcp_server_cost_info: Optional[MCPServerCostInfo] + """ + Cost per query for the MCP server tool call + """ + class StandardLoggingVectorStoreRequest(TypedDict, total=False): """ diff --git a/schema.prisma b/schema.prisma index 9b0fbbaa8f2..32c50e7ffeb 100644 --- a/schema.prisma +++ b/schema.prisma @@ -175,6 +175,7 @@ model LiteLLM_MCPServerTable { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + mcp_info Json? @default("{}") } // Generate Tokens for Proxy diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py new file mode 100644 index 00000000000..6fed5628072 --- /dev/null +++ b/tests/mcp_tests/test_mcp_logging.py @@ -0,0 +1,245 @@ +import os +import sys +import pytest +import asyncio +import json +from unittest.mock import AsyncMock, patch + + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path +import litellm +from litellm.types.utils import StandardLoggingPayload +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._experimental.mcp_server.server import ( + mcp_server_tool_call, +) +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, +) +from mcp.types import Tool as MCPTool, CallToolResult, TextContent + + +class TestMCPLogger(CustomLogger): + def __init__(self): + self.standard_logging_payload = None + super().__init__() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + print("success event") + self.standard_logging_payload = kwargs.get("standard_logging_object", None) + print(f"Captured standard_logging_payload: {self.standard_logging_payload}") + + +@pytest.mark.asyncio +async def test_mcp_cost_tracking(): + # Create a mock tool call result + litellm.logging_callback_manager._reset_all_callbacks() + mock_result = CallToolResult( + content=[TextContent(type="text", text="Test response")], + isError=False + ) + + # Create a mock MCPClient + mock_client = AsyncMock() + mock_client.call_tool = AsyncMock(return_value=mock_result) + mock_client.list_tools = AsyncMock(return_value=[ + MCPTool( + name="add_tools", + description="Test tool", + inputSchema={"type": "object", "properties": {"test": {"type": "string"}}} + ) + ]) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + # Mock the MCPClient constructor + def mock_client_constructor(*args, **kwargs): + return mock_client + + # Initialize the server manager + local_mcp_server_manager = MCPServerManager() + + with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient', mock_client_constructor): + # Load the server config + local_mcp_server_manager.load_servers_from_config( + mcp_servers_config={ + "zapier_gmail_server": { + "url": os.getenv("ZAPIER_MCP_HTTPS_SERVER_URL"), + "mcp_info": { + "mcp_server_cost_info": { + "default_cost_per_query": 1.2, + } + } + } + } + ) + + # Set up the test logger + test_logger = TestMCPLogger() + litellm.callbacks = [test_logger] + + # Initialize the tool mapping + await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() + + # Patch the global manager in both modules where it's used + with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager', local_mcp_server_manager), \ + patch('litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager', local_mcp_server_manager): + + print("tool_name_to_mcp_server_name_mapping", local_mcp_server_manager.tool_name_to_mcp_server_name_mapping) + + # Call mcp tool + response = await mcp_server_tool_call( + name="zapier_gmail_server/add_tools", # Use prefixed name + arguments={ + "test": "test" + } + ) + + # wait 1-2 seconds for logging to be processed + await asyncio.sleep(2) + + logged_standard_logging_payload = test_logger.standard_logging_payload + print("logged_standard_logging_payload", json.dumps(logged_standard_logging_payload, indent=4)) + + # Add assertions + assert response is not None + response_list = list(response) # Convert iterable to list + assert len(response_list) == 1 + assert isinstance(response_list[0], TextContent) + assert response_list[0].text == "Test response" + + # Verify client methods were called + mock_client.__aenter__.assert_called() + mock_client.call_tool.assert_called_once() + + ###### + # verify response cost is 1.2 as set on default_cost_per_query + # Critical - the cost is tracked as $1.2 + assert logged_standard_logging_payload is not None, "Standard logging payload should not be None" + assert logged_standard_logging_payload["response_cost"] == 1.2 + + +@pytest.mark.asyncio +async def test_mcp_cost_tracking_per_tool(): + """Test that individual tool costs are tracked correctly when tool_name_to_cost_per_query is configured""" + # Create a mock tool call result + litellm.logging_callback_manager._reset_all_callbacks() + mock_result = CallToolResult( + content=[TextContent(type="text", text="Test response")], + isError=False + ) + + # Create a mock MCPClient + mock_client = AsyncMock() + mock_client.call_tool = AsyncMock(return_value=mock_result) + mock_client.list_tools = AsyncMock(return_value=[ + MCPTool( + name="expensive_tool", + description="Expensive tool", + inputSchema={"type": "object", "properties": {"data": {"type": "string"}}} + ), + MCPTool( + name="cheap_tool", + description="Cheap tool", + inputSchema={"type": "object", "properties": {"data": {"type": "string"}}} + ) + ]) + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=None) + + # Mock the MCPClient constructor + def mock_client_constructor(*args, **kwargs): + return mock_client + + # Initialize the server manager + local_mcp_server_manager = MCPServerManager() + + with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient', mock_client_constructor): + # Load the server config with per-tool costs + local_mcp_server_manager.load_servers_from_config( + mcp_servers_config={ + "test_server": { + "url": os.getenv("ZAPIER_MCP_HTTPS_SERVER_URL"), + "mcp_info": { + "mcp_server_cost_info": { + "default_cost_per_query": 0.5, # Default cost + "tool_name_to_cost_per_query": { + "expensive_tool": 5.0, # High cost tool + "cheap_tool": 0.1 # Low cost tool + } + } + } + } + } + ) + + # Set up the test logger + test_logger = TestMCPLogger() + litellm.callbacks = [test_logger] + + # Initialize the tool mapping + await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping() + + # Patch the global manager in both modules where it's used + with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.global_mcp_server_manager', local_mcp_server_manager), \ + patch('litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager', local_mcp_server_manager): + + print("tool_name_to_mcp_server_name_mapping", local_mcp_server_manager.tool_name_to_mcp_server_name_mapping) + + # Test 1: Call expensive_tool - should cost 5.0 + response1 = await mcp_server_tool_call( + name="test_server/expensive_tool", # Use prefixed name + arguments={ + "data": "test_expensive" + } + ) + + # wait for logging to be processed + await asyncio.sleep(2) + + logged_standard_logging_payload_1 = test_logger.standard_logging_payload + print("logged_standard_logging_payload_1", json.dumps(logged_standard_logging_payload_1, indent=4)) + + # Verify expensive tool cost + assert logged_standard_logging_payload_1 is not None, "Standard logging payload 1 should not be None" + assert logged_standard_logging_payload_1["response_cost"] == 5.0 + + # Reset logger for second test + test_logger.standard_logging_payload = None + + # Test 2: Call cheap_tool - should cost 0.1 + response2 = await mcp_server_tool_call( + name="test_server/cheap_tool", # Use prefixed name + arguments={ + "data": "test_cheap" + } + ) + + # wait for logging to be processed + await asyncio.sleep(2) + + logged_standard_logging_payload_2 = test_logger.standard_logging_payload + print("logged_standard_logging_payload_2", json.dumps(logged_standard_logging_payload_2, indent=4)) + + # Verify cheap tool cost + assert logged_standard_logging_payload_2 is not None, "Standard logging payload 2 should not be None" + assert logged_standard_logging_payload_2["response_cost"] == 0.1 + + # Add basic response assertions + assert response1 is not None + assert response2 is not None + + response_list_1 = list(response1) + response_list_2 = list(response2) + + assert len(response_list_1) == 1 + assert len(response_list_2) == 1 + assert isinstance(response_list_1[0], TextContent) + assert isinstance(response_list_2[0], TextContent) + assert response_list_1[0].text == "Test response" + assert response_list_2[0].text == "Test response" + + # Verify client methods were called twice + assert mock_client.call_tool.call_count == 2 \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_cost_calculator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_cost_calculator.py new file mode 100644 index 00000000000..c4904cead35 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_cost_calculator.py @@ -0,0 +1,85 @@ +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import orjson +import pytest +from fastapi import Request +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from litellm.proxy._experimental.mcp_server.cost_calculator import MCPCostCalculator + + +class TestMCPCostCalculator: + def test_calculate_mcp_tool_call_cost_none_logging_obj(self): + """Test that when litellm_logging_obj is None, it returns 0.0""" + result = MCPCostCalculator.calculate_mcp_tool_call_cost(None) + assert result == 0.0 + + def test_calculate_mcp_tool_call_cost_with_tool_specific_cost(self): + """Test that when a specific tool has a defined cost, it returns that cost""" + # Mock the litellm_logging_obj + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "mcp_tool_call_metadata": { + "name": "search_web", + "mcp_server_cost_info": { + "default_cost_per_query": 0.01, + "tool_name_to_cost_per_query": { + "search_web": 0.05, + "generate_code": 0.03 + } + } + } + } + + result = MCPCostCalculator.calculate_mcp_tool_call_cost(mock_logging_obj) + assert result == 0.05 + + def test_calculate_mcp_tool_call_cost_with_default_cost(self): + """Test that when no tool-specific cost is found, it falls back to default cost""" + # Mock the litellm_logging_obj + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "mcp_tool_call_metadata": { + "name": "unknown_tool", + "mcp_server_cost_info": { + "default_cost_per_query": 0.02, + "tool_name_to_cost_per_query": { + "search_web": 0.05 + } + } + } + } + + result = MCPCostCalculator.calculate_mcp_tool_call_cost(mock_logging_obj) + assert result == 0.02 + + def test_calculate_mcp_tool_call_cost_no_cost_configuration(self): + """Test that when no cost configuration is provided, it returns 0.0""" + # Mock the litellm_logging_obj with minimal metadata + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = { + "mcp_tool_call_metadata": { + "name": "some_tool", + "mcp_server_cost_info": {} + } + } + + result = MCPCostCalculator.calculate_mcp_tool_call_cost(mock_logging_obj) + assert result == 0.0 + + def test_calculate_mcp_tool_call_cost_empty_metadata(self): + """Test that when metadata is empty or missing, it returns 0.0""" + # Mock the litellm_logging_obj with empty model_call_details + mock_logging_obj = MagicMock() + mock_logging_obj.model_call_details = {} + + result = MCPCostCalculator.calculate_mcp_tool_call_cost(mock_logging_obj) + assert result == 0.0 + diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index a88aba85a93..b4037ae5bf5 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -10,7 +10,8 @@ import { import { InfoCircleOutlined } from "@ant-design/icons"; import { Button, TextInput } from "@tremor/react"; import { createMCPServer } from "../networking"; -import { MCPServer } from "./types"; +import { MCPServer, MCPServerCostInfo } from "./types"; +import MCPServerCostConfig from "./mcp_server_cost_config"; import { isAdminRole } from "@/utils/roles"; @@ -30,20 +31,32 @@ const CreateMCPServer: React.FC = ({ }) => { const [form] = Form.useForm(); const [isLoading, setIsLoading] = useState(false); + const [costConfig, setCostConfig] = useState({}); const handleCreate = async (formValues: Record) => { setIsLoading(true); try { - console.log(`formValues: ${JSON.stringify(formValues)}`); + // Prepare the payload with cost configuration + const payload = { + ...formValues, + mcp_info: { + server_name: formValues.alias || formValues.url, + description: formValues.description, + mcp_server_cost_info: Object.keys(costConfig).length > 0 ? costConfig : null + } + }; + + console.log(`Payload: ${JSON.stringify(payload)}`); if (accessToken != null) { const response: MCPServer = await createMCPServer( accessToken, - formValues + payload ); message.success("MCP Server created successfully"); form.resetFields(); + setCostConfig({}); setModalVisible(false); onCreateSuccess(response); } @@ -59,6 +72,7 @@ const CreateMCPServer: React.FC = ({ const handleCancel = () => { form.resetFields(); + setCostConfig({}); setModalVisible(false); }; @@ -111,7 +125,7 @@ const CreateMCPServer: React.FC = ({ className="space-y-6" >
- MCP Server Name @@ -232,6 +246,16 @@ const CreateMCPServer: React.FC = ({
+ {/* Cost Configuration Section */} +
+ +
+
-
- + + + Server Configuration + Cost Configuration + + + +
+ + + + + + + + + + + + + + + + + + +
+ Cancel + +
+
+ + + +
+ + +
+ Cancel + +
+
+ + + ); }; diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 40ea23007ab..c7845fbc987 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -45,10 +45,17 @@ export interface InputSchemaProperty { required?: string[]; } + // Define MCPServerCostInfo for cost tracking + export interface MCPServerCostInfo { + default_cost_per_query?: number | null; + } + // Define MCP provider info export interface MCPInfo { server_name: string; + description?: string; logo_url?: string; + mcp_server_cost_info?: MCPServerCostInfo | null; } // Define the structure for a single MCP tool @@ -114,6 +121,7 @@ export interface MCPServer { transport?: string | null; spec_version?: string | null; auth_type?: string | null; + mcp_info?: MCPInfo | null; created_at: string; created_by: string; updated_at: string;