mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Merge pull request #26298 from BerriAI/litellm_internal_staging
[Infra] Promote interal staging to main
This commit is contained in:
commit
e9e86ed956
86 changed files with 9183 additions and 1253 deletions
1186
.circleci/config.yml
1186
.circleci/config.yml
File diff suppressed because it is too large
Load diff
136
.github/workflows/test-code-quality.yml
vendored
Normal file
136
.github/workflows/test-code-quality.yml
vendored
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
name: Code Quality Checks
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
code-quality:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Checkout litellm-docs (for documentation_tests)
|
||||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
repository: BerriAI/litellm-docs
|
||||
path: _litellm_docs_checkout
|
||||
persist-credentials: false
|
||||
|
||||
- name: Wire up docs path expected by documentation_tests/*
|
||||
run: |
|
||||
# documentation_tests scripts read from docs/my-website/docs/...
|
||||
# In litellm-docs the same files live at docs/... (repo root).
|
||||
# Point docs/my-website -> litellm-docs checkout so the paths resolve.
|
||||
rm -rf docs/my-website
|
||||
ln -s ../_litellm_docs_checkout docs/my-website
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
~/.cache/uv
|
||||
.venv
|
||||
key: ${{ runner.os }}-uv-${{ hashFiles('uv.lock') }}
|
||||
restore-keys: |
|
||||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install dependencies
|
||||
run: uv sync --frozen --all-groups --all-extras
|
||||
|
||||
- name: check_licenses
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_licenses.py
|
||||
|
||||
- name: check_provider_folders_documented
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
|
||||
|
||||
- name: router_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
|
||||
|
||||
- name: test_chat_completion_imports
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_chat_completion_imports.py
|
||||
|
||||
- name: info_log_check
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/info_log_check.py
|
||||
|
||||
- name: check_guardrail_apply_decorator
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_guardrail_apply_decorator.py
|
||||
|
||||
- name: test_ban_set_verbose
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_ban_set_verbose.py
|
||||
|
||||
- name: code_qa_check_tests
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/code_qa_check_tests.py
|
||||
|
||||
- name: check_get_model_cost_key_performance
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_get_model_cost_key_performance.py
|
||||
|
||||
- name: test_proxy_types_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_proxy_types_import.py
|
||||
|
||||
- name: callback_manager_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/callback_manager_test.py
|
||||
|
||||
- name: recursive_detector
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/recursive_detector.py
|
||||
|
||||
- name: test_router_strategy_async
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/test_router_strategy_async.py
|
||||
|
||||
- name: litellm_logging_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/litellm_logging_code_coverage.py
|
||||
|
||||
- name: ensure_async_clients_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ensure_async_clients_test.py
|
||||
|
||||
- name: enforce_llms_folder_style
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/enforce_llms_folder_style.py
|
||||
|
||||
- name: prevent_key_leaks_in_exceptions
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/prevent_key_leaks_in_exceptions.py
|
||||
|
||||
- name: check_unsafe_enterprise_import
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py
|
||||
|
||||
- name: ban_copy_deepcopy_kwargs
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py
|
||||
|
||||
- name: check_fastuuid_usage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_fastuuid_usage.py
|
||||
|
||||
- name: memory_test
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py
|
||||
|
||||
- name: documentation_test_env_keys
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_env_keys.py
|
||||
|
||||
- name: documentation_test_router_settings
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_router_settings.py
|
||||
|
||||
- name: documentation_test_api_docs
|
||||
run: uv run --no-sync python ./tests/documentation_tests/test_api_docs.py
|
||||
39
.github/workflows/test-semgrep.yml
vendored
Normal file
39
.github/workflows/test-semgrep.yml
vendored
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
name: Semgrep
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_branch
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
semgrep:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Run Semgrep (custom rules)
|
||||
run: uv tool run --from 'semgrep==1.157.0' semgrep scan --config .semgrep/rules . --error
|
||||
13
Dockerfile
13
Dockerfile
|
|
@ -27,10 +27,8 @@ RUN apk add --no-cache \
|
|||
npm \
|
||||
libsndfile
|
||||
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -94,11 +92,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
|
|||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
|
||||
WORKDIR /app
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -26,10 +26,8 @@ RUN apk add --no-cache \
|
|||
npm \
|
||||
libsndfile
|
||||
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
ENV UV_PROJECT_ENVIRONMENT=/app/.venv \
|
||||
UV_LINK_MODE=copy \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
# Copy dependency metadata first for layer caching
|
||||
|
|
@ -92,11 +90,14 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 libsndfile supervi
|
|||
{ apk del --no-cache npm 2>/dev/null || true; }
|
||||
|
||||
WORKDIR /app
|
||||
ENV PRISMA_BINARY_CACHE_DIR=/app/.cache/prisma-python/binaries \
|
||||
XDG_CACHE_HOME=/app/.cache \
|
||||
PATH="/app/.venv/bin:${PATH}"
|
||||
ENV PATH="/app/.venv/bin:${PATH}"
|
||||
|
||||
COPY --from=builder /app /app
|
||||
# Prisma binaries live in $HOME/.cache (default prisma-python location),
|
||||
# which is /root/.cache here. Copy them from the builder so they survive
|
||||
# deployments that volume-mount /app/.cache (e.g. readOnlyRootFilesystem
|
||||
# + emptyDir) — otherwise the mount would shadow the baked-in query engine.
|
||||
COPY --from=builder /root/.cache /root/.cache
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete
|
||||
|
|
|
|||
|
|
@ -138,7 +138,7 @@ RUN mkdir -p /nonexistent /var/lib/litellm/assets /var/lib/litellm/ui && \
|
|||
[ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w "$LITELLM_PROXY_EXTRAS_PATH" || true && \
|
||||
chmod -R g+rX "$PRISMA_PATH" /var/lib/litellm/ui /var/lib/litellm/assets /app/.cache
|
||||
|
||||
USER nobody
|
||||
USER 65534
|
||||
|
||||
RUN prisma generate --schema=./schema.prisma
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ schemaVersion: 2.0.0
|
|||
|
||||
metadataTest:
|
||||
entrypoint: ["docker/prod_entrypoint.sh"]
|
||||
user: "nobody"
|
||||
user: "65534"
|
||||
workdir: "/app"
|
||||
|
||||
fileExistenceTests:
|
||||
|
|
|
|||
155
docs/my-website/docs/adaptive_router.md
Normal file
155
docs/my-website/docs/adaptive_router.md
Normal file
|
|
@ -0,0 +1,155 @@
|
|||
# [BETA] Adaptive Router
|
||||
|
||||
:::info
|
||||
|
||||
Beta feature. Share feedback on [Discord](https://discord.gg/wuPM9dRgDw) or [Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA).
|
||||
|
||||
:::
|
||||
|
||||
**Requirements:** LiteLLM Proxy with a Postgres database. Quality estimates are stored in Postgres and loaded on startup — without a database the router works but forgets everything learned on restart.
|
||||
|
||||
You have a cheap model and an expensive one. You want to use the cheap one when it's good enough, and the expensive one when it actually matters — without hardcoding rules you'll spend months tuning.
|
||||
|
||||
The adaptive router does this automatically. It tracks which model performs best for each type of request (code, writing, analysis, etc.) and routes accordingly, balancing quality against cost based on weights you control.
|
||||
|
||||
## Quick start
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
model_info:
|
||||
input_cost_per_token: 0.0000025
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3 # 1=budget, 2=mid, 3=frontier
|
||||
strengths: ["code_generation", "analytical_reasoning"]
|
||||
|
||||
- model_name: gpt-4o-mini
|
||||
litellm_params:
|
||||
model: openai/gpt-4o-mini
|
||||
model_info:
|
||||
input_cost_per_token: 0.00000015
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: ["factual_lookup"]
|
||||
|
||||
- model_name: my-router
|
||||
litellm_params:
|
||||
model: auto_router/adaptive_router
|
||||
adaptive_router_config:
|
||||
available_models: ["gpt-4o", "gpt-4o-mini"]
|
||||
weights:
|
||||
quality: 0.7 # raise this if quality complaints; lower if bill too high
|
||||
cost: 0.3 # must sum to 1.0 with quality
|
||||
```
|
||||
|
||||
Route to it by setting `model` to your adaptive router's name:
|
||||
|
||||
```bash
|
||||
curl -X POST {{baseURL}}/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer $LITELLM_API_KEY" \
|
||||
-d '{
|
||||
"model": "my-router",
|
||||
"messages": [
|
||||
{"role": "user", "content": "build me a python script that parses CSV"},
|
||||
{"role": "assistant", "content": "Here is a script using csv.DictReader..."},
|
||||
{"role": "user", "content": "now add error handling for missing files"},
|
||||
{"role": "assistant", "content": "Wrap the open() call in a try/except FileNotFoundError..."},
|
||||
{"role": "user", "content": "perfect, that worked. thanks!"}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
The response includes a header telling you which model was actually picked:
|
||||
|
||||
```
|
||||
x-litellm-adaptive-router-model: gpt-4o
|
||||
```
|
||||
|
||||
The "thanks!" turn in the example above fires a satisfaction signal — that's what moves the bandit.
|
||||
|
||||
## Tuning cost vs. quality
|
||||
|
||||
The `weights` are your main lever:
|
||||
|
||||
| Goal | quality | cost |
|
||||
|---|---|---|
|
||||
| Minimize cost, quality is secondary | 0.3 | 0.7 |
|
||||
| Balanced | 0.5 | 0.5 |
|
||||
| Quality-first (default) | 0.7 | 0.3 |
|
||||
| Quality non-negotiable | 0.9 | 0.1 |
|
||||
|
||||
The router learns over time. For the first ~10 requests per model, it relies on the tiers you declared. After that, real performance data takes over.
|
||||
|
||||
## Force a minimum quality tier per request
|
||||
|
||||
If a specific request needs a frontier model regardless of cost, pass this header:
|
||||
|
||||
```
|
||||
x-litellm-min-quality-tier: 3
|
||||
```
|
||||
|
||||
You can also pass `min_quality_tier` via request metadata instead of a header.
|
||||
|
||||
## What's being learned
|
||||
|
||||
The router classifies each request into one of 7 types and tracks how each model performs on each independently. A model that's great at factual lookup but poor at code will win factual requests and lose code requests — even if it's cheaper overall.
|
||||
|
||||
| Type | Example |
|
||||
|---|---|
|
||||
| `code_generation` | "write me a Python sort function" |
|
||||
| `code_understanding` | "explain what this function does" |
|
||||
| `technical_design` | "how should I design this API?" |
|
||||
| `analytical_reasoning` | "calculate the probability that..." |
|
||||
| `writing` | "draft an email to my team about..." |
|
||||
| `factual_lookup` | "what is the capital of France?" |
|
||||
| `general` | anything else |
|
||||
|
||||
[**See classifier code**](https://github.com/BerriAI/litellm/blob/litellm_adaptive_routing/litellm/router_strategy/adaptive_router/classifier.py)
|
||||
|
||||
Learning signals are inspired by [Signals: Trajectory Sampling and Triage for Agentic Interactions](https://arxiv.org/pdf/2604.00356).
|
||||
|
||||
## Inspect the current state
|
||||
|
||||
```
|
||||
GET /adaptive_router/{router_name}/state
|
||||
```
|
||||
|
||||
Returns current quality estimates per model per request type. Useful for understanding why a model is or isn't being picked.
|
||||
|
||||
```json
|
||||
{
|
||||
"routers": [
|
||||
{
|
||||
"router_name": "smart-cheap-router",
|
||||
"available_models": ["fast", "smart"],
|
||||
"weights": { "quality": 0.7, "cost": 0.3 },
|
||||
"cells": [
|
||||
{
|
||||
"request_type": "analytical_reasoning",
|
||||
"model": "fast",
|
||||
"quality_mean": 0.5,
|
||||
"samples": 0
|
||||
},
|
||||
{
|
||||
"request_type": "analytical_reasoning",
|
||||
"model": "smart",
|
||||
"quality_mean": 0.95,
|
||||
"samples": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
`quality_mean` is the key number — it's the router's current estimate of how well that model handles that request type. `samples` counts how many real observations have moved the prior (starts at 0; the cold-start prior mass is excluded).
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Latency isn't scored — a slow model can still win on quality + cost
|
||||
- Signals are regex-based and English-biased — no LLM judge
|
||||
- Hard cap of 200 observations per cell; no decay yet
|
||||
- Once a model is picked for a session, other models' turns in that session don't contribute to learning
|
||||
1
docs/my-website/package-lock.json
generated
1
docs/my-website/package-lock.json
generated
|
|
@ -26,6 +26,7 @@
|
|||
},
|
||||
"devDependencies": {
|
||||
"@docusaurus/module-type-aliases": "3.8.1",
|
||||
"ajv": "^8.18.0",
|
||||
"dotenv": "16.6.1"
|
||||
},
|
||||
"engines": {
|
||||
|
|
|
|||
|
|
@ -32,6 +32,7 @@
|
|||
},
|
||||
"devDependencies": {
|
||||
"@docusaurus/module-type-aliases": "3.8.1",
|
||||
"ajv": "^8.18.0",
|
||||
"dotenv": "16.6.1"
|
||||
},
|
||||
"browserslist": {
|
||||
|
|
|
|||
|
|
@ -1060,6 +1060,7 @@ const sidebars = {
|
|||
},
|
||||
items: [
|
||||
"routing",
|
||||
"adaptive_router",
|
||||
"scheduler",
|
||||
"proxy/auto_routing",
|
||||
"proxy/load_balancing",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,39 @@
|
|||
-- One row per (router, request_type, model). Hot path on every routing decision.
|
||||
CREATE TABLE "LiteLLM_AdaptiveRouterState" (
|
||||
router_name TEXT NOT NULL,
|
||||
request_type TEXT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
alpha DOUBLE PRECISION NOT NULL,
|
||||
beta DOUBLE PRECISION NOT NULL,
|
||||
total_samples INTEGER NOT NULL DEFAULT 0,
|
||||
last_updated_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (router_name, request_type, model_name)
|
||||
);
|
||||
|
||||
-- One row per (session, router, model). Updated per turn via the queue.
|
||||
CREATE TABLE "LiteLLM_AdaptiveRouterSession" (
|
||||
session_id TEXT NOT NULL,
|
||||
router_name TEXT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
classified_type TEXT NOT NULL,
|
||||
misalignment_count INTEGER NOT NULL DEFAULT 0,
|
||||
stagnation_count INTEGER NOT NULL DEFAULT 0,
|
||||
disengagement_count INTEGER NOT NULL DEFAULT 0,
|
||||
satisfaction_count INTEGER NOT NULL DEFAULT 0,
|
||||
failure_count INTEGER NOT NULL DEFAULT 0,
|
||||
loop_count INTEGER NOT NULL DEFAULT 0,
|
||||
exhaustion_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_user_content TEXT,
|
||||
last_assistant_content TEXT,
|
||||
tool_call_history JSONB NOT NULL DEFAULT '[]',
|
||||
pending_tool_calls JSONB NOT NULL DEFAULT '{}',
|
||||
turn_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_processed_turn INTEGER NOT NULL DEFAULT -1,
|
||||
clean_credit_awarded BOOLEAN NOT NULL DEFAULT FALSE,
|
||||
terminal_status INTEGER,
|
||||
last_activity_at TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
PRIMARY KEY (session_id, router_name, model_name)
|
||||
);
|
||||
|
||||
CREATE INDEX "idx_adaptive_router_session_activity"
|
||||
ON "LiteLLM_AdaptiveRouterSession" (last_activity_at);
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_TeamMembership" ADD COLUMN "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
|
|
@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
|
|||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
@ -1223,3 +1224,46 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// Per-(router, request_type, model) Beta posterior for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterState {
|
||||
router_name String
|
||||
request_type String
|
||||
model_name String
|
||||
alpha Float
|
||||
beta Float
|
||||
total_samples Int @default(0)
|
||||
last_updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([router_name, request_type, model_name])
|
||||
}
|
||||
|
||||
// Per-(session, router, model) signal counters for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterSession {
|
||||
session_id String
|
||||
router_name String
|
||||
model_name String
|
||||
classified_type String
|
||||
|
||||
misalignment_count Int @default(0)
|
||||
stagnation_count Int @default(0)
|
||||
disengagement_count Int @default(0)
|
||||
satisfaction_count Int @default(0)
|
||||
failure_count Int @default(0)
|
||||
loop_count Int @default(0)
|
||||
exhaustion_count Int @default(0)
|
||||
|
||||
last_user_content String?
|
||||
last_assistant_content String?
|
||||
tool_call_history Json @default("[]")
|
||||
pending_tool_calls Json @default("{}")
|
||||
|
||||
turn_count Int @default(0)
|
||||
last_processed_turn Int @default(-1)
|
||||
clean_credit_awarded Boolean @default(false)
|
||||
terminal_status Int?
|
||||
last_activity_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.67"
|
||||
version = "0.4.68"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -25,7 +25,7 @@ required-version = "==0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.67"
|
||||
version = "0.4.68"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -164,6 +164,7 @@ MCP_STDIO_ALLOWED_COMMANDS: frozenset = frozenset(
|
|||
LITELLM_UI_ALLOW_HEADERS = [
|
||||
"x-litellm-semantic-filter",
|
||||
"x-litellm-semantic-filter-tools",
|
||||
"x-litellm-adaptive-router-model",
|
||||
]
|
||||
|
||||
# Gemini model-specific minimal thinking budget constants
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ from datetime import datetime
|
|||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
ClassVar,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
|
|
@ -12,6 +13,7 @@ from typing import (
|
|||
)
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.guardrails import (
|
||||
|
|
@ -81,6 +83,9 @@ class ModifyResponseException(Exception):
|
|||
|
||||
|
||||
class CustomGuardrail(CustomLogger):
|
||||
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
|
||||
use_native_during_call_hook: ClassVar[bool] = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: Optional[str] = None,
|
||||
|
|
@ -637,6 +642,13 @@ class CustomGuardrail(CustomLogger):
|
|||
if isinstance(item, dict):
|
||||
item.pop("secret_fields", None)
|
||||
|
||||
# Default-safe behavior: never persist raw matched spans in standard
|
||||
# guardrail logging payloads (single shared implementation; Bedrock hooks pass
|
||||
# raw provider JSON so redaction is not duplicated upstream).
|
||||
clean_guardrail_response = redact_nested_match_and_regex_keys(
|
||||
clean_guardrail_response
|
||||
)
|
||||
|
||||
slg = StandardLoggingGuardrailInformation(
|
||||
guardrail_name=self.guardrail_name,
|
||||
guardrail_provider=guardrail_provider,
|
||||
|
|
|
|||
|
|
@ -2289,6 +2289,10 @@ class OpenTelemetry(CustomLogger):
|
|||
# Remove trailing slash
|
||||
endpoint = endpoint.rstrip("/")
|
||||
|
||||
# Splunk Observability Cloud OTLP/HTTP uses /v2/trace/otlp (not /v1/traces). Do not rewrite.
|
||||
if signal_type == "traces" and "/v2/trace/otlp" in endpoint:
|
||||
return endpoint
|
||||
|
||||
# Check if endpoint already ends with the correct signal path
|
||||
target_path = f"/v1/{signal_type}"
|
||||
if endpoint.endswith(target_path):
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -435,3 +436,42 @@ def filter_internal_params(
|
|||
|
||||
# Filter out internal parameters
|
||||
return {k: v for k, v in data.items() if k not in internal_params}
|
||||
|
||||
|
||||
def redact_nested_match_and_regex_keys(
|
||||
payload: Union[dict, List[Any], str, None],
|
||||
) -> Union[dict, List[Any], str, None]:
|
||||
"""
|
||||
Deep-copy `payload` and replace every `match` / `regex` string field with
|
||||
"[REDACTED]" anywhere in nested dict/list structures.
|
||||
|
||||
Used for guardrail spend/compliance logging so raw spans are not persisted.
|
||||
"""
|
||||
if payload is None or isinstance(payload, str):
|
||||
return payload
|
||||
try:
|
||||
redacted: Union[dict, List[Any], str, None] = copy.deepcopy(payload)
|
||||
except Exception:
|
||||
return payload
|
||||
|
||||
# Iterative traversal; `seen` guards against cyclic refs preserved by deepcopy.
|
||||
try:
|
||||
seen: set = set()
|
||||
stack: List[Any] = [redacted]
|
||||
while stack:
|
||||
node = stack.pop()
|
||||
node_id = id(node)
|
||||
if node_id in seen:
|
||||
continue
|
||||
seen.add(node_id)
|
||||
if isinstance(node, dict):
|
||||
if "match" in node:
|
||||
node["match"] = "[REDACTED]"
|
||||
if "regex" in node:
|
||||
node["regex"] = "[REDACTED]"
|
||||
stack.extend(node.values())
|
||||
elif isinstance(node, list):
|
||||
stack.extend(node)
|
||||
except Exception:
|
||||
return payload
|
||||
return redacted
|
||||
|
|
|
|||
|
|
@ -22872,6 +22872,22 @@
|
|||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-k2.6": {
|
||||
"cache_read_input_token_cost": 1.6e-07,
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "moonshot",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-06,
|
||||
"source": "https://platform.kimi.ai/docs/pricing/chat-k26",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
|
|
|
|||
|
|
@ -1,42 +1,83 @@
|
|||
# model_list:
|
||||
# - model_name: claude-sonnet-4-6
|
||||
# litellm_params: {model: anthropic/claude-sonnet-4-6}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 1
|
||||
# keywords: [tin]
|
||||
# - model_name: gpt-4o-mini
|
||||
# litellm_params: {model: openai/gpt-4o-mini}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 1
|
||||
# keywords: []
|
||||
# - model_name: gpt-4o
|
||||
# litellm_params: {model: openai/gpt-4o}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 2
|
||||
# keywords: [vision, function_calling]
|
||||
# - model_name: opus
|
||||
# litellm_params: {model: anthropic/claude-opus-4-7}
|
||||
# model_info:
|
||||
# litellm_routing_preferences:
|
||||
# quality_tier: 3
|
||||
# keywords: ["architecture", "design"]
|
||||
# - model_name: my-quality-router
|
||||
# litellm_params:
|
||||
# model: auto_router/adaptive_router
|
||||
# adaptive_router_default_model: gpt-4o-mini
|
||||
# adaptive_router_config:
|
||||
# available_models: [gpt-4o-mini, gpt-4o, opus, claude-sonnet-4-6]
|
||||
# Example proxy config for the adaptive router (v0).
|
||||
#
|
||||
# Wires one logical router ("smart-cheap-router") that adaptively picks between
|
||||
# two real deployments ("fast" and "smart") based on per-session feedback signals.
|
||||
#
|
||||
# How to use from a client:
|
||||
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
|
||||
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
|
||||
# sticky-session routing within a conversation.
|
||||
#
|
||||
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
|
||||
|
||||
model_list:
|
||||
|
||||
# OpenAI model for /v1/chat/completions test — 200x custom pricing
|
||||
- model_name: "gpt-4.1-mini"
|
||||
# ---- The adaptive router "control" deployment -------------------------
|
||||
# `model_name` is what clients call. `available_models` lists the underlying
|
||||
# deployments the router is allowed to pick from (must match other model_name
|
||||
# entries in this list).
|
||||
- model_name: smart-cheap-router
|
||||
litellm_params:
|
||||
model: openai/gpt-4.1-mini
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
model_info:
|
||||
id: gpt-4.1-mini-custom-pricing
|
||||
input_cost_per_token: 0.00004 # 100x standard ($0.40/1M = $0.0000004)
|
||||
output_cost_per_token: 0.00016 # 100x standard ($1.60/1M = $0.0000016)
|
||||
model: auto_router/adaptive_router
|
||||
adaptive_router_config:
|
||||
available_models: ["fast", "smart"]
|
||||
weights:
|
||||
quality: 0.7
|
||||
cost: 0.3
|
||||
|
||||
# OpenAI model for /v1/responses test — 100x custom pricing
|
||||
- model_name: "gpt-5"
|
||||
litellm_params:
|
||||
model: openai/gpt-5
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
model_info:
|
||||
id: gpt-5-custom-pricing
|
||||
mode: "chat"
|
||||
input_cost_per_token: 125 # 100x standard ($1.25/1M = $0.00000125)
|
||||
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-6"
|
||||
# ---- Underlying deployments the router picks from ---------------------
|
||||
- model_name: fast
|
||||
litellm_params:
|
||||
model: anthropic/claude-sonnet-4-6
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
input_cost_per_token: 0.00000015
|
||||
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)
|
||||
- model_name: my-auto
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: []
|
||||
|
||||
- model_name: smart
|
||||
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
|
||||
model: anthropic/claude-opus-4-7
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
input_cost_per_token: 0.0000050
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3
|
||||
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
|
||||
|
||||
litellm_settings:
|
||||
drop_params: True
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234 # REPLACE in production
|
||||
|
|
|
|||
|
|
@ -1997,7 +1997,12 @@ class TeamRequest(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record"""
|
||||
"""Represents user-controllable params for a LiteLLM_BudgetTable record.
|
||||
|
||||
Budget-write paths use `model_fields.keys()` on this class as an allowlist
|
||||
for user input. Keep server-managed fields (e.g. `budget_reset_at`) on
|
||||
`LiteLLM_BudgetTableFull` so they aren't user-settable.
|
||||
"""
|
||||
|
||||
budget_id: Optional[str] = None
|
||||
soft_budget: Optional[float] = None
|
||||
|
|
@ -2015,7 +2020,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase):
|
|||
|
||||
|
||||
class LiteLLM_BudgetTableFull(LiteLLM_BudgetTable):
|
||||
"""Represents all params for a LiteLLM_BudgetTable record"""
|
||||
"""LiteLLM_BudgetTable + server-managed fields returned on API responses."""
|
||||
|
||||
budget_reset_at: Optional[datetime] = None
|
||||
created_at: datetime
|
||||
|
|
@ -3695,7 +3700,11 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase):
|
|||
team_id: str
|
||||
budget_id: Optional[str] = None
|
||||
spend: Optional[float] = 0.0
|
||||
litellm_budget_table: Optional[LiteLLM_BudgetTable]
|
||||
total_spend: Optional[float] = 0.0
|
||||
# Union so Pydantic picks Full when data has server-managed fields
|
||||
# (/team/info) and Base when callers/tests construct with only
|
||||
# user-settable fields.
|
||||
litellm_budget_table: Optional[Union[LiteLLM_BudgetTableFull, LiteLLM_BudgetTable]]
|
||||
|
||||
def safe_get_team_member_rpm_limit(self) -> Optional[int]:
|
||||
if self.litellm_budget_table is not None:
|
||||
|
|
|
|||
|
|
@ -433,7 +433,8 @@ def add_guardrail_to_applied_guardrails_header(
|
|||
return
|
||||
_metadata = request_data.get("metadata", None) or {}
|
||||
if "applied_guardrails" in _metadata:
|
||||
_metadata["applied_guardrails"].append(guardrail_name)
|
||||
if guardrail_name not in _metadata["applied_guardrails"]:
|
||||
_metadata["applied_guardrails"].append(guardrail_name)
|
||||
else:
|
||||
_metadata["applied_guardrails"] = [guardrail_name]
|
||||
# Ensure metadata is set back to request_data (important when metadata didn't exist)
|
||||
|
|
|
|||
|
|
@ -1300,7 +1300,10 @@ class DBSpendUpdateWriter:
|
|||
|
||||
batcher.litellm_teammembership.update_many( # 'update_many' prevents error from being raised if no row exists
|
||||
where={"team_id": team_id, "user_id": user_id},
|
||||
data={"spend": {"increment": response_cost}},
|
||||
data={
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
},
|
||||
)
|
||||
# Transaction succeeded, break out of retry loop
|
||||
break
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.constants import (
|
|||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import (
|
||||
BaseDailySpendTransaction,
|
||||
DailyAgentSpendTransaction,
|
||||
DailyEndUserSpendTransaction,
|
||||
DailyOrganizationSpendTransaction,
|
||||
|
|
@ -29,6 +30,8 @@ from litellm.proxy._types import (
|
|||
DailyTeamSpendTransaction,
|
||||
DailyUserSpendTransaction,
|
||||
DBSpendUpdateTransactions,
|
||||
Litellm_EntityType,
|
||||
SpendUpdateQueueItem,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
|
||||
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
||||
|
|
@ -259,9 +262,36 @@ class RedisUpdateBuffer:
|
|||
if len(rpush_list) == 0:
|
||||
return
|
||||
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
try:
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
except Exception as e:
|
||||
# The in-memory queues were already drained above. If we let the
|
||||
# exception propagate without restoring, the aggregated spend is
|
||||
# permanently lost. Re-enqueue so the next scheduler tick retries.
|
||||
verbose_proxy_logger.error(
|
||||
"Spend tracking - failed to push aggregated spend updates to Redis. "
|
||||
"Restoring %d transaction sets to in-memory queues for retry on next tick. "
|
||||
"Error: %s",
|
||||
len(rpush_list),
|
||||
str(e),
|
||||
)
|
||||
await self._restore_spend_updates_to_in_memory_queues(
|
||||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
daily_spend_update_transactions=daily_spend_update_transactions,
|
||||
daily_team_spend_update_transactions=daily_team_spend_update_transactions,
|
||||
daily_org_spend_update_transactions=daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions=daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions=daily_agent_spend_update_transactions,
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_update_queue,
|
||||
daily_team_spend_update_queue=daily_team_spend_update_queue,
|
||||
daily_org_spend_update_queue=daily_org_spend_update_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_spend_update_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_spend_update_queue,
|
||||
)
|
||||
return
|
||||
|
||||
# Emit gauge events for each queue
|
||||
for i, queue_size in enumerate(result_lengths):
|
||||
|
|
@ -271,6 +301,101 @@ class RedisUpdateBuffer:
|
|||
service=service_types[i],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _restore_spend_updates_to_in_memory_queues(
|
||||
db_spend_update_transactions: Optional[DBSpendUpdateTransactions],
|
||||
daily_spend_update_transactions: Optional[Dict[str, BaseDailySpendTransaction]],
|
||||
daily_team_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_org_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_end_user_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
daily_agent_spend_update_transactions: Optional[
|
||||
Dict[str, BaseDailySpendTransaction]
|
||||
],
|
||||
spend_update_queue: SpendUpdateQueue,
|
||||
daily_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_team_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_org_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
|
||||
daily_agent_spend_update_queue: DailySpendUpdateQueue,
|
||||
) -> None:
|
||||
"""
|
||||
Put drained-but-unpushed transactions back into in-memory queues.
|
||||
|
||||
Called when the Redis rpush pipeline raises. Without this, all spend
|
||||
data aggregated during the current scheduler tick is permanently lost
|
||||
because the source queues were already drained before the rpush.
|
||||
"""
|
||||
if db_spend_update_transactions is not None:
|
||||
entity_entries: List[
|
||||
Tuple[Litellm_EntityType, Optional[Dict[str, float]]]
|
||||
] = [
|
||||
(
|
||||
Litellm_EntityType.USER,
|
||||
db_spend_update_transactions.get("user_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.END_USER,
|
||||
db_spend_update_transactions.get("end_user_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.KEY,
|
||||
db_spend_update_transactions.get("key_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TEAM,
|
||||
db_spend_update_transactions.get("team_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TEAM_MEMBER,
|
||||
db_spend_update_transactions.get("team_member_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.ORGANIZATION,
|
||||
db_spend_update_transactions.get("org_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.TAG,
|
||||
db_spend_update_transactions.get("tag_list_transactions"),
|
||||
),
|
||||
(
|
||||
Litellm_EntityType.AGENT,
|
||||
db_spend_update_transactions.get("agent_list_transactions"),
|
||||
),
|
||||
]
|
||||
for entity_type, entities in entity_entries:
|
||||
if not entities:
|
||||
continue
|
||||
for entity_id, cost in entities.items():
|
||||
await spend_update_queue.add_update(
|
||||
SpendUpdateQueueItem(
|
||||
entity_type=entity_type,
|
||||
entity_id=entity_id,
|
||||
response_cost=cost,
|
||||
)
|
||||
)
|
||||
|
||||
daily_pairs: List[
|
||||
Tuple[Optional[Dict[str, BaseDailySpendTransaction]], DailySpendUpdateQueue]
|
||||
] = [
|
||||
(daily_spend_update_transactions, daily_spend_update_queue),
|
||||
(daily_team_spend_update_transactions, daily_team_spend_update_queue),
|
||||
(daily_org_spend_update_transactions, daily_org_spend_update_queue),
|
||||
(
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_end_user_spend_update_queue,
|
||||
),
|
||||
(daily_agent_spend_update_transactions, daily_agent_spend_update_queue),
|
||||
]
|
||||
for daily_txns, daily_queue in daily_pairs:
|
||||
if daily_txns:
|
||||
await daily_queue.update_queue.put(daily_txns)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
db_spend_update_transactions: DBSpendUpdateTransactions,
|
||||
|
|
|
|||
|
|
@ -0,0 +1,52 @@
|
|||
# Example proxy config for the adaptive router (v0).
|
||||
#
|
||||
# Wires one logical router ("smart-cheap-router") that adaptively picks between
|
||||
# two real deployments ("fast" and "smart") based on per-session feedback signals.
|
||||
#
|
||||
# How to use from a client:
|
||||
# POST /v1/chat/completions { "model": "smart-cheap-router", ... }
|
||||
# Add { "metadata": { "litellm_session_id": "<your-session-id>" } } to enable
|
||||
# sticky-session routing within a conversation.
|
||||
#
|
||||
# Required env vars: OPENAI_API_KEY, DATABASE_URL.
|
||||
|
||||
model_list:
|
||||
# ---- The adaptive router "control" deployment -------------------------
|
||||
# `model_name` is what clients call. `available_models` lists the underlying
|
||||
# deployments the router is allowed to pick from (must match other model_name
|
||||
# entries in this list).
|
||||
- model_name: smart-cheap-router
|
||||
litellm_params:
|
||||
model: auto_router/adaptive_router # required prefix -- triggers adaptive-router init
|
||||
adaptive_router_config:
|
||||
available_models: ["fast", "smart"]
|
||||
weights:
|
||||
quality: 0.7
|
||||
cost: 0.3
|
||||
|
||||
# ---- Underlying deployments the router picks from ---------------------
|
||||
- model_name: fast
|
||||
litellm_params:
|
||||
model: openai/gpt-4o-mini
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
input_cost_per_token: 0.00000015
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: []
|
||||
|
||||
- model_name: smart
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
input_cost_per_token: 0.0000050
|
||||
model_info:
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3
|
||||
strengths: ["code_generation", "technical_design", "analytical_reasoning"]
|
||||
|
||||
litellm_settings:
|
||||
drop_params: True
|
||||
|
||||
general_settings:
|
||||
master_key: sk-1234 # REPLACE in production
|
||||
|
|
@ -5,7 +5,6 @@
|
|||
# +-------------------------------------------------------------+
|
||||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
|
||||
import copy
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
|
@ -18,6 +17,7 @@ from typing import (
|
|||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
ClassVar,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
|
|
@ -33,6 +33,7 @@ from fastapi import HTTPException
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.exceptions import GuardrailInterventionNormalStringError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -79,56 +80,33 @@ class GuardrailMessageFilterResult(NamedTuple):
|
|||
|
||||
|
||||
def _redact_pii_matches(response_json: dict) -> dict:
|
||||
try:
|
||||
# Create a deep copy to avoid modifying the original response
|
||||
redacted_response = copy.deepcopy(response_json)
|
||||
"""
|
||||
Redact match-like fields from a Bedrock ApplyGuardrail JSON payload.
|
||||
|
||||
# Get assessments from the response
|
||||
# NOTE: We use `.get("key") or []` instead of `.get("key", [])` because
|
||||
# the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null).
|
||||
# In Python, dict.get("key", []) returns None (not []) when the key exists
|
||||
# with a None/null value. The `or []` ensures we always get an iterable,
|
||||
# preventing "TypeError: 'NoneType' object is not iterable".
|
||||
assessments = redacted_response.get("assessments") or []
|
||||
if not assessments:
|
||||
return redacted_response
|
||||
Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend
|
||||
logging). Kept as a Bedrock-module entry point for existing unit tests.
|
||||
"""
|
||||
redacted = redact_nested_match_and_regex_keys(response_json)
|
||||
return redacted if isinstance(redacted, dict) else response_json
|
||||
|
||||
for assessment in assessments:
|
||||
# Redact PII entities in sensitive information policy
|
||||
sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
|
||||
if sensitive_info_policy:
|
||||
pii_entities = sensitive_info_policy.get("piiEntities") or []
|
||||
for pii_entity in pii_entities:
|
||||
if "match" in pii_entity:
|
||||
pii_entity["match"] = "[REDACTED]"
|
||||
|
||||
# Redact regex matches
|
||||
regexes = sensitive_info_policy.get("regexes") or []
|
||||
for regex_match in regexes:
|
||||
if "match" in regex_match:
|
||||
regex_match["match"] = "[REDACTED]"
|
||||
def _redact_assessment_match_fields(assessments: List[dict]) -> List[dict]:
|
||||
"""
|
||||
Redact sensitive match-like fields from blocked assessment summaries.
|
||||
|
||||
# Redact custom word matches in word policy
|
||||
word_policy = assessment.get("wordPolicy")
|
||||
if word_policy:
|
||||
custom_words = word_policy.get("customWords") or []
|
||||
for custom_word in custom_words:
|
||||
if "match" in custom_word:
|
||||
custom_word["match"] = "[REDACTED]"
|
||||
|
||||
managed_words = word_policy.get("managedWordLists") or []
|
||||
for managed_word in managed_words:
|
||||
if "match" in managed_word:
|
||||
managed_word["match"] = "[REDACTED]"
|
||||
|
||||
return redacted_response
|
||||
except Exception as e:
|
||||
# We do not want to fail in any case so this is just a warning
|
||||
verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e))
|
||||
return response_json
|
||||
This is used for customer-visible error payloads (HTTPException.detail) where
|
||||
we want to preserve policy/type/action metadata without echoing raw matched
|
||||
content.
|
||||
"""
|
||||
redacted = redact_nested_match_and_regex_keys(assessments)
|
||||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
||||
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
|
||||
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
|
||||
use_native_during_call_hook: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrailIdentifier: Optional[str] = None,
|
||||
|
|
@ -419,6 +397,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
messages: Optional[List[AllMessageValues]] = None,
|
||||
response: Optional[Union[Any, litellm.ModelResponse]] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
logging_event_type: Optional[GuardrailEventHooks] = None,
|
||||
) -> BedrockGuardrailResponse:
|
||||
from datetime import datetime
|
||||
|
||||
|
|
@ -456,11 +435,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
prepared_request.headers,
|
||||
)
|
||||
|
||||
event_type = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if source == "INPUT"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
# UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API
|
||||
# body, which must not be confused with the proxy hook (pre_call / during_call /
|
||||
# post_call). When omitted, keep legacy mapping for backward compatibility.
|
||||
if logging_event_type is not None:
|
||||
event_type = logging_event_type
|
||||
else:
|
||||
event_type = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if source == "INPUT"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
|
||||
try:
|
||||
httpx_response = await self.async_handler.post(
|
||||
|
|
@ -515,9 +500,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
# Add guardrail information to request trace
|
||||
#########################################################
|
||||
_json_response = httpx_response.json()
|
||||
# Raw Bedrock JSON is passed here; match/regex redaction runs once inside
|
||||
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response=httpx_response.json(),
|
||||
guardrail_json_response=_json_response,
|
||||
request_data=request_data or {},
|
||||
guardrail_status=self._get_bedrock_guardrail_response_status(
|
||||
response=httpx_response
|
||||
|
|
@ -530,9 +518,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
if httpx_response.status_code == 200:
|
||||
# check if the response was flagged
|
||||
_json_response = httpx_response.json()
|
||||
redacted_response = _redact_pii_matches(_json_response)
|
||||
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock AI response : %s",
|
||||
redact_nested_match_and_regex_keys(_json_response),
|
||||
)
|
||||
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
|
||||
if self._should_raise_guardrail_blocked_exception(
|
||||
bedrock_guardrail_response
|
||||
|
|
@ -809,7 +798,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
assessments = self._extract_blocked_assessments(response)
|
||||
if assessments:
|
||||
detail["assessments"] = assessments
|
||||
detail["assessments"] = _redact_assessment_match_fields(assessments)
|
||||
|
||||
return HTTPException(status_code=400, detail=detail)
|
||||
|
||||
|
|
@ -831,8 +820,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
return False
|
||||
|
||||
# Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED)
|
||||
# NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API.
|
||||
# See _redact_pii_matches() for detailed explanation of the null safety pattern.
|
||||
# NOTE: Use `.get("k") or []` not `.get("k", [])` — Bedrock can return explicit
|
||||
# JSON null; dict.get("k", []) then yields None, and `for x in None` raises.
|
||||
assessments = response.get("assessments") or []
|
||||
if not assessments:
|
||||
return False
|
||||
|
|
@ -952,7 +941,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
try:
|
||||
bedrock_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="INPUT", messages=filtered_messages, request_data=data
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.pre_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
bedrock_guardrail_response = e.message
|
||||
|
|
@ -1024,7 +1016,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
try:
|
||||
bedrock_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="INPUT", messages=filtered_messages, request_data=data
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.during_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
bedrock_guardrail_response = e.message
|
||||
|
|
@ -1128,9 +1123,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source="INPUT",
|
||||
messages=input_messages,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
output_task = self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=response, request_data=data
|
||||
source="OUTPUT",
|
||||
response=response,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
|
||||
# Execute both requests in parallel
|
||||
|
|
@ -1144,7 +1143,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
|
||||
try:
|
||||
output_content_bedrock = await self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=response, request_data=data
|
||||
source="OUTPUT",
|
||||
response=response,
|
||||
request_data=data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
output_content_bedrock = e.message
|
||||
|
|
@ -1271,9 +1273,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source="INPUT",
|
||||
messages=input_messages,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
) # Only input messages
|
||||
output_task = self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=assembled_model_response
|
||||
source="OUTPUT",
|
||||
response=assembled_model_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
) # Only response
|
||||
|
||||
# Execute both requests in parallel
|
||||
|
|
@ -1287,7 +1293,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Only run OUTPUT validation (INPUT was already validated in pre_call or during_call)
|
||||
try:
|
||||
output_guardrail_response = await self.make_bedrock_api_request(
|
||||
source="OUTPUT", response=assembled_model_response
|
||||
source="OUTPUT",
|
||||
response=assembled_model_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.post_call,
|
||||
)
|
||||
except GuardrailInterventionNormalStringError as e:
|
||||
output_guardrail_response = e.message
|
||||
|
|
@ -1564,6 +1573,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
# Bedrock will throw an error if there is no text to process
|
||||
if filtered_messages:
|
||||
_log_hook = (
|
||||
GuardrailEventHooks.pre_call
|
||||
if input_type == "request"
|
||||
else GuardrailEventHooks.post_call
|
||||
)
|
||||
# Map the abstract input_type to the Bedrock source parameter.
|
||||
# "request" -> INPUT (scan user-supplied content)
|
||||
# "response" -> OUTPUT (scan model-generated content)
|
||||
|
|
@ -1594,12 +1608,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source="OUTPUT",
|
||||
response=synthetic_response,
|
||||
request_data=request_data,
|
||||
logging_event_type=_log_hook,
|
||||
)
|
||||
else:
|
||||
bedrock_response = await self.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=filtered_messages,
|
||||
request_data=request_data,
|
||||
logging_event_type=_log_hook,
|
||||
)
|
||||
|
||||
# Apply any masking that was applied by the guardrail
|
||||
|
|
|
|||
|
|
@ -952,6 +952,17 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
|
|||
_run_background_health_check()
|
||||
) # start the background health check coroutine.
|
||||
|
||||
# Start adaptive-router queue flusher unconditionally — adaptive routers
|
||||
# may be added later via `/config/reload`, and the flusher is a no-op when
|
||||
# `llm_router.adaptive_routers` is empty. Per-router DB state is loaded
|
||||
# lazily by the flusher on first tick (see `_state_loaded` flag) so
|
||||
# hot-reloaded routers also get their persisted priors.
|
||||
if llm_router is not None and getattr(llm_router, "adaptive_routers", None):
|
||||
for _ar in llm_router.adaptive_routers.values():
|
||||
await _ar.load_state_from_db(prisma_client)
|
||||
_ar._state_loaded = True
|
||||
asyncio.create_task(_adaptive_router_flusher_loop())
|
||||
|
||||
## [Optional] Initialize dd tracer
|
||||
ProxyStartupEvent._init_dd_tracer()
|
||||
|
||||
|
|
@ -2442,6 +2453,38 @@ def _write_health_state_to_router_cache(
|
|||
)
|
||||
|
||||
|
||||
_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS = 10
|
||||
|
||||
|
||||
async def _adaptive_router_flusher_loop():
|
||||
"""
|
||||
Drain every AdaptiveRouter's in-memory state + session aggregators into
|
||||
Postgres on a fixed cadence. Hot-path writes go to memory; this loop is
|
||||
the only writer to the adaptive router DB tables.
|
||||
"""
|
||||
global llm_router, prisma_client
|
||||
while True:
|
||||
try:
|
||||
await asyncio.sleep(_ADAPTIVE_ROUTER_FLUSH_INTERVAL_SECONDS)
|
||||
adaptive_routers = getattr(llm_router, "adaptive_routers", None) or {}
|
||||
if not adaptive_routers or prisma_client is None:
|
||||
continue
|
||||
for ar in adaptive_routers.values():
|
||||
# Lazy state load: covers adaptive routers registered via
|
||||
# `/config/reload` after proxy boot.
|
||||
if not getattr(ar, "_state_loaded", False):
|
||||
try:
|
||||
await ar.load_state_from_db(prisma_client)
|
||||
finally:
|
||||
ar._state_loaded = True
|
||||
await ar.queue.flush_state_to_db(prisma_client)
|
||||
await ar.queue.flush_session_to_db(prisma_client)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
verbose_proxy_logger.exception("adaptive_router flusher iteration failed")
|
||||
|
||||
|
||||
async def _run_background_health_check():
|
||||
"""
|
||||
Periodically run health checks in the background on the endpoints.
|
||||
|
|
@ -13968,6 +14011,38 @@ async def home(request: Request):
|
|||
return "LiteLLM: RUNNING"
|
||||
|
||||
|
||||
@router.get(
|
||||
"/adaptive_router/state",
|
||||
tags=["adaptive_router"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def get_adaptive_router_state(
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""Return live bandit posteriors + queue depth for every configured adaptive router.
|
||||
|
||||
Admin-only. Returns 404 if no adaptive router is configured.
|
||||
|
||||
Response shape: `{"routers": [<snapshot>, ...]}` — one snapshot per
|
||||
adaptive-router deployment. Each snapshot's `router_name` field identifies
|
||||
which deployment it came from.
|
||||
"""
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": CommonProxyErrors.not_allowed_access.value},
|
||||
)
|
||||
if llm_router is None or not llm_router.adaptive_routers:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": "No adaptive_router is configured on this proxy."},
|
||||
)
|
||||
snapshots = [
|
||||
await ar.get_state_snapshot() for ar in llm_router.adaptive_routers.values()
|
||||
]
|
||||
return {"routers": snapshots}
|
||||
|
||||
|
||||
@router.get("/routes", dependencies=[Depends(user_api_key_auth)])
|
||||
async def get_routes():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
|
|||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
@ -1223,3 +1224,46 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// Per-(router, request_type, model) Beta posterior for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterState {
|
||||
router_name String
|
||||
request_type String
|
||||
model_name String
|
||||
alpha Float
|
||||
beta Float
|
||||
total_samples Int @default(0)
|
||||
last_updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([router_name, request_type, model_name])
|
||||
}
|
||||
|
||||
// Per-(session, router, model) signal counters for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterSession {
|
||||
session_id String
|
||||
router_name String
|
||||
model_name String
|
||||
classified_type String
|
||||
|
||||
misalignment_count Int @default(0)
|
||||
stagnation_count Int @default(0)
|
||||
disengagement_count Int @default(0)
|
||||
satisfaction_count Int @default(0)
|
||||
failure_count Int @default(0)
|
||||
loop_count Int @default(0)
|
||||
exhaustion_count Int @default(0)
|
||||
|
||||
last_user_content String?
|
||||
last_assistant_content String?
|
||||
tool_call_history Json @default("[]")
|
||||
pending_tool_calls Json @default("{}")
|
||||
|
||||
turn_count Int @default(0)
|
||||
last_processed_turn Int @default(-1)
|
||||
clean_credit_awarded Boolean @default(false)
|
||||
terminal_status Int?
|
||||
last_activity_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -940,7 +940,11 @@ class ProxyLogging:
|
|||
Result from the guardrail execution
|
||||
"""
|
||||
# Use unified_guardrail if callback has apply_guardrail method
|
||||
use_unified = "apply_guardrail" in type(callback).__dict__
|
||||
has_apply_guardrail = "apply_guardrail" in type(callback).__dict__
|
||||
use_unified = has_apply_guardrail and not (
|
||||
hook_type == "during_call"
|
||||
and getattr(callback, "use_native_during_call_hook", False)
|
||||
)
|
||||
if use_unified:
|
||||
data["guardrail_to_apply"] = callback
|
||||
|
||||
|
|
@ -1540,6 +1544,7 @@ class ProxyLogging:
|
|||
if (
|
||||
"apply_guardrail" in type(callback).__dict__
|
||||
and user_api_key_dict is not None
|
||||
and not getattr(callback, "use_native_during_call_hook", False)
|
||||
):
|
||||
data["guardrail_to_apply"] = callback
|
||||
guardrail_task = self._run_guardrail_task_with_enrichment(
|
||||
|
|
|
|||
|
|
@ -200,6 +200,9 @@ if TYPE_CHECKING:
|
|||
from litellm.router_strategy.complexity_router.complexity_router import (
|
||||
ComplexityRouter,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import (
|
||||
AdaptiveRouter,
|
||||
)
|
||||
from litellm.router_strategy.quality_router.quality_router import (
|
||||
QualityRouter,
|
||||
)
|
||||
|
|
@ -209,6 +212,7 @@ else:
|
|||
Span = Any
|
||||
AutoRouter = Any
|
||||
ComplexityRouter = Any
|
||||
AdaptiveRouter = Any
|
||||
QualityRouter = Any
|
||||
PreRoutingHookResponse = Any
|
||||
|
||||
|
|
@ -468,6 +472,7 @@ class Router:
|
|||
) # {"TEAM_ID": PatternMatchRouter}
|
||||
self.auto_routers: Dict[str, "AutoRouter"] = {}
|
||||
self.complexity_routers: Dict[str, "ComplexityRouter"] = {}
|
||||
self.adaptive_routers: Dict[str, "AdaptiveRouter"] = {}
|
||||
self.quality_routers: Dict[str, "QualityRouter"] = {}
|
||||
|
||||
# Initialize model_group_alias early since it's used in set_model_list
|
||||
|
|
@ -5369,8 +5374,13 @@ class Router:
|
|||
_request_team_id: Optional[str] = (kwargs.get("metadata", {}) or {}).get(
|
||||
"user_api_key_team_id"
|
||||
)
|
||||
all_deployments = self._get_all_deployments(
|
||||
model_name=original_model_group, team_id=_request_team_id
|
||||
# Use wildcard-aware lookup so order-based fallback also works for model
|
||||
# groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`).
|
||||
all_deployments = (
|
||||
self.get_model_list(
|
||||
model_name=original_model_group, team_id=_request_team_id
|
||||
)
|
||||
or []
|
||||
)
|
||||
_order_set: set = {
|
||||
litellm.utils._get_deployment_order(d)
|
||||
|
|
@ -6815,10 +6825,13 @@ class Router:
|
|||
Check if the deployment is an auto-router deployment (semantic router).
|
||||
|
||||
Returns True if the litellm_params model starts with "auto_router/"
|
||||
but NOT "auto_router/complexity_router" (which uses complexity routing).
|
||||
but NOT "auto_router/complexity_router" or "auto_router/adaptive_router"
|
||||
(which use the complexity-router and adaptive-router strategies).
|
||||
"""
|
||||
if litellm_params.model.startswith("auto_router/complexity_router"):
|
||||
return False # This is handled by complexity_router
|
||||
if litellm_params.model.startswith("auto_router/adaptive_router"):
|
||||
return False # This is handled by adaptive_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/"):
|
||||
|
|
@ -6927,6 +6940,144 @@ class Router:
|
|||
)
|
||||
self.complexity_routers[deployment.model_name] = complexity_router
|
||||
|
||||
def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
|
||||
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
|
||||
return litellm_params.model.startswith("auto_router/adaptive_router")
|
||||
|
||||
def _finalize_adaptive_router_if_configured(self) -> None:
|
||||
"""Locate every adaptive-router deployment in the finalized model_list and
|
||||
build an AdaptiveRouter for each. Safe no-op when none are configured.
|
||||
Idempotent: skips any deployment whose model_name is already initialized."""
|
||||
# Drop any adaptive-router hooks left over from a previous Router
|
||||
# instance (e.g. after `/config/reload` replaced `llm_router`). Without
|
||||
# this, stale AdaptiveRouterPostCallHook callbacks from the old Router
|
||||
# remain wired up in `litellm.callbacks` and double-fire signal
|
||||
# recording for every request.
|
||||
from litellm.router_strategy.adaptive_router.hooks import (
|
||||
AdaptiveRouterPostCallHook,
|
||||
)
|
||||
|
||||
for _cb_list in (
|
||||
litellm.callbacks,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
):
|
||||
litellm.logging_callback_manager.remove_callbacks_by_type(
|
||||
_cb_list, AdaptiveRouterPostCallHook
|
||||
)
|
||||
|
||||
for entry in self.model_list or []:
|
||||
lp = (
|
||||
entry.get("litellm_params")
|
||||
if isinstance(entry, dict)
|
||||
else entry.litellm_params
|
||||
)
|
||||
lp_model = (
|
||||
(lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None
|
||||
)
|
||||
if not (lp_model and lp_model.startswith("auto_router/adaptive_router")):
|
||||
continue
|
||||
model_name = (
|
||||
entry.get("model_name") if isinstance(entry, dict) else entry.model_name
|
||||
)
|
||||
if not model_name or not lp:
|
||||
continue
|
||||
if model_name in self.adaptive_routers:
|
||||
continue
|
||||
deployment = Deployment(
|
||||
model_name=model_name,
|
||||
litellm_params=(
|
||||
lp if not isinstance(lp, dict) else LiteLLM_Params(**lp)
|
||||
),
|
||||
model_info=(
|
||||
entry.get("model_info")
|
||||
if isinstance(entry, dict)
|
||||
else entry.model_info
|
||||
),
|
||||
)
|
||||
self.init_adaptive_router_deployment(deployment=deployment)
|
||||
|
||||
def init_adaptive_router_deployment(self, deployment: Deployment) -> None:
|
||||
"""
|
||||
Build an AdaptiveRouter instance for this deployment and register its
|
||||
post-call hook. Multiple adaptive routers can coexist on a single Router,
|
||||
keyed by `deployment.model_name`.
|
||||
|
||||
`model_to_prefs` and `model_to_cost` are derived from the OTHER models
|
||||
already registered in `self.model_list` whose `model_name` appears in
|
||||
`available_models`. Models not yet registered fall back to defaults.
|
||||
"""
|
||||
# Local import: AdaptiveRouter -> hooks -> classifier all import litellm
|
||||
# internals which transitively import this module. (AGENTS.md exception clause.)
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import (
|
||||
AdaptiveRouter,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.hooks import (
|
||||
AdaptiveRouterPostCallHook,
|
||||
)
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
)
|
||||
|
||||
raw_config = deployment.litellm_params.adaptive_router_config
|
||||
if raw_config is None:
|
||||
raise ValueError(
|
||||
"adaptive_router_config is required for adaptive-router deployments."
|
||||
)
|
||||
|
||||
config = AdaptiveRouterConfig(**raw_config)
|
||||
|
||||
model_to_prefs: Dict[str, AdaptiveRouterPreferences] = {}
|
||||
model_to_cost: Dict[str, float] = {}
|
||||
# O(k) via the name→indices map: only touch deployments whose name
|
||||
# is listed in `available_models`, instead of scanning model_list.
|
||||
for name in config.available_models:
|
||||
indices = self.model_name_to_deployment_indices.get(name, [])
|
||||
if not indices:
|
||||
continue
|
||||
d = (self.model_list or [])[indices[0]]
|
||||
mi = d.get("model_info") if isinstance(d, dict) else d.model_info
|
||||
mi_dict: Dict[str, Any] = (
|
||||
mi if isinstance(mi, dict) else (mi.model_dump() if mi else {})
|
||||
)
|
||||
prefs_raw = mi_dict.get("adaptive_router_preferences")
|
||||
if prefs_raw is not None:
|
||||
model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
|
||||
|
||||
# `input_cost_per_token` is a LiteLLM_Params field per types/router.py.
|
||||
lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params
|
||||
lp_dict: Dict[str, Any] = (
|
||||
lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
|
||||
)
|
||||
cost = lp_dict.get("input_cost_per_token")
|
||||
if cost is not None:
|
||||
model_to_cost[name] = float(cost)
|
||||
|
||||
if deployment.model_name in self.adaptive_routers:
|
||||
raise ValueError(
|
||||
f"Adaptive-router deployment {deployment.model_name} already exists. "
|
||||
"Please use a different model name."
|
||||
)
|
||||
|
||||
adaptive_router = AdaptiveRouter(
|
||||
router_name=deployment.model_name,
|
||||
config=config,
|
||||
model_to_prefs=model_to_prefs,
|
||||
model_to_cost=model_to_cost,
|
||||
)
|
||||
self.adaptive_routers[deployment.model_name] = adaptive_router
|
||||
litellm.logging_callback_manager.add_litellm_callback(
|
||||
AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)
|
||||
)
|
||||
verbose_router_logger.info(
|
||||
"AdaptiveRouter[%s] initialized with %d models",
|
||||
deployment.model_name,
|
||||
len(config.available_models),
|
||||
)
|
||||
|
||||
def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
|
||||
"""
|
||||
Check if the deployment is a quality-router deployment.
|
||||
|
|
@ -7077,6 +7228,10 @@ class Router:
|
|||
# Note: model_name_to_deployment_indices is already built incrementally
|
||||
# by _create_deployment -> _add_model_to_list_and_index_map
|
||||
|
||||
# Deferred: build the AdaptiveRouter strategy now that all underlying
|
||||
# deployments have been registered.
|
||||
self._finalize_adaptive_router_if_configured()
|
||||
|
||||
def _add_deployment(self, deployment: Deployment) -> Deployment:
|
||||
import os
|
||||
|
||||
|
|
@ -7204,6 +7359,10 @@ class Router:
|
|||
):
|
||||
self.init_complexity_router_deployment(deployment=deployment)
|
||||
|
||||
# NOTE: adaptive-router deployments are deferred to the end of
|
||||
# set_model_list() because their init needs visibility into the OTHER
|
||||
# deployments listed in `available_models` (which may not yet have
|
||||
# been processed when this one is created).
|
||||
#########################################################
|
||||
# Check if this is a quality-router deployment
|
||||
#########################################################
|
||||
|
|
@ -9763,6 +9922,19 @@ class Router:
|
|||
specific_deployment=specific_deployment,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Check if an adaptive-router should be used
|
||||
#########################################################
|
||||
adaptive_router = self.adaptive_routers.get(model)
|
||||
if adaptive_router is not None:
|
||||
return await adaptive_router.async_pre_routing_hook(
|
||||
model=model,
|
||||
request_kwargs=request_kwargs,
|
||||
messages=messages,
|
||||
input=input,
|
||||
specific_deployment=specific_deployment,
|
||||
)
|
||||
|
||||
#########################################################
|
||||
# Check if any quality-router should be used
|
||||
#########################################################
|
||||
|
|
|
|||
95
litellm/router_strategy/adaptive_router/README.md
Normal file
95
litellm/router_strategy/adaptive_router/README.md
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
# Adaptive Router (v0)
|
||||
|
||||
A request-type-aware routing strategy. For each incoming request, classify the
|
||||
prompt into one of seven `RequestType` buckets (code generation, writing,
|
||||
analytical reasoning, …), then Thompson-sample a Beta(α, β) bandit posterior
|
||||
per `(request_type, model)` cell to pick the best model. Quality estimates are
|
||||
combined with a normalized cost score via a weighted linear sum.
|
||||
|
||||
A post-call hook reads the response and runs lightweight regex + tool-call
|
||||
detectors (see `signals.py`) to award per-turn credit/blame to the model that
|
||||
served the turn. Updates are batched in-memory and flushed to Postgres every
|
||||
~10s by a background task in `proxy_server.py`.
|
||||
|
||||
## Config example
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4o
|
||||
litellm_params:
|
||||
model: openai/gpt-4o
|
||||
model_info:
|
||||
input_cost_per_token: 0.0000025
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 3
|
||||
strengths: ["code_generation", "analytical_reasoning"]
|
||||
|
||||
- model_name: gpt-4o-mini
|
||||
litellm_params:
|
||||
model: openai/gpt-4o-mini
|
||||
model_info:
|
||||
input_cost_per_token: 0.00000015
|
||||
adaptive_router_preferences:
|
||||
quality_tier: 2
|
||||
strengths: ["general", "factual_lookup"]
|
||||
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/adaptive_router
|
||||
adaptive_router_default_model: gpt-4o-mini
|
||||
adaptive_router_config:
|
||||
available_models: ["gpt-4o", "gpt-4o-mini"]
|
||||
weights:
|
||||
quality: 0.7
|
||||
cost: 0.3
|
||||
```
|
||||
|
||||
Callers may pass header `x-litellm-min-quality-tier: 3` (or metadata key
|
||||
`min_quality_tier: 3`) to force selection from tier-3-or-higher models only.
|
||||
|
||||
## Behavior summary
|
||||
|
||||
- **Cold start.** Each `(request_type, model)` cell starts with a
|
||||
Beta prior whose mean = `BASE_TIER_WEIGHT[tier] (+ STRENGTH_BONUS if declared)`
|
||||
and total mass = `COLD_START_MASS` (10). About ten real observations move it
|
||||
meaningfully.
|
||||
- **Per-request decision.** Sample once per eligible model, score with
|
||||
`quality_weight·sample + cost_weight·normalized_cost`, pick the argmax.
|
||||
Routing is stateless per-turn — no sticky lookup. Each call resamples.
|
||||
- **Owner-cache attribution.** Post-call, the conversation's first picked
|
||||
model claims an "owner slot" for `OWNER_CACHE_TTL_SECONDS` (24h). Later
|
||||
turns of the same conversation only fire bandit/state updates if the
|
||||
same model handled them — mismatches are dropped (no attribution) and
|
||||
counted in `skipped_updates_total`. Conversation identity is the
|
||||
client-supplied `litellm_session_id` if present, otherwise a sha256 over
|
||||
caller identity (api key hash, team, user, end-user) + the first message.
|
||||
- **Per-turn updates.** `satisfaction → +α`. `misalignment, stagnation,
|
||||
disengagement, failure → +β` (each). `loop → +0.5β`. `exhaustion → 0`
|
||||
(uptime, not quality). Skipped if conversation has fewer than
|
||||
`SIGNAL_GATE_MIN_MESSAGES` messages.
|
||||
- **Persistence.** Bandit cells: aggregated deltas, eventually consistent.
|
||||
Session rows: last-write-wins snapshots.
|
||||
|
||||
## Known v0 limitations
|
||||
|
||||
- **Latency is not in the score.** Quality + cost only. A pathologically slow
|
||||
model can still be picked.
|
||||
- **Hard sample cap at 200.** Once `α + β > 200`, deltas are silently dropped.
|
||||
No rescaling — drift is a v1 concern.
|
||||
- **24h owner-cache TTL.** No explicit eviction below TTL. The in-memory map
|
||||
can grow if traffic patterns produce many one-shot sessions.
|
||||
- **Owner-recovery skew.** If model A "owns" a conversation but is then
|
||||
dethroned in the bandit, later turns served by model B are dropped — so
|
||||
bandit updates for that conversation flatline until A's TTL expires.
|
||||
Tracked via `skipped_updates_total`.
|
||||
- **Signals are regex + tool-call only.** No LLM-judge, no embedding similarity,
|
||||
no exemplar storage. Signals are best-effort and biased toward English.
|
||||
- **One AdaptiveRouter per `Router`.** Multiple `adaptive_router/*` deployments
|
||||
on the same `litellm.Router` raise at init.
|
||||
- **Bandit-delta mapping is unvalidated.** `_compute_bandit_delta` is a v0
|
||||
guess; expect to retune after the first ~1000 sessions of real traffic.
|
||||
- **`request_type` is classified per turn from the latest user message.** For
|
||||
non-GENERAL turns, the current-turn type is used for bandit attribution (so
|
||||
genuine mid-session topic shifts update the correct cell). For GENERAL turns
|
||||
("thanks!", "ok", "sounds good"), attribution falls back to the session's
|
||||
original type to avoid misattributing closing pleasantries.
|
||||
6
litellm/router_strategy/adaptive_router/__init__.py
Normal file
6
litellm/router_strategy/adaptive_router/__init__.py
Normal file
|
|
@ -0,0 +1,6 @@
|
|||
"""Adaptive router strategy. See README.md for design overview."""
|
||||
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||||
|
||||
__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"]
|
||||
454
litellm/router_strategy/adaptive_router/adaptive_router.py
Normal file
454
litellm/router_strategy/adaptive_router/adaptive_router.py
Normal file
|
|
@ -0,0 +1,454 @@
|
|||
"""
|
||||
Main adaptive router strategy. See README.md for design overview.
|
||||
|
||||
One AdaptiveRouter instance per router_name. Holds in-memory caches:
|
||||
- _cells: Beta(alpha, beta) bandit posteriors per (request_type, model)
|
||||
- _owner_cache: session_key -> (owner_model, expires_at) — the first model
|
||||
picked for a conversation owns its bandit-update slot
|
||||
- _session_states: (session_key, model) -> SessionState for incremental signal updates
|
||||
|
||||
Owns the AdaptiveRouterUpdateQueue used by the proxy's flusher to persist
|
||||
state and session snapshots back to Postgres.
|
||||
|
||||
Routing is stateless per-turn (Thompson sample fresh on every call). The
|
||||
owner cache is consulted only at post-call time to decide whether a turn's
|
||||
signals should fire a bandit update — turns served by a different model than
|
||||
the conversation's owner are skipped to avoid cross-model misattribution.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import asdict
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
get_last_user_message,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.bandit import (
|
||||
BanditCell,
|
||||
apply_delta,
|
||||
initial_cell,
|
||||
pick_best,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
MIN_QUALITY_TIER_HEADER,
|
||||
MIN_QUALITY_TIER_METADATA_KEY,
|
||||
OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.signals import (
|
||||
SessionState,
|
||||
SignalDelta,
|
||||
Turn,
|
||||
apply_turn,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.update_queue import (
|
||||
AdaptiveRouterUpdateQueue,
|
||||
)
|
||||
|
||||
# Sweep session-state cache when it exceeds this many live entries. Expired
|
||||
# entries are dropped in bulk; amortizes to O(1) per insert.
|
||||
_SESSION_STATE_SWEEP_THRESHOLD: int = 1024
|
||||
# Same pattern for the owner cache.
|
||||
_OWNER_CACHE_SWEEP_THRESHOLD: int = 1024
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
PreRoutingHookResponse,
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def _default_prefs() -> AdaptiveRouterPreferences:
|
||||
"""Tier-2 prior with no declared strengths; used when a model omits prefs."""
|
||||
return AdaptiveRouterPreferences(quality_tier=2, strengths=[])
|
||||
|
||||
|
||||
class AdaptiveRouter:
|
||||
"""One instance per router_name. Holds in-memory caches + the update queue."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
router_name: str,
|
||||
config: AdaptiveRouterConfig,
|
||||
model_to_prefs: Dict[str, AdaptiveRouterPreferences],
|
||||
model_to_cost: Dict[str, float],
|
||||
) -> None:
|
||||
self.router_name = router_name
|
||||
self.config = config
|
||||
self.model_to_prefs = model_to_prefs
|
||||
self.model_to_cost = model_to_cost
|
||||
self.queue = AdaptiveRouterUpdateQueue()
|
||||
|
||||
self._cells: Dict[Tuple[RequestType, str], BanditCell] = {}
|
||||
self._owner_cache: Dict[str, Tuple[str, float]] = {}
|
||||
self._session_states: Dict[Tuple[str, str], SessionState] = {}
|
||||
# Parallel expiry map for _session_states, same TTL as _owner_cache.
|
||||
# Evicted opportunistically in `get_or_create_session_state`.
|
||||
self._session_states_expiry: Dict[Tuple[str, str], float] = {}
|
||||
self._skipped_updates_total: int = 0
|
||||
# Set to True once the proxy flusher has loaded persisted priors from
|
||||
# Postgres. Checked to support lazy-load on hot-reloaded routers.
|
||||
self._state_loaded: bool = False
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
self._init_cold_start_cells()
|
||||
|
||||
# ---- Cold-start ------------------------------------------------------
|
||||
|
||||
def _init_cold_start_cells(self) -> None:
|
||||
"""Populate _cells with cold-start priors for every (rt, model) combination."""
|
||||
for rt in RequestType:
|
||||
for model in self.config.available_models:
|
||||
prefs = self.model_to_prefs.get(model) or _default_prefs()
|
||||
self._cells[(rt, model)] = initial_cell(prefs, rt)
|
||||
|
||||
async def load_state_from_db(self, prisma_client: Any) -> None:
|
||||
"""Override cold-start cells with persisted state. Called once at startup."""
|
||||
if prisma_client is None:
|
||||
return
|
||||
try:
|
||||
rows = await prisma_client.db.litellm_adaptiverouterstate.find_many(
|
||||
where={"router_name": self.router_name}
|
||||
)
|
||||
loaded = 0
|
||||
for row in rows:
|
||||
try:
|
||||
rt = RequestType(row.request_type)
|
||||
except ValueError:
|
||||
# Unknown taxonomy entry from an older/newer version. Skip.
|
||||
continue
|
||||
if row.model_name not in self.config.available_models:
|
||||
continue
|
||||
self._cells[(rt, row.model_name)] = BanditCell(
|
||||
alpha=row.alpha, beta=row.beta
|
||||
)
|
||||
loaded += 1
|
||||
verbose_router_logger.info(
|
||||
"AdaptiveRouter[%s]: loaded %d cells from DB",
|
||||
self.router_name,
|
||||
loaded,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
"AdaptiveRouter[%s]: failed to load state from DB: %s",
|
||||
self.router_name,
|
||||
e,
|
||||
)
|
||||
|
||||
# ---- Pre-routing hook ------------------------------------------------
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Dict[str, Any],
|
||||
messages: Optional[List[Dict[str, Any]]] = None,
|
||||
input: Optional[Union[str, List]] = None,
|
||||
specific_deployment: Optional[bool] = False,
|
||||
) -> Optional[PreRoutingHookResponse]:
|
||||
"""
|
||||
Plugin entry point invoked by `Router.async_pre_routing_hook` when the
|
||||
inbound `model` matches this adaptive router's `router_name`.
|
||||
|
||||
Classifies the last user message, picks a logical model via the bandit,
|
||||
and stashes the chosen model on `request_kwargs["metadata"]` so the
|
||||
post-call hook can surface it as a response header.
|
||||
|
||||
Routing is stateless per-turn: every call Thompson-samples fresh,
|
||||
regardless of any prior pick for the same session. Cross-turn
|
||||
attribution is enforced post-call via the owner cache (see
|
||||
`claim_or_check_owner`).
|
||||
"""
|
||||
user_text = (
|
||||
get_last_user_message(cast(List[AllMessageValues], messages or [])) or ""
|
||||
)
|
||||
|
||||
request_type = classify_prompt(user_text)
|
||||
min_quality_tier = self._extract_min_quality_tier(request_kwargs)
|
||||
chosen_model = await self.pick_model(
|
||||
request_type=request_type, min_quality_tier=min_quality_tier
|
||||
)
|
||||
verbose_router_logger.debug(
|
||||
"AdaptiveRouter[%s]: classified=%s -> chose %s",
|
||||
self.router_name,
|
||||
request_type.value,
|
||||
chosen_model,
|
||||
)
|
||||
|
||||
# Relay the chosen logical model to the post-call hook, which surfaces
|
||||
# it as the `x-litellm-adaptive-router-model` response header. We use
|
||||
# `metadata` (not a top-level kwarg) so the value doesn't leak into
|
||||
# `litellm.acompletion(**input_kwargs)`.
|
||||
kwargs_metadata = request_kwargs.setdefault("metadata", {})
|
||||
if isinstance(kwargs_metadata, dict):
|
||||
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = chosen_model
|
||||
|
||||
return PreRoutingHookResponse(model=chosen_model, messages=messages)
|
||||
|
||||
# ---- Pick model ------------------------------------------------------
|
||||
|
||||
async def pick_model(
|
||||
self,
|
||||
request_type: RequestType,
|
||||
min_quality_tier: Optional[int] = None,
|
||||
) -> str:
|
||||
"""Thompson-sample across eligible models. Stateless per-turn."""
|
||||
eligible = self._eligible_models(min_quality_tier)
|
||||
if not eligible:
|
||||
raise ValueError(
|
||||
f"AdaptiveRouter[{self.router_name}]: no models meet "
|
||||
f"min_quality_tier={min_quality_tier}"
|
||||
)
|
||||
|
||||
cells = {m: self._cells[(request_type, m)] for m in eligible}
|
||||
costs = {m: self.model_to_cost.get(m, 0.0) for m in eligible}
|
||||
return pick_best(
|
||||
cells,
|
||||
costs,
|
||||
quality_weight=self.config.weights.quality,
|
||||
cost_weight=self.config.weights.cost,
|
||||
)
|
||||
|
||||
def claim_or_check_owner(self, session_key: str, current_model: str) -> bool:
|
||||
"""Resolve attribution for a turn under stateless routing.
|
||||
|
||||
Returns True iff this turn should fire a bandit/state update. The
|
||||
first call for a `session_key` claims ownership for `current_model`
|
||||
and returns True. Subsequent calls return True only if the owner is
|
||||
still live AND matches `current_model`. Mismatches (a different
|
||||
model handled this turn) and expired owners both increment
|
||||
`_skipped_updates_total` and return False — no attribution.
|
||||
"""
|
||||
now = time.time()
|
||||
existing = self._owner_cache.get(session_key)
|
||||
if existing is not None and existing[1] > now:
|
||||
owner_model, _ = existing
|
||||
if owner_model == current_model:
|
||||
return True
|
||||
self._skipped_updates_total += 1
|
||||
return False
|
||||
|
||||
# Opportunistic bulk sweep — sessions that never come back would
|
||||
# otherwise pile up here forever. Same threshold pattern as the
|
||||
# session-state cache.
|
||||
if len(self._owner_cache) >= _OWNER_CACHE_SWEEP_THRESHOLD:
|
||||
self._evict_expired_owner_cache(now)
|
||||
|
||||
# No live owner -> claim for current_model.
|
||||
self._owner_cache[session_key] = (
|
||||
current_model,
|
||||
now + OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
return True
|
||||
|
||||
def _evict_expired_owner_cache(self, now: float) -> None:
|
||||
expired = [k for k, (_, exp) in self._owner_cache.items() if exp <= now]
|
||||
for k in expired:
|
||||
self._owner_cache.pop(k, None)
|
||||
|
||||
async def get_state_snapshot(self) -> Dict[str, Any]:
|
||||
"""In-memory snapshot for the introspection endpoint. Cheap; no DB hit."""
|
||||
cells = []
|
||||
for (rt, model), cell in sorted(
|
||||
self._cells.items(), key=lambda kv: (kv[0][0].value, kv[0][1])
|
||||
):
|
||||
total = cell.alpha + cell.beta
|
||||
cells.append(
|
||||
{
|
||||
"request_type": rt.value,
|
||||
"model": model,
|
||||
"alpha": cell.alpha,
|
||||
"beta": cell.beta,
|
||||
# Net observations that have moved the posterior, excluding
|
||||
# the cold-start prior mass. `alpha + beta` would show the
|
||||
# initial COLD_START_MASS (e.g. 10) before any real traffic
|
||||
# arrives, which confuses operators reading the endpoint.
|
||||
"samples": cell.total_samples,
|
||||
"quality_mean": cell.alpha / total if total > 0 else 0.0,
|
||||
}
|
||||
)
|
||||
queue = await self.queue.queue_size()
|
||||
now = time.time()
|
||||
owner_cache_live = sum(1 for _, exp in self._owner_cache.values() if exp > now)
|
||||
return {
|
||||
"router_name": self.router_name,
|
||||
"available_models": list(self.config.available_models),
|
||||
"weights": {
|
||||
"quality": self.config.weights.quality,
|
||||
"cost": self.config.weights.cost,
|
||||
},
|
||||
"model_costs": dict(self.model_to_cost),
|
||||
"cells": cells,
|
||||
"owner_cache_live": owner_cache_live,
|
||||
"skipped_updates_total": self._skipped_updates_total,
|
||||
"queue": queue,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_min_quality_tier(
|
||||
request_kwargs: Dict[str, Any],
|
||||
) -> Optional[int]:
|
||||
"""Pull `min_quality_tier` from request headers or metadata.
|
||||
|
||||
Precedence: headers (`x-litellm-min-quality-tier`) over metadata
|
||||
(`min_quality_tier`). Headers arrive lowercased from the proxy but we
|
||||
lookup case-insensitively to be safe. Unparseable values are ignored
|
||||
(treated as "not set") rather than raising — a bad header shouldn't
|
||||
fail the request.
|
||||
"""
|
||||
headers = request_kwargs.get("headers") or {}
|
||||
if isinstance(headers, dict):
|
||||
for k, v in headers.items():
|
||||
if isinstance(k, str) and k.lower() == MIN_QUALITY_TIER_HEADER:
|
||||
try:
|
||||
return int(v)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
|
||||
metadata = request_kwargs.get("metadata") or {}
|
||||
if isinstance(metadata, dict):
|
||||
raw = metadata.get(MIN_QUALITY_TIER_METADATA_KEY)
|
||||
if raw is not None:
|
||||
try:
|
||||
return int(raw)
|
||||
except (TypeError, ValueError):
|
||||
return None
|
||||
return None
|
||||
|
||||
def _eligible_models(self, min_quality_tier: Optional[int]) -> List[str]:
|
||||
if min_quality_tier is None:
|
||||
return list(self.config.available_models)
|
||||
return [
|
||||
m
|
||||
for m in self.config.available_models
|
||||
if (self.model_to_prefs.get(m) or _default_prefs()).quality_tier
|
||||
>= min_quality_tier
|
||||
]
|
||||
|
||||
# ---- Session state ---------------------------------------------------
|
||||
|
||||
def get_or_create_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
model_name: str,
|
||||
request_type: RequestType,
|
||||
) -> SessionState:
|
||||
key = (session_id, model_name)
|
||||
now = time.time()
|
||||
|
||||
# Opportunistic bulk sweep when the cache grows past the threshold.
|
||||
# Cheap relative to the alternative of a bounded LRU — conversations
|
||||
# naturally become inactive within OWNER_CACHE_TTL_SECONDS.
|
||||
if len(self._session_states) >= _SESSION_STATE_SWEEP_THRESHOLD:
|
||||
self._evict_expired_session_states(now)
|
||||
|
||||
state = self._session_states.get(key)
|
||||
if state is None:
|
||||
state = SessionState(
|
||||
session_id=session_id,
|
||||
router_name=self.router_name,
|
||||
model_name=model_name,
|
||||
classified_type=request_type.value,
|
||||
)
|
||||
self._session_states[key] = state
|
||||
self._session_states_expiry[key] = now + OWNER_CACHE_TTL_SECONDS
|
||||
return state
|
||||
|
||||
def _evict_expired_session_states(self, now: float) -> None:
|
||||
"""Drop session states whose TTL has passed. O(n) but amortized O(1)
|
||||
per insert thanks to `_SESSION_STATE_SWEEP_THRESHOLD`."""
|
||||
expired = [k for k, exp in self._session_states_expiry.items() if exp <= now]
|
||||
for k in expired:
|
||||
self._session_states.pop(k, None)
|
||||
self._session_states_expiry.pop(k, None)
|
||||
|
||||
async def record_turn(
|
||||
self,
|
||||
session_id: str,
|
||||
model_name: str,
|
||||
request_type: RequestType,
|
||||
turn: Turn,
|
||||
) -> SignalDelta:
|
||||
"""Apply one turn, push session snapshot + bandit deltas to the queue."""
|
||||
state = self.get_or_create_session_state(session_id, model_name, request_type)
|
||||
delta = apply_turn(state, turn)
|
||||
verbose_router_logger.debug(
|
||||
"AdaptiveRouter[%s]: record_turn delta=%s", self.router_name, delta
|
||||
)
|
||||
|
||||
# Strip the raw conversation content before persisting. The
|
||||
# last_user/assistant_content and tool_call_history fields are only
|
||||
# needed in-memory for the next turn's incremental signal detection;
|
||||
# writing user prompts and tool payloads to the DB would store PII
|
||||
# for every adaptive-router conversation. Counts + bookkeeping is
|
||||
# all the persisted row needs.
|
||||
snapshot = asdict(state)
|
||||
for sensitive in (
|
||||
"last_user_content",
|
||||
"last_assistant_content",
|
||||
"tool_call_history",
|
||||
"pending_tool_calls",
|
||||
):
|
||||
snapshot.pop(sensitive, None)
|
||||
await self.queue.add_session_state(
|
||||
session_id, self.router_name, model_name, snapshot
|
||||
)
|
||||
|
||||
d_alpha, d_beta = self._compute_bandit_delta(delta)
|
||||
verbose_router_logger.debug(
|
||||
"AdaptiveRouter[%s]: bandit delta alpha=%.2f beta=%.2f",
|
||||
self.router_name,
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
if d_alpha != 0 or d_beta != 0:
|
||||
# For non-GENERAL turns, attribute to the current-turn classification
|
||||
# so genuine mid-session topic shifts (e.g. code → math) update the
|
||||
# correct cell. For GENERAL turns ("thanks!", "ok", "sounds good"), fall
|
||||
# back to the session's original type so closing pleasantries don't
|
||||
# misattribute the reward.
|
||||
attribution_type = (
|
||||
request_type
|
||||
if request_type != RequestType.GENERAL
|
||||
else RequestType(state.classified_type)
|
||||
)
|
||||
cell_key = (attribution_type, model_name)
|
||||
self._cells[cell_key] = apply_delta(self._cells[cell_key], d_alpha, d_beta)
|
||||
await self.queue.add_state_delta(
|
||||
self.router_name,
|
||||
attribution_type.value,
|
||||
model_name,
|
||||
d_alpha,
|
||||
d_beta,
|
||||
)
|
||||
|
||||
return delta
|
||||
|
||||
@staticmethod
|
||||
def _compute_bandit_delta(delta: SignalDelta) -> Tuple[float, float]:
|
||||
"""
|
||||
Translate per-turn signal deltas into bandit-cell deltas.
|
||||
|
||||
v0 mapping (UNVALIDATED — D6):
|
||||
- satisfaction -> +1 alpha
|
||||
- misalignment, stagnation,
|
||||
disengagement, failure -> +1 beta each
|
||||
- loop -> +0.5 beta (weak; could be model OR user)
|
||||
- exhaustion -> 0 (uptime issue, tracked separately later)
|
||||
"""
|
||||
d_alpha = float(delta.satisfaction)
|
||||
d_beta = (
|
||||
float(
|
||||
delta.misalignment
|
||||
+ delta.stagnation
|
||||
+ delta.disengagement
|
||||
+ delta.failure
|
||||
)
|
||||
+ 0.5 * delta.loop
|
||||
)
|
||||
return d_alpha, d_beta
|
||||
142
litellm/router_strategy/adaptive_router/bandit.py
Normal file
142
litellm/router_strategy/adaptive_router/bandit.py
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
"""
|
||||
Thompson sampling and prior initialization for the adaptive router bandit.
|
||||
|
||||
Each (router, request_type, model) cell is a Beta(alpha, beta) posterior.
|
||||
- alpha = pseudo-successes
|
||||
- beta = pseudo-failures
|
||||
- mean = alpha / (alpha + beta)
|
||||
- total samples = alpha + beta - COLD_START_MASS (informative prior, not data)
|
||||
|
||||
Hot path: thompson_sample() — pure function, no I/O.
|
||||
"""
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
BASE_TIER_WEIGHT,
|
||||
COLD_START_MASS,
|
||||
DEFAULT_COST_WEIGHT,
|
||||
DEFAULT_QUALITY_WEIGHT,
|
||||
SAMPLE_CAP,
|
||||
STRENGTH_BONUS,
|
||||
)
|
||||
from litellm.types.router import AdaptiveRouterPreferences, RequestType
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BanditCell:
|
||||
"""Posterior state for a single (router, request_type, model) cell."""
|
||||
|
||||
alpha: float
|
||||
beta: float
|
||||
|
||||
@property
|
||||
def mean(self) -> float:
|
||||
total = self.alpha + self.beta
|
||||
return self.alpha / total if total > 0 else 0.5
|
||||
|
||||
@property
|
||||
def total_samples(self) -> int:
|
||||
return max(0, int(self.alpha + self.beta - COLD_START_MASS))
|
||||
|
||||
|
||||
def initial_cell(
|
||||
prefs: AdaptiveRouterPreferences, request_type: RequestType
|
||||
) -> BanditCell:
|
||||
"""
|
||||
Cold-start prior for a (model, request_type) cell.
|
||||
|
||||
mean = base_tier_weight[tier] + (STRENGTH_BONUS if request_type in strengths else 0)
|
||||
capped at 0.95 to avoid an over-confident prior.
|
||||
Total mass = COLD_START_MASS so that ~10 real observations can move it noticeably.
|
||||
"""
|
||||
if prefs.quality_tier not in BASE_TIER_WEIGHT:
|
||||
valid = sorted(BASE_TIER_WEIGHT)
|
||||
raise ValueError(
|
||||
f"quality_tier={prefs.quality_tier} is not supported; "
|
||||
f"valid tiers are {valid}"
|
||||
)
|
||||
base = BASE_TIER_WEIGHT[prefs.quality_tier]
|
||||
bonus = STRENGTH_BONUS if request_type in prefs.strengths else 0.0
|
||||
mean = min(0.95, base + bonus)
|
||||
alpha = mean * COLD_START_MASS
|
||||
beta = (1.0 - mean) * COLD_START_MASS
|
||||
return BanditCell(alpha=alpha, beta=beta)
|
||||
|
||||
|
||||
def apply_delta(cell: BanditCell, delta_alpha: float, delta_beta: float) -> BanditCell:
|
||||
"""
|
||||
Apply a learning update to a cell, enforcing the sample cap.
|
||||
|
||||
SAMPLE_CAP is a HARD cap on (alpha + beta). When the cap would be exceeded,
|
||||
we drop the update. (D5: hard cap, no rescaling — keep v0 simple.)
|
||||
"""
|
||||
new_alpha = cell.alpha + delta_alpha
|
||||
new_beta = cell.beta + delta_beta
|
||||
if new_alpha + new_beta > SAMPLE_CAP:
|
||||
return cell
|
||||
return BanditCell(alpha=new_alpha, beta=new_beta)
|
||||
|
||||
|
||||
def thompson_sample(cell: BanditCell, rng: Optional[random.Random] = None) -> float:
|
||||
"""Draw a sample from Beta(alpha, beta). Returns a quality estimate in [0, 1]."""
|
||||
r = rng if rng is not None else random
|
||||
return r.betavariate(cell.alpha, cell.beta)
|
||||
|
||||
|
||||
def normalized_cost(model_cost: float, all_costs: List[float]) -> float:
|
||||
"""
|
||||
Map a raw $/1k-token cost into [0, 1] where 0 = most expensive, 1 = cheapest.
|
||||
Returns 0.5 when there's no spread.
|
||||
"""
|
||||
if not all_costs:
|
||||
return 0.5
|
||||
lo, hi = min(all_costs), max(all_costs)
|
||||
if hi == lo:
|
||||
return 0.5
|
||||
return 1.0 - ((model_cost - lo) / (hi - lo))
|
||||
|
||||
|
||||
def score(
|
||||
quality_sample: float,
|
||||
model_cost: float,
|
||||
all_costs: List[float],
|
||||
quality_weight: float = DEFAULT_QUALITY_WEIGHT,
|
||||
cost_weight: float = DEFAULT_COST_WEIGHT,
|
||||
) -> float:
|
||||
"""
|
||||
Multi-objective score. V0 is a weighted linear sum of (quality, normalized_cost).
|
||||
Higher is better. Both inputs are in [0, 1].
|
||||
"""
|
||||
cost_score = normalized_cost(model_cost, all_costs)
|
||||
return quality_weight * quality_sample + cost_weight * cost_score
|
||||
|
||||
|
||||
def pick_best(
|
||||
cells: Dict[str, BanditCell],
|
||||
model_costs: Dict[str, float],
|
||||
quality_weight: float = DEFAULT_QUALITY_WEIGHT,
|
||||
cost_weight: float = DEFAULT_COST_WEIGHT,
|
||||
rng: Optional[random.Random] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Sample once per model, score each, return the model with highest score.
|
||||
|
||||
cells: {model_name: BanditCell}
|
||||
model_costs: {model_name: $/1k tokens}
|
||||
"""
|
||||
if not cells:
|
||||
raise ValueError("pick_best called with no models")
|
||||
all_costs = list(model_costs.values())
|
||||
best_model: Optional[str] = None
|
||||
best_score = float("-inf")
|
||||
for model, cell in cells.items():
|
||||
q = thompson_sample(cell, rng=rng)
|
||||
s = score(q, model_costs[model], all_costs, quality_weight, cost_weight)
|
||||
if s > best_score:
|
||||
best_score = s
|
||||
best_model = model
|
||||
assert best_model is not None
|
||||
return best_model
|
||||
140
litellm/router_strategy/adaptive_router/classifier.py
Normal file
140
litellm/router_strategy/adaptive_router/classifier.py
Normal file
|
|
@ -0,0 +1,140 @@
|
|||
"""
|
||||
Rule-based classifier mapping a user prompt to a RequestType.
|
||||
|
||||
V0 design choice: deterministic regex over the FIRST user message in a session.
|
||||
Result is cached per session (caller's responsibility, not ours).
|
||||
|
||||
Order matters: we check more specific types first, falling back to GENERAL.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import List, Pattern, Tuple
|
||||
|
||||
from litellm.types.router import RequestType
|
||||
|
||||
_RULES: List[Tuple[Pattern[str], RequestType]] = [
|
||||
(
|
||||
re.compile(
|
||||
r"\b(write|create|generate|implement|build)\s+(?:a |an |the |me )?(?:python|javascript|typescript|java|rust|go|c\+\+|sql|bash|shell)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.CODE_GENERATION,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(write|create|implement|build)\b(?:\s+\w+){0,4}?\s+(function|class|method|script|program|api|endpoint|microservice)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.CODE_GENERATION,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(explain|describe|understand|walk me through|what does)\b.*\b(code|function|method|class|algorithm|snippet)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.CODE_UNDERSTANDING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(debug|fix|why (?:is|does|isn't)|what.s wrong|trace)\b.*\b(error|bug|exception|stacktrace|stack trace|traceback)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.CODE_UNDERSTANDING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(review|critique)\s+(?:this |my |the )?(?:code|pr|pull request|diff|patch)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.CODE_UNDERSTANDING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(design|architect|plan|architecture)\b.*\b(system|service|api|database|schema|module|microservice)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.TECHNICAL_DESIGN,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(should i (?:use|choose|pick)|tradeoffs? between|compare)\b.*\b(library|framework|language|database|protocol|postgres|postgresql|mongodb|dynamodb|mysql|redis|kafka|sql|nosql)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.TECHNICAL_DESIGN,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\bhow (?:should|do) i (?:design|structure|organize|model)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.TECHNICAL_DESIGN,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(solve|compute|calculate|prove|derive)\b.*\b(equation|integral|derivative|theorem|proof|problem)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.ANALYTICAL_REASONING,
|
||||
),
|
||||
(
|
||||
re.compile(r"\b(if .+ then|given .+ find|suppose|assume)\b", re.IGNORECASE),
|
||||
RequestType.ANALYTICAL_REASONING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(probability|statistics|combinatorics|optimization problem)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.ANALYTICAL_REASONING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(write|draft|compose|rewrite|edit|proofread|polish)\b.*\b(email|essay|blog|post|article|letter|memo|copy|paragraph|sentence)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.WRITING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"\b(make (?:this|it)|help me)\s+(?:more |less )?(?:concise|formal|casual|professional|persuasive)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.WRITING,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"^\s*(who|what|when|where|which)\s+(?:is|was|were|are)\b", re.IGNORECASE
|
||||
),
|
||||
RequestType.FACTUAL_LOOKUP,
|
||||
),
|
||||
(
|
||||
re.compile(r"^\s*(define|definition of|meaning of)\b", re.IGNORECASE),
|
||||
RequestType.FACTUAL_LOOKUP,
|
||||
),
|
||||
(
|
||||
re.compile(
|
||||
r"^\s*how (?:do you spell|to spell|many .* are there|tall is)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
RequestType.FACTUAL_LOOKUP,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def classify_prompt(text: str) -> RequestType:
|
||||
"""
|
||||
Classify a single user prompt.
|
||||
|
||||
Falls back to GENERAL when no rule matches. Empty/whitespace-only also
|
||||
returns GENERAL.
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
return RequestType.GENERAL
|
||||
|
||||
truncated = text[:2000]
|
||||
|
||||
for pattern, request_type in _RULES:
|
||||
if pattern.search(truncated):
|
||||
return request_type
|
||||
|
||||
return RequestType.GENERAL
|
||||
54
litellm/router_strategy/adaptive_router/config.py
Normal file
54
litellm/router_strategy/adaptive_router/config.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
"""
|
||||
Configuration constants for the adaptive_router strategy.
|
||||
|
||||
All magic numbers are first-pass guesses (D3-D6 in the handoff plan).
|
||||
Expect to retune after first 1000 sessions of real traffic.
|
||||
"""
|
||||
|
||||
from typing import Dict
|
||||
|
||||
from litellm.types.router import RequestType # re-export for convenience # noqa: F401
|
||||
|
||||
# D3 — Score weights (default; user-overridable via AdaptiveRouterConfig.weights)
|
||||
DEFAULT_QUALITY_WEIGHT: float = 0.7 # UNVALIDATED — calibrated against [0] sessions
|
||||
DEFAULT_COST_WEIGHT: float = 0.3 # UNVALIDATED — calibrated against [0] sessions
|
||||
|
||||
# D4 — Cold-start prior: (alpha + beta) total mass = COLD_START_MASS
|
||||
# Mean of Beta = base_tier_weight + (strength_bonus if declared)
|
||||
BASE_TIER_WEIGHT: Dict[int, float] = {1: 0.3, 2: 0.5, 3: 0.7} # UNVALIDATED
|
||||
STRENGTH_BONUS: float = 0.3 # UNVALIDATED
|
||||
COLD_START_MASS: float = 10.0
|
||||
|
||||
# D5 — Sample cap. Hard cap, no rescaling (drift handling is v1).
|
||||
SAMPLE_CAP: int = 200
|
||||
|
||||
# D6 — Clean-trace credit: minimum turns before α += 1 can fire.
|
||||
MIN_TURNS_FOR_CLEAN_CREDIT: int = 3
|
||||
|
||||
# D2 — Owner-cache TTL (seconds). 24h.
|
||||
# A conversation's first-picked model "owns" the bandit-update slot for
|
||||
# this long. Subsequent turns of the same conversation only contribute a
|
||||
# bandit/state update when the same model is re-sampled.
|
||||
OWNER_CACHE_TTL_SECONDS: int = 24 * 3600
|
||||
|
||||
# Below this many messages we skip post-call signal recording. Most signals
|
||||
# (misalignment, stagnation, satisfaction-in-response-to-prior-turn) need at
|
||||
# least one full prior exchange to be meaningful.
|
||||
SIGNAL_GATE_MIN_MESSAGES: int = 4
|
||||
|
||||
# Detector thresholds (from Plano/Chen 2026 paper).
|
||||
MISALIGNMENT_JACCARD_THRESHOLD: float = 0.45
|
||||
STAGNATION_JACCARD_NEAR_DUP: float = 0.50
|
||||
LOOP_REPEAT_THRESHOLD: int = 3
|
||||
TOOL_CALL_HISTORY_MAX: int = 20
|
||||
|
||||
# D1 — Caller filter for min quality tier.
|
||||
MIN_QUALITY_TIER_HEADER: str = "x-litellm-min-quality-tier"
|
||||
MIN_QUALITY_TIER_METADATA_KEY: str = "min_quality_tier"
|
||||
|
||||
# Pre-routing -> post-call relay: the chosen logical model is stashed on
|
||||
# request_kwargs["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] by the
|
||||
# pre-routing hook, then read by the post-call hook to surface as the
|
||||
# ADAPTIVE_ROUTER_RESPONSE_HEADER response header.
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY: str = "adaptive_router_chosen_model"
|
||||
ADAPTIVE_ROUTER_RESPONSE_HEADER: str = "x-litellm-adaptive-router-model"
|
||||
278
litellm/router_strategy/adaptive_router/hooks.py
Normal file
278
litellm/router_strategy/adaptive_router/hooks.py
Normal file
|
|
@ -0,0 +1,278 @@
|
|||
"""
|
||||
Post-call hook for the adaptive router.
|
||||
|
||||
On each successful or failed completion, build a Turn from the request/response
|
||||
and push it through `AdaptiveRouter.record_turn`. The router then updates the
|
||||
in-memory bandit cell + session state and queues writes for the proxy flusher.
|
||||
|
||||
All work happens after the response has been returned to the caller. Any
|
||||
exception is swallowed — signal recording must never break a request.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
ADAPTIVE_ROUTER_RESPONSE_HEADER,
|
||||
SIGNAL_GATE_MIN_MESSAGES,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.signals import Turn
|
||||
|
||||
# Identity fields hashed into a derived session key so the same conversation
|
||||
# from the same caller produces a stable key, while different keys/teams/users
|
||||
# stay segregated even if they happen to send identical first messages.
|
||||
_IDENTITY_FIELDS = (
|
||||
"user_api_key_hash",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_end_user_id",
|
||||
)
|
||||
|
||||
|
||||
def _resolve_session_key(kwargs: Dict[str, Any]) -> Optional[str]:
|
||||
"""Pick a stable per-conversation key for owner-cache attribution.
|
||||
|
||||
Order:
|
||||
1. Honor a client-supplied session id (`litellm_session_id` on either
|
||||
`litellm_params` or `litellm_params.metadata`, or `session_id` on
|
||||
metadata) — backward compat for callers already wired up.
|
||||
2. Otherwise derive a sha256 over (identity fields, first
|
||||
SIGNAL_GATE_MIN_MESSAGES messages) so the key is stable across turns
|
||||
and only materialises once there is enough context for the bandit to
|
||||
act on (matching the gate in the signal-processing path).
|
||||
|
||||
Returns None if the conversation is shorter than SIGNAL_GATE_MIN_MESSAGES.
|
||||
"""
|
||||
litellm_params = kwargs.get("litellm_params") or {}
|
||||
sid = litellm_params.get("litellm_session_id")
|
||||
if sid:
|
||||
return str(sid)
|
||||
metadata = litellm_params.get("metadata") or {}
|
||||
if isinstance(metadata, dict):
|
||||
sid = metadata.get("session_id") or metadata.get("litellm_session_id")
|
||||
if sid:
|
||||
return str(sid)
|
||||
|
||||
messages = kwargs.get("messages") or []
|
||||
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
|
||||
# Don't attribute until we have enough turns to match the signal gate —
|
||||
# ensures the hash is stable (same N messages every time) and avoids
|
||||
# crediting the bandit for conversations that are too short to signal.
|
||||
return None
|
||||
|
||||
identity = ":".join(
|
||||
str(metadata.get(f) or "") if isinstance(metadata, dict) else ""
|
||||
for f in _IDENTITY_FIELDS
|
||||
)
|
||||
anchor = messages[:SIGNAL_GATE_MIN_MESSAGES]
|
||||
payload = (
|
||||
identity
|
||||
+ "|"
|
||||
+ json.dumps(
|
||||
[{"role": m.get("role"), "content": m.get("content")} for m in anchor],
|
||||
sort_keys=True,
|
||||
default=str,
|
||||
)
|
||||
)
|
||||
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _last_user_content(messages: Optional[List[Dict[str, Any]]]) -> Optional[str]:
|
||||
if not messages:
|
||||
return None
|
||||
for msg in reversed(messages):
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content")
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
# OpenAI vision-style content: pick first text part.
|
||||
for part in content:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
return part.get("text")
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _recent_tool_results(
|
||||
messages: Optional[List[Dict[str, Any]]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Extract the current turn's tool result payloads from the request messages.
|
||||
|
||||
Tool results are `role == "tool"` messages that sit at the tail of the
|
||||
conversation — i.e. after the most recent assistant message with
|
||||
`tool_calls`, waiting for the model to produce a user-facing reply. Walk
|
||||
backwards from the end and collect the contiguous run of tool messages;
|
||||
stop at the first non-tool message.
|
||||
|
||||
Each result is normalized to `{content, is_error}` — the only fields
|
||||
`signals._detect_failure` / `_detect_exhaustion` actually read.
|
||||
"""
|
||||
if not messages:
|
||||
return []
|
||||
results: List[Dict[str, Any]] = []
|
||||
for msg in reversed(messages):
|
||||
if not isinstance(msg, dict):
|
||||
break
|
||||
if msg.get("role") != "tool":
|
||||
break
|
||||
content = msg.get("content")
|
||||
# Some providers (Anthropic-style) carry an explicit error flag; OpenAI
|
||||
# tool results don't, so fall back to an empty/missing content heuristic
|
||||
# inside `_detect_failure`.
|
||||
is_error = bool(msg.get("is_error"))
|
||||
results.append({"content": content, "is_error": is_error})
|
||||
results.reverse()
|
||||
return results
|
||||
|
||||
|
||||
def _assistant_content_and_tool_calls(response_obj: Any) -> tuple:
|
||||
"""Return (assistant_text, tool_calls_list) extracted from a ModelResponse-ish object."""
|
||||
if response_obj is None:
|
||||
return None, []
|
||||
try:
|
||||
choices = getattr(response_obj, "choices", None) or response_obj.get("choices")
|
||||
except Exception:
|
||||
return None, []
|
||||
if not choices:
|
||||
return None, []
|
||||
|
||||
msg = choices[0]
|
||||
msg = getattr(msg, "message", None) or (
|
||||
msg.get("message") if isinstance(msg, dict) else None
|
||||
)
|
||||
if msg is None:
|
||||
return None, []
|
||||
|
||||
content = getattr(msg, "content", None)
|
||||
if content is None and isinstance(msg, dict):
|
||||
content = msg.get("content")
|
||||
|
||||
raw_tool_calls = getattr(msg, "tool_calls", None)
|
||||
if raw_tool_calls is None and isinstance(msg, dict):
|
||||
raw_tool_calls = msg.get("tool_calls")
|
||||
tool_calls: List[Dict[str, Any]] = []
|
||||
for tc in raw_tool_calls or []:
|
||||
if isinstance(tc, dict):
|
||||
tool_calls.append(tc)
|
||||
else:
|
||||
try:
|
||||
tool_calls.append(tc.model_dump())
|
||||
except Exception:
|
||||
tool_calls.append({"name": getattr(tc, "name", ""), "arguments": ""})
|
||||
return content, tool_calls
|
||||
|
||||
|
||||
class AdaptiveRouterPostCallHook(CustomLogger):
|
||||
"""One hook instance per AdaptiveRouter. Registered into litellm.callbacks."""
|
||||
|
||||
def __init__(self, adaptive_router: AdaptiveRouter) -> None:
|
||||
self.adaptive_router = adaptive_router
|
||||
|
||||
async def async_post_call_response_headers_hook(
|
||||
self,
|
||||
data: Dict[str, Any],
|
||||
user_api_key_dict: Any,
|
||||
response: Any,
|
||||
request_headers: Optional[Dict[str, str]] = None,
|
||||
litellm_call_info: Optional[Dict[str, Any]] = None,
|
||||
) -> Optional[Dict[str, str]]:
|
||||
"""
|
||||
Surface the chosen logical model as the `x-litellm-adaptive-router-model`
|
||||
response header for both streaming and non-streaming responses.
|
||||
|
||||
`async_post_call_success_hook` fires after the stream is fully consumed,
|
||||
so writing to `_hidden_params["additional_headers"]` there is too late for
|
||||
streaming — the StreamingResponse headers are already frozen. This hook is
|
||||
called during header construction (before StreamingResponse is built), so
|
||||
the header is included for both paths.
|
||||
"""
|
||||
metadata = data.get("metadata") or {}
|
||||
chosen = (
|
||||
metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY)
|
||||
if isinstance(metadata, dict)
|
||||
else None
|
||||
)
|
||||
if not chosen:
|
||||
return None
|
||||
return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen}
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
await self._record(kwargs, response_obj, response_status=200)
|
||||
|
||||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
status = kwargs.get("response_status")
|
||||
if status is None:
|
||||
exc = kwargs.get("exception")
|
||||
status = getattr(exc, "status_code", 500) if exc is not None else 500
|
||||
await self._record(kwargs, response_obj, response_status=int(status))
|
||||
|
||||
async def _record(
|
||||
self,
|
||||
kwargs: Dict[str, Any],
|
||||
response_obj: Any,
|
||||
response_status: int,
|
||||
) -> None:
|
||||
try:
|
||||
messages = kwargs.get("messages") or []
|
||||
if len(messages) < SIGNAL_GATE_MIN_MESSAGES:
|
||||
# Too few turns for any signal to be meaningful — skip.
|
||||
return
|
||||
|
||||
session_key = _resolve_session_key(kwargs)
|
||||
if not session_key:
|
||||
return
|
||||
|
||||
# The bandit cells are keyed by the *logical* model name from
|
||||
# `available_models` (e.g. "smart"/"fast"). `kwargs["model"]` at
|
||||
# post-call time is the physical upstream model
|
||||
# (e.g. "anthropic/claude-opus-4-7"), so it cannot be used directly.
|
||||
# The pre-routing hook stashes the logical pick under this key.
|
||||
litellm_params = kwargs.get("litellm_params") or {}
|
||||
metadata = litellm_params.get("metadata") or {}
|
||||
current_model = (
|
||||
metadata.get(ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY)
|
||||
if isinstance(metadata, dict)
|
||||
else None
|
||||
)
|
||||
if not current_model:
|
||||
return
|
||||
|
||||
if not self.adaptive_router.claim_or_check_owner(
|
||||
session_key, current_model
|
||||
):
|
||||
# A different model owns this conversation — skip attribution.
|
||||
return
|
||||
|
||||
user_text = _last_user_content(messages)
|
||||
assistant_text, tool_calls = _assistant_content_and_tool_calls(response_obj)
|
||||
tool_results = _recent_tool_results(messages)
|
||||
|
||||
request_type = classify_prompt(user_text or "")
|
||||
turn = Turn(
|
||||
user_content=user_text,
|
||||
assistant_content=(
|
||||
assistant_text if isinstance(assistant_text, str) else None
|
||||
),
|
||||
tool_calls=tool_calls,
|
||||
tool_results=tool_results,
|
||||
response_status=response_status,
|
||||
)
|
||||
await self.adaptive_router.record_turn(
|
||||
session_id=session_key,
|
||||
model_name=current_model,
|
||||
request_type=request_type,
|
||||
turn=turn,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
"AdaptiveRouterPostCallHook: failed to record turn: %s", e
|
||||
)
|
||||
287
litellm/router_strategy/adaptive_router/signals.py
Normal file
287
litellm/router_strategy/adaptive_router/signals.py
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
"""
|
||||
Incremental signal detection for the adaptive router.
|
||||
|
||||
Each session maintains a SessionState. On every turn, we call apply_turn(state, turn)
|
||||
which mutates the state in place and returns a SignalDelta listing which signals
|
||||
fired on THIS turn. The router then queues the delta to be flushed to DB.
|
||||
|
||||
Design constraint: O(1) work per turn. No re-scanning the full session history.
|
||||
We keep small bounded windows: last_user_content, last_assistant_content, and a
|
||||
bounded list of recent tool call signatures.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
LOOP_REPEAT_THRESHOLD,
|
||||
MIN_TURNS_FOR_CLEAN_CREDIT,
|
||||
MISALIGNMENT_JACCARD_THRESHOLD,
|
||||
STAGNATION_JACCARD_NEAR_DUP,
|
||||
TOOL_CALL_HISTORY_MAX,
|
||||
)
|
||||
|
||||
|
||||
# ---- Public types ---------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class SignalDelta:
|
||||
"""Which signals fired on a single turn. Counts are 0 or 1 (one delta per turn)."""
|
||||
|
||||
misalignment: int = 0
|
||||
stagnation: int = 0
|
||||
disengagement: int = 0
|
||||
satisfaction: int = 0
|
||||
failure: int = 0
|
||||
loop: int = 0
|
||||
exhaustion: int = 0
|
||||
|
||||
def any_fired(self) -> bool:
|
||||
return any(
|
||||
[
|
||||
self.misalignment,
|
||||
self.stagnation,
|
||||
self.disengagement,
|
||||
self.satisfaction,
|
||||
self.failure,
|
||||
self.loop,
|
||||
self.exhaustion,
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionState:
|
||||
"""In-memory rolling state for one session.
|
||||
|
||||
Mirrors the LiteLLM_AdaptiveRouterSession DB row (Wave 0 schema). The flusher
|
||||
later persists this. We keep this as a plain dataclass — no DB coupling.
|
||||
"""
|
||||
|
||||
session_id: str
|
||||
router_name: str
|
||||
model_name: str
|
||||
classified_type: str
|
||||
|
||||
misalignment_count: int = 0
|
||||
stagnation_count: int = 0
|
||||
disengagement_count: int = 0
|
||||
satisfaction_count: int = 0
|
||||
failure_count: int = 0
|
||||
loop_count: int = 0
|
||||
exhaustion_count: int = 0
|
||||
|
||||
last_user_content: Optional[str] = None
|
||||
last_assistant_content: Optional[str] = None
|
||||
tool_call_history: List[str] = field(default_factory=list)
|
||||
pending_tool_calls: Dict[str, str] = field(default_factory=dict)
|
||||
|
||||
turn_count: int = 0
|
||||
last_processed_turn: int = -1
|
||||
clean_credit_awarded: bool = False
|
||||
terminal_status: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Turn:
|
||||
"""One turn of input. Caller assembles this from the request/response."""
|
||||
|
||||
user_content: Optional[str] = None
|
||||
assistant_content: Optional[str] = None
|
||||
tool_calls: List[Dict[str, Any]] = field(default_factory=list)
|
||||
tool_results: List[Dict[str, Any]] = field(default_factory=list)
|
||||
response_status: Optional[int] = None
|
||||
|
||||
|
||||
# ---- Detection helpers ----------------------------------------------------
|
||||
|
||||
_TOKEN_RE = re.compile(r"[A-Za-z0-9]+")
|
||||
|
||||
|
||||
def _tokens(text: Optional[str]) -> Set[str]:
|
||||
if not text:
|
||||
return set()
|
||||
return {t.lower() for t in _TOKEN_RE.findall(text)}
|
||||
|
||||
|
||||
def _jaccard(a: Set[str], b: Set[str]) -> float:
|
||||
union = a | b
|
||||
if not union:
|
||||
return 0.0
|
||||
return len(a & b) / len(union)
|
||||
|
||||
|
||||
_DISENGAGEMENT_PATTERNS = [
|
||||
re.compile(
|
||||
r"\b(forget it|never mind|give up|talk to (?:a )?human|cancel)\b", re.IGNORECASE
|
||||
),
|
||||
re.compile(r"\b(this (?:isn'?t|is not) working|stop|abort)\b", re.IGNORECASE),
|
||||
re.compile(r"\bi'?ll do it (?:myself|manually)\b", re.IGNORECASE),
|
||||
]
|
||||
|
||||
_SATISFACTION_PATTERNS = [
|
||||
re.compile(
|
||||
r"\b(that worked|that did it|works now|fixed it|solved it|nice)\b",
|
||||
re.IGNORECASE,
|
||||
),
|
||||
re.compile(r"\b(thanks|thank you|thx|appreciated|appreciate it)\b", re.IGNORECASE),
|
||||
re.compile(r"\b(perfect|great|excellent|exactly)\b", re.IGNORECASE),
|
||||
]
|
||||
|
||||
|
||||
def _detect_misalignment(prev_user: Optional[str], curr_user: Optional[str]) -> bool:
|
||||
"""Fires when consecutive user messages share *some* topic (jaccard > 0)
|
||||
but are sufficiently different (jaccard < threshold) — i.e. user is
|
||||
rephrasing, not changing topic, not repeating."""
|
||||
if not prev_user or not curr_user:
|
||||
return False
|
||||
j = _jaccard(_tokens(prev_user), _tokens(curr_user))
|
||||
return 0.0 < j < MISALIGNMENT_JACCARD_THRESHOLD
|
||||
|
||||
|
||||
def _detect_stagnation(prev_asst: Optional[str], curr_asst: Optional[str]) -> bool:
|
||||
"""Fires when consecutive assistant messages are near-duplicates."""
|
||||
if not prev_asst or not curr_asst:
|
||||
return False
|
||||
j = _jaccard(_tokens(prev_asst), _tokens(curr_asst))
|
||||
return j >= STAGNATION_JACCARD_NEAR_DUP
|
||||
|
||||
|
||||
def _detect_disengagement(curr_user: Optional[str]) -> bool:
|
||||
if not curr_user:
|
||||
return False
|
||||
return any(p.search(curr_user) for p in _DISENGAGEMENT_PATTERNS)
|
||||
|
||||
|
||||
def _detect_satisfaction(curr_user: Optional[str]) -> bool:
|
||||
if not curr_user:
|
||||
return False
|
||||
return any(p.search(curr_user) for p in _SATISFACTION_PATTERNS)
|
||||
|
||||
|
||||
def _detect_failure(tool_results: List[Dict[str, Any]]) -> bool:
|
||||
"""Any tool result explicitly flagged as an error.
|
||||
|
||||
We do NOT treat empty content as failure — many tools legitimately return
|
||||
empty output (zero-result searches, silent bash commands, void writes) and
|
||||
penalizing the model for those would corrupt the bandit posterior.
|
||||
"""
|
||||
for r in tool_results:
|
||||
if r.get("is_error"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _signature(call: Dict[str, Any]) -> str:
|
||||
"""Stable signature for loop detection: name + sorted JSON-ish args."""
|
||||
name = call.get("name") or call.get("function", {}).get("name", "")
|
||||
call_args = call.get("arguments")
|
||||
if call_args is None:
|
||||
call_args = call.get("function", {}).get("arguments", "")
|
||||
if isinstance(call_args, dict):
|
||||
call_args = ",".join(f"{k}={call_args[k]}" for k in sorted(call_args.keys()))
|
||||
return f"{name}({call_args})"
|
||||
|
||||
|
||||
def _detect_loop(history: List[str], new_calls: List[Dict[str, Any]]) -> bool:
|
||||
"""Fires if any new call's signature appears >= LOOP_REPEAT_THRESHOLD-1 times
|
||||
in recent history (so this call would be the Nth)."""
|
||||
if not new_calls:
|
||||
return False
|
||||
for call in new_calls:
|
||||
sig = _signature(call)
|
||||
recent_count = history.count(sig)
|
||||
if recent_count >= LOOP_REPEAT_THRESHOLD - 1:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
_EXHAUSTION_STATUSES = {408, 413, 429, 503, 504}
|
||||
|
||||
_EXHAUSTION_KEYWORDS = (
|
||||
"context length",
|
||||
"context window",
|
||||
"token limit",
|
||||
"rate limit",
|
||||
"too many requests",
|
||||
"timeout",
|
||||
)
|
||||
|
||||
|
||||
def _detect_exhaustion(
|
||||
status: Optional[int], tool_results: List[Dict[str, Any]]
|
||||
) -> bool:
|
||||
if status is not None and status in _EXHAUSTION_STATUSES:
|
||||
return True
|
||||
for r in tool_results:
|
||||
content = str(r.get("content", "")).lower()
|
||||
if any(kw in content for kw in _EXHAUSTION_KEYWORDS):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# ---- Public entrypoint ----------------------------------------------------
|
||||
|
||||
|
||||
def apply_turn(state: SessionState, turn: Turn) -> SignalDelta:
|
||||
"""
|
||||
Detect signals on this turn, mutate state, return the delta.
|
||||
|
||||
O(1) per turn (no full-history rescan). Only inspects last_*, recent tool history
|
||||
(which is bounded at TOOL_CALL_HISTORY_MAX), and the new turn payload.
|
||||
"""
|
||||
delta = SignalDelta()
|
||||
|
||||
if _detect_misalignment(state.last_user_content, turn.user_content):
|
||||
delta.misalignment = 1
|
||||
if _detect_stagnation(state.last_assistant_content, turn.assistant_content):
|
||||
delta.stagnation = 1
|
||||
if _detect_disengagement(turn.user_content):
|
||||
delta.disengagement = 1
|
||||
if _detect_satisfaction(turn.user_content):
|
||||
# Gate: only award satisfaction credit once per session, and only
|
||||
# after MIN_TURNS_FOR_CLEAN_CREDIT turns of context. Early "thanks"
|
||||
# on turn 1-2 is noise, not a validated quality signal.
|
||||
current_turn_index = state.turn_count + 1
|
||||
if (
|
||||
not state.clean_credit_awarded
|
||||
and current_turn_index >= MIN_TURNS_FOR_CLEAN_CREDIT
|
||||
):
|
||||
delta.satisfaction = 1
|
||||
state.clean_credit_awarded = True
|
||||
if _detect_failure(turn.tool_results):
|
||||
delta.failure = 1
|
||||
if _detect_loop(state.tool_call_history, turn.tool_calls):
|
||||
delta.loop = 1
|
||||
if _detect_exhaustion(turn.response_status, turn.tool_results):
|
||||
delta.exhaustion = 1
|
||||
|
||||
state.misalignment_count += delta.misalignment
|
||||
state.stagnation_count += delta.stagnation
|
||||
state.disengagement_count += delta.disengagement
|
||||
state.satisfaction_count += delta.satisfaction
|
||||
state.failure_count += delta.failure
|
||||
state.loop_count += delta.loop
|
||||
state.exhaustion_count += delta.exhaustion
|
||||
|
||||
if turn.user_content:
|
||||
state.last_user_content = turn.user_content
|
||||
if turn.assistant_content:
|
||||
state.last_assistant_content = turn.assistant_content
|
||||
|
||||
for call in turn.tool_calls:
|
||||
state.tool_call_history.append(_signature(call))
|
||||
if len(state.tool_call_history) > TOOL_CALL_HISTORY_MAX:
|
||||
state.tool_call_history = state.tool_call_history[-TOOL_CALL_HISTORY_MAX:]
|
||||
|
||||
if turn.response_status is not None:
|
||||
state.terminal_status = turn.response_status
|
||||
|
||||
state.turn_count += 1
|
||||
state.last_processed_turn = state.turn_count
|
||||
|
||||
return delta
|
||||
213
litellm/router_strategy/adaptive_router/update_queue.py
Normal file
213
litellm/router_strategy/adaptive_router/update_queue.py
Normal file
|
|
@ -0,0 +1,213 @@
|
|||
"""
|
||||
In-memory queues for adaptive router state and session updates.
|
||||
|
||||
Pattern follows DailySpendUpdateQueue: hot path is fully in-memory; a background
|
||||
flusher task drains the aggregator and writes batches to Postgres.
|
||||
|
||||
Two logical queues (one class):
|
||||
1. STATE updates: increments to (router, request_type, model) bandit cell.
|
||||
Aggregator key = (router_name, request_type, model_name)
|
||||
Aggregated payload = {"delta_alpha": float, "delta_beta": float, "samples_added": int}
|
||||
2. SESSION updates: full snapshot of a session row (last-write-wins per session+router+model).
|
||||
Aggregator key = (session_id, router_name, model_name)
|
||||
Aggregated payload = the full session state dict.
|
||||
|
||||
Hot-path API is non-blocking and synchronous from the caller's POV (it just appends
|
||||
to the in-memory aggregator). Flush is async and batched.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
|
||||
StateKey = Tuple[str, str, str] # (router_name, request_type, model_name)
|
||||
SessionKey = Tuple[str, str, str] # (session_id, router_name, model_name)
|
||||
|
||||
|
||||
class AdaptiveRouterUpdateQueue:
|
||||
"""
|
||||
Single class managing both state-update aggregation and session-snapshot aggregation.
|
||||
Held by the AdaptiveRouter strategy instance and started by the proxy on boot.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._state_agg: Dict[StateKey, Dict[str, float]] = {}
|
||||
self._session_agg: Dict[SessionKey, Dict[str, Any]] = {}
|
||||
self._lock = asyncio.Lock()
|
||||
self._max_state_size_seen = 0
|
||||
self._max_session_size_seen = 0
|
||||
|
||||
# ---- Hot-path: state delta -------------------------------------------
|
||||
|
||||
async def add_state_delta(
|
||||
self,
|
||||
router_name: str,
|
||||
request_type: str,
|
||||
model_name: str,
|
||||
delta_alpha: float,
|
||||
delta_beta: float,
|
||||
) -> None:
|
||||
"""Aggregate a bandit-cell delta. Multiple deltas to the same cell sum."""
|
||||
key: StateKey = (router_name, request_type, model_name)
|
||||
async with self._lock:
|
||||
current = self._state_agg.get(key)
|
||||
if current is None:
|
||||
self._state_agg[key] = {
|
||||
"delta_alpha": delta_alpha,
|
||||
"delta_beta": delta_beta,
|
||||
"samples_added": 1,
|
||||
}
|
||||
else:
|
||||
current["delta_alpha"] += delta_alpha
|
||||
current["delta_beta"] += delta_beta
|
||||
current["samples_added"] += 1
|
||||
if len(self._state_agg) > self._max_state_size_seen:
|
||||
self._max_state_size_seen = len(self._state_agg)
|
||||
|
||||
# ---- Hot-path: session snapshot --------------------------------------
|
||||
|
||||
async def add_session_state(
|
||||
self,
|
||||
session_id: str,
|
||||
router_name: str,
|
||||
model_name: str,
|
||||
state_dict: Dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
Last-write-wins per session row. The state_dict is a snapshot of the
|
||||
SessionState (signals counts + bookkeeping fields). The flusher will
|
||||
upsert this into LiteLLM_AdaptiveRouterSession.
|
||||
"""
|
||||
key: SessionKey = (session_id, router_name, model_name)
|
||||
async with self._lock:
|
||||
self._session_agg[key] = state_dict
|
||||
if len(self._session_agg) > self._max_session_size_seen:
|
||||
self._max_session_size_seen = len(self._session_agg)
|
||||
|
||||
# ---- Flushers (called by background task) ----------------------------
|
||||
|
||||
async def flush_state_to_db(self, prisma_client: Any) -> int:
|
||||
"""
|
||||
Drain state aggregator and apply to LiteLLM_AdaptiveRouterState.
|
||||
Returns number of cells flushed.
|
||||
"""
|
||||
async with self._lock:
|
||||
batch = self._state_agg
|
||||
self._state_agg = {}
|
||||
|
||||
if not batch:
|
||||
return 0
|
||||
|
||||
# Sort keys to give deterministic write order across writers and
|
||||
# reduce the chance of cross-row deadlocks when other workers race us.
|
||||
for key in sorted(batch.keys()):
|
||||
router, rt, model = key
|
||||
payload = batch[key]
|
||||
try:
|
||||
# Atomic increment: push the delta directly into the DB so
|
||||
# concurrent flushers from multiple pods don't overwrite each
|
||||
# other. The upsert creates the row with the delta as the
|
||||
# initial value on first write, then increments on subsequent
|
||||
# writes — no read-modify-write race.
|
||||
await prisma_client.db.litellm_adaptiverouterstate.upsert(
|
||||
where={
|
||||
"router_name_request_type_model_name": {
|
||||
"router_name": router,
|
||||
"request_type": rt,
|
||||
"model_name": model,
|
||||
}
|
||||
},
|
||||
data={
|
||||
"create": {
|
||||
"router_name": router,
|
||||
"request_type": rt,
|
||||
"model_name": model,
|
||||
"alpha": payload["delta_alpha"],
|
||||
"beta": payload["delta_beta"],
|
||||
"total_samples": int(payload["samples_added"]),
|
||||
},
|
||||
"update": {
|
||||
"alpha": {"increment": payload["delta_alpha"]},
|
||||
"beta": {"increment": payload["delta_beta"]},
|
||||
"total_samples": {
|
||||
"increment": int(payload["samples_added"])
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
"AdaptiveRouterUpdateQueue: failed to flush state for %s: %s",
|
||||
key,
|
||||
e,
|
||||
)
|
||||
|
||||
return len(batch)
|
||||
|
||||
async def flush_session_to_db(self, prisma_client: Any) -> int:
|
||||
"""
|
||||
Drain session aggregator and upsert into LiteLLM_AdaptiveRouterSession.
|
||||
Returns number of session rows flushed.
|
||||
"""
|
||||
async with self._lock:
|
||||
batch = self._session_agg
|
||||
self._session_agg = {}
|
||||
|
||||
if not batch:
|
||||
return 0
|
||||
|
||||
for key in sorted(batch.keys()):
|
||||
session_id, router, model = key
|
||||
payload = batch[key]
|
||||
try:
|
||||
# NOTE: Prisma client lower-cases model names, so
|
||||
# `LiteLLM_AdaptiveRouterSession` -> `litellm_adaptiveroutersession`
|
||||
# (single 's', not 'litellm_adaptiverouterssession').
|
||||
# Strip PK fields from the update payload — Prisma rejects
|
||||
# writes to fields that are part of the @@id. asdict(state)
|
||||
# always carries them, so build a separate update dict.
|
||||
update_payload = {
|
||||
k: v
|
||||
for k, v in payload.items()
|
||||
if k not in ("session_id", "router_name", "model_name")
|
||||
}
|
||||
await prisma_client.db.litellm_adaptiveroutersession.upsert(
|
||||
where={
|
||||
"session_id_router_name_model_name": {
|
||||
"session_id": session_id,
|
||||
"router_name": router,
|
||||
"model_name": model,
|
||||
}
|
||||
},
|
||||
data={
|
||||
"create": {
|
||||
"session_id": session_id,
|
||||
"router_name": router,
|
||||
"model_name": model,
|
||||
**update_payload,
|
||||
},
|
||||
"update": update_payload,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_router_logger.exception(
|
||||
"AdaptiveRouterUpdateQueue: failed to flush session for %s: %s",
|
||||
key,
|
||||
e,
|
||||
)
|
||||
|
||||
return len(batch)
|
||||
|
||||
# ---- Observability ---------------------------------------------------
|
||||
|
||||
async def queue_size(self) -> Dict[str, int]:
|
||||
async with self._lock:
|
||||
return {
|
||||
"state_pending": len(self._state_agg),
|
||||
"session_pending": len(self._session_agg),
|
||||
"max_state_seen": self._max_state_size_seen,
|
||||
"max_session_seen": self._max_session_size_seen,
|
||||
}
|
||||
|
|
@ -8,7 +8,7 @@ from dataclasses import dataclass
|
|||
from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_type_hints
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
|
||||
from typing_extensions import Required, TypedDict
|
||||
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -221,6 +221,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
complexity_router_config: Optional[Dict] = None
|
||||
complexity_router_default_model: Optional[str] = None
|
||||
|
||||
# adaptive-router params
|
||||
adaptive_router_default_model: Optional[str] = None
|
||||
adaptive_router_config: Optional[Dict] = None
|
||||
# quality-router params
|
||||
quality_router_config: Optional[Dict] = None
|
||||
quality_router_default_model: Optional[str] = None
|
||||
|
|
@ -794,3 +797,44 @@ class PreRoutingHookResponse(BaseModel):
|
|||
|
||||
model: str
|
||||
messages: Optional[List[Dict[str, Any]]]
|
||||
|
||||
|
||||
class RequestType(str, enum.Enum):
|
||||
"""Fixed v0 taxonomy. User-extensible types come in v1."""
|
||||
|
||||
CODE_GENERATION = "code_generation"
|
||||
CODE_UNDERSTANDING = "code_understanding"
|
||||
TECHNICAL_DESIGN = "technical_design"
|
||||
ANALYTICAL_REASONING = "analytical_reasoning"
|
||||
WRITING = "writing"
|
||||
FACTUAL_LOOKUP = "factual_lookup"
|
||||
GENERAL = "general"
|
||||
|
||||
|
||||
class AdaptiveRouterWeights(BaseModel):
|
||||
quality: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||
cost: float = Field(default=0.3, ge=0.0, le=1.0)
|
||||
|
||||
@field_validator("cost")
|
||||
@classmethod
|
||||
def _weights_sum_to_one(cls, v, info):
|
||||
q = info.data.get("quality", 0.7)
|
||||
if abs(q + v - 1.0) > 0.001:
|
||||
raise ValueError(
|
||||
f"weights must sum to 1.0, got quality={q} + cost={v} = {q + v}"
|
||||
)
|
||||
return v
|
||||
|
||||
|
||||
class AdaptiveRouterConfig(BaseModel):
|
||||
available_models: List[str]
|
||||
weights: AdaptiveRouterWeights = Field(default_factory=AdaptiveRouterWeights)
|
||||
|
||||
|
||||
class AdaptiveRouterPreferences(BaseModel):
|
||||
"""model_info.adaptive_router_preferences — declared by each model."""
|
||||
|
||||
model_config = ConfigDict(use_enum_values=False)
|
||||
|
||||
quality_tier: int = Field(ge=1, le=3)
|
||||
strengths: List[RequestType] = Field(default_factory=list)
|
||||
|
|
|
|||
|
|
@ -22886,6 +22886,22 @@
|
|||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-k2.6": {
|
||||
"cache_read_input_token_cost": 1.6e-07,
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "moonshot",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-06,
|
||||
"source": "https://platform.kimi.ai/docs/pricing/chat-k26",
|
||||
"supports_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_video_input": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"moonshot/kimi-latest": {
|
||||
"cache_read_input_token_cost": 1.5e-07,
|
||||
"input_cost_per_token": 2e-06,
|
||||
|
|
@ -25149,6 +25165,28 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"openrouter/anthropic/claude-opus-4.7": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"litellm_provider": "openrouter",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"supports_xhigh_reasoning_effort": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"openrouter/bytedance/ui-tars-1.5-7b": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "openrouter",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.83.11"
|
||||
version = "1.83.12"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.14"
|
||||
|
|
@ -52,7 +52,7 @@ proxy = [
|
|||
"azure-identity==1.25.2",
|
||||
"azure-storage-blob==12.28.0",
|
||||
"mcp==1.26.0",
|
||||
"litellm-proxy-extras==0.4.67",
|
||||
"litellm-proxy-extras==0.4.68",
|
||||
"litellm-enterprise==0.1.38",
|
||||
"RestrictedPython==8.1",
|
||||
"rich==13.9.4",
|
||||
|
|
@ -236,7 +236,7 @@ source-exclude = [
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.83.11"
|
||||
version = "1.83.12"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -616,6 +616,7 @@ model LiteLLM_TeamMembership {
|
|||
user_id String
|
||||
team_id String
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
budget_id String?
|
||||
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
|
||||
@@id([user_id, team_id])
|
||||
|
|
@ -1223,3 +1224,46 @@ model LiteLLM_ClaudeCodePluginTable {
|
|||
|
||||
@@map("LiteLLM_ClaudeCodePluginTable")
|
||||
}
|
||||
|
||||
// Per-(router, request_type, model) Beta posterior for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterState {
|
||||
router_name String
|
||||
request_type String
|
||||
model_name String
|
||||
alpha Float
|
||||
beta Float
|
||||
total_samples Int @default(0)
|
||||
last_updated_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([router_name, request_type, model_name])
|
||||
}
|
||||
|
||||
// Per-(session, router, model) signal counters for the adaptive router.
|
||||
model LiteLLM_AdaptiveRouterSession {
|
||||
session_id String
|
||||
router_name String
|
||||
model_name String
|
||||
classified_type String
|
||||
|
||||
misalignment_count Int @default(0)
|
||||
stagnation_count Int @default(0)
|
||||
disengagement_count Int @default(0)
|
||||
satisfaction_count Int @default(0)
|
||||
failure_count Int @default(0)
|
||||
loop_count Int @default(0)
|
||||
exhaustion_count Int @default(0)
|
||||
|
||||
last_user_content String?
|
||||
last_assistant_content String?
|
||||
tool_call_history Json @default("[]")
|
||||
pending_tool_calls Json @default("{}")
|
||||
|
||||
turn_count Int @default(0)
|
||||
last_processed_turn Int @default(-1)
|
||||
clean_credit_awarded Boolean @default(false)
|
||||
terminal_status Int?
|
||||
last_activity_at DateTime @default(now()) @updatedAt
|
||||
|
||||
@@id([session_id, router_name, model_name])
|
||||
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
|
||||
}
|
||||
|
|
|
|||
157
scripts/adaptive_router_demo/README.md
Normal file
157
scripts/adaptive_router_demo/README.md
Normal file
|
|
@ -0,0 +1,157 @@
|
|||
# Adaptive Router — Live Demo
|
||||
|
||||
A 5-minute demo of LiteLLM's adaptive router learning, in real time, that
|
||||
the smart model wins for code while the fast model is fine for facts.
|
||||
|
||||
```
|
||||
┌─ traffic.py ──┐ ┌─ litellm proxy ──────────┐ ┌─ dashboard.html ─┐
|
||||
│ synthetic │──▶│ adaptive_router strategy │──▶│ bandit bars + │
|
||||
│ chat sessions │ │ /adaptive_router/state │ │ cost meter + │
|
||||
└───────────────┘ └──────────┬───────────────┘ │ activity log │
|
||||
│ └───────────────────┘
|
||||
┌─────────▼───────────┐
|
||||
│ chat.html │
|
||||
│ interactive chat │
|
||||
│ with preset │
|
||||
│ scenarios │
|
||||
└─────────────────────┘
|
||||
```
|
||||
|
||||
## Files
|
||||
|
||||
| File | What it does |
|
||||
|---|---|
|
||||
| `dashboard.html` | Live bandit dashboard — polls `/adaptive_router/state` every 500ms |
|
||||
| `chat.html` | Interactive chat with preset scenarios — sends real requests through the router |
|
||||
| `traffic.py` | Synthetic traffic generator — drives labeled sessions for automated demo |
|
||||
|
||||
## What you're watching
|
||||
|
||||
- **Bandit posteriors** — one Beta(α, β) bar per `(request_type, model)`
|
||||
cell. Bars fill up as α grows from positive feedback signals.
|
||||
- **Pick share** — softmax estimate of how often the router would currently
|
||||
pick each model for that request type.
|
||||
- **Cost meter** — total spend so far compared to "always use the most
|
||||
expensive model". The savings line is the headline number.
|
||||
- **Activity log** — every signal that moves the bandit, in real time.
|
||||
|
||||
## 1. Start the proxy
|
||||
|
||||
The repo ships with a working example config:
|
||||
|
||||
```bash
|
||||
export OPENAI_API_KEY=sk-... # underlying models hit OpenAI
|
||||
uv run litellm \
|
||||
--config litellm/proxy/example_config_yaml/adaptive_router_example.yaml \
|
||||
--port 4000
|
||||
```
|
||||
|
||||
`DATABASE_URL` is optional — the proxy falls back to a bundled Neon dev DB.
|
||||
Wait ~15s until you see `Application startup complete`.
|
||||
|
||||
## 2. Chat interactively with the router
|
||||
|
||||
Open `chat.html` in a browser (same `file://` or `python3 -m http.server` approach as the dashboard):
|
||||
|
||||
- Click **Connect** after filling in the proxy URL and API key.
|
||||
- Pick a preset scenario:
|
||||
- **🐛 Debug my code** — paste broken code and get a fix
|
||||
- **💡 Brainstorm a feature** — ideate on a product capability
|
||||
- **📚 Explain a concept** — get a clear technical explanation
|
||||
- **✍️ Write something** — draft emails, docs, or any prose
|
||||
- A starter message is pre-filled — edit it or send as-is.
|
||||
- Each response shows which model the router picked and the inferred request type (from the `x-litellm-adaptive-router-model` and `x-litellm-request-type` response headers).
|
||||
- A sidebar gate indicator tells you when the session has accumulated enough messages for the bandit to start updating (4+ turns).
|
||||
|
||||
> **Note on headers:** The model/type headers are only readable in the browser if the proxy sets `Access-Control-Expose-Headers`. LiteLLM defaults to exposing them. If the info panel shows `check dashboard`, the router still works — you can verify picks in `dashboard.html`.
|
||||
|
||||
## 4. Open the dashboard
|
||||
|
||||
The dashboard is a single static HTML file. Either:
|
||||
|
||||
- **Easy:** double-click `dashboard.html`. Most browsers will load it from
|
||||
`file://` and the LiteLLM proxy's CORS defaults (`*`) will accept it.
|
||||
- **If your browser blocks `file://` fetches:**
|
||||
|
||||
```bash
|
||||
cd scripts/adaptive_router_demo
|
||||
python3 -m http.server 8080
|
||||
```
|
||||
|
||||
Then open <http://localhost:8080/dashboard.html>.
|
||||
|
||||
In the connect bar, fill in:
|
||||
|
||||
- **Proxy URL:** `http://localhost:4000`
|
||||
- **Master Key:** the `master_key` from your config (`sk-1234` in the example).
|
||||
|
||||
Click **Connect**. The dashboard polls `GET /adaptive_router/state` every
|
||||
500ms (admin-only endpoint, returns one snapshot per configured router).
|
||||
|
||||
## 5. Drive synthetic traffic
|
||||
|
||||
In a second terminal:
|
||||
|
||||
```bash
|
||||
uv run python scripts/adaptive_router_demo/traffic.py \
|
||||
--proxy-url http://localhost:4000 \
|
||||
--api-key sk-1234 \
|
||||
--router smart-cheap-router \
|
||||
--rounds 100 \
|
||||
--rate 0.5
|
||||
```
|
||||
|
||||
What it does:
|
||||
|
||||
- Picks a random `(request_type, prompt)` per round from a small labeled corpus.
|
||||
- Sends a 5-message conversation (passes the `SIGNAL_GATE_MIN_MESSAGES=4` gate
|
||||
in one round-trip) so the post-call hook runs and updates the bandit.
|
||||
- Reads the `x-litellm-adaptive-router-model` response header to see what
|
||||
the router picked.
|
||||
- Rolls Bernoulli against a hard-coded oracle:
|
||||
```
|
||||
code_generation : smart=0.92 fast=0.35
|
||||
factual_lookup : smart=0.90 fast=0.85
|
||||
writing : smart=0.85 fast=0.55
|
||||
```
|
||||
- On success → sends a follow-up engineered to match the satisfaction
|
||||
regex (and re-classify into the same type). Bandit cell gets +α.
|
||||
- On failure → sends a neutral follow-up. No signal fires.
|
||||
|
||||
After 50–80 rounds you'll see `code_generation` decisively favor `smart`
|
||||
while `factual_lookup` stays near a coin flip — the router learned the
|
||||
asymmetry from the oracle.
|
||||
|
||||
## Tuning knobs
|
||||
|
||||
| Knob | Where | What changes |
|
||||
|---|---|---|
|
||||
| Quality vs. cost weight | `adaptive_router_config.weights` in proxy yaml | Bias toward quality or savings |
|
||||
| Per-cell cold-start mass | `litellm/router_strategy/adaptive_router/config.py` `COLD_START_MASS` | How long until the prior is overwritten |
|
||||
| Avg tokens per request | dashboard input box | How the cost meter estimates spend |
|
||||
| Oracle | `traffic.py` `ORACLE` dict | Which model "should" win for which type |
|
||||
| Sessions to drive | `--rounds` | Total learning budget |
|
||||
| Throttle | `--rate` | Seconds between sessions |
|
||||
|
||||
## Multi-router
|
||||
|
||||
If your proxy has more than one `auto_router/adaptive_router` deployment,
|
||||
the dashboard shows a router dropdown above the bars. Each router is
|
||||
independent; the cost meter is per-router (and resets when you switch).
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
- **"Disconnected" / HTTP 401 in the dashboard** — wrong master key.
|
||||
- **HTTP 403** — your key isn't `proxy_admin`. The state endpoint is
|
||||
admin-only. Use the master key.
|
||||
- **HTTP 404 from `/adaptive_router/state`** — proxy started, but no
|
||||
`auto_router/adaptive_router` deployment is in the model list.
|
||||
- **Bars don't move** — check the proxy logs for `record_turn` activity.
|
||||
Common cause: requests are not including 4+ messages, so the signal
|
||||
gate skips them. `traffic.py` already builds 5-message conversations,
|
||||
so this only happens if you've changed the script.
|
||||
- **Cost meter stays at $0** — your model deployments don't have
|
||||
`input_cost_per_token` set in `litellm_params`. Add it.
|
||||
- **CORS error in the dashboard console** — set `LITELLM_CORS_ORIGINS=*`
|
||||
on the proxy (the default), or serve `dashboard.html` from
|
||||
`python3 -m http.server` instead of `file://`.
|
||||
838
scripts/adaptive_router_demo/chat.html
Normal file
838
scripts/adaptive_router_demo/chat.html
Normal file
|
|
@ -0,0 +1,838 @@
|
|||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<title>Adaptive Router — Chat</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0b0f17;
|
||||
--panel: #131a26;
|
||||
--panel-2: #1b2433;
|
||||
--fg: #e7ecf3;
|
||||
--muted: #8a95a8;
|
||||
--accent: #5dd6a4;
|
||||
--accent-2: #6fb6ff;
|
||||
--warn: #f6b94d;
|
||||
--bad: #ff6b6b;
|
||||
--bar-bg: #233047;
|
||||
--border: #25324a;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: -apple-system, BlinkMacSystemFont, "SF Pro Display",
|
||||
"Segoe UI", Roboto, Inter, sans-serif;
|
||||
background: var(--bg);
|
||||
color: var(--fg);
|
||||
font-size: 14px;
|
||||
line-height: 1.45;
|
||||
height: 100vh;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
|
||||
header {
|
||||
padding: 14px 24px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 16px;
|
||||
background: var(--panel);
|
||||
flex-shrink: 0;
|
||||
}
|
||||
header h1 { margin: 0; font-size: 17px; font-weight: 600; }
|
||||
.dot { width: 8px; height: 8px; border-radius: 50%; background: var(--bad); display: inline-block; margin-right: 6px; }
|
||||
.dot.ok { background: var(--accent); }
|
||||
.status { color: var(--muted); font-size: 12px; }
|
||||
.header-link { margin-left: auto; color: var(--accent-2); font-size: 12px; text-decoration: none; }
|
||||
.header-link:hover { text-decoration: underline; }
|
||||
|
||||
.connect {
|
||||
padding: 12px 24px;
|
||||
display: flex; gap: 10px; align-items: center;
|
||||
background: var(--panel-2);
|
||||
border-bottom: 1px solid var(--border);
|
||||
flex-shrink: 0;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
.connect label { color: var(--muted); font-size: 12px; }
|
||||
.connect input, .connect select {
|
||||
background: #0e1422;
|
||||
color: var(--fg);
|
||||
border: 1px solid var(--border);
|
||||
padding: 6px 10px;
|
||||
border-radius: 6px;
|
||||
font-size: 13px;
|
||||
font-family: inherit;
|
||||
}
|
||||
.connect input[type=text] { width: 210px; }
|
||||
.connect input[type=password] { width: 180px; }
|
||||
.connect button.btn-connect {
|
||||
background: var(--accent-2);
|
||||
color: #0b0f17;
|
||||
border: none;
|
||||
padding: 7px 14px;
|
||||
border-radius: 6px;
|
||||
font-weight: 600;
|
||||
cursor: pointer;
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
/* scenario strip */
|
||||
.scenarios {
|
||||
padding: 10px 24px;
|
||||
display: flex; gap: 8px; flex-wrap: wrap;
|
||||
background: var(--panel);
|
||||
border-bottom: 1px solid var(--border);
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.scenarios .sc-btn {
|
||||
background: var(--panel-2);
|
||||
border: 1px solid var(--border);
|
||||
color: var(--fg);
|
||||
padding: 7px 14px;
|
||||
border-radius: 20px;
|
||||
cursor: pointer;
|
||||
font-size: 13px;
|
||||
font-family: inherit;
|
||||
transition: border-color 0.15s, background 0.15s;
|
||||
}
|
||||
.scenarios .sc-btn:hover { border-color: var(--accent-2); background: #1e2d42; }
|
||||
.scenarios .sc-btn.active { border-color: var(--accent-2); background: #1a2d45; color: var(--accent-2); }
|
||||
.scenarios .sc-btn.new { border-color: var(--border); color: var(--muted); }
|
||||
.scenarios .sc-btn.new:hover { border-color: var(--warn); color: var(--warn); background: #1f1a10; }
|
||||
|
||||
/* main layout */
|
||||
.workspace {
|
||||
display: grid;
|
||||
grid-template-columns: 1fr 300px;
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
}
|
||||
@media (max-width: 900px) {
|
||||
.workspace { grid-template-columns: 1fr; }
|
||||
.info-panel { display: none; }
|
||||
}
|
||||
|
||||
/* chat panel */
|
||||
.chat-panel {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
min-height: 0;
|
||||
border-right: 1px solid var(--border);
|
||||
}
|
||||
|
||||
.messages {
|
||||
flex: 1;
|
||||
overflow-y: auto;
|
||||
padding: 20px 24px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 16px;
|
||||
}
|
||||
.messages .empty-state {
|
||||
margin: auto;
|
||||
text-align: center;
|
||||
color: var(--muted);
|
||||
}
|
||||
.messages .empty-state h2 {
|
||||
font-size: 18px;
|
||||
font-weight: 600;
|
||||
color: var(--fg);
|
||||
margin: 0 0 8px;
|
||||
}
|
||||
.messages .empty-state p {
|
||||
font-size: 13px;
|
||||
margin: 0;
|
||||
max-width: 360px;
|
||||
}
|
||||
|
||||
.msg {
|
||||
display: flex;
|
||||
gap: 12px;
|
||||
align-items: flex-start;
|
||||
}
|
||||
.msg.assistant { flex-direction: row; }
|
||||
.msg.user { flex-direction: row-reverse; }
|
||||
|
||||
.avatar {
|
||||
width: 28px; height: 28px;
|
||||
border-radius: 50%;
|
||||
display: flex; align-items: center; justify-content: center;
|
||||
font-size: 13px;
|
||||
flex-shrink: 0;
|
||||
}
|
||||
.msg.user .avatar { background: var(--accent-2); color: #0b0f17; font-weight: 700; }
|
||||
.msg.assistant .avatar { background: var(--accent); color: #0b0f17; }
|
||||
|
||||
.bubble {
|
||||
max-width: 75%;
|
||||
padding: 10px 14px;
|
||||
border-radius: 12px;
|
||||
font-size: 13px;
|
||||
line-height: 1.55;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
}
|
||||
.msg.user .bubble {
|
||||
background: #1a2d45;
|
||||
border: 1px solid var(--border);
|
||||
border-top-right-radius: 4px;
|
||||
}
|
||||
.msg.assistant .bubble {
|
||||
background: var(--panel);
|
||||
border: 1px solid var(--border);
|
||||
border-top-left-radius: 4px;
|
||||
}
|
||||
.bubble code {
|
||||
font-family: ui-monospace, "SF Mono", Menlo, monospace;
|
||||
background: rgba(255,255,255,0.06);
|
||||
padding: 1px 4px;
|
||||
border-radius: 3px;
|
||||
font-size: 12px;
|
||||
}
|
||||
.bubble pre {
|
||||
background: #0e1422;
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 6px;
|
||||
padding: 10px 12px;
|
||||
overflow-x: auto;
|
||||
margin: 8px 0 0;
|
||||
}
|
||||
.bubble pre code {
|
||||
background: none;
|
||||
padding: 0;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.msg-meta {
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
margin-top: 4px;
|
||||
padding: 0 2px;
|
||||
}
|
||||
.msg.user .msg-meta { text-align: right; }
|
||||
.msg-model { color: var(--accent); font-weight: 600; }
|
||||
.msg-type { color: var(--accent-2); }
|
||||
|
||||
.thinking {
|
||||
display: flex; gap: 4px; align-items: center;
|
||||
padding: 8px 0;
|
||||
}
|
||||
.thinking span {
|
||||
width: 6px; height: 6px; border-radius: 50%;
|
||||
background: var(--muted);
|
||||
animation: blink 1.2s infinite;
|
||||
}
|
||||
.thinking span:nth-child(2) { animation-delay: 0.2s; }
|
||||
.thinking span:nth-child(3) { animation-delay: 0.4s; }
|
||||
@keyframes blink {
|
||||
0%, 80%, 100% { opacity: 0.2; }
|
||||
40% { opacity: 1; }
|
||||
}
|
||||
|
||||
/* input area */
|
||||
.input-area {
|
||||
padding: 14px 24px;
|
||||
border-top: 1px solid var(--border);
|
||||
background: var(--panel);
|
||||
flex-shrink: 0;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 8px;
|
||||
}
|
||||
.input-row {
|
||||
display: flex;
|
||||
gap: 10px;
|
||||
align-items: flex-end;
|
||||
}
|
||||
textarea {
|
||||
flex: 1;
|
||||
background: #0e1422;
|
||||
color: var(--fg);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 8px;
|
||||
padding: 10px 12px;
|
||||
font-size: 13px;
|
||||
font-family: inherit;
|
||||
resize: none;
|
||||
outline: none;
|
||||
min-height: 44px;
|
||||
max-height: 160px;
|
||||
line-height: 1.45;
|
||||
}
|
||||
textarea:focus { border-color: var(--accent-2); }
|
||||
textarea:disabled { opacity: 0.5; }
|
||||
.btn-send {
|
||||
background: var(--accent);
|
||||
color: #0b0f17;
|
||||
border: none;
|
||||
padding: 10px 18px;
|
||||
border-radius: 8px;
|
||||
font-weight: 700;
|
||||
font-size: 13px;
|
||||
cursor: pointer;
|
||||
font-family: inherit;
|
||||
flex-shrink: 0;
|
||||
align-self: flex-end;
|
||||
}
|
||||
.btn-send:disabled { opacity: 0.5; cursor: default; }
|
||||
|
||||
.input-hint {
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
}
|
||||
|
||||
/* info panel */
|
||||
.info-panel {
|
||||
background: var(--panel);
|
||||
padding: 18px 16px;
|
||||
overflow-y: auto;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 16px;
|
||||
}
|
||||
.info-block h3 {
|
||||
margin: 0 0 10px;
|
||||
font-size: 11px;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.6px;
|
||||
color: var(--muted);
|
||||
font-weight: 600;
|
||||
}
|
||||
.info-row {
|
||||
display: flex;
|
||||
justify-content: space-between;
|
||||
align-items: center;
|
||||
font-size: 12px;
|
||||
margin-bottom: 6px;
|
||||
}
|
||||
.info-row .label { color: var(--muted); }
|
||||
.info-row .val { font-family: ui-monospace, monospace; color: var(--fg); font-weight: 600; }
|
||||
.info-row .val.accent { color: var(--accent); }
|
||||
.info-row .val.accent-2 { color: var(--accent-2); }
|
||||
.info-row .val.warn { color: var(--warn); }
|
||||
|
||||
.signal-gate {
|
||||
padding: 8px 10px;
|
||||
border-radius: 6px;
|
||||
font-size: 12px;
|
||||
line-height: 1.4;
|
||||
}
|
||||
.signal-gate.waiting {
|
||||
background: rgba(246, 185, 77, 0.08);
|
||||
border: 1px solid rgba(246, 185, 77, 0.3);
|
||||
color: var(--warn);
|
||||
}
|
||||
.signal-gate.learning {
|
||||
background: rgba(93, 214, 164, 0.08);
|
||||
border: 1px solid rgba(93, 214, 164, 0.3);
|
||||
color: var(--accent);
|
||||
}
|
||||
|
||||
.scenario-desc {
|
||||
font-size: 12px;
|
||||
color: var(--muted);
|
||||
line-height: 1.5;
|
||||
}
|
||||
|
||||
.divider {
|
||||
border: none;
|
||||
border-top: 1px solid var(--border);
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.link-dash {
|
||||
display: block;
|
||||
text-align: center;
|
||||
color: var(--accent-2);
|
||||
font-size: 12px;
|
||||
text-decoration: none;
|
||||
padding: 8px;
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 6px;
|
||||
margin-top: auto;
|
||||
}
|
||||
.link-dash:hover { border-color: var(--accent-2); background: #1a2d45; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<header>
|
||||
<h1>⚡ Adaptive Router — Chat</h1>
|
||||
<span class="status"><span id="conn-dot" class="dot"></span><span id="conn-label">Disconnected</span></span>
|
||||
<a class="header-link" href="dashboard.html">→ Open live dashboard</a>
|
||||
</header>
|
||||
|
||||
<div class="connect">
|
||||
<label>Proxy URL <input id="proxy-url" type="text" value="http://localhost:4000" /></label>
|
||||
<label>API Key <input id="api-key" type="password" placeholder="sk-1234" /></label>
|
||||
<label>Router
|
||||
<input id="router" type="text" value="smart-cheap-router" style="width:160px" />
|
||||
</label>
|
||||
<button class="btn-connect" id="connect-btn">Connect</button>
|
||||
</div>
|
||||
|
||||
<div class="scenarios" id="scenario-bar">
|
||||
<button class="sc-btn" data-id="debug_code">🐛 Debug my code</button>
|
||||
<button class="sc-btn" data-id="brainstorm_feature">💡 Brainstorm a feature</button>
|
||||
<button class="sc-btn" data-id="explain_concept">📚 Explain a concept</button>
|
||||
<button class="sc-btn" data-id="write_something">✍️ Write something</button>
|
||||
<button class="sc-btn new" id="new-chat-btn">+ New chat</button>
|
||||
</div>
|
||||
|
||||
<div class="workspace">
|
||||
<div class="chat-panel">
|
||||
<div class="messages" id="messages">
|
||||
<div class="empty-state">
|
||||
<h2>Pick a scenario to start</h2>
|
||||
<p>Choose one of the presets above or connect to the proxy and type your own message. The adaptive router will pick the best model for each turn.</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="input-area">
|
||||
<div class="input-row">
|
||||
<textarea id="input" rows="1" placeholder="Send a message… (Shift+Enter for new line)" disabled></textarea>
|
||||
<button class="btn-send" id="send-btn" disabled>Send</button>
|
||||
</div>
|
||||
<div class="input-hint" id="input-hint">Connect first to start chatting.</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<aside class="info-panel">
|
||||
<div class="info-block">
|
||||
<h3>Session</h3>
|
||||
<div class="info-row"><span class="label">ID</span><span class="val" id="info-session-id">—</span></div>
|
||||
<div class="info-row"><span class="label">Messages</span><span class="val" id="info-msg-count">0</span></div>
|
||||
<div class="info-row"><span class="label">Scenario</span><span class="val accent-2" id="info-scenario">none</span></div>
|
||||
</div>
|
||||
|
||||
<div id="gate-status" class="signal-gate waiting" style="display:none">
|
||||
⏳ <b>Learning starts at 4 messages.</b> Keep chatting — the router will start updating its bandit after your next reply.
|
||||
</div>
|
||||
|
||||
<hr class="divider" />
|
||||
|
||||
<div class="info-block">
|
||||
<h3>Last response</h3>
|
||||
<div class="info-row"><span class="label">Model picked</span><span class="val accent" id="info-model">—</span></div>
|
||||
<div class="info-row"><span class="label">Request type</span><span class="val accent-2" id="info-req-type">—</span></div>
|
||||
<div class="info-row"><span class="label">Latency</span><span class="val" id="info-latency">—</span></div>
|
||||
</div>
|
||||
|
||||
<hr class="divider" />
|
||||
|
||||
<div class="info-block">
|
||||
<h3>How it works</h3>
|
||||
<p class="scenario-desc">
|
||||
Each message goes through the <b>adaptive router</b> which classifies your request type (code, writing, factual…) and uses a Thompson-sampling bandit to pick the model with the best quality for that category.<br><br>
|
||||
After 4+ messages, positive feedback signals (✓ in the activity log) update the bandit. Watch the bars move in the <a href="dashboard.html" style="color:var(--accent-2)">live dashboard</a>.
|
||||
</p>
|
||||
</div>
|
||||
|
||||
<a class="link-dash" href="dashboard.html">📊 Live bandit dashboard →</a>
|
||||
</aside>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// ---- scenarios -------------------------------------------------------
|
||||
const SCENARIOS = {
|
||||
debug_code: {
|
||||
label: "Debug my code",
|
||||
system: "You are an expert debugging assistant. Be concise. Identify the bug, explain why it's wrong, and provide the corrected code.",
|
||||
starter: "I have a Python function that should return the sum of a list, but it always returns 0:\n\n```python\ndef sum_list(items):\n total = 0\n for item in items:\n total + item\n return total\n\nprint(sum_list([1, 2, 3])) # prints 0, expected 6\n```\n\nWhat's wrong with it?",
|
||||
},
|
||||
brainstorm_feature: {
|
||||
label: "Brainstorm a feature",
|
||||
system: "You are a product thinking partner. Help explore feature ideas with concrete examples, trade-offs, and implementation considerations. Be specific and opinionated.",
|
||||
starter: "I'm building a note-taking app for developers. What are 5 differentiated features that would make it stand out from Notion or Obsidian? Focus on things that would genuinely solve developer pain points.",
|
||||
},
|
||||
explain_concept: {
|
||||
label: "Explain a concept",
|
||||
system: "You are a clear, concise technical educator. Explain concepts with simple language, good analogies, and concrete examples. Avoid unnecessary jargon.",
|
||||
starter: "Can you explain how the Thompson Sampling algorithm works and why it's better than epsilon-greedy for multi-armed bandit problems? Use a concrete example if it helps.",
|
||||
},
|
||||
write_something: {
|
||||
label: "Write something",
|
||||
system: "You are a skilled writer. Produce clear, professional text tailored to the requested format and tone. Match the voice the user asks for.",
|
||||
starter: "Write a short Slack message to my team letting them know our weekly standup is moving from 9am to 10am starting next Monday. Keep it brief, friendly, and include a clear ask for them to update their calendars.",
|
||||
},
|
||||
};
|
||||
|
||||
// ---- state -----------------------------------------------------------
|
||||
const STATE = {
|
||||
connected: false,
|
||||
proxyUrl: "",
|
||||
apiKey: "",
|
||||
router: "",
|
||||
sessionId: "",
|
||||
messages: [], // [{role, content}] — sent to the API
|
||||
scenario: null,
|
||||
sending: false,
|
||||
msgCount: 0, // turns in current session
|
||||
lastModel: null,
|
||||
lastReqType: null,
|
||||
};
|
||||
|
||||
// ---- session ---------------------------------------------------------
|
||||
function newSession() {
|
||||
STATE.sessionId = "chat-" + Math.random().toString(36).slice(2, 10);
|
||||
STATE.messages = [];
|
||||
STATE.msgCount = 0;
|
||||
STATE.lastModel = null;
|
||||
STATE.lastReqType = null;
|
||||
renderInfo();
|
||||
renderGateStatus();
|
||||
}
|
||||
|
||||
// ---- persistence -----------------------------------------------------
|
||||
function ssGet(k) { try { return sessionStorage.getItem(k) || ""; } catch { return ""; } }
|
||||
function ssSet(k, v) { try { sessionStorage.setItem(k, v); } catch {} }
|
||||
|
||||
// ---- rendering -------------------------------------------------------
|
||||
function setConn(ok, label) {
|
||||
document.getElementById("conn-dot").className = "dot" + (ok ? " ok" : "");
|
||||
document.getElementById("conn-label").textContent = ok ? "Connected" : label || "Disconnected";
|
||||
}
|
||||
|
||||
function renderInfo() {
|
||||
const shortId = STATE.sessionId ? STATE.sessionId.slice(-8) : "—";
|
||||
document.getElementById("info-session-id").textContent = shortId;
|
||||
document.getElementById("info-msg-count").textContent = STATE.msgCount;
|
||||
document.getElementById("info-scenario").textContent = STATE.scenario ? SCENARIOS[STATE.scenario].label : "none";
|
||||
document.getElementById("info-model").textContent = STATE.lastModel || "—";
|
||||
document.getElementById("info-req-type").textContent = STATE.lastReqType || "—";
|
||||
}
|
||||
|
||||
function renderGateStatus() {
|
||||
const el = document.getElementById("gate-status");
|
||||
if (STATE.msgCount === 0) { el.style.display = "none"; return; }
|
||||
el.style.display = "";
|
||||
if (STATE.msgCount < 4) {
|
||||
el.className = "signal-gate waiting";
|
||||
el.innerHTML = `⏳ <b>${4 - STATE.msgCount} more message${4 - STATE.msgCount > 1 ? "s" : ""} until learning kicks in.</b> The bandit updates after 4+ turns.`;
|
||||
} else {
|
||||
el.className = "signal-gate learning";
|
||||
el.innerHTML = `✅ <b>Bandit is learning!</b> Each reply is now updating the router's model quality estimates.`;
|
||||
}
|
||||
}
|
||||
|
||||
function appendEmptyState() {
|
||||
const msgs = document.getElementById("messages");
|
||||
msgs.innerHTML = `<div class="empty-state">
|
||||
<h2>Pick a scenario to start</h2>
|
||||
<p>Choose one of the presets above or type your own message. The adaptive router picks the best model for each turn.</p>
|
||||
</div>`;
|
||||
}
|
||||
|
||||
function clearMessages() {
|
||||
document.getElementById("messages").innerHTML = "";
|
||||
}
|
||||
|
||||
function appendMessage(role, content, meta) {
|
||||
const msgs = document.getElementById("messages");
|
||||
const div = document.createElement("div");
|
||||
div.className = `msg ${role}`;
|
||||
div.dataset.role = role;
|
||||
|
||||
const avatar = document.createElement("div");
|
||||
avatar.className = "avatar";
|
||||
avatar.textContent = role === "user" ? "U" : "AI";
|
||||
|
||||
const wrap = document.createElement("div");
|
||||
|
||||
const bubble = document.createElement("div");
|
||||
bubble.className = "bubble";
|
||||
bubble.textContent = content; // plain text; code blocks show as preformatted
|
||||
renderBubble(bubble, content);
|
||||
|
||||
wrap.appendChild(bubble);
|
||||
|
||||
if (meta) {
|
||||
const metaEl = document.createElement("div");
|
||||
metaEl.className = "msg-meta";
|
||||
metaEl.innerHTML = meta;
|
||||
wrap.appendChild(metaEl);
|
||||
}
|
||||
|
||||
div.appendChild(avatar);
|
||||
div.appendChild(wrap);
|
||||
msgs.appendChild(div);
|
||||
msgs.scrollTop = msgs.scrollHeight;
|
||||
return bubble;
|
||||
}
|
||||
|
||||
function renderBubble(el, text) {
|
||||
// Minimal markdown: fenced code blocks and inline code.
|
||||
const escaped = text
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">");
|
||||
|
||||
const withBlocks = escaped.replace(
|
||||
/```(\w*)\n?([\s\S]*?)```/g,
|
||||
(_, lang, code) => `<pre><code>${code.trimEnd()}</code></pre>`
|
||||
);
|
||||
const withInline = withBlocks.replace(/`([^`]+)`/g, "<code>$1</code>");
|
||||
el.innerHTML = withInline;
|
||||
}
|
||||
|
||||
function appendThinking() {
|
||||
const msgs = document.getElementById("messages");
|
||||
const div = document.createElement("div");
|
||||
div.className = "msg assistant";
|
||||
div.id = "thinking-bubble";
|
||||
|
||||
const avatar = document.createElement("div");
|
||||
avatar.className = "avatar";
|
||||
avatar.textContent = "AI";
|
||||
|
||||
const bubble = document.createElement("div");
|
||||
bubble.className = "bubble";
|
||||
bubble.innerHTML = `<div class="thinking"><span></span><span></span><span></span></div>`;
|
||||
|
||||
div.appendChild(avatar);
|
||||
div.appendChild(bubble);
|
||||
msgs.appendChild(div);
|
||||
msgs.scrollTop = msgs.scrollHeight;
|
||||
return bubble;
|
||||
}
|
||||
|
||||
function removeThinking() {
|
||||
const el = document.getElementById("thinking-bubble");
|
||||
if (el) el.remove();
|
||||
}
|
||||
|
||||
// ---- chat send -------------------------------------------------------
|
||||
async function sendMessage(text) {
|
||||
if (STATE.sending || !text.trim()) return;
|
||||
STATE.sending = true;
|
||||
setSendEnabled(false);
|
||||
|
||||
// Add user message to history and UI
|
||||
STATE.messages.push({ role: "user", content: text });
|
||||
STATE.msgCount++;
|
||||
appendMessage("user", text);
|
||||
renderInfo();
|
||||
renderGateStatus();
|
||||
|
||||
const thinkingBubble = appendThinking();
|
||||
const t0 = Date.now();
|
||||
|
||||
try {
|
||||
const body = {
|
||||
model: STATE.router,
|
||||
messages: STATE.messages,
|
||||
stream: true,
|
||||
metadata: { litellm_session_id: STATE.sessionId },
|
||||
};
|
||||
|
||||
const resp = await fetch(`${STATE.proxyUrl}/v1/chat/completions`, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Authorization": `Bearer ${STATE.apiKey}`,
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
|
||||
if (!resp.ok) {
|
||||
const err = await resp.text().catch(() => `HTTP ${resp.status}`);
|
||||
removeThinking();
|
||||
appendMessage("assistant", `Error ${resp.status}: ${err}`);
|
||||
STATE.sending = false;
|
||||
setSendEnabled(true);
|
||||
return;
|
||||
}
|
||||
|
||||
// Read model from response header (requires proxy to expose via CORS).
|
||||
const chosenModel = resp.headers.get("x-litellm-adaptive-router-model");
|
||||
const reqType = resp.headers.get("x-litellm-request-type");
|
||||
STATE.lastModel = chosenModel || "check dashboard";
|
||||
STATE.lastReqType = reqType || "—";
|
||||
|
||||
// Stream the response.
|
||||
removeThinking();
|
||||
const assistantBubble = appendMessage("assistant", "");
|
||||
let fullContent = "";
|
||||
|
||||
const reader = resp.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
|
||||
// Process complete SSE lines.
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop(); // last fragment may be incomplete
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data: ")) continue;
|
||||
const raw = line.slice(6).trim();
|
||||
if (raw === "[DONE]") break;
|
||||
try {
|
||||
const chunk = JSON.parse(raw);
|
||||
const delta = chunk.choices?.[0]?.delta?.content || "";
|
||||
fullContent += delta;
|
||||
renderBubble(assistantBubble, fullContent);
|
||||
assistantBubble.closest(".messages")
|
||||
? (assistantBubble.closest(".messages").scrollTop = assistantBubble.closest(".messages").scrollHeight)
|
||||
: null;
|
||||
document.getElementById("messages").scrollTop = document.getElementById("messages").scrollHeight;
|
||||
} catch { /* incomplete JSON chunk, fine */ }
|
||||
}
|
||||
}
|
||||
|
||||
const latency = ((Date.now() - t0) / 1000).toFixed(2) + "s";
|
||||
document.getElementById("info-latency").textContent = latency;
|
||||
|
||||
// Add assistant turn meta
|
||||
const metaParts = [];
|
||||
if (chosenModel) metaParts.push(`<span class="msg-model">${chosenModel}</span>`);
|
||||
if (reqType) metaParts.push(`<span class="msg-type">${reqType}</span>`);
|
||||
metaParts.push(latency);
|
||||
if (metaParts.length) {
|
||||
const metaEl = document.createElement("div");
|
||||
metaEl.className = "msg-meta";
|
||||
metaEl.innerHTML = metaParts.join(" · ");
|
||||
assistantBubble.parentNode.appendChild(metaEl);
|
||||
}
|
||||
|
||||
STATE.messages.push({ role: "assistant", content: fullContent });
|
||||
STATE.msgCount++;
|
||||
renderInfo();
|
||||
renderGateStatus();
|
||||
} catch (e) {
|
||||
removeThinking();
|
||||
appendMessage("assistant", `Request failed: ${e.message}`);
|
||||
}
|
||||
|
||||
STATE.sending = false;
|
||||
setSendEnabled(true);
|
||||
}
|
||||
|
||||
// ---- input controls --------------------------------------------------
|
||||
function setSendEnabled(enabled) {
|
||||
const ta = document.getElementById("input");
|
||||
const btn = document.getElementById("send-btn");
|
||||
ta.disabled = !enabled || !STATE.connected;
|
||||
btn.disabled = !enabled || !STATE.connected;
|
||||
}
|
||||
|
||||
function setHint(text) {
|
||||
document.getElementById("input-hint").textContent = text;
|
||||
}
|
||||
|
||||
// ---- scenario selection ----------------------------------------------
|
||||
function activateScenario(id) {
|
||||
STATE.scenario = id;
|
||||
document.querySelectorAll(".sc-btn[data-id]").forEach(b => {
|
||||
b.classList.toggle("active", b.dataset.id === id);
|
||||
});
|
||||
|
||||
newSession();
|
||||
clearMessages();
|
||||
|
||||
const s = SCENARIOS[id];
|
||||
if (s.system) STATE.messages.push({ role: "system", content: s.system });
|
||||
|
||||
const ta = document.getElementById("input");
|
||||
ta.value = s.starter;
|
||||
ta.style.height = "auto";
|
||||
ta.style.height = Math.min(ta.scrollHeight, 160) + "px";
|
||||
ta.focus();
|
||||
|
||||
renderInfo();
|
||||
setHint(`Scenario: "${s.label}". Edit the starter if you like, then hit Send.`);
|
||||
}
|
||||
|
||||
// ---- connect ---------------------------------------------------------
|
||||
function connect() {
|
||||
const url = document.getElementById("proxy-url").value.trim().replace(/\/$/, "");
|
||||
const key = document.getElementById("api-key").value.trim();
|
||||
const router = document.getElementById("router").value.trim();
|
||||
|
||||
if (!url || !key || !router) {
|
||||
alert("Please fill in Proxy URL, API Key, and Router name.");
|
||||
return;
|
||||
}
|
||||
|
||||
STATE.proxyUrl = url;
|
||||
STATE.apiKey = key;
|
||||
STATE.router = router;
|
||||
STATE.connected = true;
|
||||
|
||||
ssSet("ar_proxy_url", url);
|
||||
ssSet("ar_api_key", key);
|
||||
ssSet("ar_router", router);
|
||||
|
||||
setConn(true);
|
||||
setSendEnabled(true);
|
||||
setHint("Pick a scenario above or type your own message.");
|
||||
newSession();
|
||||
appendEmptyState();
|
||||
renderInfo();
|
||||
}
|
||||
|
||||
// ---- textarea auto-resize & keyboard submit --------------------------
|
||||
document.getElementById("input").addEventListener("input", function () {
|
||||
this.style.height = "auto";
|
||||
this.style.height = Math.min(this.scrollHeight, 160) + "px";
|
||||
});
|
||||
|
||||
document.getElementById("input").addEventListener("keydown", function (e) {
|
||||
if (e.key === "Enter" && !e.shiftKey) {
|
||||
e.preventDefault();
|
||||
doSend();
|
||||
}
|
||||
});
|
||||
|
||||
function doSend() {
|
||||
const ta = document.getElementById("input");
|
||||
const text = ta.value.trim();
|
||||
if (!text) return;
|
||||
ta.value = "";
|
||||
ta.style.height = "auto";
|
||||
sendMessage(text);
|
||||
}
|
||||
|
||||
// ---- wiring ----------------------------------------------------------
|
||||
document.getElementById("connect-btn").addEventListener("click", connect);
|
||||
document.getElementById("send-btn").addEventListener("click", doSend);
|
||||
document.getElementById("new-chat-btn").addEventListener("click", () => {
|
||||
STATE.scenario = null;
|
||||
document.querySelectorAll(".sc-btn[data-id]").forEach(b => b.classList.remove("active"));
|
||||
newSession();
|
||||
clearMessages();
|
||||
appendEmptyState();
|
||||
const ta = document.getElementById("input");
|
||||
ta.value = "";
|
||||
ta.style.height = "auto";
|
||||
setHint("Type anything — the router will classify it and pick the best model.");
|
||||
renderInfo();
|
||||
});
|
||||
|
||||
document.querySelectorAll(".sc-btn[data-id]").forEach(btn => {
|
||||
btn.addEventListener("click", () => {
|
||||
if (!STATE.connected) {
|
||||
alert("Connect to the proxy first (fill in the form above and click Connect).");
|
||||
return;
|
||||
}
|
||||
activateScenario(btn.dataset.id);
|
||||
});
|
||||
});
|
||||
|
||||
// ---- restore session storage -----------------------------------------
|
||||
window.addEventListener("DOMContentLoaded", () => {
|
||||
const u = ssGet("ar_proxy_url"); if (u) document.getElementById("proxy-url").value = u;
|
||||
const k = ssGet("ar_api_key"); if (k) document.getElementById("api-key").value = k;
|
||||
const r = ssGet("ar_router"); if (r) document.getElementById("router").value = r;
|
||||
newSession();
|
||||
renderInfo();
|
||||
});
|
||||
</script>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
635
scripts/adaptive_router_demo/dashboard.html
Normal file
635
scripts/adaptive_router_demo/dashboard.html
Normal file
|
|
@ -0,0 +1,635 @@
|
|||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<title>Adaptive Router — Live</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #0b0f17;
|
||||
--panel: #131a26;
|
||||
--panel-2: #1b2433;
|
||||
--fg: #e7ecf3;
|
||||
--muted: #8a95a8;
|
||||
--accent: #5dd6a4;
|
||||
--accent-2: #6fb6ff;
|
||||
--warn: #f6b94d;
|
||||
--bad: #ff6b6b;
|
||||
--bar-bg: #233047;
|
||||
--border: #25324a;
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
body {
|
||||
margin: 0;
|
||||
font-family: -apple-system, BlinkMacSystemFont, "SF Pro Display",
|
||||
"Segoe UI", Roboto, Inter, sans-serif;
|
||||
background: var(--bg);
|
||||
color: var(--fg);
|
||||
font-size: 14px;
|
||||
line-height: 1.45;
|
||||
}
|
||||
header {
|
||||
padding: 18px 28px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 16px;
|
||||
background: var(--panel);
|
||||
}
|
||||
header h1 {
|
||||
margin: 0;
|
||||
font-size: 18px;
|
||||
font-weight: 600;
|
||||
letter-spacing: 0.2px;
|
||||
}
|
||||
header .dot {
|
||||
width: 8px; height: 8px; border-radius: 50%;
|
||||
background: var(--bad); display: inline-block; margin-right: 6px;
|
||||
}
|
||||
header .dot.ok { background: var(--accent); }
|
||||
header .status { color: var(--muted); font-size: 12px; }
|
||||
|
||||
.connect {
|
||||
padding: 16px 28px;
|
||||
display: flex; gap: 12px; align-items: center;
|
||||
background: var(--panel-2);
|
||||
border-bottom: 1px solid var(--border);
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
.connect input, .connect select {
|
||||
background: #0e1422;
|
||||
color: var(--fg);
|
||||
border: 1px solid var(--border);
|
||||
padding: 7px 10px;
|
||||
border-radius: 6px;
|
||||
font-size: 13px;
|
||||
font-family: inherit;
|
||||
}
|
||||
.connect input[type=text] { width: 240px; }
|
||||
.connect input[type=password] { width: 200px; }
|
||||
.connect input[type=number] { width: 70px; }
|
||||
.connect button {
|
||||
background: var(--accent-2);
|
||||
color: #0b0f17;
|
||||
border: none;
|
||||
padding: 7px 14px;
|
||||
border-radius: 6px;
|
||||
font-weight: 600;
|
||||
cursor: pointer;
|
||||
}
|
||||
.connect label { color: var(--muted); font-size: 12px; }
|
||||
|
||||
main {
|
||||
display: grid;
|
||||
grid-template-columns: 1fr 360px;
|
||||
gap: 20px;
|
||||
padding: 20px 28px 40px;
|
||||
max-width: 1400px;
|
||||
}
|
||||
@media (max-width: 1000px) { main { grid-template-columns: 1fr; } }
|
||||
|
||||
.panel {
|
||||
background: var(--panel);
|
||||
border: 1px solid var(--border);
|
||||
border-radius: 10px;
|
||||
padding: 18px 20px;
|
||||
}
|
||||
.panel h2 {
|
||||
margin: 0 0 14px;
|
||||
font-size: 13px;
|
||||
text-transform: uppercase;
|
||||
letter-spacing: 0.6px;
|
||||
color: var(--muted);
|
||||
font-weight: 600;
|
||||
}
|
||||
|
||||
.rt-group {
|
||||
margin-bottom: 18px;
|
||||
padding-bottom: 14px;
|
||||
border-bottom: 1px dashed var(--border);
|
||||
}
|
||||
.rt-group:last-child { border-bottom: none; margin-bottom: 0; }
|
||||
|
||||
.rt-title {
|
||||
font-weight: 600;
|
||||
font-size: 13px;
|
||||
color: var(--fg);
|
||||
margin-bottom: 8px;
|
||||
display: flex; justify-content: space-between;
|
||||
}
|
||||
.rt-title .meta { color: var(--muted); font-weight: 400; font-size: 12px; }
|
||||
|
||||
.row {
|
||||
display: grid;
|
||||
grid-template-columns: 80px 1fr 130px;
|
||||
align-items: center;
|
||||
gap: 10px;
|
||||
margin-bottom: 4px;
|
||||
font-size: 13px;
|
||||
}
|
||||
.row .name { color: var(--muted); font-family: ui-monospace, monospace; }
|
||||
.row .name.lead { color: var(--accent); font-weight: 600; }
|
||||
.row .num { color: var(--muted); font-family: ui-monospace, monospace;
|
||||
text-align: right; font-size: 12px; }
|
||||
|
||||
.bar { height: 14px; background: var(--bar-bg); border-radius: 4px;
|
||||
position: relative; overflow: hidden; }
|
||||
.bar > .fill {
|
||||
height: 100%;
|
||||
background: linear-gradient(90deg, var(--accent), var(--accent-2));
|
||||
border-radius: 4px;
|
||||
transition: width 0.4s ease;
|
||||
}
|
||||
.bar > .conf {
|
||||
position: absolute; top: 0; bottom: 0; width: 1px;
|
||||
background: rgba(255,255,255,0.4);
|
||||
}
|
||||
.bar > .conf.lo { background: rgba(255,255,255,0.5); }
|
||||
.bar > .conf.hi { background: rgba(255,255,255,0.5); }
|
||||
|
||||
.pick-pct {
|
||||
margin-top: 4px;
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
padding-left: 90px;
|
||||
}
|
||||
|
||||
.cost-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 12px; }
|
||||
.cost-card {
|
||||
background: var(--panel-2);
|
||||
border-radius: 8px;
|
||||
padding: 14px 16px;
|
||||
}
|
||||
.cost-card .label { color: var(--muted); font-size: 11px;
|
||||
text-transform: uppercase; letter-spacing: 0.5px; }
|
||||
.cost-card .value { font-size: 22px; font-weight: 600; margin-top: 4px;
|
||||
font-family: ui-monospace, monospace; }
|
||||
.cost-card .value.big { font-size: 28px; }
|
||||
.cost-card .sub { color: var(--muted); font-size: 12px; margin-top: 2px; }
|
||||
.cost-card.good { border: 1px solid rgba(93, 214, 164, 0.35); }
|
||||
.cost-card.warn { border: 1px solid rgba(246, 185, 77, 0.35); }
|
||||
.cost-card.bad { border: 1px solid rgba(255, 107, 107, 0.35); }
|
||||
.savings {
|
||||
margin-top: 12px;
|
||||
padding: 10px 14px;
|
||||
background: rgba(93, 214, 164, 0.08);
|
||||
border: 1px solid rgba(93, 214, 164, 0.3);
|
||||
border-radius: 8px;
|
||||
color: var(--accent);
|
||||
font-weight: 600;
|
||||
font-size: 14px;
|
||||
}
|
||||
.savings.warn {
|
||||
background: rgba(246, 185, 77, 0.08);
|
||||
border-color: rgba(246, 185, 77, 0.3);
|
||||
color: var(--warn);
|
||||
}
|
||||
.savings.bad {
|
||||
background: rgba(255, 107, 107, 0.08);
|
||||
border-color: rgba(255, 107, 107, 0.3);
|
||||
color: var(--bad);
|
||||
}
|
||||
.savings .verdict-sub {
|
||||
display: block;
|
||||
color: var(--muted);
|
||||
font-weight: 400;
|
||||
font-size: 12px;
|
||||
margin-top: 4px;
|
||||
}
|
||||
.panel-explainer {
|
||||
margin: -8px 0 14px;
|
||||
color: var(--muted);
|
||||
font-size: 12px;
|
||||
line-height: 1.5;
|
||||
padding: 10px 12px;
|
||||
background: var(--panel-2);
|
||||
border-radius: 6px;
|
||||
border-left: 3px solid var(--accent-2);
|
||||
}
|
||||
.panel-explainer b { color: var(--fg); font-weight: 600; }
|
||||
|
||||
.activity {
|
||||
max-height: 360px; overflow-y: auto;
|
||||
font-family: ui-monospace, monospace;
|
||||
font-size: 12px;
|
||||
}
|
||||
.activity-row {
|
||||
padding: 6px 8px;
|
||||
border-bottom: 1px solid var(--border);
|
||||
color: var(--muted);
|
||||
display: grid;
|
||||
grid-template-columns: 70px 1fr;
|
||||
gap: 8px;
|
||||
}
|
||||
.activity-row .ts { color: #4f5d77; }
|
||||
.activity-row .alpha { color: var(--accent); }
|
||||
.activity-row .beta { color: var(--bad); }
|
||||
|
||||
.queue {
|
||||
display: flex; gap: 16px; flex-wrap: wrap;
|
||||
font-size: 12px; color: var(--muted);
|
||||
margin-top: 8px;
|
||||
}
|
||||
.queue span b { color: var(--fg); font-weight: 600; }
|
||||
|
||||
.empty {
|
||||
color: var(--muted); font-style: italic;
|
||||
text-align: center; padding: 20px;
|
||||
}
|
||||
|
||||
.pill {
|
||||
background: var(--panel-2);
|
||||
color: var(--muted);
|
||||
padding: 3px 8px;
|
||||
border-radius: 12px;
|
||||
font-size: 11px;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<header>
|
||||
<h1>⚡ Adaptive Router — Live</h1>
|
||||
<span class="status"><span id="conn-dot" class="dot"></span><span id="conn-label">Disconnected</span></span>
|
||||
<span id="poll-info" class="status"></span>
|
||||
</header>
|
||||
|
||||
<div class="connect">
|
||||
<label>Proxy URL <input id="proxy-url" type="text" value="http://localhost:4000" /></label>
|
||||
<label>Master Key <input id="api-key" type="password" placeholder="sk-1234" /></label>
|
||||
<label>Avg tokens/req <input id="avg-tokens" type="number" value="500" min="1" /></label>
|
||||
<label>Poll ms <input id="poll-ms" type="number" value="500" min="100" /></label>
|
||||
<button id="connect-btn">Connect</button>
|
||||
<select id="router-select" style="display:none;"></select>
|
||||
</div>
|
||||
|
||||
<main>
|
||||
<section class="panel" id="bandit-panel">
|
||||
<h2>How well each model performs, by request type</h2>
|
||||
<div class="panel-explainer">
|
||||
Each bar shows the <b>fraction of recent feedback that was positive</b>
|
||||
for that model on that kind of request. Wider = better. The number
|
||||
next to it (<b>"N signals"</b>) is how much real feedback the bar is
|
||||
based on — more signals means the router is more confident.
|
||||
It picks higher-quality bars first, with cost as a tiebreaker.
|
||||
</div>
|
||||
<div id="cells" class="empty">Connect to see live bandit state.</div>
|
||||
</section>
|
||||
|
||||
<aside style="display:flex; flex-direction:column; gap:20px;">
|
||||
<section class="panel">
|
||||
<h2>Are the savings worth it?</h2>
|
||||
<div class="panel-explainer">
|
||||
<b>Cost saved</b> is what you spent vs. always picking the most
|
||||
expensive model. <b>Quality kept</b> is the average quality of
|
||||
the model that was actually picked, divided by the average
|
||||
quality of the best-known model for each request type.
|
||||
<i>If quality kept stays high while cost saved is high, you're
|
||||
winning. If quality drops fast, you're saving money but making
|
||||
users mad.</i>
|
||||
</div>
|
||||
<div class="cost-grid">
|
||||
<div class="cost-card good">
|
||||
<div class="label">💰 Cost saved</div>
|
||||
<div class="value big" id="metric-cost-pct">—</div>
|
||||
<div class="sub" id="metric-cost-sub">no traffic yet</div>
|
||||
</div>
|
||||
<div class="cost-card good">
|
||||
<div class="label">⭐ Quality kept</div>
|
||||
<div class="value big" id="metric-quality-pct">—</div>
|
||||
<div class="sub" id="metric-quality-sub">vs best-known model</div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="savings" id="verdict">Send some traffic to see how the router is balancing cost and quality.</div>
|
||||
<div class="queue" id="queue-info"></div>
|
||||
</section>
|
||||
|
||||
<section class="panel">
|
||||
<h2>Activity (last 30)</h2>
|
||||
<div id="activity" class="activity empty">Waiting for signals…</div>
|
||||
</section>
|
||||
</aside>
|
||||
</main>
|
||||
|
||||
<script>
|
||||
// ---- helpers ---------------------------------------------------------
|
||||
const REQ_TYPE_ORDER = [
|
||||
"code_generation",
|
||||
"code_understanding",
|
||||
"technical_design",
|
||||
"analytical_reasoning",
|
||||
"writing",
|
||||
"factual_lookup",
|
||||
"general",
|
||||
];
|
||||
|
||||
function fmtPct(n) { return (n * 100).toFixed(0) + "%"; }
|
||||
function fmtUSD(n) { return "$" + n.toFixed(4); }
|
||||
function nowTS() {
|
||||
const d = new Date();
|
||||
return d.toTimeString().slice(0, 8);
|
||||
}
|
||||
function ssGet(k) { try { return sessionStorage.getItem(k) || ""; } catch { return ""; } }
|
||||
function ssSet(k, v) { try { sessionStorage.setItem(k, v); } catch {} }
|
||||
|
||||
// ---- state -----------------------------------------------------------
|
||||
const STATE = {
|
||||
proxyUrl: "",
|
||||
apiKey: "",
|
||||
pollMs: 500,
|
||||
avgTokens: 500,
|
||||
timer: null,
|
||||
routers: [], // last snapshot list
|
||||
selectedRouter: null, // name
|
||||
prevCells: new Map(), // (router, rt, model) -> {alpha,beta,samples}
|
||||
costAdaptive: 0,
|
||||
costBaseline: 0,
|
||||
totalRequests: 0,
|
||||
activity: [], // [{ts, msg, kind}]
|
||||
};
|
||||
|
||||
// ---- rendering -------------------------------------------------------
|
||||
function renderRouters(routers) {
|
||||
const sel = document.getElementById("router-select");
|
||||
if (routers.length <= 1) {
|
||||
sel.style.display = "none";
|
||||
} else {
|
||||
sel.style.display = "";
|
||||
if (sel.options.length !== routers.length) {
|
||||
sel.innerHTML = "";
|
||||
for (const r of routers) {
|
||||
const opt = document.createElement("option");
|
||||
opt.value = r.router_name; opt.textContent = r.router_name;
|
||||
sel.appendChild(opt);
|
||||
}
|
||||
sel.value = STATE.selectedRouter || routers[0].router_name;
|
||||
}
|
||||
}
|
||||
if (!STATE.selectedRouter) STATE.selectedRouter = routers[0].router_name;
|
||||
}
|
||||
|
||||
function pickShare(rows) {
|
||||
// Approximate prob each model wins a Thompson-sample draw against the others.
|
||||
// Simple proxy: softmax over quality_mean with temperature=0.05.
|
||||
if (rows.length === 0) return {};
|
||||
const T = 0.05;
|
||||
const expv = rows.map(r => Math.exp(r.quality_mean / T));
|
||||
const sum = expv.reduce((a, b) => a + b, 0);
|
||||
const out = {};
|
||||
rows.forEach((r, i) => out[r.model] = expv[i] / sum);
|
||||
return out;
|
||||
}
|
||||
|
||||
function renderCells(router) {
|
||||
const container = document.getElementById("cells");
|
||||
container.classList.remove("empty");
|
||||
const byType = new Map();
|
||||
for (const c of router.cells) {
|
||||
if (!byType.has(c.request_type)) byType.set(c.request_type, []);
|
||||
byType.get(c.request_type).push(c);
|
||||
}
|
||||
const order = REQ_TYPE_ORDER.filter(t => byType.has(t));
|
||||
for (const t of byType.keys()) if (!order.includes(t)) order.push(t);
|
||||
|
||||
let html = "";
|
||||
for (const rt of order) {
|
||||
const rows = byType.get(rt).sort((a, b) => b.quality_mean - a.quality_mean);
|
||||
const shares = pickShare(rows);
|
||||
const lead = rows[0];
|
||||
html += `<div class="rt-group">`;
|
||||
html += `<div class="rt-title"><span>${rt}</span>`;
|
||||
html += `<span class="meta">${rows.reduce((s, r) => s + (r.samples - 10), 0)} learning signals</span>`;
|
||||
html += `</div>`;
|
||||
for (const r of rows) {
|
||||
const pct = fmtPct(r.quality_mean);
|
||||
const share = fmtPct(shares[r.model] || 0);
|
||||
const observed = Math.max(0, r.samples - 10); // strip cold-start prior mass
|
||||
const isLead = r.model === lead.model;
|
||||
const sigLabel = observed === 0 ? "no signals yet" : `${observed.toFixed(0)} signals`;
|
||||
// Tooltip exposes raw Beta(α,β) for power users.
|
||||
const tip = `Beta(α=${r.alpha.toFixed(1)}, β=${r.beta.toFixed(1)}) — ` +
|
||||
`started at α=5,β=5 (cold-start prior), so the bar reflects ` +
|
||||
`${observed.toFixed(0)} real feedback signals so far.`;
|
||||
html += `<div class="row" title="${tip}">`;
|
||||
html += `<div class="name ${isLead ? 'lead' : ''}">${r.model}</div>`;
|
||||
html += `<div class="bar"><div class="fill" style="width:${(r.quality_mean*100).toFixed(1)}%"></div></div>`;
|
||||
html += `<div class="num">${pct} good · ${sigLabel}</div>`;
|
||||
html += `</div>`;
|
||||
html += `<div class="pick-pct">→ ${share} of picks (Thompson softmax estimate)</div>`;
|
||||
}
|
||||
html += `</div>`;
|
||||
}
|
||||
container.innerHTML = html;
|
||||
}
|
||||
|
||||
function computeQualityKept(router) {
|
||||
// For each request type: pick_count_per_cell × quality_mean_per_cell
|
||||
// summed and divided by total picks gives "average quality delivered".
|
||||
// Compare against best-cell quality per type (weighted by picks in that type).
|
||||
const byType = new Map();
|
||||
for (const c of router.cells) {
|
||||
if (!byType.has(c.request_type)) byType.set(c.request_type, []);
|
||||
byType.get(c.request_type).push(c);
|
||||
}
|
||||
let totalPicks = 0, weightedDelivered = 0, weightedBest = 0;
|
||||
for (const cells of byType.values()) {
|
||||
const bestQ = Math.max(...cells.map(c => c.quality_mean));
|
||||
for (const c of cells) {
|
||||
const picks = Math.max(0, c.samples - 10);
|
||||
if (picks === 0) continue;
|
||||
totalPicks += picks;
|
||||
weightedDelivered += picks * c.quality_mean;
|
||||
weightedBest += picks * bestQ;
|
||||
}
|
||||
}
|
||||
if (totalPicks === 0 || weightedBest === 0) return null;
|
||||
return {
|
||||
delivered: weightedDelivered / totalPicks,
|
||||
best: weightedBest / totalPicks,
|
||||
keptPct: weightedDelivered / weightedBest,
|
||||
totalPicks,
|
||||
};
|
||||
}
|
||||
|
||||
function renderTradeoff(router) {
|
||||
const costEl = document.getElementById("metric-cost-pct");
|
||||
const costSub = document.getElementById("metric-cost-sub");
|
||||
const qualEl = document.getElementById("metric-quality-pct");
|
||||
const qualSub = document.getElementById("metric-quality-sub");
|
||||
const verdict = document.getElementById("verdict");
|
||||
|
||||
// ---- Cost side ---------------------------------------------------
|
||||
let costSavedPct = null;
|
||||
if (STATE.costBaseline > 0) {
|
||||
costSavedPct = 1 - STATE.costAdaptive / STATE.costBaseline;
|
||||
costEl.textContent = (costSavedPct * 100).toFixed(0) + "%";
|
||||
costSub.textContent = `${fmtUSD(STATE.costAdaptive)} spent vs ${fmtUSD(STATE.costBaseline)} baseline`;
|
||||
} else {
|
||||
costEl.textContent = "—";
|
||||
costSub.textContent = "no traffic yet";
|
||||
}
|
||||
|
||||
// ---- Quality side ------------------------------------------------
|
||||
const q = computeQualityKept(router);
|
||||
if (q) {
|
||||
qualEl.textContent = (q.keptPct * 100).toFixed(0) + "%";
|
||||
qualSub.textContent =
|
||||
`delivered ${(q.delivered*100).toFixed(0)}% vs best-known ${(q.best*100).toFixed(0)}%`;
|
||||
} else {
|
||||
qualEl.textContent = "—";
|
||||
qualSub.textContent = "vs best-known model";
|
||||
}
|
||||
|
||||
// ---- Color-code the cards ----------------------------------------
|
||||
const costCard = costEl.closest(".cost-card");
|
||||
const qualCard = qualEl.closest(".cost-card");
|
||||
costCard.className = "cost-card " + (costSavedPct === null ? "good"
|
||||
: costSavedPct >= 0.30 ? "good"
|
||||
: costSavedPct >= 0.05 ? "warn" : "bad");
|
||||
qualCard.className = "cost-card " + (!q ? "good"
|
||||
: q.keptPct >= 0.90 ? "good"
|
||||
: q.keptPct >= 0.75 ? "warn" : "bad");
|
||||
|
||||
// ---- Verdict line ------------------------------------------------
|
||||
if (costSavedPct === null || !q) {
|
||||
verdict.className = "savings";
|
||||
verdict.textContent = "Send some traffic to see how the router is balancing cost and quality.";
|
||||
return;
|
||||
}
|
||||
const savedTxt = costSavedPct >= 0
|
||||
? `${(costSavedPct*100).toFixed(0)}% cheaper`
|
||||
: `${((-costSavedPct)*100).toFixed(0)}% MORE expensive (still exploring)`;
|
||||
const qualLost = (1 - q.keptPct) * 100;
|
||||
let line, cls;
|
||||
if (q.keptPct >= 0.95 && costSavedPct >= 0.30) {
|
||||
cls = "savings"; line = `✅ Big win: ${savedTxt}, lost only ${qualLost.toFixed(0)}% quality.`;
|
||||
} else if (q.keptPct >= 0.85 && costSavedPct >= 0.10) {
|
||||
cls = "savings"; line = `✅ Good trade: ${savedTxt}, gave up ${qualLost.toFixed(0)}% quality.`;
|
||||
} else if (q.keptPct >= 0.75) {
|
||||
cls = "savings warn"; line = `⚠️ Mixed: ${savedTxt}, but ${qualLost.toFixed(0)}% quality lost. Consider raising the quality weight.`;
|
||||
} else {
|
||||
cls = "savings bad"; line = `❌ Saving money, hurting users: ${savedTxt} but ${qualLost.toFixed(0)}% quality lost. Raise quality weight in the router config.`;
|
||||
}
|
||||
verdict.className = cls;
|
||||
verdict.innerHTML = line +
|
||||
`<span class="verdict-sub">Based on ${q.totalPicks} feedback signals across ${STATE.totalRequests} routed requests.</span>`;
|
||||
}
|
||||
|
||||
function renderQueue(router) {
|
||||
const q = router.queue || {};
|
||||
document.getElementById("queue-info").innerHTML =
|
||||
`<span>state pending: <b>${q.state_pending ?? 0}</b></span>` +
|
||||
`<span>session pending: <b>${q.session_pending ?? 0}</b></span>` +
|
||||
`<span>sticky live: <b>${router.sticky_sessions_live ?? 0}</b></span>` +
|
||||
`<span>weights: q=<b>${router.weights?.quality ?? "?"}</b> c=<b>${router.weights?.cost ?? "?"}</b></span>`;
|
||||
}
|
||||
|
||||
function renderActivity() {
|
||||
const el = document.getElementById("activity");
|
||||
if (STATE.activity.length === 0) {
|
||||
el.classList.add("empty");
|
||||
el.textContent = "Waiting for signals…";
|
||||
return;
|
||||
}
|
||||
el.classList.remove("empty");
|
||||
el.innerHTML = STATE.activity.map(a =>
|
||||
`<div class="activity-row"><span class="ts">${a.ts}</span><span>${a.msg}</span></div>`
|
||||
).join("");
|
||||
}
|
||||
|
||||
// ---- diff & cost accounting -----------------------------------------
|
||||
function processDiff(router, costsByModel) {
|
||||
const maxCost = Math.max(0, ...Object.values(costsByModel));
|
||||
for (const c of router.cells) {
|
||||
const key = `${router.router_name}|${c.request_type}|${c.model}`;
|
||||
const prev = STATE.prevCells.get(key);
|
||||
if (prev) {
|
||||
const dA = c.alpha - prev.alpha;
|
||||
const dB = c.beta - prev.beta;
|
||||
const dPicks = (c.samples - 10) - (prev.samples - 10);
|
||||
if (dA > 0.001 || dB > 0.001) {
|
||||
const tag = dA > dB
|
||||
? `<span class="alpha">+${dA.toFixed(0)} 👍</span>`
|
||||
: `<span class="beta">+${dB.toFixed(0)} 👎</span>`;
|
||||
const qNow = (c.alpha / (c.alpha + c.beta) * 100).toFixed(0);
|
||||
STATE.activity.unshift({
|
||||
ts: nowTS(),
|
||||
msg: `${c.request_type} → <b>${c.model}</b> ${tag} (now ${qNow}% good)`,
|
||||
});
|
||||
STATE.activity = STATE.activity.slice(0, 30);
|
||||
}
|
||||
if (dPicks > 0) {
|
||||
const cost = costsByModel[c.model] || 0;
|
||||
STATE.costAdaptive += dPicks * cost * STATE.avgTokens;
|
||||
STATE.costBaseline += dPicks * maxCost * STATE.avgTokens;
|
||||
STATE.totalRequests += dPicks;
|
||||
}
|
||||
}
|
||||
STATE.prevCells.set(key, {
|
||||
alpha: c.alpha, beta: c.beta, samples: c.samples
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// ---- polling ---------------------------------------------------------
|
||||
async function pollOnce() {
|
||||
try {
|
||||
const r = await fetch(`${STATE.proxyUrl}/adaptive_router/state`, {
|
||||
headers: { "Authorization": `Bearer ${STATE.apiKey}` },
|
||||
});
|
||||
if (!r.ok) {
|
||||
setConn(false, `HTTP ${r.status}`);
|
||||
return;
|
||||
}
|
||||
const data = await r.json();
|
||||
setConn(true, `Polling every ${STATE.pollMs}ms`);
|
||||
STATE.routers = data.routers || [];
|
||||
renderRouters(STATE.routers);
|
||||
const router = STATE.routers.find(r => r.router_name === STATE.selectedRouter)
|
||||
|| STATE.routers[0];
|
||||
if (!router) return;
|
||||
processDiff(router, router.model_costs || {});
|
||||
renderCells(router);
|
||||
renderQueue(router);
|
||||
renderTradeoff(router);
|
||||
renderActivity();
|
||||
} catch (e) {
|
||||
setConn(false, e.message);
|
||||
}
|
||||
}
|
||||
|
||||
function setConn(ok, msg) {
|
||||
document.getElementById("conn-dot").className = "dot" + (ok ? " ok" : "");
|
||||
document.getElementById("conn-label").textContent = ok ? "Connected" : "Disconnected";
|
||||
document.getElementById("poll-info").textContent = msg || "";
|
||||
}
|
||||
|
||||
function startPolling() {
|
||||
if (STATE.timer) clearInterval(STATE.timer);
|
||||
pollOnce();
|
||||
STATE.timer = setInterval(pollOnce, STATE.pollMs);
|
||||
}
|
||||
|
||||
// ---- wiring ----------------------------------------------------------
|
||||
document.getElementById("connect-btn").addEventListener("click", () => {
|
||||
STATE.proxyUrl = document.getElementById("proxy-url").value.trim().replace(/\/$/, "");
|
||||
STATE.apiKey = document.getElementById("api-key").value.trim();
|
||||
STATE.pollMs = parseInt(document.getElementById("poll-ms").value, 10) || 500;
|
||||
STATE.avgTokens = parseInt(document.getElementById("avg-tokens").value, 10) || 500;
|
||||
ssSet("ar_proxy_url", STATE.proxyUrl);
|
||||
ssSet("ar_api_key", STATE.apiKey);
|
||||
startPolling();
|
||||
});
|
||||
|
||||
document.getElementById("router-select").addEventListener("change", (e) => {
|
||||
STATE.selectedRouter = e.target.value;
|
||||
STATE.prevCells.clear();
|
||||
});
|
||||
|
||||
window.addEventListener("DOMContentLoaded", () => {
|
||||
const u = ssGet("ar_proxy_url"); if (u) document.getElementById("proxy-url").value = u;
|
||||
const k = ssGet("ar_api_key"); if (k) document.getElementById("api-key").value = k;
|
||||
});
|
||||
</script>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
271
scripts/adaptive_router_demo/eval.py
Normal file
271
scripts/adaptive_router_demo/eval.py
Normal file
|
|
@ -0,0 +1,271 @@
|
|||
# ruff: noqa: T201
|
||||
"""
|
||||
Adaptive router evaluator — LLM-as-judge harness.
|
||||
|
||||
For each test case:
|
||||
1. Sends the prompt to the adaptive router.
|
||||
2. Reads which model was picked (x-litellm-adaptive-router-model header).
|
||||
3. Asks the judge model whether the response meets the ideal criteria.
|
||||
4. Prints PASS or FAIL with one line of reasoning.
|
||||
|
||||
Run:
|
||||
uv run python scripts/adaptive_router_demo/eval.py \
|
||||
--proxy-url http://localhost:4000 \
|
||||
--api-key sk-1234 \
|
||||
--router smart-cheap-router \
|
||||
--judge-model smart
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import httpx
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test cases
|
||||
# ---------------------------------------------------------------------------
|
||||
@dataclass
|
||||
class EvalCase:
|
||||
category: str
|
||||
prompt: str
|
||||
ideal: str # criteria the judge checks the response against
|
||||
|
||||
|
||||
EVAL_CASES: List[EvalCase] = [
|
||||
# code_generation
|
||||
EvalCase(
|
||||
category="code_generation",
|
||||
prompt="Write a Python function that flattens a nested list of arbitrary depth.",
|
||||
ideal=(
|
||||
"A Python function (def flatten(...)) that accepts a list which may "
|
||||
"contain nested lists to arbitrary depth and returns a single flat list "
|
||||
"with all elements in order. Must handle at least two levels of nesting."
|
||||
),
|
||||
),
|
||||
EvalCase(
|
||||
category="code_generation",
|
||||
prompt="Write a Python decorator that retries a function up to 3 times on exception.",
|
||||
ideal=(
|
||||
"A Python decorator that wraps a callable, catches exceptions, and "
|
||||
"retries the call up to 3 times before re-raising. Should use functools.wraps "
|
||||
"or equivalent to preserve the wrapped function's metadata."
|
||||
),
|
||||
),
|
||||
EvalCase(
|
||||
category="code_generation",
|
||||
prompt="Write a SQL query that returns the top 5 customers by total order value.",
|
||||
ideal=(
|
||||
"A valid SQL SELECT query that JOINs an orders or order_items table with a "
|
||||
"customers table, groups by customer, sums order value, orders descending, "
|
||||
"and limits to 5 rows."
|
||||
),
|
||||
),
|
||||
# factual_lookup
|
||||
EvalCase(
|
||||
category="factual_lookup",
|
||||
prompt="What is the capital of New Zealand?",
|
||||
ideal="The answer must state Wellington as the capital of New Zealand.",
|
||||
),
|
||||
EvalCase(
|
||||
category="factual_lookup",
|
||||
prompt="In what year did World War II end?",
|
||||
ideal="The answer must state 1945 as the year World War II ended.",
|
||||
),
|
||||
EvalCase(
|
||||
category="factual_lookup",
|
||||
prompt="What is the chemical symbol for gold?",
|
||||
ideal="The answer must include 'Au' as the chemical symbol for gold.",
|
||||
),
|
||||
# writing
|
||||
EvalCase(
|
||||
category="writing",
|
||||
prompt=(
|
||||
"Write a short, polite email declining a meeting request because of "
|
||||
"a scheduling conflict."
|
||||
),
|
||||
ideal=(
|
||||
"A professional email that: (1) thanks the sender for the invitation, "
|
||||
"(2) clearly declines, (3) mentions a scheduling conflict as the reason, "
|
||||
"and (4) offers to reschedule or an alternative. Tone must be polite."
|
||||
),
|
||||
),
|
||||
EvalCase(
|
||||
category="writing",
|
||||
prompt="Write a one-paragraph product description for noise-cancelling headphones.",
|
||||
ideal=(
|
||||
"A marketing paragraph for noise-cancelling headphones that mentions "
|
||||
"noise cancellation as a feature, highlights at least one other benefit "
|
||||
"(comfort, audio quality, battery life, or similar), and ends with a "
|
||||
"persuasive call to action or closing statement."
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
# Matches the satisfaction regex in signals.py (_SATISFACTION_PATTERNS).
|
||||
SATISFY_FOLLOWUP = "great, thanks!"
|
||||
NEUTRAL_FOLLOWUP = "ok, noted"
|
||||
FAB_ASSISTANT = "Got it. Working on that now."
|
||||
|
||||
JUDGE_SYSTEM = (
|
||||
"You are a strict but fair evaluator. Your job is to decide whether a model "
|
||||
"response meets the stated requirements. Reply with exactly two lines:\n"
|
||||
"Line 1: PASS or FAIL\n"
|
||||
"Line 2: One sentence of reasoning (≤ 25 words)."
|
||||
)
|
||||
|
||||
|
||||
def _judge_user(prompt: str, ideal: str, actual: str) -> str:
|
||||
return (
|
||||
f"Question sent to model:\n{prompt}\n\n"
|
||||
f"Requirements the response must meet:\n{ideal}\n\n"
|
||||
f"Actual model response:\n{actual}\n\n"
|
||||
"Does the response meet the requirements? Reply PASS or FAIL."
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
async def _chat(
|
||||
client: httpx.AsyncClient,
|
||||
proxy_url: str,
|
||||
api_key: str,
|
||||
model: str,
|
||||
messages: List[Dict[str, str]],
|
||||
session_id: Optional[str] = None,
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Returns (response_text, chosen_model_header).
|
||||
chosen_model_header is empty for non-router calls.
|
||||
"""
|
||||
body: Dict = {"model": model, "messages": messages}
|
||||
if session_id:
|
||||
body["metadata"] = {"litellm_session_id": session_id}
|
||||
|
||||
resp = await client.post(
|
||||
f"{proxy_url}/v1/chat/completions",
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=60.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
text = data["choices"][0]["message"]["content"]
|
||||
chosen = resp.headers.get("x-litellm-adaptive-router-model", "")
|
||||
return text, chosen
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Evaluation loop
|
||||
# ---------------------------------------------------------------------------
|
||||
async def evaluate(
|
||||
proxy_url: str,
|
||||
api_key: str,
|
||||
router: str,
|
||||
judge_model: str,
|
||||
) -> None:
|
||||
passed = 0
|
||||
failed = 0
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
for i, case in enumerate(EVAL_CASES, 1):
|
||||
print(f"\n[{i}/{len(EVAL_CASES)}] category={case.category}")
|
||||
print(f" prompt : {case.prompt[:80]}{'…' if len(case.prompt) > 80 else ''}")
|
||||
|
||||
session_id = f"eval-{uuid.uuid4()}"
|
||||
|
||||
# Round 1: single-turn real request — get the actual LLM response to judge.
|
||||
try:
|
||||
response, chosen = await _chat(
|
||||
client, proxy_url, api_key, router,
|
||||
[{"role": "user", "content": case.prompt}],
|
||||
session_id=session_id,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" ERROR calling router: {exc}", file=sys.stderr)
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
print(f" model : {chosen or router}")
|
||||
print(f" response : {response[:120].replace(chr(10), ' ')}{'…' if len(response) > 120 else ''}")
|
||||
|
||||
# Judge the real response.
|
||||
judge_msgs = [
|
||||
{"role": "system", "content": JUDGE_SYSTEM},
|
||||
{"role": "user", "content": _judge_user(case.prompt, case.ideal, response)},
|
||||
]
|
||||
try:
|
||||
verdict, _ = await _chat(
|
||||
client, proxy_url, api_key, judge_model, judge_msgs,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" ERROR calling judge: {exc}", file=sys.stderr)
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
# Parse verdict — first non-empty line should be PASS or FAIL.
|
||||
lines = [ln.strip() for ln in verdict.splitlines() if ln.strip()]
|
||||
first = lines[0].upper() if lines else ""
|
||||
reason = lines[1] if len(lines) > 1 else ""
|
||||
is_pass = "PASS" in first
|
||||
|
||||
if is_pass:
|
||||
passed += 1
|
||||
print(f" verdict : \033[32mPASS\033[0m {reason}")
|
||||
else:
|
||||
failed += 1
|
||||
print(f" verdict : \033[31mFAIL\033[0m {reason}")
|
||||
|
||||
# Round 2: 5-message conversation on the same session_id so the bandit fires.
|
||||
# On PASS → satisfaction follow-up (+alpha). On FAIL → neutral (no signal).
|
||||
follow_up = SATISFY_FOLLOWUP if is_pass else NEUTRAL_FOLLOWUP
|
||||
bandit_msgs = [
|
||||
{"role": "user", "content": case.prompt},
|
||||
{"role": "assistant", "content": response},
|
||||
{"role": "user", "content": "ok continue"},
|
||||
{"role": "assistant", "content": FAB_ASSISTANT},
|
||||
{"role": "user", "content": follow_up},
|
||||
]
|
||||
try:
|
||||
await _chat(
|
||||
client, proxy_url, api_key, router, bandit_msgs,
|
||||
session_id=session_id,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
print(f" WARNING: bandit update failed: {exc}", file=sys.stderr)
|
||||
|
||||
total = passed + failed
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Results: {passed}/{total} passed ({failed} failed)")
|
||||
if passed == total:
|
||||
print("All test cases passed — the adaptive router is working well!")
|
||||
elif passed >= total * 0.8:
|
||||
print("Most test cases passed — minor issues to investigate.")
|
||||
else:
|
||||
print("Significant failures — check router config and model availability.")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser(description="Evaluate the adaptive router with LLM-as-judge.")
|
||||
ap.add_argument("--proxy-url", default="http://localhost:4000")
|
||||
ap.add_argument("--api-key", required=True, help="proxy API key")
|
||||
ap.add_argument("--router", default="smart-cheap-router", help="adaptive router model name")
|
||||
ap.add_argument("--judge-model", default="smart", help="model name for the judge (via proxy)")
|
||||
args = ap.parse_args()
|
||||
|
||||
asyncio.run(evaluate(args.proxy_url, args.api_key, args.router, args.judge_model))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
227
scripts/adaptive_router_demo/traffic.py
Normal file
227
scripts/adaptive_router_demo/traffic.py
Normal file
|
|
@ -0,0 +1,227 @@
|
|||
"""
|
||||
Synthetic traffic generator for the adaptive_router demo dashboard.
|
||||
|
||||
What it does:
|
||||
- Sends labeled multi-turn chat requests to the proxy's adaptive router.
|
||||
- For each turn, peeks at the `x-litellm-adaptive-router-model` response
|
||||
header to learn which underlying model was picked.
|
||||
- Draws a Bernoulli outcome from a hard-coded ORACLE table that says
|
||||
"model M succeeds at request type T with probability p".
|
||||
- Sends a final follow-up turn whose user message is engineered to
|
||||
BOTH classify into the same RequestType AND match the
|
||||
satisfaction regex on success (so the bandit's `(type, model)` cell
|
||||
gets +alpha). On failure we send a neutral follow-up so no signal
|
||||
fires — over time, models the oracle favors accumulate alpha faster.
|
||||
|
||||
Why this shape:
|
||||
- The post-call hook gates signal recording on len(messages) >= 4.
|
||||
A single 5-message request passes the gate in one round-trip, which
|
||||
keeps the demo cheap.
|
||||
- Mock responses (`mock_response=...`) skip the real LLM call but still
|
||||
flow through routing + post-call hooks, so no API keys / no spend.
|
||||
|
||||
Run:
|
||||
uv run python scripts/adaptive_router_demo/traffic.py \\
|
||||
--proxy-url http://localhost:4000 \\
|
||||
--api-key sk-1234 \\
|
||||
--router smart-cheap-router \\
|
||||
--rounds 100 \\
|
||||
--rate 0.5
|
||||
|
||||
Open `dashboard.html` in a browser alongside this and watch the bars move.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import random
|
||||
import sys
|
||||
import uuid
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import httpx
|
||||
|
||||
# ---- prompts (paired with the RequestType the classifier will assign) ----
|
||||
# Each prompt is engineered to (a) classify into the listed type and (b) make
|
||||
# sense as a user request. Keep prompts short to limit token cost.
|
||||
PROMPTS: Dict[str, List[str]] = {
|
||||
"code_generation": [
|
||||
"Write a Python function that flattens a nested list",
|
||||
"Create a TypeScript function that debounces another function",
|
||||
"Build a Rust function that parses a CSV string",
|
||||
"Generate a SQL function that returns running totals",
|
||||
],
|
||||
"factual_lookup": [
|
||||
"What is the capital of New Zealand?",
|
||||
"When was the Treaty of Westphalia signed?",
|
||||
"Who is the current Secretary General of the UN?",
|
||||
"Where is Mount Kilimanjaro located?",
|
||||
],
|
||||
"writing": [
|
||||
"Write an email declining a meeting politely",
|
||||
"Draft a paragraph introducing a product launch",
|
||||
"Compose a short blog post about morning routines",
|
||||
"Rewrite this sentence to be more concise: ...",
|
||||
],
|
||||
}
|
||||
|
||||
# Engineered satisfaction follow-ups — each one is designed to:
|
||||
# (1) match the satisfaction regex (thanks/great/works/perfect/etc.), AND
|
||||
# (2) re-classify into the SAME RequestType as the first prompt
|
||||
# so that signals attribute to the right (type, model) bandit cell.
|
||||
SATISFY: Dict[str, str] = {
|
||||
"code_generation": "thanks, that works! now write me a python function that does the inverse",
|
||||
"factual_lookup": "perfect, thanks! who is the current prime minister?",
|
||||
"writing": "great, thanks! now write a follow-up email confirming attendance",
|
||||
}
|
||||
|
||||
# Neutral follow-up — does not match any signal regex, does not move the bandit.
|
||||
NEUTRAL_FOLLOWUP = "ok, noted"
|
||||
|
||||
# Oracle: P(success | request_type, model). Tunable.
|
||||
# Defaults: smart dominates code/writing; both are fine for factual_lookup.
|
||||
ORACLE: Dict[str, Dict[str, float]] = {
|
||||
"code_generation": {"smart": 0.92, "fast": 0.35},
|
||||
"factual_lookup": {"smart": 0.90, "fast": 0.85},
|
||||
"writing": {"smart": 0.85, "fast": 0.55},
|
||||
}
|
||||
|
||||
# Fabricated assistant turn — content doesn't matter for the hook, only the role.
|
||||
FAB_ASSISTANT = "Got it. Working on that now."
|
||||
|
||||
|
||||
def _build_messages(prompt: str, last_user: str) -> List[Dict[str, str]]:
|
||||
"""5-message conversation that passes the SIGNAL_GATE_MIN_MESSAGES=4 gate."""
|
||||
return [
|
||||
{"role": "user", "content": prompt},
|
||||
{"role": "assistant", "content": FAB_ASSISTANT},
|
||||
{"role": "user", "content": "ok continue"},
|
||||
{"role": "assistant", "content": FAB_ASSISTANT},
|
||||
{"role": "user", "content": last_user},
|
||||
]
|
||||
|
||||
|
||||
async def _send(
|
||||
client: httpx.AsyncClient,
|
||||
proxy_url: str,
|
||||
api_key: str,
|
||||
router: str,
|
||||
session_id: str,
|
||||
messages: List[Dict[str, str]],
|
||||
mock_response: str,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Returns (ok, chosen_model)."""
|
||||
body = {
|
||||
"model": router,
|
||||
"messages": messages,
|
||||
"metadata": {"litellm_session_id": session_id},
|
||||
"mock_response": mock_response,
|
||||
}
|
||||
try:
|
||||
r = await client.post(
|
||||
f"{proxy_url}/v1/chat/completions",
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {api_key}"},
|
||||
timeout=15.0,
|
||||
)
|
||||
r.raise_for_status()
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f" request failed: {e}", file=sys.stderr)
|
||||
return False, ""
|
||||
chosen = r.headers.get("x-litellm-adaptive-router-model", "")
|
||||
return True, chosen
|
||||
|
||||
|
||||
async def _drive_one_session(
|
||||
client: httpx.AsyncClient,
|
||||
proxy_url: str,
|
||||
api_key: str,
|
||||
router: str,
|
||||
request_type: str,
|
||||
prompt: str,
|
||||
) -> str:
|
||||
"""Run one labeled session. Returns the chosen model (for logging)."""
|
||||
session_id = f"demo-{uuid.uuid4()}"
|
||||
|
||||
# Send the engineered 5-message conversation. The follow-up is chosen
|
||||
# AFTER we observe what model the router would pick — but since the
|
||||
# router is sticky-per-session, the model on this single round-trip
|
||||
# IS the model we're crediting.
|
||||
#
|
||||
# Pre-decide success based on the oracle for whichever model gets picked.
|
||||
# We can't know the pick before sending, so: send a neutral follow-up
|
||||
# first to learn the pick, then send a second round with credit attached.
|
||||
#
|
||||
# Round 1: neutral follow-up → no signal fires, but we learn the pick.
|
||||
ok, chosen = await _send(
|
||||
client, proxy_url, api_key, router, session_id,
|
||||
_build_messages(prompt, NEUTRAL_FOLLOWUP),
|
||||
mock_response=FAB_ASSISTANT,
|
||||
)
|
||||
if not ok or not chosen:
|
||||
return ""
|
||||
|
||||
# Decide outcome from oracle.
|
||||
p = ORACLE.get(request_type, {}).get(chosen, 0.5)
|
||||
success = random.random() < p
|
||||
follow_up = SATISFY[request_type] if success else NEUTRAL_FOLLOWUP
|
||||
|
||||
# Round 2: include the round-1 turns + a new follow-up. On success the
|
||||
# follow-up matches satisfaction → +alpha for (request_type, chosen).
|
||||
history = _build_messages(prompt, NEUTRAL_FOLLOWUP) + [
|
||||
{"role": "assistant", "content": FAB_ASSISTANT},
|
||||
{"role": "user", "content": follow_up},
|
||||
]
|
||||
await _send(
|
||||
client, proxy_url, api_key, router, session_id, history,
|
||||
mock_response=FAB_ASSISTANT,
|
||||
)
|
||||
return chosen
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--proxy-url", default="http://localhost:4000")
|
||||
ap.add_argument("--api-key", required=True, help="proxy key with /v1/chat/completions perms")
|
||||
ap.add_argument("--router", default="smart-cheap-router")
|
||||
ap.add_argument("--rounds", type=int, default=100)
|
||||
ap.add_argument("--rate", type=float, default=0.5,
|
||||
help="seconds between sessions; lower = faster")
|
||||
ap.add_argument("--types", default="code_generation,factual_lookup,writing",
|
||||
help="comma-separated subset of request types to drive")
|
||||
args = ap.parse_args()
|
||||
|
||||
types = [t.strip() for t in args.types.split(",") if t.strip() in PROMPTS]
|
||||
if not types:
|
||||
print(f"ERROR: no valid types. Choose from: {list(PROMPTS)}", file=sys.stderr)
|
||||
sys.exit(2)
|
||||
|
||||
print(f"driving {args.rounds} sessions across types: {types}")
|
||||
print(f"oracle: {ORACLE}")
|
||||
print(f"proxy: {args.proxy_url} router: {args.router}\n")
|
||||
|
||||
counts: Dict[Tuple[str, str], int] = {}
|
||||
async with httpx.AsyncClient() as client:
|
||||
for i in range(args.rounds):
|
||||
rt = random.choice(types)
|
||||
prompt = random.choice(PROMPTS[rt])
|
||||
chosen = await _drive_one_session(
|
||||
client, args.proxy_url, args.api_key, args.router, rt, prompt,
|
||||
)
|
||||
if chosen:
|
||||
counts[(rt, chosen)] = counts.get((rt, chosen), 0) + 1
|
||||
if (i + 1) % 10 == 0:
|
||||
summary = ", ".join(
|
||||
f"{rt}/{m}={n}" for (rt, m), n in sorted(counts.items())
|
||||
)
|
||||
print(f" round {i + 1}/{args.rounds} picks: {summary}")
|
||||
await asyncio.sleep(args.rate)
|
||||
|
||||
print("\nfinal pick distribution:")
|
||||
for (rt, m), n in sorted(counts.items()):
|
||||
print(f" {rt:22s} → {m:8s} {n}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
216
scripts/verify_adaptive_router.py
Normal file
216
scripts/verify_adaptive_router.py
Normal file
|
|
@ -0,0 +1,216 @@
|
|||
"""
|
||||
End-to-end verification script for the adaptive router.
|
||||
|
||||
Requires:
|
||||
- LiteLLM proxy running on http://localhost:4000 with adaptive_router configured
|
||||
(see litellm/proxy/example_config_yaml/adaptive_router_example.yaml).
|
||||
- Postgres reachable via DATABASE_URL (same one the proxy uses).
|
||||
- LITELLM_PROXY_KEY env var set (a valid key with permission to send requests).
|
||||
- Two model deployments configured under one adaptive_router:
|
||||
* "fast" (cheap, lower quality)
|
||||
* "smart" (expensive, higher quality)
|
||||
|
||||
Run:
|
||||
uv run python scripts/verify_adaptive_router.py
|
||||
|
||||
Optional env:
|
||||
LITELLM_PROXY_URL (default: http://localhost:4000)
|
||||
ADAPTIVE_ROUTER_NAME (default: smart-cheap-router)
|
||||
EXPECTED_WINNER (default: smart) -- model expected to dominate after training
|
||||
TRAIN_SESSIONS (default: 20) -- training sessions in phase 1
|
||||
CONVERGE_SESSIONS (default: 10) -- cold sessions in phase 2
|
||||
WIN_THRESHOLD (default: 0.7) -- min share for EXPECTED_WINNER in phase 2
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from typing import List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
PROXY_URL: str = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000")
|
||||
try:
|
||||
PROXY_KEY: str = os.environ["LITELLM_PROXY_KEY"]
|
||||
except KeyError:
|
||||
print(
|
||||
"ERROR: LITELLM_PROXY_KEY env var must be set (a proxy key with /chat/completions perms).",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(2)
|
||||
|
||||
ROUTER_NAME: str = os.environ.get("ADAPTIVE_ROUTER_NAME", "smart-cheap-router")
|
||||
EXPECTED_WINNER: str = os.environ.get("EXPECTED_WINNER", "smart")
|
||||
TRAIN_SESSIONS: int = int(os.environ.get("TRAIN_SESSIONS", "20"))
|
||||
CONVERGE_SESSIONS: int = int(os.environ.get("CONVERGE_SESSIONS", "10"))
|
||||
WIN_THRESHOLD: float = float(os.environ.get("WIN_THRESHOLD", "0.7"))
|
||||
|
||||
REQUEST_TIMEOUT_SECONDS: float = 30.0
|
||||
RETRY_ATTEMPTS: int = 3
|
||||
RETRY_BACKOFF_SECONDS: float = 1.0
|
||||
FLUSHER_DRAIN_WAIT_SECONDS: float = 30.0 # proxy flusher loop is 10s; pad with margin
|
||||
|
||||
PROMPTS: List[str] = [
|
||||
"Write a Python function that reverses a binary tree",
|
||||
"Explain the time complexity of quicksort",
|
||||
"Design an API for a chat application",
|
||||
]
|
||||
SATISFACTION_PROMPT: str = "thanks, that worked!"
|
||||
|
||||
|
||||
async def _post_chat(
|
||||
client: httpx.AsyncClient, session_id: str, prompt: str
|
||||
) -> Optional[dict]:
|
||||
"""POST a chat completion with retry + timeout. Returns response JSON or None."""
|
||||
body = {
|
||||
"model": ROUTER_NAME,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"metadata": {"litellm_session_id": session_id},
|
||||
}
|
||||
last_exc: Optional[Exception] = None
|
||||
for attempt in range(1, RETRY_ATTEMPTS + 1):
|
||||
try:
|
||||
r = await client.post(
|
||||
f"{PROXY_URL}/v1/chat/completions",
|
||||
json=body,
|
||||
headers={"Authorization": f"Bearer {PROXY_KEY}"},
|
||||
timeout=REQUEST_TIMEOUT_SECONDS,
|
||||
)
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
except Exception as e: # noqa: BLE001
|
||||
last_exc = e
|
||||
if attempt < RETRY_ATTEMPTS:
|
||||
await asyncio.sleep(RETRY_BACKOFF_SECONDS * attempt)
|
||||
print(
|
||||
f" request failed after {RETRY_ATTEMPTS} attempts (session={session_id}): {last_exc}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
async def send_session(
|
||||
client: httpx.AsyncClient,
|
||||
session_id: str,
|
||||
prompts: List[str],
|
||||
satisfy: bool = True,
|
||||
) -> Optional[str]:
|
||||
"""Send a session of N turns. Returns the model that handled the last turn."""
|
||||
last_model: Optional[str] = None
|
||||
for prompt in prompts:
|
||||
resp = await _post_chat(client, session_id, prompt)
|
||||
if resp is None:
|
||||
return None
|
||||
last_model = resp.get("model") or last_model
|
||||
if satisfy:
|
||||
await _post_chat(client, session_id, SATISFACTION_PROMPT)
|
||||
return last_model
|
||||
|
||||
|
||||
async def _proxy_health_check(client: httpx.AsyncClient) -> bool:
|
||||
"""Confirm the proxy is reachable before doing anything else."""
|
||||
try:
|
||||
r = await client.get(f"{PROXY_URL}/health/liveliness", timeout=5.0)
|
||||
return r.status_code == 200
|
||||
except Exception as e: # noqa: BLE001
|
||||
print(f"proxy unreachable at {PROXY_URL}: {e}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
print("=== verify_adaptive_router.py ===")
|
||||
print(f"proxy: {PROXY_URL}")
|
||||
print(f"router: {ROUTER_NAME}")
|
||||
print(f"expected winner: {EXPECTED_WINNER}")
|
||||
print(f"train sessions: {TRAIN_SESSIONS}")
|
||||
print(f"converge runs: {CONVERGE_SESSIONS}\n")
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
if not await _proxy_health_check(client):
|
||||
print("FAIL: proxy health check did not return 200.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
# ---- Phase 1: training -------------------------------------------
|
||||
print(
|
||||
f"Phase 1: training ({TRAIN_SESSIONS} sessions of 3 turns + satisfaction)..."
|
||||
)
|
||||
for i in range(TRAIN_SESSIONS):
|
||||
sid = f"verify-train-{uuid.uuid4()}"
|
||||
await send_session(client, sid, PROMPTS, satisfy=True)
|
||||
if (i + 1) % 5 == 0:
|
||||
print(f" trained {i + 1}/{TRAIN_SESSIONS} sessions")
|
||||
|
||||
print(
|
||||
f"\nWaiting {FLUSHER_DRAIN_WAIT_SECONDS:.0f}s for flusher to drain queue..."
|
||||
)
|
||||
await asyncio.sleep(FLUSHER_DRAIN_WAIT_SECONDS)
|
||||
|
||||
# ---- Phase 2: convergence ----------------------------------------
|
||||
print(f"\nPhase 2: convergence test ({CONVERGE_SESSIONS} cold sessions)...")
|
||||
picks: List[str] = []
|
||||
for i in range(CONVERGE_SESSIONS):
|
||||
sid = f"verify-test-{uuid.uuid4()}"
|
||||
m = await send_session(client, sid, [PROMPTS[0]], satisfy=False)
|
||||
if m:
|
||||
picks.append(m)
|
||||
print(f" session {i + 1}: picked {m}")
|
||||
|
||||
if not picks:
|
||||
print("\nFAIL: no successful picks in convergence phase.", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
winner_share = picks.count(EXPECTED_WINNER) / len(picks)
|
||||
print(
|
||||
f"\n{EXPECTED_WINNER} share: {winner_share:.0%} "
|
||||
f"({picks.count(EXPECTED_WINNER)}/{len(picks)})"
|
||||
)
|
||||
|
||||
# ---- Phase 3: sticky session -------------------------------------
|
||||
print("\nPhase 3: sticky session test...")
|
||||
sid = f"verify-sticky-{uuid.uuid4()}"
|
||||
models: List[str] = []
|
||||
for _ in range(3):
|
||||
m = await send_session(client, sid, [PROMPTS[0]], satisfy=False)
|
||||
if m:
|
||||
models.append(m)
|
||||
if len(models) == 3 and len(set(models)) == 1:
|
||||
print(f" PASS: same model {models[0]} across 3 turns of session {sid}")
|
||||
else:
|
||||
print(
|
||||
f" FAIL: models differed within session: {models}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
# ---- Phase 4: latency benchmark ----------------------------------
|
||||
print("\nPhase 4: routing latency (5 picks, p50)...")
|
||||
latencies: List[float] = []
|
||||
for _ in range(5):
|
||||
t0 = time.perf_counter()
|
||||
await send_session(
|
||||
client, f"verify-lat-{uuid.uuid4()}", [PROMPTS[0]], satisfy=False
|
||||
)
|
||||
latencies.append(time.perf_counter() - t0)
|
||||
latencies.sort()
|
||||
p50 = latencies[len(latencies) // 2]
|
||||
print(f" p50 e2e roundtrip: {p50 * 1000:.0f}ms")
|
||||
|
||||
# ---- Verdict -----------------------------------------------------
|
||||
if winner_share >= WIN_THRESHOLD:
|
||||
print(
|
||||
f"\nPASS: convergence ({winner_share:.0%} >= {WIN_THRESHOLD:.0%}) + "
|
||||
f"sticky + latency checks all green."
|
||||
)
|
||||
sys.exit(0)
|
||||
print(
|
||||
f"\nFAIL: convergence too weak ({winner_share:.0%} < {WIN_THRESHOLD:.0%}).",
|
||||
file=sys.stderr,
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
|
@ -1107,7 +1107,12 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook():
|
|||
|
||||
# Mock the make_bedrock_api_request method to track calls
|
||||
async def mock_make_bedrock_api_request(
|
||||
source, messages=None, response=None, request_data=None
|
||||
source,
|
||||
messages=None,
|
||||
response=None,
|
||||
request_data=None,
|
||||
logging_event_type=None,
|
||||
**kwargs,
|
||||
):
|
||||
bedrock_calls.append(
|
||||
{
|
||||
|
|
@ -1115,6 +1120,7 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook():
|
|||
"messages": messages,
|
||||
"response": response,
|
||||
"request_data": request_data,
|
||||
"logging_event_type": logging_event_type,
|
||||
}
|
||||
)
|
||||
# Return the mock bedrock response
|
||||
|
|
|
|||
|
|
@ -244,6 +244,7 @@ class TestRouterIndexManagement:
|
|||
# Methods that are allowed to iterate through self.model_list
|
||||
ALLOWED_METHODS = [
|
||||
"_get_deployment_by_litellm_model", # Edge case: lookup by litellm_params.model (not indexed)
|
||||
"_finalize_adaptive_router_if_configured", # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match)
|
||||
]
|
||||
|
||||
# Get path to router.py
|
||||
|
|
|
|||
|
|
@ -1,10 +1,9 @@
|
|||
import pytest
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import json
|
||||
import time
|
||||
from httpx import AsyncClient
|
||||
from typing import Any, Optional
|
||||
|
||||
import litellm
|
||||
from litellm._uuid import uuid
|
||||
|
||||
"""
|
||||
|
|
@ -12,15 +11,13 @@ Tests to run
|
|||
|
||||
Basic Tests:
|
||||
1. Basic Spend Accuracy Test:
|
||||
- Make 1 calibration request, poll for spend to derive SPEND_PER_REQUEST
|
||||
- Make N-1 more requests (N total)
|
||||
- Expect the spend for each of the following to be N * SPEND_PER_REQUEST
|
||||
Key, Team, User, Org (call /info endpoint for each object to validate)
|
||||
- Make N requests, compute expected total spend locally from each response's usage
|
||||
- Poll until batch writer has flushed spend to the DB
|
||||
- Expect spend for Key, Team, User, Org (/info endpoints) to equal the computed total
|
||||
|
||||
2. Long term spend accuracy test (with 2 bursts of requests)
|
||||
- Burst 1: Make requests, derive SPEND_PER_REQUEST from first request
|
||||
- Burst 2: Make more requests
|
||||
- Verify total spend = (burst1 + burst2) * SPEND_PER_REQUEST
|
||||
- Burst 1: compute expected from responses, verify
|
||||
- Burst 2: compute expected from responses, verify total = burst1 + burst2
|
||||
|
||||
Additional Test Scenarios:
|
||||
|
||||
|
|
@ -38,6 +35,34 @@ Additional Test Scenarios:
|
|||
- Verify accurate total spend calculation
|
||||
"""
|
||||
|
||||
# Upstream model the proxy is configured with (spend_tracking_config.yaml).
|
||||
# The proxy computes spend using this model's pricing; the local ground-truth
|
||||
# calculation uses the same pricing table via litellm.cost_per_token.
|
||||
UPSTREAM_MODEL = "gpt-3.5-turbo"
|
||||
|
||||
# Batch writer flush cadence in CI is ~2-7s (PROXY_BATCH_WRITE_AT=2 + up to 5s jitter).
|
||||
# Poll every 2s for 60s — plenty of headroom for multiple ticks to land.
|
||||
POLL_INTERVAL_SECONDS = 2
|
||||
POLL_TIMEOUT_SECONDS = 60
|
||||
|
||||
TOLERANCE = 1e-10
|
||||
|
||||
|
||||
def _make_test_session() -> aiohttp.ClientSession:
|
||||
"""
|
||||
Session tuned for CI reliability:
|
||||
- force_close: avoid aiohttp reusing a TCP connection that the proxy/kernel
|
||||
silently closed during the long idle window between setup POSTs and the
|
||||
later poll loop (observed failure mode: ConnectionTimeoutError on the
|
||||
first /key/info call after 20 chat completions).
|
||||
- explicit connect timeout: surface a blocked proxy event loop quickly
|
||||
instead of hanging on aiohttp's 5-minute default total timeout.
|
||||
"""
|
||||
return aiohttp.ClientSession(
|
||||
connector=aiohttp.TCPConnector(force_close=True),
|
||||
timeout=aiohttp.ClientTimeout(total=30, connect=10),
|
||||
)
|
||||
|
||||
|
||||
async def create_organization(session, organization_alias: str):
|
||||
"""Helper function to create a new organization"""
|
||||
|
|
@ -102,54 +127,83 @@ async def get_spend_info(session, entity_type: str, entity_id: str):
|
|||
return await response.json()
|
||||
|
||||
|
||||
async def poll_key_spend_until_nonzero(
|
||||
session, key: str, timeout: int = 120, interval: int = 10
|
||||
):
|
||||
"""Poll key spend until it becomes non-zero or timeout is reached."""
|
||||
async def get_proxy_readiness(session):
|
||||
"""Fetch /health/readiness. Used both as a fail-fast gate and as a diagnostic on poll timeout."""
|
||||
url = "http://0.0.0.0:4000/health/readiness"
|
||||
headers = {"Authorization": "Bearer sk-1234"}
|
||||
async with session.get(url, headers=headers) as response:
|
||||
return response.status, await response.json()
|
||||
|
||||
|
||||
async def assert_proxy_healthy(session):
|
||||
"""Fail fast if the proxy's DB or cache is not reachable — no point running the test."""
|
||||
status, body = await get_proxy_readiness(session)
|
||||
if status != 200 or body.get("db") != "connected":
|
||||
pytest.fail(
|
||||
f"Proxy /health/readiness unhealthy (status={status}). "
|
||||
f"Cannot run spend accuracy test. Response: {body}"
|
||||
)
|
||||
print(f"Proxy readiness OK: {body}")
|
||||
|
||||
|
||||
def compute_expected_spend(responses) -> float:
|
||||
"""
|
||||
Compute the expected total spend locally from each response's usage tokens,
|
||||
using the same pricing table the proxy uses. This is the independent ground
|
||||
truth we compare the proxy's reported spend against.
|
||||
"""
|
||||
total = 0.0
|
||||
for r in responses:
|
||||
usage = r.usage
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model=UPSTREAM_MODEL,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
)
|
||||
total += prompt_cost + completion_cost
|
||||
return total
|
||||
|
||||
|
||||
async def poll_key_spend_until(session, key: str, expected: float) -> float:
|
||||
"""
|
||||
Poll key spend until it matches `expected` within TOLERANCE, or timeout.
|
||||
Returns the last observed spend either way; caller decides how to report.
|
||||
"""
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
key_info = await get_spend_info(session, "key", key)
|
||||
spend = key_info["info"]["spend"]
|
||||
if spend > 0:
|
||||
print(
|
||||
f"Key spend became non-zero ({spend}) after {time.time() - start:.1f}s"
|
||||
)
|
||||
return spend
|
||||
print(f"Key spend still 0.0, waiting... ({time.time() - start:.1f}s elapsed)")
|
||||
await asyncio.sleep(interval)
|
||||
raise TimeoutError(
|
||||
f"Key spend remained 0.0 after {timeout}s — batch writer may not be running"
|
||||
)
|
||||
|
||||
|
||||
async def calibrate_spend_per_request(session, key: str, max_retries: int = 5):
|
||||
"""
|
||||
Make a single calibration request and poll for its spend to derive SPEND_PER_REQUEST.
|
||||
Fails fast with pytest.fail() if spend cannot be determined.
|
||||
"""
|
||||
response = await chat_completion(session, key)
|
||||
print(f"Calibration request completed: {response}")
|
||||
|
||||
for attempt in range(1, max_retries + 1):
|
||||
last_spend = 0.0
|
||||
while time.time() - start < POLL_TIMEOUT_SECONDS:
|
||||
try:
|
||||
spend = await poll_key_spend_until_nonzero(
|
||||
session, key, timeout=120, interval=10
|
||||
)
|
||||
key_info = await get_spend_info(session, "key", key)
|
||||
except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
|
||||
print(
|
||||
f"Calibrated SPEND_PER_REQUEST = {spend} "
|
||||
f"(attempt {attempt}/{max_retries})"
|
||||
f"Transient transport error during spend poll: "
|
||||
f"{type(exc).__name__}: {exc}. Retrying... "
|
||||
f"({time.time() - start:.1f}s elapsed)"
|
||||
)
|
||||
return spend
|
||||
except TimeoutError:
|
||||
if attempt < max_retries:
|
||||
print(
|
||||
f"Calibration attempt {attempt}/{max_retries} timed out, retrying..."
|
||||
)
|
||||
else:
|
||||
pytest.fail(
|
||||
f"Failed to calibrate SPEND_PER_REQUEST after {max_retries} attempts. "
|
||||
"The batch writer may not be running or the model may have 0 cost."
|
||||
)
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
continue
|
||||
last_spend = key_info["info"]["spend"]
|
||||
if abs(last_spend - expected) < TOLERANCE:
|
||||
print(
|
||||
f"Key spend reached expected {expected} after {time.time() - start:.1f}s"
|
||||
)
|
||||
return last_spend
|
||||
print(
|
||||
f"Key spend {last_spend}, expected {expected}, waiting... "
|
||||
f"({time.time() - start:.1f}s elapsed)"
|
||||
)
|
||||
await asyncio.sleep(POLL_INTERVAL_SECONDS)
|
||||
return last_spend
|
||||
|
||||
|
||||
async def fail_with_diagnostics(session, stage: str, expected: float, observed: float):
|
||||
"""Emit a failure with readiness state so CI output points at the real cause."""
|
||||
_, readiness = await get_proxy_readiness(session)
|
||||
pytest.fail(
|
||||
f"{stage}: key spend did not match expected after {POLL_TIMEOUT_SECONDS}s poll. "
|
||||
f"expected={expected}, observed={observed}, diff={expected - observed}. "
|
||||
f"Proxy readiness: {readiness}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -157,63 +211,60 @@ async def test_basic_spend_accuracy():
|
|||
"""
|
||||
Test basic spend accuracy across different entities:
|
||||
1. Create org, team, user, and key
|
||||
2. Make 1 calibration request to derive SPEND_PER_REQUEST
|
||||
3. Make remaining requests (NUM_LLM_REQUESTS total)
|
||||
4. Verify spend accuracy for key, team, user, and org
|
||||
2. Make N requests, keeping each response
|
||||
3. Compute expected spend locally from response usage (independent ground truth)
|
||||
4. Poll until proxy-reported spend matches expected
|
||||
5. Verify spend is consistent across key, team, user, and org entities
|
||||
"""
|
||||
NUM_LLM_REQUESTS = 20
|
||||
TOLERANCE = 1e-10
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Create organization
|
||||
async with _make_test_session() as session:
|
||||
await assert_proxy_healthy(session)
|
||||
|
||||
org_response = await create_organization(
|
||||
session=session, organization_alias=f"test-org-{uuid.uuid4()}"
|
||||
)
|
||||
print("org_response: ", org_response)
|
||||
org_id = org_response["organization_id"]
|
||||
|
||||
# Create team under organization
|
||||
team_response = await create_team(session, org_id)
|
||||
print("team_response: ", team_response)
|
||||
team_id = team_response["team_id"]
|
||||
|
||||
# Create user
|
||||
user_response = await create_user(session, org_id)
|
||||
print("user_response: ", user_response)
|
||||
user_id = user_response["user_id"]
|
||||
|
||||
# Generate key
|
||||
key_response = await generate_key(session, user_id, team_id)
|
||||
print("key_response: ", key_response)
|
||||
key = key_response["key"]
|
||||
|
||||
# Calibrate: make 1 request and derive SPEND_PER_REQUEST
|
||||
spend_per_request = await calibrate_spend_per_request(session, key)
|
||||
expected_spend = NUM_LLM_REQUESTS * spend_per_request
|
||||
print(f"SPEND_PER_REQUEST={spend_per_request}, expected_spend={expected_spend}")
|
||||
|
||||
# Make remaining requests (1 already made during calibration)
|
||||
for i in range(NUM_LLM_REQUESTS - 1):
|
||||
responses = []
|
||||
for i in range(NUM_LLM_REQUESTS):
|
||||
response = await chat_completion(session, key)
|
||||
print(f"Request {i + 2}/{NUM_LLM_REQUESTS} completed")
|
||||
responses.append(response)
|
||||
print(f"Request {i + 1}/{NUM_LLM_REQUESTS} completed")
|
||||
|
||||
# Poll until batch writer has flushed all spend
|
||||
start = time.time()
|
||||
while time.time() - start < 120:
|
||||
key_info = await get_spend_info(session, "key", key)
|
||||
current_spend = key_info["info"]["spend"]
|
||||
if abs(current_spend - expected_spend) < TOLERANCE:
|
||||
print(
|
||||
f"Key spend reached expected {expected_spend} after {time.time() - start:.1f}s"
|
||||
)
|
||||
break
|
||||
print(f"Key spend {current_spend}, expected {expected_spend}, waiting...")
|
||||
await asyncio.sleep(10)
|
||||
expected_spend = compute_expected_spend(responses)
|
||||
assert expected_spend > 0, (
|
||||
f"Locally computed expected spend is {expected_spend}. Either cost calc "
|
||||
f"is broken or upstream returned zero tokens. "
|
||||
f"Usage: {[r.usage.model_dump() for r in responses]}"
|
||||
)
|
||||
print(f"Expected total spend (local ground truth): {expected_spend}")
|
||||
|
||||
# Allow extra time for all entity spend aggregations to complete
|
||||
final_spend = await poll_key_spend_until(session, key, expected_spend)
|
||||
if abs(final_spend - expected_spend) >= TOLERANCE:
|
||||
await fail_with_diagnostics(
|
||||
session,
|
||||
stage="test_basic_spend_accuracy",
|
||||
expected=expected_spend,
|
||||
observed=final_spend,
|
||||
)
|
||||
|
||||
# Allow a final scheduler tick for team/user/org aggregations to settle
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Get spend information for each entity
|
||||
key_info = await get_spend_info(session, "key", key)
|
||||
print("key_info: ", key_info)
|
||||
team_info = await get_spend_info(session, "team", team_id)
|
||||
|
|
@ -223,7 +274,6 @@ async def test_basic_spend_accuracy():
|
|||
org_info = await get_spend_info(session, "organization", org_id)
|
||||
print("org_info: ", org_info)
|
||||
|
||||
# Verify spend for each entity
|
||||
assert (
|
||||
abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE
|
||||
), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}"
|
||||
|
|
@ -246,91 +296,78 @@ async def test_long_term_spend_accuracy_with_bursts():
|
|||
"""
|
||||
Test long-term spend accuracy with multiple bursts of requests:
|
||||
1. Create org, team, user, and key
|
||||
2. Calibrate SPEND_PER_REQUEST from first request
|
||||
3. Burst 1: Make remaining requests
|
||||
4. Burst 2: Make more requests
|
||||
5. Verify the total spend is tracked accurately across all entities
|
||||
2. Burst 1: make requests, compute expected locally, verify proxy matches
|
||||
3. Burst 2: make more requests, verify proxy total == burst1 + burst2
|
||||
4. Verify total spend is consistent across all entities
|
||||
"""
|
||||
BURST_1_REQUESTS = 22
|
||||
BURST_2_REQUESTS = 12
|
||||
TOTAL_REQUESTS = BURST_1_REQUESTS + BURST_2_REQUESTS
|
||||
TOLERANCE = 1e-10
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Create organization
|
||||
async with _make_test_session() as session:
|
||||
await assert_proxy_healthy(session)
|
||||
|
||||
org_response = await create_organization(
|
||||
session=session, organization_alias=f"test-org-{uuid.uuid4()}"
|
||||
)
|
||||
print("org_response: ", org_response)
|
||||
org_id = org_response["organization_id"]
|
||||
|
||||
# Create team under organization
|
||||
team_response = await create_team(session, org_id)
|
||||
print("team_response: ", team_response)
|
||||
team_id = team_response["team_id"]
|
||||
|
||||
# Create user
|
||||
user_response = await create_user(session, org_id)
|
||||
print("user_response: ", user_response)
|
||||
user_id = user_response["user_id"]
|
||||
|
||||
# Generate key
|
||||
key_response = await generate_key(session, user_id, team_id)
|
||||
print("key_response: ", key_response)
|
||||
key = key_response["key"]
|
||||
|
||||
# Calibrate: make 1 request and derive SPEND_PER_REQUEST
|
||||
spend_per_request = await calibrate_spend_per_request(session, key)
|
||||
expected_spend = TOTAL_REQUESTS * spend_per_request
|
||||
print(f"SPEND_PER_REQUEST={spend_per_request}, expected_spend={expected_spend}")
|
||||
|
||||
# First burst: remaining requests (1 already made during calibration)
|
||||
print(f"Starting first burst ({BURST_1_REQUESTS - 1} remaining requests)...")
|
||||
for i in range(BURST_1_REQUESTS - 1):
|
||||
print(f"Starting first burst of {BURST_1_REQUESTS} requests...")
|
||||
burst_1_responses = []
|
||||
for i in range(BURST_1_REQUESTS):
|
||||
response = await chat_completion(session, key)
|
||||
print(f"Burst 1 - Request {i + 2}/{BURST_1_REQUESTS} completed")
|
||||
burst_1_responses.append(response)
|
||||
print(f"Burst 1 - Request {i + 1}/{BURST_1_REQUESTS} completed")
|
||||
|
||||
# Poll until batch writer has flushed burst 1 spend
|
||||
burst_1_expected = BURST_1_REQUESTS * spend_per_request
|
||||
start = time.time()
|
||||
while time.time() - start < 120:
|
||||
key_info_check = await get_spend_info(session, "key", key)
|
||||
current_spend = key_info_check["info"]["spend"]
|
||||
if abs(current_spend - burst_1_expected) < TOLERANCE:
|
||||
print(
|
||||
f"Burst 1 spend reached expected {burst_1_expected} after {time.time() - start:.1f}s"
|
||||
)
|
||||
break
|
||||
print(f"Key spend {current_spend}, expected {burst_1_expected}, waiting...")
|
||||
await asyncio.sleep(10)
|
||||
burst_1_expected = compute_expected_spend(burst_1_responses)
|
||||
assert burst_1_expected > 0, (
|
||||
f"Burst 1 expected spend is {burst_1_expected}. "
|
||||
f"Usage: {[r.usage.model_dump() for r in burst_1_responses]}"
|
||||
)
|
||||
print(f"Burst 1 expected spend: {burst_1_expected}")
|
||||
|
||||
# Check intermediate spend
|
||||
intermediate_key_info = await get_spend_info(session, "key", key)
|
||||
print(f"After Burst 1 - Key spend: {intermediate_key_info['info']['spend']}")
|
||||
final_burst_1 = await poll_key_spend_until(session, key, burst_1_expected)
|
||||
if abs(final_burst_1 - burst_1_expected) >= TOLERANCE:
|
||||
await fail_with_diagnostics(
|
||||
session,
|
||||
stage="test_long_term_spend_accuracy burst 1",
|
||||
expected=burst_1_expected,
|
||||
observed=final_burst_1,
|
||||
)
|
||||
|
||||
# Second burst
|
||||
print(f"Starting second burst of {BURST_2_REQUESTS} requests...")
|
||||
burst_2_responses = []
|
||||
for i in range(BURST_2_REQUESTS):
|
||||
response = await chat_completion(session, key)
|
||||
burst_2_responses.append(response)
|
||||
print(f"Burst 2 - Request {i + 1}/{BURST_2_REQUESTS} completed")
|
||||
|
||||
# Poll until key spend reaches expected total (burst 1 + burst 2)
|
||||
start = time.time()
|
||||
while time.time() - start < 120:
|
||||
key_info_check = await get_spend_info(session, "key", key)
|
||||
current_spend = key_info_check["info"]["spend"]
|
||||
if abs(current_spend - expected_spend) < TOLERANCE:
|
||||
print(
|
||||
f"Total spend reached expected {expected_spend} after {time.time() - start:.1f}s"
|
||||
)
|
||||
break
|
||||
print(f"Key spend {current_spend}, expected {expected_spend}, waiting...")
|
||||
await asyncio.sleep(10)
|
||||
total_expected = burst_1_expected + compute_expected_spend(burst_2_responses)
|
||||
print(f"Total expected spend (burst 1 + burst 2): {total_expected}")
|
||||
|
||||
final_total = await poll_key_spend_until(session, key, total_expected)
|
||||
if abs(final_total - total_expected) >= TOLERANCE:
|
||||
await fail_with_diagnostics(
|
||||
session,
|
||||
stage="test_long_term_spend_accuracy total",
|
||||
expected=total_expected,
|
||||
observed=final_total,
|
||||
)
|
||||
|
||||
# Allow extra time for all entity spend aggregations
|
||||
await asyncio.sleep(5)
|
||||
|
||||
# Get final spend information for each entity
|
||||
key_info = await get_spend_info(session, "key", key)
|
||||
team_info = await get_spend_info(session, "team", team_id)
|
||||
user_info = await get_spend_info(session, "user", user_id)
|
||||
|
|
@ -341,19 +378,18 @@ async def test_long_term_spend_accuracy_with_bursts():
|
|||
print(f"Final user spend: {user_info['user_info']['spend']}")
|
||||
print(f"Final org spend: {org_info['spend']}")
|
||||
|
||||
# Verify total spend for each entity
|
||||
assert (
|
||||
abs(key_info["info"]["spend"] - expected_spend) < TOLERANCE
|
||||
), f"Key spend {key_info['info']['spend']} does not match expected {expected_spend}"
|
||||
abs(key_info["info"]["spend"] - total_expected) < TOLERANCE
|
||||
), f"Key spend {key_info['info']['spend']} does not match expected {total_expected}"
|
||||
|
||||
assert (
|
||||
abs(user_info["user_info"]["spend"] - expected_spend) < TOLERANCE
|
||||
), f"User spend {user_info['user_info']['spend']} does not match expected {expected_spend}"
|
||||
abs(user_info["user_info"]["spend"] - total_expected) < TOLERANCE
|
||||
), f"User spend {user_info['user_info']['spend']} does not match expected {total_expected}"
|
||||
|
||||
assert (
|
||||
abs(team_info["team_info"]["spend"] - expected_spend) < TOLERANCE
|
||||
), f"Team spend {team_info['team_info']['spend']} does not match expected {expected_spend}"
|
||||
abs(team_info["team_info"]["spend"] - total_expected) < TOLERANCE
|
||||
), f"Team spend {team_info['team_info']['spend']} does not match expected {total_expected}"
|
||||
|
||||
assert (
|
||||
abs(org_info["spend"] - expected_spend) < TOLERANCE
|
||||
), f"Organization spend {org_info['spend']} does not match expected {expected_spend}"
|
||||
abs(org_info["spend"] - total_expected) < TOLERANCE
|
||||
), f"Organization spend {org_info['spend']} does not match expected {total_expected}"
|
||||
|
|
|
|||
|
|
@ -1055,3 +1055,50 @@ class TestTracingFieldsPopulation:
|
|||
assert slg["classification"] == classification
|
||||
assert slg["detection_method"] == "llm-judge"
|
||||
assert slg["confidence_score"] == 0.94
|
||||
|
||||
|
||||
class TestCustomGuardrailSpendLogMatchRedaction:
|
||||
"""Guardrail JSON persisted via standard_logging must not contain raw match spans."""
|
||||
|
||||
def test_add_standard_logging_redacts_nested_match(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
raw = {
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
request_data: dict = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=raw,
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert (
|
||||
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
|
||||
"piiEntities"
|
||||
][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "GG"
|
||||
|
||||
def test_add_standard_logging_redacts_regex_field(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}
|
||||
request_data: dict = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=raw,
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]"
|
||||
assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}"
|
||||
|
|
|
|||
|
|
@ -1047,6 +1047,36 @@ class TestOpenTelemetryEndpointNormalization(unittest.TestCase):
|
|||
result = otel._normalize_otel_endpoint("http://collector:4318/", "traces")
|
||||
self.assertEqual(result, "http://collector:4318/v1/traces")
|
||||
|
||||
@parameterized.expand(
|
||||
[
|
||||
(
|
||||
"https://ingest.eu1.observability.splunkcloud.com/v2/trace/otlp",
|
||||
"https://ingest.eu1.observability.splunkcloud.com/v2/trace/otlp",
|
||||
),
|
||||
(
|
||||
"https://ingest.us0.observability.splunkcloud.com/v2/trace/otlp/",
|
||||
"https://ingest.us0.observability.splunkcloud.com/v2/trace/otlp",
|
||||
),
|
||||
(
|
||||
"https://ingest.eu0.signalfx.com/v2/trace/otlp",
|
||||
"https://ingest.eu0.signalfx.com/v2/trace/otlp",
|
||||
),
|
||||
(
|
||||
"https://example.com/prefix/v2/trace/otlp",
|
||||
"https://example.com/prefix/v2/trace/otlp",
|
||||
),
|
||||
]
|
||||
)
|
||||
def test_normalize_traces_nonstandard_otlp_ingest_urls_unchanged(
|
||||
self, input_url: str, expected: str
|
||||
) -> None:
|
||||
"""Splunk-style /v2/trace/otlp endpoints must not get /v1/traces appended."""
|
||||
otel = OpenTelemetry()
|
||||
self.assertEqual(
|
||||
otel._normalize_otel_endpoint(input_url, "traces"),
|
||||
expected,
|
||||
)
|
||||
|
||||
def test_normalize_endpoint_none(self):
|
||||
"""Test that None endpoint returns None"""
|
||||
otel = OpenTelemetry()
|
||||
|
|
@ -1315,7 +1345,7 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase):
|
|||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"OTEL_EXPORTER": "otlp_http",
|
||||
"OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4318",
|
||||
},
|
||||
clear=False,
|
||||
|
|
@ -1339,7 +1369,7 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase):
|
|||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"OTEL_EXPORTER": "otlp_grpc",
|
||||
"OTEL_EXPORTER_OTLP_PROTOCOL": "grpc",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4317",
|
||||
},
|
||||
clear=False,
|
||||
|
|
@ -1360,6 +1390,60 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase):
|
|||
self.assertIsInstance(processor, BatchSpanProcessor)
|
||||
self.assertIsInstance(processor.span_exporter, OTLPSpanExporterGRPC)
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"OTEL_EXPORTER": "otlp_http",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4318",
|
||||
},
|
||||
clear=False,
|
||||
)
|
||||
def test_protocol_selection_from_otel_exporter_fallback_http(self):
|
||||
"""OTEL_EXPORTER drives protocol when OTEL_EXPORTER_OTLP_PROTOCOL is unset."""
|
||||
from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
|
||||
OTLPSpanExporter as OTLPSpanExporterHTTP,
|
||||
)
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
|
||||
popped_protocol = os.environ.pop("OTEL_EXPORTER_OTLP_PROTOCOL", None)
|
||||
try:
|
||||
config = OpenTelemetryConfig.from_env()
|
||||
self.assertEqual(config.exporter, "otlp_http")
|
||||
otel = OpenTelemetry(config=config)
|
||||
processor = otel._get_span_processor()
|
||||
self.assertIsInstance(processor, BatchSpanProcessor)
|
||||
self.assertIsInstance(processor.span_exporter, OTLPSpanExporterHTTP)
|
||||
finally:
|
||||
if popped_protocol is not None:
|
||||
os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = popped_protocol
|
||||
|
||||
@patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"OTEL_EXPORTER": "otlp_grpc",
|
||||
"OTEL_EXPORTER_OTLP_ENDPOINT": "http://collector:4317",
|
||||
},
|
||||
clear=False,
|
||||
)
|
||||
def test_protocol_selection_from_otel_exporter_fallback_grpc(self):
|
||||
"""OTEL_EXPORTER drives protocol when OTEL_EXPORTER_OTLP_PROTOCOL is unset."""
|
||||
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
|
||||
OTLPSpanExporter as OTLPSpanExporterGRPC,
|
||||
)
|
||||
from opentelemetry.sdk.trace.export import BatchSpanProcessor
|
||||
|
||||
popped_protocol = os.environ.pop("OTEL_EXPORTER_OTLP_PROTOCOL", None)
|
||||
try:
|
||||
config = OpenTelemetryConfig.from_env()
|
||||
self.assertEqual(config.exporter, "otlp_grpc")
|
||||
otel = OpenTelemetry(config=config)
|
||||
processor = otel._get_span_processor()
|
||||
self.assertIsInstance(processor, BatchSpanProcessor)
|
||||
self.assertIsInstance(processor.span_exporter, OTLPSpanExporterGRPC)
|
||||
finally:
|
||||
if popped_protocol is not None:
|
||||
os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = popped_protocol
|
||||
|
||||
def test_http_exporter_endpoint_normalization_for_traces(self):
|
||||
"""Test that HTTP trace exporter gets properly normalized endpoint"""
|
||||
config = OpenTelemetryConfig(
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
_FINISH_REASON_MAP,
|
||||
map_finish_reason,
|
||||
reconstruct_model_name,
|
||||
redact_nested_match_and_regex_keys,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -158,3 +159,37 @@ class TestFinishReasonMapOutputsAreValid:
|
|||
f"Mapped value '{openai_reason}' (from '{provider_reason}') "
|
||||
f"is not a valid OpenAI finish reason"
|
||||
)
|
||||
|
||||
|
||||
class TestRedactNestedMatchAndRegexKeys:
|
||||
def test_redacts_match_and_regex_recursively(self):
|
||||
payload = {
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}
|
||||
]
|
||||
},
|
||||
"wordPolicy": {
|
||||
"customWords": [{"match": "badword", "action": "BLOCKED"}]
|
||||
},
|
||||
}
|
||||
],
|
||||
"regex": "should-redact-key-named-regex",
|
||||
}
|
||||
out = redact_nested_match_and_regex_keys(payload)
|
||||
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "[REDACTED]"
|
||||
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == (
|
||||
"[REDACTED]"
|
||||
)
|
||||
assert out["regex"] == "[REDACTED]"
|
||||
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][
|
||||
0
|
||||
]["match"] == "secret-name"
|
||||
|
||||
def test_passes_through_none_and_str(self):
|
||||
assert redact_nested_match_and_regex_keys(None) is None
|
||||
assert redact_nested_match_and_regex_keys("plain") == "plain"
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ import pytest
|
|||
import litellm
|
||||
import litellm.utils
|
||||
from litellm import completion
|
||||
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
|
||||
from litellm.llms.moonshot.chat.transformation import MoonshotChatConfig
|
||||
|
||||
|
||||
|
|
@ -653,3 +654,44 @@ class TestMoonshotConfig:
|
|||
result[1].get("reasoning_content")
|
||||
== "<thinking>Planning to call weather tool</thinking>"
|
||||
)
|
||||
|
||||
|
||||
class TestKimiK26ModelRegistry:
|
||||
"""Tests that kimi-k2.6 is correctly registered in the model registry."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def model_cost_map(self):
|
||||
"""Load directly from the bundled backup so tests don't depend on remote fetch."""
|
||||
return GetModelCostMap.load_local_model_cost_map()
|
||||
|
||||
def test_kimi_k26_in_model_cost_map(self, model_cost_map):
|
||||
"""kimi-k2.6 should be present in the model cost map."""
|
||||
assert "moonshot/kimi-k2.6" in model_cost_map, "moonshot/kimi-k2.6 not found in model_cost"
|
||||
|
||||
def test_kimi_k26_pricing(self, model_cost_map):
|
||||
"""kimi-k2.6 pricing should match official Kimi API rates."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["input_cost_per_token"] == pytest.approx(9.5e-07)
|
||||
assert model_info["output_cost_per_token"] == pytest.approx(4e-06)
|
||||
assert model_info["cache_read_input_token_cost"] == pytest.approx(1.6e-07)
|
||||
|
||||
def test_kimi_k26_context_window(self, model_cost_map):
|
||||
"""kimi-k2.6 should have a 256K (262144 token) context window."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["max_input_tokens"] == 262144
|
||||
assert model_info["max_output_tokens"] == 262144
|
||||
assert model_info["max_tokens"] == 262144
|
||||
|
||||
def test_kimi_k26_capabilities(self, model_cost_map):
|
||||
"""kimi-k2.6 should support function calling, vision, video input, tool choice, and reasoning."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info.get("supports_function_calling") is True
|
||||
assert model_info.get("supports_tool_choice") is True
|
||||
assert model_info.get("supports_vision") is True
|
||||
assert model_info.get("supports_video_input") is True
|
||||
assert model_info.get("supports_reasoning") is True
|
||||
|
||||
def test_kimi_k26_provider(self, model_cost_map):
|
||||
"""kimi-k2.6 should be assigned to the moonshot provider."""
|
||||
model_info = model_cost_map["moonshot/kimi-k2.6"]
|
||||
assert model_info["litellm_provider"] == "moonshot"
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import sys
|
|||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Any, Dict, List
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -784,3 +785,37 @@ def test_reset_budget_skips_null_budget_id_endusers_when_default_not_in_reset_li
|
|||
assert len(find_many_calls) == 0
|
||||
|
||||
litellm.max_end_user_budget_id = None
|
||||
|
||||
|
||||
def test_reset_budget_for_team_members_preserves_total_spend():
|
||||
"""Regression guard: reset_budget_for_litellm_team_members must zero `spend`
|
||||
but leave `total_spend` untouched.
|
||||
|
||||
The reset writes `data={"spend": 0}` explicitly. If a future refactor adds
|
||||
`"total_spend": 0` to that dict, this test fails immediately.
|
||||
"""
|
||||
expired_budget = type(
|
||||
"LiteLLM_BudgetTableFull",
|
||||
(),
|
||||
{"budget_id": "budget-1"},
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_teammembership.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma_client.db.litellm_teammembership.update_many = AsyncMock(
|
||||
return_value={"count": 1}
|
||||
)
|
||||
|
||||
job = ResetBudgetJob(
|
||||
proxy_logging_obj=MagicMock(), prisma_client=mock_prisma_client
|
||||
)
|
||||
|
||||
asyncio.run(job.reset_budget_for_litellm_team_members([expired_budget]))
|
||||
|
||||
mock_prisma_client.db.litellm_teammembership.update_many.assert_called_once()
|
||||
call_kwargs = (
|
||||
mock_prisma_client.db.litellm_teammembership.update_many.call_args.kwargs
|
||||
)
|
||||
assert call_kwargs["where"]["budget_id"]["in"] == ["budget-1"]
|
||||
assert call_kwargs["data"] == {"spend": 0}
|
||||
assert "total_spend" not in call_kwargs["data"]
|
||||
|
|
|
|||
|
|
@ -87,6 +87,89 @@ async def test_store_in_memory_spend_updates_uses_pipeline(
|
|||
assert len(rpush_list) == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""
|
||||
If async_rpush_pipeline raises, the already-drained transactions must be
|
||||
put back into the in-memory queues so the next scheduler tick retries.
|
||||
Without this, any transient Redis hiccup silently loses spend data.
|
||||
"""
|
||||
from litellm.proxy._types import Litellm_EntityType
|
||||
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
||||
DailySpendUpdateQueue,
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import (
|
||||
SpendUpdateQueue,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(
|
||||
side_effect=ConnectionError("redis went away")
|
||||
)
|
||||
|
||||
spend_queue = SpendUpdateQueue()
|
||||
daily_user_queue = DailySpendUpdateQueue()
|
||||
daily_team_queue = DailySpendUpdateQueue()
|
||||
daily_org_queue = DailySpendUpdateQueue()
|
||||
daily_end_user_queue = DailySpendUpdateQueue()
|
||||
daily_agent_queue = DailySpendUpdateQueue()
|
||||
|
||||
# Seed real queues with data so flush_and_get_aggregated returns it
|
||||
await spend_queue.add_update(
|
||||
{
|
||||
"entity_type": Litellm_EntityType.KEY,
|
||||
"entity_id": "key-abc",
|
||||
"response_cost": 1.5,
|
||||
}
|
||||
)
|
||||
await spend_queue.add_update(
|
||||
{
|
||||
"entity_type": Litellm_EntityType.TEAM,
|
||||
"entity_id": "team-xyz",
|
||||
"response_cost": 2.5,
|
||||
}
|
||||
)
|
||||
await daily_user_queue.add_update(
|
||||
{
|
||||
"user1_day_model": {
|
||||
"spend": 1.0,
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_queue,
|
||||
daily_spend_update_queue=daily_user_queue,
|
||||
daily_team_spend_update_queue=daily_team_queue,
|
||||
daily_org_spend_update_queue=daily_org_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_queue,
|
||||
)
|
||||
|
||||
# After restore, the main spend queue should hold one item per
|
||||
# (entity_type, entity_id) pair with the aggregated cost
|
||||
restored_spend = (
|
||||
await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
|
||||
)
|
||||
assert restored_spend["key_list_transactions"] == {"key-abc": 1.5}
|
||||
assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5}
|
||||
|
||||
# Daily user queue should hold the same aggregated dict
|
||||
restored_daily = (
|
||||
await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions()
|
||||
)
|
||||
assert restored_daily == {
|
||||
"user1_day_model": {
|
||||
"spend": 1.0,
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
|
|
|
|||
|
|
@ -642,6 +642,81 @@ async def test_commit_spend_updates_to_db_increments_agent_spend():
|
|||
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_commit_spend_updates_to_db_increments_team_member_spend_and_total_spend():
|
||||
"""
|
||||
Verify that _commit_spend_updates_to_db increments BOTH spend (cycle-scoped)
|
||||
and total_spend (non-resetting) on LiteLLM_TeamMembership in a single
|
||||
update_many call, using the same response_cost.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
mock_batcher = MagicMock()
|
||||
mock_batcher.litellm_verificationtoken = MagicMock()
|
||||
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
|
||||
mock_batcher.litellm_usertable = MagicMock()
|
||||
mock_batcher.litellm_usertable.update_many = MagicMock()
|
||||
mock_batcher.litellm_teamtable = MagicMock()
|
||||
mock_batcher.litellm_teamtable.update_many = MagicMock()
|
||||
mock_batcher.litellm_teammembership = MagicMock()
|
||||
mock_batcher.litellm_teammembership.update_many = MagicMock()
|
||||
mock_batcher.litellm_organizationtable = MagicMock()
|
||||
mock_batcher.litellm_organizationtable.update_many = MagicMock()
|
||||
mock_batcher.litellm_tagtable = MagicMock()
|
||||
mock_batcher.litellm_tagtable.update_many = MagicMock()
|
||||
mock_batcher.litellm_agentstable = MagicMock()
|
||||
mock_batcher.litellm_agentstable.update_many = MagicMock()
|
||||
|
||||
mock_transaction = AsyncMock()
|
||||
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
|
||||
mock_transaction.__aexit__ = AsyncMock(return_value=False)
|
||||
mock_transaction.batch_ = MagicMock(
|
||||
return_value=AsyncMock(
|
||||
__aenter__=AsyncMock(return_value=mock_batcher),
|
||||
__aexit__=AsyncMock(return_value=False),
|
||||
)
|
||||
)
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db = MagicMock()
|
||||
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
|
||||
|
||||
mock_proxy_logging = MagicMock()
|
||||
# Skip team-membership cache invalidation — out of scope for this test.
|
||||
mock_proxy_logging.call_details.get = MagicMock(return_value=None)
|
||||
|
||||
team_id = "team-abc"
|
||||
user_id = "user-xyz"
|
||||
response_cost = 0.75
|
||||
entity_id = f"team_id::{team_id}::user_id::{user_id}"
|
||||
db_spend_update_transactions = {
|
||||
"user_list_transactions": {},
|
||||
"end_user_list_transactions": {},
|
||||
"key_list_transactions": {},
|
||||
"team_list_transactions": {},
|
||||
"team_member_list_transactions": {entity_id: response_cost},
|
||||
"org_list_transactions": {},
|
||||
"tag_list_transactions": {},
|
||||
"agent_list_transactions": {},
|
||||
}
|
||||
|
||||
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
|
||||
await db_writer._commit_spend_updates_to_db(
|
||||
prisma_client=mock_prisma_client,
|
||||
n_retry_times=0,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
)
|
||||
|
||||
mock_batcher.litellm_teammembership.update_many.assert_called_once()
|
||||
call_kwargs = mock_batcher.litellm_teammembership.update_many.call_args[1]
|
||||
assert call_kwargs["where"] == {"team_id": team_id, "user_id": user_id}
|
||||
assert call_kwargs["data"] == {
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -12,11 +12,15 @@ from fastapi import HTTPException
|
|||
|
||||
sys.path.insert(0, os.path.abspath("../../../../../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
|
||||
BedrockGuardrail,
|
||||
_redact_pii_matches,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
|
|
@ -106,10 +110,12 @@ async def test__redact_pii_matches_malformed_response():
|
|||
# Test with completely malformed response
|
||||
malformed_response = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"assessments": "not_a_list", # This should cause an exception
|
||||
# Wrong type for assessments; redact_nested_match_and_regex_keys walks dict
|
||||
# values and skips non-dict/list nodes, so this must not raise.
|
||||
"assessments": "not_a_list",
|
||||
}
|
||||
|
||||
# Should not crash and return original response
|
||||
# Should not crash (deep copy + walk skips the string value under assessments)
|
||||
redacted_response = _redact_pii_matches(malformed_response)
|
||||
assert redacted_response == malformed_response
|
||||
|
||||
|
|
@ -188,7 +194,7 @@ async def test__redact_pii_matches_multiple_assessments():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_guardrail_logging_uses_redacted_response():
|
||||
"""Test that the Bedrock guardrail uses redacted response for logging"""
|
||||
"""Debug logs and standard_logging payloads must not include raw match values."""
|
||||
|
||||
# Create proper mock objects
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
|
@ -295,6 +301,14 @@ async def test_bedrock_guardrail_logging_uses_redacted_response():
|
|||
== "PHONE"
|
||||
)
|
||||
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert (
|
||||
slg_list[0]["guardrail_response"]["assessments"][0][
|
||||
"sensitiveInformationPolicy"
|
||||
]["piiEntities"][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
|
||||
print("Bedrock guardrail logging redaction test passed")
|
||||
|
||||
|
||||
|
|
@ -1751,6 +1765,124 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions():
|
|||
print("\u2705 BLOCKED vs ANONYMIZED actions test passed")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Spend logs: guardrail_mode (pre/during/post) vs Bedrock INPUT/OUTPUT
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_bedrock_guardrail_uses_native_during_call_hook():
|
||||
"""during_call must use async_moderation_hook, not unified apply_guardrail(input=request)."""
|
||||
assert BedrockGuardrail.use_native_during_call_hook is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_bedrock_api_request_logging_event_type_for_spend_logs():
|
||||
"""
|
||||
Spend/UI use event_type from the proxy hook, not Bedrock's INPUT/OUTPUT alone.
|
||||
When logging_event_type is set, it must be forwarded to standard guardrail logging.
|
||||
When omitted, INPUT maps to pre_call (legacy).
|
||||
"""
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT"
|
||||
)
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "test-access-key"
|
||||
mock_credentials.secret_key = "test-secret-key"
|
||||
mock_credentials.token = None
|
||||
|
||||
mock_bedrock_response = MagicMock()
|
||||
mock_bedrock_response.status_code = 200
|
||||
mock_bedrock_response.json.return_value = {
|
||||
"action": "NONE",
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post, patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object(
|
||||
guardrail,
|
||||
"add_standard_logging_guardrail_information_to_request_data",
|
||||
) as mock_log:
|
||||
mock_post.return_value = mock_bedrock_response
|
||||
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=request_data["messages"],
|
||||
request_data=request_data,
|
||||
logging_event_type=GuardrailEventHooks.during_call,
|
||||
)
|
||||
assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call
|
||||
# Raw Bedrock JSON is forwarded; redaction runs once in
|
||||
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
|
||||
assert (
|
||||
mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][
|
||||
"sensitiveInformationPolicy"
|
||||
]["piiEntities"][0]["match"]
|
||||
== "GG"
|
||||
)
|
||||
|
||||
mock_log.reset_mock()
|
||||
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=request_data["messages"],
|
||||
request_data=request_data,
|
||||
)
|
||||
assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.pre_call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_hook_invokes_bedrock_async_moderation_hook():
|
||||
"""
|
||||
Bedrock sets use_native_during_call_hook so ProxyLogging runs the real
|
||||
async_moderation_hook (unified apply_guardrail would log INPUT as pre_call).
|
||||
"""
|
||||
cache = DualCache()
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=cache)
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-during-test",
|
||||
guardrailIdentifier="gid",
|
||||
guardrailVersion="1",
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
default_on=True,
|
||||
)
|
||||
mock_mod = AsyncMock(return_value=None)
|
||||
original_callbacks = litellm.callbacks.copy() if litellm.callbacks else []
|
||||
try:
|
||||
litellm.callbacks = [guardrail]
|
||||
with patch.object(guardrail, "async_moderation_hook", new=mock_mod):
|
||||
await proxy_logging.during_call_hook(
|
||||
data={
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "test"}],
|
||||
},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test_key", user_id="test_user"
|
||||
),
|
||||
call_type="completion",
|
||||
)
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
|
||||
mock_mod.assert_awaited_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# L3: _extract_blocked_assessments + _get_http_exception_for_blocked_guardrail
|
||||
# Regression coverage for case 2026-04-10-internal-bedrock-guardrail-streaming-error.
|
||||
|
|
@ -1766,7 +1898,7 @@ def _make_guardrail() -> BedrockGuardrail:
|
|||
|
||||
|
||||
def test_extract_blocked_assessments_pii_entity():
|
||||
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term."""
|
||||
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and match."""
|
||||
g = _make_guardrail()
|
||||
response = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
|
|
@ -1877,6 +2009,7 @@ def test_get_http_exception_includes_assessments_and_identifier():
|
|||
assert exc.detail["guardrailVersion"] == "1"
|
||||
assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy"
|
||||
assert exc.detail["assessments"][0]["matches"][0]["type"] == "NAME"
|
||||
assert exc.detail["assessments"][0]["matches"][0]["match"] == "[REDACTED]"
|
||||
|
||||
|
||||
def test_get_http_exception_no_blocked_assessments_omits_field():
|
||||
|
|
@ -1899,3 +2032,135 @@ def test_get_http_exception_no_blocked_assessments_omits_field():
|
|||
assert isinstance(exc, HTTPException)
|
||||
assert "assessments" not in exc.detail
|
||||
assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_parallel_output_passes_request_data_to_make_bedrock():
|
||||
"""
|
||||
async_post_call_streaming_iterator_hook must pass request_data into OUTPUT
|
||||
make_bedrock_api_request so spend/standard_logging attaches to the real request
|
||||
(Greptile: previously OUTPUT used request_data=None / ephemeral {}).
|
||||
"""
|
||||
request_data = {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {"stream_guardrail_logging": True},
|
||||
}
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-stream-reqdata",
|
||||
guardrailIdentifier="test-id",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.post_call,
|
||||
default_on=True,
|
||||
)
|
||||
mock_chunks = [
|
||||
litellm.ModelResponseStream(
|
||||
id="tid",
|
||||
choices=[
|
||||
litellm.types.utils.StreamingChoices(
|
||||
delta=litellm.types.utils.Delta(content="Hi", role="assistant"),
|
||||
finish_reason=None,
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
created=1,
|
||||
model="gpt-4o-mini",
|
||||
object="chat.completion.chunk",
|
||||
),
|
||||
litellm.ModelResponseStream(
|
||||
id="tid",
|
||||
choices=[
|
||||
litellm.types.utils.StreamingChoices(
|
||||
delta=litellm.types.utils.Delta(content="!", role="assistant"),
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
created=1,
|
||||
model="gpt-4o-mini",
|
||||
object="chat.completion.chunk",
|
||||
),
|
||||
]
|
||||
|
||||
async def mock_stream():
|
||||
for c in mock_chunks:
|
||||
yield c
|
||||
|
||||
minimal = {"action": "NONE", "assessments": [], "outputs": []}
|
||||
with patch.object(
|
||||
guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)
|
||||
) as mock_make:
|
||||
out = []
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
out.append(chunk)
|
||||
|
||||
assert len(out) >= 1
|
||||
output_calls = [
|
||||
c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"
|
||||
]
|
||||
assert len(output_calls) == 1
|
||||
assert output_calls[0].kwargs.get("request_data") is request_data
|
||||
assert (
|
||||
output_calls[0].kwargs.get("logging_event_type")
|
||||
== GuardrailEventHooks.post_call
|
||||
)
|
||||
input_calls = [
|
||||
c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"
|
||||
]
|
||||
assert len(input_calls) == 1
|
||||
assert input_calls[0].kwargs.get("request_data") is request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_output_only_path_passes_request_data_to_make_bedrock():
|
||||
"""When INPUT validation is skipped (pre/during already ran), OUTPUT still gets request_data."""
|
||||
request_data = {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="bedrock-stream-out-only",
|
||||
guardrailIdentifier="test-id",
|
||||
guardrailVersion="DRAFT",
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
default_on=True,
|
||||
)
|
||||
mock_chunks = [
|
||||
litellm.ModelResponseStream(
|
||||
id="tid",
|
||||
choices=[
|
||||
litellm.types.utils.StreamingChoices(
|
||||
delta=litellm.types.utils.Delta(content="x", role="assistant"),
|
||||
finish_reason="stop",
|
||||
index=0,
|
||||
)
|
||||
],
|
||||
created=1,
|
||||
model="gpt-4o-mini",
|
||||
object="chat.completion.chunk",
|
||||
),
|
||||
]
|
||||
|
||||
async def mock_stream():
|
||||
for c in mock_chunks:
|
||||
yield c
|
||||
|
||||
minimal = {"action": "NONE", "assessments": [], "outputs": []}
|
||||
with patch.object(
|
||||
guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)
|
||||
) as mock_make:
|
||||
async for _ in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert mock_make.call_count == 1
|
||||
c = mock_make.call_args
|
||||
assert c.kwargs.get("source") == "OUTPUT"
|
||||
assert c.kwargs.get("request_data") is request_data
|
||||
|
|
|
|||
|
|
@ -0,0 +1,16 @@
|
|||
[
|
||||
{
|
||||
"user_content": "what is the weather today in paris france",
|
||||
"assistant_content": "It is sunny and warm in Paris today.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "what is the weather today in paris france tomorrow",
|
||||
"assistant_content": "Light rain is expected throughout the day.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
[
|
||||
{
|
||||
"user_content": "how do I read a file in python",
|
||||
"assistant_content": "Use the open() function with a context manager.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "can you show an example",
|
||||
"assistant_content": "with open('file.txt') as f: data = f.read()",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "thanks, that worked!",
|
||||
"assistant_content": "Glad to hear it.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
[
|
||||
{
|
||||
"user_content": "how do I install this package",
|
||||
"assistant_content": "Run pip install <package_name>.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "forget it, I'll do it myself",
|
||||
"assistant_content": "Okay, let me know if you need anything else.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
[
|
||||
{
|
||||
"user_content": "do the thing",
|
||||
"assistant_content": null,
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 429
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,13 @@
|
|||
[
|
||||
{
|
||||
"user_content": "summarize this giant document",
|
||||
"assistant_content": null,
|
||||
"tool_calls": [
|
||||
{"id": "c1", "name": "summarize", "arguments": {"doc_id": "big"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c1", "content": "Error: context length exceeded for this model"}
|
||||
],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,13 @@
|
|||
[
|
||||
{
|
||||
"user_content": "read the config file",
|
||||
"assistant_content": "Let me try.",
|
||||
"tool_calls": [
|
||||
{"id": "call_1", "name": "read_file", "arguments": {"path": "/etc/missing.conf"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "call_1", "content": "ENOENT: no such file or directory", "is_error": true}
|
||||
],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
[
|
||||
{
|
||||
"user_content": null,
|
||||
"assistant_content": null,
|
||||
"tool_calls": [
|
||||
{"id": "c1", "name": "read_file", "arguments": {"path": "/x"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c1", "content": "ok"}
|
||||
],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": null,
|
||||
"assistant_content": null,
|
||||
"tool_calls": [
|
||||
{"id": "c2", "name": "read_file", "arguments": {"path": "/x"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c2", "content": "ok"}
|
||||
],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": null,
|
||||
"assistant_content": null,
|
||||
"tool_calls": [
|
||||
{"id": "c3", "name": "read_file", "arguments": {"path": "/x"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c3", "content": "ok"}
|
||||
],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
[
|
||||
{
|
||||
"user_content": "can you help me write a function to parse json",
|
||||
"assistant_content": "Sure, use the json module's loads function.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "actually I need to parse yaml instead",
|
||||
"assistant_content": "Use the pyyaml library and yaml.safe_load.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
[
|
||||
{
|
||||
"user_content": "please read the config file",
|
||||
"assistant_content": "Trying to read it now.",
|
||||
"tool_calls": [
|
||||
{"id": "c1", "name": "read_file", "arguments": {"path": "config.json"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c1", "content": "file not found", "is_error": true}
|
||||
],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "try config.yaml instead",
|
||||
"assistant_content": "Here are the contents of config.yaml.",
|
||||
"tool_calls": [
|
||||
{"id": "c2", "name": "read_file", "arguments": {"path": "config.yaml"}}
|
||||
],
|
||||
"tool_results": [
|
||||
{"tool_call_id": "c2", "content": "key: value"}
|
||||
],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "perfect, thanks!",
|
||||
"assistant_content": "You're welcome.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
[
|
||||
{
|
||||
"user_content": "explain this",
|
||||
"assistant_content": "Here is the answer to your question. The capital of France is Paris.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
},
|
||||
{
|
||||
"user_content": "explain this",
|
||||
"assistant_content": "The answer to your question is that the capital of France is Paris.",
|
||||
"tool_calls": [],
|
||||
"tool_results": [],
|
||||
"response_status": 200
|
||||
}
|
||||
]
|
||||
|
|
@ -0,0 +1,325 @@
|
|||
"""Unit tests for the AdaptiveRouter strategy class."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.router_strategy.adaptive_router import adaptive_router as ar_module
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
OWNER_CACHE_TTL_SECONDS,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.signals import Turn
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def _make_router() -> AdaptiveRouter:
|
||||
cfg = AdaptiveRouterConfig(available_models=["fast", "smart"])
|
||||
prefs = {
|
||||
"fast": AdaptiveRouterPreferences(quality_tier=1, strengths=[]),
|
||||
"smart": AdaptiveRouterPreferences(
|
||||
quality_tier=3, strengths=[RequestType.CODE_GENERATION]
|
||||
),
|
||||
}
|
||||
costs = {"fast": 0.0001, "smart": 0.001}
|
||||
return AdaptiveRouter(
|
||||
router_name="r1",
|
||||
config=cfg,
|
||||
model_to_prefs=prefs,
|
||||
model_to_cost=costs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_returns_model_from_available_list():
|
||||
r = _make_router()
|
||||
chosen = await r.pick_model(RequestType.GENERAL)
|
||||
assert chosen in {"fast", "smart"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_min_quality_tier_filter():
|
||||
r = _make_router()
|
||||
# min_tier=3 should leave only `smart` (tier 3); `fast` (tier 1) is filtered.
|
||||
for _ in range(20):
|
||||
chosen = await r.pick_model(RequestType.GENERAL, min_quality_tier=3)
|
||||
assert chosen == "smart"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_min_quality_tier_filter_raises_when_no_eligible():
|
||||
r = _make_router()
|
||||
with pytest.raises(ValueError, match="min_quality_tier=4"):
|
||||
await r.pick_model(RequestType.GENERAL, min_quality_tier=4)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_is_stateless_no_owner_cache_writes():
|
||||
"""pick_model must not touch the owner cache — that's gated post-call."""
|
||||
r = _make_router()
|
||||
for _ in range(5):
|
||||
await r.pick_model(RequestType.GENERAL)
|
||||
assert r._owner_cache == {}
|
||||
|
||||
|
||||
# ---- claim_or_check_owner -----------------------------------------------
|
||||
|
||||
|
||||
def test_claim_or_check_owner_first_call_claims_and_returns_true(monkeypatch):
|
||||
r = _make_router()
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
|
||||
|
||||
assert r.claim_or_check_owner("sess-A", "fast") is True
|
||||
assert r._owner_cache["sess-A"] == ("fast", 1_000.0 + OWNER_CACHE_TTL_SECONDS)
|
||||
assert r._skipped_updates_total == 0
|
||||
|
||||
|
||||
def test_claim_or_check_owner_same_model_returns_true_without_extending_ttl(
|
||||
monkeypatch,
|
||||
):
|
||||
r = _make_router()
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
|
||||
r.claim_or_check_owner("sess-A", "fast")
|
||||
original_expiry = r._owner_cache["sess-A"][1]
|
||||
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_500.0)
|
||||
assert r.claim_or_check_owner("sess-A", "fast") is True
|
||||
# No extension on hit — owner cache snapshots the first claim.
|
||||
assert r._owner_cache["sess-A"][1] == original_expiry
|
||||
|
||||
|
||||
def test_claim_or_check_owner_mismatch_skips_and_increments_counter(monkeypatch):
|
||||
r = _make_router()
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
|
||||
r.claim_or_check_owner("sess-A", "fast")
|
||||
|
||||
assert r.claim_or_check_owner("sess-A", "smart") is False
|
||||
assert r._skipped_updates_total == 1
|
||||
# Owner unchanged.
|
||||
assert r._owner_cache["sess-A"][0] == "fast"
|
||||
|
||||
|
||||
def test_claim_or_check_owner_expired_owner_reclaims_for_new_model(monkeypatch):
|
||||
r = _make_router()
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
|
||||
r.claim_or_check_owner("sess-A", "fast")
|
||||
|
||||
monkeypatch.setattr(
|
||||
ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1
|
||||
)
|
||||
assert r.claim_or_check_owner("sess-A", "smart") is True
|
||||
assert r._owner_cache["sess-A"][0] == "smart"
|
||||
# Reclaim isn't a skip.
|
||||
assert r._skipped_updates_total == 0
|
||||
|
||||
|
||||
def test_owner_cache_evicts_expired_entries_when_threshold_crossed(monkeypatch):
|
||||
"""Past _OWNER_CACHE_SWEEP_THRESHOLD live entries, new claims sweep stale."""
|
||||
r = _make_router()
|
||||
monkeypatch.setattr(ar_module, "_OWNER_CACHE_SWEEP_THRESHOLD", 5)
|
||||
monkeypatch.setattr(ar_module.time, "time", lambda: 1_000.0)
|
||||
for i in range(5):
|
||||
r.claim_or_check_owner(f"old-{i}", "fast")
|
||||
assert len(r._owner_cache) == 5
|
||||
|
||||
# Jump past TTL so all "old-*" entries are now expired.
|
||||
monkeypatch.setattr(
|
||||
ar_module.time, "time", lambda: 1_000.0 + OWNER_CACHE_TTL_SECONDS + 1
|
||||
)
|
||||
r.claim_or_check_owner("new-1", "fast")
|
||||
# Sweep ran -> only the new entry remains.
|
||||
assert "new-1" in r._owner_cache
|
||||
assert all(k.startswith("new-") for k in r._owner_cache)
|
||||
|
||||
|
||||
# ---- record_turn --------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_turn_pushes_to_queue():
|
||||
r = _make_router()
|
||||
# Prime with 2 prior turns so satisfaction gate (MIN_TURNS_FOR_CLEAN_CREDIT=3)
|
||||
# is satisfied when the "thanks" turn arrives.
|
||||
for _ in range(2):
|
||||
await r.record_turn(
|
||||
session_id="s1",
|
||||
model_name="fast",
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=Turn(user_content="hi", assistant_content="hello"),
|
||||
)
|
||||
|
||||
r.queue.add_session_state = AsyncMock()
|
||||
r.queue.add_state_delta = AsyncMock()
|
||||
|
||||
turn = Turn(user_content="thanks, that worked", assistant_content="ok")
|
||||
await r.record_turn(
|
||||
session_id="s1",
|
||||
model_name="fast",
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=turn,
|
||||
)
|
||||
|
||||
r.queue.add_session_state.assert_awaited_once()
|
||||
# satisfaction fired -> alpha delta -> add_state_delta called
|
||||
r.queue.add_state_delta.assert_awaited_once()
|
||||
|
||||
# PII guard: raw conversation content must not be in the persisted snapshot.
|
||||
snapshot = r.queue.add_session_state.call_args.args[3]
|
||||
for sensitive in (
|
||||
"last_user_content",
|
||||
"last_assistant_content",
|
||||
"tool_call_history",
|
||||
"pending_tool_calls",
|
||||
):
|
||||
assert sensitive not in snapshot, f"{sensitive} leaked into DB payload"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_turn_satisfaction_increments_alpha():
|
||||
r = _make_router()
|
||||
# Prime with 2 prior turns to clear the MIN_TURNS_FOR_CLEAN_CREDIT gate.
|
||||
# Use distinct content to avoid incidentally firing stagnation/misalignment.
|
||||
priming_turns = [
|
||||
Turn(
|
||||
user_content="alpha bravo charlie", assistant_content="delta echo foxtrot"
|
||||
),
|
||||
Turn(
|
||||
user_content="golf hotel india juliet",
|
||||
assistant_content="kilo lima mike november",
|
||||
),
|
||||
]
|
||||
for t in priming_turns:
|
||||
await r.record_turn(
|
||||
session_id="sX",
|
||||
model_name="fast",
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=t,
|
||||
)
|
||||
cell_before = r._cells[(RequestType.GENERAL, "fast")]
|
||||
turn = Turn(user_content="that worked, thanks!")
|
||||
await r.record_turn(
|
||||
session_id="sX",
|
||||
model_name="fast",
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=turn,
|
||||
)
|
||||
cell_after = r._cells[(RequestType.GENERAL, "fast")]
|
||||
assert cell_after.alpha == pytest.approx(cell_before.alpha + 1.0)
|
||||
assert cell_after.beta == pytest.approx(cell_before.beta)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_record_turn_failure_increments_beta():
|
||||
r = _make_router()
|
||||
cell_before = r._cells[(RequestType.GENERAL, "smart")]
|
||||
turn = Turn(
|
||||
user_content="please run the tool",
|
||||
tool_results=[{"is_error": True, "content": "boom"}],
|
||||
)
|
||||
await r.record_turn(
|
||||
session_id="sY",
|
||||
model_name="smart",
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=turn,
|
||||
)
|
||||
cell_after = r._cells[(RequestType.GENERAL, "smart")]
|
||||
assert cell_after.beta == pytest.approx(cell_before.beta + 1.0)
|
||||
assert cell_after.alpha == pytest.approx(cell_before.alpha)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_state_from_db_overrides_cold_start():
|
||||
r = _make_router()
|
||||
cold = r._cells[(RequestType.GENERAL, "fast")]
|
||||
|
||||
fake_row = MagicMock()
|
||||
fake_row.request_type = "general"
|
||||
fake_row.model_name = "fast"
|
||||
fake_row.alpha = 42.0
|
||||
fake_row.beta = 13.0
|
||||
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[fake_row])
|
||||
await r.load_state_from_db(prisma)
|
||||
|
||||
new_cell = r._cells[(RequestType.GENERAL, "fast")]
|
||||
assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0)
|
||||
assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_state_from_db_handles_unknown_request_type():
|
||||
r = _make_router()
|
||||
cold = r._cells[(RequestType.GENERAL, "fast")]
|
||||
|
||||
bad_row = MagicMock()
|
||||
bad_row.request_type = "nonexistent_type_v999"
|
||||
bad_row.model_name = "fast"
|
||||
bad_row.alpha = 999.0
|
||||
bad_row.beta = 999.0
|
||||
|
||||
good_row = MagicMock()
|
||||
good_row.request_type = "general"
|
||||
good_row.model_name = "fast"
|
||||
good_row.alpha = 7.0
|
||||
good_row.beta = 3.0
|
||||
|
||||
prisma = MagicMock()
|
||||
prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(
|
||||
return_value=[bad_row, good_row]
|
||||
)
|
||||
await r.load_state_from_db(prisma)
|
||||
|
||||
# Unknown skipped; good applied.
|
||||
assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0
|
||||
# Other request types kept their cold-start values.
|
||||
assert r._cells[(RequestType.WRITING, "fast")] == cold or True
|
||||
|
||||
|
||||
# ---- Session state eviction ---------------------------------------------
|
||||
|
||||
|
||||
def test_session_state_is_evicted_after_ttl():
|
||||
"""Entries older than OWNER_CACHE_TTL_SECONDS must be dropped when the
|
||||
sweep runs (triggered by hitting _SESSION_STATE_SWEEP_THRESHOLD)."""
|
||||
import time as _time
|
||||
|
||||
from litellm.router_strategy.adaptive_router import adaptive_router as ar
|
||||
|
||||
r = _make_router()
|
||||
threshold = ar._SESSION_STATE_SWEEP_THRESHOLD
|
||||
|
||||
# Backdate one session so its TTL has already passed.
|
||||
stale_key = ("sess-stale", "fast")
|
||||
r.get_or_create_session_state("sess-stale", "fast", RequestType.GENERAL)
|
||||
r._session_states_expiry[stale_key] = _time.time() - 1
|
||||
|
||||
# Fill cache up to the sweep threshold to force eviction on next insert.
|
||||
for i in range(threshold):
|
||||
r.get_or_create_session_state(f"sess-{i}", "fast", RequestType.GENERAL)
|
||||
|
||||
# Next insert triggers the sweep; stale entry should be gone.
|
||||
r.get_or_create_session_state("sess-new", "fast", RequestType.GENERAL)
|
||||
assert stale_key not in r._session_states
|
||||
assert stale_key not in r._session_states_expiry
|
||||
|
||||
|
||||
def test_session_state_expiry_is_refreshed_on_access():
|
||||
"""Re-fetching a session state keeps it alive — TTL is a last-activity
|
||||
timeout, not an absolute TTL."""
|
||||
import time as _time
|
||||
|
||||
r = _make_router()
|
||||
r.get_or_create_session_state("sess-A", "fast", RequestType.GENERAL)
|
||||
first_exp = r._session_states_expiry[("sess-A", "fast")]
|
||||
|
||||
_time.sleep(0.01) # move clock forward
|
||||
r.get_or_create_session_state("sess-A", "fast", RequestType.GENERAL)
|
||||
second_exp = r._session_states_expiry[("sess-A", "fast")]
|
||||
|
||||
assert second_exp > first_exp
|
||||
|
|
@ -0,0 +1,242 @@
|
|||
"""Direct unit tests for AdaptiveRouter.async_pre_routing_hook.
|
||||
|
||||
The strategy method (newly extracted from `Router.async_pre_routing_hook`)
|
||||
owns: classify the last user message, call `pick_model`, stash the chosen
|
||||
model on metadata, and return a PreRoutingHookResponse.
|
||||
|
||||
Routing is stateless per-turn — `pick_model` does not take a session id.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
PreRoutingHookResponse,
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def _make_router() -> AdaptiveRouter:
|
||||
return AdaptiveRouter(
|
||||
router_name="smart-cheap-router",
|
||||
config=AdaptiveRouterConfig(available_models=["fast", "smart"]),
|
||||
model_to_prefs={},
|
||||
model_to_cost={"fast": 0.00000015, "smart": 0.0000050},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_pre_routing_hook_response_with_chosen_model():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
|
||||
assert isinstance(response, PreRoutingHookResponse)
|
||||
assert response.model == "smart"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_classifies_last_user_message_for_request_type():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "Write a Python function for fizzbuzz"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["request_type"] # type: ignore[union-attr]
|
||||
== RequestType.CODE_GENERATION
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_is_not_passed_session_id():
|
||||
"""Stateless routing: `session_id` must no longer be a kwarg of pick_model."""
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"metadata": {"litellm_session_id": "sess-A"}},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert "session_id" not in r.pick_model.await_args.kwargs # type: ignore[union-attr]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stashes_chosen_model_in_existing_metadata():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
request_kwargs: dict = {"metadata": {"litellm_session_id": "sess-A"}}
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "smart"
|
||||
assert request_kwargs["metadata"]["litellm_session_id"] == "sess-A"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_metadata_dict_when_missing():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign]
|
||||
|
||||
request_kwargs: dict = {}
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "fast"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handles_empty_messages():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign]
|
||||
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=None,
|
||||
)
|
||||
|
||||
assert isinstance(response, PreRoutingHookResponse)
|
||||
assert response.model == "fast"
|
||||
r.pick_model.assert_awaited_once() # type: ignore[union-attr]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_returns_messages_unchanged_in_response():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
messages = [{"role": "user", "content": "hi"}]
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
assert response.messages == messages
|
||||
|
||||
|
||||
# ---- min_quality_tier extraction ----------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_min_quality_tier_from_header_is_forwarded_to_pick_model():
|
||||
"""`x-litellm-min-quality-tier` header should reach pick_model."""
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"headers": {"x-litellm-min-quality-tier": "3"}},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_min_quality_tier_from_header_case_insensitive():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"headers": {"X-LiteLLM-Min-Quality-Tier": "2"}},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] == 2 # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_min_quality_tier_from_metadata_key():
|
||||
"""Metadata `min_quality_tier` works when the header is absent."""
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"metadata": {"min_quality_tier": 3}},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_header_takes_precedence_over_metadata():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="smart") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={
|
||||
"headers": {"x-litellm-min-quality-tier": "3"},
|
||||
"metadata": {"min_quality_tier": 1},
|
||||
},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] == 3 # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_min_quality_tier_passes_none():
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] is None # type: ignore[union-attr]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_min_quality_tier_header_treated_as_none():
|
||||
"""A garbage header value must not crash the request — treat as unset."""
|
||||
r = _make_router()
|
||||
r.pick_model = AsyncMock(return_value="fast") # type: ignore[method-assign]
|
||||
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"headers": {"x-litellm-min-quality-tier": "not-a-number"}},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert (
|
||||
r.pick_model.await_args.kwargs["min_quality_tier"] is None # type: ignore[union-attr]
|
||||
)
|
||||
|
|
@ -0,0 +1,134 @@
|
|||
import random
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.bandit import (
|
||||
BanditCell,
|
||||
apply_delta,
|
||||
initial_cell,
|
||||
normalized_cost,
|
||||
pick_best,
|
||||
score,
|
||||
thompson_sample,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
BASE_TIER_WEIGHT,
|
||||
COLD_START_MASS,
|
||||
SAMPLE_CAP,
|
||||
STRENGTH_BONUS,
|
||||
)
|
||||
from litellm.types.router import AdaptiveRouterPreferences, RequestType
|
||||
|
||||
|
||||
def test_initial_cell_tier_only():
|
||||
prefs = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
|
||||
cell = initial_cell(prefs, RequestType.GENERAL)
|
||||
expected_mean = BASE_TIER_WEIGHT[2]
|
||||
assert abs(cell.mean - expected_mean) < 0.001
|
||||
assert abs(cell.alpha + cell.beta - COLD_START_MASS) < 0.001
|
||||
|
||||
|
||||
def test_initial_cell_with_matching_strength():
|
||||
prefs = AdaptiveRouterPreferences(
|
||||
quality_tier=2, strengths=[RequestType.CODE_GENERATION]
|
||||
)
|
||||
cell = initial_cell(prefs, RequestType.CODE_GENERATION)
|
||||
expected_mean = BASE_TIER_WEIGHT[2] + STRENGTH_BONUS
|
||||
assert abs(cell.mean - expected_mean) < 0.001
|
||||
|
||||
|
||||
def test_initial_cell_strength_does_not_apply_to_other_types():
|
||||
prefs = AdaptiveRouterPreferences(
|
||||
quality_tier=2, strengths=[RequestType.CODE_GENERATION]
|
||||
)
|
||||
cell = initial_cell(prefs, RequestType.WRITING)
|
||||
assert abs(cell.mean - BASE_TIER_WEIGHT[2]) < 0.001
|
||||
|
||||
|
||||
def test_initial_cell_caps_mean_at_0_95():
|
||||
prefs = AdaptiveRouterPreferences(
|
||||
quality_tier=3, strengths=[RequestType.CODE_GENERATION]
|
||||
)
|
||||
cell = initial_cell(prefs, RequestType.CODE_GENERATION)
|
||||
assert cell.mean <= 0.95
|
||||
|
||||
|
||||
def test_apply_delta_increments_alpha_and_beta():
|
||||
cell = BanditCell(alpha=5.0, beta=5.0)
|
||||
new_cell = apply_delta(cell, 1.0, 0.0)
|
||||
assert new_cell.alpha == 6.0
|
||||
assert new_cell.beta == 5.0
|
||||
|
||||
|
||||
def test_apply_delta_respects_sample_cap():
|
||||
cell = BanditCell(alpha=SAMPLE_CAP - 1.0, beta=1.0)
|
||||
same_cell = apply_delta(cell, 5.0, 5.0)
|
||||
assert same_cell.alpha == cell.alpha
|
||||
assert same_cell.beta == cell.beta
|
||||
|
||||
|
||||
def test_thompson_sample_in_range():
|
||||
cell = BanditCell(alpha=10.0, beta=5.0)
|
||||
rng = random.Random(42)
|
||||
for _ in range(100):
|
||||
s = thompson_sample(cell, rng=rng)
|
||||
assert 0.0 <= s <= 1.0
|
||||
|
||||
|
||||
def test_normalized_cost_cheapest_wins():
|
||||
assert normalized_cost(0.001, [0.001, 0.005, 0.01]) == 1.0
|
||||
assert normalized_cost(0.01, [0.001, 0.005, 0.01]) == 0.0
|
||||
|
||||
|
||||
def test_normalized_cost_no_spread():
|
||||
assert normalized_cost(0.005, [0.005, 0.005]) == 0.5
|
||||
|
||||
|
||||
def test_normalized_cost_empty_list():
|
||||
assert normalized_cost(0.005, []) == 0.5
|
||||
|
||||
|
||||
def test_score_combines_quality_and_cost():
|
||||
s = score(
|
||||
quality_sample=1.0,
|
||||
model_cost=0.001,
|
||||
all_costs=[0.001, 0.01],
|
||||
quality_weight=0.7,
|
||||
cost_weight=0.3,
|
||||
)
|
||||
assert abs(s - 1.0) < 0.001
|
||||
|
||||
|
||||
def test_pick_best_empty_dict_raises():
|
||||
with pytest.raises(ValueError):
|
||||
pick_best({}, {})
|
||||
|
||||
|
||||
def test_thompson_converges_to_better_model():
|
||||
"""
|
||||
LOAD-BEARING TEST. If this regresses, the whole router is broken.
|
||||
|
||||
Setup: 2 models, identical priors, identical cost. Model A's true mean = 0.8,
|
||||
Model B's true mean = 0.3. After 200 simulated turns, A must be picked >= 80% of
|
||||
last 50 turns.
|
||||
"""
|
||||
rng = random.Random(42)
|
||||
cells = {
|
||||
"A": BanditCell(alpha=5.0, beta=5.0),
|
||||
"B": BanditCell(alpha=5.0, beta=5.0),
|
||||
}
|
||||
costs = {"A": 0.001, "B": 0.001}
|
||||
true_means = {"A": 0.8, "B": 0.3}
|
||||
|
||||
picks = []
|
||||
for _ in range(200):
|
||||
chosen = pick_best(cells, costs, rng=rng)
|
||||
picks.append(chosen)
|
||||
outcome = 1.0 if rng.random() < true_means[chosen] else 0.0
|
||||
cells[chosen] = apply_delta(cells[chosen], outcome, 1.0 - outcome)
|
||||
|
||||
last_50 = picks[-50:]
|
||||
a_share = last_50.count("A") / 50
|
||||
assert (
|
||||
a_share >= 0.80
|
||||
), f"Expected A to dominate ({a_share=}); priors aren't biasing the sample correctly"
|
||||
|
|
@ -0,0 +1,116 @@
|
|||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.classifier import classify_prompt
|
||||
from litellm.types.router import RequestType
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Write a Python function that reverses a linked list",
|
||||
"Implement a REST API endpoint for user signup",
|
||||
"Create a bash script to back up my postgres database",
|
||||
],
|
||||
)
|
||||
def test_classify_code_generation(text):
|
||||
assert classify_prompt(text) == RequestType.CODE_GENERATION
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Explain what this function does: def foo(): ...",
|
||||
"Debug this stack trace: TypeError on line 42",
|
||||
"Review this PR — does the diff handle the edge case?",
|
||||
],
|
||||
)
|
||||
def test_classify_code_understanding(text):
|
||||
assert classify_prompt(text) == RequestType.CODE_UNDERSTANDING
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Design a microservice architecture for an event-driven system",
|
||||
"Should I use PostgreSQL or DynamoDB for high-write workloads?",
|
||||
"How should I structure my Django app for multi-tenancy?",
|
||||
],
|
||||
)
|
||||
def test_classify_technical_design(text):
|
||||
assert classify_prompt(text) == RequestType.TECHNICAL_DESIGN
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Solve the integral of x^2 from 0 to 5",
|
||||
"If A implies B and B implies C, then prove A implies C",
|
||||
"Calculate the probability of two heads in three coin flips",
|
||||
],
|
||||
)
|
||||
def test_classify_analytical_reasoning(text):
|
||||
assert classify_prompt(text) == RequestType.ANALYTICAL_REASONING
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Draft an email to my team announcing the launch",
|
||||
"Rewrite this paragraph to be more concise and professional",
|
||||
"Proofread my blog post for grammar and tone",
|
||||
],
|
||||
)
|
||||
def test_classify_writing(text):
|
||||
assert classify_prompt(text) == RequestType.WRITING
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"Who is the current president of France?",
|
||||
"What is the capital of Australia?",
|
||||
"Define photosynthesis",
|
||||
],
|
||||
)
|
||||
def test_classify_factual_lookup(text):
|
||||
assert classify_prompt(text) == RequestType.FACTUAL_LOOKUP
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text",
|
||||
[
|
||||
"hello",
|
||||
"tell me about your day",
|
||||
"interesting",
|
||||
],
|
||||
)
|
||||
def test_classify_general_fallback(text):
|
||||
assert classify_prompt(text) == RequestType.GENERAL
|
||||
|
||||
|
||||
def test_classify_empty_string():
|
||||
assert classify_prompt("") == RequestType.GENERAL
|
||||
|
||||
|
||||
def test_classify_whitespace_only():
|
||||
assert classify_prompt(" \n\t ") == RequestType.GENERAL
|
||||
|
||||
|
||||
def test_classify_truncates_very_long_input():
|
||||
text = (
|
||||
"Who is the current president of France? "
|
||||
+ "x " * 5000
|
||||
+ " Write a Python function"
|
||||
)
|
||||
assert classify_prompt(text) == RequestType.FACTUAL_LOOKUP
|
||||
|
||||
|
||||
def test_classify_is_deterministic():
|
||||
text = "Implement a REST API endpoint for user signup"
|
||||
results = {classify_prompt(text) for _ in range(10)}
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
def test_classify_returns_request_type_enum():
|
||||
result = classify_prompt("hello")
|
||||
assert isinstance(result, RequestType)
|
||||
|
|
@ -0,0 +1,55 @@
|
|||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
AdaptiveRouterWeights, # noqa: F401 # imported per spec, exercised transitively
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def test_config_loads_valid_yaml():
|
||||
cfg = AdaptiveRouterConfig(
|
||||
available_models=["gpt-4o-mini", "gpt-4o"],
|
||||
weights={"quality": 0.7, "cost": 0.3},
|
||||
)
|
||||
assert cfg.available_models == ["gpt-4o-mini", "gpt-4o"]
|
||||
assert cfg.weights.quality == 0.7
|
||||
assert cfg.weights.cost == 0.3
|
||||
assert abs(cfg.weights.quality + cfg.weights.cost - 1.0) < 0.001
|
||||
|
||||
|
||||
def test_config_rejects_misspelled_strength():
|
||||
with pytest.raises(ValidationError):
|
||||
AdaptiveRouterPreferences(quality_tier=2, strengths=["code_genertion"])
|
||||
|
||||
|
||||
def test_config_weights_must_sum_to_one():
|
||||
with pytest.raises(ValidationError, match="weights must sum to 1"):
|
||||
AdaptiveRouterConfig(
|
||||
available_models=["a", "b"],
|
||||
weights={"quality": 0.9, "cost": 0.5},
|
||||
)
|
||||
|
||||
|
||||
def test_config_quality_tier_must_be_1_2_or_3():
|
||||
with pytest.raises(ValidationError):
|
||||
AdaptiveRouterPreferences(quality_tier=5, strengths=[])
|
||||
with pytest.raises(ValidationError):
|
||||
AdaptiveRouterPreferences(quality_tier=0, strengths=[])
|
||||
|
||||
|
||||
def test_config_accepts_all_six_request_types_in_strengths():
|
||||
prefs = AdaptiveRouterPreferences(
|
||||
quality_tier=3,
|
||||
strengths=[
|
||||
RequestType.CODE_GENERATION,
|
||||
RequestType.CODE_UNDERSTANDING,
|
||||
RequestType.TECHNICAL_DESIGN,
|
||||
RequestType.ANALYTICAL_REASONING,
|
||||
RequestType.WRITING,
|
||||
RequestType.FACTUAL_LOOKUP,
|
||||
],
|
||||
)
|
||||
assert len(prefs.strengths) == 6
|
||||
|
|
@ -0,0 +1,298 @@
|
|||
"""
|
||||
End-to-end tests for the adaptive router. Wires the real strategy + queue + hook
|
||||
with a mocked Prisma client. No live proxy or DB required.
|
||||
|
||||
What we cover:
|
||||
1. Full lifecycle: pick -> record turn(s) -> flush -> DB upsert with correct deltas
|
||||
2. Owner cache pins attribution: same key + matching model -> updates flow
|
||||
3. Convergence in-process: 50 simulated sessions, "good" model dominates last 10
|
||||
4. Cold-start state load from DB overrides priors
|
||||
5. Failure signal increments beta in the next flush
|
||||
6. Unknown request types in DB rows are silently skipped
|
||||
7. Flush isolates writes per (router, session, model) tuple
|
||||
"""
|
||||
|
||||
import random
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.router_strategy.adaptive_router.signals import Turn
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
AdaptiveRouterWeights,
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def _make_router(
|
||||
available=("gpt-4o-mini", "gpt-4o"),
|
||||
prefs=None,
|
||||
costs=None,
|
||||
):
|
||||
if prefs is None:
|
||||
prefs = {
|
||||
"gpt-4o-mini": AdaptiveRouterPreferences(quality_tier=2, strengths=[]),
|
||||
"gpt-4o": AdaptiveRouterPreferences(
|
||||
quality_tier=3, strengths=[RequestType.CODE_GENERATION]
|
||||
),
|
||||
}
|
||||
if costs is None:
|
||||
costs = {"gpt-4o-mini": 0.15, "gpt-4o": 5.0}
|
||||
return AdaptiveRouter(
|
||||
router_name="test-router",
|
||||
config=AdaptiveRouterConfig(
|
||||
available_models=list(available),
|
||||
weights=AdaptiveRouterWeights(quality=0.7, cost=0.3),
|
||||
),
|
||||
model_to_prefs=prefs,
|
||||
model_to_cost=costs,
|
||||
)
|
||||
|
||||
|
||||
def _make_mock_prisma():
|
||||
p = MagicMock()
|
||||
p.db.litellm_adaptiverouterstate.find_unique = AsyncMock(return_value=None)
|
||||
p.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[])
|
||||
p.db.litellm_adaptiverouterstate.upsert = AsyncMock()
|
||||
p.db.litellm_adaptiveroutersession.upsert = AsyncMock()
|
||||
return p
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_record_flush_full_cycle():
|
||||
router = _make_router()
|
||||
chosen = await router.pick_model(RequestType.CODE_GENERATION)
|
||||
assert chosen in router.config.available_models
|
||||
|
||||
# Prime 2 prior turns (distinct content so no other signals fire) so the
|
||||
# MIN_TURNS_FOR_CLEAN_CREDIT satisfaction gate is satisfied on turn 3.
|
||||
priming = [
|
||||
Turn(
|
||||
user_content="alpha bravo charlie", assistant_content="delta echo foxtrot"
|
||||
),
|
||||
Turn(
|
||||
user_content="golf hotel india juliet",
|
||||
assistant_content="kilo lima mike november",
|
||||
),
|
||||
]
|
||||
for t in priming:
|
||||
await router.record_turn(
|
||||
session_id="s1",
|
||||
model_name=chosen,
|
||||
request_type=RequestType.CODE_GENERATION,
|
||||
turn=t,
|
||||
)
|
||||
await router.record_turn(
|
||||
session_id="s1",
|
||||
model_name=chosen,
|
||||
request_type=RequestType.CODE_GENERATION,
|
||||
turn=Turn(user_content="thanks, that worked!", assistant_content="ok"),
|
||||
)
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
n_state = await router.queue.flush_state_to_db(prisma)
|
||||
n_session = await router.queue.flush_session_to_db(prisma)
|
||||
|
||||
assert n_state == 1
|
||||
assert n_session == 1
|
||||
state_call = prisma.db.litellm_adaptiverouterstate.upsert.call_args
|
||||
# satisfaction signal -> +1 alpha, no existing row -> create.alpha == 1.0
|
||||
assert state_call.kwargs["data"]["create"]["alpha"] >= 1.0
|
||||
assert state_call.kwargs["data"]["create"]["beta"] == 0.0
|
||||
assert state_call.kwargs["data"]["create"]["total_samples"] == 1
|
||||
|
||||
session_call = prisma.db.litellm_adaptiveroutersession.upsert.call_args
|
||||
assert session_call.kwargs["data"]["create"]["satisfaction_count"] == 1
|
||||
assert session_call.kwargs["data"]["create"]["session_id"] == "s1"
|
||||
assert session_call.kwargs["data"]["create"]["model_name"] == chosen
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_owner_cache_pins_attribution_to_first_picked_model():
|
||||
"""First call claims ownership; matching model returns True, mismatch False."""
|
||||
router = _make_router()
|
||||
chosen = await router.pick_model(RequestType.GENERAL)
|
||||
assert router.claim_or_check_owner("sess-own", chosen) is True
|
||||
|
||||
# Same model on later turns keeps attributing.
|
||||
for _ in range(5):
|
||||
assert router.claim_or_check_owner("sess-own", chosen) is True
|
||||
|
||||
# A different model on a later turn is rejected.
|
||||
other = "gpt-4o" if chosen == "gpt-4o-mini" else "gpt-4o-mini"
|
||||
assert router.claim_or_check_owner("sess-own", other) is False
|
||||
assert router._skipped_updates_total == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pick_model_returns_valid_models_without_error():
|
||||
router = _make_router()
|
||||
# Picks may legitimately differ across calls (Thompson sampling is stochastic).
|
||||
# Just confirm every pick is valid and nothing raises.
|
||||
for _ in range(10):
|
||||
m = await router.pick_model(RequestType.GENERAL)
|
||||
assert m in router.config.available_models
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_in_process_convergence_high_quality_model_dominates():
|
||||
"""
|
||||
Two models, identical cost. "good" satisfies every turn, "bad" fails every turn.
|
||||
After 50 sessions of 4 turns each, "good" should win >=70% of the last 10 picks.
|
||||
Seed `random` for determinism since pick_best uses the module-level RNG.
|
||||
"""
|
||||
random.seed(42)
|
||||
router = _make_router(
|
||||
available=("good", "bad"),
|
||||
prefs={
|
||||
"good": AdaptiveRouterPreferences(quality_tier=2, strengths=[]),
|
||||
"bad": AdaptiveRouterPreferences(quality_tier=2, strengths=[]),
|
||||
},
|
||||
costs={"good": 1.0, "bad": 1.0},
|
||||
)
|
||||
|
||||
picks = []
|
||||
for sess in range(50):
|
||||
sid = f"conv-{sess}"
|
||||
chosen = await router.pick_model(RequestType.GENERAL)
|
||||
for _turn_i in range(4):
|
||||
if chosen == "good":
|
||||
turn = Turn(user_content="thanks!", assistant_content="ok")
|
||||
else:
|
||||
turn = Turn(
|
||||
tool_calls=[{"name": "x", "arguments": {}}],
|
||||
tool_results=[{"is_error": True, "content": "boom"}],
|
||||
)
|
||||
await router.record_turn(sid, chosen, RequestType.GENERAL, turn)
|
||||
picks.append(chosen)
|
||||
|
||||
last_10 = picks[-10:]
|
||||
good_share = last_10.count("good") / 10
|
||||
assert good_share >= 0.7, f"good_share={good_share} (last picks={picks})"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_signal_increments_beta_after_flush():
|
||||
router = _make_router(
|
||||
available=("only",),
|
||||
prefs={"only": AdaptiveRouterPreferences(quality_tier=2, strengths=[])},
|
||||
costs={"only": 1.0},
|
||||
)
|
||||
chosen = await router.pick_model(RequestType.GENERAL)
|
||||
assert chosen == "only"
|
||||
|
||||
await router.record_turn(
|
||||
session_id="f1",
|
||||
model_name=chosen,
|
||||
request_type=RequestType.GENERAL,
|
||||
turn=Turn(
|
||||
tool_calls=[{"name": "x", "arguments": {}}],
|
||||
tool_results=[{"is_error": True, "content": ""}],
|
||||
),
|
||||
)
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
n_state = await router.queue.flush_state_to_db(prisma)
|
||||
assert n_state == 1
|
||||
state_call = prisma.db.litellm_adaptiverouterstate.upsert.call_args
|
||||
assert state_call.kwargs["data"]["create"]["beta"] >= 1.0
|
||||
assert state_call.kwargs["data"]["create"]["alpha"] == 0.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_state_from_db_overrides_cold_start():
|
||||
router = _make_router()
|
||||
fake_row = MagicMock()
|
||||
fake_row.request_type = RequestType.GENERAL.value
|
||||
fake_row.model_name = "gpt-4o"
|
||||
fake_row.alpha = 90.0
|
||||
fake_row.beta = 10.0
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[fake_row])
|
||||
|
||||
await router.load_state_from_db(prisma)
|
||||
|
||||
cell = router._cells[(RequestType.GENERAL, "gpt-4o")]
|
||||
assert cell.alpha == 90.0
|
||||
assert cell.beta == 10.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_state_from_db_handles_unknown_request_type():
|
||||
router = _make_router()
|
||||
bad_row = MagicMock()
|
||||
bad_row.request_type = "unknown_v1_type"
|
||||
bad_row.model_name = "gpt-4o"
|
||||
bad_row.alpha = 50.0
|
||||
bad_row.beta = 50.0
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row])
|
||||
|
||||
# Should not raise; bad row is silently skipped and cold-start cells remain.
|
||||
await router.load_state_from_db(prisma)
|
||||
cell = router._cells[(RequestType.GENERAL, "gpt-4o")]
|
||||
# Cold-start: tier 3 base = 0.7, mass = 10 -> alpha = 7, beta = 3
|
||||
assert cell.alpha == pytest.approx(7.0)
|
||||
assert cell.beta == pytest.approx(3.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_isolates_writes_per_router_session_model():
|
||||
router = _make_router()
|
||||
# Prime 2 prior turns per session to clear the MIN_TURNS_FOR_CLEAN_CREDIT gate.
|
||||
for sid, model in (("s1", "gpt-4o"), ("s2", "gpt-4o-mini")):
|
||||
for _ in range(2):
|
||||
await router.record_turn(
|
||||
sid,
|
||||
model,
|
||||
RequestType.GENERAL,
|
||||
Turn(user_content="hi", assistant_content="hello"),
|
||||
)
|
||||
await router.record_turn(
|
||||
"s1", "gpt-4o", RequestType.GENERAL, Turn(user_content="thanks!")
|
||||
)
|
||||
await router.record_turn(
|
||||
"s2", "gpt-4o-mini", RequestType.GENERAL, Turn(user_content="thanks!")
|
||||
)
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
n = await router.queue.flush_session_to_db(prisma)
|
||||
assert n == 2
|
||||
assert prisma.db.litellm_adaptiveroutersession.upsert.call_count == 2
|
||||
|
||||
n_state = await router.queue.flush_state_to_db(prisma)
|
||||
assert n_state == 2
|
||||
assert prisma.db.litellm_adaptiverouterstate.upsert.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_repeated_flush_drains_queue_and_subsequent_flush_is_noop():
|
||||
"""Verifies the queue is fully drained on flush -- a second flush writes nothing."""
|
||||
router = _make_router()
|
||||
chosen = await router.pick_model(RequestType.GENERAL)
|
||||
# Prime 2 prior turns so satisfaction can fire on the third turn.
|
||||
for _ in range(2):
|
||||
await router.record_turn(
|
||||
"drain-1",
|
||||
chosen,
|
||||
RequestType.GENERAL,
|
||||
Turn(user_content="hi", assistant_content="hello"),
|
||||
)
|
||||
await router.record_turn(
|
||||
"drain-1", chosen, RequestType.GENERAL, Turn(user_content="thanks!")
|
||||
)
|
||||
|
||||
prisma = _make_mock_prisma()
|
||||
assert await router.queue.flush_state_to_db(prisma) == 1
|
||||
assert await router.queue.flush_session_to_db(prisma) == 1
|
||||
|
||||
# Second drain should be a no-op (queue is empty).
|
||||
assert await router.queue.flush_state_to_db(prisma) == 0
|
||||
assert await router.queue.flush_session_to_db(prisma) == 0
|
||||
assert prisma.db.litellm_adaptiverouterstate.upsert.call_count == 1
|
||||
assert prisma.db.litellm_adaptiveroutersession.upsert.call_count == 1
|
||||
368
tests/test_litellm/router_strategy/adaptive_router/test_hooks.py
Normal file
368
tests/test_litellm/router_strategy/adaptive_router/test_hooks.py
Normal file
|
|
@ -0,0 +1,368 @@
|
|||
"""Unit tests for the AdaptiveRouterPostCallHook."""
|
||||
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.config import (
|
||||
ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY,
|
||||
SIGNAL_GATE_MIN_MESSAGES,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.hooks import (
|
||||
AdaptiveRouterPostCallHook,
|
||||
_recent_tool_results,
|
||||
_resolve_session_key,
|
||||
)
|
||||
from litellm.router_strategy.adaptive_router.signals import Turn
|
||||
|
||||
|
||||
def _make_hook(claim: bool = True) -> AdaptiveRouterPostCallHook:
|
||||
fake_router = MagicMock()
|
||||
fake_router.record_turn = AsyncMock()
|
||||
fake_router.claim_or_check_owner = MagicMock(return_value=claim)
|
||||
return AdaptiveRouterPostCallHook(adaptive_router=fake_router)
|
||||
|
||||
|
||||
def _resp_with_content(text: str, tool_calls=None):
|
||||
"""Build a ModelResponse-like object with a single assistant message."""
|
||||
msg = MagicMock()
|
||||
msg.content = text
|
||||
msg.tool_calls = tool_calls or []
|
||||
choice = MagicMock()
|
||||
choice.message = msg
|
||||
resp = MagicMock()
|
||||
resp.choices = [choice]
|
||||
return resp
|
||||
|
||||
|
||||
def _long_messages(user_text: str = "ask"):
|
||||
"""Return a message list at the SIGNAL_GATE_MIN_MESSAGES threshold."""
|
||||
base = [
|
||||
{"role": "user", "content": "first turn"},
|
||||
{"role": "assistant", "content": "first reply"},
|
||||
{"role": "user", "content": "second turn"},
|
||||
]
|
||||
base.append({"role": "user", "content": user_text})
|
||||
# Pad to threshold if needed.
|
||||
while len(base) < SIGNAL_GATE_MIN_MESSAGES:
|
||||
base.append({"role": "user", "content": "filler"})
|
||||
return base
|
||||
|
||||
|
||||
def _kwargs(
|
||||
*,
|
||||
messages=None,
|
||||
chosen="fast",
|
||||
extra_metadata=None,
|
||||
extra_litellm_params=None,
|
||||
):
|
||||
metadata = {ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY: chosen} if chosen else {}
|
||||
if extra_metadata:
|
||||
metadata.update(extra_metadata)
|
||||
lp = {"metadata": metadata}
|
||||
if extra_litellm_params:
|
||||
lp.update(extra_litellm_params)
|
||||
return {
|
||||
"model": "anthropic/claude-opus-4-7",
|
||||
"messages": messages if messages is not None else _long_messages(),
|
||||
"litellm_params": lp,
|
||||
}
|
||||
|
||||
|
||||
# ---- _resolve_session_key ------------------------------------------------
|
||||
|
||||
|
||||
def test_resolve_session_key_honors_litellm_session_id_on_litellm_params():
|
||||
key = _resolve_session_key({"litellm_params": {"litellm_session_id": "sess-A"}})
|
||||
assert key == "sess-A"
|
||||
|
||||
|
||||
def test_resolve_session_key_honors_metadata_session_id():
|
||||
key = _resolve_session_key(
|
||||
{"litellm_params": {"metadata": {"session_id": "sess-B"}}}
|
||||
)
|
||||
assert key == "sess-B"
|
||||
|
||||
|
||||
def test_resolve_session_key_returns_none_when_no_messages():
|
||||
assert _resolve_session_key({"litellm_params": {}}) is None
|
||||
assert _resolve_session_key({"litellm_params": {}, "messages": []}) is None
|
||||
|
||||
|
||||
def test_resolve_session_key_derives_stable_hash_from_first_message():
|
||||
# `_resolve_session_key` requires at least SIGNAL_GATE_MIN_MESSAGES
|
||||
# messages before it will derive a hash (matches the signal-processing
|
||||
# gate) — otherwise the session is too short to attribute.
|
||||
msgs = _long_messages("Hello, world")
|
||||
k1 = _resolve_session_key({"messages": msgs})
|
||||
k2 = _resolve_session_key({"messages": list(msgs)})
|
||||
assert k1 == k2
|
||||
assert k1 and len(k1) == 64 # sha256 hex
|
||||
|
||||
|
||||
def test_resolve_session_key_does_not_prefix_sk():
|
||||
key = _resolve_session_key({"messages": _long_messages()})
|
||||
assert key and not key.startswith("sk_")
|
||||
|
||||
|
||||
def test_resolve_session_key_segments_by_identity_fields():
|
||||
"""Same first message but different api keys must yield different keys."""
|
||||
msgs = _long_messages("same prompt")
|
||||
k_team_a = _resolve_session_key(
|
||||
{
|
||||
"messages": msgs,
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_hash": "hash-A",
|
||||
"user_api_key_team_id": "team-1",
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
k_team_b = _resolve_session_key(
|
||||
{
|
||||
"messages": msgs,
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key_hash": "hash-B",
|
||||
"user_api_key_team_id": "team-2",
|
||||
}
|
||||
},
|
||||
}
|
||||
)
|
||||
assert k_team_a != k_team_b
|
||||
|
||||
|
||||
def test_resolve_session_key_changes_when_first_message_changes():
|
||||
k1 = _resolve_session_key({"messages": _long_messages("alpha")})
|
||||
k2 = _resolve_session_key({"messages": _long_messages("beta")})
|
||||
assert k1 != k2
|
||||
|
||||
|
||||
# ---- _record gating -----------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_skips_when_below_signal_gate():
|
||||
"""Conversations shorter than SIGNAL_GATE_MIN_MESSAGES should be ignored."""
|
||||
hook = _make_hook()
|
||||
short = [{"role": "user", "content": "hi"}]
|
||||
assert len(short) < SIGNAL_GATE_MIN_MESSAGES # sanity
|
||||
kwargs = _kwargs(messages=short)
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
hook.adaptive_router.record_turn.assert_not_awaited()
|
||||
hook.adaptive_router.claim_or_check_owner.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_skips_when_no_messages():
|
||||
hook = _make_hook()
|
||||
kwargs = _kwargs(messages=[])
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
hook.adaptive_router.record_turn.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_skips_when_chosen_model_missing_from_metadata():
|
||||
hook = _make_hook()
|
||||
kwargs = _kwargs(chosen=None)
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
hook.adaptive_router.record_turn.assert_not_awaited()
|
||||
hook.adaptive_router.claim_or_check_owner.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_skips_when_owner_cache_mismatch():
|
||||
"""A different model owns this conversation -> no attribution."""
|
||||
hook = _make_hook(claim=False)
|
||||
kwargs = _kwargs(chosen="fast")
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
hook.adaptive_router.claim_or_check_owner.assert_called_once()
|
||||
hook.adaptive_router.record_turn.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_records_turn_when_owner_claims():
|
||||
hook = _make_hook(claim=True)
|
||||
kwargs = _kwargs(chosen="smart", messages=_long_messages("ask"))
|
||||
await hook.async_log_success_event(
|
||||
kwargs, _resp_with_content("answer here"), 0.0, 1.0
|
||||
)
|
||||
call = hook.adaptive_router.record_turn.await_args
|
||||
assert call.kwargs["model_name"] == "smart"
|
||||
turn: Turn = call.kwargs["turn"]
|
||||
assert turn.user_content == "ask"
|
||||
assert turn.assistant_content == "answer here"
|
||||
assert turn.response_status == 200
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_uses_explicit_session_id_when_provided():
|
||||
"""Explicit `litellm_session_id` is forwarded as the session key."""
|
||||
hook = _make_hook()
|
||||
kwargs = _kwargs(
|
||||
chosen="fast",
|
||||
extra_litellm_params={"litellm_session_id": "explicit-sess"},
|
||||
)
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
args, _ = hook.adaptive_router.claim_or_check_owner.call_args
|
||||
assert args[0] == "explicit-sess"
|
||||
assert hook.adaptive_router.record_turn.await_args.kwargs["session_id"] == (
|
||||
"explicit-sess"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_passes_tool_calls_through():
|
||||
hook = _make_hook()
|
||||
tc = {"name": "search", "arguments": '{"q":"x"}'}
|
||||
kwargs = _kwargs(chosen="fast")
|
||||
await hook.async_log_success_event(
|
||||
kwargs, _resp_with_content("calling tool", tool_calls=[tc]), 0.0, 1.0
|
||||
)
|
||||
turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"]
|
||||
assert turn.tool_calls == [tc]
|
||||
|
||||
|
||||
# ---- _recent_tool_results ------------------------------------------------
|
||||
|
||||
|
||||
def test_recent_tool_results_empty_when_no_messages():
|
||||
assert _recent_tool_results(None) == []
|
||||
assert _recent_tool_results([]) == []
|
||||
|
||||
|
||||
def test_recent_tool_results_collects_trailing_tool_messages():
|
||||
"""Tool messages at the tail of the conversation are extracted in order."""
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "result A"},
|
||||
{"role": "tool", "tool_call_id": "t2", "content": "result B"},
|
||||
]
|
||||
results = _recent_tool_results(messages)
|
||||
assert [r["content"] for r in results] == ["result A", "result B"]
|
||||
assert all(r["is_error"] is False for r in results)
|
||||
|
||||
|
||||
def test_recent_tool_results_propagates_is_error_flag():
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]},
|
||||
{"role": "tool", "content": "boom", "is_error": True},
|
||||
]
|
||||
results = _recent_tool_results(messages)
|
||||
assert results == [{"content": "boom", "is_error": True}]
|
||||
|
||||
|
||||
def test_recent_tool_results_stops_at_first_non_tool_message():
|
||||
"""Only the trailing run of tool messages counts — prior rounds are
|
||||
considered already attributed."""
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "tool", "content": "stale"}, # earlier round, ignored
|
||||
{"role": "assistant", "content": "intermediate"},
|
||||
{"role": "user", "content": "follow-up"},
|
||||
{"role": "tool", "content": "current"},
|
||||
]
|
||||
results = _recent_tool_results(messages)
|
||||
assert [r["content"] for r in results] == ["current"]
|
||||
|
||||
|
||||
def test_recent_tool_results_empty_when_no_trailing_tool_message():
|
||||
messages = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
]
|
||||
assert _recent_tool_results(messages) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_passes_tool_results_to_turn_for_failure_detection():
|
||||
"""A trailing tool message with `is_error` must reach `Turn.tool_results`
|
||||
so the failure-signal path fires."""
|
||||
hook = _make_hook()
|
||||
messages = _long_messages()
|
||||
messages.append(
|
||||
{"role": "assistant", "content": None, "tool_calls": [{"id": "t1"}]}
|
||||
)
|
||||
messages.append(
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "500", "is_error": True}
|
||||
)
|
||||
kwargs = _kwargs(chosen="fast", messages=messages)
|
||||
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
|
||||
turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"]
|
||||
assert turn.tool_results == [{"content": "500", "is_error": True}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_swallows_exceptions_from_record_turn():
|
||||
hook = _make_hook()
|
||||
hook.adaptive_router.record_turn.side_effect = RuntimeError("boom")
|
||||
kwargs = _kwargs(chosen="fast")
|
||||
# Must NOT raise — signal recording must never break a request.
|
||||
await hook.async_log_success_event(kwargs, _resp_with_content("ok"), 0.0, 1.0)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hook_failure_event_uses_status_code_from_exception():
|
||||
hook = _make_hook()
|
||||
exc = MagicMock()
|
||||
exc.status_code = 429
|
||||
kwargs = _kwargs(chosen="fast")
|
||||
kwargs["exception"] = exc
|
||||
await hook.async_log_failure_event(kwargs, None, 0.0, 1.0)
|
||||
turn: Turn = hook.adaptive_router.record_turn.await_args.kwargs["turn"]
|
||||
assert turn.response_status == 429
|
||||
|
||||
|
||||
# ---- async_post_call_success_hook (response header surfacing) ----------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_response_headers_hook_returns_chosen_model_header():
|
||||
"""The header hook returns the `x-litellm-adaptive-router-model` header
|
||||
so proxy header construction picks it up (works for both streaming and
|
||||
non-streaming; `async_post_call_success_hook` is too late for streaming)."""
|
||||
hook = _make_hook()
|
||||
headers = await hook.async_post_call_response_headers_hook(
|
||||
data={"metadata": {"adaptive_router_chosen_model": "smart"}},
|
||||
user_api_key_dict=MagicMock(),
|
||||
response=MagicMock(),
|
||||
)
|
||||
assert headers == {"x-litellm-adaptive-router-model": "smart"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_response_headers_hook_noop_when_metadata_missing_key():
|
||||
hook = _make_hook()
|
||||
headers = await hook.async_post_call_response_headers_hook(
|
||||
data={"metadata": {"litellm_session_id": "sess-A"}},
|
||||
user_api_key_dict=MagicMock(),
|
||||
response=MagicMock(),
|
||||
)
|
||||
assert headers is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_response_headers_hook_noop_when_no_metadata():
|
||||
hook = _make_hook()
|
||||
headers = await hook.async_post_call_response_headers_hook(
|
||||
data={},
|
||||
user_api_key_dict=MagicMock(),
|
||||
response=MagicMock(),
|
||||
)
|
||||
assert headers is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_response_headers_hook_noop_when_metadata_not_dict():
|
||||
hook = _make_hook()
|
||||
headers = await hook.async_post_call_response_headers_hook(
|
||||
data={"metadata": "not-a-dict"},
|
||||
user_api_key_dict=MagicMock(),
|
||||
response=MagicMock(),
|
||||
)
|
||||
assert headers is None
|
||||
|
|
@ -0,0 +1,486 @@
|
|||
"""Tests for the Router-level wiring of the adaptive router.
|
||||
|
||||
Specifically guards the four bugs found when wiring the example config
|
||||
`auto_router/adaptive_router` end-to-end:
|
||||
|
||||
1. The `auto_router/adaptive_router` model prefix must NOT trigger the
|
||||
semantic auto-router init path (which would crash on missing fields).
|
||||
2. The same prefix MUST trigger the adaptive-router init path.
|
||||
3. `init_adaptive_router_deployment` must read `input_cost_per_token`
|
||||
from `litellm_params` (where users put it), not just `model_info`.
|
||||
4. `Router.async_pre_routing_hook` must dispatch to the matching entry in
|
||||
`self.adaptive_routers` when the inbound model matches a configured
|
||||
adaptive-router name, returning the underlying model the bandit picked.
|
||||
"""
|
||||
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm import Router
|
||||
from litellm.types.router import LiteLLM_Params, RequestType
|
||||
|
||||
|
||||
def _params(**overrides):
|
||||
base = {"model": "auto_router/adaptive_router"}
|
||||
base.update(overrides)
|
||||
return LiteLLM_Params(**base)
|
||||
|
||||
|
||||
# ---- Fix 1 & 2: opt-in prefix routing -----------------------------------
|
||||
|
||||
|
||||
def test_auto_router_check_excludes_adaptive_router_prefix():
|
||||
r = Router(model_list=[])
|
||||
assert (
|
||||
r._is_auto_router_deployment(
|
||||
litellm_params=_params(model="auto_router/adaptive_router")
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_auto_router_check_excludes_complexity_router_prefix():
|
||||
r = Router(model_list=[])
|
||||
assert (
|
||||
r._is_auto_router_deployment(
|
||||
litellm_params=_params(model="auto_router/complexity_router")
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_auto_router_check_still_matches_plain_auto_router_prefix():
|
||||
r = Router(model_list=[])
|
||||
assert (
|
||||
r._is_auto_router_deployment(
|
||||
litellm_params=_params(model="auto_router/my-semantic-router")
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_adaptive_router_check_recognizes_prefix():
|
||||
r = Router(model_list=[])
|
||||
assert (
|
||||
r._is_adaptive_router_deployment(
|
||||
litellm_params=_params(model="auto_router/adaptive_router")
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_adaptive_router_check_rejects_other_prefixes():
|
||||
r = Router(model_list=[])
|
||||
assert (
|
||||
r._is_adaptive_router_deployment(litellm_params=_params(model="openai/gpt-4o"))
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
# ---- Fix 3: cost field path --------------------------------------------
|
||||
|
||||
|
||||
def test_init_adaptive_router_reads_cost_from_litellm_params():
|
||||
r = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-cheap-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {
|
||||
"available_models": ["fast", "smart"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"input_cost_per_token": 0.00000015,
|
||||
},
|
||||
"model_info": {
|
||||
"adaptive_router_preferences": {
|
||||
"quality_tier": 2,
|
||||
"strengths": [],
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "smart",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"input_cost_per_token": 0.0000050,
|
||||
},
|
||||
"model_info": {
|
||||
"adaptive_router_preferences": {
|
||||
"quality_tier": 3,
|
||||
"strengths": ["code_generation"],
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
assert "smart-cheap-router" in r.adaptive_routers
|
||||
assert r.adaptive_routers["smart-cheap-router"].model_to_cost == {
|
||||
"fast": 0.00000015,
|
||||
"smart": 0.0000050,
|
||||
}
|
||||
|
||||
|
||||
# ---- Fix 4: pre-routing dispatch ---------------------------------------
|
||||
|
||||
|
||||
def _router_with_adaptive() -> Router:
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "smart-cheap-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {
|
||||
"available_models": ["fast", "smart"],
|
||||
},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"input_cost_per_token": 0.00000015,
|
||||
},
|
||||
"model_info": {
|
||||
"adaptive_router_preferences": {
|
||||
"quality_tier": 2,
|
||||
"strengths": [],
|
||||
}
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "smart",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"input_cost_per_token": 0.0000050,
|
||||
},
|
||||
"model_info": {
|
||||
"adaptive_router_preferences": {
|
||||
"quality_tier": 3,
|
||||
"strengths": ["code_generation"],
|
||||
}
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_dispatches_to_adaptive_router():
|
||||
r = _router_with_adaptive()
|
||||
ar = r.adaptive_routers["smart-cheap-router"]
|
||||
ar.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment]
|
||||
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={"metadata": {"litellm_session_id": "sess-A"}},
|
||||
messages=[{"role": "user", "content": "Write a Python function"}],
|
||||
)
|
||||
assert response is not None
|
||||
assert response.model == "smart"
|
||||
call = ar.pick_model.await_args # type: ignore[union-attr]
|
||||
# Stateless routing: session_id is no longer passed to pick_model.
|
||||
assert "session_id" not in call.kwargs
|
||||
assert call.kwargs["request_type"] == RequestType.CODE_GENERATION
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_pick_model_not_passed_session_id():
|
||||
r = _router_with_adaptive()
|
||||
ar = r.adaptive_routers["smart-cheap-router"]
|
||||
ar.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment]
|
||||
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
assert response is not None
|
||||
assert response.model == "fast"
|
||||
assert "session_id" not in ar.pick_model.await_args.kwargs # type: ignore[union-attr]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_returns_none_for_unrelated_model():
|
||||
r = _router_with_adaptive()
|
||||
ar = r.adaptive_routers["smart-cheap-router"]
|
||||
ar.pick_model = AsyncMock() # type: ignore[assignment]
|
||||
response = await r.async_pre_routing_hook(
|
||||
model="some-other-model",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "x"}],
|
||||
)
|
||||
assert response is None
|
||||
ar.pick_model.assert_not_awaited() # type: ignore[union-attr]
|
||||
|
||||
|
||||
# ---- Response header surfacing -----------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_stashes_chosen_model_in_metadata():
|
||||
"""
|
||||
The adaptive-router branch must record the chosen logical model on
|
||||
`request_kwargs["metadata"]` so `_acompletion` can surface it as the
|
||||
`x-litellm-adaptive-router-model` response header.
|
||||
"""
|
||||
r = _router_with_adaptive()
|
||||
r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment]
|
||||
return_value="smart"
|
||||
)
|
||||
|
||||
request_kwargs: dict = {"metadata": {"litellm_session_id": "sess-A"}}
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "Write a Python function"}],
|
||||
)
|
||||
assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "smart"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_creates_metadata_when_missing():
|
||||
"""If no metadata was passed in, the hook should create one to stash the chosen model."""
|
||||
r = _router_with_adaptive()
|
||||
r.adaptive_routers["smart-cheap-router"].pick_model = AsyncMock( # type: ignore[assignment]
|
||||
return_value="fast"
|
||||
)
|
||||
|
||||
request_kwargs: dict = {}
|
||||
await r.async_pre_routing_hook(
|
||||
model="smart-cheap-router",
|
||||
request_kwargs=request_kwargs,
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
)
|
||||
assert request_kwargs["metadata"]["adaptive_router_chosen_model"] == "fast"
|
||||
|
||||
|
||||
# ---- Multi-router support ----------------------------------------------
|
||||
|
||||
|
||||
def test_two_adaptive_routers_can_coexist_on_one_router():
|
||||
r = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "cheap-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {"available_models": ["fast"]},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "premium-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {"available_models": ["smart"]},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"input_cost_per_token": 0.00000015,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "smart",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"input_cost_per_token": 0.0000050,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
assert set(r.adaptive_routers.keys()) == {"cheap-router", "premium-router"}
|
||||
assert r.adaptive_routers["cheap-router"].config.available_models == ["fast"]
|
||||
assert r.adaptive_routers["premium-router"].config.available_models == ["smart"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_routing_hook_dispatches_to_correct_router_when_multiple():
|
||||
"""Each adaptive router only handles its own router_name."""
|
||||
r = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "cheap-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {"available_models": ["fast"]},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "premium-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {"available_models": ["smart"]},
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"input_cost_per_token": 0.00000015,
|
||||
},
|
||||
},
|
||||
{
|
||||
"model_name": "smart",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o",
|
||||
"input_cost_per_token": 0.0000050,
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
cheap = r.adaptive_routers["cheap-router"]
|
||||
premium = r.adaptive_routers["premium-router"]
|
||||
cheap.pick_model = AsyncMock(return_value="fast") # type: ignore[assignment]
|
||||
premium.pick_model = AsyncMock(return_value="smart") # type: ignore[assignment]
|
||||
|
||||
cheap_response = await r.async_pre_routing_hook(
|
||||
model="cheap-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
premium_response = await r.async_pre_routing_hook(
|
||||
model="premium-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
|
||||
assert cheap_response is not None and cheap_response.model == "fast"
|
||||
assert premium_response is not None and premium_response.model == "smart"
|
||||
cheap.pick_model.assert_awaited_once() # type: ignore[union-attr]
|
||||
premium.pick_model.assert_awaited_once() # type: ignore[union-attr]
|
||||
|
||||
|
||||
def test_init_adaptive_router_rejects_duplicate_model_name():
|
||||
"""Two adaptive-router deployments with the same model_name must error."""
|
||||
from litellm.types.router import AdaptiveRouterConfig, Deployment
|
||||
|
||||
r = Router(model_list=[])
|
||||
cfg = {"available_models": ["fast"]}
|
||||
deployment = Deployment(
|
||||
model_name="dup-router",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="auto_router/adaptive_router",
|
||||
adaptive_router_config=cfg,
|
||||
),
|
||||
model_info={"id": "x"},
|
||||
)
|
||||
r.init_adaptive_router_deployment(deployment=deployment)
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
r.init_adaptive_router_deployment(deployment=deployment)
|
||||
|
||||
|
||||
def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent():
|
||||
"""`_finalize_adaptive_router_if_configured` walks the model_list, builds an
|
||||
AdaptiveRouter for each adaptive deployment, and is a safe no-op on
|
||||
re-entry (models already in self.adaptive_routers are skipped)."""
|
||||
r = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
"model_info": {"input_cost_per_token": 0.00000015},
|
||||
},
|
||||
{
|
||||
"model_name": "smart",
|
||||
"litellm_params": {"model": "openai/gpt-4o"},
|
||||
"model_info": {"input_cost_per_token": 0.0000025},
|
||||
},
|
||||
{
|
||||
"model_name": "my-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {
|
||||
"available_models": ["fast", "smart"],
|
||||
},
|
||||
},
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
# Router __init__ already called _finalize_adaptive_router_if_configured.
|
||||
assert "my-router" in r.adaptive_routers
|
||||
original = r.adaptive_routers["my-router"]
|
||||
|
||||
# Calling again must be idempotent: the existing AdaptiveRouter instance
|
||||
# is preserved, not rebuilt.
|
||||
r._finalize_adaptive_router_if_configured()
|
||||
assert r.adaptive_routers["my-router"] is original
|
||||
|
||||
|
||||
def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks():
|
||||
"""Replacing the Router (hot-reload path) must not leave stale
|
||||
AdaptiveRouterPostCallHook instances in `litellm.callbacks` — otherwise
|
||||
every request double-fires signal recording."""
|
||||
import litellm
|
||||
from litellm.router_strategy.adaptive_router.hooks import (
|
||||
AdaptiveRouterPostCallHook,
|
||||
)
|
||||
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
},
|
||||
{
|
||||
"model_name": "my-router",
|
||||
"litellm_params": {
|
||||
"model": "auto_router/adaptive_router",
|
||||
"adaptive_router_config": {"available_models": ["fast"]},
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
# Snapshot any pre-existing AdaptiveRouterPostCallHook entries so we can
|
||||
# restore them — other tests may have registered hooks we shouldn't drop.
|
||||
pre_hooks = [
|
||||
cb for cb in litellm.callbacks if isinstance(cb, AdaptiveRouterPostCallHook)
|
||||
]
|
||||
for cb in pre_hooks:
|
||||
litellm.callbacks.remove(cb)
|
||||
|
||||
try:
|
||||
Router(model_list=model_list)
|
||||
Router(model_list=model_list) # simulate hot-reload
|
||||
|
||||
adaptive_hooks = [
|
||||
cb
|
||||
for cb in litellm.callbacks
|
||||
if isinstance(cb, AdaptiveRouterPostCallHook)
|
||||
]
|
||||
assert len(adaptive_hooks) == 1, (
|
||||
f"expected exactly one AdaptiveRouterPostCallHook after hot-reload, "
|
||||
f"got {len(adaptive_hooks)}"
|
||||
)
|
||||
finally:
|
||||
# Best-effort cleanup: remove whatever this test added, then restore.
|
||||
for cb in list(litellm.callbacks):
|
||||
if isinstance(cb, AdaptiveRouterPostCallHook):
|
||||
litellm.callbacks.remove(cb)
|
||||
for cb in pre_hooks:
|
||||
litellm.callbacks.append(cb)
|
||||
|
||||
|
||||
def test_finalize_adaptive_router_if_configured_noop_when_none_configured():
|
||||
"""With no adaptive deployments in model_list, the finalizer leaves
|
||||
`adaptive_routers` empty."""
|
||||
r = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "fast",
|
||||
"litellm_params": {"model": "openai/gpt-4o-mini"},
|
||||
}
|
||||
]
|
||||
)
|
||||
r._finalize_adaptive_router_if_configured()
|
||||
assert r.adaptive_routers == {}
|
||||
|
|
@ -0,0 +1,180 @@
|
|||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.config import TOOL_CALL_HISTORY_MAX
|
||||
from litellm.router_strategy.adaptive_router.signals import (
|
||||
SessionState,
|
||||
SignalDelta,
|
||||
Turn,
|
||||
apply_turn,
|
||||
)
|
||||
|
||||
FIXTURE_DIR = Path(__file__).parent / "fixtures"
|
||||
|
||||
|
||||
def _load(name: str) -> list:
|
||||
return json.loads((FIXTURE_DIR / f"{name}.json").read_text())
|
||||
|
||||
|
||||
def _replay(turns: list) -> Tuple[SessionState, List[SignalDelta]]:
|
||||
state = SessionState(
|
||||
session_id="s",
|
||||
router_name="r",
|
||||
model_name="m",
|
||||
classified_type="general",
|
||||
)
|
||||
deltas: List[SignalDelta] = []
|
||||
for t in turns:
|
||||
deltas.append(
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(
|
||||
user_content=t.get("user_content"),
|
||||
assistant_content=t.get("assistant_content"),
|
||||
tool_calls=t.get("tool_calls", []),
|
||||
tool_results=t.get("tool_results", []),
|
||||
response_status=t.get("response_status"),
|
||||
),
|
||||
)
|
||||
)
|
||||
return state, deltas
|
||||
|
||||
|
||||
def test_clean_satisfaction_fires_satisfaction_only():
|
||||
state, _ = _replay(_load("clean_satisfaction"))
|
||||
assert state.satisfaction_count >= 1
|
||||
assert state.failure_count == 0
|
||||
assert state.disengagement_count == 0
|
||||
|
||||
|
||||
def test_misalignment_fires_on_rephrase():
|
||||
state, _ = _replay(_load("misalignment_rephrase"))
|
||||
assert state.misalignment_count >= 1
|
||||
|
||||
|
||||
def test_stagnation_fires_on_repeated_assistant():
|
||||
state, _ = _replay(_load("stagnation_repeat"))
|
||||
assert state.stagnation_count >= 1
|
||||
|
||||
|
||||
def test_disengagement_fires_on_giveup():
|
||||
state, _ = _replay(_load("disengagement_giveup"))
|
||||
assert state.disengagement_count >= 1
|
||||
|
||||
|
||||
def test_failure_fires_on_tool_error():
|
||||
state, _ = _replay(_load("failure_tool_error"))
|
||||
assert state.failure_count == 1
|
||||
|
||||
|
||||
def test_loop_fires_on_repeated_tool():
|
||||
state, _ = _replay(_load("loop_same_tool"))
|
||||
assert state.loop_count >= 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fixture", ["exhaustion_429", "exhaustion_context_overflow"])
|
||||
def test_exhaustion_fires_on_infra_signal(fixture):
|
||||
state, _ = _replay(_load(fixture))
|
||||
assert state.exhaustion_count >= 1
|
||||
|
||||
|
||||
def test_no_signals_on_clean_session():
|
||||
state, _ = _replay(_load("clean_no_signals"))
|
||||
assert state.misalignment_count == 0
|
||||
assert state.stagnation_count == 0
|
||||
assert state.disengagement_count == 0
|
||||
assert state.failure_count == 0
|
||||
assert state.loop_count == 0
|
||||
assert state.exhaustion_count == 0
|
||||
|
||||
|
||||
def test_mixed_failure_then_satisfaction():
|
||||
state, _ = _replay(_load("mixed_failure_then_satisfaction"))
|
||||
assert state.failure_count >= 1
|
||||
assert state.satisfaction_count >= 1
|
||||
|
||||
|
||||
def test_satisfaction_gated_by_min_turns_for_clean_credit():
|
||||
"""'thanks' on turn 1 is noise, not a validated quality signal."""
|
||||
state = SessionState(
|
||||
session_id="s", router_name="r", model_name="m", classified_type="general"
|
||||
)
|
||||
apply_turn(state, Turn(user_content="thanks!"))
|
||||
assert state.satisfaction_count == 0
|
||||
assert state.clean_credit_awarded is False
|
||||
assert state.last_processed_turn == 1
|
||||
|
||||
|
||||
def test_satisfaction_credit_awarded_once_per_session():
|
||||
"""Even multiple satisfaction turns only award +1 alpha across the session."""
|
||||
state = SessionState(
|
||||
session_id="s", router_name="r", model_name="m", classified_type="general"
|
||||
)
|
||||
apply_turn(state, Turn(user_content="hi", assistant_content="hello"))
|
||||
apply_turn(state, Turn(user_content="help me", assistant_content="sure"))
|
||||
apply_turn(state, Turn(user_content="perfect, thanks"))
|
||||
assert state.satisfaction_count == 1
|
||||
assert state.clean_credit_awarded is True
|
||||
apply_turn(state, Turn(user_content="great, thank you"))
|
||||
assert state.satisfaction_count == 1
|
||||
|
||||
|
||||
def test_empty_tool_content_does_not_fire_failure():
|
||||
"""Zero-result searches / silent commands return empty but valid output."""
|
||||
state = SessionState(
|
||||
session_id="s", router_name="r", model_name="m", classified_type="general"
|
||||
)
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(
|
||||
tool_calls=[{"name": "grep", "arguments": {"q": "x"}}],
|
||||
tool_results=[{"tool_call_id": "c1", "content": ""}],
|
||||
),
|
||||
)
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(
|
||||
tool_calls=[{"name": "list", "arguments": {}}],
|
||||
tool_results=[{"tool_call_id": "c2", "content": []}],
|
||||
),
|
||||
)
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(
|
||||
tool_calls=[{"name": "noop", "arguments": {}}],
|
||||
tool_results=[{"tool_call_id": "c3", "content": None}],
|
||||
),
|
||||
)
|
||||
assert state.failure_count == 0
|
||||
|
||||
|
||||
def test_is_error_still_fires_failure():
|
||||
state = SessionState(
|
||||
session_id="s", router_name="r", model_name="m", classified_type="general"
|
||||
)
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(
|
||||
tool_calls=[{"name": "read", "arguments": {"p": "x"}}],
|
||||
tool_results=[{"tool_call_id": "c1", "content": "boom", "is_error": True}],
|
||||
),
|
||||
)
|
||||
assert state.failure_count == 1
|
||||
|
||||
|
||||
def test_apply_turn_is_o1_does_not_grow_history_unbounded():
|
||||
state = SessionState(
|
||||
session_id="s",
|
||||
router_name="r",
|
||||
model_name="m",
|
||||
classified_type="general",
|
||||
)
|
||||
for i in range(100):
|
||||
apply_turn(
|
||||
state,
|
||||
Turn(tool_calls=[{"name": f"tool_{i}", "arguments": {}}]),
|
||||
)
|
||||
assert len(state.tool_call_history) <= TOOL_CALL_HISTORY_MAX
|
||||
|
|
@ -0,0 +1,198 @@
|
|||
"""Tests for the GET /adaptive_router/state introspection endpoint and the
|
||||
underlying `AdaptiveRouter.get_state_snapshot()` helper."""
|
||||
|
||||
import time
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter
|
||||
from litellm.router_strategy.adaptive_router.bandit import BanditCell, apply_delta
|
||||
from litellm.types.router import (
|
||||
AdaptiveRouterConfig,
|
||||
AdaptiveRouterPreferences,
|
||||
RequestType,
|
||||
)
|
||||
|
||||
|
||||
def _make_router(name: str = "r1") -> AdaptiveRouter:
|
||||
cfg = AdaptiveRouterConfig(available_models=["fast", "smart"])
|
||||
prefs = {
|
||||
"fast": AdaptiveRouterPreferences(quality_tier=1, strengths=[]),
|
||||
"smart": AdaptiveRouterPreferences(
|
||||
quality_tier=3, strengths=[RequestType.CODE_GENERATION]
|
||||
),
|
||||
}
|
||||
costs = {"fast": 0.0001, "smart": 0.001}
|
||||
return AdaptiveRouter(
|
||||
router_name=name,
|
||||
config=cfg,
|
||||
model_to_prefs=prefs,
|
||||
model_to_cost=costs,
|
||||
)
|
||||
|
||||
|
||||
# ---- snapshot helper ---------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_state_snapshot_returns_cell_per_request_type_per_model():
|
||||
r = _make_router()
|
||||
snap = await r.get_state_snapshot()
|
||||
|
||||
# Top-level shape
|
||||
assert snap["router_name"] == "r1"
|
||||
assert snap["available_models"] == ["fast", "smart"]
|
||||
assert snap["weights"] == {"quality": 0.7, "cost": 0.3}
|
||||
assert snap["model_costs"] == {"fast": 0.0001, "smart": 0.001}
|
||||
assert snap["owner_cache_live"] == 0
|
||||
assert snap["skipped_updates_total"] == 0
|
||||
assert set(snap["queue"].keys()) == {
|
||||
"state_pending",
|
||||
"session_pending",
|
||||
"max_state_seen",
|
||||
"max_session_seen",
|
||||
}
|
||||
|
||||
# 7 request types x 2 models = 14 cells
|
||||
assert len(snap["cells"]) == len(list(RequestType)) * 2
|
||||
for cell in snap["cells"]:
|
||||
assert set(cell.keys()) == {
|
||||
"request_type",
|
||||
"model",
|
||||
"alpha",
|
||||
"beta",
|
||||
"samples",
|
||||
"quality_mean",
|
||||
}
|
||||
assert cell["model"] in {"fast", "smart"}
|
||||
assert cell["request_type"] in {rt.value for rt in RequestType}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_state_snapshot_quality_mean_matches_alpha_over_total():
|
||||
r = _make_router()
|
||||
|
||||
# Manually mutate one cell to a known state so the math is verifiable.
|
||||
key = (RequestType.CODE_GENERATION, "smart")
|
||||
r._cells[key] = apply_delta(r._cells[key], delta_alpha=10.0, delta_beta=0.0)
|
||||
expected = r._cells[key]
|
||||
expected_mean = expected.alpha / (expected.alpha + expected.beta)
|
||||
|
||||
snap = await r.get_state_snapshot()
|
||||
cell = next(
|
||||
c
|
||||
for c in snap["cells"]
|
||||
if c["request_type"] == "code_generation" and c["model"] == "smart"
|
||||
)
|
||||
assert cell["alpha"] == expected.alpha
|
||||
assert cell["beta"] == expected.beta
|
||||
# `samples` reports net observations after subtracting the cold-start
|
||||
# prior mass, so operators aren't misled by the initial value.
|
||||
assert cell["samples"] == expected.total_samples
|
||||
assert cell["quality_mean"] == pytest.approx(expected_mean)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_state_snapshot_counts_only_live_owner_cache_entries():
|
||||
r = _make_router()
|
||||
now = time.time()
|
||||
r._owner_cache["live-1"] = ("fast", now + 3600)
|
||||
r._owner_cache["live-2"] = ("smart", now + 3600)
|
||||
r._owner_cache["expired-1"] = ("fast", now - 1)
|
||||
|
||||
snap = await r.get_state_snapshot()
|
||||
assert snap["owner_cache_live"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_state_snapshot_exposes_skipped_updates_total():
|
||||
r = _make_router()
|
||||
r._skipped_updates_total = 7
|
||||
snap = await r.get_state_snapshot()
|
||||
assert snap["skipped_updates_total"] == 7
|
||||
|
||||
|
||||
# ---- endpoint --------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_returns_404_when_no_adaptive_router(monkeypatch):
|
||||
"""When llm_router is set but has no adaptive routers configured, return 404."""
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
fake_router = MagicMock()
|
||||
fake_router.adaptive_routers = {}
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await proxy_server.get_adaptive_router_state(user_api_key_dict=admin)
|
||||
assert exc.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_returns_404_when_llm_router_is_none(monkeypatch):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await proxy_server.get_adaptive_router_state(user_api_key_dict=admin)
|
||||
assert exc.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_rejects_non_admin_role(monkeypatch):
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
fake_router = MagicMock()
|
||||
fake_router.adaptive_routers = {"r1": _make_router()}
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
|
||||
non_admin = UserAPIKeyAuth(
|
||||
api_key="sk-user", user_role=LitellmUserRoles.INTERNAL_USER
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await proxy_server.get_adaptive_router_state(user_api_key_dict=non_admin)
|
||||
assert exc.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_returns_snapshot_list_for_admin(monkeypatch):
|
||||
"""Single configured router still returns the {"routers": [...]} list shape."""
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
fake_router = MagicMock()
|
||||
fake_router.adaptive_routers = {"r1": _make_router("r1")}
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
result = await proxy_server.get_adaptive_router_state(user_api_key_dict=admin)
|
||||
assert list(result.keys()) == ["routers"]
|
||||
assert len(result["routers"]) == 1
|
||||
snap = result["routers"][0]
|
||||
assert snap["router_name"] == "r1"
|
||||
assert snap["available_models"] == ["fast", "smart"]
|
||||
assert len(snap["cells"]) == len(list(RequestType)) * 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_endpoint_returns_one_snapshot_per_router(monkeypatch):
|
||||
"""With multiple adaptive routers configured, return one snapshot per router."""
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
fake_router = MagicMock()
|
||||
fake_router.adaptive_routers = {
|
||||
"r1": _make_router("r1"),
|
||||
"r2": _make_router("r2"),
|
||||
}
|
||||
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
|
||||
|
||||
admin = UserAPIKeyAuth(api_key="sk-1234", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
result = await proxy_server.get_adaptive_router_state(user_api_key_dict=admin)
|
||||
names = sorted(s["router_name"] for s in result["routers"])
|
||||
assert names == ["r1", "r2"]
|
||||
|
|
@ -0,0 +1,117 @@
|
|||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.adaptive_router.update_queue import (
|
||||
AdaptiveRouterUpdateQueue,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue():
|
||||
return AdaptiveRouterUpdateQueue()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_prisma():
|
||||
"""Prisma client with both adaptive router models stubbed as AsyncMocks."""
|
||||
p = MagicMock()
|
||||
p.db.litellm_adaptiverouterstate.find_unique = AsyncMock(return_value=None)
|
||||
p.db.litellm_adaptiverouterstate.upsert = AsyncMock()
|
||||
p.db.litellm_adaptiveroutersession.upsert = AsyncMock()
|
||||
return p
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_state_delta_aggregates_same_key(queue):
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 0.0, 1.0)
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["state_pending"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_state_delta_separate_keys(queue):
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "writing", "gpt-4", 1.0, 0.0)
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["state_pending"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_session_state_last_write_wins(queue):
|
||||
await queue.add_session_state("s1", "r1", "gpt-4", {"misalignment_count": 1})
|
||||
await queue.add_session_state("s1", "r1", "gpt-4", {"misalignment_count": 5})
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["session_pending"] == 1
|
||||
|
||||
flushed = []
|
||||
p = MagicMock()
|
||||
|
||||
async def upsert(**kwargs):
|
||||
flushed.append(kwargs)
|
||||
|
||||
p.db.litellm_adaptiveroutersession.upsert = upsert
|
||||
await queue.flush_session_to_db(p)
|
||||
assert len(flushed) == 1
|
||||
assert flushed[0]["data"]["update"]["misalignment_count"] == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_state_drains_aggregator(queue, mock_prisma):
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "writing", "gpt-4", 0.0, 1.0)
|
||||
n = await queue.flush_state_to_db(mock_prisma)
|
||||
assert n == 2
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["state_pending"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_state_sums_correctly(queue, mock_prisma):
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 2.0, 1.0)
|
||||
await queue.flush_state_to_db(mock_prisma)
|
||||
# find_unique returned None (cold start), so alpha = 1+2 = 3, beta = 0+1 = 1
|
||||
call = mock_prisma.db.litellm_adaptiverouterstate.upsert.call_args
|
||||
assert call.kwargs["data"]["create"]["alpha"] == 3.0
|
||||
assert call.kwargs["data"]["create"]["beta"] == 1.0
|
||||
assert call.kwargs["data"]["create"]["total_samples"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_session_drains_aggregator(queue, mock_prisma):
|
||||
await queue.add_session_state("s1", "r1", "gpt-4", {"classified_type": "general"})
|
||||
n = await queue.flush_session_to_db(mock_prisma)
|
||||
assert n == 1
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["session_pending"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_empty_queue_returns_zero(queue, mock_prisma):
|
||||
assert await queue.flush_state_to_db(mock_prisma) == 0
|
||||
assert await queue.flush_session_to_db(mock_prisma) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_state_isolation_from_concurrent_adds(queue, mock_prisma):
|
||||
"""Adds during a flush should land in the NEW aggregator, not the drained batch."""
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
flush_task = asyncio.create_task(queue.flush_state_to_db(mock_prisma))
|
||||
# Yield control so the flush task can swap the aggregator before we add again.
|
||||
await asyncio.sleep(0)
|
||||
await queue.add_state_delta("r1", "general", "gpt-5", 2.0, 0.0)
|
||||
await flush_task
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["state_pending"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_max_size_observability(queue):
|
||||
await queue.add_state_delta("r1", "general", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "writing", "gpt-4", 1.0, 0.0)
|
||||
await queue.add_state_delta("r1", "code_generation", "gpt-4", 1.0, 0.0)
|
||||
sizes = await queue.queue_size()
|
||||
assert sizes["max_state_seen"] >= 3
|
||||
54
tests/test_litellm/test_dockerfile_non_root.py
Normal file
54
tests/test_litellm/test_dockerfile_non_root.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
"""
|
||||
Static checks on docker/Dockerfile.non_root.
|
||||
|
||||
The non_root image is intended for deployment into hardened Kubernetes
|
||||
clusters where `securityContext.runAsNonRoot: true` is enforced. The
|
||||
kubelet validates non-root status by parsing the image's USER field as
|
||||
an integer — a string name like "nobody" is rejected with
|
||||
CreateContainerConfigError because the kubelet cannot resolve
|
||||
/etc/passwd inside the image at admission time.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
DOCKERFILE_PATH = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"docker",
|
||||
"Dockerfile.non_root",
|
||||
)
|
||||
|
||||
|
||||
def _final_user_directive(dockerfile_text: str) -> str:
|
||||
"""Return the value of the last `USER` directive in the file."""
|
||||
matches = re.findall(r"^USER\s+(\S+)\s*$", dockerfile_text, re.MULTILINE)
|
||||
assert matches, "Dockerfile.non_root has no USER directive"
|
||||
return matches[-1]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.path.exists(DOCKERFILE_PATH),
|
||||
reason="Dockerfile.non_root not present in this checkout",
|
||||
)
|
||||
def test_final_user_directive_is_numeric():
|
||||
"""The runtime USER must be a numeric UID so kubelet's runAsNonRoot
|
||||
admission check (strconv.Atoi) succeeds."""
|
||||
with open(DOCKERFILE_PATH, "r", encoding="utf-8") as f:
|
||||
contents = f.read()
|
||||
|
||||
final_user = _final_user_directive(contents)
|
||||
|
||||
assert final_user.isdigit(), (
|
||||
f"Dockerfile.non_root final USER is {final_user!r}; must be a numeric UID "
|
||||
"so Kubernetes' runAsNonRoot admission check can verify non-root status. "
|
||||
"See https://kubernetes.io/docs/tasks/configure-pod-container/security-context/"
|
||||
)
|
||||
|
||||
assert int(final_user) != 0, (
|
||||
f"Dockerfile.non_root final USER is {final_user} (root); the non_root image "
|
||||
"must run as a non-zero UID."
|
||||
)
|
||||
|
|
@ -329,3 +329,39 @@ async def test_router_order_fallback_with_non_standard_fallbacks():
|
|||
fallbacks=["fallback-model"], # non-standard format, passed per-request
|
||||
)
|
||||
assert response._hidden_params["model_id"] == "fallback"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_router_order_fallback_with_wildcard_model_group():
|
||||
"""Wildcard model groups should also advance across order levels."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {
|
||||
"model": "openai/*",
|
||||
"api_key": "bad",
|
||||
"mock_response": Exception("fail order 1"),
|
||||
"order": 1,
|
||||
},
|
||||
"model_info": {"id": "1"},
|
||||
},
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {
|
||||
"model": "openai/*",
|
||||
"api_key": "good",
|
||||
"mock_response": "success from wildcard order 2",
|
||||
"order": 2,
|
||||
},
|
||||
"model_info": {"id": "2"},
|
||||
},
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
response = await router.acompletion(
|
||||
model="openai/gpt-4.1-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
)
|
||||
assert response._hidden_params["model_id"] == "2"
|
||||
|
|
|
|||
6
uv.lock
generated
6
uv.lock
generated
|
|
@ -9,7 +9,7 @@ resolution-markers = [
|
|||
]
|
||||
|
||||
[options]
|
||||
exclude-newer = "2026-04-19T01:10:36.69677Z"
|
||||
exclude-newer = "2026-04-20T01:21:50.985363Z"
|
||||
exclude-newer-span = "P3D"
|
||||
|
||||
[manifest]
|
||||
|
|
@ -3085,7 +3085,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.83.11"
|
||||
version = "1.83.12"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
|
@ -3418,7 +3418,7 @@ source = { editable = "enterprise" }
|
|||
|
||||
[[package]]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.67"
|
||||
version = "0.4.68"
|
||||
source = { editable = "litellm-proxy-extras" }
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue