mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge remote-tracking branch 'upstream/litellm_internal_staging' into litellm_post_call_non_streaming
This commit is contained in:
commit
62c2c553d7
76 changed files with 6658 additions and 688 deletions
|
|
@ -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
|
||||
|
|
@ -3042,10 +3051,19 @@ jobs:
|
|||
- ui/litellm-dashboard/node_modules
|
||||
- run:
|
||||
name: Build UI from source
|
||||
# Prior version used `cp -r out/ ../../litellm/proxy/_experimental/out/`.
|
||||
# GNU cp (used on CircleCI's Ubuntu image) interprets that as "copy the
|
||||
# source directory as a child of the destination" when the destination
|
||||
# already exists — silently creating `_experimental/out/out/` instead of
|
||||
# replacing the served bundle. The proxy continued serving whatever was
|
||||
# checked into `_experimental/out/*`, so this job was effectively testing
|
||||
# the pre-build bundle on every run. Replace-and-move guarantees the
|
||||
# freshly built bundle is what the proxy actually serves.
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm run build
|
||||
cp -r out/ ../../litellm/proxy/_experimental/out/
|
||||
rm -rf ../../litellm/proxy/_experimental/out
|
||||
mv out ../../litellm/proxy/_experimental/out
|
||||
# Restructure HTML so extensionless routes work (login.html -> login/index.html)
|
||||
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name 'index.html' | while read -r f; do
|
||||
d="${f%.html}"; mkdir -p "$d"; mv "$f" "$d/index.html"
|
||||
|
|
|
|||
|
|
@ -15,29 +15,21 @@ COPY --from=uvbin /uv /usr/local/bin/uv
|
|||
COPY --from=uvbin /uvx /usr/local/bin/uvx
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache \
|
||||
python3 \
|
||||
python3-dev \
|
||||
clang \
|
||||
llvm \
|
||||
lld \
|
||||
gcc \
|
||||
linux-headers \
|
||||
build-base \
|
||||
bash \
|
||||
coreutils \
|
||||
curl \
|
||||
openssl \
|
||||
openssl-dev \
|
||||
nodejs \
|
||||
npm \
|
||||
libsndfile && break || sleep 5; \
|
||||
apk add --no-cache \
|
||||
python3 \
|
||||
python3-dev \
|
||||
gcc \
|
||||
bash \
|
||||
coreutils \
|
||||
curl \
|
||||
openssl \
|
||||
libsndfile \
|
||||
nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
NVM_DIR=/root/.nvm \
|
||||
PATH="/root/.nvm/versions/node/v20.20.2/bin:/app/.venv/bin:${PATH}" \
|
||||
PATH="/app/.venv/bin:${PATH}" \
|
||||
LITELLM_NON_ROOT=true \
|
||||
PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
PRISMA_CLI_BINARY_TARGETS="debian-openssl-3.0.x" \
|
||||
|
|
@ -49,7 +41,8 @@ COPY enterprise/pyproject.toml enterprise/
|
|||
COPY litellm-proxy-extras/pyproject.toml litellm-proxy-extras/
|
||||
|
||||
# Install third-party dependencies (cached unless pyproject.toml/uv.lock change)
|
||||
RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
||||
uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
--extra extra_proxy \
|
||||
|
|
@ -62,38 +55,12 @@ COPY . .
|
|||
# Set non-root flag for build time consistency
|
||||
ENV LITELLM_NON_ROOT=true
|
||||
|
||||
# Build Admin UI once and stage the static output for the runtime image.
|
||||
# NOTE: .npmrc files (which may set ignore-scripts=true and min-release-age=3d)
|
||||
# are temporarily renamed during npm install/ci so they don't block lifecycle
|
||||
# scripts needed by the build. This is safe because npm ci installs from
|
||||
# package-lock.json with pinned versions + integrity hashes.
|
||||
# Stage the pre-built Admin UI from the checked-in Next.js static export.
|
||||
# _experimental/out/ is regenerated as part of the release runbook.
|
||||
# Restructure extensionless routes (foo.html -> foo/index.html) to match the layout
|
||||
# proxy_server.py expects, and drop a readiness marker.
|
||||
RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
|
||||
([ -f /app/.npmrc ] && mv /app/.npmrc /app/.npmrc.bak || true) && \
|
||||
NVM_VERSION="v0.40.4" && \
|
||||
NVM_CHECKSUM="4b7412c49960c7d31e8df72da90c1fb5b8cccb419ac99537b737028d497aba4f" && \
|
||||
NODE_VERSION="v20.20.2" && \
|
||||
NVM_SCRIPT="/tmp/install-nvm.sh" && \
|
||||
curl -fsSL "https://raw.githubusercontent.com/nvm-sh/nvm/${NVM_VERSION}/install.sh" -o "$NVM_SCRIPT" && \
|
||||
echo "${NVM_CHECKSUM} ${NVM_SCRIPT}" | sha256sum -c - && \
|
||||
bash "$NVM_SCRIPT" && \
|
||||
export NVM_DIR="$HOME/.nvm" && \
|
||||
. "$NVM_DIR/nvm.sh" && \
|
||||
nvm install "${NODE_VERSION}" && \
|
||||
nvm use "${NODE_VERSION}" && \
|
||||
npm install -g npm@11.12.1 && \
|
||||
npm install -g node-gyp@12.2.0 && \
|
||||
ln -sf "$(npm root -g)/node-gyp" "$(npm root -g)/npm/node_modules/node-gyp" && \
|
||||
npm cache clean --force && \
|
||||
cd /app/ui/litellm-dashboard && \
|
||||
if [ -f "/app/enterprise/enterprise_ui/enterprise_colors.json" ]; then \
|
||||
cp /app/enterprise/enterprise_ui/enterprise_colors.json ./ui_colors.json; \
|
||||
fi && \
|
||||
([ -f .npmrc ] && mv .npmrc .npmrc.bak || true) && \
|
||||
npm ci --no-audit --no-fund && \
|
||||
([ -f .npmrc.bak ] && mv .npmrc.bak .npmrc || true) && \
|
||||
([ -f /app/.npmrc.bak ] && mv /app/.npmrc.bak /app/.npmrc || true) && \
|
||||
npm run build && \
|
||||
cp -r /app/ui/litellm-dashboard/out/* /var/lib/litellm/ui/ && \
|
||||
cp -r /app/litellm/proxy/_experimental/out/. /var/lib/litellm/ui/ && \
|
||||
cp /app/litellm/proxy/logo.jpg /var/lib/litellm/assets/logo.jpg && \
|
||||
( cd /var/lib/litellm/ui && \
|
||||
for html_file in *.html; do \
|
||||
|
|
@ -103,10 +70,10 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
|
|||
mv "$html_file" "$folder_name/index.html"; \
|
||||
fi; \
|
||||
done && \
|
||||
touch .litellm_ui_ready ) && \
|
||||
cd /app/ui/litellm-dashboard && rm -rf ./out
|
||||
touch .litellm_ui_ready )
|
||||
|
||||
RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
||||
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
|
||||
if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
||||
uv sync --frozen --no-default-groups --no-editable \
|
||||
--extra proxy \
|
||||
--extra proxy-runtime \
|
||||
|
|
@ -123,10 +90,7 @@ RUN if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
|
|||
--python python3; \
|
||||
fi
|
||||
|
||||
RUN mkdir -p /app/.cache/npm && \
|
||||
prisma generate --schema=./schema.prisma && \
|
||||
prisma --version && \
|
||||
prisma migrate diff --from-empty --to-schema-datamodel ./schema.prisma --script > /dev/null 2>&1 || true
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
|
@ -137,33 +101,11 @@ WORKDIR /app
|
|||
USER root
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk upgrade --no-cache && break || sleep 5; \
|
||||
apk upgrade --no-cache && break || sleep 5; \
|
||||
done && \
|
||||
for i in 1 2 3; do \
|
||||
apk add --no-cache python3 bash openssl tzdata nodejs npm supervisor libsndfile && break || sleep 5; \
|
||||
done && \
|
||||
apk upgrade --no-cache nodejs && \
|
||||
npm install -g npm@11.12.1 tar@7.5.11 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \
|
||||
GLOBAL="$(npm root -g)" && \
|
||||
find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "glob" -path "*/node_modules/glob" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/glob" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "brace-expansion" -path "*/node_modules/@isaacs/brace-expansion" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/@isaacs/brace-expansion" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "minimatch" -path "*/node_modules/minimatch" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/minimatch" "$d"; \
|
||||
done && \
|
||||
find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \
|
||||
rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \
|
||||
done && \
|
||||
find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \
|
||||
sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \
|
||||
npm cache clean --force && \
|
||||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
apk add --no-cache python3 bash openssl tzdata supervisor libsndfile nodejs && break || sleep 5; \
|
||||
done
|
||||
|
||||
COPY --from=builder /app /app
|
||||
COPY --from=builder /var/lib/litellm/ui /var/lib/litellm/ui
|
||||
|
|
@ -179,15 +121,10 @@ ENV PATH="/app/.venv/bin:${PATH}" \
|
|||
PRISMA_SKIP_POSTINSTALL_GENERATE=1 \
|
||||
PRISMA_HIDE_UPDATE_MESSAGE=1 \
|
||||
PRISMA_ENGINES_CHECKSUM_IGNORE_MISSING=1 \
|
||||
NPM_CONFIG_CACHE=/app/.cache/npm \
|
||||
NPM_CONFIG_PREFER_OFFLINE=true \
|
||||
PRISMA_OFFLINE_MODE=true
|
||||
|
||||
RUN sed -i 's/\r$//' docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && \
|
||||
chmod +x docker/entrypoint.sh docker/prod_entrypoint.sh && \
|
||||
mkdir -p /nonexistent /.npm /var/lib/litellm/assets /var/lib/litellm/ui /tmp/.npm && \
|
||||
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent /.npm /tmp/.npm && \
|
||||
RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
|
||||
chown -R nobody:nogroup /app /var/lib/litellm/ui /var/lib/litellm/assets /nonexistent && \
|
||||
PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \
|
||||
chown -R nobody:nogroup "$PRISMA_PATH" && \
|
||||
LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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`.
|
||||
|
|
|
|||
95
docs/my-website/docs/proxy/agentic_loop_hook.md
Normal file
95
docs/my-website/docs/proxy/agentic_loop_hook.md
Normal 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`
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -536,6 +536,7 @@ const sidebars = {
|
|||
description: "Modify requests, responses, and more",
|
||||
items: [
|
||||
"proxy/call_hooks",
|
||||
"proxy/agentic_loop_hook",
|
||||
"proxy/rules",
|
||||
]
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
14
litellm/integrations/compression_interception/__init__.py
Normal file
14
litellm/integrations/compression_interception/__init__.py
Normal 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",
|
||||
]
|
||||
399
litellm/integrations/compression_interception/handler.py
Normal file
399
litellm/integrations/compression_interception/handler.py
Normal 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
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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)}"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
158
litellm/llms/scaleway/audio_transcription/transformation.py
Normal file
158
litellm/llms/scaleway/audio_transcription/transformation.py
Normal 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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1950,7 +1950,7 @@
|
|||
"responses": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_transcriptions": true,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
"batches": false,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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={
|
||||
|
|
|
|||
|
|
@ -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={
|
||||
|
|
|
|||
|
|
@ -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 = (
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
21
litellm/router_strategy/quality_router/__init__.py
Normal file
21
litellm/router_strategy/quality_router/__init__.py
Normal 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",
|
||||
]
|
||||
74
litellm/router_strategy/quality_router/config.py
Normal file
74
litellm/router_strategy/quality_router/config.py
Normal 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")
|
||||
446
litellm/router_strategy/quality_router/quality_router.py
Normal file
446
litellm/router_strategy/quality_router/quality_router.py
Normal 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,
|
||||
)
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
27
litellm/types/integrations/compression_interception.py
Normal file
27
litellm/types/integrations/compression_interception.py
Normal 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]]
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
},
|
||||
|
|
|
|||
|
|
@ -1968,7 +1968,7 @@
|
|||
"responses": true,
|
||||
"embeddings": false,
|
||||
"image_generations": false,
|
||||
"audio_transcriptions": false,
|
||||
"audio_transcriptions": true,
|
||||
"audio_speech": false,
|
||||
"moderations": false,
|
||||
"batches": false,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
},
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
||||
|
|
|
|||
|
|
@ -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]"
|
||||
)
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
1033
tests/test_litellm/router_strategy/test_quality_router.py
Normal file
1033
tests/test_litellm/router_strategy/test_quality_router.py
Normal file
File diff suppressed because it is too large
Load diff
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue