scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

Merge remote-tracking branch 'origin/main' into litellm_redis_durable_spend_log_buffer

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

# Conflicts:
#	litellm/proxy/utils.py
This commit is contained in:
yassin 2026-09-19 23:21:34 +00:00
commit 0c1841affc
189 changed files with 13816 additions and 2324 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
.github/e2e-stack/assert_tests_ran.py vendored
View file

@ -1,3 +1,5 @@
import os
import re
import sys
import xml.etree.ElementTree as ET
from pathlib import Path
@ -15,6 +17,12 @@ def main() -> int:
_ = sys.stdout.write("::error::could not read the test execution report\n")
return 1
cases: Final = tuple(report.iter("testcase"))
expected_count: Final = os.environ.get("E2E_REQUIRED_TEST_COUNT")
if expected_count is not None and (
len(cases) != int(expected_count) or any(case.find("skipped") is not None for case in cases)
):
_ = sys.stdout.write("::error::required test count was not met or a required case was skipped\n")
return 1
passed: Final = frozenset(
case.get("file") for case in cases if all(case.find(tag) is None for tag in ("skipped", "failure", "error"))
)
@ -38,6 +46,13 @@ def main() -> int:
if case.get("file") != path or all(case.find(tag) is None for tag in ("failure", "error")):
continue
_ = sys.stdout.write(f" failed: {case.get('classname', '')}::{case.get('name', '')}\n")
for prop in case.findall("./properties/property"):
name = prop.get("name", "")
value = prop.get("value", "")
if name in ("oauth_failure_phase", "oauth_exception_type", "oauth_frame") and re.fullmatch(
r"[A-Za-z0-9_.:<>-]{1,240}", value
):
_ = sys.stdout.write(f" {name}: {value}\n")
if (
selected
and not missing

1
.github/e2e-stack/select_tests.py vendored
View file

@ -5,6 +5,7 @@ from typing import Final
SELECTABLE: Final = re.compile(r"^tests/e2e/([A-Za-z0-9_.-]+/)*test_[A-Za-z0-9_.-]+\.py$")
UNSUPPORTED: Final = re.compile(
r"^tests/e2e/(ui|claude_code|load)/"
r"|^tests/e2e/mcp/test_mcp_oauth_happy_path_e2e\.py$"
r"|^tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e\.py$"
r"|^tests/e2e/batches/test_managed_files_enforcement_e2e\.py$"
r"|^tests/e2e/guardrails/test_presidio_masking_e2e\.py$"

177
.github/workflows/test-mcp-oauth-e2e.yml vendored Normal file
View file

@ -0,0 +1,177 @@
name: MCP OAuth happy path
on:
pull_request:
paths:
- '.github/workflows/test-mcp-oauth-e2e.yml'
- '.github/e2e-stack/**'
- 'tests/e2e/*.py'
- 'tests/e2e/pytest.ini'
- 'tests/e2e/idp_realm.json'
- 'tests/e2e/mcp/**'
- 'litellm/experimental_mcp_client/**'
- 'litellm/proxy/_experimental/mcp_server/**'
- 'litellm/proxy/auth/**'
- 'litellm/proxy/management_endpoints/*sso*.py'
- 'litellm/proxy/management_endpoints/sso/**'
- 'litellm/proxy/common_utils/encrypt_decrypt_utils.py'
- 'litellm/proxy/proxy_server.py'
- 'litellm/proxy/schema.prisma'
- 'ui/litellm-dashboard/src/app/connect/**'
- 'ui/litellm-dashboard/src/app/mcp/oauth/**'
- 'pyproject.toml'
- 'uv.lock'
workflow_dispatch:
permissions: {}
concurrency:
group: mcp-oauth-${{ github.ref }}
cancel-in-progress: true
jobs:
oauth:
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.head.repo.full_name == github.repository
runs-on: ubuntu-latest
environment: e2e-changed
timeout-minutes: 45
permissions:
contents: read
id-token: write
services:
postgres:
image: postgres:16.6
env:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: dbpassword9090
POSTGRES_DB: litellm
ports:
- 5432:5432
options: >-
--health-cmd "pg_isready -U litellm"
--health-interval 5s
--health-timeout 5s
--health-retries 10
env:
DATABASE_HOST: 127.0.0.1
DATABASE_PORT: '5432'
DATABASE_USER: litellm
DATABASE_PASSWORD: dbpassword9090
DATABASE_NAME: litellm
DATABASE_URL: postgresql://litellm:dbpassword9090@127.0.0.1:5432/litellm
E2E_KEYCLOAK_URL: http://127.0.0.1:8081
E2E_KEYCLOAK_ADMIN_USER: admin
E2E_KEYCLOAK_ADMIN_PASSWORD: e2e-ephemeral-idp-not-a-secret
E2E_FIXTURE_MODE: live
E2E_PROVIDER_CACHE: '0'
E2E_MCP_OAUTH_LIVE: '1'
E2E_REQUIRED_TEST_COUNT: '4'
steps:
- name: Checkout the tested source
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
ref: ${{ github.event.pull_request.head.sha || github.sha }}
persist-credentials: false
- name: Require and materialize the upstream login
env:
STORAGE_STATE: ${{ secrets.E2E_LINEAR_STORAGE_STATE_B64 }}
run: |
umask 077
python3 - <<'PY'
import base64
import json
import os
import secrets
from pathlib import Path
encoded = os.environ.get("STORAGE_STATE", "")
if not encoded:
raise SystemExit("E2E_LINEAR_STORAGE_STATE_B64 is required; capture and provision a test-account login")
state = json.loads(base64.b64decode(encoded, validate=True))
if not isinstance(state, dict) or not state.get("cookies"):
raise SystemExit("The captured login must contain browser cookies")
directory = Path(os.environ["RUNNER_TEMP"]) / "mcp-oauth-private"
directory.mkdir(mode=0o700)
path = directory / "linear-state.json"
path.write_text(json.dumps(state))
with open(os.environ["GITHUB_ENV"], "a") as output:
output.write(f"E2E_LINEAR_STORAGE_STATE={path}\n")
for name in ("LITELLM_MASTER_KEY", "LITELLM_SALT_KEY"):
value = "sk-e2e-" + secrets.token_hex(24)
print(f"::add-mask::{value}")
output.write(f"{name}={value}\n")
PY
- uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
with:
python-version: '3.13'
- uses: ./.github/actions/setup-uv-with-retries
with:
version: '0.10.9'
- uses: ./.github/actions/cache-cargo-build
- name: Install the frozen E2E environment
run: |
.github/scripts/uv_sync_with_retries.sh --frozen --extra proxy --extra proxy-runtime --extra extra_proxy --group ci --group proxy-dev --group e2e-dev
uv run --no-sync python scripts/prisma_generate_if_needed.py
uv run --no-sync playwright install --with-deps chromium
- name: Configure license access
id: aws
uses: aws-actions/configure-aws-credentials@e7f100cf4c008499ea8adda475de1042d6975c7b # v6.2.0
with:
role-to-assume: ${{ vars.E2E_AWS_ROLE_TO_ASSUME }}
aws-region: us-east-1
role-session-name: mcp-oauth-${{ github.run_id }}
role-duration-seconds: 900
output-env-credentials: false
output-credentials: true
- name: Load the E2E license
env:
AWS_ACCESS_KEY_ID: ${{ steps.aws.outputs.aws-access-key-id }}
AWS_SECRET_ACCESS_KEY: ${{ steps.aws.outputs.aws-secret-access-key }}
AWS_SESSION_TOKEN: ${{ steps.aws.outputs.aws-session-token }}
AWS_DEFAULT_REGION: us-east-1
run: |
license="$(aws secretsmanager get-secret-value --secret-id litellm-e2e-changed-license --query SecretString --output text)"
test -n "${license}"
echo "::add-mask::${license}"
echo "LITELLM_LICENSE=${license}" >> "${GITHUB_ENV}"
- uses: actions/setup-node@49933ea5288caeca8642d1e84afbd3f7d6820020 # v4.4.0
with:
node-version-file: ui/litellm-dashboard/.nvmrc
- name: Build the gateway consent UI at the tested commit
run: |
cd ui/litellm-dashboard
../../scripts/with_dashboard_node.sh npm ci
../../scripts/with_dashboard_node.sh npm run build
mkdir -p ../../litellm/proxy/_experimental/out
cp -r out/. ../../litellm/proxy/_experimental/out/
find ../../litellm/proxy/_experimental/out -name '*.html' ! -name index.html | while read -r page; do
mkdir -p "${page%.html}"
mv "${page}" "${page%.html}/index.html"
done
- name: Prepare the isolated database and IdP
run: |
umask 077
bash .github/e2e-stack/start-idp.sh
uv run --no-sync python migrations/run.py > "${RUNNER_TEMP}/mcp-oauth-private/migrations.log" 2>&1
- name: Run every required OAuth variant without retries
run: |
umask 077
uv run --no-sync pytest -c tests/e2e/pytest.ini tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0 --tb=short -o junit_family=xunit1 \
--junitxml="${RUNNER_TEMP}/mcp-oauth-private/results.xml" \
> "${RUNNER_TEMP}/mcp-oauth-private/pytest.log" 2>&1
- name: Report JUnit results and reject skipped or missing cases
if: always()
run: |
uv run --no-sync python .github/e2e-stack/assert_tests_ran.py \
"${RUNNER_TEMP}/mcp-oauth-private/results.xml" tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py
- name: Remove private login and logs
if: always()
run: |
docker rm -f e2e-keycloak >/dev/null 2>&1 || true
rm -rf "${RUNNER_TEMP}/mcp-oauth-private"

9
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py
View file

@ -60,6 +60,11 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
"""Helper function to save email settings to general_settings in db"""
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name="general_settings", changed_keys={"email_settings": settings}
)
try:
verbose_proxy_logger.debug(
f"Saving email settings to general_settings: {settings}"
@ -168,6 +173,8 @@ async def update_event_settings(
await _save_email_settings(prisma_client, settings_dict)
return {"message": "Email event settings updated successfully"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error updating email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))
@ -197,6 +204,8 @@ async def reset_event_settings(
await _save_email_settings(prisma_client, default_settings)
return {"message": "Email event settings reset to defaults"}
except HTTPException:
raise
except Exception as e:
verbose_proxy_logger.exception(f"Error resetting email settings: {str(e)}")
raise HTTPException(status_code=500, detail=str(e))

116
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -20,6 +20,7 @@ from typing import (
)
from uuid import NAMESPACE_URL, uuid5
import httpx
from fastapi import HTTPException
from pydantic import ValidationError
@ -34,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
)
from openai.types.file_deleted import FileDeleted
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import (
build_list_page,
@ -59,6 +61,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_content_type_from_file_object,
get_model_id_from_unified_batch_id,
get_original_file_id,
is_litellm_executed_batch,
map_raw_file_ids_to_unified,
normalize_mime_type_for_provider,
resolve_managed_output_file_model_name,
@ -75,6 +78,7 @@ from litellm.types.llms.openai import ( # pyright: ignore[reportAttributeAccess
CreateFileRequest,
FileListPage,
FileObject,
HttpxBinaryResponseContent,
OpenAIFileObject,
ResponsesAPIResponse,
)
@ -86,10 +90,6 @@ from litellm.types.utils import (
SpecialEnums,
)
if TYPE_CHECKING:
from litellm.types.llms.openai import HttpxBinaryResponseContent
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
from prisma.models import (
@ -204,6 +204,19 @@ def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableAct
return prisma_client.db.litellm_managedobjecttable
def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets
"Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
)
return MappingProxyType(
{
key: value
for key in ("storage_backend", "storage_url")
if isinstance(value := hidden_params.get(key), str)
}
)
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Class variables or attributes
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
@ -226,6 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
user_api_key_dict: UserAPIKeyAuth,
) -> None:
verbose_logger.info(f"Storing LiteLLM Managed File object with id={file_id} in cache")
storage_metadata: Final = _storage_metadata_of(file_object)
if file_object is not None:
litellm_managed_file_object = LiteLLM_ManagedFileTable(
unified_file_id=file_id,
@ -235,6 +249,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
created_by=resolve_resource_owner_id(user_api_key_dict),
team_id=user_api_key_dict.team_id,
updated_by=user_api_key_dict.user_id,
storage_backend=storage_metadata.get("storage_backend"),
storage_url=storage_metadata.get("storage_url"),
)
await self.internal_usage_cache.async_set_cache(
key=file_id,
@ -262,14 +278,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_object_json = file_object.model_dump_json()
db_data["file_object"] = file_object_json
update_data["file_object"] = file_object_json
# Extract storage metadata from hidden params if present
hidden_params = getattr(file_object, "_hidden_params", {}) or {}
if "storage_backend" in hidden_params:
db_data["storage_backend"] = hidden_params["storage_backend"]
update_data["storage_backend"] = hidden_params["storage_backend"]
if "storage_url" in hidden_params:
db_data["storage_url"] = hidden_params["storage_url"]
update_data["storage_url"] = hidden_params["storage_url"]
db_data.update(storage_metadata)
update_data.update(storage_metadata)
verbose_logger.debug(
f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
@ -314,6 +324,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
create_if_missing: bool = True,
batch_processed: bool = False,
) -> None:
"""Persist a managed object row, caching it and upserting it in the DB.
@ -328,6 +339,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
row absent from the table is left absent rather than created with the
observer as its creator, because created_by and team_id are written from
whoever calls the create branch.
batch_processed is set by callers that have already billed the batch
themselves, so CheckBatchCost skips the row instead of billing it twice.
It is written only in the upsert create branch.
"""
verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
litellm_managed_object = LiteLLM_ManagedObjectTable(
@ -379,6 +394,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
**attribution_columns,
"batch_processed": batch_processed,
},
"update": update_columns,
},
@ -1343,6 +1359,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes
) -> LLMResponseTypes:
if isinstance(response, LiteLLMBatch):
decoded_batch_id: Final = _is_base64_encoded_unified_file_id(response.id)
if decoded_batch_id and is_litellm_executed_batch(decoded_batch_id):
return response
## Check if unified_file_id is in the response
unified_file_id = response._hidden_params.get("unified_file_id") # managed file id
unified_batch_id = response._hidden_params.get("unified_batch_id") # managed batch id
@ -1794,24 +1813,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Check if file deletion should be blocked due to batch references
await self._check_file_deletion_allowed(file_id)
# file_id = convert_b64_uid_to_unified_uid(file_id)
model_file_id_mapping = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
if specific_model_file_id_mapping:
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**{k: v for k, v in filtered_data.items() if k != "_litellm_internal_model_credentials"},
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
await self._delete_storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
else:
await self._delete_provider_files(file_id, litellm_parent_otel_span, llm_router, data)
await self.delete_unified_file_id(file_id, litellm_parent_otel_span)
@ -1820,16 +1826,53 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
prom_logger.record_managed_file_deleted(result="success")
return FileDeleted(id=file_id, object="file", deleted=True)
async def _delete_storage_backend_content(self, storage_backend_name: str, storage_url: str) -> None:
try:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
raise HTTPException(status_code=400, detail=f"Cannot delete the stored file content: {e}") from e
await storage_backend.delete_file(storage_url)
async def _delete_provider_files(
self,
file_id: str,
litellm_parent_otel_span: Span | None,
llm_router: Router,
data: Mapping[str, object],
) -> None:
model_file_id_mapping: Final = await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
specific_model_file_id_mapping: Final = model_file_id_mapping.get(file_id)
if not specific_model_file_id_mapping:
return
filtered_data: Final = {
k: v for k, v in data.items() if k not in ("model", "file_id", "_litellm_internal_model_credentials")
}
for model_id, model_file_id in specific_model_file_id_mapping.items():
credentials = llm_router.get_deployment_credentials_with_provider(model_id=model_id)
delete_data = {
**filtered_data,
**(
{"_litellm_internal_model_credentials": MappingProxyType(dict(credentials))}
if credentials is not None
else {}
),
}
await llm_router.afile_delete(model=model_id, file_id=model_file_id, **delete_data)
async def afile_content(
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
llm_router: Router,
**data: Dict,
) -> "HttpxBinaryResponseContent":
) -> HttpxBinaryResponseContent:
"""
Get the content of a file from first model that has it
"""
managed_file: Final = await self.get_unified_file_id(file_id, litellm_parent_otel_span)
if managed_file is not None and managed_file.storage_backend and managed_file.storage_url:
return await self._storage_backend_content(managed_file.storage_backend, managed_file.storage_url)
model_file_id_mapping = data.pop("model_file_id_mapping", None)
model_file_id_mapping = model_file_id_mapping or await self.get_model_file_id_mapping(
[file_id], litellm_parent_otel_span
@ -1859,6 +1902,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
else:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
async def _storage_backend_content(self, storage_backend_name: str, storage_url: str) -> HttpxBinaryResponseContent:
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
content: Final = await storage_backend.download_file(storage_url)
return HttpxBinaryResponseContent(response=httpx.Response(status_code=httpx.codes.OK, content=content))
async def _convert_storage_files_to_base64(
self,
messages: List[AllMessageValues],
@ -1889,16 +1937,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# File is stored in a storage backend, download and convert to base64
try:
from litellm.llms.base_llm.files.storage_backend_factory import (
get_storage_backend,
)
storage_backend_name = db_file.storage_backend
storage_url = db_file.storage_url
# Get storage backend (uses same env vars as callback)
try:
storage_backend = get_storage_backend(storage_backend_name)
storage_backend = get_storage_backend(storage_backend_name, prisma_client=self.prisma_client)
except ValueError as e:
verbose_logger.warning(
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-enterprise"
version = "0.1.68"
version = "0.1.69"
description = "Package for LiteLLM Enterprise features"
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.1.68"
version = "0.1.69"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-enterprise==",

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260915000000_add_autorouter_savings_estimate_coverage/migration.sql Normal file
View file

@ -0,0 +1,5 @@
ALTER TABLE "LiteLLM_AutoRouterSession"
ADD COLUMN IF NOT EXISTS "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
ADD COLUMN IF NOT EXISTS "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}';

36
litellm-proxy-extras/litellm_proxy_extras/migrations/20260915010000_add_autorouter_baseline_state/migration.sql Normal file
View file

@ -0,0 +1,36 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineComparison" (
"scope" TEXT PRIMARY KEY,
"api_key" TEXT NOT NULL,
"session_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"initial_equivalent" BOOLEAN NOT NULL,
"revision" BIGINT NOT NULL DEFAULT 0,
"published_revision" BIGINT NOT NULL DEFAULT 0,
"history" TEXT,
"attempted_at" TIMESTAMP(3),
"retired" BOOLEAN NOT NULL DEFAULT FALSE,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_scope"
ON "LiteLLM_AutoRouterBaselineComparison" ("api_key", "session_id", "router_name");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_updated"
ON "LiteLLM_AutoRouterBaselineComparison" ("updated_at");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_dirty"
ON "LiteLLM_AutoRouterBaselineComparison" ("attempted_at", "updated_at", "scope")
WHERE NOT "retired" AND "revision" <> "published_revision";
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterBaselineObservation" (
"request_id" TEXT PRIMARY KEY,
"scope" TEXT NOT NULL,
"started_at" DOUBLE PRECISION NOT NULL,
"revision" BIGINT NOT NULL,
"data" TEXT NOT NULL,
"publication" TEXT,
"conflicted" BOOLEAN NOT NULL DEFAULT FALSE
);
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_order"
ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "started_at", "request_id");
CREATE INDEX IF NOT EXISTS "idx_autorouter_baseline_event_revision"
ON "LiteLLM_AutoRouterBaselineObservation" ("scope", "revision", "started_at");

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260918000000_add_managed_file_content_table/migration.sql Normal file
View file

@ -0,0 +1,8 @@
-- CreateTable
CREATE TABLE IF NOT EXISTS "LiteLLM_ManagedFileContentTable" (
"id" TEXT NOT NULL,
"content" BYTEA NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_ManagedFileContentTable_pkey" PRIMARY KEY ("id")
);

40
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID
@ -1545,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession {
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
model LiteLLM_AutoRouterBaselineComparison {
scope String @id
api_key String
session_id String
router_name String
initial_equivalent Boolean
revision BigInt @default(0)
published_revision BigInt @default(0)
history String?
attempted_at DateTime?
retired Boolean @default(false)
updated_at DateTime @default(now())
@@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope")
@@index([updated_at], map: "idx_autorouter_baseline_updated")
}
model LiteLLM_AutoRouterBaselineObservation {
request_id String @id
scope String
started_at Float
revision BigInt
data String
publication String?
conflicted Boolean @default(false)
@@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order")
@@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision")
}
model LiteLLM_AutoRouterSession {
api_key String
session_id String
@ -1571,6 +1607,10 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.99"
version = "0.4.100"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
readme = "README.md"
requires-python = ">=3.9"
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
module-root = ""
[tool.commitizen]
version = "0.4.99"
version = "0.4.100"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

1
litellm-rust/Cargo.lock generated
View file

@ -3046,7 +3046,6 @@ checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147"
dependencies = [
"base64 0.22.1",
"bytes",
"futures-channel",
"futures-core",
"futures-util",
"h2 0.4.15",

2
litellm-rust/Cargo.toml
View file

@ -34,7 +34,7 @@ pyo3 = "0.29.2"
pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] }
pythonize = "0.29.0"
rand = "0.8"
reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "rustls-tls", "http2", "stream"] }
reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] }
rstest = "0.26.1"
rstest_reuse = "0.7.0"
rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] }

10
litellm-rust/clippy.toml Normal file
View file

@ -0,0 +1,10 @@
# The Tokio runtime is reached only through `host-python/src/execution.rs`, whose fork gate
# must see every entry. Going around it makes a fork-after-use hang instead of raising.
disallowed-methods = [
{ path = "pyo3_async_runtimes::tokio::get_runtime", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::future_into_py_with_locals", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::local_future_into_py", reason = "use litellm_host_python::run_async / run_async_value" },
{ path = "pyo3_async_runtimes::tokio::run", reason = "use litellm_host_python::run_sync / run_sync_value" },
{ path = "pyo3_async_runtimes::tokio::run_until_complete", reason = "use litellm_host_python::run_sync / run_sync_value" },
]

102
litellm-rust/crates/host-python/src/execution.rs
View file

@ -4,6 +4,7 @@ use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use std::time::Duration;
use crate::fork_gate::{ForkGate, Refused, RuntimeAlreadyStarted};
use crate::{Pythonized, panic_to_pyerr, release_gil};
use futures_util::FutureExt;
use pyo3::exceptions::PyRuntimeError;
@ -12,6 +13,67 @@ use serde::Serialize;
use tokio::runtime::{Handle, Runtime};
use tokio::time::{self, MissedTickBehavior};
pyo3::create_exception!(
_native,
ForkedAfterNativeRuntimeStarted,
PyRuntimeError,
"This process was forked after the native runtime started. Runtime threads do not survive fork(), so native routes cannot run here."
);
pyo3::create_exception!(
_native,
ProcessReservedForForking,
PyRuntimeError,
"This process was reserved for forking workers, so native routes cannot run here."
);
static FORK_GATE: ForkGate = ForkGate::new();
/// Whether this process has started the Tokio runtime.
pub fn runtime_started() -> bool {
FORK_GATE.started(std::process::id())
}
/// Declares that this process exists to fork workers, so it must never start the runtime.
/// Fails if it already has. Workers are unaffected: the reservation is keyed by pid.
pub fn reserve_process_for_forking() -> Result<(), RuntimeAlreadyStarted> {
FORK_GATE.reserve(std::process::id())
}
/// The only door to the Tokio runtime: every route reaches it through this module, which is
/// what lets the gate speak for the whole extension. `clippy.toml` disallows going around it.
fn enter_runtime() -> PyResult<()> {
FORK_GATE
.enter(std::process::id())
.map_err(|refused| match refused {
Refused::ReservedForForking => ProcessReservedForForking::new_err(
"this process is reserved for forking workers and cannot run native routes; \
move the call into a worker, after the fork",
),
Refused::ForkedAfterStart => ForkedAfterNativeRuntimeStarted::new_err(
"this process was forked after the native runtime started, and runtime threads \
do not survive fork(); start workers with spawn or forkserver, or fork before \
the first native call",
),
})
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn runtime() -> PyResult<&'static Runtime> {
enter_runtime()?;
Ok(pyo3_async_runtimes::tokio::get_runtime())
}
#[expect(clippy::disallowed_methods, reason = "this is the gated door")]
fn future_into_py<F, T>(py: Python<'_>, future: F) -> PyResult<Bound<'_, PyAny>>
where
F: Future<Output = PyResult<T>> + Send + 'static,
T: for<'py> IntoPyObject<'py> + Send + 'static,
{
enter_runtime()?;
pyo3_async_runtimes::tokio::future_into_py(py, future)
}
pub fn run_sync<T, E, F>(
py: Python<'_>,
future: F,
@ -22,12 +84,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
run_sync_on(
py,
pyo3_async_runtimes::tokio::get_runtime(),
future,
map_error,
)
run_sync_on(py, runtime()?, future, map_error)
}
pub fn run_sync_value<T, F>(py: Python<'_>, future: F) -> PyResult<T>
@ -35,7 +92,7 @@ where
T: Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
run_sync_value_on(py, pyo3_async_runtimes::tokio::get_runtime(), future)
run_sync_value_on(py, runtime()?, future)
}
fn run_sync_value_on<T, F>(py: Python<'_>, runtime: &Runtime, future: F) -> PyResult<T>
@ -83,7 +140,7 @@ where
E: Send + 'static,
F: Future<Output = Result<T, E>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move {
future_into_py(py, async move {
let result = catch_future_panic(future).await?;
let result = map_core_result(result, map_error)?;
Ok(Pythonized(result))
@ -95,7 +152,7 @@ where
T: for<'py> IntoPyObject<'py> + Send + 'static,
F: Future<Output = PyResult<T>> + Send + 'static,
{
pyo3_async_runtimes::tokio::future_into_py(py, async move { catch_future_panic(future).await? })
future_into_py(py, async move { catch_future_panic(future).await? })
}
pub fn poll_async_value<T, F>(py: Python<'_>, future: Pin<&mut F>) -> PyResult<Poll<T>>
@ -103,8 +160,9 @@ where
T: Send,
F: Future<Output = PyResult<T>> + Send,
{
let runtime = runtime()?;
let result = release_gil(py, || {
let _runtime = pyo3_async_runtimes::tokio::get_runtime().enter();
let _runtime = runtime.enter();
std::panic::catch_unwind(AssertUnwindSafe(|| {
future.poll(&mut Context::from_waker(Waker::noop()))
}))
@ -286,27 +344,25 @@ mod tests {
}
#[pyfunction]
fn runtime_worker_count() -> usize {
pyo3_async_runtimes::tokio::get_runtime()
.metrics()
.num_workers()
fn runtime_worker_count() -> PyResult<usize> {
Ok(runtime()?.metrics().num_workers())
}
#[pyfunction]
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> bool {
fn runtime_is_responsive(_py: Python<'_>, expected_completions: usize) -> PyResult<bool> {
let completion_deadline = Instant::now() + Duration::from_secs(2);
while ASYNC_PROBE_COMPLETED.load(Ordering::SeqCst) < expected_completions {
if Instant::now() >= completion_deadline {
return false;
return Ok(false);
}
thread::sleep(Duration::from_millis(1));
}
let (heartbeat_tx, heartbeat_rx) = mpsc::sync_channel(1);
pyo3_async_runtimes::tokio::get_runtime().spawn(async move {
runtime()?.spawn(async move {
let _ = heartbeat_tx.send(());
});
heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok()
Ok(heartbeat_rx.recv_timeout(Duration::from_secs(2)).is_ok())
}
fn extract_bool(py: Python<'_>, result: PyResult<Py<PyAny>>) -> bool {
@ -317,6 +373,16 @@ mod tests {
.expect("result should convert")
}
#[rstest]
fn reaching_the_runtime_marks_the_process_as_started(
#[from(initialized_python)] python: &InitializedPython,
) {
python.attach(|py| {
run_sync_value(py, async { Ok(()) }).unwrap();
assert!(runtime_started());
});
}
#[rstest]
fn inline_poll_releases_gil_and_enters_runtime(
#[from(initialized_python)] python: &InitializedPython,

139
litellm-rust/crates/host-python/src/fork_gate.rs Normal file
View file

@ -0,0 +1,139 @@
use std::sync::atomic::{AtomicU32, Ordering};
const UNSET: u32 = 0;
/// Decides which process may use the Tokio runtime. Its worker threads do not survive
/// `fork()`: a child forked after they started hangs on its first native call. The gate turns
/// both halves of that hazard into errors, keyed by pid so a fork needs no hook to be seen:
/// a process reserved for forking can never start the runtime, and a child of a process that
/// did start it is refused instead of hanging.
pub(crate) struct ForkGate {
runtime_pid: AtomicU32,
fork_only_pid: AtomicU32,
}
#[derive(Debug, PartialEq, Eq)]
pub(crate) enum Refused {
ReservedForForking,
ForkedAfterStart,
}
#[derive(Debug, PartialEq, Eq)]
pub struct RuntimeAlreadyStarted;
impl ForkGate {
pub(crate) const fn new() -> Self {
Self {
runtime_pid: AtomicU32::new(UNSET),
fork_only_pid: AtomicU32::new(UNSET),
}
}
/// Claims the runtime for `pid`. Claim first, then look for a reservation: `reserve` does
/// the mirror image, so when the two race at least one of them sees the other.
pub(crate) fn enter(&self, pid: u32) -> Result<(), Refused> {
match self
.runtime_pid
.compare_exchange(UNSET, pid, Ordering::SeqCst, Ordering::SeqCst)
{
Err(owner) if owner != pid => return Err(Refused::ForkedAfterStart),
_ => {}
}
if self.fork_only_pid.load(Ordering::SeqCst) == pid {
// Nothing was started, so the workers forked from here must still find it unclaimed.
let _ =
self.runtime_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(Refused::ReservedForForking);
}
Ok(())
}
/// Reserves `pid` for forking. Reserve first, then look for a started runtime: `enter` does
/// the mirror image, so when the two race at least one of them sees the other. A refused
/// reservation leaves the gate exactly as it was, so a process already running the runtime
/// keeps refusing the children it forks.
pub(crate) fn reserve(&self, pid: u32) -> Result<(), RuntimeAlreadyStarted> {
self.fork_only_pid.store(pid, Ordering::SeqCst);
if self.runtime_pid.load(Ordering::SeqCst) == pid {
let _ =
self.fork_only_pid
.compare_exchange(pid, UNSET, Ordering::SeqCst, Ordering::SeqCst);
return Err(RuntimeAlreadyStarted);
}
Ok(())
}
pub(crate) fn started(&self, pid: u32) -> bool {
self.runtime_pid.load(Ordering::SeqCst) == pid
}
}
#[cfg(test)]
mod tests {
use super::*;
const MASTER: u32 = 100;
const WORKER: u32 = 101;
#[test]
fn unreserved_process_starts_the_runtime_and_stays_started() {
let gate = ForkGate::new();
assert!(!gate.started(MASTER));
assert_eq!(gate.enter(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
}
#[test]
fn reserved_process_can_never_start_the_runtime() {
let gate = ForkGate::new();
assert_eq!(gate.reserve(MASTER), Ok(()));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert_eq!(gate.enter(MASTER), Err(Refused::ReservedForForking));
assert!(!gate.started(MASTER));
}
#[test]
fn workers_forked_from_a_reserved_process_start_their_own_runtime() {
let gate = ForkGate::new();
gate.reserve(MASTER).unwrap();
gate.enter(MASTER).unwrap_err();
assert_eq!(gate.enter(WORKER), Ok(()));
assert!(gate.started(WORKER));
}
#[test]
fn reserving_after_the_runtime_started_is_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
}
#[test]
fn a_refused_reservation_leaves_the_runtime_claimed_and_its_children_refused() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.reserve(MASTER), Err(RuntimeAlreadyStarted));
assert_eq!(gate.enter(MASTER), Ok(()));
assert!(gate.started(MASTER));
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
}
#[test]
fn child_forked_after_the_runtime_started_is_refused_instead_of_hanging() {
let gate = ForkGate::new();
gate.enter(MASTER).unwrap();
assert_eq!(gate.enter(WORKER), Err(Refused::ForkedAfterStart));
assert!(!gate.started(WORKER));
assert_eq!(gate.enter(MASTER), Ok(()));
}
}

8
litellm-rust/crates/host-python/src/lib.rs
View file

@ -8,6 +8,7 @@ mod argument;
mod callable;
mod driver;
mod execution;
mod fork_gate;
mod gil;
mod handle;
mod marshal;
@ -18,7 +19,12 @@ pub use adapter::{
pub use argument::lookup;
pub use callable::wrap_failure;
pub use driver::run_call;
pub use execution::{poll_async_value, run_async, run_async_value, run_sync, run_sync_value};
pub use execution::{
ForkedAfterNativeRuntimeStarted, ProcessReservedForForking, poll_async_value,
reserve_process_for_forking, run_async, run_async_value, run_sync, run_sync_value,
runtime_started,
};
pub use fork_gate::RuntimeAlreadyStarted;
pub use gil::{release_count, release_gil};
pub use handle::{Execution, ExecutionBody, ExecutionStep};
pub use marshal::{Pythonized, from_py, from_py_argument, panic_to_pyerr, to_py};

18
litellm-rust/crates/python-bridge/src/diagnostics.rs
View file

@ -1,5 +1,5 @@
use litellm_host_python::release_count;
use pyo3::{prelude::*, types::PyDict};
use litellm_host_python::{release_count, runtime_started};
use pyo3::{exceptions::PyRuntimeError, prelude::*, types::PyDict};
#[pyfunction]
pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
@ -8,6 +8,20 @@ pub(crate) fn gil_stats(py: Python<'_>) -> PyResult<Py<PyAny>> {
Ok(stats.into_any().unbind())
}
/// True once this process has started the native runtime, which does not survive `fork()`.
#[pyfunction]
pub(crate) fn process_state_started() -> bool {
runtime_started()
}
/// Declares that this process only forks workers: from now on every native route raises here,
/// so the runtime can never start. Raises if it already has. Forked workers are unaffected.
#[pyfunction]
pub(crate) fn reserve_process_for_forking() -> PyResult<()> {
litellm_host_python::reserve_process_for_forking()
.map_err(|_| PyRuntimeError::new_err("the native runtime already started in this process"))
}
#[cfg(feature = "panic-test")]
#[pyfunction]
pub(crate) fn _panic_for_test() {

8
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -13,7 +13,7 @@ mod _native {
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
#[pymodule_export]
use crate::diagnostics::gil_stats;
use crate::diagnostics::{gil_stats, process_state_started, reserve_process_for_forking};
#[pymodule_export]
use crate::errors::{RustBridgeDeclined, RustUpstreamError};
#[pymodule_export]
@ -30,6 +30,8 @@ mod _native {
use crate::routes::responses::ResponsesWebSocketConnection;
#[pymodule_export]
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
}
use pyo3::prelude::*;
@ -50,6 +52,8 @@ mod tests {
let mut expected = vec![
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ocr",
"aocr",
"transcription",
@ -62,6 +66,8 @@ mod tests {
"ResponsesWebSocketConnection",
"TokenCounter",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",
];
expected.sort_unstable();

12
litellm-rust/crates/python-bridge/src/routes/responses.rs
View file

@ -25,7 +25,7 @@ impl ResponsesWebSocketConnection {
) -> PyResult<Bound<'py, PyAny>> {
let headers = marshal_headers(headers)?;
let timeout = optional_timeout(timeout_seconds);
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout)
.await
.map_err(responses_error_to_pyerr)?;
@ -35,7 +35,7 @@ impl ResponsesWebSocketConnection {
fn send_text<'py>(&self, py: Python<'py>, text: String) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner
.send_text(text)
.await
@ -45,14 +45,14 @@ impl ResponsesWebSocketConnection {
fn recv_text<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.recv_text().await.map_err(responses_error_to_pyerr)
})
}
fn close<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let inner = self.inner.clone();
pyo3_async_runtimes::tokio::future_into_py(py, async move {
litellm_host_python::run_async_value(py, async move {
inner.close().await.map_err(responses_error_to_pyerr)
})
}
@ -68,6 +68,10 @@ mod tests {
use tokio_tungstenite::{accept_async, tungstenite::Message};
#[test]
#[expect(
clippy::disallowed_methods,
reason = "the test server shares the routes' runtime"
)]
fn responses_websocket_connection_round_trips_through_python() {
Python::initialize();
let runtime = pyo3_async_runtimes::tokio::get_runtime();

1
litellm/constants.py
View file

@ -1697,6 +1697,7 @@ LOGIN_THROTTLE_NOT_BLOCKED: Final = (0, 0)
LITELLM_PROXY_ADMIN_NAME: Final = "default_user_id"
LITELLM_PROXY_BUDGET_NAME: Final = "litellm-proxy-budget"
GLOBAL_PROXY_SPEND_CACHE_KEY: Final = f"{LITELLM_PROXY_ADMIN_NAME}:spend"
LITELLM_EXECUTED_BATCH_CONCURRENCY: Final = max(1, int(os.getenv("LITELLM_EXECUTED_BATCH_CONCURRENCY", "4")))
########################### CLI SSO AUTHENTICATION CONSTANTS ###########################
LITELLM_CLI_SOURCE_IDENTIFIER: Final = "litellm-cli"

85
litellm/integrations/otel/model/payloads.py
View file

@ -7,9 +7,11 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from enum import Enum
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final, cast
from typing import TYPE_CHECKING, ClassVar, Final, Literal, cast
from urllib.parse import urlsplit
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.otel.model.metadata import RequestContext, RequestIdentity
from litellm.integrations.otel.model.semconv import (
GenAIOperation,
@ -25,6 +27,7 @@ from litellm.integrations.otel.model.utils import (
as_float,
as_int,
as_str,
as_str_mapping,
as_str_tuple,
)
@ -424,7 +427,7 @@ class LLMCallSpanData:
# plain ``.get`` — no repeated ``isinstance`` guards.
raw_response: Final = payload.get("response")
response: Final = cast(Mapping[str, object], raw_response if isinstance(raw_response, dict) else {})
choices_out: Final = _dicts(response.get("choices"))
choices_out: Final = _dicts(response.get("choices")) or _responses_choices(response)
# ``finish_reasons`` is metadata, not content, so derive it from
# ``choices_out`` before gating. The raw message/choice bodies are only
# retained when content capture is enabled (see ``capture_span_content``);
@ -703,6 +706,84 @@ def _finish_reasons(choices: tuple[Mapping[str, object], ...]) -> tuple[str, ...
return tuple(r for c in choices if (r := as_str(c.get("finish_reason"))))
class _ToolFunction(TypedDict):
name: ReadOnly[str]
arguments: ReadOnly[str]
class _ToolCall(TypedDict):
id: ReadOnly[str]
type: ReadOnly[Literal["function"]]
function: ReadOnly[_ToolFunction]
class _AssistantMessage(TypedDict):
role: ReadOnly[str]
content: ReadOnly[str | None]
refusal: ReadOnly[str | None]
tool_calls: ReadOnly[tuple[_ToolCall, ...] | None]
class _Choice(TypedDict):
message: ReadOnly[_AssistantMessage]
finish_reason: ReadOnly[str | None]
_RESPONSES_TOOL_CALL_TYPES: Final = frozenset({"function_call", "custom_tool_call"})
def _responses_choices(response: Mapping[str, object]) -> tuple[_Choice, ...]:
"""A Responses API ``output`` folded into one chat-shaped assistant choice."""
items: Final = _dicts(response.get("output"))
messages: Final = tuple(item for item in items if item.get("type") == "message")
parts: Final = tuple(part for item in messages for part in _dicts(item.get("content")))
tool_calls: Final = tuple(
_responses_tool_call(item) for item in items if item.get("type") in _RESPONSES_TOOL_CALL_TYPES
)
if not messages and not tool_calls:
return ()
message: Final[_AssistantMessage] = {
"role": next((role for item in messages if (role := as_str(item.get("role")))), "assistant"),
"content": _responses_parts_text(parts, "output_text", "text"),
"refusal": _responses_parts_text(parts, "refusal", "refusal"),
"tool_calls": tool_calls or None,
}
choice: Final[_Choice] = {"message": message, "finish_reason": _responses_finish_reason(response, bool(tool_calls))}
return (choice,)
def _responses_parts_text(parts: tuple[Mapping[str, object], ...], part_type: str, field: str) -> str | None:
texts: Final = tuple(
text for part in parts if part.get("type") == part_type if (text := as_str(part.get(field))) is not None
)
return "".join(texts) if texts else None
def _responses_tool_call(item: Mapping[str, object]) -> _ToolCall:
custom: Final = item.get("type") == "custom_tool_call"
function: Final[_ToolFunction] = {
"name": as_str(item.get("name")) or "",
"arguments": as_str(item.get("input" if custom else "arguments")) or "",
}
tool_call: Final[_ToolCall] = {
"id": as_str(item.get("call_id")) or as_str(item.get("id")) or "",
"type": "function",
"function": function,
}
return tool_call
def _responses_finish_reason(response: Mapping[str, object], has_tool_calls: bool) -> str | None:
status: Final = as_str(response.get("status"))
if status == "completed":
return "tool_calls" if has_tool_calls else "stop"
if status != "incomplete":
return None
details: Final = as_str_mapping(response.get("incomplete_details"))
reason: Final = details.get("reason") if details is not None else None
return "content_filter" if reason == "content_filter" else "length"
def _parse_error(payload: StandardLoggingPayload) -> SpanError | None:
"""A ``SpanError`` for a failed request, or ``None`` on success."""
if payload.get("status") != "failure":

85
litellm/litellm_core_utils/litellm_logging.py
View file

@ -212,6 +212,7 @@ if TYPE_CHECKING:
from litellm.integrations.otel.model.config import ExporterSpec, OpenTelemetryV2Config
from litellm.litellm_core_utils.llm_cost_calc.utils import BilledTokenRates
from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector
from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation
try:
from litellm_enterprise.enterprise_callbacks.callback_controls import (
EnterpriseCallbackControls,
@ -501,6 +502,8 @@ class Logging(LiteLLMLoggingBaseClass):
litellm_request_debug: bool = False
streamed_anthropic_message_id: str | None = None
classifier_input: Mapping[str, JsonValue] | None = None
baseline_cache_context: "BaselineCacheContext | None" = None
baseline_observation: "CapturedBaselineObservation | None" = None
def __init__(
self,
@ -508,7 +511,7 @@ class Logging(LiteLLMLoggingBaseClass):
messages,
stream,
call_type,
start_time,
start_time: datetime.datetime,
litellm_call_id: str,
function_id: str,
litellm_trace_id: str | None = None,
@ -2181,6 +2184,7 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result,
start_time,
end_time,
build_logging_payload: bool = True,
):
"""Resolve hidden params, compute response cost, and emit the standard logging payload."""
hidden_params: Final = getattr(logging_result, "_hidden_params", {})
@ -2205,6 +2209,9 @@ class Logging(LiteLLMLoggingBaseClass):
else:
self.model_call_details["response_cost"] = self._response_cost_calculator(result=logging_result)
if not build_logging_payload:
return
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
logging_result, start_time, end_time
)
@ -2215,6 +2222,19 @@ class Logging(LiteLLMLoggingBaseClass):
if standard_logging_payload is not None:
emit_standard_logging_payload(standard_logging_payload)
async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None:
if self.baseline_cache_context is None:
return
from litellm.proxy.hooks.autorouter_baseline_cache import finalize_baseline_cache
await finalize_baseline_cache(self, response_obj)
async def invalidate_baseline_cache_estimate(self, reason: str, *, completed: bool = False) -> None:
"""Invalidate uncertain attempts; retire the reservation at logical completion."""
from litellm.proxy.hooks.autorouter_baseline_cache import invalidate_baseline_cache
await invalidate_baseline_cache(self, reason, completed=completed)
def _build_standard_logging_payload(
self, init_response_obj: object, start_time: Any, end_time: Any
) -> StandardLoggingPayload | None:
@ -2266,6 +2286,7 @@ class Logging(LiteLLMLoggingBaseClass):
end_time=None,
cache_hit=None,
standard_logging_object: StandardLoggingPayload | None = None,
build_logging_payload: bool = True,
):
try:
if start_time is None:
@ -2303,6 +2324,7 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result=logging_result,
start_time=start_time,
end_time=end_time,
build_logging_payload=build_logging_payload,
)
elif standard_logging_object is not None:
self.model_call_details["standard_logging_object"] = standard_logging_object
@ -3051,8 +3073,17 @@ class Logging(LiteLLMLoggingBaseClass):
result=result,
cache_hit=cache_hit,
standard_logging_object=kwargs.get("standard_logging_object", None),
build_logging_payload=self.baseline_cache_context is None,
)
if self.stream is not True and self.baseline_cache_context is not None:
await self._prepare_baseline_cache_estimate(result)
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
result, start_time, end_time
)
if (prepared_payload := self.model_call_details.get("standard_logging_object")) is not None:
emit_standard_logging_payload(prepared_payload)
## BUILD COMPLETE STREAMED RESPONSE
if "async_complete_streaming_response" in self.model_call_details:
return # break out of this.
@ -3097,6 +3128,8 @@ class Logging(LiteLLMLoggingBaseClass):
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
await self._prepare_baseline_cache_estimate(complete_streaming_response)
## STANDARDIZED LOGGING PAYLOAD
try:
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
@ -3125,6 +3158,7 @@ class Logging(LiteLLMLoggingBaseClass):
# Only build standard_logging_object if not already built by
# _success_handler_helper_fn
if self.model_call_details.get("standard_logging_object") is None:
await self._prepare_baseline_cache_estimate(result)
## STANDARDIZED LOGGING PAYLOAD
self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload(
result, start_time, end_time
@ -3631,6 +3665,8 @@ class Logging(LiteLLMLoggingBaseClass):
"""
Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
"""
if self.baseline_cache_context is not None:
await self.invalidate_baseline_cache_estimate("failed_request")
await self.special_failure_handlers(exception=exception)
if not self.should_run_logging(event_type="async_failure"): # prevent double logging
return
@ -5528,6 +5564,10 @@ class StandardLoggingPayloadSetup:
for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
clean_metadata[key] = metadata[key]
recorded_guardrails: Final = metadata.get("applied_guardrails")
if applied_guardrails and isinstance(recorded_guardrails, list):
clean_metadata["applied_guardrails"] = list(dict.fromkeys([*applied_guardrails, *recorded_guardrails]))
user_api_key: Final = metadata.get("user_api_key")
if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key):
clean_metadata["user_api_key_hash"] = user_api_key
@ -6149,6 +6189,8 @@ def _autorouter_savings_for_payload(
model_id: str | None,
usage_object: Mapping[str, object] | None,
cost_breakdown: Mapping[str, object] | None,
baseline_usage: Usage | None = None,
baseline_provenance: Literal["observed_initial", "modeled"] | None = None,
) -> float | None:
"""The auto-router savings figure for the payload, or ``None`` when there is none.
@ -6167,6 +6209,8 @@ def _autorouter_savings_for_payload(
model_id=model_id,
usage_object=usage_object,
cost_breakdown=cost_breakdown,
baseline_usage=baseline_usage,
baseline_provenance=baseline_provenance,
)
except Exception as e: # noqa: BLE001 # a savings figure must never fail request logging
verbose_logger.debug("autorouter savings skipped on logging payload: %s", e)
@ -6343,13 +6387,18 @@ def get_standard_logging_object_payload(
model_name = response_model_name
request_cost_breakdown: Final = cost_breakdown_with_guardrail(logging_obj.cost_breakdown, guardrail_cost)
autorouter_savings: Final = _autorouter_savings_for_payload(
request_metadata=metadata,
model=model_name,
custom_llm_provider=custom_llm_provider,
model_id=_model_id,
usage_object=usage_dict,
cost_breakdown=request_cost_breakdown,
captured_baseline: Final = logging_obj.baseline_observation
autorouter_savings: Final = (
None
if status != "success" or cache_hit or logging_obj.baseline_cache_context is not None
else _autorouter_savings_for_payload(
request_metadata=metadata,
model=model_name,
custom_llm_provider=custom_llm_provider,
model_id=_model_id,
usage_object=usage_dict,
cost_breakdown=request_cost_breakdown,
)
)
payload: Final[StandardLoggingPayload] = StandardLoggingPayload(
@ -6396,6 +6445,26 @@ def get_standard_logging_object_payload(
response_cost=response_cost,
cost_breakdown=request_cost_breakdown,
autorouter_savings=autorouter_savings,
autorouter_savings_estimate=(
{
"version": 3,
"status": "unknown",
"reason": "pending_projection",
} # mutable-ok: spend-log JSON serialization requires plain mappings
if captured_baseline is not None
else (
{ # mutable-ok: spend-log JSON serialization requires plain mappings
"version": 1,
"status": "estimated" if autorouter_savings is not None else "unknown",
"reason": "uncached_usage" if autorouter_savings is not None else "baseline_unavailable",
}
if metadata.get("routing_decision")
else None
)
),
autorouter_baseline_observation=(
captured_baseline.model_dump_json() if captured_baseline is not None else None
),
total_tokens=usage_dict.get("total_tokens", 0),
prompt_tokens=usage_dict.get("prompt_tokens", 0),
completion_tokens=usage_dict.get("completion_tokens", 0),

8
litellm/litellm_core_utils/redact_messages.py
View file

@ -128,6 +128,8 @@ def _redact_responses_api_output(output_items):
for content_part in output_item.content:
if getattr(content_part, "text", None) is not None:
content_part.text = REDACTED_BY_LITELLM
if getattr(content_part, "refusal", None) is not None:
content_part.refusal = REDACTED_BY_LITELLM
# Redact reasoning items in output array
if hasattr(output_item, "type") and output_item.type == "reasoning":
@ -138,6 +140,8 @@ def _redact_responses_api_output(output_items):
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
output_item.arguments = REDACTED_BY_LITELLM
if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"):
output_item.input = REDACTED_BY_LITELLM
def _redact_responses_api_output_dict(output_items, redacted_str: str):
@ -153,6 +157,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
for content_item in output_item["content"]:
if isinstance(content_item, dict) and content_item.get("text") is not None:
content_item["text"] = redacted_str
if isinstance(content_item, dict) and content_item.get("refusal") is not None:
content_item["refusal"] = redacted_str
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
for summary_item in output_item["summary"]:
@ -161,6 +167,8 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
if output_item.get("type") == "function_call" and "arguments" in output_item:
output_item["arguments"] = redacted_str
if output_item.get("type") == "custom_tool_call" and "input" in output_item:
output_item["input"] = redacted_str
def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[str, object]:

11
litellm/llms/anthropic/chat/handler.py
View file

@ -4,7 +4,8 @@ Calling + translation logic for anthropic's `/v1/messages` endpoint
import copy
import json
from collections.abc import Callable
from collections.abc import Callable, Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
import httpx
@ -31,7 +32,6 @@ from litellm.types.llms.anthropic import (
ContentBlockStop,
MessageBlockDelta,
MessageStartBlock,
UsageDelta,
)
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
@ -557,6 +557,7 @@ class ModelResponseIterator:
self.tool_index = -1
self.json_mode = json_mode
self.speed = speed
self._cumulative_usage: Mapping[str, object] = MappingProxyType({})
# rewritten-name -> caller's original. Built per-request from the
# forward map in AnthropicConfig._build_request_tool_name_maps; only
# contains entries we actually rewrote, so a tool legitimately named
@ -631,10 +632,12 @@ class ModelResponseIterator:
return True
return False
def _handle_usage(self, anthropic_usage_chunk: dict | UsageDelta) -> Usage:
def _handle_usage(self, anthropic_usage_chunk: Mapping[str, object]) -> Usage:
# message_delta usage is cumulative but may omit fields reported at message_start.
self._cumulative_usage = MappingProxyType({**self._cumulative_usage, **anthropic_usage_chunk})
reasoning_content: Final = "".join(self.reasoning_content_chunks) if self.reasoning_content_chunks else None
usage: Final = AnthropicConfig().calculate_usage(
usage_object=cast(dict, anthropic_usage_chunk),
usage_object=self._cumulative_usage,
reasoning_content=reasoning_content,
speed=self.speed,
)

20
litellm/llms/anthropic/count_tokens/handler.py
View file

@ -4,9 +4,11 @@ Anthropic CountTokens API handler.
Uses httpx for HTTP requests instead of the Anthropic SDK.
"""
from typing import Any, Final
from collections.abc import Mapping
from typing import Final
import httpx
from pydantic import JsonValue, TypeAdapter
import litellm
from litellm._logging import verbose_logger
@ -16,6 +18,8 @@ from litellm.llms.anthropic.count_tokens.transformation import (
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
_COUNT_RESPONSE: Final = TypeAdapter(dict[str, JsonValue])
class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
"""
@ -27,13 +31,14 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
async def handle_count_tokens_request(
self,
model: str,
messages: list[dict[str, Any]],
messages: list[dict[str, JsonValue]],
api_key: str,
api_base: str | None = None,
timeout: float | httpx.Timeout | None = None,
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
) -> dict[str, Any]:
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
) -> dict[str, JsonValue]:
"""
Handle a CountTokens request using httpx.
@ -52,7 +57,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
"""
try:
# Validate the request
self.validate_request(model, messages)
self.validate_request(model, messages, system=system, tools=tools)
verbose_logger.debug("Processing Anthropic CountTokens request for model: %s", model)
@ -62,6 +67,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
messages=messages,
tools=tools,
system=system,
optional_params=optional_params,
)
verbose_logger.debug("Transformed request: %s", request_body)
@ -97,7 +103,7 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig):
message=error_text,
)
anthropic_response: Final = response.json()
anthropic_response: Final = _COUNT_RESPONSE.validate_json(response.content)
verbose_logger.debug("Anthropic response: %s", anthropic_response)

56
litellm/llms/anthropic/count_tokens/transformation.py
View file

@ -4,10 +4,17 @@ Anthropic CountTokens API transformation logic.
This module handles the transformation of requests to Anthropic's CountTokens API format.
"""
from typing import Any, Final
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final
from pydantic import JsonValue, TypeAdapter
from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION
_COUNT_REQUEST: Final = TypeAdapter(dict[str, JsonValue])
COUNT_TOKEN_OPTION_NAMES: Final = ("thinking", "tool_choice", "output_config")
class AnthropicCountTokensConfig:
"""
@ -31,27 +38,31 @@ class AnthropicCountTokensConfig:
def transform_request_to_count_tokens(
self,
model: str,
messages: list[dict[str, Any]],
tools: list[dict[str, Any]] | None = None,
system: Any | None = None,
) -> dict[str, Any]:
messages: list[dict[str, JsonValue]],
tools: list[dict[str, JsonValue]] | None = None,
system: JsonValue = None,
optional_params: Mapping[str, JsonValue] | None = None,
) -> dict[str, JsonValue]: # mutable-ok: provider transport requires JSON dictionaries
"""
Transform request to Anthropic CountTokens format.
Includes optional system and tools fields for accurate token counting.
"""
request: Final[dict[str, Any]] = {
"model": model,
"messages": messages,
}
if system is not None:
request["system"] = system
if tools is not None:
request["tools"] = tools
return request
options: Final[Mapping[str, JsonValue]] = optional_params or MappingProxyType({})
return _COUNT_REQUEST.validate_python(
MappingProxyType(
{
"model": model,
"messages": messages,
**MappingProxyType(
{key: value for key, value in (("system", system), ("tools", tools)) if value is not None}
),
**MappingProxyType(
{key: value for key, value in options.items() if key in COUNT_TOKEN_OPTION_NAMES}
),
}
)
)
def get_required_headers(self, api_key: str) -> dict[str, str]:
"""
@ -76,7 +87,14 @@ class AnthropicCountTokensConfig:
headers, _ = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key)
return headers
def validate_request(self, model: str, messages: list[dict[str, Any]]) -> None:
def validate_request(
self,
model: str,
messages: Sequence[Mapping[str, JsonValue]],
*,
system: JsonValue = None,
tools: list[dict[str, JsonValue]] | None = None,
) -> None:
"""
Validate the incoming count tokens request.
@ -90,7 +108,7 @@ class AnthropicCountTokensConfig:
if not model:
raise ValueError("model parameter is required")
if not messages:
if not messages and not system and not tools:
raise ValueError("messages parameter is required")
if not isinstance(messages, list):

462
litellm/llms/anthropic/prompt_cache_prediction.py
View file

@ -1,10 +1,11 @@
from __future__ import annotations
import asyncio
import hashlib
import json
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from dataclasses import dataclass, field
from itertools import accumulate, groupby
from types import MappingProxyType
from typing import Annotated, Final, Literal, Protocol, TypeAlias
@ -14,9 +15,14 @@ from pydantic import BaseModel, ConfigDict, Field, JsonValue, StrictInt, TypeAda
import litellm
from litellm.llms.anthropic.common_utils import AnthropicModelInfo, is_anthropic_oauth_key
from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import DEFAULT_ANTHROPIC_API_VERSION
from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES
from litellm.llms.anthropic.experimental_pass_through.messages.transformation import (
DEFAULT_ANTHROPIC_API_VERSION,
AnthropicMessagesConfig,
)
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import ModelResponse
from litellm.utils import supports_thinking_cache_preservation
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
_HEADERS: Final = TypeAdapter(dict[str, str])
@ -100,10 +106,7 @@ _Block: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult, Field(discriminato
class _Message(_StrictModel):
role: Literal["user", "assistant"]
content: str | Annotated[tuple[_Block, ...], Field(strict=False)]
def blocks(self) -> tuple[_Text | _ToolUse | _ToolResult, ...]:
return (_Text(type="text", text=self.content),) if isinstance(self.content, str) else tuple(self.content)
content: Annotated[str, Field(min_length=1, pattern=r"\S")] | Annotated[tuple[_Block, ...], Field(strict=False)]
class _Tool(_StrictModel):
@ -113,10 +116,7 @@ class _Tool(_StrictModel):
type: Literal["custom"] | None = None
class _Request(_StrictModel):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
class _RequestOptions(_StrictModel):
model: str | None = None
max_tokens: int | None = None
stream: bool | None = None
@ -127,6 +127,289 @@ class _Request(_StrictModel):
metadata: Mapping[str, JsonValue] | None = None
class _Request(_RequestOptions):
messages: tuple[_Message, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_ResultText, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_Tool, ...], Field(strict=False)] | None = None
class _Thinking(_StrictModel):
type: Literal["thinking"]
thinking: str
signature: str = Field(min_length=1)
_PlanBlock: TypeAlias = Annotated[_Text | _ToolUse | _ToolResult | _Thinking, Field(discriminator="type")]
class _PlanMessage(_StrictModel):
role: Literal["user", "assistant", "system"]
content: str | Annotated[tuple[_PlanBlock, ...], Field(strict=False)]
class _PlanTool(_Tool):
cache_control: _CacheControl | None = None
class _PlanRequest(_RequestOptions):
messages: tuple[_PlanMessage, ...] = Field(min_length=1, strict=False)
system: str | Annotated[tuple[_Text, ...], Field(strict=False)] | None = None
tools: Annotated[tuple[_PlanTool, ...], Field(strict=False)] | None = None
cache_control: _CacheControl | None = None
thinking: Mapping[str, JsonValue] | None = None
tool_choice: Mapping[str, JsonValue] | None = None
output_config: Mapping[str, JsonValue] | None = None
speed: Literal["fast", "standard"] | None = None
service_tier: Literal["auto", "standard_only"] | None = None
@dataclass(frozen=True, slots=True)
class CacheBoundary:
fingerprint: str
prefix_body: Mapping[str, JsonValue] = field(repr=False)
ttl_seconds: int
lookback_fingerprints: tuple[str, ...]
content_fingerprint: str = ""
lookback_content_fingerprints: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class PromptCachePlan:
full_body: Mapping[str, JsonValue] = field(repr=False)
breakpoints: tuple[CacheBoundary, ...]
@dataclass(frozen=True, slots=True)
class UnsupportedCachePlan:
reason: Literal[
"unsupported_prompt_shape",
"conflicting_cache_ttl",
"too_many_cache_breakpoints",
"invalid_cache_ttl_order",
"unsupported_thinking_cache_semantics",
"token_count_unavailable",
"inconsistent_prefix_token_count",
]
@dataclass(frozen=True, slots=True)
class CountedBreakpoint:
fingerprint: str
ttl_seconds: int
prefix_tokens: int
lookback_fingerprints: tuple[str, ...]
content_fingerprint: str = ""
lookback_content_fingerprints: tuple[str, ...] = ()
@dataclass(frozen=True, slots=True)
class CountedPromptCachePlan:
total_tokens: int
breakpoints: tuple[CountedBreakpoint, ...]
@dataclass(frozen=True, slots=True)
class _Position:
section: Literal["tools", "system", "messages"]
message_index: int
role: str
block: Mapping[str, JsonValue]
marker: _CacheControl | None
def _content_blocks(content: JsonValue) -> tuple[Mapping[str, JsonValue], ...]:
if isinstance(content, str):
return (MappingProxyType({"type": "text", "text": content}),)
return tuple(_JSON_OBJECT.validate_python(block) for block in content) if isinstance(content, list) else ()
def _position(
section: Literal["tools", "system", "messages"],
message_index: int,
role: str,
block: Mapping[str, JsonValue],
) -> _Position:
control: Final = block.get("cache_control")
return _Position(
section,
message_index,
role,
MappingProxyType({key: value for key, value in block.items() if key != "cache_control"}),
_CacheControl.model_validate(control) if control is not None else None,
)
def _positions(body: Mapping[str, JsonValue]) -> tuple[_Position, ...]:
tools: Final = body.get("tools")
messages: Final = body.get("messages")
return (
*tuple(
_position("tools", -1, "", _JSON_OBJECT.validate_python(tool))
for tool in (tools if isinstance(tools, list) else ())
),
*tuple(_position("system", -1, "", block) for block in _content_blocks(body.get("system"))),
*tuple(
_position("messages", message_index, str(message.get("role")), block)
for message_index, raw_message in enumerate(messages if isinstance(messages, list) else ())
for message in (_JSON_OBJECT.validate_python(raw_message),)
for block in _content_blocks(message.get("content"))
),
)
def _prefix_body(
body: Mapping[str, JsonValue],
positions: tuple[_Position, ...],
last_index: int,
) -> Mapping[str, JsonValue]:
prefix: Final = positions[: last_index + 1]
sections: Final = MappingProxyType(
{
section: _count_objects(tuple(position.block for position in prefix if position.section == section))
for section in ("tools", "system")
if any(position.section == section for position in prefix)
}
)
messages: Final = tuple(
MappingProxyType(
_JSON_OBJECT.validate_python(
MappingProxyType(
{"role": group[0].role, "content": _count_objects(tuple(position.block for position in group))}
)
)
)
for _, values in groupby(
(position for position in prefix if position.section == "messages"),
key=lambda position: position.message_index,
)
for group in (tuple(values),)
)
return MappingProxyType(
_JSON_OBJECT.validate_python(
MappingProxyType(
{
**MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body}),
**sections,
"messages": _count_objects(messages),
}
)
)
)
def _position_group(position: _Position, index: int) -> tuple[str, int, str | int]:
block_type: Final = position.block.get("type")
return (
position.section,
position.message_index,
block_type if isinstance(block_type, str) and block_type in ("tool_use", "tool_result") else index,
)
def _chain_digest(previous: str, current: str) -> str:
return _digest((previous, current))
def _cacheable_position(position: _Position) -> bool:
block_type: Final = position.block.get("type")
if block_type == "thinking":
return False
text: Final = position.block.get("text")
return block_type != "text" or (isinstance(text, str) and bool(text.strip()))
def _entry_fingerprint(fingerprint: str, ttl_seconds: int) -> str:
return _digest(("native-cache-prefix-v2", fingerprint, ttl_seconds))
def parse_cache_plan(body: Mapping[str, JsonValue]) -> PromptCachePlan | UnsupportedCachePlan:
try:
request: Final = _PlanRequest.model_validate(body)
positions: Final = _positions(body)
except ValidationError:
return UnsupportedCachePlan("unsupported_prompt_shape")
explicit: Final = tuple(
(index, position.marker) for index, position in enumerate(positions) if position.marker is not None
)
automatic_index: Final = next(
(index for index in reversed(range(len(positions))) if _cacheable_position(positions[index])), None
)
automatic_existing: Final = next((marker for index, marker in explicit if index == automatic_index), None)
if (
request.cache_control is not None
and automatic_existing is not None
and automatic_existing != request.cache_control
):
return UnsupportedCachePlan("conflicting_cache_ttl")
automatic: Final = (
((automatic_index, request.cache_control),)
if (request.cache_control is not None and automatic_index is not None and automatic_existing is None)
else ()
)
markers: Final = tuple(sorted((*explicit, *automatic), key=lambda value: value[0]))
if len(markers) > 4:
return UnsupportedCachePlan("too_many_cache_breakpoints")
ttls: Final = tuple(3600 if marker.ttl == "1h" else 300 for _, marker in markers)
if any(first < second for first, second in zip(ttls, ttls[1:])):
return UnsupportedCachePlan("invalid_cache_ttl_order")
settings: Final = MappingProxyType(
{
key: body[key]
for key in ("thinking", "output_config", "speed")
if key in body and not (key == "speed" and body[key] == "standard")
}
)
hashes: Final = tuple(
accumulate(
(
_digest(
(
position.section,
position.message_index,
position.role,
position.block,
body.get("tool_choice") if position.section == "messages" else None,
)
)
for position in positions
),
_chain_digest,
initial=_digest(settings),
)
)[1:]
groups: Final = tuple(
tuple(index for index, _ in values)
for _, values in groupby(
enumerate(positions),
key=lambda item: _position_group(item[1], item[0]),
)
)
return PromptCachePlan(
full_body=MappingProxyType(dict(body)),
breakpoints=tuple(
CacheBoundary(
fingerprint=_entry_fingerprint(hashes[index], ttl),
prefix_body=_prefix_body(body, positions, index),
ttl_seconds=ttl,
lookback_fingerprints=tuple(
_entry_fingerprint(hashes[earlier], ttl)
for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:])
for earlier in reversed(group)
if earlier <= index
),
content_fingerprint=hashes[index],
lookback_content_fingerprints=tuple(
hashes[earlier]
for group in reversed(tuple(group for group in groups if group[0] <= index)[-20:])
for earlier in reversed(group)
if earlier <= index
),
)
for (index, _), ttl in zip(markers, ttls)
),
)
@dataclass(frozen=True, slots=True)
class PromptPrefix:
prefix_body: Mapping[str, JsonValue]
@ -137,68 +420,28 @@ class PromptPrefix:
def _digest(value: object) -> str:
return hashlib.sha256(
json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()
json.dumps(value, default=_json_object, separators=(",", ":"), ensure_ascii=False).encode()
).hexdigest()
def _next_digest(previous: str, boundary: tuple[int, str, Mapping[str, JsonValue]]) -> str:
return _digest((previous, boundary))
def _json_object(value: object) -> dict[str, JsonValue]: # mutable-ok: JSON serialization requires a dictionary
return _JSON_OBJECT.validate_python(value)
def parse_prompt(body: Mapping[str, JsonValue]) -> PromptPrefix | None:
try:
request: Final = _Request.model_validate(body)
blocks: Final = tuple(message.blocks() for message in request.messages)
_Request.model_validate(body)
except ValidationError:
return None
markers: Final = tuple(
(message_index, block_index, block.cache_control)
for message_index, message_blocks in enumerate(blocks)
for block_index, block in enumerate(message_blocks)
if block.cache_control is not None
)
if len(markers) != 1:
plan: Final = parse_cache_plan(body)
if isinstance(plan, UnsupportedCachePlan) or len(plan.breakpoints) != 1:
return None
message_end, block_end, marker = markers[0]
normalized: Final = _JSON_OBJECT.validate_python(request.model_dump(mode="json", exclude_none=True))
context: Final = MappingProxyType({key: normalized[key] for key in ("system", "tools") if key in normalized})
boundaries: Final = tuple(
(
message_index,
request.messages[message_index].role,
_JSON_OBJECT.validate_python(
block.model_dump(mode="json", exclude=MappingProxyType({"cache_control": True}), exclude_none=True)
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
)
hashes: Final = tuple(
accumulate(boundaries, _next_digest, initial=_digest((_JSON_OBJECT.validate_python(context), marker.ttl)))
)[1:]
prefix_messages: Final = tuple(
_Message(
role=request.messages[message_index].role,
content=tuple(
block
for block_index, block in enumerate(message_blocks)
if message_index < message_end or block_index <= block_end
),
)
for message_index, message_blocks in enumerate(blocks[: message_end + 1])
)
prefix: Final = plan.breakpoints[0]
return PromptPrefix(
prefix_body=MappingProxyType(
_JSON_OBJECT.validate_python(
_Request(messages=prefix_messages, system=request.system, tools=request.tools).model_dump(
mode="json", exclude_none=True
)
)
),
fingerprint=hashes[-1],
fingerprints=tuple(reversed(hashes[-20:])),
ttl_seconds=3600 if marker.ttl == "1h" else 300,
prefix_body=prefix.prefix_body,
fingerprint=prefix.fingerprint,
fingerprints=prefix.lookback_fingerprints,
ttl_seconds=prefix.ttl_seconds,
)
@ -246,6 +489,9 @@ class _CountBody(BaseModel):
messages: Sequence[Mapping[str, JsonValue]]
tools: Sequence[Mapping[str, JsonValue]] | None = None
system: str | Sequence[Mapping[str, JsonValue]] | None = None
thinking: Mapping[str, JsonValue] | None = None
tool_choice: Mapping[str, JsonValue] | None = None
output_config: Mapping[str, JsonValue] | None = None
class _CountResult(BaseModel):
@ -262,16 +508,36 @@ def _count_objects(
return [dict(value) for value in values] # mutable-ok: serialize read-only inputs at the provider API boundary
async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
native: Final = _CountBody.model_validate(body)
def _messages_url(model: str, api_key: str, api_base: str | None) -> str:
return AnthropicMessagesConfig().get_complete_url( # pyright: ignore[reportUnknownMemberType] # canonical native URL owner takes legacy JSON arguments
api_base=api_base,
api_key=api_key,
model=model,
optional_params=_JSON_OBJECT.validate_python(MappingProxyType({})),
litellm_params=_JSON_OBJECT.validate_python(MappingProxyType({})),
)
async def count_prompt_tokens(
model: str,
api_key: str,
body: Mapping[str, JsonValue],
api_base: str | None = None,
) -> int | None:
try:
native: Final = _CountBody.model_validate(body)
count_url: Final = _messages_url(model, api_key, api_base) + "/count_tokens"
result: Final = _CountResult.model_validate(
await _counter.handle_count_tokens_request(
model=model,
messages=_count_objects(native.messages),
tools=_count_objects(native.tools) if native.tools is not None else None,
system=native.system,
system=_JSON_OBJECT.validate_python(MappingProxyType({"system": native.system}))["system"],
api_key=api_key,
api_base=count_url,
optional_params=_JSON_OBJECT.validate_python(
MappingProxyType({key: body[key] for key in COUNT_TOKEN_OPTION_NAMES if key in body})
),
timeout=15.0,
)
)
@ -280,10 +546,55 @@ async def count_prompt_tokens(model: str, api_key: str, body: Mapping[str, JsonV
return result.input_tokens
async def count_cache_plan(
model: str,
api_key: str,
plan: PromptCachePlan,
token_counter: TokenCounter = count_prompt_tokens,
) -> CountedPromptCachePlan | UnsupportedCachePlan:
if any(position.block.get("type") == "thinking" for position in _positions(plan.full_body)):
if not supports_thinking_cache_preservation(model, "anthropic"):
return UnsupportedCachePlan("unsupported_thinking_cache_semantics")
total: Final = await token_counter(model, api_key, plan.full_body)
if total is None:
return UnsupportedCachePlan("token_count_unavailable")
counts: Final = tuple(
await asyncio.gather(*(token_counter(model, api_key, marker.prefix_body) for marker in plan.breakpoints))
)
if any(value is None for value in counts):
return UnsupportedCachePlan("token_count_unavailable")
known: Final = tuple(value for value in counts if value is not None)
if any(value < 0 for value in (total, *known)) or any(
first > second for first, second in zip(known, (*known[1:], total))
):
return UnsupportedCachePlan("inconsistent_prefix_token_count")
return CountedPromptCachePlan(
total,
tuple(
CountedBreakpoint(
marker.fingerprint,
marker.ttl_seconds,
count,
marker.lookback_fingerprints,
marker.content_fingerprint,
marker.lookback_content_fingerprints,
)
for marker, count in zip(plan.breakpoints, known)
),
)
@dataclass(frozen=True, slots=True)
class NativePredictionTarget:
model: str
api_key: str
api_key: str = field(repr=False)
api_base: str | None = None
def supported_baseline_recipient(target: NativePredictionTarget, wire: httpx.Request) -> bool:
return wire.headers.get("x-api-key") == target.api_key and wire.url == httpx.URL(
_messages_url(target.model, target.api_key, target.api_base)
)
@dataclass(frozen=True, slots=True)
@ -297,11 +608,26 @@ class UnsupportedPredictionTarget:
def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
return _resolve_prediction_target(params, allow_configured_endpoint=False)
def resolve_baseline_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget | UnsupportedPredictionTarget:
return _resolve_prediction_target(params, allow_configured_endpoint=True)
def _resolve_prediction_target(
params: LiteLLM_Params,
*,
allow_configured_endpoint: bool,
) -> NativePredictionTarget | UnsupportedPredictionTarget:
configured_options: Final = frozenset(params.model_dump(exclude_defaults=True, exclude_none=True))
if configured_options - _DEPLOYMENT_OPTIONS:
return UnsupportedPredictionTarget("unsupported_deployment_configuration")
api_base: Final = AnthropicModelInfo.get_api_base(params.api_base)
if api_base not in ("https://api.anthropic.com", "https://api.anthropic.com/v1/messages"):
if not allow_configured_endpoint and api_base not in (
"https://api.anthropic.com",
"https://api.anthropic.com/v1/messages",
):
return UnsupportedPredictionTarget("unsupported_provider_endpoint")
try:
model, provider, _, _ = litellm.get_llm_provider(
@ -314,7 +640,7 @@ def resolve_prediction_target(params: LiteLLM_Params) -> NativePredictionTarget
api_key: Final = AnthropicModelInfo.get_api_key(params.api_key)
if api_key is None or not _supported_provider_key(api_key):
return UnsupportedPredictionTarget("unsupported_provider_credentials")
return NativePredictionTarget(model=model, api_key=api_key)
return NativePredictionTarget(model=model, api_key=api_key, api_base=api_base)
def _supported_provider_key(api_key: str) -> bool:

9
litellm/llms/azure/chat/gpt_transformation.py
View file

@ -280,10 +280,17 @@ class AzureOpenAIConfig(BaseConfig):
ordered_messages: Final = system_messages_first(messages) if litellm.openai_system_messages_first else messages
stripped_messages: Final = drop_tool_reference_parts_from_tool_messages(ordered_messages)
azure_messages: Final = convert_to_azure_openai_messages(hoist_images_from_tool_messages(stripped_messages))
request_params: Final = MappingProxyType(
{
key: value
for key, value in optional_params.items()
if key != "tool_choice" or optional_params.get("tools") or optional_params.get("functions")
}
)
return {
"model": model,
"messages": azure_messages,
**optional_params,
**request_params,
**sanitized_tools_update(optional_params),
}

40
litellm/llms/base_llm/files/litellm_db_storage_backend.py Normal file
View file

@ -0,0 +1,40 @@
from typing import TYPE_CHECKING, Final
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.repositories.managed_file_content_repository import ManagedFileContentRepository
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
LITELLM_DB_STORAGE_BACKEND_NAME: Final = "litellm_db"
LITELLM_DB_STORAGE_URL_PREFIX: Final = f"{LITELLM_DB_STORAGE_BACKEND_NAME}://"
def storage_url_to_row_id(storage_url: str) -> str:
if not storage_url.startswith(LITELLM_DB_STORAGE_URL_PREFIX):
raise ValueError(f"Not a {LITELLM_DB_STORAGE_BACKEND_NAME} storage url: {storage_url}")
return storage_url.removeprefix(LITELLM_DB_STORAGE_URL_PREFIX)
class LiteLLMDbStorageBackend(BaseFileStorageBackend):
def __init__(self, prisma_client: "PrismaClient") -> None:
self._contents = ManagedFileContentRepository(prisma_client)
async def upload_file(
self,
file_content: bytes,
filename: str,
content_type: str,
path_prefix: str | None = None,
file_naming_strategy: str = "uuid",
) -> str:
return f"{LITELLM_DB_STORAGE_URL_PREFIX}{await self._contents.store(file_content)}"
async def download_file(self, storage_url: str) -> bytes:
content: Final = await self._contents.load(storage_url_to_row_id(storage_url))
if content is None:
raise ValueError(f"No stored file content for {storage_url}")
return content
async def delete_file(self, storage_url: str) -> None:
await self._contents.delete(storage_url_to_row_id(storage_url))

26
litellm/llms/base_llm/files/storage_backend_factory.py
View file

@ -6,32 +6,46 @@ based on the backend type. Backends use the same configuration as their correspo
callbacks (e.g., azure_storage uses the same env vars as AzureBlobStorageLogger).
"""
from typing import TYPE_CHECKING
from litellm._logging import verbose_logger
from .azure_blob_storage_backend import AzureBlobStorageBackend
from .litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME, LiteLLMDbStorageBackend
from .storage_backend import BaseFileStorageBackend
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def get_storage_backend(backend_type: str) -> BaseFileStorageBackend:
def get_storage_backend(backend_type: str, prisma_client: "PrismaClient | None" = None) -> BaseFileStorageBackend:
"""
Factory function to create a storage backend instance.
Backends are configured using the same environment variables as their
corresponding callbacks. For example, "azure_storage" uses the same
env vars as AzureBlobStorageLogger.
env vars as AzureBlobStorageLogger. "litellm_db" stores file bytes in the
proxy's own database and needs the connected Prisma client.
Args:
backend_type: Backend type identifier (e.g., "azure_storage")
backend_type: Backend type identifier (e.g., "azure_storage", "litellm_db")
prisma_client: The proxy's database client, required by "litellm_db"
Returns:
BaseFileStorageBackend: Instance of the appropriate storage backend
Raises:
ValueError: If backend_type is not supported
ValueError: If backend_type is not supported, or "litellm_db" is asked for without a database
"""
verbose_logger.debug("Creating storage backend: type=%s", backend_type)
if backend_type == "azure_storage":
return AzureBlobStorageBackend()
else:
raise ValueError(f"Unsupported storage backend type: {backend_type}. Supported types: azure_storage")
if backend_type == LITELLM_DB_STORAGE_BACKEND_NAME:
if prisma_client is None:
raise ValueError(f"Storage backend {LITELLM_DB_STORAGE_BACKEND_NAME} requires a database-connected proxy")
return LiteLLMDbStorageBackend(prisma_client)
raise ValueError(
f"Unsupported storage backend type: {backend_type}. "
f"Supported types: azure_storage, {LITELLM_DB_STORAGE_BACKEND_NAME}"
)

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

@ -2163,6 +2163,8 @@ class BaseLLMHTTPHandler:
e=e, litellm_params=litellm_params_dict
)
if should_retry and not hit_max_attempt:
if logging_obj.baseline_cache_context is not None:
await logging_obj.invalidate_baseline_cache_estimate("retried_request")
verbose_logger.debug(
"Anthropic /v1/messages: invalid thinking signature; "
"stripping thinking blocks and retrying (attempt %s/%s).",

10
litellm/model_prices_and_context_window_backup.json
View file

@ -14510,6 +14510,7 @@
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
@ -14547,6 +14548,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14698,6 +14700,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14727,6 +14730,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14759,6 +14763,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14796,6 +14801,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14831,6 +14837,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14869,6 +14876,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14986,6 +14994,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -15027,6 +15036,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,

15
litellm/models/autorouter_session.py
View file

@ -8,6 +8,8 @@ maintains per (api_key, session_id, router_name).
from collections.abc import Mapping
from datetime import datetime
from pydantic import Field
from litellm.types.llms.base import LiteLLMPydanticObjectBase
@ -22,18 +24,25 @@ class LiteLLM_AutoRouterSession(LiteLLMPydanticObjectBase):
turns: int
spend: float
saved_spend: float
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
savings_estimated_baseline_models: Mapping[str, int] = Field(default_factory=dict)
classifier_cost: float
tier_turns: Mapping[str, int]
baseline_models: Mapping[str, int]
@property
def baseline_model(self) -> str | None:
"""The baseline most of this session's turns were priced against, or None when no turn recorded one.
"""The baseline most covered turns were priced against, or None when none were estimated.
A router reconfigured mid-session leaves turns priced against two baselines; the row keeps both
counts, and the label is the one that priced the most money-carrying turns rather than whatever the
router is configured with now.
"""
if not self.baseline_models:
if not self.savings_estimated_baseline_models:
return None
return max(self.baseline_models, key=lambda model: (self.baseline_models[model], model))
return max(
self.savings_estimated_baseline_models,
key=lambda model: (self.savings_estimated_baseline_models[model], model),
)

6
litellm/proxy/_types.py
View file

@ -13,6 +13,7 @@ from pydantic import (
ConfigDict,
Field,
Json,
JsonValue,
PositiveInt,
field_validator,
model_validator,
@ -123,6 +124,7 @@ class SupportedDBObjectType(str, enum.Enum):
MODEL_COST_MAP = "model_cost_map"
TOOLS = "tools"
CONFIG_OVERRIDES = "config_overrides"
WEBSEARCH_INTERCEPTION_SETTINGS = "websearch_interception_settings"
def __str__(self):
return str(self.value)
@ -3940,6 +3942,7 @@ class SpendLogsRouterMetadata(TypedDict):
class SpendLogsMetadata(TypedDict):
autorouter_baseline_observation: ReadOnly[str | None]
"""
Specific metadata k,v pairs logged to spendlogs for easier cost tracking
"""
@ -3980,7 +3983,8 @@ class SpendLogsMetadata(TypedDict):
original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks
cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
compression_savings: CompressionSavingsMetadata | None
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
autorouter_savings: ReadOnly[float | None]
autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None]
litellm_gateway_injected_cache: ReadOnly[str | None]
router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model
azure_spillover: ReadOnly[AzureSpillover | None] # None = Azure did not report spillover

60
litellm/proxy/auth/auth_checks.py
View file

@ -4012,6 +4012,64 @@ async def get_org_object(
return _org_obj
def _last_known_org_cache_key(org_id: str) -> str:
return f"org_id:{org_id}:with_budget:last_known"
async def _keep_last_known_org(
org: LiteLLM_OrganizationTable, org_id: str, user_api_key_cache: UserApiKeyCache
) -> None:
cache_key: Final = _last_known_org_cache_key(org_id)
held_locally: Final = await user_api_key_cache.async_get_cache(
key=cache_key, local_only=True, model_type=LiteLLM_OrganizationTable
)
if held_locally is not None:
return
await user_api_key_cache.async_set_cache(
key=cache_key,
value=org,
model_type=LiteLLM_OrganizationTable,
ttl=get_management_object_ttl(user_api_key_cache),
)
async def get_org_object_for_request(
org_id: str,
prisma_client: PrismaClient,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> LiteLLM_OrganizationTable | None:
try:
org: Final = await get_org_object(
org_id=org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
include_budget_table=True,
)
except OrganizationNotFoundError:
return None
except Exception as e: # noqa: BLE001 # only a DB outage may fail auth here, anything else degrades to no org limits
if not PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(e):
verbose_proxy_logger.debug("org lookup failed, continuing without org limits", exc_info=True)
return None
last_known_org: Final = await user_api_key_cache.async_get_cache(
key=_last_known_org_cache_key(org_id),
model_type=LiteLLM_OrganizationTable,
)
if last_known_org is not None:
return last_known_org
if PrismaDBExceptionHandler.should_allow_request_on_db_unavailable():
return None
raise
if org is None:
return None
await _keep_last_known_org(org, org_id, user_api_key_cache)
return org
async def _get_resources_from_access_groups(
access_group_ids: Sequence[str],
resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"],
@ -5680,7 +5738,7 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget):
if max_budget is None or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend

52
litellm/proxy/auth/user_api_key_auth.py
View file

@ -58,6 +58,7 @@ from litellm.proxy.auth.auth_checks import (
get_jwt_key_mapping_object,
get_key_end_user_budget_id,
get_object_permission,
get_org_object_for_request,
get_project_object,
get_team_membership,
get_team_object,
@ -2611,6 +2612,47 @@ def _token_can_vouch_for_team(valid_token: UserAPIKeyAuth, lookup_error: BaseExc
return PrismaDBExceptionHandler.should_allow_request_on_db_unavailable()
async def _inherit_org_identity(
user_api_key_auth_obj: UserAPIKeyAuth,
team_object: LiteLLM_TeamTableCachedObj | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
parent_otel_span: Span | None,
proxy_logging_obj: ProxyLogging | None,
) -> None:
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
already_populated: Final = any(
value is not None
for value in (
user_api_key_auth_obj.organization_alias,
user_api_key_auth_obj.organization_max_budget,
user_api_key_auth_obj.organization_tpm_limit,
user_api_key_auth_obj.organization_rpm_limit,
user_api_key_auth_obj.organization_metadata,
)
)
if user_api_key_auth_obj.org_id is None or already_populated or prisma_client is None:
return
org_object: Final = await get_org_object_for_request(
org_id=user_api_key_auth_obj.org_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if org_object is None:
return
user_api_key_auth_obj.organization_alias = org_object.organization_alias
user_api_key_auth_obj.organization_metadata = org_object.metadata
budget: Final = org_object.litellm_budget_table
if budget is None:
return
user_api_key_auth_obj.organization_max_budget = budget.max_budget
user_api_key_auth_obj.organization_tpm_limit = budget.tpm_limit
user_api_key_auth_obj.organization_rpm_limit = budget.rpm_limit
def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, object]) -> bool:
return master_key is None and not any(
general_settings.get(flag, False)
@ -2849,8 +2891,14 @@ async def _run_centralized_common_checks(
user_object=user_object,
)
if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None:
user_api_key_auth_obj.org_id = team_object.organization_id
await _inherit_org_identity(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
# common_checks identifies admin via user_object, not the token
# (non_proxy_admin_allowed_routes_check). JWT admin shortcut and

177
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -7,10 +7,12 @@
import asyncio
import os
from collections.abc import Mapping
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final, cast
from typing import TYPE_CHECKING, Any, Final, Literal, cast
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_proxy_logger
@ -18,6 +20,15 @@ from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE,
LiteLLMExecutedBatchRunner,
ManagedBatchStore,
batch_error,
executed_batch_runner_lost,
litellm_executed_provider_for,
resolve_litellm_executed_provider,
)
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
log_llm_api_exception,
@ -47,16 +58,87 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_model_id_from_unified_batch_id,
get_models_from_unified_file_id,
get_original_file_id,
is_litellm_executed_batch,
prepare_data_with_credentials,
update_batch_in_database,
validate_managed_id_requirement,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_metadata
from litellm.proxy.route_llm_request import raise_if_required_body_param_missing
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy, is_known_model
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.router import Router
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from prisma.models import LiteLLM_ManagedObjectTable
router: Final = APIRouter()
_METADATA_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object])
def _request_tags(data: Mapping[str, object]) -> tuple[str, ...] | None:
metadata: Final = data.get("litellm_metadata")
if metadata is None:
return None
return request_tags_from_metadata(_METADATA_ADAPTER.validate_python(metadata))
def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyLogging) -> LiteLLMExecutedBatchRunner:
from litellm.proxy.proxy_server import general_settings, prisma_client
managed_files: Final = proxy_logging_obj.get_proxy_hook("managed_files")
if prisma_client is None or not isinstance(managed_files, ManagedBatchStore):
raise batch_error(
400,
"LiteLLM-executed batches need a database: set DATABASE_URL so LiteLLM can keep the batch and its files",
)
return LiteLLMExecutedBatchRunner(
llm_router=llm_router,
prisma_client=prisma_client,
managed_files=managed_files,
batches=ManagedBatchRepository(prisma_client),
proxy_logging_obj=proxy_logging_obj,
general_settings=general_settings,
)
async def _batch_from_database(
batch_id: str,
unified_batch_id: str | Literal[False],
executed_batch: bool,
managed_files_obj: object,
prisma_client: PrismaClient | None,
llm_router: Router | None,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple["LiteLLM_ManagedObjectTable | None", LiteLLMBatch | None]:
row, batch = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
)
updated_at: Final[object] = getattr(row, "updated_at", None)
if not executed_batch or batch is None or llm_router is None or not isinstance(updated_at, datetime):
return row, batch
if not executed_batch_runner_lost(batch.status, updated_at):
return row, batch
runner: Final = _litellm_executed_batch_runner(llm_router, proxy_logging_obj)
return row, await runner.fail_abandoned(batch, user_api_key_dict)
async def _raise_when_input_file_must_be_managed(model: str, credentials: Mapping[str, object]) -> None:
if await litellm_executed_provider_for(credentials) is None:
return
raise batch_error(
400,
f"Batches for {model} run inside LiteLLM, so the input file must be a LiteLLM managed file: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
def _raise_not_found_when_openai_fallback_unservable(
@ -101,6 +183,24 @@ async def _resolve_managed_input_file_storage_url(input_file_id: str) -> "str |
return db_file.storage_url or None
async def _create_provider_batch_for_managed_file(
llm_router: Router,
create_batch_data: LiteLLMBatchCreateRequest,
input_file_id: str,
unified_file_id: str,
) -> LiteLLMBatch:
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
request: Final[LiteLLMBatchCreateRequest] = {
**create_batch_data,
"input_file_id": resolved_storage_url or input_file_id,
"disable_fallbacks": True,
}
response: Final = await llm_router.acreate_batch(**request)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
return response
@router.post(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -296,24 +396,35 @@ async def create_batch(
await authorize_model_for_key(model_id=model, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
_create_batch_data["model"] = model
resolved_storage_url: Final = await _resolve_managed_input_file_storage_url(input_file_id)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url
if llm_router is None:
raise HTTPException(
status_code=500,
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
)
_create_batch_data.update(disable_fallbacks=True) # pyright: ignore[reportCallIssue] # router flag
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id
executed_provider: Final = await resolve_litellm_executed_provider(
llm_router, model, user_api_key_dict.team_id
)
response = (
await _litellm_executed_batch_runner(llm_router, proxy_logging_obj).create(
create_request=_create_batch_data,
unified_input_file_id=input_file_id,
model=model,
provider=executed_provider,
user_api_key_dict=user_api_key_dict,
request_tags=_request_tags(_create_batch_data),
)
if executed_provider is not None
else await _create_provider_batch_for_managed_file(
llm_router, _create_batch_data, input_file_id, unified_file_id
)
)
else:
# Check if model specified via header/query/body param
model_param: Final = (
data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model")
_create_batch_data.get("model")
or request.query_params.get("model")
or request.headers.get("x-litellm-model")
)
# SCENARIO 2 & 3: Model from header/query OR custom_llm_provider fallback
@ -325,6 +436,7 @@ async def create_batch(
user_api_key_dict=user_api_key_dict,
operation_context="batch creation",
)
await _raise_when_input_file_must_be_managed(model_param, credentials)
prepare_data_with_credentials(
data=_create_batch_data,
@ -486,23 +598,26 @@ async def retrieve_batch(
managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files")
from litellm.proxy.proxy_server import prisma_client
db_batch_object, response = await get_batch_from_database(
executed_batch: Final = isinstance(unified_batch_id, str) and is_litellm_executed_batch(unified_batch_id)
db_batch_object, response = await _batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
executed_batch=executed_batch,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
llm_router=llm_router,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
)
if executed_batch and response is None:
raise batch_error(404, f"No batch found with id '{batch_id}'.")
# If batch is in a terminal state, return immediately.
# Include "complete" (DB-normalized form of "completed").
if response is not None and response.status in [
"completed",
"complete",
"failed",
"cancelled",
"expired",
]:
if response is not None and (
response.status in ("completed", "complete", "failed", "cancelled", "expired") or executed_batch
):
# Call hooks and return
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
@ -978,6 +1093,17 @@ async def cancel_batch(
proxy_config=proxy_config,
)
unified_model_id: Final = get_model_id_from_unified_batch_id(unified_batch_id) if unified_batch_id else None
if unified_model_id is not None:
resolved_unified_model: Final = (
llm_router.resolve_model_name_from_model_id(unified_model_id) if llm_router is not None else None
)
await authorize_model_for_key(
model_id=resolved_unified_model or unified_model_id,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
# SCENARIO 1: Batch ID is encoded with model info
if model_from_id is not None:
credentials: Final = await get_authorized_credentials_for_model(
@ -1009,6 +1135,12 @@ async def cancel_batch(
)
# SCENARIO 2: target_model_names based routing
elif unified_batch_id and is_litellm_executed_batch(unified_batch_id):
if llm_router is None:
raise batch_error(500, "LLM Router not initialized. Ensure models added to proxy.")
response = await _litellm_executed_batch_runner( # rebind-ok: each cancel path sets the route's response
llm_router, proxy_logging_obj
).cancel(batch_id, user_api_key_dict)
elif unified_batch_id:
if llm_router is None:
raise HTTPException(
@ -1022,11 +1154,6 @@ async def cancel_batch(
status_code=400,
detail={"error": "Invalid LiteLLM managed batch ID. Missing model_id."},
)
await authorize_model_for_key(
model_id=llm_router.resolve_model_name_from_model_id(model_id_from_batch) or model_id_from_batch,
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
)
data["model"] = model_id_from_batch
data["batch_id"] = get_batch_id_from_unified_batch_id(unified_batch_id)
response = await llm_router.acancel_batch(**data)

716
litellm/proxy/batches_endpoints/litellm_executed_batches.py Normal file
View file

@ -0,0 +1,716 @@
import asyncio
import json
import time
from collections.abc import Awaitable, Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from itertools import pairwise
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, runtime_checkable
import httpx
from openai.types.batch import Errors
from openai.types.batch_error import BatchError
from openai.types.batch_request_counts import BatchRequestCounts
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid as uuid_module
from litellm.constants import LITELLM_EXECUTED_BATCH_CONCURRENCY
from litellm.integrations.prometheus import PrometheusLogger
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_utils import is_request_body_safe
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
from litellm.proxy.openai_files_endpoints.common_utils import LITELLM_EXECUTED_BATCH_ID_PREFIX
from litellm.proxy.openai_files_endpoints.storage_backend_service import StorageBackendFileService
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.managed_batch_repository import ManagedBatchRepository
from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import LITELLM_EXECUTED_BATCH_PROVIDERS, ExtractedFileData, LiteLLMBatch, LlmProviders
if TYPE_CHECKING:
from prisma import types as prisma_types
from litellm.router import Router
BatchEndpoint: TypeAlias = Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions", "/v1/responses"]
BatchStatus: TypeAlias = Literal[
"in_progress", "finalizing", "completed", "failed", "cancelling", "cancelled", "expired"
]
TERMINAL_BATCH_STATUSES: Final[frozenset[str]] = frozenset({"completed", "failed", "cancelled", "expired"})
_STOP_STATUSES: Final[frozenset[str]] = TERMINAL_BATCH_STATUSES | frozenset({"cancelling"})
_BATCH_ENDPOINT_ADAPTER: Final[TypeAdapter[BatchEndpoint]] = TypeAdapter(BatchEndpoint)
_CANCEL_POLL_SECONDS: Final = 1.0
_HEARTBEAT_SECONDS: Final = 30.0
_STALE_AFTER_SECONDS: Final = 180.0
_FILES_API_PROBE_TIMEOUT_SECONDS: Final = 5.0
_COMPLETION_WINDOW_SECONDS: Final = 24 * 60 * 60
_RUNNER_LOST_MESSAGE: Final = "the proxy replica running this batch stopped before it finished; resubmit the batch"
_EXPIRED_MESSAGE: Final = "This request could not be executed before the completion window expired."
_ROUTER_METHODS: Final[Mapping[BatchEndpoint, str]] = MappingProxyType(
{
"/v1/chat/completions": "acompletion",
"/v1/completions": "atext_completion",
"/v1/embeddings": "aembedding",
"/v1/responses": "aresponses",
}
)
_CANCELLING_TRANSITIONS: Final[Mapping[BatchStatus, BatchStatus]] = MappingProxyType(
{"completed": "cancelled", "expired": "cancelled", "in_progress": "cancelling", "finalizing": "cancelling"}
)
LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE: Final = (
"upload it through POST /v1/files with purpose=batch and either the x-litellm-model header or the "
"target_model_names form field naming the model, so LiteLLM keeps the file and runs the batch itself"
)
_RUNNING_BATCHES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong references keep running batch tasks alive
_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({})
_NO_HEADERS: Final[Mapping[str, str]] = MappingProxyType({})
class _ErrorDetail(TypedDict):
message: ReadOnly[str]
type: ReadOnly[str]
param: ReadOnly[None]
code: ReadOnly[None]
class _ErrorBody(TypedDict):
error: ReadOnly[_ErrorDetail]
class _ResultResponse(TypedDict):
status_code: ReadOnly[int]
request_id: ReadOnly[str]
body: ReadOnly[Mapping[str, object]]
class _LineError(TypedDict):
code: ReadOnly[str]
message: ReadOnly[str]
class _ResultLine(TypedDict):
id: ReadOnly[str]
custom_id: ReadOnly[str]
response: ReadOnly[_ResultResponse | None]
error: ReadOnly[_LineError | None]
class BatchInputLine(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
custom_id: str
method: Literal["POST"]
url: str
body: Mapping[str, object]
@dataclass(frozen=True, slots=True)
class InvalidBatchInput:
line_number: int | None
reason: str
def describe(self) -> str:
return f"line {self.line_number}: {self.reason}" if self.line_number is not None else self.reason
@dataclass(frozen=True, slots=True)
class RowOutcome:
custom_id: str
status_code: int
body: Mapping[str, object]
succeeded: bool
@dataclass(frozen=True, slots=True)
class ExpiredRow:
custom_id: str
@dataclass(frozen=True, slots=True)
class _BatchRun:
unified_batch_id: str
llm_batch_id: str
model: str
endpoint: BatchEndpoint
lines: tuple[BatchInputLine, ...]
user_api_key_dict: UserAPIKeyAuth
request_tags: tuple[str, ...]
deadline: float
@runtime_checkable
class ManagedBatchStore(Protocol):
def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: ...
async def get_unified_file_id(
self, file_id: str, litellm_parent_otel_span: object | None = None
) -> LiteLLM_ManagedFileTable | None: ...
async def store_unified_object_id(
self,
unified_object_id: str,
file_object: LiteLLMBatch,
litellm_parent_otel_span: object | None,
model_object_id: str,
file_purpose: Literal["batch", "fine-tune", "response"],
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None = None,
persist_attribution: bool = False,
batch_processed: bool = False,
) -> None: ...
class _StorageBackendFactory(Protocol):
def __call__(self, backend_type: str, prisma_client: PrismaClient | None = None) -> BaseFileStorageBackend: ...
class _ResultFileUploader(Protocol):
def __call__(
self,
file_data: Mapping[str, object],
target_storage: str,
target_model_names: Sequence[str],
purpose: OpenAIFilesPurpose,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None = None,
) -> Awaitable[OpenAIFileObject]: ...
@runtime_checkable
class _RouterCall(Protocol):
def __call__(self, **params: object) -> Awaitable[object]: ... # kwargs-ok: the request body is passed as keywords
def litellm_executed_provider_of(credentials: Mapping[str, object]) -> str | None:
explicit_provider: Final = credentials.get("custom_llm_provider")
provider: Final = (
explicit_provider if isinstance(explicit_provider, str) else _provider_of(credentials.get("model"))
)
return provider if provider in LITELLM_EXECUTED_BATCH_PROVIDERS else None
class _HttpGetter(Protocol):
async def get(
self, url: str, *, headers: dict[str, str] | None = None, timeout: float | httpx.Timeout | None = None
) -> httpx.Response: ...
class FilesApiProbe(Protocol):
async def __call__(self, api_base: str, api_key: str | None) -> bool: ...
class BodyRejection(Protocol):
def __call__(self, body: Mapping[str, object], /) -> str | None: ...
async def upstream_lacks_files_api(api_base: str, api_key: str | None, http_client: _HttpGetter | None = None) -> bool:
client: Final = http_client or get_async_httpx_client(llm_provider=LlmProviders.HOSTED_VLLM)
try:
response: Final = await client.get(
f"{api_base.rstrip('/')}/files",
headers=(
{"Authorization": f"Bearer {api_key}"} # mutable-ok: AsyncHTTPHandler.get wants a plain dict
if api_key
else None
),
timeout=_FILES_API_PROBE_TIMEOUT_SECONDS,
)
except httpx.HTTPError:
return False
return response.status_code == httpx.codes.NOT_FOUND
def _upstream_of(credentials: Mapping[str, object], provider: str) -> tuple[str, str | None] | None:
model: Final = credentials.get("model")
api_base: Final = credentials.get("api_base")
api_key: Final = credentials.get("api_key")
if not isinstance(model, str):
return None
try:
_, _, resolved_api_key, resolved_api_base = litellm.get_llm_provider(
model=model,
custom_llm_provider=provider,
api_base=api_base if isinstance(api_base, str) else None,
api_key=api_key if isinstance(api_key, str) else None,
)
except Exception: # noqa: BLE001 # get_llm_provider raises on a model it cannot map, which means nothing to probe
return None
return None if resolved_api_base is None else (resolved_api_base, resolved_api_key)
async def litellm_executed_provider_for(
credentials: Mapping[str, object], lacks_files_api: FilesApiProbe = upstream_lacks_files_api
) -> str | None:
provider: Final = litellm_executed_provider_of(credentials)
if provider is None:
return None
upstream: Final = _upstream_of(credentials, provider)
if upstream is None:
return None
return provider if await lacks_files_api(*upstream) else None
async def resolve_litellm_executed_provider(
llm_router: "Router",
model: str,
team_id: str | None,
lacks_files_api: FilesApiProbe = upstream_lacks_files_api,
) -> str | None:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=model, team_id=team_id)
return None if credentials is None else await litellm_executed_provider_for(credentials, lacks_files_api)
def _provider_of(model: object) -> str | None:
if not isinstance(model, str):
return None
try:
return litellm.get_llm_provider(model=model)[1]
except Exception: # noqa: BLE001 # get_llm_provider raises on an unknown model, which means no provider
return None
def _validation_reason(error: ValidationError) -> str:
return "; ".join(
f"{'.'.join(str(part) for part in item['loc'])}: {item['msg']}" if item["loc"] else item["msg"]
for item in error.errors()
)
def _accept_every_body(_body: Mapping[str, object]) -> str | None:
return None
def _parse_line(
line_number: int, raw: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection
) -> BatchInputLine | InvalidBatchInput:
try:
line: Final = BatchInputLine.model_validate_json(raw)
except ValidationError as e:
return InvalidBatchInput(line_number, _validation_reason(e))
if line.url != endpoint:
return InvalidBatchInput(line_number, f"url {line.url!r} does not match the batch endpoint {endpoint!r}")
if line.body.get("stream"):
return InvalidBatchInput(line_number, "streaming requests are not supported in a batch")
rejection: Final = reject_body(line.body)
if rejection is not None:
return InvalidBatchInput(line_number, rejection)
return line
def parse_batch_input(
content: bytes, endpoint: BatchEndpoint, reject_body: BodyRejection = _accept_every_body
) -> tuple[BatchInputLine, ...] | InvalidBatchInput:
raw_lines: Final = tuple((number, raw) for number, raw in enumerate(content.splitlines(), start=1) if raw.strip())
if not raw_lines:
return InvalidBatchInput(None, "the input file has no requests")
parsed: Final = tuple(_parse_line(number, raw, endpoint, reject_body) for number, raw in raw_lines)
first_invalid: Final = next((item for item in parsed if isinstance(item, InvalidBatchInput)), None)
if first_invalid is not None:
return first_invalid
lines: Final = tuple(item for item in parsed if isinstance(item, BatchInputLine))
custom_ids: Final = sorted(line.custom_id for line in lines)
duplicate: Final = next((first for first, second in pairwise(custom_ids) if first == second), None)
if duplicate is not None:
return InvalidBatchInput(None, f"custom_id {duplicate!r} is used more than once")
return lines
def batch_error(status_code: int, message: str) -> ProxyException:
error_type: Final = "invalid_request_error" if status_code < 500 else ProxyErrorTypes.internal_server_error.value
return ProxyException(message=message, type=error_type, param=None, code=status_code)
def _validate_endpoint(endpoint: object) -> BatchEndpoint:
try:
return _BATCH_ENDPOINT_ADAPTER.validate_python(endpoint)
except ValidationError:
raise batch_error(400, f"endpoint {endpoint!r} is not supported for a LiteLLM-executed batch")
def _status_code_of(error: Exception) -> int:
status_code: Final[object] = getattr(error, "status_code", None)
return status_code if isinstance(status_code, int) else 500
def _error_body(error: Exception) -> _ErrorBody:
body: Final[_ErrorBody] = {
"error": {"message": str(error), "type": type(error).__name__, "param": None, "code": None}
}
return body
def _line_response(outcome: RowOutcome | ExpiredRow) -> _ResultResponse | None:
if isinstance(outcome, ExpiredRow):
return None
response: Final[_ResultResponse] = {
"status_code": outcome.status_code,
"request_id": f"req_{uuid_module.uuid4().hex[:24]}",
"body": outcome.body,
}
return response
def _line_error(outcome: RowOutcome | ExpiredRow) -> _LineError | None:
if isinstance(outcome, RowOutcome):
return None
error: Final[_LineError] = {"code": "batch_expired", "message": _EXPIRED_MESSAGE}
return error
def _result_line(outcome: RowOutcome | ExpiredRow) -> _ResultLine:
line: Final[_ResultLine] = {
"id": f"batch_req_{uuid_module.uuid4().hex[:24]}",
"custom_id": outcome.custom_id,
"response": _line_response(outcome),
"error": _line_error(outcome),
}
return line
def _dump(response: object) -> Mapping[str, object]:
if isinstance(response, BaseModel):
return response.model_dump(mode="json")
raise TypeError(f"Batch rows must return a single response object, got {type(response).__name__}")
def _resolve_transition(current_status: str, requested: BatchStatus) -> BatchStatus:
if current_status != "cancelling":
return requested
return _CANCELLING_TRANSITIONS.get(requested, requested)
def executed_batch_runner_lost(status: str, updated_at: datetime) -> bool:
if status in TERMINAL_BATCH_STATUSES:
return False
return (datetime.now(timezone.utc) - updated_at).total_seconds() > _STALE_AFTER_SECONDS
class _StopWatch:
def __init__(self, load_status: Callable[[], Awaitable[str | None]], interval_seconds: float) -> None:
self._load_status = load_status
self._interval_seconds = interval_seconds
self._checked_at = float("-inf")
self._stopped = False
async def stopped(self) -> bool:
if self._stopped:
return True
now: Final = time.monotonic()
if now - self._checked_at < self._interval_seconds:
return False
self._checked_at = now
self._stopped = await self._load_status() in _STOP_STATUSES
return self._stopped
class LiteLLMExecutedBatchRunner:
def __init__(
self,
llm_router: "Router",
prisma_client: PrismaClient,
managed_files: ManagedBatchStore,
batches: ManagedBatchRepository,
proxy_logging_obj: ProxyLogging,
general_settings: Mapping[str, object],
concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY,
heartbeat_seconds: float = _HEARTBEAT_SECONDS,
completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS,
storage_backend_factory: _StorageBackendFactory = get_storage_backend,
upload_result_file: _ResultFileUploader = StorageBackendFileService.upload_file_to_storage_backend,
) -> None:
self.llm_router = llm_router
self.prisma_client = prisma_client
self.managed_files = managed_files
self.batches = batches
self.proxy_logging_obj = proxy_logging_obj
self.general_settings = general_settings
self.concurrency = concurrency
self.heartbeat_seconds = heartbeat_seconds
self.completion_window_seconds = completion_window_seconds
self.storage_backend_factory = storage_backend_factory
self.upload_result_file = upload_result_file
async def create(
self,
create_request: LiteLLMBatchCreateRequest,
unified_input_file_id: str,
model: str,
provider: str,
user_api_key_dict: UserAPIKeyAuth,
request_tags: Sequence[str] | None,
) -> LiteLLMBatch:
endpoint: Final = _validate_endpoint(create_request.get("endpoint"))
content: Final = await self._download_input(unified_input_file_id, user_api_key_dict)
parsed: Final = parse_batch_input(content, endpoint, self._body_rejection(model))
if isinstance(parsed, InvalidBatchInput):
raise batch_error(400, f"Invalid batch input file: {parsed.describe()}")
llm_batch_id: Final = f"{LITELLM_EXECUTED_BATCH_ID_PREFIX}{uuid_module.uuid4().hex}"
model_id: Final = next(iter(self.llm_router.get_model_ids(model_name=model)), model)
unified_batch_id: Final = self.managed_files.get_unified_batch_id(batch_id=llm_batch_id, model_id=model_id)
now: Final = time.time()
created_at: Final = int(now)
batch: Final = LiteLLMBatch(
id=unified_batch_id,
object="batch",
endpoint=endpoint,
input_file_id=unified_input_file_id,
completion_window="24h",
status="validating",
created_at=created_at,
expires_at=created_at + int(self.completion_window_seconds),
metadata=create_request.get("metadata"),
model=model,
request_counts=BatchRequestCounts(completed=0, failed=0, total=len(parsed)),
)
await self.managed_files.store_unified_object_id(
unified_object_id=unified_batch_id,
file_object=batch,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
model_object_id=llm_batch_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=request_tags,
persist_attribution=True,
batch_processed=True,
)
_record_batch_created(model, provider, user_api_key_dict)
run: Final = _BatchRun(
unified_batch_id=unified_batch_id,
llm_batch_id=llm_batch_id,
model=model,
endpoint=endpoint,
lines=parsed,
user_api_key_dict=user_api_key_dict,
request_tags=tuple(request_tags or ()),
deadline=now + self.completion_window_seconds,
)
task: Final = asyncio.create_task(self._run(run))
_RUNNING_BATCHES.add(task)
task.add_done_callback(_RUNNING_BATCHES.discard)
return batch
async def cancel(self, unified_batch_id: str, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
current: Final = await self.batches.load_batch(unified_batch_id)
if current is None:
raise batch_error(404, f"Batch {unified_batch_id} not found")
if current.status in TERMINAL_BATCH_STATUSES:
raise batch_error(400, f"Cannot cancel a batch with status '{current.status}'")
if current.status == "cancelling":
return current
cancelling: Final = current.model_copy(
update=MappingProxyType({"status": "cancelling", "cancelling_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(cancelling, unchanged, user_api_key_dict.user_id):
return cancelling
return await self.cancel(unified_batch_id, user_api_key_dict)
async def fail_abandoned(self, batch: LiteLLMBatch, user_api_key_dict: UserAPIKeyAuth) -> LiteLLMBatch:
error: Final = BatchError(message=_RUNNER_LOST_MESSAGE, code="runner_lost")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
failed: Final = batch.model_copy(
update=MappingProxyType({"status": "failed", "failed_at": int(time.time()), "errors": errors})
)
untouched: Final[prisma_types.DateTimeFilter] = {
"lt": datetime.now(timezone.utc) - timedelta(seconds=_STALE_AFTER_SECONDS)
}
still_abandoned: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {
"status": batch.status,
"updated_at": untouched,
}
if await self.batches.compare_and_set(failed, still_abandoned, user_api_key_dict.user_id):
return failed
return await self.batches.load_batch(batch.id) or batch
def _body_rejection(self, model: str) -> BodyRejection:
def reject(body: Mapping[str, object]) -> str | None:
try:
is_request_body_safe(
request_body=dict(body), # mutable-ok: is_request_body_safe takes a dict
general_settings=dict(self.general_settings), # mutable-ok: is_request_body_safe takes a dict
llm_router=self.llm_router,
model=model,
)
except ValueError as e:
return str(e)
return None
return reject
async def _download_input(self, unified_input_file_id: str, user_api_key_dict: UserAPIKeyAuth) -> bytes:
stored: Final = await self.managed_files.get_unified_file_id(
unified_input_file_id, litellm_parent_otel_span=user_api_key_dict.parent_otel_span
)
if stored is None or not stored.storage_backend or not stored.storage_url:
raise batch_error(
400,
f"LiteLLM does not hold the content of input file {unified_input_file_id}: "
f"{LITELLM_EXECUTED_BATCH_UPLOAD_GUIDANCE}",
)
try:
backend: Final = self.storage_backend_factory(stored.storage_backend, prisma_client=self.prisma_client)
return await backend.download_file(stored.storage_url)
except ValueError as e:
raise batch_error(400, str(e))
async def _run(self, run: _BatchRun) -> None:
heartbeat: Final = asyncio.create_task(self._heartbeat(run))
try:
await self._execute(run)
except Exception as e: # noqa: BLE001 # whatever fails, the batch must end up marked failed
verbose_proxy_logger.exception("LiteLLM-executed batch %s failed: %s", run.unified_batch_id, e)
error: Final = BatchError(message=str(e), code="internal_error")
errors: Final = Errors(data=[error], object="list") # mutable-ok: Errors.data is typed as a list
try:
await self._advance(run, "failed", MappingProxyType({"errors": errors}))
except Exception as advance_error: # noqa: BLE001 # a failed status write is logged, never raised
verbose_proxy_logger.exception(
"LiteLLM-executed batch %s could not be marked failed: %s", run.unified_batch_id, advance_error
)
finally:
heartbeat.cancel()
async def _heartbeat(self, run: _BatchRun) -> None:
while True:
await asyncio.sleep(self.heartbeat_seconds)
try:
await self._touch(run)
except Exception as e: # noqa: BLE001 # a missed beat is logged and the next one retries
verbose_proxy_logger.warning("LiteLLM-executed batch %s heartbeat failed: %s", run.unified_batch_id, e)
async def _touch(self, run: _BatchRun) -> None:
await self.batches.touch(run.unified_batch_id, run.user_api_key_dict.user_id)
async def _execute(self, run: _BatchRun) -> None:
await self._advance(run, "in_progress")
watch: Final = _StopWatch(lambda: self.batches.load_status(run.unified_batch_id), _CANCEL_POLL_SECONDS)
semaphore: Final = asyncio.Semaphore(self.concurrency)
results: Final = await asyncio.gather(*(self._run_row(run, line, watch, semaphore) for line in run.lines))
outcomes: Final = tuple(outcome for outcome in results if outcome is not None)
if await self._advance(run, "finalizing") is None:
return
succeeded: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, RowOutcome) and outcome.succeeded
)
failed: Final = tuple(
outcome for outcome in outcomes if isinstance(outcome, ExpiredRow) or not outcome.succeeded
)
output_file_id: Final = await self._upload_results(run, "output", succeeded)
error_file_id: Final = await self._upload_results(run, "error", failed)
request_counts: Final = BatchRequestCounts(completed=len(succeeded), failed=len(failed), total=len(run.lines))
final_status: Final[BatchStatus] = (
"expired" if any(isinstance(outcome, ExpiredRow) for outcome in outcomes) else "completed"
)
await self._advance(
run,
final_status,
MappingProxyType(
{"output_file_id": output_file_id, "error_file_id": error_file_id, "request_counts": request_counts}
),
)
async def _run_row(
self, run: _BatchRun, line: BatchInputLine, watch: _StopWatch, semaphore: asyncio.Semaphore
) -> RowOutcome | ExpiredRow | None:
async with semaphore:
if await watch.stopped():
return None
remaining: Final = run.deadline - time.time()
if remaining <= 0:
return ExpiredRow(custom_id=line.custom_id)
try:
return await asyncio.wait_for(self._row_outcome(run, line), timeout=remaining)
except asyncio.TimeoutError:
return ExpiredRow(custom_id=line.custom_id)
async def _row_outcome(self, run: _BatchRun, line: BatchInputLine) -> RowOutcome:
try:
body: Final = await self._dispatch(run, line)
except Exception as e: # noqa: BLE001 # a provider error becomes the row's error line, never a crashed batch
return RowOutcome(
custom_id=line.custom_id, status_code=_status_code_of(e), body=_error_body(e), succeeded=False
)
return RowOutcome(custom_id=line.custom_id, status_code=200, body=body, succeeded=True)
async def _dispatch(self, run: _BatchRun, line: BatchInputLine) -> Mapping[str, object]:
params: Final = MappingProxyType(
{**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True}
)
return _dump(await self._router_call(run.endpoint)(**params))
def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall:
method: Final[object] = getattr(self.llm_router, _ROUTER_METHODS[endpoint], None)
if not isinstance(method, _RouterCall):
raise TypeError(f"the router has no callable for {endpoint}")
return method
def _row_metadata(self, run: _BatchRun) -> dict[str, object]: # mutable-ok: router updates metadata in place
return { # mutable-ok: the router updates request metadata in place
**LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(run.user_api_key_dict),
"user_api_key": run.user_api_key_dict.api_key,
"user_api_end_user_max_budget": run.user_api_key_dict.end_user_max_budget,
"tags": list(run.request_tags), # mutable-ok: litellm types request tags as a list
"batch_id": run.unified_batch_id,
}
async def _upload_results(
self, run: _BatchRun, kind: Literal["output", "error"], outcomes: Sequence[RowOutcome | ExpiredRow]
) -> str | None:
if not outcomes:
return None
content: Final = "".join(f"{json.dumps(_result_line(outcome))}\n" for outcome in outcomes).encode()
file_data: Final[ExtractedFileData] = {
"filename": f"{run.llm_batch_id}_{kind}.jsonl",
"content": content,
"content_type": "application/jsonl",
"headers": _NO_HEADERS,
}
file_object: Final = await self.upload_result_file(
file_data=file_data,
target_storage=LITELLM_DB_STORAGE_BACKEND_NAME,
target_model_names=(run.model,),
purpose="batch_output",
proxy_logging_obj=self.proxy_logging_obj,
user_api_key_dict=run.user_api_key_dict,
prisma_client=self.prisma_client,
)
return file_object.id
async def _advance(
self, run: _BatchRun, requested: BatchStatus, fields: Mapping[str, object] = _NO_FIELDS
) -> BatchStatus | None:
current: Final = await self.batches.load_batch(run.unified_batch_id)
if current is None:
raise RuntimeError(f"Batch {run.unified_batch_id} is no longer stored")
if current.status in TERMINAL_BATCH_STATUSES:
return None
status: Final = _resolve_transition(current.status, requested)
updated: Final = current.model_copy(
update=MappingProxyType({**fields, "status": status, f"{status}_at": int(time.time())})
)
unchanged: Final[prisma_types.LiteLLM_ManagedObjectTableWhereInput] = {"status": current.status}
if await self.batches.compare_and_set(updated, unchanged, run.user_api_key_dict.user_id):
return status
return await self._advance(run, requested, fields)
def _record_batch_created(model: str, provider: str, user_api_key_dict: UserAPIKeyAuth) -> None:
prometheus_logger: Final = PrometheusLogger.get_instance()
if prometheus_logger is None:
return
prometheus_logger.record_managed_batch_created(
model=model,
api_provider=provider,
user=user_api_key_dict.user_id or "",
user_email=user_api_key_dict.user_email or "",
api_key_alias=user_api_key_dict.key_alias or "",
)

67
litellm/proxy/client/cli/commands/statusline_script.py
View file

@ -32,6 +32,7 @@ import unicodedata
import urllib.error
import urllib.request
from collections.abc import Callable, Mapping
from math import isfinite
from pathlib import Path
from types import MappingProxyType
from typing import IO, Final, NamedTuple, Protocol
@ -43,6 +44,7 @@ FETCH_TIMEOUT_SECONDS: Final = 3
BAR_WIDTH: Final = 24
BAR_FULL: Final = "\u2588"
BAR_EMPTY: Final = "\u2591"
SEPARATOR: Final = " \u00b7 "
TRANSCRIPT_SCAN_LIMIT_BYTES: Final = 4 * 1024 * 1024
CLAUDE_BASE_URL_ENV_KEYS: Final = ("ANTHROPIC_BASE_URL",)
CLAUDE_API_KEY_ENV_KEYS: Final = ("ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_KEY")
@ -63,8 +65,11 @@ class Session(NamedTuple):
router_name: str
last_model: str
spend: float
baseline_spend: float
baseline_spend: float | None
baseline_model: str | None
turns: int | None = None
savings_estimated_turns: int | None = None
savings_estimated_actual_spend: float | None = None
class Credentials(NamedTuple):
@ -205,17 +210,38 @@ def _session_from_payload(payload: Mapping[str, object]) -> Session | None:
router_name: Final = printable(payload.get("router_name"))
last_model: Final = printable(payload.get("last_model"))
spend: Final = payload.get("spend")
baseline_spend: Final = payload.get("baseline_spend")
baseline_spend: Final = payload.get("savings_estimated_baseline_spend", payload.get("baseline_spend"))
turns: Final = payload.get("turns")
estimated_turns: Final = payload.get("savings_estimated_turns")
estimated_actual: Final = payload.get("savings_estimated_actual_spend")
if not router_name or not last_model:
return None
if not isinstance(spend, (int, float)) or not isinstance(baseline_spend, (int, float)):
if not isinstance(spend, (int, float)) or isinstance(spend, bool) or not isfinite(spend):
return None
if baseline_spend is not None and (
not isinstance(baseline_spend, (int, float)) or isinstance(baseline_spend, bool) or not isfinite(baseline_spend)
):
return None
return Session(
router_name=router_name,
last_model=last_model,
spend=float(spend),
baseline_spend=float(baseline_spend),
baseline_spend=float(baseline_spend) if baseline_spend is not None else None,
baseline_model=printable(payload.get("baseline_model")) or None,
turns=turns if isinstance(turns, int) and not isinstance(turns, bool) and turns >= 0 else None,
savings_estimated_turns=(
estimated_turns
if isinstance(estimated_turns, int) and not isinstance(estimated_turns, bool) and estimated_turns >= 0
else (0 if estimated_turns is not None else None)
),
savings_estimated_actual_spend=(
float(estimated_actual)
if isinstance(estimated_actual, (int, float))
and not isinstance(estimated_actual, bool)
and isfinite(estimated_actual)
and estimated_actual >= 0
else None
),
)
@ -314,15 +340,36 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo
return f"{code}{text}{RESET}" if use_color else text
routed: Final = paint(BOLD, f"Routed to: {model}")
if session is None or session.baseline_model is None or session.baseline_spend <= 0:
if session is None:
return routed
if session.savings_estimated_turns == 0 or session.baseline_spend is None:
return f"{routed}{SEPARATOR}Savings unavailable"
if session.baseline_model is None or session.baseline_spend <= 0:
return routed
if session.savings_estimated_turns is not None and (
session.savings_estimated_actual_spend is None
or session.turns is None
or session.savings_estimated_turns > session.turns
):
return f"{routed}{SEPARATOR}Savings unavailable"
compared_spend: Final = (
session.savings_estimated_actual_spend
if session.savings_estimated_turns is not None and session.savings_estimated_actual_spend is not None
else session.spend
)
coverage: Final = (
f"{SEPARATOR}{session.savings_estimated_turns} of {session.turns} turns estimated"
if session.savings_estimated_turns is not None
else ""
)
reference: Final = baseline_label(session.baseline_model, config_dir)
pct: Final = (session.baseline_spend - session.spend) / session.baseline_spend * 100
delta: Final = paint(LITELLM_COLOR, f"{'-' if pct >= 0 else '+'}{abs(round(pct))}% vs {reference}")
peak: Final = max(session.spend, session.baseline_spend)
pct: Final = round((session.baseline_spend - compared_spend) / session.baseline_spend * 100)
sign: Final = "-" if pct > 0 else "+" if pct < 0 else ""
delta: Final = paint(LITELLM_COLOR, f"{sign}{abs(pct)}% vs {reference}")
peak: Final = max(compared_spend, session.baseline_spend)
label_width: Final = max(_display_width(session.router_name), _display_width(reference))
rows: Final = (
(session.router_name, session.spend, LITELLM_COLOR),
(session.router_name, compared_spend, LITELLM_COLOR),
(reference, session.baseline_spend, BASELINE_COLOR),
)
lines: Final = (
@ -331,7 +378,7 @@ def render(model: str, session: Session | None, config_dir: Path, use_color: boo
f"{paint(DIM, f'${amount:.2f}')}"
for label, amount, color in rows
)
return "\n".join((f"{routed} {delta}", *lines))
return "\n".join((f"{routed} {delta}{coverage}", *lines))
def color_enabled(env: Mapping[str, str]) -> bool:

8
litellm/proxy/common_request_processing.py
View file

@ -3749,6 +3749,14 @@ class ProxyBaseLLMRequestProcessing:
"async_streaming_data_generator: error closing response stream: %s",
e,
)
logging_obj: Final = request_data.get("litellm_logging_obj")
if (
not stream_completed
and isinstance(logging_obj, LiteLLMLoggingObj)
and logging_obj.baseline_cache_context is not None
and logging_obj.model_call_details.get("prompt_cache_response_complete") is not True
):
await logging_obj.invalidate_baseline_cache_estimate("incomplete_response", completed=True)
@staticmethod
async def async_streaming_data_generator(

21
litellm/proxy/config_resolvers/settings_store.py
View file

@ -60,9 +60,7 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def rejected_writes(self, incoming: Mapping[str, JsonValue]) -> tuple[str, ...]:
return tuple(
sorted(
key for key, value in incoming.items() if self.owned_by_config(key) and value != self._yaml_values[key]
)
sorted(key for key, value in incoming.items() if self.owned_by_config(key) and value != self.get(key))
)
def shadowed_db_keys(self) -> tuple[str, ...]:
@ -74,8 +72,13 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def apply_db_row(self, row: DbRow, db_row: Mapping[str, JsonValue]) -> None:
previous_row: Final = self._database_rows.get(row, _EMPTY_VALUES)
changed: Final = frozenset(
key
for key in (*previous_row, *db_row)
if previous_row.get(key, ABSENT) != db_row.get(key, ABSENT) # pyright: ignore[reportUnknownArgumentType] # JsonValue vs Absent compare
)
self._database_rows = MappingProxyType({**self._database_rows, row: MappingProxyType(dict(db_row))})
self._clear_runtime_keys(frozenset((*previous_row, *db_row)))
self._clear_runtime_keys(changed)
def resolved(self) -> Mapping[str, JsonValue]:
return MappingProxyType(dict(self))
@ -130,6 +133,9 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def __len__(self) -> int:
return sum(1 for _ in self)
def __bool__(self) -> bool:
return any(True for _ in self)
def _clear_runtime(self) -> None:
self._runtime_values = _EMPTY_VALUES
self._deleted_runtime_keys = frozenset()
@ -160,7 +166,12 @@ class SettingsStore(MutableMapping[str, JsonValue]):
def _db_value_is_shadowed(self, key: str) -> bool:
db_value: Final = self._db_value(key)
return not isinstance(db_value, Absent) and db_value is not None and db_value != self.get(key)
return (
not isinstance(db_value, Absent)
and db_value is not None
and db_value != self.get(key)
and db_value != self.config_value(key)
)
def _resolution_for(self, key: str) -> Resolved:
yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT)

43
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -26,6 +26,7 @@ from typing import TYPE_CHECKING, Final, NamedTuple
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
from litellm.proxy.db.create_views import SupportsExecuteRaw
if TYPE_CHECKING:
from litellm.proxy._types import SpendLogsPayload
@ -75,6 +76,9 @@ SELECT
COALESCE(SUM(total_tokens), 0)::bigint AS total_tokens,
COALESCE(SUM(spend), 0)::float8 AS spend,
COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend,
COALESCE(SUM(savings_estimated_turns), 0)::int AS savings_estimated_turns,
COALESCE(SUM(savings_estimated_actual_spend), 0)::float8 AS savings_estimated_actual_spend,
COALESCE(SUM(savings_estimated_saved_spend), 0)::float8 AS savings_estimated_saved_spend,
COALESCE(SUM(classifier_cost), 0)::float8 AS classifier_cost,
COALESCE(SUM(classifier_cost_recorded_turns), 0)::int AS classifier_cost_recorded_turns,
COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds
@ -104,6 +108,9 @@ class AutoRouterTurnTransaction:
cache_touched: bool
tier: str | None = None
baseline_model: str | None = None
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
class TurnCacheFacts(NamedTuple):
@ -215,13 +222,18 @@ def build_autorouter_turn_transaction(
turn_at: Final = _turn_time_utc(str(payload.get("startTime") or ""))
if turn_at is None:
return None
from litellm.proxy.spend_tracking.savings import classifier_cost_from_decision
from litellm.proxy.spend_tracking.savings import (
classifier_cost_from_decision,
recorded_estimated_autorouter_savings,
)
usage_object_raw: Final = metadata.get("usage_object")
cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None)
tier_raw: Final = routing_decision.get("tier")
baseline_raw: Final = routing_decision.get("savings_baseline_model")
classifier_cost: Final = classifier_cost_from_decision(routing_decision)
actual_spend: Final = float(payload.get("spend") or 0.0) + (classifier_cost or 0.0)
estimated_savings: Final = recorded_estimated_autorouter_savings(metadata)
return AutoRouterTurnTransaction(
api_key=api_key,
session_id=bounded_session_id(session_id),
@ -232,13 +244,16 @@ def build_autorouter_turn_transaction(
model=model,
turn_at=turn_at,
total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0),
spend=float(payload.get("spend") or 0.0) + (classifier_cost or 0.0),
spend=actual_spend,
saved_spend=saved_spend,
classifier_cost=classifier_cost or 0.0,
covered=cache.covered,
cache_hit=cache.read_tokens > 0,
cache_ttl_seconds=cache.write_ttl_seconds,
cache_touched=cache.touched,
savings_estimated_turns=int(estimated_savings is not None),
savings_estimated_actual_spend=actual_spend if estimated_savings is not None else 0.0,
savings_estimated_saved_spend=estimated_savings if estimated_savings is not None else 0.0,
)
@ -263,6 +278,10 @@ _BASELINE: Final = f"{_p('baseline_model')}::text"
_BASELINE_DELTA: Final = (
f"(CASE WHEN {_BASELINE} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_BASELINE}, 1) END)"
)
_ESTIMATED_BASELINE: Final = f"{_p('savings_estimated_turns')}::int = 1 AND {_BASELINE} IS NOT NULL"
_ESTIMATED_BASELINE_DELTA: Final = (
f"(CASE WHEN {_ESTIMATED_BASELINE} THEN jsonb_build_object({_BASELINE}, 1) ELSE '{{}}'::jsonb END)"
)
_IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at"
_SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}"
@ -281,7 +300,8 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t (
same_model_turns, same_model_hits, first_visit_turns, first_visit_hits,
return_turns, return_hits, return_expired_misses, return_within_ttl_misses,
ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns,
baseline_models
baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend,
savings_estimated_baseline_models
)
VALUES (
{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp,
@ -292,13 +312,18 @@ VALUES (
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END),
(CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END),
{_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8,
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA}
{_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA},
{_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8,
{_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA}
)
ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
turns = t.turns + 1,
total_tokens = t.total_tokens + EXCLUDED.total_tokens,
spend = t.spend + EXCLUDED.spend,
saved_spend = t.saved_spend + EXCLUDED.saved_spend,
savings_estimated_turns = t.savings_estimated_turns + EXCLUDED.savings_estimated_turns,
savings_estimated_actual_spend = t.savings_estimated_actual_spend + EXCLUDED.savings_estimated_actual_spend,
savings_estimated_saved_spend = t.savings_estimated_saved_spend + EXCLUDED.savings_estimated_saved_spend,
classifier_cost = t.classifier_cost + EXCLUDED.classifier_cost,
classifier_cost_recorded_turns = t.classifier_cost_recorded_turns + 1,
covered_turns = t.covered_turns + EXCLUDED.covered_turns,
@ -331,6 +356,10 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET
baseline_models = (CASE WHEN {_BASELINE} IS NOT NULL
THEN t.baseline_models || jsonb_build_object({_BASELINE}, COALESCE((t.baseline_models ->> {_BASELINE})::int, 0) + 1)
ELSE t.baseline_models END),
savings_estimated_baseline_models = (CASE WHEN {_ESTIMATED_BASELINE}
THEN t.savings_estimated_baseline_models || jsonb_build_object(
{_BASELINE}, COALESCE((t.savings_estimated_baseline_models ->> {_BASELINE})::int, 0) + 1)
ELSE t.savings_estimated_baseline_models END),
first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at),
last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at)
"""
@ -348,6 +377,10 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float
return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS)
async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None:
await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction))
async def _upsert_turn_with_retry(
prisma_client: PrismaClient,
transaction: AutoRouterTurnTransaction,
@ -355,7 +388,7 @@ async def _upsert_turn_with_retry(
) -> None:
for attempt in range(n_retry_times + 1):
try:
await prisma_client.db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction))
await write_autorouter_turn(prisma_client.db, transaction)
except DB_RETRY_SAFE_ERROR_TYPES:
if attempt >= n_retry_times:
raise

640
litellm/proxy/db/baseline_accounting.py Normal file
View file

@ -0,0 +1,640 @@
from __future__ import annotations
import asyncio
import json
from collections.abc import AsyncIterator, Callable, Sequence
from datetime import datetime, timedelta
from functools import reduce
from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, model_validator
from typing_extensions import Self
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.autorouter_session_rollup import (
AutoRouterTurnTransaction,
write_autorouter_turn,
)
from litellm.proxy.db.create_views import SupportsRawQueries
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
DailySpendEntity,
SpendRow,
build_bulk_upsert,
merge_by_conflict_key,
)
from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper
from litellm.proxy.spend_tracking.baseline_accounting import (
BaselineEstimate,
BaselineHistory,
BaselineObservation,
advance_baseline_history,
)
from litellm.proxy.spend_tracking.savings import BaselineCosts, BaselineCostSnapshot, price_baseline_comparison
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class DailyBaselineTarget(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
entity: DailySpendEntity
entity_id: str | None
class DailyBaselineAttribution(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
date: str
api_key: str
model: str | None = None
custom_llm_provider: str | None = None
model_group: str | None = None
endpoint: str | None = None
mcp_namespaced_tool_name: str | None = None
targets: tuple[DailyBaselineTarget, ...] = ()
def adjustment(self, target: DailyBaselineTarget, savings_delta: float, request_id: str) -> SpendRow:
table: Final = DAILY_SPEND_TABLES[target.entity]
return MappingProxyType(
{
"date": self.date,
"api_key": self.api_key,
"model": self.model,
"custom_llm_provider": self.custom_llm_provider,
"model_group": self.model_group,
"endpoint": self.endpoint,
"mcp_namespaced_tool_name": self.mcp_namespaced_tool_name,
table.entity_id_column: target.entity_id,
"request_id": request_id,
"autorouter_savings_spend": savings_delta,
}
)
class BaselineAccountingRecord(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
scope: str = Field(pattern=r"^autorouter-baseline:v3:[a-f0-9]{64}$")
api_key: str = Field(min_length=1)
session_id: str = Field(min_length=1, max_length=256)
router_name: str = Field(min_length=1)
baseline_model: str = Field(min_length=1)
observation: BaselineObservation
pricing: BaselineCostSnapshot
turn: AutoRouterTurnTransaction | None
daily: DailyBaselineAttribution | None
@model_validator(mode="after")
def consistent_turn(self) -> Self:
turn: Final = self.turn
if turn is not None and (
(turn.api_key, turn.session_id, turn.router_name, turn.baseline_model)
!= (self.api_key, self.session_id, self.router_name, self.baseline_model)
or turn.spend != self.pricing.actual_spend + self.pricing.classifier_cost
or any(
(
turn.saved_spend,
turn.savings_estimated_turns,
turn.savings_estimated_actual_spend,
turn.savings_estimated_saved_spend,
)
)
):
raise ValueError("Baseline observation must own an unestimated turn with matching scope and actual cost")
return self
class BaselinePublication(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
version: Literal[3] = 3
comparison_id: str
comparison_started_at: float
status: Literal["estimated", "unknown"]
reason: str
provenance: Literal["observed_identical", "modeled"] | None = None
actual_spend: float | None = None
baseline_spend: float | None = None
input_tokens: int | None = None
cache_read_input_tokens: int | None = None
cache_creation_5m_input_tokens: int | None = None
cache_creation_1h_input_tokens: int | None = None
@property
def costs(self) -> BaselineCosts | None:
if self.status != "estimated" or self.actual_spend is None or self.baseline_spend is None:
return None
return BaselineCosts(self.actual_spend, self.baseline_spend)
def baseline_publication(
record: BaselineAccountingRecord, estimate: BaselineEstimate, first_at: float
) -> BaselinePublication:
costs: Final = price_baseline_comparison(record.pricing, estimate.usage, estimate.provenance)
details: Final = estimate.usage.prompt_tokens_details if estimate.usage is not None else None
writes: Final = details.cache_creation_token_details if details is not None else None
return BaselinePublication(
comparison_id=record.scope,
comparison_started_at=first_at,
status="estimated" if costs is not None else "unknown",
reason=estimate.reason if costs is not None or estimate.usage is None else "pricing_unavailable",
provenance=estimate.provenance if costs is not None else None,
actual_spend=costs.actual if costs is not None else None,
baseline_spend=costs.baseline if costs is not None else None,
input_tokens=details.text_tokens if details is not None else None,
cache_read_input_tokens=details.cached_tokens if details is not None else None,
cache_creation_5m_input_tokens=writes.ephemeral_5m_input_tokens if writes is not None else None,
cache_creation_1h_input_tokens=writes.ephemeral_1h_input_tokens if writes is not None else None,
)
class _Comparison(BaseModel):
revision: int
published_revision: int
initial_equivalent: bool
retired: bool
history: str | None
class _StoredRecord(BaseModel):
data: str
publication: str | None
conflicted: bool
started_at: float
class _Change(BaseModel):
request_id: str
publication: BaselinePublication
api_key: str
session_id: str
router_name: str
baseline_model: str
covered_delta: int
actual_delta: float
savings_delta: float
daily: DailyBaselineAttribution | None
class _TransactionManager(Protocol):
async def __aenter__(self) -> SupportsRawQueries: ...
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
class _TransactionalDatabase(Protocol):
def tx(self, *, timeout: timedelta) -> _TransactionManager: ...
_COMPARISONS: Final = TypeAdapter(tuple[_Comparison, ...])
_RECORDS: Final = TypeAdapter(tuple[_StoredRecord, ...])
_HISTORY: Final = TypeAdapter(BaselineHistory)
_PAGE_TIMESTAMPS: Final = 128
_TRANSACTION_TIMEOUT: Final = timedelta(seconds=10)
_CREATE_COMPARISON: Final = """
INSERT INTO "LiteLLM_AutoRouterBaselineComparison"
(scope, api_key, session_id, router_name, initial_equivalent)
VALUES ($1, $2, $3, $4, NOT EXISTS (
SELECT 1 FROM "LiteLLM_AutoRouterSession"
WHERE api_key = $2 AND session_id = $3 AND router_name = $4
)) ON CONFLICT (scope) DO NOTHING
"""
_LOCK_COMPARISON: Final = """
SELECT revision, published_revision, initial_equivalent, retired, history
FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope = $1 FOR UPDATE
"""
_INSERT_RECORD: Final = """
INSERT INTO "LiteLLM_AutoRouterBaselineObservation"
(request_id, scope, started_at, revision, data)
VALUES ($1, $2, $3::float8, $4::bigint, $5)
ON CONFLICT (request_id) DO NOTHING
"""
_MARK_CONFLICT: Final = """
UPDATE "LiteLLM_AutoRouterBaselineObservation"
SET conflicted = TRUE, revision = $4::bigint
WHERE request_id = $1 AND scope = $2 AND data <> $3 AND NOT conflicted
"""
_READ_PAGE: Final = """
WITH times AS (
SELECT DISTINCT started_at FROM "LiteLLM_AutoRouterBaselineObservation"
WHERE scope = $1 AND revision > $2::bigint
AND ($3::float8 IS NULL OR started_at > $3::float8)
AND ($5::float8 IS NULL OR (
started_at >= $5::float8 AND publication::jsonb->>'status' = 'estimated'
))
ORDER BY started_at LIMIT $4::int
)
SELECT data, publication, conflicted, started_at
FROM "LiteLLM_AutoRouterBaselineObservation"
WHERE scope = $1 AND revision > $2::bigint
AND started_at IN (SELECT started_at FROM times)
AND ($5::float8 IS NULL OR publication::jsonb->>'status' = 'estimated')
ORDER BY started_at, request_id
"""
_UPDATE_LOGS: Final = """
WITH changes AS (
SELECT request_id, publication::jsonb AS publication
FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
)
UPDATE "LiteLLM_SpendLogs" AS logs
SET metadata = (COALESCE(logs.metadata::jsonb, '{}'::jsonb) - 'autorouter_baseline_observation') || jsonb_build_object(
'autorouter_savings_estimate', changes.publication,
'autorouter_savings', CASE WHEN changes.publication->>'status' = 'estimated' THEN
(changes.publication->>'baseline_spend')::float8 - (changes.publication->>'actual_spend')::float8
ELSE NULL END
)
FROM changes WHERE logs.request_id = changes.request_id
"""
_UPDATE_PUBLICATIONS: Final = """
UPDATE "LiteLLM_AutoRouterBaselineObservation" AS observations
SET publication = x.publication::text
FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb)
WHERE observations.request_id = x.request_id
"""
_UPDATE_SESSIONS: Final = """
WITH changes AS (
SELECT * FROM jsonb_to_recordset($1::jsonb) AS x(
api_key text, session_id text, router_name text, baseline_model text,
covered_delta int, actual_delta float8, savings_delta float8
)
), totals AS (
SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta,
SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta
FROM changes GROUP BY api_key, session_id, router_name
), models AS (
SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas
FROM (
SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta
FROM changes GROUP BY api_key, session_id, router_name, baseline_model
) grouped GROUP BY api_key, session_id, router_name
)
UPDATE "LiteLLM_AutoRouterSession" AS session
SET saved_spend = session.saved_spend + totals.savings_delta,
savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta,
savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta,
savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta,
savings_estimated_baseline_models = (
SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM (
SELECT key, SUM(value::int)::int AS value FROM (
SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models)
UNION ALL SELECT * FROM jsonb_each_text(models.deltas)
) combined GROUP BY key HAVING SUM(value::int) > 0
) counts
)
FROM totals JOIN models USING (api_key, session_id, router_name)
WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id
AND session.router_name = totals.router_name
"""
def _primary_transaction(client: PrismaClient) -> _TransactionManager:
primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db))
return primary.tx(timeout=_TRANSACTION_TIMEOUT)
def _serialized(model: BaseModel) -> str:
return json.dumps(model.model_dump(mode="json"), sort_keys=True, separators=(",", ":"))
def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, new: BaselinePublication) -> _Change:
previous: Final = old.costs if old is not None else None
current: Final = new.costs
return _Change(
request_id=record.observation.request_id,
publication=new,
api_key=record.api_key,
session_id=record.session_id,
router_name=record.router_name,
baseline_model=record.baseline_model,
covered_delta=int(current is not None) - int(previous is not None),
actual_delta=(current.actual if current is not None else 0.0)
- (previous.actual if previous is not None else 0.0),
savings_delta=(current.savings if current is not None else 0.0)
- (previous.savings if previous is not None else 0.0),
daily=record.daily,
)
def _project_group(
previous: tuple[BaselineHistory, tuple[_Change, ...]], stored: Sequence[_StoredRecord]
) -> tuple[BaselineHistory, tuple[_Change, ...]]:
history, prior_changes = previous
records: Final = tuple(BaselineAccountingRecord.model_validate_json(item.data) for item in stored)
observations: Final = tuple(
record.observation.model_copy(
update=MappingProxyType(
{"outcome": "uncertain", "baseline_equivalent": False, "reason": "conflicting_observation"}
)
)
if row.conflicted
else record.observation
for record, row in zip(records, stored)
)
advanced, estimates = advance_baseline_history(history, observations)
publications: Final = tuple(
baseline_publication(
record, estimate, advanced.first_at if advanced.first_at is not None else observations[0].started_at
)
for record, estimate in zip(records, estimates)
)
changes: Final = tuple(
_change(record, old, publication)
for record, row, publication in zip(records, stored, publications)
for old in (BaselinePublication.model_validate_json(row.publication) if row.publication else None,)
if publication != old
)
return advanced, (*prior_changes, *changes)
async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None:
if not changes:
return
serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":"))
await db.execute_raw(_UPDATE_LOGS, serialized)
await db.execute_raw(_UPDATE_SESSIONS, serialized)
for entity, table in DAILY_SPEND_TABLES.items():
if adjustments := tuple(
change.daily.adjustment(target, change.savings_delta, change.request_id)
for change in changes
if change.daily is not None and change.savings_delta != 0
for target in change.daily.targets
if target.entity == entity
):
statement, values = build_bulk_upsert(table, merge_by_conflict_key(table, adjustments))
await db.execute_raw(statement, *values)
await db.execute_raw(_UPDATE_PUBLICATIONS, serialized)
class BaselineAccountingStore:
def __init__(self, transaction: Callable[[], _TransactionManager]) -> None:
self.transaction: Final = transaction
@classmethod
def for_client(cls, client: PrismaClient) -> BaselineAccountingStore:
def transaction() -> _TransactionManager:
return _primary_transaction(client)
return cls(transaction)
async def append(
self, record: BaselineAccountingRecord
) -> Literal["recorded", "retired", "conflict", "unavailable"]:
try:
async with self.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
await db.execute_raw(
_CREATE_COMPARISON, record.scope, record.api_key, record.session_id, record.router_name
)
rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, record.scope)))
if not rows:
return "unavailable"
revision: Final = rows[0].revision + 1
data: Final = _serialized(record)
inserted: Final = await db.execute_raw(
_INSERT_RECORD,
record.observation.request_id,
record.scope,
record.observation.started_at,
revision,
data,
)
if inserted and record.turn is not None:
await write_autorouter_turn(db, record.turn)
conflicted: Final = (
0
if inserted
else await db.execute_raw(
_MARK_CONFLICT, record.observation.request_id, record.scope, data, revision
)
)
canonical: Final = (
_RECORDS.validate_python(
tuple(
await db.query_raw(
'SELECT data, publication, conflicted, started_at FROM "LiteLLM_AutoRouterBaselineObservation" '
"WHERE request_id=$1 AND scope=$2",
record.observation.request_id,
record.scope,
)
)
)
if not inserted
else ()
)
if not inserted and not canonical:
return "conflict"
if rows[0].retired:
await _publish(
db,
(
_change(
BaselineAccountingRecord.model_validate_json(canonical[0].data)
if canonical
else record,
BaselinePublication.model_validate_json(canonical[0].publication)
if canonical and canonical[0].publication is not None
else None,
BaselinePublication(
comparison_id=record.scope,
comparison_started_at=canonical[0].started_at
if canonical
else record.observation.started_at,
status="unknown",
reason="comparison_retired",
),
),
),
)
return "retired"
if inserted or conflicted:
await self._withdraw(
db, record.scope, canonical[0].started_at if canonical else record.observation.started_at
)
await db.execute_raw(
'UPDATE "LiteLLM_AutoRouterBaselineComparison" SET revision = $2::bigint, '
"updated_at = CURRENT_TIMESTAMP, attempted_at = NULL WHERE scope = $1",
record.scope,
revision,
)
return "recorded"
except Exception: # noqa: BLE001 # accounting failure must not change inference or actual billing
verbose_proxy_logger.warning("Auto-router baseline observation could not be persisted")
return "unavailable"
async def _pages(
self, db: SupportsRawQueries, scope: str, after_revision: int, withdraw_from: float | None = None
) -> AsyncIterator[tuple[_StoredRecord, ...]]:
cursor: float | None = None
while page := _RECORDS.validate_python(
tuple(await db.query_raw(_READ_PAGE, scope, after_revision, cursor, _PAGE_TIMESTAMPS, withdraw_from))
):
yield page
cursor = page[-1].started_at # rebind-ok: keyset pagination advances after each complete timestamp group
async def _withdraw(self, db: SupportsRawQueries, scope: str, started_at: float) -> None:
async for page in self._pages(db, scope, 0, withdraw_from=started_at):
await _publish(
db,
tuple(
_change(
BaselineAccountingRecord.model_validate_json(row.data),
previous,
BaselinePublication(
comparison_id=scope,
comparison_started_at=min(previous.comparison_started_at, started_at),
status="unknown",
reason="pending_projection",
),
)
for row in page
if row.publication is not None
for previous in (BaselinePublication.model_validate_json(row.publication),)
),
)
async def retire_before(self, cutoff: datetime, batch_size: int, timeout_ms: int) -> None:
async with self.transaction() as db:
await db.execute_raw(f"SET LOCAL statement_timeout = {max(1, timeout_ms)}")
await db.execute_raw(f"SET LOCAL lock_timeout = {max(1, timeout_ms)}")
await db.execute_raw(
'WITH expired AS (SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison" '
"WHERE NOT retired AND updated_at < $1::timestamptz ORDER BY updated_at "
"LIMIT $2::int FOR UPDATE SKIP LOCKED) "
'UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison '
"SET retired=TRUE, history=NULL FROM expired WHERE comparison.scope=expired.scope",
cutoff,
batch_size,
)
await db.execute_raw(
'DELETE FROM "LiteLLM_AutoRouterBaselineObservation" WHERE request_id IN ('
'SELECT event.request_id FROM "LiteLLM_AutoRouterBaselineObservation" AS event '
'JOIN "LiteLLM_AutoRouterBaselineComparison" AS comparison USING (scope) '
"WHERE comparison.retired AND comparison.updated_at < $1::timestamptz "
"LIMIT $2::int)",
cutoff,
batch_size,
)
async def project(self, scope: str) -> Literal["published", "unchanged", "unavailable"]:
try:
async with self.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 5000")
await db.execute_raw("SET LOCAL lock_timeout = 1000")
rows: Final = _COMPARISONS.validate_python(tuple(await db.query_raw(_LOCK_COMPARISON, scope)))
if not rows or rows[0].retired or rows[0].revision == rows[0].published_revision:
return "unchanged"
missing_log: Final = await db.query_raw(
'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" AS observation '
'WHERE scope=$1 AND publication IS NULL AND NOT EXISTS (SELECT 1 FROM "LiteLLM_SpendLogs" AS log '
"WHERE log.request_id=observation.request_id) LIMIT 1",
scope,
)
if missing_log:
return "unavailable"
state: Final = rows[0]
checkpoint: Final = (
_HISTORY.validate_json(state.history)
if state.history is not None
else BaselineHistory(equivalent=state.initial_equivalent)
)
changed: Final = await db.query_raw(
'SELECT 1 FROM "LiteLLM_AutoRouterBaselineObservation" '
"WHERE scope = $1 AND revision > $2::bigint AND started_at <= $3::float8 LIMIT 1",
scope,
state.published_revision,
checkpoint.last_at,
)
history = BaselineHistory(equivalent=state.initial_equivalent) if changed else checkpoint
async for page in self._pages(db, scope, 0 if changed else state.published_revision):
history, updates = reduce(
_project_group,
(tuple(group) for _, group in groupby(page, key=lambda item: item.started_at)),
(history, ()),
)
await _publish(db, updates)
await db.execute_raw(
'UPDATE "LiteLLM_AutoRouterBaselineComparison" '
"SET published_revision = revision, history = $2 WHERE scope = $1",
scope,
_HISTORY.dump_json(history).decode(),
)
return "published"
except Exception: # noqa: BLE001 # rollback leaves the durable revision dirty for a later flush
verbose_proxy_logger.warning("Auto-router baseline projection remains pending")
return "unavailable"
class _Scope(BaseModel):
scope: str
_SCOPES: Final = TypeAdapter(tuple[_Scope, ...])
_CLAIM_DIRTY: Final = """
WITH candidates AS (
SELECT scope FROM "LiteLLM_AutoRouterBaselineComparison"
WHERE NOT retired AND revision <> published_revision
AND (attempted_at IS NULL OR attempted_at < CURRENT_TIMESTAMP - INTERVAL '30 seconds')
ORDER BY attempted_at NULLS FIRST, updated_at, scope LIMIT 32 FOR UPDATE SKIP LOCKED
)
UPDATE "LiteLLM_AutoRouterBaselineComparison" AS comparison
SET attempted_at = CURRENT_TIMESTAMP FROM candidates
WHERE comparison.scope = candidates.scope RETURNING comparison.scope
"""
async def _flush_records(
store: BaselineAccountingStore, records: Sequence[BaselineAccountingRecord]
) -> tuple[BaselineAccountingRecord, ...]:
slots: Final = asyncio.Semaphore(4)
async def append(record: BaselineAccountingRecord) -> bool:
async with slots:
return await store.append(record) == "unavailable"
failed: Final = await asyncio.gather(*(append(record) for record in records))
return tuple(record for record, retry in zip(records, failed) if retry)
async def flush_baseline_accounting(client: PrismaClient) -> None:
from litellm.proxy.utils import request_spend_log_flush
store: Final = BaselineAccountingStore.for_client(client)
async with client.baseline_accounting_lock:
batch: Final = tuple(client.baseline_accounting_transactions[:32])
client.baseline_accounting_transactions = client.baseline_accounting_transactions[
32:
] # rebind-ok: drain under lock
more_queued: Final = bool(client.baseline_accounting_transactions)
try:
remaining: Final = await asyncio.wait_for(_flush_records(store, batch), timeout=5)
except (Exception, asyncio.CancelledError) as error: # noqa: BLE001 # unknown acknowledgements can be replayed safely
async with client.baseline_accounting_lock:
client.baseline_accounting_transactions.extend(batch)
if isinstance(error, asyncio.CancelledError):
raise
return
async with client.baseline_accounting_lock:
client.baseline_accounting_transactions.extend(remaining)
if more_queued and len(remaining) < len(batch):
request_spend_log_flush(client)
try:
async with store.transaction() as db:
await db.execute_raw("SET LOCAL statement_timeout = 1000")
scopes: Final = _SCOPES.validate_python(tuple(await db.query_raw(_CLAIM_DIRTY)))
slots: Final = asyncio.Semaphore(4)
async def project(item: _Scope) -> str:
async with slots:
return await store.project(item.scope)
outcomes: Final = await asyncio.wait_for(asyncio.gather(*(project(item) for item in scopes)), timeout=5)
if len(scopes) == 32 and "published" in outcomes:
request_spend_log_flush(client)
except Exception: # noqa: BLE001 # durable dirty comparisons remain eligible after the retry interval
verbose_proxy_logger.warning("Auto-router baseline projection will retry on a later spend flush")

32
litellm/proxy/db/daily_spend_bulk_upsert.py
View file

@ -14,6 +14,8 @@ from itertools import groupby
from types import MappingProxyType
from typing import Final, Literal
from pydantic import TypeAdapter
DailySpendEntity = Literal["user", "team", "org", "tag", "end_user", "agent"]
SqlValue = str | int | float | None
@ -43,6 +45,36 @@ DAILY_SPEND_TABLES: Final[Mapping[DailySpendEntity, DailySpendTable]] = MappingP
}
)
_ENTITY_INPUT_KEYS: Final[Mapping[DailySpendEntity, str]] = MappingProxyType(
{
"user": "user",
"team": "team_id",
"org": "organization_id",
"end_user": "end_user",
"agent": "agent_id",
"tag": "request_tags",
}
)
_TAGS: Final = TypeAdapter(tuple[str, ...])
def daily_spend_entity_ids(payload: Mapping[str, object], entity: DailySpendEntity) -> tuple[str | None, ...]:
key: Final = _ENTITY_INPUT_KEYS[entity]
if key not in payload:
return ()
value: Final = payload[key]
if entity == "tag":
if value is None:
return ()
tags: Final = _TAGS.validate_json(value) if isinstance(value, str) else _TAGS.validate_python(value)
return tuple(dict.fromkeys(tags))
if value is None:
return (None,) if entity == "user" else ()
if not isinstance(value, str) or (entity == "end_user" and not value):
return ()
return (value,)
# The unique constraint's columns after the entity id, in constraint order. A NULL can
# never match itself in a unique index, so every one of these is normalized to '': the
# conflict target has to be NULL-free or the row is re-inserted on every single flush.

146
litellm/proxy/db/db_spend_update_writer.py
View file

@ -18,6 +18,7 @@ from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload
from urllib.parse import quote, unquote
from pydantic import TypeAdapter
from typing_extensions import LiteralString, ReadOnly, TypedDict
import litellm
@ -51,6 +52,7 @@ from litellm.proxy.common_utils.user_api_key_cache import project_cache_key
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
daily_spend_entity_ids,
merge_by_conflict_key,
)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
@ -82,6 +84,8 @@ from litellm.repositories.prisma_protocols import BatchTable
from litellm.types.utils import CallTypes
if TYPE_CHECKING:
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution
from litellm.proxy.utils import PrismaClient, ProxyLogging
else:
PrismaClient = Any
@ -89,6 +93,7 @@ else:
RESPONSES_SESSION_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
_SPEND_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
def _org_member_transaction_key(org_id: str, user_id: str) -> str:
@ -579,25 +584,31 @@ class DBSpendUpdateWriter:
metadata_raw: Final = payload.get("metadata")
if not metadata_raw:
return
metadata: Final = json.loads(metadata_raw)
if not isinstance(metadata, dict) or not metadata.get("routing_decision"):
metadata: Final = _SPEND_METADATA_ADAPTER.validate_json(metadata_raw)
routing_decision: Final = metadata.get("routing_decision")
if not isinstance(routing_decision, Mapping) or not routing_decision:
return
from litellm.proxy.db.autorouter_session_rollup import (
build_autorouter_turn_transaction,
)
usage_object_raw: Final = metadata.get("usage_object")
cost_breakdown: Final = metadata.get("cost_breakdown")
savings_estimate: Final = metadata.get("autorouter_savings_estimate")
savings_spend: Final = compute_savings_spend(
model=payload.get("model"),
custom_llm_provider=payload.get("custom_llm_provider"),
compression_saved_tokens=0,
gateway_injected_cache=marks_gateway_injection(metadata, payload.get("model_id")),
routing_decision=metadata.get("routing_decision"),
routing_decision=routing_decision,
usage_object=usage_object_raw if isinstance(usage_object_raw, dict) else None,
model_id=payload.get("model_id"),
llm_router=get_llm_router,
cost_breakdown=metadata.get("cost_breakdown"),
cost_breakdown=cost_breakdown if isinstance(cost_breakdown, Mapping) else None,
recorded_autorouter_savings=metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=(
savings_estimate if isinstance(savings_estimate, Mapping) else None
),
billed_at=payload.get("endTime"),
)
transaction: Final = build_autorouter_turn_transaction(
@ -605,6 +616,11 @@ class DBSpendUpdateWriter:
metadata=metadata,
saved_spend=savings_spend.autorouter,
)
try:
if await self._enqueue_baseline_accounting(payload, metadata, transaction, prisma_client):
return
except Exception: # noqa: BLE001 # optional baseline capture must preserve the original actual-spend rollup
verbose_proxy_logger.warning("Auto-router baseline observation was unavailable; actual turn retained")
if transaction is None:
return
async with prisma_client._autorouter_turn_transactions_lock:
@ -612,6 +628,95 @@ class DBSpendUpdateWriter:
except Exception as e: # noqa: BLE001 # a metrics enqueue must never fail the spend write
verbose_proxy_logger.debug("_enqueue_autorouter_turn_transaction error (non-blocking): %s", e)
async def _enqueue_baseline_accounting(
self,
payload: SpendLogsPayload,
metadata: Mapping[str, object],
turn: "AutoRouterTurnTransaction | None",
prisma_client: "PrismaClient",
) -> bool:
from litellm.proxy.db.baseline_accounting import (
BaselineAccountingRecord,
)
from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation
from litellm.proxy.spend_tracking.savings import baseline_cost_snapshot
serialized: Final = metadata.get("autorouter_baseline_observation")
if not isinstance(serialized, str):
return False
captured: Final = CapturedBaselineObservation.model_validate_json(serialized)
if captured.api_key != payload["api_key"] or captured.session_id != payload["session_id"]:
return False
decision: Final = _SPEND_METADATA_ADAPTER.validate_python(
metadata.get("routing_decision") or MappingProxyType({})
)
breakdown: Final = _SPEND_METADATA_ADAPTER.validate_python(
metadata.get("cost_breakdown") or MappingProxyType({})
)
daily: Final = await self._baseline_daily_attribution(payload, prisma_client)
record: Final = BaselineAccountingRecord(
scope=captured.scope,
api_key=captured.api_key,
session_id=captured.session_id,
router_name=captured.router_name,
baseline_model=captured.baseline_model,
observation=captured.observation.model_copy(update=MappingProxyType({"request_id": payload["request_id"]})),
pricing=baseline_cost_snapshot(captured.model, captured.prices, payload["spend"], breakdown, decision),
turn=turn,
daily=daily,
)
async with prisma_client.baseline_accounting_lock:
if len(prisma_client.baseline_accounting_transactions) >= 10000:
verbose_proxy_logger.warning("Auto-router baseline observation queue is full")
return False
prisma_client.baseline_accounting_transactions.append(record)
from litellm.proxy.utils import request_spend_log_flush
request_spend_log_flush(prisma_client)
return True
async def _baseline_daily_attribution(
self,
payload: SpendLogsPayload,
prisma_client: "PrismaClient",
) -> "DailyBaselineAttribution | None":
from litellm.proxy.db.baseline_accounting import DailyBaselineAttribution, DailyBaselineTarget
normalized: Final = cast(SpendLogsPayload, MappingProxyType({**payload, "end_user_id": payload["end_user"]}))
bases: Final = tuple(
zip(
DAILY_SPEND_TABLES,
await asyncio.gather(
*(
self._common_add_spend_log_transaction_to_daily_transaction( # pyright: ignore[reportUnknownMemberType] # legacy payload union; this caller supplies a validated spend payload
normalized,
prisma_client,
"request_tags" if entity == "tag" else entity,
)
for entity in DAILY_SPEND_TABLES
)
),
)
)
base: Final = next((base for _, base in bases if base is not None), None)
if base is None:
return None
return DailyBaselineAttribution(
date=base["date"],
api_key=base["api_key"],
model=base.get("model"),
custom_llm_provider=base.get("custom_llm_provider"),
model_group=base.get("model_group"),
endpoint=base.get("endpoint"),
mcp_namespaced_tool_name=base.get("mcp_namespaced_tool_name"),
targets=tuple(
DailyBaselineTarget(entity=entity, entity_id=identity)
for entity, values in bases
if values is not None
for identity in daily_spend_entity_ids(payload, entity)
),
)
def _enqueue_tool_registry_upsert(
self,
kwargs: dict | None,
@ -2322,21 +2427,13 @@ class DBSpendUpdateWriter:
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
) -> BaseDailySpendTransaction | None:
common_expected_keys: Final = ["startTime", "api_key"]
if type == "user":
expected_keys = ["user", *common_expected_keys]
elif type == "team":
expected_keys = ["team_id", *common_expected_keys]
elif type == "org":
expected_keys = ["organization_id", *common_expected_keys]
elif type == "request_tags":
expected_keys = ["request_tags", *common_expected_keys]
elif type == "end_user":
expected_keys = ["end_user_id", *common_expected_keys]
elif type == "agent":
expected_keys = ["agent_id", *common_expected_keys]
else:
raise ValueError(f"Invalid type: {type}")
entity: Final = "tag" if type == "request_tags" else type
identity_payload: Final = (
MappingProxyType({**payload, "end_user": payload.get("end_user_id")}) if type == "end_user" else payload
)
if not daily_spend_entity_ids(identity_payload, entity):
return None
expected_keys: Final = ("startTime", "api_key")
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
"Missing expected keys: %s, in payload, skipping from daily_user_spend_transactions", expected_keys
@ -2399,6 +2496,7 @@ class DBSpendUpdateWriter:
usage_object=usage_obj,
cost_breakdown=_metadata.get("cost_breakdown"),
recorded_autorouter_savings=_metadata.get("autorouter_savings"),
recorded_autorouter_savings_estimate=_metadata.get("autorouter_savings_estimate"),
billed_at=payload.get("endTime"),
)
timed_duration_ms: Final = _timed_request_duration_ms(payload, request_status, is_internal_call)
@ -2597,14 +2695,10 @@ class DBSpendUpdateWriter:
verbose_proxy_logger.debug("request_tags is None for request. Skipping incrementing tag spend.")
return
request_tags: Sequence[str] = []
if isinstance(payload["request_tags"], str):
request_tags = json.loads(payload["request_tags"])
elif isinstance(payload["request_tags"], list):
request_tags = payload["request_tags"]
else:
raise ValueError(f"Invalid request_tags: {payload['request_tags']}")
request_tags: Final = daily_spend_entity_ids(payload, "tag")
for tag in request_tags:
if tag is None:
continue
endpoint_str = base_daily_transaction.get("endpoint") or ""
daily_transaction_key = f"{tag}_{base_daily_transaction['date']}_{payload['api_key']}_{payload['model']}_{payload['custom_llm_provider']}_{endpoint_str}"
daily_transaction = DailyTagSpendTransaction(

11
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -549,6 +549,17 @@ class SpendLogCleanup:
Prune auto-router session rollup rows, which carry their own retention horizon.
"""
session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
from litellm.proxy.db.baseline_accounting import BaselineAccountingStore
if remaining_ms := self._remaining_timeout_ms(deadline)():
try:
await BaselineAccountingStore.for_client(prisma_client).retire_before(
session_cutoff,
self.batch_size,
remaining_ms,
)
except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job
verbose_proxy_logger.warning("Auto-router baseline retention remains pending")
sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
return (sessions_result,)

2
litellm/proxy/hooks/__init__.py
View file

@ -2,6 +2,7 @@ import os
from typing import Final, Literal
from . import *
from .autorouter_baseline_cache import AutoRouterBaselineCache
from .cache_control_check import _PROXY_CacheControlCheck
from .litellm_skills import SkillsInjectionHook
from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler
@ -25,6 +26,7 @@ PROXY_HOOKS: Final = {
"max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler,
"sensitive_data_routing": _PROXY_SensitiveDataRoutingHandler,
"prompt_cache_prediction": PromptCacheObserver,
"autorouter_baseline_cache": AutoRouterBaselineCache,
}
## FEATURE FLAG HOOKS ##

344
litellm/proxy/hooks/autorouter_baseline_cache.py Normal file
View file

@ -0,0 +1,344 @@
from __future__ import annotations
import asyncio
import hashlib
import json
import time
from collections.abc import Callable, Mapping
from dataclasses import dataclass, replace
from datetime import datetime
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
import httpx
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
get_litellm_metadata_from_kwargs, # pyright: ignore[reportUnknownVariableType] # legacy metadata boundary validated below
)
from litellm.llms.anthropic.prompt_cache_prediction import (
CountedPromptCachePlan,
NativePredictionTarget,
TokenCounter,
UnsupportedCachePlan,
UnsupportedPredictionTarget,
count_cache_plan,
count_prompt_tokens,
parse_cache_plan,
resolve_baseline_prediction_target,
supported_baseline_recipient,
supported_prediction_headers,
)
from litellm.proxy.spend_tracking.baseline_accounting import BaselineObservation
from litellm.proxy.spend_tracking.savings import (
_effective_model_info, # pyright: ignore[reportPrivateUsage] # existing deployment-price owner
_proxy_llm_router, # pyright: ignore[reportPrivateUsage] # existing optional proxy-router owner
)
from litellm.types.router import BaselineRouteStamp
from litellm.types.utils import CallTypes, ModelInfo, Usage
from litellm.utils import get_prompt_cache_min_tokens
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
_METADATA: Final = TypeAdapter(Mapping[str, object])
_PRICES: Final[TypeAdapter[ModelInfo | None]] = TypeAdapter(ModelInfo | None)
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
_COUNT_TIMEOUT: Final = 3.0
_MAX_COUNTS: Final = 4096
class CapturedBaselineObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
scope: str
api_key: str
session_id: str
router_name: str
baseline_model: str
model: str
prices: ModelInfo | None
observation: BaselineObservation
@dataclass(frozen=True, slots=True)
class BaselineCacheContext:
collector: AutoRouterBaselineCache
capture: CapturedBaselineObservation
target: NativePredictionTarget | UnsupportedPredictionTarget
baseline_deployment_id: str
invalidated: str | None = None
class _Metadata(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
route: BaselineRouteStamp = Field(alias="_autorouter_baseline_route")
user_api_key_hash: str = Field(min_length=1)
session_id: str | None = None
class _WireEvent(BaseModel):
model_config = ConfigDict(strict=True, arbitrary_types_allowed=True)
httpx_response: httpx.Response
api_call_start_time: datetime
completion_start_time: datetime
custom_llm_provider: str
stream: bool = False
prompt_cache_response_complete: bool = False
class _ResponseUsage(BaseModel):
model_config = ConfigDict(strict=True, from_attributes=True)
usage: Usage | None = None
def _digest(value: object) -> str:
return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
class AutoRouterBaselineCache(CustomLogger):
def __init__(
self,
prisma_client: PrismaClient | None,
router: Callable[[], Router | None] = _proxy_llm_router,
token_counter: TokenCounter | None = None,
clock: Callable[[], float] = time.time,
) -> None:
super().__init__() # pyright: ignore[reportUnknownMemberType] # legacy callback constructor
self.router: Final = router
self.token_counter: Final = token_counter
self.clock: Final = clock
self.count_slots: Final = asyncio.Semaphore(8)
self.counts: Mapping[str, tuple[int, float]] = MappingProxyType({})
async def async_pre_call_deployment_hook(self, kwargs: Mapping[str, object], call_type: CallTypes | None) -> None:
from litellm.litellm_core_utils.litellm_logging import Logging
logging_obj: Final = kwargs.get("litellm_logging_obj")
if not isinstance(logging_obj, Logging) or call_type != CallTypes.anthropic_messages:
return
try:
metadata: Final = _METADATA.validate_python(
get_litellm_metadata_from_kwargs(
{"litellm_params": kwargs} # mutable-ok: legacy metadata owner requires a dictionary
)
)
if metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
return
if logging_obj.baseline_cache_context is not None:
await invalidate_baseline_cache(logging_obj, "retried_request")
return
request: Final = _Metadata.model_validate(metadata)
session: Final = kwargs.get("litellm_session_id") or request.session_id or logging_obj.litellm_session_id
if not isinstance(session, str) or not session or len(session) > 256:
return
router: Final = self.router()
deployment: Final = router.get_deployment(request.route.baseline_deployment_id) if router else None
if deployment is None:
return
target: Final = resolve_baseline_prediction_target(deployment.litellm_params)
prices: Final = _PRICES.validate_python(
_effective_model_info(router, request.route.baseline_deployment_id, request.route.baseline_model)
)
scope: Final = "autorouter-baseline:v3:" + _digest(
(
request.user_api_key_hash,
session,
request.route.router_name,
request.route.baseline_deployment_id,
deployment.litellm_params.model_dump(mode="json"),
prices,
)
)
started: Final = logging_obj.start_time.timestamp()
capture: Final = CapturedBaselineObservation(
scope=scope,
api_key=request.user_api_key_hash,
session_id=session,
router_name=request.route.router_name,
baseline_model=request.route.baseline_model,
model=target.model if isinstance(target, NativePredictionTarget) else request.route.baseline_model,
prices=prices,
observation=BaselineObservation(
request_id=logging_obj.litellm_call_id,
started_at=started,
available_at=started,
outcome="uncertain",
baseline_equivalent=False,
reason="incomplete_response",
),
)
logging_obj.baseline_cache_context = BaselineCacheContext(
self, capture, target, request.route.baseline_deployment_id
)
except Exception: # noqa: BLE001 # optional observation cannot fail inference
verbose_proxy_logger.warning("Auto-router baseline observation could not be initialized")
async def _count(self, target: NativePredictionTarget, body: Mapping[str, JsonValue]) -> int | None:
key: Final = _digest((target.model, target.api_key, target.api_base, _JSON_BODY.validate_python(body)))
now: Final = self.clock()
cached: Final = self.counts.get(key)
if cached is not None and cached[1] > now:
return cached[0]
async with self.count_slots:
tokens: Final = (
await self.token_counter(target.model, target.api_key, body)
if self.token_counter is not None
else await count_prompt_tokens(target.model, target.api_key, body, api_base=target.api_base)
)
if tokens is None or tokens < 0:
return None
retained: Final = tuple((k, v) for k, v in self.counts.items() if v[1] > now and k != key)[-(_MAX_COUNTS - 1) :]
self.counts = MappingProxyType(dict((*retained, (key, (tokens, now + 3600)))))
return tokens
async def plan(
self, target: NativePredictionTarget, wire: httpx.Request, body: Mapping[str, JsonValue], usage: Usage | None
) -> tuple[CountedPromptCachePlan | None, str | None]:
if not supported_prediction_headers(wire.headers):
return None, "unsupported_request_headers"
plan: Final = parse_cache_plan(body)
if isinstance(plan, UnsupportedCachePlan):
return None, plan.reason
details: Final = usage.prompt_tokens_details if usage is not None else None
if (
not plan.breakpoints
and details is not None
and ((details.cached_tokens or 0) + (details.cache_creation_tokens or 0))
):
return None, "implicit_cache_without_breakpoints"
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int | None:
return await self._count(target, body)
try:
counted: Final = await asyncio.wait_for(
count_cache_plan(target.model, target.api_key, plan, token_counter=count), timeout=_COUNT_TIMEOUT
)
return (None, counted.reason) if isinstance(counted, UnsupportedCachePlan) else (counted, None)
except TimeoutError:
return None, "token_count_timeout"
except Exception: # noqa: BLE001 # token counting cannot fail a completed request
return None, "token_count_unavailable"
async def invalidate_baseline_cache(logging_obj: Logging, reason: str, *, completed: bool = False) -> None:
context: Final = logging_obj.baseline_cache_context
if context is not None:
logging_obj.baseline_cache_context = replace(
context, invalidated=reason
) # rebind-ok: request-owned retry marker
logging_obj.baseline_observation = context.capture.model_copy(
update=MappingProxyType(
{ # rebind-ok: capture uncertainty for failure logging
"observation": context.capture.observation.model_copy(
update=MappingProxyType(
{
"available_at": max(context.capture.observation.started_at, context.collector.clock()),
"reason": reason,
}
)
),
}
)
)
async def finalize_baseline_cache(logging_obj: Logging, response_obj: object) -> None:
context: Final = logging_obj.baseline_cache_context
if context is None:
return
try:
capture: Final = await _capture(context, logging_obj, response_obj)
if logging_obj.baseline_cache_context is context:
logging_obj.baseline_observation = capture # rebind-ok: attach only to the captured request owner
except Exception: # noqa: BLE001 # observation failures must preserve inference and billing
await invalidate_baseline_cache(logging_obj, "observation_unavailable")
async def _capture(
context: BaselineCacheContext, logging_obj: Logging, response_obj: object
) -> CapturedBaselineObservation:
original: Final = context.capture.observation
details: Final = _METADATA.validate_python(logging_obj.model_call_details)
if details.get("cache_hit") is True:
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType({"outcome": "response_cache", "reason": "response_cache_hit"})
)
}
)
)
event: Final = _WireEvent.model_validate(details)
wire: Final = event.httpx_response.request
usage: Final = _ResponseUsage.model_validate(response_obj).usage
complete: Final = (
event.custom_llm_provider == "anthropic"
and event.httpx_response.status_code == 200
and (not event.stream or event.prompt_cache_response_complete)
)
started: Final = original.started_at
available: Final = event.completion_start_time.timestamp()
if context.invalidated or not complete or not started <= available <= context.collector.clock():
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType(
{
"available_at": max(started, context.collector.clock()),
"reason": context.invalidated or "incomplete_response",
}
)
)
}
)
)
target: Final = context.target
if isinstance(target, UnsupportedPredictionTarget) or not supported_baseline_recipient(target, wire):
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": original.model_copy(
update=MappingProxyType(
{
"available_at": available,
"reason": target.reason
if isinstance(target, UnsupportedPredictionTarget)
else "unsupported_baseline_recipient",
}
)
)
}
)
)
body: Final = _JSON_BODY.validate_json(wire.content)
same: Final = (
logging_obj.get_router_model_id() == context.baseline_deployment_id and body.get("model") == target.model
)
plan, reason = await context.collector.plan(target, wire, body, usage)
minimum: Final = get_prompt_cache_min_tokens(target.model)
return context.capture.model_copy(
update=MappingProxyType(
{
"observation": BaselineObservation(
request_id=original.request_id,
started_at=started,
available_at=available,
outcome="complete",
baseline_equivalent=same,
usage=usage,
plan=plan,
minimum_cache_tokens=minimum,
reason=reason,
)
}
)
)

42
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -556,6 +556,9 @@ class _SessionAggRow(BaseModel):
total_tokens: int
spend: float
saved_spend: float
savings_estimated_turns: int = 0
savings_estimated_actual_spend: float = 0.0
savings_estimated_saved_spend: float = 0.0
classifier_cost: float
classifier_cost_recorded_turns: int
session_seconds: float
@ -582,9 +585,19 @@ def _cache_bucket(turns: int, hits: int) -> AutoRouterCacheBucket:
return AutoRouterCacheBucket(turns=turns, hits=hits, hit_rate_pct=_pct(hits, turns))
def _savings_cohort(
turns: int, estimated_turns: int, actual_spend: float, saved_spend: float
) -> tuple[float | None, float | None]:
if turns > 0 and estimated_turns == 0:
return None, None
return saved_spend, actual_spend + saved_spend
def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
return_misses: Final = row.return_turns - row.return_hits
baseline_spend: Final = row.spend + row.saved_spend
saved_spend, baseline_spend = _savings_cohort(
row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend
)
sessions: Final = row.sessions
return AutoRouterBenchmarkTotals(
sessions=sessions,
@ -593,11 +606,15 @@ def _benchmark_totals(row: _SessionAggRow) -> AutoRouterBenchmarkTotals:
avg_session_seconds=row.session_seconds / sessions if sessions else 0.0,
avg_tokens_per_session=row.total_tokens / sessions if sessions else 0.0,
spend=row.spend,
saved_spend=row.saved_spend,
savings_estimated_turns=row.savings_estimated_turns,
savings_estimated_actual_spend=row.savings_estimated_actual_spend,
saved_spend=saved_spend,
classifier_cost=row.classifier_cost if row.classifier_cost_recorded_turns == row.turns else None,
baseline_spend=baseline_spend,
saved_pct=_pct(row.saved_spend, baseline_spend),
saved_per_session=row.saved_spend / sessions if sessions else 0.0,
saved_pct=_pct(saved_spend, baseline_spend) if saved_spend is not None and baseline_spend is not None else None,
saved_per_session=(row.savings_estimated_saved_spend / sessions if sessions else 0.0)
if row.savings_estimated_turns == row.turns
else None,
cache=AutoRouterCacheStats(
coverage_pct=_pct(row.covered_turns, row.turns),
hit_rate_pct=_pct(row.cache_hits, row.covered_turns),
@ -627,6 +644,8 @@ def _benchmark_group(row: _SessionAggRow) -> AutoRouterBenchmarkGroup:
avg_tokens_per_session=totals.avg_tokens_per_session,
spend=totals.spend,
saved_spend=totals.saved_spend,
savings_estimated_turns=totals.savings_estimated_turns,
savings_estimated_actual_spend=totals.savings_estimated_actual_spend,
classifier_cost=totals.classifier_cost,
baseline_spend=totals.baseline_spend,
saved_pct=totals.saved_pct,
@ -658,6 +677,9 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow:
total_tokens=sum(row.total_tokens for row in rows),
spend=sum(row.spend for row in rows),
saved_spend=sum(row.saved_spend for row in rows),
savings_estimated_turns=sum(row.savings_estimated_turns for row in rows),
savings_estimated_actual_spend=sum(row.savings_estimated_actual_spend for row in rows),
savings_estimated_saved_spend=sum(row.savings_estimated_saved_spend for row in rows),
classifier_cost=sum(row.classifier_cost for row in rows),
classifier_cost_recorded_turns=sum(row.classifier_cost_recorded_turns for row in rows),
session_seconds=sum(row.session_seconds for row in rows),
@ -807,6 +829,9 @@ async def get_auto_router_session(
raise HTTPException(
status_code=404, detail=f"No auto-routed turns recorded for session {session_id!r} under this key"
)
saved_spend, baseline_spend = _savings_cohort(
row.turns, row.savings_estimated_turns, row.savings_estimated_actual_spend, row.savings_estimated_saved_spend
)
return AutoRouterSessionResponse(
session_id=session_id,
router_name=row.router_name,
@ -814,10 +839,13 @@ async def get_auto_router_session(
turns=row.turns,
last_model=row.last_model,
spend=row.spend,
saved_spend=row.saved_spend,
baseline_spend=row.spend + row.saved_spend,
savings_estimated_turns=row.savings_estimated_turns,
savings_estimated_actual_spend=row.savings_estimated_actual_spend,
saved_spend=saved_spend,
baseline_spend=baseline_spend if row.savings_estimated_turns == row.turns else None,
savings_estimated_baseline_spend=baseline_spend,
baseline_model=row.baseline_model,
baseline_models=row.baseline_models,
baseline_models=row.savings_estimated_baseline_models,
)

5
litellm/proxy/management_endpoints/coordination_redis_endpoints.py
View file

@ -364,6 +364,11 @@ async def update_coordination_redis_settings(
settings: Final = _merge_over_saved(request.settings, saved_settings or {})
_validated_params(settings)
from litellm.proxy.proxy_server import proxy_config
proxy_config.reject_config_owned_writes(
section_name=_GENERAL_SETTINGS_PARAM_NAME, changed_keys={_COORDINATION_REDIS_KEY: settings}
)
general_settings: Final = await _read_general_settings()
before_settings: Final = general_settings.get(_COORDINATION_REDIS_KEY)
action: Final[AUDIT_ACTIONS] = "updated" if isinstance(before_settings, dict) else "created"

6
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -38,6 +38,7 @@ if TYPE_CHECKING:
FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500
BATCH_CREATE_HIDDEN_PARAM: Final = "batch_create"
LITELLM_EXECUTED_BATCH_ID_PREFIX: Final = "litellm_batch_"
def validate_file_list_limit(limit: int | None) -> None:
@ -179,6 +180,11 @@ def get_batch_id_from_unified_batch_id(file_id: str) -> str:
return re.split(r"[;,]", batch_id, maxsplit=1)[0]
def is_litellm_executed_batch(decoded_unified_batch_id: str) -> bool:
_, marker, batch_id = decoded_unified_batch_id.partition("llm_batch_id:")
return bool(marker) and batch_id.startswith(LITELLM_EXECUTED_BATCH_ID_PREFIX)
def encode_file_id_with_model(file_id: str, model: str, id_type: Literal["file", "batch"] = "file") -> str:
"""
Encode a file/batch ID with model routing information.

107
litellm/proxy/openai_files_endpoints/files_endpoints.py
View file

@ -7,7 +7,7 @@
import asyncio
import traceback
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Any, BinaryIO, Final, TypedDict, cast, get_args
import httpx
@ -32,10 +32,15 @@ from litellm.litellm_core_utils.cloud_storage_security import (
is_managed_cloud_storage_uri,
)
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.llms.base_llm.managed_resources.isolation import build_list_page
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.batches_endpoints.litellm_executed_batches import (
litellm_executed_provider_of,
resolve_litellm_executed_provider,
)
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.common_utils.http_parsing_utils import (
_read_request_body,
@ -67,6 +72,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
add_internal_model_credentials,
apply_team_provider_credentials,
authorize_model_for_key,
encode_file_id_with_model,
extract_file_creation_params,
get_authorized_credentials_for_model,
@ -86,7 +92,7 @@ from litellm.proxy.openai_files_endpoints.general_upload_validation import (
coerce_optional_str_list_setting,
raise_upload_validation_failure,
)
from litellm.proxy.utils import ProxyLogging, is_known_model
from litellm.proxy.utils import PrismaClient, ProxyLogging, is_known_model
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.router import Router
from litellm.types.llms.openai import (
@ -99,6 +105,64 @@ from litellm.types.llms.openai import (
router: Final = APIRouter()
def _names_a_litellm_executed_provider(llm_router: Router, candidate: str, team_id: str | None) -> bool:
credentials: Final = llm_router.get_deployment_credentials_with_provider(model_id=candidate, team_id=team_id)
return credentials is not None and litellm_executed_provider_of(credentials) is not None
async def _litellm_executed_batch_input_model(
llm_router: Router | None,
purpose: OpenAIFilesPurpose,
model: str | None,
target_model_names_list: Sequence[str],
user_api_key_dict: UserAPIKeyAuth,
explicit_storage: str | None,
) -> str | None:
if llm_router is None:
return None
candidates: Final = (model,) if model is not None else tuple(target_model_names_list)
team_id: Final = user_api_key_dict.team_id
await asyncio.gather(
*(
authorize_model_for_key(model_id=candidate, llm_router=llm_router, user_api_key_dict=user_api_key_dict)
for candidate in candidates
if _names_a_litellm_executed_provider(llm_router, candidate, team_id)
)
)
if explicit_storage is not None:
return None
providers: Final = await asyncio.gather(
*(resolve_litellm_executed_provider(llm_router, candidate, team_id) for candidate in candidates)
)
executed: Final = tuple(
candidate for candidate, provider in zip(candidates, providers, strict=True) if provider is not None
)
if not executed:
return None
if purpose != "batch":
raise ProxyException(
message=(
f"The server behind {', '.join(executed)} has no Files API, so LiteLLM keeps only batch input "
f"files for it and runs the batch itself: upload with purpose=batch; got purpose={purpose}"
),
type="invalid_request_error",
param="purpose",
code=400,
)
if len(candidates) == 1:
return executed[0]
raise ProxyException(
message=(
f"LiteLLM runs batches for {', '.join(executed)} itself and keeps their input files, so a batch "
f"input file can target only that one model; got target_model_names={', '.join(candidates)}"
),
type="invalid_request_error",
param="target_model_names",
code=400,
)
_MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None)
_LISTED_FILES_ADAPTER: Final = TypeAdapter(list[OpenAIFileObject])
@ -244,30 +308,41 @@ async def route_create_file(
5. Else -> use custom_llm_provider with files_settings
"""
# Handle custom storage backend
if target_storage and target_storage != "default":
explicit_storage: Final = target_storage if target_storage and target_storage != "default" else None
if explicit_storage == LITELLM_DB_STORAGE_BACKEND_NAME:
raise ProxyException(
message=(
f"target_storage={LITELLM_DB_STORAGE_BACKEND_NAME} is not a storage a caller can pick: LiteLLM "
"chooses it on its own for the batch input files of a model whose batches it runs itself, so "
"upload with purpose=batch and name that model instead of target_storage"
),
type="invalid_request_error",
param="target_storage",
code=400,
)
executed_model: Final = await _litellm_executed_batch_input_model(
llm_router, purpose, model, target_model_names_list, user_api_key_dict, explicit_storage
)
storage: Final = explicit_storage or (LITELLM_DB_STORAGE_BACKEND_NAME if executed_model is not None else None)
if storage is not None:
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
from litellm.proxy.openai_files_endpoints.storage_backend_service import (
StorageBackendFileService,
)
from litellm.proxy.proxy_server import prisma_client
# Extract file data
file_data: Final = extract_file_data(cast(Any, _create_file_request.get("file")))
# Use storage backend service to handle upload
file_object: Final = await StorageBackendFileService.upload_file_to_storage_backend(
file_data=file_data,
target_storage=target_storage,
target_model_names=target_model_names_list,
return await StorageBackendFileService.upload_file_to_storage_backend(
file_data=extract_file_data(cast(Any, _create_file_request.get("file"))),
target_storage=storage,
target_model_names=(executed_model,) if executed_model is not None else target_model_names_list,
purpose=purpose,
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
return file_object
# NEW: Handle model-based routing (no DB required)
if model is not None:
# Get credentials from model_list via router
@ -848,7 +923,7 @@ async def get_file_content(
# Check if file is stored in a storage backend (check DB)
if hasattr(managed_files_obj, "prisma_client") and getattr(managed_files_obj, "prisma_client", None):
prisma_client: Final = getattr(managed_files_obj, "prisma_client")
prisma_client: Final[PrismaClient] = getattr(managed_files_obj, "prisma_client")
db_file: Final = await ManagedFileRepository(prisma_client).table.find_first(
where={"unified_file_id": file_id}
)
@ -863,7 +938,7 @@ async def get_file_content(
try:
# Get storage backend (uses same env vars as callback)
storage_backend: Final = get_storage_backend(storage_backend_name)
storage_backend: Final = get_storage_backend(storage_backend_name, prisma_client=prisma_client)
file_content: Final = await storage_backend.download_file(storage_url)
# Return file content

40
litellm/proxy/openai_files_endpoints/storage_backend_service.py
View file

@ -7,15 +7,16 @@ storage backends (e.g., Azure Blob Storage) and managing associated metadata.
import base64
import time
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from typing import Any, Final, cast
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid as uuid_module
from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend
from litellm.llms.base_llm.files.storage_backend_factory import get_storage_backend
from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.types.llms.openai import OpenAIFileObject, OpenAIFilesPurpose
from litellm.types.utils import SpecialEnums
@ -35,21 +36,23 @@ class StorageBackendFileService:
async def upload_file_to_storage_backend(
file_data: Mapping[str, Any],
target_storage: str,
target_model_names: list[str],
target_model_names: Sequence[str],
purpose: OpenAIFilesPurpose,
proxy_logging_obj: ProxyLogging,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient | None = None,
) -> OpenAIFileObject:
"""
Upload a file to a storage backend and create a file object.
Args:
file_data: File data dictionary from extract_file_data()
target_storage: Storage backend name (e.g., "azure_storage")
target_storage: Storage backend name (e.g., "azure_storage", "litellm_db")
target_model_names: List of model names for managed files
purpose: File purpose (e.g., "user_data", "batch")
proxy_logging_obj: Proxy logging object for accessing hooks
user_api_key_dict: User API key authentication data
prisma_client: The proxy's database client, required by the "litellm_db" backend
Returns:
OpenAIFileObject: Created file object with storage metadata
@ -59,7 +62,7 @@ class StorageBackendFileService:
"""
# Get storage backend instance
try:
storage_backend: Final = get_storage_backend(target_storage)
storage_backend: Final = get_storage_backend(target_storage, prisma_client=prisma_client)
except ValueError as e:
raise ProxyException(
message=str(e),
@ -103,8 +106,9 @@ class StorageBackendFileService:
storage_url=storage_url,
)
# Store in managed files if target_model_names provided
if target_model_names:
if not target_model_names:
return file_object
try:
await StorageBackendFileService._store_in_managed_files(
file_object=file_object,
file_data=file_data,
@ -114,9 +118,25 @@ class StorageBackendFileService:
proxy_logging_obj=proxy_logging_obj,
user_api_key_dict=user_api_key_dict,
)
except Exception:
await StorageBackendFileService._discard_orphaned_content(storage_backend, storage_url, target_storage)
raise
return file_object
@staticmethod
async def _discard_orphaned_content(
storage_backend: BaseFileStorageBackend, storage_url: str, target_storage: str
) -> None:
try:
await storage_backend.delete_file(storage_url)
except Exception as e: # noqa: BLE001 # the metadata failure is what surfaces; a failed cleanup is only logged
verbose_proxy_logger.warning(
"Could not delete orphaned content at %s on %s after its metadata write failed: %s",
storage_url,
target_storage,
e,
)
@staticmethod
def _create_file_object_with_storage_metadata(
file_content: bytes,
@ -164,7 +184,7 @@ class StorageBackendFileService:
@staticmethod
def _create_unified_file_id(
file_type: str,
target_model_names: list[str],
target_model_names: Sequence[str],
file_id: str,
) -> str:
"""
@ -194,7 +214,7 @@ class StorageBackendFileService:
async def _store_in_managed_files(
file_object: OpenAIFileObject,
file_data: Mapping[str, Any],
target_model_names: list[str],
target_model_names: Sequence[str],
target_storage: str,
storage_url: str,
proxy_logging_obj: ProxyLogging,

5
litellm/proxy/proxy_cli.py
View file

@ -589,6 +589,11 @@ class ProxyInitializationHelpers:
gunicorn_options["certfile"] = ssl_certfile_path
gunicorn_options["keyfile"] = ssl_keyfile_path
# The master preloads the app and then forks every worker, so native routes are
# forbidden in it: their runtime threads would not survive the fork.
from litellm.rust_bridge.fork_guard import reserve_process_for_forking
reserve_process_for_forking("the gunicorn master")
start_query_engine_reaper()
StandaloneApplication(app=app, options=gunicorn_options).run() # Run gunicorn

111
litellm/proxy/proxy_server.py
View file

@ -4787,6 +4787,22 @@ def _swap_in_model_cost_map(new_model_cost_map: dict) -> int:
return adopt_model_cost_map(new_model_cost_map)
def _websearch_handler_params(stored: Mapping[str, object]) -> dict[str, object]:
"""
Translate stored web search interception settings into handler kwargs.
Drops ``enabled``, which gates the callback rather than configuring it, and
drops an ``enabled_providers`` that is not a non-empty list so the handler
applies its own default. An empty list otherwise matches no provider at all,
and a bare string is iterated one character at a time.
"""
params: Final = {key: value for key, value in stored.items() if key != "enabled"}
providers: Final = params.get("enabled_providers")
if not isinstance(providers, list) or not providers:
params.pop("enabled_providers", None)
return params
def should_load_db_object(object_type: str | SupportedDBObjectType) -> bool:
"""
Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
@ -4888,6 +4904,7 @@ class ProxyConfig:
def __init__(self) -> None:
self.config: Mapping[str, object] = MappingProxyType({})
self._last_semantic_filter_config: dict[str, object] | None = None
self._last_websearch_interception_config: dict[str, object] | None = None
self._last_hashicorp_vault_config: dict[str, object] | None = None
self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache
self._cyberark_boot_env: dict[str, str | None] | None = None # mutable-ok: deployment env snapshot, set once
@ -5117,6 +5134,15 @@ class ProxyConfig:
store.apply_db_row(cast(DbRow, section_name), wrote_section)
await invalidate_config_param(section_name)
def reject_config_owned_deletes(self, *, section_name: str, keys: tuple[str, ...]) -> None:
"""Refuse a delete of a setting the config file owns; unlike a write, the value never makes it allowed."""
store: Final = self._settings_stores.get(cast(Section, section_name))
if store is None:
return
owned: Final = tuple(sorted(key for key in keys if store.owned_by_config(key)))
if owned:
self._raise_config_owned(section_name=section_name, rejected=owned, store=store)
def reject_config_owned_writes(self, *, section_name: str, changed_keys: Mapping[str, JsonValue]) -> None:
"""Refuse a write to a setting the config file owns, rather than storing a value that never applies."""
store: Final = self._settings_stores.get(cast(Section, section_name))
@ -5125,6 +5151,9 @@ class ProxyConfig:
rejected: Final = store.rejected_writes(changed_keys)
if not rejected:
return
self._raise_config_owned(section_name=section_name, rejected=rejected, store=store)
def _raise_config_owned(self, *, section_name: str, rejected: tuple[str, ...], store: SettingsStore) -> None:
subject: Final = (
f"key '{rejected[0]}' is" if len(rejected) == 1 else f"keys {', '.join(repr(key) for key in rejected)} are"
)
@ -7697,6 +7726,9 @@ class ProxyConfig:
if self._should_load_db_object(object_type="semantic_filter_settings"):
await self._init_semantic_filter_settings_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type=SupportedDBObjectType.WEBSEARCH_INTERCEPTION_SETTINGS):
await self.init_websearch_interception_settings_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="config_overrides"):
await self._init_hashicorp_vault_config_override(prisma_client=prisma_client)
await self._init_cyberark_config_override(prisma_client=prisma_client)
@ -7775,6 +7807,66 @@ class ProxyConfig:
except Exception as e:
verbose_proxy_logger.exception("Error initializing semantic filter settings from DB: %s", e)
async def init_websearch_interception_settings_in_db(self, prisma_client: PrismaClient):
"""
Initialize web search interception settings from database.
Called periodically (approximately every 10 seconds) by background task to hot-reload settings across all pods.
"""
import json
import litellm
from litellm.integrations.websearch_interception.handler import (
WebSearchInterceptionLogger,
)
try:
config_record: Final = await get_config_param(prisma_client, "litellm_settings")
if config_record is None or config_record.param_value is None:
return
litellm_settings = config_record.param_value
if isinstance(litellm_settings, str):
litellm_settings = json.loads(litellm_settings)
websearch_config: Final = litellm_settings.get("websearch_interception_params", None)
if not isinstance(websearch_config, Mapping):
return
if "enabled" not in websearch_config and self._last_websearch_interception_config is None:
verbose_proxy_logger.debug(
"Web search interception: stored settings carry no 'enabled' flag and none were applied "
"before, so litellm_settings.callbacks keeps ownership of the callback."
)
return
enabled: Final = bool(coerce_bool(websearch_config.get("enabled", True)))
registered: Final = bool(
litellm.logging_callback_manager.get_custom_loggers_for_type(WebSearchInterceptionLogger)
)
if self._last_websearch_interception_config == websearch_config and registered == enabled:
return
replacement: Final = (
WebSearchInterceptionLogger.from_config_yaml(_websearch_handler_params(websearch_config))
if enabled
else None
)
litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, WebSearchInterceptionLogger)
if replacement is not None:
litellm.logging_callback_manager.add_litellm_callback(replacement)
verbose_proxy_logger.info("Web search interception reinitialized from DB")
else:
verbose_proxy_logger.info("Web search interception disabled")
self._last_websearch_interception_config = dict(websearch_config)
except Exception as e:
verbose_proxy_logger.exception("Error initializing web search interception settings from DB: %s", e)
async def _init_sso_settings_in_db(self, prisma_client: PrismaClient):
"""
Initialize SSO settings from database into the router on startup.
@ -9684,10 +9776,12 @@ class ProxyStartupEvent:
user_api_key_cache: UserApiKeyCache,
):
"""Initialize JWT auth on startup"""
if general_settings.get("litellm_jwtauth", None) is not None:
for k, v in general_settings["litellm_jwtauth"].items():
if isinstance(v, str) and v.startswith("os.environ/"):
general_settings["litellm_jwtauth"][k] = get_secret(v)
declared_jwtauth: Final = general_settings.get("litellm_jwtauth", None)
if declared_jwtauth is not None:
resolved_jwtauth: Final = {
key: (get_secret(value) if isinstance(value, str) and value.startswith("os.environ/") else value)
for key, value in declared_jwtauth.items()
}
# ``user_config_file_path`` is set by ``ProxyConfig._get_config_from_file``
# during startup. Threading it through lets an operator-
# configured ``custom_validate: s3://...`` resolve through
@ -9695,7 +9789,7 @@ class ProxyStartupEvent:
# file context) hit the gate and refuse remote loads.
litellm_jwtauth = LiteLLM_JWTAuth(
config_file_path=user_config_file_path,
**general_settings["litellm_jwtauth"],
**resolved_jwtauth,
)
else:
litellm_jwtauth = LiteLLM_JWTAuth()
@ -17665,9 +17759,12 @@ async def get_config_general_settings(
detail={"error": f"Field name={field_name} is not set"},
)
declared: Final = (
settings.config_value(field_name) if settings.owned_by_config(field_name) else settings[field_name]
)
field_value = _redact_general_setting_value(
field_name,
settings[field_name],
declared,
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN,
)
if field_name == "plugins" and isinstance(field_value, list):
@ -18041,6 +18138,8 @@ async def delete_config_general_settings(
detail={"error": f"Invalid field={data.field_name} passed in."},
)
proxy_config.reject_config_owned_deletes(section_name="general_settings", keys=(data.field_name,))
## get general settings from db
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
where={"param_name": "general_settings"}

9
litellm/proxy/public_endpoints/public_endpoints.py
View file

@ -23,6 +23,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.utils import get_custom_url
from litellm.repositories.table_repositories import ClaudeCodePluginRepository
from litellm.router_strategy.complexity_router.fuse_presets import FusePresetCatalog, get_fuse_presets
from litellm.types.agents import AgentCard
from litellm.types.mcp import MCPPublicServer
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
@ -424,6 +425,14 @@ async def get_complexity_scorer_defaults() -> ComplexityScorerDefaults:
)
@router.get(
"/public/complexity_router/fuse_presets",
response_model=FusePresetCatalog,
)
async def get_public_fuse_presets() -> FusePresetCatalog:
return get_fuse_presets()
@router.get(
"/public/litellm_model_cost_map",
tags=["public", "model management"],

40
litellm/proxy/schema.prisma
View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID
@ -1545,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession {
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
model LiteLLM_AutoRouterBaselineComparison {
scope String @id
api_key String
session_id String
router_name String
initial_equivalent Boolean
revision BigInt @default(0)
published_revision BigInt @default(0)
history String?
attempted_at DateTime?
retired Boolean @default(false)
updated_at DateTime @default(now())
@@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope")
@@index([updated_at], map: "idx_autorouter_baseline_updated")
}
model LiteLLM_AutoRouterBaselineObservation {
request_id String @id
scope String
started_at Float
revision BigInt
data String
publication String?
conflicted Boolean @default(false)
@@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order")
@@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision")
}
model LiteLLM_AutoRouterSession {
api_key String
session_id String
@ -1571,6 +1607,10 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")

348
litellm/proxy/spend_tracking/baseline_accounting.py Normal file
View file

@ -0,0 +1,348 @@
"""Pure, chronological cache accounting for the recorded baseline comparison.
Observation collection, pricing and durable publication belong to their existing
owners. Replaying these values in event order is independent of callback order.
"""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from itertools import groupby
from math import isfinite
from types import MappingProxyType
from typing import Final, Literal
from pydantic import BaseModel, ConfigDict, Field
from litellm.llms.anthropic.prompt_cache_prediction import CountedBreakpoint, CountedPromptCachePlan
from litellm.types.utils import CacheCreationTokenDetails, PromptTokensDetailsWrapper, Usage
MAX_CACHE_TTL: Final = 3600
MAX_CACHE_ENTRIES: Final = 1024
class BaselineObservation(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
version: Literal[3] = 3
request_id: str = Field(min_length=1)
started_at: float = Field(allow_inf_nan=False, ge=0)
available_at: float = Field(allow_inf_nan=False, ge=0)
outcome: Literal["complete", "uncertain", "response_cache"]
baseline_equivalent: bool
usage: Usage | None = None
plan: CountedPromptCachePlan | None = None
minimum_cache_tokens: int = Field(default=0, ge=0)
reason: str | None = None
@dataclass(frozen=True, slots=True)
class BaselineEstimate:
request_id: str
reason: str
provenance: Literal["observed_identical", "modeled"] | None = None
usage: Usage | None = None
@dataclass(frozen=True, slots=True)
class CacheEntry:
fingerprint: str
content_fingerprint: str
tokens: int
ttl_seconds: int
available_at: float
expires_at: float
uncertain: bool = False
@dataclass(frozen=True, slots=True)
class BaselineHistory:
first_at: float | None = None
last_at: float | None = None
equivalent: bool = True
uncertain_before: float = 0.0
entries: tuple[CacheEntry, ...] = ()
blocked_until: float = 0.0
def _complete_usage(usage: Usage | None) -> bool:
if usage is None or usage.prompt_tokens < 0 or usage.completion_tokens < 0:
return False
details: Final = usage.prompt_tokens_details
if details is None:
return False
values: Final = (details.text_tokens, details.cached_tokens, details.cache_creation_tokens)
if any(value is None or value < 0 for value in values):
return False
split: Final = details.cache_creation_token_details
writes: Final = details.cache_creation_tokens or 0
return (
usage.total_tokens == usage.prompt_tokens + usage.completion_tokens
and sum(value or 0 for value in values) == usage.prompt_tokens
and (
writes == 0
or (
split is not None
and split.ephemeral_5m_input_tokens is not None
and split.ephemeral_1h_input_tokens is not None
and min(split.ephemeral_5m_input_tokens, split.ephemeral_1h_input_tokens) >= 0
and split.ephemeral_5m_input_tokens + split.ephemeral_1h_input_tokens == writes
)
)
)
def _valid_plan(plan: CountedPromptCachePlan | None) -> bool:
if plan is None or plan.total_tokens < 0 or len(plan.breakpoints) > 4:
return False
return all(
marker.fingerprint
and marker.content_fingerprint
and marker.fingerprint in marker.lookback_fingerprints
and marker.content_fingerprint in marker.lookback_content_fingerprints
and marker.ttl_seconds in (300, 3600)
and 0 <= marker.prefix_tokens <= plan.total_tokens
for marker in plan.breakpoints
) and all(
left.prefix_tokens <= right.prefix_tokens and left.ttl_seconds >= right.ttl_seconds
for left, right in zip(plan.breakpoints, plan.breakpoints[1:])
)
def _markers(observation: BaselineObservation) -> tuple[CountedBreakpoint, ...]:
return (
tuple(
marker
for marker in observation.plan.breakpoints
if marker.prefix_tokens >= observation.minimum_cache_tokens
)
if observation.plan is not None
else ()
)
def _matches(entry: CacheEntry, markers: tuple[CountedBreakpoint, ...], started: float) -> bool:
return entry.available_at <= started < entry.expires_at and any(
entry.fingerprint in marker.lookback_fingerprints
and entry.tokens <= marker.prefix_tokens
and entry.ttl_seconds == marker.ttl_seconds
for marker in markers
)
def _ambiguous(entry: CacheEntry, markers: tuple[CountedBreakpoint, ...], started: float) -> bool:
return entry.available_at <= started < entry.expires_at and any(
entry.content_fingerprint in marker.lookback_content_fingerprints
and (entry.uncertain or entry.ttl_seconds != marker.ttl_seconds)
for marker in markers
)
def _usage_with_cache(usage: Usage, total: int, read: int, write_5m: int, write_1h: int) -> Usage:
writes: Final = write_5m + write_1h
original_details: Final = usage.prompt_tokens_details or PromptTokensDetailsWrapper()
details: Final = original_details.model_copy(
deep=True,
update=MappingProxyType(
{
"text_tokens": total - read - writes,
"cached_tokens": read,
"cache_creation_tokens": writes,
"cache_write_tokens": writes,
"cache_creation_token_details": CacheCreationTokenDetails(
ephemeral_5m_input_tokens=write_5m,
ephemeral_1h_input_tokens=write_1h,
),
}
),
)
return Usage.model_validate(
{ # mutable-ok: Usage only runs its normalizing constructor for a plain dictionary
**usage.model_dump(),
"prompt_tokens": total,
"total_tokens": total + usage.completion_tokens,
"prompt_tokens_details": details,
"cache_read_input_tokens": read,
"cache_creation_input_tokens": writes,
},
)
def _estimate(history: BaselineHistory, observation: BaselineObservation, equivalent: bool) -> BaselineEstimate:
if observation.outcome != "complete" or not _complete_usage(observation.usage):
return BaselineEstimate(observation.request_id, observation.reason or observation.outcome)
usage: Final = observation.usage
if usage is None:
return BaselineEstimate(observation.request_id, "missing_usage")
if equivalent and observation.baseline_equivalent:
return BaselineEstimate(
observation.request_id, "identical_baseline_path", "observed_identical", usage.model_copy(deep=True)
)
if observation.started_at < history.blocked_until:
return BaselineEstimate(observation.request_id, "concurrent_uncertainty")
plan: Final = observation.plan
if not _valid_plan(plan) or plan is None:
return BaselineEstimate(observation.request_id, observation.reason or "unsupported_cache_plan")
markers: Final = _markers(observation)
if any(_ambiguous(entry, markers, observation.started_at) for entry in history.entries):
return BaselineEstimate(observation.request_id, "cache_ttl_changed")
read: Final = max(
(
entry.tokens
for entry in history.entries
if not entry.uncertain and _matches(entry, markers, observation.started_at)
),
default=0,
)
end: Final = markers[-1].prefix_tokens if markers else 0
if read < end and observation.started_at < history.uncertain_before + max(marker.ttl_seconds for marker in markers):
return BaselineEstimate(observation.request_id, "history_unavailable")
one_hour: Final = max(
(marker.prefix_tokens for marker in markers if marker.ttl_seconds == 3600 and marker.prefix_tokens > read),
default=read,
)
expired: Final = any(
entry.expires_at <= observation.started_at
and any(entry.fingerprint in marker.lookback_fingerprints for marker in markers)
for entry in history.entries
)
reason: Final = (
"cache_prefix_available"
if read
else "cache_prefix_expired"
if expired
else "cache_prefix_cold"
if markers
else "below_cache_minimum"
if plan.breakpoints
else "no_cache_breakpoints"
)
return BaselineEstimate(
observation.request_id,
reason,
"modeled",
_usage_with_cache(usage, plan.total_tokens, read, end - one_hour, one_hour - read),
)
def _writes(history: BaselineHistory, observation: BaselineObservation) -> tuple[CacheEntry, ...]:
if (
observation.outcome != "complete"
or observation.started_at < history.blocked_until
or not _complete_usage(observation.usage)
or not _valid_plan(observation.plan)
):
return ()
markers: Final = _markers(observation)
ambiguous: Final = tuple(entry for entry in history.entries if _ambiguous(entry, markers, observation.started_at))
hit: Final = (
max(
(
entry
for entry in history.entries
if not entry.uncertain and _matches(entry, markers, observation.started_at)
),
key=lambda entry: entry.tokens,
default=None,
)
if not ambiguous
else None
)
refresh: Final = (
(
CacheEntry(
hit.fingerprint,
hit.content_fingerprint,
hit.tokens,
hit.ttl_seconds,
observation.available_at,
observation.started_at + hit.ttl_seconds,
),
)
if hit is not None and all(marker.fingerprint != hit.fingerprint for marker in markers)
else ()
)
return (
*refresh,
*(
CacheEntry(
marker.fingerprint,
marker.content_fingerprint,
marker.prefix_tokens,
marker.ttl_seconds,
observation.available_at,
observation.started_at + max((marker.ttl_seconds, *(entry.ttl_seconds for entry in ambiguous))),
uncertain=bool(ambiguous),
)
for marker in markers
),
)
def _entry_key(entry: CacheEntry) -> tuple[str, str, int, int, bool]:
return entry.fingerprint, entry.content_fingerprint, entry.tokens, entry.ttl_seconds, entry.uncertain
def _compact_entries(entries: tuple[CacheEntry, ...], started: float) -> tuple[CacheEntry, ...]:
ordered: Final = sorted((entry for entry in entries if entry.expires_at >= started - MAX_CACHE_TTL), key=_entry_key)
return tuple(
retained
for _, values in groupby(ordered, key=_entry_key)
for group in (tuple(values),)
for retained in (
max(
(entry for entry in group if entry.available_at <= started),
key=lambda entry: entry.expires_at,
default=None,
),
*(entry for entry in group if entry.available_at > started),
)
if retained is not None
)
def advance_baseline_history(
history: BaselineHistory,
simultaneous: Sequence[BaselineObservation],
) -> tuple[BaselineHistory, tuple[BaselineEstimate, ...]]:
"""Apply one request-start timestamp; ties cannot manufacture initial equality.
The storage owner groups and orders observations before calling this function.
Equal timestamps are evaluated against the same preceding cache snapshot.
"""
if not simultaneous:
return history, ()
started: Final = simultaneous[0].started_at
valid_order: Final = (
isfinite(started)
and all(item.started_at == started and item.available_at >= started for item in simultaneous)
and (history.last_at is None or started > history.last_at)
)
if not valid_order:
return history, tuple(BaselineEstimate(item.request_id, "invalid_observation_order") for item in simultaneous)
first: Final = started if history.first_at is None else history.first_at
uncertain: Final = max(history.uncertain_before, first)
relevant: Final = tuple(item for item in simultaneous if item.outcome != "response_cache")
equivalent: Final = history.equivalent and all(item.baseline_equivalent for item in relevant)
before: Final = BaselineHistory(
first, history.last_at, equivalent, uncertain, history.entries, history.blocked_until
)
estimates: Final = tuple(_estimate(before, item, equivalent) for item in simultaneous)
invalidated: Final = any(
item.outcome != "complete" or not _complete_usage(item.usage) or not _valid_plan(item.plan) for item in relevant
)
blocked: Final = max((history.blocked_until, *(item.available_at for item in relevant if invalidated)))
entries: Final = _compact_entries(
() if invalidated else (*history.entries, *(entry for item in relevant for entry in _writes(before, item))),
started,
)
overflow: Final = len(entries) > MAX_CACHE_ENTRIES
return BaselineHistory(
first_at=first,
last_at=started,
equivalent=equivalent,
uncertain_before=max(started, blocked) if invalidated or overflow else uncertain,
entries=() if overflow else entries,
blocked_until=blocked,
), estimates

266
litellm/proxy/spend_tracking/savings.py
View file

@ -10,7 +10,11 @@ have been aggregated across models.
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Final, NamedTuple
from math import isclose, isfinite
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, NamedTuple
from pydantic import BaseModel, ConfigDict, Field
import litellm
from litellm._logging import verbose_proxy_logger
@ -65,7 +69,7 @@ def _resolve_model(model: str | None, custom_llm_provider: str | None) -> _Model
return None
try:
resolved_model, provider, _, _ = litellm.get_llm_provider(model=model, custom_llm_provider=custom_llm_provider)
except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to zero savings
except Exception as e: # noqa: BLE001 # get_llm_provider raises for unroutable names; degrade to an unavailable estimate
verbose_proxy_logger.debug(
"savings: cannot resolve provider for model=%s custom_llm_provider=%s (%s)", model, custom_llm_provider, e
)
@ -118,6 +122,68 @@ class PricingBasis(NamedTuple):
_STANDARD_RATES: Final = PricingBasis()
class BaselineCostSnapshot(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
model: str
provider: str
prices: ModelInfo | None
basis: PricingBasis = _STANDARD_RATES
actual_spend: float = Field(allow_inf_nan=False, ge=0)
actual_token_cost: float | None = Field(default=None, allow_inf_nan=False, ge=0)
classifier_cost: float = Field(default=0.0, allow_inf_nan=False, ge=0)
def baseline_cost_snapshot(
model: str,
prices: ModelInfo | None,
actual_spend: float,
cost_breakdown: Mapping[str, object] | None,
routing_decision: Mapping[str, object] | None,
) -> BaselineCostSnapshot:
return BaselineCostSnapshot(
model=model,
provider="anthropic",
prices=prices,
actual_spend=actual_spend,
basis=_pricing_basis(cost_breakdown),
actual_token_cost=_recorded_token_cost(cost_breakdown),
classifier_cost=classifier_cost_from_decision(routing_decision) or 0.0,
)
class BaselineCosts(NamedTuple):
actual: float
baseline: float
@property
def savings(self) -> float:
return self.baseline - self.actual
def price_baseline_comparison(
snapshot: BaselineCostSnapshot,
baseline_usage: Usage | None,
provenance: Literal["observed_identical", "modeled"] | None,
) -> BaselineCosts | None:
if baseline_usage is None or provenance is None:
return None
actual: Final = snapshot.actual_spend + snapshot.classifier_cost
if provenance == "observed_identical":
return BaselineCosts(actual=actual, baseline=snapshot.actual_spend)
if snapshot.prices is None or snapshot.actual_token_cost is None:
return None
token_cost: Final = _cost_of_usage(
_ModelIdentity(snapshot.model, snapshot.provider), baseline_usage, snapshot.prices, snapshot.basis
)
if token_cost is None or not isfinite(token_cost) or token_cost < 0:
return None
baseline: Final = snapshot.actual_spend + token_cost - snapshot.actual_token_cost
if not isfinite(baseline) or baseline < 0:
return None
return BaselineCosts(actual=actual, baseline=baseline)
def _pricing_basis(cost_breakdown: Mapping[str, object] | None) -> PricingBasis:
"""The basis recorded on a request, defaulting to standard rates when absent.
@ -225,56 +291,16 @@ def _baseline_cache_rate_keys(baseline_info: ModelInfo | None) -> tuple[bool, bo
)
def _baseline_usage(usage: Usage, conversation_continuing: bool, baseline_info: ModelInfo | None = None) -> Usage:
"""The same request as a single-model baseline would have met it.
The baseline is one model serving every turn, so whether it had this prompt cached
is simply whether the conversation was already underway. On a continuing
conversation it wrote the prompt on an earlier turn and would only read it now, so
the cache tokens move into the read bucket and whatever this request paid to write
counts against the saving; that write is what switching models costs.
On a conversation's first turn nothing was cached anywhere, for any model. The
baseline would have written the same prompt, so the cache buckets stay where they are
and both arms carry the write at their own rates, unless the baseline has no rate for
a bucket, in which case those tokens are its plain input. Charging the write to this case
too, which is all a single rollup row can support, understates a first turn to a
few percent of its value and can render a profitable route as a loss.
A continuing turn that mostly read from cache is the third case: the selected model
was already warm, so it is the one that has been serving this conversation and the
baseline's cache holds exactly what its does. The tokens written are the turn's own
growth, new to every model, and the baseline would have paid to write them too.
Moving them would forgive the baseline a write it really owes and shrink the
reported saving. "Mostly read" rather than "read anything" on purpose: a switch onto
a model holding a small prefix of this prompt still writes most of it, and must keep
counting that write against the saving.
Only the cache buckets move. Every other field the request was priced on travels
through untouched, audio and image and video counts among them, because the baseline
is this same request served by a model that happened to be warm; naming the fields to
keep instead would price the baseline on a request that never ran, and would go stale
the next time a priced field is added.
"""
def _baseline_usage(usage: Usage, baseline_info: ModelInfo | None = None) -> Usage:
cache_read, cache_creation = _cache_token_split(usage)
details: Final = usage.prompt_tokens_details
if details is None or (cache_read <= 0 and cache_creation <= 0):
return usage
# The tokens this request paid to write move into the cached count and the creation
# charge is dropped: on one model that cache was already warm, so the baseline would
# have read them rather than paying to create them. The 5m/1h breakdown goes with
# them; left behind it re-charges the write.
warm: Final = conversation_continuing and cache_creation > 0 and cache_read <= cache_creation
reads = cache_read + cache_creation if warm else cache_read
writes = 0 if warm else cache_creation
prices_reads, prices_writes = _baseline_cache_rate_keys(baseline_info)
reads = reads if prices_reads else 0
writes = writes if prices_writes else 0
reads: Final = cache_read if prices_reads else 0
writes: Final = cache_creation if prices_writes else 0
if (reads, writes) == (cache_read, cache_creation):
return usage
other_modalities: Final = sum(
(getattr(details, field, 0) or 0) for field in ("audio_tokens", "image_tokens", "video_tokens")
)
@ -309,64 +335,47 @@ def compute_autorouter_savings(
cost_breakdown: Mapping[str, object] | None = None,
baseline_deployment_id: str | None = None,
selected_deployment_id: str | None = None,
) -> float:
"""Net dollars the router saved, or cost, by serving this request on ``selected_model``.
Signed on purpose. Switching models leaves the new one with a cold cache, so the
request pays a cache-creation charge that staying on one model would not have
incurred; when that charge outweighs the cheaper rates, routing lost money and the
dashboard has to be able to say so. Zero when both sides resolve to the same
deployment, or when either cannot be resolved or priced.
Only one side of this subtraction is a counterfactual. What the request cost on the
model that served it is a number the operator was actually billed, and the cost
calculator already wrote it down, so ``cost_breakdown`` is read rather than
re-derived. Recomputing it means restating every pricing dimension the biller
applied, and each one omitted is a silent disagreement with the ``spend`` column
beside it; a request billed at a priority tier recomputed at standard rates reads as
half its real cost.
The baseline has no such record, since it never ran, so it is priced through the same
cost engine on the basis the biller used for this request. An operator running that
one model instead of the router would have sent this request to the same tier and the
same region, because both are properties of the request and the deployment's
contract, not of which model the router happened to pick.
``conversation_continuing`` says whether the baseline would already have had this
prompt cached. It defaults to True because that is the conservative reading: a
request whose shape the router could not determine is charged the write and
under-claims rather than inflating a savings figure.
"""
# No provider argument for the baseline on purpose: it arrives from the routing
# metadata as a single self-describing string, already qualified by the auto-router,
# so there is no second field that could disagree with it.
baseline_usage: Usage | None = None,
baseline_provenance: Literal["observed_initial", "modeled"] | None = None,
) -> float | None:
"""Price established baseline usage; conversation shape cannot establish cache warmth."""
baseline: Final = _resolve_model(baseline_model, None)
selected: Final = _resolve_model(selected_model, selected_provider)
if baseline is None or selected is None:
return 0.0
same_target: Final = (
baseline_deployment_id == selected_deployment_id
if baseline_deployment_id and selected_deployment_id
else baseline == selected
)
if same_target:
return 0.0
return None
if baseline_usage is None and any(_cache_token_split(usage)):
return None
basis: Final = _pricing_basis(cost_breakdown)
effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline)
modeled_usage: Final = baseline_usage if baseline_usage is not None else usage
baseline_cost: Final = _cost_of_usage(
baseline,
_baseline_usage(usage, conversation_continuing, effective_baseline_info),
effective_baseline_info,
basis,
baseline, _baseline_usage(modeled_usage, effective_baseline_info), effective_baseline_info, basis
)
recorded_selected_cost: Final = _recorded_token_cost(cost_breakdown)
selected_cost: Final = (
recorded_selected_cost
if recorded_selected_cost is not None
else _cost_of_usage(selected, usage, selected_info, basis)
)
# Falls back to pricing the request only when the biller recorded nothing, which is
# every row written before the breakdown carried its basis.
selected_cost = _recorded_token_cost(cost_breakdown)
if selected_cost is None:
selected_cost = _cost_of_usage(selected, usage, selected_info, basis)
if baseline_cost is None or selected_cost is None:
return 0.0
return baseline_cost - selected_cost
return None
if baseline_provenance == "observed_initial":
same_prices: Final = effective_baseline_info == (
selected_info if selected_info is not None else _model_info(selected)
)
equivalent: Final = (
baseline_usage is not None
and baseline_usage == usage
and baseline == selected
and bool(baseline_deployment_id)
and baseline_deployment_id == selected_deployment_id
and same_prices
and recorded_selected_cost is not None
and isclose(baseline_cost, recorded_selected_cost, rel_tol=1e-9, abs_tol=1e-12)
)
return 0.0 if equivalent else None
difference: Final = baseline_cost - selected_cost
return difference if isfinite(difference) else None
def _usage_from_spend_log(usage_object: Mapping[str, object] | None) -> Usage | None:
@ -463,11 +472,23 @@ def _proxy_llm_router() -> "Router | None":
def _numeric_savings(value: object) -> float | None:
"""``value`` as a recorded savings figure, or ``None`` when it is not one."""
if isinstance(value, bool) or not isinstance(value, (int, float)):
if isinstance(value, bool) or not isinstance(value, (int, float)) or not isfinite(value):
return None
return float(value)
def recorded_estimated_autorouter_savings(metadata: Mapping[str, object]) -> float | None:
estimate: Final = metadata.get("autorouter_savings_estimate")
if (
not isinstance(estimate, Mapping)
or type(estimate.get("version")) is not int
or estimate.get("version") not in (1, 2, 3)
or estimate.get("status") != "estimated"
):
return None
return _numeric_savings(metadata.get("autorouter_savings"))
def classifier_cost_from_decision(routing_decision: Mapping[str, object] | None) -> float | None:
"""The LLM-classifier cost a routing decision recorded, or ``None`` when it holds none.
@ -490,22 +511,10 @@ def autorouter_savings_for_request(
model_id: str | None = None,
llm_router: "Callable[[], Router | None] | None" = None,
cost_breakdown: Mapping[str, object] | None = None,
baseline_usage: Usage | None = None,
baseline_provenance: Literal["observed_initial", "modeled"] | None = None,
) -> float | None:
"""Auto-router savings for one request, net of the classifier call that routed it,
or ``None`` when the driver is off.
``None`` and ``0.0`` are different facts: ``None`` means this request cannot carry a
figure at all (no routing decision, no baseline, unusable usage), while ``0.0`` is a
real figure for a routed request whose baseline resolved to the served deployment.
Never raises: pricing failures inside degrade to zero, and the driver-off cases
return ``None``, so this is safe on the logging path where a raise would fail the
request's logging.
The classifier deduction lives here, at the figure's one computation owner, rather
than in any reader: the stamped ``autorouter_savings`` is then already net, so the
session rollup, the daily tables and every logging consumer agree without each
re-deriving the deduction, and the recorded-figure-wins path cannot deduct twice.
"""
"""Return net savings for established usage, or None when the estimate is unavailable."""
usage: Final = _usage_from_spend_log(usage_object)
if usage is None or not model:
return None
@ -522,15 +531,16 @@ def autorouter_savings_for_request(
selected_model=model,
selected_provider=custom_llm_provider,
usage=usage,
# Absent means the router never recorded a shape, which is the conservative
# reading: charge the cache write rather than claim a first turn's saving.
conversation_continuing=decision.get("conversation_continuing") is not False,
selected_info=_effective_model_info(router_instance, model_id, model or ""),
baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""),
cost_breakdown=cost_breakdown,
baseline_deployment_id=baseline_id,
selected_deployment_id=model_id,
baseline_usage=baseline_usage,
baseline_provenance=baseline_provenance,
)
if gross is None:
return None
classifier_cost: Final = classifier_cost_from_decision(decision)
return gross if classifier_cost is None else gross - classifier_cost
@ -542,6 +552,8 @@ def autorouter_savings_for_logging_payload(
model_id: str | None,
usage_object: Mapping[str, object] | None,
cost_breakdown: Mapping[str, object] | None,
baseline_usage: Usage | None = None,
baseline_provenance: Literal["observed_initial", "modeled"] | None = None,
) -> float | None:
"""The figure the logging payload records for a request, or ``None`` when none should be.
@ -561,6 +573,8 @@ def autorouter_savings_for_logging_payload(
model_id=model_id,
llm_router=_proxy_llm_router,
cost_breakdown=cost_breakdown,
baseline_usage=baseline_usage,
baseline_provenance=baseline_provenance,
)
@ -575,6 +589,7 @@ def compute_savings_spend(
llm_router: "Callable[[], Router | None] | None" = None,
cost_breakdown: Mapping[str, object] | None = None,
recorded_autorouter_savings: object = None,
recorded_autorouter_savings_estimate: Mapping[str, object] | None = None,
billed_at: datetime | str | None = None,
) -> SavingsSpend:
"""
@ -604,11 +619,9 @@ def compute_savings_spend(
figure is normally the smaller of the two, being a subset of the same requests, but
not always: a request that only writes cache and never reads it has negative net
savings, and dropping such a request from the attributed figure can lift it above
the total. Auto-router savings compare the
served ``model`` against the counterfactual baseline the router recorded on
its ``routing_decision``, and are zero unless the two differ. That record
also says whether the conversation was already underway, which is what tells
a mid-conversation switch from a first turn.
the total. Auto-router savings compare established baseline usage against the
recorded selected-model cost. Versioned unknown estimates contribute no dollars
to this subtotal and are excluded from the separately reported coverage cohort.
``llm_router`` is passed as a provider rather than a router because every spend write
calls this and only auto-routed ones need one, so looking it up eagerly at the call
@ -653,10 +666,21 @@ def compute_savings_spend(
# The figure the logging path recorded wins, before the usage gate on purpose: a row
# whose usage no longer parses still carries the number computed when it did.
recorded_savings: Final = _numeric_savings(recorded_autorouter_savings)
recorded_savings: Final = (
recorded_estimated_autorouter_savings(
MappingProxyType(
{
"autorouter_savings": recorded_autorouter_savings,
"autorouter_savings_estimate": recorded_autorouter_savings_estimate,
}
)
)
if recorded_autorouter_savings_estimate is not None
else _numeric_savings(recorded_autorouter_savings)
)
autorouter: Final = (
recorded_savings
if recorded_savings is not None
if recorded_savings is not None or recorded_autorouter_savings_estimate is not None
else autorouter_savings_for_request(
model=model,
custom_llm_provider=custom_llm_provider,

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

@ -9,7 +9,7 @@ from functools import reduce
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable
from pydantic import BaseModel
from pydantic import BaseModel, JsonValue
import litellm
from litellm._logging import verbose_proxy_logger
@ -137,7 +137,15 @@ def _get_router_metadata_for_spend_log(
)
_STAMPED_METADATA_KEYS: Final = frozenset(("router_metadata", "azure_spillover"))
_STAMPED_METADATA_KEYS: Final = frozenset(
(
"router_metadata",
"azure_spillover",
"autorouter_savings",
"autorouter_savings_estimate",
"autorouter_baseline_observation",
)
)
def _get_spend_logs_metadata(
@ -156,6 +164,8 @@ def _get_spend_logs_metadata(
cost_breakdown: CostBreakdown | None = None,
litellm_call_id: str | None = None,
autorouter_savings: float | None = None,
autorouter_savings_estimate: Mapping[str, JsonValue] | None = None,
autorouter_baseline_observation: str | None = None,
router_metadata: SpendLogsRouterMetadata | None = None,
azure_spillover: AzureSpillover | None = None,
) -> SpendLogsMetadata:
@ -196,6 +206,8 @@ def _get_spend_logs_metadata(
cost_breakdown=None,
compression_savings=None,
autorouter_savings=autorouter_savings,
autorouter_savings_estimate=autorouter_savings_estimate,
autorouter_baseline_observation=autorouter_baseline_observation,
litellm_gateway_injected_cache=None,
litellm_call_id=litellm_call_id,
router_metadata=router_metadata,
@ -207,7 +219,12 @@ def _get_spend_logs_metadata(
# Filter the metadata dictionary to include only the specified keys
clean_metadata: Final = SpendLogsMetadata(
**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS},
**MappingProxyType(
{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS}
),
autorouter_savings=autorouter_savings,
autorouter_savings_estimate=autorouter_savings_estimate,
autorouter_baseline_observation=autorouter_baseline_observation,
router_metadata=router_metadata,
azure_spillover=azure_spillover,
)
@ -231,7 +248,6 @@ def _get_spend_logs_metadata(
clean_metadata["cold_storage_object_key"] = cold_storage_object_key
clean_metadata["litellm_overhead_time_ms"] = litellm_overhead_time_ms
clean_metadata["cost_breakdown"] = cost_breakdown
clean_metadata["autorouter_savings"] = autorouter_savings
clean_metadata["litellm_call_id"] = litellm_call_id
return clean_metadata
@ -660,6 +676,16 @@ def get_logging_payload(
autorouter_savings=(
standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None
),
autorouter_savings_estimate=(
standard_logging_payload.get("autorouter_savings_estimate")
if standard_logging_payload is not None
else None
),
autorouter_baseline_observation=(
standard_logging_payload.get("autorouter_baseline_observation")
if standard_logging_payload is not None
else None
),
litellm_call_id=litellm_call_id,
router_metadata=_get_router_metadata_for_spend_log(
metadata=metadata,

146
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -6,6 +6,7 @@ from collections import Counter
from collections.abc import Mapping, MutableMapping, Sequence
from types import MappingProxyType
from typing import (
Annotated,
Final,
NamedTuple,
Protocol,
@ -477,6 +478,72 @@ class MCPToolSearchSettingsResponse(SettingsResponse):
"""Response model for native MCP tool search settings"""
class WebSearchInterceptionSettings(BaseModel):
"""Configuration for server-side web search interception"""
enabled: bool = Field(
default=False,
description="Serve web search tool calls from a configured search tool instead of passing them upstream",
)
enabled_providers: list[str] = Field(
default_factory=list,
description="LLM providers to intercept for (e.g. 'bedrock', 'vertex_ai'). Empty intercepts Bedrock only.",
)
search_tool_name: str | None = Field(
default=None,
description="Name of the configured search tool to run searches through. Empty uses the first one available.",
)
max_agentic_loops: int | None = Field(
default=None,
ge=1,
description="How many follow-up model calls one intercepted request may chain. Empty applies the default of 3.",
)
class WebSearchInterceptionSettingsResponse(SettingsResponse):
"""Response model for web search interception settings"""
def _with_websearch_enabled_resolved(config: Mapping[str, object]) -> dict[str, object]:
"""
Answer with the stored flag when there is one, and only otherwise with what
this process is running.
A stored flag is the cluster's own answer, so it is the same on every pod and
is safe for the page to send back on save. Deriving the answer from this
process instead would report off on a pod that has not polled yet, and the
next save would persist that as a cluster-wide off. Without a stored flag the
only available answer is local: litellm_settings.callbacks activates
interception without storing one, and a write through the generic config
endpoint can drop the flag from a block that is still live. Reporting the
field default there would claim the feature is off while it serves.
"""
from litellm.integrations.websearch_interception.handler import (
WebSearchInterceptionLogger,
)
litellm_settings: Final[Mapping[str, object]] = _as_settings_section(config.get("litellm_settings"))
stored: Final[Mapping[str, object]] = _as_settings_section(litellm_settings.get("websearch_interception_params"))
if "enabled" in stored:
return dict(config)
resolved: Final = {
**stored,
"enabled": bool(litellm.logging_callback_manager.get_custom_loggers_for_type(WebSearchInterceptionLogger)),
}
return {
**config,
"litellm_settings": {**litellm_settings, "websearch_interception_params": resolved},
}
def _as_settings_section(value: object) -> Mapping[str, object]:
return cast("Mapping[str, object]", value) if isinstance(value, Mapping) else MappingProxyType({})
@router.get(
"/get/allowed_ips",
tags=["Budget & Spend Tracking"],
@ -875,7 +942,13 @@ async def update_default_team_member_budget(teams: list[NewUserRequestTeam], use
async def _update_litellm_setting(
settings: DefaultInternalUserParams | DefaultTeamSSOParams | MCPSemanticFilterSettings | MCPToolSearchSettings,
settings: (
DefaultInternalUserParams
| DefaultTeamSSOParams
| MCPSemanticFilterSettings
| MCPToolSearchSettings
| WebSearchInterceptionSettings
),
settings_key: str,
success_message: str,
user_api_key_dict: UserAPIKeyAuth,
@ -1399,6 +1472,77 @@ async def update_mcp_semantic_filter_settings(
return result
@router.get(
"/get/websearch_interception_settings",
tags=["Settings"],
dependencies=[Depends(user_api_key_auth)],
response_model=WebSearchInterceptionSettingsResponse,
)
async def get_websearch_interception_settings(
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Get web search interception configuration.
Returns the current settings plus their schema, for the Admin UI to render.
"""
from litellm.proxy.proxy_server import prisma_client, proxy_config
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": "Database not connected. Please connect a database."},
)
config: Final = await proxy_config.get_config()
return await _get_settings_with_schema(
settings_key="websearch_interception_params",
settings_class=WebSearchInterceptionSettings,
config=_with_websearch_enabled_resolved(config),
)
@router.patch(
"/update/websearch_interception_settings",
tags=["Settings"],
dependencies=[Depends(user_api_key_auth)],
)
async def update_websearch_interception_settings(
settings: WebSearchInterceptionSettings,
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Update web search interception settings in database.
Settings will be picked up by all pods within approximately 10 seconds via background polling.
"""
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
raise HTTPException(
status_code=403,
detail="Only proxy admins can update web search interception settings.",
)
result: Final = await _update_litellm_setting(
settings=settings,
settings_key="websearch_interception_params",
success_message=(
"Web search interception settings updated successfully. "
"Changes will be applied across all pods within 10 seconds."
),
user_api_key_dict=user_api_key_dict,
)
try:
from litellm.proxy.proxy_server import prisma_client, proxy_config
if prisma_client is not None:
await proxy_config.init_websearch_interception_settings_in_db(prisma_client=prisma_client)
except Exception as e:
verbose_proxy_logger.warning("Failed to reinitialize web search interception settings immediately: %s", e)
return result
@router.get(
"/get/mcp_tool_search_settings",
tags=["Settings"], # mutable-ok: FastAPI's route decorator only accepts a list

36
litellm/proxy/utils.py
View file

@ -249,6 +249,7 @@ if TYPE_CHECKING:
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.models.team import LiteLLM_TeamTableCachedObj
from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction
from litellm.proxy.db.baseline_accounting import BaselineAccountingRecord
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
from litellm.repositories.prisma_protocols import TableActions
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
@ -471,12 +472,16 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje
class _UpstreamStreamBoundary(Generic[_T]):
__slots__ = ("_upstream", "failure")
__slots__ = ("_source", "_upstream", "failure")
def __init__(self, upstream: AsyncIterable[_T]) -> None:
self._source: Final = upstream
self._upstream: Final = upstream.__aiter__()
self.failure: BaseException | None = None
def __getattr__(self, name: str) -> object:
return getattr(self._source, name)
def __aiter__(self) -> "_UpstreamStreamBoundary[_T]":
return self
@ -3037,6 +3042,10 @@ class ProxyLogging:
Otherwise, returns None and the original exception is used.
"""
logging_obj: Final[object] = request_data.get("litellm_logging_obj") # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] # legacy request data is narrowed to Logging below
if isinstance(logging_obj, Logging) and logging_obj.baseline_cache_context is not None:
await logging_obj.invalidate_baseline_cache_estimate("failed_request", completed=True)
### ALERTING ###
await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail")
if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception):
@ -4190,6 +4199,10 @@ class PrismaClient:
http_client: "HttpConfig | None" = None,
):
## init logging object
self.baseline_accounting_transactions: list[
BaselineAccountingRecord
] = [] # mutable-ok: locked background queue
self.baseline_accounting_lock: Final = asyncio.Lock()
self.proxy_logging_obj = proxy_logging_obj
self.token_auth: DatabaseTokenAuth | None = resolve_database_token_auth()
verbose_proxy_logger.debug("Creating Prisma Client..")
@ -7204,7 +7217,15 @@ async def _total_queued_spend_transactions(prisma_client: PrismaClient) -> int:
autorouter_queue_size: Final = len(prisma_client.autorouter_turn_transactions)
from litellm.proxy.db.shadow_eval_funnel import pending_shadow_eval_funnel_events
return spend_queue_size + tool_queue_size + autorouter_queue_size + pending_shadow_eval_funnel_events()
async with prisma_client.baseline_accounting_lock:
baseline_queue_size: Final = len(prisma_client.baseline_accounting_transactions)
return (
spend_queue_size
+ tool_queue_size
+ autorouter_queue_size
+ baseline_queue_size
+ pending_shadow_eval_funnel_events()
)
async def update_daily_tag_spend(
@ -7263,7 +7284,10 @@ async def update_spend_logs_job(
This job is triggered based on queue size rather than time.
Pops the batch once, writes spend logs, then runs guardrail usage tracking.
"""
from litellm.proxy.db.baseline_accounting import flush_baseline_accounting
if await _total_queued_spend_transactions(prisma_client) == 0:
await flush_baseline_accounting(prisma_client)
return
async with prisma_client.spend_log_write_lock:
await _run_spend_logs_job(prisma_client, db_writer_client, proxy_logging_obj)
@ -7274,6 +7298,8 @@ async def _run_spend_logs_job(
db_writer_client: AsyncHTTPHandler | None,
proxy_logging_obj: ProxyLogging,
) -> None:
from litellm.proxy.db.baseline_accounting import flush_baseline_accounting
n_retry_times: Final = 3
MAX_LOGS_PER_INTERVAL: Final = 10000
@ -7331,6 +7357,8 @@ async def _run_spend_logs_job(
tool_tracking_err,
)
await flush_baseline_accounting(prisma_client)
async with prisma_client._autorouter_turn_transactions_lock:
autorouter_turns_to_process: Final = prisma_client.autorouter_turn_transactions[:MAX_LOGS_PER_INTERVAL]
remaining_autorouter_turns: Final = prisma_client.autorouter_turn_transactions[
@ -7473,7 +7501,9 @@ async def _monitor_spend_logs_queue(
proxy_logging_obj=proxy_logging_obj,
)
else:
# Exponential backoff when no logs to process
from litellm.proxy.db.baseline_accounting import flush_baseline_accounting
await flush_baseline_accounting(prisma_client)
current_interval = min(current_interval * backoff_multiplier, max_backoff)
if await _wait_for_spend_log_flush_request(flush_requested, current_interval):

48
litellm/repositories/managed_batch_repository.py Normal file
View file

@ -0,0 +1,48 @@
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final
from litellm.repositories.table_repositories import PrismaTableRepository
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from prisma import models as prisma_models
def _batch_of(blob: object) -> LiteLLMBatch:
return LiteLLMBatch.model_validate_json(blob) if isinstance(blob, str) else LiteLLMBatch.model_validate(blob)
class ManagedBatchRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedObjectTable"]):
table_name = "litellm_managedobjecttable"
async def load_batch(self, unified_batch_id: str) -> LiteLLMBatch | None:
row: Final = await self._find_row(unified_batch_id)
return None if row is None or not row.file_object else _batch_of(row.file_object)
async def load_status(self, unified_batch_id: str) -> str | None:
row: Final = await self._find_row(unified_batch_id)
return row.status if row is not None else None
async def compare_and_set(
self, batch: LiteLLMBatch, unchanged: Mapping[str, object], updated_by: str | None
) -> bool:
updated_rows: Final = await self.table.update_many(
where={"unified_object_id": batch.id, **unchanged}, # mutable-ok: prisma filters are plain dicts
data={ # mutable-ok: prisma payloads are plain dicts
"file_object": batch.model_dump_json(),
"status": batch.status,
"updated_by": updated_by,
},
)
return updated_rows > 0
async def touch(self, unified_batch_id: str, updated_by: str | None) -> None:
await self.table.update_many(
where={"unified_object_id": unified_batch_id}, # mutable-ok: prisma filters are plain dicts
data={"updated_by": updated_by}, # mutable-ok: prisma payloads are plain dicts
)
async def _find_row(self, unified_batch_id: str) -> "prisma_models.LiteLLM_ManagedObjectTable | None":
return await self.table.find_first(
where={"unified_object_id": unified_batch_id} # mutable-ok: prisma filters are plain dicts
)

32
litellm/repositories/managed_file_content_repository.py Normal file
View file

@ -0,0 +1,32 @@
from typing import TYPE_CHECKING, Final
from litellm.repositories.table_repositories import PrismaTableRepository
if TYPE_CHECKING:
from prisma import models as prisma_models
class ManagedFileContentRepository(PrismaTableRepository["prisma_models.LiteLLM_ManagedFileContentTable"]):
table_name = "litellm_managedfilecontenttable"
async def store(self, content: bytes) -> str:
from prisma import Base64
row: Final = await self.table.create(
data={"content": Base64.encode(content)} # mutable-ok: prisma payloads are plain dicts
)
return row.id
async def load(self, row_id: str) -> bytes | None:
row: Final[prisma_models.LiteLLM_ManagedFileContentTable | None] = await self.table.find_unique(
where={"id": row_id} # mutable-ok: prisma filters are plain dicts
)
return None if row is None else row.content.decode()
async def delete(self, row_id: str) -> None:
from prisma.errors import RecordNotFoundError
try:
await self.table.delete(where={"id": row_id}) # mutable-ok: prisma filters are plain dicts
except RecordNotFoundError:
return

14
litellm/router.py
View file

@ -13723,6 +13723,20 @@ class Router:
to the deployment that actually served the request. Every attempt therefore
writes or clears, never just writes.
"""
from litellm.types.router import BaselineRouteStamp
baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None
baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None
router_name: Final = routing_decision.get("router_model_name") if routing_decision else None
Router._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key="_autorouter_baseline_route",
value=(
BaselineRouteStamp(router_name, baseline_model, baseline_id)
if router_name and baseline_model and baseline_id
else None
),
)
Router._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key="routing_decision",

21
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -75,6 +75,7 @@ from litellm.types.utils import (
AUTOROUTER_CLASSIFIER_CALL_ORIGIN,
ModelResponse,
RoutingDecisionCause,
StandardLoggingHeuristicV2Forecast,
StandardLoggingRoutingDecision,
StandardLoggingRoutingDecisionTierBoundaries,
)
@ -1043,6 +1044,7 @@ class ClassificationOutcome(NamedTuple):
capability_forecast: CapabilityClassifierForecast | None = None
llm_v2_forecast: LLMV2Decision | None = None
jev_verdict: JevVerdict | None = None
heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None
def _with_signal(outcome: ClassificationOutcome, signal: str | None) -> ClassificationOutcome:
@ -1075,6 +1077,8 @@ def _with_classifier_forecast(
decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome
) -> StandardLoggingRoutingDecision:
"""Attach validated forecasts and their applied policy to the routing decision."""
if outcome.heuristic_v2_forecast is not None:
return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast}
if outcome.jev_verdict is not None:
forecasted_decision: Final[StandardLoggingRoutingDecision] = {
**decision,
@ -1772,6 +1776,7 @@ class ComplexityRouter(CustomLogger):
conversation_continuing: bool = True,
tier_litellm_params: Mapping[str, object] | None = None,
context_escalation_original_tier: ComplexityTier | str | None = None,
heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None,
) -> StandardLoggingRoutingDecision:
"""Assemble the per-request provenance record for this router's decision.
@ -1831,7 +1836,9 @@ class ComplexityRouter(CustomLogger):
masked_tier_litellm_params: Final = mask_credentials_in_payload(tier_litellm_params)
if isinstance(masked_tier_litellm_params, Mapping):
decision["tier_litellm_params"] = masked_tier_litellm_params
return decision
return (
decision if heuristic_v2_forecast is None else {**decision, "heuristic_v2_forecast": heuristic_v2_forecast}
)
async def aclassify(
self,
@ -1888,6 +1895,15 @@ class ComplexityRouter(CustomLogger):
score=None,
signals=(f"request-type:{request_type.value}", *probability_signals),
cause="heuristic_v2",
heuristic_v2_forecast=StandardLoggingHeuristicV2Forecast(
probabilities={
candidate.value: prediction.probabilities[index]
for index, candidate in enumerate(TIER_SEVERITY_ORDER, start=1)
},
threshold=predictor.routing_threshold,
predicted_tier=tier.value,
request_type=request_type.value,
),
)
async def _classify_heuristic_first(
@ -3553,6 +3569,7 @@ class ComplexityRouter(CustomLogger):
context_escalation_original_tier=(
decision.get("context_escalation_original_tier") if decision is not None else None
),
heuristic_v2_forecast=decision.get("heuristic_v2_forecast") if decision is not None else None,
)
from litellm.types.router import PreRoutingHookResponse as HookResponse
@ -3732,6 +3749,7 @@ class ComplexityRouter(CustomLogger):
conversation_continuing=bool(decision.get("conversation_continuing", True)),
tier_litellm_params=self._litellm_params_for_model(candidate_tier, new_model),
context_escalation_original_tier=decision.get("context_escalation_original_tier"),
heuristic_v2_forecast=decision.get("heuristic_v2_forecast"),
)
return response.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict
@ -3776,6 +3794,7 @@ class ComplexityRouter(CustomLogger):
conversation_continuing=bool(decision.get("conversation_continuing", True)),
tier_litellm_params=self._litellm_params_for_model(None, default_model),
context_escalation_original_tier=decision.get("context_escalation_original_tier"),
heuristic_v2_forecast=decision.get("heuristic_v2_forecast"),
)
return response.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict

100
litellm/router_strategy/complexity_router/fuse_presets.json Normal file
View file

@ -0,0 +1,100 @@
{
"version": "2026-09-17-v1",
"models": [
{
"id": "gpt-6-astra-v1",
"label": "GPT-6 Astra",
"model": "gpt-6-astra",
"text": "OpenAI model for demanding end-to-end work, including reasoning, coding, research, and document tasks",
"sources": ["https://developers.openai.com/api/docs/models/gpt-6-astra"]
},
{
"id": "gpt-5.6-sol-v1",
"label": "GPT-5.6 Sol",
"model": "gpt-5.6-sol",
"text": "OpenAI model for complex professional work, supporting reasoning and tool calling",
"sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-sol"]
},
{
"id": "gpt-5.6-luna-v1",
"label": "GPT-5.6 Luna",
"model": "gpt-5.6-luna",
"text": "OpenAI model for high-volume workloads, supporting reasoning and tool calling",
"sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-luna"]
},
{
"id": "gpt-5.6-terra-v1",
"label": "GPT-5.6 Terra",
"model": "gpt-5.6-terra",
"text": "OpenAI general-purpose model supporting reasoning, text and image input, and tool calling",
"sources": ["https://developers.openai.com/api/docs/models/gpt-5.6-terra"]
},
{
"id": "claude-haiku-4-5-v1",
"label": "Claude Haiku 4.5",
"model": "claude-haiku-4-5",
"text": "Anthropic latency-focused model supporting text and image input, tool use, and extended thinking",
"sources": ["https://platform.claude.com/docs/en/models/haiku-4-5/overview"]
},
{
"id": "claude-sonnet-5-v1",
"label": "Claude Sonnet 5",
"model": "claude-sonnet-5",
"text": "Anthropic model balancing speed and capability, with adaptive thinking and tool use",
"sources": ["https://platform.claude.com/docs/en/models/sonnet-5/overview"]
},
{
"id": "claude-opus-5-v1",
"label": "Claude Opus 5",
"model": "claude-opus-5",
"text": "Anthropic model for complex agentic coding and enterprise work, with adaptive thinking",
"sources": ["https://platform.claude.com/docs/en/models/opus-5/overview"]
},
{
"id": "claude-fable-5-v1",
"label": "Claude Fable 5",
"model": "claude-fable-5",
"text": "Anthropic model for demanding reasoning and long-running agent tasks, with always-on adaptive thinking",
"sources": ["https://platform.claude.com/docs/en/models/fable-5/introducing-claude-fable-5-and-claude-mythos-5"]
},
{
"id": "claude-fable-5-1-v1",
"label": "Claude Fable 5.1",
"model": "claude-fable-5-1",
"text": "Anthropic model for demanding reasoning, long-running agentic coding, and multistep research, with always-on adaptive thinking",
"sources": ["https://platform.claude.com/docs/en/models/fable-5-1/overview"]
}
],
"harnesses": [
{
"id": "unspecified-v1",
"label": "Unspecified runtime",
"text": "Agent runtime is unspecified. Assess the task using the supplied context without assuming repository access, runnable tests, network access, additional tools, or a step, time, or spending budget",
"sources": ["https://code.claude.com/docs/en/how-claude-code-works", "https://mini-swe-agent.com/latest/faq/"]
},
{
"id": "claude-code-v1",
"label": "Claude Code",
"text": "Claude Code supplies an agent loop with context management and configured tools. Available actions depend on the session's tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget",
"sources": ["https://code.claude.com/docs/en/how-claude-code-works"]
},
{
"id": "codex-cli-v1",
"label": "Codex CLI",
"text": "Codex CLI supplies a terminal-based coding agent. File operations, command execution, and integrations depend on the session's tools, permissions, and sandbox. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget",
"sources": ["https://learn.chatgpt.com/docs/codex/cli", "https://learn.chatgpt.com/codex/permissions"]
},
{
"id": "opencode-v1",
"label": "OpenCode",
"text": "OpenCode supplies a configurable agent runtime. Available actions depend on the selected agent, tools, permissions, and execution environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget",
"sources": ["https://opencode.ai/docs/agents/"]
},
{
"id": "mini-swe-agent-v1",
"label": "mini-SWE-agent",
"text": "The standard mini-SWE-agent setup uses a bash-only action interface and separate command executions. Available commands and resources depend on its configured environment. The runtime name alone does not establish repository access, runnable tests, network access, additional tools, or a step, time, or spending budget",
"sources": ["https://mini-swe-agent.com/latest/faq/"]
}
]
}

52
litellm/router_strategy/complexity_router/fuse_presets.py Normal file
View file

@ -0,0 +1,52 @@
from functools import lru_cache
from importlib.resources import files
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, StringConstraints
ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)]
class FuseModelPreset(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
id: str
label: str
text: ProfileText
sources: tuple[str, ...] = Field(min_length=1)
model: str
class FuseHarnessPreset(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
id: str
label: str
text: ProfileText
sources: tuple[str, ...] = Field(min_length=1)
class FusePresetCatalog(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True)
version: str
models: tuple[FuseModelPreset, ...]
harnesses: tuple[FuseHarnessPreset, ...]
@lru_cache(maxsize=1)
def get_fuse_presets() -> FusePresetCatalog:
return FusePresetCatalog.model_validate_json(
files(__package__).joinpath("fuse_presets.json").read_text(encoding="utf-8")
)
def resolve_fuse_profile(text: str | None, preset_id: str | None, kind: Literal["model", "harness"]) -> str | None:
if preset_id is None:
return text
catalog: Final = get_fuse_presets()
presets: Final = catalog.models if kind == "model" else catalog.harnesses
preset: Final = next((entry for entry in presets if entry.id == preset_id), None)
if preset is None:
return None
return text if text is not None else preset.text

37
litellm/router_strategy/complexity_router/llm_v2.py
View file

@ -7,15 +7,15 @@ from dataclasses import dataclass
from sys import float_info
from typing import Annotated, Final, Literal, TypeAlias
from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter
from pydantic import BaseModel, ConfigDict, Field, StrictFloat, StringConstraints, TypeAdapter, model_validator
from typing_extensions import ReadOnly, TypedDict
from litellm.llms.base_llm.base_utils import (
type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # legacy output validated below
)
from litellm.router_strategy.complexity_router.fuse_presets import ProfileText, resolve_fuse_profile
ShortText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=512)]
ProfileText: TypeAlias = Annotated[str, StringConstraints(strip_whitespace=True, min_length=1, max_length=4000)]
class _SolverProfile(TypedDict):
@ -139,20 +139,41 @@ class LLMV2Config(BaseModel):
efficient_tier: str = "SIMPLE"
capable_tier: str = "REASONING"
efficient_profile: ProfileText
capable_profile: ProfileText
harness: ProfileText
efficient_profile: ProfileText | None = None
capable_profile: ProfileText | None = None
harness: ProfileText | None = None
efficient_profile_preset: str | None = None
capable_profile_preset: str | None = None
harness_preset: str | None = None
max_quality_gap: float = Field(ge=0.0, le=1.0, description="Maximum estimated success loss allowed for efficient.")
max_output_tokens: int = Field(default=1024, ge=1)
response_format: Literal["json_schema", "json_object"] = "json_schema"
calibration: LLMV2Calibration | None = None
@model_validator(mode="after")
def validate_profiles(self) -> LLMV2Config:
self._profile_texts()
return self
def _profile_texts(self) -> tuple[str, str, str]:
efficient: Final = resolve_fuse_profile(self.efficient_profile, self.efficient_profile_preset, "model")
capable: Final = resolve_fuse_profile(self.capable_profile, self.capable_profile_preset, "model")
harness: Final = resolve_fuse_profile(self.harness, self.harness_preset, "harness")
if efficient is None:
raise ValueError("efficient_profile requires text or a known efficient_profile_preset")
if capable is None:
raise ValueError("capable_profile requires text or a known capable_profile_preset")
if harness is None:
raise ValueError("harness requires text or a known harness_preset")
return efficient, capable, harness
def system_prompt(self, efficient_model: str, capable_model: str) -> str:
efficient, capable, harness = self._profile_texts()
profiles: Final[_SolverProfiles] = {
"prompt_version": LLM_V2_PROMPT_VERSION,
"harness": self.harness,
"efficient": {"model": efficient_model, "profile": self.efficient_profile},
"capable": {"model": capable_model, "profile": self.capable_profile},
"harness": harness,
"efficient": {"model": efficient_model, "profile": efficient},
"capable": {"model": capable_model, "profile": capable},
}
schema: Final = (
"\n\nResponse JSON schema:\n" + json.dumps(LLMV2Verdict.model_json_schema())

8
litellm/rust_bridge/_native.pyi
View file

@ -9,6 +9,8 @@ from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMe
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
class ForkedAfterNativeRuntimeStarted(RuntimeError): ...
class ProcessReservedForForking(RuntimeError): ...
def ocr(
request: LiteLLMOcrRequest,
@ -101,8 +103,12 @@ class TokenCounter:
def acount_request(self, body: bytes) -> Future[dict[str, object]]: ...
def gil_stats() -> dict[str, int]: ...
def process_state_started() -> bool: ...
def reserve_process_for_forking() -> None: ...
__all__ = [
"ForkedAfterNativeRuntimeStarted",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",
"RustUpstreamError",
@ -116,5 +122,7 @@ __all__ = [
"gil_stats",
"messages",
"ocr",
"process_state_started",
"reserve_process_for_forking",
"transcription",
]

47
litellm/rust_bridge/fork_guard.py Normal file
View file

@ -0,0 +1,47 @@
"""Fork safety of the Rust extension.
Its runtime threads do not survive ``fork``, so a child forked after the first native call
cannot run native routes: it raises ``ForkedAfterNativeRuntimeStarted`` instead of hanging.
Fork before the first native call, or start workers with ``spawn`` / ``forkserver``.
A process whose job is to fork workers (the gunicorn master under ``preload``) reserves itself:
from then on any native route called in it raises ``ProcessReservedForForking`` at the call
site, so the runtime can never start there. Workers forked from it are unaffected.
"""
from __future__ import annotations
from typing import Final
from litellm.rust_bridge.loader import get_native_bridge
class NativeStateStartedBeforeFork(RuntimeError):
pass
class _NeverRaised(RuntimeError):
"""Stands in for a native exception when the extension is unavailable or predates it."""
_native: Final = get_native_bridge()
ForkedAfterNativeRuntimeStarted: Final[type[RuntimeError]] = getattr(
_native, "ForkedAfterNativeRuntimeStarted", _NeverRaised
)
ProcessReservedForForking: Final[type[RuntimeError]] = getattr(_native, "ProcessReservedForForking", _NeverRaised)
def reserve_process_for_forking(where: str) -> None:
"""Forbid native routes in this process. Raises if one already ran here."""
native: Final = get_native_bridge()
reserve: Final = getattr(native, "reserve_process_for_forking", None)
if not callable(reserve):
return
try:
reserve()
except RuntimeError as error:
raise NativeStateStartedBeforeFork(
f"The LiteLLM Rust extension already ran a native route in {where}, and its runtime "
"threads do not survive fork(). Move the native call (warm-up, health check, "
"import-time initialization) into the worker, after the fork."
) from error

1
litellm/types/llms/openai.py
View file

@ -512,6 +512,7 @@ class CreateBatchRequest(TypedDict, total=False):
class LiteLLMBatchCreateRequest(CreateBatchRequest, total=False):
model: str
disable_fallbacks: ReadOnly[bool]
class RetrieveBatchRequest(TypedDict, total=False):

38
litellm/types/management_endpoints/auto_router_endpoints.py
View file

@ -199,13 +199,20 @@ class AutoRouterBenchmarkTotals(BaseModel):
description="Recorded LLM classifier cost already included in spend; null when any session turns predate "
"subtotal recording, and zero for an empty window"
)
saved_spend: float = Field(
description="Signed dollars saved versus each router's savings baseline (derived from its hardest "
"tier, or the configured override), from the same per-request savings record the usage tab reads"
savings_estimated_turns: int = Field(
description="Turns covered by the current savings estimator; legacy estimates are excluded"
)
savings_estimated_actual_spend: float = Field(
description="Actual spend, including classifier cost, for covered turns only"
)
saved_spend: float | None = Field(
description="Signed savings for covered turns only; null when traffic has no current estimates"
)
baseline_spend: float | None = Field(description="Estimated single-model cost for covered turns only")
saved_pct: float | None = Field(description="Covered savings over covered baseline spend, as a percentage")
saved_per_session: float | None = Field(
description="Average session savings; unavailable unless every turn is covered"
)
baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost")
saved_pct: float = Field(description="saved_spend over baseline_spend, as a percentage")
saved_per_session: float
cache: AutoRouterCacheStats
@ -236,16 +243,27 @@ class AutoRouterSessionResponse(BaseModel):
turns: int = Field(description="Auto-routed turns the rollup has recorded for this session so far")
last_model: str = Field(description="The deployment model the most recent turn was routed to")
spend: float = Field(description="What the session's routed traffic actually cost, classifier calls included")
saved_spend: float = Field(description="Estimated savings against the baseline, net of classifier cost")
baseline_spend: float = Field(description="spend plus saved_spend: the estimated single-model cost")
savings_estimated_turns: int = Field(
description="Turns covered by the current savings estimator; legacy estimates are excluded"
)
savings_estimated_actual_spend: float = Field(
description="Actual spend, including classifier cost, for covered turns only"
)
saved_spend: float | None = Field(description="Estimated savings for covered turns only, net of classifier cost")
baseline_spend: float | None = Field(
description="Estimated single-model cost; unavailable unless every turn is covered"
)
savings_estimated_baseline_spend: float | None = Field(
description="Estimated single-model cost for covered turns only"
)
baseline_model: str | None = Field(
description="The savings baseline most of this session's turns were priced against, recorded turn by "
description="The savings baseline most covered turns were priced against, recorded turn by "
"turn, so it still names the counterfactual after the router is reconfigured or removed. None when no "
"turn recorded one: rows from before the baseline was recorded, and adaptive and quality routers, "
"which derive no baseline and so report no savings"
)
baseline_models: Mapping[str, int] = Field(
description="Turns priced against each baseline model; more than one entry means the router's "
description="Covered turns priced against each baseline model; more than one entry means the router's "
"baseline changed mid-session and baseline_spend mixes both"
)

7
litellm/types/router.py
View file

@ -1057,6 +1057,13 @@ class TaggedPreRoutingStrategy(Generic[_PreRoutingStrategyT_co]):
strategy: _PreRoutingStrategyT_co
@dataclass(frozen=True, slots=True)
class BaselineRouteStamp:
router_name: str
baseline_model: str
baseline_deployment_id: str
@dataclass(frozen=True, slots=True)
class ConsumedRequestTagsStamp:
"""The model group a tagged router rewrote to, plus the request tags spent selecting it."""

16
litellm/types/utils.py
View file

@ -146,6 +146,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
supports_assistant_prefill: bool | None
supports_prompt_caching: bool | None
supports_prompt_cache_breakpoint: ReadOnly[bool | None]
supports_thinking_cache_preservation: ReadOnly[bool | None]
supports_computer_use: bool | None
supports_audio_input: bool | None
supports_embedding_image_input: bool | None
@ -2974,6 +2975,13 @@ LLM_AS_A_JUDGE_GUARDRAIL_CALL_ORIGIN: Final[InternalCallOrigin] = "llm_as_a_judg
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN: Final[InternalCallOrigin] = "background_response_cost_poll"
class StandardLoggingHeuristicV2Forecast(TypedDict):
probabilities: ReadOnly[Mapping[str, float]]
threshold: ReadOnly[float]
predicted_tier: ReadOnly[str]
request_type: ReadOnly[str]
class StandardLoggingRoutingDecision(TypedDict, total=False):
"""Per-request provenance for a pre-routing strategy (auto-router) decision."""
@ -2992,6 +3000,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
classifier_cost: float
classifier_probabilities: ReadOnly[Mapping[str, float]]
classifier_confidence: ReadOnly[float]
heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast]
classifier_crux: str # writable-ok: added only when a capability verdict is available
classifier_primary_rule: str # writable-ok: added only when a capability verdict is available
classifier_capability_boundary: str # writable-ok: added only when a capability verdict is available
@ -3037,6 +3046,7 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
"classifier_cost",
"classifier_probabilities",
"classifier_confidence",
"heuristic_v2_forecast",
"classifier_primary_rule",
"classifier_capability_boundary",
"classifier_p_solve",
@ -3450,7 +3460,9 @@ class StandardLoggingPayload(ClassifierAudit):
stream: bool | None
response_cost: float
cost_breakdown: CostBreakdown | None # Detailed cost breakdown
autorouter_savings: ReadOnly[float | None] # None = not an auto-routed caller request; 0.0 is a real figure
autorouter_savings: ReadOnly[float | None]
autorouter_savings_estimate: ReadOnly[Mapping[str, JsonValue] | None]
autorouter_baseline_observation: ReadOnly[str | None]
response_cost_failure_debug_info: StandardLoggingModelCostFailureDebugInformation | None
status: StandardLoggingPayloadStatus
status_fields: StandardLoggingPayloadStatusFields
@ -4178,6 +4190,8 @@ FILE_CONTENT_STREAMING_PROVIDERS: Final[frozenset[str]] = frozenset(
{*OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, LlmProviders.VERTEX_AI.value}
)
LITELLM_EXECUTED_BATCH_PROVIDERS: Final[frozenset[str]] = frozenset({LlmProviders.HOSTED_VLLM.value})
ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "litellm_proxy", "vertex_ai"]
LIST_BATCHES_SUPPORTED_PROVIDERS: Final[frozenset[str]] = frozenset(get_args(ListBatchesSupportedProvider))

11
litellm/utils.py
View file

@ -1881,6 +1881,7 @@ def client(original_function):
# Type assertion: logging_obj is guaranteed to be non-None after function_setup
assert logging_obj is not None, "logging_obj should not be None after function_setup"
kwargs["litellm_logging_obj"] = logging_obj
modified_kwargs: Final = await async_pre_call_deployment_hook(kwargs, call_type)
if modified_kwargs is not None:
kwargs = modified_kwargs
@ -2848,6 +2849,14 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None
)
def supports_thinking_cache_preservation(model: str, custom_llm_provider: str | None = None) -> bool:
return _supports_factory(
model=model,
custom_llm_provider=custom_llm_provider,
key="supports_thinking_cache_preservation",
)
def supports_computer_use(model: str, custom_llm_provider: str | None = None) -> bool:
"""
Check if the given model supports computer use and return a boolean value.
@ -5822,6 +5831,7 @@ def _get_model_info_helper(
supports_assistant_prefill=None,
supports_prompt_caching=None,
supports_prompt_cache_breakpoint=None,
supports_thinking_cache_preservation=None,
supports_computer_use=None,
supports_pdf_input=None,
)
@ -6094,6 +6104,7 @@ def _get_model_info_helper(
supports_assistant_prefill=_model_info.get("supports_assistant_prefill", None),
supports_prompt_caching=_model_info.get("supports_prompt_caching", None),
supports_prompt_cache_breakpoint=_model_info.get("supports_prompt_cache_breakpoint", None),
supports_thinking_cache_preservation=_model_info.get("supports_thinking_cache_preservation", None),
supports_audio_input=_model_info.get("supports_audio_input", None),
supports_audio_output=_model_info.get("supports_audio_output", None),
supports_pdf_input=_model_info.get("supports_pdf_input", None),

10
model_prices_and_context_window.json
View file

@ -14510,6 +14510,7 @@
"supports_native_structured_output": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_sampling_params": false,
@ -14547,6 +14548,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14698,6 +14700,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14727,6 +14730,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14759,6 +14763,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14796,6 +14801,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14831,6 +14837,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14869,6 +14876,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -14986,6 +14994,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,
@ -15027,6 +15036,7 @@
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_thinking_cache_preservation": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_native_structured_output": true,

3
model_prices_and_context_window.schema.json
View file

@ -911,6 +911,9 @@
"supports_system_messages": {
"type": "boolean"
},
"supports_thinking_cache_preservation": {
"type": "boolean"
},
"supports_tool_choice": {
"type": "boolean"
},

5
pyproject.toml
View file

@ -71,8 +71,8 @@ proxy = [
"mcp>=2.2.0,<3",
"httpx2>=2.5.0,<3",
"pydantic>=2.12.0,<3",
"litellm-proxy-extras==0.4.99",
"litellm-enterprise==0.1.68",
"litellm-proxy-extras==0.4.100",
"litellm-enterprise==0.1.69",
"RestrictedPython>=8.5,<9.0",
"rich>=13.9.4,<14.0",
"InquirerPy>=0.3.4,<1.0",
@ -302,6 +302,7 @@ editable-profile = "dev"
include = [
"litellm/proxy/_experimental/out/**",
"litellm/router_strategy/complexity_router/artifacts/*.json",
"litellm/router_strategy/complexity_router/fuse_presets.json",
"litellm/proxy/client/cli/commands/codex_base_instructions.md",
]
exclude = [

40
schema.prisma
View file

@ -1107,6 +1107,12 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
@@index([team_id, created_at(sort: Desc)])
}
model LiteLLM_ManagedFileContentTable {
id String @id @default(uuid())
content Bytes
created_at DateTime @default(now())
}
model LiteLLM_ManagedVectorStoreTable {
id String @id @default(uuid())
unified_resource_id String @unique // The base64 encoded unified vector store ID
@ -1545,6 +1551,36 @@ model LiteLLM_AdaptiveRouterSession {
@@index([last_activity_at], map: "idx_adaptive_router_session_activity")
}
model LiteLLM_AutoRouterBaselineComparison {
scope String @id
api_key String
session_id String
router_name String
initial_equivalent Boolean
revision BigInt @default(0)
published_revision BigInt @default(0)
history String?
attempted_at DateTime?
retired Boolean @default(false)
updated_at DateTime @default(now())
@@index([api_key, session_id, router_name], map: "idx_autorouter_baseline_scope")
@@index([updated_at], map: "idx_autorouter_baseline_updated")
}
model LiteLLM_AutoRouterBaselineObservation {
request_id String @id
scope String
started_at Float
revision BigInt
data String
publication String?
conflicted Boolean @default(false)
@@index([scope, started_at, request_id], map: "idx_autorouter_baseline_event_order")
@@index([scope, revision, started_at], map: "idx_autorouter_baseline_event_revision")
}
model LiteLLM_AutoRouterSession {
api_key String
session_id String
@ -1571,6 +1607,10 @@ model LiteLLM_AutoRouterSession {
total_tokens BigInt @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
savings_estimated_baseline_models Json @default("{}")
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
tier_turns Json @default("{}")

40
scripts/check_test_quality.py
View file

@ -60,6 +60,13 @@ TQ007 A module global that a conftest saves before every test and restores aft
names are read from the keys the conftest assigns directly and from whatever the
save loop iterates, including a module-level tuple or dict it names rather than
spells out.
TQ009 A child interpreter spawned as `subprocess.run([sys.executable, ...])` without
`-I`/`-P` as its first flag. Without isolation the child's sys.path leads with
the working directory, so a source checkout shadows the installed package and
the child tests a different `litellm` than the parent imported -- TQ003 is the
same working-directory hazard seen from the child's side. Use
tests.test_litellm_rust.support.child_interpreter.run_child_interpreter, which
also asserts the child resolved the same `litellm.__file__` as the parent.
Every rule is suppressible with `# test-quality-ok: <reason>` on the reported
line, following the repo's `*-ok: <reason>` convention. A suppression without a
@ -140,6 +147,9 @@ SKIP_CALLS: Final = frozenset(("pytest.skip", "skip"))
CONFTEST_NAME: Final = "conftest.py"
SDK_MODULE: Final = "litellm"
SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_call", "call"))
INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P"))
CREDENTIAL_NAME_RE: Final = re.compile(
r"(?:API_KEY|_KEY|TOKEN|SECRET|PASSWORD|CREDENTIAL|DATABASE_URL|ACCESS_KEY_ID)$"
)
@ -709,6 +719,35 @@ def _snapshotted_names(tree: ast.Module) -> Iterator[tuple[str, int]]:
yield from _string_members(iterable)
def iter_child_interpreter_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
for node in ast.walk(tree):
if not (isinstance(node, ast.Call) and node.args):
continue
if _dotted_name(node.func).rsplit(".", 1)[-1] not in SUBPROCESS_SPAWNS:
continue
argv: Final = node.args[0]
if not isinstance(argv, (ast.List, ast.Tuple)) or not argv.elts:
continue
if _dotted_name(argv.elts[0]) != "sys.executable":
continue
isolated: Final = (
len(argv.elts) > 1
and isinstance(argv.elts[1], ast.Constant)
and argv.elts[1].value in INTERPRETER_ISOLATION_FLAGS
)
if isolated:
continue
yield Violation(
path,
node.lineno,
"TQ009",
"child interpreter spawned without -I/-P; the working directory lands on sys.path "
"and a source checkout can shadow the installed package, use "
"tests.test_litellm_rust.support.child_interpreter.run_child_interpreter or pass -I "
f"(suppress: `# {SUPPRESSION_TOKEN}: <reason>`)",
)
def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]:
if path.name != CONFTEST_NAME:
return
@ -746,6 +785,7 @@ def check_file(path: Path) -> tuple[Violation, ...]:
*iter_credential_skip_violations(path, tree),
*iter_conftest_inventory_violations(path, tree),
*iter_internal_patch_violations(path, tree),
*iter_child_interpreter_violations(path, tree),
)
if violation.line not in skip
)

22
terraform/provider/docs/index.md
View file

@ -43,22 +43,22 @@ resource "litellm_team" "dev_team" {
The LiteLLM provider supports the following resources:
* [`litellm_model`](./resources/model) - Manage LiteLLM model configurations
* [`litellm_team`](./resources/team) - Manage teams and their permissions
* [`litellm_team_member`](./resources/team_member) - Manage team member configurations
* [`litellm_team_member_add`](./resources/team_member_add) - Add members to teams
* [`litellm_key`](./resources/key) - Manage API keys
* [`litellm_mcp_server`](./resources/mcp_server) - Manage MCP (Model Context Protocol) servers
* [`litellm_credential`](./resources/credential) - Manage credentials for various providers
* [`litellm_vector_store`](./resources/vector_store) - Manage vector stores
* [`litellm_jwt_key_mapping`](./resources/jwt_key_mapping) - Map JWT claim values to virtual keys
* `litellm_model` - Manage LiteLLM model configurations
* `litellm_team` - Manage teams and their permissions
* `litellm_team_member` - Manage team member configurations
* `litellm_team_member_add` - Add members to teams
* `litellm_key` - Manage API keys
* `litellm_mcp_server` - Manage MCP (Model Context Protocol) servers
* `litellm_credential` - Manage credentials for various providers
* `litellm_vector_store` - Manage vector stores
* `litellm_jwt_key_mapping` - Map JWT claim values to virtual keys
## Available Data Sources
The LiteLLM provider supports the following data sources:
* [`litellm_credential`](./data-sources/credential) - Retrieve credential information
* [`litellm_vector_store`](./data-sources/vector_store) - Retrieve vector store information
* `litellm_credential` - Retrieve credential information
* `litellm_vector_store` - Retrieve vector store information
## Authentication

3
test-quality-budget.json
View file

@ -22,5 +22,8 @@
},
"TQ008": {
"limit": 10993
},
"TQ009": {
"limit": 59
}
}

28
tests/code_coverage_tests/test_e2e_changed_gate.py
View file

@ -226,3 +226,31 @@ def test_an_unusable_secret_is_named_without_printing_its_value(
assert unprintable not in result.stderr
assert result.stdout == ""
assert not env_path.exists()
@pytest.mark.parametrize("phase", ("setup", "call", "teardown"))
def test_oauth_failure_diagnostics_do_not_publish_private_payloads(tmp_path: Path, phase: str) -> None:
suite: Final = ET.Element("testsuite")
case: Final = ET.SubElement(suite, "testcase", file=SELECTED[0])
private: Final = "private-token-in-exception-message"
failure: Final = ET.SubElement(case, "failure", message=private)
failure.text = private
properties: Final = ET.SubElement(case, "properties")
for name, value in (
("oauth_failure_phase", phase),
("oauth_exception_type", "AssertionError"),
("oauth_frame", "oauth_gateway.py:120:start"),
("oauth_frame", f"injected\\n{private}"),
("unrelated_property", private),
):
_ = ET.SubElement(properties, "property", name=name, value=value)
report: Final = tmp_path / "report.xml"
ET.ElementTree(suite).write(report)
result: Final = subprocess.run(
[sys.executable, "-I", str(GATE), str(report), SELECTED[0]], capture_output=True, text=True
)
assert result.returncode == 1
assert f"oauth_failure_phase: {phase}" in result.stdout
assert "oauth_exception_type: AssertionError" in result.stdout
assert "oauth_frame: oauth_gateway.py:120:start" in result.stdout
assert private not in result.stdout + result.stderr

8
tests/e2e/AGENTS.md
View file

@ -14,7 +14,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
- `quota_management/` - quota enforcement and accounting, one subfolder per behavior: `ratelimit/` (rpm/tpm blocks, window reset, pacing headers on live traffic), `budgets/` (budget definition, enforcement, and reset windows: key, team, tag, soft, multi-window), and `spend_tracking/` (spend logging and cost attribution on `/spend/*`)
- `management/` - key/team/user/organization management routes: create/update/delete persistence via the info routes, team membership, and llm-only-key route denials (API surface; not Playwright)
- `a2a/` - the A2A (agent-to-agent) surface: admin registration via `/v1/agents`, proxy-fronted card discovery at `/.well-known/agent-card.json`, and JSON-RPC `message/send` invocation, driving agents backed by the litellm completion bridge (a real provider) and asserting protocol-version normalization (0.3 vs 1.0)
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion lists and executes the server's tools with the stored per-user token
- `mcp/` - the MCP server surface over api_key auth against the real Datadog remote MCP server (see "MCP suite: real Datadog only" below); plus the gateway-managed OAuth (authorization_code) path exercised through `/chat/completions` in `test_mcp_chat_completion_oauth_e2e.py` and direct MCP protocol operations in `test_mcp_oauth_happy_path_e2e.py`, the one behavior Datadog's static-header auth cannot reach, seeding the per-user upstream token via the interactive authorize dance driven with the mcp SDK's own OAuth client (headless-browser consent from a saved session) and asserting the completion or protocol call lists and executes the server's tools with the stored per-user token
- `logging/` - logging-integration delivery (datadog and friends)
- `security/` - secret handling and log-leak protection
- `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory regression test (`test_reliability_memory_e2e.py`: a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak)
@ -26,14 +26,14 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family
## MCP suite: real Datadog only
Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite
Every test under `tests/e2e/mcp/` must exercise the proxy against the real Datadog remote MCP server, except the two Linear OAuth tests `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Do not add a compose service, FastMCP fixture, mock upstream, or any other fake MCP host for this suite
- Register via `register_datadog_mcp` in `tests/e2e/mcp/datadog_mcp.py` (or extend that helper if you need a different `toolsets=` / `allowed_tools` slice of the same Datadog endpoint). That posts `/v1/mcp/server` with `url=datadog_mcp_url(...)` and static headers `DD-API-KEY` / `DD-APPLICATION-KEY` from the process env
- Auth is Datadog's documented CI/header path, not a browser OAuth authorize/token dance. Hard-fail when `DD_API_KEY` or `DD_APP_KEY` is missing (`assert_dd_mcp_creds`); never skip for a missing fake upstream
- Prefer calling real Datadog tools that prove the product path (e.g. `search_datadog_logs` for list/call and permission denials). Seed a unique marker (`e2e-datadog-mcp-*`) in a chat completion when you need a log the tool can find; dual-read with `dd_logs` from conftest when delivery matters
- Delete the MCP server (and any keys) through `resources.defer` the same way every other suite tears down
- If a new MCP behavior cannot be covered with Datadog's tool surface, say so in the PR and get agreement before inventing another upstream; the default is always Datadog
- The one standing exception is `test_mcp_chat_completion_oauth_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so it cannot exercise gateway-managed OAuth or per-user token seeding in any form. That test drives a real Linear MCP server instead; it is still a real remote upstream, so the no-mock, no-fixture rule above holds unchanged
- The two standing exceptions are `test_mcp_chat_completion_oauth_e2e.py` and `test_mcp_oauth_happy_path_e2e.py`. Datadog authenticates with the static `DD-API-KEY` / `DD-APPLICATION-KEY` headers and exposes no authorize/token dance at all, so these tests drive a real Linear MCP server instead; they are still real remote upstreams, so the no-mock, no-fixture rule above holds unchanged. The direct OAuth test also uses the existing live provider edge to inspect forwarded headers without replay, and owns a separate source-built gateway for cold restarts
## Lay the pattern down in a class
@ -152,7 +152,7 @@ MCPs - endpoint features with the protocol op as the variant
mcp.<operation>.<auth_family>.<assertion>
operation : list_tools | call_tool | list_resources | read_resource | list_prompts | get_prompt
auth_family : none | api_key | bearer | oauth
assertion : succeeds | denied_without_permission
assertion : succeeds | denied_without_permission | persists_across_processes
e.g. mcp.call_tool.oauth.succeeds
```

53
tests/e2e/CONTRIBUTING.md
View file

@ -248,3 +248,56 @@ The semantic header set is `content-type`, `accept`, `anthropic-version`, `anthr
Excluded transport and telemetry headers are `host`, `content-length`, `connection`, `accept-encoding`, `user-agent`, `traceparent`, `tracestate`, `x-request-id`, `x-client-request-id` and `x-stainless-*`. Inbound transfer-encoding is unsupported; send JSON with content-length framing. The destination represents host identity and the relay carries original body bytes. Replay does not verify credentials, SDK timeout/retry behavior, transport performance, model availability or stateful remote IDs. Live relay uses original request bytes and header values, never the stored identity
Strict replay harness regression tests live in `tests/code_coverage_tests/test_provider_replay_harness.py`. The CircleCI `provider_replay_harness` job runs them alongside the existing legacy harness files with `--noconftest -o pythonpath=tests/e2e`; they need only synthetic HTTP providers and temporary fixture storage
## MCP OAuth happy path
`test_mcp_oauth_happy_path_e2e.py` runs one shared scenario with four variants:
aggregate gateway SSO and explicitly configured per-server JWT, each directly
against Linear and through the live provider edge. The edge forwards to real
Linear without replay and compares the forwarded bearer to the encrypted
canonical user/server credential. This observes the forwarding boundary, not
Linear's internal logs. Direct variants independently exercise discovery
Use the existing database preparation, Prisma generation and Keycloak setup.
Build and stage the dashboard from the tested checkout as in the UI runner.
Provide `DATABASE_URL`, `LITELLM_MASTER_KEY`, `LITELLM_SALT_KEY`, `LITELLM_LICENSE`,
and the `E2E_KEYCLOAK_*` settings. Capture a test-account Linear login using
`mcp/linear_session_capture.py` and set `E2E_LINEAR_STORAGE_STATE` to that private
file. The test workspace must contain a team. Do not publish browser state or
raw test/proxy output
```bash
E2E_MCP_OAUTH_LIVE=1 E2E_FIXTURE_MODE=live E2E_PROVIDER_CACHE=0 \
uv run --no-sync pytest tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py \
--rootdir=. --reruns 0
```
The test starts and restarts its own source-built proxy on a free loopback port,
retaining its database and SSO client but no Redis or process-local cache. It
does not restart an existing proxy or clear shared databases. Gateway login,
consent, immediate list/call and post-restart reconnect must all succeed. The
aggregate client never injects a gateway header; the explicitly labeled JWT
variant configures `x-litellm-api-key` for the first consent and reconnects with
only its gateway JWT after restart
`.github/workflows/test-mcp-oauth-e2e.yml` automatically requests a run for
same-repository pull requests changing MCP, gateway authentication/SSO, consent
UI, dependencies or the relevant E2E harness/workflow paths. It retains manual
`workflow_dispatch` for targeted verification. The four cases run in the
protected `e2e-changed` environment after its normal deployment approval;
reviewers should approve and inspect this separate OAuth check when it appears.
Fork pull requests do not run this credentialed job; use a reviewed
same-repository branch for their verification. The workflow's path-filtered
check is not configured here as a globally required branch-protection check.
Provision `E2E_LINEAR_STORAGE_STATE_B64` as a secret there and retain the existing E2E license/AWS role configuration. A missing or
expired session fails the job; collection, deselection and skips are not passes.
The generic changed-test job excludes this file because it requires an owned
proxy and consent UI. No LLM call is needed
Coverage remains limited to authorization-code OAuth over HTTP. M2M, OBO,
PKCE passthrough, static/BYOK, ID-JAG, forwarding, SigV4 and stdio are outside this
scenario; consult the registry and LIT-3559 for their existing coverage and gaps.
LIT-4506 owns broader isolation/failure regressions. LIT-7737 retains ownership
of dependency/Python compatibility and its matrix; this test reuses its delivered
environment and does not change dependency constraints or compatibility gates

158
tests/e2e/batches/test_batches_e2e.py
View file

@ -1160,62 +1160,151 @@ def _vllm_params(api_base: str, api_key: str | None, model_id: str) -> LiteLLMPa
)
class TestHostedVllmBatch:
"""hosted_vllm file upload + batch create (OpenAI-compatible path, LIT-3266).
HOSTED_VLLM_DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
HOSTED_VLLM_BAD_LINE_CUSTOM_ID = "req-bad"
hosted_vllm is in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, so /v1/files
and /v1/batches route through the OpenAI handler against the deployment's
api_base. Skipped for now: it needs a live vLLM (or OpenAI-compatible) server
exposing the files/batches APIs (HOSTED_VLLM_API_BASE), which the e2e
environment does not currently provision.
def _hosted_vllm_deployment(client: BatchClient, resources: ResourceManager) -> str:
api_base = os.environ.get("HOSTED_VLLM_API_BASE")
if api_base is None:
pytest.skip("set HOSTED_VLLM_API_BASE (the live vLLM server this deployment targets)")
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
model_id = (os.environ.get("HOSTED_VLLM_MODEL") or HOSTED_VLLM_DEFAULT_MODEL).strip()
proxy_name = batch_model_name("hosted-vllm-batch")
model_row_id = client.create_model(proxy_name, _vllm_params(api_base, api_key, model_id))
resources.defer(lambda: client.delete_model(model_row_id))
return proxy_name
def _upload_hosted_vllm_input(
client: BatchClient, content: bytes, *, proxy_name: str, key: str, upload_route: str
) -> Result[FileObject]:
if upload_route == "model_query":
return client.upload_file(content=content, form=FileUploadForm(purpose="batch"), model=proxy_name, key=key)
return client.upload_file(
content=content, form=FileUploadForm(purpose="batch", target_model_names=proxy_name), key=key
)
def _jsonl_with_a_failing_line(model: str) -> bytes:
bad_line = {
"custom_id": HOSTED_VLLM_BAD_LINE_CUSTOM_ID,
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": model, "messages": [{"role": "user", "content": "ping"}], "max_tokens": -1},
}
return render_jsonl(model) + (json.dumps(bad_line) + "\n").encode()
def _download_managed_file(client: BatchClient, file_id: str, *, key: str) -> list[str]:
downloaded = client.proxy.transport.download(
f"/v1/files/{file_id}/content", headers=client.proxy.transport.bearer(key)
)
assert downloaded.status_code == 200, (
f"file content must be 200, got {downloaded.status_code}: {downloaded.body[:300]}"
)
return downloaded.body.strip().splitlines()
class TestHostedVllmBatch:
"""hosted_vllm file upload + batch execution (LIT-5739).
vLLM implements neither /v1/files nor /v1/batches, so LiteLLM keeps the batch
input in its own database, runs every line through the deployment's
/v1/chat/completions itself, and serves the batch plus its output and error
files from that database under the creating key. Needs a live vLLM server
(HOSTED_VLLM_API_BASE), which the default e2e stack does not provision, so
the cases skip without it.
"""
@pytest.mark.skip(
reason="hosted_vllm batch/files needs a live vLLM server (HOSTED_VLLM_API_BASE) "
"not provisioned in the e2e environment; re-enable when available (LIT-3266)"
)
@pytest.mark.parametrize("upload_route", ["target_model_names", "model_query"])
@pytest.mark.covers(
"llm.batches.hosted_vllm.basic.nonstream.works",
"llm.files.hosted_vllm.upload.nonstream.works",
exercised_on=["batches", "files"],
)
def test_unified_file_and_batch_create(
self, client: BatchClient, resources: ResourceManager
def test_batch_runs_to_completion_with_a_downloadable_output(
self, client: BatchClient, resources: ResourceManager, upload_route: str
) -> None:
api_base = os.environ["HOSTED_VLLM_API_BASE"]
api_key = (os.environ.get("HOSTED_VLLM_API_KEY") or "").strip() or None
model_id = (
os.environ.get("HOSTED_VLLM_MODEL") or "meta-llama/Llama-3.2-3B-Instruct"
).strip()
proxy_name = batch_model_name("hosted-vllm-batch")
model_row_id = client.create_model(
proxy_name, _vllm_params(api_base, api_key, model_id)
)
resources.defer(lambda: client.delete_model(model_row_id))
proxy_name = _hosted_vllm_deployment(client, resources)
key = resources.key()
file = unwrap(
client.upload_file(
content=render_jsonl(model_id),
form=FileUploadForm(purpose="batch", target_model_names=proxy_name),
key=key,
_upload_hosted_vllm_input(
client, render_jsonl(proxy_name), proxy_name=proxy_name, key=key, upload_route=upload_route
)
)
resources.defer(lambda: cleanup_file(client, file.id, key=key))
assert_file_object(file, provider="hosted_vllm")
assert is_managed_id(file.id), f"hosted_vllm batch input must stay in LiteLLM, got file id {file.id!r}"
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key))
assert batch.id, f"hosted_vllm create returned no batch id: {created.body[:200]}"
assert batch.status in CREATED_BATCH_STATUSES, (
f"hosted_vllm batch has non-transitional status {batch.status!r}"
)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True))
assert is_managed_id(batch.id), f"hosted_vllm batch must be LiteLLM-managed, got {batch.id!r}"
assert batch.status in CREATED_BATCH_STATUSES, f"hosted_vllm batch has non-transitional status {batch.status!r}"
assert_batch_object(batch)
finished = _poll_until_terminal(client, batch.id, key)
assert finished.status == "completed", f"hosted_vllm batch ended {finished.status!r}: {finished.errors!r}"
assert finished.output_file_id, "completed hosted_vllm batch has no output_file_id"
assert finished.error_file_id is None, f"all lines succeeded but error_file_id={finished.error_file_id!r}"
output_lines = _download_managed_file(client, finished.output_file_id, key=key)
assert len(output_lines) == 1, f"one input line must yield one output line, got {output_lines!r}"
first_line = BatchOutputLine.model_validate_json(output_lines[0])
assert first_line.custom_id == "req-1", f"output line lost its custom_id: {output_lines[0][:300]}"
assert first_line.response.status_code == 200, f"batch output line reports failure: {output_lines[0][:400]}"
assert first_line.response.body is not None and first_line.response.body.choices, (
"batch output line has no choices"
)
rows = client.proxy.poll_logs_for_key(
key, predicate=lambda found: any(row.call_type == "acompletion" for row in found)
)
line_rows = [row for row in rows if row.call_type == "acompletion"]
assert line_rows, f"the batch line's chat call was not logged under the creating key: {rows!r}"
assert all(row.custom_llm_provider == "hosted_vllm" for row in line_rows), (
f"batch line rows must be attributed to hosted_vllm: {line_rows!r}"
)
@pytest.mark.covers("llm.batches.hosted_vllm.basic.nonstream.works", exercised_on=["batches", "files"])
def test_failing_line_lands_in_the_error_file_not_the_batch_status(
self, client: BatchClient, resources: ResourceManager
) -> None:
proxy_name = _hosted_vllm_deployment(client, resources)
key = resources.key()
file = unwrap(
_upload_hosted_vllm_input(
client,
_jsonl_with_a_failing_line(proxy_name),
proxy_name=proxy_name,
key=key,
upload_route="target_model_names",
)
)
resources.defer(lambda: cleanup_file(client, file.id, key=key))
created = client.create_batch(body=BatchCreateBody(input_file_id=file.id), key=key)
require_successful_call(created)
batch = BatchObject.model_validate_json(created.body)
resources.defer(lambda: cleanup_batch(client, batch.id, key=key, delete_output_files=True))
finished = _poll_until_terminal(client, batch.id, key)
assert finished.status == "completed", f"a failing line must not fail the batch, got {finished.status!r}"
assert finished.output_file_id, "the good line must still produce an output file"
assert finished.error_file_id, "the failing line must produce an error file"
output_lines = _download_managed_file(client, finished.output_file_id, key=key)
error_lines = _download_managed_file(client, finished.error_file_id, key=key)
assert [BatchOutputLine.model_validate_json(line).custom_id for line in output_lines] == ["req-1"]
assert len(error_lines) == 1, f"one failing line must yield one error line, got {error_lines!r}"
error_line = BatchOutputLine.model_validate_json(error_lines[0])
assert error_line.custom_id == HOSTED_VLLM_BAD_LINE_CUSTOM_ID
assert error_line.response.status_code == 400, f"error line must carry the provider's 4xx: {error_lines[0][:400]}"
BATCH_TERMINAL_STATUSES = frozenset({"completed", "failed", "expired", "cancelled"})
FAILED_BATCH_POLL_SECONDS = 120.0
@ -1443,6 +1532,7 @@ class BatchOutputResponse(BaseModel):
class BatchOutputLine(BaseModel):
custom_id: str | None = None
response: BatchOutputResponse

17
tests/e2e/conftest.py
View file

@ -17,6 +17,7 @@ import functools
import os
from collections.abc import Generator, Iterator
from datetime import datetime, timezone
from pathlib import Path
from types import MappingProxyType
from typing import Final
@ -28,6 +29,7 @@ from e2e_config import (
FIXTURE_DIR,
FIXTURE_MODE_RAW,
MANAGED_FILES_OPT_IN_ENV,
MCP_OAUTH_LIVE_OPT_IN_ENV,
PROMPT_CACHING_OPT_IN_ENV,
PROXY_BASE_URL,
REDIS_CHAOS_OPT_IN_ENV,
@ -56,6 +58,7 @@ OPT_IN_MARKERS: Final = MappingProxyType(
"prompt_caching_stack": PROMPT_CACHING_OPT_IN_ENV,
"redis_chaos": REDIS_CHAOS_OPT_IN_ENV,
"cli_determinism": CLI_DETERMINISM_OPT_IN_ENV,
"mcp_oauth_live": MCP_OAUTH_LIVE_OPT_IN_ENV,
}
)
@ -132,6 +135,11 @@ def pytest_configure(config: pytest.Config) -> None:
"redis_chaos: load test that pauses the proxy's Redis outright mid-run; needs a proxy booted from "
"gateway/redis_chaos_ci_config.yml on the same host, and is deselected unless E2E_REDIS_CHAOS is set",
)
config.addinivalue_line(
"markers",
"mcp_oauth_live: real Linear OAuth consent via a captured browser session; deselected unless "
"E2E_MCP_OAUTH_LIVE is set",
)
def pytest_sessionstart(session: pytest.Session) -> None:
@ -213,6 +221,8 @@ def pytest_runtest_setup(item: pytest.Item) -> None:
LIVE_PROVIDER_REQUIRED.set(item.get_closest_marker("provider_live") is not None)
if item.get_closest_marker("e2e") is None:
return
if isinstance(item, pytest.Function) and "oauth_gateway" in item.fixturenames:
return
reason = _proxy_fail_reason()
if reason is not None:
pytest.fail(reason)
@ -236,6 +246,13 @@ def pytest_runtest_makereport(
"""Stash the call-phase outcome so teardown can tell a passed test from a
failed one without re-deriving it."""
report = yield
if item.get_closest_marker("mcp_oauth_live") is not None and call.excinfo is not None:
# Publish code locations only, never exception messages, source text or locals.
item.user_properties.append(("oauth_failure_phase", report.when))
item.user_properties.append(("oauth_exception_type", call.excinfo.type.__name__))
for entry in call.excinfo.traceback:
item.user_properties.append(("oauth_frame", f"{Path(entry.path).name}:{entry.lineno + 1}:{entry.name}"))
report.user_properties = list(item.user_properties)
if report.when == "call":
item.stash[_CALL_PASSED] = report.passed
return report

8
tests/e2e/coverage_registry/mcp.yaml
View file

@ -71,6 +71,14 @@
assertions: [succeeds]
source: "db.py user_oauth_credential lookup"
rationale: OAuth2 token passthrough; per-user credential storage
- id: mcp.call_tool.oauth.persists_across_processes
module: mcp
tier: P1
operation: call_tool
auth_family: oauth
assertions: [persists_across_processes]
source: "outbound_credentials/per_user_oauth_store.py V2PerUserTokenStore"
rationale: Stored per-user token survives a verified restart of an owned gateway with no Redis cache
- id: mcp.list_tools.none.succeeds
module: mcp
tier: P1

1
tests/e2e/coverage_registry/mgmt.yaml
View file

@ -72,7 +72,6 @@
- {id: mgmt.callback.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "callback_management_endpoints.py", rationale: "Callback config (smoke)"}
- {id: mgmt.cost_tracking.estimate.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "cost_tracking_settings.py", rationale: "Cost estimate (smoke)"}
- {id: mgmt.router_settings.update.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "router_settings_endpoints.py", rationale: "Router config (smoke)"}
- {id: mgmt.config.allowed_ip.changed_key_only, module: mgmt, tier: P2, surface: api, assertions: [persists], source: "proxy_setting_endpoints.py:496", rationale: "An allowed-IP change leaves unrelated file settings out of the DB row"}
- {id: mgmt.jwt_key_mapping.new.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "jwt_key_mapping_endpoints.py", rationale: "JWT->key mapping (smoke)"}
- {id: mgmt.compliance.gdpr.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "compliance_endpoints.py", rationale: "GDPR ops (smoke)"}
- {id: mgmt.tool_management.list.happy_path, module: mgmt, tier: P2, surface: api, assertions: [happy_path], source: "tool_management_endpoints.py", rationale: "Tool inventory (smoke)"}

2
tests/e2e/e2e_config.py
View file

@ -52,6 +52,7 @@ CHEAP_OPENAI_MODEL = os.environ.get("E2E_CHEAP_OPENAI_MODEL", "gpt-5.5")
LINEAR_MCP_URL = os.environ.get("E2E_LINEAR_MCP_URL", "https://mcp.linear.app/mcp")
LINEAR_STORAGE_STATE = os.environ.get("E2E_LINEAR_STORAGE_STATE", "")
LINEAR_READONLY_TOOL: Final = "list_teams" # as listed by tools/list on mcp.linear.app when PR #33787 landed
# Jaeger query API of the compose stack's OTEL trace destination (the `jaeger`
# service in docker-compose.yml maps it to host 16686). Trace-completeness tests
@ -144,6 +145,7 @@ MANAGED_FILES_OPT_IN_ENV = "E2E_MANAGED_FILES_STACK"
PROMPT_CACHING_OPT_IN_ENV = "E2E_PROMPT_CACHING_STACK"
REDIS_CHAOS_OPT_IN_ENV = "E2E_REDIS_CHAOS"
CLI_DETERMINISM_OPT_IN_ENV = "E2E_CLI_DETERMINISM"
MCP_OAUTH_LIVE_OPT_IN_ENV: Final = "E2E_MCP_OAUTH_LIVE"
ANOMALY_SESSIONS = int(os.environ.get("E2E_ANOMALY_SESSIONS", "6"))
ANOMALY_TURNS_PER_SESSION = int(os.environ.get("E2E_ANOMALY_TURNS_PER_SESSION", "6"))
ANOMALY_TURN_ATTEMPTS = int(os.environ.get("E2E_ANOMALY_TURN_ATTEMPTS", "3"))

1
tests/e2e/gateway/stage_mirror_ci_config.yml
View file

@ -1,5 +1,4 @@
general_settings:
max_parallel_requests: 100
proxy_batch_write_at: 5
enable_jwt_auth: true
litellm_jwtauth:

4
tests/e2e/idp.py
View file

@ -435,7 +435,7 @@ def _signal_process_group(process_id: int, signum: int) -> bool:
return True
def _stop_process_group(child: subprocess.Popen[bytes]) -> None:
def stop_process_group(child: subprocess.Popen[bytes]) -> None:
_signal_process_group(child.pid, signal.SIGTERM)
deadline: Final = time.monotonic() + 5
while _process_group_exists(child.pid):
@ -476,7 +476,7 @@ def run_oidc_profile(proxy_url: str, command: list[str]) -> int:
try:
return child.wait()
finally:
_stop_process_group(child)
stop_process_group(child)
if __name__ == "__main__":

149
tests/e2e/management/test_config_misc_endpoints_e2e.py
View file

@ -7,13 +7,18 @@ so a read-back reflects the change. Router settings, which mutate global proxy
state, are exercised with a benign, self-restoring change so a shared proxy is left
as it was found.
Cache settings and the Vault config override are deliberately not covered here.
Both routes reconfigure the whole proxy: /cache/settings persists what it receives
into a row that outranks the YAML cache_params and is re-applied on a timer, and
/config_overrides/hashicorp_vault swaps the process-wide secret manager. Neither can
be exercised safely against the shared proxy the suites run on, so they need an
isolated proxy before a test lands. Do not add a read-then-write-back test for
either one.
Cache settings, the Vault config override and the allowed-IP routes are deliberately
not covered here. All three reconfigure the whole proxy: /cache/settings persists what
it receives into a row that outranks the YAML cache_params and is re-applied on a timer,
/config_overrides/hashicorp_vault swaps the process-wide secret manager, and
/add/allowed_ip mutates the live general_settings["allowed_ips"] that
auth_utils._check_valid_ip reads, so the first call locks every other client out of the
shared proxy. The allowlist is an exact string match with no CIDR support, and no route
reports the caller's address as the proxy sees it, so a test cannot allowlist itself
first; /delete/allowed_ip sits behind the same auth dependency, so the cleanup is locked
out too and the proxy stays poisoned for the rest of the build. None of the three can be
exercised safely against the shared proxy the suites run on, so they need an isolated
proxy before a test lands. Do not add a read-then-write-back test for any of them.
"""
from __future__ import annotations
@ -21,10 +26,9 @@ from __future__ import annotations
import math
import time
from collections.abc import Callable
from typing import Final
import pytest
from pydantic import BaseModel, JsonValue, RootModel
from pydantic import BaseModel
from e2e_config import unique_marker
from e2e_http import NoBody, Success, unwrap, unwrap_status
@ -188,7 +192,7 @@ class JwtKeyMappingResponse(BaseModel):
class RouterSettingsPatch(BaseModel):
num_retries: int
retry_after: int
class ConfigUpdateBody(BaseModel):
@ -199,39 +203,8 @@ class ConfigUpdateResponse(BaseModel):
message: str
class AllowedIpBody(BaseModel):
ip: str
class ConfigFieldInfoParams(BaseModel):
field_name: str
class ConfigFieldInfoResponse(BaseModel):
field_name: str
field_value: JsonValue
source: str
editable: bool
class ConfigListParams(BaseModel):
config_type: str
class ConfigListEntry(BaseModel):
field_name: str
field_value: JsonValue
stored_in_db: bool | None
source: str
editable: bool
class ConfigListResponse(RootModel[list[ConfigListEntry]]):
pass
class RouterCurrentValues(BaseModel):
num_retries: int | None = None
retry_after: int | None = None
class RouterSettingsResponse(BaseModel):
@ -493,17 +466,25 @@ class TestRouterSettings:
) -> None:
"""/config/update is the only write path for router_settings (there is no
dedicated router-settings write route). The change is restored on teardown so
the shared proxy keeps its original retry policy."""
original = self._read_num_retries(client)
assert original is not None, "GET /router/settings did not report num_retries; cannot prove a change"
resources.defer(lambda: self._write_num_retries(client, original))
the shared proxy keeps its original retry policy.
target = original + 5
retry_after is the subject because it satisfies all three constraints at once:
no lane's config file declares it, so the database owns it and the write is not
refused as config-owned; it is in RUNTIME_UPDATABLE_ROUTER_SETTINGS, so
/config/update accepts it; and it is in ROUTER_SETTINGS_FIELDS backed by an
always-set Router attribute, so GET /router/settings reports it for the
read-back. Bumping it by one second is the smallest change that proves the
round-trip without slowing a concurrent test that hits a retry."""
original = self._read_retry_after(client)
assert original is not None, "GET /router/settings did not report retry_after; cannot prove a change"
resources.defer(lambda: self._write_retry_after(client, original))
target = original + 1
response = unwrap(
client.proxy.transport.post(
"/config/update",
headers=client.proxy.transport.master,
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=target)),
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=target)),
response_type=ConfigUpdateResponse,
)
)
@ -513,20 +494,20 @@ class TestRouterSettings:
_ = _poll(
client,
lambda: True if self._read_num_retries(client) == target else None,
f"GET /router/settings never reported num_retries {target} after /config/update",
lambda: True if self._read_retry_after(client) == target else None,
f"GET /router/settings never reported retry_after {target} after /config/update",
)
self._write_num_retries(client, original)
self._write_retry_after(client, original)
restored = _poll(
client,
lambda: original if self._read_num_retries(client) == original else None,
f"GET /router/settings never returned to the original num_retries {original} after the restore",
lambda: original if self._read_retry_after(client) == original else None,
f"GET /router/settings never returned to the original retry_after {original} after the restore",
)
assert restored == original, f"router num_retries left at {restored}, expected the original {original}"
assert restored == original, f"router retry_after left at {restored}, expected the original {original}"
@staticmethod
def _read_num_retries(client: ManagementClient) -> int | None:
def _read_retry_after(client: ManagementClient) -> int | None:
return unwrap(
client.proxy.transport.get(
"/router/settings",
@ -534,72 +515,20 @@ class TestRouterSettings:
params=NoBody(),
response_type=RouterSettingsResponse,
)
).current_values.num_retries
).current_values.retry_after
@staticmethod
def _write_num_retries(client: ManagementClient, value: int) -> None:
def _write_retry_after(client: ManagementClient, value: int) -> None:
_ = unwrap(
client.proxy.transport.post(
"/config/update",
headers=client.proxy.transport.master,
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(num_retries=value)),
json=ConfigUpdateBody(router_settings=RouterSettingsPatch(retry_after=value)),
response_type=ConfigUpdateResponse,
)
)
class TestConfigPersistence:
@pytest.mark.covers("mgmt.config.allowed_ip.changed_key_only")
def test_add_allowed_ip_does_not_store_unrelated_config_value(
self, client: ManagementClient, resources: ResourceManager
) -> None:
allowed_ip: Final = "127.0.0.1"
added: Final = unwrap(
client.proxy.transport.post(
"/add/allowed_ip",
headers=client.proxy.transport.master,
json=AllowedIpBody(ip=allowed_ip),
response_type=ConfigUpdateResponse,
)
)
resources.defer(
lambda: unwrap(
client.proxy.transport.post(
"/delete/allowed_ip",
headers=client.proxy.transport.master,
json=AllowedIpBody(ip=allowed_ip),
response_type=ConfigUpdateResponse,
)
)
)
assert added.message == f"IP {allowed_ip} address added successfully"
listed: Final = unwrap(
client.proxy.transport.get(
"/config/list",
headers=client.proxy.transport.master,
params=ConfigListParams(config_type="general_settings"),
response_type=ConfigListResponse,
)
)
unrelated: Final = next(entry for entry in listed.root if entry.field_name == "max_parallel_requests")
assert unrelated.stored_in_db is not True
assert unrelated.source == "config"
assert unrelated.editable is False
field_info: Final = unwrap(
client.proxy.transport.get(
"/config/field/info",
headers=client.proxy.transport.master,
params=ConfigFieldInfoParams(field_name="max_parallel_requests"),
response_type=ConfigFieldInfoResponse,
)
)
assert field_info.source == "config"
assert field_info.editable is False
assert field_info.field_value == unrelated.field_value
class TestMcpServerSubmission:
@pytest.mark.covers("mgmt.mcp_server.register.happy_path")
def test_register_submits_pending_server(self, client: ManagementClient, resources: ResourceManager) -> None:

198
tests/e2e/mcp/oauth_chat_client.py
View file

@ -18,7 +18,7 @@ import asyncio
import re
import time
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Final
from urllib.parse import parse_qsl
import httpx
@ -26,11 +26,21 @@ import httpx2
import pytest
from e2e_config import PROXY_BASE_URL, REQUEST_TIMEOUT
from e2e_http import AuthHeaders, NoBody, unwrap
from idp import Identity
from mcp import ClientSession
from mcp.client.auth import OAuthClientProvider
from mcp.client.streamable_http import streamable_http_client
from mcp.shared.auth import AuthorizationCodeResult, OAuthClientInformationFull, OAuthClientMetadata, OAuthToken
from models import ChatBody, ChatResponse, McpServerCreateBody, McpServerInfo
from mcp.types import TextContent
from models import (
ChatBody,
ChatResponse,
McpOauthUserCredentialStatus,
McpServerCreateBody,
McpServerInfo,
McpServerUserCredentialListResponse,
McpServerUserCredentialRow,
)
from proxy_client import ProxyClient
if TYPE_CHECKING:
@ -44,8 +54,8 @@ OAUTH_CLIENT_REDIRECT_URI = "http://127.0.0.1:53682/e2e/callback"
BROWSER_CONSENT_TIMEOUT = 60.0
def _mcp_url(alias: str) -> str:
return f"{PROXY_BASE_URL}/{alias}/mcp"
def _mcp_url(alias: str, base_url: str = PROXY_BASE_URL) -> str:
return f"{base_url.rstrip('/')}/{alias}/mcp"
class InMemoryTokenStorage:
@ -69,7 +79,13 @@ class InMemoryTokenStorage:
self._client_info = client_info
async def _browser_follow_authorize(start_url: str, storage_state_path: str) -> tuple[str, str | None]:
async def _browser_follow_authorize(
start_url: str,
storage_state_path: str,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> tuple[str, str | None]:
"""Play the browser's role for a real upstream whose authorize endpoint
serves an interactive consent page (Linear). A headless Chromium primed
with a human's saved Linear session opens the gateway authorize URL and
@ -85,6 +101,9 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
def _note_request(request: object) -> None:
url = getattr(request, "url", "")
host = httpx.URL(url).host
if not allow_upstream_consent and (host == "linear.app" or host.endswith(".linear.app")):
captured["upstream_consent"] = "seen"
if url.startswith(OAUTH_CLIENT_REDIRECT_URI) and "url" not in captured:
captured["url"] = url
@ -96,7 +115,7 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
context = await browser.new_context(storage_state=storage_state_path)
await context.route(re.compile(re.escape(OAUTH_CLIENT_REDIRECT_URI) + r".*"), _swallow_redirect)
page = await context.new_page()
page.on("request", _note_request)
context.on("request", _note_request)
page.on("framenavigated", lambda frame: trail.append(frame.url.split("?", 1)[0]))
await page.goto(start_url, wait_until="domcontentloaded")
deadline = time.monotonic() + BROWSER_CONSENT_TIMEOUT
@ -105,8 +124,29 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
await page.wait_for_load_state("networkidle", timeout=8000)
except Exception: # noqa: BLE001 - a busy consent page never idles; fall through and try to advance it
pass
if "url" in captured:
if "upstream_consent" in captured or "url" in captured:
break
if await page.locator("#username").count() and identity is not None:
await page.locator("#username").fill(identity.username)
await page.locator("#password").fill(identity.password)
await page.locator("#kc-login").click()
continue
if "/ui/connect" in page.url and server_alias is not None:
card = page.locator("div.cursor-pointer").filter(has=page.get_by_text(server_alias, exact=True))
if await card.count() != 1:
await asyncio.sleep(0.5)
continue
connect = card.get_by_text("Connect", exact=True)
if await connect.count():
await connect.click()
continue
if not await card.locator("svg.text-success").count():
await asyncio.sleep(0.5)
continue
finish = page.get_by_role("button", name="Finish connecting", exact=True)
if await finish.count() and await finish.is_enabled():
await finish.click()
continue
control = page.locator(
'button[name="action"][value="approve"], button:has-text("Authorize"), '
'button:has-text("Allow"), button:has-text("@"), a:has-text("@")'
@ -118,27 +158,44 @@ async def _browser_follow_authorize(start_url: str, storage_state_path: str) ->
final_url = page.url
await browser.close()
# A redirect chain can finish inside goto/networkidle before the loop checks the page.
assert "upstream_consent" not in captured, "cold reconnect required upstream consent"
landing = captured.get("url")
assert landing is not None, (
f"consent flow never reached {OAUTH_CLIENT_REDIRECT_URI}; "
f"final={final_url.split('?', 1)[0]!r}; trail={trail[-6:]}"
)
params = dict(parse_qsl(httpx.URL(landing).query.decode()))
assert "code" in params, f"client redirect_uri carried no code: {landing}"
assert "code" in params, "client redirect_uri carried no authorization code"
return params["code"], params.get("state")
def _oauth_provider(url: str, storage: InMemoryTokenStorage, storage_state_path: str) -> OAuthClientProvider:
def _oauth_provider(
url: str,
storage: InMemoryTokenStorage,
storage_state_path: str | None,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OAuthClientProvider:
"""The SDK's real OAuth machinery (RFC 9728/8414 discovery, RFC 7591 DCR,
PKCE, token exchange) with the browser leg driven by Playwright against the
upstream's consent screen."""
code_holder: dict[str, str | None] = {} # mutable-ok: hand-off between the two SDK callbacks
async def redirect_handler(authorize_url: str) -> None:
code, state = await _browser_follow_authorize(authorize_url, storage_state_path)
async def _reject_redirect(_: str) -> None:
raise AssertionError("gateway demanded a fresh upstream consent; stored per-user token was not reused")
async def _follow_redirect(authorize_url: str) -> None:
assert storage_state_path is not None
code, state = await _browser_follow_authorize(
authorize_url, storage_state_path, identity, server_alias, allow_upstream_consent
)
code_holder["code"] = code
code_holder["state"] = state
redirect_handler: Final = _reject_redirect if storage_state_path is None else _follow_redirect
async def callback_handler() -> AuthorizationCodeResult:
code = code_holder.get("code")
assert code is not None, "callback_handler ran before the authorize redirect completed"
@ -167,24 +224,45 @@ class _HeaderInjectingTransport(httpx2.AsyncBaseTransport):
store the upstream token for from the key on the token exchange, exactly
like a production MCP host configured with a LiteLLM key header."""
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str]) -> None:
def __init__(self, inner: httpx2.AsyncBaseTransport, headers: dict[str, str], gateway_url: str) -> None:
self._inner = inner
self._headers = headers
self._gateway_url = httpx2.URL(gateway_url)
@staticmethod
def _port(url: httpx2.URL) -> int | None:
if url.port is not None:
return url.port
return {"http": 80, "https": 443}.get(url.scheme)
async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
same_origin: Final = (
request.url.scheme == self._gateway_url.scheme
and request.url.host == self._gateway_url.host
and self._port(request.url) == self._port(self._gateway_url)
)
if same_origin:
for name, value in self._headers.items():
if name not in request.headers:
request.headers[name] = value
else:
for name, value in self._headers.items():
if request.headers.get(name) == value:
del request.headers[name]
return await self._inner.handle_async_request(request)
async def aclose(self) -> None:
await self._inner.aclose()
def _oauth_http_client(headers: dict[str, str], auth: OAuthClientProvider) -> httpx2.AsyncClient:
def _oauth_http_client(
headers: dict[str, str], auth: OAuthClientProvider, gateway_url: str = PROXY_BASE_URL
) -> httpx2.AsyncClient:
return httpx2.AsyncClient(
headers=headers,
auth=auth,
timeout=httpx2.Timeout(REQUEST_TIMEOUT),
follow_redirects=True,
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers),
transport=_HeaderInjectingTransport(httpx2.AsyncHTTPTransport(), headers, gateway_url),
)
@ -199,6 +277,43 @@ async def _seed_via_dance(
return tuple(sorted(tool.name for tool in listed.tools))
@dataclass(frozen=True, slots=True)
class OauthToolRun:
tools: tuple[str, ...]
is_error: bool
text: str
async def _list_and_call(
url: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
gateway_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
server_alias: str | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
async with _oauth_http_client(
headers,
_oauth_provider(url, storage, storage_state_path, identity, server_alias, allow_upstream_consent),
gateway_url,
) as http_client:
async with streamable_http_client(url, http_client=http_client) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
listed: Final = await session.list_tools()
result: Final = await session.call_tool(tool, arguments)
text: Final = "".join(content.text for content in result.content if isinstance(content, TextContent))
return OauthToolRun(
tools=tuple(sorted(tool_item.name for tool_item in listed.tools)),
is_error=result.is_error,
text=text,
)
@dataclass(frozen=True, slots=True)
class ChatMcpClient:
proxy: ProxyClient
@ -252,6 +367,53 @@ class ChatMcpClient:
f"last error: {last_error!r}"
)
def list_and_call(
self,
alias: str,
headers: dict[str, str],
storage: InMemoryTokenStorage,
storage_state_path: str | None,
tool: str,
arguments: dict[str, str],
base_url: str = PROXY_BASE_URL,
identity: Identity | None = None,
allow_upstream_consent: bool = True,
) -> OauthToolRun:
return asyncio.run(
_list_and_call(
f"{base_url.rstrip('/')}/mcp" if identity is not None else _mcp_url(alias, base_url),
headers,
storage,
storage_state_path,
tool,
arguments,
base_url,
identity,
alias,
allow_upstream_consent,
)
)
def server_user_credentials(self, server_id: str) -> tuple[McpServerUserCredentialRow, ...]:
return unwrap(
self.proxy.transport.get(
f"/v1/mcp/server/{server_id}/user-credentials",
headers=self.proxy.transport.master,
params=NoBody(),
response_type=McpServerUserCredentialListResponse,
)
).root
def revoke_user_token(self, server_id: str, headers: AuthHeaders) -> None:
_ = unwrap(
self.proxy.transport.delete(
f"/v1/mcp/server/{server_id}/oauth-user-credential",
headers=headers,
json=NoBody(),
response_type=McpOauthUserCredentialStatus,
)
)
def chat_with_mcp(self, headers: AuthHeaders, body: ChatBody) -> ChatResponse:
"""POST /chat/completions carrying the LiteLLM key in `headers` (either
ingress form) with an MCP server attached in `body.tools`. The gateway

198
tests/e2e/mcp/oauth_gateway.py Normal file
View file

@ -0,0 +1,198 @@
"""An owned, source-built OAuth gateway with cold restarts and credential observations.
Only this child process is restarted. Its database and SSO client survive while
its process-local caches do not; Redis is deliberately absent from its config.
The optional live edge measures headers without recording credentials or bodies.
"""
from __future__ import annotations
import os
import socket
import subprocess
import sys
import threading
import time
from collections.abc import Callable, Mapping
from contextlib import ExitStack
from dataclasses import dataclass, field
from pathlib import Path
from typing import Final
import psycopg
from e2e_http import NoBody
from idp import Keycloak, stop_process_group
from proxy_client import ProxyClient, build_proxy_client
from psycopg.rows import class_row
from pydantic import BaseModel, SecretStr, TypeAdapter, ValidationError
INHERITED_ENV_PREFIXES: Final = ("REDIS_", "MICROSOFT_", "GOOGLE_", "GENERIC_", "PROXY_")
class StoredOAuth(BaseModel):
type: str
access_token: SecretStr
@dataclass(frozen=True, slots=True)
class CredentialRow:
credential_b64: str = field(repr=False)
def stored_oauth(user_id: str, server_id: str) -> StoredOAuth:
"""Read the encrypted credential because management APIs omit the plaintext token."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
with psycopg.Connection[CredentialRow].connect(
os.environ["DATABASE_URL"], row_factory=class_row(CredentialRow)
) as conn:
row: Final = conn.execute(
'SELECT credential_b64 FROM "LiteLLM_MCPUserCredentials" WHERE user_id = %s AND server_id = %s',
(user_id, server_id),
).fetchone()
assert row is not None, "canonical user/server has no persisted credential"
plaintext: Final = decrypt_value_helper(
row.credential_b64, "e2e_mcp_oauth", exception_type="debug", return_original_value=False
)
assert plaintext is not None, "persisted credential must decrypt with the gateway salt"
assert plaintext != row.credential_b64, "persisted credential must be encrypted"
try:
credential: Final = StoredOAuth.model_validate_json(plaintext)
except ValidationError:
raise AssertionError("decrypted credential is not an OAuth payload") from None
assert credential.type == "oauth2"
assert bool(credential.access_token.get_secret_value()), "stored upstream token is empty"
return credential
class RpcMethod(BaseModel):
method: str = ""
@dataclass(slots=True)
class OAuthObservation:
gateway_token: str = field(default="", repr=False)
_seen: tuple[tuple[str, str, bool], ...] = field(default=(), init=False, repr=False)
_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False)
def observe(self, url: str, headers: Mapping[str, str], body: bytes | None) -> None:
if body is None or not url.endswith("/mcp"):
return
try:
operation: Final = RpcMethod.model_validate_json(body).method
except ValidationError:
return
if operation not in ("tools/list", "tools/call"):
return
received: Final = headers.get("authorization", "")
gateway_leaked: Final = any(
value in (self.gateway_token, f"Bearer {self.gateway_token}") for value in headers.values()
)
with self._lock:
self._seen = (*self._seen, (operation, received, gateway_leaked))
def assert_forwarded(self, expected: StoredOAuth) -> None:
with self._lock:
snapshot: Final = self._seen
self._seen = ()
assert {item[0] for item in snapshot} == {"tools/list", "tools/call"}, "missing upstream observations"
expected_header: Final = f"Bearer {expected.access_token.get_secret_value()}"
assert all(item[1] == expected_header for item in snapshot), "upstream bearer did not match the stored token"
assert all(not item[2] for item in snapshot), "gateway bearer leaked to the upstream"
def available_port() -> int:
with socket.socket() as listener:
listener.bind(("127.0.0.1", 0))
return TypeAdapter(tuple[str, int]).validate_python(listener.getsockname())[1]
@dataclass(slots=True)
class OAuthGateway:
base_url: str
proxy: ProxyClient
_environment: Mapping[str, str] = field(repr=False)
_command: tuple[str, ...] = field(repr=False)
_log_path: Path
_child: subprocess.Popen[bytes] | None = field(default=None, init=False, repr=False)
def start(self) -> None:
with self._log_path.open("ab") as log:
self._child = subprocess.Popen(
self._command,
env=self._environment,
stdout=log,
stderr=log,
start_new_session=True,
)
deadline: Final = time.monotonic() + 120
while time.monotonic() < deadline:
assert self._child.poll() is None, "owned OAuth gateway exited; inspect its private log"
result = self.proxy.transport.probe("/health/liveliness", params=NoBody())
if result.status_code == 200:
return
time.sleep(0.5)
raise AssertionError("owned OAuth gateway did not become ready")
def stop(self) -> None:
if self._child is not None:
stop_process_group(self._child)
assert self._child.poll() is not None, "old gateway process is still alive"
def restart(self) -> None:
assert self._child is not None
previous: Final = self._child.pid
self.stop()
self.start()
assert self._child.pid != previous, "gateway restart did not create a new process"
def owned_gateway(idp: Keycloak, directory: Path, cleanup: ExitStack) -> OAuthGateway:
for name in ("DATABASE_URL", "LITELLM_LICENSE", "LITELLM_SALT_KEY", "LITELLM_MASTER_KEY"):
assert os.environ.get(name), f"{name} is required for the owned OAuth gateway"
port: Final = available_port()
base_url: Final = f"http://127.0.0.1:{port}"
def defer(callback: Callable[[], object]) -> None:
cleanup.callback(callback)
browser: Final = idp.browser_client(callback_url=f"{base_url}/sso/callback", defer=defer)
config: Final = directory / "oauth-gateway.yaml"
config.write_text(
"model_list: []\n"
"general_settings:\n"
" master_key: os.environ/LITELLM_MASTER_KEY\n"
" database_url: os.environ/DATABASE_URL\n"
" enable_jwt_auth: true\n"
" litellm_jwtauth:\n"
" user_id_jwt_field: sub\n"
" user_email_jwt_field: email\n"
" team_ids_jwt_field: groups\n"
" user_id_upsert: true\n"
)
environment: Final = {
**{key: value for key, value in os.environ.items() if not key.startswith(INHERITED_ENV_PREFIXES)},
**browser.environment(idp.discovery()),
"PROXY_BASE_URL": base_url,
"JWT_PUBLIC_KEY_URL": idp.jwks_url,
"JWT_ISSUER": idp.issuer,
"JWT_AUDIENCE": "litellm-e2e",
"DISABLE_SCHEMA_UPDATE": "true",
"STORE_MODEL_IN_DB": "True",
"PYTHONPATH": str(Path(__file__).resolve().parents[3]),
}
gateway: Final = OAuthGateway(
base_url=base_url,
proxy=build_proxy_client(
base_url=base_url,
control_plane_base_url=base_url,
replica_urls=(base_url,),
master_key=os.environ["LITELLM_MASTER_KEY"],
),
_environment=environment,
_command=(sys.executable, "-m", "litellm.proxy.proxy_cli", "--config", str(config), "--port", str(port)),
_log_path=directory / "oauth-gateway.log",
)
cleanup.callback(gateway.stop)
gateway.start()
return gateway

207
tests/e2e/mcp/test_mcp_oauth_happy_path_e2e.py Normal file
View file

@ -0,0 +1,207 @@
"""Real OAuth consent, immediate MCP operations and cold-restart persistence.
Aggregate SSO uses the SDK's normal authentication. The per-server variant is
explicitly a configured two-header client, not an Authorization-only OAuth host.
The observed variants forward to the same real Linear upstream and compare its
bearer at the forwarding boundary; direct variants retain unmodified discovery.
"""
from __future__ import annotations
import os
from collections.abc import Iterator
from contextlib import ExitStack
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal
import pytest
from e2e_config import LINEAR_MCP_URL, LINEAR_READONLY_TOOL, LINEAR_STORAGE_STATE, unique_marker
from e2e_http import AuthHeaders, NoBody, get_external, unwrap
from idp import Identity, Keycloak
from lifecycle import ResourceManager
from models import (
McpOauthCredentials,
McpServerCreateBody,
ObjectPermission,
TeamMemberAddBody,
TeamMemberEntry,
TeamUpdateBody,
)
from oauth_chat_client import ChatMcpClient, InMemoryTokenStorage, OauthToolRun, build_chat_client
from oauth_gateway import OAuthGateway, OAuthObservation, owned_gateway, stored_oauth
from provider_edge import LiveEdge, start_provider_edge
from proxy_client import ProxyClient
from pydantic import BaseModel, ValidationError
pytestmark = [pytest.mark.e2e, pytest.mark.mcp_oauth_live, pytest.mark.provider_live]
class OAuthMetadata(BaseModel):
authorization_endpoint: str
token_endpoint: str
registration_endpoint: str
class LinearTeam(BaseModel):
id: str
name: str
class LinearTeams(BaseModel):
teams: tuple[LinearTeam, ...]
def assert_tool_result(run: OauthToolRun, tool: str) -> None:
assert tool in run.tools
assert run.is_error is False
try:
result: Final = LinearTeams.model_validate_json(run.text)
except ValidationError:
raise AssertionError("list_teams did not return the expected teams payload") from None
assert result.teams, "the test workspace must contain at least one team"
assert all(team.id and team.name for team in result.teams), "team results must contain identifiers and names"
@pytest.fixture(scope="module")
def oauth_gateway(idp: Keycloak, tmp_path_factory: pytest.TempPathFactory) -> Iterator[OAuthGateway]:
assert LINEAR_STORAGE_STATE and Path(LINEAR_STORAGE_STATE).is_file(), (
"E2E_LINEAR_STORAGE_STATE must name a captured Linear login; see mcp/linear_session_capture.py"
)
assert os.environ.get("E2E_FIXTURE_MODE", "live") == "live", "OAuth acceptance cannot use replay"
with ExitStack() as cleanup:
yield owned_gateway(idp, tmp_path_factory.mktemp("mcp-oauth"), cleanup)
@pytest.fixture(scope="module")
def proxy(oauth_gateway: OAuthGateway) -> ProxyClient:
return oauth_gateway.proxy
@pytest.fixture(scope="module")
def client(proxy: ProxyClient) -> ChatMcpClient:
return build_chat_client(proxy)
class TestMcpOauthHappyPath:
@pytest.mark.covers("mcp.list_tools.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.succeeds")
@pytest.mark.covers("mcp.call_tool.oauth.persists_across_processes")
@pytest.mark.parametrize("route", ("aggregate_sso", "explicit_header_jwt"))
@pytest.mark.parametrize("observed", (False, True), ids=("direct", "observed"))
def test_consent_list_call_and_cold_restart(
self,
client: ChatMcpClient,
resources: ResourceManager,
jwt_identity: Identity,
idp: Keycloak,
oauth_gateway: OAuthGateway,
route: Literal["aggregate_sso", "explicit_header_jwt"],
observed: bool,
) -> None:
alias: Final = f"e2elinear{unique_marker()}"
tool: Final = f"{alias}-{LINEAR_READONLY_TOOL}"
token: Final = idp.access_token(jwt_identity)
observation: Final = OAuthObservation(gateway_token=token)
edge: Final = (
start_provider_edge(
LiveEdge(observe_request=observation.observe),
mounts=MappingProxyType(
{"linear": "https://mcp.linear.app", ".well-known": "https://mcp.linear.app/.well-known"}
),
)
if observed
else None
)
if edge is not None:
resources.defer(edge.shutdown)
metadata: Final = (
unwrap(
get_external(
"https://mcp.linear.app/.well-known/oauth-authorization-server",
response_type=OAuthMetadata,
)
)
if observed
else None
)
created: Final = client.create_server(
McpServerCreateBody(
alias=alias,
server_name=alias,
url=f"{edge.edge.api_base('linear')}/mcp" if edge is not None else LINEAR_MCP_URL,
transport="http",
allow_all_keys=False,
auth_type="oauth2",
oauth2_flow="authorization_code",
per_server_oauth_discovery=route == "explicit_header_jwt",
authorization_url=metadata.authorization_endpoint if metadata else None,
token_url=metadata.token_endpoint if metadata else None,
registration_url=metadata.registration_endpoint if metadata else None,
credentials=McpOauthCredentials(upstream_resource=LINEAR_MCP_URL) if observed else None,
)
)
resources.defer(lambda: client.delete_server(created.server_id))
assert client.server_user_credentials(created.server_id) == (), (
"scenario must start without upstream credentials"
)
unwrap(
client.proxy.transport.post(
"/team/member_add",
headers=client.proxy.transport.master,
json=TeamMemberAddBody(
team_id=jwt_identity.group, member=TeamMemberEntry(user_id=jwt_identity.user_id, role="user")
),
response_type=NoBody,
)
)
client.proxy.update_team(
TeamUpdateBody(
team_id=jwt_identity.group,
object_permission=ObjectPermission(mcp_servers=[created.server_id]),
)
)
headers: Final = {"x-litellm-api-key": f"Bearer {token}"} if route == "explicit_header_jwt" else {}
resources.defer(
lambda: client.revoke_user_token(
created.server_id,
AuthHeaders(authorization=f"Bearer {idp.access_token(jwt_identity)}"),
)
)
identity: Final = jwt_identity if route == "aggregate_sso" else None
first: Final = client.list_and_call(
alias,
headers,
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
)
assert_tool_result(first, tool)
credentials: Final = client.server_user_credentials(created.server_id)
assert len(credentials) == 1
assert credentials[0].user_id == jwt_identity.user_id
assert credentials[0].credential_type == "oauth2"
first_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(first_stored_oauth)
oauth_gateway.restart()
fresh_token: Final = idp.access_token(jwt_identity)
observation.gateway_token = fresh_token
second: Final = client.list_and_call(
alias,
{"Authorization": f"Bearer {fresh_token}"} if identity is None else {},
InMemoryTokenStorage(),
LINEAR_STORAGE_STATE if identity is not None else None,
tool,
{},
base_url=oauth_gateway.base_url,
identity=identity,
allow_upstream_consent=False,
)
assert_tool_result(second, tool)
second_stored_oauth: Final = stored_oauth(jwt_identity.user_id, created.server_id)
if observed:
observation.assert_forwarded(second_stored_oauth)

32
tests/e2e/models.py
View file

@ -192,7 +192,7 @@ class ImageUrl(BaseModel):
class TextContentPart(BaseModel):
type: str = "text"
text: str
cache_control: "CacheControl | None" = None
cache_control: CacheControl | None = None
class ImageContentPart(BaseModel):
@ -572,6 +572,10 @@ class McpInfo(BaseModel):
logo_url: str | None = None
class McpOauthCredentials(BaseModel):
upstream_resource: str
class McpServerCreateBody(BaseModel):
"""POST /v1/mcp/server. For a gateway-managed OAuth server, `auth_type` is
`oauth2` and `oauth2_flow` is `authorization_code`; the upstream endpoints
@ -584,8 +588,11 @@ class McpServerCreateBody(BaseModel):
allow_all_keys: bool = True
auth_type: str | None = None
oauth2_flow: Literal["client_credentials", "authorization_code"] | None = None
per_server_oauth_discovery: bool | None = None
authorization_url: str | None = None
token_url: str | None = None
registration_url: str | None = None
credentials: McpOauthCredentials | None = None
server_name: str | None = None
description: str | None = None
mcp_info: McpInfo | None = None
@ -625,6 +632,26 @@ class McpServerListResponse(RootModel[list[McpServerRow]]):
"""GET /v1/mcp/server answers with a bare array of servers."""
class McpServerUserCredentialRow(BaseModel):
user_id: str
credential_type: Literal["oauth2", "byok"]
expires_at: str | None = None
connected_at: str | None = None
updated_at: str
class McpServerUserCredentialListResponse(RootModel[tuple[McpServerUserCredentialRow, ...]]):
"""GET /v1/mcp/server/{server_id}/user-credentials answers with a bare array."""
class McpOauthUserCredentialStatus(BaseModel):
server_id: str
has_credential: bool
expires_at: str | None = None
is_expired: bool = False
connected_at: str | None = None
class ToolsetTool(BaseModel):
server_id: str
tool_name: str
@ -1172,8 +1199,9 @@ class TeamNewResponse(BaseModel):
class TeamUpdateBody(BaseModel):
team_id: str
team_alias: str
team_alias: str | None = None
models: list[str] | None = None
object_permission: ObjectPermission | None = None
class TeamInfoParams(BaseModel):

12
tests/e2e/provider_edge.py
View file

@ -46,7 +46,7 @@ import os
import re
import threading
from collections import deque
from collections.abc import Generator, Mapping, Sequence
from collections.abc import Callable, Generator, Mapping, Sequence
from contextlib import closing, contextmanager
from dataclasses import dataclass, field, replace
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@ -538,7 +538,7 @@ class ReplayEdge:
@dataclass(frozen=True, slots=True)
class LiveEdge:
pass
observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None
type EdgeBackend = RecordEdge | ReplayEdge | LiveEdge | CacheEdge
@ -787,10 +787,13 @@ def _handle_record(
def _handle_live(
method: str, url: str, headers: Mapping[str, str], body: bytes | None, timeout: float,
cache: CacheEdge | None = None, mount: str = "", test_key: str | None = None,
observe_request: Callable[[str, Mapping[str, str], bytes | None], None] | None = None,
) -> EdgeOutcome:
forwarded: Final = {
name: value for name, value in headers.items() if name.lower() not in _REQUEST_DROPPED_HEADERS
}
if observe_request is not None:
observe_request(url, forwarded, body)
head: Final = (
forward_stream(method, url, headers=forwarded, body=body, timeout=timeout)
if cache is None else cache.forward(mount, method, url, forwarded, body, timeout, test_key=test_key)
@ -868,9 +871,10 @@ def handle_edge_request(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
backend, mount, test_key,
)
case LiveEdge():
case LiveEdge(observe_request=observe_request):
return _handle_live(
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout
method, _upstream_url(upstream_base, upstream_path, split.query), headers, body, timeout,
observe_request=observe_request,
)
case RecordEdge():
return _handle_record(

Some files were not shown because too many files have changed in this diff Show more