Merge branch 'litellm_internal_staging' into litellm_non-root-dockerfile-optimization-31b6

This commit is contained in:
Yuneng Jiang 2026-04-20 16:42:32 -07:00
commit 0ecc89c355
No known key found for this signature in database
75 changed files with 6621 additions and 597 deletions

View file

@ -439,7 +439,14 @@ jobs:
auth:
username: ${DOCKERHUB_USERNAME}
password: ${DOCKERHUB_PASSWORD}
- image: cimg/postgres:16.0
environment:
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
POSTGRES_DB: litellm_test
working_directory: ~/project
environment:
DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test"
steps:
- checkout
@ -463,12 +470,14 @@ jobs:
paths:
- ./.venv
key: v2-dependencies-{{ checksum "uv.lock" }}-{{ checksum ".circleci/config.yml" }}
- wait_for_service:
url: tcp://localhost:5432
timeout: "60"
- run:
name: Run prisma ./docker/entrypoint.sh
name: Seed DB schema via prisma db push
command: |
set +e
chmod +x docker/entrypoint.sh
./docker/entrypoint.sh
uv run --no-sync litellm --skip_server_startup --use_prisma_db_push
set -e
- run:
name: Generate Prisma Client

View file

@ -10,6 +10,7 @@ Supported Providers:
- Vertex AI (`vertex_ai/`, `vertex_ai_beta/`)
- Bedrock (`bedrock/`, `bedrock/invoke/`, `bedrock/converse`) ([All models bedrock supports prompt caching on](https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-caching.html))
- Deepseek API (`deepseek/`)
- xAI (`xai/`)
For the supported providers, LiteLLM follows the OpenAI prompt caching usage object format:

View file

