scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(agentic)!: surface a failed agentic follow-up call instead of the pre-loop tool call (#49227)

Browse source 
* fix(agentic): surface a failed agentic follow-up call instead of the pre-loop tool call

* fix(agentic): surface failed follow-ups from the OpenAI chat dispatcher and update pinned integration tests

* fix(agentic): end a capped agentic turn on chat and responses instead of a retried 500

* fix(openai): run agentic hooks after the first call's error handling so a follow-up error keeps its own type and message

* fix(agentic loop): end capped turns on every chat path and bill the first call when a follow-up fails

A turn refused at max_agentic_loops now drops the internal tool calls of every hook that claims calls in it, while a call the client declared survives. The hosted_vllm and openai chat paths now carry the loop depth, fingerprints and max_agentic_loops into the follow-up, so they stop at the cap instead of looping forever.

A failed follow-up call now records the first call's usage and cost on the request's failure log before the error reaches the client. A code interpreter sandbox that is missing or fails becomes the tool call's output for the model to handle, so the turn completes without a retry.

On native streaming /v1/messages a hook or follow-up failure now reaches the client as an Anthropic error event with the mapped type, and compresr declares its retrieve tool so the stream holds it back. Comments and docstrings the branch added against the repo rule are removed

* fix(agentic loop): bill retried attempts and clean up what failed or capped turns return

The router's retry loop now carries each failed attempt's paid first call into the request's billing, so a request whose follow-up keeps failing bills every first call on its failure row, and one that recovers bills the failed attempt's first call on its success row. A failure row no longer adds unpriced estimated tokens on top of attempts that were billed, so its tokens match its spend

A sandbox failure now answers the model with a fixed sentence and marks the code_interpreter_call as failed, without the sandbox's error text. Code interpreter on /v1/responses with Claude now runs its loop, because the bridged inner call no longer strips the hook's own markers

Dropped internal tool calls leave the kept calls numbered from 0, also on a replayed stream. A capped native /v1/messages stream now says the loop hit its limit and that max_agentic_loops raises it, instead of asking the client to retry

The agentic refusal and hook error logs no longer include the model or the exception text, and the agentic loop helpers take typed messages, params, and read-only fingerprints

* fix(agentic loop): keep a recovered retry's cost breakdown on the response it returns

A retried agentic request that recovers no longer adds the failed attempt's first call to its success row. The proxy keeps the follow-up's own row for that request id, so the added spend never showed up in the stored row, and that duplicate row is tracked separately. When every attempt fails, the failure row still bills each attempt's first call

Pricing a first call for the failure row no longer replaces the request's cost breakdown. Before, a recovered retry answered with x-litellm-response-cost-original, -input, and -output from the failed attempt's first call. Now its headers and logged breakdown match a request that succeeds on its first try

* fix(agentic): stop retrying a request whose agentic follow-up failed and bill every round it served

An error raised by an agentic follow-up after the first call was served now carries a typed mark, so the
router's retry loops, the mid-stream retry, and litellm's num_retries wrapper raise it as is instead of
rerunning the whole request. A failure of the first call itself still retries, and fallbacks still run.

The failure row bills every round the loop served before the failure, at any depth, and a later attempt
that succeeds bills what the failed attempts paid on its success row, with the cost breakdown and the
response cost headers adding up. The error response's x-litellm-response-cost reports the same amount
its failure row bills. Failure rows outside agentic hooks keep the prompt-token estimate they had before.

A refused tool call now matches by id among the refused ids, or by name among refused calls that carry
no id, so a name-only refusal next to an id refusal is no longer kept.

* fix(code_interpreter): delete a session sandbox whose code run failed so the next call starts a fresh one

* fix(code_interpreter): delete only the session sandbox that failed, never its replacement

* test(code_interpreter): type the session sandbox test helpers

* fix(proxy): report what a failed agentic follow-up's failure row bills on /v1/messages errors

* fix(agentic): keep the no-retry mark through error mapping and bill only failed follow-ups on success

An agentic loop error that is not a provider error, such as the streaming chat cap or a crash in a hook's
plan build, now keeps its no-retry mark when litellm maps it to an APIConnectionError, so the router no
longer reruns every paid round. The mark is stored on the exception, so a built-in error can carry it too.

A crash in a hook's plan build on /v1/messages and /v1/responses bills the first call on its failure row.

A recovered attempt adds a failed attempt's usage to its success row only when that attempt was a failed
agentic follow-up. Other failed attempts, such as a dropped stream, log their own row as on main.

A code interpreter sandbox that fails while running the code now says the code may have run but its
result is lost, instead of saying the code did not run.

* test(responses): bill a retried MCP request each attempt on its own row again

Round 4 stopped adding a non-agentic failed attempt to the success row, so this test goes back to main's assertions: the success row bills only its own rounds and the failed attempt's first round stays on its own failure row

* fix(spend): name a failed follow-up attempt the success record already bills

With the usage meter on, a failed agentic follow-up attempt gets its own charge-eligible record, and the request's success record also carries that attempt's usage and cost once the success row bills it. The success record now lists the attempt's record ids in included_attempt_ids, so a reader can tell the success record already covers it. Nothing changes when the meter is off

* fix(spend): count a billed failed follow-up's tokens on the success row

* fix(spend): bill a failed follow-up on a fallback served from cache

* fix(spend): show a billed failed follow-up in a cached fallback's cost breakdown

* fix(spend): keep a failed follow-up's billing state off the request-merged call details

* fix(spend): record the cost breakdown of a successful agentic follow-up

The chat agentic loop returned the follow-up response with the cost its own internal call had already stamped, so the request's logging reused that cost and never recorded cost_breakdown on the spend row. The loop now drops that stamped cost unless the follow-up was served from the cache, and the request prices the response itself

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-03 18:00:33 -08:00 • committed by GitHub
parent d5ca88900c
commit b9ad61d4f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 3471 additions and 806 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

128
litellm/integrations/code_interpreter_interception/handler.py
View file

@ -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,

36
litellm/litellm_core_utils/agentic_followup_failure.py Normal file
View file

@ -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))

