diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 861ffd91bfa..2f3800c592f 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -13,6 +13,7 @@ from typing import ( from typing_extensions import ReadOnly, TypedDict +from litellm.types.integrations.custom_logger import CHAT_COMPLETION_AGENTIC_SURFACE from litellm.types.llms.openai import ResponsesAPIResponse if TYPE_CHECKING: @@ -317,6 +318,7 @@ class ResponsesToCompletionBridgeHandler: result: Final = await aresponses( **request_data, aresponses=True, + _agentic_loop_api_surface=CHAT_COMPLETION_AGENTIC_SURFACE, ) from litellm.types.utils import ModelResponse diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index f17bb827641..c7ef1c4b321 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -70,6 +70,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac # Class variables or attributes server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() wants_failed_prompt_token_estimate: ClassVar[bool] = False + wants_chat_completion_hooks_on_bridged_chat: ClassVar[bool] = False enforces_request_content: bool = False """ diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index b386c020a1b..a6d0ca69557 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -12,7 +12,7 @@ import uuid from collections.abc import AsyncIterator, Callable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, TypeVar, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypedDict, TypeVar, cast from pydantic import TypeAdapter from typing_extensions import Never, ReadOnly @@ -54,6 +54,7 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, as_converted_stream, + bridged_from_chat_completion, ) from litellm.types.integrations.websearch_interception import ( AnthropicSearchQuery, @@ -286,6 +287,19 @@ def _kept_allowed_tools_tool_choice(optional_params: Mapping[str, object]) -> Ma return allowed_tools_tool_choice_kept_on_follow_up(optional_params.get("tool_choice")) +def _is_litellm_web_search_function_responses(tool: Mapping[str, object]) -> bool: + return tool.get("type") == "function" and tool.get("name") == LITELLM_WEB_SEARCH_TOOL_NAME + + +def _kept_on_the_responses_agentic_loop( + kwargs: Mapping[str, object], tools: Sequence[Mapping[str, object]] +) -> Mapping[str, object]: + chat_side_runs_the_search_loop: Final = any(_is_litellm_web_search_function_responses(tool) for tool in tools) + if not bridged_from_chat_completion(kwargs) or chat_side_runs_the_search_loop: + return kwargs + return {key: value for key, value in kwargs.items() if key != "_agentic_loop_api_surface"} + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -298,6 +312,8 @@ class WebSearchInterceptionLogger(CustomLogger): 4. Returns final response """ + wants_chat_completion_hooks_on_bridged_chat: ClassVar[bool] = True + def __init__( self, enabled_providers: list[LlmProviders | str] | None = None, @@ -563,7 +579,7 @@ class WebSearchInterceptionLogger(CustomLogger): ] converted_kwargs: Final = { - **kwargs, + **_kept_on_the_responses_agentic_loop(kwargs, tools), "tools": converted_tools, **_tool_choice_naming_litellm_web_search(kwargs, is_web_search_tool_responses), } diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 93799c72209..1b87b8e8537 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -1,14 +1,15 @@ # this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api import json -from collections.abc import Awaitable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from itertools import chain from types import MappingProxyType from typing import Final, TypeVar from pydantic import TypeAdapter, ValidationError -from litellm._internal_context import internal_sub_call +from litellm._internal_context import billed_separately, internal_sub_call, is_internal_call from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.agentic_followup_failure import mark_agentic_followup_failure, rounds_served_before @@ -24,10 +25,13 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( ) from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model from litellm.litellm_core_utils.get_litellm_params import AGENTIC_LOOP_TOOL_CHOICE_KEY -from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.litellm_logging import UncachedPrice from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict, object_items from litellm.llms.base_llm.base_model_iterator import MockResponseIterator +from litellm.responses.mcp.parked_stream_dispatch import SERVED_FROM_CACHE +from litellm.responses.mcp.round_billing import billed_for_every_round, summed_cost_breakdown from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, @@ -61,6 +65,20 @@ _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset( ) +def _first_round_bills_itself(logging_obj: object) -> bool: + return ( + isinstance(logging_obj, LiteLLMLoggingObject) + and logging_obj.bridged_chat_request is not None + and logging_obj.stream is not True + and not is_internal_call.get() + and logging_obj.async_success_log_enqueued + ) + + +def _follow_ups_bill_themselves(logging_obj: object, request_kwargs: Mapping[str, object]) -> bool: + return _first_round_bills_itself(logging_obj) and not converted_stream_requested(request_kwargs) + + async def bill_first_round_on_failure( dispatch: Awaitable[_DispatchResult], *, logging_obj: object, response: object, stream: bool ) -> _DispatchResult: @@ -70,7 +88,7 @@ async def bill_first_round_on_failure( if stream or not isinstance(logging_obj, LiteLLMLoggingObject): mark_agentic_followup_failure(error, rounds_served_before(error)) raise - this_round: Final = logging_obj.served_round(response) + this_round: Final = None if _first_round_bills_itself(logging_obj) else logging_obj.served_round(response) served_rounds: Final = (*rounds_served_before(error), *(() if this_round is None else (this_round,))) mark_agentic_followup_failure(error, served_rounds) logging_obj.record_served_rounds_for_failure(served_rounds) @@ -92,6 +110,10 @@ def _chat_build_plan_overridden(callback: CustomLogger) -> bool: ) +def runs_its_own_chat_completion_loop(callback: CustomLogger) -> bool: + return callback.wants_chat_completion_hooks_on_bridged_chat + + def _post_hook_overridden(callback: CustomLogger) -> bool: return ( type(callback).async_post_agentic_loop_response_hook is not CustomLogger.async_post_agentic_loop_response_hook @@ -100,6 +122,7 @@ def _post_hook_overridden(callback: CustomLogger) -> bool: _STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) _HOOK_RESPONSE: Final = TypeAdapter(object) +_PRICE_OF_EVERY_ROUND: Final = "agentic_loop_price_of_every_round" def _plan_response_override(plan: AgenticLoopPlan) -> object | None: @@ -165,6 +188,9 @@ def _converted_stream_replay( return response if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject): return None + every_round: Final = _price_of_every_round(response) + if every_round is not None: + logging_obj.record_every_round_breakdown(every_round.cost_breakdown) return CustomStreamWrapper( completion_stream=MockResponseIterator(model_response=response), @@ -172,6 +198,7 @@ def _converted_stream_replay( custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, stream_options=_stashed_stream_options(kwargs), + response_cost=None if every_round is None else every_round.cost, ) @@ -345,6 +372,61 @@ def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: } +def _followup_billing_scope(logging_obj: object, request_kwargs: Mapping[str, object]) -> AbstractContextManager[None]: + if _follow_ups_bill_themselves(logging_obj, request_kwargs): + return billed_separately() + return internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObject) else None) + + +def _billed_with_the_round_it_follows(round_served: object, follow_up: object, logging_obj: object) -> object: + if not ( + isinstance(round_served, ModelResponse) + and isinstance(follow_up, ModelResponse) + and isinstance(logging_obj, LiteLLMLoggingObject) + and logging_obj.bridged_chat_request is not None + ): + return follow_up + if not (is_internal_call.get() or logging_obj.unpark_success_log()): + return follow_up + first_price: Final = _price_of_round(round_served, logging_obj) + final_price: Final = _price_of_round(follow_up, logging_obj) + first: Final = _priced(round_served, first_price) + final: Final = _priced(follow_up, final_price) + breakdown: Final = summed_cost_breakdown(first, first_price.cost_breakdown, final, final_price.cost_breakdown) + logging_obj.record_every_round_breakdown(breakdown) + billed: Final = billed_for_every_round(first, first_price.cost_breakdown, final, final_price.cost_breakdown) + billed_hidden_params: Final = get_hidden_params(billed) or {} + total_cost: Final = billed_hidden_params.get("response_cost") + set_hidden_params( + billed, + { + **billed_hidden_params, + _PRICE_OF_EVERY_ROUND: UncachedPrice( + cost=float(total_cost) if isinstance(total_cost, (int, float)) else None, cost_breakdown=breakdown + ), + }, + ) + return billed + + +def _price_of_every_round(response: ModelResponse) -> UncachedPrice | None: + every_round: Final = (get_hidden_params(response) or {}).get(_PRICE_OF_EVERY_ROUND) + return every_round if isinstance(every_round, UncachedPrice) else None + + +def _price_of_round(round_response: ModelResponse, logging_obj: LiteLLMLoggingObject) -> UncachedPrice: + every_round: Final = _price_of_every_round(round_response) + if every_round is not None: + return every_round + return SERVED_FROM_CACHE if served_from_cache(round_response) else logging_obj.uncached_price(round_response) + + +def _priced(round_response: ModelResponse, price: UncachedPrice) -> ModelResponse: + priced: Final = round_response.model_copy() + set_hidden_params(priced, {**(get_hidden_params(round_response) or {}), "response_cost": price.cost}) + return priced + + def _leave_pricing_to_the_request(response: object) -> None: hidden_params: Final = get_hidden_params(response) if hidden_params is not None and not served_from_cache(response): @@ -353,6 +435,7 @@ def _leave_pricing_to_the_request(response: object) -> None: async def _execute_chat_completion_agentic_plan( *, + response: ModelResponse, plan: AgenticLoopPlan, callback: CustomLogger, model: str, @@ -392,7 +475,7 @@ async def _execute_chat_completion_agentic_plan( ) try: - with internal_sub_call(logging_obj.litellm_call_id if isinstance(logging_obj, LiteLLMLoggingObject) else None): + with _followup_billing_scope(logging_obj, kwargs): response_followup = await litellm.acompletion( model=full_model_name, messages=patch.messages, @@ -416,15 +499,16 @@ async def _execute_chat_completion_agentic_plan( str(e), ) _leave_pricing_to_the_request(response_followup) + billed: Final = _billed_with_the_round_it_follows(response, response_followup, logging_obj) replay: Final = _converted_stream_replay( - response_followup, + billed, kwargs=kwargs, depth=depth, model=model, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) - return response_followup if replay is None else replay + return billed if replay is None else replay finally: try: await callback.async_agentic_loop_cleanup_hook(plan=plan, kwargs=kwargs) @@ -547,6 +631,7 @@ async def maybe_run_chat_completion_agentic_loop( continue return await _execute_chat_completion_agentic_plan( + response=response, plan=plan, callback=callback, model=model, @@ -560,14 +645,42 @@ async def maybe_run_chat_completion_agentic_loop( fingerprint=fingerprint, ) - return _converted_stream_replay( - response, - kwargs=request_kwargs, - depth=depth, - model=model, - custom_llm_provider=custom_llm_provider, - logging_obj=logging_obj, + bridged_chat_turn: Final = ( + await run_chat_completion_agentic_hooks( + response=response, + model=model, + messages=messages, + optional_params=optional_params, + kwargs=dict(request_kwargs), + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + stream=stream, + asking_only=runs_its_own_chat_completion_loop, + ) + if isinstance(logging_obj, LiteLLMLoggingObject) and logging_obj.bridged_chat_request is not None + else None ) + match bridged_chat_turn: + case CustomStreamWrapper(): + return bridged_chat_turn + case ModelResponse(): + return _end_refused_turn( + bridged_chat_turn, + kwargs=request_kwargs, + depth=depth, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + case _: + return _converted_stream_replay( + response, + kwargs=request_kwargs, + depth=depth, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) async def run_chat_completion_agentic_hooks( @@ -580,13 +693,14 @@ async def run_chat_completion_agentic_hooks( logging_obj: LiteLLMLoggingObject, custom_llm_provider: str, stream: bool, + asking_only: Callable[[CustomLogger], bool] | None = None, ) -> object | None: callbacks: Final = _custom_loggers(logging_obj) depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) tools: Final = _request_tools(optional_params) hook_kwargs: Final = {**kwargs, "custom_llm_provider": custom_llm_provider} - for callback in callbacks: + for callback in callbacks if asking_only is None else tuple(filter(asking_only, callbacks)): try: should_run, tool_calls = await callback.async_should_run_chat_completion_agentic_loop( response=response, @@ -664,6 +778,7 @@ async def run_chat_completion_agentic_hooks( continue return await _execute_chat_completion_agentic_plan( + response=response, plan=plan, callback=callback, model=model, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2f421c4d536..ed276eebbc8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2426,7 +2426,7 @@ class Logging(LiteLLMLoggingBaseClass): ) result_hidden_params: Final = getattr(priced_result, HIDDEN_PARAMS_ATTR, None) or MappingProxyType({}) - if isinstance(priced_result, (BaseModel, HttpxBinaryResponseContent)) and hasattr( + if isinstance(priced_result, (BaseModel, HttpxBinaryResponseContent, litellm.CustomStreamWrapper)) and hasattr( priced_result, "_hidden_params" ): hidden_params: Final = result_hidden_params @@ -3003,6 +3003,9 @@ class Logging(LiteLLMLoggingBaseClass): self.record_partial_usage_for_failure(spend.usage, spend.cost) self.cost_breakdown = spend.cost_breakdown + def record_every_round_breakdown(self, cost_breakdown: CostBreakdown | None) -> None: + self.cost_breakdown = cost_breakdown + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" @@ -3082,11 +3085,8 @@ class Logging(LiteLLMLoggingBaseClass): """A billed response the router discards after its success log was parked (a safeguard refusal it falls back from) is owed by whichever attempt ends the request, instead of being logged in its place.""" parked: Final = self.deferred_success_result - if parked is None or self.enqueue_deferred_logging is None: + if parked is None or not self.unpark_success_log(): return False - self.deferred_success_result = None - self.enqueue_deferred_logging = None - self._async_success_scheduled = False served: Final = self.served_round(parked[0]) if served is None: return True @@ -3097,6 +3097,14 @@ class Logging(LiteLLMLoggingBaseClass): self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, served) return True + def unpark_success_log(self) -> bool: + if self.deferred_success_result is None or self.enqueue_deferred_logging is None: + return False + self.deferred_success_result = None + self.enqueue_deferred_logging = None + self._async_success_scheduled = False + return True + def mark_stash_billed_by_continuation(self) -> None: self._stash_billed_by_continuation = True diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 46fb41608cd..392609b7ef7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -480,6 +480,7 @@ class CustomStreamWrapper: _response_headers: dict | httpx.Headers | None = None, count_prompt_tokens: Callable[[], int] | None = None, is_mock: bool = False, + response_cost: float | None = None, ): self.model = model self.make_call = make_call @@ -534,6 +535,7 @@ class CustomStreamWrapper: self._hidden_params: dict[str, object] = { "model_id": (_model_info.get("id", None)), "api_base": _api_base, + **({} if response_cost is None else {"response_cost": response_cost}), } # returned as x-litellm-model-id response header in proxy self._hidden_params["additional_headers"] = process_response_headers( @@ -574,7 +576,7 @@ class CustomStreamWrapper: # this cache must be removed. self._base_hidden_params: dict[str, object] = { **self._hidden_params, - "response_cost": None, + "response_cost": response_cost, } self._post_streaming_hooks: list | None = None diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 90ba74fcd4f..7831ee61e3e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( bill_first_round_on_failure, holds_a_call_the_choice_excludes, run_chat_completion_agentic_hooks, + runs_its_own_chat_completion_loop, with_caller_tool_choice, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params @@ -169,6 +170,7 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, AgenticLoopSafetyError, AgenticLoopToolChoiceError, + bridged_from_chat_completion, converted_stream_requested, is_interception_internal_key, ) @@ -3438,6 +3440,7 @@ class BaseLLMHTTPHandler: ) agentic_kwargs: Final = dict(litellm_params) + chat_completion_loop_follows: Final = bridged_from_chat_completion(agentic_kwargs) final_response: Final = await self._call_agentic_completion_hooks( response=initial_response, model=model, @@ -3449,10 +3452,15 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, kwargs=agentic_kwargs, api_surface="responses", + skip_chat_completion_loops=chat_completion_loop_follows, ) result: Final = final_response if final_response is not None else initial_response - if converted_stream_requested(agentic_kwargs) and not agentic_kwargs.get("_agentic_loop_depth"): + if ( + not chat_completion_loop_follows + and converted_stream_requested(agentic_kwargs) + and not agentic_kwargs.get("_agentic_loop_depth") + ): return self._wrap_responses_response_as_fake_stream( result=result, model=model, @@ -6397,6 +6405,7 @@ class BaseLLMHTTPHandler: custom_llm_provider: str, kwargs: dict, api_surface: str = "anthropic_messages", + skip_chat_completion_loops: bool = False, ) -> Any | None: """ Call agentic completion hooks for all custom loggers (Anthropic Messages API). @@ -6422,7 +6431,9 @@ class BaseLLMHTTPHandler: } for callback in callbacks: - if not isinstance(callback, CustomLogger): + if not isinstance(callback, CustomLogger) or ( + skip_chat_completion_loops and runs_its_own_chat_completion_loop(callback) + ): continue should_run: bool = False diff --git a/litellm/responses/main.py b/litellm/responses/main.py index db5ac5a264b..96155da88c7 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -78,6 +78,7 @@ from litellm.responses.mcp.round_billing import ( ) from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.types.integrations.custom_logger import CHAT_COMPLETION_AGENTIC_SURFACE from litellm.types.llms.openai import ( PromptObject, Reasoning, @@ -1161,6 +1162,14 @@ def _resolve_prompt_swapped_provider( ) +def _record_model_for_the_responses_agentic_loop( + logging_obj: LiteLLMLoggingObj, model: str, request_kwargs: Mapping[str, object] +) -> None: + if request_kwargs.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE: + return + logging_obj.record_agentic_loop_params(AgenticLoopParams(model=model)) + + def _apply_prompt_management_to_responses_call( input: str | ResponseInputParam, tools: Iterable[ToolParam] | None, @@ -1606,7 +1615,8 @@ def responses( local_vars=local_vars, use_chat_completions_api=use_chat_completions_api, ) - litellm_logging_obj.record_agentic_loop_params(AgenticLoopParams(model=model)) + request_kwargs: Final = cast("Mapping[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict + _record_model_for_the_responses_agentic_loop(litellm_logging_obj, model, request_kwargs) from litellm.litellm_core_utils.core_helpers import attribute_or_none @@ -1731,7 +1741,6 @@ def responses( if current_reasoning is None and reasoning_effort is not None else current_reasoning ) - routing_kwargs: Final = cast("Mapping[str, object]", kwargs) # cast-ok: kwargs is the untyped request dict deployment_reasoning: Final = cast( # cast-ok: the resolver hands a Reasoning back in the shape it was given Reasoning | None, resolve_deployment_reasoning_effort( @@ -1739,9 +1748,9 @@ def responses( model=model, custom_llm_provider=custom_llm_provider, deployment_model_info=deployment_model_info, - model_group=model_group_from_metadata(routing_kwargs.get("litellm_metadata")), + model_group=model_group_from_metadata(request_kwargs.get("litellm_metadata")), drop_params=litellm.drop_params is True - or normalize_drop_params(routing_kwargs.get("drop_params")) is True, + or normalize_drop_params(request_kwargs.get("drop_params")) is True, ).value, ) response_api_optional_params: Final[ResponsesAPIOptionalRequestParams] = ( diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 127b3d61753..b541b036d70 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -50,6 +50,10 @@ def converted_stream_requested(params: Mapping[str, object]) -> bool: return any(bool(params.get(key)) for key in CONVERTED_STREAM_KEYS) +def bridged_from_chat_completion(params: Mapping[str, object]) -> bool: + return params.get("_agentic_loop_api_surface") == CHAT_COMPLETION_AGENTIC_SURFACE + + def as_converted_stream( params: Mapping[str, object], prefix: str ) -> dict[str, object]: # mutable-ok: deployment hooks hand back the kwargs dict their caller keeps editing diff --git a/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py b/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py index e5ef356b0f3..3b6682ea67e 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py @@ -11,11 +11,21 @@ Also covers the max_agentic_loops knob on websearch_interception_params, from config.yaml through to the settings the loop actually reads. """ +import asyncio import json import time +from collections.abc import Callable, Iterator, Mapping +from contextlib import suppress +from datetime import datetime, timezone +from types import MappingProxyType +from typing import ClassVar, Final, TypedDict from unittest.mock import MagicMock, patch +import httpx import pytest +import respx +from pydantic import TypeAdapter +from typing_extensions import ReadOnly import litellm from litellm.exceptions import AuthenticationError, RateLimitError @@ -27,14 +37,19 @@ from litellm.integrations.websearch_interception.handler import ( from litellm.integrations.websearch_interception.tools import get_litellm_web_search_tool from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.litellm_core_utils.model_cost_table import changed_in_place from litellm.litellm_core_utils.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, ) +from litellm.llms.base_llm.anthropic_messages.transformation import BaseAnthropicMessagesConfig from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost from litellm.secret_managers.main import get_secret from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import CallTypes from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -1164,3 +1179,757 @@ class TestRebuiltStreamIsWellFormed: assert result_block["tool_use_id"] == "srvtoolu_01" assert result_block["content"][0]["url"] == "https://example.com/h100" + + +OPENAI_COMPATIBLE_API_BASE: Final = "https://openai-compatible.example/v1" +SEARCHES_THE_MODEL_ASKS_FOR: Final = DEFAULT_MAX_AGENTIC_LOOPS + 2 + + +def _openai_reply_asking_for_another_search_until_enough_results(request: httpx.Request) -> httpx.Response: + searches_done: Final = sum(1 for message in json.loads(request.content)["messages"] if message["role"] == "tool") + search_call: Final = { + "id": f"call_{searches_done}", + "type": "function", + "function": { + "name": INTERNAL_TOOL_NAME, + "arguments": json.dumps({"query": f"world cup winner, source {searches_done}"}), + }, + } + message: Final = ( + {"role": "assistant", "content": "Argentina won the 2022 World Cup."} + if searches_done >= SEARCHES_THE_MODEL_ASKS_FOR + else {"role": "assistant", "content": None, "tool_calls": [search_call]} + ) + return httpx.Response( + 200, + json={ + "id": f"chatcmpl-{searches_done}", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.6", + "choices": [ + {"index": 0, "message": message, "finish_reason": "tool_calls" if "tool_calls" in message else "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_cap", "loops_allowed"), + [({"max_agentic_loops": 1}, 1), ({}, DEFAULT_MAX_AGENTIC_LOOPS)], + ids=["max_agentic_loops=1", "default cap"], +) +async def test_openai_chat_completion_stops_searching_at_the_agentic_loop_cap( + request_cap: dict[str, int], + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"])]) + litellm.in_memory_llm_clients_cache.flush_cache() + openai_route: Final = respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/chat/completions").mock( + side_effect=_openai_reply_asking_for_another_search_until_enough_results + ) + + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + **request_cap, + ) + + assert openai_route.call_count == loops_allowed + 1 + assert perplexity_search.call_count == loops_allowed + assert {call.request.headers["authorization"] for call in openai_route.calls} == {"Bearer sk-test"} + + +def _round_input_tokens(searches_done: int) -> int: + return 10 + 1000 * searches_done + + +def _responses_reply_asking_for_searches_until(searches_wanted: int) -> Callable[[httpx.Request], httpx.Response]: + def reply(request: httpx.Request) -> httpx.Response: + request_input: Final = json.loads(request.content)["input"] + searches_done: Final = sum(1 for item in request_input if item.get("type") == "function_call_output") + output_item: Final = ( + { + "type": "message", + "id": f"msg_{searches_done}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Argentina won the 2022 World Cup.", "annotations": []}], + } + if searches_done >= searches_wanted + else { + "type": "function_call", + "id": f"fc_{searches_done}", + "call_id": f"call_{searches_done}", + "name": INTERNAL_TOOL_NAME, + "arguments": json.dumps({"query": f"world cup winner, source {searches_done}"}), + "status": "completed", + } + ) + return httpx.Response( + 200, + json={ + "id": f"resp_{searches_done}", + "object": "response", + "created_at": 0, + "status": "completed", + "model": "gpt-5.6", + "output": [output_item], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": { + "input_tokens": _round_input_tokens(searches_done), + "output_tokens": 5, + "total_tokens": _round_input_tokens(searches_done) + 5, + }, + }, + ) + + return reply + + +_responses_reply_asking_for_another_search_until_enough_results: Final = _responses_reply_asking_for_searches_until( + SEARCHES_THE_MODEL_ASKS_FOR +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_cap", "loops_allowed"), + [({"max_agentic_loops": 1}, 1), ({}, DEFAULT_MAX_AGENTIC_LOOPS)], + ids=["max_agentic_loops=1", "default cap"], +) +async def test_openai_chat_completion_bridged_to_the_responses_api_stops_searching_at_the_agentic_loop_cap( + request_cap: dict[str, int], + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"])]) + litellm.in_memory_llm_clients_cache.flush_cache() + responses_route: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + **request_cap, + ) + + assert responses_route.call_count == loops_allowed + 1 + assert perplexity_search.call_count == loops_allowed + assert all("_agentic_loop_api_surface" not in json.loads(call.request.content) for call in responses_route.calls) + + +@pytest.mark.asyncio +async def test_responses_pinned_openai_chat_completion_keeps_its_search_follow_ups_on_the_responses_api( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"])]) + litellm.in_memory_llm_clients_cache.flush_cache() + responses_route: Final = respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + chat_route: Final = respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/chat/completions").respond(404) + + response = await litellm.acompletion( + model="openai/responses/responses-only-model", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=SEARCHES_THE_MODEL_ASKS_FOR, + ) + + assert chat_route.call_count == 0 + assert responses_route.call_count == SEARCHES_THE_MODEL_ASKS_FOR + 1 + assert response.choices[0].message.content == "Argentina won the 2022 World Cup." + + +class _LoggedRow(TypedDict): + call_type: ReadOnly[str] + prompt_tokens: ReadOnly[int] + response_cost: ReadOnly[float] + cost_breakdown: ReadOnly[Mapping[str, object] | None] + + +class _ModelSpendRowRecorder(CustomLogger): + def __init__(self, rows_expected: int) -> None: + super().__init__() + self.prompt_tokens: Final[list[int]] = [] + self.response_costs: Final[list[float]] = [] + self.breakdown_totals: Final[list[object]] = [] + self._rows_expected: Final = rows_expected + self.all_rows_logged: Final = asyncio.Event() + self.streamed_reply_logged: Final = asyncio.Event() + + def _record(self, kwargs: dict[str, object]) -> None: + row: Final = TypeAdapter(_LoggedRow).validate_python(kwargs["standard_logging_object"]) + if row["call_type"] != "acompletion": + return + self.prompt_tokens.append(row["prompt_tokens"]) + self.response_costs.append(row["response_cost"]) + self.breakdown_totals.append((row["cost_breakdown"] or {}).get("total_cost")) + if kwargs.get("stream") is True: + self.streamed_reply_logged.set() + if len(self.prompt_tokens) >= self._rows_expected: + self.all_rows_logged.set() + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + +def _every_round_input_tokens(loops_allowed: int) -> int: + return sum(_round_input_tokens(searches_done) for searches_done in range(loops_allowed + 1)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loops_allowed", [1, 2]) +async def test_responses_pinned_openai_chat_completion_bills_every_search_round_it_served( + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + recorder: Final = _ModelSpendRowRecorder(rows_expected=loops_allowed + 1) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + responses_route: Final = respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=loops_allowed, + ) + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(recorder.all_rows_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert responses_route.call_count == loops_allowed + 1 + assert sum(recorder.prompt_tokens) == _every_round_input_tokens(loops_allowed) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loops_allowed", [1, 2]) +async def test_streamed_responses_pinned_openai_chat_completion_bills_every_search_round_once( + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + recorder: Final = _ModelSpendRowRecorder(rows_expected=2) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + stream: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=loops_allowed, + stream=True, + ) + assert isinstance(stream, litellm.CustomStreamWrapper) + _ = [chunk async for chunk in stream] + await asyncio.wait_for(recorder.streamed_reply_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert sum(recorder.prompt_tokens) == _every_round_input_tokens(loops_allowed) + + +def _responses_reply_failing_after(searches_that_succeed: int): + def reply(request: httpx.Request) -> httpx.Response: + request_input: Final = json.loads(request.content)["input"] + searches_done: Final = sum(1 for item in request_input if item.get("type") == "function_call_output") + if searches_done >= searches_that_succeed: + return httpx.Response(400, json={"error": {"message": "rejected", "type": "invalid_request_error"}}) + return _responses_reply_asking_for_another_search_until_enough_results(request) + + return reply + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non_streamed", "streamed"]) +@pytest.mark.parametrize("rounds_served", [1, 2]) +async def test_responses_pinned_openai_chat_completion_bills_each_served_round_once_when_a_follow_up_fails( + rounds_served: int, + stream: bool, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + recorder: Final = _ModelSpendRowRecorder(rows_expected=2 if stream else 2 * rounds_served + 1) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_failing_after(rounds_served) + ) + + with pytest.raises(litellm.BadRequestError): + await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=SEARCHES_THE_MODEL_ASKS_FOR, + stream=stream, + ) + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(recorder.all_rows_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert sum(recorder.prompt_tokens) == sum(_round_input_tokens(served) for served in range(rounds_served)) + + +class _ProxyParksTheRequestSuccessLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.request_logging_obj: Logging | None = None + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if self.request_logging_obj is not None or not isinstance(logging_obj, Logging): + return None + logging_obj.defer_async_logging = True + self.request_logging_obj = logging_obj + return None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rounds_served", [1, 2]) +async def test_responses_pinned_openai_chat_completion_bills_its_parked_first_round_when_a_follow_up_fails( + rounds_served: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + parker: Final = _ProxyParksTheRequestSuccessLog() + recorder: Final = _ModelSpendRowRecorder(rows_expected=1) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), parker, recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_failing_after(rounds_served) + ) + + with pytest.raises(litellm.BadRequestError): + await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=SEARCHES_THE_MODEL_ASKS_FOR, + ) + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(recorder.all_rows_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert parker.request_logging_obj is not None + assert sum(recorder.prompt_tokens) == sum(_round_input_tokens(served) for served in range(rounds_served)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loops_allowed", [1, 2]) +async def test_responses_pinned_openai_chat_completion_parks_one_row_billing_every_round_until_the_proxy_releases_it( + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + parker: Final = _ProxyParksTheRequestSuccessLog() + recorder: Final = _ModelSpendRowRecorder(rows_expected=loops_allowed + 1) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), parker, recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + response: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=loops_allowed, + ) + await GLOBAL_LOGGING_WORKER.flush() + rows_before_the_proxy_releases_the_log: Final = tuple(recorder.prompt_tokens) + request_logging_obj: Final = parker.request_logging_obj + assert request_logging_obj is not None + prepared: Final = prepare_response_cost(request_logging_obj, response, datetime.now(timezone.utc)) + release: Final = request_logging_obj.enqueue_deferred_logging + assert release is not None + release() + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(recorder.all_rows_logged.wait(), timeout=1) + await GLOBAL_LOGGING_WORKER.flush() + + assert rows_before_the_proxy_releases_the_log == () + assert prepared is not None + assert prepared.usage.prompt_tokens == _every_round_input_tokens(loops_allowed) + assert recorder.prompt_tokens == [_every_round_input_tokens(loops_allowed)] + + +ROUND_INPUT_TOKENS_UNDER_THE_LONG_CONTEXT_TIER: Final = 200_000 +ROUND_OUTPUT_TOKENS: Final = 5 +TEST_OWNED_TIER_PRICES: Final = MappingProxyType( + { + "input_cost_per_token": 1e-06, + "input_cost_per_token_above_272k_tokens": 3e-06, + "output_cost_per_token": 2e-06, + "output_cost_per_token_above_272k_tokens": 5e-06, + } +) + + +@pytest.fixture +def gpt_5_6_at_test_owned_tier_prices(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + vendor_entry: Final = litellm.model_cost["gpt-5.6"] + assert isinstance(vendor_entry, dict) + unpriced: Final = {key: value for key, value in vendor_entry.items() if "cost" not in key} + monkeypatch.setitem(litellm.model_cost, "gpt-5.6", {**unpriced, **TEST_OWNED_TIER_PRICES}) + changed_in_place() + yield + changed_in_place() + + +def _long_context_round_reply(request: httpx.Request) -> httpx.Response: + reply: Final = _responses_reply_asking_for_another_search_until_enough_results(request) + return httpx.Response( + 200, + json={ + **reply.json(), + "usage": { + "input_tokens": ROUND_INPUT_TOKENS_UNDER_THE_LONG_CONTEXT_TIER, + "output_tokens": ROUND_OUTPUT_TOKENS, + "total_tokens": ROUND_INPUT_TOKENS_UNDER_THE_LONG_CONTEXT_TIER + ROUND_OUTPUT_TOKENS, + }, + }, + ) + + +def _gpt_5_6_price(rounds: int) -> float: + prompt_cost, completion_cost = litellm.cost_per_token( + model="gpt-5.6", + custom_llm_provider="openai", + prompt_tokens=ROUND_INPUT_TOKENS_UNDER_THE_LONG_CONTEXT_TIER * rounds, + completion_tokens=ROUND_OUTPUT_TOKENS * rounds, + ) + return prompt_cost + completion_cost + + +def _price_of_rounds_priced_one_by_one(rounds: int) -> float: + one_round: Final = _gpt_5_6_price(1) + assert _gpt_5_6_price(rounds) != pytest.approx(rounds * one_round) + return rounds * one_round + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loops_allowed", [1, 2]) +@pytest.mark.usefixtures("gpt_5_6_at_test_owned_tier_prices") +async def test_responses_pinned_openai_chat_completion_prices_each_round_of_its_one_row_at_that_rounds_size( + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + expected_cost: Final = _price_of_rounds_priced_one_by_one(loops_allowed + 1) + parker: Final = _ProxyParksTheRequestSuccessLog() + recorder: Final = _ModelSpendRowRecorder(rows_expected=1) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), parker, recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock(side_effect=_long_context_round_reply) + + response: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=loops_allowed, + ) + request_logging_obj: Final = parker.request_logging_obj + assert request_logging_obj is not None + prepared: Final = prepare_response_cost(request_logging_obj, response, datetime.now(timezone.utc)) + release: Final = request_logging_obj.enqueue_deferred_logging + assert release is not None + release() + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(recorder.all_rows_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert prepared is not None + assert prepared.amount == pytest.approx(expected_cost) + assert recorder.response_costs == [pytest.approx(expected_cost)] + assert recorder.breakdown_totals == [pytest.approx(expected_cost)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("loops_allowed", [1, 2]) +@pytest.mark.usefixtures("gpt_5_6_at_test_owned_tier_prices") +async def test_streamed_responses_pinned_openai_chat_completion_prices_each_round_at_that_rounds_size( + loops_allowed: int, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + expected_cost: Final = _price_of_rounds_priced_one_by_one(loops_allowed + 1) + recorder: Final = _ModelSpendRowRecorder(rows_expected=2) + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), recorder] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock(side_effect=_long_context_round_reply) + + stream: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + max_agentic_loops=loops_allowed, + stream=True, + ) + assert isinstance(stream, litellm.CustomStreamWrapper) + _ = [chunk async for chunk in stream] + await asyncio.wait_for(recorder.streamed_reply_logged.wait(), timeout=5) + await GLOBAL_LOGGING_WORKER.flush() + + assert sum(recorder.response_costs) == pytest.approx(expected_cost) + assert recorder.breakdown_totals == pytest.approx(recorder.response_costs) + + +class _AnswersOnTheResponsesApi(CustomLogger): + async def async_should_run_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return kwargs.get("_agentic_loop_api_surface") == "responses", {} + + async def async_build_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + anthropic_messages_provider_config: BaseAnthropicMessagesConfig | None, + anthropic_messages_optional_request_params: dict, + logging_obj: Logging, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + assert isinstance(response, ResponsesAPIResponse) + return AgenticLoopPlan( + response_override=response.model_copy( + update={ + "output": [ + { + "type": "message", + "id": "msg_handled", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Handled on the Responses API.", "annotations": []} + ], + } + ] + } + ) + ) + + +@pytest.mark.asyncio +async def test_responses_pinned_openai_chat_completion_still_runs_agentic_hooks_that_only_handle_the_responses_api( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + monkeypatch.setattr( + litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"]), _AnswersOnTheResponsesApi()] + ) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + response: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + tools=[{"type": "web_search"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + ) + + assert response.choices[0].message.content == "Handled on the Responses API." + + +class _AnswersOnBothApis(_AnswersOnTheResponsesApi): + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return isinstance(response, litellm.ModelResponse), {} + + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + optional_params: dict, + logging_obj: Logging, + stream: bool, + kwargs: dict, + ) -> AgenticLoopPlan: + assert isinstance(response, litellm.ModelResponse) + return AgenticLoopPlan( + response_override=litellm.ModelResponse( + model=response.model, + choices=[ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "Handled on the chat API."}, + } + ], + ) + ) + + +class _DeclinesOnTheChatApi(_AnswersOnTheResponsesApi): + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return False, {} + + +class _WantsChatHooksOnBridgedChat(_AnswersOnBothApis): + wants_chat_completion_hooks_on_bridged_chat: ClassVar[bool] = True + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("perplexity_search") +@pytest.mark.parametrize( + ("callback", "answer"), + [ + (_AnswersOnBothApis(), "Handled on the Responses API."), + (_DeclinesOnTheChatApi(), "Handled on the Responses API."), + (_WantsChatHooksOnBridgedChat(), "Handled on the chat API."), + ], +) +async def test_responses_pinned_openai_chat_completion_asks_the_chat_hook_only_of_callbacks_that_opt_in( + callback: CustomLogger, answer: str, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "callbacks", [callback]) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_another_search_until_enough_results + ) + + response: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + ) + + assert response.choices[0].message.content == answer + + +async def _answer_text(response: litellm.ModelResponse | litellm.CustomStreamWrapper) -> str | None: + if isinstance(response, litellm.CustomStreamWrapper): + return "".join([chunk.choices[0].delta.content or "" async for chunk in response if chunk.choices]) + return response.choices[0].message.content + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non-streamed", "streamed"]) +async def test_responses_pinned_openai_chat_completion_runs_the_search_its_web_search_options_ask_for( + stream: bool, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + perplexity_search: respx.Route, +): + monkeypatch.setattr(litellm, "callbacks", [WebSearchInterceptionLogger(enabled_providers=["openai"])]) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(f"{OPENAI_COMPATIBLE_API_BASE}/responses").mock( + side_effect=_responses_reply_asking_for_searches_until(1) + ) + + response: Final = await litellm.acompletion( + model="openai/responses/gpt-5.6", + messages=[{"role": "user", "content": "Who won the last World Cup?"}], + web_search_options={}, + allowed_openai_params=["web_search_options"], + api_key="sk-test", + api_base=OPENAI_COMPATIBLE_API_BASE, + stream=stream, + ) + + assert perplexity_search.call_count == 1 + assert await _answer_text(response) == "Argentina won the 2022 World Cup." diff --git a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py index 7d41cbd6628..0fdfef1b842 100644 --- a/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py +++ b/tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py @@ -1069,6 +1069,7 @@ class TestAgenticFollowUpKeepsTheProviderPrefix: model: str, custom_llm_provider: str, patched_model: str | None = None ) -> ModelResponse: response: Final = await _execute_chat_completion_agentic_plan( + response=_tool_call_model_response(), plan=AgenticLoopPlan( run_agentic_loop=True, request_patch=AgenticLoopRequestPatch( diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 1936fc1d2df..72afb840766 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -2700,6 +2700,83 @@ def test_openai_config_hoists_a_later_developer_message_only_for_a_model_off_ope assert [message["role"] for message in messages] == expected_roles +class _ChatCompletionOnlyAgenticLoop(CustomLogger): + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: list[dict], + tools: list[dict] | None, + stream: bool, + custom_llm_provider: str, + kwargs: dict, + ) -> tuple[bool, dict]: + return isinstance(response, ModelResponse) and response.choices[0].finish_reason == "tool_calls", {} + + async def async_run_chat_completion_agentic_loop( + self, + tools: dict, + model: str, + messages: list[dict], + response: object, + optional_params: dict, + logging_obj: object, + stream: bool, + kwargs: dict, + ) -> ModelResponse: + return ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "order shipped"}}] + ) + + +@pytest.mark.asyncio +async def test_openai_runs_a_callback_that_only_implements_the_chat_completion_agentic_loop( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "callbacks", [_ChatCompletionOnlyAgenticLoop()]) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://openai-compatible.example/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-lookup", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_lookup", + "type": "function", + "function": {"name": "lookup_order", "arguments": '{"order_id": "42"}'}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + ) + ) + + response: Final = await litellm.acompletion( + model="openai/gpt-5.6", + messages=[{"role": "user", "content": "Where is order 42?"}], + tools=[{"type": "function", "function": {"name": "lookup_order", "parameters": {"type": "object"}}}], + api_key="sk-test", + api_base="https://openai-compatible.example/v1", + ) + + assert response.choices[0].message.content == "order shipped" + + @pytest.mark.asyncio async def test_sync_batch_list_given_an_async_client_returns_a_page_the_caller_can_await( http_backend: ModuleType,