fix(spend): bill web search interception follow-ups once under the outer request (#49762)

* fix(spend): bill web search interception follow-ups once under the outer request

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(spend): keep web search rounds inside interception follow-ups billed

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(spend): keep the outer request as parent for nested interception follow-ups

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* revert(spend): drop parent override for nested interception follow-ups

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/_internal_context.py

@@ -52,6 +52,18 @@ def internal_sub_call(parent_request_id: str | None = None) -> Generator[None]:
         internal_parent_request_id.reset(parent_token)
 
 
+@contextmanager
+def billed_separately() -> Generator[None]:
+    """A call that writes its own spend row even inside a sub-call whose parent bills the model response."""
+    token: Final = is_internal_call.set(False)
+    parent_token: Final = internal_parent_request_id.set(None)
+    try:
+        yield
+    finally:
+        is_internal_call.reset(token)
+        internal_parent_request_id.reset(parent_token)
+
+
 @contextmanager
 def post_response_phase() -> Generator[None]:
     """Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns."""

litellm/integrations/websearch_interception/handler.py

@@ -17,6 +17,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast
 from typing_extensions import Never, ReadOnly
 
 import litellm
+from litellm._internal_context import billed_separately
 from litellm._logging import verbose_logger
 from litellm.anthropic_interface import messages as anthropic_messages
 from litellm.constants import LITELLM_WEB_SEARCH_TOOL_NAME
@@ -1628,19 +1629,20 @@ class WebSearchInterceptionLogger(CustomLogger):
             search_kwargs: Final = MappingProxyType(
                 {**configured_search_kwargs, **parent_correlation.as_search_kwargs()}
             )
-            result: Final = (
-                await litellm.asearch(
-                    query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
+            with billed_separately():
+                result: Final = (
+                    await litellm.asearch(
+                        query=query_arg, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs
+                    )
+                    if search_metadata is None
+                    else await litellm.asearch(
+                        query=query_arg,
+                        search_provider=search_provider,
+                        litellm_metadata=search_metadata,
+                        **_NO_ASEARCH_NAMED,
+                        **search_kwargs,
+                    )
                 )
-                if search_metadata is None
-                else await litellm.asearch(
-                    query=query_arg,
-                    search_provider=search_provider,
-                    litellm_metadata=search_metadata,
-                    **_NO_ASEARCH_NAMED,
-                    **search_kwargs,
-                )
-            )
 
             # Format using transformation function
             search_result_text: Final = WebSearchTransformation.format_search_response(result)

litellm/litellm_core_utils/chat_completion_agentic_loop.py

@@ -8,6 +8,7 @@ from typing import Final
 
 from pydantic import TypeAdapter, ValidationError
 
+from litellm._internal_context import internal_sub_call
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
@@ -204,12 +205,13 @@ async def _execute_chat_completion_agentic_plan(
     )
 
     try:
-        response_followup = await litellm.acompletion(
-            model=full_model_name,
-            messages=patch.messages,
-            **optional_params_for_followup,
-            **kwargs_for_followup,
-        )
+        with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObject) else None):
+            response_followup = await litellm.acompletion(
+                model=full_model_name,
+                messages=patch.messages,
+                **optional_params_for_followup,
+                **kwargs_for_followup,
+            )
         if _post_hook_overridden(callback):
             try:
                 response_followup = await callback.async_post_agentic_loop_response_hook(

litellm/llms/custom_httpx/llm_http_handler.py

@@ -37,6 +37,7 @@ import litellm
 import litellm.litellm_core_utils
 import litellm.types
 import litellm.types.utils
+from litellm._internal_context import internal_sub_call
 from litellm._logging import redact_string, verbose_logger
 from litellm.anthropic_beta_headers_manager import (
     aload_beta_headers_config,
@@ -6252,12 +6253,13 @@ class BaseLLMHTTPHandler:
         )
 
         try:
-            response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses(
-                model=patch.model or model,
-                input=patch.messages,  # pyright: ignore[reportArgumentType]  # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime
-                **optional_params,
-                **kwargs_for_followup,
-            )
+            with internal_sub_call(logging_obj.litellm_call_id):
+                response: ResponsesAPIResponse | BaseResponsesAPIStreamingIterator = await litellm.aresponses(
+                    model=patch.model or model,
+                    input=patch.messages,  # pyright: ignore[reportArgumentType]  # pre-existing mismatch surfaced by the Router import; patch messages are valid response input at runtime
+                    **optional_params,
+                    **kwargs_for_followup,
+                )
 
             if callback is not None:
                 try:
@@ -6346,6 +6348,7 @@ class BaseLLMHTTPHandler:
         max_loops: int,
         fingerprints: list[str],
         fingerprint: str,
+        logging_obj: Optional["LiteLLMLoggingObj"] = None,
     ) -> ModelResponse | CustomStreamWrapper:
         patch: Final = plan.request_patch or AgenticLoopRequestPatch()
         if patch.messages is None:
@@ -6385,12 +6388,14 @@ class BaseLLMHTTPHandler:
             fingerprint=fingerprint,
         )
 
-        return await litellm.acompletion(
-            model=full_model_name,
-            messages=patch.messages,
-            **optional_params_for_followup,
-            **kwargs_for_followup,
-        )
+        with internal_sub_call(logging_obj.litellm_call_id if logging_obj is not None else None):
+            response: Final = await litellm.acompletion(
+                model=full_model_name,
+                messages=patch.messages,
+                **optional_params_for_followup,
+                **kwargs_for_followup,
+            )
+        return response
 
     def _maybe_wrap_in_fake_stream(
         self,
@@ -6739,6 +6744,7 @@ class BaseLLMHTTPHandler:
                     messages=messages,
                     optional_params=optional_params,
                     kwargs=kwargs_with_provider,
+                    logging_obj=logging_obj,
                     custom_llm_provider=custom_llm_provider,
                     depth=depth,
                     max_loops=max_loops,

tests/integration/_support/interception_vendor.py

@@ -27,7 +27,7 @@ JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
 _ITEMS: Final = TypeAdapter(list[dict[str, JsonValue]])
 _CHAT_TARGET: Final = re.compile(r"^/openai/deployments/([^/]+)/chat/completions$")
 
-Mode: TypeAlias = Literal["search", "plain", "search-down", "followup-fails", "missing"]
+Mode: TypeAlias = Literal["search", "multi-search", "plain", "search-down", "followup-fails", "missing"]
 
 
 def deployment(mode: Mode) -> str:
@@ -42,8 +42,8 @@ def plain(name: str) -> str:
     return f"plain-{name}"
 
 
-def query(name: str) -> str:
-    return f"q-{name}"
+def query(name: str, turn: int = 1) -> str:
+    return f"q-{name}" if turn == 1 else f"q-{name}-{turn}"
 
 
 def _minted(prefix: str) -> str:
@@ -98,9 +98,20 @@ def _chat(name: str, body: Mapping[str, JsonValue]) -> Reply:
             content_type="text/event-stream", chunks=chat_stream(_minted(f"chatcmpl-{name}-s-"), "gpt-4o", plain(name))
         )
     messages: Final = _ITEMS.validate_python(body["messages"])
-    if any(message.get("role") == "tool" for message in messages):
+    tool_results: Final = tuple(message for message in messages if message.get("role") == "tool")
+    if tool_results:
         if name.endswith("-followup-fails"):
             return _azure_error(500, FOLLOWUP_FAILURE)
+        if name.endswith("-multi-search") and len(tool_results) == 1:
+            call: Final = {
+                "id": f"call_{name[4:36]}-2",
+                "type": "function",
+                "function": {"name": SEARCH_TOOL, "arguments": json.dumps({"query": query(name, turn=2)})},
+            }
+            message: Final = JSON_OBJECT.validate_python({"role": "assistant", "content": None, "tool_calls": [call]})
+            return _chat_completion(
+                _minted(f"chatcmpl-{name}-2-"), message, "tool_calls", FOLLOWUP_USAGE
+            )
         return _chat_completion(
             _minted(f"chatcmpl-{name}-2-"), {"role": "assistant", "content": answer(name)}, "stop", FOLLOWUP_USAGE
         )
