diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ed29d49fc29..78a1ad68d90 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1552,7 +1552,7 @@ class Usage(SafeAttributeModel, CompletionUsage): completion_tokens_details: Optional[ Union[CompletionTokensDetailsWrapper, dict] ] = None, - server_tool_use: Optional[ServerToolUse] = None, + server_tool_use: Optional[Union[ServerToolUse, dict]] = None, cost: Optional[float] = None, **params, ): @@ -1653,6 +1653,9 @@ class Usage(SafeAttributeModel, CompletionUsage): prompt_tokens_details=_prompt_tokens_details or None, ) + if isinstance(server_tool_use, dict): + server_tool_use = ServerToolUse(**server_tool_use) + if server_tool_use is not None: self.server_tool_use = server_tool_use else: # maintain openai compatibility in usage object if possible diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index a04f6407e4b..c43291566b6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,17 +1,14 @@ -import json import os import sys -from unittest.mock import MagicMock import pytest -from fastapi.testclient import TestClient import litellm from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) from litellm.types.llms.openai import FileSearchTool, WebSearchOptions -from litellm.types.utils import ModelInfo, ModelResponse, StandardBuiltInToolsParams +from litellm.types.utils import ModelResponse, StandardBuiltInToolsParams sys.path.insert( 0, os.path.abspath("../../..") @@ -139,6 +136,22 @@ def test_get_cost_for_anthropic_web_search(): assert cost > 0.0 +def test_get_cost_for_anthropic_web_search_with_server_tool_use_dict(): + """ + Anthropic-compatible passthrough responses can construct Usage from a raw + usage payload. Ensure dict server_tool_use values are normalized before + built-in tool cost tracking reads server_tool_use.web_search_requests. + """ + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage(server_tool_use={"web_search_requests": 1}) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert StandardBuiltInToolCostTracking.response_object_includes_web_search_call( + response_object=None, usage=usage + ) + + @pytest.mark.parametrize( "model", ["gemini/gemini-2.0-flash-001", "gemini-2.0-flash-001"] ) diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index c146847f391..5960840838c 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -1,13 +1,9 @@ -import asyncio import os import sys -from typing import Optional -from unittest.mock import AsyncMock, patch import pytest sys.path.insert(0, os.path.abspath("../..")) -import json from litellm.types.utils import HiddenParams @@ -75,6 +71,26 @@ def test_usage_dump(): assert new_usage.prompt_tokens_details.web_search_requests == 1 +def test_usage_converts_server_tool_use_dict(): + from litellm.types.utils import ServerToolUse, Usage + + usage = Usage( + completion_tokens=2, + prompt_tokens=1, + total_tokens=3, + server_tool_use={"web_search_requests": 4, "tool_search_requests": 1}, + ) + + assert isinstance(usage.server_tool_use, ServerToolUse) + assert usage.server_tool_use.web_search_requests == 4 + assert usage.server_tool_use.tool_search_requests == 1 + + round_trip = Usage(**usage.model_dump()) + assert isinstance(round_trip.server_tool_use, ServerToolUse) + assert round_trip.server_tool_use.web_search_requests == 4 + assert round_trip.server_tool_use.tool_search_requests == 1 + + def test_usage_completion_tokens_details_text_tokens(): from litellm.types.utils import Usage