@ -8,6 +8,7 @@ The function keeps high-relevance and recent context, replaces low-relevance con
```python
import litellm
from litellm.types.utils import CallTypes
messages = [
{"role": "system", "content": "You are a coding assistant."},
@ -19,6 +20,7 @@ messages = [
compressed = litellm.compress(
messages=messages,
model="gpt-4o",
call_type=CallTypes.completion,
compression_trigger=1000,
compression_target=500,
)
@ -45,6 +47,7 @@ response = litellm.completion(
- `messages` (`List[dict]`, required): input conversation messages
- `model` (`str`, required): model name used for token counting
- `call_type` (`CallTypes`, default `CallTypes.completion`): the LiteLLM call type whose message schema these messages follow. Supported values: `CallTypes.completion` / `CallTypes.acompletion` (OpenAI chat-completions shape) and `CallTypes.anthropic_messages` (Anthropic Messages shape)
- `compression_trigger` (`int`, default `200000`): compress only if input token count exceeds this
- `compression_target` (`Optional[int]`, default `70% of compression_trigger`): desired post-compression token budget
- `embedding_model` (`Optional[str]`): if set, combines BM25 + embedding relevance scoring
@ -70,6 +73,28 @@ args = json.loads(tool_call.function.arguments)
full_content = compressed["cache"][args["key"]]
```
## Server-side Callback Loop (`/v1/messages`)
You can enable callback-based compression interception to make retrieval loops
transparent for Anthropic Messages calls:
```yaml
litellm_settings:
callbacks: ["compression_interception"]
compression_interception_params:
enabled: true
compression_trigger: 10000
compression_target: 7000
```
With this enabled, LiteLLM runs the following server-side flow:
1. Compresses inbound messages before the first provider call.
2. Injects the `litellm_content_retrieve` tool.
3. Detects retrieval `tool_use` blocks in the model response.
4. Resolves retrieval keys from the compression cache.
5. Reruns the model via agentic loop and returns the final answer.
## Performance
Benchmarked on [SWE-bench Lite](https://huggingface.co/datasets/princeton-nlp/SWE-bench_Lite_bm25_27K) (real GitHub issues with ~27k tokens of BM25-retrieved repo context per problem).

View file

@ -60,3 +60,44 @@ curl http://localhost:4000/chat/completions \
## Supported features
Scaleway provider supports all features in [Generative APIs reference documentation ↗](https://www.scaleway.com/en/developers/api/generative-apis/), such as streaming, structured outputs and tool calling.
## Audio transcription
Scaleway's `/audio/transcriptions` endpoint is OpenAI-compatible and works with Whisper models.
### Python SDK
```python
import os
from litellm import transcription
os.environ["SCW_SECRET_KEY"] = "your-scaleway-secret-key"
with open("speech.mp3", "rb") as audio_file:
response = transcription(
model="scaleway/whisper-large-v3",
file=audio_file,
)
print(response.text)
```
### Proxy config
```yaml
model_list:
- model_name: scaleway-whisper
litellm_params:
model: scaleway/whisper-large-v3
api_key: "os.environ/SCW_SECRET_KEY"
```
### Proxy request
```bash
curl http://localhost:4000/v1/audio/transcriptions \
-H "Authorization: Bearer YOUR_LITELLM_MASTER_KEY" \
-F model="scaleway-whisper" \
-F file="@speech.mp3"
```
Supported optional params: `language`, `prompt`, `response_format`, `temperature`, `timestamp_granularities`.

View file

@ -0,0 +1,95 @@
# Agentic Loop Hook
Build a `CustomLogger` callback that intercepts a model response, fulfills tool calls server-side, and reruns the model — transparently to the caller.
:::info Supported call types
- `async` only (sync calls do not trigger the hook)
- Non-streaming only (streaming responses cannot be inspected for tool calls)
- Works on both `/v1/messages` and `/v1/chat/completions`
:::
## Implement the callback
Override two methods on `CustomLogger`:
```python
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
MY_TOOL = "my_tool"
class MyToolCallback(CustomLogger):
async def async_should_run_agentic_loop(
self, response, model, messages, tools, stream, custom_llm_provider, kwargs
):
# Return (True, context_dict) if there are tool calls to handle
content = getattr(response, "content", None) or []
calls = [b for b in content if isinstance(b, dict)
and b.get("type") == "tool_use" and b.get("name") == MY_TOOL]
if not calls:
return False, {}
return True, {"tool_calls": calls}
async def async_build_agentic_loop_plan(
self, tools, model, messages, response,
anthropic_messages_provider_config,
anthropic_messages_optional_request_params,
logging_obj, stream, kwargs,
):
calls = tools["tool_calls"]
results = [f"result for {c['input']}" for c in calls] # your logic here
follow_up = messages + [
{"role": "assistant", "content": [
{"type": "tool_use", "id": c["id"], "name": c["name"], "input": c["input"]}
for c in calls
]},
{"role": "user", "content": [
{"type": "tool_result", "tool_use_id": c["id"], "content": results[i]}
for i, c in enumerate(calls)
]},
]
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=AgenticLoopRequestPatch(messages=follow_up),
)
```
For `/v1/chat/completions`, override `async_build_chat_completion_agentic_loop_plan` instead — same idea, `optional_params` replaces `anthropic_messages_optional_request_params`.
## Register it
```python
import litellm
litellm.callbacks = [MyToolCallback()]
```
Or in `config.yaml`:
```yaml
litellm_settings:
callbacks: ["my_module.MyToolCallback"]
```
## `AgenticLoopPlan` fields
| Field | Effect |
|---|---|
| `run_agentic_loop=True` + `request_patch` | Reruns the model with the patched request |
| `response_override` | Returns this value directly to the caller (no rerun) |
| `terminate=True` | Stops the loop, returns the current response |
| `run_agentic_loop=False` (default) | Skips; next callback is checked |
`AgenticLoopRequestPatch` accepts: `model`, `messages`, `tools`, `max_tokens`, `optional_params`, `kwargs`.
## Loop safety
- Default max reruns: `3` — override per-request with `kwargs["max_agentic_loops"]`
- Identical tool-call fingerprints abort the loop automatically
- Current depth is in `kwargs["_agentic_loop_depth"]`
## Examples in this repo
- `litellm/integrations/compression_interception/handler.py`
- `litellm/integrations/websearch_interception/handler.py`

View file

@ -8,6 +8,22 @@ Reduce costs by up to 90% by using LiteLLM to auto-inject prompt caching checkpo
<Image img={require('../../img/auto_prompt_caching.png')} style={{ width: '800px', height: 'auto' }} />
Supported Providers (`cache_control` marker):
- Anthropic API (`anthropic/`)
- AWS Bedrock - Claude (`bedrock/`)
- Vertex AI - Claude and Gemini (`vertex_ai/`)
- Google AI Studio - Gemini (`gemini/`)
- Azure AI - Claude (`azure_ai/`)
- OpenRouter - Claude, Gemini, MiniMax, GLM, z-ai routes (`openrouter/`)
- Databricks - Claude (`databricks/`)
- DashScope / Qwen (`dashscope/`)
- MiniMax (`minimax/`)
- Z.ai / GLM (`zai/`)
Provider Managed (automatic, no marker needed):
- OpenAI (`openai/`)
- DeepSeek (`deepseek/`)
- xAI (`xai/`)
## How it works

View file

@ -536,6 +536,7 @@ const sidebars = {
description: "Modify requests, responses, and more",
items: [
"proxy/call_hooks",
"proxy/agentic_loop_hook",
"proxy/rules",
]
},

View file

@ -148,6 +148,7 @@ _custom_logger_compatible_callbacks_literal = Literal[
"vantage",
"posthog",
"levo",
"compression_interception",
]
cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None
logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None

View file

@ -1,9 +1,9 @@
"""
Main compress() function orchestrates BM25/embedding scoring, message stubbing,
and retrieval tool injection.
Main compress() function normalizes input messages, orchestrates BM25/embedding
scoring, message stubbing, and retrieval tool injection.
"""
from typing import Any, Dict, List, Optional, Set, Union, cast
from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast
from litellm.caching.dual_cache import DualCache
from litellm.compression.message_stubbing import (
@ -15,27 +15,196 @@ from litellm.compression.retrieval_tool import build_retrieval_tool
from litellm.compression.scoring.bm25 import bm25_score_messages
from litellm.litellm_core_utils.token_counter import token_counter
from litellm.types.compression import CompressedResult
from litellm.types.utils import AllMessageValues, Message
from litellm.types.utils import CallTypes
# CallTypes that produce Anthropic-shaped messages (structured content blocks).
# Everything else is treated as OpenAI chat-completions shape.
_ANTHROPIC_CALL_TYPES = frozenset({CallTypes.anthropic_messages.value})
# CallTypes that are valid targets for compression. Compression operates on
# message-shaped inputs, so we only accept call types whose payload is a list
# of role/content messages.
_SUPPORTED_CALL_TYPES = frozenset(
{
CallTypes.completion.value,
CallTypes.acompletion.value,
CallTypes.anthropic_messages.value,
}
)
def _normalize_call_type(call_type: Union[CallTypes, str]) -> str:
"""Return the string value for a ``CallTypes`` enum or a raw string."""
if isinstance(call_type, CallTypes):
return call_type.value
return call_type
def _is_anthropic_call_type(call_type: str) -> bool:
return call_type in _ANTHROPIC_CALL_TYPES
def _build_retrieval_tools(keys: List[str], call_type: str) -> List[dict]:
"""
Build retrieval tool definitions in the target request schema.
- Chat-completions call types: keep OpenAI function-tool schema.
- Anthropic messages call type: remap to Anthropic's custom tool schema.
"""
if not keys:
return []
openai_tools = [build_retrieval_tool(keys)]
if not _is_anthropic_call_type(call_type):
return openai_tools
# Lazy import to avoid introducing provider transformation imports during
# module import for non-Anthropic call paths.
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
anthropic_tools, _mcp_servers = AnthropicConfig()._map_tools(openai_tools)
return cast(List[dict], anthropic_tools)
def _content_to_text(content: Any) -> str:
"""
Convert OpenAI/Anthropic message content blocks to plain text.
Text extraction policy:
- Include text-bearing fields only (`text` blocks + string values).
- For `tool_result`, expand into nested `content` items.
- Ignore non-textual blocks (images/documents/tool metadata/thinking metadata).
Implemented iteratively (stack-based) to avoid unbounded recursion.
"""
parts: List[str] = []
stack: List[Any] = [content]
while stack:
item = stack.pop()
if isinstance(item, str):
parts.append(item)
elif isinstance(item, list):
# Push list items in reverse order so they are processed left-to-right.
for element in reversed(item):
stack.append(element)
elif isinstance(item, dict):
item_type = item.get("type")
if item_type == "text":
parts.append(str(item.get("text", "")))
elif item_type == "tool_result":
stack.append(item.get("content", ""))
return " ".join(parts)
def _normalize_messages_for_compression(
messages: List[dict],
call_type: str,
) -> Tuple[List[dict], List[dict]]:
"""
Normalize each original message to a text-surrogate content for scoring.
Returns:
(normalized_messages, original_messages_copy)
"""
if call_type not in _SUPPORTED_CALL_TYPES:
raise ValueError(
f"Unsupported call_type={call_type!r} for compression. "
f"Expected one of: {sorted(_SUPPORTED_CALL_TYPES)}."
)
original_messages: List[Dict[str, Any]] = [dict(m) for m in messages]
normalized_messages: List[dict] = []
for msg in original_messages:
normalized_messages.append(
{
**msg,
"content": _content_to_text(msg.get("content", "")),
}
)
return normalized_messages, original_messages
def _extract_last_user_message(messages: List[dict]) -> str:
"""Return the text content of the last user message."""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content", "")
if isinstance(content, str):
return content
if isinstance(content, list):
parts = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
parts.append(part.get("text", ""))
elif isinstance(part, str):
parts.append(part)
return " ".join(parts)
return _content_to_text(msg.get("content", ""))
return ""
def _extract_tool_use_ids(content: Any) -> List[str]:
if not isinstance(content, list):
return []
tool_use_ids: List[str] = []
for part in content:
if not isinstance(part, dict):
continue
if part.get("type") != "tool_use":
continue
tool_use_id = part.get("id")
if isinstance(tool_use_id, str) and tool_use_id:
tool_use_ids.append(tool_use_id)
return tool_use_ids
def _extract_tool_result_ids(content: Any) -> Set[str]:
if not isinstance(content, list):
return set()
tool_result_ids: Set[str] = set()
for part in content:
if not isinstance(part, dict):
continue
if part.get("type") != "tool_result":
continue
tool_use_id = part.get("tool_use_id")
if isinstance(tool_use_id, str) and tool_use_id:
tool_result_ids.add(tool_use_id)
return tool_result_ids
def _extract_anthropic_tool_exchange_spans(
messages: List[dict],
) -> Tuple[List[Set[int]], Optional[str]]:
"""
Return atomic 2-message spans for Anthropic tool exchanges.
Each assistant message containing `tool_use` must be immediately followed by a
user message containing matching `tool_result` blocks for all tool_use ids.
"""
spans: List[Set[int]] = []
i = 0
while i < len(messages):
current = messages[i]
if current.get("role") != "assistant":
i += 1
continue
tool_use_ids = _extract_tool_use_ids(current.get("content"))
if not tool_use_ids:
i += 1
continue
if i + 1 >= len(messages):
return [], "invalid_anthropic_tool_sequence"
next_msg = messages[i + 1]
if next_msg.get("role") != "user":
return [], "invalid_anthropic_tool_sequence"
tool_result_ids = _extract_tool_result_ids(next_msg.get("content"))
if not tool_result_ids:
return [], "invalid_anthropic_tool_sequence"
for tool_use_id in tool_use_ids:
if tool_use_id not in tool_result_ids:
return [], "invalid_anthropic_tool_sequence"
spans.append({i, i + 1})
i += 2
return spans, None
def _get_protected_indices(messages: List[dict]) -> List[int]:
"""
Return indices of messages that must never be compressed:
@ -87,9 +256,98 @@ def _combine_scores(
return [bm25_weight * b + emb_weight * e for b, e in zip(norm_bm25, norm_emb)]
def _select_kept_indices_for_budget(
normalized_messages: List[dict],
original_messages: List[dict],
combined_scores: List[float],
compression_target: int,
model: str,
initial_kept_indices: Set[int],
tool_exchange_spans: List[Set[int]],
) -> Tuple[Set[int], Dict[int, dict]]:
kept_indices = set(initial_kept_indices)
current_tokens = 0
for i in kept_indices:
current_tokens += token_counter(
model=model,
text=cast(str, normalized_messages[i].get("content", "") or ""),
)
# Fill token budget from highest-scoring units.
# A unit is either:
# 1) a single message index, or
# 2) an Anthropic tool-exchange span that must be kept/dropped atomically.
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
span_id_by_index: Dict[int, int] = {}
for span_id, span in enumerate(tool_exchange_spans):
for idx in span:
span_id_by_index[idx] = span_id
# Build single-message candidate units (non-span messages).
candidate_units: List[Tuple[float, Tuple[int, ...], bool]] = []
for idx in range(len(normalized_messages)):
if idx in span_id_by_index or idx in kept_indices:
continue
candidate_units.append((combined_scores[idx], (idx,), True))
# Build span candidate units (atomic keep/drop for tool exchanges).
for span in tool_exchange_spans:
span_indices = tuple(sorted(span))
if any(idx in kept_indices for idx in span_indices):
continue
span_score = max(combined_scores[idx] for idx in span_indices)
candidate_units.append((span_score, span_indices, False))
# Sort by descending relevance score.
candidate_units.sort(key=lambda item: item[0], reverse=True)
for _score, indices, can_truncate in candidate_units:
if any(idx in kept_indices for idx in indices):
continue
msg_tokens = 0
for idx in indices:
msg_tokens += token_counter(
model=model,
text=cast(str, normalized_messages[idx].get("content", "") or ""),
)
remaining = compression_target - current_tokens
if remaining <= 0:
break # budget exhausted
if current_tokens + msg_tokens <= compression_target:
# Fits entirely
kept_indices.update(indices)
current_tokens += msg_tokens
elif can_truncate and len(indices) == 1 and remaining >= 100:
# Too large to fit whole single message, but we have budget — truncate it.
idx = indices[0]
truncated = truncate_message(original_messages[idx], remaining)
truncated_tokens = token_counter(
model=model,
text=truncated.get("content", "") or "",
)
truncated_overrides[idx] = truncated
kept_indices.add(idx)
current_tokens += truncated_tokens
return kept_indices, truncated_overrides
def _get_dropped_tool_span_indices(
kept_indices: Set[int], tool_exchange_spans: List[Set[int]]
) -> Set[int]:
dropped_tool_span_indices: Set[int] = set()
for span in tool_exchange_spans:
if not any(idx in kept_indices for idx in span):
dropped_tool_span_indices.update(span)
return dropped_tool_span_indices
def compress(
messages: List[dict],
model: str,
call_type: Union[CallTypes, str] = CallTypes.completion,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
@ -108,6 +366,12 @@ def compress(
Parameters:
messages: The conversation messages to (potentially) compress.
model: The LLM model name used for token counting.
call_type: The LiteLLM call type whose message schema these messages
follow. Supported values:
- ``CallTypes.completion`` / ``CallTypes.acompletion`` OpenAI
chat-completions shape (default)
- ``CallTypes.anthropic_messages`` Anthropic Messages shape
(structured content blocks + atomic tool exchanges)
compression_trigger: Only compress if input exceeds this token count.
compression_target: Target token count after compression.
Defaults to ``compression_trigger // 2``.
@ -122,29 +386,37 @@ def compress(
A ``CompressedResult`` dict containing compressed messages, token
counts, a cache of original content, and the retrieval tool definition.
"""
call_type_str = _normalize_call_type(call_type)
normalized_messages, original_messages = _normalize_messages_for_compression(
messages=messages,
call_type=call_type_str,
)
if compression_target is None:
compression_target = compression_trigger * 7 // 10
original_tokens = token_counter(
model=model, messages=cast(List[Union[AllMessageValues, Message]], messages)
model=model,
messages=cast(List[Any], original_messages),
)
# Pass through if below trigger
if original_tokens <= compression_trigger:
return CompressedResult(
messages=messages,
messages=original_messages,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
compression_ratio=0.0,
cache={},
tools=[],
compression_skipped_reason="below_trigger",
)
# Extract query for relevance scoring
query = _extract_last_user_message(messages)
query = _extract_last_user_message(normalized_messages)
# Score each message
bm25_scores = bm25_score_messages(query, messages)
bm25_scores = bm25_score_messages(query, normalized_messages)
if embedding_model:
from litellm.compression.scoring.embedding_scorer import (
@ -153,7 +425,7 @@ def compress(
emb_scores = embedding_score_messages(
query,
messages,
normalized_messages,
model=embedding_model,
cache=compression_cache,
embedding_model_params=embedding_model_params,
@ -162,85 +434,69 @@ def compress(
else:
combined_scores = bm25_scores
# Sort message indices by score descending
ranked_indices = sorted(
range(len(messages)),
key=lambda i: combined_scores[i],
reverse=True,
)
# Protected messages are never compressed
protected_indices = _get_protected_indices(messages)
protected_indices = _get_protected_indices(normalized_messages)
kept_indices: Set[int] = set(protected_indices)
# Count tokens for protected messages
current_tokens = 0
for i in kept_indices:
current_tokens += token_counter(
model=model, text=messages[i].get("content", "") or ""
tool_exchange_spans: List[Set[int]] = []
if _is_anthropic_call_type(call_type_str):
tool_exchange_spans, tool_sequence_error = (
_extract_anthropic_tool_exchange_spans(original_messages)
)
# Fill token budget from highest-scoring messages.
# For each candidate (ranked by relevance):
# - If it fits entirely → keep it as-is.
# - If it doesn't fit but there's meaningful remaining budget → truncate it
# to fill as much of the budget as possible.
# - Otherwise → stub it (pointer only, content goes to cache).
# Multiple messages may be truncated so we preserve partial content from
# several high-scoring messages rather than fully stubbing all but one.
truncated_overrides: Dict[int, dict] = {} # idx -> truncated message dict
for idx in ranked_indices:
if idx in kept_indices:
continue
msg_content = messages[idx].get("content", "") or ""
msg_tokens = token_counter(model=model, text=msg_content)
remaining = compression_target - current_tokens
if remaining <= 0:
break # budget exhausted
if current_tokens + msg_tokens <= compression_target:
# Fits entirely
kept_indices.add(idx)
current_tokens += msg_tokens
elif remaining >= 100:
# Too large to fit whole, but we have budget — truncate it.
truncated = truncate_message(messages[idx], remaining)
truncated_tokens = token_counter(
model=model,
text=truncated.get("content", "") or "",
if tool_sequence_error is not None:
return CompressedResult(
messages=original_messages,
original_tokens=original_tokens,
compressed_tokens=original_tokens,
compression_ratio=0.0,
cache={},
tools=[],
compression_skipped_reason=tool_sequence_error,
)
truncated_overrides[idx] = truncated
kept_indices.add(idx)
current_tokens += truncated_tokens
for span in tool_exchange_spans:
# If any message in the span is protected, keep the whole span.
if any(idx in kept_indices for idx in span):
kept_indices.update(span)
kept_indices, truncated_overrides = _select_kept_indices_for_budget(
normalized_messages=normalized_messages,
original_messages=original_messages,
combined_scores=combined_scores,
compression_target=compression_target,
model=model,
initial_kept_indices=kept_indices,
tool_exchange_spans=tool_exchange_spans,
)
# Build compressed messages and cache
compressed_messages: List[dict] = []
cache: Dict[str, str] = {}
used_keys: Set[str] = set()
dropped_tool_span_indices = _get_dropped_tool_span_indices(
kept_indices=kept_indices, tool_exchange_spans=tool_exchange_spans
)
for i, msg in enumerate(messages):
for i, msg in enumerate(original_messages):
if i in dropped_tool_span_indices:
continue
if i in kept_indices:
# Use the truncated version if we made one, otherwise the original
compressed_messages.append(truncated_overrides.get(i, msg))
else:
key = extract_key(msg, fallback_index=i, used_keys=used_keys)
content = msg.get("content", "")
if isinstance(content, list):
content = " ".join(
p.get("text", "") if isinstance(p, dict) else str(p)
for p in content
)
key = extract_key(
normalized_messages[i], fallback_index=i, used_keys=used_keys
)
content = _content_to_text(msg.get("content", ""))
cache[key] = content
compressed_messages.append(stub_message(msg, key))
# Build retrieval tool
tools = [build_retrieval_tool(list(cache.keys()))] if cache else []
# Build retrieval tool in the target request schema
tools = _build_retrieval_tools(list(cache.keys()), call_type=call_type_str)
compressed_tokens = token_counter(
model=model,
messages=cast(List[Union[AllMessageValues, Message]], compressed_messages),
messages=cast(List[Any], compressed_messages),
)
return CompressedResult(

View file

@ -0,0 +1,14 @@
"""
Compression Interception Module
Provides server-side prompt compression + retrieval tool fulfillment for
Anthropic Messages agentic loops.
"""
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
__all__ = [
"CompressionInterceptionLogger",
]

View file

@ -0,0 +1,399 @@
"""
Compression Interception Handler
CustomLogger that compresses inbound Anthropic Messages requests and fulfills
litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
"""
import time
import uuid
from typing import Any, Dict, List, Optional, Tuple, cast
from litellm._logging import verbose_logger
from litellm.compression import compress
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.integrations.compression_interception import (
CompressionInterceptionConfig,
)
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.utils import CallTypes
LITELLM_CONTENT_RETRIEVE_TOOL_NAME = "litellm_content_retrieve"
_CACHE_TTL_SECONDS = 15 * 60
class CompressionInterceptionLogger(CustomLogger):
"""
CustomLogger that implements transparent prompt compression + retrieval loops.
Flow:
1. Compress inbound /v1/messages requests in pre-call hook.
2. Inject litellm_content_retrieve tool and persist compressed cache by call_id.
3. Detect retrieval tool_use blocks in first model response.
4. Build typed rerun plan with tool_result blocks from the compressed cache.
"""
def __init__(
self,
enabled: bool = True,
compression_trigger: int = 200_000,
compression_target: Optional[int] = None,
embedding_model: Optional[str] = None,
embedding_model_params: Optional[Dict[str, Any]] = None,
):
super().__init__()
self.enabled = enabled
self.compression_trigger = compression_trigger
self.compression_target = compression_target
self.embedding_model = embedding_model
self.embedding_model_params = embedding_model_params
self._compression_cache_by_call_id: Dict[str, Tuple[Dict[str, str], float]] = {}
@classmethod
def from_config_yaml(
cls, config: CompressionInterceptionConfig
) -> "CompressionInterceptionLogger":
return cls(
enabled=bool(config.get("enabled", True)),
compression_trigger=int(config.get("compression_trigger", 200_000)),
compression_target=config.get("compression_target"),
embedding_model=config.get("embedding_model"),
embedding_model_params=config.get("embedding_model_params"),
)
@staticmethod
def initialize_from_proxy_config(
litellm_settings: Dict[str, Any],
callback_specific_params: Dict[str, Any],
) -> "CompressionInterceptionLogger":
compression_params: CompressionInterceptionConfig = {}
if "compression_interception_params" in litellm_settings:
compression_params = litellm_settings["compression_interception_params"]
elif "compression_interception" in callback_specific_params:
compression_params = callback_specific_params["compression_interception"]
return CompressionInterceptionLogger.from_config_yaml(compression_params)
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
if not self.enabled:
return None
if call_type is not None and call_type != CallTypes.anthropic_messages:
return None
if int(kwargs.get("_agentic_loop_depth", 0) or 0) > 0:
return None
messages = kwargs.get("messages")
model = kwargs.get("model")
if not isinstance(messages, list) or not isinstance(model, str):
return None
if self._has_retrieval_tool(kwargs.get("tools")):
return None
self._prune_expired_cache()
compressed = compress( # type: ignore
messages=messages,
model=model,
call_type=CallTypes.anthropic_messages,
compression_trigger=self.compression_trigger,
compression_target=self.compression_target,
embedding_model=self.embedding_model,
embedding_model_params=self.embedding_model_params,
)
cache = cast(Dict[str, str], compressed.get("cache", {}))
skip_reason = cast(Optional[str], compressed.get("compression_skipped_reason"))
compressed_tools = cast(List[Dict[str, Any]], compressed.get("tools", []))
# Only mutate kwargs when compression actually produced a result.
# If compression was a no-op (below trigger, invalid tool sequence, etc.),
# leave ``messages`` and ``tools`` untouched — injecting an empty
# ``tools: []`` onto a request that originally had no tools breaks
# Anthropic Messages requests.
if cache:
kwargs["messages"] = compressed["messages"]
if compressed_tools:
kwargs["tools"] = self._merge_tools(
existing_tools=cast(
Optional[List[Dict[str, Any]]], kwargs.get("tools")
),
compressed_tools=compressed_tools,
)
call_id = cast(Optional[str], kwargs.get("litellm_call_id"))
if not call_id:
call_id = str(uuid.uuid4())
kwargs["litellm_call_id"] = call_id
self._compression_cache_by_call_id[call_id] = (cache, time.time())
verbose_logger.debug(
"CompressionInterception: compressed request [call_id=%s original=%d compressed=%d cached_keys=%d]",
call_id,
compressed.get("original_tokens"),
compressed.get("compressed_tokens"),
len(cache),
)
elif skip_reason is not None:
verbose_logger.debug(
"CompressionInterception: compression skipped [reason=%s original=%d compressed=%d]",
skip_reason,
compressed.get("original_tokens"),
compressed.get("compressed_tokens"),
)
return kwargs
async def async_should_run_agentic_loop(
self,
response: Any,
model: str,
messages: List[Dict],
tools: Optional[List[Dict]],
stream: bool,
custom_llm_provider: str,
kwargs: Dict,
) -> Tuple[bool, Dict]:
if not self.enabled:
return False, {}
if not self._has_retrieval_tool(tools):
return False, {}
tool_calls, thinking_blocks = self._extract_retrieval_tool_calls(
response=response
)
if not tool_calls:
return False, {}
return True, {
"tool_calls": tool_calls,
"thinking_blocks": thinking_blocks,
"tool_type": "compression_retrieval",
}
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
self._prune_expired_cache()
tool_calls = cast(List[Dict[str, Any]], tools.get("tool_calls", []))
thinking_blocks = cast(List[Dict[str, Any]], tools.get("thinking_blocks", []))
call_id = self._resolve_call_id(logging_obj=logging_obj, kwargs=kwargs)
cache = self._get_cache(call_id=call_id)
retrieval_results = [
self._resolve_retrieval_content(tc, cache) for tc in tool_calls
]
assistant_message = {
"role": "assistant",
"content": thinking_blocks
+ [
{
"type": "tool_use",
"id": tc.get("id"),
"name": tc.get("name", LITELLM_CONTENT_RETRIEVE_TOOL_NAME),
"input": tc.get("input", {}),
}
for tc in tool_calls
],
}
user_message = {
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": tool_calls[i].get("id"),
"content": retrieval_results[i],
}
for i in range(len(tool_calls))
],
}
follow_up_messages = messages + [assistant_message, user_message]
max_tokens = cast(
Optional[int],
anthropic_messages_optional_request_params.get("max_tokens")
or kwargs.get("max_tokens"),
)
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
full_model_name = model
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = cast(str, agentic_params.get("model", model))
request_patch = AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
max_tokens=max_tokens,
optional_params=optional_params_without_max_tokens,
kwargs=self._prepare_followup_kwargs(kwargs=kwargs),
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "compression_retrieval", "call_id": call_id or ""},
)
def _prune_expired_cache(self) -> None:
now = time.time()
self._compression_cache_by_call_id = {
call_id: (cache, created_at)
for call_id, (
cache,
created_at,
) in self._compression_cache_by_call_id.items()
if now - created_at <= _CACHE_TTL_SECONDS
}
def _get_cache(self, call_id: Optional[str]) -> Dict[str, str]:
if not call_id:
return {}
cache_entry = self._compression_cache_by_call_id.get(call_id)
if cache_entry is None:
return {}
return cache_entry[0]
def _resolve_call_id(
self, logging_obj: Any, kwargs: Dict[str, Any]
) -> Optional[str]:
if logging_obj is not None:
logging_call_id = getattr(logging_obj, "litellm_call_id", None)
if isinstance(logging_call_id, str) and logging_call_id:
return logging_call_id
kwargs_call_id = kwargs.get("litellm_call_id")
return cast(
Optional[str], kwargs_call_id if isinstance(kwargs_call_id, str) else None
)
def _resolve_retrieval_content(
self, tool_call: Dict[str, Any], cache: Dict[str, str]
) -> str:
raw_input = tool_call.get("input", {})
key = ""
if isinstance(raw_input, dict):
key = str(raw_input.get("key", "") or "")
if not key:
return "No retrieval key provided."
if key in cache:
return cache[key]
return f"[compressed content key '{key}' not found]"
def _extract_retrieval_tool_calls(
self, response: Any
) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
if isinstance(response, dict):
content = response.get("content", [])
else:
content = getattr(response, "content", []) or []
if not isinstance(content, list):
return [], []
tool_calls: List[Dict[str, Any]] = []
thinking_blocks: List[Dict[str, Any]] = []
for block in content:
if isinstance(block, dict):
block_type = block.get("type")
block_name = block.get("name")
if block_type in ("thinking", "redacted_thinking"):
thinking_blocks.append(block)
if (
block_type == "tool_use"
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
tool_calls.append(
{
"id": block.get("id"),
"type": "tool_use",
"name": block_name,
"input": block.get("input", {}),
}
)
else:
block_type = getattr(block, "type", None)
block_name = getattr(block, "name", None)
if block_type == "thinking":
thinking_blocks.append(
{
"type": "thinking",
"thinking": getattr(block, "thinking", ""),
"signature": getattr(block, "signature", ""),
}
)
elif block_type == "redacted_thinking":
thinking_blocks.append(
{
"type": "redacted_thinking",
"data": getattr(block, "data", ""),
}
)
if (
block_type == "tool_use"
and block_name == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
tool_calls.append(
{
"id": getattr(block, "id", None),
"type": "tool_use",
"name": block_name,
"input": getattr(block, "input", {}) or {},
}
)
return tool_calls, thinking_blocks
def _prepare_followup_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
internal_keys = {"litellm_logging_obj"}
return {
k: v
for k, v in kwargs.items()
if not k.startswith("_compression_interception") and k not in internal_keys
}
def _has_retrieval_tool(self, tools: Any) -> bool:
if not isinstance(tools, list):
return False
for tool in tools:
if not isinstance(tool, dict):
continue
function = tool.get("function")
if tool.get("type") == "function" and isinstance(function, dict):
if function.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME:
return True
if (
tool.get("type") == "custom"
and tool.get("name") == LITELLM_CONTENT_RETRIEVE_TOOL_NAME
):
return True
return False
def _merge_tools(
self,
existing_tools: Optional[List[Dict[str, Any]]],
compressed_tools: List[Dict[str, Any]],
) -> List[Dict[str, Any]]:
merged = list(existing_tools or [])
if self._has_retrieval_tool(merged):
return merged
merged.extend(compressed_tools)
return merged

View file

@ -20,6 +20,7 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
from litellm.types.integrations.argilla import ArgillaItem
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
from litellm.types.prompts.init_prompts import PromptSpec
from litellm.types.integrations.custom_logger import AgenticLoopPlan
from litellm.types.utils import (
AdapterCompletionStreamWrapper,
CallTypes,
@ -239,7 +240,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
@ -676,6 +677,26 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
"""
Build a typed rerun plan for Anthropic Messages agentic loops.
Override this method to separate callback decision/tool execution from
follow-up request execution (handled by BaseLLMHTTPHandler).
"""
return AgenticLoopPlan(run_agentic_loop=False)
async def async_should_run_chat_completion_agentic_loop(
self,
response: Any,
@ -707,6 +728,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
pass
async def async_build_chat_completion_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
optional_params: Dict,
logging_obj: "LiteLLMLoggingObj",
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
"""
Build a typed rerun plan for chat-completions agentic loops.
"""
return AgenticLoopPlan(run_agentic_loop=False)
# Useful helpers for custom logger classes
def truncate_standard_logging_payload_content(

View file

@ -51,6 +51,7 @@ if TYPE_CHECKING:
else:
AsyncIOScheduler = Any
class PrometheusLogger(CustomLogger):
# Class variables or attributes
@ -991,9 +992,7 @@ class PrometheusLogger(CustomLogger):
amount: float = 1.0,
) -> None:
_labels = prometheus_label_factory(
supported_enum_labels=self.get_labels_for_metric(
metric_name=metric_name
),
supported_enum_labels=self.get_labels_for_metric(metric_name=metric_name),
enum_values=enum_values,
label_context=label_context,
)
@ -1118,7 +1117,9 @@ class PrometheusLogger(CustomLogger):
user_api_key = hash_token(user_api_key)
label_context = PrometheusLabelFactoryContext(enum_values) #amortized per request.
label_context = PrometheusLabelFactoryContext(
enum_values
) # amortized per request.
# increment total LLM requests and spend metric
self._increment_top_level_request_and_spend_metrics(
@ -3490,7 +3491,9 @@ def _prometheus_labels_from_context(
}
if UserAPIKeyLabelNames.END_USER.value in filtered_labels:
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = ctx.get_resolved_end_user()
filtered_labels[UserAPIKeyLabelNames.END_USER.value] = (
ctx.get_resolved_end_user()
)
for sk, val in ctx._custom_by_sanitized_key.items():
if sk in supported_enum_labels:

View file

@ -51,8 +51,7 @@ class PrometheusLabelFactoryContext:
self.enum_values = enum_values
enum_dict = enum_values.model_dump()
self._sanitized_enum: Dict[str, Optional[str]] = {
k: _sanitize_prometheus_label_value(v)
for k, v in enum_dict.items()
k: _sanitize_prometheus_label_value(v) for k, v in enum_dict.items()
}
self._custom_by_sanitized_key: Dict[str, Optional[str]] = {}
if enum_values.custom_metadata_labels is not None:

View file

@ -28,6 +28,10 @@ from litellm.integrations.websearch_interception.transformation import (
from litellm.types.integrations.websearch_interception import (
WebSearchInterceptionConfig,
)
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
@ -573,6 +577,35 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs=kwargs,
)
async def async_build_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
tool_calls = tools["tool_calls"]
thinking_blocks = tools.get("thinking_blocks", [])
request_patch = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
thinking_blocks=thinking_blocks,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs,
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "websearch", "response_format": "anthropic"},
)
async def async_run_chat_completion_agentic_loop(
self,
tools: Dict,
@ -608,6 +641,33 @@ class WebSearchInterceptionLogger(CustomLogger):
response_format=response_format,
)
async def async_build_chat_completion_agentic_loop_plan(
self,
tools: Dict,
model: str,
messages: List[Dict],
response: Any,
optional_params: Dict,
logging_obj: Any,
stream: bool,
kwargs: Dict,
) -> AgenticLoopPlan:
tool_calls = tools["tool_calls"]
response_format = tools.get("response_format", "openai")
request_patch = await self._build_chat_completion_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
optional_params=optional_params,
kwargs=kwargs,
response_format=response_format,
)
return AgenticLoopPlan(
run_agentic_loop=True,
request_patch=request_patch,
metadata={"tool_type": "websearch", "response_format": response_format},
)
@staticmethod
def _resolve_max_tokens(
optional_params: Dict,
@ -672,7 +732,48 @@ class WebSearchInterceptionLogger(CustomLogger):
stream: bool,
kwargs: Dict,
) -> Any:
"""Execute litellm.search() and make follow-up request"""
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch = await self._build_anthropic_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
thinking_blocks=thinking_blocks,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs,
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
optional_params = dict(anthropic_messages_optional_request_params)
optional_params.update(request_patch.optional_params)
max_tokens = request_patch.max_tokens
if max_tokens is None:
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
else:
optional_params.pop("max_tokens", None)
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
return await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=request_patch.messages,
model=request_patch.model or model,
**optional_params,
**request_patch.kwargs,
)
async def _build_anthropic_request_patch(
self,
model: str,
messages: List[Dict],
tool_calls: List[Dict],
thinking_blocks: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
kwargs: Dict,
) -> AgenticLoopRequestPatch:
"""Execute litellm.search() and build follow-up request patch."""
# Extract search queries from tool_use blocks
search_tasks = []
@ -721,20 +822,8 @@ class WebSearchInterceptionLogger(CustomLogger):
thinking_blocks=thinking_blocks,
)
# Make follow-up request with search results
# Type cast: user_message is a Dict for Anthropic format (default response_format)
follow_up_messages = messages + [assistant_message, cast(Dict, user_message)]
verbose_logger.debug(
"WebSearchInterception: Making follow-up request with search results"
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
)
verbose_logger.debug(
f"WebSearchInterception: Last message (tool_result): {user_message}"
)
# Correlation context for structured logging
_call_id = getattr(logging_obj, "litellm_call_id", None) or kwargs.get(
"litellm_call_id", "unknown"
@ -742,61 +831,41 @@ class WebSearchInterceptionLogger(CustomLogger):
full_model_name = model # safe default before try block
# Use anthropic_messages.acreate for follow-up request
try:
max_tokens = self._resolve_max_tokens(
anthropic_messages_optional_request_params, kwargs
)
max_tokens = self._resolve_max_tokens(
anthropic_messages_optional_request_params, kwargs
)
verbose_logger.debug(
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
)
verbose_logger.debug(
f"WebSearchInterception: Using max_tokens={max_tokens} for follow-up request"
)
# Create a copy of optional params without max_tokens (since we pass it explicitly)
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
optional_params_without_max_tokens = {
k: v
for k, v in anthropic_messages_optional_request_params.items()
if k != "max_tokens"
}
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
kwargs_for_followup = self._prepare_followup_kwargs(kwargs)
# Get model from logging_obj.model_call_details["agentic_loop_params"]
# This preserves the full model name with provider prefix (e.g., "bedrock/invoke/...")
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = agentic_params.get("model", model)
verbose_logger.debug(
f"WebSearchInterception: Using model name: {full_model_name}"
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
final_response = await anthropic_messages.acreate(
max_tokens=max_tokens,
messages=follow_up_messages,
model=full_model_name,
**optional_params_without_max_tokens,
**kwargs_for_followup,
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
)
verbose_logger.debug(
f"WebSearchInterception: Final response: {final_response}"
)
return final_response
except Exception as e:
verbose_logger.exception(
"WebSearchInterception: Follow-up request failed "
"[call_id=%s model=%s messages=%d searches=%d]: %s",
_call_id,
full_model_name,
len(follow_up_messages),
len(final_search_results),
str(e),
)
raise
full_model_name = agentic_params.get("model", model)
verbose_logger.debug(
"WebSearchInterception: Built anthropic request patch "
"[call_id=%s model=%s messages=%d searches=%d]",
_call_id,
full_model_name,
len(follow_up_messages),
len(final_search_results),
)
return AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
max_tokens=max_tokens,
optional_params=optional_params_without_max_tokens,
kwargs=kwargs_for_followup,
)
async def _execute_search(self, query: str) -> str:
"""Execute a single web search using router's search tools"""
@ -883,7 +952,36 @@ class WebSearchInterceptionLogger(CustomLogger):
kwargs: Dict,
response_format: str = "openai",
) -> Any:
"""Execute litellm.search() and make follow-up chat completion request"""
"""Legacy path: execute search + build patch + run follow-up call."""
request_patch = await self._build_chat_completion_request_patch(
model=model,
messages=messages,
tool_calls=tool_calls,
optional_params=optional_params,
kwargs=kwargs,
response_format=response_format,
)
if request_patch.messages is None:
raise ValueError("WebSearchInterception: missing follow-up messages")
params = dict(optional_params)
params.update(request_patch.optional_params)
return await litellm.acompletion(
model=request_patch.model or model,
messages=request_patch.messages,
**params,
**request_patch.kwargs,
)
async def _build_chat_completion_request_patch( # noqa: PLR0915
self,
model: str,
messages: List[Dict],
tool_calls: List[Dict],
optional_params: Dict,
kwargs: Dict,
response_format: str = "openai",
) -> AgenticLoopRequestPatch:
"""Execute litellm.search() and build chat-completion rerun patch."""
# Extract search queries from tool_calls
search_tasks = []
@ -963,74 +1061,56 @@ class WebSearchInterceptionLogger(CustomLogger):
f"WebSearchInterception: Follow-up messages count: {len(follow_up_messages)}"
)
# Use litellm.acompletion for follow-up request
try:
# Remove internal parameters that shouldn't be passed to follow-up request
internal_params = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
# Remove internal parameters that shouldn't be passed to follow-up request
internal_params = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception") and k not in internal_params
}
full_model_name = model
if "custom_llm_provider" in kwargs:
custom_llm_provider = kwargs["custom_llm_provider"]
if not model.startswith(custom_llm_provider) and "/" not in model:
full_model_name = f"{custom_llm_provider}/{model}"
verbose_logger.debug(
"WebSearchInterception: Built chat completion request patch model=%s messages=%d",
full_model_name,
len(follow_up_messages),
)
tools_param = optional_params.get("tools")
optional_params_clean = {
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception")
and k not in internal_params
}
}
if tools_param is not None:
optional_params_clean["tools"] = tools_param
# Get full model name from kwargs
full_model_name = model
if "custom_llm_provider" in kwargs:
custom_llm_provider = kwargs["custom_llm_provider"]
# Reconstruct full model name with provider prefix if needed
if not model.startswith(custom_llm_provider):
# Check if model already has a provider prefix
if "/" not in model:
full_model_name = f"{custom_llm_provider}/{model}"
verbose_logger.debug(
f"WebSearchInterception: Using model name: {full_model_name}"
)
# Prepare tools for follow-up request (same as original)
tools_param = optional_params.get("tools")
# Remove tools and extra_body from optional_params to avoid issues
# extra_body often contains internal LiteLLM params that shouldn't be forwarded
optional_params_clean = {
k: v
for k, v in optional_params.items()
if k
not in {
"tools",
"extra_body",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
}
final_response = await litellm.acompletion(
model=full_model_name,
messages=follow_up_messages,
tools=tools_param,
**optional_params_clean,
**kwargs_for_followup,
)
verbose_logger.debug(
f"WebSearchInterception: Follow-up request completed, response type: {type(final_response)}"
)
return final_response
except Exception as e:
verbose_logger.exception(
f"WebSearchInterception: Follow-up request failed: {str(e)}"
)
raise
return AgenticLoopRequestPatch(
model=full_model_name,
messages=follow_up_messages,
optional_params=optional_params_clean,
kwargs=kwargs_for_followup,
)
async def _create_empty_search_result(self) -> str:
"""Create an empty search result for tool calls without queries"""

View file

@ -296,6 +296,15 @@ def get_supported_openai_params( # noqa: PLR0915
return OVHCloudAudioTranscriptionConfig().get_supported_openai_params(
model=model
)
elif custom_llm_provider == "scaleway":
if request_type == "transcription":
from litellm.llms.scaleway.audio_transcription.transformation import (
ScalewayAudioTranscriptionConfig,
)
return ScalewayAudioTranscriptionConfig().get_supported_openai_params(
model=model
)
elif custom_llm_provider == "elevenlabs":
if request_type == "transcription":
from litellm.llms.elevenlabs.audio_transcription.transformation import (

View file

@ -684,7 +684,7 @@ def generic_cost_per_token( # noqa: PLR0915
- cache_creation
- image_tokens
)
# Clamp to zero: inconsistent streaming usage
# Clamp to zero: inconsistent streaming usage
if text_tokens < 0:
text_tokens = 0
prompt_tokens_details["text_tokens"] = text_tokens

View file

@ -34,6 +34,7 @@ from litellm.types.llms.anthropic import (
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionRequest,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
)
@ -67,6 +68,32 @@ class AnthropicMessagesHandler(BaseTranslation):
super().__init__()
self.adapter = LiteLLMAnthropicMessagesAdapter()
def _translate_to_openai(self, data: dict) -> ChatCompletionRequest:
"""Translate Anthropic request to OpenAI chat completion format."""
(
chat_completion_compatible_request,
_tool_name_mapping,
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
)
return chat_completion_compatible_request
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert Anthropic messages request data to OpenAI-spec structured messages.
Uses the Anthropic-to-OpenAI adapter to translate message format.
"""
messages = data.get("messages")
if messages is None:
return None
chat_completion_compatible_request = self._translate_to_openai(data)
result = cast(
List[AllMessageValues],
chat_completion_compatible_request.get("messages", []),
)
return result if result else None
async def process_input_messages(
self,
data: dict,
@ -82,13 +109,7 @@ class AnthropicMessagesHandler(BaseTranslation):
skip_system = effective_skip_system_message_for_guardrail(guardrail_to_apply)
(
chat_completion_compatible_request,
_tool_name_mapping,
) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai(
# Use a shallow copy to avoid mutating request data (pop on litellm_metadata).
anthropic_message_request=cast(AnthropicMessagesRequest, data.copy())
)
chat_completion_compatible_request = self._translate_to_openai(data)
structured_messages = cast(
List[AllMessageValues],
@ -103,8 +124,6 @@ class AnthropicMessagesHandler(BaseTranslation):
chat_completion_compatible_request.get("tools", [])
)
task_mappings: List[Tuple[int, Optional[int]]] = []
# Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# Step 1: Extract all text content and images
for msg_idx, message in enumerate(messages):

View file

@ -0,0 +1,322 @@
"""
Agentic Streaming Iterator for Anthropic Messages
Wraps the raw SSE byte stream from the Anthropic pass-through endpoint,
yields every chunk to the caller (preserving real streaming), collects
all bytes, and on stream exhaustion rebuilds the full Anthropic response
to run through agentic completion hooks. If an agentic hook fires, the
follow-up response is chained as Phase 2 of the same iterator.
"""
import json
from typing import Any, AsyncIterator, Dict, List, Optional, cast
from litellm._logging import verbose_logger
# ---------------------------------------------------------------------------
# SSE parsing helpers (module-level to keep the class lean)
# ---------------------------------------------------------------------------
def _parse_sse_events(raw: bytes) -> List[tuple]:
"""Return a list of (event_type, parsed_data_dict) from raw SSE bytes."""
text = raw.decode("utf-8", errors="replace")
lines = text.split("\n")
events: List[tuple] = []
current_event_type: Optional[str] = None
for line in lines:
stripped = line.strip()
if stripped.startswith("event:"):
current_event_type = stripped[len("event:") :].strip()
continue
if not stripped.startswith("data:"):
continue
data_str = stripped[len("data:") :].strip()
try:
data = json.loads(data_str)
except (json.JSONDecodeError, ValueError):
continue
event_type = current_event_type or data.get("type", "")
current_event_type = None
events.append((event_type, data))
return events
def _handle_message_start(data: Dict, response: Dict) -> None:
msg = data.get("message", {})
response["id"] = msg.get("id", response["id"])
response["model"] = msg.get("model", response["model"])
response["role"] = msg.get("role", response["role"])
usage = msg.get("usage", {})
if usage:
response["usage"]["input_tokens"] = usage.get("input_tokens", 0)
for key in ("cache_creation_input_tokens", "cache_read_input_tokens"):
if key in usage:
response["usage"][key] = usage[key]
def _handle_content_block_start(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", len(content_blocks))
block = data.get("content_block", {})
block_type = block.get("type", "text")
_BLOCK_TEMPLATES: Dict[str, Dict] = {
"text": {"type": "text", "text": ""},
"thinking": {"type": "thinking", "thinking": "", "signature": ""},
"redacted_thinking": {
"type": "redacted_thinking",
"data": block.get("data", ""),
},
}
if block_type == "tool_use":
content_blocks[idx] = {
"type": "tool_use",
"id": block.get("id", ""),
"name": block.get("name", ""),
"input": {},
"_partial_json": "",
}
elif block_type in _BLOCK_TEMPLATES:
content_blocks[idx] = dict(_BLOCK_TEMPLATES[block_type])
else:
content_blocks[idx] = dict(block)
def _handle_content_block_delta(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
delta = data.get("delta", {})
delta_type = delta.get("type", "")
block = content_blocks.get(idx)
if block is None:
return
if delta_type == "text_delta":
block["text"] = block.get("text", "") + delta.get("text", "")
elif delta_type == "input_json_delta":
block["_partial_json"] = block.get("_partial_json", "") + delta.get(
"partial_json", ""
)
elif delta_type == "thinking_delta":
block["thinking"] = block.get("thinking", "") + delta.get("thinking", "")
elif delta_type == "signature_delta":
block["signature"] = delta.get("signature", block.get("signature", ""))
def _handle_content_block_stop(data: Dict, content_blocks: Dict[int, Dict]) -> None:
idx = data.get("index", 0)
block = content_blocks.get(idx)
if block and block.get("type") == "tool_use":
partial = block.pop("_partial_json", "")
if partial:
try:
block["input"] = json.loads(partial)
except (json.JSONDecodeError, ValueError):
block["input"] = {"_raw": partial}
def _handle_message_delta(data: Dict, response: Dict) -> None:
delta = data.get("delta", {})
if "stop_reason" in delta:
response["stop_reason"] = delta["stop_reason"]
if "stop_sequence" in delta:
response["stop_sequence"] = delta["stop_sequence"]
usage = data.get("usage", {})
if usage.get("output_tokens") is not None:
response["usage"]["output_tokens"] = usage["output_tokens"]
for key in (
"input_tokens",
"cache_creation_input_tokens",
"cache_read_input_tokens",
):
if key in usage:
response["usage"][key] = usage[key]
class AgenticAnthropicStreamingIterator:
"""
Two-phase async iterator that enables agentic hooks on streaming
Anthropic Messages pass-through responses.
Phase 1: Yield raw SSE bytes from the upstream response while
accumulating them. When the inner iterator is exhausted,
rebuild the full Anthropic response dict and call agentic hooks.
Phase 2: If an agentic hook fires and returns a follow-up response
(streaming or non-streaming), yield those bytes to the caller.
"""
def __init__(
self,
completion_stream: AsyncIterator,
http_handler: Any,
model: str,
messages: List[Dict],
anthropic_messages_provider_config: Any,
anthropic_messages_optional_request_params: Dict,
logging_obj: Any,
custom_llm_provider: str,
kwargs: Dict,
):
self._inner = completion_stream.__aiter__()
self._http_handler = http_handler
self._model = model
self._messages = messages
self._anthropic_messages_provider_config = anthropic_messages_provider_config
self._anthropic_messages_optional_request_params = (
anthropic_messages_optional_request_params
)
self._logging_obj = logging_obj
self._custom_llm_provider = custom_llm_provider
self._kwargs = kwargs
self._collected_bytes: List[bytes] = []
self._stream_exhausted = False
self._hook_processing_done = False
self._follow_up_iterator: Optional[AsyncIterator] = None
def __aiter__(self):
return self
async def __anext__(self) -> bytes:
# Phase 1: yield from upstream, collect bytes
if not self._stream_exhausted:
try:
chunk = await self._inner.__anext__()
self._collected_bytes.append(chunk)
return chunk
except StopAsyncIteration:
self._stream_exhausted = True
await self._process_agentic_hooks()
# Fall through to Phase 2
# Phase 2: yield from follow-up stream if one was created
if self._follow_up_iterator is not None:
chunk = await self._follow_up_iterator.__anext__()
return chunk
raise StopAsyncIteration
async def _process_agentic_hooks(self) -> None:
"""Rebuild the Anthropic response from collected SSE bytes and call hooks."""
if self._hook_processing_done:
return
self._hook_processing_done = True
if not self._collected_bytes:
return
try:
rebuilt = self._rebuild_anthropic_response_from_sse(self._collected_bytes)
if rebuilt is None:
verbose_logger.debug(
"AgenticStreamingIterator: Could not rebuild response from SSE bytes"
)
return
[
(
f"{b.get('type')}({b.get('name', '')})"
if b.get("type") == "tool_use"
else b.get("type")
)
for b in rebuilt.get("content", [])
]
result = await self._http_handler._call_agentic_completion_hooks(
response=rebuilt,
model=self._model,
messages=self._messages,
anthropic_messages_provider_config=self._anthropic_messages_provider_config,
anthropic_messages_optional_request_params=self._anthropic_messages_optional_request_params,
logging_obj=self._logging_obj,
stream=True,
custom_llm_provider=self._custom_llm_provider,
kwargs=self._kwargs,
)
if result is None:
return
if hasattr(result, "__aiter__"):
self._follow_up_iterator = result.__aiter__()
elif isinstance(result, dict):
from litellm.llms.anthropic.experimental_pass_through.messages.fake_stream_iterator import (
FakeAnthropicMessagesStreamIterator,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
fake = FakeAnthropicMessagesStreamIterator(
response=cast(AnthropicMessagesResponse, result)
)
self._follow_up_iterator = fake.__aiter__()
else:
verbose_logger.warning(
"AgenticStreamingIterator: Unexpected result type from hooks: %s",
type(result).__name__,
)
except Exception as e:
_call_id = 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),
)
@staticmethod
def _rebuild_anthropic_response_from_sse(
raw_bytes: List[bytes],
) -> Optional[Dict[str, Any]]:
"""
Parse collected SSE bytes into an Anthropic Messages response dict.
Processes SSE events in order:
- message_start -> envelope (id, model, role, usage)
- content_block_start -> new content block
- content_block_delta -> accumulate text/json/thinking deltas
- content_block_stop -> finalize block
- message_delta -> stop_reason, output usage
- message_stop -> end
"""
events = _parse_sse_events(b"".join(raw_bytes))
response: Dict[str, Any] = {
"id": "",
"type": "message",
"role": "assistant",
"model": "",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
content_blocks: Dict[int, Dict[str, Any]] = {}
saw_message_start = False
for event_type, data in events:
if event_type == "message_start":
saw_message_start = True
_handle_message_start(data, response)
elif event_type == "content_block_start":
_handle_content_block_start(data, content_blocks)
elif event_type == "content_block_delta":
_handle_content_block_delta(data, content_blocks)
elif event_type == "content_block_stop":
_handle_content_block_stop(data, content_blocks)
elif event_type == "message_delta":
_handle_message_delta(data, response)
if not saw_message_start:
return None
for idx in sorted(content_blocks.keys()):
block = content_blocks[idx]
block.pop("_partial_json", None)
response["content"].append(block)
return response

View file

@ -5,6 +5,7 @@ if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues
class BaseTranslation(ABC):
@ -101,6 +102,16 @@ class BaseTranslation(ABC):
"""
return responses_so_far
def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]:
"""
Convert request data to OpenAI-spec structured messages.
Override in subclasses for format-specific conversion.
Returns None if no convertible content is found.
"""
return None
def extract_request_tool_names(self, data: dict) -> List[str]:
"""
Extract tool names from the request body for allowlist/policy checks.

View file

@ -78,6 +78,10 @@ from litellm.types.containers.main import (
DeleteContainerResult,
)
from litellm.types.files import TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
AgenticLoopPlan,
AgenticLoopRequestPatch,
)
from litellm.types.llms.anthropic_messages.anthropic_response import (
AnthropicMessagesResponse,
)
@ -2047,7 +2051,23 @@ class BaseLLMHTTPHandler:
request_body=request_body,
litellm_logging_obj=logging_obj,
)
initial_response = completion_stream
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
)
initial_response = AgenticAnthropicStreamingIterator(
completion_stream=completion_stream,
http_handler=self,
model=model,
messages=messages,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
return initial_response
else:
initial_response = anthropic_messages_provider_config.transform_anthropic_messages_response(
model=model,
@ -2055,7 +2075,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
)
# Call agentic completion hooks
# Call agentic completion hooks (non-streaming path only)
final_response = await self._call_agentic_completion_hooks(
response=initial_response,
model=model,
@ -2063,7 +2083,7 @@ class BaseLLMHTTPHandler:
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream or False,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
@ -4516,6 +4536,167 @@ class BaseLLMHTTPHandler:
return stream, data
return stream, data
@staticmethod
def _get_agentic_loop_settings(kwargs: Dict) -> Tuple[int, int, List[str]]:
depth = int(kwargs.get("_agentic_loop_depth", 0) or 0)
max_loops = int(kwargs.get("max_agentic_loops", 3) or 3)
fingerprints = list(kwargs.get("_agentic_loop_fingerprints", []) or [])
return depth, max(max_loops, 1), fingerprints
@staticmethod
def _check_agentic_loop_safety(
tool_calls: Any,
fingerprints: List[str],
depth: int,
max_loops: int,
model: str,
) -> str:
"""
Evaluate agentic-loop safety guards (fingerprint cycle / max depth).
Raises ValueError on abort. Returns the current fingerprint on success.
These checks must not be swallowed by the per-callback ``except Exception``
block that wraps callback dispatch they are bounded-loop / cycle-break
safety rails and must abort the agentic dispatch when they trip.
"""
fingerprint = BaseLLMHTTPHandler._fingerprint_agentic_tools(tool_calls)
if fingerprint in fingerprints:
raise ValueError(
"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}"
)
return fingerprint
@staticmethod
def _fingerprint_agentic_tools(tools: Dict) -> str:
try:
return json.dumps(tools, sort_keys=True, default=str)
except Exception:
return str(tools)
async def _execute_anthropic_agentic_plan(
self,
plan: AgenticLoopPlan,
model: str,
messages: List[Dict],
anthropic_messages_optional_request_params: Dict,
logging_obj: "LiteLLMLoggingObj",
kwargs: Dict,
depth: int,
max_loops: int,
fingerprints: List[str],
fingerprint: str,
stream: bool = False,
) -> Any:
from litellm.anthropic_interface import messages as anthropic_messages
patch = plan.request_patch or AgenticLoopRequestPatch()
if patch.messages is None:
raise ValueError("Agentic loop plan missing patched messages")
full_model_name = model
if logging_obj is not None:
agentic_params = logging_obj.model_call_details.get(
"agentic_loop_params", {}
)
full_model_name = cast(str, agentic_params.get("model", model))
optional_params = dict(anthropic_messages_optional_request_params)
optional_params.update(patch.optional_params)
if patch.tools is not None:
optional_params["tools"] = patch.tools
max_tokens = patch.max_tokens
if max_tokens is None:
max_tokens = cast(Optional[int], optional_params.pop("max_tokens", None))
else:
optional_params.pop("max_tokens", None)
if max_tokens is None:
max_tokens = cast(int, kwargs.get("max_tokens", 1024))
internal_keys = {"litellm_logging_obj"}
kwargs_for_followup = {
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_keys
and k not in optional_params
}
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]
return await anthropic_messages.acreate(
**{
"max_tokens": max_tokens,
"messages": patch.messages,
"model": patch.model or full_model_name,
"stream": stream,
**optional_params,
**kwargs_for_followup,
}
)
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,
) -> Any:
patch = plan.request_patch or AgenticLoopRequestPatch()
if patch.messages is None:
raise ValueError("Agentic loop plan missing patched messages")
full_model_name = patch.model or model
if "/" not in full_model_name:
full_model_name = f"{custom_llm_provider}/{full_model_name}"
optional_params_for_followup = 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 = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup = {
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
}
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]
return await litellm.acompletion(
model=full_model_name,
messages=patch.messages,
**optional_params_for_followup,
**kwargs_for_followup,
)
async def _call_agentic_completion_hooks(
self,
response: Any,
@ -4541,45 +4722,111 @@ class BaseLLMHTTPHandler:
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
tools = anthropic_messages_optional_request_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
for callback in callbacks:
if not isinstance(callback, CustomLogger):
continue
should_run: bool = False
tool_calls: Any = None
try:
if isinstance(callback, CustomLogger):
# First: Check if agentic loop should run
(
should_run,
tool_calls,
) = await callback.async_should_run_agentic_loop(
response=response,
# First: Check if agentic loop should run. Wrap in try/except
# to shield from buggy user callbacks — a callback crash should
# not abort the whole request.
(
should_run,
tool_calls,
) = await callback.async_should_run_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in "
"async_should_run_agentic_loop [call_id=%s model=%s]: %s",
_call_id,
model,
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,
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_agentic_loop_plan
is not CustomLogger.async_build_agentic_loop_plan
)
if not build_plan_overridden:
return await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
tools=tools,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs=kwargs_with_provider,
)
if should_run:
# Second: Execute agentic loop
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = (
custom_llm_provider
)
agentic_response = await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
# First hook that runs agentic loop wins
return agentic_response
plan = await callback.async_build_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=anthropic_messages_provider_config,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_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 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_anthropic_agentic_plan(
plan=plan,
model=model,
messages=messages,
anthropic_messages_optional_request_params=anthropic_messages_optional_request_params,
logging_obj=logging_obj,
kwargs=kwargs_with_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
stream=stream,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
@ -4653,52 +4900,104 @@ class BaseLLMHTTPHandler:
callbacks = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
tools = optional_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
# Check if callback has the chat completion agentic loop method
if not hasattr(
callback, "async_should_run_chat_completion_agentic_loop"
):
continue
if not isinstance(callback, CustomLogger):
continue
if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"):
continue
# First: Check if agentic loop should run
(
should_run,
tool_calls,
) = await callback.async_should_run_chat_completion_agentic_loop(
response=response,
should_run: bool = False
tool_calls: Any = 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,
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,
tools=tools,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
kwargs=kwargs_with_provider,
)
if should_run:
# Second: Execute agentic loop
# Add custom_llm_provider to kwargs so the agentic loop can reconstruct the full model name
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = (
custom_llm_provider
)
agentic_response = (
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
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,
custom_llm_provider=custom_llm_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
except Exception as e:
verbose_logger.exception(
f"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: {str(e)}"

View file

@ -294,9 +294,7 @@ class Authenticator:
access_token_url = os.getenv(
"GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL
)
client_id = os.getenv(
"GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID
)
client_id = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID)
for attempt in range(max_attempts):
try:

View file

@ -48,6 +48,17 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert chat completions request data to OpenAI-spec structured messages.
Messages are already in OpenAI format, so this is a simple extraction.
"""
messages = data.get("messages")
if messages is None:
return None
return cast(List[AllMessageValues], messages)
async def process_input_messages(
self,
data: dict,
@ -68,9 +79,6 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
tool_calls_to_check: List[ChatCompletionToolParam] = []
text_task_mappings: List[Tuple[int, Optional[int]]] = []
tool_call_task_mappings: List[Tuple[int, int]] = []
# text_task_mappings: Track (message_index, content_index) for each text
# content_index is None for string content, int for list content
# tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call
# Step 1: Extract all text content, images, and tool calls
for msg_idx, message in enumerate(messages):
@ -92,12 +100,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
inputs["images"] = images_to_check
if tool_calls_to_check:
inputs["tool_calls"] = tool_calls_to_check # type: ignore
if messages:
msg_list = cast(List[AllMessageValues], messages)
structured_messages = self.get_structured_messages(data)
if structured_messages:
inputs["structured_messages"] = (
openai_messages_without_system(msg_list)
openai_messages_without_system(structured_messages)
if skip_system
else msg_list
else structured_messages
)
# Pass tools (function definitions) to the guardrail
tools = data.get("tools")

View file

@ -43,6 +43,7 @@ from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionToolCallChunk,
ChatCompletionToolParam,
)
@ -70,6 +71,24 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]:
"""
Convert Responses API request data to OpenAI-spec structured messages.
Transforms `input` (string or ResponseInputParam) and optional
`instructions` into chat completion messages.
"""
input_data = data.get("input")
if input_data is None:
return None
messages = (
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input_data,
responses_api_request=data,
)
)
return cast(List[AllMessageValues], messages) if messages else None
async def process_input_messages(
self,
data: dict,
@ -86,12 +105,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if input_data is None:
return data
structured_messages = (
LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=input_data,
responses_api_request=data,
)
)
structured_messages = self.get_structured_messages(data)
# Handle simple string input
if isinstance(input_data, str):

View file

@ -0,0 +1,158 @@
"""
Support for Scaleway's OpenAI-compatible `/v1/audio/transcriptions` endpoint.
API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path-audio-create-an-audio-transcription
"""
from typing import List, Optional, Union
import httpx
from litellm.litellm_core_utils.audio_utils.utils import process_audio_file
from litellm.llms.base_llm.audio_transcription.transformation import (
AudioTranscriptionRequestData,
BaseAudioTranscriptionConfig,
)
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIAudioTranscriptionOptionalParams,
)
from litellm.types.utils import FileTypes, TranscriptionResponse
class ScalewayAudioTranscriptionException(BaseLLMException):
pass
class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig):
def get_supported_openai_params(
self, model: str
) -> List[OpenAIAudioTranscriptionOptionalParams]:
return [
"language",
"prompt",
"response_format",
"temperature",
"timestamp_granularities",
]
def map_openai_params(
self,
non_default_params: dict,
optional_params: dict,
model: str,
drop_params: bool,
) -> dict:
supported_params = self.get_supported_openai_params(model)
for k, v in non_default_params.items():
if k in supported_params:
optional_params[k] = v
return optional_params
def get_complete_url(
self,
api_base: Optional[str],
api_key: Optional[str],
model: str,
optional_params: dict,
litellm_params: dict,
stream: Optional[bool] = None,
) -> str:
api_base = (
"https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/")
)
return f"{api_base}/audio/transcriptions"
def get_error_class(
self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
) -> BaseLLMException:
return ScalewayAudioTranscriptionException(
message=error_message,
status_code=status_code,
headers=headers,
)
def validate_environment(
self,
headers: dict,
model: str,
messages: List[AllMessageValues],
optional_params: dict,
litellm_params: dict,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
) -> dict:
if api_key is None:
api_key = get_secret_str("SCW_SECRET_KEY")
if not api_key:
raise ScalewayAudioTranscriptionException(
message=(
"Scaleway API key not found. Pass `api_key=...` or set the "
"SCW_SECRET_KEY environment variable."
),
status_code=401,
headers={},
)
default_headers = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
}
default_headers.update(headers or {})
return default_headers
def transform_audio_transcription_request(
self,
model: str,
audio_file: FileTypes,
optional_params: dict,
litellm_params: dict,
) -> AudioTranscriptionRequestData:
processed_audio = process_audio_file(audio_file)
form_fields: dict = {"model": model}
for key in self.get_supported_openai_params(model):
value = optional_params.get(key)
if value is not None:
form_fields[key] = value
files = {
"file": (
processed_audio.filename,
processed_audio.file_content,
processed_audio.content_type,
)
}
return AudioTranscriptionRequestData(data=form_fields, files=files)
def transform_audio_transcription_response(
self,
raw_response: httpx.Response,
) -> TranscriptionResponse:
content_type = (raw_response.headers.get("content-type") or "").lower()
if "application/json" not in content_type:
return TranscriptionResponse(text=raw_response.text)
try:
response_json = raw_response.json()
except Exception:
raise ScalewayAudioTranscriptionException(
message=raw_response.text,
status_code=raw_response.status_code,
headers=raw_response.headers,
)
text = response_json.get("text") or ""
response = TranscriptionResponse(text=text)
if "segments" in response_json:
response["segments"] = response_json["segments"]
if "language" in response_json:
response["language"] = response_json["language"]
response._hidden_params = response_json
return response

