From 2a94bb3835f1c3376e4ae948cf028c3ca409007a Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 04:08:45 +0000 Subject: [PATCH] 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 | 12 +++ .../websearch_interception/handler.py | 26 ++--- .../chat_completion_agentic_loop.py | 14 +-- litellm/llms/custom_httpx/llm_http_handler.py | 30 +++--- .../_support/interception_vendor.py | 32 +++++- ...t_interception_converted_stream_options.py | 98 ++++++++++++++++++- 6 files changed, 175 insertions(+), 37 deletions(-) diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index 40d3377bf66..f2f7869876d 100644 --- a/litellm/_internal_context.py +++ b/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.""" diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 2bb89248be4..d41647fe5f1 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/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) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 2ba47cbb400..1f29b5d7a4d 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/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( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 437944f8e03..879e4e02e3e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/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, diff --git a/tests/integration/_support/interception_vendor.py b/tests/integration/_support/interception_vendor.py index 3ca40087441..b46eb998f7a 100644 --- a/tests/integration/_support/interception_vendor.py +++ b/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 ) diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 0ca62d2f014..5a7855102c0 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/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(