tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 14:23:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #17078 from BerriAI/litellm_add_search_logging

Browse Source 
Add search API logging and cost tracking in LiteLLM Proxy
This commit is contained in:
Sameer Kankute
2025-11-25 23:59:41 +05:30
committed by GitHub
parent 2e50db81a5 67d69d12b0
commit 59bcf079fb
7 changed files with 310 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/cost_calculator.py
+51
View File
@@ -1031,6 +1031,57 @@ def completion_cost( # noqa: PLR0915
billed_units.get("search_units") or 1
) # cohere charges per request by default.
completion_tokens = search_units
elif (
call_type == CallTypes.search.value
or call_type == CallTypes.asearch.value
):
from litellm.search import search_provider_cost_per_query
# Extract number_of_queries from optional_params or default to 1
number_of_queries = 1
if optional_params is not None:
# Check if query is a list (multiple queries)
query = optional_params.get("query")
if isinstance(query, list):
number_of_queries = len(query)
elif query is not None:
number_of_queries = 1
search_model = model or ""
if custom_llm_provider and "/" not in search_model:
# If model is like "tavily-search", construct "tavily/search" for cost lookup
search_model = f"{custom_llm_provider}/search"
prompt_cost, completion_cost_result = search_provider_cost_per_query(
model=search_model,
custom_llm_provider=custom_llm_provider,
number_of_queries=number_of_queries,
optional_params=optional_params,
)
# Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost)
_final_cost = prompt_cost + completion_cost_result
# Apply discount
original_cost = _final_cost
_final_cost, discount_percent, discount_amount = _apply_cost_discount(
base_cost=_final_cost,
custom_llm_provider=custom_llm_provider,
)
# Store cost breakdown in logging object if available
_store_cost_breakdown_in_logging_obj(
litellm_logging_obj=litellm_logging_obj,
prompt_tokens_cost_usd_dollar=prompt_cost,
completion_tokens_cost_usd_dollar=completion_cost_result,
cost_for_built_in_tools_cost_usd_dollar=0.0,
total_cost_usd_dollar=_final_cost,
original_cost=original_cost,
discount_percent=discount_percent,
discount_amount=discount_amount,
)
return _final_cost
elif call_type == CallTypes.arealtime.value and isinstance(
completion_response, LiteLLMRealtimeStreamLoggingObject
):
litellm/litellm_core_utils/litellm_logging.py
+5
View File
@@ -70,6 +70,7 @@ from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_logging,
)
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.containers.main import ContainerObject
from litellm.types.llms.openai import (
@@ -1298,6 +1299,7 @@ class Logging(LiteLLMLoggingBaseClass):
OpenAIFileObject,
LiteLLMRealtimeStreamLoggingObject,
OpenAIModerationResponse,
"SearchResponse",
],
cache_hit: Optional[bool] = None,
litellm_model_name: Optional[str] = None,
@@ -1710,8 +1712,11 @@ class Logging(LiteLLMLoggingBaseClass):
or isinstance(logging_result, LiteLLMRealtimeStreamLoggingObject)
or isinstance(logging_result, OpenAIModerationResponse)
or isinstance(logging_result, OCRResponse) # OCR
or isinstance(logging_result, SearchResponse) # Search API
or isinstance(logging_result, dict)
and logging_result.get("object") == "vector_store.search_results.page"
or isinstance(logging_result, dict)
and logging_result.get("object") == "search" # Search API (dict format)
or isinstance(logging_result, VideoObject)
or isinstance(logging_result, ContainerObject)
or (self.call_type == CallTypes.call_mcp_tool.value)
litellm/proxy/_types.py
+2
View File
@@ -348,6 +348,8 @@ class LiteLLMRoutes(enum.Enum):
# search
"/search",
"/v1/search",
"/search/{search_tool_name}",
"/v1/search/{search_tool_name}",
# OCR
"/ocr",
"/v1/ocr",
litellm/proxy/proxy_server.py
+27 -9
View File
@@ -2010,6 +2010,16 @@ class ProxyConfig:
)
print(f"\033[32m {search_tool_name} ({search_provider})\033[0m") # noqa
# Handle os.environ/ variables in litellm_params
litellm_params = search_tool.get("litellm_params", {})
if litellm_params:
for k, v in litellm_params.items():
if isinstance(v, str) and v.startswith("os.environ/"):
_v = v.replace("os.environ/", "")
v = get_secret(_v)
litellm_params[k] = v
search_tool["litellm_params"] = litellm_params
# Cast to SearchToolTypedDict for type safety
try:
search_tool_typed: SearchToolTypedDict = SearchToolTypedDict(**search_tool) # type: ignore
@@ -3761,6 +3771,7 @@ class ProxyConfig:
async def _init_search_tools_in_db(self, prisma_client: PrismaClient):
"""
Initialize search tools from database into the router on startup.
Only updates router if there are tools in the database, otherwise preserves config-loaded tools.
"""
global llm_router
@@ -3778,17 +3789,24 @@ class ProxyConfig:
f"Loading {len(search_tools)} search tool(s) from database into router"
)
if llm_router is not None:
# Add search tools to the router
await SearchAPIRouter.update_router_search_tools(
router_instance=llm_router, search_tools=search_tools
)
verbose_proxy_logger.info(
f"Successfully loaded {len(search_tools)} search tool(s) into router"
)
# Only update router if there are tools in the database
# This prevents overwriting config-loaded tools with an empty list
if len(search_tools) > 0:
if llm_router is not None:
# Add search tools to the router
await SearchAPIRouter.update_router_search_tools(
router_instance=llm_router, search_tools=search_tools
)
verbose_proxy_logger.info(
f"Successfully loaded {len(search_tools)} search tool(s) into router"
)
else:
verbose_proxy_logger.debug(
"Router not initialized yet, search tools will be added when router is created"
)
else:
verbose_proxy_logger.debug(
"Router not initialized yet, search tools will be added when router is created"
"No search tools found in database, keeping config-loaded search tools (if any)"
)
except Exception as e:
litellm/proxy/search_endpoints/endpoints.py
+21
View File
@@ -131,6 +131,27 @@ async def search(
if search_tool_name is not None:
data["search_tool_name"] = search_tool_name
if "search_tool_name" in data and data["search_tool_name"]:
data["model"] = data["search_tool_name"]
if llm_router is not None and hasattr(llm_router, "search_tools"):
search_tool_name_value = data["search_tool_name"]
matching_tools = [
tool for tool in llm_router.search_tools
if tool.get("search_tool_name") == search_tool_name_value
]
if matching_tools:
search_tool = matching_tools[0]
search_provider = search_tool.get("litellm_params", {}).get("search_provider")
if search_provider:
data["custom_llm_provider"] = search_provider
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["model_group"] = search_tool_name_value
# Process request using ProxyBaseLLMRequestProcessing
processor = ProxyBaseLLMRequestProcessing(data=data)
try:
litellm/proxy/spend_tracking/spend_tracking_utils.py
+2 -1
View File
@@ -228,7 +228,8 @@ def get_logging_payload( # noqa: PLR0915
if call_type in ["ocr", "aocr"]:
usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict)
else:
usage = cast(dict, response_obj).get("usage", None) or {}
# Use response_obj_dict instead of response_obj to avoid calling .get() on Pydantic models
usage = response_obj_dict.get("usage", None) or {}
if isinstance(usage, litellm.Usage):
usage = dict(usage)
tests/proxy_unit_tests/test_search_api_logging.py
+202
View File
@@ -0,0 +1,202 @@
"""
Test search API logging and cost tracking in proxy.
Tests that search API requests are properly logged to LiteLLM_SpendLogs
with correct fields populated (call_type, model, custom_llm_provider,
model_group, spend, etc.)
"""
import asyncio
import os
import sys
import time
from datetime import datetime
from unittest.mock import AsyncMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import Router
from litellm.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
from litellm.proxy.spend_tracking.spend_management_endpoints import view_spend_logs
from litellm.proxy.utils import ProxyLogging, hash_token, update_spend
from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult
@pytest.fixture
def prisma_client():
from litellm.proxy import proxy_server
from litellm.proxy.proxy_cli import append_query_params
from litellm.proxy.utils import PrismaClient
params = {"connection_limit": 100, "pool_timeout": 60}
database_url = os.getenv("DATABASE_URL")
if database_url is None:
pytest.skip("DATABASE_URL not set")
modified_url = append_query_params(database_url, params)
os.environ["DATABASE_URL"] = modified_url
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
proxy_server.litellm_proxy_budget_name = (
f"litellm-proxy-budget-{time.time()}"
)
proxy_server.user_custom_key_generate = None
return prisma_client
@pytest.mark.asyncio
async def test_search_api_logging_and_cost_tracking(prisma_client):
"""
Test that search API requests are logged with correct fields and cost tracking.
Verifies:
1. Search request creates a spend log entry
2. call_type is set to "asearch"
3. model is set to search_tool_name
4. custom_llm_provider is set correctly
5. model_group is set to search_tool_name
6. spend is calculated and logged
"""
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Setup router with search tool
search_tool_name = "tavily-search"
search_provider = "tavily"
router = Router(model_list=[])
router.search_tools = [
{
"search_tool_name": search_tool_name,
"litellm_params": {
"search_provider": search_provider,
},
}
]
setattr(litellm.proxy.proxy_server, "llm_router", router)
# Generate a test API key
from litellm.proxy.management_endpoints.key_management_endpoints import generate_key_fn
from litellm.proxy._types import GenerateKeyRequest
from litellm.proxy._types import LitellmUserRoles
user_api_key_dict = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="test_user",
)
key_request = GenerateKeyRequest(models=[], duration=None)
key_response = await generate_key_fn(
data=key_request, user_api_key_dict=user_api_key_dict
)
generated_key = key_response.key
user_id = key_response.user_id
# Create mock search response
mock_search_result = SearchResult(
title="Test Result",
url="https://example.com",
snippet="Test snippet",
)
mock_search_response = SearchResponse(
object="search",
results=[mock_search_result],
)
# Mock the search function to return our mock response
with patch("litellm.search.main.asearch", new_callable=AsyncMock) as mock_asearch:
mock_asearch.return_value = mock_search_response
# Setup proxy logging
user_api_key_cache = DualCache()
proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache)
setattr(litellm.proxy.proxy_server, "proxy_logging_obj", proxy_logging_obj)
# Call the track_cost_callback directly to simulate what happens after a search
proxy_db_logger = _ProxyDBLogger()
# Simulate the kwargs that would be passed from the search endpoint
request_id = "search_test_123"
kwargs = {
"call_type": "asearch",
"model": search_tool_name,
"custom_llm_provider": search_provider,
"litellm_call_id": request_id, # Set request_id in kwargs
"litellm_params": {
"metadata": {
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
"model_group": search_tool_name,
}
},
"metadata": {
"user_api_key": hash_token(generated_key),
"user_api_key_user_id": user_id,
"model_group": search_tool_name,
},
"response_cost": 0.008, # Mock cost for tavily search
}
# Set id on the response object
mock_search_response.id = request_id
await proxy_db_logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=mock_search_response,
start_time=datetime.now(),
end_time=datetime.now(),
)
# Wait for async operations
await asyncio.sleep(2)
await update_spend(
prisma_client=prisma_client,
db_writer_client=None,
proxy_logging_obj=proxy_logging_obj,
)
# Query spend logs
spend_logs = await view_spend_logs(
request_id=request_id,
user_api_key_dict=UserAPIKeyAuth(api_key=generated_key),
)
# Verify spend log was created
assert len(spend_logs) == 1, f"Expected 1 spend log, got {len(spend_logs)}"
spend_log = spend_logs[0]
# Verify all fields are populated correctly
assert spend_log.request_id == request_id
assert spend_log.call_type == "asearch"
assert spend_log.model == search_tool_name
assert spend_log.custom_llm_provider == search_provider
assert spend_log.model_group == search_tool_name
assert spend_log.spend == 0.008
# API key should be hashed (either the generated key or the one from metadata)
assert spend_log.api_key != "" # Should be populated
# Note: user field may be empty if not set in the request, but user_id should be in metadata
assert spend_log.metadata.get("user_api_key_user_id") == user_id or spend_log.user == user_id
print(f"✅ Search API logging test passed!")
print(f" - call_type: {spend_log.call_type}")
print(f" - model: {spend_log.model}")
print(f" - custom_llm_provider: {spend_log.custom_llm_provider}")
print(f" - model_group: {spend_log.model_group}")
print(f" - spend: {spend_log.spend}")