View file

@ -79,7 +79,9 @@ class BasePassthroughUtils:
for header_name, header_value in request_headers.items():
if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX):
# Strip the 'x-pass-' prefix and normalize to lowercase
actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower()
actual_header_name = header_name[
len(PASS_THROUGH_HEADER_PREFIX) :
].lower()
if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any(
actual_header_name.startswith(p)
for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES

View file

@ -1950,7 +1950,7 @@
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_transcriptions": true,
"audio_speech": false,
"moderations": false,
"batches": false,

View file

@ -22,11 +22,21 @@ model_list:
output_cost_per_token: 10 # 100x standard ($10.00/1M = $0.00001)
# Anthropic model for /v1/messages test — 100x custom pricing
- model_name: "claude-sonnet-4-20250514"
- model_name: "claude-sonnet-4-6"
litellm_params:
model: anthropic/claude-sonnet-4-20250514
model: anthropic/claude-sonnet-4-6
api_key: os.environ/ANTHROPIC_API_KEY
model_info:
id: claude-sonnet-4-custom-pricing
input_cost_per_token: 0.0003 # 100x standard ($0.000003)
output_cost_per_token: 0.0015 # 100x standard ($0.000015)
output_cost_per_token: 0.0015 # 100x standard ($0.000015)
- model_name: my-auto
litellm_params:
model: auto_router/complexity_router
complexity_router_config:
tiers:
SIMPLE: "gpt-4.1-mini"
COMPLEX: claude-sonnet-4-6
tier_boundaries:
simple_medium: 0.30
complexity_router_default_model: small-model

View file

@ -3126,9 +3126,7 @@ async def _virtual_key_max_budget_alert_check(
alert_email_config: Optional[Dict[str, List[str]]] = (
_merge_budget_alert_email_configs(
global_cfg=litellm.default_key_max_budget_alert_emails,
per_key_cfg=(valid_token.metadata or {}).get(
"max_budget_alert_emails"
),
per_key_cfg=(valid_token.metadata or {}).get("max_budget_alert_emails"),
)
)
@ -3138,7 +3136,9 @@ async def _virtual_key_max_budget_alert_check(
(int(k) for k in alert_email_config if k.isdigit()),
default=None,
)
if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0):
if min_pct is None or valid_token.spend < valid_token.max_budget * (
min_pct / 100.0
):
return
call_info = CallInfo(
@ -3164,8 +3164,7 @@ async def _virtual_key_max_budget_alert_check(
else:
# Old path: existing single 80% threshold — completely unchanged
alert_threshold = (
valid_token.max_budget
* EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
if (

View file

@ -37,6 +37,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915
if isinstance(value, list):
imported_list: List[Any] = []
for callback in value: # ["presidio", <my-custom-callback>]
if isinstance(callback, str) and callback == "compression_interception":
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
compression_interception_obj = (
CompressionInterceptionLogger.initialize_from_proxy_config(
litellm_settings=litellm_settings,
callback_specific_params=callback_specific_params,
)
)
imported_list.append(compression_interception_obj)
continue
# check if callback is a custom logger compatible callback
if isinstance(callback, str):
callback = LoggingCallbackManager._add_custom_callback_generic_api_str(

View file

@ -306,7 +306,9 @@ def _health_check_deployment_is_wildcard(litellm_params: dict) -> bool:
return "*" in _deployment_model_string_for_health_check(litellm_params)
def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) -> Optional[int]:
def _resolve_health_check_max_tokens(
model_info: dict, litellm_params: dict
) -> Optional[int]:
"""
Pick max_tokens for the health check request.
@ -341,10 +343,7 @@ def _resolve_health_check_max_tokens(model_info: dict, litellm_params: dict) ->
return int(tokens_reasoning)
if not is_reasoning and tokens_non_reasoning is not None:
return int(tokens_non_reasoning)
if (
is_reasoning
and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None
):
if is_reasoning and BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING is not None:
return int(BACKGROUND_HEALTH_CHECK_MAX_TOKENS_REASONING)
if BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None:

View file

@ -1121,14 +1121,14 @@ async def _db_health_readiness_check():
return db_health_cache
except Exception as e:
db_health_cache = {"status": "disconnected", "last_updated": datetime.now()}
PrismaDBExceptionHandler.handle_db_exception(e)
if PrismaDBExceptionHandler.is_database_transport_error(e):
try:
verbose_proxy_logger.warning(
"_db_health_readiness_check: health_check failed, attempting reconnect"
)
await prisma_client.disconnect()
await prisma_client.connect()
await prisma_client.attempt_db_reconnect(
reason="health_readiness_check"
)
await prisma_client.health_check()
verbose_proxy_logger.info(
"_db_health_readiness_check: reconnect succeeded"

View file

@ -1570,9 +1570,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
user_api_key_project_id = standard_logging_metadata.get(
"user_api_key_project_id"
)
user_api_key_end_user_id = kwargs.get(
"user"
) or standard_logging_metadata.get("user_api_key_end_user_id")
user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get(
"user_api_key_end_user_id"
)
model_group = get_model_group_from_litellm_kwargs(kwargs)
# Get total tokens from response

View file

@ -1,5 +1,6 @@
import asyncio
import copy
import re
import time
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@ -28,6 +29,14 @@ _SPECIAL_HEADERS_CACHE = frozenset(
v.value.lower() for v in SpecialHeaders._member_map_.values()
)
# Matches any header of the form x-<something>-session-id (case-insensitive).
# Excludes the two explicit litellm headers which are handled with higher priority.
_GENERIC_SESSION_ID_HEADER_RE = re.compile(r"^x-.+-session-id$", re.IGNORECASE)
_EXPLICIT_SESSION_HEADERS = frozenset({"x-litellm-trace-id", "x-litellm-session-id"})
# Session-id values must be non-empty strings of alphanumerics, hyphens, or underscores
# (covers UUIDs and most common session-id formats).
_SESSION_ID_VALUE_RE = re.compile(r"^[a-zA-Z0-9_\-]{8,}$")
def _sanitize_for_log(value: Any) -> str:
"""
@ -115,13 +124,43 @@ def _get_metadata_variable_name(request: Request) -> str:
return "metadata"
def _extract_generic_session_id_from_headers(
normalized: Dict[str, str],
) -> Optional[str]:
"""
Scan a normalised (lower-cased keys) header dict for any header that looks
like ``x-<vendor>-session-id`` and whose value is a plausible session/trace
identifier (alphanumeric + hyphens/underscores, at least 8 chars).
The two explicit LiteLLM headers (``x-litellm-trace-id`` /
``x-litellm-session-id``) are excluded here because they are handled with
higher priority by the caller.
Example: ``x-claude-code-session-id: e96634a3-fa28-4083-b354-55542e2dca01``
"""
for key, value in normalized.items():
if (
key not in _EXPLICIT_SESSION_HEADERS
and _GENERIC_SESSION_ID_HEADER_RE.match(key)
and isinstance(value, str)
and _SESSION_ID_VALUE_RE.match(value)
):
return value
return None
def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str]:
"""
Extract chain id for call chaining from request headers.
x-litellm-trace-id and x-litellm-session-id are interchangeable; when both
are present, x-litellm-trace-id takes precedence. Header keys are matched
case-insensitively so this works with raw header dicts from any transport.
Priority order:
1. ``x-litellm-trace-id`` (explicit, highest priority)
2. ``x-litellm-session-id`` (explicit)
3. Any ``x-<vendor>-session-id`` header whose value looks like a session id
(alphanumeric / UUID, at least 8 chars). E.g. ``x-claude-code-session-id``.
Header keys are matched case-insensitively so this works with raw header
dicts from any transport.
Used by MCP (and other paths that have raw_headers but no Request) to set
litellm_trace_id/litellm_session_id for spend logs and logging consistency.
@ -129,8 +168,10 @@ def get_chain_id_from_headers(headers: Optional[Dict[str, str]]) -> Optional[str
if not headers:
return None
normalized = {k.lower(): v for k, v in headers.items() if isinstance(k, str)}
return normalized.get("x-litellm-trace-id") or normalized.get(
"x-litellm-session-id"
return (
normalized.get("x-litellm-trace-id")
or normalized.get("x-litellm-session-id")
or _extract_generic_session_id_from_headers(normalized)
)
@ -649,10 +690,8 @@ class LiteLLMProxyRequestSetup:
#########################################################################################
agent_id_from_header = headers.get("x-litellm-agent-id")
# x-litellm-trace-id and x-litellm-session-id are interchangeable for call chaining
chain_id = headers.get("x-litellm-trace-id") or headers.get(
"x-litellm-session-id"
)
# Explicit litellm headers take precedence; fall back to any x-*-session-id header.
chain_id = get_chain_id_from_headers(dict(headers))
if agent_id_from_header:
metadata_from_headers["agent_id"] = agent_id_from_header

View file

@ -2120,9 +2120,7 @@ async def delete_user(
for m in all_target_memberships:
if not m.organization_id:
continue
target_org_ids_by_user.setdefault(m.user_id, set()).add(
m.organization_id
)
target_org_ids_by_user.setdefault(m.user_id, set()).add(m.organization_id)
# check that all teams passed exist
for user_id in data.user_ids:
@ -2141,9 +2139,7 @@ async def delete_user(
# Org-admin may only delete users whose entire org membership is
# within their admin scope. A target with ANY org outside the
# caller's scope (or no org at all) requires PROXY_ADMIN.
if not target_org_ids or not target_org_ids.issubset(
caller_admin_org_ids
):
if not target_org_ids or not target_org_ids.issubset(caller_admin_org_ids):
raise HTTPException(
status_code=403,
detail={

View file

@ -1078,10 +1078,7 @@ async def organization_member_update(
LitellmUserRoles.PROXY_ADMIN.value,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value,
):
if (
user_api_key_dict.user_role
!= LitellmUserRoles.PROXY_ADMIN.value
):
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
raise HTTPException(
status_code=403,
detail={

View file

@ -1570,8 +1570,7 @@ async def update_team( # noqa: PLR0915
current_org_id = getattr(existing_team_row, "organization_id", None)
if (
data.organization_id != current_org_id
and user_api_key_dict.user_role
!= LitellmUserRoles.PROXY_ADMIN.value
and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
):
# Is the caller org_admin of the destination org?
caller_memberships = (

View file

@ -577,6 +577,12 @@ class ProxyInitializationHelpers:
help="Exit with error if database migration fails on startup.",
envvar="ENFORCE_PRISMA_MIGRATION_CHECK",
)
@click.option(
"--reload",
is_flag=True,
default=False,
help="Enable uvicorn hot reload (dev only). Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
)
def run_server( # noqa: PLR0915
host,
port,
@ -618,6 +624,7 @@ def run_server( # noqa: PLR0915
keepalive_timeout,
max_requests_before_restart,
enforce_prisma_migration_check: bool,
reload: bool,
):
if setup:
from litellm.setup_wizard import run_setup_wizard
@ -954,6 +961,9 @@ def run_server( # noqa: PLR0915
if loop_type:
uvicorn_args["loop"] = loop_type
if reload:
uvicorn_args["reload"] = True
uvicorn.run(
**uvicorn_args,
workers=num_workers,

View file

@ -200,12 +200,16 @@ if TYPE_CHECKING:
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from litellm.router_strategy.quality_router.quality_router import (
QualityRouter,
)
Span = Union[_Span, Any]
else:
Span = Any
AutoRouter = Any
ComplexityRouter = Any
QualityRouter = Any
PreRoutingHookResponse = Any
@ -464,6 +468,7 @@ class Router:
) # {"TEAM_ID": PatternMatchRouter}
self.auto_routers: Dict[str, "AutoRouter"] = {}
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
self.quality_routers: Dict[str, "QualityRouter"] = {}
# Initialize model_group_alias early since it's used in set_model_list
self.model_group_alias: Dict[str, Union[str, RouterModelGroupAliasItem]] = (
@ -5884,7 +5889,7 @@ class Router:
response = await response
## PROCESS RESPONSE HEADERS
response = await self.set_response_headers(
response=response, model_group=model_group
response=response, model_group=model_group, request_kwargs=kwargs
)
return response
@ -6814,6 +6819,8 @@ class Router:
"""
if litellm_params.model.startswith("auto_router/complexity_router"):
return False # This is handled by complexity_router
if litellm_params.model.startswith("auto_router/quality_router"):
return False # This is handled by quality_router
if litellm_params.model.startswith("auto_router/"):
return True
return False
@ -6920,6 +6927,58 @@ class Router:
)
self.complexity_routers[deployment.model_name] = complexity_router
def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""
Check if the deployment is a quality-router deployment.
Returns True if the litellm_params model starts with "auto_router/quality_router".
"""
if litellm_params.model.startswith("auto_router/quality_router"):
return True
return False
def init_quality_router_deployment(self, deployment: Deployment):
"""
Initialize the quality-router deployment.
Resolves the default model from either `quality_router_default_model` or
`quality_router_config["default_model"]`, then instantiates the
QualityRouter and stores it in `self.quality_routers`.
"""
# Import here to mirror the AutoRouter / ComplexityRouter init pattern
# and avoid circular imports.
from litellm.router_strategy.quality_router.quality_router import (
QualityRouter,
)
quality_router_config: Optional[dict] = (
deployment.litellm_params.quality_router_config
)
default_model: Optional[str] = (
deployment.litellm_params.quality_router_default_model
)
if default_model is None and quality_router_config:
default_model = quality_router_config.get("default_model")
if default_model is None:
raise ValueError(
"quality_router_default_model is required for quality-router deployments, "
"or set default_model in quality_router_config. Please configure it in the litellm_params"
)
quality_router: QualityRouter = QualityRouter(
model_name=deployment.model_name,
default_model=default_model,
litellm_router_instance=self,
quality_router_config=quality_router_config,
)
if deployment.model_name in self.quality_routers:
raise ValueError(
f"Quality-router deployment {deployment.model_name} already exists. Please use a different model name."
)
self.quality_routers[deployment.model_name] = quality_router
def deployment_is_active_for_environment(self, deployment: Deployment) -> bool:
"""
Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments
@ -6966,6 +7025,11 @@ class Router:
self.model_id_to_deployment_index_map = {} # Reset the index
self.model_name_to_deployment_indices = {} # Reset the model_name index
self.team_model_to_deployment_indices = {} # Reset the team_model index
# Reset per-strategy router registries so hot-reload doesn't leave
# stale routers pointing at the old model_list.
self.quality_routers = {}
self.complexity_routers = {}
self.auto_routers = {}
self._invalidate_model_group_info_cache()
self._invalidate_access_groups_cache()
# we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works
@ -7140,6 +7204,12 @@ class Router:
):
self.init_complexity_router_deployment(deployment=deployment)
#########################################################
# Check if this is a quality-router deployment
#########################################################
if self._is_quality_router_deployment(litellm_params=deployment.litellm_params):
self.init_quality_router_deployment(deployment=deployment)
return deployment
def _initialize_deployment_for_pass_through(
@ -8143,7 +8213,10 @@ class Router:
return returned_dict
async def set_response_headers(
self, response: Any, model_group: Optional[str] = None
self,
response: Any,
model_group: Optional[str] = None,
request_kwargs: Optional[dict] = None,
) -> Any:
"""
Add the most accurate rate limit headers for a given model response.
@ -8164,6 +8237,45 @@ class Router:
additional_headers = response._hidden_params["additional_headers"] # type: ignore
# Lift QualityRouter routing decision into response headers for
# transparency. The decision is stashed in request_kwargs.metadata
# by QualityRouter.async_pre_routing_hook.
metadata = (
(request_kwargs.get("metadata") or {})
if isinstance(request_kwargs, dict)
else {}
)
decision = (
metadata.get("quality_router_decision")
if isinstance(metadata, dict)
else None
)
if isinstance(decision, dict):
# Only emit headers for fields that have a meaningful value.
# `complexity_tier` and `matched_keyword` are mutually exclusive
# (the keyword path short-circuits classification), so each
# request emits one or the other but not both.
if decision.get("routed_model") is not None:
additional_headers["x-litellm-quality-router-model"] = str(
decision["routed_model"]
)
if decision.get("quality_tier") is not None:
additional_headers["x-litellm-quality-router-tier"] = str(
decision["quality_tier"]
)
if decision.get("routed_via") is not None:
additional_headers["x-litellm-quality-router-via"] = str(
decision["routed_via"]
)
if decision.get("matched_keyword") is not None:
additional_headers["x-litellm-quality-router-keyword"] = str(
decision["matched_keyword"]
)
if decision.get("complexity_tier") is not None:
additional_headers["x-litellm-quality-router-complexity"] = str(
decision["complexity_tier"]
)
if (
"x-ratelimit-remaining-tokens" not in additional_headers
and "x-ratelimit-remaining-requests" not in additional_headers
@ -8708,8 +8820,6 @@ class Router:
and self.routing_strategy == "latency-based-routing"
):
_settings_to_return[var] = self.lowestlatency_logger.routing_args.json()
elif var == "routing_strategy_args":
_settings_to_return[var] = None
return _settings_to_return
def update_settings(self, **kwargs):
@ -9620,7 +9730,7 @@ class Router:
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional[PreRoutingHookResponse]:
@ -9653,6 +9763,18 @@ class Router:
specific_deployment=specific_deployment,
)
#########################################################
# Check if any quality-router should be used
#########################################################
if model in self.quality_routers:
return await self.quality_routers[model].async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=messages,
input=input,
specific_deployment=specific_deployment,
)
return None
def get_available_deployment(

View file

@ -82,11 +82,34 @@ class AutoRouter(CustomLogger):
)
return auto_router_routes
@staticmethod
def _extract_text_from_messages(messages: List[Dict[str, Any]]) -> str:
"""
Extract text content from the last user message for routing.
Handles tool-call conversations (where the last message may be an
assistant or tool message with non-string content) and multimodal
messages (where content is a list of content blocks).
"""
for msg in reversed(messages):
if msg.get("role") == "user":
content = msg.get("content")
if content is None:
return ""
if isinstance(content, list):
return " ".join(
block.get("text", "")
for block in content
if isinstance(block, dict) and block.get("type") == "text"
)
return str(content)
return ""
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, str]]] = None,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
@ -120,8 +143,7 @@ class AutoRouter(CustomLogger):
auto_sync=self.auto_sync_value,
)
user_message: Dict[str, str] = messages[-1]
message_content: str = user_message.get("content", "")
message_content = self._extract_text_from_messages(messages)
route_choice: Optional[Union[RouteChoice, List[RouteChoice]]] = self.routelayer(
text=message_content
)

View file

@ -332,45 +332,68 @@ class ComplexityRouter(CustomLogger):
f"No model configured for tier {tier_key} and no default_model set"
)
async def async_pre_routing_hook(
def _resolve_messages(
self,
model: str,
messages: Optional[List[Dict[str, Any]]],
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
) -> Optional[List[Dict[str, Any]]]:
"""
Pre-routing hook called before the routing decision.
Resolve messages from the request, converting from other formats if needed.
Classifies the request by complexity and returns the appropriate model.
Args:
model: The original model name requested.
request_kwargs: The request kwargs.
messages: The messages in the request.
input: Optional input for embeddings.
specific_deployment: Whether a specific deployment was requested.
Returns:
PreRoutingHookResponse with the routed model, or None if no routing needed.
Uses the guardrail translation handler dispatch to convert Responses API
``input`` (or other non-chat-completions formats) into OpenAI-spec messages.
"""
from litellm.types.router import PreRoutingHookResponse
if messages:
return messages
if messages is None or len(messages) == 0:
verbose_router_logger.debug(
"ComplexityRouter: No messages provided, skipping routing"
)
return None
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
from litellm.llms import load_guardrail_translation_mappings
from litellm.types.utils import CallTypes
# Extract the last user message and the last system prompt
mappings = load_guardrail_translation_mappings()
call_type: Optional[CallTypes] = None
# 1. Try route-based inference from proxy metadata
route = request_kwargs.get("litellm_metadata", {}).get(
"user_api_key_request_route"
)
if route:
call_types_list = get_call_types_for_route(route)
if call_types_list:
for ct in call_types_list:
if ct in mappings:
call_type = ct
break
# 2. Fallback: try each mapped handler until one produces messages
handlers_to_try: List[Any] = []
if call_type is not None and call_type in mappings:
handlers_to_try.append(mappings[call_type]())
else:
handlers_to_try.extend(handler_cls() for handler_cls in mappings.values())
for handler in handlers_to_try:
structured = handler.get_structured_messages(request_kwargs)
if structured:
return [
msg if isinstance(msg, dict) else msg.model_dump() # type: ignore
for msg in structured
]
return None
@staticmethod
def _extract_user_message_and_system_prompt(
messages: List[Dict[str, Any]],
) -> Tuple[Optional[str], Optional[str]]:
"""Extract the last user message text and last system prompt from messages."""
user_message: Optional[str] = None
system_prompt: Optional[str] = None
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content") or ""
# content may be a list of content parts (e.g. [{"type": "text", "text": "..."}])
if isinstance(content, list):
text_parts = [
part.get("text", "")
@ -383,6 +406,52 @@ class ComplexityRouter(CustomLogger):
user_message = content
elif role == "system" and system_prompt is None:
system_prompt = content
if user_message is not None and system_prompt is not None:
break
return user_message, system_prompt
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
"""
Pre-routing hook called before the routing decision.
Classifies the request by complexity and returns the appropriate model.
Supports chat completions (messages), Responses API (input), and other
formats via the guardrail translation handler dispatch.
Args:
model: The original model name requested.
request_kwargs: The request kwargs.
messages: The messages in the request.
input: Optional input for Responses API or embeddings.
specific_deployment: Whether a specific deployment was requested.
Returns:
PreRoutingHookResponse with the routed model, or None if no routing needed.
"""
from litellm.types.router import PreRoutingHookResponse
resolved_messages = self._resolve_messages(messages, request_kwargs)
if not resolved_messages:
verbose_router_logger.debug(
"ComplexityRouter: No messages could be resolved, skipping routing"
)
return None
# Determine whether the original request used messages directly
has_original_messages = messages is not None and len(messages) > 0
user_message, system_prompt = self._extract_user_message_and_system_prompt(
resolved_messages
)
if user_message is None:
verbose_router_logger.debug(
@ -391,13 +460,10 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=self.config.default_model
or self.get_model_for_tier(ComplexityTier.MEDIUM),
messages=messages,
messages=messages if has_original_messages else None,
)
# Classify the request
tier, score, signals = self.classify(user_message, system_prompt)
# Get the model for this tier
routed_model = self.get_model_for_tier(tier)
verbose_router_logger.info(
@ -407,5 +473,5 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
messages=messages if has_original_messages else None,
)

View file

@ -0,0 +1,21 @@
"""
Quality-tier auto-router.
Re-uses the ComplexityRouter's classification to decide a request's complexity,
then maps that complexity to an admin-configured quality tier and resolves the
target model from each candidate's `model_info.litellm_routing_preferences`.
"""
from .config import (
DEFAULT_COMPLEXITY_TO_QUALITY,
QualityRouterConfig,
RoutingPreferences,
)
from .quality_router import QualityRouter
__all__ = [
"QualityRouter",
"QualityRouterConfig",
"RoutingPreferences",
"DEFAULT_COMPLEXITY_TO_QUALITY",
]

View file

@ -0,0 +1,74 @@
"""
Configuration models for the QualityRouter.
"""
from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict, Field
# Default mapping from ComplexityTier name (string) to quality tier (int).
# Higher tier = higher capability requirement.
DEFAULT_COMPLEXITY_TO_QUALITY: Dict[str, int] = {
"SIMPLE": 1,
"MEDIUM": 2,
"COMPLEX": 3,
"REASONING": 4,
}
class QualityRouterConfig(BaseModel):
"""Configuration for the QualityRouter."""
available_models: List[str] = Field(
default_factory=list,
description=(
"List of candidate model names this router may route to. Each model "
"must declare its quality_tier in model_info.litellm_routing_preferences."
),
)
default_model: Optional[str] = Field(
default=None,
description="Fallback model when no quality tier resolves.",
)
complexity_to_quality: Dict[str, int] = Field(
default_factory=lambda: DEFAULT_COMPLEXITY_TO_QUALITY.copy(),
description="Mapping from ComplexityTier name to quality tier (int).",
)
model_config = ConfigDict(extra="allow")
class RoutingPreferences(BaseModel):
"""Per-deployment routing preferences declared on model_info."""
quality_tier: int = Field(
...,
description="The quality tier this deployment satisfies.",
)
keywords: List[str] = Field(
default_factory=list,
description=(
"Substring keywords (case-insensitive) that, when present in the "
"user message, route the request to this deployment. See `order` "
"for explicit collision handling, otherwise ties fall through to "
"(highest quality_tier, then cheapest model_info.input_cost_per_token)."
),
)
order: Optional[int] = Field(
default=None,
description=(
"Explicit priority used to break ties between deployments at the "
"same quality tier. Lower values win. Applies both to keyword "
"collisions and to picking between multiple deployments at the "
"same quality_tier. Tiebreak order is "
"(quality_tier DESC, order ASC, input_cost_per_token ASC, "
"model_name ASC) — quality always wins first, then explicit "
"order, then price."
),
)
model_config = ConfigDict(extra="allow")

View file

@ -0,0 +1,446 @@
"""
Quality-tier Auto Router.
Routes a request to a model at a target quality tier. The quality tier is
inferred by re-using the existing ComplexityRouter's classification, then
mapped through an admin-configured `complexity_to_quality` table. Each
candidate model declares its own `quality_tier` in
`model_info.litellm_routing_preferences`.
Optional keyword override: deployments may also declare `keywords` in
`litellm_routing_preferences`. If any declared keyword appears in the user
message (case-insensitive substring match), the router short-circuits the
complexity-classification flow and routes to the matching deployment. When
multiple deployments match, ties are broken by (highest quality_tier first,
then cheapest `model_info.input_cost_per_token`).
"""
import math
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
from litellm._logging import verbose_router_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from .config import QualityRouterConfig, RoutingPreferences
if TYPE_CHECKING:
from litellm.router import Router
from litellm.types.router import PreRoutingHookResponse
else:
Router = Any
PreRoutingHookResponse = Any
class QualityRouter(CustomLogger):
"""
Routes requests to a model at a target quality tier, with an optional
keyword override.
"""
def __init__(
self,
model_name: str,
litellm_router_instance: "Router",
default_model: Optional[str] = None,
quality_router_config: Optional[Dict[str, Any]] = None,
):
self.model_name = model_name
self.litellm_router_instance = litellm_router_instance
if quality_router_config:
self.config = QualityRouterConfig(**quality_router_config)
else:
self.config = QualityRouterConfig()
# Explicit default_model arg overrides anything in the config dict.
if default_model:
self.config.default_model = default_model
# Internal scorer — re-use the existing rule-based classifier.
self._scorer = ComplexityRouter(
model_name=f"{model_name}::scorer",
litellm_router_instance=litellm_router_instance,
)
# Per-model indices populated alongside the tier index. `_model_keywords`
# stores keywords lowercased so we can substring-match against the
# lowercased user message in O(total-keyword-count). `_model_quality`,
# `_model_cost`, and `_model_order` drive tiebreaking — `_model_order`
# is the explicit priority (lower wins, unset = +inf).
self._model_keywords: Dict[str, List[str]] = {}
self._model_quality: Dict[str, int] = {}
self._model_cost: Dict[str, Optional[float]] = {}
self._model_order: Dict[str, Optional[int]] = {}
# Tier → models index. Built lazily on first access so the QualityRouter
# deployment does NOT need to appear after all its referenced models in
# the config — when `_build_tier_index` runs eagerly in `__init__`, the
# router instance's `model_list` is still being assembled incrementally
# by `_create_deployment`, and any `available_models` defined AFTER the
# router entry in config.yaml would silently be reported as missing.
self._tier_to_models_cache: Optional[Dict[int, List[str]]] = None
verbose_router_logger.debug(
f"QualityRouter initialized for {model_name} with "
f"available_models={self.config.available_models}, "
f"default_model={self.config.default_model}"
)
@property
def _tier_to_models(self) -> Dict[int, List[str]]:
"""Lazy tier→models index; built on first access."""
if self._tier_to_models_cache is None:
self._tier_to_models_cache = self._build_tier_index()
return self._tier_to_models_cache
def _get_routing_preferences(self, deployment: Any) -> Optional[Dict[str, Any]]:
"""
Extract litellm_routing_preferences from a deployment, handling both
dict-shaped and Pydantic-object-shaped deployments.
"""
# Dict-shaped deployment.
if isinstance(deployment, dict):
model_info = deployment.get("model_info") or {}
if isinstance(model_info, dict):
return model_info.get("litellm_routing_preferences")
# Pydantic ModelInfo nested in a dict.
return getattr(model_info, "litellm_routing_preferences", None)
# Pydantic-object deployment.
model_info = getattr(deployment, "model_info", None)
if model_info is None:
return None
if isinstance(model_info, dict):
return model_info.get("litellm_routing_preferences")
return getattr(model_info, "litellm_routing_preferences", None)
def _get_deployment_input_cost(self, deployment: Any) -> Optional[float]:
"""
Extract `input_cost_per_token` from a deployment's model_info.
Returns None when not declared None is treated as "infinite cost"
for the cheapest-tiebreak ordering, so unpriced models lose ties to
priced ones. (Admins who want a model to win on price must declare it.)
"""
if isinstance(deployment, dict):
model_info = deployment.get("model_info") or {}
else:
model_info = getattr(deployment, "model_info", None) or {}
if isinstance(model_info, dict):
cost = model_info.get("input_cost_per_token")
else:
cost = getattr(model_info, "input_cost_per_token", None)
if cost is None:
return None
try:
return float(cost)
except (TypeError, ValueError):
return None
def _get_deployment_model_name(self, deployment: Any) -> Optional[str]:
"""Extract `model_name` from a dict- or object-shaped deployment."""
if isinstance(deployment, dict):
return deployment.get("model_name")
return getattr(deployment, "model_name", None)
def _build_tier_index(self) -> Dict[int, List[str]]:
"""
Build {quality_tier: [model_name, ...]} for every model in
`available_models`, plus side indices `_model_keywords`,
`_model_quality`, and `_model_cost`. Raises if any listed model is
missing `litellm_routing_preferences`.
"""
model_list = getattr(self.litellm_router_instance, "model_list", None) or []
available = set(self.config.available_models)
# Track which available models we've matched so we can error on missing.
seen: Dict[str, bool] = {name: False for name in available}
tier_to_models: Dict[int, List[str]] = {}
for deployment in model_list:
name = self._get_deployment_model_name(deployment)
if name is None or name not in available:
continue
raw_prefs = self._get_routing_preferences(deployment)
if raw_prefs is None:
raise ValueError(
f"QualityRouter: model '{name}' is listed in available_models "
f"but has no model_info.litellm_routing_preferences"
)
# Validate via the Pydantic model so we get a clear error for
# missing quality_tier, wrong types, etc. This also means
# `RoutingPreferences` is the single source of truth for the
# accepted shape — readers relied on raw dicts before.
try:
if isinstance(raw_prefs, RoutingPreferences):
prefs = raw_prefs
elif isinstance(raw_prefs, dict):
prefs = RoutingPreferences(**raw_prefs)
else:
# A Pydantic object of some other shape — coerce via its dict.
prefs = RoutingPreferences(
**(
raw_prefs.model_dump()
if hasattr(raw_prefs, "model_dump")
else dict(raw_prefs)
)
)
except Exception as e:
raise ValueError(
f"QualityRouter: model '{name}' has invalid "
f"litellm_routing_preferences: {e}"
) from e
tier_int = int(prefs.quality_tier)
tier_to_models.setdefault(tier_int, []).append(name)
self._model_keywords[name] = [str(k).lower() for k in prefs.keywords if k]
self._model_quality[name] = tier_int
self._model_cost[name] = self._get_deployment_input_cost(deployment)
self._model_order[name] = prefs.order
seen[name] = True
missing = [name for name, found in seen.items() if not found]
if missing:
raise ValueError(
f"QualityRouter: the following available_models are not present in "
f"the router's model_list (or are missing routing preferences): {missing}"
)
# Sort each tier's model list so `_resolve_model_for_quality_tier`
# (which picks index [0]) honors (order ASC, cost ASC, name ASC).
# Quality is moot within a single tier; keep parity with the keyword
# tiebreak by ordering on (order, cost, name) here.
for models in tier_to_models.values():
models.sort(key=lambda n: (self._order_key(n), self._cost_key(n), n))
return tier_to_models
def _order_key(self, model_name: str) -> float:
"""`order` lookup as a float — unset becomes +inf so explicit wins."""
order = self._model_order.get(model_name)
return float(order) if order is not None else math.inf
def _cost_key(self, model_name: str) -> float:
"""`input_cost_per_token` as a float — unset becomes +inf."""
cost = self._model_cost.get(model_name)
return float(cost) if cost is not None else math.inf
def _keyword_override(self, user_message: str) -> Optional[Tuple[str, str]]:
"""
Find a deployment whose declared keywords appear in `user_message`.
Returns (model_name, matched_keyword) or None when no keyword matches.
When multiple deployments match, sorts by:
1. quality_tier DESC (best quality always wins first)
2. `order` ASC (explicit priority unset = +inf so explicit wins
within the same tier)
3. input_cost_per_token ASC (unpriced = +inf so priced wins)
4. model_name ASC (deterministic stability)
"""
# Touch the lazy index so `_model_keywords` / `_model_quality` /
# `_model_cost` / `_model_order` are populated.
_ = self._tier_to_models
text = user_message.lower()
matches: List[Tuple[str, str]] = [] # (model_name, matched_keyword)
for model_name, keywords in self._model_keywords.items():
for kw in keywords:
if kw and kw in text:
matches.append((model_name, kw))
break # one match per model is enough
if not matches:
return None
def sort_key(match: Tuple[str, str]) -> Tuple[int, float, float, str]:
name = match[0]
quality = self._model_quality.get(name, 0)
order_val = self._order_key(name)
cost = self._model_cost.get(name)
cost_val = cost if cost is not None else math.inf
# Negate quality so higher tier sorts first under ASC sort.
return (-quality, order_val, cost_val, name)
matches.sort(key=sort_key)
return matches[0]
def _resolve_model_for_quality_tier(self, tier: int) -> str:
"""
Resolve a quality tier to a concrete model name.
Strategy:
1. Exact tier match first model registered at that tier.
2. Round UP to the next higher tier that has a model (closer to a
request we might lack capacity for).
3. Round DOWN to the closest lower tier that has a model (degrade
gracefully instead of jumping straight to `default_model`,
which may be off-tier).
4. Fall back to `config.default_model`.
5. Otherwise raise.
"""
tier_index = self._tier_to_models
if tier in tier_index and tier_index[tier]:
return tier_index[tier][0]
# Round up.
higher_tiers = sorted(t for t in tier_index if t > tier)
for t in higher_tiers:
if tier_index[t]:
return tier_index[t][0]
# Round down — closest lower tier first.
lower_tiers = sorted((t for t in tier_index if t < tier), reverse=True)
for t in lower_tiers:
if tier_index[t]:
return tier_index[t][0]
if self.config.default_model:
return self.config.default_model
raise ValueError(
f"QualityRouter: no model available for quality tier {tier} and "
f"no default_model configured"
)
def _stash_decision(
self,
request_kwargs: Optional[Dict[str, Any]],
decision: Dict[str, Any],
) -> None:
"""
Stash the routing decision in request_kwargs.metadata so the Router can
lift it into response headers (`x-litellm-quality-router-*`). The same
dict object flows from here through to `make_call.set_response_headers`.
"""
if request_kwargs is None:
return
metadata = request_kwargs.setdefault("metadata", {})
if isinstance(metadata, dict):
metadata["quality_router_decision"] = decision
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: Dict,
messages: Optional[List[Dict[str, Any]]] = None,
input: Optional[Union[str, List]] = None,
specific_deployment: Optional[bool] = False,
) -> Optional["PreRoutingHookResponse"]:
"""Try keyword override first; fall back to complexity-tier routing."""
from litellm.types.router import PreRoutingHookResponse
if messages is None or len(messages) == 0:
verbose_router_logger.debug(
"QualityRouter: No messages provided, skipping routing"
)
return None
# Extract last user message and last system prompt — same rules as
# ComplexityRouter.async_pre_routing_hook.
user_message: Optional[str] = None
system_prompt: Optional[str] = None
for msg in reversed(messages):
role = msg.get("role", "")
content = msg.get("content") or ""
if isinstance(content, list):
text_parts = [
part.get("text", "")
for part in content
if isinstance(part, dict) and part.get("type") == "text"
]
content = " ".join(text_parts).strip()
if isinstance(content, str) and content:
if role == "user" and user_message is None:
user_message = content
elif role == "system" and system_prompt is None:
system_prompt = content
if user_message is None:
verbose_router_logger.debug(
"QualityRouter: No user message found, routing to default model"
)
if not self.config.default_model:
raise ValueError(
"QualityRouter: no user message and no default_model configured"
)
return PreRoutingHookResponse(
model=self.config.default_model,
messages=messages,
)
# Try keyword override first — it short-circuits complexity classification.
keyword_match = self._keyword_override(user_message)
if keyword_match is not None:
routed_model, matched_keyword = keyword_match
verbose_router_logger.info(
f"QualityRouter: keyword override matched='{matched_keyword}' "
f"routed_model={routed_model} "
f"(quality_tier={self._model_quality.get(routed_model)}, "
f"input_cost_per_token={self._model_cost.get(routed_model)})"
)
self._stash_decision(
request_kwargs,
{
"router_model_name": self.model_name,
"routed_model": routed_model,
"routed_via": "keyword",
"matched_keyword": matched_keyword,
"quality_tier": self._model_quality.get(routed_model),
"complexity_tier": None,
},
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
)
# No keyword match → complexity classification flow.
complexity_tier, score, signals = self._scorer.classify(
user_message, system_prompt
)
complexity_name = (
complexity_tier.value
if hasattr(complexity_tier, "value")
else str(complexity_tier)
)
quality_tier = self.config.complexity_to_quality.get(complexity_name)
if quality_tier is None:
raise ValueError(
f"QualityRouter: complexity tier '{complexity_name}' not present "
f"in complexity_to_quality mapping {self.config.complexity_to_quality}"
)
routed_model = self._resolve_model_for_quality_tier(int(quality_tier))
verbose_router_logger.info(
f"QualityRouter: complexity={complexity_name}, score={score:.3f}, "
f"signals={signals}, quality_tier={quality_tier}, "
f"routed_model={routed_model}"
)
self._stash_decision(
request_kwargs,
{
"router_model_name": self.model_name,
"routed_model": routed_model,
"routed_via": "quality_tier",
"matched_keyword": None,
"quality_tier": int(quality_tier),
"complexity_tier": complexity_name,
},
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
)

View file

@ -2,7 +2,14 @@
Type definitions for litellm.compress().
"""
from typing import Dict, List, TypedDict
import sys
if sys.version_info >= (3, 11):
from typing import Dict, List, NotRequired, TypedDict
else:
from typing import Dict, List, TypedDict
from typing_extensions import NotRequired
class CompressedResult(TypedDict):
@ -12,3 +19,4 @@ class CompressedResult(TypedDict):
compression_ratio: float # fraction reduced, e.g. 0.6 means 60% reduction
cache: Dict[str, str] # key -> original content (for retrieval tool responses)
tools: List[dict] # [litellm_content_retrieve tool definition]
compression_skipped_reason: NotRequired[str]

View file

@ -0,0 +1,27 @@
"""
Type definitions for Compression Interception integration.
"""
from typing import Any, Dict, Optional, TypedDict
class CompressionInterceptionConfig(TypedDict, total=False):
"""
Configuration parameters for CompressionInterceptionLogger.
Used in proxy_config.yaml under litellm_settings:
litellm_settings:
compression_interception_params:
enabled: true
compression_trigger: 100000
compression_target: 70000
embedding_model: "text-embedding-3-small"
embedding_model_params:
dimensions: 512
"""
enabled: bool
compression_trigger: int
compression_target: Optional[int]
embedding_model: Optional[str]
embedding_model_params: Optional[Dict[str, Any]]

View file

@ -1,6 +1,6 @@
from typing import Optional
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from pydantic import BaseModel, Field
class StandardCustomLoggerInitParams(BaseModel):
@ -9,3 +9,29 @@ class StandardCustomLoggerInitParams(BaseModel):
"""
turn_off_message_logging: Optional[bool] = False
class AgenticLoopRequestPatch(BaseModel):
"""
Patch returned by callbacks to request a follow-up LLM call.
"""
model: Optional[str] = None
messages: Optional[List[Dict[str, Any]]] = None
tools: Optional[List[Dict[str, Any]]] = None
max_tokens: Optional[int] = None
optional_params: Dict[str, Any] = Field(default_factory=dict)
kwargs: Dict[str, Any] = Field(default_factory=dict)
class AgenticLoopPlan(BaseModel):
"""
Typed callback response for agentic-loop reruns.
"""
run_agentic_loop: bool = False
request_patch: Optional[AgenticLoopRequestPatch] = None
response_override: Optional[Any] = None
terminate: bool = False
stop_reason: Optional[str] = None
metadata: Dict[str, Any] = Field(default_factory=dict)

View file

@ -784,7 +784,7 @@ class UserAPIKeyLabelValues:
org_id: Optional[str] = None
org_alias: Optional[str] = None
#Added for test compatibility.
# Added for test compatibility.
def __init__(self, **kwargs: Any) -> None:
"""
Match former Pydantic behavior: unknown keys are ignored; ``api_key_hash`` maps to

View file

@ -221,6 +221,10 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
complexity_router_config: Optional[Dict] = None
complexity_router_default_model: Optional[str] = None
# quality-router params
quality_router_config: Optional[Dict] = None
quality_router_default_model: Optional[str] = None
# Batch/File API Params
s3_bucket_name: Optional[str] = None
s3_encryption_key_id: Optional[str] = None

View file

@ -3290,6 +3290,7 @@ class LlmProviders(str, Enum):
MANUS = "manus"
WANDB = "wandb"
OVHCLOUD = "ovhcloud"
SCALEWAY = "scaleway"
LEMONADE = "lemonade"
AMAZON_NOVA = "amazon_nova"
A2A_AGENT = "a2a_agent"

View file

@ -8472,6 +8472,12 @@ class ProviderConfigManager:
)
return OVHCloudAudioTranscriptionConfig()
elif litellm.LlmProviders.SCALEWAY == provider:
from litellm.llms.scaleway.audio_transcription.transformation import (
ScalewayAudioTranscriptionConfig,
)
return ScalewayAudioTranscriptionConfig()
elif litellm.LlmProviders.MISTRAL == provider:
from litellm.llms.mistral.audio_transcription.transformation import (
MistralAudioTranscriptionConfig,

View file

@ -33302,6 +33302,7 @@
"output_cost_per_token": 1.5e-05,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_web_search": true
@ -33317,6 +33318,7 @@
"output_cost_per_token": 1.5e-05,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_web_search": true
@ -33332,6 +33334,7 @@
"output_cost_per_token": 2.5e-05,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_web_search": true
@ -33347,6 +33350,7 @@
"output_cost_per_token": 2.5e-05,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_web_search": true
@ -33362,6 +33366,7 @@
"output_cost_per_token": 1.5e-05,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": false,
"supports_tool_choice": true,
"supports_web_search": true
@ -33378,6 +33383,7 @@
"output_cost_per_token": 5e-07,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33395,6 +33401,7 @@
"output_cost_per_token": 5e-07,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33411,6 +33418,7 @@
"output_cost_per_token": 4e-06,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33427,6 +33435,7 @@
"output_cost_per_token": 4e-06,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33443,6 +33452,7 @@
"output_cost_per_token": 4e-06,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33459,6 +33469,7 @@
"output_cost_per_token": 5e-07,
"source": "https://x.ai/api#pricing",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": false,
"supports_tool_choice": true,
@ -33474,38 +33485,41 @@
"output_cost_per_token": 1.5e-05,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"xai/grok-4-fast-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_128k_tokens": 4e-07,
"litellm_provider": "xai",
"max_input_tokens": 2000000.0,
"max_output_tokens": 2000000.0,
"max_tokens": 2000000.0,
"mode": "chat",
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_128k_tokens": 4e-07,
"output_cost_per_token": 5e-07,
"output_cost_per_token_above_128k_tokens": 1e-06,
"cache_read_input_token_cost": 5e-08,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_web_search": true
},
"xai/grok-4-fast-non-reasoning": {
"cache_read_input_token_cost": 5e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_128k_tokens": 4e-07,
"litellm_provider": "xai",
"max_input_tokens": 2000000.0,
"max_output_tokens": 2000000.0,
"cache_read_input_token_cost": 5e-08,
"max_tokens": 2000000.0,
"mode": "chat",
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_128k_tokens": 4e-07,
"output_cost_per_token": 5e-07,
"output_cost_per_token_above_128k_tokens": 1e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_web_search": true
},
@ -33521,6 +33535,7 @@
"output_cost_per_token_above_128k_tokens": 3e-05,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_web_search": true
},
@ -33536,6 +33551,7 @@
"output_cost_per_token_above_128k_tokens": 3e-05,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_web_search": true
},
@ -33553,6 +33569,7 @@
"source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -33573,6 +33590,7 @@
"source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -33593,6 +33611,7 @@
"source": "https://docs.x.ai/docs/models/grok-4-1-fast-reasoning",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_tool_choice": true,
@ -33613,6 +33632,7 @@
"source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -33632,6 +33652,7 @@
"source": "https://docs.x.ai/docs/models/grok-4-1-fast-non-reasoning",
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -33648,6 +33669,7 @@
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -33664,6 +33686,7 @@
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
@ -33696,6 +33719,7 @@
"output_cost_per_token": 6e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": true
@ -33724,6 +33748,7 @@
"output_cost_per_token": 1.5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
@ -33738,6 +33763,7 @@
"output_cost_per_token": 1.5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
@ -33752,6 +33778,7 @@
"output_cost_per_token": 1.5e-06,
"source": "https://docs.x.ai/docs/models",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true
},

