mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
444 lines
16 KiB
Python
444 lines
16 KiB
Python
import os
|
|
import pytest
|
|
import asyncio
|
|
from typing import Optional
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
|
|
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,
|
|
set_auth_context,
|
|
)
|
|
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
|
|
MCPServerManager,
|
|
)
|
|
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
|
|
from litellm.types.mcp import MCPPostCallResponseObject
|
|
from litellm.types.utils import HiddenParams
|
|
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}")
|
|
|
|
|
|
def _set_authorized_user(server_ids):
|
|
"""Configure auth context with permission to call the specified servers."""
|
|
server_list = list(server_ids)
|
|
user_auth = UserAPIKeyAuth(
|
|
api_key="test",
|
|
user_id="test_user",
|
|
object_permission=LiteLLM_ObjectPermissionTable(
|
|
object_permission_id="mcp-test-permissions",
|
|
mcp_servers=server_list,
|
|
),
|
|
)
|
|
set_auth_context(user_api_key_auth=user_auth, mcp_servers=server_list)
|
|
|
|
|
|
@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 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
|
|
await 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,
|
|
),
|
|
):
|
|
|
|
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())
|
|
|
|
print(
|
|
"tool_name_to_mcp_server_name_mapping",
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping,
|
|
)
|
|
|
|
# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"add_tools"
|
|
] = "zapier_gmail_server"
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"zapier_gmail_server-add_tools"
|
|
] = "zapier_gmail_server"
|
|
|
|
# Call mcp tool
|
|
response = await mcp_server_tool_call(
|
|
name="zapier_gmail_server-add_tools", # Use correct prefixed name with - separator
|
|
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", logged_standard_logging_payload)
|
|
|
|
# Add assertions
|
|
assert response is not None
|
|
# Handle CallToolResult - access .content for the list of content items
|
|
if isinstance(response, CallToolResult):
|
|
response_list = response.content
|
|
else:
|
|
response_list = list(
|
|
response
|
|
) # Convert iterable to list for backward compatibility
|
|
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.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 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
|
|
await 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()
|
|
|
|
# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"expensive_tool"
|
|
] = "test_server"
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"test_server-expensive_tool"
|
|
] = "test_server"
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["cheap_tool"] = (
|
|
"test_server"
|
|
)
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"test_server-cheap_tool"
|
|
] = "test_server"
|
|
|
|
# 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,
|
|
),
|
|
):
|
|
|
|
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())
|
|
|
|
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 correct prefixed name with - separator
|
|
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", logged_standard_logging_payload_1
|
|
)
|
|
|
|
# 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 correct prefixed name with - separator
|
|
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", logged_standard_logging_payload_2
|
|
)
|
|
|
|
# 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.content)
|
|
response_list_2 = list(response2.content)
|
|
|
|
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
|
|
|
|
|
|
class MCPLoggerHook(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}")
|
|
|
|
async def async_post_mcp_tool_call_hook(
|
|
self, kwargs, response_obj: MCPPostCallResponseObject, start_time, end_time
|
|
) -> Optional[MCPPostCallResponseObject]:
|
|
print("post mcp tool call response_obj", response_obj)
|
|
# update the MCPPostCallResponseObject with the response_cost
|
|
response_obj.hidden_params.response_cost = 1.42
|
|
return response_obj
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mcp_tool_call_hook():
|
|
# 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 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
|
|
await local_mcp_server_manager.load_servers_from_config(
|
|
mcp_servers_config={
|
|
"zapier_gmail_server": {
|
|
"url": os.getenv("ZAPIER_MCP_HTTPS_SERVER_URL"),
|
|
}
|
|
}
|
|
)
|
|
|
|
# Set up the test logger
|
|
test_logger = MCPLoggerHook()
|
|
litellm.callbacks = [test_logger]
|
|
|
|
# Initialize the tool mapping
|
|
await local_mcp_server_manager._initialize_tool_name_to_mcp_server_name_mapping()
|
|
|
|
# Manually add the tool mapping to ensure it's available (since mocking might not capture it properly)
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping["add_tools"] = (
|
|
"zapier_gmail_server"
|
|
)
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping[
|
|
"zapier_gmail_server-add_tools"
|
|
] = "zapier_gmail_server"
|
|
|
|
# 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,
|
|
),
|
|
):
|
|
|
|
_set_authorized_user(local_mcp_server_manager.get_all_mcp_server_ids())
|
|
|
|
print(
|
|
"tool_name_to_mcp_server_name_mapping",
|
|
local_mcp_server_manager.tool_name_to_mcp_server_name_mapping,
|
|
)
|
|
|
|
# Call mcp tool using the correct separator format (- not /)
|
|
response = await mcp_server_tool_call(
|
|
name="zapier_gmail_server-add_tools", # Use correct prefixed name with - separator
|
|
arguments={"test": "test"},
|
|
)
|
|
|
|
# wait 1-2 seconds for logging to be processed
|
|
await asyncio.sleep(2)
|
|
|
|
# check logged standard logging payload
|
|
logged_standard_logging_payload = test_logger.standard_logging_payload
|
|
print("logged_standard_logging_payload", logged_standard_logging_payload)
|
|
assert (
|
|
logged_standard_logging_payload is not None
|
|
), "Standard logging payload should not be None"
|
|
assert logged_standard_logging_payload["response_cost"] == 1.42
|