From 67d69d12b059777701d96e4ad60592f333e0ffc2 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 25 Nov 2025 17:14:59 +0530 Subject: [PATCH] Add cost tracking and logging support --- litellm/cost_calculator.py | 51 +++++ litellm/litellm_core_utils/litellm_logging.py | 5 + litellm/proxy/_types.py | 2 + litellm/proxy/proxy_server.py | 36 +++- litellm/proxy/search_endpoints/endpoints.py | 21 ++ .../spend_tracking/spend_tracking_utils.py | 3 +- .../test_search_api_logging.py | 202 ++++++++++++++++++ 7 files changed, 310 insertions(+), 10 deletions(-) create mode 100644 tests/proxy_unit_tests/test_search_api_logging.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0f5195e31a..9ef26d23ce 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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 ): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9c4a7e3876..5b4fa6b7e2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8a8bbdfe2e..9bfd118aba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0c0ced82ba..1330774286 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/proxy/search_endpoints/endpoints.py b/litellm/proxy/search_endpoints/endpoints.py index da5388f389..8cfc8bb510 100644 --- a/litellm/proxy/search_endpoints/endpoints.py +++ b/litellm/proxy/search_endpoints/endpoints.py @@ -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: diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 08e5624885..59f712e5b6 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -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) diff --git a/tests/proxy_unit_tests/test_search_api_logging.py b/tests/proxy_unit_tests/test_search_api_logging.py new file mode 100644 index 0000000000..7ac22e51ef --- /dev/null +++ b/tests/proxy_unit_tests/test_search_api_logging.py @@ -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}") +