View file

@ -1968,7 +1968,7 @@
"responses": true,
"embeddings": false,
"image_generations": false,
"audio_transcriptions": false,
"audio_transcriptions": true,
"audio_speech": false,
"moderations": false,
"batches": false,

View file

@ -33,6 +33,7 @@ from dataclasses import asdict, dataclass, field
from typing import Optional
import litellm
from litellm.types.utils import CallTypes
# ---------------------------------------------------------------------------
# Problem definitions (HumanEval-style)
@ -880,6 +881,7 @@ def eval_problem(
result = litellm.compress(
messages=messages,
model=model,
call_type=CallTypes.completion,
compression_trigger=compression_trigger,
embedding_model=embedding_model,
)

View file

@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import litellm # noqa: E402
from litellm.compression import compress as litellm_compress # noqa: E402
from litellm.types.utils import CallTypes # noqa: E402
# ---------------------------------------------------------------------------
# Prompts
@ -445,7 +446,7 @@ def eval_instance(
compress_kwargs: dict = {
"messages": messages,
"model": model,
"input_type": "openai_chat_completions",
"call_type": CallTypes.completion,
"compression_trigger": compression_trigger,
"embedding_model": embedding_model,
}

View file

@ -72,7 +72,7 @@ def test_batch_completions_models():
def test_batch_completion_models_all_responses():
try:
responses = batch_completion_models_all_responses(
models=["gemini/gemini-2.5-flash-lite", "claude-3-haiku-20240307"],
models=["gemini/gemini-2.5-flash-lite", "claude-haiku-4-5-20251001"],
messages=[{"role": "user", "content": "write a poem"}],
max_tokens=10,
)

View file

@ -142,7 +142,7 @@ def trade(model_name: str) -> List[Trade]: # type: ignore
@pytest.mark.parametrize(
"model", ["claude-3-haiku-20240307", "anthropic.claude-3-haiku-20240307-v1:0"]
"model", ["claude-haiku-4-5-20251001", "anthropic.claude-3-haiku-20240307-v1:0"]
)
@pytest.mark.flaky(retries=6, delay=10)
def test_function_call_parsing(model):

View file

@ -47,7 +47,7 @@ def get_current_weather(location, unit="fahrenheit"):
[
"gpt-3.5-turbo-1106",
"mistral/mistral-large-latest",
"claude-3-haiku-20240307",
"claude-haiku-4-5-20251001",
"gemini/gemini-2.5-flash-lite",
"anthropic.claude-3-sonnet-20240229-v1:0",
],
@ -275,7 +275,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message
"anthropic.claude-3-sonnet-20240229-v1:0",
"bedrock",
),
("claude-3-haiku-20240307", "anthropic"),
("claude-haiku-4-5-20251001", "anthropic"),
],
)
@pytest.mark.parametrize(

View file

@ -1509,7 +1509,7 @@ def test_router_fallbacks_with_wildcard_model_name():
{
"model_name": "claude-3-haiku",
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
"mock_response": "Hi this is claude!",
},
@ -1555,7 +1555,7 @@ def test_fallbacks_with_different_messages():
{
"model_name": "claude-3-haiku",
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
},

View file

@ -1727,7 +1727,7 @@ def test_openai_chat_completion_complete_response_call():
"model",
[
"gpt-3.5-turbo",
"claude-3-haiku-20240307",
"claude-haiku-4-5-20251001",
"o1",
],
)
@ -2247,7 +2247,7 @@ def streaming_and_function_calling_format_tests(idx, chunk):
[
# "gpt-3.5-turbo",
# "anthropic.claude-3-sonnet-20240229-v1:0",
"claude-3-haiku-20240307",
"claude-haiku-4-5-20251001",
],
)
def test_streaming_and_function_calling(model):