@@ -153,9 +164,20 @@ def _responses(body: Mapping[str, JsonValue]) -> Reply:
             content_type="text/event-stream", chunks=responses_stream(_minted(f"resp_{name}_s_"), name, plain(name))
         )
     items: Final = _ITEMS.validate_python(body["input"]) if isinstance(body.get("input"), list) else []
-    if any(item.get("type") == "function_call_output" for item in items):
+    function_call_outputs: Final = tuple(item for item in items if item.get("type") == "function_call_output")
+    if function_call_outputs:
         if name.endswith("-followup-fails"):
             return _azure_error(500, FOLLOWUP_FAILURE)
+        if name.endswith("-multi-search") and len(function_call_outputs) == 1:
+            call: Final = {
+                "type": "function_call",
+                "id": f"fc_{name[4:36]}-2",
+                "call_id": f"call_{name[4:36]}-2",
+                "name": SEARCH_TOOL,
+                "arguments": json.dumps({"query": query(name, turn=2)}),
+                "status": "completed",
+            }
+            return _response(_minted(f"resp_{name}_2_"), name, (call,), RESPONSES_FOLLOWUP_USAGE)
         return _response(
             _minted(f"resp_{name}_2_"), name, (_text_item(f"{name}_2", answer(name)),), RESPONSES_FOLLOWUP_USAGE
         )

tests/integration/streaming/test_interception_converted_stream_options.py

@@ -5,6 +5,7 @@ import uuid
 from collections.abc import Callable, Generator, Iterator, Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
+from hashlib import sha256
 from pathlib import Path
 from queue import SimpleQueue
 from types import MappingProxyType
@@ -216,10 +217,13 @@ def scenario(rig: _Rig) -> Iterator[Scenario]:
         yield opened
 
 
-def _deploy(rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY) -> _Deployment:
+def _deploy(
+    rig: _Rig, scenario: Scenario, mode: iv.Mode, *, api_key: str = iv.AZURE_KEY, priced: bool = False
+) -> _Deployment:
     name: Final = iv.deployment(mode)
+    pricing: Final = {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002} if priced else {}
     model: Final = scenario.model(
-        model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION
+        model=f"azure/{name}", api_base=rig.wire.url, api_key=api_key, api_version=iv.API_VERSION, **pricing
     )
     return _Deployment(model, name)
 
@@ -325,6 +329,46 @@ def _assert_response_landed_once(model: str, identity: str) -> None:
     assert row["status"] == "success" and rv.same_response(string_value(row["request_id"]), identity), row
 
 
+def _assert_single_model_spend_for_key(key: str, call_type: str, *, expected_searches: int = 1) -> None:
+    digest: Final = sha256(key.encode()).hexdigest()
+    observed: Final = eventually(
+        lambda: (
+            read_rows(
+                "SELECT status, spend, request_id, call_type, model_group, session_id "
+                'FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND status = %s',
+                (digest, "success"),
+            ),
+            read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (digest,)),
+        ),
+        lambda values: (
+            bool(values[0])
+            and (
+                len(values[0]) > expected_searches + 1
+                or (
+                    len(values[0]) == expected_searches + 1
+                    and len(values[1]) == 1
+                    and values[1][0]["spend"] is not None
+                    and float(str(values[1][0]["spend"]))
+                    == pytest.approx(sum(float(str(row["spend"])) for row in values[0]))
+                )
+            )
+        ),
+        seconds=70,
+    )
+    rows, balance = observed
+    model_rows: Final = [row for row in rows if row["call_type"] == call_type]
+    search_rows: Final = [row for row in rows if row["call_type"] == "asearch"]
+    spend: Final = sum(float(str(row["spend"])) for row in rows)
+    assert (
+        len(rows) == expected_searches + 1
+        and len(model_rows) == 1
+        and len(search_rows) == expected_searches
+        and float(str(model_rows[0]["spend"])) > 0
+        and all(float(str(row["spend"])) > 0 for row in search_rows)
+        and float(str(balance[0]["spend"])) == pytest.approx(spend)
+    ), (rows, balance)
+
+
 def _openai(rig: _Rig) -> openai.OpenAI:
     return openai.OpenAI(base_url=f"{rig.base_url}/v1", api_key=rig.gateway.key, max_retries=0)
 
