From aee4a6e62238f00f8f8e2f8ebe76f90ec3b5d91f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 15 Jan 2026 15:55:57 -0800 Subject: [PATCH] fix _get_total_tokens_from_usage --- .../hooks/parallel_request_limiter_v3.py | 2 +- ...ase_anthropic_messages_tool_search_test.py | 277 ++++++++++++++++++ 2 files changed, 278 insertions(+), 1 deletion(-) create mode 100644 tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 4d17cca22ad..b5bbb4237c1 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1236,7 +1236,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return pipeline_operations def _get_total_tokens_from_usage( - self, usage: Any | None, rate_limit_type: Literal["output", "input", "total"] + self, usage: Optional[Any], rate_limit_type: Literal["output", "input", "total"] ) -> int: """ Get total tokens from response usage for rate limiting. diff --git a/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py b/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py new file mode 100644 index 00000000000..f1c1b4d4c0c --- /dev/null +++ b/tests/pass_through_unit_tests/base_anthropic_messages_tool_search_test.py @@ -0,0 +1,277 @@ +""" +Base test class for Anthropic Messages API tool search E2E tests. + +Tests that tool search works correctly via litellm.anthropic.messages interface +by making actual API calls and validating that tool search discovers deferred tools. + +Reference: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool +""" + +import json +import os +import sys +from abc import ABC, abstractmethod +from typing import Any, Dict, List + +sys.path.insert(0, os.path.abspath("../../..")) + +import pytest +import litellm + + +# Sample tools for tool search testing +def get_deferred_tools() -> List[Dict[str, Any]]: + """ + Returns a list of tools with defer_loading: true. + These tools should only be discovered via tool search. + """ + return [ + { + "name": "get_weather", + "description": "Get the current weather for a location", + "input_schema": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "The city and state, e.g. San Francisco, CA" + } + }, + "required": ["location"] + }, + "defer_loading": True + }, + { + "name": "get_stock_price", + "description": "Get the current stock price for a ticker symbol", + "input_schema": { + "type": "object", + "properties": { + "ticker": { + "type": "string", + "description": "The stock ticker symbol, e.g. AAPL" + } + }, + "required": ["ticker"] + }, + "defer_loading": True + }, + { + "name": "search_web", + "description": "Search the web for information", + "input_schema": { + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "The search query" + } + }, + "required": ["query"] + }, + "defer_loading": True + }, + ] + + +def get_tool_search_tool_regex() -> Dict[str, Any]: + """Returns the tool search tool using regex variant.""" + return { + "type": "tool_search_tool_regex_20251119", + "name": "tool_search_tool_regex" + } + + +def get_tool_search_tool_bm25() -> Dict[str, Any]: + """Returns the tool search tool using BM25 variant.""" + return { + "type": "tool_search_tool_bm25_20251119", + "name": "tool_search_tool_bm25" + } + + +class BaseAnthropicMessagesToolSearchTest(ABC): + """ + Base test class for tool search E2E tests across different providers. + + Subclasses must implement: + - get_model(): Returns the model string to use for tests + """ + + @abstractmethod + def get_model(self) -> str: + """ + Returns the model string to use for tests. + + Examples: + - "anthropic/claude-sonnet-4-20250514" + - "vertex_ai/claude-sonnet-4@20250514" + - "bedrock/invoke/anthropic.claude-sonnet-4-20250514-v1:0" + """ + pass + + def get_tools_with_tool_search(self) -> List[Dict[str, Any]]: + """ + Returns tools list with tool search tool and deferred tools. + """ + return [get_tool_search_tool_regex()] + get_deferred_tools() + + @pytest.mark.asyncio + async def test_tool_search_basic_request(self): + """ + E2E test: Basic tool search request should succeed. + + This validates that the tool search beta header is being sent correctly + and the request doesn't error out. + """ + litellm._turn_on_debug() + + tools = self.get_tools_with_tool_search() + messages = [ + { + "role": "user", + "content": "What's the weather in San Francisco?" + } + ] + + response = await litellm.anthropic.messages.acreate( + model=self.get_model(), + messages=messages, + tools=tools, + max_tokens=1024, + ) + + print(f"Response: {json.dumps(response, indent=2, default=str)}") + + # Validate response structure + assert "content" in response, "Response should contain content" + assert "usage" in response, "Response should contain usage" + + # The model should either respond with text or use a tool + content = response.get("content", []) + assert len(content) > 0, "Response should have content" + + @pytest.mark.asyncio + async def test_tool_search_discovers_tool(self): + """ + E2E test: Tool search should discover and use a deferred tool. + + This validates that when the user asks about weather, the model + discovers the get_weather tool via tool search and attempts to use it. + """ + litellm._turn_on_debug() + + tools = self.get_tools_with_tool_search() + messages = [ + { + "role": "user", + "content": "I need to know the current weather in New York City. Please use the appropriate tool." + } + ] + + response = await litellm.anthropic.messages.acreate( + model=self.get_model(), + messages=messages, + tools=tools, + max_tokens=1024, + ) + + print(f"Response: {json.dumps(response, indent=2, default=str)}") + + content = response.get("content", []) + + # Check if the model used tool_use (either tool_search or get_weather) + tool_uses = [block for block in content if block.get("type") == "tool_use"] + + print(f"Tool uses: {json.dumps(tool_uses, indent=2, default=str)}") + + # The model should attempt to use tools when asked about weather + # It might use tool_search first, or directly use get_weather if discovered + if response.get("stop_reason") == "tool_use": + assert len(tool_uses) > 0, "Expected tool_use blocks when stop_reason is tool_use" + + @pytest.mark.asyncio + async def test_tool_search_streaming(self): + """ + E2E test: Tool search should work with streaming responses. + """ + litellm._turn_on_debug() + + tools = self.get_tools_with_tool_search() + messages = [ + { + "role": "user", + "content": "What's the weather like in Tokyo?" + } + ] + + response = await litellm.anthropic.messages.acreate( + model=self.get_model(), + messages=messages, + tools=tools, + max_tokens=1024, + stream=True, + ) + + # Collect all chunks + chunks = [] + async for chunk in response: + if isinstance(chunk, bytes): + chunk_str = chunk.decode("utf-8") + for line in chunk_str.split("\n"): + if line.startswith("data: "): + try: + json_data = json.loads(line[6:]) + chunks.append(json_data) + print(f"Chunk: {json.dumps(json_data, indent=2, default=str)}") + except json.JSONDecodeError: + pass + elif isinstance(chunk, dict): + chunks.append(chunk) + print(f"Chunk: {json.dumps(chunk, indent=2, default=str)}") + + # Should have received chunks + assert len(chunks) > 0, "Expected to receive streaming chunks" + + # Should have message_start + message_starts = [c for c in chunks if c.get("type") == "message_start"] + assert len(message_starts) > 0, "Expected message_start in streaming response" + + @pytest.mark.asyncio + async def test_tool_search_with_multiple_deferred_tools(self): + """ + E2E test: Tool search should work with multiple deferred tools. + + This validates that the model can discover the appropriate tool + from a larger catalog of deferred tools. + """ + litellm._turn_on_debug() + + tools = self.get_tools_with_tool_search() + messages = [ + { + "role": "user", + "content": "What's the stock price of Apple (AAPL)?" + } + ] + + response = await litellm.anthropic.messages.acreate( + model=self.get_model(), + messages=messages, + tools=tools, + max_tokens=1024, + ) + + print(f"Response: {json.dumps(response, indent=2, default=str)}") + + # Validate response + assert "content" in response, "Response should contain content" + + content = response.get("content", []) + tool_uses = [block for block in content if block.get("type") == "tool_use"] + + # If the model decides to use a tool, it should be related to stocks + if tool_uses: + tool_names = [t.get("name") for t in tool_uses] + print(f"Tools used: {tool_names}") +