View file

@ -77,7 +77,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest):
@property
def model_config(self) -> Dict[str, Any]:
return {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
}
@ -86,7 +86,7 @@ class TestAnthropicDirectAPI(BaseAnthropicMessagesTest):
"""
This is the model name that is expected to be in the logging payload
"""
return "claude-3-haiku-20240307"
return "claude-haiku-4-5-20251001"
class TestAnthropicBedrockAPI(BaseAnthropicMessagesTest):
@ -140,7 +140,7 @@ async def test_anthropic_messages_streaming_with_bad_request():
response = await litellm.anthropic.messages.acreate(
messages=[{"role": "user", "content": "hi"}],
api_key=os.getenv("ANTHROPIC_API_KEY"),
model="claude-3-haiku-20240307",
model="claude-haiku-4-5-20251001",
max_tokens=100,
stream=True,
)
@ -168,7 +168,7 @@ async def test_anthropic_messages_router_streaming_with_bad_request():
{
"model_name": "claude-special-alias",
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
}
@ -205,7 +205,7 @@ async def test_anthropic_messages_litellm_router_non_streaming():
{
"model_name": "claude-special-alias",
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
}
@ -243,7 +243,7 @@ async def test_anthropic_messages_litellm_router_routing_strategy():
{
"model_name": "claude-special-alias",
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
}
@ -341,7 +341,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking():
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "Here's a joke for you!"}],
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
@ -355,7 +355,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking():
{
"model_name": MODEL_GROUP,
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
}
@ -419,7 +419,7 @@ async def test_anthropic_messages_litellm_router_latency_metadata_tracking():
assert "model_info" in litellm_metadata
# Verify other call parameters
assert call_kwargs["model"] == "claude-3-haiku-20240307"
assert call_kwargs["model"] == "claude-haiku-4-5-20251001"
assert call_kwargs["messages"] == messages
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["metadata"] == {"user_id": "hello"}
@ -459,7 +459,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging():
{
"model_name": MODEL_GROUP,
"litellm_params": {
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
},
}
@ -496,7 +496,7 @@ async def test_anthropic_messages_litellm_router_non_streaming_with_logging():
assert test_custom_logger.logged_standard_logging_payload["response"] is not None
assert (
test_custom_logger.logged_standard_logging_payload["model"]
== "claude-3-haiku-20240307"
== "claude-haiku-4-5-20251001"
)
# check logged usage + spend
@ -543,7 +543,7 @@ async def test_anthropic_messages_with_extra_headers():
"text": "Why did the chicken cross the road? To get to the other side!",
}
],
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
@ -556,7 +556,7 @@ async def test_anthropic_messages_with_extra_headers():
response = await litellm.anthropic.messages.acreate(
messages=messages,
api_key=api_key,
model="claude-3-haiku-20240307",
model="claude-haiku-4-5-20251001",
max_tokens=100,
client=mock_client,
provider_specific_header={
@ -689,7 +689,7 @@ async def test_anthropic_messages_with_thinking():
"text": "Why did the chicken cross the road? To get to the other side!",
}
],
"model": "claude-3-haiku-20240307",
"model": "claude-haiku-4-5-20251001",
"stop_reason": "end_turn",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
@ -702,7 +702,7 @@ async def test_anthropic_messages_with_thinking():
response = await litellm.anthropic.messages.acreate(
messages=messages,
api_key=api_key,
model="claude-3-haiku-20240307",
model="claude-haiku-4-5-20251001",
max_tokens=100,
client=mock_client,
thinking={"budget_tokens": 100},
@ -717,7 +717,7 @@ async def test_anthropic_messages_with_thinking():
request_body = json.loads(call_kwargs.get("data", {}))
print("REQUEST BODY", request_body)
assert request_body["max_tokens"] == 100
assert request_body["model"] == "claude-3-haiku-20240307"
assert request_body["model"] == "claude-haiku-4-5-20251001"
assert request_body["messages"] == messages
assert request_body["thinking"] == {"budget_tokens": 100}

View file

@ -0,0 +1,364 @@
"""
Unit tests for Compression Interception Handler.
"""
from unittest.mock import MagicMock
import pytest
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
from litellm.types.utils import CallTypes
def test_initialize_from_proxy_config():
"""Test initialization from proxy config with litellm_settings."""
litellm_settings = {
"compression_interception_params": {
"enabled": True,
"compression_trigger": 1234,
"compression_target": 789,
}
}
logger = CompressionInterceptionLogger.initialize_from_proxy_config(
litellm_settings=litellm_settings,
callback_specific_params={},
)
assert logger.enabled is True
assert logger.compression_trigger == 1234
assert logger.compression_target == 789
@pytest.mark.asyncio
async def test_pre_call_hook_compresses_messages_and_injects_tool(monkeypatch):
"""Test pre-call hook compresses and stores per-call cache."""
logger = CompressionInterceptionLogger()
compressed_result = {
"messages": [{"role": "user", "content": "stubbed"}],
"original_tokens": 12000,
"compressed_tokens": 5000,
"compression_ratio": 0.58,
"cache": {"auth.py": "full file content"},
"tools": [
{
"type": "function",
"function": {
"name": "litellm_content_retrieve",
"parameters": {
"type": "object",
"properties": {"key": {"type": "string"}},
},
},
}
],
}
def _fake_compress(**kwargs):
return compressed_result
# The handler does ``from litellm.compression import compress`` at module
# scope, so we must patch the binding on the handler module — patching
# ``litellm.compress`` has no effect on the already-bound reference.
monkeypatch.setattr(
"litellm.integrations.compression_interception.handler.compress",
_fake_compress,
)
kwargs = {
"model": "bedrock/us.anthropic.claude-sonnet-4-5",
"messages": [{"role": "user", "content": "very large context"}],
"tools": [
{
"type": "function",
"function": {"name": "existing_tool", "parameters": {"type": "object"}},
}
],
}
result = await logger.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.anthropic_messages
)
assert result is not None
assert result["messages"] == compressed_result["messages"]
tool_names = [t.get("function", {}).get("name") for t in result["tools"]]
assert "existing_tool" in tool_names
assert "litellm_content_retrieve" in tool_names
assert result["litellm_call_id"] in logger._compression_cache_by_call_id
@pytest.mark.asyncio
async def test_pre_call_hook_below_trigger_does_not_inject_empty_tools(monkeypatch):
"""
When compression is a no-op (below trigger / invalid tool sequence), the
hook must NOT replace ``messages`` or inject an empty ``tools: []`` onto
a request that originally had no tools Anthropic Messages rejects
``tools: []``.
"""
logger = CompressionInterceptionLogger()
original_messages = [{"role": "user", "content": "short prompt"}]
def _fake_compress_noop(**kwargs):
return {
"messages": original_messages,
"original_tokens": 42,
"compressed_tokens": 42,
"compression_ratio": 0.0,
"cache": {},
"tools": [],
"compression_skipped_reason": "below_trigger",
}
monkeypatch.setattr(
"litellm.integrations.compression_interception.handler.compress",
_fake_compress_noop,
)
kwargs = {
"model": "bedrock/us.anthropic.claude-sonnet-4-5",
"messages": original_messages,
}
result = await logger.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.anthropic_messages
)
assert result is not None
# Original request had no ``tools`` — skipped compression must leave it that way.
assert "tools" not in result
# Cache must not be populated for a no-op.
assert result.get("litellm_call_id") not in logger._compression_cache_by_call_id
@pytest.mark.asyncio
async def test_should_run_agentic_loop_detects_retrieval_tool_use():
"""Test should-run hook returns tool calls for retrieval tool_use blocks."""
logger = CompressionInterceptionLogger()
response = {
"content": [
{
"type": "tool_use",
"id": "toolu_123",
"name": "litellm_content_retrieve",
"input": {"key": "auth.py"},
}
]
}
should_run, tools_dict = await logger.async_should_run_agentic_loop(
response=response,
model="bedrock/claude",
messages=[],
tools=[
{
"type": "function",
"function": {
"name": "litellm_content_retrieve",
"parameters": {"type": "object"},
},
}
],
stream=False,
custom_llm_provider="bedrock",
kwargs={},
)
assert should_run is True
assert len(tools_dict["tool_calls"]) == 1
assert tools_dict["tool_calls"][0]["input"]["key"] == "auth.py"
@pytest.mark.asyncio
async def test_build_agentic_loop_plan_returns_request_patch():
"""Callback should return typed patch with tool_result content."""
logger = CompressionInterceptionLogger()
call_id = "call_123"
logger._compression_cache_by_call_id[call_id] = (
{"auth.py": "full auth file"},
9999999999.0,
)
logging_obj = MagicMock()
logging_obj.litellm_call_id = call_id
logging_obj.model_call_details = {
"agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"}
}
plan = await logger.async_build_agentic_loop_plan(
tools={
"tool_calls": [
{
"id": "toolu_abc",
"type": "tool_use",
"name": "litellm_content_retrieve",
"input": {"key": "auth.py"},
}
]
},
model="claude-3-5-sonnet",
messages=[{"role": "user", "content": "read auth.py"}],
response=None,
anthropic_messages_provider_config=None,
anthropic_messages_optional_request_params={
"max_tokens": 1024,
"tools": [{"name": "litellm_content_retrieve"}],
},
logging_obj=logging_obj,
stream=False,
kwargs={
"temperature": 0.1,
"_compression_interception_internal": True,
"litellm_logging_obj": object(),
},
)
assert plan.run_agentic_loop is True
assert plan.request_patch is not None
assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet"
assert plan.request_patch.max_tokens == 1024
assert plan.request_patch.messages is not None
assert len(plan.request_patch.messages) == 3
tool_result_content = plan.request_patch.messages[-1]["content"][0]["content"]
assert tool_result_content == "full auth file"
assert "_compression_interception_internal" not in plan.request_patch.kwargs
assert "litellm_logging_obj" not in plan.request_patch.kwargs
assert plan.request_patch.kwargs["temperature"] == 0.1
assert "max_tokens" not in plan.request_patch.optional_params
@pytest.mark.asyncio
async def test_should_run_agentic_loop_with_custom_type_tools():
"""Test that async_should_run_agentic_loop returns True when tools contain
litellm_content_retrieve as a custom-typed tool (e.g. Claude Code tool list)
and the model response includes a matching tool_use block."""
logger = CompressionInterceptionLogger()
# Exact tools payload produced by Claude Code litellm_content_retrieve is
# the final entry and uses type="custom" (not type="function").
tools = [
{
"name": "Agent",
"description": "Launch a new agent to handle complex, multi-step tasks.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"description": {"type": "string"},
"prompt": {"type": "string"},
},
"required": ["description", "prompt"],
"additionalProperties": False,
},
},
{
"name": "AskUserQuestion",
"description": "Use this tool when you need to ask the user questions.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {
"questions": {"type": "array", "items": {"type": "object"}},
},
"required": ["questions"],
"additionalProperties": False,
},
},
{
"name": "Bash",
"description": "Executes a given bash command and returns its output.",
"input_schema": {
"$schema": "https://json-schema.org/draft/2020-12/schema",
"type": "object",
"properties": {"command": {"type": "string"}},
"required": ["command"],
"additionalProperties": False,
},
},
{
"name": "litellm_content_retrieve",
"description": "Retrieve the full content of a file or message that was compressed to save tokens.",
"input_schema": {
"type": "object",
"properties": {
"key": {
"type": "string",
"description": "The identifier of the content to retrieve",
"enum": [
"message_0",
"HA_UPTIME_ROUTER_SPEC.md",
"message_159",
"message_160",
],
}
},
"required": ["key"],
},
"type": "custom",
},
]
response = {
"content": [
{
"type": "tool_use",
"id": "toolu_abc",
"name": "litellm_content_retrieve",
"input": {"key": "message_0"},
}
]
}
should_run, tools_dict = await logger.async_should_run_agentic_loop(
response=response,
model="claude-3-5-sonnet",
messages=[],
tools=tools,
stream=False,
custom_llm_provider="anthropic",
kwargs={},
)
assert should_run is True
assert tools_dict["tool_type"] == "compression_retrieval"
assert len(tools_dict["tool_calls"]) == 1
assert tools_dict["tool_calls"][0]["input"]["key"] == "message_0"
@pytest.mark.asyncio
async def test_build_agentic_loop_plan_missing_key_fallback():
"""Missing cache keys should produce deterministic fallback content."""
logger = CompressionInterceptionLogger()
logging_obj = MagicMock()
logging_obj.litellm_call_id = "missing_call"
logging_obj.model_call_details = {"agentic_loop_params": {}}
plan = await logger.async_build_agentic_loop_plan(
tools={
"tool_calls": [
{
"id": "toolu_missing",
"type": "tool_use",
"name": "litellm_content_retrieve",
"input": {"key": "not_found.py"},
}
]
},
model="claude-3-5-sonnet",
messages=[{"role": "user", "content": "read file"}],
response=None,
anthropic_messages_provider_config=None,
anthropic_messages_optional_request_params={},
logging_obj=logging_obj,
stream=False,
kwargs={},
)
assert plan.request_patch is not None
assert (
plan.request_patch.messages[-1]["content"][0]["content"]
== "[compressed content key 'not_found.py' not found]"
)