@@ -402,6 +446,31 @@ def test_raw_chat_stream_without_stream_options_answers_without_a_usage_chunk(ri
     _assert_search_turn(rig, deployment, identity)
 
 
+def test_chat_stream_through_web_search_has_single_spend_owner(rig: _Rig, scenario: Scenario) -> None:
+    deployment: Final = _deploy(rig, scenario, "search", priced=True)
+    key: Final = scenario.key(models=[deployment.model])
+    stream: Final = _post(rig, "/v1/chat/completions", _chat_body(deployment), key=key)
+    identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
+    _assert_search_turn(rig, deployment, identity)
+    _assert_single_model_spend_for_key(key, "acompletion")
+
+
+def test_multi_round_chat_stream_through_web_search_bills_each_search(
+    rig: _Rig, scenario: Scenario
+) -> None:
+    deployment: Final = _deploy(rig, scenario, "multi-search", priced=True)
+    key: Final = scenario.key(models=[deployment.model])
+    stream: Final = _post(rig, "/v1/chat/completions", _chat_body(deployment), key=key)
+    identity: Final = _assert_restreamed(stream, iv.answer(deployment.name), usage=None)
+    received: Final = _upstream(rig, deployment)
+    _assert_upstream(received, model_calls=3, searches=2)
+    assert [search["query"] for search in received.searches] == [
+        iv.query(deployment.name),
+        iv.query(deployment.name, turn=2),
+    ], received.searches
+    _assert_single_model_spend_for_key(key, "acompletion", expected_searches=2)
+
+
 def test_openai_sdk_responses_stream_through_web_search_keeps_stream_options_off_azure(
     rig: _Rig, scenario: Scenario
 ) -> None:
@@ -454,6 +523,31 @@ def test_raw_responses_stream_without_stream_options_answers(rig: _Rig, scenario
     _assert_response_landed_once(deployment.model, identity)
 
 
+def test_responses_stream_through_web_search_has_single_spend_owner(rig: _Rig, scenario: Scenario) -> None:
+    deployment: Final = _deploy(rig, scenario, "search", priced=True)
+    key: Final = scenario.key(models=[deployment.model])
+    stream: Final = _post(rig, "/v1/responses", _responses_body(deployment), key=key)
+    _assert_responses_restreamed(stream, iv.answer(deployment.name))
+    _assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1)
+    _assert_single_model_spend_for_key(key, "aresponses")
+
+
+def test_multi_round_responses_stream_through_web_search_bills_each_search(
+    rig: _Rig, scenario: Scenario
+) -> None:
+    deployment: Final = _deploy(rig, scenario, "multi-search", priced=True)
+    key: Final = scenario.key(models=[deployment.model])
+    stream: Final = _post(rig, "/v1/responses", _responses_body(deployment), key=key)
+    _assert_responses_restreamed(stream, iv.answer(deployment.name))
+    received: Final = _upstream(rig, deployment)
+    _assert_upstream(received, model_calls=3, searches=2)
+    assert [search["query"] for search in received.searches] == [
+        iv.query(deployment.name),
+        iv.query(deployment.name, turn=2),
+    ], received.searches
+    _assert_single_model_spend_for_key(key, "aresponses", expected_searches=2)
+
+
 def test_openai_sdk_chat_without_streaming_through_web_search_answers_in_json(rig: _Rig, scenario: Scenario) -> None:
     deployment: Final = _deploy(rig, scenario, "search")
     completion: Final = _openai(rig).chat.completions.create(