diff --git a/litellm/integrations/code_interpreter_interception/handler.py b/litellm/integrations/code_interpreter_interception/handler.py index 69b528a29f8..b035732afa6 100644 --- a/litellm/integrations/code_interpreter_interception/handler.py +++ b/litellm/integrations/code_interpreter_interception/handler.py @@ -10,6 +10,7 @@ import json import time import uuid from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypedDict, runtime_checkable from pydantic import ValidationError @@ -73,7 +74,7 @@ class CodeInterpreterLogOutput(TypedDict): class CodeInterpreterCall(TypedDict): id: str type: Literal["code_interpreter_call"] - status: Literal["completed"] + status: Literal["completed", "failed"] code: str container_id: str | None outputs: list[CodeInterpreterLogOutput] @@ -141,6 +142,39 @@ class _SupportsOutput(Protocol): _CachedContainer: TypeAlias = tuple[ContainerHandle, SandboxToolParams | None, float, str | None] +@dataclass(frozen=True, slots=True) +class _ReadySandbox: + container: ContainerHandle + params: SandboxToolParams | None + + +@dataclass(frozen=True, slots=True) +class _SandboxUnavailable: + output: str + + +_Sandbox: TypeAlias = _ReadySandbox | _SandboxUnavailable + + +@dataclass(frozen=True, slots=True) +class _ToolRun: + output: str + status: Literal["completed", "failed"] + + +_SANDBOX_UNAVAILABLE_OUTPUT: Final = "The code did not run because the code execution sandbox is unavailable." +_SANDBOX_LOST_RESULT_OUTPUT: Final = ( + "The code execution sandbox failed while running the code, so the code may have run but its result is lost." +) + + +def _sandbox_unavailable(step: str, error: Exception, output: str = _SANDBOX_UNAVAILABLE_OUTPUT) -> _SandboxUnavailable: + verbose_logger.warning( + "CodeInterpreterInterception: %s failed, answering the tool call with the failure: %s", step, error + ) + return _SandboxUnavailable(output=output) + + def _output_item_type(item: object) -> object: if isinstance(item, dict): item_mapping: Final[dict[str, object]] = item @@ -189,6 +223,14 @@ def _extract_identity(kwargs: Mapping[str, object]) -> str: return identity if isinstance(identity, str) else "" +def _bridged_from_responses(kwargs: Mapping[str, object], call_type: CallTypes | None) -> bool: + route: Final = getattr(kwargs.get("litellm_logging_obj"), "call_type", None) + return call_type in (CallTypes.completion, CallTypes.acompletion) and route in ( + CallTypes.responses, + CallTypes.aresponses, + ) + + def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> SandboxToolParams | None: if sandbox_tool_name is None: return None @@ -261,7 +303,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: CallTypes | None ) -> dict | None: - if not kwargs.get("_agentic_loop_depth"): + if not kwargs.get("_agentic_loop_depth") and not _bridged_from_responses(kwargs, call_type): kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None) kwargs.pop(_SANDBOX_KEY, None) self._strip_interception_metadata(kwargs) @@ -451,16 +493,16 @@ class CodeInterpreterInterceptionLogger(CustomLogger): sandbox_key: Final = self._extract_sandbox_key(kwargs) is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY)) identity: Final = _extract_identity(kwargs) if is_session else None - container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) + sandbox: Final = await self._acquire_sandbox(cache_key=sandbox_key, identity=identity) try: - container_id: Final = self._container_id(container) + container_id: Final = self._sandbox_container_id(sandbox) input_list: Final = self._normalize_messages(messages) code_interpreter_calls: Final[list[CodeInterpreterCall]] = [] for tool_call in tool_calls: arguments = tool_call.get("arguments", "") code = self._parse_code(arguments) - stdout = await self._run_tool_call(container=container, params=params, arguments=arguments) + tool_run = await self._sandbox_tool_output(sandbox=sandbox, arguments=arguments) input_list.append( { "type": "function_call", @@ -473,22 +515,23 @@ class CodeInterpreterInterceptionLogger(CustomLogger): { "type": "function_call_output", "call_id": tool_call.get("call_id"), - "output": stdout, + "output": tool_run.output, } ) code_interpreter_calls.append( { "id": f"ci_{uuid.uuid4().hex}", "type": "code_interpreter_call", - "status": "completed", + "status": tool_run.status, "code": code, "container_id": container_id, - "outputs": ([{"type": "logs", "logs": stdout}] if stdout else []), + "outputs": ([{"type": "logs", "logs": tool_run.output}] if tool_run.output else []), } ) except Exception: await self._delete_container_for_cache_key(sandbox_key) raise + await self._evict_sandbox_if_a_run_failed(sandbox_key, sandbox, code_interpreter_calls) optional_params: Final = anthropic_messages_optional_request_params request_patch: Final = AgenticLoopRequestPatch( @@ -526,14 +569,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger): sandbox_key: Final = self._extract_sandbox_key(kwargs) is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY)) identity: Final = _extract_identity(kwargs) if is_session else None - container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity) + sandbox: Final = await self._acquire_sandbox(cache_key=sandbox_key, identity=identity) try: - container_id: Final = self._container_id(container) + container_id: Final = self._sandbox_container_id(sandbox) tool_results: Final = [ await self._build_chat_completion_tool_result( - container=container, - params=params, + sandbox=sandbox, tool_call=tool_call, container_id=container_id, ) @@ -544,6 +586,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger): raise tool_messages: Final = [result[0] for result in tool_results] code_interpreter_calls: Final = [result[1] for result in tool_results] + await self._evict_sandbox_if_a_run_failed(sandbox_key, sandbox, code_interpreter_calls) request_patch: Final = AgenticLoopRequestPatch( model=model, @@ -569,9 +612,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger): ) @staticmethod - def _container_id(container: ContainerHandle) -> str | None: - container_id: Final[object] = getattr(container, "id", None) - return container_id if isinstance(container_id, str) else None + def _sandbox_container_id(sandbox: _Sandbox) -> str | None: + match sandbox: + case _ReadySandbox(container=container): + container_id: Final[object] = getattr(container, "id", None) + return container_id if isinstance(container_id, str) else None + case _SandboxUnavailable(): + return None @staticmethod def _agentic_tool_calls(tools: dict[str, object]) -> list[CodeExecutionToolCall]: @@ -588,28 +635,27 @@ class CodeInterpreterInterceptionLogger(CustomLogger): async def _build_chat_completion_tool_result( self, - container: ContainerHandle, - params: SandboxToolParams | None, + sandbox: _Sandbox, tool_call: CodeExecutionToolCall, container_id: str | None, ) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]: arguments: Final = tool_call.get("arguments", "") code: Final = self._parse_code(arguments) - stdout: Final = await self._run_tool_call(container=container, params=params, arguments=arguments) + tool_run: Final = await self._sandbox_tool_output(sandbox=sandbox, arguments=arguments) tool_call_id: Final = tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex return ( { "role": "tool", "tool_call_id": tool_call_id, - "content": stdout, + "content": tool_run.output, }, { "id": f"ci_{uuid.uuid4().hex}", "type": "code_interpreter_call", - "status": "completed", + "status": tool_run.status, "code": code, "container_id": container_id, - "outputs": [{"type": "logs", "logs": stdout}] if stdout else [], + "outputs": [{"type": "logs", "logs": tool_run.output}] if tool_run.output else [], }, ) @@ -694,19 +740,55 @@ class CodeInterpreterInterceptionLogger(CustomLogger): except (json.JSONDecodeError, TypeError, AttributeError): return "" - async def _run_tool_call(self, container: ContainerHandle, params: SandboxToolParams | None, arguments: str) -> str: + async def _evict_sandbox_if_a_run_failed( + self, sandbox_key: str | None, sandbox: _Sandbox, calls: Sequence[CodeInterpreterCall] + ) -> None: + if not sandbox_key or not isinstance(sandbox, _ReadySandbox): + return + if not any(call["status"] == "failed" for call in calls): + return + cached: Final = self._container_cache.get(sandbox_key) + if cached is None or cached[0] is not sandbox.container: + return + self._container_cache.pop(sandbox_key) + await self._delete_container(container=cached[0], params=cached[1]) + + async def _sandbox_tool_output(self, sandbox: _Sandbox, arguments: str) -> _ToolRun: + match sandbox: + case _SandboxUnavailable(output=output): + return _ToolRun(output=output, status="failed") + case _ReadySandbox(container=container, params=params): + match await self._run_tool_call(container=container, params=params, arguments=arguments): + case _SandboxUnavailable(output=output): + return _ToolRun(output=output, status="failed") + case str() as output: + return _ToolRun(output=output, status="completed") + + async def _run_tool_call( + self, container: ContainerHandle, params: SandboxToolParams | None, arguments: str + ) -> str | _SandboxUnavailable: try: code: Final = json.loads(arguments).get("code", "") if arguments else "" except (json.JSONDecodeError, TypeError): return "[invalid tool arguments: could not parse code]" - result: Final = await self._run_code(container=container, params=params, code=code) + try: + result: Final = await self._run_code(container=container, params=params, code=code) + except Exception as e: + return _sandbox_unavailable("running the code", e, _SANDBOX_LOST_RESULT_OUTPUT) if getattr(result, "error", None): error: Final = result.error message: Final = error.get("value") or error.get("name") if isinstance(error, dict) else str(error) return f"[execution error] {message}" return getattr(result, "stdout", "") or "" + async def _acquire_sandbox(self, cache_key: str | None, identity: str | None) -> _Sandbox: + try: + container, params = await self._get_or_create_container(cache_key=cache_key, identity=identity) + except Exception as e: + return _sandbox_unavailable("creating the sandbox", e) + return _ReadySandbox(container=container, params=params) + async def _get_or_create_container( self, cache_key: str | None, diff --git a/litellm/litellm_core_utils/agentic_followup_failure.py b/litellm/litellm_core_utils/agentic_followup_failure.py new file mode 100644 index 00000000000..318429d64b5 --- /dev/null +++ b/litellm/litellm_core_utils/agentic_followup_failure.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass +from typing import Final, TypeAlias + +from litellm.types.utils import Usage + +ServedRound: TypeAlias = tuple[Usage, float] + +_MARK_ATTRIBUTE: Final = "_litellm_agentic_followup_mark" + + +@dataclass(frozen=True, slots=True) +class _FollowUpMark: + served_rounds: tuple[ServedRound, ...] + + +def _mark_of(error: BaseException) -> _FollowUpMark | None: + mark: Final[object] = getattr(error, _MARK_ATTRIBUTE, None) + return mark if isinstance(mark, _FollowUpMark) else None + + +def mark_agentic_followup_failure(error: BaseException, served_rounds: tuple[ServedRound, ...]) -> None: + setattr(error, _MARK_ATTRIBUTE, _FollowUpMark(served_rounds)) + + +def is_agentic_followup_failure(error: BaseException) -> bool: + return _mark_of(error) is not None + + +def rounds_served_before(error: BaseException) -> tuple[ServedRound, ...]: + mark: Final = _mark_of(error) + return () if mark is None else mark.served_rounds + + +def carry_agentic_followup_mark(source: BaseException, mapped: BaseException) -> None: + if mapped is not source and is_agentic_followup_failure(source): + mark_agentic_followup_failure(mapped, rounds_served_before(source)) diff --git a/litellm/litellm_core_utils/agentic_loop_refusal.py b/litellm/litellm_core_utils/agentic_loop_refusal.py new file mode 100644 index 00000000000..f3ee15dffdd --- /dev/null +++ b/litellm/litellm_core_utils/agentic_loop_refusal.py @@ -0,0 +1,107 @@ +import asyncio +from collections.abc import Awaitable, Iterable, Mapping +from itertools import chain +from types import MappingProxyType +from typing import Final + +from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ( + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, + Choices, + ModelResponse, +) + + +def _strings(values: Iterable[object]) -> frozenset[str]: + return frozenset(value for value in values if isinstance(value, str)) + + +def _claimed_calls(tool_calls: object) -> tuple[object, ...]: + wrapped: Final = message_field(tool_calls, "tool_calls") + return parts_of(tool_calls if wrapped is None else wrapped) + + +async def _calls_a_gate_claims(gate: Awaitable[tuple[bool, object]]) -> tuple[object, ...]: + try: + should_run, tool_calls = await gate + except Exception as e: + verbose_logger.debug("LiteLLM.AgenticLoopRefused: skipping a gate that failed with %s", type(e).__name__) + return () + return _claimed_calls(tool_calls) if should_run else () + + +async def every_claimed_tool_call( + refused_tool_calls: object, other_gates: Iterable[Awaitable[tuple[bool, object]]] +) -> Mapping[str, tuple[object, ...]]: + claims: Final = await asyncio.gather(*(_calls_a_gate_claims(gate) for gate in other_gates)) + return MappingProxyType({"tool_calls": (*_claimed_calls(refused_tool_calls), *chain.from_iterable(claims))}) + + +def _call_ids(call: object) -> frozenset[str]: + return _strings((message_field(call, "id"), message_field(call, "call_id"))) + + +def refused_tool_call_identifiers(tool_calls: object) -> tuple[frozenset[str], frozenset[str]]: + calls: Final = _claimed_calls(tool_calls) + ids: Final = frozenset(chain.from_iterable(_call_ids(call) for call in calls)) + names: Final = _strings(message_field(call, "name") for call in calls if not _call_ids(call)) + return ids, names + + +def is_refused_call(call_id: object, name: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool: + return (isinstance(call_id, str) and call_id in refused_ids) or name in refused_names + + +def is_refused_item( + item: object, item_type: str, id_field: str, refused_ids: frozenset[str], refused_names: frozenset[str] +) -> bool: + if message_field(item, "type") != item_type: + return False + return is_refused_call(message_field(item, id_field), message_field(item, "name"), refused_ids, refused_names) + + +def finalize_refused_responses_turn(response: ResponsesAPIResponse, tool_calls: object) -> ResponsesAPIResponse: + refused_ids, refused_names = refused_tool_call_identifiers(tool_calls) + kept: Final = [ + item + for item in parts_of(message_field(response, "output")) + if not is_refused_item(item, "function_call", "call_id", refused_ids, refused_names) + ] + return response.model_copy(update={"output": kept}) + + +def _is_refused_chat_tool_call(call: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool: + name: Final = message_field(message_field(call, "function"), "name") + return is_refused_call(message_field(call, "id"), name, refused_ids, refused_names) + + +def _at_position( + call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall, position: int +) -> ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall: + return call if message_field(call, "index") is None else call.model_copy(update={"index": position}) + + +def _finalized_choice(choice: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> object: + if not isinstance(choice, Choices) or not choice.message.tool_calls: + return choice + surviving: Final = ( + call for call in choice.message.tool_calls if not _is_refused_chat_tool_call(call, refused_ids, refused_names) + ) + kept: Final = [_at_position(call, position) for position, call in enumerate(surviving)] + finish_reason: Final = "stop" if not kept and choice.finish_reason == "tool_calls" else choice.finish_reason + return choice.model_copy( + update={ + "message": choice.message.model_copy(update={"tool_calls": kept or None}), + "finish_reason": finish_reason, + } + ) + + +def finalize_refused_chat_turn(response: ModelResponse, tool_calls: object) -> ModelResponse: + refused_ids, refused_names = refused_tool_call_identifiers(tool_calls) + return response.model_copy( + update={"choices": [_finalized_choice(choice, refused_ids, refused_names) for choice in response.choices]} + ) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index 1f29b5d7a4d..ce98d2cbb01 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -1,29 +1,35 @@ # 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 Mapping +from collections.abc import Awaitable, Iterator, Mapping, Sequence from itertools import chain from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar 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_failure import mark_agentic_followup_failure, rounds_served_before from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs +from litellm.litellm_core_utils.agentic_loop_refusal import every_claimed_tool_call, finalize_refused_chat_turn from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +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.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopLimitError, AgenticLoopPlan, AgenticLoopRequestPatch, + AgenticLoopSafetyError, converted_stream_requested, is_interception_internal_key, stashed_stream_options, @@ -31,6 +37,8 @@ from litellm.types.integrations.custom_logger import ( from litellm.types.utils import ModelResponse from litellm.utils import CustomStreamWrapper +_DispatchResult: Final = TypeVar("_DispatchResult") + _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset( ( "acompletion", @@ -44,22 +52,41 @@ _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset( ) +async def bill_first_round_on_failure( + dispatch: Awaitable[_DispatchResult], *, logging_obj: object, response: object, stream: bool +) -> _DispatchResult: + try: + return await dispatch + except Exception as error: + 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) + 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) + raise + + def _gate_overridden(callback: CustomLogger) -> bool: - base: Final = CustomLogger.async_should_run_agentic_loop - func: Final = type(callback).async_should_run_agentic_loop - return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + return type(callback).async_should_run_agentic_loop is not CustomLogger.async_should_run_agentic_loop def _build_plan_overridden(callback: CustomLogger) -> bool: - base: Final = CustomLogger.async_build_agentic_loop_plan - func: Final = type(callback).async_build_agentic_loop_plan - return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + return type(callback).async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan + + +def _chat_build_plan_overridden(callback: CustomLogger) -> bool: + return ( + type(callback).async_build_chat_completion_agentic_loop_plan + is not CustomLogger.async_build_chat_completion_agentic_loop_plan + ) def _post_hook_overridden(callback: CustomLogger) -> bool: - base: Final = CustomLogger.async_post_agentic_loop_response_hook - func: Final = type(callback).async_post_agentic_loop_response_hook - return getattr(func, "__func__", func) is not getattr(base, "__func__", base) + return ( + type(callback).async_post_agentic_loop_response_hook is not CustomLogger.async_post_agentic_loop_response_hook + ) _STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -76,14 +103,14 @@ def _coerce_int(value: object, default: int) -> int: return int(value) if isinstance(value, (int, str)) else default -def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]: +def _agentic_loop_settings(kwargs: Mapping[str, object]) -> tuple[int, int, tuple[str, ...]]: depth: Final = _coerce_int(kwargs.get("_agentic_loop_depth"), 0) configured: Final = validated_max_agentic_loops( kwargs.get("max_agentic_loops"), field="litellm_params.max_agentic_loops" ) max_loops: Final = DEFAULT_MAX_AGENTIC_LOOPS if configured is None else configured raw_fingerprints: Final = kwargs.get("_agentic_loop_fingerprints") - fingerprints: Final = [str(fp) for fp in raw_fingerprints] if isinstance(raw_fingerprints, list) else [] + fingerprints: Final = tuple(str(fp) for fp in raw_fingerprints) if isinstance(raw_fingerprints, list) else () return depth, max_loops, fingerprints @@ -96,16 +123,16 @@ def _fingerprint_tools(tool_calls: object) -> str: def _check_agentic_loop_safety( tool_calls: object, - fingerprints: list[str], + fingerprints: Sequence[str], depth: int, max_loops: int, model: str, ) -> str: fingerprint: Final = _fingerprint_tools(tool_calls) if fingerprint in fingerprints: - raise ValueError("Agentic loop detected repeated tool-call fingerprint; aborting rerun") + raise AgenticLoopSafetyError("Agentic loop detected repeated tool-call fingerprint; aborting rerun") if depth >= max_loops: - raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}") + raise AgenticLoopLimitError(f"Exceeded max_agentic_loops={max_loops} for model={model}") return fingerprint @@ -134,6 +161,83 @@ def _converted_stream_replay( ) +def _end_refused_turn( + finalized: ModelResponse, + *, + kwargs: Mapping[str, object], + depth: int, + model: str, + custom_llm_provider: str, + logging_obj: object, +) -> ModelResponse | CustomStreamWrapper: + replay: Final = _converted_stream_replay( + finalized, + kwargs=kwargs, + depth=depth, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + return finalized if replay is None else replay + + +def _request_tools(optional_params: Mapping[str, object]) -> list[dict[str, object]]: + return [tool for tool in object_items(optional_params.get("tools")) if is_str_keyed_dict(tool)] + + +def _custom_loggers(logging_obj: object) -> tuple[CustomLogger, ...]: + import litellm + + dynamic: Final = getattr(logging_obj, "dynamic_success_callbacks", None) or [] + return tuple(callback for callback in (*litellm.callbacks, *dynamic) if isinstance(callback, CustomLogger)) + + +async def _refused_chat_turn( + callbacks: Sequence[CustomLogger], + refusing: CustomLogger, + refused_tool_calls: object, + *, + response: ModelResponse, + model: str, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + custom_llm_provider: str, + kwargs: Mapping[str, object], +) -> ModelResponse: + surface_kwargs: Final = { + **kwargs, + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + "custom_llm_provider": custom_llm_provider, + } + + def other_gates() -> Iterator[Awaitable[tuple[bool, object]]]: + for callback in callbacks: + if callback is refusing: + continue + yield callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(kwargs), + ) + if _gate_overridden(callback): + yield callback.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=False, + custom_llm_provider=custom_llm_provider, + kwargs=dict(surface_kwargs), + ) + + claimed: Final = await every_claimed_tool_call(refused_tool_calls, other_gates()) + return finalize_refused_chat_turn(response, claimed) + + def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]: metadata: Final = kwargs_for_followup.get("litellm_metadata") return MappingProxyType( @@ -164,6 +268,12 @@ def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]: } +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): + hidden_params.pop("response_cost", None) + + async def _execute_chat_completion_agentic_plan( *, plan: AgenticLoopPlan, @@ -175,7 +285,7 @@ async def _execute_chat_completion_agentic_plan( custom_llm_provider: str, depth: int, max_loops: int, - fingerprints: list[str], + fingerprints: Sequence[str], fingerprint: str, ) -> object: import litellm @@ -226,6 +336,7 @@ async def _execute_chat_completion_agentic_plan( model, str(e), ) + _leave_pricing_to_the_request(response_followup) replay: Final = _converted_stream_replay( response_followup, kwargs=kwargs, @@ -252,23 +363,18 @@ async def maybe_run_chat_completion_agentic_loop( *, response: ModelResponse, model: str, - messages: list, - optional_params: dict, + messages: list[dict[str, object]], + optional_params: dict[str, object], kwargs: dict[str, object], logging_obj: object, custom_llm_provider: str, stream: bool, ) -> ModelResponse | CustomStreamWrapper | None: - import litellm - - callbacks: Final = litellm.callbacks + (getattr(logging_obj, "dynamic_success_callbacks", None) or []) + callbacks: Final = _custom_loggers(logging_obj) depth, max_loops, fingerprints = _agentic_loop_settings(kwargs) - tools: Final = optional_params.get("tools", []) + tools: Final = _request_tools(optional_params) for callback in callbacks: - if not isinstance(callback, CustomLogger): - continue - if not _gate_overridden(callback): continue @@ -297,29 +403,39 @@ async def maybe_run_chat_completion_agentic_loop( if not should_run: continue - fingerprint = _check_agentic_loop_safety( - tool_calls=tool_calls, - fingerprints=fingerprints, - depth=depth, - max_loops=max_loops, - model=model, - ) - try: - if not _build_plan_overridden(callback): - return await callback.async_run_agentic_loop( - tools=tool_calls, + fingerprint = _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + except AgenticLoopSafetyError as e: + if stream: + raise + verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) + return _end_refused_turn( + await _refused_chat_turn( + callbacks, + callback, + tool_calls, + response=response, model=model, messages=messages, - response=response, - anthropic_messages_provider_config=None, - anthropic_messages_optional_request_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=hook_kwargs, - ) + tools=tools, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ), + kwargs=kwargs, + depth=depth, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) - plan = await callback.async_build_agentic_loop_plan( + if not _build_plan_overridden(callback): + return await callback.async_run_agentic_loop( tools=tool_calls, model=model, messages=messages, @@ -331,31 +447,38 @@ async def maybe_run_chat_completion_agentic_loop( kwargs=hook_kwargs, ) - if plan.response_override is not None: - return plan.response_override - if plan.terminate: - return response - if not plan.run_agentic_loop: - continue + plan = await callback.async_build_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=hook_kwargs, + ) - return await _execute_chat_completion_agentic_plan( - plan=plan, - callback=callback, - model=model, - optional_params=optional_params, - kwargs=kwargs, - logging_obj=logging_obj, - custom_llm_provider=custom_llm_provider, - depth=depth, - max_loops=max_loops, - fingerprints=fingerprints, - fingerprint=fingerprint, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s", - str(e), - ) + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + return response + if not plan.run_agentic_loop: + continue + + return await _execute_chat_completion_agentic_plan( + plan=plan, + callback=callback, + model=model, + optional_params=optional_params, + kwargs=kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) return _converted_stream_replay( response, @@ -365,3 +488,113 @@ async def maybe_run_chat_completion_agentic_loop( custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, ) + + +async def run_chat_completion_agentic_hooks( + *, + response: ModelResponse, + model: str, + messages: list[dict[str, object]], + optional_params: dict[str, object], + kwargs: dict[str, object], + logging_obj: LiteLLMLoggingObject, + custom_llm_provider: str, + stream: bool, +) -> 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: + try: + should_run, tool_calls = await callback.async_should_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + except Exception: + verbose_logger.exception( + "LiteLLM.AgenticHookError: Exception in async_should_run_chat_completion_agentic_loop" + ) + continue + + if not should_run: + continue + + try: + fingerprint = _check_agentic_loop_safety( + tool_calls=tool_calls, + fingerprints=fingerprints, + depth=depth, + max_loops=max_loops, + model=model, + ) + except AgenticLoopSafetyError as e: + if stream: + raise + verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) + return await _refused_chat_turn( + callbacks, + callback, + tool_calls, + response=response, + model=model, + messages=messages, + tools=tools, + custom_llm_provider=custom_llm_provider, + kwargs=kwargs, + ) + + if not _chat_build_plan_overridden(callback): + return await callback.async_run_chat_completion_agentic_loop( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=hook_kwargs, + ) + + plan = await callback.async_build_chat_completion_agentic_loop_plan( + tools=tool_calls, + model=model, + messages=messages, + response=response, + optional_params=optional_params, + logging_obj=logging_obj, + stream=stream, + kwargs=hook_kwargs, + ) + + if plan.response_override is not None: + return plan.response_override + if plan.terminate: + verbose_logger.debug( + "Agentic chat loop terminated by callback=%s reason=%s", type(callback).__name__, plan.stop_reason + ) + return response + if not plan.run_agentic_loop: + continue + + return await _execute_chat_completion_agentic_plan( + plan=plan, + callback=callback, + model=model, + optional_params=optional_params, + kwargs=hook_kwargs, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + depth=depth, + max_loops=max_loops, + fingerprints=fingerprints, + fingerprint=fingerprint, + ) + + return None diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 422a2ef1219..6d39f272f3e 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -11,6 +11,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import _ENABLE_SECRET_REDACTION, redact_string, verbose_logger +from litellm.litellm_core_utils.agentic_followup_failure import carry_agentic_followup_mark from litellm.litellm_core_utils.bug_report import ( bug_report_notice, build_bug_report, @@ -2902,6 +2903,7 @@ def exception_type( request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), # stub the request ) except Exception as e: + carry_agentic_followup_mark(source=original_exception, mapped=e) # LOGGING exception_logging( logger_fn=None, @@ -2927,6 +2929,7 @@ def exception_type( model="", ) setattr(raised_exc, "litellm_response_headers", litellm_response_headers) + carry_agentic_followup_mark(source=original_exception, mapped=raised_exc) raise raised_exc diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index d6d9be4e132..d0ac5c00513 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -33,6 +33,9 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( ) PROVIDER_AFFINITY_HEADER_KWARG_KEY: Final = "provider_affinity_header" +AGENTIC_LOOP_CONTROL_KWARGS_KEYS: Final = frozenset( + {"_agentic_loop_depth", "_agentic_loop_fingerprints", "max_agentic_loops"} +) OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS: Final = frozenset( { "token_exchange_endpoint", @@ -94,6 +97,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( | ANTHROPIC_WIF_KWARGS_KEYS | OPENAI_WIF_KWARGS_KEYS | OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS + | AGENTIC_LOOP_CONTROL_KWARGS_KEYS | frozenset(CustomPricingLiteLLMParams.model_fields) ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6e79f8e56ef..0fe520103e2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -54,6 +54,7 @@ from litellm.constants import ( REDACTED_BY_LITELLM, ) from litellm.cost_calculator import ( + BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, declared_cost_map_pricing, @@ -77,6 +78,7 @@ from litellm.integrations.langtrace import langtrace_trace_endpoint from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.agentic_followup_failure import ServedRound from litellm.litellm_core_utils.classifier_logging import ( classifier_audit_fields, classifier_input_snapshot, @@ -944,6 +946,11 @@ class Logging(LiteLLMLoggingBaseClass): self._sync_success_scheduled: bool = False self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None self._stream_failure_claimed: bool = False + self._stash_owed_on_success: bool = False + self._stash_billed_by_continuation: bool = False + self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = () + self._failed_attempt_ids_owed_on_success: tuple[str, ...] = () + self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = () self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging": @@ -2320,12 +2327,19 @@ class Logging(LiteLLMLoggingBaseClass): return 0.0 def uncached_price(self, result: object) -> UncachedPrice: + return self._price_without_recording(result, cache_hit=False, ignore_stamped_cost=True) + + def _price_without_recording( + self, result: object, *, cache_hit: bool | None, ignore_stamped_cost: bool + ) -> UncachedPrice: previous_cost_breakdown: Final = self.cost_breakdown previous_billed_token_rates: Final = self.billed_token_rates debug_key: Final = "response_cost_failure_debug_information" previous_failure_debug_info: Final = self.model_call_details.get(debug_key) try: - cost: Final = self.response_cost_calculator(result=result, cache_hit=False, ignore_stamped_cost=True) + cost: Final = self.response_cost_calculator( + result=result, cache_hit=cache_hit, ignore_stamped_cost=ignore_stamped_cost + ) return UncachedPrice(cost=cost, cost_breakdown=self.cost_breakdown) finally: self.cost_breakdown = previous_cost_breakdown @@ -2777,27 +2791,43 @@ class Logging(LiteLLMLoggingBaseClass): own usage while the request's failure log still bills every attempt.""" usage: Final = self._model_call_details_view.get("combined_usage_object") response_cost: Final = self._model_call_details_view.get("response_cost") - if isinstance(usage, Usage) and type(self).usage_capture_callback is not None: + billed_by_continuation: Final = self._stash_billed_by_continuation + owed_on_success: Final = self._stash_owed_on_success + self._stash_billed_by_continuation = False + self._stash_owed_on_success = False + captured_attempt_ids: Final = ( self.capture_failed_attempt_usage(usage) + if isinstance(usage, Usage) and type(self).usage_capture_callback is not None + else () + ) self.model_call_details.pop("_accounting_capture", None) self.model_call_details.pop("_accounting_pricing_snapshot", None) self.model_call_details.pop("_accounting_usage_observed", None) self.model_call_details.pop("combined_usage_object", None) self.model_call_details.pop("response_cost", None) - if isinstance(usage, Usage): - self.model_call_details["current_attempt_api_call_start_time"] = None - self.model_call_details["failed_attempts_usage"] = ( - *self.model_call_details.get("failed_attempts_usage", ()), - (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), - ) + if not isinstance(usage, Usage): + return + attempt: Final = (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0) + self.model_call_details["current_attempt_api_call_start_time"] = None + self.model_call_details["failed_attempts_usage"] = ( + *self.model_call_details.get("failed_attempts_usage", ()), + attempt, + ) + if not owed_on_success or billed_by_continuation: + return + self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt) + self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids) - def capture_failed_attempt_usage(self, usage: Usage) -> None: + def mark_stash_billed_by_continuation(self) -> None: + self._stash_billed_by_continuation = True + + def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]: from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.proxy.spend_tracking.usage_capture import CapturedUsage capture: Final = type(self).usage_capture_callback if capture is None: - return + return () attempt: Final = self.attempt_scoped_copy() attempt.prepared_cost = None attempt.finalized_cost = None @@ -2832,11 +2862,13 @@ class Logging(LiteLLMLoggingBaseClass): datetime.datetime.now(datetime.timezone.utc), False, ) - if isinstance(captured, CapturedUsage): - self.model_call_details["_accounting_prior_captures"] = ( - *self.model_call_details.get("_accounting_prior_captures", ()), - captured, - ) + if not isinstance(captured, CapturedUsage): + return () + self.model_call_details["_accounting_prior_captures"] = ( + *self.model_call_details.get("_accounting_prior_captures", ()), + captured, + ) + return tuple(write.event.event_id for write in captured.writes) def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" @@ -2844,20 +2876,17 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["continued_partial_cost"] = response_cost self.model_call_details["continued_response_id"] = response_id - def _add_continued_partial_spend(self, response: object) -> None: + def _add_continued_partial_spend(self, response: object) -> Usage | None: partial_usage: Final = self.model_call_details.get("continued_partial_usage") partial_cost: Final = self.model_call_details.get("continued_partial_cost") own_usage: Final = getattr(response, "usage", None) own_cost: Final = self.model_call_details.get("response_cost") if not (isinstance(partial_usage, Usage) and isinstance(partial_cost, float) and isinstance(own_usage, Usage)): - return - from litellm.cost_calculator import BaseTokenUsageProcessor - - self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=[own_usage, partial_usage] - ) + return own_usage if isinstance(own_usage, Usage) else None + combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=[own_usage, partial_usage]) + self.model_call_details["combined_usage_object"] = combined_usage if not isinstance(own_cost, (int, float)): - return + return combined_usage self.model_call_details["response_cost"] = own_cost + partial_cost if self.cost_breakdown is not None: self.cost_breakdown = self.cost_breakdown | CostBreakdown( @@ -2867,6 +2896,27 @@ class Logging(LiteLLMLoggingBaseClass): "mid_stream_partial_cost": partial_cost, }, ) + return combined_usage + + def _add_failed_attempts_spend(self, own_usage: Usage | None) -> None: + own_cost: Final = self.model_call_details.get("response_cost") + if own_usage is None or not isinstance(own_cost, (int, float)): + return + failed_attempts: Final = self._failed_attempts_owed_on_success + if not failed_attempts: + return + failed_attempts_cost: Final = sum(cost for _, cost in failed_attempts) + self.model_call_details.pop("failed_attempts_usage", None) + self.failed_attempts_billed_on_success = failed_attempts + self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts + if self._failed_attempt_ids_owed_on_success: + self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success + self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)] + ) + self.model_call_details["response_cost"] = own_cost + failed_attempts_cost + if self.cost_breakdown is not None: + self.cost_breakdown = _with_failed_attempts_cost(self.cost_breakdown, own_cost, failed_attempts_cost) def claim_stream_failure_report(self) -> bool: if self._stream_failure_claimed: @@ -2874,6 +2924,24 @@ class Logging(LiteLLMLoggingBaseClass): self._stream_failure_claimed = True return True + def served_round(self, response: object) -> ServedRound | None: + usage: Final = get_usage_object(completion_response=response) + if usage is None: + return None + price: Final = self._price_without_recording(response, cache_hit=None, ignore_stamped_cost=False) + return usage, price.cost or 0.0 + + def record_served_rounds_for_failure(self, served_rounds: tuple[ServedRound, ...]) -> None: + if not served_rounds: + return + self._stash_owed_on_success = True + self.record_partial_usage_for_failure( + served_rounds[0][0] + if len(served_rounds) == 1 + else BaseTokenUsageProcessor.combine_usage_objects(usage_objects=[usage for usage, _ in served_rounds]), + sum(cost for _, cost in served_rounds), + ) + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) @@ -3178,9 +3246,11 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.finalized_cost.amount elif self.model_call_details.get("cache_hit") is True: self.model_call_details["response_cost"] = 0.0 + self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result)) elif "response_cost" in hidden_params: self.model_call_details["response_cost"] = hidden_params["response_cost"] self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"]) + self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result)) elif (existing_cost := self.model_call_details.get("response_cost")) is not None and ( existing_cost != 0 or self.call_type == CallTypes.pass_through.value # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped ): @@ -3190,6 +3260,7 @@ class Logging(LiteLLMLoggingBaseClass): pass else: self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result) + self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result)) if not build_logging_payload: return None @@ -3753,7 +3824,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.response_cost_calculator( result=complete_streaming_response if provider_transcription is None else provider_transcription ) - self._add_continued_partial_spend(complete_streaming_response) + self._add_failed_attempts_spend(self._add_continued_partial_spend(complete_streaming_response)) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) logged_streaming_response: Final = self._client_facing_streaming_response( self._as_served_chat_response(complete_streaming_response) @@ -4349,7 +4420,7 @@ class Logging(LiteLLMLoggingBaseClass): self.model, ) self.model_call_details["response_cost"] = None - self._add_continued_partial_spend(complete_streaming_response) + self._add_failed_attempts_spend(self._add_continued_partial_spend(complete_streaming_response)) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) @@ -7842,6 +7913,20 @@ def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown: ) +def _with_failed_attempts_cost(breakdown: CostBreakdown, own_cost: float, failed_attempts_cost: float) -> CostBreakdown: + return breakdown | CostBreakdown( + total_cost=own_cost + failed_attempts_cost, + additional_costs={**breakdown.get("additional_costs", {}), "failed_attempts_cost": failed_attempts_cost}, + ) + + +def _cached_cost_breakdown(served: CostBreakdown | None, failed_attempts: tuple[ServedRound, ...]) -> CostBreakdown: + cached: Final = _cache_hit_cost_breakdown(served) + if not failed_attempts: + return cached + return _with_failed_attempts_cost(cached, 0.0, sum(cost for _, cost in failed_attempts)) + + def _finalized_cost_of(logging_obj: object) -> FinalizedCost | None: """Duck-typed logging objects (and test doubles) carry no finalized cost.""" finalized: Final = getattr(logging_obj, "finalized_cost", None) @@ -8035,7 +8120,7 @@ def get_standard_logging_object_payload( if finalized_cost is not None else cost_breakdown_with_guardrail( ( - _cache_hit_cost_breakdown(uncached_price.cost_breakdown) + _cached_cost_breakdown(uncached_price.cost_breakdown, logging_obj.failed_attempts_billed_on_success) if uncached_price is not None else logging_obj.cost_breakdown ), diff --git a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py index 25d4d7b3ce1..0aa30998557 100644 --- a/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py @@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +from litellm.types.integrations.custom_logger import AgenticLoopLimitError, AgenticLoopSafetyError if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -35,10 +36,16 @@ SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b'data: {"type": "error", "error": {"type": "api_error", "message": ' b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n' ) +AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES: Final = ( + b"event: error\n" + b'data: {"type": "error", "error": {"type": "api_error", "message": ' + b'"The agentic loop reached its limit before this turn finished. ' + b'Raise max_agentic_loops to allow more loops."}}\n\n' +) def is_server_fulfilled_tool_leak_error(chunk: object) -> bool: - return chunk == SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES + return chunk in (SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES) async def _anext_or_none(iterator: AsyncIterator) -> bytes | None: @@ -216,6 +223,7 @@ class AgenticAnthropicStreamingIterator: self._follow_up_chunk_task: asyncio.Task | None = None self._replay_index = 0 self._error_emitted = False + self._loop_limit_reached = False @property def has_buffered_provider_output(self) -> bool: @@ -288,6 +296,8 @@ class AgenticAnthropicStreamingIterator: "tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client", self._model, ) + if self._loop_limit_reached: + return AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES if self._replay_index < len(self._collected_bytes): @@ -393,14 +403,9 @@ class AgenticAnthropicStreamingIterator: "AgenticStreamingIterator: Unexpected result type from hooks: %s", type(result).__name__, ) - except Exception as e: - _call_id: Final = getattr(self._logging_obj, "litellm_call_id", "unknown") - verbose_logger.exception( - "AgenticStreamingIterator: Error in agentic hook processing [call_id=%s model=%s]: %s", - _call_id, - self._model, - str(e), - ) + except AgenticLoopSafetyError as e: + self._loop_limit_reached = isinstance(e, AgenticLoopLimitError) + verbose_logger.warning("AgenticStreamingIterator: agentic loop refused [rail=%s]", type(e).__name__) @staticmethod def _rebuild_anthropic_response_from_sse( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f17dc6dd10..4a10db31606 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -47,6 +47,12 @@ from litellm.anthropic_beta_headers_manager import ( from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs +from litellm.litellm_core_utils.agentic_loop_refusal import ( + every_claimed_tool_call, + finalize_refused_responses_turn, + is_refused_item, + refused_tool_call_identifiers, +) from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, @@ -56,11 +62,11 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SUBTITLE_RESPONSE_FORMATS, synthesize_subtitle_document, ) -from litellm.litellm_core_utils.core_helpers import ( - as_str_mapping, - qualify_agentic_followup_model, - set_provider_response_headers_in_hidden_params, +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + bill_first_round_on_failure, + run_chat_completion_agentic_hooks, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header @@ -150,6 +156,7 @@ from litellm.types.extraction import ExtractionResponse from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, + AgenticLoopLimitError, AgenticLoopPlan, AgenticLoopRequestPatch, AgenticLoopSafetyError, @@ -6005,7 +6012,7 @@ class BaseLLMHTTPHandler: @staticmethod def _check_agentic_loop_safety( tool_calls: object, - fingerprints: list[str], + fingerprints: Sequence[str], depth: int, max_loops: int, model: str, @@ -6024,7 +6031,7 @@ class BaseLLMHTTPHandler: if fingerprint in fingerprints: raise AgenticLoopSafetyError("Agentic loop detected repeated tool-call fingerprint; aborting rerun") if depth >= max_loops: - raise AgenticLoopSafetyError(f"Exceeded max_agentic_loops={max_loops} for model={model}") + raise AgenticLoopLimitError(f"Exceeded max_agentic_loops={max_loops} for model={model}") return fingerprint @staticmethod @@ -6034,26 +6041,6 @@ class BaseLLMHTTPHandler: except Exception: return str(tools) - @staticmethod - def _refused_agentic_tool_identifiers(tool_calls: object) -> tuple[frozenset[str], frozenset[str]]: - """ - Collect the ids and names of the tool calls a safety rail just refused. - - Callbacks hand back either a bare list of tool calls or a dict wrapping - that list under ``tool_calls``, and both the anthropic and responses - shapes carry an ``id`` (or ``call_id``) plus a ``name``. - """ - calls: Final = tool_calls.get("tool_calls") if isinstance(tool_calls, dict) else tool_calls - if not isinstance(calls, list): - return frozenset(), frozenset() - dict_calls: Final = (call for call in calls if isinstance(call, dict)) - fields: Final = tuple((call.get("id"), call.get("call_id"), call.get("name")) for call in dict_calls) - ids: Final = frozenset( - value for call_id, caller_id, _ in fields for value in (call_id, caller_id) if isinstance(value, str) - ) - names: Final = frozenset(name for _, _, name in fields if isinstance(name, str)) - return ids, names - @staticmethod def _is_refused_tool_use_block(block: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool: """ @@ -6064,22 +6051,13 @@ class BaseLLMHTTPHandler: share a name with a refused one. The name is only consulted for tool call shapes that arrive without an id. """ - if not isinstance(block, dict) or block.get("type") != "tool_use": - return False - block_id: Final = block.get("id") - if isinstance(block_id, str) and refused_ids: - return block_id in refused_ids - return block.get("name") in refused_names + return is_refused_item(block, "tool_use", "id", refused_ids, refused_names) @staticmethod def _can_replace_turn_with_terminal_response(stream: bool, api_surface: str) -> bool: """ Whether a refused rerun can still be answered with a finalized turn. - Only the anthropic messages surface can. The responses surface carries a - pydantic model the finalizer does not rewrite, so it keeps raising, which - is what every surface did before this path learned to end the turn. - The messages and responses call sites pass ``stream=False``, because interception converts an intercepted stream to non-streaming before the loop runs and rebuilds the SSE stream from the finalized turn @@ -6088,7 +6066,7 @@ class BaseLLMHTTPHandler: finalized turn would reach the client as a second message rather than as a replacement. """ - return not stream and api_surface == "anthropic_messages" + return not stream and api_surface in ("anthropic_messages", "responses") @staticmethod def _finalize_refused_agentic_response(response: object, tool_calls: object) -> object: @@ -6105,10 +6083,12 @@ class BaseLLMHTTPHandler: one is still in the response the stop reason stays ``tool_use`` so the client knows to answer it. """ + if isinstance(response, ResponsesAPIResponse): + return finalize_refused_responses_turn(response, tool_calls) if not isinstance(response, dict): return response - refused_ids, refused_names = BaseLLMHTTPHandler._refused_agentic_tool_identifiers(tool_calls) + refused_ids, refused_names = refused_tool_call_identifiers(tool_calls) finalized: Final = dict(response) content: Final = finalized.get("content") if isinstance(content, list): @@ -6135,7 +6115,7 @@ class BaseLLMHTTPHandler: kwargs: dict, depth: int, max_loops: int, - fingerprints: list[str], + fingerprints: Sequence[str], fingerprint: str, stream: bool = False, callback: Optional["CustomLogger"] = None, @@ -6176,7 +6156,7 @@ class BaseLLMHTTPHandler: kwargs_for_followup.update(patch.kwargs) kwargs_for_followup["_agentic_loop_depth"] = depth + 1 kwargs_for_followup["max_agentic_loops"] = max_loops - kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint] + kwargs_for_followup["_agentic_loop_fingerprints"] = [*fingerprints, fingerprint] response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( **{ @@ -6215,7 +6195,7 @@ class BaseLLMHTTPHandler: kwargs: dict, depth: int, max_loops: int, - fingerprints: list[str], + fingerprints: Sequence[str], fingerprint: str, callback: Optional["CustomLogger"] = None, ) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator: @@ -6336,67 +6316,6 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, ) - async def _execute_chat_completion_agentic_plan( - self, - plan: AgenticLoopPlan, - model: str, - messages: list[dict], - optional_params: dict, - kwargs: dict, - custom_llm_provider: str, - depth: int, - 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: - raise ValueError("Agentic loop plan missing patched messages") - - full_model_name: Final = qualify_agentic_followup_model(patch.model, model, custom_llm_provider) - - optional_params_for_followup: Final = dict(optional_params) - optional_params_for_followup.update(patch.optional_params) - if patch.tools is not None: - optional_params_for_followup["tools"] = patch.tools - - internal_params: Final = { - "_websearch_interception", - "acompletion", - "litellm_logging_obj", - "custom_llm_provider", - "model_alias_map", - "stream_response", - "custom_prompt_dict", - } - kwargs_for_followup: Final = build_agentic_followup_kwargs( - request_kwargs=MappingProxyType( - { - k: v - for k, v in kwargs.items() - if not k.startswith("_websearch_interception") - and not k.startswith("_compression_interception") - and k not in internal_params - } - ), - patch_kwargs=patch.kwargs, - request_params=frozenset((*optional_params_for_followup, "model", "messages")), - depth=depth, - max_loops=max_loops, - fingerprints=fingerprints, - fingerprint=fingerprint, - ) - - 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, response: _ResponseT, @@ -6513,26 +6432,38 @@ class BaseLLMHTTPHandler: except AgenticLoopSafetyError as e: if not self._can_replace_turn_with_terminal_response(stream, api_surface): raise - verbose_logger.warning( - "LiteLLM.AgenticLoopRefused: ending turn [call_id=%s model=%s]: %s", - logging_obj.litellm_call_id, - model, - str(e), - ) + verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__) return self._maybe_wrap_in_fake_stream( - self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), + self._finalize_refused_agentic_response( + response=response, + tool_calls=await every_claimed_tool_call( + tool_calls, + ( + other.async_should_run_agentic_loop( + response=response, + model=model, + messages=messages, + tools=tools, + stream=stream, + custom_llm_provider=custom_llm_provider, + kwargs=hook_kwargs, + ) + for other in callbacks + if isinstance(other, CustomLogger) and other is not callback + ), + ), + ), logging_obj, api_surface, ) - try: - kwargs_with_provider = hook_kwargs.copy() - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider - build_plan_overridden = ( - callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan - ) - if not build_plan_overridden: - agentic_result: object = await callback.async_run_agentic_loop( + kwargs_with_provider = {**hook_kwargs, "custom_llm_provider": custom_llm_provider} + build_plan_overridden = ( + callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan + ) + if not build_plan_overridden: + agentic_result: object = await bill_first_round_on_failure( + callback.async_run_agentic_loop( tools=tool_calls, model=model, messages=messages, @@ -6542,10 +6473,15 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=stream, kwargs=kwargs_with_provider, - ) - return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface) + ), + logging_obj=logging_obj, + response=response, + stream=stream, + ) + return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface) - plan = await callback.async_build_agentic_loop_plan( + plan = await bill_first_round_on_failure( + callback.async_build_agentic_loop_plan( tools=tool_calls, model=model, messages=messages, @@ -6555,30 +6491,35 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, stream=stream, kwargs=kwargs_with_provider, + ), + logging_obj=logging_obj, + response=response, + stream=stream, + ) + + if plan.response_override is not None: + return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface) + if plan.terminate: + verbose_logger.debug( + "Agentic loop terminated by callback=%s reason=%s", + callback.__class__.__name__, + plan.stop_reason, ) + return self._maybe_wrap_in_fake_stream( + await callback.async_post_agentic_loop_response_hook( + response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), + plan=plan, + kwargs=kwargs_with_provider, + ), + logging_obj, + api_surface, + ) + if not plan.run_agentic_loop: + continue - if plan.response_override is not None: - return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface) - if plan.terminate: - verbose_logger.debug( - "Agentic loop terminated by callback=%s reason=%s", - callback.__class__.__name__, - plan.stop_reason, - ) - return self._maybe_wrap_in_fake_stream( - await callback.async_post_agentic_loop_response_hook( - response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls), - plan=plan, - kwargs=kwargs_with_provider, - ), - logging_obj, - api_surface, - ) - if not plan.run_agentic_loop: - continue - - if api_surface == "responses": - return await self._execute_responses_agentic_plan( + if api_surface == "responses": + return await bill_first_round_on_failure( + self._execute_responses_agentic_plan( plan=plan, model=model, response_api_optional_request_params=anthropic_messages_optional_request_params, @@ -6589,10 +6530,15 @@ class BaseLLMHTTPHandler: fingerprints=fingerprints, fingerprint=fingerprint, callback=callback, - ) + ), + logging_obj=logging_obj, + response=response, + stream=stream, + ) - return self._maybe_wrap_in_fake_stream( - await self._execute_anthropic_agentic_plan( + return self._maybe_wrap_in_fake_stream( + await bill_first_round_on_failure( + self._execute_anthropic_agentic_plan( plan=plan, model=model, messages=messages, @@ -6606,17 +6552,13 @@ class BaseLLMHTTPHandler: stream=stream, callback=callback, ), - logging_obj, - api_surface, - ) - except Exception as e: - _call_id = getattr(logging_obj, "litellm_call_id", "unknown") - verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in agentic completion hooks [call_id=%s model=%s]: %s", - _call_id, - model, - str(e), - ) + logging_obj=logging_obj, + response=response, + stream=stream, + ), + logging_obj, + api_surface, + ) # Check if we need to convert response to fake stream # This happens when: @@ -6648,111 +6590,23 @@ class BaseLLMHTTPHandler: Returns the response from agentic loop, or None if no hook runs. """ - from litellm._logging import verbose_logger - from litellm.integrations.custom_logger import CustomLogger - - callbacks: Final = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) - tools: Final = optional_params.get("tools", []) - depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs) - - for callback in callbacks: - if not isinstance(callback, CustomLogger): - continue - if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): - continue - - should_run: bool = False - tool_calls: object = None - try: - ( - should_run, - tool_calls, - ) = await callback.async_should_run_chat_completion_agentic_loop( - response=response, - model=model, - messages=messages, - tools=tools, - stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=kwargs, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in async_should_run_chat_completion_agentic_loop: %s", - str(e), - ) - continue - - if not should_run: - continue - - # Safety guards must run OUTSIDE the callback try/except — they are - # bounded-loop / cycle-break rails that must propagate to the caller. - fingerprint = self._check_agentic_loop_safety( - tool_calls=tool_calls, - fingerprints=fingerprints, - depth=depth, - max_loops=max_loops, + agentic_response: Final = await bill_first_round_on_failure( + run_chat_completion_agentic_hooks( + response=response, model=model, - ) - - try: - kwargs_with_provider = kwargs.copy() if kwargs else {} - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider - build_plan_overridden = ( - callback.__class__.async_build_chat_completion_agentic_loop_plan - is not CustomLogger.async_build_chat_completion_agentic_loop_plan - ) - if not build_plan_overridden: - return await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - - plan = await callback.async_build_chat_completion_agentic_loop_plan( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - - if plan.response_override is not None: - return plan.response_override - if plan.terminate: - verbose_logger.debug( - "Agentic chat loop terminated by callback=%s reason=%s", - callback.__class__.__name__, - plan.stop_reason, - ) - return response - if not plan.run_agentic_loop: - continue - - return await self._execute_chat_completion_agentic_plan( - plan=plan, - model=model, - 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, - fingerprints=fingerprints, - fingerprint=fingerprint, - ) - except Exception as e: - verbose_logger.exception("LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s", e) + messages=messages, + optional_params=optional_params, + kwargs=kwargs or {}, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + stream=stream, + ), + logging_obj=logging_obj, + response=response, + stream=stream, + ) + if agentic_response is not None: + return agentic_response # Check if we need to convert response to fake stream for chat completions # This happens when: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e22ef9112e..87c944a642a 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -30,6 +30,10 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RETRIES from litellm.exceptions import UnsupportedParamsError from litellm.files.types import FileContentStreamingResult +from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + bill_first_round_on_failure, + run_chat_completion_agentic_hooks, +) from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing @@ -663,67 +667,23 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): Returns the response from agentic loop, or None if no hook runs. """ - from litellm._logging import verbose_logger - from litellm.integrations.custom_logger import CustomLogger - - callbacks: Final = litellm.callbacks + (logging_obj.dynamic_success_callbacks or []) - # Avoid logging full callback objects to prevent leaking sensitive data - verbose_logger.debug("LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks)) - tools: Final = optional_params.get("tools", []) - # Avoid logging full tools payloads; they may contain sensitive parameters - verbose_logger.debug( - "LiteLLM.AgenticHooks: tools_count=%s", - len(tools) if isinstance(tools, list) else 1 if tools else 0, + if not isinstance(response, ModelResponse): + return None + return await bill_first_round_on_failure( + run_chat_completion_agentic_hooks( + response=response, + model=model, + messages=messages, + optional_params=optional_params, + kwargs=litellm_params or {}, + logging_obj=logging_obj, + custom_llm_provider=litellm_params.get("custom_llm_provider", "openai"), + stream=stream, + ), + logging_obj=logging_obj, + response=response, + stream=stream, ) - # Get custom_llm_provider from litellm_params - custom_llm_provider: Final = litellm_params.get("custom_llm_provider", "openai") - - for callback in callbacks: - try: - if isinstance(callback, CustomLogger): - # Check if the callback has the chat completion agentic loop methods - if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"): - continue - - # First: Check if agentic loop should run (using chat completion method) - ( - should_run, - tool_calls, - ) = await callback.async_should_run_chat_completion_agentic_loop( - response=response, - model=model, - messages=messages, - tools=tools, - stream=stream, - custom_llm_provider=custom_llm_provider, - kwargs=litellm_params, - ) - - if should_run: - # Second: Execute agentic loop - kwargs_with_provider = litellm_params.copy() if litellm_params else {} - kwargs_with_provider["custom_llm_provider"] = custom_llm_provider - - # For OpenAI Chat Completions, use the chat completion agentic loop method - agentic_response: object = await callback.async_run_chat_completion_agentic_loop( - tools=tool_calls, - model=model, - messages=messages, - response=response, - optional_params=optional_params, - logging_obj=logging_obj, - stream=stream, - kwargs=kwargs_with_provider, - ) - # First hook that runs agentic loop wins - return agentic_response - - except Exception as e: - verbose_logger.exception( - "LiteLLM.AgenticHookError: Exception in agentic completion hooks for OpenAI: %s", e - ) - - return None def mock_streaming( self, @@ -1074,8 +1034,28 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): hidden_params={"headers": headers}, _response_headers=headers, ) + except openai.UnprocessableEntityError as e: + if litellm.drop_params is True or drop_params is True: + data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) + else: + raise e + except Exception as e: + exception_response: object = getattr(e, "response", None) + status_code = getattr(e, "status_code", 500) + exception_body = getattr(e, "body", None) + error_headers = getattr(e, "headers", None) + if error_headers is None and exception_response: + error_headers = getattr(exception_response, "headers", None) + message = getattr(e, "message", str(e)) - # Call agentic completion hooks (e.g., for websearch_interception) + raise OpenAIError( + status_code=status_code, + message=message, + response=upstream_error_response(e), + headers=error_headers, + body=exception_body, + ) + else: agentic_response = await self._call_agentic_completion_hooks_openai( response=final_response_obj, model=model, @@ -1098,29 +1078,6 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): ) return final_response_obj - except openai.UnprocessableEntityError as e: - ## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800 - if litellm.drop_params is True or drop_params is True: - data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True) - else: - raise e - # e.message - except Exception as e: - exception_response: object = getattr(e, "response", None) - status_code = getattr(e, "status_code", 500) - exception_body = getattr(e, "body", None) - error_headers = getattr(e, "headers", None) - if error_headers is None and exception_response: - error_headers = getattr(exception_response, "headers", None) - message = getattr(e, "message", str(e)) - - raise OpenAIError( - status_code=status_code, - message=message, - response=upstream_error_response(e), - headers=error_headers, - body=exception_body, - ) def streaming( self, diff --git a/litellm/main.py b/litellm/main.py index 86da08a23ab..794b3cef729 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -81,11 +81,13 @@ from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_for_health_check, ) from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + bill_first_round_on_failure, maybe_run_chat_completion_agentic_loop, ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.get_litellm_params import ( + AGENTIC_LOOP_CONTROL_KWARGS_KEYS, AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, PROVIDER_AFFINITY_HEADER_KWARG_KEY, @@ -954,26 +956,31 @@ async def acompletion( # no equivalent: every provider already funnels through one shared # handler where the loop is dispatched. if isinstance(response, litellm.ModelResponse): - looped: Final = await maybe_run_chat_completion_agentic_loop( - response=response, - model=model, - messages=messages, - optional_params={ - k: v - for k, v in completion_kwargs.items() - if v is not None - and k - not in ( - "model", - "messages", - "stream", - "acompletion", - "deployment_id", - ) - }, - kwargs=kwargs, + looped: Final = await bill_first_round_on_failure( + maybe_run_chat_completion_agentic_loop( + response=response, + model=model, + messages=messages, + optional_params={ + k: v + for k, v in completion_kwargs.items() + if v is not None + and k + not in ( + "model", + "messages", + "stream", + "acompletion", + "deployment_id", + ) + }, + kwargs=kwargs, + logging_obj=kwargs.get("litellm_logging_obj"), + custom_llm_provider=custom_llm_provider, + stream=bool(stream), + ), logging_obj=kwargs.get("litellm_logging_obj"), - custom_llm_provider=custom_llm_provider, + response=response, stream=bool(stream), ) if looped is not None: @@ -6169,6 +6176,7 @@ def completion( *ANTHROPIC_WIF_KWARGS_KEYS, *OPENAI_WIF_KWARGS_KEYS, *OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, + *AGENTIC_LOOP_CONTROL_KWARGS_KEYS, "github_copilot_auth_type", "github_copilot_host", "github_copilot_user_session", diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a0c43b9a8c5..312bdbc55ab 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -16,6 +16,7 @@ from litellm.anthropic_interface.exceptions import ( AnthropicExceptionMapping, ) from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.tool_fields import first_validation_error from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler from litellm.llms.anthropic.messages.context_management import ( @@ -49,6 +50,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( ) from litellm.proxy.common_utils.token_count_log import rejected_token_count from litellm.proxy.common_utils.token_count_settings import count_source_fields +from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, should_auto_drop_params_for_agentic_cli, @@ -251,7 +253,9 @@ async def anthropic_response( call_id=base_llm_response_processor.litellm_call_id, model_id=model_id, version=version, - response_cost=0, + response_cost=( + failure_row_response_cost(base_llm_response_processor.data) if is_agentic_followup_failure(e) else 0 + ), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), request_data=data, timeout=getattr(answered, "timeout", None), diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 647b19d5597..5ae6c628324 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -55,6 +55,7 @@ from litellm.constants import ( from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.otel.runtime import phase_event from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.bug_report import ( allowlisted, bug_report_notice, @@ -132,6 +133,7 @@ from litellm.proxy.common_utils.sse_keepalive import ( from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash +from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( endpoint_unsupported_by_model, @@ -4402,6 +4404,7 @@ class ProxyBaseLLMRequestProcessing: """Raises ProxyException (OpenAI API compatible) if an exception is raised""" log_llm_api_exception(e, self.litellm_call_id) billing_logging_obj: Final = _REQUEST_FIELDS.validate_python(self.data).get("litellm_logging_obj") + failed_in_agentic_follow_up: Final = is_agentic_followup_failure(e) # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -4438,7 +4441,9 @@ class ProxyBaseLLMRequestProcessing: model_id=model_id, version=version, response_cost=( - _litellm_logging_obj.finalized_cost.amount + failure_row_response_cost(self.data) + if failed_in_agentic_follow_up + else _litellm_logging_obj.finalized_cost.amount if isinstance(_litellm_logging_obj, LiteLLMLoggingObj) and _litellm_logging_obj.finalized_cost is not None else 0 diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 4d910cde603..9b9e0949a1a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -22,7 +22,7 @@ import json import time from collections import Counter, OrderedDict from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeGuard +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, TypeGuard from urllib.parse import urlparse import httpx @@ -507,6 +507,8 @@ class _CompressionResult: class CompresrGuardrail(CustomGuardrail): + server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({COMPRESR_RETRIEVE_TOOL_NAME}) + def __init__( self, api_base: str | None = None, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 8a015d396d0..031bdfa19cd 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -225,6 +225,32 @@ def _guardrail_info_lifted_from_litellm_metadata( return {"standard_logging_guardrail_information": guardrail_info} +def _failure_row_guardrail_information(request_data: Mapping[str, object]) -> object: + metadata: Final = _str_keyed_fields(request_data.get("metadata")) + lifted: Final = _guardrail_info_lifted_from_litellm_metadata( + _str_keyed_fields(request_data.get("litellm_metadata")), + guardrail_info_already_set="standard_logging_guardrail_information" in metadata, + ) + return {**metadata, **lifted}.get("standard_logging_guardrail_information") + + +def failure_row_response_cost(request_data: Mapping[str, object]) -> float: + stashed_cost: Final = request_data.get("response_cost") + recovered_stream_cost: Final = ( + max(float(stashed_cost), 0.0) + if isinstance(request_data.get("combined_usage_object"), litellm.Usage) + and isinstance(stashed_cost, (int, float)) + else 0.0 + ) + finalized_failure: Final = _str_keyed_fields(request_data.get("standard_logging_object")) + finalized_amount: Final = finalized_failure.get("response_cost") + return ( + float(finalized_amount) + if finalized_failure.get("cost_finalized") is True and isinstance(finalized_amount, (int, float)) + else recovered_stream_cost + guardrail_information_cost(_failure_row_guardrail_information(request_data)) + ) + + def _redaction_opt_in_lifted_from_litellm_metadata( litellm_metadata: ObjectMapping, *, caller_set_headers: bool ) -> ObjectMapping: @@ -577,19 +603,7 @@ class _ProxyDBLogger(CustomLogger): # recovered cost onto request_data (the usage rides along in # ``combined_usage_object`` for the token columns), so attribute the # real partial spend to this failure row instead of zero. - recovered_stream_cost: Final = ( - max(float(request_data.get("response_cost") or 0.0), 0.0) - if isinstance(request_data.get("combined_usage_object"), litellm.Usage) - else 0.0 - ) - finalized_failure: Final = _str_keyed_fields(request_data.get("standard_logging_object")) - finalized_amount: Final = finalized_failure.get("response_cost") - recovered_response_cost: Final = ( - float(finalized_amount) - if finalized_failure.get("cost_finalized") is True and isinstance(finalized_amount, (int, float)) - else recovered_stream_cost - + guardrail_information_cost(existing_metadata.get("standard_logging_guardrail_information")) - ) + recovered_response_cost: Final = failure_row_response_cost(request_data) failure_snapshot_with_late_guardrails: Final = _failure_snapshot_with_guardrails_recorded_after_it( request_data.get("standard_logging_object"), diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b1544b1e042..dfb8d966078 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -649,10 +649,12 @@ def get_logging_payload( usage = _usage # A request that failed mid-stream has no usable response_obj usage, and a stream - # continued on a fallback only has the fallback's. combined_usage_object holds the - # usage of every call the client's answer took, so it wins in both cases. + # continued on a fallback or a success that bills a failed follow-up only has its own. + # combined_usage_object holds the usage of every call the row bills, so it wins in each case. _combined_usage: Final = kwargs.get("combined_usage_object") - if (not usage or "continued_partial_usage" in kwargs) and isinstance(_combined_usage, litellm.Usage): + if ( + not usage or "continued_partial_usage" in kwargs or "failed_attempts_billed_on_success" in kwargs + ) and isinstance(_combined_usage, litellm.Usage): usage = _combined_usage.model_dump() id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 08f717ee81c..e28ebf439b5 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -110,6 +110,7 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( replace_complexity_router_headers, response_total_token_count, ) +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.batch_utils import ( batch_create_routing, @@ -4044,6 +4045,8 @@ class PythonCore(DeprecatedRouterMembers): if continuation_output is not None else None ) + if cut_partial is not None: + model_response.logging_obj.mark_stash_billed_by_continuation() if continue_after_content: emitted_tokens: Final = int( getattr(complete_response_object_usage, "completion_tokens", 0) or 0 @@ -7373,6 +7376,8 @@ class PythonCore(DeprecatedRouterMembers): last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see for attempt in range(attempted_retries_for_request(initial_kwargs), budget): trigger = _mid_stream_retry_trigger(last_error) + if is_agentic_followup_failure(trigger): + return _StreamRetriesExhausted(last_error) if not policy_applies and not self._stream_should_retry( trigger, healthy_deployments, all_deployments, initial_kwargs ): @@ -9602,7 +9607,7 @@ class PythonCore(DeprecatedRouterMembers): response = add_retry_headers_to_response(response=response, attempted_retries=0, max_retries=None) return response except Exception as e: - if is_guardrail_intervention(e): + if is_guardrail_intervention(e) or is_agentic_followup_failure(e): raise current_attempt = None original_exception = e @@ -9697,6 +9702,7 @@ class PythonCore(DeprecatedRouterMembers): # Update retry tracking metadata before each retry attempt _metadata["attempted_retries"] = current_attempt + 1 _metadata["max_retries"] = num_retries + carry_superseded_attempt_usage(kwargs) # if the function call is successful, no exception will be raised and we'll break out of the loop response = await self.make_call(original_function, *args, **kwargs) if coroutine_checker.is_async_callable(response): # async errors are often returned as coroutines @@ -9710,6 +9716,8 @@ class PythonCore(DeprecatedRouterMembers): return response except Exception as e: + if is_agentic_followup_failure(e): + raise # Always track the latest error so we raise the most # recent exception instead of the first one. original_exception = e diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 77e75dea030..14ed719a880 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -74,14 +74,13 @@ class AgenticLoopSafetyError(ValueError): Covers both rails: the bounded-loop cap (``max_agentic_loops``) and the repeated tool-call fingerprint cycle break. Subclasses ``ValueError`` so callers that already catch the broader type keep working. - - Only the anthropic messages loop raises this today. The chat completions - loop in ``litellm_core_utils/chat_completion_agentic_loop.py`` still raises - a plain ``ValueError`` from its own copy of the same rails, so catching - this type alone will not cover that surface until it is moved over. """ +class AgenticLoopLimitError(AgenticLoopSafetyError): + pass + + class StandardCustomLoggerInitParams(LiteLLMBaseModel): """ Params for initializing a CustomLogger. diff --git a/litellm/utils.py b/litellm/utils.py index 9f49e632395..ec7bb041fb5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -83,6 +83,7 @@ from litellm.constants import ( PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( CHAT_ALLOWED_TOOLS_LAYOUT, parse_chat_allowed_tools_tool_choice, @@ -2312,7 +2313,7 @@ def client(original_function): is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_completion_litellm_router_call + num_retries and not is_completion_litellm_router_call and not is_agentic_followup_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2346,7 +2347,7 @@ def client(original_function): is_responses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_responses_litellm_router_call + num_retries and not is_responses_litellm_router_call and not is_agentic_followup_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2645,7 +2646,10 @@ def client(original_function): is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_acompletion_litellm_router_call + and not isinstance(e, ImportError) + and not is_agentic_followup_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops @@ -2676,7 +2680,10 @@ def client(original_function): is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_aresponses_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_aresponses_litellm_router_call + and not isinstance(e, ImportError) + and not is_agentic_followup_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops diff --git a/tests/integration/streaming/test_interception_converted_stream_options.py b/tests/integration/streaming/test_interception_converted_stream_options.py index 5a7855102c0..6fc38dfb440 100644 --- a/tests/integration/streaming/test_interception_converted_stream_options.py +++ b/tests/integration/streaming/test_interception_converted_stream_options.py @@ -848,60 +848,36 @@ def test_wrong_azure_key_surfaces_azures_401(rig: _Rig, scenario: Scenario) -> N assert len(_upstream(rig, deployment).model_calls) == 1 -def _tool_call(call: Mapping[str, JsonValue]) -> tuple[str, Mapping[str, JsonValue]]: - function: Final = object_value(call["function"]) - return string_value(function["name"]), iv.JSON_OBJECT.validate_json(string_value(function["arguments"])) +def _assert_failed_once(call_id: str, *, spend: float) -> None: + rows: Final = eventually( + lambda: read_rows('SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s', (call_id,)), + lambda found: len(found) >= 1, + seconds=70, + ) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(rows[0]["spend"]) == pytest.approx(spend), rows -def _stream_choices(stream: _Stream) -> Iterator[Mapping[str, JsonValue]]: - for frame in stream.frames: - yield from rv.ITEMS.validate_python(frame.get("choices") or []) - - -def _streamed_tool_calls(stream: _Stream) -> Iterator[tuple[str, Mapping[str, JsonValue]]]: - for choice in _stream_choices(stream): - yield from map( - _tool_call, rv.ITEMS.validate_python(object_value(choice.get("delta") or {}).get("tool_calls") or []) - ) - - -def _handed_back(deployment: _Deployment) -> list[tuple[str, Mapping[str, JsonValue]]]: - return [(iv.SEARCH_TOOL, {"query": iv.query(deployment.name)})] - - -def _assert_succeeded_once(model: str, identity: str) -> None: - rows: Final = _spend_rows(model, 1) - assert [string_value(row["request_id"]) for row in rows if row["status"] == "success"] == [identity], rows - - -def test_failed_follow_up_call_hands_the_search_tool_call_back_with_its_usage(rig: _Rig, scenario: Scenario) -> None: - deployment: Final = _deploy(rig, scenario, "followup-fails") - stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})) - assert stream.status == 200 and stream.content == "", stream.text - assert list(_streamed_tool_calls(stream)) == _handed_back(deployment), stream.text - finished: Final = [choice["finish_reason"] for choice in _stream_choices(stream) if choice.get("finish_reason")] - assert finished == ["tool_calls"], stream.text - assert [_totals(usage) for usage in stream.usages] == [iv.FIRST_USAGE["total_tokens"]], stream.text - (identity,) = stream.ids +@pytest.mark.parametrize("stream", [True, False], ids=["stream", "json"]) +def test_failed_follow_up_call_answers_with_the_providers_error_instead_of_the_search_tool_call( + rig: _Rig, scenario: Scenario, stream: bool +) -> None: + deployment: Final = _deploy( + rig, scenario, "followup-fails", priced=True + ) + extra: Final = {**_NO_CACHE, "stream_options": dict(_USAGE)} if stream else _NO_CACHE + body: Final = _chat_body(deployment, stream=stream, extra=extra) + response: Final = rig.gateway.request("POST", "/v1/chat/completions", body) + assert response.status_code == 500 and iv.FOLLOWUP_FAILURE in response.text, response.text + assert iv.SEARCH_TOOL not in response.text, response.text _assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1) - _assert_succeeded_once(deployment.model, identity) + _assert_failed_once( + response.headers["x-litellm-call-id"], + spend=iv.FIRST_USAGE["prompt_tokens"] * 0.001 + iv.FIRST_USAGE["completion_tokens"] * 0.002, + ) assert rig.gateway.chat(_deploy(rig, scenario, "plain").model)["object"] == "chat.completion" -def test_failed_follow_up_call_on_a_json_request_hands_the_search_tool_call_back(rig: _Rig, scenario: Scenario) -> None: - deployment: Final = _deploy(rig, scenario, "followup-fails") - response: Final = rig.gateway.request("POST", "/v1/chat/completions", _chat_body(deployment, stream=False)) - assert response.status_code == 200, response.text - body: Final = iv.JSON_OBJECT.validate_json(response.content) - (choice,) = rv.ITEMS.validate_python(body["choices"]) - calls: Final = rv.ITEMS.validate_python(object_value(choice["message"])["tool_calls"]) - assert [_tool_call(call) for call in calls] == _handed_back(deployment), body - assert choice["finish_reason"] == "tool_calls", body - assert _totals(object_value(body["usage"])) == iv.FIRST_USAGE["total_tokens"], body - _assert_upstream(_upstream(rig, deployment), model_calls=2, searches=1) - _assert_succeeded_once(deployment.model, string_value(body["id"])) - - def test_search_outage_still_restreams_the_follow_up_answer(rig: _Rig, scenario: Scenario) -> None: deployment: Final = _deploy(rig, scenario, "search-down") stream: Final = _chat(rig, _chat_body(deployment, extra={**_NO_CACHE, "stream_options": dict(_USAGE)})) diff --git a/tests/unit/integrations/code_interpreter_interception/test_handler.py b/tests/unit/integrations/code_interpreter_interception/test_handler.py index c037b6be80d..f4f6a127fc7 100644 --- a/tests/unit/integrations/code_interpreter_interception/test_handler.py +++ b/tests/unit/integrations/code_interpreter_interception/test_handler.py @@ -5,16 +5,24 @@ All sandbox dependencies are injected: a FakeSandbox stands in for the real e2b config and records how it is called. """ +import asyncio import time +from collections.abc import Mapping from typing import Final +import httpx import pytest +import respx + +import litellm from litellm.integrations.code_interpreter_interception.handler import ( _INTERCEPTION_ACTIVE_KEY as _ACTIVE_KEY, ) from litellm.integrations.code_interpreter_interception.handler import ( _SANDBOX_KEY, + _SANDBOX_LOST_RESULT_OUTPUT, + _SANDBOX_UNAVAILABLE_OUTPUT, _SESSION_SCOPED_KEY, LITELLM_CODE_EXECUTION_TOOL_NAME, CodeInterpreterInterceptionLogger, @@ -23,6 +31,7 @@ from litellm.llms.base_llm.sandbox.transformation import CodeExecutionResult from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, CODE_INTERPRETER_STREAM_OPTIONS_KEY, + AgenticLoopPlan, NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, is_interception_internal_key, ) @@ -831,11 +840,7 @@ async def test_post_hook_delete_is_idempotent_across_loop_levels(): @pytest.mark.asyncio -async def test_build_plan_deletes_sandbox_when_execution_raises(): - """If sandbox execution raises before a plan is built (e.g. E2B aborts - output over its cap), the cached sandbox must be deleted before re-raising, - otherwise a caller can leak paid containers until the prune TTL.""" - +async def test_a_failed_code_run_still_deletes_its_sandbox_once_the_loop_ends(): class RaisingSandbox(FakeSandbox): async def arun_code(self, *, container, code, **kwargs): raise ValueError("output exceeded cap") @@ -843,13 +848,11 @@ async def test_build_plan_deletes_sandbox_when_execution_raises(): sandbox = RaisingSandbox() logger = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) - with pytest.raises(ValueError, match="exceeded cap"): - await _build_plan(logger, sandbox, call_id="k1") + plan = await _build_plan(logger, sandbox, call_id="k1") + await logger.async_agentic_loop_cleanup_hook(plan=plan, kwargs={}) - assert len(sandbox.create_calls) == 1, "the sandbox must have been created" - assert len(sandbox.delete_calls) == 1, ( - "a build failure must delete the cached sandbox so it does not keep running and billing" - ) + assert len(sandbox.create_calls) == 1 + assert len(sandbox.delete_calls) == 1 assert "sbxkey1" not in logger._container_cache @@ -1448,3 +1451,215 @@ async def test_per_identity_cap_evicts_lru_session(): "the entry with the oldest last_accessed timestamp must be evicted first (LRU)" ) assert len(sandbox.delete_calls) == 1, "evicted sandbox must be deleted, not just removed from cache" + + +class _SandboxThatCannotStart(FakeSandbox): + async def acreate_sandbox(self, **kwargs): + raise ConnectionError("sandbox provider unreachable") + + +class _SandboxThatCannotRunCode(FakeSandbox): + async def arun_code(self, *, container, code, **kwargs): + raise TimeoutError("sandbox run timed out") + + +_SANDBOX_FAILURES: Final = ( + pytest.param(None, _SANDBOX_UNAVAILABLE_OUTPUT, id="no_sandbox"), + pytest.param(_SandboxThatCannotStart(), _SANDBOX_UNAVAILABLE_OUTPUT, id="create_fails"), + pytest.param(_SandboxThatCannotRunCode(), _SANDBOX_LOST_RESULT_OUTPUT, id="run_fails"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("sandbox", "expected_output"), _SANDBOX_FAILURES) +async def test_a_sandbox_failure_answers_the_responses_tool_call_with_a_fixed_failure(sandbox, expected_output): + plan: Final = await _build_plan(CodeInterpreterInterceptionLogger(sandbox_config=sandbox), sandbox) + + outputs: Final = [message for message in _iter_messages(plan) if message.get("type") == "function_call_output"] + assert plan.run_agentic_loop is True + assert [(output["call_id"], output["output"]) for output in outputs] == [("c1", expected_output)] + assert [call["status"] for call in plan.metadata["code_interpreter_calls"]] == ["failed"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("sandbox", "expected_output"), _SANDBOX_FAILURES) +async def test_a_sandbox_failure_answers_the_chat_tool_call_with_a_fixed_failure(sandbox, expected_output): + logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + plan: Final = await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + {"id": "call_1", "name": LITELLM_CODE_EXECUTION_TOOL_NAME, "arguments": '{"code":"print(40 + 2)"}'} + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response={"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]}, + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id="k1"), + stream=False, + kwargs={ + "litellm_call_id": "k1", + _ACTIVE_KEY: True, + _SANDBOX_KEY: "sbxkey1", + "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE, + }, + ) + + tool_messages: Final = [message for message in _iter_messages(plan) if message.get("role") == "tool"] + assert plan.run_agentic_loop is True + assert [(message["tool_call_id"], message["content"]) for message in tool_messages] == [ + ("call_1", expected_output) + ] + assert [call["status"] for call in plan.metadata["code_interpreter_calls"]] == ["failed"] + + +_SESSION_PLAN_SURFACES: Final = ( + pytest.param({}, FakeResponse(output=[_function_call_item()]), id="responses"), + pytest.param( + {_ACTIVE_KEY: True, "_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE}, + {"choices": [{"message": {"tool_calls": [_chat_function_call_item()]}}]}, + id="chat", + ), +) + + +async def _build_session_plan( + logger: CodeInterpreterInterceptionLogger, + request_id: str, + surface_kwargs: Mapping[str, object] | None = None, + response: object | None = None, +) -> AgenticLoopPlan: + return await logger.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_1", + "call_id": "c1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "arguments": '{"code":"print(1)"}', + } + ] + }, + model="gpt-5", + messages=[{"role": "user", "content": "x"}], + response=response or FakeResponse(output=[_function_call_item()]), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={"tools": []}, + logging_obj=FakeLogging(litellm_call_id=request_id), + stream=False, + kwargs={ + "litellm_call_id": request_id, + _SANDBOX_KEY: "conv-1", + _SESSION_SCOPED_KEY: True, + **(surface_kwargs or {}), + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("surface_kwargs", "response"), _SESSION_PLAN_SURFACES) +async def test_a_session_sandbox_whose_code_run_fails_is_not_reused_by_the_next_request( + surface_kwargs: Mapping[str, object], response: object +) -> None: + sandbox: Final = _SandboxThatCannotRunCode() + logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + await _build_session_plan(logger, "req1", surface_kwargs, response) + await _build_session_plan(logger, "req2", surface_kwargs, response) + + assert len(sandbox.create_calls) == 2 + assert len(sandbox.delete_calls) == 2 + + +class _SandboxWhoseFirstContainerDies(FakeSandbox): + def __init__(self) -> None: + super().__init__() + self.handles: list[FakeHandle] = [] + self.first_run_started = asyncio.Event() + self.first_run_may_finish = asyncio.Event() + + async def acreate_sandbox(self, **kwargs: object) -> FakeHandle: + self.create_calls.append(kwargs) + handle: Final = FakeHandle(sandbox_id=f"sbx_{len(self.handles)}") + self.handles.append(handle) + return handle + + async def arun_code(self, *, container: FakeHandle, code: str, **kwargs: object) -> CodeExecutionResult: + self.run_calls.append({"container": container, "code": code}) + if len(self.run_calls) == 1: + self.first_run_started.set() + await self.first_run_may_finish.wait() + if container is self.handles[0]: + raise TimeoutError("sandbox expired") + return CodeExecutionResult(stdout=self.stdout) + + +@pytest.mark.asyncio +async def test_a_late_failure_on_a_dead_session_sandbox_keeps_its_healthy_replacement() -> None: + sandbox: Final = _SandboxWhoseFirstContainerDies() + logger: Final = CodeInterpreterInterceptionLogger(sandbox_config=sandbox) + + slow_request: Final = asyncio.create_task(_build_session_plan(logger, "req-slow")) + await sandbox.first_run_started.wait() + await _build_session_plan(logger, "req-fails-fast") + await _build_session_plan(logger, "req-gets-replacement") + sandbox.first_run_may_finish.set() + await slow_request + await _build_session_plan(logger, "req-after") + + assert [call["container"] for call in sandbox.delete_calls] == [sandbox.handles[0]] + assert [call["container"] for call in sandbox.run_calls[-2:]] == [sandbox.handles[1], sandbox.handles[1]] + + +def _anthropic_message(content: list[dict[str, object]], stop_reason: str) -> dict[str, object]: + return { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": stop_reason, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + + +@pytest.mark.asyncio +async def test_a_responses_request_bridged_to_claude_runs_the_code_and_ends_without_the_internal_call( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + sandbox: Final = FakeSandbox(stdout="42") + monkeypatch.setattr(litellm, "callbacks", [CodeInterpreterInterceptionLogger(sandbox_config=sandbox)]) + provider: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + side_effect=[ + httpx.Response( + 200, + json=_anthropic_message( + [ + { + "type": "tool_use", + "id": "toolu_1", + "name": LITELLM_CODE_EXECUTION_TOOL_NAME, + "input": {"code": "print(40 + 2)"}, + } + ], + "tool_use", + ), + ), + httpx.Response(200, json=_anthropic_message([{"type": "text", "text": "The answer is 42."}], "end_turn")), + ] + ) + + response: Final = await litellm.aresponses( + model="anthropic/claude-sonnet-5-5", + input="compute 40 + 2", + api_key="sk-test", + tools=[{"type": "code_interpreter", "container": {"type": "auto"}}], + ) + + assert provider.call_count == 2 + assert [run["code"] for run in sandbox.run_calls] == ["print(40 + 2)"] + assert not [item for item in response.output if getattr(item, "type", None) == "function_call"] 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 351ca9beb4d..e5ef356b0f3 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 @@ -12,6 +12,7 @@ config.yaml through to the settings the loop actually reads. """ import json +import time from unittest.mock import MagicMock, patch import pytest @@ -24,6 +25,8 @@ from litellm.integrations.websearch_interception.handler import ( WebSearchInterceptionLogger, ) 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.agentic_loop_settings import DEFAULT_MAX_AGENTIC_LOOPS from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( FakeAnthropicMessagesStreamIterator, @@ -31,6 +34,7 @@ from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.secret_managers.main import get_secret +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -38,6 +42,7 @@ from litellm.types.integrations.custom_logger import ( ) INTERNAL_TOOL_NAME = "litellm_web_search" +CODE_EXECUTION_TOOL_NAME = "litellm_code_execution" @pytest.fixture(autouse=True) @@ -99,6 +104,29 @@ def _tool_use_names(response: dict) -> list[str]: return [block.get("name") for block in response["content"] if block.get("type") == "tool_use"] +def _responses_turn(*function_calls: tuple[str, str]) -> ResponsesAPIResponse: + return ResponsesAPIResponse( + id="resp_1", + created_at=0, + status="completed", + model="claude-sonnet-4-5", + object="response", + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Let me check one more source.", "annotations": []}], + }, + *( + {"type": "function_call", "id": f"fc_{call_id}", "call_id": call_id, "name": name, "arguments": "{}"} + for call_id, name in function_calls + ), + ], + ) + + class _InterceptingCallback(CustomLogger): """ Stands in for the websearch interceptor: asks for another loop whenever the @@ -152,14 +180,46 @@ class _InterceptingCallback(CustomLogger): return response -def _logging_obj(callback: CustomLogger, converted_stream: bool = False) -> MagicMock: +class _ResponsesInterceptingCallback(_InterceptingCallback): + async def async_should_run_agentic_loop( + self, response, model, messages, tools, stream, custom_llm_provider, kwargs + ): + calls = [ + {"id": item.call_id, "call_id": item.call_id, "name": item.name} + for item in response.output + if item.type == "function_call" and item.name == INTERNAL_TOOL_NAME + ] + return bool(calls), {"tool_calls": calls} + + +def _logging_obj(*callbacks: CustomLogger, converted_stream: bool = False) -> MagicMock: logging_obj = MagicMock() logging_obj.model_call_details = {"websearch_interception_converted_stream": converted_stream} - logging_obj.dynamic_success_callbacks = [callback] + logging_obj.dynamic_success_callbacks = list(callbacks) logging_obj.litellm_call_id = "call-abc" return logging_obj +class _CodeExecutionGate(CustomLogger): + async def async_should_run_agentic_loop( + self, response, model, messages, tools, stream, custom_llm_provider, kwargs + ): + calls = ( + [ + {"id": block["id"], "name": block["name"]} + for block in response["content"] + if block.get("type") == "tool_use" and block.get("name") == CODE_EXECUTION_TOOL_NAME + ] + if isinstance(response, dict) + else [ + {"call_id": item.call_id, "name": item.name} + for item in response.output + if item.type == "function_call" and item.name == CODE_EXECUTION_TOOL_NAME + ] + ) + return bool(calls), {"tool_calls": calls} + + async def _run_hooks( handler: BaseLLMHTTPHandler, callback: CustomLogger, @@ -314,6 +374,48 @@ class TestCappedLoopReturnsTerminalResponse: assert _tool_use_names(result) == ["get_weather"] assert result["stop_reason"] == "tool_use" + @pytest.mark.asyncio + async def test_every_hooks_internal_call_is_dropped_and_the_clients_survives(self): + response = _response_asking_for_another_search() + response["content"] += [ + {"id": "toolu_code_1", "type": "tool_use", "name": CODE_EXECUTION_TOOL_NAME, "input": {}}, + {"id": "toolu_client_1", "type": "tool_use", "name": "get_weather", "input": {}}, + ] + + result = await self.handler._call_agentic_completion_hooks( + response=response, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj(self.callback, _CodeExecutionGate()), + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + ) + + assert _tool_use_names(result) == ["get_weather"] + assert self.callback.plan_calls == 0 + + @pytest.mark.asyncio + async def test_responses_surface_drops_every_hooks_internal_call_and_keeps_the_clients(self): + result = await self.handler._call_agentic_completion_hooks( + response=_responses_turn( + ("call_search", INTERNAL_TOOL_NAME), ("call_code", CODE_EXECUTION_TOOL_NAME), ("call_client", "lookup") + ), + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=_logging_obj(_ResponsesInterceptingCallback(), _CodeExecutionGate()), + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + api_surface="responses", + ) + + assert [getattr(item, "call_id", None) for item in result.output] == [None, "call_client"] + def test_only_the_refused_tool_calls_are_dropped(self): """ A block is matched on the id the rail refused, not on the tool name, so a @@ -347,6 +449,18 @@ class TestCappedLoopReturnsTerminalResponse: assert _tool_use_names(result) == [] assert result["stop_reason"] == "end_turn" + def test_a_name_only_refusal_keeps_the_native_search_blocks_sharing_that_name(self): + response = _response_asking_for_another_search() + response["content"][-1] = {"type": "tool_use", "name": "web_search", "input": {}} + + result = BaseLLMHTTPHandler._finalize_refused_agentic_response( + response=response, + tool_calls={"tool_calls": [{"name": "web_search", "input": {}}]}, + ) + + assert _block_types(result) == ["server_tool_use", "web_search_tool_result", "text"] + assert result["stop_reason"] == "end_turn" + @pytest.mark.asyncio async def test_streaming_caller_is_left_to_its_existing_behavior(self): """ @@ -365,19 +479,34 @@ class TestCappedLoopReturnsTerminalResponse: assert self.callback.plan_calls == 0 @pytest.mark.asyncio - async def test_responses_surface_is_left_to_its_existing_behavior(self): - """ - The responses surface carries a pydantic model rather than the anthropic - dict this finalizer rewrites, so it keeps raising instead of being handed - a response that was never actually finalized. - """ - with pytest.raises(AgenticLoopSafetyError): - await _run_hooks( - self.handler, - self.callback, - kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, - api_surface="responses", - ) + async def test_responses_surface_drops_the_refused_function_call_and_keeps_the_answer(self): + response = _responses_turn(("call_search", INTERNAL_TOOL_NAME), ("call_client", "lookup")) + callback = _ResponsesInterceptingCallback() + + result = await _run_hooks( + self.handler, + callback, + kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + response=response, + api_surface="responses", + ) + + assert isinstance(result, ResponsesAPIResponse) + assert [(item.type, getattr(item, "name", None)) for item in result.output] == [ + ("message", None), + ("function_call", "lookup"), + ] + assert callback.plan_calls == 0 + + def test_responses_refusal_by_call_id_alone_keeps_the_calls_it_did_not_list(self): + response = _responses_turn(("call_search", INTERNAL_TOOL_NAME), ("call_unlisted", INTERNAL_TOOL_NAME)) + + result = BaseLLMHTTPHandler._finalize_refused_agentic_response( + response=response, + tool_calls={"tool_calls": [{"call_id": "call_search", "name": INTERNAL_TOOL_NAME}]}, + ) + + assert [getattr(item, "call_id", None) for item in result.output] == [None, "call_unlisted"] @pytest.mark.asyncio async def test_non_dict_response_is_returned_untouched(self): @@ -509,6 +638,156 @@ def _response_asking_for_searches(*queries: str) -> dict: } +def _request_logging_obj(call_type: str, *callbacks: CustomLogger) -> Logging: + logging_obj = Logging( + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + stream=False, + call_type=call_type, + start_time=time.time(), + litellm_call_id="call-abc", + function_id="f", + dynamic_success_callbacks=list(callbacks), + ) + logging_obj.update_environment_variables( + model="claude-sonnet-4-5", user="", optional_params={}, litellm_params={}, custom_llm_provider="anthropic" + ) + return logging_obj + + +def _assert_first_call_billed(logging_obj: Logging, *, input_tokens: int, output_tokens: int) -> None: + price = litellm.model_cost["claude-sonnet-4-5"] + assert logging_obj.model_call_details["combined_usage_object"].total_tokens == input_tokens + output_tokens + assert logging_obj.model_call_details["response_cost"] == pytest.approx( + input_tokens * price["input_cost_per_token"] + output_tokens * price["output_cost_per_token"] + ) + + +class TestFailedFollowUpCallReachesTheClient: + def setup_method(self): + self.handler = BaseLLMHTTPHandler() + self.logger = WebSearchInterceptionLogger(enabled_providers=["anthropic"]) + + @staticmethod + async def _rate_limited_follow_up(**call_kwargs): + raise RateLimitError( + "UserByModelByMinuteUncachedInputTokens", llm_provider="anthropic", model="claude-sonnet-4-5" + ) + + @staticmethod + async def _search(query, kwargs=None, rich=None): + found = SearchResult(title="Result", url="https://example.com", snippet="A result.", date=None) + return ("Title: Result\nURL: https://example.com", SearchResponse(results=[found])) + + @pytest.mark.asyncio + @pytest.mark.parametrize("converted_stream", (False, True), ids=("non_streaming", "streaming_client")) + async def test_follow_up_rate_limit_is_raised_to_the_caller(self, monkeypatch, converted_stream): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._rate_limited_follow_up) + + with ( + patch.object(self.logger, "_execute_search", side_effect=self._search), + pytest.raises(RateLimitError, match="UserByModelByMinuteUncachedInputTokens"), + ): + await self.handler._call_agentic_completion_hooks( + response=_response_asking_for_searches("who won the world cup"), + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]}, + logging_obj=_logging_obj(self.logger, converted_stream=converted_stream), + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + @pytest.mark.asyncio + async def test_follow_up_rate_limit_bills_the_first_call_on_the_failure_log(self, monkeypatch): + monkeypatch.setattr("litellm.anthropic_interface.messages.acreate", self._rate_limited_follow_up) + logging_obj = _request_logging_obj("anthropic_messages", self.logger) + + with ( + patch.object(self.logger, "_execute_search", side_effect=self._search), + pytest.raises(RateLimitError), + ): + await self.handler._call_agentic_completion_hooks( + response=_response_asking_for_searches("who won the world cup"), + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]}, + logging_obj=logging_obj, + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + _assert_first_call_billed(logging_obj, input_tokens=10, output_tokens=5) + + @pytest.mark.asyncio + async def test_a_crashed_plan_build_is_not_retried_and_bills_the_first_call(self): + logging_obj = _request_logging_obj("anthropic_messages", self.logger) + + with ( + patch.object(self.logger, "async_build_agentic_loop_plan", side_effect=ValueError("the plan crashed")), + pytest.raises(ValueError, match="the plan crashed") as raised, + ): + await self.handler._call_agentic_completion_hooks( + response=_response_asking_for_searches("who won the world cup"), + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={"tools": [get_litellm_web_search_tool()]}, + logging_obj=logging_obj, + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3, WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, + ) + + assert is_agentic_followup_failure(raised.value) + _assert_first_call_billed(logging_obj, input_tokens=10, output_tokens=5) + + @pytest.mark.asyncio + async def test_responses_follow_up_rate_limit_bills_the_first_call_on_the_failure_log(self, monkeypatch): + monkeypatch.setattr(litellm, "aresponses", self._rate_limited_follow_up) + callback = _ResponsesInterceptingCallback() + logging_obj = _request_logging_obj("aresponses", callback) + first_call = _responses_turn(("call_search", INTERNAL_TOOL_NAME)).model_copy( + update={"usage": {"input_tokens": 13, "output_tokens": 6, "total_tokens": 19}} + ) + + with pytest.raises(RateLimitError): + await self.handler._call_agentic_completion_hooks( + response=first_call, + model="claude-sonnet-4-5", + messages=[{"role": "user", "content": "who won the world cup"}], + anthropic_messages_provider_config=MagicMock(), + anthropic_messages_optional_request_params={}, + logging_obj=logging_obj, + stream=False, + custom_llm_provider="anthropic", + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3}, + api_surface="responses", + ) + + _assert_first_call_billed(logging_obj, input_tokens=13, output_tokens=6) + + @pytest.mark.asyncio + async def test_responses_follow_up_rate_limit_is_raised_to_the_caller(self, monkeypatch): + monkeypatch.setattr(litellm, "aresponses", self._rate_limited_follow_up) + callback = _ResponsesInterceptingCallback() + + with pytest.raises(RateLimitError, match="UserByModelByMinuteUncachedInputTokens"): + await _run_hooks( + self.handler, + callback, + kwargs={"_agentic_loop_depth": 0, "max_agentic_loops": 3}, + response=_responses_turn(("call_search", INTERNAL_TOOL_NAME)), + api_surface="responses", + ) + + assert callback.plan_calls == 1 + + class TestFailedSearchEndsTheTurn: """ A search that failed used to come back to the client as an empty successful diff --git a/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py index b2d5225070c..ee4c4d32478 100644 --- a/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py +++ b/tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py @@ -371,8 +371,7 @@ class TestLoggingObjExcludedFromFollowUp: class TestFollowUpErrorScenarios: """Regression tests: the agentic loop must surface errors properly and - not silently swallow them (except at the _call_agentic_completion_hooks - level which intentionally catches to return the initial response).""" + not silently swallow them.""" @pytest.mark.asyncio async def test_followup_400_raises(self): 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 6167e91c37f..d26da679246 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 @@ -21,6 +21,8 @@ they stay out of the body even without it. """ import datetime +import json +from collections.abc import Mapping, Sequence from typing import Any, Dict, Final, List, Optional, Tuple from unittest.mock import AsyncMock, MagicMock, patch @@ -28,11 +30,13 @@ import pytest from openai._legacy_response import LegacyAPIResponse import litellm +from litellm.caching.caching import Cache from litellm.integrations.code_interpreter_interception.handler import ( CodeInterpreterInterceptionLogger, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.chat_completion_agentic_loop import ( + _execute_chat_completion_agentic_plan, maybe_run_chat_completion_agentic_loop, ) from litellm.litellm_core_utils.litellm_logging import Logging @@ -43,6 +47,7 @@ from litellm.types.integrations.custom_logger import ( WEBSEARCH_INTERCEPTION_PREFIX, AgenticLoopPlan, AgenticLoopRequestPatch, + AgenticLoopSafetyError, ) from litellm.types.utils import ( ChatCompletionMessageToolCall, @@ -50,6 +55,7 @@ from litellm.types.utils import ( Function, Message, ModelResponse, + Usage, ) # The internal control fields that must never reach a provider request body. @@ -260,8 +266,6 @@ class _GateOnlyLogger(CustomLogger): class _RecordingGateLogger(_GateOnlyLogger): - """Also keeps every follow-up response the loop hands to the post hook.""" - def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None: super().__init__(plan=plan, tool_calls=tool_calls) self.followups: tuple[ModelResponse, ...] = () @@ -441,72 +445,544 @@ async def test_dispatcher_followup_does_not_repeat_a_request_param_found_in_requ assert call_kwargs["api_base"] == "https://a" +_REFUSED_GATE_TOOL_CALLS: Final = {"tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}]} +_REFUSING_RAILS: Final = ( + {"_agentic_loop_depth": 3, "max_agentic_loops": 3}, + { + "_agentic_loop_depth": 0, + "max_agentic_loops": 3, + "_agentic_loop_fingerprints": [json.dumps(_REFUSED_GATE_TOOL_CALLS, sort_keys=True, default=str)], + }, +) + + +def _tool_call(call_id: str, name: str) -> ChatCompletionMessageToolCall: + return ChatCompletionMessageToolCall(id=call_id, type="function", function=Function(name=name, arguments="{}")) + + +async def _run_refused_rerun( + response: ModelResponse, + rail_kwargs: Dict[str, object], + stream: bool = False, + gate_tool_calls: Optional[Dict[str, object]] = _REFUSED_GATE_TOOL_CALLS, + other_hooks_tool_calls: tuple[Dict[str, object], ...] = (), +) -> Any: + litellm.callbacks = [ + _GateOnlyLogger(plan=AgenticLoopPlan(run_agentic_loop=True), tool_calls=tool_calls) + for tool_calls in (gate_tool_calls, *other_hooks_tool_calls) + ] + acompletion_mock = AsyncMock() + with patch.object(litellm, "acompletion", acompletion_mock): + try: + return await maybe_run_chat_completion_agentic_loop( + response=response, + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs=dict(rail_kwargs), + logging_obj=_LoggingStub(), + custom_llm_provider="openai", + stream=stream, + ) + finally: + acompletion_mock.assert_not_awaited() + + @pytest.mark.asyncio -async def test_dispatcher_raises_when_depth_reaches_max_agentic_loops( - restore_callbacks, +@pytest.mark.parametrize("rail_kwargs", _REFUSING_RAILS, ids=["max_agentic_loops", "repeated_fingerprint"]) +async def test_refused_rerun_ends_the_turn_without_the_internal_tool_call( + restore_callbacks, rail_kwargs: Dict[str, object] ): - """depth >= max_agentic_loops -> ValueError mentioning max_agentic_loops, - before any follow-up call is attempted.""" - logger = _GateOnlyLogger( - plan=AgenticLoopPlan(run_agentic_loop=True), + result = await _run_refused_rerun(_tool_call_model_response(), rail_kwargs) + + assert isinstance(result, ModelResponse) + assert not result.choices[0].message.tool_calls + assert result.choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_refused_rerun_keeps_the_clients_own_tool_call(restore_callbacks): + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content="Running it.", + tool_calls=[_tool_call("call_abc", "litellm_code_execution"), _tool_call("call_client", "lookup")], + ), + ) + ] + ) + + result = await _run_refused_rerun(response, _REFUSING_RAILS[0]) + + assert [call.id for call in result.choices[0].message.tool_calls] == ["call_client"] + assert result.choices[0].finish_reason == "tool_calls" + assert result.choices[0].message.content == "Running it." + + +def _indexed_tool_call(call_id: str, name: str, index: int) -> ChatCompletionMessageToolCall: + return ChatCompletionMessageToolCall( + id=call_id, type="function", function=Function(name=name, arguments="{}"), index=index + ) + + +def _turn_with_the_internal_call_between_the_clients() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + _indexed_tool_call("call_weather", "get_weather", 0), + _indexed_tool_call("call_abc", "litellm_code_execution", 1), + _indexed_tool_call("call_time", "get_time", 2), + ], + ), + ) + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rail_kwargs", _REFUSING_RAILS, ids=["max_agentic_loops", "repeated_fingerprint"]) +async def test_refused_rerun_renumbers_the_kept_tool_calls_from_zero(restore_callbacks, rail_kwargs: Dict[str, object]): + result = await _run_refused_rerun(_turn_with_the_internal_call_between_the_clients(), rail_kwargs) + + assert [(call.id, call.index) for call in result.choices[0].message.tool_calls] == [ + ("call_weather", 0), + ("call_time", 1), + ] + + +@pytest.mark.asyncio +async def test_refused_rerun_on_a_converted_stream_replays_the_kept_tool_calls_from_index_zero( + restore_callbacks, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr( + litellm, + "callbacks", + [_GateOnlyLogger(plan=AgenticLoopPlan(run_agentic_loop=True), tool_calls=_REFUSED_GATE_TOOL_CALLS)], + ) + logging_obj: Final = Logging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-refused-stream", + function_id="fn-refused-stream", + ) + + result: Final = await maybe_run_chat_completion_agentic_loop( + response=_turn_with_the_internal_call_between_the_clients(), + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + kwargs={**_REFUSING_RAILS[1], f"{CODE_INTERPRETER_INTERCEPTION_PREFIX}_converted_stream": True}, + logging_obj=logging_obj, + custom_llm_provider="openai", + stream=False, + ) + + assert isinstance(result, CustomStreamWrapper) + deltas: Final = [chunk.choices[0].delta async for chunk in result if chunk.choices] + streamed_calls: Final = [call for delta in deltas for call in delta.tool_calls or []] + assert [(call.id, call.index) for call in streamed_calls] == [("call_weather", 0), ("call_time", 1)] + + +@pytest.mark.asyncio +async def test_refused_rerun_drops_every_hooks_internal_call_and_keeps_the_clients(restore_callbacks): + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + _tool_call("call_abc", "litellm_code_execution"), + _tool_call("call_search", "litellm_web_search"), + _tool_call("call_client", "lookup"), + ], + ), + ) + ] + ) + + result = await _run_refused_rerun( + response, + _REFUSING_RAILS[0], + other_hooks_tool_calls=({"tool_calls": [{"id": "call_search", "name": "litellm_web_search"}]},), + ) + + assert [call.id for call in result.choices[0].message.tool_calls] == ["call_client"] + + +@pytest.mark.asyncio +async def test_refused_rerun_drops_a_call_claimed_only_by_name_next_to_one_claimed_by_id(restore_callbacks): + response: Final = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + _tool_call("call_abc", "litellm_code_execution"), + _tool_call("call_search", "litellm_web_search"), + _tool_call("call_client", "lookup"), + ], + ), + ) + ] + ) + + result: Final = await _run_refused_rerun( + response, + _REFUSING_RAILS[0], + gate_tool_calls={"tool_calls": [{"id": "call_abc"}]}, + other_hooks_tool_calls=({"tool_calls": [{"name": "litellm_web_search"}]},), + ) + + assert [call.id for call in result.choices[0].message.tool_calls] == ["call_client"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response, gate_tool_calls", + [(_plain_model_response(), _REFUSED_GATE_TOOL_CALLS), (_tool_call_model_response(), None)], + ids=["turn_has_no_tool_calls", "gate_names_no_tool_call"], +) +async def test_refused_rerun_returns_a_turn_it_cannot_match_unchanged( + restore_callbacks, response: ModelResponse, gate_tool_calls: Optional[Dict[str, object]] +): + result = await _run_refused_rerun(response, _REFUSING_RAILS[0], gate_tool_calls=gate_tool_calls) + + assert result.choices[0].model_dump() == response.choices[0].model_dump() + + +@pytest.mark.asyncio +async def test_refused_rerun_on_a_live_stream_still_raises(restore_callbacks): + with pytest.raises(AgenticLoopSafetyError, match="max_agentic_loops"): + await _run_refused_rerun(_tool_call_model_response(), _REFUSING_RAILS[0], stream=True) + + +@pytest.mark.asyncio +async def test_dispatcher_surfaces_a_failed_follow_up_instead_of_the_tool_call_turn(): + logger: Final = _GateOnlyLogger( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=_patched_messages(), kwargs={"mock_response": "litellm.RateLimitError"} + ), + ), + tool_calls={"tool_calls": [{"id": "call_abc"}]}, + ) + + with pytest.raises(litellm.RateLimitError): + await _dispatch_tool_call_turn(_tool_call_model_response(), logger) + + assert logger.cleanup_calls == 1 + + +class _PlanRecordingLogger(_GateOnlyLogger): + def __init__(self, plan: AgenticLoopPlan, tool_calls: Dict[str, Any]) -> None: + super().__init__(plan=plan, tool_calls=tool_calls) + self.logging_objs: List[Any] = [] + + async def async_build_agentic_loop_plan( + self, + tools: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict[str, Any], + logging_obj: Any, + stream: bool, + kwargs: Dict[str, Any], + ) -> AgenticLoopPlan: + self.logging_objs.append(logging_obj) + return self._plan + + +@pytest.mark.asyncio +async def test_acompletion_bills_the_first_call_on_the_failure_log_when_the_follow_up_fails(restore_callbacks): + logger: Final = _PlanRecordingLogger( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=_patched_messages(), kwargs={"mock_response": "litellm.RateLimitError"} + ), + ), tool_calls={"tool_calls": [{"id": "call_abc"}]}, ) litellm.callbacks = [logger] + first_call: Final = ModelResponse( + choices=_tool_call_model_response().choices, usage=Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8) + ) + mock_client: Final = MagicMock() + mock_client.post = AsyncMock(return_value=_raw_response_for(first_call)) - acompletion_mock = AsyncMock() - with patch.object(litellm, "acompletion", acompletion_mock): - with pytest.raises(ValueError, match="max_agentic_loops"): - await maybe_run_chat_completion_agentic_loop( - response=_tool_call_model_response(), - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - optional_params={}, - kwargs={"_agentic_loop_depth": 3, "max_agentic_loops": 3}, - logging_obj=_LoggingStub(), - custom_llm_provider="openai", - stream=False, - ) + with pytest.raises(litellm.RateLimitError): + await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + api_key="sk-test", + client=mock_client, + num_retries=0, + ) - acompletion_mock.assert_not_awaited() + (logging_obj,) = logger.logging_objs + price: Final = litellm.model_cost["gpt-4o-mini"] + assert logging_obj.model_call_details["combined_usage_object"].total_tokens == 8 + assert logging_obj.model_call_details["response_cost"] == pytest.approx( + 5 * price["input_cost_per_token"] + 3 * price["output_cost_per_token"] + ) + + +async def _agentic_turn_with_follow_up( + logger: _PlanRecordingLogger, monkeypatch: pytest.MonkeyPatch +) -> tuple[ModelResponse, Logging]: + monkeypatch.setattr(litellm, "callbacks", [logger]) + mock_client: Final = MagicMock() + mock_client.post = AsyncMock(return_value=_raw_response_for(_tool_call_model_response())) + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + api_key="sk-test", + client=mock_client, + num_retries=0, + ) + assert isinstance(response, ModelResponse) + (logging_obj,) = logger.logging_objs + return response, logging_obj + + +def _follow_up_logger() -> _PlanRecordingLogger: + return _PlanRecordingLogger( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages(), kwargs={"mock_response": "42"}), + ), + tool_calls={"tool_calls": [{"id": "call_abc"}]}, + ) @pytest.mark.asyncio -async def test_dispatcher_raises_on_repeated_tool_call_fingerprint(restore_callbacks): - """A tool_calls fingerprint already present in _agentic_loop_fingerprints -> - ValueError about the repeated fingerprint (cycle guard), with no follow-up - call.""" - import json +async def test_acompletion_records_the_cost_breakdown_of_a_successful_follow_up(monkeypatch): + response, logging_obj = await _agentic_turn_with_follow_up(_follow_up_logger(), monkeypatch) - # The dispatcher fingerprints the whole value the gate returns as its second - # tuple element, so the seeded fingerprint must mirror that dict exactly. - gate_tool_calls = {"tool_calls": [{"id": "call_abc", "name": "litellm_code_execution"}]} - fingerprint = json.dumps(gate_tool_calls, sort_keys=True, default=str) + response_cost: Final = response._hidden_params["response_cost"] + assert response_cost > 0 + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["total_cost"] == pytest.approx(response_cost) - logger = _GateOnlyLogger( - plan=AgenticLoopPlan(run_agentic_loop=True), - tool_calls=gate_tool_calls, + +@pytest.mark.asyncio +async def test_acompletion_keeps_a_cached_follow_up_free(monkeypatch): + monkeypatch.setattr(litellm, "cache", Cache(type="local")) + await litellm.acompletion(model="openai/gpt-4o-mini", messages=_patched_messages(), mock_response="42") + + response, _ = await _agentic_turn_with_follow_up(_follow_up_logger(), monkeypatch) + + assert response._hidden_params["cache_hit"] is True + assert response._hidden_params["response_cost"] == 0 + + +def _tool_call_turn(call_id: str, usage: Usage) -> ModelResponse: + return ModelResponse( + model="gpt-4o-mini", + choices=[ + Choices( + finish_reason="tool_calls", + message=Message(role="assistant", content=None, tool_calls=[_tool_call(call_id, "litellm_web_search")]), + ) + ], + usage=usage, ) - litellm.callbacks = [logger] - acompletion_mock = AsyncMock() - with patch.object(litellm, "acompletion", acompletion_mock): - with pytest.raises(ValueError, match="fingerprint"): - await maybe_run_chat_completion_agentic_loop( - response=_tool_call_model_response(), + +class _TwoRoundLogger(CustomLogger): + def __init__(self, middle_round: ModelResponse) -> None: + super().__init__() + self._follow_ups: Final = (middle_round, "litellm.RateLimitError") + self.logging_objs: tuple[object, ...] = () + + async def async_should_run_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + assert isinstance(response, ModelResponse) + calls: Final = response.choices[0].message.tool_calls or [] + return bool(calls), {"tool_calls": [{"id": call.id} for call in calls]} + + async def async_build_agentic_loop_plan( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> AgenticLoopPlan: + follow_up: Final = self._follow_ups[len(self.logging_objs)] + self.logging_objs = (*self.logging_objs, logging_obj) + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch(messages=_patched_messages(), kwargs={"mock_response": follow_up}), + ) + + +@pytest.mark.asyncio +async def test_acompletion_bills_every_served_round_on_the_failure_log_when_a_later_follow_up_fails( + restore_callbacks, monkeypatch: pytest.MonkeyPatch +): + first_usage: Final = Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8) + middle_usage: Final = Usage(prompt_tokens=11, completion_tokens=4, total_tokens=15) + logger: Final = _TwoRoundLogger(middle_round=_tool_call_turn("call_next", middle_usage)) + monkeypatch.setattr(litellm, "callbacks", [logger]) + mock_client: Final = MagicMock() + mock_client.post = AsyncMock(return_value=_raw_response_for(_tool_call_turn("call_abc", first_usage))) + + with pytest.raises(litellm.RateLimitError): + await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + api_key="sk-test", + client=mock_client, + num_retries=0, + ) + + first_logging_obj, _ = logger.logging_objs + assert isinstance(first_logging_obj, Logging) + price: Final = litellm.model_cost["gpt-4o-mini"] + billed_usage: Final = first_logging_obj.model_call_details["combined_usage_object"] + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == (5 + 11, 3 + 4) + assert first_logging_obj.model_call_details["response_cost"] == pytest.approx( + (5 + 11) * price["input_cost_per_token"] + (3 + 4) * price["output_cost_per_token"] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +async def test_a_follow_up_failing_with_a_builtin_error_reaches_the_caller_unchanged(stream: bool): + from litellm.litellm_core_utils.chat_completion_agentic_loop import bill_first_round_on_failure + + sandbox_error: Final = RuntimeError("sandbox crashed") + + async def failing_follow_up() -> ModelResponse: + raise sandbox_error + + with pytest.raises(RuntimeError) as raised: + await bill_first_round_on_failure( + failing_follow_up(), + logging_obj=Logging( model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], - optional_params={}, - kwargs={ - "_agentic_loop_depth": 0, - "max_agentic_loops": 3, - "_agentic_loop_fingerprints": [fingerprint], - }, - logging_obj=_LoggingStub(), - custom_llm_provider="openai", - stream=False, - ) + stream=stream, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), + litellm_call_id="call-builtin-error", + function_id="fn-builtin-error", + ), + response=_tool_call_turn("call_abc", Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8)), + stream=stream, + ) - acompletion_mock.assert_not_awaited() + assert raised.value is sandbox_error + + +class _RunOnlyLoopLogger(CustomLogger): + def __init__(self, result: ModelResponse) -> None: + super().__init__() + self._result = result + + async def async_should_run_agentic_loop( + self, + response: Any, + model: str, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + stream: bool, + custom_llm_provider: str, + kwargs: Dict[str, Any], + ) -> Tuple[bool, Dict[str, Any]]: + return True, {"tool_calls": [{"id": "call_abc"}]} + + async def async_run_agentic_loop( + self, + tools: Dict[str, Any], + model: str, + messages: List[Dict[str, Any]], + response: Any, + anthropic_messages_provider_config: Any, + anthropic_messages_optional_request_params: Dict[str, Any], + logging_obj: Any, + stream: bool, + kwargs: Dict[str, Any], + ) -> Any: + return self._result + + +class _LoggingWithCallbacks(_LoggingStub): + def __init__(self, callbacks: Tuple[CustomLogger, ...]) -> None: + self.dynamic_success_callbacks = list(callbacks) + + +async def _dispatch_tool_call_turn( + response: ModelResponse, *callbacks: CustomLogger +) -> ModelResponse | CustomStreamWrapper | None: + return await maybe_run_chat_completion_agentic_loop( + response=response, + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is 6*7?"}], + optional_params={}, + kwargs={}, + logging_obj=_LoggingWithCallbacks(callbacks), + custom_llm_provider="openai", + stream=False, + ) + + +@pytest.mark.asyncio +async def test_dispatcher_returns_the_tool_call_turn_when_the_plan_terminates(): + response: Final = _tool_call_model_response() + terminating: Final = _GateOnlyLogger( + plan=AgenticLoopPlan(terminate=True), tool_calls={"tool_calls": [{"id": "call_abc"}]} + ) + + assert await _dispatch_tool_call_turn(response, terminating) is response + + +@pytest.mark.asyncio +async def test_dispatcher_moves_past_a_plan_that_declines_the_loop(): + override: Final = _plain_model_response("from the next callback") + tool_calls: Final = {"tool_calls": [{"id": "call_abc"}]} + declined: Final = _GateOnlyLogger(plan=AgenticLoopPlan(run_agentic_loop=False), tool_calls=tool_calls) + overriding: Final = _GateOnlyLogger(plan=AgenticLoopPlan(response_override=override), tool_calls=tool_calls) + + assert await _dispatch_tool_call_turn(_tool_call_model_response(), declined, overriding) is override + + +@pytest.mark.asyncio +async def test_dispatcher_returns_a_run_only_loops_result(): + loop_result: Final = _plain_model_response() + + assert await _dispatch_tool_call_turn(_tool_call_model_response(), _RunOnlyLoopLogger(loop_result)) is loop_result @pytest.mark.asyncio @@ -562,3 +1038,55 @@ async def test_converted_stream_replays_the_final_answer_with_the_requested_usag if include_usage else [] ) + + +class TestAgenticFollowUpKeepsTheProviderPrefix: + @staticmethod + async def _run_followup( + model: str, custom_llm_provider: str, patched_model: str | None = None + ) -> ModelResponse: + response: Final = await _execute_chat_completion_agentic_plan( + plan=AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[{"role": "user", "content": "hi"}], model=patched_model + ), + ), + callback=CustomLogger(), + model=model, + optional_params={"mock_response": "ok from followup"}, + kwargs={}, + logging_obj=None, + custom_llm_provider=custom_llm_provider, + depth=0, + max_loops=2, + fingerprints=[], + fingerprint="fp", + ) + assert isinstance(response, ModelResponse) + return response + + @pytest.mark.asyncio + async def test_a_sub_path_model_still_resolves_a_provider(self) -> None: + response: Final = await self._run_followup("mantle/anthropic.claude-sonnet-5", "bedrock") + + assert response.choices[0].message.content == "ok from followup" + + @pytest.mark.asyncio + async def test_an_ordinary_model_still_resolves_a_provider(self) -> None: + response: Final = await self._run_followup("gpt-4o-mini", "openai") + + assert response.choices[0].message.content == "ok from followup" + + @pytest.mark.asyncio + async def test_a_hook_that_patches_another_providers_model_reaches_that_provider(self) -> None: + response: Final = await self._run_followup("gpt-4o", "openai", patched_model="anthropic/claude-sonnet-4-5") + + assert response.model == "claude-sonnet-4-5" + assert response._hidden_params["custom_llm_provider"] == "anthropic" + + @pytest.mark.asyncio + async def test_a_hook_echoing_an_org_style_model_stays_on_the_request_provider(self) -> None: + response: Final = await self._run_followup("openai/gpt-4o", "openrouter", patched_model="openai/gpt-4o") + + assert response._hidden_params["custom_llm_provider"] == "openrouter" diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 13e181b2b92..cf9ad4f91a3 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -14451,3 +14451,45 @@ def test_record_native_bedrock_guardrail_bills_each_call_its_own_trace_when_call (first_cost, second_cost) = costs assert second_cost == pytest.approx(2 * first_cost) assert shared_metadata == {"tag": "x"} + + +_FAILED_ATTEMPT: Final = (Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), 0.5) + + +async def _success_row_cost_after_a_failed_attempt(stash: Callable[[LitellmLogging], None]) -> float: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="recovered-attempt", + function_id="recovered-attempt", + ) + logging_obj.update_environment_variables( + model="gpt-4o-mini", user="", optional_params={}, litellm_params={"metadata": {}}, custom_llm_provider="openai" + ) + stash(logging_obj) + logging_obj.carry_failed_attempt_usage() + now: Final = datetime.datetime.now() + await logging_obj.async_success_handler( + ModelResponse(model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)), + now, + now, + ) + return logging_obj.model_call_details["standard_logging_object"]["response_cost"] + + +@pytest.mark.asyncio +async def test_a_recovered_attempt_bills_a_failed_agentic_follow_up_but_not_a_failure_that_logged_its_own_row() -> None: + own_cost: Final = await _success_row_cost_after_a_failed_attempt(lambda logging_obj: None) + after_a_dropped_stream: Final = await _success_row_cost_after_a_failed_attempt( + lambda logging_obj: logging_obj.record_partial_usage_for_failure(*_FAILED_ATTEMPT) + ) + after_a_failed_follow_up: Final = await _success_row_cost_after_a_failed_attempt( + lambda logging_obj: logging_obj.record_served_rounds_for_failure((_FAILED_ATTEMPT,)) + ) + + assert own_cost > 0 + assert after_a_dropped_stream == pytest.approx(own_cost) + assert after_a_failed_follow_up == pytest.approx(own_cost + _FAILED_ATTEMPT[1]) diff --git a/tests/unit/llms/anthropic/messages/native/test_agentic_streaming_iterator.py b/tests/unit/llms/anthropic/messages/native/test_agentic_streaming_iterator.py index 90a6ae3a574..7db81b7841e 100644 --- a/tests/unit/llms/anthropic/messages/native/test_agentic_streaming_iterator.py +++ b/tests/unit/llms/anthropic/messages/native/test_agentic_streaming_iterator.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, MagicMock import pytest - +import litellm from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, @@ -21,6 +21,10 @@ from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( _handle_message_start, _parse_sse_events, ) +from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.utils import ProxyLogging # --------------------------------------------------------------------------- @@ -88,6 +92,14 @@ def _build_simple_text_stream() -> List[bytes]: return chunks +def _sse_frames(wire: str) -> list[tuple[str, str]]: + return [ + (frame.split("\n")[0].removeprefix("event: "), frame.split("data: ", 1)[1]) + for frame in wire.split("\n\n") + if frame.startswith("event: ") and "data: " in frame + ] + + def _build_tool_use_stream() -> List[bytes]: """Produce SSE bytes for a response with a tool_use block.""" chunks = [] @@ -727,8 +739,7 @@ class TestAgenticStreamingIteratorPhase2: class TestAgenticStreamingIteratorErrorHandling: @pytest.mark.asyncio - async def test_should_swallow_hook_errors(self): - """Errors in hook processing should be swallowed; Phase 1 chunks are still yielded.""" + async def test_a_hook_failure_reaches_the_client_after_the_streamed_message(self): chunks = _build_simple_text_stream() mock_stream = MockAsyncStream(chunks) @@ -751,11 +762,15 @@ class TestAgenticStreamingIteratorErrorHandling: ) collected = [] - async for chunk in iterator: - collected.append(chunk) - # All Phase 1 chunks should still have been yielded - assert len(collected) == len(chunks) + async def _drain(): + async for chunk in iterator: + collected.append(chunk) + + with pytest.raises(RuntimeError, match="hook exploded"): + await _drain() + + assert collected == chunks @pytest.mark.asyncio async def test_should_handle_empty_stream(self): @@ -924,21 +939,54 @@ class TestAgenticStreamingIteratorHoldBack: assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == phase2_chunks @pytest.mark.asyncio - async def test_should_error_instead_of_replaying_server_fulfilled_tool_use_when_hook_crashes(self): - """A hook crash must not replay the buffered retrieval tool_use: that is the unknown-tool bug.""" + async def test_a_rate_limited_follow_up_reaches_the_client_as_its_anthropic_error_event(self): chunks = _build_tool_use_stream() mock_handler = MagicMock() - mock_handler._call_agentic_completion_hooks = AsyncMock(side_effect=RuntimeError("hook exploded")) + mock_handler._call_agentic_completion_hooks = AsyncMock( + side_effect=litellm.RateLimitError( + message="follow-up rate limited", llm_provider="anthropic", model="claude-sonnet-4-20250514" + ) + ) - iterator = _build_hold_back_iterator(MockAsyncStream(chunks), mock_handler) + wire = "".join( + [ + frame if isinstance(frame, str) else frame.decode() + async for frame in ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=_build_hold_back_iterator(MockAsyncStream(chunks), mock_handler), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + request_data={"model": "claude-sonnet-4-20250514"}, + proxy_logging_obj=ProxyLogging(user_api_key_cache=MagicMock()), + ) + ] + ) - collected = [] - async for chunk in iterator: - collected.append(chunk) + error_events = [json.loads(data) for event, data in _sse_frames(wire) if event == "error"] + assert [event["error"]["type"] for event in error_events] == ["rate_limit_error"] + assert "follow-up rate limited" in error_events[0]["error"]["message"] + assert "litellm_content_retrieve" not in wire - assert [c for c in collected if c != STREAM_SSE_KEEPALIVE_PING_BYTES] == [SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] - assert b"litellm_content_retrieve" not in b"".join(collected) + @pytest.mark.asyncio + async def test_a_capped_loop_tells_the_client_to_raise_max_agentic_loops_instead_of_retrying(self): + mock_handler = MagicMock() + mock_handler._call_agentic_completion_hooks = AsyncMock( + side_effect=lambda **_: BaseLLMHTTPHandler._check_agentic_loop_safety( + tool_calls=[{"name": "litellm_content_retrieve"}], + fingerprints=[], + depth=1, + max_loops=1, + model="claude-sonnet-4-20250514", + ) + ) + + iterator = _build_hold_back_iterator(MockAsyncStream(_build_tool_use_stream()), mock_handler) + wire = b"".join([chunk async for chunk in iterator if chunk != STREAM_SSE_KEEPALIVE_PING_BYTES]).decode() + + error_messages = [json.loads(data)["error"]["message"] for event, data in _sse_frames(wire) if event == "error"] + assert len(error_messages) == 1 + assert "max_agentic_loops" in error_messages[0] + assert "retry" not in error_messages[0].lower() + assert "litellm_content_retrieve" not in wire @pytest.mark.asyncio async def test_should_error_instead_of_replaying_when_no_hook_fires_on_tool_use(self): diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 01760572f33..dd0e5762845 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5,7 +5,7 @@ import json import logging import threading import time -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, Mock, patch @@ -68,6 +68,7 @@ from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, ) +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( @@ -3614,20 +3615,40 @@ class TestServerFulfilledToolsInRequest: == frozenset() ) - def test_interception_callbacks_declare_their_retrieval_tools(self): + @pytest.mark.parametrize("hook", ["compression", "headroom", "compresr"]) + def test_a_retrieval_hook_holds_back_a_stream_carrying_its_retrieve_tool(self, hook): from litellm.integrations.compression_interception.handler import ( LITELLM_CONTENT_RETRIEVE_TOOL_NAME, CompressionInterceptionLogger, ) + from litellm.proxy.guardrails.guardrail_hooks.compresr.compresr import ( + COMPRESR_RETRIEVE_TOOL_NAME, + CompresrGuardrail, + ) from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( HEADROOM_RETRIEVE_TOOL_NAME, HeadroomGuardrail, ) - assert HeadroomGuardrail.server_fulfilled_tool_names == frozenset({HEADROOM_RETRIEVE_TOOL_NAME}) - assert CompressionInterceptionLogger.server_fulfilled_tool_names == frozenset( - {LITELLM_CONTENT_RETRIEVE_TOOL_NAME} - ) + callback, retrieve_tool_name = { + "compression": (CompressionInterceptionLogger(), LITELLM_CONTENT_RETRIEVE_TOOL_NAME), + "headroom": ( + HeadroomGuardrail(api_base="https://headroom.test", api_key="hr-test", guardrail_name="headroom"), + HEADROOM_RETRIEVE_TOOL_NAME, + ), + "compresr": ( + CompresrGuardrail(api_base="https://compresr.test", api_key="cr-test", guardrail_name="compresr"), + COMPRESR_RETRIEVE_TOOL_NAME, + ), + }[hook] + tools = [ + {"name": "Bash", "input_schema": {"type": "object"}}, + {"name": retrieve_tool_name, "input_schema": {"type": "object"}}, + ] + + assert BaseLLMHTTPHandler._server_fulfilled_tools_in_request( + logging_obj=self._logging_obj_with([callback]), tools=tools + ) == frozenset({retrieve_tool_name}) def _make_stub_direct_vector_store_config(response): @@ -5341,6 +5362,165 @@ async def test_responses_agentic_followup_sends_the_plans_request_param_over_a_s assert followup_calls[0]["prompt_cache_key"] == "from-plan-params" +class _PlannedChatLoop(CustomLogger): + def __init__(self, plan: AgenticLoopPlan, claimed_call_name: str = "litellm_web_search") -> None: + super().__init__() + self._plan = plan + self._claimed_call_name = claimed_call_name + self.logging_objs: list[object] = [] + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + claimed: Final = [ + {"id": call.id, "name": call.function.name} + for call in _tool_calls_of(response) + if call.function.name == self._claimed_call_name + ] + return True, {"tool_calls": claimed or [{"id": "call_search", "name": self._claimed_call_name}]} + + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> AgenticLoopPlan: + self.logging_objs.append(logging_obj) + return self._plan + + +class _RunOnlyChatLoop(CustomLogger): + def __init__(self, result: ModelResponse) -> None: + super().__init__() + self._result = result + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + return True, {"tool_calls": [{"id": "call_search", "name": "litellm_web_search"}]} + + async def async_run_chat_completion_agentic_loop( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> object: + return self._result + + +def _tool_calls_of(response: object) -> list: + if not isinstance(response, ModelResponse) or not response.choices: + return [] + return response.choices[0].message.tool_calls or [] + + +async def _run_chat_completion_agentic_hooks( + response: ModelResponse, *callbacks: CustomLogger, kwargs: Mapping[str, object] = MappingProxyType({}) +) -> object: + logging_obj: Final = Mock(litellm_call_id="call-1", dynamic_success_callbacks=list(callbacks), model_call_details={}) + return await BaseLLMHTTPHandler()._call_agentic_chat_completion_hooks( + response=response, + model="gpt-4o-mini", + messages=[{"role": "user", "content": "what is new in litellm?"}], + optional_params={}, + logging_obj=logging_obj, + stream=False, + custom_llm_provider="openai", + kwargs=dict(kwargs), + ) + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_hooks_surface_a_failed_follow_up_instead_of_the_tool_call_turn(): + rate_limited_follow_up: Final = AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[{"role": "user", "content": "here are the search results"}], + kwargs={"mock_response": "litellm.RateLimitError"}, + ), + ) + + with pytest.raises(litellm.RateLimitError): + await _run_chat_completion_agentic_hooks(ModelResponse(), _PlannedChatLoop(rate_limited_follow_up)) + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_hooks_return_the_tool_call_turn_when_the_plan_terminates(): + response: Final = ModelResponse() + + assert await _run_chat_completion_agentic_hooks(response, _PlannedChatLoop(AgenticLoopPlan(terminate=True))) is response + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_hooks_move_past_a_plan_that_declines_the_loop(): + override: Final = ModelResponse() + declined: Final = _PlannedChatLoop(AgenticLoopPlan(run_agentic_loop=False)) + overriding: Final = _PlannedChatLoop(AgenticLoopPlan(response_override=override)) + + assert await _run_chat_completion_agentic_hooks(ModelResponse(), declined, overriding) is override + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_hooks_return_a_run_only_loops_result(): + loop_result: Final = ModelResponse() + + assert await _run_chat_completion_agentic_hooks(ModelResponse(), _RunOnlyChatLoop(loop_result)) is loop_result + + +@pytest.mark.asyncio +async def test_chat_completion_agentic_hooks_at_the_cap_drop_every_hooks_internal_call_and_keep_the_clients(): + response: Final = ModelResponse( + choices=[ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": f"call_{name}", "type": "function", "function": {"name": name, "arguments": "{}"}} + for name in ("litellm_web_search", "litellm_code_execution", "lookup") + ], + }, + } + ] + ) + searching: Final = _PlannedChatLoop(AgenticLoopPlan(run_agentic_loop=True)) + running_code: Final = _PlannedChatLoop( + AgenticLoopPlan(run_agentic_loop=True), claimed_call_name="litellm_code_execution" + ) + + result: Final = await _run_chat_completion_agentic_hooks( + response, searching, running_code, kwargs={"_agentic_loop_depth": 1, "max_agentic_loops": 1} + ) + + assert [call.id for call in _tool_calls_of(result)] == ["call_lookup"] + + @pytest.mark.asyncio async def test_chat_completion_agentic_followup_does_not_repeat_request_params_from_plan_kwargs(monkeypatch): """A plan whose kwargs repeat a request param, or the explicitly passed model, must not crash the chat follow-up with a duplicate keyword""" @@ -5364,17 +5544,17 @@ async def test_chat_completion_agentic_followup_does_not_repeat_request_params_f ), ) - response: Final = await BaseLLMHTTPHandler()._execute_chat_completion_agentic_plan( - plan=plan, + response: Final = await BaseLLMHTTPHandler()._call_agentic_chat_completion_hooks( + response=ModelResponse(), model="gpt-5", messages=[{"role": "user", "content": "x"}], optional_params={"temperature": 0.2}, - kwargs=dict(request_kwargs), + logging_obj=Mock( + litellm_call_id="call-1", dynamic_success_callbacks=[_PlannedChatLoop(plan)], model_call_details={} + ), + stream=False, custom_llm_provider="openai", - depth=0, - max_loops=3, - fingerprints=[], - fingerprint="fp", + kwargs=dict(request_kwargs), ) assert response == "followup-response" @@ -5384,6 +5564,138 @@ async def test_chat_completion_agentic_followup_does_not_repeat_request_params_f assert followup_calls[0]["model"] == "openai/gpt-5" +def _a_search_tool_call_turn(call_number: int) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": f"chatcmpl-{call_number}", + "object": "chat.completion", + "created": 0, + "model": "served-model", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": f"call-{call_number}", + "type": "function", + "function": {"name": "litellm_web_search", "arguments": '{"query": "weather"}'}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + }, + ) + + +class _SearchAgainLoop(CustomLogger): + def __init__(self, client: AsyncHTTPHandler) -> None: + super().__init__() + self._client = client + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + searches: Final = [ + {"id": call.id, "name": call.function.name} + for call in response.choices[0].message.tool_calls or [] + if call.function.name == "litellm_web_search" + ] + return bool(searches), {"tool_calls": searches} + + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> AgenticLoopPlan: + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[*messages, {"role": "user", "content": "here are the search results"}], + kwargs={"client": self._client}, + ), + ) + + +@pytest.mark.asyncio +async def test_hosted_vllm_bills_the_first_call_on_the_failure_log_when_the_follow_up_fails(monkeypatch): + loop: Final = _PlannedChatLoop( + AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[{"role": "user", "content": "here are the search results"}], + kwargs={"mock_response": "litellm.RateLimitError"}, + ), + ) + ) + monkeypatch.setattr(litellm, "callbacks", [loop]) + + with pytest.raises(litellm.RateLimitError): + await litellm.acompletion( + model="hosted_vllm/served-model", + api_base="https://fake-vllm.test/v1", + api_key="transport-only", + client=AsyncHTTPHandler(transport=httpx.MockTransport(lambda request: _a_search_tool_call_turn(1))), + messages=[{"role": "user", "content": "search the web for the weather"}], + input_cost_per_token=0.001, + output_cost_per_token=0.002, + num_retries=0, + max_retries=0, + ) + + (logging_obj,) = loop.logging_objs + assert logging_obj.model_call_details["combined_usage_object"].total_tokens == 8 + assert logging_obj.model_call_details["response_cost"] == pytest.approx(5 * 0.001 + 3 * 0.002) + + +@pytest.mark.asyncio +async def test_hosted_vllm_ends_the_turn_at_max_agentic_loops_without_the_internal_search_call(monkeypatch): + upstream_calls: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + + def always_search(request: httpx.Request) -> httpx.Response: + upstream_calls.put_nowait(request) + return _a_search_tool_call_turn(upstream_calls.qsize()) + + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(always_search)) + monkeypatch.setattr(litellm, "callbacks", [_SearchAgainLoop(client)]) + + response: Final = await asyncio.wait_for( + litellm.acompletion( + model="hosted_vllm/served-model", + api_base="https://fake-vllm.test/v1", + api_key="transport-only", + client=client, + messages=[{"role": "user", "content": "search the web for the weather"}], + max_agentic_loops=1, + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + assert upstream_calls.qsize() == 2 + assert not response.choices[0].message.tool_calls + + _UPSTREAM_HEADERS: Final = {"x-request-id": "req_upstream", "x-ratelimit-remaining-requests": "41"} @@ -6016,63 +6328,6 @@ async def _search_through(transport: httpx.MockTransport, is_async: bool) -> obj return await litellm.vector_stores.asearch(**search_kwargs, client=client) -class TestAgenticFollowUpKeepsTheProviderPrefix: - @staticmethod - async def _run_followup( - model: str, custom_llm_provider: str, patched_model: str | None = None - ) -> ModelResponse: - from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler - from litellm.types.integrations.custom_logger import ( - AgenticLoopPlan, - AgenticLoopRequestPatch, - ) - - response: Final = await BaseLLMHTTPHandler()._execute_chat_completion_agentic_plan( - plan=AgenticLoopPlan( - run_agentic_loop=True, - request_patch=AgenticLoopRequestPatch( - messages=[{"role": "user", "content": "hi"}], model=patched_model - ), - ), - model=model, - messages=[{"role": "user", "content": "hi"}], - optional_params={"mock_response": "ok from followup"}, - kwargs={}, - custom_llm_provider=custom_llm_provider, - depth=0, - max_loops=2, - fingerprints=[], - fingerprint="fp", - ) - assert isinstance(response, ModelResponse) - return response - - @pytest.mark.asyncio - async def test_a_sub_path_model_still_resolves_a_provider(self) -> None: - response: Final = await self._run_followup("mantle/anthropic.claude-sonnet-5", "bedrock") - - assert response.choices[0].message.content == "ok from followup" - - @pytest.mark.asyncio - async def test_an_ordinary_model_still_resolves_a_provider(self) -> None: - response: Final = await self._run_followup("gpt-4o-mini", "openai") - - assert response.choices[0].message.content == "ok from followup" - - @pytest.mark.asyncio - async def test_a_hook_that_patches_another_providers_model_reaches_that_provider(self) -> None: - response: Final = await self._run_followup("gpt-4o", "openai", patched_model="anthropic/claude-sonnet-4-5") - - assert response.model == "claude-sonnet-4-5" - assert response._hidden_params["custom_llm_provider"] == "anthropic" - - @pytest.mark.asyncio - async def test_a_hook_echoing_an_org_style_model_stays_on_the_request_provider(self) -> None: - response: Final = await self._run_followup("openai/gpt-4o", "openrouter", patched_model="openai/gpt-4o") - - assert response._hidden_params["custom_llm_provider"] == "openrouter" - - @pytest.mark.asyncio @pytest.mark.parametrize("stream", [False, True]) @pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) diff --git a/tests/unit/llms/openai/test_openai.py b/tests/unit/llms/openai/test_openai.py index 77a7c3b83fd..f61cfba9a10 100644 --- a/tests/unit/llms/openai/test_openai.py +++ b/tests/unit/llms/openai/test_openai.py @@ -3,7 +3,7 @@ import gzip import importlib import json import os -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping, Sequence from types import ModuleType from typing import Final, Literal, Protocol from unittest.mock import MagicMock, Mock, patch @@ -19,6 +19,8 @@ from pydantic import BaseModel, ConfigDict, JsonValue import litellm from litellm import create_thread, get_thread from litellm.caching.caching import Cache, LiteLLMCacheType +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch from litellm.llms.openai.openai import ( AssistantEventHandler, AsyncAssistantEventHandler, @@ -225,6 +227,333 @@ async def test_acompletion_returns_prompt_cache_details_over_injected_transport( assert prompt_tokens_details.cached_tokens == 7 + +class _SearchingChatGate(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.runs = 0 + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + return True, {"tool_calls": [{"id": "call-1"}]} + + +class _FailingFollowUpChatLoop(_SearchingChatGate): + def __init__(self, error: Exception) -> None: + super().__init__() + self.error = error + self.logging_objs: list[object] = [] + + async def async_run_chat_completion_agentic_loop( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> object: + self.runs += 1 + self.logging_objs.append(logging_obj) + raise self.error + + +class _RaisingAgenticGate(CustomLogger): + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + raise ValueError("gate bug") + + +class _AnsweringChatLoop(_SearchingChatGate): + async def async_run_chat_completion_agentic_loop( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> object: + self.runs += 1 + return litellm.ModelResponse( + choices=[ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "after the search"}} + ] + ) + + +def _respond_with_a_search_tool_call(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "chatcmpl-search", + "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-1", + "type": "function", + "function": {"name": "litellm_web_search", "arguments": '{"query": "weather"}'}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, + }, + ) + + +async def _acompletion_through(http_client: httpx.AsyncClient) -> object: + return await asyncio.wait_for( + litellm.acompletion( + model="openai/gpt-5.6", + api_key="transport-only", + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + messages=[{"role": "user", "content": "search the web for the weather"}], + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + +@pytest.mark.asyncio +async def test_acompletion_surfaces_a_failed_agentic_follow_up_instead_of_the_first_turn( + monkeypatch: pytest.MonkeyPatch, +): + follow_up_error: Final = litellm.RateLimitError( + message="follow-up rate limited", llm_provider="openai", model="gpt-5.6" + ) + loop: Final = _FailingFollowUpChatLoop(follow_up_error) + monkeypatch.setattr(litellm, "callbacks", [loop]) + + async with httpx.AsyncClient(transport=httpx.MockTransport(_respond_with_a_search_tool_call)) as http_client: + with pytest.raises(litellm.RateLimitError) as raised: + await _acompletion_through(http_client) + + assert raised.value.message == follow_up_error.message + assert loop.runs == 1 + + +@pytest.mark.asyncio +async def test_a_failed_agentic_follow_up_bills_the_first_call_on_the_failure_log(monkeypatch: pytest.MonkeyPatch): + loop: Final = _FailingFollowUpChatLoop( + litellm.RateLimitError(message="follow-up rate limited", llm_provider="openai", model="gpt-5.6") + ) + monkeypatch.setattr(litellm, "callbacks", [loop]) + + async with httpx.AsyncClient(transport=httpx.MockTransport(_respond_with_a_search_tool_call)) as http_client: + with pytest.raises(litellm.RateLimitError): + await _acompletion_through(http_client) + + (logging_obj,) = loop.logging_objs + price: Final = litellm.model_cost["gpt-5.6"] + assert logging_obj.model_call_details["combined_usage_object"].total_tokens == 8 + assert logging_obj.model_call_details["response_cost"] == pytest.approx( + 5 * price["input_cost_per_token"] + 3 * price["output_cost_per_token"] + ) + + +@pytest.mark.asyncio +async def test_acompletion_does_not_resend_the_first_call_when_the_agentic_follow_up_answers_422( + monkeypatch: pytest.MonkeyPatch, +): + follow_up_error: Final = litellm.UnprocessableEntityError( + message="follow-up rejected", + model="gpt-5.6", + llm_provider="openai", + response=httpx.Response(422, request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")), + ) + loop: Final = _FailingFollowUpChatLoop(follow_up_error) + monkeypatch.setattr(litellm, "callbacks", [loop]) + monkeypatch.setattr(litellm, "drop_params", True) + + async with httpx.AsyncClient(transport=httpx.MockTransport(_respond_with_a_search_tool_call)) as http_client: + with pytest.raises(litellm.UnprocessableEntityError) as raised: + await _acompletion_through(http_client) + + assert raised.value.message == follow_up_error.message + assert loop.runs == 1 + + +def _rejecting_the_user_param(outbound: asyncio.Queue[dict[str, object]]) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + body: Final = json.loads(request.content) + outbound.put_nowait(body) + if "user" in body: + rejection: Final = json.dumps({"detail": [{"loc": ["body", "user"], "msg": "extra fields not permitted"}]}) + return httpx.Response(422, json={"error": {"message": rejection}}) + return httpx.Response( + 200, + json={ + "id": "chatcmpl-resent", + "object": "chat.completion", + "created": 0, + "model": "gpt-5.6", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "done"}}], + }, + ) + + return httpx.MockTransport(respond) + + +async def _acompletion_with_a_user_param(http_client: httpx.AsyncClient) -> object: + return await asyncio.wait_for( + litellm.acompletion( + model="openai/gpt-5.6", + api_key="transport-only", + client=AsyncOpenAI(api_key="transport-only", http_client=http_client), + messages=[{"role": "user", "content": "hi"}], + user="end-user-1", + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + +@pytest.mark.asyncio +async def test_acompletion_resends_without_the_param_a_422_rejects_when_drop_params_is_set( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "drop_params", True) + outbound: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + + async with httpx.AsyncClient(transport=_rejecting_the_user_param(outbound)) as http_client: + response: Final = await _acompletion_with_a_user_param(http_client) + + assert response.choices[0].message.content == "done" + assert outbound.get_nowait()["user"] == "end-user-1" + assert "user" not in outbound.get_nowait() + assert outbound.empty() + + +@pytest.mark.asyncio +async def test_acompletion_raises_a_422_without_resending_when_drop_params_is_off(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + outbound: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + + async with httpx.AsyncClient(transport=_rejecting_the_user_param(outbound)) as http_client: + with pytest.raises(litellm.BadRequestError, match="extra fields not permitted"): + await _acompletion_with_a_user_param(http_client) + + assert outbound.get_nowait()["user"] == "end-user-1" + assert outbound.empty() + + +@pytest.mark.asyncio +async def test_acompletion_skips_a_failing_agentic_gate_and_runs_the_next_hook(monkeypatch: pytest.MonkeyPatch): + loop: Final = _AnsweringChatLoop() + monkeypatch.setattr(litellm, "callbacks", [_RaisingAgenticGate(), loop]) + + async with httpx.AsyncClient(transport=httpx.MockTransport(_respond_with_a_search_tool_call)) as http_client: + response = await _acompletion_through(http_client) + + assert loop.runs == 1 + assert response.choices[0].message.content == "after the search" + + +class _SearchAgainLoop(CustomLogger): + def __init__(self, client: AsyncOpenAI) -> None: + super().__init__() + self._client = client + + async def async_should_run_chat_completion_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + searches: Final = [ + {"id": call.id, "name": call.function.name} + for call in response.choices[0].message.tool_calls or [] + if call.function.name == "litellm_web_search" + ] + return bool(searches), {"tool_calls": searches} + + async def async_build_chat_completion_agentic_loop_plan( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + optional_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> AgenticLoopPlan: + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[*messages, {"role": "user", "content": "here are the search results"}], + kwargs={"client": self._client}, + ), + ) + + +@pytest.mark.asyncio +async def test_acompletion_ends_the_turn_at_max_agentic_loops_without_the_internal_search_call( + monkeypatch: pytest.MonkeyPatch, +): + upstream_calls: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + + def always_search(request: httpx.Request) -> httpx.Response: + upstream_calls.put_nowait(request) + return _respond_with_a_search_tool_call(request) + + async with httpx.AsyncClient(transport=httpx.MockTransport(always_search)) as http_client: + client: Final = AsyncOpenAI(api_key="transport-only", http_client=http_client) + monkeypatch.setattr(litellm, "callbacks", [_SearchAgainLoop(client)]) + response: Final = await asyncio.wait_for( + litellm.acompletion( + model="openai/gpt-5.6", + api_key="transport-only", + client=client, + messages=[{"role": "user", "content": "search the web for the weather"}], + max_agentic_loops=1, + num_retries=0, + max_retries=0, + ), + timeout=10, + ) + + assert upstream_calls.qsize() == 2 + assert not response.choices[0].message.tool_calls + + _REASONING_TOOL_CALL_REPLY: Final = { "id": "chatcmpl-fake-stream", "object": "chat.completion", diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index dccd1e877fc..80e7e793f6e 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -2,6 +2,7 @@ Test for anthropic_endpoints/endpoints.py, focusing on handling dictionary objects in streaming responses """ +import datetime import json import logging import unittest @@ -249,6 +250,86 @@ class TestBlockedResponseUsage: assert response["content"] == [{"type": "text", "text": _BLOCK_TEXT}] +class TestAnthropicResponseReportsWhatAFailedAgenticFollowUpBills: + async def _error_response_cost(self, raise_the_follow_up_error: bool) -> float: + import litellm + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + from litellm.litellm_core_utils.chat_completion_agentic_loop import bill_first_round_on_failure + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + from litellm.types.utils import ModelResponse, Usage + + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alert_types = [] + follow_up_error: Final = litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4o") + + async def failing_follow_up() -> ModelResponse: + raise follow_up_error + + async def fail_after_the_first_call(self: ProxyBaseLLMRequestProcessing, **kwargs: object) -> None: + _, logging_obj = await self.common_processing_pre_call_logic( + request=kwargs["request"], + general_settings={}, + user_api_key_dict=kwargs["user_api_key_dict"], + proxy_logging_obj=proxy_logging, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="anthropic_messages", + ) + logging_obj.model_call_details.update(first_api_call_start_time=datetime.datetime.now()) + with pytest.raises(litellm.RateLimitError): + await bill_first_round_on_failure( + failing_follow_up(), + logging_obj=logging_obj, + response=ModelResponse( + model="gpt-4o", usage=Usage(prompt_tokens=559, completion_tokens=113, total_tokens=672) + ), + stream=False, + ) + raise ( + follow_up_error + if raise_the_follow_up_error + else litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4o") + ) + + request: Final = MagicMock(spec=Request) + request.headers = {} + with ( + patch.object( + ep, + "read_request_body", + new=AsyncMock( + return_value={ + "model": "gpt-4o", + "max_tokens": 64, + "messages": [{"role": "user", "content": "what is 6*7?"}], + } + ), + ), + patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fail_after_the_first_call), + patch.object(proxy_server, "proxy_logging_obj", proxy_logging), + ): + response: Final = await ep.anthropic_response( + fastapi_response=Response(), + request=request, + user_api_key_dict=UserAPIKeyAuth(request_route="/v1/messages"), + ) + + assert response.status_code == 429 + return float(response.headers["x-litellm-response-cost"]) + + @pytest.mark.asyncio + async def test_the_error_response_reports_the_served_first_call_its_failure_row_bills(self) -> None: + import litellm + + price: Final = litellm.model_cost["gpt-4o"] + + assert await self._error_response_cost(raise_the_follow_up_error=True) == pytest.approx( + 559 * price["input_cost_per_token"] + 113 * price["output_cost_per_token"] + ) + assert await self._error_response_cost(raise_the_follow_up_error=False) == 0 + + class TestProxyExceptionAnthropicEnvelope: @pytest.mark.asyncio async def test_anthropic_response_maps_proxy_exception_to_anthropic_envelope(self): diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index d5771dff626..4e113ce03a4 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -4418,6 +4418,81 @@ def test_get_logging_payload_uses_recovered_combined_usage_on_failure(): assert payload["total_tokens"] == 31 +def _fallback_after_a_failed_follow_up( + cache_hit: bool, request_params: Mapping[str, object] = MappingProxyType({}) +): + from litellm.litellm_core_utils.hidden_params import set_hidden_param + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import ModelResponse, Usage + + now: Final = datetime.datetime.now(timezone.utc) + logger: Final = Logging( + model="fallback-model", + messages=[], + stream=False, + call_type="acompletion", + start_time=now, + litellm_call_id="caller", + function_id="function", + ) + logger.update_environment_variables( + model="fallback-model", + custom_llm_provider="openai", + litellm_params={"metadata": {"user_api_key": "sk-test"}}, + optional_params=dict(request_params), + ) + fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4)) + set_hidden_param(fallback, "response_cost", 0.1) + logger.record_served_rounds_for_failure(((Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4),)) + logger.carry_failed_attempt_usage() + logger.model_call_details["cache_hit"] = cache_hit + prepared: Final = logger.prepare_response_cost(fallback, now) + assert prepared is not None + return logger, fallback, prepared, now + + +@pytest.mark.parametrize(("cache_hit", "own_cost"), [(False, 0.1), (True, 0.0)], ids=["fresh", "cached"]) +def test_success_row_bills_a_failed_follow_ups_cost_and_tokens(cache_hit: bool, own_cost: float): + logger, fallback, _, now = _fallback_after_a_failed_follow_up(cache_hit) + + payload: Final = get_logging_payload( + kwargs=logger.model_call_details, response_obj=fallback, start_time=now, end_time=now + ) + + assert payload["spend"] == pytest.approx(own_cost + 0.4) + assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (3 + 5, 1 + 2, 4 + 7) + + +@pytest.mark.parametrize(("cache_hit", "own_cost"), [(False, 0.1), (True, 0.0)], ids=["fresh", "cached"]) +def test_request_params_cannot_change_what_a_failed_follow_up_bills(cache_hit: bool, own_cost: float): + forged_attempt: Final = [{"prompt_tokens": 1000, "completion_tokens": 1000, "total_tokens": 2000}, -100.0] + logger, fallback, _, now = _fallback_after_a_failed_follow_up( + cache_hit, + { + "stash_billed_by_continuation": True, + "failed_attempts_owed_on_success": [forged_attempt], + "failed_attempts_billed_on_success": [forged_attempt], + }, + ) + + payload: Final = get_logging_payload( + kwargs=logger.model_call_details, response_obj=fallback, start_time=now, end_time=now + ) + + assert payload["spend"] == pytest.approx(own_cost + 0.4) + assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (8, 3, 11) + + +def test_cached_fallback_cost_breakdown_adds_up_to_the_billed_failed_follow_up(): + _, _, prepared, _ = _fallback_after_a_failed_follow_up(cache_hit=True) + + breakdown: Final = prepared.payload["cost_breakdown"] + + assert breakdown is not None + assert breakdown["total_cost"] == pytest.approx(prepared.payload["response_cost"]) == pytest.approx(0.4) + assert breakdown["additional_costs"]["failed_attempts_cost"] == pytest.approx(0.4) + + def test_get_logging_payload_failure_without_recovered_usage_is_zero(): """A failure with no recovered usage keeps zero token counts, so the combined-usage override never invents tokens for ordinary failures. diff --git a/tests/unit/proxy/spend_tracking/test_usage_capture.py b/tests/unit/proxy/spend_tracking/test_usage_capture.py index 32710e594ee..0c59473114a 100644 --- a/tests/unit/proxy/spend_tracking/test_usage_capture.py +++ b/tests/unit/proxy/spend_tracking/test_usage_capture.py @@ -544,6 +544,32 @@ async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_succe assert successful["included_attempt_ids"] == [] assert successful["usage"]["total_tokens"] == 3 assert successful["charge_eligible"] is True + + +@pytest.mark.asyncio +async def test_success_that_bills_a_failed_follow_up_names_that_attempt_as_included() -> None: + from litellm.litellm_core_utils.hidden_params import set_hidden_param + from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime + from litellm.types.utils import ModelResponse, Usage + + database: Final = StoredUsage() + runtime: Final = UsageCaptureRuntime() + runtime.start(database.insert, UsageCaptureSettings()) + logger: Final = _logger() + fallback: Final = ModelResponse(id="fallback", usage=Usage(prompt_tokens=3, completion_tokens=1, total_tokens=4)) + set_hidden_param(fallback, "response_cost", 0.1) + try: + logger.record_served_rounds_for_failure(((Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.4),)) + logger.carry_failed_attempt_usage() + logger.prepare_response_cost(fallback, NOW) + logger.capture_usage_before_logging(fallback, NOW, NOW) + finally: + await runtime.stop(1) + envelopes: Final = {row.event.event_id: json.loads(row.envelope) for row in database.rows} + assert envelopes["caller-attempt-0"]["usage"]["total_tokens"] == 7 + assert envelopes["caller-attempt-0"]["charge_eligible"] is True + assert envelopes["fallback"]["usage"]["total_tokens"] == 4 + 7 + assert envelopes["fallback"]["included_attempt_ids"] == ["caller-attempt-0"] assert database.rows[2].event.outcome is Outcome.SUCCEEDED diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 41e7ecc9428..13f6f840cb7 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4293,6 +4293,62 @@ class TestHandleLLMApiExceptionPricesWhatTheServerDispatched: assert resent >= dispatched_after_trim +class TestHandleLLMApiExceptionReportsWhatAFailedAgenticFollowUpBills: + async def _error_response_cost(self, raise_the_follow_up_error: bool) -> float: + from litellm.litellm_core_utils.chat_completion_agentic_loop import bill_first_round_on_failure + from litellm.types.utils import ModelResponse, Usage + + processing_obj: Final = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o", "messages": [{"role": "user", "content": "what is 6*7?"}]} + ) + caller: Final = ProxyUserAPIKeyAuth(request_route="/v1/chat/completions") + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alert_types = [] + _, logging_obj = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=caller, + proxy_logging_obj=proxy_logging, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + ) + logging_obj.model_call_details.update(first_api_call_start_time=datetime.datetime.now()) + follow_up_error: Final = litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4o") + + async def failing_follow_up() -> ModelResponse: + raise follow_up_error + + with pytest.raises(litellm.RateLimitError): + await bill_first_round_on_failure( + failing_follow_up(), + logging_obj=logging_obj, + response=ModelResponse( + model="gpt-4o", usage=Usage(prompt_tokens=559, completion_tokens=113, total_tokens=672) + ), + stream=False, + ) + with pytest.raises(ProxyException) as raised: + await processing_obj.handle_llm_api_exception( + e=follow_up_error + if raise_the_follow_up_error + else litellm.RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4o"), + user_api_key_dict=caller, + proxy_logging_obj=proxy_logging, + ) + return float(raised.value.headers["x-litellm-response-cost"]) + + @pytest.mark.asyncio + async def test_the_error_response_reports_the_served_first_call_its_failure_row_bills(self) -> None: + price: Final = litellm.model_cost["gpt-4o"] + + assert await self._error_response_cost(raise_the_follow_up_error=True) == pytest.approx( + 559 * price["input_cost_per_token"] + 113 * price["output_cost_per_token"] + ) + assert await self._error_response_cost(raise_the_follow_up_error=False) == 0 + + @pytest.mark.asyncio class TestHandleLLMApiExceptionDictDetail: """ diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index db0db97cd88..d28619c0938 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -3814,6 +3814,56 @@ def test_a_retry_failure_bills_input_only_for_dispatched_attempts( assert lifted["response_cost"] == 0.25 +@pytest.mark.asyncio +async def test_an_agentic_follow_up_failure_row_bills_each_served_first_call_without_an_estimate() -> None: + import litellm + from litellm.litellm_core_utils.chat_completion_agentic_loop import bill_first_round_on_failure + from litellm.proxy.utils import _failure_fields_to_lift + from litellm.types.utils import ModelResponse, Usage + + logging_obj: Final = Logging( + model=FAILURE_USAGE_MODEL, + messages=ONE_USER_MESSAGE, + stream=False, + call_type="acompletion", + start_time=real_datetime.datetime.now(), + litellm_call_id="call-agentic-follow-up", + function_id="fn-agentic-follow-up", + ) + logging_obj.update_environment_variables( + model=FAILURE_USAGE_MODEL, user="", optional_params={}, litellm_params={"custom_llm_provider": "openai"} + ) + + async def failing_follow_up() -> ModelResponse: + raise litellm.RateLimitError(message="rate limited", llm_provider="openai", model=FAILURE_USAGE_MODEL) + + async def serve_first_call_then_fail_its_follow_up(usage: Usage) -> float: + logging_obj.pre_call(input=ONE_USER_MESSAGE, api_key="sk-test") + with pytest.raises(litellm.RateLimitError): + await bill_first_round_on_failure( + failing_follow_up(), + logging_obj=logging_obj, + response=ModelResponse(model=FAILURE_USAGE_MODEL, usage=usage), + stream=False, + ) + return logging_obj.model_call_details["response_cost"] + + primary_cost: Final = await serve_first_call_then_fail_its_follow_up( + Usage(prompt_tokens=559, completion_tokens=113, total_tokens=672) + ) + logging_obj.carry_failed_attempt_usage() + backup_cost: Final = await serve_first_call_then_fail_its_follow_up( + Usage(prompt_tokens=31, completion_tokens=5, total_tokens=36) + ) + + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}) + + usage: Final = lifted["combined_usage_object"] + assert (usage.prompt_tokens, usage.completion_tokens) == (559 + 31, 113 + 5) + assert primary_cost > 0 and backup_cost > 0 + assert lifted["response_cost"] == pytest.approx(primary_cost + backup_cost) + + @pytest.mark.parametrize( "model_call_details, dispatched", [ diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index ec0be896095..037004ab999 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -540,13 +540,9 @@ async def test_mcp_allowed_tools_filtering(): ] # First filter by allowed tools - filtered_tools_with_duplicates = ( - LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( - mcp_tools=mock_mcp_tools_with_duplicates, - mcp_tools_with_litellm_proxy=cast( - List[ToolParam], mcp_tool_config_with_duplicates - ), - ) + filtered_tools_with_duplicates = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( + mcp_tools=mock_mcp_tools_with_duplicates, + mcp_tools_with_litellm_proxy=cast(List[ToolParam], mcp_tool_config_with_duplicates), ) # Then deduplicate the filtered tools @@ -605,11 +601,9 @@ async def test_mcp_allowed_tools_filtering(): }, ] - filtered_tools_multiple = ( - LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( - mcp_tools=mock_mcp_tools_from_server, - mcp_tools_with_litellm_proxy=cast(List[ToolParam], multiple_mcp_configs), - ) + filtered_tools_multiple = LiteLLM_Proxy_MCP_Handler._filter_mcp_tools_by_allowed_tools( + mcp_tools=mock_mcp_tools_from_server, + mcp_tools_with_litellm_proxy=cast(List[ToolParam], multiple_mcp_configs), ) # Should return union of all allowed tools (3 unique tools) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3cc97f4507b..f815b107648 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -13,7 +13,7 @@ import os import sys import threading import warnings -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timedelta, timezone from pathlib import Path from types import CodeType, FrameType, MappingProxyType, SimpleNamespace @@ -43,6 +43,7 @@ from litellm.integrations.otel.plumbing.context import set_request_root_span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.anthropic.messages.native.streaming_iterator import closing_superseded_attempt from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( + AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES, SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, ) from litellm.llms.bedrock.common_utils import BedrockError @@ -82,6 +83,7 @@ from litellm.router_utils.fallback_event_handlers import ( from litellm.router_utils.prompt_caching_cache import PromptCachingCache from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.scheduler import FlowItem +from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.router_api import RoutingContext as RouterContext from litellm.types.router import ( @@ -3652,14 +3654,28 @@ class _CutSSE(httpx.AsyncByteStream): raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") -_BACKUP_SSE: Final = "".join( - f"data: {json.dumps(chunk)}\n\n" - for chunk in ( - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]}, - {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", - "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}}, +_BACKUP_SSE: Final = ( + "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [ + {"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None} + ], + }, + { + "id": "b", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}, + }, + ) ) ) + "data: [DONE]\n\n" @@ -3845,6 +3861,255 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t ) +_AGENTIC_FIRST_CALL_USAGE: Final = {"prompt_tokens": 50, "completion_tokens": 7, "total_tokens": 57} +_AGENTIC_FIRST_CALL_BODY: Final = { + "id": "chatcmpl-first", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_abc", "type": "function", "function": {"name": "litellm_code_execution", "arguments": "{}"}} + ], + }, + } + ], + "usage": _AGENTIC_FIRST_CALL_USAGE, +} + + +_CRASH_THE_PLAN: Final = "crash the plan" + + +class _ScriptedAgenticFollowUps(CustomLogger): + def __init__(self, follow_ups: tuple[str, ...]) -> None: + super().__init__() + self._follow_ups: Final = follow_ups + self.plans_built = 0 + self.shared_call_succeeded: Final = asyncio.Event() + + async def async_should_run_agentic_loop( + self, + response: object, + model: str, + messages: Sequence[Mapping[str, object]], + tools: Sequence[Mapping[str, object]] | None, + stream: bool, + custom_llm_provider: str, + kwargs: Mapping[str, object], + ) -> tuple[bool, dict[str, object]]: + assert isinstance(response, litellm.ModelResponse) + return bool(response.choices[0].message.tool_calls), {"tool_calls": [{"id": "call_abc"}]} + + async def async_build_agentic_loop_plan( + self, + tools: Mapping[str, object], + model: str, + messages: Sequence[Mapping[str, object]], + response: object, + anthropic_messages_provider_config: object, + anthropic_messages_optional_request_params: Mapping[str, object], + logging_obj: object, + stream: bool, + kwargs: Mapping[str, object], + ) -> AgenticLoopPlan: + follow_up: Final = self._follow_ups[min(self.plans_built, len(self._follow_ups) - 1)] + self.plans_built += 1 + if follow_up == _CRASH_THE_PLAN: + raise ValueError("the agentic plan build crashed") + return AgenticLoopPlan( + run_agentic_loop=True, + request_patch=AgenticLoopRequestPatch( + messages=[{"role": "user", "content": "what is 6*7?"}, {"role": "assistant", "content": "42"}], + kwargs={"mock_response": follow_up}, + ), + ) + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + if kwargs.get("litellm_call_id") == "shared-agentic-call": + self.shared_call_succeeded.set() + + +_AGENTIC_BACKUP_USAGE: Final = {"prompt_tokens": 20, "completion_tokens": 3, "total_tokens": 23} +_AGENTIC_BACKUP_BODY: Final = { + "id": "chatcmpl-backup", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "42"}}], + "usage": _AGENTIC_BACKUP_USAGE, +} +_AGENTIC_RATE_LIMITED: Final = { + "error": {"message": "Rate limit reached", "type": "requests", "code": "rate_limit_exceeded"} +} + + +async def _route_agentic_request( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + follow_ups: tuple[str, ...], + model: str = "m", + first_call_replies: tuple[httpx.Response, ...] = (), + fallbacks: tuple[str, ...] = (), +) -> tuple[object, LiteLLMLogging, respx.Route, _ScriptedAgenticFollowUps]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + hook: Final = _ScriptedAgenticFollowUps(follow_ups) + monkeypatch.setattr(litellm, "callbacks", [hook]) + first_calls: Final = ( + respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=list(first_call_replies)) + if first_call_replies + else respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_AGENTIC_FIRST_CALL_BODY) + ) + ) + respx_mock.post("https://backup.test/v1/chat/completions").mock( + return_value=httpx.Response(200, json=_AGENTIC_BACKUP_BODY) + ) + router: Final = Router( + model_list=[ + {"model_name": "m", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}}, + { + "model_name": "backup", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-test", + "api_base": "https://backup.test/v1", + }, + }, + ], + num_retries=2, + retry_after=0, + fallbacks=[{"m": list(fallbacks)}] if fallbacks else [], + ) + messages: Final = [{"role": "user", "content": "what is 6*7?"}] + shared_logging, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model=model, + messages=messages, + litellm_call_id="shared-agentic-call", + ) + try: + outcome: object = await router.acompletion( + model=model, messages=messages, litellm_call_id="shared-agentic-call", litellm_logging_obj=shared_logging + ) + except (litellm.RateLimitError, litellm.APIConnectionError) as error: + outcome = error + return outcome, shared_logging, first_calls, hook + + +def _gpt_4o_mini_cost(usage: Mapping[str, int]) -> float: + return sum( + litellm.cost_per_token( + model="gpt-4o-mini", prompt_tokens=usage["prompt_tokens"], completion_tokens=usage["completion_tokens"] + ) + ) + + +@pytest.mark.asyncio +async def test_agentic_request_whose_follow_up_fails_is_not_retried_and_bills_its_one_first_call( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + from litellm.proxy.utils import _failure_fields_to_lift + + outcome, shared_logging, first_calls, _ = await _route_agentic_request( + monkeypatch, respx_mock, ("litellm.RateLimitError",) + ) + + lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": shared_logging}) + assert isinstance(outcome, litellm.RateLimitError) + assert outcome.status_code == 429 + assert first_calls.call_count == 1 + assert lifted["combined_usage_object"].prompt_tokens == _AGENTIC_FIRST_CALL_USAGE["prompt_tokens"] + assert lifted["combined_usage_object"].completion_tokens == _AGENTIC_FIRST_CALL_USAGE["completion_tokens"] + assert lifted["response_cost"] == pytest.approx(_gpt_4o_mini_cost(_AGENTIC_FIRST_CALL_USAGE)) + + +@pytest.mark.asyncio +async def test_agentic_request_whose_plan_build_crashes_after_the_first_call_is_not_retried( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + outcome, _, first_calls, hook = await _route_agentic_request(monkeypatch, respx_mock, (_CRASH_THE_PLAN,)) + + assert isinstance(outcome, litellm.APIConnectionError) + assert first_calls.call_count == 1 + assert hook.plans_built == 1 + + +@pytest.mark.asyncio +async def test_agentic_request_whose_first_call_fails_is_still_retried( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + outcome, _, first_calls, hook = await _route_agentic_request( + monkeypatch, + respx_mock, + ("42",), + first_call_replies=( + httpx.Response(429, json=_AGENTIC_RATE_LIMITED), + httpx.Response(200, json=_AGENTIC_FIRST_CALL_BODY), + ), + ) + + assert isinstance(outcome, litellm.ModelResponse) + assert first_calls.call_count == 2 + assert hook.plans_built == 1 + + +async def _successful_agentic_request_billing( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + follow_ups: tuple[str, ...], + model: str, + fallbacks: tuple[str, ...] = (), +) -> tuple[Mapping[str, object], Mapping[str, str | None]]: + from litellm.proxy.common_request_processing import cost_headers + + outcome, shared_logging, _, hook = await _route_agentic_request( + monkeypatch, respx_mock, follow_ups, model=model, fallbacks=fallbacks + ) + await asyncio.wait_for(hook.shared_call_succeeded.wait(), timeout=10) + assert isinstance(outcome, litellm.ModelResponse) + success: Final = shared_logging.model_call_details["standard_logging_object"] + return success, cost_headers(success["response_cost"], shared_logging) + + +@pytest.mark.asyncio +async def test_fallback_after_a_failed_agentic_follow_up_bills_the_primarys_first_call_on_the_success_row( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + first_try, first_try_headers = await _successful_agentic_request_billing( + monkeypatch, respx_mock, ("42",), model="backup" + ) + recovered, recovered_headers = await _successful_agentic_request_billing( + monkeypatch, respx_mock, ("litellm.RateLimitError",), model="m", fallbacks=("backup",) + ) + + first_call_cost: Final = _gpt_4o_mini_cost(_AGENTIC_FIRST_CALL_USAGE) + assert first_try["response_cost"] == pytest.approx(_gpt_4o_mini_cost(_AGENTIC_BACKUP_USAGE)) + assert "failed_attempts_cost" not in first_try["cost_breakdown"].get("additional_costs", {}) + assert recovered["response_cost"] == pytest.approx(first_try["response_cost"] + first_call_cost) + assert recovered["prompt_tokens"] == first_try["prompt_tokens"] + _AGENTIC_FIRST_CALL_USAGE["prompt_tokens"] + assert ( + recovered["completion_tokens"] + == first_try["completion_tokens"] + _AGENTIC_FIRST_CALL_USAGE["completion_tokens"] + ) + assert recovered["cost_breakdown"]["total_cost"] == pytest.approx(recovered["response_cost"]) + assert recovered["cost_breakdown"]["additional_costs"]["failed_attempts_cost"] == pytest.approx(first_call_cost) + assert float(recovered_headers["x-litellm-response-cost"]) == pytest.approx(recovered["response_cost"]) + assert {key: value for key, value in recovered_headers.items() if key != "x-litellm-response-cost"} == { + key: value for key, value in first_try_headers.items() if key != "x-litellm-response-cost" + } + + def _responses_sse(*events: Mapping[str, object]) -> str: return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) @@ -13898,7 +14163,11 @@ def test_get_search_tool_support_reports_each_routable_deployment_of_the_group() @pytest.mark.parametrize("provider", ["chatgpt", "github_copilot"]) @pytest.mark.parametrize("explicit_provider", [False, True]) async def test_search_tool_catalog_lookup_needs_no_credentials( - provider: str, explicit_provider: bool, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter + provider: str, + explicit_provider: bool, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, ) -> None: monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path / "chatgpt")) monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path / "copilot")) @@ -17566,13 +17835,16 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool ) -> None: - models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ( - "anthropic/claude-opus-5", other_model + models: Final = (other_model, "anthropic/claude-opus-5") if reverse else ("anthropic/claude-opus-5", other_model) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "api_key": "fake-key"}, + } + for model in models + ] ) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "api_key": "fake-key"}, - } for model in models]) result: Final = router.get_model_group_info("fast-group") @@ -17930,10 +18202,7 @@ class _AnthropicMessagesFallbackByteStream: def _anthropic_messages_overloaded_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' - ) + return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' def _anthropic_messages_invalid_request_error_chunk() -> bytes: @@ -18749,6 +19018,7 @@ def test_anthropic_stream_forwards_ping_live_direct_call(): def test_anthropic_stream_error_is_gateway_verdict_direct_call(): assert _anthropic_stream_error_is_gateway_verdict(SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES) is True + assert _anthropic_stream_error_is_gateway_verdict(AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES) is True assert _anthropic_stream_error_is_gateway_verdict(_anthropic_messages_overloaded_error_chunk()) is False assert _anthropic_stream_error_is_gateway_verdict(_anthropic_messages_ping_chunk()) is False @@ -24683,16 +24953,25 @@ class TestAutoRouterTraceProvenance: monkeypatch.setattr(litellm, "_async_success_callback", [logger, CallbackComplete()]) classifier: Final = ObservedClassifier() - router: Final = Router(model_list=[ - {**self._marker("blue"), "litellm_params": { - "model": "auto_router/complexity_router", - "complexity_router_config": { - "classifier_type": "custom", "classifier_plugin": classifier, - "tiers": {"SIMPLE": "answer"}, "session_affinity": False, "deployment_affinity": False, + router: Final = Router( + model_list=[ + { + **self._marker("blue"), + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "classifier_type": "custom", + "classifier_plugin": classifier, + "tiers": {"SIMPLE": "answer"}, + "session_affinity": False, + "deployment_affinity": False, + }, + }, }, - }}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], plugins=[classifier]) + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + plugins=[classifier], + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale, "user_api_key_team_id": "team"}} with logger.tracer.start_as_current_span("request") as root: @@ -24777,20 +25056,22 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]" ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), - { - "model_name": "traced-router", - "model_info": {"id": "definition-semantic"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "tags": ["blue"], - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", + router: Final = Router( + model_list=[ + self._marker("blue"), + { + "model_name": "traced-router", + "model_info": {"id": "definition-semantic"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "tags": ["blue"], + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, }, - }, - ]) + ] + ) with logger.start_phase_span("route traced-router"): response: Final = await router.async_pre_routing_hook( model="traced-router", request_kwargs={"metadata": {"tags": ["blue"]}} @@ -24810,26 +25091,42 @@ class TestAutoRouterTraceProvenance: ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - respx_mock.post("https://embedding.example/v1/embeddings").respond(200, json={ - "object": "list", - "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - router: Final = Router(model_list=[ - { - "model_name": "internal-semantic-router", - "model_info": {"id": "semantic-definition", "team_id": "trace-team", "team_public_model_name": "public-router"}, - "litellm_params": { - "model": "auto_router/semantic-router", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - }, + respx_mock.post("https://embedding.example/v1/embeddings").respond( + 200, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, }, - {"model_name": "embed", "litellm_params": {"model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", "api_key": "test"}}, - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ]) + ) + router: Final = Router( + model_list=[ + { + "model_name": "internal-semantic-router", + "model_info": { + "id": "semantic-definition", + "team_id": "trace-team", + "team_public_model_name": "public-router", + }, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + }, + }, + { + "model_name": "embed", + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + }, + }, + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ] + ) request: Final = {"metadata": {"user_api_key_team_id": "trace-team"}} with logger.tracer.start_as_current_span("request"): await router.async_get_available_deployment(