scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

[Feat] Add MCP Cost Tracking (#12385)

Browse source 
* 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
This commit is contained in:
Ishaan Jaff 2025-07-07 21:08:10 -07:00 • committed by GitHub
parent 85bc1065a0
commit 5cad0dd94b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 696 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

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

14
litellm/cost_calculator.py
View file

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

47
litellm/litellm_core_utils/litellm_logging.py
View file

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

50
litellm/proxy/_experimental/mcp_server/cost_calculator.py Normal file
View file

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

8
litellm/proxy/_experimental/mcp_server/db.py
View file

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

9
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

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

22
litellm/proxy/_experimental/mcp_server/server.py
View file

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

5
litellm/proxy/_types.py
View file

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

1
litellm/proxy/schema.prisma
View file

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

15
litellm/types/mcp.py
View file

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

2
litellm/types/mcp_server/mcp_server_manager.py
View file

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

13
litellm/types/utils.py
View file

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

1
schema.prisma
View file

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

245
tests/mcp_tests/test_mcp_logging.py Normal file
View file

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

85
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_cost_calculator.py Normal file
View file

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

32
ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx
View file

@ -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<CreateMCPServerProps> = ({
}) => {
const [form] = Form.useForm();
const [isLoading, setIsLoading] = useState(false);
const [costConfig, setCostConfig] = useState<MCPServerCostInfo>({});
const handleCreate = async (formValues: Record<string, any>) => {
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<CreateMCPServerProps> = ({
const handleCancel = () => {
form.resetFields();
setCostConfig({});
setModalVisible(false);
};
@ -111,7 +125,7 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
className="space-y-6"
>
<div className="grid grid-cols-1 gap-6">
<Form.Item
<Form.Item
label={
<span className="text-sm font-medium text-gray-700 flex items-center">
MCP Server Name
@ -232,6 +246,16 @@ const CreateMCPServer: React.FC<CreateMCPServerProps> = ({
</Form.Item>
</div>
{/* Cost Configuration Section */}
<div className="mt-8 pt-6 border-t border-gray-200">
<MCPServerCostConfig
value={costConfig}
onChange={setCostConfig}
accessToken={accessToken}
disabled={false}
/>
</div>
<div className="flex items-center justify-end space-x-3 pt-6 border-t border-gray-100">
<Button
variant="secondary"

83
ui/litellm-dashboard/src/components/mcp_tools/mcp_server_cost_config.tsx Normal file
View file

@ -0,0 +1,83 @@
import React from "react";
import { Tooltip, InputNumber } from "antd";
import { InfoCircleOutlined, DollarOutlined } from "@ant-design/icons";
import { Card, Title, Text } from "@tremor/react";
import { MCPServerCostInfo } from "./types";
interface MCPServerCostConfigProps {
value?: MCPServerCostInfo;
onChange?: (value: MCPServerCostInfo) => void;
serverId?: string;
serverUrl?: string;
accessToken: string | null;
disabled?: boolean;
}
const MCPServerCostConfig: React.FC<MCPServerCostConfigProps> = ({
value = {},
onChange,
serverId,
serverUrl,
accessToken,
disabled = false
}) => {
const handleDefaultCostChange = (defaultCost: number | null) => {
const updated = {
...value,
default_cost_per_query: defaultCost
};
onChange?.(updated);
};
return (
<Card>
<div className="space-y-6">
<div className="flex items-center gap-2 mb-4">
<DollarOutlined className="text-green-600" />
<Title>Cost Configuration</Title>
<Tooltip title="Configure costs for this MCP server's tool calls. These costs will be tracked when the server's tools are used.">
<InfoCircleOutlined className="text-gray-400" />
</Tooltip>
</div>
<div className="space-y-4">
<div>
<label className="block text-sm font-medium text-gray-700 mb-2">
Default Cost per Query ($)
<Tooltip title="Default cost charged for each tool call to this server.">
<InfoCircleOutlined className="ml-1 text-gray-400" />
</Tooltip>
</label>
<InputNumber
min={0}
step={0.0001}
precision={4}
placeholder="0.0000"
value={value.default_cost_per_query}
onChange={handleDefaultCostChange}
disabled={disabled}
style={{ width: '200px' }}
addonBefore="$"
/>
<Text className="block mt-1 text-gray-500 text-sm">
Set a default cost for all tool calls to this server
</Text>
</div>
</div>
{value.default_cost_per_query && (
<div className="mt-6 p-4 bg-blue-50 border border-blue-200 rounded-lg">
<Text className="text-blue-800 font-medium">Cost Summary:</Text>
<div className="mt-2 space-y-1">
<Text className="text-blue-700">
• Default cost: ${value.default_cost_per_query.toFixed(4)} per query
</Text>
</div>
</div>
)}
</div>
</Card>
);
};
export default MCPServerCostConfig;

126
ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx
View file

@ -1,8 +1,9 @@
import React from "react";
import React, { useState, useEffect } from "react";
import { Form, Select, Button as AntdButton, message } from "antd";
import { Button, TextInput } from "@tremor/react";
import { MCPServer } from "./types";
import { Button, TextInput, TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react";
import { MCPServer, MCPServerCostInfo } from "./types";
import { updateMCPServer } from "../networking";
import MCPServerCostConfig from "./mcp_server_cost_config";
interface MCPServerEditProps {
mcpServer: MCPServer;
@ -13,11 +14,30 @@ interface MCPServerEditProps {
const MCPServerEdit: React.FC<MCPServerEditProps> = ({ mcpServer, accessToken, onCancel, onSuccess }) => {
const [form] = Form.useForm();
const [costConfig, setCostConfig] = useState<MCPServerCostInfo>({});
// Initialize cost config from existing server data
useEffect(() => {
if (mcpServer.mcp_info?.mcp_server_cost_info) {
setCostConfig(mcpServer.mcp_info.mcp_server_cost_info);
}
}, [mcpServer]);
const handleSave = async (values: Record<string, any>) => {
if (!accessToken) return;
try {
const updated = await updateMCPServer(accessToken, { ...values, server_id: mcpServer.server_id });
// Prepare the payload with cost configuration
const payload = {
...values,
server_id: mcpServer.server_id,
mcp_info: {
server_name: values.alias || values.url,
description: values.description,
mcp_server_cost_info: Object.keys(costConfig).length > 0 ? costConfig : null
}
};
const updated = await updateMCPServer(accessToken, payload);
message.success("MCP Server updated successfully");
onSuccess(updated);
} catch (error: any) {
@ -26,41 +46,69 @@ const MCPServerEdit: React.FC<MCPServerEditProps> = ({ mcpServer, accessToken, o
};
return (
<Form form={form} onFinish={handleSave} initialValues={mcpServer} layout="vertical">
<Form.Item label="MCP Server Name" name="alias">
<TextInput />
</Form.Item>
<Form.Item label="Description" name="description">
<TextInput />
</Form.Item>
<Form.Item label="MCP Server URL" name="url" rules={[{ required: true, message: "Please enter a server URL" }]}>
<TextInput />
</Form.Item>
<Form.Item label="Transport Type" name="transport" rules={[{ required: true }]}>
<Select>
<Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
<Select.Option value="http">HTTP</Select.Option>
</Select>
</Form.Item>
<Form.Item label="Authentication" name="auth_type" rules={[{ required: true }]}>
<Select>
<Select.Option value="none">None</Select.Option>
<Select.Option value="api_key">API Key</Select.Option>
<Select.Option value="bearer_token">Bearer Token</Select.Option>
<Select.Option value="basic">Basic Auth</Select.Option>
</Select>
</Form.Item>
<Form.Item label="MCP Version" name="spec_version" rules={[{ required: true }]}>
<Select>
<Select.Option value="2025-03-26">2025-03-26 (Latest)</Select.Option>
<Select.Option value="2024-11-05">2024-11-05</Select.Option>
</Select>
</Form.Item>
<div className="flex justify-end gap-2">
<AntdButton onClick={onCancel}>Cancel</AntdButton>
<Button type="submit">Save Changes</Button>
</div>
</Form>
<TabGroup>
<TabList className="grid w-full grid-cols-2">
<Tab>Server Configuration</Tab>
<Tab>Cost Configuration</Tab>
</TabList>
<TabPanels className="mt-6">
<TabPanel>
<Form form={form} onFinish={handleSave} initialValues={mcpServer} layout="vertical">
<Form.Item label="MCP Server Name" name="alias">
<TextInput />
</Form.Item>
<Form.Item label="Description" name="description">
<TextInput />
</Form.Item>
<Form.Item label="MCP Server URL" name="url" rules={[{ required: true, message: "Please enter a server URL" }]}>
<TextInput />
</Form.Item>
<Form.Item label="Transport Type" name="transport" rules={[{ required: true }]}>
<Select>
<Select.Option value="sse">Server-Sent Events (SSE)</Select.Option>
<Select.Option value="http">HTTP</Select.Option>
</Select>
</Form.Item>
<Form.Item label="Authentication" name="auth_type" rules={[{ required: true }]}>
<Select>
<Select.Option value="none">None</Select.Option>
<Select.Option value="api_key">API Key</Select.Option>
<Select.Option value="bearer_token">Bearer Token</Select.Option>
<Select.Option value="basic">Basic Auth</Select.Option>
</Select>
</Form.Item>
<Form.Item label="MCP Version" name="spec_version" rules={[{ required: true }]}>
<Select>
<Select.Option value="2025-03-26">2025-03-26 (Latest)</Select.Option>
<Select.Option value="2024-11-05">2024-11-05</Select.Option>
</Select>
</Form.Item>
<div className="flex justify-end gap-2">
<AntdButton onClick={onCancel}>Cancel</AntdButton>
<Button type="submit">Save Changes</Button>
</div>
</Form>
</TabPanel>
<TabPanel>
<div className="space-y-6">
<MCPServerCostConfig
value={costConfig}
onChange={setCostConfig}
serverId={mcpServer.server_id}
serverUrl={mcpServer.url}
accessToken={accessToken}
disabled={false}
/>
<div className="flex justify-end gap-2">
<AntdButton onClick={onCancel}>Cancel</AntdButton>
<Button onClick={() => form.submit()}>Save Changes</Button>
</div>
</div>
</TabPanel>
</TabPanels>
</TabGroup>
);
};

8
ui/litellm-dashboard/src/components/mcp_tools/types.tsx
View file

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