diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 816311e401e..7452a4a20ef 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -63,8 +63,10 @@ def get_cost_for_web_search_request(custom_llm_provider: str, usage: "Usage", mo return cost_per_web_search_request_vertex_ai(usage=usage, model_info=model_info) elif custom_llm_provider == "azure_ai": + from litellm.litellm_core_utils.llm_cost_calc.utils import get_web_search_requests_from_usage + azure_ai_model_key: Final[str] = model_info.get("key", "") if model_info else "" - if "claude" not in azure_ai_model_key.lower(): + if "claude" not in azure_ai_model_key.lower() and get_web_search_requests_from_usage(usage) is None: return None from .anthropic.cost_calculation import get_cost_for_anthropic_web_search diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index 3e76181c1ec..24b57c3e93c 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -992,6 +992,283 @@ AZURE_AI_CLAUDE_HAIKU_4_5_RESPONSES_WEB_SEARCH: Final = CostTrackingTestCase( ) +# azure_ai/claude-sonnet-4-5 (public row has no search price, the deployment owns input, output and search rates) +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$UNIQUE_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer b1c2d3e4f5a6"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + rollups=True, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_web_search_two_requests_stream"}], + "stream": True, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer c2d3e4f5a6b1"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_zero_search_rate_messages_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.0), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_zero_search_rate_messages_two_requests"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer d3e4f5a6b1c2"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0039, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.0, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_NO_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_messages_no_server_tool_use", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/messages", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "max_tokens": 412, + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_messages_no_server_tool_use"}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer e4f5a6b1c2d3"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412}, + }, + ), + expected=ExactExpected( + spend=0.0039, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.0, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_three_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_three_requests"}], + "stream": False, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}], + "allowed_openai_params": ["web_search_options"], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer f5a6b1c2d3e4"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 3}}, + }, + ), + expected=ExactExpected( + spend=0.0339, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.03, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_two_requests_stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="azure_ai/claude-sonnet-4-5", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "azure-ai-claude-sonnet-4-5-deployment_rates_chat_web_search_two_requests_stream"}], + "stream": True, + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + "allowed_openai_params": ["web_search_options"], + }, + response=SseResponse( + content_type="text/event-stream", + frames=sse_frames( + { + "type": "message_start", + "message": { + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1840}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "scripted answer a6b1c2d3e4f5"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn"}, + "usage": {"output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + {"type": "message_stop"}, + ), + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + +AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_RESPONSES_WEB_SEARCH: Final = CostTrackingTestCase( + name="azure-ai-claude-sonnet-4-5-deployment_rates_responses_web_search_two_requests", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="azure_ai/claude-sonnet-4-5", + endpoint="/v1/responses", + deployment=Deployment(input_cost_per_token=1e-06, output_cost_per_token=5e-06, search_cost_per_query=0.01), + request={ + "model": "$MODEL", + "input": "azure-ai-claude-sonnet-4-5-deployment_rates_responses_web_search_two_requests", + "tools": [{"type": "web_search"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "msg_$REQUEST_ID", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted answer b1c2d3e4f5a7"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 1840, "output_tokens": 412, "server_tool_use": {"web_search_requests": 2}}, + }, + ), + expected=ExactExpected( + spend=0.0239, + input_cost=0.00184, + output_cost=0.00206, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.02, + ), +) + + # groq/qwen/qwen3.8-27b GROQ_QWEN_3_8_JSON: Final = CostTrackingTestCase( name="groq-qwen-3.8-json", @@ -1155,6 +1432,13 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH, AZURE_AI_CLAUDE_HAIKU_4_5_CHAT_WEB_SEARCH_STREAM, AZURE_AI_CLAUDE_HAIKU_4_5_RESPONSES_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_WEB_SEARCH_ZERO_RATE, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_MESSAGES_NO_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_CHAT_WEB_SEARCH_STREAM, + AZURE_AI_CLAUDE_SONNET_4_5_DEPLOYMENT_RESPONSES_WEB_SEARCH, GROQ_QWEN_3_8_JSON, GROQ_QWEN_3_8_STREAM_X_GROQ_RECOUNT, COHERE_COMMAND_A_V2_TOKENS, diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 7a75320a70f..9bac5cbdcef 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -310,6 +310,17 @@ def register_scenario_deployment( if case.deployment is not None else {} ), + **( + { + "search_context_cost_per_query": { + "search_context_size_low": case.deployment.search_cost_per_query, + "search_context_size_medium": case.deployment.search_cost_per_query, + "search_context_size_high": case.deployment.search_cost_per_query, + } + } + if case.deployment is not None and case.deployment.search_cost_per_query is not None + else {} + ), **( {"vertex_credentials": _vertex_service_account_json(control_url)} if case.rates.litellm_provider.startswith("vertex_ai") diff --git a/tests/integration/cost_calculation/cost_map.py b/tests/integration/cost_calculation/cost_map.py index f8db4eb1a04..5e211d25a94 100644 --- a/tests/integration/cost_calculation/cost_map.py +++ b/tests/integration/cost_calculation/cost_map.py @@ -620,6 +620,12 @@ COST_MAP_ENTRIES: Final[dict[str, dict[str, JsonValue]]] = { }, "web_search_billing_unit": "per_query", }, + "azure_ai/claude-sonnet-4-5": { + "litellm_provider": "azure_ai", + "mode": "chat", + "input_cost_per_token": 3e-06, + "output_cost_per_token": 1.5e-05, + }, "groq/qwen/qwen3.8-27b": { "litellm_provider": "groq", "mode": "chat", diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 9304100f3a5..4990ef2a5e4 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -105,6 +105,7 @@ class Deployment(BaseModel): base_model: str | None = None input_cost_per_token: float | None = None output_cost_per_token: float | None = None + search_cost_per_query: float | None = None class WavUpload(BaseModel): diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 759f686f32e..320a4eff371 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -600,6 +600,9 @@ _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 ("azure_ai/claude-sonnet-4-5", 3, 0.01, 0.03), ("azure_ai/claude-haiku-4-5", 1, 0.01, 0.01), ("azure_ai/claude-haiku-4-5", 2, 0.0, 0.0), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 2, 0.01, 0.02), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 3, 0.01, 0.03), + ("4d696a77-1451-438a-85be-9750c7fb4eb7", 2, 0.0, 0.0), ], ) def test_azure_ai_claude_web_search_billed_per_server_tool_use_request( @@ -635,7 +638,7 @@ def test_azure_ai_claude_web_search_billed_per_server_tool_use_request( def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: from litellm.llms import get_cost_for_web_search_request - from litellm.types.utils import ModelInfo, SearchContextCostPerQuery, ServerToolUse, Usage + from litellm.types.utils import ModelInfo, SearchContextCostPerQuery, Usage model_info = ModelInfo( key="azure_ai/gpt-5.4-mini-2026-03-17", @@ -648,11 +651,6 @@ def test_azure_ai_non_claude_web_search_keeps_generic_fallback() -> None: output_cost_per_token=5e-6, search_context_cost_per_query=SearchContextCostPerQuery(search_context_size_medium=0.01), ) - usage = Usage( - prompt_tokens=1840, - completion_tokens=412, - total_tokens=2252, - server_tool_use=ServerToolUse(web_search_requests=2), - ) + usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252) assert get_cost_for_web_search_request(custom_llm_provider="azure_ai", usage=usage, model_info=model_info) is None