mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(websearch_interception): pass metadata to search calls for cost attribution
The websearch interception handler executes search API calls via litellm.asearch() but does not forward the request metadata (API key, user, team, etc.) from the original request. This means the proxy's spend tracking callbacks cannot attribute search costs to the correct API key or user budget. This change: - Adds an optional `metadata` parameter to `_execute_search()` - Extracts metadata from kwargs in both `_execute_agentic_loop()` and `_execute_chat_completion_agentic_loop()` - Forwards it to `litellm.asearch()` so that the `@client` decorator's logging pipeline can attribute costs correctly Search provider costs (e.g. google_pse/search at $0.005/query) are already defined in model_prices_and_context_window_backup.json and calculated by search/cost_calculator.py. This fix closes the gap so those costs actually reach the proxy's spend tracking system.
This commit is contained in:
parent
63425b4cb4
commit
ba87b6239c
1 changed files with 35 additions and 5 deletions
|
|
@ -668,6 +668,11 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
) -> Any:
|
||||
"""Execute litellm.search() and make follow-up request"""
|
||||
|
||||
# Extract metadata for cost attribution
|
||||
_metadata = kwargs.get("metadata") or kwargs.get("litellm_params", {}).get(
|
||||
"metadata"
|
||||
)
|
||||
|
||||
# Extract search queries from tool_use blocks
|
||||
search_tasks = []
|
||||
for tool_call in tool_calls:
|
||||
|
|
@ -676,7 +681,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Queuing search for query='{query}'"
|
||||
)
|
||||
search_tasks.append(self._execute_search(query))
|
||||
search_tasks.append(
|
||||
self._execute_search(query, metadata=_metadata)
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Tool call {tool_call['id']} has no query"
|
||||
|
|
@ -792,8 +799,18 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
raise
|
||||
|
||||
async def _execute_search(self, query: str) -> str:
|
||||
"""Execute a single web search using router's search tools"""
|
||||
async def _execute_search(
|
||||
self, query: str, metadata: Optional[Dict] = None
|
||||
) -> str:
|
||||
"""Execute a single web search using router's search tools.
|
||||
|
||||
Args:
|
||||
query: The search query string.
|
||||
metadata: Optional metadata dict from the original request. When
|
||||
provided, it is forwarded to ``litellm.asearch`` so that proxy
|
||||
callbacks (e.g. spend tracking) can attribute the search cost
|
||||
to the correct API key / user / team.
|
||||
"""
|
||||
try:
|
||||
# Import router from proxy_server
|
||||
try:
|
||||
|
|
@ -851,7 +868,13 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Executing search for '{query}' using provider '{search_provider}'"
|
||||
)
|
||||
result = await litellm.asearch(query=query, search_provider=search_provider)
|
||||
search_kwargs: Dict[str, Any] = {
|
||||
"query": query,
|
||||
"search_provider": search_provider,
|
||||
}
|
||||
if metadata:
|
||||
search_kwargs["metadata"] = metadata
|
||||
result = await litellm.asearch(**search_kwargs)
|
||||
|
||||
# Format using transformation function
|
||||
search_result_text = WebSearchTransformation.format_search_response(result)
|
||||
|
|
@ -879,6 +902,11 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
) -> Any:
|
||||
"""Execute litellm.search() and make follow-up chat completion request"""
|
||||
|
||||
# Extract metadata for cost attribution
|
||||
_metadata = kwargs.get("metadata") or kwargs.get("litellm_params", {}).get(
|
||||
"metadata"
|
||||
)
|
||||
|
||||
# Extract search queries from tool_calls
|
||||
search_tasks = []
|
||||
for tool_call in tool_calls:
|
||||
|
|
@ -897,7 +925,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Queuing search for query='{query}'"
|
||||
)
|
||||
search_tasks.append(self._execute_search(query))
|
||||
search_tasks.append(
|
||||
self._execute_search(query, metadata=_metadata)
|
||||
)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
f"WebSearchInterception: Tool call {tool_call.get('id')} has no query"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue