diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 80d43d87dc3..c3ea605dd9e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1540,6 +1540,11 @@ class ServerToolUse(BaseModel): web_search_requests: Optional[int] = None tool_search_requests: Optional[int] = None + def __getitem__(self, key: str) -> Optional[int]: + if key not in self.__class__.model_fields: + raise KeyError(key) + return getattr(self, key) + class Usage(SafeAttributeModel, CompletionUsage): _cache_creation_input_tokens: int = PrivateAttr( @@ -1570,7 +1575,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, ): @@ -1671,6 +1676,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/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index f8b6fbde3dc..1114b3df0c2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -321,6 +321,44 @@ class TestAzureAnthropicCostCalculation: assert call_kwargs["model"] == "azure_ai/claude-sonnet-4-5_gb_20250929" assert call_kwargs["custom_llm_provider"] == "azure_ai" + def test_passthrough_logging_sets_response_cost_with_server_tool_use_dict(self): + from litellm.types.utils import Choices, Message, ModelResponse + + logging_obj = self._create_mock_logging_obj(model="claude-3-7-sonnet-20250219") + logging_obj.get_router_model_id.return_value = None + logging_obj.litellm_params = {} + + response = ModelResponse( + id="test-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message(content="test", role="assistant"), + ) + ], + created=1234567890, + model="claude-3-7-sonnet-20250219", + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "server_tool_use": {"web_search_requests": 1}, + }, + ) + + kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload( + litellm_model_response=response, + model="claude-3-7-sonnet-20250219", + kwargs={}, + start_time=datetime.now(), + end_time=datetime.now(), + logging_obj=logging_obj, + ) + + assert "response_cost" in kwargs + assert kwargs["response_cost"] > 0 + class TestAnthropicBatchPassthroughCostTracking: """Test cases for Anthropic batch passthrough cost tracking functionality""" diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index c146847f391..a4074ccdaaa 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,48 @@ def test_usage_dump(): assert new_usage.prompt_tokens_details.web_search_requests == 1 +def test_usage_server_tool_use_dict_is_coerced_and_round_trips(): + from litellm.types.utils import ServerToolUse, Usage + + current_usage = Usage( + completion_tokens=1, + prompt_tokens=1, + total_tokens=2, + server_tool_use={"web_search_requests": 1}, + ) + + assert isinstance(current_usage.server_tool_use, ServerToolUse) + assert current_usage.server_tool_use.web_search_requests == 1 + + new_usage = Usage(**current_usage.model_dump()) + assert isinstance(new_usage.server_tool_use, ServerToolUse) + assert new_usage.server_tool_use.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["web_search_requests"] == 4 + assert usage.server_tool_use.tool_search_requests == 1 + with pytest.raises(KeyError): + usage.server_tool_use["unknown_metric"] + + 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["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