View file

@ -4,7 +4,7 @@ Unit tests for WebSearch Interception Handler
Tests the WebSearchInterceptionLogger class and helper functions.
"""
from unittest.mock import MagicMock, Mock
from unittest.mock import AsyncMock, MagicMock, Mock
import pytest
@ -69,6 +69,61 @@ async def test_async_should_run_agentic_loop():
assert tools_dict == {}
@pytest.mark.asyncio
async def test_async_build_agentic_loop_plan_returns_request_patch():
"""Callback should return a typed patch for base handler reruns."""
logger = WebSearchInterceptionLogger(enabled_providers=["bedrock"])
logger._execute_search = AsyncMock( # type: ignore
return_value="Title: LiteLLM\nURL: docs\nSnippet: test"
)
tools_dict = {
"tool_calls": [
{
"id": "toolu_123",
"type": "tool_use",
"name": "litellm_web_search",
"input": {"query": "what is litellm"},
}
],
"response_format": "anthropic",
}
logging_obj = MagicMock()
logging_obj.model_call_details = {
"agentic_loop_params": {"model": "bedrock/invoke/claude-3-5-sonnet"}
}
kwargs = {
"temperature": 0.2,
"_websearch_interception_converted_stream": True,
"litellm_logging_obj": object(),
}
plan = await logger.async_build_agentic_loop_plan(
tools=tools_dict,
model="claude-3-5-sonnet",
messages=[{"role": "user", "content": "search LiteLLM"}],
response=None,
anthropic_messages_provider_config=None,
anthropic_messages_optional_request_params={
"max_tokens": 1024,
"tools": [{"name": "litellm_web_search"}],
},
logging_obj=logging_obj,
stream=False,
kwargs=kwargs,
)
assert plan.run_agentic_loop is True
assert plan.request_patch is not None
assert plan.request_patch.model == "bedrock/invoke/claude-3-5-sonnet"
assert plan.request_patch.max_tokens == 1024
assert plan.request_patch.messages is not None
assert len(plan.request_patch.messages) == 3
assert "_websearch_interception_converted_stream" not in plan.request_patch.kwargs
assert "litellm_logging_obj" not in plan.request_patch.kwargs
assert plan.request_patch.kwargs["temperature"] == 0.2
@pytest.mark.asyncio
async def test_internal_flags_filtered_from_followup_kwargs():
"""Test that internal _websearch_interception flags are filtered from follow-up request kwargs.

View file

