fix(cost): bill azure_ai claude web searches under deployment-owned pricing ids (#46135)

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-13 03:59:12 +00:00 · 2026-10-12 19:40:53 -07:00 · 2026-10-12 19:40:53 -07:00 · 24bbf8f0e1
commit 24bbf8f0e1
parent ffec84f918
6 changed files with 310 additions and 8 deletions
--- 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

--- 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,
--- 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")
--- 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",
--- 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):
--- 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