107
litellm/litellm_core_utils/agentic_loop_refusal.py Normal file
View file

@ -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]}
)

371
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -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

3
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -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

4
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -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)
)

137
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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
),

23
litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py
View file

@ -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(

374
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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:

125
litellm/llms/openai/openai.py
View file

@ -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,

46
litellm/main.py
View file

@ -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",

6
litellm/proxy/anthropic_endpoints/endpoints.py
View file

@ -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),

7
litellm/proxy/common_request_processing.py
View file

@ -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

4
litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py
View file

@ -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,

40
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

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

8
litellm/proxy/spend_tracking/spend_tracking_utils.py
View file

@ -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)

10
litellm/router_python_core.py
View file

@ -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

9
litellm/types/integrations/custom_logger.py
View file

@ -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.

15
litellm/utils.py
View file

@ -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

72
tests/integration/streaming/test_interception_converted_stream_options.py
View file

@ -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)}))

237
tests/unit/integrations/code_interpreter_interception/test_handler.py
View file

@ -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"]

309
tests/unit/integrations/websearch_interception/test_websearch_agentic_loop_cap.py
View file

@ -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

3
tests/unit/integrations/websearch_interception/test_websearch_thinking_constraint.py
View file

@ -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):

630
tests/unit/litellm_core_utils/test_chat_completion_agentic_loop.py
View file

@ -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"

42
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -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])

80
tests/unit/llms/anthropic/messages/native/test_agentic_streaming_iterator.py
View file

@ -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):

395
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -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"])

331
tests/unit/llms/openai/test_openai.py
View file

@ -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",

81
tests/unit/proxy/anthropic_endpoints/test_endpoints.py
View file

@ -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):

75
tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -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.

26
tests/unit/proxy/spend_tracking/test_usage_capture.py
View file

@ -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

56
tests/unit/proxy/test_common_request_processing.py
View file

@ -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:
"""

50
tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py
View file

@ -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",
[

18
tests/unit/responses/mcp/test_aresponses_api_with_mcp.py
View file

@ -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)

419
tests/unit/test_router/test_router.py
View file

@ -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(