@ -0,0 +1,792 @@
"""
Tests for AgenticAnthropicStreamingIterator and SSE rebuild helpers.
"""
import json
import os
import sys
from typing import Any, Dict, List, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
import pytest
sys.path.insert(0, os.path.abspath("../../../../.."))
from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import (
AgenticAnthropicStreamingIterator,
_handle_content_block_delta,
_handle_content_block_start,
_handle_content_block_stop,
_handle_message_delta,
_handle_message_start,
_parse_sse_events,
)
# ---------------------------------------------------------------------------
# Helpers to build SSE byte payloads
# ---------------------------------------------------------------------------
def _sse_event(event_type: str, data: dict) -> bytes:
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n".encode()
def _build_simple_text_stream() -> List[bytes]:
"""Produce SSE bytes for a simple text response (no tool calls)."""
chunks = []
chunks.append(
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_123",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 0},
},
},
)
)
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Hello, world!"},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
chunks.append(
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 5},
},
)
)
chunks.append(_sse_event("message_stop", {"type": "message_stop"}))
return chunks
def _build_tool_use_stream() -> List[bytes]:
"""Produce SSE bytes for a response with a tool_use block."""
chunks = []
chunks.append(
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_tool_456",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 50, "output_tokens": 0},
},
},
)
)
# thinking block
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {
"type": "thinking",
"thinking": "",
"signature": "",
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {
"type": "thinking_delta",
"thinking": "I need to retrieve...",
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "signature_delta", "signature": "sig_abc"},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0})
)
# tool_use block
chunks.append(
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 1,
"content_block": {
"type": "tool_use",
"id": "toolu_001",
"name": "litellm_content_retrieve",
"input": {},
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 1,
"delta": {
"type": "input_json_delta",
"partial_json": '{"key": "section_',
},
},
)
)
chunks.append(
_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": 1,
"delta": {"type": "input_json_delta", "partial_json": '1"}'},
},
)
)
chunks.append(
_sse_event("content_block_stop", {"type": "content_block_stop", "index": 1})
)
chunks.append(
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "tool_use"},
"usage": {"output_tokens": 20},
},
)
)
chunks.append(_sse_event("message_stop", {"type": "message_stop"}))
return chunks
# ---------------------------------------------------------------------------
# Mock async stream
# ---------------------------------------------------------------------------
class MockAsyncStream:
"""Async iterator that yields a list of byte chunks."""
def __init__(self, chunks: List[bytes]):
self._chunks = list(chunks)
self._idx = 0
def __aiter__(self):
return self
async def __anext__(self) -> bytes:
if self._idx >= len(self._chunks):
raise StopAsyncIteration
chunk = self._chunks[self._idx]
self._idx += 1
return chunk
# ---------------------------------------------------------------------------
# Tests for _parse_sse_events
# ---------------------------------------------------------------------------
class TestParseSSEEvents:
def test_should_parse_single_event(self):
raw = _sse_event(
"message_start", {"type": "message_start", "message": {"id": "1"}}
)
events = _parse_sse_events(raw)
assert len(events) == 1
assert events[0][0] == "message_start"
assert events[0][1]["message"]["id"] == "1"
def test_should_parse_multiple_events(self):
raw = b"".join(_build_simple_text_stream())
events = _parse_sse_events(raw)
event_types = [e[0] for e in events]
assert "message_start" in event_types
assert "content_block_start" in event_types
assert "content_block_delta" in event_types
assert "content_block_stop" in event_types
assert "message_delta" in event_types
assert "message_stop" in event_types
def test_should_skip_malformed_json(self):
raw = b"event: message_start\ndata: {invalid json}\n\n"
events = _parse_sse_events(raw)
assert len(events) == 0
def test_should_handle_empty_bytes(self):
events = _parse_sse_events(b"")
assert events == []
# ---------------------------------------------------------------------------
# Tests for _handle_* helpers
# ---------------------------------------------------------------------------
class TestHandleMessageStart:
def test_should_populate_envelope(self):
response: Dict[str, Any] = {
"id": "",
"model": "",
"role": "assistant",
"usage": {"input_tokens": 0, "output_tokens": 0},
}
data = {
"message": {
"id": "msg_abc",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {
"input_tokens": 42,
"cache_creation_input_tokens": 100,
},
}
}
_handle_message_start(data, response)
assert response["id"] == "msg_abc"
assert response["model"] == "claude-sonnet-4-20250514"
assert response["usage"]["input_tokens"] == 42
assert response["usage"]["cache_creation_input_tokens"] == 100
class TestHandleContentBlockStart:
def test_should_create_text_block(self):
blocks: Dict[int, Dict] = {}
data = {"index": 0, "content_block": {"type": "text", "text": ""}}
_handle_content_block_start(data, blocks)
assert blocks[0] == {"type": "text", "text": ""}
def test_should_create_tool_use_block(self):
blocks: Dict[int, Dict] = {}
data = {
"index": 1,
"content_block": {
"type": "tool_use",
"id": "toolu_x",
"name": "my_tool",
"input": {},
},
}
_handle_content_block_start(data, blocks)
assert blocks[1]["type"] == "tool_use"
assert blocks[1]["name"] == "my_tool"
assert blocks[1]["_partial_json"] == ""
def test_should_create_thinking_block(self):
blocks: Dict[int, Dict] = {}
data = {
"index": 0,
"content_block": {"type": "thinking", "thinking": "", "signature": ""},
}
_handle_content_block_start(data, blocks)
assert blocks[0]["type"] == "thinking"
class TestHandleContentBlockDelta:
def test_should_accumulate_text(self):
blocks = {0: {"type": "text", "text": "Hello"}}
_handle_content_block_delta(
{"index": 0, "delta": {"type": "text_delta", "text": " World"}},
blocks,
)
assert blocks[0]["text"] == "Hello World"
def test_should_accumulate_json(self):
blocks = {0: {"type": "tool_use", "_partial_json": '{"key":'}}
_handle_content_block_delta(
{
"index": 0,
"delta": {"type": "input_json_delta", "partial_json": '"val"}'},
},
blocks,
)
assert blocks[0]["_partial_json"] == '{"key":"val"}'
def test_should_ignore_missing_block(self):
blocks: Dict[int, Dict] = {}
_handle_content_block_delta(
{"index": 99, "delta": {"type": "text_delta", "text": "x"}},
blocks,
)
assert 99 not in blocks
class TestHandleContentBlockStop:
def test_should_parse_tool_input_json(self):
blocks = {
0: {
"type": "tool_use",
"input": {},
"_partial_json": '{"key": "section_1"}',
}
}
_handle_content_block_stop({"index": 0}, blocks)
assert blocks[0]["input"] == {"key": "section_1"}
assert "_partial_json" not in blocks[0]
def test_should_handle_invalid_json_gracefully(self):
blocks = {
0: {
"type": "tool_use",
"input": {},
"_partial_json": "not valid json",
}
}
_handle_content_block_stop({"index": 0}, blocks)
assert blocks[0]["input"] == {"_raw": "not valid json"}
class TestHandleMessageDelta:
def test_should_set_stop_reason_and_usage(self):
response: Dict[str, Any] = {
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
}
_handle_message_delta(
{
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 15},
},
response,
)
assert response["stop_reason"] == "end_turn"
assert response["usage"]["output_tokens"] == 15
# ---------------------------------------------------------------------------
# Tests for _rebuild_anthropic_response_from_sse
# ---------------------------------------------------------------------------
class TestRebuildAnthropicResponse:
def test_should_rebuild_simple_text_response(self):
raw_bytes = _build_simple_text_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["id"] == "msg_123"
assert result["model"] == "claude-sonnet-4-20250514"
assert result["stop_reason"] == "end_turn"
assert len(result["content"]) == 1
assert result["content"][0]["type"] == "text"
assert result["content"][0]["text"] == "Hello, world!"
assert result["usage"]["input_tokens"] == 10
assert result["usage"]["output_tokens"] == 5
def test_should_rebuild_tool_use_response(self):
raw_bytes = _build_tool_use_stream()
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["id"] == "msg_tool_456"
assert result["stop_reason"] == "tool_use"
assert len(result["content"]) == 2
thinking = result["content"][0]
assert thinking["type"] == "thinking"
assert thinking["thinking"] == "I need to retrieve..."
assert thinking["signature"] == "sig_abc"
tool = result["content"][1]
assert tool["type"] == "tool_use"
assert tool["id"] == "toolu_001"
assert tool["name"] == "litellm_content_retrieve"
assert tool["input"] == {"key": "section_1"}
def test_should_return_none_without_message_start(self):
raw_bytes = [
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text"},
},
)
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is None
def test_should_handle_empty_bytes(self):
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[]
)
assert result is None
def test_should_handle_multi_event_chunks(self):
"""When multiple SSE events arrive in a single bytes chunk."""
combined = b"".join(_build_simple_text_stream())
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
[combined]
)
assert result is not None
assert result["content"][0]["text"] == "Hello, world!"
def test_should_preserve_cache_usage_fields(self):
raw_bytes = [
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_cache",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {
"input_tokens": 100,
"cache_creation_input_tokens": 50,
"cache_read_input_tokens": 30,
},
},
},
),
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 10},
},
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["usage"]["cache_creation_input_tokens"] == 50
assert result["usage"]["cache_read_input_tokens"] == 30
def test_should_handle_redacted_thinking_block(self):
raw_bytes = [
_sse_event(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_redact",
"model": "claude-sonnet-4-20250514",
"role": "assistant",
"usage": {"input_tokens": 5},
},
},
),
_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "redacted_thinking", "data": "abc123"},
},
),
_sse_event(
"content_block_stop",
{"type": "content_block_stop", "index": 0},
),
_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn"},
"usage": {"output_tokens": 1},
},
),
_sse_event("message_stop", {"type": "message_stop"}),
]
result = AgenticAnthropicStreamingIterator._rebuild_anthropic_response_from_sse(
raw_bytes
)
assert result is not None
assert result["content"][0]["type"] == "redacted_thinking"
# ---------------------------------------------------------------------------
# Tests for AgenticAnthropicStreamingIterator (Phase 1 / Phase 2)
# ---------------------------------------------------------------------------
class TestAgenticStreamingIteratorPhase1:
@pytest.mark.asyncio
async def test_should_yield_all_chunks_when_no_hook_fires(self):
"""When hooks return None, the wrapper should yield all original chunks."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[{"role": "user", "content": "hi"}],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert len(collected) == len(chunks)
for orig, got in zip(chunks, collected):
assert orig == got
mock_handler._call_agentic_completion_hooks.assert_awaited_once()
@pytest.mark.asyncio
async def test_should_pass_rebuilt_response_to_hooks(self):
"""The rebuilt dict passed to hooks should match the original stream content."""
chunks = _build_tool_use_stream()
mock_stream = MockAsyncStream(chunks)
captured_response = {}
async def mock_hooks(**kwargs):
captured_response.update(kwargs["response"])
return None
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = mock_hooks
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
async for _ in iterator:
pass
assert captured_response["id"] == "msg_tool_456"
assert captured_response["stop_reason"] == "tool_use"
assert captured_response["content"][1]["name"] == "litellm_content_retrieve"
class TestAgenticStreamingIteratorPhase2:
@pytest.mark.asyncio
async def test_should_chain_follow_up_async_iterator(self):
"""When hooks return an async iterator, Phase 2 should yield from it."""
phase1_chunks = _build_simple_text_stream()
phase2_chunks = [b"follow-up-chunk-1", b"follow-up-chunk-2"]
mock_stream = MockAsyncStream(phase1_chunks)
follow_up = MockAsyncStream(phase2_chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=follow_up)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert len(collected) == len(phase1_chunks) + len(phase2_chunks)
assert collected[-2:] == phase2_chunks
@pytest.mark.asyncio
async def test_should_convert_dict_response_to_fake_stream(self):
"""When hooks return a dict, it should be wrapped in FakeAnthropicMessagesStreamIterator."""
phase1_chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(phase1_chunks)
fake_response = {
"id": "msg_followup",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-20250514",
"content": [{"type": "text", "text": "follow-up answer"}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 100, "output_tokens": 20},
}
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
return_value=fake_response
)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
# Phase 1 chunks + Phase 2 fake-stream chunks
assert len(collected) > len(phase1_chunks)
# The follow-up chunks should contain the text from the dict response
phase2_bytes = b"".join(collected[len(phase1_chunks) :])
assert b"follow-up answer" in phase2_bytes
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."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(
side_effect=RuntimeError("hook exploded")
)
mock_logging = MagicMock()
mock_logging.litellm_call_id = "test_call_123"
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=mock_logging,
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
# All Phase 1 chunks should still have been yielded
assert len(collected) == len(chunks)
@pytest.mark.asyncio
async def test_should_handle_empty_stream(self):
"""An empty upstream stream should not crash."""
mock_stream = MockAsyncStream([])
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
collected = []
async for chunk in iterator:
collected.append(chunk)
assert collected == []
# hooks should not be called since no bytes were collected
mock_handler._call_agentic_completion_hooks.assert_not_awaited()
@pytest.mark.asyncio
async def test_should_pass_stream_true_to_hooks(self):
"""The wrapper should always pass stream=True to hooks."""
chunks = _build_simple_text_stream()
mock_stream = MockAsyncStream(chunks)
mock_handler = MagicMock()
mock_handler._call_agentic_completion_hooks = AsyncMock(return_value=None)
iterator = AgenticAnthropicStreamingIterator(
completion_stream=mock_stream,
http_handler=mock_handler,
model="claude-sonnet-4-20250514",
messages=[],
anthropic_messages_provider_config=MagicMock(),
anthropic_messages_optional_request_params={},
logging_obj=MagicMock(),
custom_llm_provider="anthropic",
kwargs={},
)
async for _ in iterator:
pass
call_kwargs = mock_handler._call_agentic_completion_hooks.call_args
assert call_kwargs.kwargs["stream"] is True

View file

@ -74,6 +74,34 @@ def test_prepare_fake_stream_request():
assert result_data["messages"] == [{"role": "user", "content": "Hello"}]
def test_get_agentic_loop_settings_defaults_and_overrides():
handler = BaseLLMHTTPHandler()
depth, max_loops, fingerprints = handler._get_agentic_loop_settings(kwargs={})
assert depth == 0
assert max_loops == 3
assert fingerprints == []
depth, max_loops, fingerprints = handler._get_agentic_loop_settings(
kwargs={
"_agentic_loop_depth": 2,
"max_agentic_loops": 7,
"_agentic_loop_fingerprints": ["fp-1", "fp-2"],
}
)
assert depth == 2
assert max_loops == 7
assert fingerprints == ["fp-1", "fp-2"]
def test_fingerprint_agentic_tools_is_deterministic():
handler = BaseLLMHTTPHandler()
tools_a = {"tool_calls": [{"id": "1", "input": {"q": "abc"}, "name": "web_search"}]}
tools_b = {"tool_calls": [{"name": "web_search", "input": {"q": "abc"}, "id": "1"}]}
assert handler._fingerprint_agentic_tools(tools_a) == handler._fingerprint_agentic_tools(tools_b)
@pytest.mark.asyncio
async def test_async_anthropic_messages_handler_extra_headers():
"""

View file

@ -891,6 +891,61 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
assert result == responses_so_far
class TestGetStructuredMessages:
"""Test the get_structured_messages method."""
def test_should_return_messages_from_chat_completions_request(self):
"""Test that messages are returned from a chat completions request."""
handler = OpenAIChatCompletionsHandler()
data = {
"messages": [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
]
}
result = handler.get_structured_messages(data)
assert result is not None
assert len(result) == 2
assert result[0]["role"] == "system"
assert result[1]["role"] == "user"
def test_should_return_none_when_no_messages(self):
"""Test that None is returned when no messages key exists."""
handler = OpenAIChatCompletionsHandler()
data = {"model": "gpt-4"}
result = handler.get_structured_messages(data)
assert result is None
def test_should_return_none_for_none_messages(self):
"""Test that None is returned when messages is explicitly None."""
handler = OpenAIChatCompletionsHandler()
data = {"messages": None}
result = handler.get_structured_messages(data)
assert result is None
def test_should_handle_multimodal_content(self):
"""Test that messages with multimodal content are returned."""
handler = OpenAIChatCompletionsHandler()
data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/image.png"},
},
],
}
]
}
result = handler.get_structured_messages(data)
assert result is not None
assert len(result) == 1
assert isinstance(result[0]["content"], list)
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])

View file

@ -995,3 +995,63 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
# Should return the responses
assert result == responses_so_far
class TestGetStructuredMessages:
"""Test the get_structured_messages method for Responses API handler."""
def test_should_convert_string_input_to_messages(self):
"""Test that a simple string input is converted to OpenAI messages."""
handler = OpenAIResponsesHandler()
data = {"input": "What is the capital of France?"}
result = handler.get_structured_messages(data)
assert result is not None
assert len(result) >= 1
found_user = False
for msg in result:
if isinstance(msg, dict) and msg.get("role") == "user":
found_user = True
break
assert found_user, f"Expected a user message, got: {result}"
def test_should_convert_list_input_to_messages(self):
"""Test that list input (ResponseInputParam) is converted to OpenAI messages."""
handler = OpenAIResponsesHandler()
data = {
"input": [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
{"role": "user", "content": "How are you?"},
]
}
result = handler.get_structured_messages(data)
assert result is not None
assert len(result) >= 3
def test_should_include_instructions_as_system_message(self):
"""Test that instructions are included as a system message."""
handler = OpenAIResponsesHandler()
data = {
"input": "Roll a d20",
"instructions": "You are a helpful dungeon master.",
}
result = handler.get_structured_messages(data)
assert result is not None
has_system = any(
isinstance(msg, dict) and msg.get("role") == "system" for msg in result
)
assert has_system, f"Expected system message from instructions, got: {result}"
def test_should_return_none_when_no_input(self):
"""Test that None is returned when input key is missing."""
handler = OpenAIResponsesHandler()
data = {"model": "gpt-4o"}
result = handler.get_structured_messages(data)
assert result is None
def test_should_return_none_for_none_input(self):
"""Test that None is returned when input is explicitly None."""
handler = OpenAIResponsesHandler()
data = {"input": None}
result = handler.get_structured_messages(data)
assert result is None

View file

@ -0,0 +1,240 @@
import os
from unittest.mock import MagicMock
import httpx
import pytest
from litellm.llms.scaleway.audio_transcription.transformation import (
ScalewayAudioTranscriptionConfig,
ScalewayAudioTranscriptionException,
)
from litellm.types.utils import TranscriptionResponse
# ---------------------------------------------------------------------------
# get_complete_url
# ---------------------------------------------------------------------------
def test_scaleway_get_complete_url_default_base():
"""With no api_base supplied, Scaleway's Generative API endpoint is used."""
url = ScalewayAudioTranscriptionConfig().get_complete_url(
api_base=None,
api_key="fake",
model="whisper-large-v3",
optional_params={},
litellm_params={},
)
assert url == "https://api.scaleway.ai/v1/audio/transcriptions"
def test_scaleway_get_complete_url_custom_base_strips_trailing_slash():
"""Caller-supplied api_base is respected; trailing slash is normalized."""
url = ScalewayAudioTranscriptionConfig().get_complete_url(
api_base="https://custom.example.com/v1/",
api_key="fake",
model="whisper-large-v3",
optional_params={},
litellm_params={},
)
assert url == "https://custom.example.com/v1/audio/transcriptions"
# ---------------------------------------------------------------------------
# validate_environment
# ---------------------------------------------------------------------------
def test_scaleway_validate_environment_explicit_api_key():
headers = ScalewayAudioTranscriptionConfig().validate_environment(
headers={},
model="whisper-large-v3",
messages=[],
optional_params={},
litellm_params={},
api_key="explicit-key",
)
assert headers["Authorization"] == "Bearer explicit-key"
assert headers["accept"] == "application/json"
def test_scaleway_validate_environment_reads_scw_secret_key(monkeypatch):
monkeypatch.setenv("SCW_SECRET_KEY", "env-secret")
headers = ScalewayAudioTranscriptionConfig().validate_environment(
headers={},
model="whisper-large-v3",
messages=[],
optional_params={},
litellm_params={},
)
assert headers["Authorization"] == "Bearer env-secret"
def test_scaleway_validate_environment_explicit_api_key_wins_over_env(monkeypatch):
"""Caller-supplied api_key must win over the SCW_SECRET_KEY env var."""
monkeypatch.setenv("SCW_SECRET_KEY", "env-secret")
headers = ScalewayAudioTranscriptionConfig().validate_environment(
headers={},
model="whisper-large-v3",
messages=[],
optional_params={},
litellm_params={},
api_key="explicit-wins",
)
assert headers["Authorization"] == "Bearer explicit-wins"
# ---------------------------------------------------------------------------
# transform_audio_transcription_request
# ---------------------------------------------------------------------------
def _open_test_audio():
"""Shared helper: open the repo's canonical speech fixture."""
wav_path = os.path.join(
os.path.dirname(__file__),
"../../../..",
"tests",
"llm_translation",
"gettysburg.wav",
)
return open(wav_path, "rb")
def test_scaleway_transform_request_builds_multipart_with_supported_params():
with _open_test_audio() as audio_file:
result = (
ScalewayAudioTranscriptionConfig().transform_audio_transcription_request(
model="whisper-large-v3",
audio_file=audio_file,
optional_params={
"language": "en",
"temperature": 0.0,
"response_format": "verbose_json",
},
litellm_params={},
)
)
assert isinstance(result.data, dict)
assert result.data["model"] == "whisper-large-v3"
assert result.data["language"] == "en"
assert result.data["temperature"] == 0.0
assert result.data["response_format"] == "verbose_json"
assert result.files is not None
assert "file" in result.files
assert len(result.files["file"]) == 3 # (filename, content, content_type)
def test_scaleway_transform_request_drops_unsupported_params():
"""Only params in get_supported_openai_params() should land in the form."""
with _open_test_audio() as audio_file:
result = (
ScalewayAudioTranscriptionConfig().transform_audio_transcription_request(
model="whisper-large-v3",
audio_file=audio_file,
optional_params={
"language": "en",
"stream": True, # not supported
"diarize": True, # not supported
},
litellm_params={},
)
)
assert "stream" not in result.data
assert "diarize" not in result.data
assert result.data["language"] == "en"
# ---------------------------------------------------------------------------
# transform_audio_transcription_response
# ---------------------------------------------------------------------------
def test_scaleway_transform_response_parses_text():
mock_response = MagicMock(spec=httpx.Response)
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = {"text": "Four score and seven years ago"}
response = (
ScalewayAudioTranscriptionConfig().transform_audio_transcription_response(
mock_response
)
)
assert isinstance(response, TranscriptionResponse)
assert response.text == "Four score and seven years ago"
def test_scaleway_transform_response_preserves_segments_and_language():
mock_response = MagicMock(spec=httpx.Response)
mock_response.headers = {"content-type": "application/json"}
mock_response.json.return_value = {
"text": "hello world",
"language": "en",
"segments": [
{"text": "hello", "start": 0.0, "end": 0.5},
{"text": "world", "start": 0.6, "end": 1.1},
],
}
response = (
ScalewayAudioTranscriptionConfig().transform_audio_transcription_response(
mock_response
)
)
assert response.text == "hello world"
assert response["language"] == "en"
assert len(response["segments"]) == 2
def test_scaleway_transform_response_raises_typed_exception_on_non_json():
"""Malformed upstream body must raise the Scaleway-typed exception so
error handlers downstream can classify it as a Scaleway failure."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.json.side_effect = ValueError("not json")
mock_response.headers = {"content-type": "application/json"}
mock_response.text = "upstream 502 bad gateway"
mock_response.status_code = 502
with pytest.raises(ScalewayAudioTranscriptionException):
ScalewayAudioTranscriptionConfig().transform_audio_transcription_response(
mock_response
)
def test_scaleway_transform_response_returns_plain_text_for_non_json_content_type():
"""When Scaleway responds with text/srt/vtt (response_format="text" etc.),
the content-type is not application/json return the body as plain text
rather than exploding on .json()."""
mock_response = MagicMock(spec=httpx.Response)
mock_response.headers = {"content-type": "text/plain; charset=utf-8"}
mock_response.text = "Four score and seven years ago"
response = (
ScalewayAudioTranscriptionConfig().transform_audio_transcription_response(
mock_response
)
)
assert isinstance(response, TranscriptionResponse)
assert response.text == "Four score and seven years ago"
def test_scaleway_validate_environment_raises_when_no_key(monkeypatch):
"""Missing credential should fail fast with a typed exception rather than
silently emitting 'Bearer None'."""
monkeypatch.delenv("SCW_SECRET_KEY", raising=False)
with pytest.raises(ScalewayAudioTranscriptionException) as excinfo:
ScalewayAudioTranscriptionConfig().validate_environment(
headers={},
model="whisper-large-v3",
messages=[],
optional_params={},
litellm_params={},
)
assert "SCW_SECRET_KEY" in str(excinfo.value)

View file

@ -1,14 +1,17 @@
import sys
import os
from types import SimpleNamespace
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
from litellm.proxy.common_utils.callback_utils import (
initialize_callbacks_on_proxy,
get_remaining_tokens_and_requests_from_request_data,
normalize_callback_names,
)
import litellm
from unittest.mock import patch
from litellm.proxy.common_utils.callback_utils import process_callback
@ -84,3 +87,35 @@ def test_normalize_callback_names_lowercases_strings():
"s3",
"custom_callback",
]
def test_initialize_callbacks_on_proxy_instantiates_compression_interception(
monkeypatch,
):
dummy_callback = object()
monkeypatch.setitem(
sys.modules,
"litellm.proxy.proxy_server",
SimpleNamespace(prisma_client=None),
)
monkeypatch.setattr(
"litellm.integrations.compression_interception.handler.CompressionInterceptionLogger.initialize_from_proxy_config",
lambda litellm_settings, callback_specific_params: dummy_callback,
)
original_callbacks = (
list(litellm.callbacks) if isinstance(litellm.callbacks, list) else []
)
litellm.callbacks = []
try:
initialize_callbacks_on_proxy(
value=["compression_interception"],
premium_user=False,
config_file_path=".",
litellm_settings={"compression_interception_params": {"enabled": True}},
callback_specific_params={},
)
assert dummy_callback in litellm.callbacks
assert "compression_interception" not in litellm.callbacks
finally:
litellm.callbacks = original_callbacks

View file

@ -9,6 +9,7 @@ sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
import httpx
import pytest
from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, PrismaError
@ -110,220 +111,127 @@ async def test_db_health_prisma_client_none():
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",
"transport_error",
[
PrismaError(),
httpx.ConnectError("All connection attempts failed"),
ClientNotConnectedError(),
HTTPClientClosedError(),
PrismaError("Can't reach database server"),
],
)
async def test_db_health_error_flag_off_raises_no_reconnect(prisma_error):
async def test_db_health_transport_error_never_raises(transport_error):
"""
When health_check raises and allow_requests_on_db_unavailable is False,
handle_db_exception re-raises immediately. The reconnect path is never
reached, so disconnect/connect are never called.
Regression test for the /health/readiness 503 loop bug.
handle_db_exception() used to re-raise inside _db_health_readiness_check,
turning any DB outage into a 503 "Service Unhealthy" response that never
recovered. Transport errors (ClientNotConnectedError, httpx.ConnectError,
etc.) must return {"status": "disconnected"} never raise.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=prisma_error)
mock_prisma.disconnect = AsyncMock()
mock_prisma.health_check = AsyncMock(side_effect=transport_error)
mock_prisma.attempt_db_reconnect = AsyncMock(return_value=False)
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
),
):
with pytest.raises(Exception) as exc_info:
await _db_health_readiness_check()
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert exc_info.value is prisma_error
mock_prisma.disconnect.assert_not_called()
assert _health_endpoints_module.db_health_cache["status"] == "disconnected"
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",
"transport_error",
[
PrismaError("Can't reach database server"),
httpx.ConnectError("All connection attempts failed"),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_error_flag_on_reconnect_succeeds(prisma_error):
async def test_db_health_transport_error_reconnect_succeeds(transport_error):
"""
When health_check raises, allow_requests_on_db_unavailable is True,
and the reconnect cycle (disconnect -> connect -> health_check) succeeds,
return 'connected' and update the cache.
When health_check raises a transport error and attempt_db_reconnect
succeeds, the second health_check passes and we return 'connected'.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=[prisma_error, None])
mock_prisma.disconnect = AsyncMock()
mock_prisma.connect = AsyncMock()
mock_prisma.health_check = AsyncMock(side_effect=[transport_error, None])
mock_prisma.attempt_db_reconnect = AsyncMock(return_value=True)
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert result["status"] == "connected"
mock_prisma.disconnect.assert_called_once()
mock_prisma.connect.assert_called_once()
mock_prisma.attempt_db_reconnect.assert_called_once_with(
reason="health_readiness_check"
)
assert mock_prisma.health_check.call_count == 2
@pytest.mark.asyncio
@pytest.mark.parametrize(
"prisma_error",
"transport_error",
[
PrismaError("Can't reach database server"),
httpx.ConnectError("All connection attempts failed"),
ClientNotConnectedError(),
HTTPClientClosedError(),
],
)
async def test_db_health_error_flag_on_reconnect_fails(prisma_error):
async def test_db_health_transport_error_reconnect_fails(transport_error):
"""
When health_check raises, allow_requests_on_db_unavailable is True,
but the reconnect also fails, return 'disconnected' instead of raising.
This respects the flag's intent: keep serving even without a DB.
When health_check raises a transport error and attempt_db_reconnect also
fails, return 'disconnected' without raising.
"""
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=prisma_error)
mock_prisma.disconnect = AsyncMock()
mock_prisma.connect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.disconnect.assert_called_once()
mock_prisma.connect.assert_called_once()
@pytest.mark.asyncio
async def test_db_health_non_transport_error_flag_off_raises():
"""
When health_check raises a non-transport error and
allow_requests_on_db_unavailable is False, handle_db_exception
re-raises before reaching the is_database_transport_error guard.
Cache is still invalidated before the re-raise.
"""
non_transport_error = PrismaError("UniqueViolationError")
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=non_transport_error)
mock_prisma.disconnect = AsyncMock()
mock_prisma.connect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": False},
),
):
with pytest.raises(PrismaError):
await _db_health_readiness_check()
assert _health_endpoints_module.db_health_cache["status"] == "disconnected"
mock_prisma.disconnect.assert_not_called()
mock_prisma.connect.assert_not_called()
@pytest.mark.asyncio
async def test_db_health_non_transport_error_flag_on_skips_reconnect():
"""
When health_check raises a non-transport error (e.g. data-layer) and
allow_requests_on_db_unavailable is True, handle_db_exception swallows
the exception, then is_database_transport_error returns False so the
reconnect cycle is skipped. Returns 'disconnected' without calling
disconnect/connect.
"""
non_transport_error = PrismaError("UniqueViolationError")
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=non_transport_error)
mock_prisma.disconnect = AsyncMock()
mock_prisma.connect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.disconnect.assert_not_called()
mock_prisma.connect.assert_not_called()
@pytest.mark.asyncio
async def test_db_health_reconnect_disconnect_fails():
"""
When disconnect() itself raises during the reconnect cycle,
the inner except catches it and returns 'disconnected'.
connect() and the second health_check() are never called.
"""
transport_error = ClientNotConnectedError()
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=transport_error)
mock_prisma.disconnect = AsyncMock(side_effect=RuntimeError("already closed"))
mock_prisma.connect = AsyncMock()
mock_prisma.attempt_db_reconnect = AsyncMock(
side_effect=RuntimeError("reconnect failed")
)
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.proxy_server.general_settings",
{"allow_requests_on_db_unavailable": True},
),
):
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.disconnect.assert_called_once()
mock_prisma.connect.assert_not_called()
@pytest.mark.asyncio
async def test_db_health_non_transport_error_returns_disconnected():
"""
When health_check raises a non-transport error (e.g. data-layer error),
is_database_transport_error returns False so reconnect is skipped.
Returns 'disconnected' without raising and without calling attempt_db_reconnect.
"""
non_transport_error = PrismaError("UniqueViolationError")
mock_prisma = MagicMock()
mock_prisma.health_check = AsyncMock(side_effect=non_transport_error)
mock_prisma.attempt_db_reconnect = AsyncMock()
_health_endpoints_module.db_health_cache = {
"status": "connected",
"last_updated": datetime.now() - timedelta(seconds=20),
}
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
result = await _db_health_readiness_check()
assert result["status"] == "disconnected"
mock_prisma.attempt_db_reconnect.assert_not_called()
@pytest.mark.asyncio

View file

@ -1791,6 +1791,57 @@ def test_add_litellm_metadata_from_request_headers_both_headers_trace_id_precede
assert data["litellm_trace_id"] == "trace-value"
def test_add_litellm_metadata_from_request_headers_generic_session_id_header():
"""A generic x-<vendor>-session-id header is used when no explicit litellm header is set."""
headers = {"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"}
data = {"metadata": {}}
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=headers, data=data, _metadata_variable_name="metadata"
)
assert data["metadata"]["session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01"
assert data["litellm_session_id"] == "e96634a3-fa28-4083-b354-55542e2dca01"
assert data["litellm_trace_id"] == "e96634a3-fa28-4083-b354-55542e2dca01"
def test_add_litellm_metadata_from_request_headers_explicit_header_beats_generic():
"""Explicit x-litellm-trace-id wins over a generic x-*-session-id header."""
headers = {
"x-litellm-trace-id": "explicit-trace-id-value",
"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01",
}
data = {"metadata": {}}
LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers(
headers=headers, data=data, _metadata_variable_name="metadata"
)
assert data["litellm_session_id"] == "explicit-trace-id-value"
assert data["litellm_trace_id"] == "explicit-trace-id-value"
def test_get_chain_id_from_headers_generic_vendor_session_id():
"""get_chain_id_from_headers picks up any x-<vendor>-session-id with a valid value."""
from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
assert (
get_chain_id_from_headers(
{"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01"}
)
== "e96634a3-fa28-4083-b354-55542e2dca01"
)
# Short / non-alphanumeric values should be ignored
assert get_chain_id_from_headers({"x-foo-session-id": "short"}) is None
assert get_chain_id_from_headers({"x-foo-session-id": "has spaces!!"}) is None
# Explicit headers still take precedence
assert (
get_chain_id_from_headers(
{
"x-litellm-trace-id": "explicit-id-value",
"x-claude-code-session-id": "e96634a3-fa28-4083-b354-55542e2dca01",
}
)
== "explicit-id-value"
)
def test_get_internal_user_header_from_mapping_returns_expected_header():
mappings = [
{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"},

View file

@ -12,7 +12,148 @@ sys.path.insert(
from litellm.router_strategy.auto_router.auto_router import AutoRouter
pytestmark = pytest.mark.skip(reason="Skipping auto router tests - beta feature")
pytestmark_skip_beta = pytest.mark.skip(
reason="Skipping auto router tests - beta feature"
)
class TestExtractTextFromMessages:
"""Tests for AutoRouter._extract_text_from_messages (no semantic_router dependency)."""
def test_should_extract_content_from_simple_user_message(self):
messages = [{"role": "user", "content": "Hello world"}]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "Hello world"
def test_should_extract_last_user_message_from_tool_call_conversation(self):
messages = [
{"role": "user", "content": "What's the weather in NYC?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_abc123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"location": "NYC"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_abc123",
"content": "72°F and sunny",
},
{"role": "user", "content": "Now tell me about London"},
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "Now tell me about London"
def test_should_find_user_message_when_last_message_is_assistant_with_tool_calls(
self,
):
messages = [
{"role": "user", "content": "What's the weather?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_abc",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "What's the weather?"
def test_should_find_user_message_when_last_message_is_tool_response(self):
messages = [
{"role": "user", "content": "What's the weather?"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_abc",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
{
"role": "tool",
"tool_call_id": "call_abc",
"content": "72°F and sunny",
},
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "What's the weather?"
def test_should_handle_multimodal_content_list(self):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/img.png"},
},
],
}
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "What's in this image?"
def test_should_handle_multimodal_content_with_multiple_text_blocks(self):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "First part"},
{"type": "text", "text": "Second part"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/img.png"},
},
],
}
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == "First part Second part"
def test_should_return_empty_string_when_user_content_is_none(self):
messages = [{"role": "user", "content": None}]
result = AutoRouter._extract_text_from_messages(messages)
assert result == ""
def test_should_return_empty_string_when_no_user_messages(self):
messages = [
{"role": "system", "content": "You are a helpful assistant"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_abc",
"type": "function",
"function": {"name": "get_weather", "arguments": "{}"},
}
],
},
]
result = AutoRouter._extract_text_from_messages(messages)
assert result == ""
def test_should_return_empty_string_for_empty_messages_list(self):
result = AutoRouter._extract_text_from_messages([])
assert result == ""
@pytest.fixture
@ -41,6 +182,7 @@ def mock_route_choice():
return mock_choice
@pytestmark_skip_beta
class TestAutoRouter:
"""Test class for AutoRouter methods."""

View file

@ -7,7 +7,7 @@ Tests the rule-based complexity scoring and tier assignment logic.
import os
import sys
from typing import Dict, List
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import pytest
@ -828,3 +828,222 @@ class TestRouterComplexityDeploymentMethods:
)
router.init_complexity_router_deployment(deployment)
assert "auto_router/complexity_router/test-router" in router.complexity_routers
class TestAsyncPreRoutingHookMultiFormat:
"""Test async_pre_routing_hook with multiple input formats."""
@pytest.mark.asyncio
async def test_should_route_with_chat_completions_messages(self, complexity_router):
"""Test routing with standard chat completions messages."""
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={},
messages=[{"role": "user", "content": "What is 2+2?"}],
)
assert result is not None
assert result.model is not None
assert result.messages is not None
@pytest.mark.asyncio
async def test_should_route_with_responses_api_string_input(
self, complexity_router
):
"""Test routing with Responses API string input via handler dispatch."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.types.utils import CallTypes
mock_mappings = {CallTypes.responses: OpenAIResponsesHandler}
with patch(
"litellm.llms.load_guardrail_translation_mappings",
return_value=mock_mappings,
):
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={"input": "What is the capital of France?"},
messages=None,
input="What is the capital of France?",
)
assert result is not None
assert result.model is not None
# messages should be None since the original request didn't have messages
assert result.messages is None
@pytest.mark.asyncio
async def test_should_route_with_responses_api_list_input(self, complexity_router):
"""Test routing with Responses API list input via handler dispatch."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.types.utils import CallTypes
mock_mappings = {CallTypes.responses: OpenAIResponsesHandler}
list_input = [
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi there!"},
{
"role": "user",
"content": "Write a Python function to sort a list using merge sort",
},
]
with patch(
"litellm.llms.load_guardrail_translation_mappings",
return_value=mock_mappings,
):
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={"input": list_input},
messages=None,
input=list_input,
)
assert result is not None
assert result.model is not None
assert result.messages is None
@pytest.mark.asyncio
async def test_should_use_route_based_inference(self, complexity_router):
"""Test that route-based call type inference is used when available."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.types.utils import CallTypes
mock_mappings = {CallTypes.responses: OpenAIResponsesHandler}
with patch(
"litellm.llms.load_guardrail_translation_mappings",
return_value=mock_mappings,
):
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={
"input": "Roll 2d4+1",
"litellm_metadata": {
"user_api_key_request_route": "/v1/responses",
},
},
messages=None,
)
assert result is not None
assert result.model is not None
@pytest.mark.asyncio
async def test_should_return_none_when_no_messages_or_input(
self, complexity_router
):
"""Test that None is returned when neither messages nor input is available."""
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={},
messages=None,
input=None,
)
assert result is None
@pytest.mark.asyncio
async def test_should_prefer_original_messages_over_conversion(
self, complexity_router
):
"""Test that original messages are used when both messages and input are available."""
messages = [{"role": "user", "content": "What is 2+2?"}]
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={"input": "This should be ignored"},
messages=messages,
)
assert result is not None
assert result.messages == messages
@pytest.mark.asyncio
async def test_should_include_instructions_in_classification(
self, complexity_router
):
"""Test that Responses API instructions influence classification via system message."""
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
)
from litellm.types.utils import CallTypes
mock_mappings = {CallTypes.responses: OpenAIResponsesHandler}
with patch(
"litellm.llms.load_guardrail_translation_mappings",
return_value=mock_mappings,
):
result = await complexity_router.async_pre_routing_hook(
model="test-model",
request_kwargs={
"input": "Write merge sort",
"instructions": "You are an expert Python developer. Use advanced algorithms and optimize for performance.",
},
messages=None,
)
assert result is not None
assert result.model is not None
class TestExtractUserMessageAndSystemPrompt:
"""Test the _extract_user_message_and_system_prompt static method."""
def test_should_extract_user_message(self):
"""Test extraction of the last user message."""
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
{"role": "assistant", "content": "Hi!"},
{"role": "user", "content": "How are you?"},
]
user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(
messages
)
assert user_msg == "How are you?"
assert sys_prompt == "You are helpful."
def test_should_handle_no_user_message(self):
"""Test when there is no user message."""
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "assistant", "content": "Hi!"},
]
user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(
messages
)
assert user_msg is None
assert sys_prompt == "You are helpful."
def test_should_handle_multipart_content(self):
"""Test extraction from multipart content messages."""
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Describe this image"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/img.png"},
},
],
}
]
user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(
messages
)
assert user_msg == "Describe this image"
assert sys_prompt is None
def test_should_handle_empty_messages(self):
"""Test with empty messages list."""
user_msg, sys_prompt = ComplexityRouter._extract_user_message_and_system_prompt(
[]
)
assert user_msg is None
assert sys_prompt is None

File diff suppressed because it is too large Load diff

View file

@ -3,6 +3,7 @@ Unit tests for litellm.compress().
"""
import os
import importlib
import pytest
@ -12,6 +13,10 @@ from litellm.compression.scoring.embedding_scorer import embedding_score_message
from litellm.compression.content_detection import detect_content_type
from litellm.compression.message_stubbing import extract_key, stub_message
from litellm.compression.retrieval_tool import build_retrieval_tool
from litellm.types.utils import CallTypes
CALL_TYPE = CallTypes.completion
ANTHROPIC_CALL_TYPE = CallTypes.anthropic_messages
# ---------------------------------------------------------------------------
@ -149,7 +154,7 @@ def test_retrieval_tool_description_lists_keys():
def test_compress_below_trigger_passthrough():
messages = [{"role": "user", "content": "hello"}]
result = litellm.compress(messages, model="gpt-4o")
result = litellm.compress(messages, model="gpt-4o", call_type=CALL_TYPE)
assert result["messages"] == messages
assert result["cache"] == {}
assert result["tools"] == []
@ -178,6 +183,7 @@ def test_compress_above_trigger():
result = litellm.compress(
big_messages,
model="gpt-4o",
call_type=CALL_TYPE,
compression_trigger=1000,
compression_target=500,
)
@ -189,13 +195,62 @@ def test_compress_above_trigger():
assert result["tools"][0]["function"]["name"] == "litellm_content_retrieve"
def test_compress_anthropic_list_content_is_boundary_stable():
messages = [
{"role": "system", "content": [{"type": "text", "text": "System prompt"}]},
{
"role": "user",
"content": [
{"type": "text", "text": "# a.py\n" + "alpha " * 2000},
{
"type": "image_url",
"image_url": {"url": "https://example.com/a.png"},
},
],
},
{
"role": "user",
"content": [
{"type": "text", "text": "# b.py\n" + "beta " * 2000},
{
"type": "image_url",
"image_url": {"url": "https://example.com/b.png"},
},
],
},
{
"role": "user",
"content": [{"type": "text", "text": "Fix alpha bug in a.py"}],
},
]
result = litellm.compress(
messages=messages,
model="claude-sonnet-4-20250514",
call_type=ANTHROPIC_CALL_TYPE,
compression_trigger=1000,
compression_target=500,
)
assert result["compressed_tokens"] < result["original_tokens"]
assert len(result["messages"]) == len(messages)
assert [m["role"] for m in result["messages"]] == [m["role"] for m in messages]
assert len(result["cache"]) > 0
assert len(result["tools"]) == 1
assert result["tools"][0]["type"] == "custom"
assert result["tools"][0]["name"] == "litellm_content_retrieve"
assert "input_schema" in result["tools"][0]
def test_compress_preserves_system_message():
messages = [
{"role": "system", "content": "System prompt. " * 500},
{"role": "user", "content": "Large file content. " * 5000},
{"role": "user", "content": "Fix the bug"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
)
assert result["messages"][0]["role"] == "system"
assert "System prompt" in result["messages"][0]["content"]
@ -205,7 +260,9 @@ def test_compress_preserves_last_user_message():
{"role": "user", "content": "Big context " * 5000},
{"role": "user", "content": "Fix the bug in auth.py"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
)
last_user = [m for m in result["messages"] if m["role"] == "user"][-1]
assert "Fix the bug in auth.py" in last_user["content"]
@ -216,7 +273,9 @@ def test_compress_preserves_last_assistant_message():
{"role": "assistant", "content": "I'll help with that. " * 2000},
{"role": "user", "content": "Now fix the bug"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
)
assistant_msgs = [m for m in result["messages"] if m["role"] == "assistant"]
assert len(assistant_msgs) >= 1
# The last assistant message should be preserved (not stubbed)
@ -229,7 +288,9 @@ def test_cache_keys_match_stubs():
{"role": "user", "content": "# auth.py\n" + "code " * 5000},
{"role": "user", "content": "Fix it"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
)
if result["tools"]:
tool_desc = result["tools"][0]["function"]["description"]
for key in result["cache"]:
@ -242,11 +303,75 @@ def test_compress_default_target():
{"role": "user", "content": "content " * 5000},
{"role": "user", "content": "query"},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=2000)
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=2000
)
# Should have compressed — target = 1000
assert result["compressed_tokens"] <= result["original_tokens"]
def test_compress_nested_tool_result_extracts_text_only():
messages = [
{"role": "system", "content": [{"type": "text", "text": "System rules"}]},
{
"role": "user",
"content": [
{"type": "text", "text": "prefix"},
{
"type": "tool_result",
"tool_use_id": "toolu_1",
"content": [
{"type": "text", "text": "nested text fragment"},
{
"type": "image_url",
"image_url": {
"url": "https://example.com/secret-tool.png",
},
},
],
},
{
"type": "image_url",
"image_url": {"url": "https://example.com/top.png"},
},
{"type": "text", "text": " " + ("irrelevant " * 3000)},
],
},
{
"role": "user",
"content": [{"type": "text", "text": "final query that must remain"}],
},
]
result = litellm.compress(
messages=messages,
model="claude-sonnet-4-20250514",
call_type=ANTHROPIC_CALL_TYPE,
compression_trigger=500,
compression_target=100,
)
cached_text = " ".join(result["cache"].values())
assert "nested text fragment" in cached_text
assert "https://example.com/secret-tool.png" not in cached_text
assert "https://example.com/top.png" not in cached_text
def test_compress_default_call_type_is_completion():
result = litellm.compress(
messages=[
{"role": "user", "content": "Large context " * 4000},
{"role": "user", "content": "query"},
],
model="gpt-4o",
compression_trigger=1000,
compression_target=500,
)
assert result["compressed_tokens"] <= result["original_tokens"]
assert isinstance(result["tools"], list)
def test_compress_forwards_embedding_model_params(monkeypatch):
captured = {}
@ -269,6 +394,7 @@ def test_compress_forwards_embedding_model_params(monkeypatch):
{"role": "user", "content": "Fix auth"},
],
model="gpt-4o",
call_type=CALL_TYPE,
compression_trigger=1000,
embedding_model="text-embedding-3-small",
embedding_model_params={"api_base": "https://example-embeddings.test"},
@ -326,6 +452,7 @@ def test_embedding_scorer():
{"role": "user", "content": "Fix auth"},
],
model="gpt-4o",
call_type=CALL_TYPE,
compression_trigger=1000,
embedding_model="text-embedding-3-small",
)
@ -346,8 +473,9 @@ def test_simple_compression(final_user_message, expected_content):
{"role": "user", "content": "Unrelated cooking recipes " * 2000},
{"role": "user", "content": final_user_message},
]
result = litellm.compress(messages, model="gpt-4o", compression_trigger=1000)
print(result["messages"])
result = litellm.compress(
messages, model="gpt-4o", call_type=CALL_TYPE, compression_trigger=1000
)
if expected_content == "Unrelated cooking recipes ":
assert "Unrelated cooking recipes " in result["messages"][1]["content"]
assert "Authentication code " not in result["messages"][0]["content"]
@ -356,3 +484,184 @@ def test_simple_compression(final_user_message, expected_content):
assert "Unrelated cooking recipes " not in result["messages"][1]["content"]
else:
raise ValueError(f"Unexpected expected_content: {expected_content}")
def test_compress_anthropic_drops_irrelevant_tool_exchange_span(monkeypatch):
compress_module = importlib.import_module("litellm.compression.compress")
def fake_bm25_score_messages(query, messages):
assert "final query" in query
assert len(messages) == 5
# Prefer idx=0 and de-prioritize the tool exchange span (idx=1,2)
return [0.95, 0.01, 0.02, 0.8, 1.0]
def fake_token_counter(model, messages=None, text=None):
if messages is not None:
return 1000
if text is None:
return 0
if "final query" in text:
return 50
if "assistant_tail" in text:
return 20
if "other_blob" in text:
return 220
if "tool_payload_relevant" in text:
return 200
if text == "":
return 1
return 10
monkeypatch.setattr(
compress_module, "bm25_score_messages", fake_bm25_score_messages
)
monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)
messages = [
{"role": "user", "content": "other_blob " * 300},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_drop",
"name": "litellm_content_retrieve",
"input": {"key": "message_1"},
}
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_drop",
"content": [{"type": "text", "text": "tool_payload_relevant"}],
}
],
},
{"role": "assistant", "content": "assistant_tail"},
{"role": "user", "content": "final query"},
]
result = litellm.compress(
messages=messages,
model="claude-sonnet-4-20250514",
call_type=ANTHROPIC_CALL_TYPE,
compression_trigger=100,
compression_target=280,
)
# idx=1,2 should be dropped atomically (no orphan tool blocks left behind)
assert len(result["messages"]) == 3
assert result["messages"][0]["role"] == "user"
assert "other_blob" in result["messages"][0]["content"]
assert result["messages"][1]["content"] == "assistant_tail"
assert result["messages"][2]["content"] == "final query"
assert result["cache"] == {}
def test_compress_anthropic_keeps_relevant_tool_exchange_span(monkeypatch):
compress_module = importlib.import_module("litellm.compression.compress")
def fake_bm25_score_messages(query, messages):
assert "final query" in query
assert len(messages) == 5
# Prefer the tool exchange span over idx=0
return [0.05, 0.01, 0.92, 0.8, 1.0]
def fake_token_counter(model, messages=None, text=None):
if messages is not None:
return 1000
if text is None:
return 0
if "final query" in text:
return 50
if "assistant_tail" in text:
return 20
if "other_blob" in text:
return 220
if "tool_payload_relevant" in text:
return 200
if text == "":
return 1
return 10
monkeypatch.setattr(
compress_module, "bm25_score_messages", fake_bm25_score_messages
)
monkeypatch.setattr(compress_module, "token_counter", fake_token_counter)
messages = [
{"role": "user", "content": "other_blob " * 300},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_keep",
"name": "litellm_content_retrieve",
"input": {"key": "message_1"},
}
],
},
{
"role": "user",
"content": [
{
"type": "tool_result",
"tool_use_id": "toolu_keep",
"content": [{"type": "text", "text": "tool_payload_relevant"}],
}
],
},
{"role": "assistant", "content": "assistant_tail"},
{"role": "user", "content": "final query"},
]
result = litellm.compress(
messages=messages,
model="claude-sonnet-4-20250514",
call_type=ANTHROPIC_CALL_TYPE,
compression_trigger=100,
compression_target=280,
)
assert len(result["messages"]) == 5
assert result["messages"][1]["role"] == "assistant"
assert result["messages"][2]["role"] == "user"
# idx=0 should be compressed instead
assert "litellm_content_retrieve" in result["messages"][0]["content"]
assert len(result["cache"]) == 1
def test_compress_anthropic_malformed_tool_sequence_passes_through():
messages = [
{"role": "user", "content": "other_blob " * 300},
{
"role": "assistant",
"content": [
{
"type": "tool_use",
"id": "toolu_broken",
"name": "litellm_content_retrieve",
"input": {"key": "message_1"},
}
],
},
{"role": "user", "content": [{"type": "text", "text": "missing tool_result"}]},
{"role": "user", "content": "final query"},
]
result = litellm.compress(
messages=messages,
model="claude-sonnet-4-20250514",
call_type=ANTHROPIC_CALL_TYPE,
compression_trigger=100,
compression_target=280,
)
assert result["messages"] == messages
assert result["cache"] == {}
assert result["tools"] == []
assert result["compression_skipped_reason"] == "invalid_anthropic_tool_sequence"