scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165

merge(litellm_internal_staging): sync latest staging and reconcile lint budgets
Some checks failed
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
Details

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
shivam 2026-07-31 00:29:59 +00:00
commit f7eaeed00e
172 changed files with 11116 additions and 2043 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/pull_request_template.md vendored
View file

@ -40,6 +40,7 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
The proof must be completely e2e with no mocks, using, for example, actual LLM calls costing real $. `pytest` commands are not enough
For bug fixes: show reproduction before the fix and passing behavior after
Include the commit hash each proof was captured at, for both the before and the after runs
If the change applies to all three LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), include proof for every single one of them, not just one
For new features: show the feature working end-to-end
For UI changes: include before/after screenshots -->

2
CLAUDE.md
View file

@ -61,6 +61,8 @@ Do not add `Co-Authored-By: Claude` or any Claude attribution to commit messages
When working on a PR, keep the PR description in sync with new commits being made
Replies/rebuttals to AI PR review bots must be 15-25 word human-readable replies
Monkeypatching attributes of a class to do testing is an anti-pattern. Prefer dependency-injecting things into classes. That way, at unit test time, you can pass a mocked dependency in
Do not put names of customers or customer company names in code, PR descriptions, issue bodies, etc. This means never mention literally any company name. Especially if you're about to say a sentence mentioning that the reason the PR exists was a feature/model/bug fix/etc. requested by a company. That's the indication that you should replace that company name with "the customer". e.g. not "Model request from Acme (Pylon #1234)" but "Model request from a customer (Pylon #1234)". This is because the codebase is public. The only exception is for publicly known providers or vendors such as OpenAI, Anthropic, AWS Bedrock, etc. only IF we're adding support for that provider/vendor in general and NOT if that PR or whatnot was a request by one of them, and they're actually one of our customers.

24
basedpyright-code-budget.json
View file

@ -6,7 +6,7 @@
"limit": 2645
},
"reportAssignmentType": {
"limit": 328
"limit": 329
},
"reportAttributeAccessIssue": {
"limit": 516
@ -18,7 +18,7 @@
"limit": 59
},
"reportDeprecated": {
"limit": 324
"limit": 325
},
"reportDuplicateImport": {
"limit": 42
@ -54,13 +54,13 @@
"limit": 0
},
"reportMissingParameterType": {
"limit": 5845
"limit": 5869
},
"reportMissingTypeArgument": {
"limit": 15846
"limit": 15861
},
"reportMissingTypeStubs": {
"limit": 41
"limit": 82
},
"reportOperatorIssue": {
"limit": 0
@ -72,7 +72,7 @@
"limit": 0
},
"reportOptionalMemberAccess": {
"limit": 1073
"limit": 1079
},
"reportOptionalOperand": {
"limit": 0
@ -90,7 +90,7 @@
"limit": 12
},
"reportReturnType": {
"limit": 217
"limit": 219
},
"reportTypedDictNotRequiredAccess": {
"limit": 27
@ -102,16 +102,16 @@
"limit": 45498
},
"reportUnknownLambdaType": {
"limit": 109
"limit": 113
},
"reportUnknownMemberType": {
"limit": 40458
"limit": 40477
},
"reportUnknownParameterType": {
"limit": 20302
"limit": 20338
},
"reportUnknownVariableType": {
"limit": 32026
"limit": 32047
},
"reportUnnecessaryCast": {
"limit": 177
@ -138,7 +138,7 @@
"limit": 206
},
"reportUnusedImport": {
"limit": 1001
"limit": 1003
},
"reportUnusedVariable": {
"limit": 1297

523
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_genai_otel/grafana_dashboard.json Normal file
View file

@ -0,0 +1,523 @@
{
"annotations": {
"list": []
},
"editable": true,
"fiscalYearStartMonth": 0,
"graphTooltip": 0,
"links": [],
"panels": [
{
"type": "stat",
"title": "Requests",
"gridPos": {
"h": 4,
"w": 6,
"x": 0,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"decimals": 0,
"color": {
"mode": "fixed",
"fixedColor": "blue"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 1
},
{
"type": "stat",
"title": "Spend",
"description": "LiteLLM's computed cost for the selected window, from gen_ai.usage.cost",
"gridPos": {
"h": 4,
"w": 6,
"x": 6,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "currencyUSD",
"decimals": 4,
"color": {
"mode": "fixed",
"fixedColor": "green"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 2
},
{
"type": "stat",
"title": "Tokens",
"gridPos": {
"h": 4,
"w": 6,
"x": 12,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"decimals": 0,
"color": {
"mode": "fixed",
"fixedColor": "purple"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "sum(increase(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range]))"
}
],
"id": 3
},
{
"type": "stat",
"title": "p95 request duration",
"gridPos": {
"h": 4,
"w": 6,
"x": 18,
"y": 0
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"decimals": 2,
"color": {
"mode": "fixed",
"fixedColor": "orange"
}
},
"overrides": []
},
"options": {
"reduceOptions": {
"calcs": [
"lastNotNull"
],
"fields": "",
"values": false
},
"colorMode": "background",
"graphMode": "none"
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"instant": true,
"expr": "histogram_quantile(0.95, sum by (le) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__range])))"
}
],
"id": 4
},
{
"type": "timeseries",
"title": "Request rate by model",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 4
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "reqpm",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
}
],
"id": 5
},
{
"type": "timeseries",
"title": "Spend rate by model",
"description": "USD per hour, derived from the gen_ai.usage.cost histogram",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 4
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "currencyUSD",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "sum by (gen_ai_request_model) (rate(gen_ai_usage_cost_USD_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 3600"
}
],
"id": 6
},
{
"type": "timeseries",
"title": "Tokens per minute by model and type",
"description": "gen_ai.client.token.usage split by the gen_ai.token.type attribute",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 12
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "short",
"custom": {
"lineWidth": 2,
"fillOpacity": 8,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}} {{gen_ai_token_type}}",
"expr": "sum by (gen_ai_request_model, gen_ai_token_type) (rate(gen_ai_client_token_usage_sum{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])) * 60"
}
],
"id": 7
},
{
"type": "timeseries",
"title": "p95 request duration by model",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 12
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_operation_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 8
},
{
"type": "timeseries",
"title": "p95 time to first token (streaming)",
"description": "gen_ai.server.time_to_first_token, recorded only for streaming requests",
"gridPos": {
"h": 8,
"w": 12,
"x": 0,
"y": 20
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_server_time_to_first_token_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 9
},
{
"type": "timeseries",
"title": "p95 provider generation time",
"description": "gen_ai.client.response.duration, upstream generation time excluding LiteLLM overhead",
"gridPos": {
"h": 8,
"w": 12,
"x": 12,
"y": 20
},
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"fieldConfig": {
"defaults": {
"unit": "s",
"custom": {
"lineWidth": 2,
"fillOpacity": 0,
"showPoints": "never"
}
},
"overrides": []
},
"options": {
"legend": {
"displayMode": "list",
"placement": "bottom"
},
"tooltip": {
"mode": "multi",
"sort": "desc"
}
},
"targets": [
{
"refId": "A",
"editorMode": "code",
"legendFormat": "{{gen_ai_request_model}}",
"expr": "histogram_quantile(0.95, sum by (le, gen_ai_request_model) (rate(gen_ai_client_response_duration_seconds_bucket{service_name=~\"$service\", gen_ai_request_model=~\"$model\"}[$__rate_interval])))"
}
],
"id": 10
}
],
"preload": false,
"refresh": "30s",
"schemaVersion": 42,
"tags": [
"litellm",
"genai",
"opentelemetry"
],
"templating": {
"list": [
{
"name": "datasource",
"label": "Prometheus",
"type": "datasource",
"query": "prometheus",
"current": {},
"hide": 0
},
{
"name": "service",
"label": "Service",
"type": "query",
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"query": "label_values(gen_ai_client_operation_duration_seconds_count, service_name)",
"refresh": 2,
"includeAll": true,
"multi": true,
"current": {
"text": "All",
"value": "$__all"
}
},
{
"name": "model",
"label": "Model",
"type": "query",
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"query": "label_values(gen_ai_client_operation_duration_seconds_count{service_name=~\"$service\"}, gen_ai_request_model)",
"refresh": 2,
"includeAll": true,
"multi": true,
"current": {
"text": "All",
"value": "$__all"
}
}
]
},
"time": {
"from": "now-1h",
"to": "now"
},
"timepicker": {},
"timezone": "browser",
"title": "LiteLLM GenAI (OpenTelemetry)",
"uid": "litellm-genai-otel",
"weekStart": ""
}

35
cookbook/litellm_proxy_server/grafana_dashboard/dashboard_genai_otel/readme.md Normal file
View file

@ -0,0 +1,35 @@
# LiteLLM GenAI dashboard (OpenTelemetry metrics)
Dashboard for the `gen_ai.*` metrics the OpenTelemetry v2 integration emits, as opposed to the `litellm_*` Prometheus metrics the other dashboards in this folder chart.
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source. Panels: request count, spend, token count, p95 duration, request rate by model, spend rate per hour by model, tokens per minute split by input and output, p95 duration by model, p95 time to first token, and p95 provider generation time. Template variables for data source, service, and model.
## Pre-requisites
Metrics are off by default. In the proxy environment:
```shell
LITELLM_OTEL_V2=true
LITELLM_OTEL_INTEGRATION_ENABLE_METRICS=true
OTEL_EXPORTER="otlp_http"
OTEL_ENDPOINT="<your OTLP endpoint>"
```
You also need the metric attribute filter, or the panels will plot flat lines at zero. LiteLLM's default attribute set includes per-request fields, so nearly every request lands in its own time series with a single sample, and `rate()` has nothing to compute over:
```yaml title="config.yaml"
callback_settings:
otel:
attributes:
include_list:
- gen_ai.operation.name
- gen_ai.system
- gen_ai.request.model
- gen_ai.framework
```
See [Grafana Cloud](https://docs.litellm.ai/docs/observability/grafana_cloud) for the full setup, and [OpenTelemetry v2](https://docs.litellm.ai/docs/observability/opentelemetry_v2#metrics) for the metric reference.
## Note on Grafana's AI Observability integration
Grafana Cloud ships prebuilt GenAI dashboards that query these same metric names, so they look like a drop-in alternative to this one. They are not: twenty of their twenty-two panels filter on `telemetry_sdk_name="openlit"`, a label LiteLLM does not carry and cannot be configured to add, so those panels stay empty.

4
cookbook/litellm_proxy_server/grafana_dashboard/readme.md
View file

@ -2,6 +2,10 @@
This folder contains the `json` for creating Grafana Dashboards
## [LiteLLM GenAI Dashboard (OpenTelemetry)](./dashboard_genai_otel)
Charts the `gen_ai.*` metrics from the OpenTelemetry v2 integration: spend, tokens, request rate, and latency percentiles by model. Separate from the dashboards below, which chart the `litellm_*` Prometheus metrics.
## [LiteLLM v2 Dashboard](./dashboard_v2)
<img width="1316" alt="grafana_1" src="https://github.com/user-attachments/assets/d0df802d-0cb9-4906-a679-941c547789ab">

106
litellm-proxy-extras/litellm_proxy_extras/replica_identity.py Normal file
View file

@ -0,0 +1,106 @@
"""Optional post-migration step that raises Postgres REPLICA IDENTITY to FULL.
Logical-replication consumers (Neon / lakehouse sync and similar) need FULL
replica identity to reconstruct the old row of an UPDATE or DELETE. Prisma
leaves every table it creates at the Postgres default, so the setting has to be
re-applied by hand after each migration run. Setting
``LITELLM_SET_REPLICA_IDENTITY_FULL`` makes every migration run re-assert it.
The statement goes through the Prisma CLI rather than a Postgres driver because
``litellm-proxy-extras`` has no runtime dependencies, while the CLI is already
required for the migrations themselves.
"""
import subprocess
import tempfile
from pathlib import Path
from litellm_proxy_extras._logging import logger
REPLICA_IDENTITY_FULL_ENV_VAR = "LITELLM_SET_REPLICA_IDENTITY_FULL"
REPLICA_IDENTITY_FULL_SQL = r"""
DO $$
DECLARE
target regclass;
BEGIN
SET LOCAL lock_timeout = '5s';
FOR target IN
SELECT c.oid::regclass
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind = 'r'
AND c.relreplident <> 'f'
AND n.nspname = ANY (current_schemas(false))
AND c.relname LIKE 'LiteLLM\_%'
LOOP
BEGIN
EXECUTE format('ALTER TABLE %s REPLICA IDENTITY FULL', target);
EXCEPTION WHEN lock_not_available THEN
RAISE WARNING 'REPLICA IDENTITY FULL skipped for %: table busy, retrying next run', target;
END;
END LOOP;
END
$$;
"""
def apply_replica_identity_full(
schema_path: str,
prisma_command: str,
prisma_env: dict[str, str],
) -> bool:
"""Set REPLICA IDENTITY FULL on every LiteLLM table that is not already FULL.
Never raises. Replication metadata is not needed to serve requests, so
every failure mode is reported and stepped over rather than taking down a
migration run that already succeeded: a database that refuses the ALTER
(most often because the runtime user does not own the tables), a missing
or unrunnable Prisma CLI, a read-only temp directory, or a timeout.
Returns True when the statement was applied, False when it failed.
"""
logger.info("Applying REPLICA IDENTITY FULL to LiteLLM tables")
try:
with tempfile.TemporaryDirectory(prefix="litellm_replica_identity_") as tmp_dir:
sql_path = Path(tmp_dir) / "replica_identity_full.sql"
sql_path.write_text(REPLICA_IDENTITY_FULL_SQL)
subprocess.run(
[
prisma_command,
"db",
"execute",
"--file",
str(sql_path),
"--schema",
schema_path,
],
timeout=60,
check=True,
capture_output=True,
text=True,
env=prisma_env,
)
except subprocess.CalledProcessError as e:
logger.error(
"Failed to set REPLICA IDENTITY FULL. Logical replication "
"consumers may reject updates to these tables. Grant table "
"ownership to the migration user, or apply "
"`ALTER TABLE ... REPLICA IDENTITY FULL` by hand. Error: %s",
e.stderr,
)
return False
except subprocess.TimeoutExpired:
logger.error("Timed out setting REPLICA IDENTITY FULL on LiteLLM tables")
return False
except OSError as e:
logger.error(
"Could not run the REPLICA IDENTITY FULL statement. Logical "
"replication consumers may reject updates to these tables. "
"Error: %s",
e,
)
return False
logger.info("REPLICA IDENTITY FULL applied to LiteLLM tables")
return True

46
litellm-proxy-extras/litellm_proxy_extras/utils.py
View file

@ -10,6 +10,10 @@ from pathlib import Path
from typing import Optional
from litellm_proxy_extras._logging import logger
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
def str_to_bool(value: Optional[str]) -> bool:
@ -676,6 +680,39 @@ class ProxyExtrasDBManager:
finally:
os.chdir(original_dir)
@staticmethod
def apply_replica_identity_full_if_requested() -> bool:
"""
Re-assert REPLICA IDENTITY FULL on LiteLLM's tables when the operator
opted in via LITELLM_SET_REPLICA_IDENTITY_FULL.
Prisma leaves new tables at the Postgres default, which logical
replication consumers reject, so the setting has to be re-applied after
every migration run rather than once by hand.
Returns:
bool: True if the setting was applied, False if it was not
requested or could not be applied.
"""
if not str_to_bool(os.getenv(REPLICA_IDENTITY_FULL_ENV_VAR)):
return False
try:
schema_path = ProxyExtrasDBManager._get_prisma_dir() + "/schema.prisma"
prisma_command = _get_prisma_command()
prisma_env = _get_prisma_env()
except OSError as e:
logger.error(
"Could not resolve the migrations directory for the REPLICA "
"IDENTITY FULL step, skipping it. Error: %s",
e,
)
return False
return apply_replica_identity_full(
schema_path=schema_path,
prisma_command=prisma_command,
prisma_env=prisma_env,
)
@staticmethod
def setup_database(
use_migrate: bool = False, use_v2_resolver: bool = False
@ -694,6 +731,15 @@ class ProxyExtrasDBManager:
Returns:
bool: True if setup was successful, False otherwise
"""
migrated = ProxyExtrasDBManager._run_migrations(
use_migrate=use_migrate, use_v2_resolver=use_v2_resolver
)
if migrated:
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
return migrated
@staticmethod
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
if use_v2_resolver:
logger.info("Using v2 migration resolver (--use_v2_migration_resolver)")
return ProxyExtrasDBManager._setup_database_v2(use_migrate=use_migrate)

64
litellm/_redis.py
View file

@ -61,23 +61,51 @@ def _get_redis_kwargs():
return available_args
def _get_redis_url_kwargs(client=None):
def _init_arg_names(cls: type) -> frozenset[str]:
"""Every ``__init__`` parameter accepted anywhere in a class's MRO.
Keyword-only parameters are included, and the MRO is walked because redis-py splits a
connection's parameters between ``AbstractConnection`` and its concrete subclasses.
"""
return frozenset(
name
for klass in inspect.getmro(cls)
if klass is not object
for spec in (inspect.getfullargspec(klass.__init__),)
for name in spec.args + spec.kwonlyargs
)
def _get_redis_url_kwargs(client: Optional[type] = None) -> tuple[str, ...]:
"""Connection kwargs that redis-py forwards from ``from_url`` down to the connection.
``from_url`` is declared as ``(cls, url, **kwargs)``, so introspecting it yields no
connection kwargs at all. What it really does is hand its kwargs to the connection
class, so that class's signature is the allowlist.
Taking the client's signature instead would be wrong in both directions: it omits
nothing useful, but it admits client-only parameters such as
``single_connection_client`` and ``auto_close_connection_pool``, plus the ``ssl_*``
family that only ``SSLConnection`` accepts. Those reach ``AbstractConnection`` and
raise ``TypeError`` the first time a connection is created. TLS on a url config is
selected by the ``rediss://`` scheme, which picks ``SSLConnection`` on its own.
"""
if client is None:
client = redis.Redis.from_url
arg_spec = inspect.getfullargspec(redis.Redis.from_url)
client = redis.Redis
connection_cls = async_redis.Connection if client is async_redis.Redis else redis.Connection
exclude_args = frozenset(
{
"self",
"connection_pool",
"retry",
}
)
# Only allow primitive arguments
exclude_args = {
"self",
"connection_pool",
"retry",
}
include_args = ("url", "max_connections")
include_args = ["url"]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
return available_args
return tuple(x for x in _init_arg_names(connection_cls) if x not in exclude_args) + include_args
def _get_redis_cluster_kwargs(client=None):
@ -614,7 +642,7 @@ def get_redis_async_client(
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
if connection_pool is not None:
return async_redis.Redis(connection_pool=connection_pool)
args = _get_redis_url_kwargs(client=async_redis.Redis.from_url)
args = _get_redis_url_kwargs(client=async_redis.Redis)
url_kwargs = {}
for arg in redis_kwargs:
if arg in args:
@ -662,10 +690,10 @@ def get_redis_connection_pool(
return None
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
pool_kwargs = {
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
"url": redis_kwargs["url"],
}
allowed_args = _get_redis_url_kwargs(client=async_redis.Redis)
pool_kwargs = {k: v for k, v in redis_kwargs.items() if k in allowed_args and k != "max_connections"}
pool_kwargs["timeout"] = REDIS_CONNECTION_POOL_TIMEOUT
pool_kwargs["url"] = redis_kwargs["url"]
if "max_connections" in redis_kwargs:
try:
pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])

1
litellm/a2a_protocol/main.py
View file

@ -568,6 +568,7 @@ def _build_streaming_logging_obj(
logging_obj.custom_llm_provider = "a2a_agent"
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "a2a_agent"
logging_obj.model_call_details["call_type"] = logging_obj.call_type
if agent_id:
logging_obj.model_call_details["agent_id"] = agent_id

299
litellm/batches/batch_utils.py
View file

@ -1,5 +1,6 @@
import json
from typing import Any, Iterator, List, Literal, Optional, Tuple
from dataclasses import dataclass
from typing import Any, Iterable, Iterator, List, Literal, Optional, Tuple
import litellm
from litellm._logging import verbose_logger
@ -24,20 +25,20 @@ async def calculate_batch_cost_and_usage(
deployment-specific pricing (e.g. input_cost_per_token_batches)
is used instead of the global cost map.
"""
batch_cost = _batch_cost_calculator(
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return batch_cost, batch_usage, [model_name]
return _aggregate_batch_cost_usage_models(
entries=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
file_content_dictionary=file_content_dictionary,
model_name=model_name,
model_info=model_info,
)
batch_usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
)
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name, custom_llm_provider)
return batch_cost, batch_usage, batch_models
async def _handle_completed_batch(
@ -46,7 +47,9 @@ async def _handle_completed_batch(
model_name: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
"""Helper function to process a completed batch and handle logging
"""Fetch a completed batch's output file and aggregate its cost, usage, and
models in a single pass over the JSONL lines, so the parsed file content is
never materialized in memory.
Args:
batch: The batch object
@ -54,75 +57,109 @@ async def _handle_completed_batch(
model_name: Optional model name
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
"""
# Get batch results
file_content_dictionary = await _get_batch_output_file_content_as_dictionary(
batch, custom_llm_provider, litellm_params=litellm_params
)
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
# Calculate costs and usage
batch_cost = _batch_cost_calculator(
custom_llm_provider=custom_llm_provider,
file_content_dictionary=file_content_dictionary,
model_name=model_name,
)
batch_usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_name=model_name,
)
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name, custom_llm_provider)
return batch_cost, batch_usage, batch_models
def _get_batch_models_from_file_content(
file_content_dictionary: List[dict],
model_name: Optional[str] = None,
custom_llm_provider: str = "openai",
) -> List[str]:
"""
Get the models from the file content
"""
if model_name:
return [model_name]
batch_models = []
for _item in file_content_dictionary:
if _batch_response_was_successful(_item, custom_llm_provider):
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
_model = _response_body.get("model")
if _model:
batch_models.append(_model)
return batch_models
def _batch_cost_calculator(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> float:
"""
Calculate the cost of a batch based on the output file id
"""
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost)
return batch_cost
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
_get_file_content_as_dictionary(file_content), model_name
)
return batch_cost, batch_usage, [model_name]
# For other providers, use the existing logic
total_cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content_dictionary,
return _aggregate_batch_cost_usage_models(
entries=_iter_batch_input_entries(file_content),
custom_llm_provider=custom_llm_provider,
model_name=model_name,
model_info=model_info,
)
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
@dataclass(frozen=True, slots=True)
class _BatchOutputLineStats:
cost: float
prompt_tokens: int
completion_tokens: int
total_tokens: int
cache_read_tokens: int
cache_creation_tokens: int
model: Optional[str]
def _iter_successful_output_line_stats(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str],
model_info: Optional[ModelInfo],
) -> Iterator[_BatchOutputLineStats]:
from litellm.cost_calculator import batch_cost_calculator
for entry in entries:
if not _batch_response_was_successful(entry, custom_llm_provider):
continue
response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
prompt_details = _parse_prompt_tokens_details(usage)
raw_model = response_body.get("model")
response_model = raw_model if isinstance(raw_model, str) and raw_model else None
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
if custom_llm_provider == "bedrock" and model_name:
cost_model = model_name
else:
cost_model = response_model or model_name or ""
prompt_cost, completion_cost = batch_cost_calculator(
usage=usage,
model=cost_model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
line_cost = prompt_cost + completion_cost
else:
line_cost = litellm.completion_cost(
completion_response=response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
yield _BatchOutputLineStats(
cost=line_cost,
prompt_tokens=usage.prompt_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens,
cache_read_tokens=prompt_details["cache_hit_tokens"],
cache_creation_tokens=prompt_details["cache_creation_tokens"],
model=response_model,
)
def _aggregate_batch_cost_usage_models(
entries: Iterable[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
"""Aggregate cost, usage, and models from batch output entries in a single
pass, holding one small stats record per line instead of the parsed file."""
line_stats = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
cache_token_params = {
key: tokens
for key, tokens in (
("cache_read_input_tokens", sum(stats.cache_read_tokens for stats in line_stats)),
("cache_creation_input_tokens", sum(stats.cache_creation_tokens for stats in line_stats)),
)
if tokens > 0
}
batch_usage = Usage(
total_tokens=sum(stats.total_tokens for stats in line_stats),
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
**cache_token_params,
)
batch_models = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
total_cost = sum((stats.cost for stats in line_stats), 0.0)
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
return total_cost, batch_usage, batch_models
def calculate_vertex_ai_batch_cost_and_usage(
@ -193,13 +230,13 @@ def calculate_vertex_ai_batch_cost_and_usage(
)
async def _get_batch_output_file_content_as_dictionary(
async def _fetch_batch_output_file_content(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
litellm_params: Optional[dict] = None,
) -> List[dict]:
) -> bytes:
"""
Get the batch output file content as a list of dictionaries
Fetch the batch output file and return its raw JSONL bytes
Args:
batch: The batch object
@ -212,9 +249,6 @@ async def _get_batch_output_file_content_as_dictionary(
_is_base64_encoded_unified_file_id,
)
if custom_llm_provider == "vertex_ai":
raise ValueError("Vertex AI does not support file content retrieval")
if batch.output_file_id is None:
raise ValueError("Output file id is None cannot retrieve file content")
@ -240,7 +274,7 @@ async def _get_batch_output_file_content_as_dictionary(
file_content_kwargs.update(credentials)
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
return _get_file_content_as_dictionary(_file_content.content)
return _file_content.content
def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
@ -270,6 +304,8 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
"vertex_project",
"vertex_location",
"vertex_credentials",
"gcs_bucket_name",
"bucket_name",
"timeout",
"max_retries",
]
@ -284,17 +320,7 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
"""
Get the file content as a list of dictionaries from JSON Lines format
"""
try:
_file_content_str = file_content.decode("utf-8")
# Split by newlines and parse each line as a separate JSON object
json_objects = []
for line in _file_content_str.strip().split("\n"):
if line: # Skip empty lines
json_objects.append(json.loads(line))
verbose_logger.debug("json_objects=%s", json.dumps(json_objects, indent=4))
return json_objects
except Exception as e:
raise e
return list(_iter_batch_input_entries(file_content))
def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
@ -361,101 +387,6 @@ def _count_entry_tokens(
return 0
def _get_batch_job_cost_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> float:
"""
Get the cost of a batch job from the file content
"""
from litellm.cost_calculator import batch_cost_calculator
try:
total_cost: float = 0.0
# parse the file content as json
verbose_logger.debug("file_content_dictionary=%s", json.dumps(file_content_dictionary, indent=4))
for _item in file_content_dictionary:
if _batch_response_was_successful(_item, custom_llm_provider):
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
# Bedrock batch output lines report a short internal model id
# (e.g. "claude-sonnet-4-6") that is not in the cost map; use the
# deployment model name for pricing when available.
if custom_llm_provider == "bedrock" and model_name:
model = model_name
else:
model = _response_body.get("model") or model_name or ""
prompt_cost, completion_cost = batch_cost_calculator(
usage=usage,
model=model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
total_cost += prompt_cost + completion_cost
else:
total_cost += litellm.completion_cost(
completion_response=_response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
except Exception as e:
verbose_logger.error("error in _get_batch_job_cost_from_file_content", e)
raise e
def _get_batch_job_total_usage_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_name: Optional[str] = None,
) -> Usage:
"""
Get the tokens of a batch job from the file content
"""
if (
custom_llm_provider == "vertex_ai"
and model_name
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
):
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
return batch_usage
# For other providers, use the existing logic
total_tokens: int = 0
prompt_tokens: int = 0
completion_tokens: int = 0
cache_read_tokens: int = 0
cache_creation_tokens: int = 0
for _item in file_content_dictionary:
if _batch_response_was_successful(_item, custom_llm_provider):
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
usage: Usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
total_tokens += usage.total_tokens
prompt_tokens += usage.prompt_tokens
completion_tokens += usage.completion_tokens
prompt_details = _parse_prompt_tokens_details(usage)
cache_read_tokens += prompt_details["cache_hit_tokens"]
cache_creation_tokens += prompt_details["cache_creation_tokens"]
cache_token_params = {
key: tokens
for key, tokens in (
("cache_read_input_tokens", cache_read_tokens),
("cache_creation_input_tokens", cache_creation_tokens),
)
if tokens > 0
}
return Usage(
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
**cache_token_params,
)
def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
"""Token-count a ``prompt`` / ``input`` field that the OpenAI batch
schema allows in four shapes:

1
litellm/caching/caching_handler.py
View file

@ -517,6 +517,7 @@ class LLMCachingHandler:
cached_result=final_embedding_cached_response,
is_async=True,
is_embedding=True,
custom_llm_provider=custom_llm_provider,
)
self._async_log_cache_hit_on_callbacks(
logging_obj=logging_obj,

117
litellm/caching/redis_cache.py
View file

@ -17,7 +17,8 @@ import json
import time
from collections.abc import Awaitable, Callable, Sequence
from datetime import timedelta
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast
from contextvars import ContextVar
from typing import TYPE_CHECKING, Any, List, Optional, Tuple, TypeVar, Union, cast
import litellm
from litellm._logging import print_verbose, verbose_logger
@ -168,24 +169,97 @@ class RedisCircuitBreaker:
self._state = self.CLOSED
_RedisCallResult = TypeVar("_RedisCallResult")
_swallowed_redis_failures: ContextVar[int] = ContextVar("litellm_swallowed_redis_failures", default=0)
@functools.lru_cache(maxsize=1)
def _redis_health_error_types() -> tuple[type, ...]:
"""Exception types that mean the Redis backend itself is unhealthy.
Command and data errors say nothing about connectivity: an INCR against a non-numeric
value or an undecodable cached entry is a request problem, and counting those would let
a caller trip the shared breaker on demand, dropping rate limiting to per-process
counters that spreading traffic across replicas can outrun.
Imported lazily because this module is reachable from a base ``import litellm`` while
redis is not a base dependency.
"""
from redis.exceptions import BusyLoadingError, ClusterDownError
from redis.exceptions import ConnectionError as RedisConnectionError
from redis.exceptions import TimeoutError as RedisTimeoutError
return (RedisConnectionError, RedisTimeoutError, BusyLoadingError, ClusterDownError, OSError, asyncio.TimeoutError)
def _is_redis_health_failure(exc: BaseException) -> bool:
"""True when ``exc`` indicates Redis is unreachable rather than the request being bad."""
try:
return isinstance(exc, _redis_health_error_types())
except ImportError:
return True
def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
"""Record a Redis failure that the calling method is about to swallow.
The marker is a ContextVar rather than a counter on the breaker because breakers are
shared by every concurrent caller. A plain shared counter cannot tell "my call failed"
from "some other in-flight call failed", so a success overlapping someone else's
failure would be discarded and a Redis that is answering would still be evicted.
asyncio gives each task its own copy of the context, so this is per-call.
"""
if not _is_redis_health_failure(exc):
return
breaker.record_failure()
_swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
async def _run_under_circuit_breaker(
breaker: RedisCircuitBreaker,
name: str,
call: Callable[[], Awaitable[_RedisCallResult]],
) -> _RedisCallResult:
"""Run one Redis coroutine under a circuit breaker.
Shared by the method decorator and the Lua script executor so both feed the same
health signal. Success is recorded only when nothing failed while ``call`` ran,
because several Redis methods catch their own connection errors and return a default.
"""
if breaker.is_open():
raise Exception(f"Redis circuit breaker is open — skipping {name}")
swallowed_before = _swallowed_redis_failures.get()
try:
result = await call()
except Exception as e:
if _is_redis_health_failure(e):
breaker.record_failure()
raise
if _swallowed_redis_failures.get() == swallowed_before:
breaker.record_success()
return result
def _redis_circuit_breaker_guard(method): # type: ignore
"""
Decorator for RedisCache async methods.
Checks the circuit breaker before each call; records success/failure after.
Does not apply to ping/disconnect/test_connection (health/teardown must always run).
A returning method is not proof of a healthy Redis: several methods catch their own
connection errors and return a default so callers degrade rather than fail. Counting
those as successes reset the failure streak on every request, so the breaker could
never open and Redis was never taken out of the pool. Success is therefore recorded
only when no failure was registered while the method ran.
"""
@functools.wraps(method)
async def wrapper(self, *args, **kwargs): # type: ignore
if self._circuit_breaker.is_open():
raise Exception(f"Redis circuit breaker is open — skipping {method.__name__}")
try:
result = await method(self, *args, **kwargs)
self._circuit_breaker.record_success()
return result
except Exception:
self._circuit_breaker.record_failure()
raise
return await _run_under_circuit_breaker(
self._circuit_breaker, method.__name__, lambda: method(self, *args, **kwargs)
)
return wrapper
@ -551,13 +625,16 @@ class RedisCache(BaseCache):
)
async def run_script(keys: Sequence[str], args: Sequence[Any], client: Any = None) -> Any:
executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
)
if executor is None:
executor = self._register_script_for_current_loop(script)
litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
return await executor(keys=keys, args=args, client=client)
async def execute() -> object:
executor: Optional[Callable[..., Awaitable[Any]]] = litellm.in_memory_llm_clients_cache.get_cache(
key=script_cache_key
)
if executor is None:
executor = self._register_script_for_current_loop(script)
litellm.in_memory_llm_clients_cache.set_cache(key=script_cache_key, value=executor)
return await executor(keys=keys, args=args, client=client)
return await _run_under_circuit_breaker(self._circuit_breaker, "run_script", execute)
return run_script
@ -674,6 +751,7 @@ class RedisCache(BaseCache):
str(e),
value,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
async def _pipeline_helper(
self,
@ -758,6 +836,7 @@ class RedisCache(BaseCache):
str(e),
cache_value,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
async def _set_cache_sadd_helper(
self,
@ -842,6 +921,7 @@ class RedisCache(BaseCache):
str(e),
value,
)
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def batch_cache_write(self, key, value, **kwargs):
@ -1106,6 +1186,7 @@ class RedisCache(BaseCache):
)
)
print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {str(e)}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
@_redis_circuit_breaker_guard
async def async_batch_get_cache(
@ -1177,6 +1258,7 @@ class RedisCache(BaseCache):
)
)
verbose_logger.error(f"Error occurred in async batch get cache - {str(e)}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
return key_value_dict
def sync_ping(self) -> bool:
@ -1432,6 +1514,7 @@ class RedisCache(BaseCache):
return ttl
except Exception as e:
verbose_logger.debug(f"Redis TTL Error: {e}")
_record_swallowed_redis_failure(self._circuit_breaker, e)
return None
@_redis_circuit_breaker_guard

8
litellm/integrations/custom_logger.py
View file

@ -519,7 +519,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
This log gets called after the MCP tool call is made.
Useful if you want to modiy the standard logging payload after the MCP tool call is made.
Useful if you want to modify the standard logging payload after the MCP tool call is made.
To change what the caller sends back to the MCP client, mutate ``response_obj``
in place: every call site discards the returned object, because the
dispatcher unwraps it to ``mcp_tool_call_response`` (a raw content list, not
a ``CallToolResult``) which the tool-call paths cannot forward. Guardrails
that mask or reject tool output should use ``post_mcp_call`` instead.
"""
return None

1
litellm/integrations/opentelemetry.py
View file

@ -118,6 +118,7 @@ TOKEN_TYPE_ATTRIBUTE: str = "gen_ai.token.type"
VALID_METRIC_ATTRIBUTE_NAMES: FrozenSet[str] = frozenset(
(
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",

29
litellm/integrations/otel/README.md
View file

@ -222,7 +222,29 @@ lives in [`plumbing/`](./plumbing):
otherwise the operator's globally configured `MeterProvider` is reused so its
readers/exporters receive them alongside the server metrics, and one is built
and registered as the global only when none is set (mirroring how V2 owns trace
export).
export). A **failed** call records `gen_ai.client.operation.duration` too,
carrying the semconv `error.type` (the mapped provider exception's class name),
so the histogram covers the whole traffic and failure-rate panels are buildable;
the other five instruments describe a completed generation and are skipped
rather than filled with a fabricated zero. `error.type` is stamped after the
cardinality filter, so an `otel.attributes` list cannot merge the failure series
back into the success series. A proxy-gate rejection (auth / rate limit) records
nothing, for the same reason it gets no span: no upstream call happened.
Both paths cap their attributes at `METRIC_ATTRIBUTE_CEILING` before the
operator's own `otel.attributes` filter runs, so the filter can narrow the set
but never widen it. The ceiling is what keeps series count bounded by the
deployment's own key/team/user/deployment count instead of by its traffic: a
label value that moves per request mints a time series per request, which both
bills per request on a hosted backend and leaves a histogram that cannot be
aggregated. So client-supplied and per-request metadata (`requester_metadata`,
`spend_logs_metadata`, `user_api_key_end_user_id`, `requester_ip_address`) is
metric-ineligible and stays on the span, where cardinality is free, and the
`hidden_params` label carries only `model_id`, the deployment identity a
per-deployment panel joins on. `api_base` is excluded despite naming the same
deployment, because it is a documented per-call parameter and so is caller-chosen
in SDK use. Because the shared validator accepts every span attribute name, a
filter that names a metric-ineligible one logs a warning once when the filter
resolves rather than silently emitting nothing for it.
- [`events.py`](./plumbing/events.py) — GenAI client events. Gated on
`enable_events` (`LITELLM_OTEL_INTEGRATION_ENABLE_EVENTS`), a failed LLM call
records the semconv `gen_ai.client.operation.exception` log event at severity
@ -268,7 +290,10 @@ lives in [`plumbing/`](./plumbing):
- **A new attribute vocabulary for a backend**: add a mapper in `mappers/`
(a class with a `map(data) -> AttributeMap` method, typically built from
`key -> extractor` tables) and register it in `mappers/__init__._MAPPER_BY_NAME`.
`key -> extractor` tables) and register it in `mappers/__init__._PLAIN_MAPPERS`.
If it spells declared tool definitions out per index, register it in
`_TOOL_DEFINITION_MAPPERS` instead and take the shared attribute budget in its
constructor, so the family stays bounded span-wide rather than per vocabulary.
- **A new integration**: add a preset in `presets/` that returns an
`OpenTelemetryV2Config`, and register it in `presets/__init__.PRESET_BY_CALLBACK`.
For admin-owned per-key/team destinations, add an adapter mapping the named

25
litellm/integrations/otel/logger.py
View file

@ -275,13 +275,29 @@ class OpenTelemetryV2(CustomLogger):
self._record_metrics(kwargs, response_obj, start_time, end_time)
def _record_metrics(self, kwargs, response_obj, start_time, end_time) -> None:
"""Record the GenAI metrics for a successful LLM call. Best-effort: a
recording failure (e.g. a malformed payload) must never break the span
close or the request itself."""
"""Record the GenAI metrics for a successful LLM call."""
self._guarded_record(lambda recorder: recorder.record(kwargs, response_obj, start_time, end_time))
def _record_failure_metrics(self, kwargs, start_time, end_time) -> None:
"""Record the GenAI metrics for a failed LLM call, so the duration
histogram covers the whole traffic rather than only what survived.
A synthetic proxy-gate log (auth / rate-limit rejection) is skipped for the
same reason it gets no span: no upstream call happened, so its duration is
not a GenAI operation's duration and would pull the histogram down."""
if LLMCallEvent.from_dict(kwargs).is_no_upstream_call:
return
self._guarded_record(lambda recorder: recorder.record_failure(kwargs, start_time, end_time))
def _guarded_record(self, record: "Callable[[GenAIMetricRecorder], None]") -> None:
"""Run one metric recording. Best-effort: a recording failure (e.g. a
malformed payload) must never break the span close or the request itself. A
misconfigured attribute filter is operator-fixable, so it is surfaced once
at ERROR instead of being swallowed."""
if self._metrics_recorder is None:
return
try:
self._metrics_recorder.record(kwargs, response_obj, start_time, end_time)
record(self._metrics_recorder)
except ValueError as exc:
if not self._metric_filter_error_logged:
verbose_logger.error(
@ -298,6 +314,7 @@ class OpenTelemetryV2(CustomLogger):
if self._emit_mcp_list_tools(kwargs, start_time, end_time):
return
self._close_llm_call(kwargs, start_time, end_time)
self._record_failure_metrics(kwargs, start_time, end_time)
def _seed_identity_baggage(self, identity: RequestIdentity, model: str | None, context: Context) -> Context:
"""Seed authenticated request-identity Baggage onto ``context`` so the Baggage processor

30
litellm/integrations/otel/mappers/__init__.py
View file

@ -18,13 +18,19 @@ from litellm.integrations.otel.mappers.langfuse import LangfuseMapper
from litellm.integrations.otel.mappers.langtrace import LangtraceMapper
from litellm.integrations.otel.mappers.legacy import LegacyMapper
from litellm.integrations.otel.mappers.openinference import OpenInferenceMapper
from litellm.integrations.otel.mappers.utils import tool_attr_budget
from litellm.integrations.otel.mappers.weave import WeaveMapper
# Registry keyed by ``config.mapper_names`` entries.
_MAPPER_BY_NAME: dict[str, Callable[[], AttributeMapper]] = {
# Registries keyed by ``config.mapper_names`` entries, split by whether the
# vocabulary spells declared tool definitions out per index. Those share one
# span-wide attribute ceiling, so resolution has to know how many of them are
# active before it can build them.
_TOOL_DEFINITION_MAPPERS: dict[str, Callable[[int], AttributeMapper]] = {
"genai": GenAIMapper,
"legacy": LegacyMapper,
"openinference": OpenInferenceMapper,
}
_PLAIN_MAPPERS: dict[str, Callable[[], AttributeMapper]] = {
"langfuse": LangfuseMapper,
"weave": WeaveMapper,
"langtrace": LangtraceMapper,
@ -33,13 +39,19 @@ _MAPPER_BY_NAME: dict[str, Callable[[], AttributeMapper]] = {
def resolve_mappers(names: Iterable[str]) -> list[AttributeMapper]:
"""Resolve mapper names to instances. Unknown names raise ``ValueError``."""
out: list[AttributeMapper] = []
for name in names:
factory = _MAPPER_BY_NAME.get(name)
if factory is None:
raise ValueError(f"unknown mapper name {name!r}; known: {sorted(_MAPPER_BY_NAME)}")
out.append(factory())
return out
ordered = tuple(names)
for name in ordered:
if name not in _TOOL_DEFINITION_MAPPERS and name not in _PLAIN_MAPPERS:
known = sorted((*_TOOL_DEFINITION_MAPPERS, *_PLAIN_MAPPERS))
raise ValueError(f"unknown mapper name {name!r}; known: {known}")
# Distinct vocabularies each write the tool family under their own keys, so
# the ceiling is split by how many of them are configured. Repeating a name
# rewrites the same keys, so only distinct ones count.
budget = tool_attr_budget(len({*ordered} & _TOOL_DEFINITION_MAPPERS.keys()))
return [
_TOOL_DEFINITION_MAPPERS[name](budget) if name in _TOOL_DEFINITION_MAPPERS else _PLAIN_MAPPERS[name]()
for name in ordered
]
__all__ = [

28
litellm/integrations/otel/mappers/genai.py
View file

@ -11,10 +11,11 @@ from typing import Callable
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import (
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
collect,
drop_none,
output_messages,
serialize_messages,
tool_definition_attrs,
)
from litellm.integrations.otel.model.payloads import (
GuardrailSpanData,
@ -135,6 +136,9 @@ class GenAIMapper:
LiteLLM.SERVICE_CALL_TYPE: lambda d: d.call_type,
}
def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None:
self._tool_attr_budget = tool_attr_budget
def map(self, data: SpanData) -> AttributeMap:
match data:
case LLMCallSpanData():
@ -150,18 +154,18 @@ class GenAIMapper:
case _:
return {}
@classmethod
def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
attrs = collect(cls._LLM_CALL_ATTRS, data)
attrs.update(
drop_none(
{
f"gen_ai.tool.{idx}.{suffix}": extract(tool)
for idx, tool in enumerate(data.tools)
for suffix, extract in cls._TOOL_ATTRS.items()
}
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
attrs = collect(self._LLM_CALL_ATTRS, data)
if data.tools:
attrs[LiteLLM.TOOLS_DECLARED] = len(data.tools)
attrs.update(
tool_definition_attrs(
lambda idx, suffix: f"gen_ai.tool.{idx}.{suffix}",
data.tools,
self._TOOL_ATTRS,
self._tool_attr_budget,
)
)
)
return attrs
@classmethod

25
litellm/integrations/otel/mappers/legacy.py
View file

@ -12,7 +12,11 @@ Like ``GenAIMapper``, each span kind declares its schema as a flat
from typing import Callable, Final
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, SpanData
from litellm.integrations.otel.mappers.utils import collect, drop_none
from litellm.integrations.otel.mappers.utils import (
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
collect,
tool_definition_attrs,
)
from litellm.integrations.otel.model.payloads import (
LLMCallSpanData,
ServiceSpanData,
@ -63,6 +67,9 @@ class LegacyMapper:
_LEGACY_ERROR: lambda d: d.error.message if d.error is not None and d.error.message else None,
}
def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None:
self._tool_attr_budget = tool_attr_budget
def map(self, data: SpanData) -> AttributeMap:
match data:
case LLMCallSpanData():
@ -72,16 +79,14 @@ class LegacyMapper:
case _:
return {}
@classmethod
def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
attrs = collect(cls._LLM_CALL_ATTRS, data)
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
attrs = collect(self._LLM_CALL_ATTRS, data)
attrs.update(
drop_none(
{
f"llm.request.functions.{idx}.{suffix}": extract(tool)
for idx, tool in enumerate(data.tools)
for suffix, extract in cls._TOOL_ATTRS.items()
}
tool_definition_attrs(
lambda idx, suffix: f"llm.request.functions.{idx}.{suffix}",
data.tools,
self._TOOL_ATTRS,
self._tool_attr_budget,
)
)
return attrs

32
litellm/integrations/otel/mappers/openinference.py
View file

@ -13,9 +13,11 @@ from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue, Span
from litellm.integrations.otel.mappers.utils import (
collect,
drop_none,
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
json_if,
message_content,
output_messages,
tool_definition_attrs,
)
from litellm.integrations.otel.model.payloads import (
LLMCallSpanData,
@ -70,6 +72,9 @@ class OpenInferenceMapper:
),
}
def __init__(self, tool_attr_budget: int = MAX_TOOL_DEFINITION_ATTRS_PER_SPAN) -> None:
self._tool_attr_budget = tool_attr_budget
def map(self, data: SpanData) -> AttributeMap:
match data:
case LLMCallSpanData():
@ -77,14 +82,13 @@ class OpenInferenceMapper:
case _:
return {}
@classmethod
def _llm_call(cls, data: LLMCallSpanData) -> AttributeMap:
def _llm_call(self, data: LLMCallSpanData) -> AttributeMap:
return {
**collect(cls._LLM_CALL_ATTRS, data),
**collect(cls._BLOB_ATTRS, data),
**cls._messages("llm.input_messages", "input.value", data.messages_in),
**cls._messages("llm.output_messages", "output.value", output_messages(data)),
**cls._tools(data),
**collect(self._LLM_CALL_ATTRS, data),
**collect(self._BLOB_ATTRS, data),
**self._messages("llm.input_messages", "input.value", data.messages_in),
**self._messages("llm.output_messages", "output.value", output_messages(data)),
**self._tools(data),
}
@staticmethod
@ -108,12 +112,10 @@ class OpenInferenceMapper:
attrs[value_key] = json.dumps([{"role": role, "content": content} for role, content in parsed])
return attrs
@classmethod
def _tools(cls, data: LLMCallSpanData) -> AttributeMap:
return drop_none(
{
f"llm.tools.{idx}.{suffix}": extract(tool)
for idx, tool in enumerate(data.tools)
for suffix, extract in cls._TOOL_ATTRS.items()
}
def _tools(self, data: LLMCallSpanData) -> AttributeMap:
return tool_definition_attrs(
lambda idx, suffix: f"llm.tools.{idx}.{suffix}",
data.tools,
self._TOOL_ATTRS,
self._tool_attr_budget,
)

51
litellm/integrations/otel/mappers/utils.py
View file

@ -6,10 +6,34 @@ they live in one place.
"""
import json
from typing import Callable, Mapping, Sequence
from typing import Callable, Final, Mapping, Sequence
from litellm.integrations.otel.mappers.base import AttributeMap, AttrValue
from litellm.integrations.otel.model.payloads import LLMCallSpanData
from litellm.integrations.otel.model.payloads import LLMCallSpanData, ToolDefinition
DEFAULT_SPAN_ATTRIBUTE_LIMIT: Final = 128
"""The OTel SDK's default per-span attribute count limit."""
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN: Final = DEFAULT_SPAN_ATTRIBUTE_LIMIT // 4
"""Span-wide ceiling on attributes spent spelling out declared tool definitions.
Tool definitions are an unbounded attribute family: one entry per declared
tool, per field, per active vocabulary. Agentic clients declare hundreds, which
overruns the span attribute limit. That limit evicts oldest-first, so an
uncapped family silently destroys the core ``gen_ai.*`` attributes written
before it.
The ceiling is span-wide rather than per-mapper because several vocabularies
can be active at once and each spells the same tools out under its own keys, so
a per-mapper allowance multiplies by the number of vocabularies and reaches the
limit again. Reserving a quarter of the span for tool detail leaves the rest to
core telemetry no matter how many vocabularies are configured.
"""
def tool_attr_budget(vocabularies: int) -> int:
"""Split the span-wide tool-definition ceiling across active vocabularies."""
return MAX_TOOL_DEFINITION_ATTRS_PER_SPAN // max(vocabularies, 1)
def drop_none(values: Mapping[str, AttrValue | None]) -> AttributeMap:
@ -17,6 +41,29 @@ def drop_none(values: Mapping[str, AttrValue | None]) -> AttributeMap:
return {k: v for k, v in values.items() if v is not None}
def tool_definition_attrs(
key_for: Callable[[int, str], str],
tools: Sequence[ToolDefinition],
extractors: Mapping[str, Callable[[ToolDefinition], AttrValue | None]],
attr_budget: int,
) -> AttributeMap:
"""Per-index attributes for as many tools as ``attr_budget`` affords.
``key_for`` builds a vocabulary's key from the tool's index and the field
name, so each mapper keeps its own naming while sharing the budget. One tool
always keeps its detail, so the family stays legible even when many
vocabularies split the ceiling.
"""
max_tools = max(attr_budget // max(len(extractors), 1), 1)
return drop_none(
{
key_for(idx, suffix): extract(tool)
for idx, tool in enumerate(tools[:max_tools])
for suffix, extract in extractors.items()
}
)
def collect(table: Mapping[str, Callable], source: object) -> AttributeMap:
"""Apply an extractor table to ``source``, dropping ``None`` results."""
return drop_none({key: extract(source) for key, extract in table.items()})

73
litellm/integrations/otel/model/semconv.py
View file

@ -6,17 +6,30 @@ without a semconv equivalent lives under the ``litellm.*`` vendor namespace.
from enum import Enum
from typing import Final
from litellm._logging import verbose_logger
class GenAIOperation(str, Enum):
"""Values for ``gen_ai.operation.name``."""
"""Values for ``gen_ai.operation.name``.
The first block is the convention's own vocabulary. The ``LITELLM_`` members
are vendor values for operations the convention names nothing for; its note
on this attribute directs instrumentation to use a system-specific name in
exactly that case, the same allowance :func:`resolve_provider` relies on for
unmapped providers. They stay under the ``litellm.`` prefix so a value the
convention adds later can never collide with one of ours.
"""
CHAT = "chat"
TEXT_COMPLETION = "text_completion"
EMBEDDINGS = "embeddings"
GENERATE_CONTENT = "generate_content"
RETRIEVAL = "retrieval" # vector-store search / RAG query spans
CREATE_AGENT = "create_agent" # reserved for future agent spans
INVOKE_AGENT = "invoke_agent" # reserved for future agent spans
INVOKE_AGENT = "invoke_agent" # agent (A2A) message spans
EXECUTE_TOOL = "execute_tool" # MCP tool-call spans
LITELLM_VECTOR_STORE_MANAGEMENT = "litellm.vector_store_management"
LITELLM_VECTOR_STORE_FILE_MANAGEMENT = "litellm.vector_store_file_management"
class GenAIProvider(str, Enum):
@ -49,11 +62,17 @@ class MCPMethod(str, Enum):
class GenAI:
"""Canonical OTel GenAI span-attribute keys."""
"""Canonical OTel GenAI attribute keys.
``SYSTEM`` is the one exception: the convention deprecated it in favor of
``PROVIDER_NAME``, and it survives here only so already-shipped series keep
resolving for consumers that query it. Nothing new should use it.
"""
# request
OPERATION_NAME: Final = "gen_ai.operation.name"
PROVIDER_NAME: Final = "gen_ai.provider.name"
SYSTEM: Final = "gen_ai.system"
REQUEST_MODEL: Final = "gen_ai.request.model"
REQUEST_TEMPERATURE: Final = "gen_ai.request.temperature"
REQUEST_TOP_P: Final = "gen_ai.request.top_p"
@ -233,6 +252,7 @@ class LiteLLM:
# ``litellm_params.model``), distinct from the user-facing ``gen_ai.request.model``.
PROVIDER_MODEL: Final = "litellm.provider.model"
REQUEST_STREAMING: Final = "litellm.request.streaming"
TOOLS_DECLARED: Final = "litellm.request.tools.declared"
GUARDRAIL_NAME: Final = "litellm.guardrail.name"
GUARDRAIL_MODE: Final = "litellm.guardrail.mode"
GUARDRAIL_STATUS: Final = "litellm.guardrail.status"
@ -315,6 +335,35 @@ _OPERATION_BY_CALL_TYPE: dict[str, GenAIOperation] = {
"responses": GenAIOperation.CHAT,
"aresponses": GenAIOperation.CHAT,
"call_mcp_tool": GenAIOperation.EXECUTE_TOOL,
"vector_store_search": GenAIOperation.RETRIEVAL,
"avector_store_search": GenAIOperation.RETRIEVAL,
"query": GenAIOperation.RETRIEVAL,
"aquery": GenAIOperation.RETRIEVAL,
"send_message": GenAIOperation.INVOKE_AGENT,
"asend_message": GenAIOperation.INVOKE_AGENT,
"asend_message_streaming": GenAIOperation.INVOKE_AGENT,
"vector_store_create": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_create": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_list": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_list": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_update": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_update": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_delete": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"avector_store_delete": GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT,
"vector_store_file_create": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_create": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_list": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_list": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_retrieve": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_content": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_content": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_update": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_update": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"vector_store_file_delete": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
"avector_store_file_delete": GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT,
}
@ -331,7 +380,21 @@ def resolve_provider(custom_llm_provider: str | None) -> str:
def resolve_operation(call_type: str | None) -> GenAIOperation:
"""Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value."""
"""Map a litellm ``call_type`` to a ``gen_ai.operation.name`` value.
An unmapped call type still falls back to ``chat`` so every series keeps an
operation label, but it logs at debug rather than falling through silently:
a new call type mislabelled as ``chat`` mixes its latency and cost into
everyone's chat charts, which is invisible until someone reads the numbers.
"""
if not call_type:
return GenAIOperation.CHAT
return _OPERATION_BY_CALL_TYPE.get(call_type.lower(), GenAIOperation.CHAT)
mapped = _OPERATION_BY_CALL_TYPE.get(call_type.lower())
if mapped is not None:
return mapped
verbose_logger.debug(
"otel: call_type %r has no gen_ai.operation.name mapping; labelling it %r. Add it to _OPERATION_BY_CALL_TYPE.",
call_type,
GenAIOperation.CHAT.value,
)
return GenAIOperation.CHAT

189
litellm/integrations/otel/plumbing/metrics.py
View file

@ -1,6 +1,6 @@
"""GenAI client metrics: the six ``gen_ai.client.*`` histograms plus the
recorder that builds attributes, applies the shared cardinality filter, and
records a request's metrics in the success path.
records a request's metrics on both the success and the failure path.
The instrument names/units/descriptions and the recording + timing math mirror
the v1 :mod:`litellm.integrations.opentelemetry` integration so both engines emit
@ -10,11 +10,12 @@ identical metrics. The attribute cardinality filter is reused from v1 by import
from dataclasses import dataclass
from datetime import datetime
from typing import Any, FrozenSet, Mapping, Optional
from typing import Any, Final, FrozenSet, Mapping, Optional, TypeAlias
from opentelemetry.metrics import Histogram, Meter
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.opentelemetry import (
METRIC_METADATA_KEYS,
TOKEN_TYPE_ATTRIBUTE,
@ -22,11 +23,34 @@ from litellm.integrations.opentelemetry import (
_resolve_metric_attribute_filter,
)
from litellm.integrations.otel.model.metadata import time_to_first_chunk_seconds
from litellm.integrations.otel.model.semconv import Metric, resolve_operation
from litellm.integrations.otel.model.semconv import (
Error,
GenAI,
Metric,
resolve_operation,
resolve_provider,
)
from litellm.integrations.otel.model.utils import to_seconds
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
def _provider_attributes(custom_llm_provider: object) -> Mapping[str, str]:
"""The provider labels for one call's metrics.
``gen_ai.provider.name`` carries the semconv-mapped value; the deprecated
``gen_ai.system`` spelling is dual-emitted with the raw litellm provider
string it has always carried, so a dashboard already querying it keeps
matching. A call with no provider gets neither label: a placeholder value
would mint a permanent series that no operator can act on.
"""
if not isinstance(custom_llm_provider, str) or not custom_llm_provider:
return {}
return {
GenAI.PROVIDER_NAME: resolve_provider(custom_llm_provider),
GenAI.SYSTEM: custom_llm_provider,
}
@dataclass(frozen=True)
class GenAIMetrics:
operation_duration: Histogram
@ -72,8 +96,83 @@ def create_genai_metrics(meter: Meter) -> GenAIMetrics:
)
# A metric datapoint's attributes. Values are the strings the recorder builds, except
# the request model, which is whatever the caller passed and may be absent.
MetricAttributes: TypeAlias = Mapping[str, "str | None"]
ERROR_TYPE_FALLBACK: Final = "_OTHER"
# Every attribute a metric datapoint may carry, on either path. A label value that
# is unique per request is a new time series that will never be written to again, so
# this set is what keeps the series count bounded by the deployment's own
# key/team/user/deployment count rather than by its traffic. Each entry is a fixed
# enum or an operator-provisioned identifier.
#
# Deliberately excluded is everything the *client* supplies or that moves per
# request: ``metadata.requester_metadata`` and ``metadata.spend_logs_metadata`` (both
# free-form from the request body), ``metadata.user_api_key_end_user_id`` (the body's
# ``user`` field), and ``metadata.requester_ip_address``. Those stay on the span,
# where cardinality is free and where they already are.
# ``metadata.user_api_key_user_email`` is left out too: it is bounded, but it is PII
# duplicating the user id already here.
#
# This is a CEILING, applied before the operator's own include/exclude filter, so an
# operator can narrow it but never widen it back to an unbounded attribute.
METRIC_ATTRIBUTE_CEILING: Final[frozenset[str]] = frozenset(
(
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",
"metadata.user_api_key_hash",
"metadata.user_api_key_alias",
"metadata.user_api_key_team_id",
"metadata.user_api_key_team_alias",
"metadata.user_api_key_org_id",
"metadata.user_api_key_user_id",
"hidden_params",
)
)
# The only ``hidden_params`` field that becomes part of the ``hidden_params`` label.
# The object as a whole is per-request by construction -- ``response_cost``,
# ``litellm_overhead_time_ms``, ``cache_key``, ``usage_object`` and the provider's
# ``additional_headers`` rate-limit counters all move on every call -- so dumping it
# whole made one series per request out of every instrument.
#
# ``model_id`` is the router's own deployment id, so it is bounded by the deployment
# list and is what a per-deployment panel joins on. ``api_base`` is deliberately NOT
# here even though it names the same thing: it is a documented per-call parameter, so
# in SDK use it is chosen by the caller rather than provisioned by the operator, and a
# caller varying it would put the per-request cardinality straight back.
BOUNDED_HIDDEN_PARAM_KEYS: Final[tuple[str, ...]] = ("model_id",)
def resolve_error_type(kwargs: Mapping[str, Any]) -> str:
"""The ``error.type`` value for a failed request.
Bounded by construction: the mapped provider exception's class name (the same
``error_information.error_class`` the failure span stamps), else the provider
status code, else the raw exception's class name, else ``_OTHER`` — the value
the convention reserves for a failure the instrumentation cannot classify. The
exception *message* is unbounded and never becomes a label; it stays on the
span and its exception event, where high cardinality is free.
"""
std_log = kwargs.get("standard_logging_object")
info = getattr(std_log, "error_information", None) or (std_log or {}).get("error_information") or {}
error_class = info.get("error_class") or info.get("error_code")
if error_class:
return str(error_class)
exception = kwargs.get("exception")
if exception is not None:
return type(exception).__name__
return ERROR_TYPE_FALLBACK
class GenAIMetricRecorder:
"""Records the six GenAI histograms for one successful LLM call.
"""Records the six GenAI histograms for one successful LLM call, and the
duration histogram alone for one failed LLM call (see :meth:`record_failure`).
The cardinality filter is resolved lazily on the first record: the proxy
populates ``callback_settings.otel.attributes`` after the logger is built, so
@ -96,7 +195,7 @@ class GenAIMetricRecorder:
start_time: datetime,
end_time: datetime,
) -> None:
common_attrs = self._filter_attributes(self._common_attributes(kwargs))
common_attrs = self._filter_attributes(self._bounded_attributes(kwargs))
duration_s = (end_time - start_time).total_seconds()
self._metrics.operation_duration.record(duration_s, attributes=common_attrs)
@ -110,17 +209,48 @@ class GenAIMetricRecorder:
self._record_time_per_output_token(kwargs, response_obj, end_time, duration_s, common_attrs)
self._record_response_duration(kwargs, end_time, common_attrs)
def record_failure(
self,
kwargs: Mapping[str, Any],
start_time: datetime,
end_time: datetime,
) -> None:
"""Record the one metric a failed request can honestly report: the
operation's duration, tagged with ``error.type``.
The other five instruments all describe a completed generation and have
nothing to measure here. litellm hands the failure callback no
``response_obj`` at all, so there is no usage to split into input/output
tokens and no completion-token count to divide generation time by; it also
zeroes ``response_cost`` on failure. Recording them anyway would put a
fabricated zero into series that dashboards average.
The attribute set is :data:`METRIC_ATTRIBUTE_CEILING`, the same cap the
success path uses. A failure needs no provider spend, so a caller who can put
a unique value into a client-supplied attribute could mint one histogram
series per request for free; the cap is what makes that impossible on either
path.
``error.type`` is stamped after both filters, exactly like
``gen_ai.token.type``, so an operator's include/exclude list cannot strip
the discriminator and silently merge failures back into the success series.
"""
attributes = {
**self._filter_attributes(self._bounded_attributes(kwargs)),
Error.TYPE: resolve_error_type(kwargs),
}
self._metrics.operation_duration.record((end_time - start_time).total_seconds(), attributes=attributes)
# ------------------------------------------------------------------ #
# Attribute building + cardinality filter
# ------------------------------------------------------------------ #
def _common_attributes(self, kwargs: Mapping[str, Any]) -> dict:
params = kwargs.get("litellm_params") or {}
provider = params.get("custom_llm_provider", "Unknown")
common_attrs: dict = {
"gen_ai.operation.name": resolve_operation(kwargs.get("call_type")).value,
"gen_ai.system": provider,
"gen_ai.request.model": kwargs.get("model"),
GenAI.OPERATION_NAME: resolve_operation(kwargs.get("call_type")).value,
**_provider_attributes(params.get("custom_llm_provider")),
GenAI.REQUEST_MODEL: kwargs.get("model"),
"gen_ai.framework": "litellm",
}
@ -136,11 +266,25 @@ class GenAIMetricRecorder:
common_attrs[f"metadata.{key}"] = str(value)
hidden_params = getattr(std_log, "hidden_params", None) or (std_log or {}).get("hidden_params", {})
if hidden_params:
common_attrs["hidden_params"] = safe_dumps(hidden_params)
bounded_hidden_params = {
key: hidden_params[key]
for key in BOUNDED_HIDDEN_PARAM_KEYS
if isinstance(hidden_params, Mapping) and hidden_params.get(key) is not None
}
if bounded_hidden_params:
common_attrs["hidden_params"] = safe_dumps(bounded_hidden_params)
return common_attrs
def _bounded_attributes(self, kwargs: Mapping[str, Any]) -> MetricAttributes:
"""The datapoint attributes, capped at :data:`METRIC_ATTRIBUTE_CEILING`.
The cap runs BEFORE the operator's include/exclude filter so the filter can
only narrow it. An operator who names an excluded attribute in an include
list gets nothing for it rather than reintroducing an unbounded label.
"""
return {k: v for k, v in self._common_attributes(kwargs).items() if k in METRIC_ATTRIBUTE_CEILING}
def _ensure_filter(self) -> None:
if self._filter_resolved:
return
@ -157,8 +301,29 @@ class GenAIMetricRecorder:
# without reconstructing the recorder.
self._include, self._exclude = _resolve_metric_attribute_filter(attributes)
self._filter_resolved = True
self._warn_about_metric_ineligible_names()
def _filter_attributes(self, attrs: dict) -> dict:
def _warn_about_metric_ineligible_names(self) -> None:
"""Say so when the operator's filter names an attribute the ceiling removes.
The shared validator accepts every span attribute name, so a name that is
legal on a span but metric-ineligible would otherwise be a silent no-op: an
``include_list`` naming it emits nothing for it and an ``exclude_list`` naming
it looks like it worked. Logged once, when the filter resolves, rather than
per request.
"""
named = (self._include or frozenset()) | (self._exclude or frozenset())
ineligible = sorted(named - METRIC_ATTRIBUTE_CEILING - {TOKEN_TYPE_ATTRIBUTE})
if ineligible:
verbose_logger.warning(
"OTel metrics: %s cannot be a metric attribute and is being ignored; it varies "
"per request or is client-supplied, so it would make one time series per request. "
"It is still on the span. Metric attributes are limited to: %s",
", ".join(ineligible),
", ".join(sorted(METRIC_ATTRIBUTE_CEILING)),
)
def _filter_attributes(self, attrs: MetricAttributes) -> MetricAttributes:
self._ensure_filter()
if self._include is not None:
return {k: v for k, v in attrs.items() if k in self._include}

2
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4710,6 +4710,7 @@ class StandardLoggingPayloadSetup:
applied_guardrails=applied_guardrails,
mcp_tool_call_metadata=mcp_tool_call_metadata,
vector_store_request_metadata=vector_store_request_metadata,
routing_decision=None,
usage_object=usage_object,
requester_custom_headers=None,
cold_storage_object_key=None,
@ -5549,6 +5550,7 @@ def get_standard_logging_metadata(
applied_guardrails=None,
mcp_tool_call_metadata=None,
vector_store_request_metadata=None,
routing_decision=None,
usage_object=None,
requester_custom_headers=None,
user_api_key_request_route=None,

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

@ -1626,7 +1626,7 @@ class OpenAIFilesAPI(BaseLLM):
openai_client: AsyncOpenAI,
) -> OpenAIFileObject:
response = await openai_client.files.create(**create_file_data) # type: ignore[arg-type]
return OpenAIFileObject(**response.model_dump())
return OpenAIFileObject.model_validate(response.model_dump())
def create_file(
self,
@ -1662,7 +1662,7 @@ class OpenAIFilesAPI(BaseLLM):
create_file_data=create_file_data, openai_client=openai_client
)
response = cast(OpenAI, openai_client).files.create(**create_file_data) # type: ignore[arg-type]
return OpenAIFileObject(**response.model_dump())
return OpenAIFileObject.model_validate(response.model_dump())
async def afile_content(
self,
@ -1986,7 +1986,7 @@ class OpenAIBatchesAPI(BaseLLM):
openai_client: AsyncOpenAI,
) -> LiteLLMBatch:
response = await openai_client.batches.create(**create_batch_data) # type: ignore[arg-type]
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
def create_batch(
self,
@ -2023,7 +2023,7 @@ class OpenAIBatchesAPI(BaseLLM):
)
response = cast(OpenAI, openai_client).batches.create(**create_batch_data) # type: ignore[arg-type]
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
async def aretrieve_batch(
self,
@ -2032,7 +2032,7 @@ class OpenAIBatchesAPI(BaseLLM):
) -> LiteLLMBatch:
verbose_logger.debug("retrieving batch, args= %s", retrieve_batch_data)
response = await openai_client.batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type]
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
def retrieve_batch(
self,
@ -2068,7 +2068,7 @@ class OpenAIBatchesAPI(BaseLLM):
retrieve_batch_data=retrieve_batch_data, openai_client=openai_client
)
response = cast(OpenAI, openai_client).batches.retrieve(**retrieve_batch_data) # type: ignore[arg-type]
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
async def acancel_batch(
self,
@ -2077,7 +2077,7 @@ class OpenAIBatchesAPI(BaseLLM):
) -> LiteLLMBatch:
verbose_logger.debug("async cancelling batch, args= %s", cancel_batch_data)
response = await openai_client.batches.cancel(**cancel_batch_data)
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
def cancel_batch(
self,
@ -2117,7 +2117,7 @@ class OpenAIBatchesAPI(BaseLLM):
if not isinstance(openai_client, OpenAI):
raise ValueError("OpenAI client is not an instance of OpenAI. Make sure you passed a sync OpenAI client.")
response = openai_client.batches.cancel(**cancel_batch_data)
return LiteLLMBatch(**response.model_dump())
return LiteLLMBatch.model_validate(response.model_dump())
async def alist_batches(
self,
@ -2477,9 +2477,9 @@ class OpenAIAssistantsAPI(BaseLLM):
response_obj: Optional[OpenAIMessage] = None
if getattr(thread_message, "status", None) is None:
thread_message.status = "completed"
response_obj = OpenAIMessage(**thread_message.dict())
response_obj = OpenAIMessage.model_validate(thread_message.dict())
else:
response_obj = OpenAIMessage(**thread_message.dict())
response_obj = OpenAIMessage.model_validate(thread_message.dict())
return response_obj
# fmt: off
@ -2556,9 +2556,9 @@ class OpenAIAssistantsAPI(BaseLLM):
response_obj: Optional[OpenAIMessage] = None
if getattr(thread_message, "status", None) is None:
thread_message.status = "completed"
response_obj = OpenAIMessage(**thread_message.dict())
response_obj = OpenAIMessage.model_validate(thread_message.dict())
else:
response_obj = OpenAIMessage(**thread_message.dict())
response_obj = OpenAIMessage.model_validate(thread_message.dict())
return response_obj
async def async_get_messages(

8
litellm/llms/openai/responses/transformation.py
View file

@ -280,7 +280,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
raw_response_headers = dict(raw_response.headers)
processed_headers = process_response_headers(raw_response_headers)
try:
response = ResponsesAPIResponse(**raw_response_json)
response = ResponsesAPIResponse.model_validate(raw_response_json)
except Exception:
verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct")
response = ResponsesAPIResponse.model_construct(**raw_response_json)
@ -506,7 +506,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
raise OpenAIError(message=raw_response.text, status_code=raw_response.status_code)
raw_response_headers = dict(raw_response.headers)
processed_headers = process_response_headers(raw_response_headers)
response = ResponsesAPIResponse(**raw_response_json)
response = ResponsesAPIResponse.model_validate(raw_response_json)
response._hidden_params["additional_headers"] = processed_headers
response._hidden_params["headers"] = raw_response_headers
@ -588,7 +588,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
raw_response_headers = dict(raw_response.headers)
processed_headers = process_response_headers(raw_response_headers)
response = ResponsesAPIResponse(**raw_response_json)
response = ResponsesAPIResponse.model_validate(raw_response_json)
response._hidden_params["additional_headers"] = processed_headers
response._hidden_params["headers"] = raw_response_headers
@ -647,7 +647,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig):
processed_headers = process_response_headers(raw_response_headers)
try:
response = ResponsesAPIResponse(**raw_response_json)
response = ResponsesAPIResponse.model_validate(raw_response_json)
except Exception:
verbose_logger.debug(f"Error constructing ResponsesAPIResponse: {raw_response_json}, using model_construct")
response = ResponsesAPIResponse.model_construct(**raw_response_json)

80
litellm/model_prices_and_context_window_backup.json
View file

@ -3454,22 +3454,16 @@
},
"azure_ai/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.5e-07,
"cache_read_input_token_cost_priority": 1.5e-07,
"cache_read_input_token_cost_above_272k_tokens_priority": 3e-07,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_above_272k_tokens": 1.5e-06,
"input_cost_per_token_priority": 1.5e-06,
"input_cost_per_token_above_272k_tokens_priority": 3e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"output_cost_per_token_above_272k_tokens": 6.75e-06,
"output_cost_per_token_priority": 9e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.35e-05,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-mini",
"supported_endpoints": [
"/v1/chat/completions",
@ -3500,22 +3494,16 @@
},
"azure_ai/gpt-5.4-mini-2026-03-17": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.5e-07,
"cache_read_input_token_cost_priority": 1.5e-07,
"cache_read_input_token_cost_above_272k_tokens_priority": 3e-07,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_above_272k_tokens": 1.5e-06,
"input_cost_per_token_priority": 1.5e-06,
"input_cost_per_token_above_272k_tokens_priority": 3e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"output_cost_per_token_above_272k_tokens": 6.75e-06,
"output_cost_per_token_priority": 9e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.35e-05,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-mini",
"supported_endpoints": [
"/v1/chat/completions",
@ -3546,22 +3534,16 @@
},
"azure_ai/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_priority": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_priority": 4e-07,
"input_cost_per_token_above_272k_tokens_priority": 8e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"output_cost_per_token_above_272k_tokens": 1.875e-06,
"output_cost_per_token_priority": 2.5e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.75e-06,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-nano",
"supported_endpoints": [
"/v1/chat/completions",
@ -3592,22 +3574,16 @@
},
"azure_ai/gpt-5.4-nano-2026-03-17": {
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_priority": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_priority": 4e-07,
"input_cost_per_token_above_272k_tokens_priority": 8e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"output_cost_per_token_above_272k_tokens": 1.875e-06,
"output_cost_per_token_priority": 2.5e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.75e-06,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-nano",
"supported_endpoints": [
"/v1/chat/completions",
@ -7201,7 +7177,7 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7236,7 +7212,7 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7271,7 +7247,7 @@
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7306,7 +7282,7 @@
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -16703,8 +16679,8 @@
"input_cost_per_token": 6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://fireworks.ai/pricing",
@ -16717,8 +16693,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -16733,8 +16709,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17077,8 +17053,8 @@
"input_cost_per_token": 6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://fireworks.ai/pricing",
@ -17091,8 +17067,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17107,8 +17083,8 @@
"input_cost_per_token": 2e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17123,8 +17099,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17139,8 +17115,8 @@
"input_cost_per_token": 1.9e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -42501,8 +42477,8 @@
"input_cost_per_token": 2e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -42517,8 +42493,8 @@
"input_cost_per_token": 1.9e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",

231
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -1,6 +1,6 @@
import re
from datetime import datetime, timezone
from typing import Dict, List, Optional, Set, Tuple, cast
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Set, Tuple, cast
from fastapi import HTTPException
from starlette.datastructures import Headers
@ -30,6 +30,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent
)
from litellm.proxy._types import (
UI_TEAM_ID,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
ProxyException,
SpecialHeaders,
@ -43,13 +44,27 @@ from litellm.proxy.auth.user_api_key_auth import (
user_api_key_auth,
)
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
from litellm.proxy.common_utils.user_api_key_cache import get_management_object_ttl
from litellm.proxy.common_utils.user_api_key_cache import (
USER_NO_MCP_PERMISSION_SENTINEL,
get_management_object_ttl,
user_object_permission_id_cache_key,
)
from litellm.repositories.table_repositories import (
AgentsRepository,
MCPServerRepository,
)
from litellm.repositories.user_repository import UserRepository
from litellm.types.mcp_server.mcp_server_manager import MCPServer
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
def _as_list(values: Sequence[str] | None) -> list[str] | None: # mutable-ok: resolver returns a list
"""Widen a read-only allowlist back to the mutable list the resolver's own contract returns,
preserving the ``None`` that means "no restriction"."""
return None if values is None else list(values)
def _parse_mcp_server_names_from_path(path: str, mcp_servers_header: Optional[List[str]] = None) -> Optional[List[str]]:
"""Resolve the single MCP server name a cold-start passthrough bypass may
@ -1408,6 +1423,15 @@ class MCPRequestHandler:
f"Applied agent intersection filter. Final allowed servers: {allowed_mcp_servers}"
)
#########################################################
# Apply the internal user's own ceiling (the entitlement attached to the human)
#########################################################
capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(
allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source
)
allowed_mcp_servers = list(capped)
has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts
#########################################################
# Apply org-level ceiling if org_id is set
#########################################################
@ -1831,6 +1855,12 @@ class MCPRequestHandler:
# No team restrictions → use key restrictions
allowed_tools = cast(List[str], key_tools)
allowed_tools = _as_list(
await MCPRequestHandler._apply_user_tool_ceiling(
allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source
)
)
return await MCPRequestHandler._apply_agent_and_org_tool_ceilings(
allowed_tools, server_id, user_api_key_auth, keyless_source=keyless_source
)
@ -2376,6 +2406,203 @@ class MCPRequestHandler:
verbose_logger.warning(f"Failed to get allowed MCP servers for end_user: {str(e)}")
return []
@staticmethod
async def _get_user_object_permission(
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> LiteLLM_ObjectPermissionTable | None:
"""The internal user's OWN object_permission: the entitlement attached to the HUMAN rather
than to the credential they authenticated with.
A key's object_permission is the credential's scope and a team's is the group's; this one
answers "which MCP servers and tools is this person entitled to", independent of how many keys
they hold. Caches the ``user_id -> object_permission_id`` mapping (with a sentinel for "no
entitlement") exactly as the agent path does, then reuses the shared ``object_permission_id``
cache, so a warm request reads no rows.
``None`` means the human places NO ceiling: no user row, or a row naming no permission. The
two fault classes are deliberately NOT collapsed into that: a user row we cannot read leaves
us unable to say whether they are entitled at all, which is exactly the state before this
level existed, so it places no ceiling; a row that NAMES a permission we cannot read is a
KNOWN entitlement with unknown contents, so it raises and the caller denies.
"""
from litellm.proxy.auth.auth_checks import get_object_permission
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if not user_api_key_auth or not user_api_key_auth.user_id:
return None
if prisma_client is None:
verbose_logger.debug("prisma_client is None")
return None
user_id = user_api_key_auth.user_id
object_permission_id = await MCPRequestHandler._user_object_permission_id(user_id, prisma_client)
if object_permission_id is None:
return None
object_permission = await get_object_permission(
object_permission_id=object_permission_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if object_permission is None:
raise ValueError(
f"user {user_id!r} names object_permission_id {object_permission_id!r} which could not be loaded"
)
return object_permission
@staticmethod
async def _user_object_permission_id(user_id: str, prisma_client: "PrismaClient") -> str | None:
"""The permission row this human's user row links to, or None when they link none.
Caches the link (with a sentinel for "links none") so a human without an entitlement costs no
DB read per MCP request. Anything other than an id string is treated as a cache MISS rather
than carried into the permission lookup, and a read that fails answers None: not knowing
whether someone is entitled is the state that existed before this level, so it places no
ceiling. Only a link we DID resolve can make the caller deny.
"""
from litellm.proxy.proxy_server import user_api_key_cache
cache_key = user_object_permission_id_cache_key(user_id)
try:
cached: object = await user_api_key_cache.async_get_cache(key=cache_key)
if cached == USER_NO_MCP_PERMISSION_SENTINEL:
return None
if isinstance(cached, str) and cached:
return cached
user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id})
linked: object = getattr(user_row, "object_permission_id", None) if user_row is not None else None
object_permission_id = linked if isinstance(linked, str) and linked else None
await user_api_key_cache.async_set_cache(
key=cache_key,
value=object_permission_id or USER_NO_MCP_PERMISSION_SENTINEL,
ttl=get_management_object_ttl(user_api_key_cache),
)
return object_permission_id
except Exception as e: # noqa: BLE001 # unknown whether entitled at all: no ceiling, as before
verbose_logger.warning(f"MCP user entitlement: link for {user_id!r} unresolved, no ceiling: {str(e)}")
return None
@staticmethod
async def _get_allowed_mcp_servers_for_user(
user_api_key_auth: UserAPIKeyAuth | None = None,
) -> Sequence[str] | None:
"""The MCP servers the internal user is entitled to, as server ids.
``[]`` means this human places no restriction (allow-all from this level); ``None`` means the
ceiling is UNRESOLVED, which the caller denies on. Servers named only under
``mcp_tool_permissions`` count as entitled, exactly as they do for a key or a team, so
granting one tool never requires naming its server twice.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
try:
object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth)
if object_permissions is None:
return []
direct_mcp_servers = global_mcp_server_manager.expand_permission_list(object_permissions.mcp_servers or [])
access_group_servers = await MCPRequestHandler._get_mcp_servers_from_access_groups(
object_permissions.mcp_access_groups or []
)
tool_perm_servers = list(
global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).keys()
)
return list(set(direct_mcp_servers + access_group_servers + tool_perm_servers))
except Exception as e: # noqa: BLE001 # any resolution fault is an unresolved ceiling, never "no ceiling"
verbose_logger.warning(f"Failed to get allowed MCP servers for user: {str(e)}")
return None
@staticmethod
async def _apply_user_server_ceiling(
allowed_mcp_servers: Sequence[str],
user_api_key_auth: UserAPIKeyAuth | None = None,
*,
keyless_source: bool = False,
) -> tuple[tuple[str, ...], bool]:
"""Narrow a resolved server list by the internal user's own entitlement.
Returns the capped list and whether this human restricted it at all; the caller needs the
second value because an org list may only CAP a lower-level restriction, never replace one, so
a user ceiling has to be visible to the org step.
RAISES when the entitlement is known but unreadable, which the resolver's own handler turns
into deny-all. That is the point of the level: dropping a ceiling we know exists is exactly the
silent widening it is there to prevent.
"""
if keyless_source:
return tuple(allowed_mcp_servers), False
entitled = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth)
if entitled is None:
raise ValueError(
f"MCP user ceiling unresolvable for user_id="
f"{user_api_key_auth.user_id if user_api_key_auth else None!r}"
)
if not entitled:
return tuple(allowed_mcp_servers), False
capped = tuple(server for server in allowed_mcp_servers if server in set(entitled))
verbose_logger.debug(f"Applied user ceiling filter. Final allowed servers: {capped}")
return capped, True
@staticmethod
async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool:
"""Whether this human's own entitlement bounds their MCP access at all.
True when they are entitled to a specific set of servers, and also when that entitlement is
UNRESOLVED — a caller uses this to decide whether it may skip the resolver, and skipping it on
a transient fault would widen access.
"""
entitled_servers = await MCPRequestHandler._get_allowed_mcp_servers_for_user(user_api_key_auth)
return entitled_servers is None or len(entitled_servers) > 0
@staticmethod
async def _apply_user_tool_ceiling(
allowed_tools: Sequence[str] | None,
server_id: str,
user_api_key_auth: UserAPIKeyAuth | None = None,
*,
keyless_source: bool = False,
) -> Sequence[str] | None:
"""Narrow a key/team tool allowlist by the internal user's own tool entitlement.
The human's entitlement can only ever narrow: a user naming tools on ``server_id`` intersects
(and becomes the allowlist when no lower level restricts), while a user naming none places no
restriction. Returns ``[]`` (deny every tool on this server) when the entitlement cannot be
resolved, because the caller's own except-handler treats a raise as allow-all for key auth.
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
if keyless_source:
return allowed_tools
try:
object_permissions = await MCPRequestHandler._get_user_object_permission(user_api_key_auth)
except Exception as e: # noqa: BLE001 # an unresolved human entitlement must deny, not widen
verbose_logger.warning(f"MCP user tool ceiling unresolvable, denying tools on {server_id!r}: {str(e)}")
return []
if object_permissions is None or not object_permissions.mcp_tool_permissions:
return allowed_tools
user_tools = global_mcp_server_manager.expand_tool_permissions(object_permissions.mcp_tool_permissions).get(
server_id
)
if user_tools is None:
return allowed_tools
if allowed_tools is None:
return list(user_tools)
return list(set(allowed_tools) & set(user_tools))
# Sentinel stored in cache when an agent has no object_permission, so we
# don't re-query the DB on every MCP request for that agent.
_AGENT_NO_PERMISSION_SENTINEL = "__agent_no_mcp_permission__"

9
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
View file

@ -1778,10 +1778,12 @@ async def token_endpoint(
@router.post("/authorize/complete")
async def authorize_complete(request: Request, flow: str = Form(...)):
async def authorize_complete(request: Request, flow: str = Form(...), delivery: str | None = Form(None)):
"""Finish an aggregate connect flow: mint the gateway authorization code for the
signed-in user and redirect back to the DCR client. POST plus the per-flow HttpOnly
cookie set at /authorize; an anonymous or bad-flow request just 400s."""
signed-in user and hand it back to the DCR client, by 303 redirect (default) or, for
a loopback client on a different machine, as a copyable callback URL
(``delivery=manual``). POST plus the per-flow HttpOnly cookie set at /authorize; an
anonymous or bad-flow request just 400s."""
from litellm.proxy.proxy_server import user_api_key_cache # noqa: PLC0415 # circular import at module load
return await complete_connect_flow(
@ -1789,6 +1791,7 @@ async def authorize_complete(request: Request, flow: str = Form(...)):
flow_handle=flow,
session_user_id=_session_cookie_user_id(request),
cache=user_api_key_cache,
delivery=delivery,
)

70
litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py
View file

@ -39,6 +39,7 @@ from __future__ import annotations
import hashlib
import hmac
import html
import secrets
from base64 import urlsafe_b64encode
from collections.abc import Mapping
@ -47,7 +48,7 @@ from typing import Awaitable, Callable, Literal, TypeVar
from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse
from fastapi import HTTPException, Request
from fastapi.responses import JSONResponse, RedirectResponse, Response
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response
from pydantic import BaseModel, ConfigDict, Field, ValidationError
from typing_extensions import assert_never
@ -94,6 +95,13 @@ server-side session store, and the sealed value never appears in a URL)."""
CONNECT_FLOW_TTL_SECONDS = 600
GATEWAY_AUTH_CODE_TTL_SECONDS = 120
MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS = 300
"""Lifetime of a code the user delivers by hand (headless/remote client, LIT-4863 class):
copy-pasting a callback URL from a laptop browser to an SSH session is slower than a
browser redirect, so manual-delivery codes get 5 minutes instead of 2, still well under
the 10-minute ceiling RFC 6749 section 4.1.2 recommends. Single-use and PKCE binding are
unchanged, so the longer window only extends how long the legitimate holder has to paste
it, not what an observer could do with it."""
_CLAIM_TTL_BUFFER_SECONDS = 60
_USED_CODE_CACHE_PREFIX = "mcp_gateway_dcr_code_used:"
_USED_FLOW_CACHE_PREFIX = "mcp_gateway_dcr_flow_used:"
@ -390,6 +398,7 @@ async def complete_connect_flow(
flow_handle: str,
session_user_id: str | None,
cache: DualCache,
delivery: str | None = None,
) -> Response:
"""The deliberate finish step of the connect flow: mint the gateway authorization
code and send the browser back to the client.
@ -399,7 +408,24 @@ async def complete_connect_flow(
into the flow: a link crafted by another party dies here with ``access_denied``
instead of minting a code for the victim's identity. The flow is single-use (an atomic
claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in.
``delivery`` chooses how the code reaches the client. Default (absent or
``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"``
renders the callback URL on a page instead, for a client whose redirect URI is a
loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box,
container): the 303 would dereference the browser machine's loopback and the code
would never arrive, so the user carries it over by pasting the URL into the client or
fetching it from the client machine's terminal. Manual delivery is honored only for
loopback redirect URIs; a routable redirect URI works from any browser by
construction, so those flows always redirect. The user who sees the page is exactly
the user the 303 would have carried the code to, and the same user already sees the
code today in the dead redirect's address bar, so the page exposes the code to no new
party. Unknown ``delivery`` values are rejected rather than defaulted: a client that
asked for manual delivery and got a dead redirect instead would silently lose its
code.
"""
if delivery not in (None, "redirect", "manual"):
return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'")
sealed_flow = request.cookies.get(_flow_cookie_name(flow_handle))
if sealed_flow is None:
return _oauth_error(400, "invalid_request", "unknown or expired connect flow")
@ -417,6 +443,8 @@ async def complete_connect_flow(
f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
):
return _oauth_error(400, "invalid_request", "this connect flow was already completed; restart the connection")
manual_delivery = delivery == "manual" and is_loopback_redirect_host(urlparse(flow.redirect_uri))
code_ttl = MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS if manual_delivery else GATEWAY_AUTH_CODE_TTL_SECONDS
code = _seal(
GATEWAY_AUTH_CODE_PREFIX,
_GatewayAuthCode(
@ -426,16 +454,46 @@ async def complete_connect_flow(
code_challenge=flow.code_challenge,
jti=secrets.token_urlsafe(24),
iat=int(now.timestamp()),
exp=int(now.timestamp()) + GATEWAY_AUTH_CODE_TTL_SECONDS,
exp=int(now.timestamp()) + code_ttl,
),
)
params = {"code": code, **({"state": flow.state} if flow.state else {})}
response = RedirectResponse(_append_query_params(flow.redirect_uri, params), status_code=303)
callback_url = _append_query_params(flow.redirect_uri, params)
response: Response = (
_manual_delivery_response(callback_url) if manual_delivery else RedirectResponse(callback_url, status_code=303)
)
path, secure = _cookie_path_and_secure(request)
response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax")
return response
def _manual_delivery_response(callback_url: str) -> Response:
"""The manual code-delivery page: the callback URL the 303 would have followed,
rendered for the user to carry to the machine the client actually runs on (paste into
the client's prompt, or fetch with curl from that machine's terminal). Served
no-store because the body holds a live single-use code, and the URL is HTML-escaped
because it is client-influenced. The page renders the URL as data only, never as a
ready-to-paste shell command: no single quoting of an attacker-influenced string is
correct across POSIX shells, cmd.exe, and PowerShell (cmd.exe ignores single quotes
and percent-expands inside double quotes), so any command string this page suggested
would be wrong for some shell the user might paste it into."""
safe_url = html.escape(callback_url, quote=True)
minutes = MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS // 60
body = (
"<html><head><title>Finish connecting</title></head><body>"
"<h2>Almost done</h2>"
"<p>Your MCP client runs on a different machine, so this browser cannot deliver the"
" authorization code to it. On the machine where the client runs, paste this URL into"
" the client's prompt (Claude Code accepts the pasted callback URL), or pass it as the"
" quoted argument of a curl command from that machine's terminal:</p>"
f'<p><input type="text" value="{safe_url}" readonly size="100" onclick="this.select()"></p>'
f"<p>The code is single-use and expires in {minutes} minutes. You can close this window"
" once the client confirms it is connected.</p>"
"</body></html>"
)
return HTMLResponse(body, headers=TOKEN_NO_CACHE_HEADERS)
def _pkce_verifier_matches(code_verifier: str, code_challenge: str) -> bool:
"""RFC 7636 S256 verification, total over hostile input. The comparison is over bytes
so a non-ASCII ``code_challenge`` (which reaches here unvalidated from the client's
@ -601,9 +659,11 @@ async def _authorization_code_grant(
if failure is not None:
return _reload_failure_response(failure)
# Atomic single-use claim is the gate: on a concurrent double-redeem exactly one caller
# wins, and a claim that cannot be recorded fails closed.
# wins, and a claim that cannot be recorded fails closed. The marker's TTL derives from
# the code's own remaining lifetime so it outlives whichever lifetime the code was minted with.
if not await guard.claim(
f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}", GATEWAY_AUTH_CODE_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS
f"{_USED_CODE_CACHE_PREFIX}{parsed.jti}",
parsed.exp - int(now.timestamp()) + _CLAIM_TTL_BUFFER_SECONDS,
):
return _oauth_error(400, "invalid_grant", "the authorization code was already used")
return _session_token_pair(SessionPrincipal(user_id=parsed.user_id, client_id=client_id), keys, now)

101
litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py
View file

@ -13,11 +13,22 @@ payload (name + arguments) so we just build the tool_call.
from typing import TYPE_CHECKING, Any, Dict, Optional
from fastapi import HTTPException
from mcp.types import Tool as MCPTool
from litellm._logging import verbose_proxy_logger
from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.proxy._experimental.mcp_server.utils import (
json_string_leaves,
json_unrewritable_labels,
mcp_content_item_text,
mcp_tool_result_content_list,
mcp_tool_result_structured_content,
set_mcp_tool_result_structured_content,
with_json_string_leaves,
with_mcp_content_item_text,
)
from litellm.types.llms.openai import (
ChatCompletionToolParam,
ChatCompletionToolParamFunctionChunk,
@ -92,7 +103,93 @@ class MCPGuardrailTranslationHandler(BaseTranslation):
user_api_key_dict: Optional[Any] = None,
request_data: Optional[dict] = None,
) -> Any:
verbose_proxy_logger.debug(
"MCP Guardrail: Output processing not implemented for MCP tools",
"""Scan the text content of an MCP tool result and write masked text back.
The content list is rewritten in place (only the entries the guardrail
actually changed) rather than returned as a new result: the same object is
already referenced by the logging payload captured before this hook runs,
so a copy would leave the unmasked text in the spend log / span. A
guardrail that rejects the result raises, and the exception propagates to
the caller.
``structuredContent`` is scanned and masked too, in the same
``apply_guardrail`` call: it is serialized to the client alongside
``content``, so a value living only there would otherwise reach the
client unscanned.
"""
content = mcp_tool_result_content_list(response)
text_blocks = (
tuple(
(index, text) for index, item in enumerate(content) if (text := mcp_content_item_text(item)) is not None
)
if content is not None
else ()
)
structured = mcp_tool_result_structured_content(response)
structured_leaves = json_string_leaves(structured) if structured is not None else ()
structured_labels = json_unrewritable_labels(structured) if structured is not None else ()
if structured_leaves is None or structured_labels is None:
raise HTTPException(
status_code=400,
detail={
"error": (
"Content blocked: MCP tool result structuredContent is nested too deeply to be scanned "
"by the configured guardrail"
)
},
)
if not text_blocks and not structured_leaves and not structured_labels:
verbose_proxy_logger.debug("MCP Guardrail: tool result has no scannable text, nothing to do")
return response
originals = (
tuple(text for _, text in text_blocks) + tuple(text for _, text in structured_leaves) + structured_labels
)
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=GenericGuardrailAPIInputs(texts=list(originals)),
request_data=request_data if request_data is not None else {},
input_type="response",
logging_obj=litellm_logging_obj,
)
masked_texts = guardrailed_inputs.get("texts") if guardrailed_inputs else None
if masked_texts is None:
return response
if len(masked_texts) != len(originals):
verbose_proxy_logger.warning(
"MCP Guardrail: guardrail returned %d texts for %d tool result texts; leaving the result unmasked",
len(masked_texts),
len(originals),
)
return response
split = len(text_blocks)
if content is not None:
for (index, original), masked in zip(text_blocks, masked_texts[:split]):
if masked != original:
content[index] = with_mcp_content_item_text(content[index], masked)
label_start = split + len(structured_leaves)
if any(masked != original for original, masked in zip(structured_labels, masked_texts[label_start:])):
raise HTTPException(
status_code=400,
detail={
"error": (
"Content blocked: MCP tool result matched a masking rule on a non-rewritable field "
"(a structuredContent key or numeric value), which cannot be redacted without changing "
"the payload contract"
)
},
)
structured_replacements = {
path: masked
for (path, original), masked in zip(structured_leaves, masked_texts[split:label_start])
if masked != original
}
if structured_replacements:
set_mcp_tool_result_structured_content(
response, with_json_string_leaves(structured, structured_replacements)
)
return response

7
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -2411,6 +2411,11 @@ class MCPServerManager:
and not is_admitted_subject
and _user_has_admin_view(user_api_key_auth)
and not has_explicit_object_permission
# An entitlement attached to the HUMAN binds them whatever their role: it is the
# person's scope, not the credential's, so an admin role is not a waiver of it. An
# UNRESOLVED entitlement also skips the shortcut, so the resolver denies rather than
# handing over the whole registry on a transient fault.
and not await MCPRequestHandler._user_places_mcp_ceiling(user_api_key_auth)
):
verbose_logger.debug("Admin user without explicit object_permission - returning all servers")
return list(self.get_registry().keys())
@ -5322,7 +5327,7 @@ class MCPServerManager:
]
}
)
db_mcp_servers = [LiteLLM_MCPServerTable(**r.model_dump()) for r in raw_rows]
db_mcp_servers = [LiteLLM_MCPServerTable.model_validate(r.model_dump()) for r in raw_rows]
verbose_logger.info(f"Found {len(db_mcp_servers)} MCP servers in database")
previous_registry = self.registry

44
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
View file

@ -12,6 +12,11 @@ import httpx
from fastapi import APIRouter, Depends, HTTPException, Query, Request, status
from litellm._logging import verbose_logger
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
ModifyResponseException,
)
from litellm.proxy._experimental.mcp_server.exceptions import (
MCPServerListError,
MCPUpstreamAuthError,
@ -33,6 +38,8 @@ from litellm.proxy.auth.ip_address_utils import IPAddressUtils
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
if TYPE_CHECKING:
from mcp.types import CallToolResult
from litellm.proxy._experimental.mcp_server.db import OAuthCredentialPayload
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.types.mcp import MCPAuth
@ -51,6 +58,13 @@ router = APIRouter(
tags=["mcp"],
)
_MCP_GUARDRAIL_REJECTIONS = (
BlockedPiiEntityError,
GuardrailRaisedException,
ModifyResponseException,
HTTPException,
)
def _connection_error_message(exc: BaseException) -> str:
if isinstance(exc, httpx.LocalProtocolError):
@ -99,9 +113,17 @@ if MCP_AVAILABLE:
end_time: datetime,
user_api_key_auth: UserAPIKeyAuth | None = None,
request_data: Mapping[str, object] | None = None,
) -> None:
) -> "CallToolResult":
"""Fire post-call logging, returning the tool result to send to the client.
``post_mcp_call`` guardrails already ran on ``execute_mcp_tool``'s return
path, so the result arriving here is the guardrailed one. A guardrail
rejection raised by a native ``async_post_mcp_tool_call_hook`` is still
re-raised rather than swallowed as a logging failure, which would return
the unguarded result.
"""
if logging_obj is None:
return
return result
logging_results = await asyncio.gather(
_fire_mcp_tool_call_logging(
logging_obj,
@ -113,11 +135,13 @@ if MCP_AVAILABLE:
),
return_exceptions=True,
)
logging_error = logging_results[0]
if isinstance(logging_error, asyncio.CancelledError):
raise logging_error
if isinstance(logging_error, BaseException):
verbose_logger.warning("MCP tool call logging failed (continuing): %s", logging_error)
outcome = logging_results[0]
if isinstance(outcome, (asyncio.CancelledError, *_MCP_GUARDRAIL_REJECTIONS)):
raise outcome
if isinstance(outcome, BaseException):
verbose_logger.warning("MCP tool call logging failed (continuing): %s", outcome)
return result
return outcome
def _relay_upstream_auth_http_exception(e: MCPUpstreamAuthError, request: Request) -> HTTPException:
"""Convert a client-forwarded pass-through upstream 401 into an HTTPException that preserves the
@ -196,7 +220,7 @@ if MCP_AVAILABLE:
raw_headers=virtual_raw_headers,
litellm_logging_obj=virtual_logging_obj,
)
await _safe_fire_mcp_tool_call_logging(
return await _safe_fire_mcp_tool_call_logging(
virtual_logging_obj,
result,
_tool_start_time,
@ -204,7 +228,6 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_dict,
request_data=data,
)
return result
def _get_server_auth_header(
server,
@ -998,7 +1021,7 @@ if MCP_AVAILABLE:
litellm_logging_obj=data.get("litellm_logging_obj"),
requested_server_id=canonical_server_id,
)
await _safe_fire_mcp_tool_call_logging(
return await _safe_fire_mcp_tool_call_logging(
logging_obj,
result,
_tool_start_time,
@ -1006,7 +1029,6 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_dict,
request_data=data,
)
return result
except MCPMissingUserEnvVarsError as e:
verbose_logger.info(
"MCP tool call missing per-user env vars: server_id=%s missing=%s",

53
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -2910,7 +2910,38 @@ if MCP_AVAILABLE:
local_content = await _handle_local_mcp_tool(original_tool_name, arguments)
response = CallToolResult(content=cast(Any, local_content), isError=False)
return response
return await _run_post_mcp_call_guardrails(
result=response,
litellm_logging_obj=litellm_logging_obj,
user_api_key_auth=user_api_key_auth,
request_data=kwargs,
)
async def _run_post_mcp_call_guardrails(
result: CallToolResult,
litellm_logging_obj: LiteLLMLoggingObj | None,
user_api_key_auth: UserAPIKeyAuth | None,
request_data: Mapping[str, object],
) -> CallToolResult:
"""Run ``post_mcp_call`` guardrails over an executed tool result.
Lives on ``execute_mcp_tool``'s return path rather than inside
``_fire_mcp_tool_call_logging`` so enforcement never depends on logging
being configured, and so every dispatch route gets it: the MCP protocol
handler, the REST endpoint, and tool search all funnel through here.
A guardrail that rejects the result raises, matching ``pre_mcp_call``.
"""
from litellm.proxy.proxy_server import proxy_logging_obj
if proxy_logging_obj is None:
return result
return await proxy_logging_obj.post_mcp_call_hook(
response=result,
request_data=(
litellm_logging_obj.model_call_details if litellm_logging_obj is not None else dict(request_data)
),
user_api_key_dict=user_api_key_auth,
)
_MCP_CREDENTIAL_REQUEST_FIELDS = frozenset(
{
@ -2929,8 +2960,14 @@ if MCP_AVAILABLE:
end_time: datetime,
user_api_key_auth: UserAPIKeyAuth | None = None,
request_data: Mapping[str, object] | None = None,
) -> None:
"""Fire post-call logging for an executed MCP tool call.
) -> CallToolResult:
"""Fire post-call logging for an executed MCP tool call, returning the result to send.
The returned result is what the caller must forward to the client: a
``post_mcp_call`` guardrail may rewrite the tool output (e.g. mask
sensitive values) or reject it, in which case its exception propagates.
Guardrails run before the success/failure logging so the masked text, not
the raw one, is what gets logged.
A result with ``isError=True`` is logged as a failure (``status="failure"``
payload, so OTel marks the span ERROR) while the HTTP wire behavior stays
@ -2946,6 +2983,8 @@ if MCP_AVAILABLE:
stripped before the dict is handed to ``post_call_failure_hook``
callbacks.
"""
from litellm.proxy.proxy_server import proxy_logging_obj
logging_obj.post_call(original_response=result)
await logging_obj.async_post_mcp_tool_call_hook(
kwargs=logging_obj.model_call_details,
@ -2957,7 +2996,7 @@ if MCP_AVAILABLE:
error_message = extract_mcp_tool_result_error_message(result)
if error_message is None:
await logging_obj.async_success_handler(result=result, start_time=start_time, end_time=end_time)
return
return result
logging_obj.has_run_logging(event_type="sync_success")
logging_obj.has_run_logging(event_type="async_success")
@ -2966,8 +3005,7 @@ if MCP_AVAILABLE:
await logging_obj.async_failure_handler(tool_error, "", start_time, end_time)
if user_api_key_auth is None:
return
from litellm.proxy.proxy_server import proxy_logging_obj
return result
if proxy_logging_obj:
sanitized_request_data = {
@ -2979,6 +3017,7 @@ if MCP_AVAILABLE:
user_api_key_dict=user_api_key_auth,
route="/mcp/call_tool",
)
return result
@client
async def call_mcp_tool(
@ -3062,7 +3101,7 @@ if MCP_AVAILABLE:
raise
if litellm_logging_obj:
await _fire_mcp_tool_call_logging(
response = await _fire_mcp_tool_call_logging(
logging_obj=litellm_logging_obj,
result=response,
start_time=start_time,

160
litellm/proxy/_experimental/mcp_server/utils.py
View file

@ -4,6 +4,7 @@ MCP Server Utilities
import json
import re
from collections.abc import MutableMapping, MutableSequence
from typing import (
Any,
Dict,
@ -434,6 +435,56 @@ def extract_mcp_tool_result_error_message(result: object) -> Optional[str]:
return "MCP tool call returned isError=true"
def mcp_tool_result_content_list(result: object) -> MutableSequence[object] | None: # mutable-ok: see below
"""The mutable content list of an MCP tool result, or ``None`` when it has none.
Deliberately mutable: a guardrail masking the result rewrites entries in place,
because the logging payload captured before the guardrail runs references this
same list, so handing back a copy would leave the unmasked text in the spend log
and the OTel span.
Accepts both ``mcp.types.CallToolResult`` objects and their dict
equivalents, duck-typed so the ``mcp`` package is not required.
"""
content: object = result.get("content") if isinstance(result, Mapping) else getattr(result, "content", None)
if isinstance(content, MutableSequence):
return content
return None
def mcp_content_item_text(item: object) -> str | None:
"""The ``text`` of a rewritable MCP content item, or ``None``.
Only mappings and Pydantic-style models report a text, because those are the
only shapes ``with_mcp_content_item_text`` can rewrite; a caller therefore
never reads text it would be unable to write back (e.g. masked by a
guardrail). Non-text content (images, embedded resources) has no ``text``
and is reported as ``None``.
"""
text: object
if isinstance(item, Mapping):
text = item.get("text")
elif callable(getattr(item, "model_copy", None)):
text = getattr(item, "text", None)
else:
return None
return text if isinstance(text, str) else None
def with_mcp_content_item_text(item: object, text: str) -> object:
"""A copy of an MCP content item carrying ``text`` instead of its own.
Only meaningful for items ``mcp_content_item_text`` returned a text for; any
other item is returned unchanged.
"""
if isinstance(item, Mapping):
return {**item, "text": text}
model_copy = getattr(item, "model_copy", None)
if callable(model_copy):
return model_copy(update={"text": text})
return item
TOOL_DISPLAY_NAME_PATTERN = re.compile(r"^[a-zA-Z0-9_-]+$")
@ -618,3 +669,112 @@ def merge_mcp_headers(
merged.update({str(k): str(v) for k, v in static_headers.items()})
return merged or None
# Local rather than litellm.constants: this module deliberately imports no litellm
# package, so pulling one in for a single integer would drag in litellm/__init__.
MAX_STRUCTURED_CONTENT_SCAN_DEPTH = 100
JSONLeafPath = tuple[str | int, ...]
def _flatten_leaf_groups(
groups: Iterable[tuple[tuple[JSONLeafPath, str], ...] | None],
) -> tuple[tuple[JSONLeafPath, str], ...] | None:
"""Concatenate child leaf groups, propagating the too-deep sentinel."""
materialized = tuple(groups)
if any(group is None for group in materialized):
return None
return tuple(leaf for group in materialized if group is not None for leaf in group)
def json_string_leaves(value: object, path: JSONLeafPath = ()) -> tuple[tuple[JSONLeafPath, str], ...] | None:
"""Depth-first, deterministically ordered string leaves of a JSON value.
Returns ``None`` when the value is nested past ``MAX_STRUCTURED_CONTENT_SCAN_DEPTH``,
so the caller blocks rather than letting deeper values through unscanned; an
empty tuple means there was simply nothing to scan. A sentinel rather than an
exception because this module is reloaded by tests (see the note above the
environment-backed constants), which would give a custom exception class a new
identity and let it escape a caller's ``except``.
"""
if len(path) > MAX_STRUCTURED_CONTENT_SCAN_DEPTH:
return None
if isinstance(value, str):
return ((path, value),)
if isinstance(value, dict):
return _flatten_leaf_groups(json_string_leaves(item, (*path, key)) for key, item in value.items())
if isinstance(value, list):
return _flatten_leaf_groups(json_string_leaves(item, (*path, index)) for index, item in enumerate(value))
return ()
def with_json_string_leaves(
value: object,
replacements: Mapping[JSONLeafPath, str],
path: JSONLeafPath = (),
) -> object:
"""Rebuild a JSON value with the guardrail's rewritten string leaves."""
if isinstance(value, str):
return replacements.get(path, value)
if isinstance(value, dict):
return {key: with_json_string_leaves(item, replacements, (*path, key)) for key, item in value.items()}
if isinstance(value, list):
return [with_json_string_leaves(item, replacements, (*path, index)) for index, item in enumerate(value)]
return value
def json_unrewritable_labels(value: object, path_depth: int = 0) -> tuple[str, ...] | None:
"""Strings in a JSON value that carry meaning but cannot be rewritten.
Dictionary keys and non-string scalars: masking either would change the
payload's contract rather than redact a value, so a caller scans these and
blocks on a match instead of rewriting, matching what the content filter
already does for MCP tool call arguments. ``None`` means the value is nested
past the scan depth, same contract as ``json_string_leaves``.
"""
if path_depth > MAX_STRUCTURED_CONTENT_SCAN_DEPTH:
return None
if isinstance(value, bool) or value is None or isinstance(value, str):
return ()
if isinstance(value, (int, float)):
return (str(value),)
if isinstance(value, dict):
own = tuple(key for key in value if isinstance(key, str))
nested = tuple(json_unrewritable_labels(item, path_depth + 1) for item in value.values())
if any(group is None for group in nested):
return None
return own + tuple(label for group in nested if group is not None for label in group)
if isinstance(value, list):
nested = tuple(json_unrewritable_labels(item, path_depth + 1) for item in value)
if any(group is None for group in nested):
return None
return tuple(label for group in nested if group is not None for label in group)
return ()
def mcp_tool_result_structured_content(result: object) -> object:
"""The ``structuredContent`` of an MCP tool result, or ``None`` when it has none."""
if isinstance(result, Mapping):
return result.get("structuredContent")
return getattr(result, "structuredContent", None)
def set_mcp_tool_result_structured_content(result: object, value: object) -> bool:
"""Replace ``structuredContent`` in place; ``False`` when the shape does not carry it.
In place for the same reason the content list is: the logging payload captured
before the guardrail ran references this object, so a copy would leave the
unmasked value in the spend log and the OTel span.
"""
if isinstance(result, MutableMapping):
result["structuredContent"] = value
return True
if not hasattr(result, "structuredContent"):
return False
try:
setattr(result, "structuredContent", value) # attribute name is fixed by the MCP result shape
return True
except (AttributeError, TypeError, ValueError):
return False

13
litellm/proxy/_types.py
View file

@ -57,6 +57,7 @@ from litellm.types.utils import (
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingRoutingDecision,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
@ -2464,16 +2465,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"is active as a reminder that hard enforcement is relaxed."
),
)
skip_user_budget_on_team_key: bool | None = Field(
None,
description=(
"If True, restores the legacy behavior where a user's personal "
"max_budget is NOT enforced when their key belongs to a team; only "
"the team (and team-member) budgets apply. Defaults to False, meaning "
"the user's personal max_budget is always enforced regardless of "
"whether the key belongs to a team (see GitHub issue #12905)."
),
)
user_url_validation: Optional[bool] = Field(
None,
description=(
@ -2786,6 +2777,7 @@ class UserInfoV2Response(LiteLLMPydanticObjectBase):
updated_at: Optional[datetime] = None
sso_user_id: Optional[str] = None
teams: List[str] = [] # Just team IDs, not full team objects
object_permission: LiteLLM_ObjectPermissionTable | None = None
from litellm.models.config import LiteLLM_Config as LiteLLM_Config # noqa: E402
@ -3316,6 +3308,7 @@ class SpendLogsMetadata(TypedDict):
applied_guardrails: Optional[List[str]]
mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall]
vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]]
routing_decision: StandardLoggingRoutingDecision | None
guardrail_information: Optional[List[StandardLoggingGuardrailInformation]]
eval_information: Optional[Any]
status: StandardLoggingPayloadStatus

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

@ -74,6 +74,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
get_management_object_ttl,
object_permission_cache_key,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.guardrails.tool_name_extraction import (
@ -485,6 +486,14 @@ MODEL_DISCOVERY_ROUTES = frozenset(
}
)
BUDGET_ENFORCED_SIDE_EFFECT_ROUTES = frozenset(
{
"/health",
"/health/services",
"/health/test_connection",
}
)
async def common_checks(
request_body: dict,
@ -531,8 +540,10 @@ async def common_checks(
request=request,
)
if route in MODEL_DISCOVERY_ROUTES:
skip_budget_checks = True
skip_all_budget_checks = skip_budget_checks or (
route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES
and (route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_api_route(route=route))
)
# 1. If team is blocked
if team_object is not None and team_object.blocked is True:
@ -606,7 +617,7 @@ async def common_checks(
project_object=project_object,
_model=_model,
llm_router=llm_router,
skip_budget_checks=skip_budget_checks,
skip_budget_checks=skip_all_budget_checks,
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
@ -615,7 +626,7 @@ async def common_checks(
_reject_clientside_metadata_tags_check(general_settings, request_body, route)
# If this is a free model, skip all budget checks
if not skip_budget_checks:
if not skip_all_budget_checks:
# Key metadata.tags are injected into request_body here so the tag budget
# check can read them; this mutation must run before the gathered checks.
if valid_token is not None:
@ -632,31 +643,28 @@ async def common_checks(
)
async def _user_max_budget_check() -> None:
if user_object is None or user_object.max_budget is None:
return
skip_for_team = (
general_settings.get("skip_user_budget_on_team_key") is True
and team_object is not None
and team_object.team_id is not None
)
if skip_for_team:
return
from litellm.proxy.proxy_server import get_current_spend
# 4.1 personal budget, if personal key
if (
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.max_budget is not None
):
from litellm.proxy.proxy_server import get_current_spend
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
user_budget = user_object.max_budget
user_spend = await get_current_spend(
counter_key=f"spend:user:{user_object.user_id}",
fallback_spend=user_object.spend or 0.0,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
)
if math.isfinite(user_budget) and user_spend >= user_budget:
raise litellm.BudgetExceededError(
current_cost=user_spend,
max_budget=user_budget,
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
entity_type=Litellm_EntityType.USER.value,
entity_id=user_object.user_id,
)
# Each scope reads a distinct counter key with no cross-scope ordering
# dependency, so the per-scope Redis-first reads run concurrently instead
@ -715,7 +723,7 @@ async def common_checks(
raise budget_error
_enforce_user_param_check(general_settings, request, request_body, route)
_global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route)
_global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route)
_guardrail_modification_check(request_body, team_object)
# 10 [OPTIONAL] Organization RBAC checks
@ -2434,7 +2442,7 @@ class ExperimentalUIJWTToken:
if decrypted_token is None:
return None
try:
return UserAPIKeyAuth(**json.loads(decrypted_token))
return UserAPIKeyAuth.model_validate(json.loads(decrypted_token))
except Exception as e:
raise Exception(f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}")
@ -2553,7 +2561,7 @@ async def get_key_object(
code=status.HTTP_401_UNAUTHORIZED,
)
_response = UserAPIKeyAuth(**_valid_token.model_dump(exclude_none=True))
_response = UserAPIKeyAuth.model_validate(_valid_token.model_dump(exclude_none=True))
# Load object_permission if object_permission_id exists but object_permission is not loaded
if _response.object_permission_id and not _response.object_permission:
@ -2609,7 +2617,7 @@ async def get_object_permission(
raise Exception("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
# check if in cache
key = "object_permission_id:{}".format(object_permission_id)
key = object_permission_cache_key(object_permission_id)
deserialized_perm = await user_api_key_cache.async_get_cache(
key=key,
model_type=LiteLLM_ObjectPermissionTable,

19
litellm/proxy/auth/oauth2_proxy_hook.py
View file

@ -1,4 +1,5 @@
from typing import Any, Dict, FrozenSet
from collections.abc import Mapping
from typing import Dict, FrozenSet, List, Union
from fastapi import Request
@ -83,21 +84,17 @@ async def handle_oauth2_proxy_request(request: Request) -> UserAPIKeyAuth:
"(signature-validated) instead of header-trust."
)
auth_data: Dict[str, Any] = {}
for key, header in oauth2_config_mappings.items():
value = request.headers.get(header)
if not value:
continue
if key == "models":
auth_data[key] = [model.strip() for model in value.split(",")]
else:
auth_data[key] = value
auth_data: Mapping[str, Union[str, List[str]]] = {
key: [model.strip() for model in value.split(",")] if key == "models" else value
for key, header in oauth2_config_mappings.items()
if (value := request.headers.get(header))
}
verbose_proxy_logger.debug(
"Auth data before creating UserAPIKeyAuth object: keys=%s",
list(auth_data.keys()),
)
user_api_key_auth = UserAPIKeyAuth(**auth_data)
user_api_key_auth = UserAPIKeyAuth.model_validate(auth_data)
verbose_proxy_logger.debug(
"UserAPIKeyAuth object created with keys: %s",
list(user_api_key_auth.__fields_set__),

2
litellm/proxy/auth/resolvers/store.py
View file

@ -118,7 +118,7 @@ class IdentityStore:
if from_db is None:
raise KeyNotFoundError(hashed_token)
key = UserAPIKeyAuth(**from_db.model_dump(exclude_none=True))
key = UserAPIKeyAuth.model_validate(from_db.model_dump(exclude_none=True))
if key.object_permission_id and not key.object_permission:
try:

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

@ -2498,7 +2498,6 @@ async def _reserve_budget_after_common_checks(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=general_settings.get("skip_user_budget_on_team_key") is True,
fail_closed_budget_enforcement=general_settings.get("fail_closed_budget_enforcement") is True,
)

22
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -150,6 +150,28 @@ class UserApiKeyCache(DualCache):
return await super().async_set_cache_pipeline(cache_list=normalized, local_only=local_only, **kwargs)
#: Value cached under ``user_object_permission_id_cache_key`` when the user links no permission row,
#: so a human without an entitlement costs no DB read per request. Lives beside the key builder
#: because it is part of the same cache protocol: a reader that knows the key must know this value.
USER_NO_MCP_PERMISSION_SENTINEL = "__user_no_mcp_permission__"
def user_object_permission_id_cache_key(user_id: str) -> str:
"""Cache key for the ``user_id -> object_permission_id`` link.
Lives here rather than next to either user because two modules own the two halves: the MCP auth
resolver writes it on read, and ``/user/update`` deletes it after changing the link. A key format
duplicated across those two drifts silently, and the failure is an entitlement change that never
takes effect.
"""
return f"user_object_permission_id:{user_id}"
def object_permission_cache_key(object_permission_id: str) -> str:
"""Cache key ``get_object_permission`` stores a permission row under."""
return f"object_permission_id:{object_permission_id}"
def get_management_object_ttl(cache: DualCache) -> float:
"""
In-memory TTL for management-object cache writes (keys, teams, users, budgets, ...).

16
litellm/proxy/db/prisma_client.py
View file

@ -834,6 +834,21 @@ class PrismaManager:
dname = os.path.dirname(os.path.dirname(abspath))
return dname
@staticmethod
def _apply_replica_identity_full_if_requested() -> None:
"""
`prisma db push` bypasses litellm-proxy-extras, so the opt-in
REPLICA IDENTITY FULL step has to be driven from here too.
litellm-proxy-extras is an optional install, so this is a no-op when it
is absent.
"""
try:
from litellm_proxy_extras.utils import ProxyExtrasDBManager
except ImportError:
return
ProxyExtrasDBManager.apply_replica_identity_full_if_requested()
@staticmethod
def setup_database(use_migrate: bool = False, use_v2_resolver: bool = False) -> bool:
"""
@ -880,6 +895,7 @@ class PrismaManager:
timeout=60,
check=True,
)
PrismaManager._apply_replica_identity_full_if_requested()
return True
except subprocess.TimeoutExpired:
verbose_proxy_logger.warning(f"Attempt {attempt + 1} timed out")

1
litellm/proxy/guardrails/guardrail_hooks/presidio.py
View file

@ -75,6 +75,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
GuardrailEventHooks.post_call,
GuardrailEventHooks.logging_only,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.post_mcp_call,
]
# Class variables or attributes

7
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -804,6 +804,8 @@ class UnifiedLLMGuardrails(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
guardrail_to_apply: Union[CustomGuardrail, None] = None,
buffer_until_moderated_default: bool = False,
) -> AsyncGenerator[Any, None]:
"""
Passes the entire stream to the guardrail
@ -824,7 +826,8 @@ class UnifiedLLMGuardrails(CustomLogger):
# litellm.integrations.custom_guardrail.
from litellm.integrations.custom_guardrail import ModifyResponseException
guardrail_to_apply: CustomGuardrail = request_data.pop("guardrail_to_apply", None)
if guardrail_to_apply is None:
guardrail_to_apply = request_data.pop("guardrail_to_apply", None)
# Get streaming configuration. Resolution order (later wins): default
# < guardrail attribute < guardrail_config dict < this callback's
@ -852,7 +855,7 @@ class UnifiedLLMGuardrails(CustomLogger):
# release the original chunks are replayed as-is, so a
# content-rewriting guardrail (e.g. PII masking) would leak
# unredacted content. Guarded below via mask_response_content.
buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", False)
buffer_until_moderated = _streaming_flag("streaming_buffer_until_moderated", buffer_until_moderated_default)
if (
buffer_until_moderated

38
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -28,6 +28,7 @@ from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_str_from_messages,
)
@ -2447,13 +2448,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
data["litellm_proxy_rate_limit_response"] = response
# Mirror into metadata so streaming success logging can find
# it via ``kwargs["litellm_params"]["metadata"]``.
self._stash_value_in_metadata_channels(
self._stash_value_in_internal_metadata(
data=data,
key=RATE_LIMIT_RESPONSE_KEY,
value=response,
)
if parallel_slot_id is not None:
self._stash_value_in_metadata_channels(
self._stash_value_in_internal_metadata(
data=data,
key=MAX_PARALLEL_SLOT_ACQUIRED_KEY,
value={
@ -2533,7 +2534,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requested_model=requested_model,
)
else:
self._stash_value_in_metadata_channels(
self._stash_value_in_internal_metadata(
data=data,
key=RATE_LIMIT_DESCRIPTORS_KEY,
value=descriptors,
@ -2566,7 +2567,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
data["litellm_proxy_rate_limit_response"] = tpm_response
# Keep the metadata stash in sync when this is the
# first snapshot written.
self._stash_value_in_metadata_channels(
self._stash_value_in_internal_metadata(
data=data,
key=RATE_LIMIT_RESPONSE_KEY,
value=tpm_response,
@ -2803,19 +2804,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return merged
@staticmethod
def _stash_value_in_metadata_channels(
def _stash_value_in_internal_metadata(
data: Dict[str, Any],
key: str,
value: Any,
) -> None:
for channel in ("metadata", "litellm_metadata"):
existing = data.get(channel)
if isinstance(existing, dict):
existing[key] = value
elif channel == "metadata":
# ``litellm_metadata`` is owned by the router; don't conjure
# it here.
data[channel] = {key: value}
# Writes only the proxy-internal bucket. Routes that own
# ``litellm_metadata`` (Responses, /v1/messages, batches, files) expose
# ``metadata`` as a provider request parameter, so creating or adding to
# it here would forward internal state upstream.
_, metadata_bucket = get_or_create_metadata_bucket(data)
metadata_bucket[key] = value
@classmethod
def _stash_reservation_in_data(
@ -2831,11 +2830,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
scopes_payload: Optional[List[List[str]]] = [[k, v] for k, v in reserved_scopes] if reserved_scopes else None
cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_TOKENS_KEY, value=estimated_tokens)
cls._stash_value_in_internal_metadata(data=data, key=TPM_RESERVED_TOKENS_KEY, value=estimated_tokens)
if reserved_model:
cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_MODEL_KEY, value=reserved_model)
cls._stash_value_in_internal_metadata(data=data, key=TPM_RESERVED_MODEL_KEY, value=reserved_model)
if scopes_payload is not None:
cls._stash_value_in_metadata_channels(data=data, key=TPM_RESERVED_SCOPES_KEY, value=scopes_payload)
cls._stash_value_in_internal_metadata(data=data, key=TPM_RESERVED_SCOPES_KEY, value=scopes_payload)
@staticmethod
def _lookup_stashed_value(
@ -2858,9 +2857,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
return candidate
litellm_params = kwargs.get("litellm_params")
if isinstance(litellm_params, dict):
lp_metadata = litellm_params.get("metadata")
if isinstance(lp_metadata, dict):
candidate = lp_metadata.get(key)
for channel in ("litellm_metadata", "metadata"):
lp_metadata = litellm_params.get(channel)
if isinstance(lp_metadata, dict) and lp_metadata.get(key) is not None:
return lp_metadata[key]
if candidate is None and isinstance(standard_logging_metadata, dict):
candidate = standard_logging_metadata.get(key)
return candidate

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -154,6 +154,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS = (
"applied_guardrails",
"applied_policies",
"policy_sources",
"routing_decision",
"pillar_response_headers",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
@ -201,6 +202,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = (
"applied_guardrails",
"applied_policies",
"policy_sources",
"routing_decision",
"standard_logging_object",
"proxy_server_request",
"secret_fields",

4
litellm/proxy/management_endpoints/common_utils.py
View file

@ -216,7 +216,7 @@ async def _user_has_admin_privileges(
teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_obj.teams}})
for team in teams:
team_obj = LiteLLM_TeamTable(**team.model_dump())
team_obj = LiteLLM_TeamTable.model_validate(team.model_dump())
if _is_user_team_admin(user_api_key_dict=user_api_key_dict, team_obj=team_obj):
return True
@ -288,7 +288,7 @@ async def _team_admin_can_invite_user(
for team in teams
if _is_user_team_admin(
user_api_key_dict=user_api_key_dict,
team_obj=LiteLLM_TeamTable(**team.model_dump()),
team_obj=LiteLLM_TeamTable.model_validate(team.model_dump()),
)
]
if not admin_team_ids:

105
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -45,6 +45,14 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
generate_key_helper_fn,
prepare_metadata_fields,
)
from litellm.proxy.common_utils.user_api_key_cache import (
object_permission_cache_key,
user_object_permission_id_cache_key,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
handle_update_object_permission_common,
)
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
from litellm.repositories.organization_repository import OrganizationRepository
@ -401,7 +409,7 @@ async def new_user(
- duration: Optional[str] - Duration for the key auto-created on `/user/new`. Default is None.
- key_alias: Optional[str] - Alias for the key auto-created on `/user/new`. Default is None.
- sso_user_id: Optional[str] - The id of the user in the SSO provider.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1"], "mcp_servers": ["github"], "mcp_tool_permissions": {"github": ["list_issues"]}}. The MCP grants act as a ceiling on every key this user holds. IF null or {} then no object permission.
- prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts.
- organizations: List[str] - List of organization id's the user is a member of
- budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].
@ -466,6 +474,10 @@ async def new_user(
data_json = data.json() # type: ignore
data_json = _update_internal_new_user_params(data_json, data)
# Persist the requested grants as their own row and link it, mirroring key/team creation.
# generate_key_helper_fn only forwards object_permission_id, so without this the entitlement
# the caller sent would be dropped on the floor.
data_json = await _set_object_permission(data_json=data_json, prisma_client=prisma_client)
_hash_password_in_dict(data_json)
teams = data.teams
if teams is None:
@ -852,9 +864,12 @@ async def _check_user_info_v2_access(
if prisma_client is None:
return None
# Helper: fetch the target user row (reused across branches)
# Helper: fetch the target user row (reused across branches). object_permission is included so
# callers can read the user's MCP/vector-store entitlements without a second round trip.
async def _fetch_target_user():
return await UserRepository(prisma_client).table.find_unique(where={"user_id": target_user_id})
return await UserRepository(prisma_client).table.find_unique(
where={"user_id": target_user_id}, include={"object_permission": True}
)
# Rule 1: Proxy admins — fetch and return the target row directly
if _user_has_admin_view(user_api_key_dict):
@ -972,6 +987,7 @@ async def user_info_v2(
updated_at=user_data.get("updated_at"),
sso_user_id=user_data.get("sso_user_id"),
teams=user_data.get("teams") or [],
object_permission=user_data.get("object_permission"),
)
except Exception as e:
verbose_proxy_logger.exception(
@ -1207,6 +1223,48 @@ async def _invalidate_user_spend_counter_if_changed(
await _invalidate_spend_counter(counter_key=f"spend:user:{non_default_values['user_id']}")
def _clears_object_permission(user_request: UpdateUserRequest) -> bool:
"""Whether the caller explicitly asked to remove this user's object_permission.
Distinguishes "sent nothing" from "sent an empty grant set". Only the latter clears; an omitted
field must leave an existing entitlement alone.
"""
if "object_permission" not in (user_request.fields_set() if hasattr(user_request, "fields_set") else set()):
return False
sent = user_request.object_permission
return sent is None or not sent.model_dump(exclude_unset=True, exclude_none=True)
async def _invalidate_cached_user_entitlement(user_id: str | None, object_permission_ids: tuple[str, ...]) -> None:
"""Drop the cache entries an entitlement change makes stale.
All three kinds are needed: a permission row is cached under its own id (so re-reading the same
link still yields the OLD grants), the ``user_id -> object_permission_id`` link is cached
separately (so a user who previously had NO entitlement keeps its "none" sentinel), and the user
row itself is cached whole. Leaving any behind means an admin revoking a tool keeps serving it
until the management-object TTL expires.
Both the outgoing and incoming permission ids are passed, because a clear leaves no incoming id
at all and an upsert may mint a new row; invalidating only one of the two leaves the other's
grants live.
Each deletion is isolated: one that fails must not skip the others, or a single unreachable key
would silently leave the rest of a revocation in place. Best-effort overall, exactly as the caches
are everywhere else, since one we cannot clear still expires on its own.
"""
from litellm.proxy.proxy_server import user_api_key_cache
keys = (
*(object_permission_cache_key(permission_id) for permission_id in dict.fromkeys(object_permission_ids)),
*((user_object_permission_id_cache_key(user_id), user_id) if user_id is not None else ()),
)
for key in keys:
try:
await user_api_key_cache.async_delete_cache(key=key)
except Exception as e: # noqa: BLE001 # a cache we cannot clear still expires; never fail the write
verbose_proxy_logger.warning(f"Failed to invalidate cached entitlement key {key!r}: {str(e)}")
async def _update_single_user_helper(
user_request: UpdateUserRequest,
user_api_key_dict: UserAPIKeyAuth,
@ -1259,9 +1317,15 @@ async def _update_single_user_helper(
)
_is_self_update = _target_user_id is not None and user_api_key_dict.user_id == _target_user_id
if _is_self_update and user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value:
_protected_fields = ("max_budget", "soft_budget", "spend")
# object_permission is a CEILING on what this human may reach, so a self-write is an
# escalation path: sending an empty grant list means "no restriction" and would lift a
# restriction an admin placed on them. Checked against the fields the caller actually SENT,
# because `_update_internal_user_params` drops empty values, and `object_permission: {}` is
# precisely the clear-my-own-ceiling case this must refuse.
_sent_fields = user_request.fields_set() if hasattr(user_request, "fields_set") else set()
_protected_fields = ("max_budget", "soft_budget", "spend", "object_permission")
for _field in _protected_fields:
if _field in non_default_values:
if _field in non_default_values or _field in _sent_fields:
raise HTTPException(
status_code=403,
detail={
@ -1282,6 +1346,22 @@ async def _update_single_user_helper(
# Reject NaN/±inf spend before it can reach the DB / spend counter.
validate_finite_spend(non_default_values.get("spend"))
# Upsert the grants into their own row and link it, mirroring /key/update and /team/update.
# This also removes object_permission from the payload, which is not a column on the user table.
if "object_permission" in non_default_values:
object_permission_id = await handle_update_object_permission_common(
data_json=non_default_values,
existing_object_permission_id=getattr(existing_user_row, "object_permission_id", None),
prisma_client=prisma_client,
)
if object_permission_id is not None:
non_default_values["object_permission_id"] = object_permission_id
elif _clears_object_permission(user_request):
# An explicit `{}` or null means "no object permission", which the merge-based upsert cannot
# express: merging an empty grant set over the existing row leaves every grant in place. So
# the link is dropped instead, which is what makes the documented clear actually clear.
non_default_values["object_permission_id"] = None
# Perform the update
response: dict[str, Any] | None = None
@ -1326,6 +1406,19 @@ async def _update_single_user_helper(
await _invalidate_user_spend_counter_if_changed(non_default_values)
if "object_permission_id" in non_default_values:
await _invalidate_cached_user_entitlement(
user_id=non_default_values.get("user_id"),
object_permission_ids=tuple(
permission_id
for permission_id in (
getattr(existing_user_row, "object_permission_id", None),
non_default_values.get("object_permission_id"),
)
if isinstance(permission_id, str)
),
)
if response is None:
raise HTTPException(
status_code=400,
@ -1407,7 +1500,7 @@ async def user_update(
- team_id: Optional[str] - [DEPRECATED PARAM] The team id of the user. Default is None.
- duration: Optional[str] - [NOT IMPLEMENTED].
- key_alias: Optional[str] - [NOT IMPLEMENTED].
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"]}. IF null or {} then no object permission.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - internal user-specific object permission. Example - {"vector_stores": ["vector_store_1"], "mcp_servers": ["github"], "mcp_tool_permissions": {"github": ["list_issues"]}}. The MCP grants act as a ceiling on every key this user holds. IF null or {} then no object permission.
- prompts: Optional[List[str]] - List of allowed prompts for the user. If specified, the user will only be able to use these specific prompts.
- budget_limits: Optional[list] - List of concurrent budget windows for the user. Each window specifies a budget_limit, time_period, and optional budget_duration. Example - [{"budget_limit": 10.0, "time_period": "1d"}, {"budget_limit": 50.0, "time_period": "7d"}].

4
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -459,7 +459,7 @@ if MCP_AVAILABLE:
payload_dict: dict[str, Any] = loaded
try:
return MCPServer(**payload_dict)
return MCPServer.model_validate(payload_dict)
except Exception as e:
verbose_proxy_logger.debug(f"Invalid temporary MCP server payload in Redis cache: {str(e)}")
return None
@ -704,7 +704,7 @@ if MCP_AVAILABLE:
except AttributeError:
payload_dict = payload.dict() # type: ignore[attr-defined]
payload_dict["credentials"] = inherited_credentials
return NewMCPServerRequest(**payload_dict)
return NewMCPServerRequest.model_validate(payload_dict)
def _build_temporary_mcp_server_record(
payload: NewMCPServerRequest,

8
litellm/proxy/management_helpers/utils.py
View file

@ -308,7 +308,7 @@ async def add_new_member(
)
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
if _returned_user is not None:
returned_user = LiteLLM_UserTable(**_returned_user.model_dump())
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif new_member.user_email is not None:
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
## user email is not unique acc. to prisma schema -> future improvement
@ -323,11 +323,11 @@ async def add_new_member(
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user") # type: ignore
if _returned_user is not None:
returned_user = LiteLLM_UserTable(**_returned_user.model_dump())
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
elif len(existing_user_row) == 1:
user_info = existing_user_row[0]
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id)
returned_user = LiteLLM_UserTable(**user_info.model_dump())
returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump())
elif len(existing_user_row) > 1:
raise HTTPException(
status_code=400,
@ -354,7 +354,7 @@ async def add_new_member(
include={"litellm_budget_table": True},
)
returned_team_membership = LiteLLM_TeamMembership(**_returned_team_membership.model_dump())
returned_team_membership = LiteLLM_TeamMembership.model_validate(_returned_team_membership.model_dump())
if returned_user is None:
raise Exception("Unable to update user table with membership information!")

39
litellm/proxy/proxy_server.py
View file

@ -5398,7 +5398,7 @@ class ProxyConfig:
# decrypt values
for k, v in _litellm_params.items():
_litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v)
_litellm_params = LiteLLM_Params(**_litellm_params)
_litellm_params = LiteLLM_Params.model_validate(_litellm_params)
else:
verbose_proxy_logger.error(
@ -5429,7 +5429,7 @@ class ProxyConfig:
# decrypt values
for k, v in _litellm_params.items():
_litellm_params[k] = self._resolve_db_litellm_param(key=k, value=v)
_litellm_params = LiteLLM_Params(**_litellm_params)
_litellm_params = LiteLLM_Params.model_validate(_litellm_params)
else:
verbose_proxy_logger.error(
f"Invalid model added to proxy db. Invalid litellm params. litellm_params={_litellm_params}"
@ -11827,6 +11827,22 @@ def _sort_models(
return all_models
def _is_auto_router_model(model: Mapping[str, object]) -> bool:
"""
True for any auto-router deployment, i.e. every `auto_router/*` strategy
(semantic, complexity, adaptive, quality).
Router._is_auto_router_deployment is deliberately narrower; it answers "is this the
*semantic* auto-router strategy" and returns False for the complexity and adaptive
prefixes, so it is not reusable here.
"""
litellm_params = model.get("litellm_params")
if not isinstance(litellm_params, Mapping):
return False
litellm_model = litellm_params.get("model")
return isinstance(litellm_model, str) and litellm_model.startswith("auto_router/")
def _paginate_models_response(
all_models: List[Dict[str, Any]],
page: int,
@ -12121,6 +12137,15 @@ async def model_info_v2(
"asc",
description="Sort order. Options: asc, desc",
),
exclude_auto_routers: bool | None = fastapi.Query(
False,
description=(
"Omit auto-router deployments (litellm model prefixed `auto_router/`). "
"They select among deployments rather than being deployments themselves, so a "
"caller rendering a deployment list can leave them out. Defaults to false, so "
"existing callers are unaffected"
),
),
):
"""
Paginated model metadata for proxy deployments (pricing, provider, team access).
@ -12288,6 +12313,11 @@ async def model_info_v2(
user_api_key_dict=user_api_key_dict,
)
# `is True` because direct-call tests bypass FastAPI, so the Query default arrives as a
# truthy sentinel object rather than False.
if exclude_auto_routers is True:
all_models = [m for m in all_models if not _is_auto_router_model(m)]
# Update total count to include agents
search_total_count = len(all_models)
@ -13063,7 +13093,7 @@ def _get_model_group_info(
_model_group_info = llm_router.get_model_group_info(model_group=model)
if _model_group_info is not None:
model_groups.append(ModelGroupInfoProxy(**_model_group_info.model_dump()))
model_groups.append(ModelGroupInfoProxy.model_validate(_model_group_info.model_dump()))
else:
model_group_info = ModelGroupInfoProxy(
model_group=model,
@ -14782,7 +14812,7 @@ async def update_config_general_settings(
)
try:
ConfigGeneralSettings(**{data.field_name: data.field_value})
ConfigGeneralSettings.model_validate({data.field_name: data.field_value})
except Exception:
raise HTTPException(
status_code=400,
@ -15228,7 +15258,6 @@ async def get_config_list(
"forward_client_headers_to_llm_api": {"type": "Boolean"},
"mcp_required_fields": {"type": "List"},
"cancel_on_disconnect": {"type": "Boolean"},
"skip_user_budget_on_team_key": {"type": "Boolean"},
"disable_auto_add_proxy_admin_to_teams": {"type": "Boolean"},
}

6
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -155,7 +155,6 @@ async def reserve_budget_for_request(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
fail_closed_budget_enforcement: bool = False,
) -> Optional[dict]:
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
@ -175,7 +174,6 @@ async def reserve_budget_for_request(
proxy_logging_obj=proxy_logging_obj,
end_user_id=end_user_id,
end_user_object=end_user_object,
skip_user_budget_on_team_key=skip_user_budget_on_team_key,
)
if not counters:
return None
@ -333,7 +331,6 @@ async def _get_budget_counters(
proxy_logging_obj: ProxyLogging,
end_user_id: Optional[str] = None,
end_user_object: Optional[Any] = None,
skip_user_budget_on_team_key: bool = False,
) -> List[_BudgetCounter]:
counters: List[_BudgetCounter] = []
@ -382,9 +379,8 @@ async def _get_budget_counters(
)
)
is_team_key = team_object is not None and team_object.team_id is not None
if (
not (is_team_key and skip_user_budget_on_team_key)
(team_object is None or team_object.team_id is None)
and user_object is not None
and user_object.user_id is not None
and user_object.max_budget is not None

88
litellm/proxy/utils.py
View file

@ -109,6 +109,7 @@ from litellm.litellm_core_utils.core_helpers import coerce_token_limit
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.llms import load_guardrail_translation_mappings
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import (
AlertType,
@ -172,6 +173,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionRe
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
if TYPE_CHECKING:
from mcp.types import CallToolResult
from opentelemetry.trace import Span as _Span
from prisma.client import TransactionManager
@ -185,6 +187,8 @@ else:
unified_guardrail = UnifiedLLMGuardrails()
NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages})
def print_verbose(print_statement):
"""
@ -1760,6 +1764,20 @@ class ProxyLogging:
cache[sig] = caps
return caps
@staticmethod
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
from litellm.litellm_core_utils.api_route_to_call_types import (
get_call_types_for_route,
)
route = user_api_key_dict.request_route
if not route:
return False
call_types = get_call_types_for_route(route)
if not call_types:
return False
return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
@staticmethod
def has_post_call_response_headers_callbacks() -> bool:
return ProxyLogging._callback_capabilities().has_post_call_response_headers
@ -2470,6 +2488,59 @@ class ProxyLogging:
if raised:
raise raised[0]
async def post_mcp_call_hook(
self,
response: "CallToolResult",
request_data: Mapping[str, Any],
user_api_key_dict: UserAPIKeyAuth | None = None,
) -> "CallToolResult":
"""
Run guardrails configured for ``post_mcp_call`` against an MCP tool result.
The MCP counterpart of ``post_call_success_hook``: guardrails that
implement ``apply_guardrail`` see the tool result's text through the
unified guardrail seam (``MCPGuardrailTranslationHandler``), so a text
guardrail can mask sensitive values in the result without any MCP-specific
code of its own. Guardrails that instead implement
``async_post_mcp_tool_call_hook`` are dispatched by
``Logging.async_post_mcp_tool_call_hook`` and are not run here.
A guardrail that rejects the result raises, and the exception propagates
(matching the inbound ``pre_mcp_call`` behavior) rather than being
swallowed into an unguarded result.
"""
caps = ProxyLogging._callback_capabilities()
if not caps.has_guardrail:
return response
handler_cls = load_guardrail_translation_mappings().get(CallTypes.call_mcp_tool)
if handler_cls is None:
verbose_proxy_logger.debug("MCP guardrail translation handler unavailable; skipping post_mcp_call hook")
return response
for callback in caps.resolved_callbacks:
if not isinstance(callback, CustomGuardrail):
continue
if "apply_guardrail" not in type(callback).__dict__:
continue
if (
callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_mcp_call)
is not True
):
continue
response = await self._run_guardrail_with_metrics(
callback,
handler_cls().process_output_response(
response=response,
guardrail_to_apply=callback,
litellm_logging_obj=request_data.get("litellm_logging_obj"),
user_api_key_dict=user_api_key_dict,
request_data=request_data,
),
"post_mcp_call",
)
return response
async def post_call_response_headers_hook(
self,
data: dict,
@ -2668,6 +2739,7 @@ class ProxyLogging:
request_data = _check_and_merge_model_level_guardrails(data=request_data, llm_router=llm_router)
current_response = response
stream_needs_translation = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict)
for resolved_callback, kind in caps.iterator_overrides:
if isinstance(resolved_callback, CustomGuardrail):
@ -2676,7 +2748,18 @@ class ProxyLogging:
is not True
):
continue
if kind == "override":
effective_kind = (
"apply_guardrail"
if (
kind == "override"
and stream_needs_translation
and isinstance(resolved_callback, CustomGuardrail)
and resolved_callback.uses_apply_guardrail_interface()
and not resolved_callback.mask_response_content
)
else kind
)
if effective_kind == "override":
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
resolved_callback.async_post_call_streaming_iterator_hook(
@ -2687,13 +2770,14 @@ class ProxyLogging:
)
else:
# kind == "apply_guardrail": route through unified_guardrail
request_data["guardrail_to_apply"] = resolved_callback
current_response = self._wrap_streaming_iterator_with_enrichment(
resolved_callback,
unified_guardrail.async_post_call_streaming_iterator_hook(
user_api_key_dict=user_api_key_dict,
request_data=request_data,
response=current_response,
guardrail_to_apply=resolved_callback,
buffer_until_moderated_default=(kind == "override"),
),
)

54
litellm/repositories/base_repository.py
View file

@ -3,38 +3,58 @@ Base repository class with common functionality.
"""
from abc import ABC, abstractmethod
from typing import Any, Dict, Generic, List, Optional, Type, TypeVar
from collections.abc import Iterable, Mapping, Sequence
from typing import Any, Dict, Generic, List, Optional, Protocol, Tuple, Type, TypeVar, Union, runtime_checkable
from pydantic import BaseModel
T = TypeVar("T", bound=BaseModel)
def _record_to_dict(record: Any) -> Dict[str, Any]:
if isinstance(record, dict):
return record
if hasattr(record, "model_dump") and callable(record.model_dump):
@runtime_checkable
class SupportsModelDump(Protocol):
def model_dump(self) -> Dict[str, object]: ...
@runtime_checkable
class SupportsDict(Protocol):
def dict(self) -> Dict[str, object]: ...
DbRecord = Union[
Mapping[str, object],
SupportsModelDump,
SupportsDict,
Sequence[Tuple[str, object]],
]
def record_to_dict(record: DbRecord) -> Mapping[str, object]:
"""Project a database record into a mapping of column name to value."""
if isinstance(record, SupportsModelDump):
return record.model_dump()
if hasattr(record, "dict") and callable(record.dict):
if isinstance(record, SupportsDict):
return record.dict()
return dict(record)
if isinstance(record, Mapping):
return record
return {key: value for key, value in record}
class BaseRepository(ABC, Generic[T]):
"""Abstract base class for all repositories."""
def __init__(self, prisma_client: Any):
def __init__(self, prisma_client: Any): # any-ok: PrismaClient is an untyped runtime wrapper
self._prisma_client = prisma_client
@property
def prisma_client(self) -> Any:
def prisma_client(self) -> Any: # any-ok: PrismaClient is an untyped runtime wrapper
if self._prisma_client is None:
raise RuntimeError("No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys")
return self._prisma_client
@property
@abstractmethod
def table(self) -> Any:
def table(self) -> Any: # any-ok: Prisma table actions are reached through the untyped client wrapper
"""Return the Prisma table for this repository."""
...
@ -44,21 +64,15 @@ class BaseRepository(ABC, Generic[T]):
"""Return the domain model class for this repository."""
...
def _to_model(self, record: Any) -> Optional[T]:
def _to_model(self, record: Optional[DbRecord]) -> Optional[T]:
"""Convert a database record to a domain model."""
if record is None:
return None
return self.model_class(**_record_to_dict(record))
return self.model_class.model_validate(record_to_dict(record))
def _to_model_list(self, records: List[Any]) -> List[T]:
def _to_model_list(self, records: Iterable[Optional[DbRecord]]) -> List[T]:
"""Convert a list of database records to domain models."""
result: List[T] = []
for r in records:
if r is not None:
model = self._to_model(r)
if model is not None:
result.append(model)
return result
return [model for record in records if record is not None and (model := self._to_model(record)) is not None]
async def find_by_id(self, id_value: str, id_field: str = "id") -> Optional[T]:
"""Find a record by its primary key."""

6
litellm/repositories/organization_repository.py
View file

@ -26,10 +26,8 @@ class OrganizationRepository(BaseRepository[LiteLLM_OrganizationTable]):
async def find_by_alias(self, organization_alias: str) -> Optional[LiteLLM_OrganizationTable]:
"""Find an organization by alias."""
records = await self.table.find_many(where={"organization_alias": organization_alias})
if records:
return self._to_model(records[0])
return None
organizations = await self.find_many(where={"organization_alias": organization_alias})
return organizations[0] if organizations else None
async def create_organization(
self,

9
litellm/repositories/project_repository.py
View file

@ -24,15 +24,12 @@ class ProjectRepository(BaseRepository[LiteLLM_ProjectTable]):
async def find_by_alias(self, project_alias: str) -> Optional[LiteLLM_ProjectTable]:
"""Find a project by alias."""
records = await self.table.find_many(where={"project_alias": project_alias})
if records:
return self._to_model(records[0])
return None
projects = await self.find_many(where={"project_alias": project_alias})
return projects[0] if projects else None
async def find_by_team_id(self, team_id: str) -> List[LiteLLM_ProjectTable]:
"""Find all projects belonging to a team."""
records = await self.table.find_many(where={"team_id": team_id})
return self._to_model_list(records)
return await self.find_many(where={"team_id": team_id})
async def create_project(
self,

56
litellm/repositories/team_repository.py
View file

@ -3,55 +3,59 @@ Team repository for database operations on LiteLLM_TeamTable.
"""
import json
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
from pydantic import TypeAdapter
from litellm.models.team import LiteLLM_TeamTable, Member
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.base_repository import (
BaseRepository,
DbRecord,
record_to_dict,
)
if TYPE_CHECKING:
from prisma import Prisma
_MEMBERS_WITH_ROLES_ADAPTER = TypeAdapter(list[Member])
_JSON_ENCODED_TEAM_FIELDS = (
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
"budget_limits",
"members_with_roles",
)
class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
"""Repository for team database operations."""
@property
def table(self) -> Any:
def table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
return self.prisma_client.db.litellm_teamtable
@property
def deleted_table(self) -> Any:
def deleted_table(self) -> Any: # any-ok: PrismaClient.db is an untyped runtime wrapper
return self.prisma_client.db.litellm_deletedteamtable
@property
def model_class(self) -> Type[LiteLLM_TeamTable]:
return LiteLLM_TeamTable
def _to_model(self, record: Any) -> Optional[LiteLLM_TeamTable]:
def _to_model(self, record: Optional[DbRecord]) -> Optional[LiteLLM_TeamTable]:
"""Convert a database record to a Team model."""
if record is None:
return None
data = record.dict() if hasattr(record, "dict") else dict(record)
data = {
field: json.loads(value) if field in _JSON_ENCODED_TEAM_FIELDS and isinstance(value, str) else value
for field, value in record_to_dict(record).items()
}
json_fields = [
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
"budget_limits",
"members_with_roles",
]
for field in json_fields:
if isinstance(data.get(field), str):
data[field] = json.loads(data[field])
return LiteLLM_TeamTable(**data)
return LiteLLM_TeamTable.model_validate(data)
async def get_members_with_roles_locked(self, tx: "Prisma", team_id: str) -> List[Member]:
"""Return the team's members_with_roles, locking the row FOR UPDATE.
@ -103,8 +107,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
organization_id: Optional[str] = None,
admins: Optional[List[str]] = None,
members: Optional[List[str]] = None,
members_with_roles: Optional[Dict[str, Any]] = None,
metadata: Optional[Dict[str, Any]] = None,
members_with_roles: Optional[Mapping[str, object]] = None,
metadata: Optional[Mapping[str, object]] = None,
max_budget: Optional[float] = None,
soft_budget: Optional[float] = None,
models: Optional[List[str]] = None,
@ -115,7 +119,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
object_permission_id: Optional[str] = None,
) -> LiteLLM_TeamTable:
"""Create a new team."""
data: Dict[str, Any] = {"team_id": team_id}
data: Dict[str, object] = {"team_id": team_id}
if team_alias is not None:
data["team_alias"] = team_alias
if organization_id is not None:
@ -154,8 +158,8 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
organization_id: Optional[str] = None,
admins: Optional[List[str]] = None,
members: Optional[List[str]] = None,
members_with_roles: Optional[Dict[str, Any]] = None,
metadata: Optional[Dict[str, Any]] = None,
members_with_roles: Optional[Mapping[str, object]] = None,
metadata: Optional[Mapping[str, object]] = None,
max_budget: Optional[float] = None,
soft_budget: Optional[float] = None,
models: Optional[List[str]] = None,
@ -167,7 +171,7 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
object_permission_id: Optional[str] = None,
) -> Optional[LiteLLM_TeamTable]:
"""Update a team."""
data: Dict[str, Any] = {}
data: Dict[str, object] = {}
if team_alias is not None:
data["team_alias"] = team_alias
if organization_id is not None:
@ -228,9 +232,9 @@ class TeamRepository(BaseRepository[LiteLLM_TeamTable]):
return team
def _build_archive_data(self, team: LiteLLM_TeamTable) -> Dict[str, Any]:
def _build_archive_data(self, team: LiteLLM_TeamTable) -> Dict[str, object]:
"""Build archive data dict with only columns that exist in LiteLLM_DeletedTeamTable."""
data: Dict[str, Any] = {"team_id": team.team_id}
data: Dict[str, object] = {"team_id": team.team_id}
if team.team_alias is not None:
data["team_alias"] = team.team_alias
if team.organization_id is not None:

58
litellm/repositories/verification_token_repository.py
View file

@ -3,14 +3,18 @@ VerificationToken repository for database operations on LiteLLM_VerificationToke
"""
import json
from collections.abc import Iterator, Mapping
from collections.abc import Mapping
from datetime import datetime
from typing import TYPE_CHECKING, Any, Protocol
from typing import TYPE_CHECKING, Any
from litellm.models.verification_token import (
LiteLLM_VerificationToken,
)
from litellm.repositories.base_repository import BaseRepository
from litellm.repositories.base_repository import (
BaseRepository,
DbRecord,
record_to_dict,
)
if TYPE_CHECKING:
from prisma.models import (
@ -19,11 +23,17 @@ if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
class _DictConvertible(Protocol):
def dict(self) -> dict[str, object]: ...
def __iter__(self) -> Iterator[tuple[str, object]]: ...
_JSON_ENCODED_TOKEN_FIELDS = (
"aliases",
"config",
"permissions",
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
"budget_limits",
"litellm_budget_table",
)
class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
@ -46,31 +56,21 @@ class VerificationTokenRepository(BaseRepository[LiteLLM_VerificationToken]):
def model_class(self) -> type[LiteLLM_VerificationToken]:
return LiteLLM_VerificationToken
def _to_model(self, record: _DictConvertible | None) -> LiteLLM_VerificationToken | None:
def _to_model(self, record: DbRecord | None) -> LiteLLM_VerificationToken | None:
"""Convert a database record to a VerificationToken model."""
if record is None:
return None
data = record.dict() if hasattr(record, "dict") else dict(record)
json_fields = [
"aliases",
"config",
"permissions",
"metadata",
"model_spend",
"model_max_budget",
"router_settings",
"budget_limits",
"litellm_budget_table",
]
for field in json_fields:
value = data.get(field)
if isinstance(value, str):
data[field] = json.loads(value)
if data.get("org_id") is None and data.get("organization_id") is not None:
data["org_id"] = data["organization_id"]
decoded = {
field: json.loads(value) if field in _JSON_ENCODED_TOKEN_FIELDS and isinstance(value, str) else value
for field, value in record_to_dict(record).items()
}
organization_id = decoded.get("organization_id")
data = (
decoded
if decoded.get("org_id") is not None or organization_id is None
else {**decoded, "org_id": organization_id}
)
return LiteLLM_VerificationToken.model_validate(data)

19
litellm/responses/mcp/litellm_proxy_mcp_handler.py
View file

@ -795,20 +795,33 @@ class LiteLLM_Proxy_MCP_Handler:
proxy_logging_obj=proxy_logging_obj,
)
if proxy_logging_obj:
result = await proxy_logging_obj.post_mcp_call_hook(
response=result,
request_data=(
litellm_logging_obj.model_call_details
if litellm_logging_obj
else {"mcp_tool_name": tool_name}
),
user_api_key_dict=user_api_key_auth,
)
if litellm_logging_obj:
try:
litellm_logging_obj.post_call(original_response=result)
end_time = datetime.now()
await litellm_logging_obj.async_post_mcp_tool_call_hook(
kwargs=litellm_logging_obj.model_call_details,
response_obj=result,
start_time=start_time,
end_time=end_time,
end_time=datetime.now(),
)
except Exception:
verbose_logger.exception("Failed to run post-call logging for MCP tool call %s", tool_name)
try:
await litellm_logging_obj.async_success_handler(
result=result,
start_time=start_time,
end_time=end_time,
end_time=datetime.now(),
)
except Exception:
verbose_logger.exception("Failed to log MCP tool call success for %s", tool_name)

70
litellm/router.py
View file

@ -71,6 +71,7 @@ from litellm.litellm_core_utils.core_helpers import (
_get_parent_otel_span_from_kwargs,
coerce_token_limit,
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
)
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
@ -210,8 +211,10 @@ from litellm.types.utils import (
from litellm.types.utils import ModelInfo
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import (
PROMPT_QUOTING_ROUTING_DECISION_FIELDS,
ModelResponseStream,
StandardLoggingPayload,
StandardLoggingRoutingDecision,
Usage,
)
from litellm.utils import (
@ -11158,6 +11161,7 @@ class Router:
router_strategy = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs)
if router_strategy is None:
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
return None
pre_routing_hook_response = await router_strategy.async_pre_routing_hook(
@ -11167,6 +11171,10 @@ class Router:
input=input,
specific_deployment=specific_deployment,
)
self._record_routing_decision(
request_kwargs=request_kwargs,
routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None),
)
# `model` (the alias, e.g. "smart-router") is never the deployment actually
# called - apply the alias's own litellm_params (besides `model` itself,
@ -11185,6 +11193,68 @@ class Router:
return pre_routing_hook_response
@staticmethod
def _record_routing_decision(
request_kwargs: dict,
routing_decision: StandardLoggingRoutingDecision | None,
) -> None:
"""Make the request's metadata describe THIS routing attempt, and only this one.
Fallbacks re-enter the hook with the same `request_kwargs`, so an attempt that
picks a plain model group after an auto-router group failed must clear the
earlier decision; leaving it would attribute the first router's tier and cause
to the deployment that actually served the request. Every attempt therefore
writes or clears, never just writes.
"""
if routing_decision is None:
for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")):
if isinstance(bucket, dict):
bucket.pop("routing_decision", None)
return
# `get_or_create_metadata_bucket` is the single owner of "which dict holds
# proxy-internal metadata": it picks `litellm_metadata` when present (so the
# decision never lands in the `metadata` dict that routes like /v1/messages
# forward to the provider) and replaces a non-dict value rather than silently
# skipping the write.
_, metadata_bucket = get_or_create_metadata_bucket(request_kwargs)
metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed(
request_kwargs=request_kwargs, routing_decision=routing_decision
)
@staticmethod
def _redact_prompt_text_if_needed(
request_kwargs: Mapping[str, Any],
routing_decision: StandardLoggingRoutingDecision,
) -> StandardLoggingRoutingDecision:
"""Drop verbatim prompt text from the record when message logging is redacted.
An operator who turns message logging off has said prompt content must not reach
the logs, so the fields that quote the prompt (the matched keywords, and the
signals that name them) are omitted. Derived values are kept, because a tier, a
cause, a score or an escalation flag aggregates the prompt rather than
reproducing any of it, and dropping them would leave the row unexplainable for
no privacy gain. Applied here rather than in each strategy so a strategy added
later cannot bypass it.
"""
from litellm.litellm_core_utils.redact_messages import (
should_redact_message_logging,
)
if not should_redact_message_logging(
{
"litellm_params": request_kwargs,
"standard_callback_dynamic_params": request_kwargs.get("standard_callback_dynamic_params"),
}
):
return routing_decision
kept = {
field: value
for field, value in routing_decision.items()
if field not in PROMPT_QUOTING_ROUTING_DECISION_FIELDS
}
return cast(StandardLoggingRoutingDecision, kept) # cast-ok: dropping optional keys preserves the type
def get_available_deployment(
self,
model: str,

13
litellm/router_strategy/adaptive_router/adaptive_router.py
View file

@ -23,6 +23,7 @@ from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_last_user_message,
)
from litellm.types.utils import StandardLoggingRoutingDecision
from litellm.router_strategy.adaptive_router.bandit import (
BanditCell,
apply_delta,
@ -193,7 +194,17 @@ class AdaptiveRouter:
if isinstance(kwargs_metadata, dict):
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = chosen_model
return PreRoutingHookResponse(model=chosen_model, messages=messages)
return PreRoutingHookResponse(
model=chosen_model,
messages=messages,
routing_decision=StandardLoggingRoutingDecision(
router_model_name=self.router_name,
router_type="adaptive",
routed_model=chosen_model,
cause="bandit",
request_type=request_type.value,
),
)
# ---- Pick model ------------------------------------------------------

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

@ -19,7 +19,7 @@ import asyncio
import random
import re
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Literal, Union, cast
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Union, cast
from pydantic import BaseModel
@ -27,7 +27,12 @@ from litellm._logging import verbose_router_logger
from litellm.constants import RETURN_RAW_MODEL_NAME_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.base_utils import type_to_response_format_param
from litellm.types.utils import ModelResponse
from litellm.types.utils import (
ModelResponse,
RoutingDecisionCause,
StandardLoggingRoutingDecision,
StandardLoggingRoutingDecisionTierBoundaries,
)
from .config import (
DEFAULT_CODE_KEYWORDS,
@ -135,6 +140,27 @@ class DimensionScore:
self.signal = signal
class KeywordOverride(NamedTuple):
"""A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired."""
tier: ComplexityTier
matched_keyword: str | None
class ClassificationOutcome(NamedTuple):
"""What the classifier decided and which mechanism actually produced it.
`cause` reflects the path that ran, not the configured classifier_type: an LLM
classifier that fails falls back to the heuristic scorer and reports it.
`score` is None on the LLM path, which produces a tier label and no score.
"""
tier: ComplexityTier
score: float | None
signals: tuple[str, ...]
cause: Literal["heuristic_scorer", "reasoning_override", "llm_classifier"]
class ComplexityRouter(CustomLogger):
"""
Complexity router that classifies requests and routes to appropriate models.
@ -256,6 +282,7 @@ class ComplexityRouter(CustomLogger):
def _score_keyword_match(
self,
text: str,
disclosable_text: str,
keywords: list[str],
name: str,
signal_label: str,
@ -264,6 +291,15 @@ class ComplexityRouter(CustomLogger):
) -> tuple[DimensionScore, int]:
"""Score based on keyword matches using word boundary matching.
Scoring reads `text`, which for most dimensions includes the system prompt.
The signal names only the terms that also appear in `disclosable_text`, the
caller's own message: signals are persisted to the request's spend log, which
the caller can read, so naming a term matched solely in the system prompt would
let a caller recover configured terms from a prompt it cannot see. Terms it did
not supply are reported as a count instead, which explains the score without
disclosing anything. `disclosable_text` is required rather than defaulted so a
future dimension has to state which text it is willing to quote.
Returns:
Tuple of (DimensionScore, match_count) so callers can reuse the count.
"""
@ -272,18 +308,13 @@ class ComplexityRouter(CustomLogger):
matches = [kw for kw in keywords if self._keyword_matches(text, kw)]
match_count = len(matches)
if match_count < low_threshold:
return DimensionScore(name, score_none, None), match_count
if match_count >= high_threshold:
return (
DimensionScore(name, score_high, f"{signal_label} ({', '.join(matches[:3])})"),
match_count,
)
if match_count >= low_threshold:
return (
DimensionScore(name, score_low, f"{signal_label} ({', '.join(matches[:3])})"),
match_count,
)
return DimensionScore(name, score_none, None), match_count
disclosable = [kw for kw in matches if self._keyword_matches(disclosable_text, kw)]
detail = ", ".join(disclosable[:3]) if disclosable else f"{match_count} matches"
score = score_high if match_count >= high_threshold else score_low
return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count
def _score_multi_step(self, text: str) -> DimensionScore:
"""Score based on multi-step patterns."""
@ -300,8 +331,19 @@ class ComplexityRouter(CustomLogger):
return DimensionScore("questionComplexity", 0, None)
def classify(self, prompt: str, system_prompt: str | None = None) -> tuple[ComplexityTier, float, list[str]]:
"""Classify a prompt by complexity, discarding which rule decided the tier.
Kept for callers that only need the tier and score; `_score_and_classify` is the
single computation behind both, so the two can never disagree.
"""
Classify a prompt by complexity.
tier, score, signals, _cause = self._score_and_classify(prompt, system_prompt)
return tier, score, list(signals)
def _score_and_classify(
self, prompt: str, system_prompt: str | None = None
) -> tuple[ComplexityTier, float, tuple[str, ...], Literal["heuristic_scorer", "reasoning_override"]]:
"""
Classify a prompt by complexity, reporting whether the score chose the tier.
Args:
prompt: The user's prompt/message.
@ -327,6 +369,7 @@ class ComplexityRouter(CustomLogger):
# Score all dimensions, capturing match counts where needed
code_score, _ = self._score_keyword_match(
full_text,
user_text,
self.code_keywords,
"codePresence",
"code",
@ -334,6 +377,7 @@ class ComplexityRouter(CustomLogger):
(0, 0.5, 1.0),
)
reasoning_score, reasoning_match_count = self._score_keyword_match(
user_text,
user_text,
self.reasoning_keywords,
"reasoningMarkers",
@ -343,6 +387,7 @@ class ComplexityRouter(CustomLogger):
)
technical_score, _ = self._score_keyword_match(
full_text,
user_text,
self.technical_keywords,
"technicalTerms",
"technical",
@ -351,6 +396,7 @@ class ComplexityRouter(CustomLogger):
)
simple_score, _ = self._score_keyword_match(
full_text,
user_text,
self.simple_keywords,
"simpleIndicators",
"simple",
@ -378,48 +424,112 @@ class ComplexityRouter(CustomLogger):
# Check for reasoning override (2+ reasoning markers)
# Reuse match count from _score_keyword_match to avoid scanning twice
if reasoning_match_count >= 2:
return ComplexityTier.REASONING, weighted_score, signals
return ComplexityTier.REASONING, weighted_score, tuple(signals), "reasoning_override"
# Map score to tier
boundaries = self.config.tier_boundaries
simple_medium = boundaries.get("simple_medium", 0.15)
medium_complex = boundaries.get("medium_complex", 0.35)
complex_reasoning = boundaries.get("complex_reasoning", 0.60)
if weighted_score < simple_medium:
boundaries = self._effective_tier_boundaries()
if weighted_score < boundaries["simple_medium"]:
tier = ComplexityTier.SIMPLE
elif weighted_score < medium_complex:
elif weighted_score < boundaries["medium_complex"]:
tier = ComplexityTier.MEDIUM
elif weighted_score < complex_reasoning:
elif weighted_score < boundaries["complex_reasoning"]:
tier = ComplexityTier.COMPLEX
else:
tier = ComplexityTier.REASONING
return tier, weighted_score, signals
return tier, weighted_score, tuple(signals), "heuristic_scorer"
def _effective_tier_boundaries(self) -> StandardLoggingRoutingDecisionTierBoundaries:
"""The tier boundaries in effect, with the documented defaults filled in.
Shared by score-to-tier mapping and the per-request routing decision snapshot,
so a logged decision always reflects the boundaries that actually applied.
"""
boundaries = self.config.tier_boundaries
return StandardLoggingRoutingDecisionTierBoundaries(
simple_medium=boundaries.get("simple_medium", 0.15),
medium_complex=boundaries.get("medium_complex", 0.35),
complex_reasoning=boundaries.get("complex_reasoning", 0.60),
)
def _build_routing_decision(
self,
*,
routed_model: str,
cause: RoutingDecisionCause,
tier: ComplexityTier | None = None,
score: float | None = None,
signals: tuple[str, ...] | None = None,
matched_keyword: str | None = None,
escalation_keyword: str | None = None,
escalated: bool = False,
classifier_model: str | None = None,
) -> StandardLoggingRoutingDecision:
"""Assemble the per-request provenance record for this router's decision.
Optional facts are omitted rather than set to None, so a spend log row only
carries the keys that applied to its path. `tier_boundaries` rides with
`score` because the score is only interpretable against the boundaries that
mapped it to a tier.
"""
decision = StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="complexity",
routed_model=routed_model,
cause=cause,
)
if tier is not None:
decision["tier"] = tier.value
if score is not None:
decision["score"] = score
decision["tier_boundaries"] = self._effective_tier_boundaries()
if signals:
# Stored as a list because this record is serialized to JSON for the spend
# log and read back as an array by the dashboard; a sequence type that only
# happens to survive the serializer would make the wire shape depend on it.
decision["signals"] = list(signals)
if matched_keyword is not None:
decision["matched_keyword"] = matched_keyword
if escalation_keyword is not None:
# Two separate facts: the caller asked to escalate, and whether the tier
# actually moved. A request that escalates from an already-highest tier has
# nowhere to go, so it records the keyword with escalated=False rather than
# dropping the ask (which reads as an ordinary route) or claiming a bump
# that never happened. Every path reports both the same way.
decision["escalation_keyword"] = escalation_keyword
decision["escalated"] = escalated
if classifier_model is not None:
decision["classifier_model"] = classifier_model
return decision
async def aclassify(
self,
prompt: str,
system_prompt: str | None = None,
request_kwargs: dict[str, Any] | None = None,
) -> tuple[ComplexityTier, float, list[str]]:
) -> ClassificationOutcome:
"""
Classify a prompt by complexity, using the LLM classifier when configured.
Falls back to the local heuristic scorer if classifier_type is "heuristic",
or if the LLM call fails, times out, or returns an unparseable response.
The outcome's `cause` reports which path actually classified the request.
"""
if self.config.classifier_type != "llm" or self.config.classifier_llm_config is None:
return self.classify(prompt, system_prompt)
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
try:
tier = await self._classify_with_llm(prompt, system_prompt, request_kwargs)
return tier, 1.0, [f"llm-classifier:{tier.value}"]
return ClassificationOutcome(
tier=tier, score=None, signals=(f"llm-classifier:{tier.value}",), cause="llm_classifier"
)
except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the heuristic scorer
verbose_router_logger.warning(
f"ComplexityRouter: LLM classifier failed ({e}), falling back to heuristic scoring"
)
return self.classify(prompt, system_prompt)
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
async def _classify_with_llm(
self,
@ -699,16 +809,16 @@ class ComplexityRouter(CustomLogger):
}
return best_model
def _escalation_triggered(self, user_message: str) -> bool:
"""Whether the prompt asks to escalate to a stronger model.
def _matched_escalation_keyword(self, user_message: str) -> str | None:
"""The escalation keyword the prompt contains, or None when escalation is off.
Matching is a case-sensitive substring test so the default "LITELLM ESCALATE"
only fires on the deliberate, shouted form and not on incidental lowercase
mentions of the word (e.g. "how do I escalate this ticket").
"""
if not self.escalation_keywords:
return False
return any(keyword in user_message for keyword in self.escalation_keywords)
return None
return next((keyword for keyword in self.escalation_keywords if keyword in user_message), None)
def _tier_for_model(self, model: str) -> ComplexityTier | None:
"""Return the most-severe configured tier whose pool contains this model."""
@ -746,7 +856,7 @@ class ComplexityRouter(CustomLogger):
return pinned_model
return self.get_model_for_tier(escalated_tier)
def _lexical_tier_override(self, user_message: str) -> ComplexityTier | None:
def _lexical_tier_override(self, user_message: str) -> KeywordOverride | None:
"""When keyword_tier_rules match literally, the most-severe matched tier wins.
Escalating to the highest tier (rather than the first rule in the list) keeps
@ -757,12 +867,15 @@ class ComplexityRouter(CustomLogger):
if not rules:
return None
text = user_message.lower()
matched_tiers = [
rule.tier for rule in rules if any(self._keyword_matches(text, keyword) for keyword in rule.keywords)
matches = [
KeywordOverride(tier=rule.tier, matched_keyword=matched_keyword)
for rule in rules
if (matched_keyword := next((kw for kw in rule.keywords if self._keyword_matches(text, kw)), None))
is not None
]
if not matched_tiers:
if not matches:
return None
return max(matched_tiers, key=TIER_SEVERITY_ORDER.index)
return max(matches, key=lambda match: TIER_SEVERITY_ORDER.index(match.tier))
def _get_or_create_semantic_routelayer(self) -> SemanticRouter:
"""Build (once) a SemanticRouter with one route per tier, utterances = that tier's keywords."""
@ -867,7 +980,7 @@ class ComplexityRouter(CustomLogger):
except ValueError:
return None
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: dict) -> ComplexityTier | None:
async def _resolve_keyword_tier_override(self, user_message: str, request_kwargs: dict) -> KeywordOverride | None:
"""Resolve a keyword_tier_rule override, semantically or lexically per config.
Returns None (no override -> fall through to the scorer) not only when no rule
@ -879,12 +992,17 @@ class ComplexityRouter(CustomLogger):
if not self.config.semantic_keyword_matching:
return self._lexical_tier_override(user_message)
try:
return await self._semantic_tier_override(user_message, request_kwargs)
semantic_tier = await self._semantic_tier_override(user_message, request_kwargs)
except Exception as e: # noqa: BLE001 -- embedding call can fail many ways (timeout, provider/network/parse error); any failure must fall back to scoring, never fail the request
verbose_router_logger.warning(
f"ComplexityRouter: semantic keyword matching failed ({e}), falling back to complexity scoring"
)
return None
if semantic_tier is None:
return None
# A semantic match is a similarity hit against the rule's utterances, not a
# literal keyword, so there is no single matched keyword to report.
return KeywordOverride(tier=semantic_tier, matched_keyword=None)
def _resolve_messages(
self,
@ -1003,6 +1121,7 @@ class ComplexityRouter(CustomLogger):
pinned_model = await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
if isinstance(pinned_model, str):
routed_model: str | None = pinned_model
pin_escalation_keyword: str | None = None
if self.escalation_keywords:
resolved_messages = self._resolve_messages(messages, request_kwargs)
user_message = (
@ -1010,7 +1129,9 @@ class ComplexityRouter(CustomLogger):
if resolved_messages
else None
)
if user_message is not None and self._escalation_triggered(user_message):
if user_message is not None:
pin_escalation_keyword = self._matched_escalation_keyword(user_message)
if pin_escalation_keyword is not None:
routed_model = self._escalated_pin(pinned_model)
if routed_model is not None:
# Refresh the TTL on every hit so an active session doesn't lose its
@ -1028,7 +1149,8 @@ class ComplexityRouter(CustomLogger):
kwargs_metadata = request_kwargs.setdefault("metadata", {})
if isinstance(kwargs_metadata, dict):
kwargs_metadata[ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY] = routed_model
cause = "session_affinity_escalation" if routed_model != pinned_model else "session_affinity_pin"
escalated = routed_model != pinned_model
cause: RoutingDecisionCause = "session_affinity_escalation" if escalated else "session_affinity_pin"
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause={cause}, routed_model={routed_model}"
)
@ -1036,6 +1158,12 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
),
)
response = await self._classify_and_route(
@ -1106,29 +1234,45 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(routed_model=routed_model, cause="default_fallback"),
)
escalate = self._escalation_triggered(user_message)
escalation_keyword = self._matched_escalation_keyword(user_message)
override_tier = await self._resolve_keyword_tier_override(user_message, request_kwargs)
if override_tier is not None:
routed_tier = self._escalate_tier(override_tier) if escalate else override_tier
override = await self._resolve_keyword_tier_override(user_message, request_kwargs)
if override is not None:
routed_tier = self._escalate_tier(override.tier) if escalation_keyword is not None else override.tier
keyword_escalated = routed_tier != override.tier
routed_model = await self._pick_model_for_tier(routed_tier, messages, resolved_messages, request_kwargs)
base_cause = "semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match"
cause = f"{base_cause}+escalation" if escalate else base_cause
keyword_cause: RoutingDecisionCause = (
"semantic_keyword_match" if self.config.semantic_keyword_matching else "literal_keyword_match"
)
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause={cause}, "
f"ComplexityRouter: routing decision cause={keyword_cause}, escalated={keyword_escalated}, "
f"tier={routed_tier.value}, routed_model={routed_model}"
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=keyword_cause,
tier=routed_tier,
matched_keyword=override.matched_keyword,
escalation_keyword=escalation_keyword,
escalated=keyword_escalated,
),
)
tier, score, signals = await self.aclassify(user_message, system_prompt, request_kwargs)
if escalate:
outcome = await self.aclassify(user_message, system_prompt, request_kwargs)
tier, score, signals = outcome.tier, outcome.score, outcome.signals
classified_tier = tier
if escalation_keyword is not None:
tier = self._escalate_tier(tier)
signals = [*signals, "escalation"]
escalated = tier != classified_tier
if escalated:
signals = (*signals, "escalation")
score_repr = f"{score:.3f}" if score is not None else "n/a"
if self.config.adaptive:
routed_model = self._soft_floor_pick(tier, user_message, request_kwargs)
adaptive = self._ensure_adaptive_router()
@ -1138,18 +1282,33 @@ class ComplexityRouter(CustomLogger):
chosen_key = getattr(self, "_adaptive_chosen_model_key", "adaptive_router_chosen_model")
kwargs_metadata[chosen_key] = routed_model
verbose_router_logger.info(
f"ComplexityRouter[adaptive]: routing decision cause=complexity_scorer, "
f"tier={tier.value}, score={score:.3f}, "
f"ComplexityRouter[adaptive]: routing decision cause={outcome.cause}, "
f"tier={tier.value}, score={score_repr}, "
f"signals={signals}, routed_model={routed_model}"
)
else:
routed_model = await self._pick_model_for_tier(tier, messages, resolved_messages, request_kwargs)
verbose_router_logger.info(
f"ComplexityRouter: routing decision cause=complexity_scorer, tier={tier.value}, "
f"score={score:.3f}, signals={signals}, routed_model={routed_model}"
f"ComplexityRouter: routing decision cause={outcome.cause}, tier={tier.value}, "
f"score={score_repr}, signals={signals}, routed_model={routed_model}"
)
classifier_model = (
self.config.classifier_llm_config.model
if outcome.cause == "llm_classifier" and self.config.classifier_llm_config is not None
else None
)
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=outcome.cause,
tier=tier,
score=score,
signals=signals,
escalation_keyword=escalation_keyword,
escalated=escalated,
classifier_model=classifier_model,
),
)

27
litellm/router_strategy/quality_router/quality_router.py
View file

@ -23,6 +23,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.router_strategy.complexity_router.complexity_router import (
ComplexityRouter,
)
from litellm.types.utils import StandardLoggingRoutingDecision
from .config import QualityRouterConfig, RoutingPreferences
@ -357,6 +358,12 @@ class QualityRouter(CustomLogger):
return PreRoutingHookResponse(
model=self.config.default_model,
messages=messages,
routing_decision=StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="quality",
routed_model=self.config.default_model,
cause="default_fallback",
),
)
# Try keyword override first — it short-circuits complexity classification.
@ -380,9 +387,20 @@ class QualityRouter(CustomLogger):
"complexity_tier": None,
},
)
routing_decision = StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="quality",
routed_model=routed_model,
cause="keyword",
matched_keyword=matched_keyword,
)
keyword_quality_tier = self._model_quality.get(routed_model)
if keyword_quality_tier is not None:
routing_decision["tier"] = str(keyword_quality_tier)
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
routing_decision=routing_decision,
)
# No keyword match → complexity classification flow.
@ -419,4 +437,13 @@ class QualityRouter(CustomLogger):
return PreRoutingHookResponse(
model=routed_model,
messages=messages,
routing_decision=StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="quality",
routed_model=routed_model,
cause="quality_tier",
tier=str(int(quality_tier)),
score=score,
signals=list(signals),
),
)

1
litellm/types/guardrails.py
View file

@ -1043,6 +1043,7 @@ class GuardrailEventHooks(str, Enum):
logging_only = "logging_only"
pre_mcp_call = "pre_mcp_call"
during_mcp_call = "during_mcp_call"
post_mcp_call = "post_mcp_call"
realtime_input_transcription = "realtime_input_transcription"

3
litellm/types/router.py
View file

@ -28,7 +28,7 @@ from .completion import CompletionRequest
from .embedding import EmbeddingRequest
from .llms.openai import OpenAIFileObject
from .search import SearchProvider
from .utils import CustomPricingLiteLLMParams, ModelResponse
from .utils import CustomPricingLiteLLMParams, ModelResponse, StandardLoggingRoutingDecision
class ConfigurableClientsideParamsCustomAuth(TypedDict):
@ -839,6 +839,7 @@ class PreRoutingHookResponse(BaseModel):
model: str
messages: Optional[List[Dict[str, Any]]]
routing_decision: StandardLoggingRoutingDecision | None = None
_PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True)

68
litellm/types/utils.py
View file

@ -2675,6 +2675,73 @@ class StandardLoggingPromptManagementMetadata(TypedDict):
prompt_integration: str
class StandardLoggingRoutingDecisionTierBoundaries(TypedDict):
"""Snapshot of the complexity scorer's tier boundaries at decision time, so a
historical spend log row stays explainable after the router config changes."""
simple_medium: float
medium_complex: float
complex_reasoning: float
RoutingDecisionCause = Literal[
"heuristic_scorer",
# The scorer found 2+ reasoning markers and forced REASONING regardless of score.
# A distinct cause rather than a marker inside `signals`, because it is the fact
# that tells a reader the score did NOT choose the tier; encoding it as free text
# meant anything that filtered `signals` silently changed what the row claimed.
"reasoning_override",
"llm_classifier",
"literal_keyword_match",
"semantic_keyword_match",
"session_affinity_pin",
"session_affinity_escalation",
"default_fallback",
"keyword",
"quality_tier",
"bandit",
]
class StandardLoggingRoutingDecision(TypedDict, total=False):
"""Per-request provenance for a pre-routing strategy (auto-router) decision."""
router_model_name: str
router_type: Literal["complexity", "adaptive", "quality"]
routed_model: str
cause: RoutingDecisionCause
tier: str
request_type: str
score: float
signals: Sequence[str]
matched_keyword: str
escalation_keyword: str
classifier_model: str
escalated: bool
tier_boundaries: StandardLoggingRoutingDecisionTierBoundaries
# Fields whose values quote the caller's prompt. Dropped when an operator turns message
# logging off. Every other field aggregates the prompt without reproducing it and is kept,
# so a redacted row stays explainable. `test_every_routing_decision_field_is_classified`
# fails if a field is added to the record without being placed in one set or the other.
PROMPT_QUOTING_ROUTING_DECISION_FIELDS: FrozenSet[str] = frozenset({"signals", "matched_keyword", "escalation_keyword"})
DERIVED_ROUTING_DECISION_FIELDS: FrozenSet[str] = frozenset(
{
"router_model_name",
"router_type",
"routed_model",
"cause",
"tier",
"request_type",
"score",
"classifier_model",
"escalated",
"tier_boundaries",
}
)
class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
"""
Specific metadata k,v pairs logged to integration for easier cost tracking and prompt management
@ -2688,6 +2755,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata):
prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata]
mcp_tool_call_metadata: Optional[StandardLoggingMCPToolCall]
vector_store_request_metadata: Optional[List[StandardLoggingVectorStoreRequest]]
routing_decision: StandardLoggingRoutingDecision | None
applied_guardrails: Optional[List[str]]
usage_object: Optional[dict]
cold_storage_object_key: Optional[str] # S3/GCS object key for cold storage retrieval

88
model_prices_and_context_window.json
View file

@ -3454,22 +3454,16 @@
},
"azure_ai/gpt-5.4-mini": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.5e-07,
"cache_read_input_token_cost_priority": 1.5e-07,
"cache_read_input_token_cost_above_272k_tokens_priority": 3e-07,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_above_272k_tokens": 1.5e-06,
"input_cost_per_token_priority": 1.5e-06,
"input_cost_per_token_above_272k_tokens_priority": 3e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"output_cost_per_token_above_272k_tokens": 6.75e-06,
"output_cost_per_token_priority": 9e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.35e-05,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-mini",
"supported_endpoints": [
"/v1/chat/completions",
@ -3500,22 +3494,16 @@
},
"azure_ai/gpt-5.4-mini-2026-03-17": {
"cache_read_input_token_cost": 7.5e-08,
"cache_read_input_token_cost_above_272k_tokens": 1.5e-07,
"cache_read_input_token_cost_priority": 1.5e-07,
"cache_read_input_token_cost_above_272k_tokens_priority": 3e-07,
"input_cost_per_token": 7.5e-07,
"input_cost_per_token_above_272k_tokens": 1.5e-06,
"input_cost_per_token_priority": 1.5e-06,
"input_cost_per_token_above_272k_tokens_priority": 3e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 4.5e-06,
"output_cost_per_token_above_272k_tokens": 6.75e-06,
"output_cost_per_token_priority": 9e-06,
"output_cost_per_token_above_272k_tokens_priority": 1.35e-05,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-mini",
"supported_endpoints": [
"/v1/chat/completions",
@ -3546,22 +3534,16 @@
},
"azure_ai/gpt-5.4-nano": {
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_priority": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_priority": 4e-07,
"input_cost_per_token_above_272k_tokens_priority": 8e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"output_cost_per_token_above_272k_tokens": 1.875e-06,
"output_cost_per_token_priority": 2.5e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.75e-06,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-nano",
"supported_endpoints": [
"/v1/chat/completions",
@ -3592,22 +3574,16 @@
},
"azure_ai/gpt-5.4-nano-2026-03-17": {
"cache_read_input_token_cost": 2e-08,
"cache_read_input_token_cost_above_272k_tokens": 4e-08,
"cache_read_input_token_cost_priority": 4e-08,
"cache_read_input_token_cost_above_272k_tokens_priority": 8e-08,
"input_cost_per_token": 2e-07,
"input_cost_per_token_above_272k_tokens": 4e-07,
"input_cost_per_token_priority": 4e-07,
"input_cost_per_token_above_272k_tokens_priority": 8e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 400000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.25e-06,
"output_cost_per_token_above_272k_tokens": 1.875e-06,
"output_cost_per_token_priority": 2.5e-06,
"output_cost_per_token_above_272k_tokens_priority": 3.75e-06,
"source": "https://ai.azure.com/catalog/models/gpt-5.4-nano",
"supported_endpoints": [
"/v1/chat/completions",
@ -7201,7 +7177,7 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7236,7 +7212,7 @@
"cache_read_input_token_cost": 7.5e-08,
"input_cost_per_token": 7.5e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7271,7 +7247,7 @@
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -7306,7 +7282,7 @@
"cache_read_input_token_cost": 2e-08,
"input_cost_per_token": 2e-07,
"litellm_provider": "azure",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -16703,8 +16679,8 @@
"input_cost_per_token": 6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://fireworks.ai/pricing",
@ -16717,8 +16693,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -16733,8 +16709,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17077,8 +17053,8 @@
"input_cost_per_token": 6e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://fireworks.ai/pricing",
@ -17091,8 +17067,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17107,8 +17083,8 @@
"input_cost_per_token": 2e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17123,8 +17099,8 @@
"input_cost_per_token": 9.5e-07,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -17139,8 +17115,8 @@
"input_cost_per_token": 1.9e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -24364,7 +24340,7 @@
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_priority": 1.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -24410,7 +24386,7 @@
"input_cost_per_token_batches": 3.75e-07,
"input_cost_per_token_priority": 1.5e-06,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -24454,7 +24430,7 @@
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_batches": 1e-07,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -24497,7 +24473,7 @@
"input_cost_per_token_flex": 1e-07,
"input_cost_per_token_batches": 1e-07,
"litellm_provider": "openai",
"max_input_tokens": 1050000,
"max_input_tokens": 272000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
@ -42622,8 +42598,8 @@
"input_cost_per_token": 2e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",
@ -42638,8 +42614,8 @@
"input_cost_per_token": 1.9e-06,
"litellm_provider": "fireworks_ai",
"max_input_tokens": 262144,
"max_output_tokens": 262144,
"max_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 8e-06,
"source": "https://docs.fireworks.ai/serverless/pricing",

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.95.0"
version = "1.96.0"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.15"
@ -302,7 +302,7 @@ members = ["enterprise", "litellm-proxy-extras"]
profile = "black"
[tool.commitizen]
version = "1.95.0"
version = "1.96.0"
version_files = [
"pyproject.toml:^version",
]

40
ruff-strict-budget.json
View file

@ -1,6 +1,6 @@
{
"ANN001": {
"limit": 3094
"limit": 3118
},
"ANN002": {
"limit": 69
@ -24,7 +24,7 @@
"limit": 130
},
"ANN401": {
"limit": 2012
"limit": 2017
},
"ASYNC230": {
"limit": 14
@ -33,7 +33,7 @@
"limit": 4
},
"B006": {
"limit": 186
"limit": 188
},
"B008": {
"limit": 505
@ -42,7 +42,7 @@
"limit": 84
},
"B010": {
"limit": 191
"limit": 194
},
"B018": {
"limit": 5
@ -60,7 +60,7 @@
"limit": 4
},
"BLE001": {
"limit": 2895
"limit": 2899
},
"C401": {
"limit": 11
@ -81,7 +81,7 @@
"limit": 4
},
"C901": {
"limit": 312
"limit": 314
},
"D419": {
"limit": 9
@ -123,7 +123,7 @@
"limit": 52
},
"I001": {
"limit": 269
"limit": 271
},
"LOG015": {
"limit": 8
@ -180,7 +180,7 @@
"limit": 34
},
"PLR1714": {
"limit": 257
"limit": 261
},
"PLR1730": {
"limit": 10
@ -189,7 +189,7 @@
"limit": 4
},
"PLW0127": {
"limit": 42
"limit": 43
},
"PLW0133": {
"limit": 4
@ -222,7 +222,7 @@
"limit": 38
},
"RET504": {
"limit": 714
"limit": 716
},
"RUF010": {
"limit": 874
@ -237,7 +237,7 @@
"limit": 41
},
"RUF022": {
"limit": 84
"limit": 85
},
"RUF023": {
"limit": 5
@ -261,7 +261,7 @@
"limit": 24
},
"SIM101": {
"limit": 59
"limit": 61
},
"SIM102": {
"limit": 324
@ -273,7 +273,7 @@
"limit": 6
},
"SIM114": {
"limit": 109
"limit": 111
},
"SIM115": {
"limit": 5
@ -288,7 +288,7 @@
"limit": 4
},
"SIM210": {
"limit": 10
"limit": 11
},
"SIM211": {
"limit": 4
@ -306,10 +306,10 @@
"limit": 9
},
"TID251": {
"limit": 2651
"limit": 2653
},
"TRY002": {
"limit": 546
"limit": 547
},
"TRY004": {
"limit": 98
@ -324,7 +324,7 @@
"limit": 883
},
"UP006": {
"limit": 12145
"limit": 12168
},
"UP007": {
"limit": 2526
@ -348,7 +348,7 @@
"limit": 5
},
"UP032": {
"limit": 625
"limit": 626
},
"UP034": {
"limit": 4
@ -360,9 +360,9 @@
"limit": 4
},
"UP037": {
"limit": 103
"limit": 105
},
"UP045": {
"limit": 17777
"limit": 17805
}
}

27
tests/batches_tests/test_batch_custom_pricing.py
View file

@ -12,8 +12,7 @@ import litellm
import pytest
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_cost_from_file_content,
_aggregate_batch_cost_usage_models,
calculate_batch_cost_and_usage,
)
from litellm.cost_calculator import batch_cost_calculator
@ -113,28 +112,12 @@ def test_batch_cost_calculator_uses_custom_model_info():
), f"Expected completion cost {expected_completion}, got {completion_cost}"
def test_get_batch_job_cost_from_file_content_uses_custom_model_info():
"""_get_batch_job_cost_from_file_content should thread model_info to completion_cost."""
def test_aggregate_batch_cost_uses_custom_model_info():
"""_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)
expected = (10 * 0.00125) + (5 * 0.005)
assert cost == pytest.approx(
expected
), f"Expected total cost {expected}, got {cost}"
def test_batch_cost_calculator_func_uses_custom_model_info():
"""_batch_cost_calculator should thread model_info."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
cost = _batch_cost_calculator(
file_content_dictionary=file_content,
cost, _, _ = _aggregate_batch_cost_usage_models(
entries=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)

6
tests/batches_tests/test_batch_rate_limits.py
View file

@ -913,7 +913,7 @@ async def test_batch_logging_azure_credentials_regression():
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.batches.batch_utils import (
_extract_file_access_credentials,
_get_batch_output_file_content_as_dictionary,
_fetch_batch_output_file_content,
_handle_completed_batch,
)
from litellm.types.llms.openai import Batch, HttpxBinaryResponseContent
@ -996,7 +996,7 @@ async def test_batch_logging_azure_credentials_regression():
with patch(
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
result = await _get_batch_output_file_content_as_dictionary(
result = await _fetch_batch_output_file_content(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,
@ -1092,7 +1092,7 @@ async def test_batch_logging_azure_credentials_regression():
)
# Call without litellm_params (should still work for OpenAI)
result = await _get_batch_output_file_content_as_dictionary(
result = await _fetch_batch_output_file_content(
batch=mock_batch,
custom_llm_provider="openai",
litellm_params=None,

22
tests/batches_tests/test_batches_logging_unit_tests.py
View file

@ -19,10 +19,8 @@ import litellm
from litellm import create_batch, create_file
from litellm._logging import verbose_logger
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_aggregate_batch_cost_usage_models,
_get_file_content_as_dictionary,
_get_batch_job_cost_from_file_content,
_get_batch_job_total_usage_from_file_content,
_get_batch_job_usage_from_response_body,
_get_response_from_batch_job_output_file,
_batch_response_was_successful,
@ -139,9 +137,10 @@ def test_get_file_content_as_dictionary(sample_file_content):
def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
usage = _get_batch_job_total_usage_from_file_content(
sample_file_content_dict, custom_llm_provider="openai"
)
with patch("litellm.completion_cost", return_value=0.0):
_, usage, _ = _aggregate_batch_cost_usage_models(
entries=sample_file_content_dict, custom_llm_provider="openai"
)
assert usage.total_tokens == 62 # 30 + 32
assert usage.prompt_tokens == 42 # 20 + 22
assert usage.completion_tokens == 20 # 10 + 10
@ -157,8 +156,8 @@ async def test_batch_cost_calculator(sample_file_content_dict):
so we expect the cost to be 0.5 * 2 = 1.0
"""
with patch("litellm.completion_cost", return_value=0.5):
cost = _batch_cost_calculator(
file_content_dictionary=sample_file_content_dict,
cost, _, _ = _aggregate_batch_cost_usage_models(
entries=sample_file_content_dict,
custom_llm_provider="openai",
)
assert cost == 1.0 # 0.5 * 2 successful responses
@ -278,9 +277,12 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
created_at=1234567890,
)
sample_file_content_bytes = "\n".join(
json.dumps(row) for row in sample_file_content_dict
).encode()
with patch(
"litellm.batches.batch_utils._get_batch_output_file_content_as_dictionary",
new=AsyncMock(return_value=sample_file_content_dict),
"litellm.batches.batch_utils._fetch_batch_output_file_content",
new=AsyncMock(return_value=sample_file_content_bytes),
):
cost, usage, models = await _handle_completed_batch(
batch=batch, custom_llm_provider="openai"

3
tests/code_coverage_tests/recursive_detector.py
View file

@ -57,6 +57,9 @@ IGNORE_FUNCTIONS = [
"_filter_mcp_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the MCP call at the cap.
"_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap.
"_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap.
"json_string_leaves", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned.
"with_json_string_leaves", # transitively bounded: only runs on a tree json_string_leaves already walked under the cap.
"json_unrewritable_labels", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); returns the None sentinel at the cap so the caller blocks.
]

2
tests/e2e/coverage_registry/quota_management.yaml
View file

@ -23,7 +23,7 @@
- {id: quota_management.budget.key.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: key, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes key spend after the window; a blocked key serves again"}
- {id: quota_management.budget.team.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_duration zeroes a team's spend after the window; every key on the team serves again"}
- {id: quota_management.budget.organization.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: organization, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An org budget resets after its window; keys under the org serve again"}
- {id: quota_management.budget.internal_user.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An internal user's budget resets after its window; their personal and team-member keys serve again"}
- {id: quota_management.budget.internal_user.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: internal_user, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "An internal user's budget resets after its window; their personal keys serve again"}
- {id: quota_management.budget.team_member.resets_after_window, module: quota_management, tier: P1, behavior: budget, variant: team_member, assertions: [resets_after_window], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Member per-team budget reset keeps advancing window after window"}
- {id: quota_management.budget.key_multi_window.blocks_then_resets, module: quota_management, tier: P1, behavior: budget, variant: key_multi_window, assertions: [blocks_then_resets], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "budget_limits enforce within a short window and serve again in the next"}
- {id: quota_management.budget.key_multi_window.resets_windows_independently, module: quota_management, tier: P2, behavior: budget, variant: key_multi_window, assertions: [resets_windows_independently], exercised_on: [chat_completions], source: "proxy/common_utils/reset_budget_job.py", rationale: "Each window of a multi-window budget resets on its own schedule"}

24
tests/e2e/quota_management/budgets/test_budget_enforcement_e2e.py
View file

@ -79,9 +79,14 @@ class TestBudgetBlocksPerLevel:
)
@pytest.mark.covers("quota_management.budget.internal_user.blocks_over_limit")
def test_user_budget_enforced_across_all_their_keys(
def test_user_budget_enforced_across_their_personal_keys(
self, client: BudgetClient, resources: ResourceManager
) -> None:
"""A user's max_budget follows the person across their personal keys, so a
second untouched key is not a fresh allowance. It stops at the team
boundary: the same user's team-scoped key is governed by the team and
team-member budgets, both uncapped here, so it is the control that must
keep serving while the personal keys are refused."""
user_id = client.create_user(max_budget=TINY_CAP)
resources.defer(lambda: client.delete_user(user_id))
first_key = client.generate_key(user_id=user_id)
@ -95,12 +100,17 @@ class TestBudgetBlocksPerLevel:
resources.defer(lambda: client.delete_key(team_key))
_assert_blocked_429(client, first_key)
for label, key in (("second personal key", second_key), ("team-member key", team_key)):
result = _chat(client, key)
assert is_budget_block(result) and result.status_code == 429, (
f"the {label} of a user over budget must get the same 429 budget_exceeded, "
f"got {result.status_code}: {result.body[:200]}"
)
second = _chat(client, second_key)
assert is_budget_block(second) and second.status_code == 429, (
f"the second personal key of a user over budget must get the same 429 budget_exceeded, "
f"got {second.status_code}: {second.body[:200]}"
)
team_result = _chat(client, team_key)
assert not is_budget_block(team_result), (
f"the team-scoped key of a user over their personal budget must keep serving; "
f"got {team_result.status_code}: {team_result.body[:200]}"
)
require_successful_call(team_result)
@pytest.mark.covers("quota_management.budget.end_user.blocks_over_limit")
def test_end_user_budget_blocks_attributed_calls(

15
tests/e2e/quota_management/budgets/test_budget_reset_e2e.py
View file

@ -102,21 +102,6 @@ class TestBudgetResetPerLevel:
_drive_to_block(client, key)
_poll_until_serves_again(client, key)
@pytest.mark.covers("quota_management.budget.internal_user.resets_after_window")
def test_team_member_key_user_budget_resets_after_window(
self, client: BudgetClient, resources: ResourceManager
) -> None:
user_id = client.create_user(max_budget=TINY_CAP, budget_duration=WINDOW)
resources.defer(lambda: client.delete_user(user_id))
team_id = client.create_team(alias=f"e2e-user-team-reset-{unique_marker()}")
resources.defer(lambda: client.delete_team(team_id))
client.add_team_member(team_id, user_id, max_budget_in_team=100.0)
key = client.generate_key(team_id=team_id, user_id=user_id)
resources.defer(lambda: client.delete_key(key))
_drive_to_block(client, key)
_poll_until_serves_again(client, key)
class TestKeyBudgetResetAcrossKeyKinds:
"""The tiny max_budget and its 30s window sit on the key itself while the user,

159
tests/proxy_migration_tests/test_replica_identity_full.py Normal file
View file

@ -0,0 +1,159 @@
"""Coverage for the opt-in REPLICA IDENTITY FULL post-migration step.
The DB-backed tests run against the same Postgres the migration suite uses, in
a throwaway schema so they cannot disturb the migrated tables.
"""
import os
import uuid
import pytest
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.utils import ProxyExtrasDBManager
psycopg = pytest.importorskip("psycopg")
requires_db = pytest.mark.skipif(
"DATABASE_URL" not in os.environ,
reason="requires a postgres database (DATABASE_URL)",
)
def _base_url() -> str:
return os.environ["DATABASE_URL"].split("?")[0]
def _replica_identities(schema: str) -> dict:
with psycopg.connect(_base_url(), autocommit=True) as conn:
rows = conn.execute(
"SELECT c.relname, c.relreplident FROM pg_class c "
"JOIN pg_namespace n ON n.oid = c.relnamespace "
"WHERE n.nspname = %s AND c.relkind = 'r'",
(schema,),
).fetchall()
return dict(rows)
@pytest.fixture
def scratch_schema(monkeypatch):
"""A schema holding two LiteLLM tables and one foreign table, all at the default."""
schema = f"replica_identity_{uuid.uuid4().hex[:8]}"
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'CREATE SCHEMA "{schema}"')
conn.execute(
f'CREATE TABLE "{schema}"."LiteLLM_ScratchTable" (id TEXT PRIMARY KEY, note TEXT)'
)
conn.execute(f'CREATE TABLE "{schema}"."LiteLLM_ScratchSibling" (id TEXT PRIMARY KEY)')
conn.execute(f'CREATE TABLE "{schema}"."ScratchForeignTable" (id TEXT PRIMARY KEY)')
monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}")
yield schema
with psycopg.connect(_base_url(), autocommit=True) as conn:
conn.execute(f'DROP SCHEMA "{schema}" CASCADE')
@requires_db
def test_applies_full_to_litellm_tables_only(scratch_schema, monkeypatch):
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True
identities = _replica_identities(scratch_schema)
assert identities["LiteLLM_ScratchTable"] == "f"
assert identities["LiteLLM_ScratchSibling"] == "f"
assert identities["ScratchForeignTable"] == "d"
@requires_db
def test_a_locked_table_does_not_block_the_others(scratch_schema, monkeypatch):
"""ALTER TABLE needs an exclusive lock, so a table busy with a long read has
to be skipped for the next run instead of stalling every other table behind it."""
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
with psycopg.connect(_base_url()) as holder:
holder.execute(f'SELECT * FROM "{scratch_schema}"."LiteLLM_ScratchTable"')
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True
identities = _replica_identities(scratch_schema)
assert identities["LiteLLM_ScratchTable"] == "d"
assert identities["LiteLLM_ScratchSibling"] == "f"
@requires_db
def test_leaves_tables_alone_when_not_requested(scratch_schema, monkeypatch):
monkeypatch.delenv(REPLICA_IDENTITY_FULL_ENV_VAR, raising=False)
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False
assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "d"
@requires_db
def test_is_idempotent_across_runs(scratch_schema, monkeypatch):
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is True
assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "f"
@requires_db
def test_reports_failure_without_raising(scratch_schema, monkeypatch):
"""A run that cannot execute the statement must not take the migration down."""
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
monkeypatch.setattr(
ProxyExtrasDBManager,
"_get_prisma_dir",
staticmethod(lambda: "/nonexistent/prisma/dir"),
)
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False
assert _replica_identities(scratch_schema)["LiteLLM_ScratchTable"] == "d"
def test_reports_an_unrunnable_prisma_cli_without_raising(tmp_path):
"""A deployment without the Prisma CLI on PATH must still finish its
migration run instead of dying on the optional replication step."""
assert (
apply_replica_identity_full(
schema_path=str(tmp_path / "schema.prisma"),
prisma_command=str(tmp_path / "no-such-prisma"),
prisma_env={},
)
is False
)
def test_setup_database_applies_after_a_successful_migration_run(monkeypatch):
applied = []
monkeypatch.setattr(
ProxyExtrasDBManager, "_run_migrations", staticmethod(lambda **kwargs: True)
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"apply_replica_identity_full_if_requested",
staticmethod(lambda: applied.append(True)),
)
assert ProxyExtrasDBManager.setup_database(use_migrate=True) is True
assert applied == [True]
def test_setup_database_skips_replica_identity_when_migrations_fail(monkeypatch):
applied = []
monkeypatch.setattr(
ProxyExtrasDBManager, "_run_migrations", staticmethod(lambda **kwargs: False)
)
monkeypatch.setattr(
ProxyExtrasDBManager,
"apply_replica_identity_full_if_requested",
staticmethod(lambda: applied.append(True)),
)
assert ProxyExtrasDBManager.setup_database(use_migrate=True) is False
assert applied == []

11
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -219,8 +219,8 @@ async def test_aaauser_personal_budgets(key_ownership):
"""
Set a personal budget on a user
User budget is enforced regardless of key ownership (personal or team).
Both cases should raise BudgetExceededError when the user is over budget.
- have it only apply when key belongs to user -> raises BudgetExceededError
- if key belongs to team, have key respect team budget -> allows call to go through
"""
import asyncio
import time
@ -278,9 +278,12 @@ async def test_aaauser_personal_budgets(key_ownership):
== valid_token
)
with pytest.raises(ProxyException) as exc_info:
if key_ownership == "user_key":
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
else:
await user_api_key_auth(request=request, api_key="Bearer " + user_key)
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
@pytest.mark.asyncio

29
tests/test_litellm/a2a_protocol/test_main.py
View file

@ -104,3 +104,32 @@ async def test_streaming_trace_id_prefers_logging_trace_id():
pass
assert captured["extra_headers"]["X-LiteLLM-Trace-Id"] == "trace-from-logging"
def test_streaming_logging_obj_carries_call_type_into_model_call_details():
"""The streaming logging object is built by hand rather than through
``update_environment_variables``, which is the only place ``call_type`` normally
reaches ``model_call_details``. Callbacks read the call type from there, so
without this the streamed turn arrives at every logger with no call type at all
and OTel's GenAI metrics label it ``chat`` instead of ``invoke_agent``."""
from a2a.compat.v0_3.types import MessageSendParams, SendStreamingMessageRequest
from litellm.a2a_protocol.main import _build_streaming_logging_obj
request = SendStreamingMessageRequest(
id="rpc-call-type",
params=MessageSendParams(
message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": "hi"}]}
),
)
logging_obj = _build_streaming_logging_obj(
request=request,
agent_name="some-agent",
agent_id=None,
litellm_params=None,
metadata=None,
proxy_server_request=None,
)
assert logging_obj.model_call_details["call_type"] == "asend_message_streaming"

433
tests/test_litellm/batches/test_batch_utils.py
View file

@ -14,10 +14,13 @@ maps (litellm.completion_cost, batch_cost_calculator), the tokenizer
deterministic stand-ins so the arithmetic under test is the only variable.
"""
import json
import os
import sys
import httpx
import pytest
import respx
sys.path.insert(0, os.path.abspath("../../../.."))
@ -200,29 +203,34 @@ def test_estimate_tokens_never_zero_for_short_rows():
# =========================================================================== #
# _get_batch_models_from_file_content (output file)
# _aggregate_batch_cost_usage_models: models (output file)
# =========================================================================== #
def test_output_models_uses_model_name_override():
# model_name short-circuits: content is ignored entirely.
assert bu._get_batch_models_from_file_content([_success_row(model="ignored")], model_name="forced-model") == [
"forced-model"
]
def test_output_models_uses_model_name_override(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
_, _, models = bu._aggregate_batch_cost_usage_models(
entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
)
assert models == ["forced-model"]
def test_output_models_collects_from_successful_only():
def test_output_models_collects_from_successful_only(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [
_success_row(model="gpt-4o"),
_failed_row(model="should-be-skipped"),
_success_row(model="claude-3"),
]
assert bu._get_batch_models_from_file_content(rows) == ["gpt-4o", "claude-3"]
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert models == ["gpt-4o", "claude-3"]
def test_output_models_skips_successful_without_model():
def test_output_models_skips_successful_without_model(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [{"response": {"status_code": 200, "body": {}}}]
assert bu._get_batch_models_from_file_content(rows) == []
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert models == []
# =========================================================================== #
@ -235,6 +243,8 @@ def test_extract_credentials_only_known_keys():
"api_key": "sk-1",
"api_base": "https://b",
"vertex_project": "proj",
"gcs_bucket_name": "my-bucket",
"bucket_name": "my-alias-bucket",
"model": "gpt-4o", # not a credential key
"unrelated": "x",
}
@ -242,6 +252,8 @@ def test_extract_credentials_only_known_keys():
"api_key": "sk-1",
"api_base": "https://b",
"vertex_project": "proj",
"gcs_bucket_name": "my-bucket",
"bucket_name": "my-alias-bucket",
}
@ -261,6 +273,8 @@ def test_extract_credentials_all_supported_keys():
"vertex_project",
"vertex_location",
"vertex_credentials",
"gcs_bucket_name",
"bucket_name",
"timeout",
"max_retries",
}
@ -372,17 +386,18 @@ def test_count_entry_uses_model_name_fallback(monkeypatch):
# =========================================================================== #
# _get_batch_job_total_usage_from_file_content (output usage aggregation)
# _aggregate_batch_cost_usage_models: usage (output usage aggregation)
# =========================================================================== #
def test_total_usage_sums_successful_only():
def test_total_usage_sums_successful_only(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [
_success_row(usage=_usage(10, 5)), # 15
_failed_row(), # excluded
_success_row(usage=_usage(20, 10)), # 30
]
usage = bu._get_batch_job_total_usage_from_file_content(rows)
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
30,
15,
@ -391,7 +406,9 @@ def test_total_usage_sums_successful_only():
def test_total_usage_empty_is_zero():
usage = bu._get_batch_job_total_usage_from_file_content([])
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
assert cost == 0.0
assert models == []
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
0,
0,
@ -400,7 +417,7 @@ def test_total_usage_empty_is_zero():
# =========================================================================== #
# _get_batch_job_cost_from_file_content (cost maps mocked)
# _aggregate_batch_cost_usage_models: cost (cost maps mocked)
# =========================================================================== #
@ -419,7 +436,7 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
_success_row(usage=_usage(20, 10)),
]
total = bu._get_batch_job_cost_from_file_content(rows, custom_llm_provider="openai")
total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert total == 1.0 # 2 successful * 0.5
assert len(calls) == 2 # failed row not costed
@ -435,8 +452,8 @@ def test_cost_from_content_model_info_path(monkeypatch):
_success_row(usage=_usage(20, 10)),
]
total = bu._get_batch_job_cost_from_file_content(
rows,
total, _, _ = bu._aggregate_batch_cost_usage_models(
entries=rows,
custom_llm_provider="openai",
model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path
)
@ -444,32 +461,65 @@ def test_cost_from_content_model_info_path(monkeypatch):
assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2)
def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
"""A one-shot generator: any implementation that iterates the entries twice
(e.g. separate cost and usage passes) sees nothing on the second pass and
returns wrong totals for at least one of cost/usage/models."""
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5)
one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
assert cost == 1.0
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
assert models == ["gpt-4o", "gpt-4o"]
# =========================================================================== #
# _batch_cost_calculator (dispatch: vertex-disable-transform vs generic)
# calculate_batch_cost_and_usage (dispatch: vertex-disable-transform vs generic)
# =========================================================================== #
def test_batch_cost_calculator_generic_path(monkeypatch):
monkeypatch.setattr(bu, "_get_batch_job_cost_from_file_content", lambda **kw: 4.2)
assert bu._batch_cost_calculator([], custom_llm_provider="openai", model_name="gpt-4o") == 4.2
def test_batch_cost_calculator_vertex_disable_transform_path(monkeypatch):
@pytest.mark.asyncio
async def test_calculate_vertex_disable_transform_path(monkeypatch):
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model: (9.9, Usage()),
lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)),
)
# generic path must NOT be taken
monkeypatch.setattr(
bu,
"_get_batch_job_cost_from_file_content",
"_aggregate_batch_cost_usage_models",
lambda **kw: pytest.fail("generic path should not run"),
)
cost = bu._batch_cost_calculator([], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001")
cost, usage, models = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001"
)
assert cost == 9.9
assert usage.total_tokens == 3
assert models == ["gemini-2.0-flash-001"]
@pytest.mark.asyncio
async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
"""Without a model_name the raw-vertex path cannot price lines; the generic
aggregation path must run even with the disable flag set."""
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model: pytest.fail("raw vertex path should not run"),
)
cost, usage, models = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=[], custom_llm_provider="vertex_ai"
)
assert cost == 0.0
assert usage.total_tokens == 0
assert models == []
# =========================================================================== #
@ -579,24 +629,19 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
@pytest.mark.asyncio
async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 2.5)
monkeypatch.setattr(
bu,
"_get_batch_job_total_usage_from_file_content",
lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
cost, usage, models = await bu.calculate_batch_cost_and_usage(
file_content_dictionary=rows, custom_llm_provider="openai"
)
assert cost == 2.5
assert usage.total_tokens == 15
assert models == ["gpt-4o"] # real _get_batch_models_from_file_content
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert models == ["gpt-4o"]
# =========================================================================== #
# _get_batch_output_file_content_as_dictionary (file fetch + credential merge)
# _fetch_batch_output_file_content (file fetch + credential merge)
# =========================================================================== #
@ -615,16 +660,217 @@ def _batch(output_file_id):
)
def _vertex_openai_row(custom_id, model, prompt_tokens, completion_tokens):
return {
"id": f"batch_req_{custom_id}",
"custom_id": custom_id,
"response": {
"status_code": 200,
"request_id": custom_id,
"body": {
"id": f"chatcmpl-{custom_id}",
"object": "chat.completion",
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "ok"},
"finish_reason": "stop",
}
],
"usage": _usage(prompt_tokens, completion_tokens),
},
},
"error": None,
}
def _vertex_jsonl(rows):
return "\n".join(json.dumps(row) for row in rows).encode()
@pytest.mark.asyncio
async def test_output_file_content_vertex_raises():
with pytest.raises(ValueError, match="Vertex AI does not support"):
await bu._get_batch_output_file_content_as_dictionary(_batch("of"), custom_llm_provider="vertex_ai")
async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch):
import litellm.files.main as files_main
rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)]
captured: dict = {}
async def fake_afile_content(**kw):
captured.update(kw)
return type("R", (), {"content": _vertex_jsonl(rows)})()
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
result = await bu._fetch_batch_output_file_content(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={
"vertex_project": "proj-1",
"vertex_location": "us-central1",
"vertex_credentials": "/path/to/creds.json",
"gcs_bucket_name": "litellm-bucket",
"model": "vertex_ai/gemini-3.6-flash",
},
)
assert bu._get_file_content_as_dictionary(result) == rows
assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
assert captured["custom_llm_provider"] == "vertex_ai"
assert captured["vertex_project"] == "proj-1"
assert captured["vertex_location"] == "us-central1"
assert captured["vertex_credentials"] == "/path/to/creds.json"
assert captured["gcs_bucket_name"] == "litellm-bucket"
assert "model" not in captured
@pytest.mark.asyncio
async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monkeypatch):
import base64
import litellm.files.main as files_main
captured: dict = {}
async def fake_afile_content(**kw):
captured.update(kw)
return type("R", (), {"content": b'{"a": 1}'})()
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
unified_id = (
"litellm_proxy:application/jsonl;unified_id,uuid-1;target_model_names,vertex-model;"
"llm_output_file_id,gs://litellm-bucket/output/predictions.jsonl;llm_output_file_model_id,model-1"
)
encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="vertex_ai")
assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
assert captured["custom_llm_provider"] == "vertex_ai"
def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens):
return {
"request": {
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
"labels": {"litellm_custom_id": custom_id},
},
"status": "",
"response": {
"candidates": [
{
"content": {"role": "model", "parts": [{"text": "ok"}]},
"finishReason": "STOP",
}
],
"usageMetadata": {
"promptTokenCount": prompt_tokens,
"candidatesTokenCount": completion_tokens,
"totalTokenCount": prompt_tokens + completion_tokens,
},
"modelVersion": "gemini-3.6-flash",
},
"processed_time": "2026-07-30T00:00:00.000000+00:00",
}
@pytest.fixture
def respx_interceptable_httpx_client(monkeypatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.in_memory_llm_clients_cache.flush_cache()
yield
litellm.in_memory_llm_clients_cache.flush_cache()
@pytest.mark.asyncio
@respx.mock
async def test_output_file_content_vertex_managed_uri_accepted_by_real_validation(respx_interceptable_httpx_client):
managed_output_uri = (
"gs://litellm-bucket/litellm-vertex-files/publishers/google/models/"
"gemini-3.6-flash/abc-123/prediction-model/predictions.jsonl"
)
rows = [
_vertex_predictions_row("request-1", 10, 5),
_vertex_predictions_row("request-2", 20, 10),
]
route = respx.get(url__regex=r"https://storage\.googleapis\.com/storage/v1/b/litellm-bucket/o/.*").mock(
return_value=httpx.Response(200, content=_vertex_jsonl(rows))
)
file_content = await bu._fetch_batch_output_file_content(
_batch(managed_output_uri),
custom_llm_provider="vertex_ai",
litellm_params={
"api_key": "test-token",
"vertex_project": "proj-1",
"vertex_location": "us-central1",
"gcs_bucket_name": "litellm-bucket",
},
)
result = bu._get_file_content_as_dictionary(file_content)
assert route.call_count == 1
request = route.calls.last.request
assert request.url.raw_path == (
b"/storage/v1/b/litellm-bucket/o/"
b"litellm-vertex-files%2Fpublishers%2Fgoogle%2Fmodels%2Fgemini-3.6-flash"
b"%2Fabc-123%2Fprediction-model%2Fpredictions.jsonl?alt=media"
)
assert [row["custom_id"] for row in result] == ["request-1", "request-2"]
assert all(row["response"]["status_code"] == 200 for row in result)
assert all(row["response"]["body"]["model"] == "gemini-3.6-flash" for row in result)
assert [row["response"]["body"]["usage"]["prompt_tokens"] for row in result] == [10, 20]
assert [row["response"]["body"]["usage"]["completion_tokens"] for row in result] == [5, 10]
@pytest.mark.asyncio
@respx.mock
async def test_output_file_content_vertex_foreign_bucket_rejected_by_real_validation():
with pytest.raises(Exception, match="does not match the configured storage bucket"):
await bu._fetch_batch_output_file_content(
_batch("gs://attacker-bucket/litellm-vertex-files/x/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={
"api_key": "test-token",
"vertex_project": "proj-1",
"vertex_location": "us-central1",
"gcs_bucket_name": "litellm-bucket",
},
)
assert respx.mock.calls.call_count == 0
@pytest.mark.asyncio
async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monkeypatch):
import litellm.files.main as files_main
rows = [
_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5),
_vertex_openai_row("request-2", "gemini-3.6-flash", 20, 10),
]
async def fake_afile_content(**kw):
return type("R", (), {"content": _vertex_jsonl(rows)})()
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
cost, usage, models = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
)
assert cost > 0
assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06)
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
assert models == ["gemini-3.6-flash", "gemini-3.6-flash"]
@pytest.mark.asyncio
async def test_output_file_content_no_output_file_id_raises():
with pytest.raises(ValueError, match="Output file id is None"):
await bu._get_batch_output_file_content_as_dictionary(_batch(None), custom_llm_provider="openai")
await bu._fetch_batch_output_file_content(_batch(None), custom_llm_provider="openai")
@pytest.mark.asyncio
@ -641,13 +887,13 @@ async def test_output_file_content_fetches_and_parses(monkeypatch):
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
monkeypatch.setattr(cu, "_is_base64_encoded_unified_file_id", lambda fid: False)
result = await bu._get_batch_output_file_content_as_dictionary(
result = await bu._fetch_batch_output_file_content(
_batch("file-out"),
custom_llm_provider="azure",
litellm_params={"api_key": "sk-az", "api_base": "https://az", "model": "x"},
)
assert result == [{"a": 1}, {"b": 2}]
assert result == b'{"a": 1}\n{"b": 2}'
# afile_content received the file id + extracted credentials (not "model").
assert captured["file_id"] == "file-out"
assert captured["custom_llm_provider"] == "azure"
@ -676,13 +922,13 @@ async def test_output_file_content_unified_file_id_extraction(monkeypatch):
lambda fid: "litellm_proxy;llm_output_file_id,real-file-99;rest",
)
await bu._get_batch_output_file_content_as_dictionary(_batch("encoded-blob"), custom_llm_provider="openai")
await bu._fetch_batch_output_file_content(_batch("encoded-blob"), custom_llm_provider="openai")
assert captured["file_id"] == "real-file-99"
# =========================================================================== #
# _handle_completed_batch (async orchestrator: fetch -> cost/usage/models)
# _handle_completed_batch (async orchestrator: fetch -> single-pass aggregate)
# =========================================================================== #
@ -690,49 +936,48 @@ async def test_output_file_content_unified_file_id_extraction(monkeypatch):
async def test_handle_completed_batch_orchestration(monkeypatch):
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
async def fake_get_content(batch, custom_llm_provider, litellm_params=None):
return rows
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl(rows)
monkeypatch.setattr(bu, "_get_batch_output_file_content_as_dictionary", fake_get_content)
monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 3.3)
monkeypatch.setattr(
bu,
"_get_batch_job_total_usage_from_file_content",
lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3)
cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
assert cost == 3.3
assert usage.total_tokens == 15
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert models == ["gpt-4o"]
# =========================================================================== #
# Remaining branch: vertex usage disable-transform path.
#
# NOTE: the error path of _get_batch_job_cost_from_file_content (its `raise e`)
# is intentionally NOT tested: the preceding line logs via
# `verbose_logger.error("...", e)`, which passes the exception as a logging
# format-arg with no placeholder and itself raises TypeError under
# logging.raiseExceptions, masking the original error. Asserting that masked
# behavior would lock a source bug; left uncovered on purpose.
# =========================================================================== #
@pytest.mark.asyncio
async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch):
raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}]
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
return _vertex_jsonl(raw_rows)
def test_total_usage_vertex_disable_transform_path(monkeypatch):
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
monkeypatch.setattr(
bu,
"calculate_vertex_ai_batch_cost_and_usage",
lambda content, model: (
0.0,
Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
),
seen: dict = {}
def fake_vertex_calc(content, model):
seen["content"] = content
seen["model"] = model
return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)
monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc)
cost, usage, models = await bu._handle_completed_batch(
_batch("gs://litellm-bucket/output/predictions.jsonl"),
custom_llm_provider="vertex_ai",
model_name="gemini-x",
)
usage = bu._get_batch_job_total_usage_from_file_content([], custom_llm_provider="vertex_ai", model_name="gemini-x")
assert cost == 7.7
assert usage.total_tokens == 3
assert models == ["gemini-x"]
assert seen["content"] == raw_rows
assert seen["model"] == "gemini-x"
def _anthropic_usage(input_tokens, output_tokens, cache_creation=0, cache_read=0):
@ -832,46 +1077,54 @@ def test_bedrock_cost_uses_deployment_model_name():
"recordId": "1",
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
}
cost = bu._get_batch_job_cost_from_file_content(
file_content_dictionary=[row],
cost, _, models = bu._aggregate_batch_cost_usage_models(
entries=[row],
custom_llm_provider="bedrock",
model_name="us.anthropic.claude-sonnet-4-6",
model_info={},
)
assert cost > 0
assert models == ["us.anthropic.claude-sonnet-4-6"]
def test_anthropic_total_usage_sums_succeeded_only():
def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
import litellm.cost_calculator as cc
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
rows = [
_anthropic_succeeded_row(usage=_anthropic_usage(10, 5)),
_anthropic_errored_row(),
_anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)),
]
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="anthropic")
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145)
def test_anthropic_total_usage_aggregates_cache_token_details():
def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
import litellm.cost_calculator as cc
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
rows = [
_anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
_anthropic_errored_row(),
_anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)),
]
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="anthropic")
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert usage.prompt_tokens_details.cached_tokens == 8700
assert usage.prompt_tokens_details.cache_creation_tokens == 2300
assert usage.cache_read_input_tokens == 8700
assert usage.cache_creation_input_tokens == 2300
def test_total_usage_without_cache_tokens_has_no_prompt_details():
def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
rows = [
{
"custom_id": "req-1",
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
}
]
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="openai")
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
assert usage.prompt_tokens_details is None
@ -884,8 +1137,8 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
_anthropic_errored_row(),
]
total = bu._get_batch_job_cost_from_file_content(
rows,
total, _, _ = bu._aggregate_batch_cost_usage_models(
entries=rows,
custom_llm_provider="anthropic",
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
)
@ -910,8 +1163,8 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"),
)
total = bu._get_batch_job_cost_from_file_content(
[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
total, _, _ = bu._aggregate_batch_cost_usage_models(
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
)
assert total == pytest.approx(0.3)
@ -920,12 +1173,16 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
assert seen[0]["usage"].prompt_tokens == 10
def test_anthropic_batch_models_collected_from_succeeded_rows():
def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch):
import litellm.cost_calculator as cc
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
rows = [
_anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"),
_anthropic_errored_row(),
]
assert bu._get_batch_models_from_file_content(rows, None, "anthropic") == ["claude-sonnet-4-5-20250929"]
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
assert models == ["claude-sonnet-4-5-20250929"]
@pytest.mark.asyncio

38
tests/test_litellm/caching/test_caching_handler.py
View file

@ -558,6 +558,44 @@ async def test_embedding_cache_falls_back_to_token_counter_for_legacy_entries():
assert response.usage.prompt_tokens > 0
@pytest.mark.asyncio
async def test_embedding_cache_hit_sets_custom_llm_provider_on_logging_obj():
"""A full embedding cache hit must stamp the resolved provider onto the logging
obj so spend logs record the provider instead of None/unknown."""
from litellm.types.utils import CallTypes
llm_caching_handler = LLMCachingHandler(
original_function=MagicMock(),
request_kwargs={},
start_time=datetime.now(),
)
cached_result = [
{
"embedding": [-0.025, -0.019],
"index": 0,
"object": "embedding",
"model": "text-embedding-3-small",
"prompt_tokens": 5,
}
]
logging_obj = _build_logging_obj(CallTypes.aembedding.value, stream=False)
logging_obj.async_success_handler = AsyncMock()
response, cache_hit = llm_caching_handler._process_async_embedding_cached_response(
final_embedding_cached_response=None,
cached_result=cached_result,
kwargs={"model": "text-embedding-3-small", "input": "hello world"},
logging_obj=logging_obj,
start_time=datetime.now(),
model="text-embedding-3-small",
)
assert cache_hit
assert logging_obj.model_call_details["custom_llm_provider"] == "openai"
def test_request_kwargs_does_not_retain_logging_obj():
"""
The caching handler lives on logging_obj._llm_caching_handler, so keeping

186
tests/test_litellm/caching/test_redis_cache.py
View file

@ -59,9 +59,24 @@ def test_delete_cache_applies_namespace(namespace, monkeypatch, redis_no_ping):
@pytest.mark.asyncio
async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping):
monkeypatch.setenv("REDIS_HOST", "my-fake-host")
redis_cache = RedisCache(socket_timeout=1.0)
@pytest.mark.parametrize(
"redis_config",
[
pytest.param({"host": "my-fake-host"}, id="host_port"),
pytest.param({"url": "redis://my-fake-host:6379"}, id="url"),
],
)
async def test_redis_client_init_with_socket_timeout(monkeypatch, redis_no_ping, redis_config):
"""socket_timeout has to reach the connection however Redis was configured.
A url config used to drop every connection kwarg, so redis-py was left with
socket_timeout (and socket_connect_timeout, which falls back to it) unset. A
Redis host that drops packets instead of refusing them then blocks each caller
indefinitely, and the circuit breaker never trips because no call ever returns.
"""
monkeypatch.delenv("REDIS_URL", raising=False)
monkeypatch.delenv("REDIS_HOST", raising=False)
redis_cache = RedisCache(socket_timeout=1.0, **redis_config)
assert redis_cache.redis_kwargs["socket_timeout"] == 1.0
client = redis_cache.init_async_client()
assert client is not None
@ -428,3 +443,168 @@ def test_delete_cache_namespaces_key(namespace, expected, monkeypatch, redis_no_
redis_cache.redis_client = mock_client
redis_cache.delete_cache(key="k")
mock_client.delete.assert_called_once_with(expected)
def _closed_port() -> int:
"""A port with nothing listening, so Redis calls fail fast and deterministically."""
import socket
with socket.socket() as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"call_method",
[
pytest.param(lambda c: c.async_get_cache("lit4930"), id="async_get_cache"),
pytest.param(lambda c: c.async_batch_get_cache(["lit4930"]), id="async_batch_get_cache"),
pytest.param(lambda c: c.async_set_cache("lit4930", "v"), id="async_set_cache"),
pytest.param(lambda c: c.async_get_ttl("lit4930"), id="async_get_ttl"),
],
)
async def test_circuit_breaker_opens_when_method_swallows_redis_failure(redis_no_ping, call_method):
"""A guarded method that swallows its own Redis error must still count as a failure.
These methods catch connection errors and return a default so callers degrade instead
of failing, which is correct. But that returns cleanly through the circuit breaker
guard, and counting it as a success reset the failure streak on every call, so the
breaker could never open. An unreachable Redis then stayed in the pool and every
request kept paying the full socket timeout on it.
"""
from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD
cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5)
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD):
await call_method(cache)
with pytest.raises(Exception, match="circuit breaker is open"):
await call_method(cache)
@pytest.mark.asyncio
async def test_circuit_breaker_success_still_resets_the_failure_streak(redis_no_ping):
"""A reachable Redis must keep the breaker closed, however many earlier calls failed.
The guard now records success only when nothing failed while the method ran, so this
pins the other half of that contract: a call that genuinely reaches Redis has to clear
the streak, or a healthy Redis would eventually be evicted from the pool.
"""
from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD
cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5)
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1):
await cache.async_get_cache("lit4930")
assert cache._circuit_breaker.is_open() is False
reachable_redis = AsyncMock()
reachable_redis.get.return_value = None
with patch.object(cache, "init_async_client", return_value=reachable_redis):
await cache.async_get_cache("lit4930")
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD - 1):
await cache.async_get_cache("lit4930")
assert cache._circuit_breaker.is_open() is False, "one success must clear the streak"
@pytest.mark.asyncio
async def test_circuit_breaker_covers_lua_script_execution(redis_no_ping):
"""Lua script execution must feed the breaker like every other Redis call.
The v3 rate limiter issues all of its Redis traffic through async_register_script, so
leaving that path unguarded meant the coordination calls during an outage never
counted toward taking Redis out of the pool and kept paying a full socket timeout
each, which is the traffic the outage hurts most.
"""
from litellm.constants import REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD
cache = RedisCache(host="127.0.0.1", port=_closed_port(), socket_timeout=0.5)
run_script = cache.async_register_script("return 1")
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD):
with pytest.raises(Exception):
await run_script(keys=["lit4930"], args=[1])
with pytest.raises(Exception, match="circuit breaker is open"):
await run_script(keys=["lit4930"], args=[1])
@pytest.mark.asyncio
async def test_concurrent_success_is_not_cancelled_by_another_calls_failure():
"""One caller's failure must not discard a different caller's success.
A breaker is shared by every concurrent caller, so tracking "did this call fail" on the
breaker itself cannot tell my failure from someone else's. A Redis that is still
answering would then be evicted from the pool by unrelated in-flight failures, which is
the opposite of the outage this guard exists to handle.
"""
from redis.exceptions import ConnectionError as RedisConnectionError
from litellm.caching.redis_cache import (
RedisCircuitBreaker,
_record_swallowed_redis_failure,
_run_under_circuit_breaker,
)
breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60)
# The failure has to land after both calls are already in flight, which is the only
# ordering where a shared counter confuses the two. Failing before the healthy call
# starts would leave its snapshot correct and prove nothing.
async def swallows_a_failure():
await asyncio.sleep(0.02)
_record_swallowed_redis_failure(breaker, RedisConnectionError("redis unreachable"))
return None
async def succeeds_while_the_other_fails():
await asyncio.sleep(0.05)
return "ok"
rounds = breaker.failure_threshold + 1
for _ in range(rounds):
await asyncio.gather(
_run_under_circuit_breaker(breaker, "failing", swallows_a_failure),
_run_under_circuit_breaker(breaker, "healthy", succeeds_while_the_other_fails),
)
assert breaker._failure_count < breaker.failure_threshold, "the healthy call must clear the streak"
assert breaker.is_open() is False, "a Redis answering every round must stay in the pool"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error, opens_breaker",
[
pytest.param("ConnectionError", True, id="connection_refused_is_unhealthy"),
pytest.param("TimeoutError", True, id="timeout_is_unhealthy"),
pytest.param("BusyLoadingError", True, id="loading_is_unhealthy"),
pytest.param("ResponseError", False, id="wrong_type_command_is_not"),
pytest.param("DataError", False, id="bad_data_is_not"),
],
)
async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker):
"""Command and data errors must not count against Redis health.
They say nothing about connectivity, and a caller able to provoke them (an INCR against
a non-numeric value, say) could otherwise trip the shared breaker on demand and drop
rate limiting to per-process counters, which spreading traffic across replicas outruns.
"""
import redis.exceptions
from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker
breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60)
raised = getattr(redis.exceptions, error)("boom")
async def failing_call():
raise raised
for _ in range(breaker.failure_threshold + 1):
with pytest.raises(Exception):
await _run_under_circuit_breaker(breaker, "op", failing_call)
assert breaker.is_open() is opens_breaker

135
tests/test_litellm/integrations/otel/test_otel_v2_emitter.py
View file

@ -17,6 +17,9 @@ from litellm.integrations.otel.plumbing import context as ctx_mod # noqa: E402
from litellm.integrations.otel.plumbing import providers # noqa: E402
from litellm.integrations.otel.emitter import SpanEmitter # noqa: E402
from litellm.integrations.otel.emitter import stamp_error # noqa: E402
from litellm.integrations.otel.mappers.utils import ( # noqa: E402
MAX_TOOL_DEFINITION_ATTRS_PER_SPAN,
)
from litellm.integrations.otel.model.payloads import ( # noqa: E402
GuardrailSpanData,
LLMCallSpanData,
@ -305,3 +308,135 @@ def test_guardrail_success_span_is_unset():
)
(span,) = exporter.get_finished_spans()
assert span.status.status_code is StatusCode.UNSET
def _tools_payload(count):
"""A request declaring ``count`` tools, in the chat-completion shape."""
return _payload(
model_parameters={
"temperature": 0.7,
"tools": [
{
"type": "function",
"function": {
"name": f"tool_{i}",
"description": f"description for tool {i}",
"parameters": {"type": "object", "properties": {}},
},
}
for i in range(count)
],
}
)
def test_many_tools_do_not_evict_core_attributes():
"""Tool definitions must never crowd core telemetry off the span.
An agentic client declares hundreds of tools. Spelling each one out as
per-index attributes overruns the OTel SDK's 128-attribute span limit,
which evicts oldest-first and so destroys the ``gen_ai.*`` attributes
written before it. Capping the tool family keeps the core intact.
"""
engine, exporter = _engine()
data = LLMCallSpanData.from_standard_logging_payload(_tools_payload(127))
engine.emit(SpanRole.LLM_CALL, data)
(span,) = exporter.get_finished_spans()
a = span.attributes
assert a[GenAI.REQUEST_MODEL] == "gpt-4o"
assert a[GenAI.PROVIDER_NAME] == "openai"
assert a[GenAI.USAGE_INPUT_TOKENS] == 10
assert a[GenAI.USAGE_OUTPUT_TOKENS] == 5
assert a[GenAI.RESPONSE_FINISH_REASONS] == ("stop",)
assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002
assert a["gen_ai.usage.prompt_tokens"] == 10
assert span.dropped_attributes == 0
assert a[LiteLLM.TOOLS_DECLARED] == 127
assert a["gen_ai.tool.0.name"] == "tool_0"
assert "gen_ai.tool.126.name" not in a
assert "llm.request.functions.126.name" not in a
def test_tool_definitions_kept_in_full_below_the_cap():
"""A handful of tools keeps full per-index detail in both vocabularies."""
engine, exporter = _engine()
data = LLMCallSpanData.from_standard_logging_payload(_tools_payload(3))
engine.emit(SpanRole.LLM_CALL, data)
(span,) = exporter.get_finished_spans()
a = span.attributes
assert a[LiteLLM.TOOLS_DECLARED] == 3
for idx in range(3):
assert a[f"gen_ai.tool.{idx}.name"] == f"tool_{idx}"
assert a[f"gen_ai.tool.{idx}.description"] == f"description for tool {idx}"
assert a[f"gen_ai.tool.{idx}.parameters"]
assert a[f"llm.request.functions.{idx}.name"] == f"tool_{idx}"
def _tool_span(mapper_names, tool_count):
"""The exported LLM-call span for ``mapper_names`` and ``tool_count`` tools."""
cfg = OpenTelemetryV2Config(
exporter="in_memory",
legacy_compat=True,
mapper_names=list(mapper_names),
)
provider, exporter = providers.in_memory_provider(cfg)
engine = SpanEmitter(providers.get_tracer(provider, "litellm-test"), cfg)
engine.emit(
SpanRole.LLM_CALL,
LLMCallSpanData.from_standard_logging_payload(_tools_payload(tool_count)),
)
(span,) = exporter.get_finished_spans()
return span
def _tool_definition_keys(attributes):
return [
key
for key in attributes
if key.startswith(("gen_ai.tool.", "llm.request.functions.", "llm.tools."))
]
@pytest.mark.parametrize(
"mapper_names",
[
["genai"],
["genai", "openinference"],
["genai", "openinference", "langfuse", "weave", "langtrace"],
],
)
def test_tool_definitions_stay_within_one_span_wide_budget(mapper_names):
"""Every supported composition has to leave core telemetry on the span.
Each vocabulary spells the same tools out under its own keys, so an
allowance handed to each mapper separately multiplies by the number of
configured vocabularies and reaches the attribute limit again. Arize and
Phoenix already layer OpenInference on top of the default two, and every
vendor vocabulary can be listed at once. One budget shared across them all
is what keeps the total bounded.
"""
span = _tool_span(mapper_names, 127)
a = span.attributes
assert span.dropped_attributes == 0
assert a[GenAI.REQUEST_MODEL] == "gpt-4o"
assert a[GenAI.PROVIDER_NAME] == "openai"
assert a[GenAI.USAGE_INPUT_TOKENS] == 10
assert a[GenAI.USAGE_OUTPUT_TOKENS] == 5
assert a[f"{LiteLLM.COST_PREFIX}total"] == 0.002
assert a[LiteLLM.TOOLS_DECLARED] == 127
emitted = _tool_definition_keys(a)
assert emitted, "some tool detail should survive in every composition"
assert len(emitted) <= MAX_TOOL_DEFINITION_ATTRS_PER_SPAN
def test_vendor_tool_definitions_are_truncated_not_dropped():
"""The OpenInference vocabulary keeps its leading tools and loses the tail."""
a = _tool_span(["genai", "openinference"], 127).attributes
assert a["llm.tools.0.tool.name"] == "tool_0"
assert a["llm.tools.0.tool.json_schema"]
assert "llm.tools.126.tool.name" not in a

532
tests/test_litellm/integrations/otel/test_otel_v2_metrics.py
View file

@ -12,9 +12,17 @@ raises out of ``GenAIMetricRecorder.record`` -- asserted directly at the recorde
layer -- and the logger turns that raise into a single ERROR ("metrics disabled")
plus a quiet no-op for the rest of the process, asserted at the logger layer so
the misconfig never breaks a request nor spams a log line per request.
The failure path is driven the same way, through the real
``OpenTelemetryV2.async_log_failure_event``: a failed call records
``gen_ai.client.operation.duration`` and nothing else, tagged with ``error.type``,
and a success driven through the same reader keeps a datapoint whose attributes are
byte-for-byte what it had before the failure path existed -- the guard for every
dashboard already querying that histogram.
"""
import asyncio
import json
from datetime import datetime, timedelta
import pytest
@ -25,6 +33,9 @@ from opentelemetry.sdk.metrics import MeterProvider # noqa: E402
from opentelemetry.sdk.metrics.export import InMemoryMetricReader # noqa: E402
import litellm # noqa: E402
from litellm.constants import ( # noqa: E402
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
)
from litellm.integrations.otel.logger import OpenTelemetryV2 # noqa: E402
from litellm.integrations.otel.model.config import ( # noqa: E402
OpenTelemetryV2Config,
@ -57,15 +68,17 @@ ALL_METRICS = frozenset(
TOKEN_TYPE = "gen_ai.token.type"
MODEL_KEY = "gen_ai.request.model"
OPERATION_KEY = "gen_ai.operation.name"
PROVIDER_NAME_KEY = "gen_ai.provider.name"
SYSTEM_KEY = "gen_ai.system"
# Each is a member of VALID_METRIC_ATTRIBUTE_NAMES and is stamped on the metric
# by default (proven by the no-filter test below).
HIGH_CARDINALITY_KEYS = (
# Keys inside the ceiling that an operator's filter must still be able to remove.
# Every one is bounded, so it survives the ceiling and only the operator's own
# exclude_list takes it off; that is what makes the filter tests non-vacuous.
FILTERABLE_KEYS = (
"hidden_params",
"metadata.user_api_key_hash",
"metadata.requester_ip_address",
"metadata.requester_metadata",
"metadata.applied_guardrails",
"metadata.user_api_key_team_id",
)
PROMPT_TOKENS = 137
@ -73,18 +86,25 @@ COMPLETION_TOKENS = 89
RESPONSE_COST = 0.0023
def _build_call(stream: bool = True):
def _build_call(
stream: bool = True,
provider: str | None = "openai",
call_type: str = "completion",
):
"""A captured success-call (kwargs, response_obj, start, end) that exercises
every one of the six metrics: usage for token.usage, response_cost for cost,
streaming + timing for the response-time histograms."""
streaming + timing for the response-time histograms.
``provider=None`` omits ``custom_llm_provider`` entirely, reproducing a call
litellm could not attribute to a provider."""
start = datetime(2026, 6, 12, 12, 0, 0)
api_call_start = start + timedelta(seconds=0.1)
completion_start = start + timedelta(seconds=0.5)
end = start + timedelta(seconds=1.0)
kwargs = {
"model": "gpt-4o-mini",
"call_type": "completion",
"litellm_params": {"custom_llm_provider": "openai"},
"call_type": call_type,
"litellm_params": ({"custom_llm_provider": provider} if provider is not None else {}),
"optional_params": {"stream": stream},
"response_cost": RESPONSE_COST,
"api_call_start_time": api_call_start,
@ -93,6 +113,7 @@ def _build_call(stream: bool = True):
"standard_logging_object": {
"metadata": {
"user_api_key_hash": "hash-abc123",
"user_api_key_team_id": "team-1",
"requester_ip_address": "10.0.0.7",
"requester_metadata": {"team": "alpha", "tier": "gold"},
"applied_guardrails": ["pii", "toxicity"],
@ -131,7 +152,7 @@ def _metrics_by_name(reader):
return out
def _drive_success(reader, callback_settings_attributes=None):
def _drive_success(reader, callback_settings_attributes=None, **call_overrides):
"""Construct a metrics-on logger, optionally populate callback_settings AFTER
construction (mirroring the proxy ordering), run the real success hook."""
logger = _logger(reader, enable_metrics=True)
@ -141,7 +162,7 @@ def _drive_success(reader, callback_settings_attributes=None):
"otel": {"attributes": callback_settings_attributes}
}
try:
kwargs, response_obj, start, end = _build_call()
kwargs, response_obj, start, end = _build_call(**call_overrides)
asyncio.run(logger.async_log_success_event(kwargs, response_obj, start, end))
finally:
litellm.callback_settings = previous
@ -205,14 +226,14 @@ def test_metrics_off_by_default_records_nothing():
def test_exclude_list_strips_high_cardinality_across_metrics():
"""exclude_list set AFTER construction (the proxy path) removes every
high-cardinality key from more than one metric while the low-cardinality
model attribute survives."""
"""exclude_list set AFTER construction (the proxy path) removes every listed
key from more than one metric while the low-cardinality model attribute
survives."""
metrics = _drive_success(
InMemoryMetricReader(),
callback_settings_attributes={"exclude_list": list(HIGH_CARDINALITY_KEYS)},
callback_settings_attributes={"exclude_list": list(FILTERABLE_KEYS)},
)
excluded = set(HIGH_CARDINALITY_KEYS)
excluded = set(FILTERABLE_KEYS)
for name in (OPERATION_DURATION, TOKEN_USAGE):
points = metrics[name]
@ -241,18 +262,226 @@ def test_include_list_allows_only_listed_attributes():
assert set(dp.attributes.keys()) - {TOKEN_TYPE} == allowed
def test_no_filter_keeps_high_cardinality_keys():
"""Backward compatibility: without an attributes config every high-cardinality
key the call carries is still stamped, so the filter tests above prove a real
removal rather than a key that was never present."""
def test_no_filter_still_keeps_the_filterable_keys():
"""Without an attributes config every key the filter tests remove is present,
so those tests prove a real removal rather than a key that was never there."""
metrics = _drive_success(InMemoryMetricReader())
expected = set(HIGH_CARDINALITY_KEYS)
expected = set(FILTERABLE_KEYS)
for name in (OPERATION_DURATION, TOKEN_USAGE):
for dp in metrics[name]:
assert expected.issubset(set(dp.attributes.keys()))
def test_a_metric_ineligible_filter_name_is_reported_not_silently_dropped(caplog):
"""Naming a metric-ineligible attribute in a filter has to say so.
The shared validator accepts every span attribute name, so an operator can put
one in an ``include_list``, get nothing for it, and have no way to tell that from
a value that happened to be absent. The ceiling is deliberate, but silent is what
makes it a support ticket.
"""
with caplog.at_level("WARNING"):
_drive_success(
InMemoryMetricReader(),
callback_settings_attributes={
"include_list": [MODEL_KEY, "metadata.requester_ip_address"]
},
)
reported = [
r.getMessage().split(" cannot be a metric attribute")[0].removeprefix("OTel metrics: ")
for r in caplog.records
if r.levelname == "WARNING" and "cannot be a metric attribute" in r.getMessage()
]
assert reported == ["metadata.requester_ip_address"], reported
def test_two_calls_differing_only_per_request_share_one_series():
"""The whole point of the ceiling: metric cardinality must not grow with traffic.
Every field here moves on every real request -- the response cost, the call id,
the cache key, the provider's remaining-rate-limit headers -- and each one used
to reach the datapoint inside a single ``hidden_params`` label. A unique label
value is a new time series, so each of the six instruments minted one series per
request, which is both a Grafana Cloud bill proportional to traffic and a
histogram that cannot be aggregated. Identical attribute sets is what "one
series" means to the SDK.
"""
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
for index, cost in enumerate((RESPONSE_COST, RESPONSE_COST * 3)):
kwargs, response_obj, start, end = _build_call()
kwargs["response_cost"] = cost
kwargs["standard_logging_object"]["hidden_params"] = {
"model_id": "m-1",
# A documented per-call parameter, so it varies here on purpose: the same
# deployment reached under a caller-chosen base must not split the series.
"api_base": f"https://proxy-{index}.example.com/v1",
"litellm_call_id": f"call-{index}",
"cache_key": f"cache-{index}",
"response_cost": cost,
"litellm_overhead_time_ms": 1.5 + index,
"usage_object": {"prompt_tokens": index, "completion_tokens": index},
"additional_headers": {"x_ratelimit_remaining_requests": 100 - index},
}
asyncio.run(logger.async_log_success_event(kwargs, response_obj, start, end))
for name in ALL_METRICS:
attribute_sets = {
tuple(sorted((k, v) for k, v in dp.attributes.items() if k != TOKEN_TYPE))
for dp in _metrics_by_name(reader)[name]
}
assert len(attribute_sets) == 1, f"{name} split into {len(attribute_sets)} series across 2 requests"
def test_hidden_params_label_carries_only_bounded_deployment_fields():
"""``hidden_params`` survives the ceiling, but only as the deployment identity.
``model_id`` is the router's deployment id, bounded by the deployment list, and is
what a per-deployment dashboard reads. Everything else in the object is
per-request or caller-chosen and belongs on the span, which already carries it.
``api_base`` is excluded despite naming the same deployment: it is a documented
per-call parameter, so a caller varying it would restore the per-request
cardinality this cap exists to remove.
"""
kwargs, response_obj, start, end = _build_call()
kwargs["standard_logging_object"]["hidden_params"] = {
"model_id": "m-1",
"api_base": "https://api.openai.com/v1",
"litellm_call_id": "abc",
"cache_key": "ck-1",
"response_cost": RESPONSE_COST,
}
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
asyncio.run(logger.async_log_success_event(kwargs, response_obj, start, end))
label = _metrics_by_name(reader)[OPERATION_DURATION][0].attributes["hidden_params"]
assert json.loads(label) == {"model_id": "m-1"}
def test_success_attributes_are_capped_at_the_ceiling():
"""The success path carries exactly the ceiling, no client-supplied attributes.
The fixture deliberately sets every excluded key, so this asserts a real removal
rather than keys that were never present.
"""
kwargs, response_obj, start, end = _build_call()
metadata = kwargs["standard_logging_object"]["metadata"]
metadata.update(
{
"spend_logs_metadata": {"cost_center": "abc"},
"user_api_key_end_user_id": "end-user-1",
"user_api_key_user_email": "someone@example.com",
}
)
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
asyncio.run(logger.async_log_success_event(kwargs, response_obj, start, end))
metrics = _metrics_by_name(reader)
for name in ALL_METRICS:
for dp in metrics[name]:
leaked = set(dp.attributes) - set(BOUNDED_KEYS) - {TOKEN_TYPE}
assert not leaked, f"{name} leaked {leaked}"
def test_provider_is_labelled_with_semconv_provider_name():
"""Every recorded point carries gen_ai.provider.name holding the semconv
provider value (bedrock -> aws.bedrock), the key the GenAI convention and the
dashboards built on it query. The deprecated gen_ai.system spelling alone is
unreadable to them."""
metrics = _drive_success(InMemoryMetricReader(), provider="bedrock")
for name in ALL_METRICS:
points = metrics[name]
assert points, f"{name} was not recorded"
for dp in points:
assert dp.attributes[PROVIDER_NAME_KEY] == "aws.bedrock"
def test_deprecated_gen_ai_system_is_dual_emitted_verbatim():
"""gen_ai.system keeps its raw litellm provider value alongside the new key
for one release, so a dashboard already filtering on it keeps matching. Its
value must not be swapped for the mapped one, which would break exactly the
queries the dual emission exists to protect."""
metrics = _drive_success(InMemoryMetricReader(), provider="bedrock")
for dp in metrics[OPERATION_DURATION]:
assert dp.attributes[SYSTEM_KEY] == "bedrock"
assert dp.attributes[PROVIDER_NAME_KEY] == "aws.bedrock"
def test_no_provider_attribute_when_provider_is_absent():
"""A call litellm could not attribute to a provider carries no provider label
at all. A placeholder value ("Unknown") would mint a permanent series that
aggregates every unattributable request and that no operator can act on."""
metrics = _drive_success(InMemoryMetricReader(), provider=None)
for name in ALL_METRICS:
points = metrics[name]
assert points, f"{name} was not recorded"
for dp in points:
keys = set(dp.attributes.keys())
assert PROVIDER_NAME_KEY not in keys
assert SYSTEM_KEY not in keys
assert "Unknown" not in set(dp.attributes.values())
def test_vector_store_search_is_not_labelled_as_chat():
"""A vector-store search records under gen_ai.operation.name=retrieval, so its
latency and cost stay out of the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type="avector_store_search")
for name in (OPERATION_DURATION, TOKEN_COST):
for dp in metrics[name]:
assert dp.attributes[OPERATION_KEY] == "retrieval"
@pytest.mark.parametrize(
"call_type,expected",
[
("avector_store_create", "litellm.vector_store_management"),
("avector_store_delete", "litellm.vector_store_management"),
("avector_store_file_create", "litellm.vector_store_file_management"),
("avector_store_file_list", "litellm.vector_store_file_management"),
],
)
def test_vector_store_management_is_not_labelled_as_chat(call_type, expected):
"""Store and file management reach the recorder through the same success hook as a
completion, so leaving them unmapped kept billing- and latency-relevant admin calls
inside the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type=call_type)
for dp in metrics[OPERATION_DURATION]:
assert dp.attributes[OPERATION_KEY] == expected
@pytest.mark.parametrize("call_type", ["asend_message", "asend_message_streaming"])
def test_agent_message_is_not_labelled_as_chat(call_type):
"""An A2A agent send records under gen_ai.operation.name=invoke_agent, streamed or
not. The streaming iterator dispatches the same success handlers under its own
``asend_message_streaming`` call type, so an unmapped streaming spelling puts every
streamed agent turn's latency and cost back into the chat series."""
metrics = _drive_success(InMemoryMetricReader(), call_type=call_type)
for name in (OPERATION_DURATION, TOKEN_COST):
for dp in metrics[name]:
assert dp.attributes[OPERATION_KEY] == "invoke_agent"
def test_provider_name_is_filterable():
"""gen_ai.provider.name is a member of the metric-attribute allowlist, so an
operator can include or exclude it; an unlisted name raises instead."""
metrics = _drive_success(
InMemoryMetricReader(),
callback_settings_attributes={"include_list": [PROVIDER_NAME_KEY]},
)
for dp in metrics[OPERATION_DURATION]:
assert set(dp.attributes.keys()) == {PROVIDER_NAME_KEY}
def test_metrics_reach_operator_configured_global_provider(monkeypatch):
"""Regression: with no meter provider injected, the six gen_ai.client.*
histograms must record through the operator's globally configured
@ -331,3 +560,262 @@ def test_token_type_rejected_from_either_list(attributes, monkeypatch):
# the specific reason so dropping that guard (and falling through to "unknown
# attribute name") is caught.
assert "discriminator" in str(exc_info.value)
# --- failure path ------------------------------------------------------------ #
ERROR_TYPE = "error.type"
ERROR_CLASS = "RateLimitError"
FAILURE_DURATION_S = 1.0
# Attributes a failure datapoint must never carry. Each is either supplied by the
# caller (so a caller could mint a fresh series per request, and a failure costs
# them no provider spend) or varies per request, or is PII duplicating an id that
# is already on the series.
UNBOUNDED_KEYS = (
"metadata.requester_metadata",
"metadata.requester_ip_address",
"metadata.spend_logs_metadata",
"metadata.user_api_key_end_user_id",
"metadata.user_api_key_user_email",
)
# The exact set a datapoint may carry on either path: the operation, the
# operator-provisioned identity, and the deployment that served it.
BOUNDED_KEYS = (
"hidden_params",
"gen_ai.operation.name",
"gen_ai.provider.name",
"gen_ai.system",
"gen_ai.request.model",
"gen_ai.framework",
"metadata.user_api_key_hash",
"metadata.user_api_key_alias",
"metadata.user_api_key_team_id",
"metadata.user_api_key_team_alias",
"metadata.user_api_key_org_id",
"metadata.user_api_key_user_id",
)
def _build_failure(
*,
error_information=None,
exception=None,
no_upstream_call=False,
):
"""A captured failure-call ``(kwargs, start, end)``.
Mirrors what litellm actually hands ``async_log_failure_event``: no
``response_obj`` at all, but the streaming timings and the recovered
``response_cost`` a mid-stream failure still carries -- so routing the failure
path through the full success recorder would show up here as extra series
rather than passing unnoticed. The metadata carries both the bounded identity
keys and every caller-supplied / per-request key, so the allowlist test below
proves a real removal rather than a key that was never there.
"""
start = datetime(2026, 6, 12, 12, 0, 0)
api_call_start = start + timedelta(seconds=0.1)
completion_start = start + timedelta(seconds=0.5)
end = start + timedelta(seconds=FAILURE_DURATION_S)
standard_logging_object = {
"status": "failure",
"metadata": {
"user_api_key_hash": "hash-abc123",
"user_api_key_alias": "alias-abc",
"user_api_key_team_id": "team-1",
"user_api_key_team_alias": "team-alpha",
"user_api_key_org_id": "org-1",
"user_api_key_user_id": "user-1",
"user_api_key_user_email": "user@example.com",
"user_api_key_end_user_id": "end-user-42",
"requester_ip_address": "10.0.0.7",
"requester_metadata": {"trace": "caller-supplied-unique-value"},
"spend_logs_metadata": {"ticket": "caller-supplied-unique-value"},
},
"hidden_params": {
"litellm_call_id": "abc",
"model_id": "m-1",
"api_base": "https://api.openai.com/v1",
},
}
if error_information is not None:
standard_logging_object["error_information"] = error_information
kwargs = {
"model": "gpt-4o-mini",
"call_type": "completion",
"litellm_params": {"custom_llm_provider": "openai"},
"optional_params": {"stream": True},
"response_cost": RESPONSE_COST,
"api_call_start_time": api_call_start,
"completion_start_time": completion_start,
"end_time": end,
"standard_logging_object": standard_logging_object,
}
if exception is not None:
kwargs["exception"] = exception
if no_upstream_call:
kwargs[LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL] = True
return kwargs, start, end
def _drive_failure(reader, callback_settings_attributes=None, **failure_kwargs):
logger = _logger(reader, enable_metrics=True)
previous = litellm.callback_settings
if callback_settings_attributes is not None:
litellm.callback_settings = {"otel": {"attributes": callback_settings_attributes}}
try:
kwargs, start, end = _build_failure(**failure_kwargs)
asyncio.run(logger.async_log_failure_event(kwargs, None, start, end))
finally:
litellm.callback_settings = previous
return _metrics_by_name(reader)
def test_failure_records_only_the_duration_histogram():
"""A failed call contributes to gen_ai.client.operation.duration -- before this
existed a failure recorded nothing at all, so the histogram measured only the
traffic that survived. It contributes to nothing else: the other five
instruments describe a completed generation, and the call carries a streaming
timing pair and a recovered response_cost that would light four of them up if
the failure were routed through the success recorder."""
metrics = _drive_failure(
InMemoryMetricReader(),
error_information={"error_class": ERROR_CLASS, "error_code": "429"},
)
assert set(metrics.keys()) == {OPERATION_DURATION}
points = metrics[OPERATION_DURATION]
assert len(points) == 1
assert points[0].count == 1
assert points[0].sum == pytest.approx(FAILURE_DURATION_S)
assert points[0].attributes[ERROR_TYPE] == ERROR_CLASS
def test_success_and_failure_are_separable_and_success_attributes_unchanged():
"""The pooled histogram stays queryable per outcome, and the existing
dashboards keep working.
A success and a failure through one reader must land on two distinct series --
one with error.type, one without -- so a failure-rate panel is expressible and
an operator can still get success-only latency by filtering error.type="". The
success datapoint's attribute map must be byte-for-byte the map a success-only
run produces, which is what stops the new attribute from leaking onto the
series every current query reads."""
baseline_reader = InMemoryMetricReader()
baseline = _drive_success(baseline_reader)
baseline_points = baseline[OPERATION_DURATION]
assert len(baseline_points) == 1
baseline_attributes = dict(baseline_points[0].attributes)
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
ok_kwargs, response_obj, ok_start, ok_end = _build_call()
asyncio.run(logger.async_log_success_event(ok_kwargs, response_obj, ok_start, ok_end))
bad_kwargs, bad_start, bad_end = _build_failure(error_information={"error_class": ERROR_CLASS})
asyncio.run(logger.async_log_failure_event(bad_kwargs, None, bad_start, bad_end))
points = metrics = _metrics_by_name(reader)[OPERATION_DURATION]
assert len(points) == 2, f"success and failure collapsed into {len(points)} series: {metrics}"
succeeded = [dp for dp in points if ERROR_TYPE not in dp.attributes]
failed = [dp for dp in points if dp.attributes.get(ERROR_TYPE) == ERROR_CLASS]
assert len(succeeded) == 1 and len(failed) == 1
assert dict(succeeded[0].attributes) == baseline_attributes
def test_failure_attributes_are_a_bounded_allowlist():
"""A failure datapoint carries exactly the bounded allowlist plus error.type.
A failed request needs no provider spend, so nothing rate-limits a caller who
puts a unique value into an attribute they control and mints one histogram
series per request. The same payload is driven through the success path first,
which does carry those keys, so this asserts a real removal on the failure path
rather than keys that were never present. The exact-set assertion is the guard
against the natural refactor of "just reuse _common_attributes"."""
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
kwargs, start, end = _build_failure(error_information={"error_class": ERROR_CLASS})
usage = {"usage": {"prompt_tokens": 1, "completion_tokens": 1}}
asyncio.run(logger.async_log_success_event(kwargs, usage, start, end))
asyncio.run(logger.async_log_failure_event(kwargs, None, start, end))
points = _metrics_by_name(reader)[OPERATION_DURATION]
succeeded = next(dp for dp in points if ERROR_TYPE not in dp.attributes)
failed = next(dp for dp in points if ERROR_TYPE in dp.attributes)
supplied = set(kwargs["standard_logging_object"]["metadata"])
missing = {key for key in UNBOUNDED_KEYS if key.removeprefix("metadata.") not in supplied}
assert not missing, f"fixture never carried {missing}, so the exclusion below proves nothing"
leaked = set(UNBOUNDED_KEYS) & set(failed.attributes)
assert not leaked, f"failure datapoint leaked unbounded attributes: {leaked}"
assert set(failed.attributes) == set(BOUNDED_KEYS) | {ERROR_TYPE}
assert json.loads(failed.attributes["hidden_params"]) == {"model_id": "m-1"}
def test_operator_filter_can_still_narrow_the_failure_allowlist():
"""The allowlist is a ceiling, not a floor: an exclude_list an operator sets
still removes a listed key from the failure series."""
metrics = _drive_failure(
InMemoryMetricReader(),
callback_settings_attributes={"exclude_list": ["metadata.user_api_key_hash"]},
error_information={"error_class": ERROR_CLASS},
)
attributes = metrics[OPERATION_DURATION][0].attributes
assert "metadata.user_api_key_hash" not in attributes
assert attributes[ERROR_TYPE] == ERROR_CLASS
assert attributes[MODEL_KEY] == "gpt-4o-mini"
@pytest.mark.parametrize(
"failure_kwargs, expected",
[
({"error_information": {"error_class": ERROR_CLASS, "error_code": "429"}}, ERROR_CLASS),
({"error_information": {"error_code": "429"}}, "429"),
({"exception": ValueError("boom")}, "ValueError"),
({}, "_OTHER"),
],
ids=["error_class", "error_code_only", "exception_fallback", "unclassifiable"],
)
def test_error_type_is_bounded_and_falls_back(failure_kwargs, expected):
"""error.type is always a bounded value: the mapped exception's class name, the
provider status code, the raw exception's class name, or the semconv _OTHER
fallback. Never the exception message, which is unbounded."""
metrics = _drive_failure(InMemoryMetricReader(), **failure_kwargs)
assert metrics[OPERATION_DURATION][0].attributes[ERROR_TYPE] == expected
def test_include_list_cannot_strip_error_type():
"""error.type is a structural discriminator like gen_ai.token.type: an
include_list that does not mention it must not merge the failure series back
into the success series, so it is stamped after the filter runs."""
metrics = _drive_failure(
InMemoryMetricReader(),
callback_settings_attributes={"include_list": [MODEL_KEY]},
error_information={"error_class": ERROR_CLASS},
)
attributes = metrics[OPERATION_DURATION][0].attributes
assert dict(attributes) == {MODEL_KEY: "gpt-4o-mini", ERROR_TYPE: ERROR_CLASS}
def test_proxy_gate_rejection_records_no_duration():
"""A synthetic proxy-gate failure log (auth / rate-limit rejection) never made
an upstream call, so its wall time is not a GenAI operation's duration; it is
skipped for the same reason it gets no span. Recording it would pull the
histogram toward the proxy's own latency.
Both failures go through one reader so the assertion is that exactly the
upstream one landed, rather than the vacuous "nothing was recorded" a
failure path that records nothing at all would also satisfy."""
reader = InMemoryMetricReader()
logger = _logger(reader, enable_metrics=True)
gate_kwargs, gate_start, gate_end = _build_failure(
error_information={"error_class": "AuthenticationError"},
no_upstream_call=True,
)
asyncio.run(logger.async_log_failure_event(gate_kwargs, None, gate_start, gate_end))
upstream_kwargs, upstream_start, upstream_end = _build_failure(error_information={"error_class": ERROR_CLASS})
asyncio.run(logger.async_log_failure_event(upstream_kwargs, None, upstream_start, upstream_end))
points = _metrics_by_name(reader)[OPERATION_DURATION]
assert [dp.attributes[ERROR_TYPE] for dp in points] == [ERROR_CLASS]
assert points[0].count == 1

102
tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py
View file

@ -1,8 +1,13 @@
"""Tests for the OTel v2 sources of truth: span registry, semconv keys, config,
and the typed StandardLoggingPayload adapter. These need no OTel SDK."""
import logging
import re
from pathlib import Path
import pytest
import litellm
from litellm.integrations.otel import (
BAGGAGE_PROMOTED_KEYS,
DB,
@ -208,6 +213,103 @@ def test_operation_resolution():
assert resolve_operation("call_mcp_tool") is GenAIOperation.EXECUTE_TOOL
@pytest.mark.parametrize("call_type", ["vector_store_search", "avector_store_search"])
def test_vector_store_search_is_a_retrieval_operation(call_type):
"""A vector-store search is a retrieval, so its duration and cost must not
land in the chat series that dashboards read latency off."""
assert resolve_operation(call_type) is GenAIOperation.RETRIEVAL
assert resolve_operation(call_type).value == "retrieval"
@pytest.mark.parametrize("call_type", ["query", "aquery"])
def test_rag_query_is_a_retrieval_operation(call_type):
"""``/rag/query`` reaches the same recorder as a vector-store search and is the
same operation, so it must not be the one retrieval surface left reading as chat."""
assert resolve_operation(call_type) is GenAIOperation.RETRIEVAL
@pytest.mark.parametrize(
"call_type",
[
f"{prefix}vector_store_{verb}"
for verb in ("create", "retrieve", "list", "update", "delete")
for prefix in ("", "a")
],
)
def test_vector_store_management_is_not_chat(call_type):
"""The store lifecycle calls are not GenAI client operations and the convention
names nothing for them, so they take a vendor value rather than defaulting into
the chat series."""
assert resolve_operation(call_type) is GenAIOperation.LITELLM_VECTOR_STORE_MANAGEMENT
assert resolve_operation(call_type).value == "litellm.vector_store_management"
@pytest.mark.parametrize(
"call_type",
[
f"{prefix}vector_store_file_{verb}"
for verb in ("create", "list", "retrieve", "content", "update", "delete")
for prefix in ("", "a")
],
)
def test_vector_store_file_management_is_not_chat(call_type):
"""The file operations are a distinct REST resource from the store lifecycle, so
they get their own vendor value instead of sharing one bucket."""
assert resolve_operation(call_type) is GenAIOperation.LITELLM_VECTOR_STORE_FILE_MANAGEMENT
assert resolve_operation(call_type).value == "litellm.vector_store_file_management"
def test_vendor_operation_values_are_namespaced():
"""A vendor value must stay under the ``litellm.`` prefix: an unprefixed invented
name could collide with a value the convention adds later, silently changing what
a conformant consumer thinks it is reading."""
vendor = [op for op in GenAIOperation if op.name.startswith("LITELLM_")]
assert vendor, "no vendor operation values defined"
assert all(op.value.startswith("litellm.") for op in vendor)
@pytest.mark.parametrize("call_type", ["send_message", "asend_message", "asend_message_streaming"])
def test_agent_message_is_an_invoke_agent_operation(call_type):
"""An agent (A2A) message send is an agent invocation, not a chat completion.
The streaming spelling counts: ``_build_streaming_logging_obj`` in
``litellm/a2a_protocol/main.py`` stamps ``asend_message_streaming`` on the
logging object the streaming iterator dispatches success handlers with, so a
missing entry sends every streamed agent turn into the chat series. There is
no sync spelling because A2A streaming is async-only.
"""
assert resolve_operation(call_type) is GenAIOperation.INVOKE_AGENT
assert resolve_operation(call_type).value == "invoke_agent"
def test_every_call_type_the_a2a_package_stamps_is_an_agent_operation():
"""Pins the map to the call types the A2A code actually stamps on its logging
objects. A new spelling added there without a map entry fails here instead of
quietly landing in the chat series, which is how the streaming one was missed."""
a2a_package = Path(litellm.__file__).parent / "a2a_protocol"
stamped = {
call_type
for source in a2a_package.rglob("*.py")
for call_type in re.findall(r'call_type="([^"]+)"', source.read_text())
}
assert stamped, "no call_type literals found in litellm/a2a_protocol"
unmapped = {
call_type: resolve_operation(call_type).value
for call_type in stamped
if resolve_operation(call_type) is not GenAIOperation.INVOKE_AGENT
}
assert not unmapped, f"add these to _OPERATION_BY_CALL_TYPE: {unmapped}"
def test_unmapped_call_type_falls_back_to_chat_loudly(caplog):
"""The fallback still labels the series ``chat`` so it is never unlabelled,
but it says so at debug: a silent default is how retrieval and agent calls
ended up in the chat charts in the first place."""
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
assert resolve_operation("some_future_call_type") is GenAIOperation.CHAT
assert any("some_future_call_type" in record.getMessage() for record in caplog.records)
# --- MCP tool-call (source of truth #1/#2/#3) ------------------------------- #

76
tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_kimi_model_metadata.py Normal file
View file

@ -0,0 +1,76 @@
"""
Regression test for Fireworks Kimi K2.5 / K2.6 / K2.7 context and output limits.
Fireworks publishes a 262144-token context window for every Kimi K2.5, K2.6 and
K2.7 model, but caps generation well below that. A previous bulk edit had flattened
max_output_tokens/max_tokens to 262144 (equal to the context window), which let the
pre-call context-window check admit requests asking for a full 262144-token
completion that Fireworks then rejects. These assertions pin the corrected per-alias
limits so a future bulk edit can't silently flatten them again.
"""
import json
from importlib.resources import files
import pytest
CONTEXT_WINDOW = 262144
OUTPUT_LIMIT = 32768
KIMI_ALIASES = (
"fireworks_ai/kimi-k2p5",
"fireworks_ai/kimi-k2p6",
"fireworks_ai/kimi-k2p6-fast",
"fireworks_ai/kimi-k2p7-code",
"fireworks_ai/kimi-k2p7-code-fast",
"fireworks_ai/accounts/fireworks/models/kimi-k2p5",
"fireworks_ai/accounts/fireworks/models/kimi-k2p6",
"fireworks_ai/accounts/fireworks/models/kimi-k2p7-code",
"fireworks_ai/accounts/fireworks/routers/kimi-k2p6-fast",
"fireworks_ai/accounts/fireworks/routers/kimi-k2p7-code-fast",
)
@pytest.fixture(scope="module")
def use_local_model_cost_map():
monkeypatch = pytest.MonkeyPatch()
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
import litellm
from litellm.utils import _invalidate_model_cost_lowercase_map
original_model_cost = litellm.model_cost
litellm.model_cost = json.loads(
files("litellm")
.joinpath("model_prices_and_context_window_backup.json")
.read_text(encoding="utf-8")
)
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
try:
yield litellm
finally:
litellm.model_cost = original_model_cost
litellm.get_model_info.cache_clear()
_invalidate_model_cost_lowercase_map()
monkeypatch.undo()
@pytest.mark.parametrize("alias", KIMI_ALIASES)
def test_fireworks_kimi_raw_cost_entry_limits(use_local_model_cost_map, alias):
entry = use_local_model_cost_map.model_cost[alias]
assert entry["litellm_provider"] == "fireworks_ai"
assert entry["max_input_tokens"] == CONTEXT_WINDOW
assert entry["max_output_tokens"] == OUTPUT_LIMIT
assert entry["max_tokens"] == OUTPUT_LIMIT
assert entry["max_output_tokens"] < entry["max_input_tokens"]
@pytest.mark.parametrize("alias", KIMI_ALIASES)
def test_fireworks_kimi_get_model_info_limits(use_local_model_cost_map, alias):
model_info = use_local_model_cost_map.get_model_info(model=alias)
assert model_info["max_input_tokens"] == CONTEXT_WINDOW
assert model_info["max_output_tokens"] == OUTPUT_LIMIT
assert model_info["max_tokens"] == OUTPUT_LIMIT

343
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -7645,3 +7645,346 @@ class TestSessionBearerEgressScrub:
assert oauth2 is None
assert "authorization" not in {k.lower() for k in raw}
assert per_server == {"github": {"Authorization": "Bearer gh_injected_upstream"}}
# ---------------------------------------------------------------------------
# Internal-user (human) MCP entitlement tests
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
class TestUserMCPEntitlement:
"""The entitlement attached to the HUMAN, read at both list time and tool-call time.
A key's object_permission scopes the credential and a team's scopes the group; the user's own
scopes the person, so it must cap every key they hold and every tool those keys may invoke.
"""
def _auth(self, user_id: str = "human-1", **kwargs) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="sk-test", user_id=user_id, **kwargs)
def _perm(self, *, servers=None, access_groups=None, tool_permissions=None) -> LiteLLM_ObjectPermissionTable:
return LiteLLM_ObjectPermissionTable(
object_permission_id="perm-human-1",
mcp_servers=servers if servers is not None else [],
mcp_access_groups=access_groups if access_groups is not None else [],
mcp_tool_permissions=tool_permissions,
)
@contextlib.contextmanager
def _entitled(self, perm):
"""Patch the human's entitlement lookup. ``perm`` may be a permission row, None, or an
exception instance to raise (an entitlement that cannot be resolved)."""
side_effect = perm if isinstance(perm, Exception) else None
with patch.object(
MCPRequestHandler,
"_get_user_object_permission",
new_callable=AsyncMock,
return_value=None if side_effect else perm,
side_effect=side_effect,
) as patched:
yield patched
@contextlib.contextmanager
def _key_and_team_servers(self, key_servers, team_servers):
with (
patch.object(
MCPRequestHandler,
"_get_allowed_mcp_servers_for_key",
new_callable=AsyncMock,
return_value=key_servers,
),
patch.object(
MCPRequestHandler,
"_get_allowed_mcp_servers_for_team",
new_callable=AsyncMock,
return_value=team_servers,
),
patch.object(
MCPRequestHandler,
"_get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
),
patch.object(
MCPRequestHandler,
"_get_key_access_group_mcp_server_extras",
new_callable=AsyncMock,
return_value=[],
),
):
yield
async def test_entitlement_caps_the_servers_the_key_reaches(self):
"""The key grants two servers; the human is entitled to one, so only that one resolves."""
with self._key_and_team_servers(["srv-a", "srv-b"], []):
with self._entitled(self._perm(servers=["srv-a"])):
result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth())
assert result == ["srv-a"]
async def test_entitlement_never_widens_the_key(self):
"""A human entitled to a server their key does not grant still cannot reach it: the level is a
ceiling, so it intersects rather than unions."""
with self._key_and_team_servers(["srv-a"], []):
with self._entitled(self._perm(servers=["srv-a", "srv-elsewhere"])):
result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth())
assert result == ["srv-a"]
async def test_no_entitlement_places_no_ceiling(self):
"""A human with no entitlement row leaves the key/team result untouched."""
with self._key_and_team_servers(["srv-a", "srv-b"], []):
with self._entitled(None):
result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth())
assert sorted(result) == ["srv-a", "srv-b"]
async def test_unresolvable_entitlement_denies_every_server(self):
"""A KNOWN entitlement whose contents cannot be read must deny, not fall back to the key's
wider scope."""
with self._key_and_team_servers(["srv-a", "srv-b"], []):
with self._entitled(ValueError("permission row unreadable")):
result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth())
assert result == []
async def test_entitlement_caps_the_tools_the_key_reaches(self):
"""Tool-level: the key allows three tools on the server, the human is entitled to one."""
key_perm = self._perm(tool_permissions={"srv-a": ["read", "write", "delete"]})
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=key_perm):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})):
result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth())
assert result == ["read"]
async def test_entitlement_alone_restricts_tools_on_an_otherwise_unrestricted_key(self):
"""An unrestricted key (no tool permissions of its own) is still bound by the human's tools."""
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})):
result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth())
assert result == ["read"]
async def test_entitlement_on_another_server_does_not_restrict_this_one(self):
"""Tool grants are per server: naming tools on srv-b places no bound on srv-a."""
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(self._perm(tool_permissions={"srv-b": ["read"]})):
result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth())
assert result is None
async def test_unresolvable_entitlement_denies_every_tool(self):
"""Fail closed on the tool axis too. The caller's own except-handler treats a raise as
allow-all for key auth, so the ceiling must return the empty allowlist itself."""
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(ValueError("permission row unreadable")):
result = await MCPRequestHandler.get_allowed_tools_for_server("srv-a", self._auth())
assert result == []
async def test_tool_call_is_rejected_at_call_time(self):
"""The end-to-end contract: a tool the human is not entitled to is refused when INVOKED, not
merely hidden from the advertised list."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
server = MCPServer(
server_id="srv-a",
name="srv-a",
server_name="srv-a",
url="https://srv-a.example.com",
transport=MCPTransport.http,
)
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})):
await global_mcp_server_manager.check_tool_permission_for_key_team(
tool_name="read", server=server, user_api_key_auth=self._auth()
)
with pytest.raises(HTTPException) as exc:
await global_mcp_server_manager.check_tool_permission_for_key_team(
tool_name="delete", server=server, user_api_key_auth=self._auth()
)
assert exc.value.status_code == 403
async def test_keyless_admitted_source_is_not_capped_by_the_user_level(self):
"""A gateway-admitted human resolves as a UNION over their own grants plus their teams', and
their own grants ARE the user source there. Re-applying them as a ceiling per source would
make one team's narrower scope silently bound another's, so the level is skipped."""
with self._key_and_team_servers(["srv-a", "srv-b"], []):
with self._entitled(self._perm(servers=["srv-a"])) as lookup:
result = await MCPRequestHandler.get_allowed_mcp_servers(self._auth(), keyless_source=True)
assert sorted(result) == ["srv-a", "srv-b"]
lookup.assert_not_awaited()
async def test_keyless_admitted_source_tools_are_not_capped_by_the_user_level(self):
with patch.object(MCPRequestHandler, "_get_key_object_permission", return_value=None):
with patch.object(
MCPRequestHandler, "_get_team_object_permission", new_callable=AsyncMock, return_value=None
):
with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})) as lookup:
result = await MCPRequestHandler.get_allowed_tools_for_server(
"srv-a", self._auth(), keyless_source=True
)
assert result is None
lookup.assert_not_awaited()
async def test_servers_named_only_under_tool_permissions_are_entitled(self):
"""Granting one tool on a server entitles the human to that server, so an admin never has to
name it twice."""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPTransport
from litellm.types.mcp_server.mcp_server_manager import MCPServer
global_mcp_server_manager.registry["srv-a"] = MCPServer(
server_id="srv-a",
name="srv-a",
server_name="srv-a",
url="https://srv-a.example.com",
transport=MCPTransport.http,
)
try:
with self._entitled(self._perm(tool_permissions={"srv-a": ["read"]})):
with patch.object(
MCPRequestHandler,
"_get_mcp_servers_from_access_groups",
new_callable=AsyncMock,
return_value=[],
):
result = await MCPRequestHandler._get_allowed_mcp_servers_for_user(self._auth())
finally:
global_mcp_server_manager.registry.pop("srv-a", None)
assert result == ["srv-a"]
async def test_places_ceiling_is_true_when_unresolvable(self):
"""``_user_places_mcp_ceiling`` gates the admin shortcut that hands over the whole registry, so
an entitlement it cannot resolve must still count as a ceiling."""
with self._entitled(ValueError("boom")):
assert await MCPRequestHandler._user_places_mcp_ceiling(self._auth()) is True
with self._entitled(None):
assert await MCPRequestHandler._user_places_mcp_ceiling(self._auth()) is False
@pytest.mark.asyncio
class TestGetUserObjectPermission:
"""Resolution of the ``user_id -> object_permission_id -> grants`` chain."""
def _prisma_with_user(self, user_row):
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
return prisma_client
async def test_resolves_through_the_shared_permission_cache(self):
from litellm.caching.dual_cache import DualCache
user_row = MagicMock()
user_row.object_permission_id = "perm-1"
prisma_client = self._prisma_with_user(user_row)
auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-shared")
expected = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
return_value=expected,
) as mock_get_perm,
):
assert await MCPRequestHandler._get_user_object_permission(auth) is expected
assert mock_get_perm.await_args.kwargs["object_permission_id"] == "perm-1"
# The user_id -> object_permission_id link is cached, so the user row is read once.
prisma_client.db.litellm_usertable.find_unique.reset_mock()
await MCPRequestHandler._get_user_object_permission(auth)
prisma_client.db.litellm_usertable.find_unique.assert_not_called()
async def test_caches_a_sentinel_for_a_human_with_no_entitlement(self):
"""A human without an entitlement is the common case and must cost no DB read per request."""
from litellm.caching.dual_cache import DualCache
user_row = MagicMock()
user_row.object_permission_id = None
prisma_client = self._prisma_with_user(user_row)
auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-no-perm")
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch("litellm.proxy.auth.auth_checks.get_object_permission", new_callable=AsyncMock) as mock_get_perm,
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
assert await MCPRequestHandler._get_user_object_permission(auth) is None
mock_get_perm.assert_not_awaited()
prisma_client.db.litellm_usertable.find_unique.assert_awaited_once()
async def test_missing_user_row_places_no_ceiling(self):
"""Whether this human is entitled at all is unknown when their row is absent, which is the
state before the level existed, so it must not deny."""
from litellm.caching.dual_cache import DualCache
prisma_client = self._prisma_with_user(None)
auth = UserAPIKeyAuth(api_key="sk-test", user_id="ghost")
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
async def test_unreadable_user_row_places_no_ceiling(self):
from litellm.caching.dual_cache import DualCache
prisma_client = MagicMock()
prisma_client.db.litellm_usertable.find_unique = AsyncMock(side_effect=Exception("db down"))
auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-db-down")
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
):
assert await MCPRequestHandler._get_user_object_permission(auth) is None
async def test_named_but_unreadable_permission_raises(self):
"""A KNOWN entitlement with unknown contents is indeterminate: it must surface so the callers
can deny rather than serve the wider key scope."""
from litellm.caching.dual_cache import DualCache
user_row = MagicMock()
user_row.object_permission_id = "perm-gone"
prisma_client = self._prisma_with_user(user_row)
auth = UserAPIKeyAuth(api_key="sk-test", user_id="human-dangling")
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma_client),
patch("litellm.proxy.proxy_server.user_api_key_cache", DualCache()),
patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
patch(
"litellm.proxy.auth.auth_checks.get_object_permission",
new_callable=AsyncMock,
return_value=None,
),
):
with pytest.raises(ValueError):
await MCPRequestHandler._get_user_object_permission(auth)
async def test_no_user_id_places_no_ceiling(self):
assert await MCPRequestHandler._get_user_object_permission(UserAPIKeyAuth(api_key="sk-test")) is None
assert await MCPRequestHandler._get_user_object_permission(None) is None

304
tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py
View file

@ -1,11 +1,14 @@
"""Tests for the MCP guardrail translation handler."""
import pytest
from mcp.types import CallToolResult, ImageContent, TextContent
from litellm.exceptions import BlockedPiiEntityError
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import (
MCPGuardrailTranslationHandler,
)
from litellm.types.utils import GenericGuardrailAPIInputs
class MockGuardrail(CustomGuardrail):
@ -80,3 +83,304 @@ async def test_process_input_messages_handles_minimal_data():
tools = guardrail.last_inputs.get("tools", [])
assert len(tools) == 1
assert tools[0]["function"]["name"] == "simple_tool"
class MaskingGuardrail(CustomGuardrail):
"""Guardrail that rewrites every scanned text, recording what it saw."""
def __init__(self, masked_texts=None, raises=None):
super().__init__(guardrail_name="masking-mcp-guardrail")
self.masked_texts = masked_texts
self.raises = raises
self.call_count = 0
self.last_inputs = None
self.last_input_type = None
self.last_request_data = None
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.call_count += 1
self.last_inputs = inputs
self.last_input_type = input_type
self.last_request_data = request_data
if self.raises is not None:
raise self.raises
if self.masked_texts is None:
return inputs
return GenericGuardrailAPIInputs(texts=list(self.masked_texts))
@pytest.mark.asyncio
async def test_process_output_response_masks_text_content():
"""Masked text returned by the guardrail must land in the tool result."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(masked_texts=["email <EMAIL_ADDRESS>", "call <PHONE_NUMBER>"])
result = CallToolResult(
content=[
TextContent(type="text", text="email jane@example.com"),
TextContent(type="text", text="call 415-555-0132"),
],
isError=False,
)
returned = await handler.process_output_response(
response=result,
guardrail_to_apply=guardrail,
request_data={"mcp_tool_name": "echo"},
)
assert guardrail.call_count == 1
assert guardrail.last_input_type == "response"
assert guardrail.last_inputs["texts"] == ["email jane@example.com", "call 415-555-0132"]
assert [item.text for item in returned.content] == ["email <EMAIL_ADDRESS>", "call <PHONE_NUMBER>"]
assert [item.text for item in result.content] == ["email <EMAIL_ADDRESS>", "call <PHONE_NUMBER>"]
@pytest.mark.asyncio
async def test_process_output_response_masks_dict_shaped_result():
"""A dict-shaped tool result (REST/JSON-RPC payload) must be masked too."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(masked_texts=["<EMAIL_ADDRESS>"])
result = {"content": [{"type": "text", "text": "jane@example.com"}], "isError": False}
returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
assert returned["content"][0]["text"] == "<EMAIL_ADDRESS>"
assert returned["content"][0]["type"] == "text"
@pytest.mark.asyncio
async def test_process_output_response_propagates_block():
"""A guardrail rejecting the tool result must not be swallowed."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(
raises=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="masking-mcp-guardrail")
)
result = CallToolResult(content=[TextContent(type="text", text="jane@example.com")], isError=False)
with pytest.raises(BlockedPiiEntityError):
await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
@pytest.mark.asyncio
async def test_process_output_response_skips_non_text_content():
"""A result carrying no text content must not be sent to the guardrail."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(masked_texts=["should not be used"])
result = CallToolResult(
content=[ImageContent(type="image", data="aGk=", mimeType="image/png")],
isError=False,
)
returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
assert guardrail.call_count == 0
assert returned is result
@pytest.mark.asyncio
async def test_process_output_response_handles_result_without_content():
"""An unexpected result shape must be passed through, not crash the tool call."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(masked_texts=["should not be used"])
returned = await handler.process_output_response(response={"error": "boom"}, guardrail_to_apply=guardrail)
assert guardrail.call_count == 0
assert returned == {"error": "boom"}
@pytest.mark.asyncio
async def test_process_output_response_leaves_result_unmasked_on_text_count_mismatch():
"""A guardrail returning the wrong number of texts must not shuffle content."""
handler = MCPGuardrailTranslationHandler()
guardrail = MaskingGuardrail(masked_texts=["<EMAIL_ADDRESS>"])
result = CallToolResult(
content=[
TextContent(type="text", text="jane@example.com"),
TextContent(type="text", text="415-555-0132"),
],
isError=False,
)
returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail)
assert [item.text for item in returned.content] == ["jane@example.com", "415-555-0132"]
class SubstitutingGuardrail(CustomGuardrail):
"""Masks one substring wherever it appears, across however many texts it is given."""
def __init__(self, needle: str, replacement: str):
super().__init__(guardrail_name="substituting-mcp-guardrail")
self.needle = needle
self.replacement = replacement
self.seen_texts: list = []
async def apply_guardrail(self, inputs, request_data, input_type, **kwargs):
self.seen_texts = list(inputs.get("texts") or [])
return GenericGuardrailAPIInputs(
texts=[text.replace(self.needle, self.replacement) for text in self.seen_texts]
)
@pytest.mark.asyncio
async def test_structured_content_is_masked_alongside_content():
"""structuredContent goes to the client too, so it must be masked, not just content."""
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="email jane@example.com")],
structuredContent={"contact": {"email": "jane@example.com"}, "balance": 42.0},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
assert returned.structuredContent == {"contact": {"email": "<EMAIL_ADDRESS>"}, "balance": 42.0}
@pytest.mark.asyncio
async def test_value_present_only_in_structured_content_is_masked():
"""The gap this closes: a sensitive value that never appears in the text content.
Scanning only content would hand it to the guardrail never, so it would reach
the client unscanned behind a result that looks inspected.
"""
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structuredContent={"records": [{"email": "jane@example.com"}]},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert "jane@example.com" in guardrail.seen_texts
assert returned.structuredContent == {"records": [{"email": "<EMAIL_ADDRESS>"}]}
assert returned.content[0].text == "lookup complete"
@pytest.mark.asyncio
async def test_structured_content_without_a_match_is_untouched():
"""Unrelated structured data keeps its values and its types."""
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structuredContent={"record_id": "C-1001", "balance": 42.0, "active": True, "note": None},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "active": True, "note": None}
@pytest.mark.asyncio
async def test_structured_content_nested_too_deeply_is_blocked():
"""Too deep to walk must block rather than pass the deeper values unscanned."""
from fastapi import HTTPException
from litellm.proxy._experimental.mcp_server.utils import MAX_STRUCTURED_CONTENT_SCAN_DEPTH
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
nested: dict = {"leaf": "jane@example.com"}
for _ in range(MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1):
nested = {"next": nested}
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structuredContent=nested,
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert exc_info.value.status_code == 400
def test_too_deep_json_returns_a_sentinel_rather_than_raising():
"""The too-deep signal must be a return value, not a custom exception.
mcp_server/utils.py is reloaded by tests that override its environment-backed
constants, which gives any exception class defined there a fresh identity and
lets it escape a caller's except clause; under xdist that surfaced as a failure
in an unrelated shard. A sentinel has no identity to lose. Asserted directly on
the helper so this pins the contract without reloading the module and leaking
that reload into other tests.
"""
from litellm.proxy._experimental.mcp_server.utils import (
MAX_STRUCTURED_CONTENT_SCAN_DEPTH,
json_string_leaves,
)
nested: dict = {"leaf": "jane@example.com"}
for _ in range(MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1):
nested = {"next": nested}
assert json_string_leaves(nested) is None
assert json_string_leaves({"a": "b"}) == ((("a",), "b"),)
@pytest.mark.asyncio
async def test_sensitive_structured_content_key_is_blocked():
"""A dict key is client-visible but not rewritable, so a match must block.
Maps keyed by an identifier are a common API shape, and renaming the key would
change the payload contract rather than redact a value; the content filter takes
the same position on MCP tool call arguments.
"""
from fastapi import HTTPException
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structuredContent={"jane@example.com": {"balance": 42.0}},
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert exc_info.value.status_code == 400
assert "non-rewritable" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_sensitive_structured_content_numeric_value_is_blocked():
"""A numeric value cannot be masked in place either, so a match must block."""
from fastapi import HTTPException
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("4155550199", "<PHONE_NUMBER>")
response = CallToolResult(
content=[TextContent(type="text", text="lookup complete")],
structuredContent={"phone": 4155550199},
isError=False,
)
with pytest.raises(HTTPException) as exc_info:
await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_clean_structured_content_keys_do_not_block():
"""Ordinary keys and numbers must pass through untouched."""
handler = MCPGuardrailTranslationHandler()
guardrail = SubstitutingGuardrail("jane@example.com", "<EMAIL_ADDRESS>")
response = CallToolResult(
content=[TextContent(type="text", text="email jane@example.com")],
structuredContent={"record_id": "C-1001", "balance": 42.0, "count": 3},
isError=False,
)
returned = await handler.process_output_response(response=response, guardrail_to_apply=guardrail)
assert returned.content[0].text == "email <EMAIL_ADDRESS>"
assert returned.structuredContent == {"record_id": "C-1001", "balance": 42.0, "count": 3}

207
tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py
View file

@ -1,6 +1,7 @@
"""Tests for the aggregate gateway DCR flow (register, authorize, complete, token)."""
import hashlib
import html
import json
from base64 import urlsafe_b64encode
from datetime import datetime, timedelta, timezone
@ -16,7 +17,10 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import (
GATEWAY_AUTH_CODE_PREFIX,
GATEWAY_AUTH_CODE_TTL_SECONDS,
GATEWAY_DCR_CLIENT_ID_PREFIX,
MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS,
_AUTH_CODE_DEBUG_KEY,
_GatewayAuthCode,
_open_sealed,
_seal,
aggregate_authorize,
aggregate_token,
@ -588,3 +592,206 @@ async def test_single_use_guard_fails_closed_when_redis_errors():
guard = _SingleUseGuard(cache)
assert await guard.claim("jti-fault", 60) is False # fail closed, not a fallback count of 1
LOOPBACK_REDIRECT_URI = "http://localhost:3118/callback"
async def _complete(redirect_uri: str, delivery, cookies=None, handle=None, session_user_id="u1"):
client_id = (await _register([redirect_uri]))["client_id"]
if cookies is None:
handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1", redirect_uri=redirect_uri))
response = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id=session_user_id,
cache=DualCache(),
delivery=delivery,
)
return client_id, response
def _callback_url_from_page(response) -> str:
import html as html_lib
import re
match = re.search(r'value="([^"]+)"', response.body.decode())
assert match is not None
return html_lib.unescape(match.group(1))
@pytest.mark.asyncio
async def test_manual_delivery_renders_pasteable_callback_url_for_loopback_client():
"""The LIT-4863 headless path: a loopback client on another machine gets the callback
URL on a page instead of a dead 303, and the code on that page is a full-fidelity
authorization code (PKCE-bound, single-use, redeemable at /token)."""
client_id, response = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
assert response.status_code == 200
assert response.headers["content-type"].startswith("text/html")
assert response.headers["cache-control"] == "no-store"
assert f"{CONNECT_FLOW_COOKIE_PREFIX}" in response.headers["set-cookie"]
callback_url = _callback_url_from_page(response)
parsed = urlparse(callback_url)
assert f"{parsed.scheme}://{parsed.netloc}{parsed.path}" == LOOPBACK_REDIRECT_URI
params = parse_qs(parsed.query)
assert params["state"] == ["client-state-123"]
code = params["code"][0]
assert code.startswith(GATEWAY_AUTH_CODE_PREFIX)
cache = DualCache()
token_response = await aggregate_token(
request=_request("/token", method="POST"),
grant_type="authorization_code",
code=code,
redirect_uri=LOOPBACK_REDIRECT_URI,
client_id=client_id,
code_verifier=CODE_VERIFIER,
refresh_token=None,
master_key=MASTER_KEY,
reload_user=_reload_user_active,
cache=cache,
)
assert token_response.status_code == 200
replay = await aggregate_token(
request=_request("/token", method="POST"),
grant_type="authorization_code",
code=code,
redirect_uri=LOOPBACK_REDIRECT_URI,
client_id=client_id,
code_verifier=CODE_VERIFIER,
refresh_token=None,
master_key=MASTER_KEY,
reload_user=_reload_user_active,
cache=cache,
)
assert json.loads(replay.body)["error"] == "invalid_grant"
@pytest.mark.asyncio
async def test_manual_delivery_code_gets_the_longer_ttl_and_redirect_code_does_not():
_, manual = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
manual_code = parse_qs(urlparse(_callback_url_from_page(manual)).query)["code"][0]
opened_manual = _open_sealed(manual_code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY)
assert opened_manual is not None
assert opened_manual.exp - opened_manual.iat == MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS
_, redirected = await _complete(LOOPBACK_REDIRECT_URI, delivery=None)
redirect_code = parse_qs(urlparse(redirected.headers["location"]).query)["code"][0]
opened_redirect = _open_sealed(redirect_code, GATEWAY_AUTH_CODE_PREFIX, _GatewayAuthCode, _AUTH_CODE_DEBUG_KEY)
assert opened_redirect is not None
assert opened_redirect.exp - opened_redirect.iat == GATEWAY_AUTH_CODE_TTL_SECONDS
@pytest.mark.asyncio
@pytest.mark.parametrize("delivery", [None, "redirect"])
async def test_loopback_client_still_redirects_when_manual_not_requested(delivery):
_, response = await _complete(LOOPBACK_REDIRECT_URI, delivery=delivery)
assert response.status_code == 303
assert response.headers["location"].startswith(LOOPBACK_REDIRECT_URI)
@pytest.mark.asyncio
async def test_manual_delivery_is_ignored_for_routable_redirect_uri():
"""A routable redirect URI works from any browser by construction, so manual is a
no-op there and the flow keeps its normal shape."""
_, response = await _complete(REDIRECT_URI, delivery="manual")
assert response.status_code == 303
assert response.headers["location"].startswith(REDIRECT_URI)
@pytest.mark.asyncio
async def test_unknown_delivery_value_is_rejected_before_the_flow_is_consumed():
"""A typo'd delivery must not burn the single-use flow: the user fixes the form and
finishes normally."""
client_id = (await _register([LOOPBACK_REDIRECT_URI]))["client_id"]
handle, cookies = _flow_cookie_from(_authorize(client_id, session_user_id="u1", redirect_uri=LOOPBACK_REDIRECT_URI))
rejected = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=DualCache(),
delivery="carrier-pigeon",
)
assert rejected.status_code == 400
assert json.loads(rejected.body)["error"] == "invalid_request"
retried = await complete_connect_flow(
request=_request("/authorize/complete", cookies=cookies, method="POST"),
flow_handle=handle,
session_user_id="u1",
cache=DualCache(),
delivery="manual",
)
assert retried.status_code == 200
@pytest.mark.asyncio
async def test_manual_delivery_page_escapes_client_influenced_values():
"""redirect_uri (and everything else on the page) is client-registered input; a quote
or tag in its path must render inert."""
hostile_uri = 'http://127.0.0.1:9/cb"><script>alert(1)</script>'
_, response = await _complete(hostile_uri, delivery="manual")
assert response.status_code == 200
body = response.body.decode()
assert "<script>alert(1)</script>" not in body
assert "&lt;script&gt;" in body
class _TtlRecordingCache(DualCache):
"""Captures the TTL of every single-use claim recorded through the in-memory arm."""
def __init__(self):
super().__init__()
self.claim_ttls: dict = {}
async def async_increment_cache(self, key, value, ttl=None, **kwargs):
self.claim_ttls[key] = ttl
return await super().async_increment_cache(key, value, ttl=ttl, **kwargs)
@pytest.mark.asyncio
async def test_used_code_marker_outlives_the_manually_delivered_code():
"""Veria review finding on the LIT-4863 change: a manual code lives 300s, but the
used-code marker was retained for the 120s redirect lifetime plus buffer, so a client
could redeem, wait out the marker, and redeem the still-valid code again. The marker's
TTL must cover the code's own remaining lifetime plus the claim buffer."""
client_id, response = await _complete(LOOPBACK_REDIRECT_URI, delivery="manual")
code = parse_qs(urlparse(_callback_url_from_page(response)).query)["code"][0]
cache = _TtlRecordingCache()
token_response = await aggregate_token(
request=_request("/token", method="POST"),
grant_type="authorization_code",
code=code,
redirect_uri=LOOPBACK_REDIRECT_URI,
client_id=client_id,
code_verifier=CODE_VERIFIER,
refresh_token=None,
master_key=MASTER_KEY,
reload_user=_reload_user_active,
cache=cache,
)
assert token_response.status_code == 200
marker_ttls = [ttl for key, ttl in cache.claim_ttls.items() if key.startswith("mcp_gateway_dcr_code_used:")]
assert len(marker_ttls) == 1
assert marker_ttls[0] >= MANUAL_DELIVERY_AUTH_CODE_TTL_SECONDS
@pytest.mark.asyncio
@pytest.mark.parametrize("redirect_uri", [LOOPBACK_REDIRECT_URI, "http://127.0.0.1:9/cb$(whoami)&calc& rem x"])
async def test_manual_delivery_page_renders_the_url_as_data_never_as_a_shell_command(redirect_uri):
"""Two review rounds proved no single command string is safe across POSIX shells,
cmd.exe, and PowerShell (single quotes are not quoting in cmd.exe; percent expands
there even inside double quotes), so the page must render the callback URL as data
only and never as a ready-to-paste command."""
_, response = await _complete(redirect_uri, delivery="manual")
assert response.status_code == 200
body = response.body.decode()
assert "<code>" not in body
assert 'curl "' not in body
assert "curl '" not in body
assert 'value="' in body

103
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -7130,6 +7130,14 @@ def _mock_mcp_logging_obj() -> MagicMock:
return logging_obj
def _mock_mcp_proxy_logging() -> MagicMock:
"""ProxyLogging stand-in whose post_mcp_call_hook passes the result through."""
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock.post_mcp_call_hook = AsyncMock(side_effect=lambda response, **_: response)
return proxy_logging_mock
def test_extract_mcp_tool_result_error_message():
from litellm.proxy._experimental.mcp_server.utils import (
extract_mcp_tool_result_error_message,
@ -7160,8 +7168,7 @@ async def test_fire_mcp_tool_call_logging_iserror_logs_failure():
from litellm.proxy._experimental.mcp_server.exceptions import MCPToolResultError
logging_obj = _mock_mcp_logging_obj()
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock = _mock_mcp_proxy_logging()
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
@ -7199,8 +7206,7 @@ async def test_fire_mcp_tool_call_logging_success_path_unchanged():
)
logging_obj = _mock_mcp_logging_obj()
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock = _mock_mcp_proxy_logging()
result = _call_tool_result(False, "all good")
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
@ -7229,8 +7235,7 @@ async def test_fire_mcp_tool_call_logging_iserror_without_auth_skips_failure_hoo
)
logging_obj = _mock_mcp_logging_obj()
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock = _mock_mcp_proxy_logging()
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
await _fire_mcp_tool_call_logging(
@ -7256,8 +7261,7 @@ async def test_fire_mcp_tool_call_logging_strips_credentials_from_failure_hook()
)
logging_obj = _mock_mcp_logging_obj()
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock = _mock_mcp_proxy_logging()
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
request_data = {
"name": "explode",
@ -7528,8 +7532,7 @@ async def test_call_mcp_tool_skips_failure_hook_for_upstream_auth_error():
transport=MCPTransport.http,
mcp_info={"server_name": "test_server"},
)
proxy_logging_mock = MagicMock()
proxy_logging_mock.post_call_failure_hook = AsyncMock()
proxy_logging_mock = _mock_mcp_proxy_logging()
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
with (
@ -7841,3 +7844,83 @@ class TestPreemptive401ModeAware:
await self._run(delegate, None, has_stored_token=False)
assert exc.value.status_code == 401
await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
@pytest.mark.asyncio
async def test_post_mcp_call_guardrails_return_the_rewritten_result():
"""The result a post_mcp_call guardrail rewrote must be what the caller sends back."""
from litellm.proxy._experimental.mcp_server.server import (
_run_post_mcp_call_guardrails,
)
logging_obj = _mock_mcp_logging_obj()
raw_result = _call_tool_result(False, "jane@example.com")
masked_result = _call_tool_result(False, "<EMAIL_ADDRESS>")
proxy_logging_mock = _mock_mcp_proxy_logging()
proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result)
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
returned = await _run_post_mcp_call_guardrails(
result=raw_result,
litellm_logging_obj=logging_obj,
user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"),
request_data={},
)
assert returned is masked_result
hook_kwargs = proxy_logging_mock.post_mcp_call_hook.await_args.kwargs
assert hook_kwargs["response"] is raw_result
assert hook_kwargs["request_data"] is logging_obj.model_call_details
@pytest.mark.asyncio
async def test_post_mcp_call_guardrails_run_without_a_logging_object():
"""Enforcement must not depend on logging being configured.
A tool call dispatched without a litellm_logging_obj (tool search, and any
caller that omits it) would otherwise skip the guardrail entirely and return
the unscanned tool output to the client.
"""
from litellm.proxy._experimental.mcp_server.server import (
_run_post_mcp_call_guardrails,
)
raw_result = _call_tool_result(False, "jane@example.com")
masked_result = _call_tool_result(False, "<EMAIL_ADDRESS>")
proxy_logging_mock = _mock_mcp_proxy_logging()
proxy_logging_mock.post_mcp_call_hook = AsyncMock(return_value=masked_result)
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
returned = await _run_post_mcp_call_guardrails(
result=raw_result,
litellm_logging_obj=None,
user_api_key_auth=UserAPIKeyAuth(api_key="test-key", user_id="test-user"),
request_data={"name": "fetch_record"},
)
assert returned is masked_result
proxy_logging_mock.post_mcp_call_hook.assert_awaited_once()
assert proxy_logging_mock.post_mcp_call_hook.await_args.kwargs["request_data"] == {"name": "fetch_record"}
@pytest.mark.asyncio
async def test_post_mcp_call_guardrails_propagate_a_block():
"""A post_mcp_call guardrail rejection must propagate instead of returning the result."""
from litellm.exceptions import BlockedPiiEntityError
from litellm.proxy._experimental.mcp_server.server import (
_run_post_mcp_call_guardrails,
)
proxy_logging_mock = _mock_mcp_proxy_logging()
proxy_logging_mock.post_mcp_call_hook = AsyncMock(
side_effect=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="presidio-mcp")
)
with patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_mock):
with pytest.raises(BlockedPiiEntityError):
await _run_post_mcp_call_guardrails(
result=_call_tool_result(False, "jane@example.com"),
litellm_logging_obj=_mock_mcp_logging_obj(),
user_api_key_auth=None,
request_data={},
)

83
tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py
View file

@ -1710,6 +1710,89 @@ class TestCallToolRestAPI:
assert captured["allowed_mcp_servers"] == [stub_server]
fire_logging.assert_awaited_once()
async def test_returns_guardrail_rewritten_tool_result(self, monkeypatch):
"""A post_mcp_call guardrail rewrite of the tool result must reach the REST caller,
not the raw result the upstream server returned."""
async def fake_contexts(user_api_key_auth):
return [user_api_key_auth]
async def fake_get_allowed_mcp_servers(*args, **kwargs):
return ["server-1"]
class StubServer:
server_id = "server-1"
alias = "server-1"
server_name = "server-1"
name = "stub"
allowed_tools = None
mcp_info = {"server_name": "stub"}
available_on_public_internet = True
auth_type = None
stub_server = StubServer()
async def fake_add_litellm_data_to_request(**kwargs):
return kwargs.get("data", {})
async def fake_execute_mcp_tool(**kwargs):
return {"content": [{"type": "text", "text": "jane@example.com"}]}
monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_allowed_mcp_servers",
fake_get_allowed_mcp_servers,
raising=False,
)
monkeypatch.setattr(
rest_endpoints.global_mcp_server_manager,
"get_mcp_server_by_id",
lambda server_id: stub_server if server_id == "server-1" else None,
raising=False,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.add_litellm_data_to_request",
fake_add_litellm_data_to_request,
raising=False,
)
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", {}, raising=False)
monkeypatch.setattr(rest_endpoints, "execute_mcp_tool", fake_execute_mcp_tool, raising=False)
masked_result = {"content": [{"type": "text", "text": "<EMAIL_ADDRESS>"}]}
monkeypatch.setattr(
rest_endpoints,
"_fire_mcp_tool_call_logging",
AsyncMock(return_value=masked_result),
raising=False,
)
request = _build_request(
path="/mcp-rest/tools/call",
method="POST",
json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}},
)
result = await rest_endpoints.call_tool_rest_api(request, user_api_key_dict=UserAPIKeyAuth())
assert result == masked_result
async def test_success_logging_guardrail_rejection_propagates(self, monkeypatch):
"""A guardrail rejecting the tool result must not be swallowed as a logging failure,
otherwise the unguarded result would still be returned to the caller."""
from litellm.exceptions import BlockedPiiEntityError
fire_logging = AsyncMock(
side_effect=BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="presidio-mcp")
)
monkeypatch.setattr(rest_endpoints, "_fire_mcp_tool_call_logging", fire_logging, raising=False)
with pytest.raises(BlockedPiiEntityError):
await rest_endpoints._safe_fire_mcp_tool_call_logging(
object(), {"result": "ok"}, datetime.now(), datetime.now()
)
fire_logging.assert_awaited_once()
@pytest.mark.parametrize("upstream_status", [401, 403])
async def test_call_tool_rest_relays_upstream_auth_failure(self, monkeypatch, upstream_status):
"""A pass-through call that hits an upstream 401/403 (surfaced by the manager as

166
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -4835,66 +4835,26 @@ async def test_common_checks_personal_user_budget_blocks_in_gather():
@pytest.mark.asyncio
async def test_user_budget_enforced_on_team_key():
"""User budget must be enforced even when the key belongs to a team.
async def test_common_checks_personal_user_budget_skipped_for_team_key():
"""A user's personal max_budget does not apply to a team-scoped key.
Previously _user_max_budget_check skipped enforcement for team keys,
letting a user with a $100 personal budget spend unlimited through a
team key. This regression test ensures that is no longer the case.
Team keys are governed by the team (and team-member) budgets only; the key
owner's personal budget is deliberately out of scope. This asserts the read
path lets a team key through even when the user is far over their personal
budget, and fails if personal enforcement is reintroduced for team keys.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0)
team = LiteLLM_TeamTable(team_id="t1", spend=0.0, max_budget=1000.0)
token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == "spend:user:u1" else 0.0
async def _no_membership(*a, **kw):
return None
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership):
with pytest.raises(litellm.BudgetExceededError) as over:
await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=team,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
valid_token=token,
request=MagicMock(spec=Request),
)
assert "User=u1" in str(over.value)
@pytest.mark.asyncio
async def test_skip_user_budget_on_team_key_flag_restores_old_behavior():
"""Setting skip_user_budget_on_team_key=True skips user budget for team keys.
This is the opt-in escape hatch that restores the legacy behavior where
user budgets were not enforced when the key belonged to a team.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=100.0)
team = LiteLLM_TeamTable(team_id="t1", max_budget=2100.0)
token = UserAPIKeyAuth(token="k1", user_id="u1", team_id="t1")
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == "spend:user:u1" else 0.0
async def _no_membership(*a, **kw):
async def _no_membership(*args, **kwargs):
return None
with patch("litellm.proxy.proxy_server.prisma_client", None), patch(
@ -4906,7 +4866,7 @@ async def test_skip_user_budget_on_team_key_flag_restores_old_behavior():
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={"skip_user_budget_on_team_key": True},
general_settings={},
route="/chat/completions",
llm_router=None,
proxy_logging_obj=MagicMock(),
@ -4916,6 +4876,114 @@ async def test_skip_user_budget_on_team_key_flag_restores_old_behavior():
assert result is True
@pytest.mark.parametrize(
"scope, route, expect_blocked",
[
("user", "/chat/completions", True),
("user", "/key/list", False),
("team", "/chat/completions", True),
("team", "/key/list", False),
("org", "/chat/completions", True),
("org", "/key/list", False),
],
)
@pytest.mark.asyncio
async def test_budget_checks_only_run_on_llm_api_routes(scope, route, expect_blocked):
"""Budgets cap spend, so they must only gate routes that can spend.
Enforcing them on management routes locked an over-budget caller out of the
Admin UI, which authenticates with a normal virtual key, leaving no way to
reach the page that raises the limit.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
over_budget_counter = {"user": "spend:user:u1", "team": "spend:team:t1", "org": "spend:org:o1"}[scope]
async def _spend_by_counter(counter_key, fallback_spend, max_budget=None, **kwargs):
return 999.0 if counter_key == over_budget_counter else 0.0
async def _no_membership(*a, **kw):
return None
org_table = MagicMock()
org_table.spend = 999.0
org_table.litellm_budget_table = MagicMock()
org_table.litellm_budget_table.max_budget = 10.0
async def _get_org(*a, **kw):
return org_table
user = LiteLLM_UserTable(user_id="u1", spend=0.0, max_budget=10.0 if scope == "user" else None)
team = LiteLLM_TeamTable(team_id="t1", max_budget=10.0) if scope == "team" else None
token = UserAPIKeyAuth(
token="k1",
user_id="u1",
team_id="t1" if scope == "team" else None,
org_id="o1" if scope == "org" else None,
)
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
async def _run():
return await common_checks(
request_body={"messages": [{"role": "user", "content": "hi"}]},
team_object=team,
user_object=user,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=proxy_logging_obj,
valid_token=token,
request=MagicMock(spec=Request),
)
with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch(
"litellm.proxy.proxy_server.get_current_spend", _spend_by_counter
), patch("litellm.proxy.auth.auth_checks.get_team_membership", _no_membership), patch(
"litellm.proxy.auth.auth_checks.get_org_object", _get_org
):
if expect_blocked:
with pytest.raises(litellm.BudgetExceededError):
await _run()
else:
assert await _run() is True
@pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"])
@pytest.mark.asyncio
async def test_spend_capable_non_llm_routes_still_enforce_budget(route):
"""These routes are not LLM API routes but still reach a provider or an
external service: /health and /health/test_connection run litellm.ahealth_check
against real deployments, and /health/services fires Slack/email/webhook sends.
Exempting them with the other management routes would let an exhausted budget
keep spending.
"""
from fastapi import Request
from litellm.proxy.auth.auth_checks import common_checks
team = LiteLLM_TeamTable(team_id="t1", spend=150.0, max_budget=100.0)
with pytest.raises(litellm.BudgetExceededError):
await common_checks(
request_body={},
team_object=team,
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=AsyncMock(),
valid_token=UserAPIKeyAuth(token="k1", team_id="t1"),
request=MagicMock(spec=Request),
)
@pytest.mark.asyncio
async def test_get_default_end_user_budget_db_fetch_returns_validated_budget(monkeypatch):
from litellm.proxy.auth.auth_checks import get_default_end_user_budget

22
tests/test_litellm/proxy/db/test_prisma_client.py
View file

@ -193,3 +193,25 @@ async def test_recreate_prisma_client_recovers_from_disconnected_client(
mock_kill.assert_not_called()
assert wrapper._original_prisma is mock_new_prisma
mock_new_prisma.connect.assert_awaited_once()
def test_db_push_applies_replica_identity_full_when_requested(monkeypatch):
"""`prisma db push` bypasses litellm-proxy-extras, so it needs its own call
into the opt-in REPLICA IDENTITY FULL step."""
from litellm.proxy.db.prisma_client import PrismaManager
from litellm_proxy_extras.replica_identity import REPLICA_IDENTITY_FULL_ENV_VAR
from litellm_proxy_extras.utils import ProxyExtrasDBManager
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
applied = []
monkeypatch.setattr(
ProxyExtrasDBManager,
"apply_replica_identity_full_if_requested",
staticmethod(lambda: applied.append(True)),
)
with patch("litellm.proxy.db.prisma_client.subprocess.run") as mock_run:
assert PrismaManager.setup_database(use_migrate=False) is True
assert mock_run.call_args[0][0][:3] == ["prisma", "db", "push"]
assert applied == [True]

85
tests/test_litellm/proxy/db/test_replica_identity.py Normal file
View file

@ -0,0 +1,85 @@
"""The opt-in REPLICA IDENTITY FULL step, without a database.
The behavior against real Postgres is covered by
tests/proxy_migration_tests/test_replica_identity_full.py; these pin the two
things that hold with no database at all: the statement handed to the Prisma
CLI, and the promise that no failure of this optional step escapes into a
migration run that already succeeded.
"""
import subprocess
from pathlib import Path
from unittest.mock import patch
import pytest
from litellm_proxy_extras.replica_identity import (
REPLICA_IDENTITY_FULL_ENV_VAR,
apply_replica_identity_full,
)
from litellm_proxy_extras.utils import ProxyExtrasDBManager
def test_hands_the_alter_statement_to_the_prisma_cli():
captured = {}
def capture(cmd, **kwargs):
captured["cmd"] = cmd
captured["sql"] = Path(cmd[cmd.index("--file") + 1]).read_text()
return subprocess.CompletedProcess(cmd, 0)
with patch(
"litellm_proxy_extras.replica_identity.subprocess.run", side_effect=capture
):
applied = apply_replica_identity_full(
schema_path="/somewhere/schema.prisma",
prisma_command="prisma",
prisma_env={"DATABASE_URL": "postgresql://x/y"},
)
assert applied is True
assert captured["cmd"][:3] == ["prisma", "db", "execute"]
assert captured["cmd"][-2:] == ["--schema", "/somewhere/schema.prisma"]
sql = captured["sql"]
assert "ALTER TABLE %s REPLICA IDENTITY FULL" in sql
assert r"c.relname LIKE 'LiteLLM\_%'" in sql
assert "c.relreplident <> 'f'" in sql
assert "lock_timeout" in sql
@pytest.mark.parametrize(
"failure",
[
subprocess.CalledProcessError(1, "prisma", stderr="must be owner of table"),
subprocess.TimeoutExpired("prisma", 60),
OSError(2, "No such file or directory"),
PermissionError(13, "Read-only file system"),
],
ids=["rejected", "timed-out", "cli-missing", "read-only-fs"],
)
def test_every_failure_is_reported_instead_of_raised(failure):
with patch(
"litellm_proxy_extras.replica_identity.subprocess.run", side_effect=failure
):
assert (
apply_replica_identity_full(
schema_path="/somewhere/schema.prisma",
prisma_command="prisma",
prisma_env={},
)
is False
)
def test_an_unusable_migrations_dir_skips_the_step_instead_of_killing_the_run(
tmp_path, monkeypatch
):
"""LITELLM_MIGRATION_DIR makes the step copy the migrations tree before it
can run, and that copy is filesystem work that can fail on its own."""
blocker = tmp_path / "blocker"
blocker.write_text("not a directory")
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(blocker / "migrations"))
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False

83
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -3165,6 +3165,89 @@ async def test_pre_call_hook_does_not_leak_internal_stash_to_request_body():
assert isinstance(metadata.get(RATE_LIMIT_DESCRIPTORS_KEY), list)
@pytest.mark.asyncio
@pytest.mark.parametrize("caller_metadata", [None, {"user_tag": "abc"}])
async def test_pre_call_hook_does_not_touch_provider_metadata_on_litellm_metadata_routes(
caller_metadata,
):
"""Regression for #35197: routes that own ``litellm_metadata`` (Responses,
/v1/messages, batches, files) send ``metadata`` to the provider, so the
limiter must never create it or write stash keys into it."""
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_LITELLM_STASH_KEYS,
RATE_LIMIT_DESCRIPTORS_KEY,
RATE_LIMIT_RESPONSE_KEY,
TPM_RESERVED_TOKENS_KEY,
)
user_api_key_dict = UserAPIKeyAuth(
api_key=hash_token("sk-responses-metadata"),
tpm_limit=1000,
rpm_limit=5,
)
local_cache = DualCache()
handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
)
async def mock_should_rate_limit(descriptors, **kwargs):
return {
"overall_code": "OK",
"statuses": [
{
"code": "OK",
"current_limit": 5,
"limit_remaining": 4,
"descriptor_key": d["key"],
"descriptor_value": d["value"],
"rate_limit_type": "requests",
}
for d in descriptors
],
}
async def mock_reserve_tpm_tokens(descriptors, estimated_tokens, **kwargs):
return {"overall_code": "OK", "statuses": []}
handler.should_rate_limit = mock_should_rate_limit
handler.reserve_tpm_tokens = mock_reserve_tpm_tokens
data: Dict[str, Any] = {
"model": "responses-model",
"input": "hello",
"litellm_metadata": {},
}
if caller_metadata is not None:
data["metadata"] = dict(caller_metadata)
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=local_cache,
data=data,
call_type="aresponses",
)
if caller_metadata is None:
assert "metadata" not in data, f"limiter created provider metadata: {data.get('metadata')!r}"
else:
assert data["metadata"] == caller_metadata
litellm_metadata = data["litellm_metadata"]
assert litellm_metadata.get(TPM_RESERVED_TOKENS_KEY)
assert isinstance(litellm_metadata.get(RATE_LIMIT_DESCRIPTORS_KEY), list)
assert litellm_metadata.get(RATE_LIMIT_RESPONSE_KEY)
leaked = [k for k in _LITELLM_STASH_KEYS if k in data]
assert not leaked, f"stash keys leaked to top level: {leaked}"
for key in _LITELLM_STASH_KEYS:
assert handler._lookup_stashed_value(
kwargs={"litellm_params": {"litellm_metadata": litellm_metadata}},
standard_logging_metadata=None,
key=key,
) == litellm_metadata.get(key)
@pytest.mark.asyncio
async def test_pre_call_hook_rejects_caller_supplied_stash_values():
"""Caller cannot pre-populate stash keys in body metadata to drive a

330
tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -2893,6 +2893,7 @@ async def test_user_info_v2_response_shape(mocker):
"updated_at",
"sso_user_id",
"teams",
"object_permission",
}
assert set(response_dict.keys()) == expected_fields
@ -3702,3 +3703,332 @@ async def test_get_user_info_for_proxy_admin_validates_keys_and_teams():
returned_key = result.keys[0]
assert returned_key["team_id"] == "team-a"
assert returned_key["models"] == []
def _object_permission_mocks(mocker, existing_object_permission_id=None):
"""Prisma double whose user row optionally already links a permission row."""
mock_prisma_client = mocker.MagicMock()
existing_user = mocker.MagicMock()
existing_user.model_dump.return_value = {
"user_id": "target-user",
"object_permission_id": existing_object_permission_id,
}
existing_user.user_id = "target-user"
existing_user.object_permission_id = existing_object_permission_id
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
return_value=existing_user
)
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = mocker.AsyncMock(
return_value=None
)
mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock(
return_value=SimpleNamespace(object_permission_id="perm-new")
)
mock_prisma_client.update_data = mocker.AsyncMock(
return_value={"user_id": "target-user"}
)
mock_prisma_client.jsonify_object = mocker.MagicMock(side_effect=lambda x: x)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
mocker.patch(
"litellm.proxy.proxy_server._invalidate_spend_counter",
new=mocker.AsyncMock(),
)
return mock_prisma_client
@pytest.mark.asyncio
async def test_user_update_persists_mcp_entitlement_and_links_it(mocker):
"""/user/update documents an object_permission param; it must actually be stored.
The grants live in their own table, so the endpoint has to upsert them and hand the user row
only the resulting object_permission_id. Passing object_permission through to the user update
would not even be a column.
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mock_prisma_client = _object_permission_mocks(mocker)
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
await _update_single_user_helper(
user_request=UpdateUserRequest(
user_id="target-user",
object_permission={
"mcp_servers": ["github"],
"mcp_tool_permissions": {"github": ["list_issues"]},
},
),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
upsert_kwargs = mock_prisma_client.db.litellm_objectpermissiontable.upsert.call_args.kwargs
created = upsert_kwargs["data"]["create"]
assert created["mcp_servers"] == ["github"]
assert json.loads(created["mcp_tool_permissions"]) == {"github": ["list_issues"]}
written = mock_prisma_client.update_data.call_args.kwargs["data"]
assert written["object_permission_id"] == "perm-new"
assert "object_permission" not in written
@pytest.mark.asyncio
async def test_user_update_invalidates_the_cached_entitlement(mocker):
"""An admin revoking a tool must take effect now, not at the end of the cache TTL.
Three entries go stale: the permission row (keyed by its own id), the user -> permission link
(which carries a "no entitlement" sentinel), and the cached user row.
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
_object_permission_mocks(mocker)
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
await _update_single_user_helper(
user_request=UpdateUserRequest(
user_id="target-user",
object_permission={"mcp_tool_permissions": {"github": []}},
),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
assert deleted == {
"object_permission_id:perm-new",
"user_object_permission_id:target-user",
"target-user",
}
@pytest.mark.asyncio
async def test_admin_can_clear_a_users_mcp_entitlement(mocker):
"""An explicit empty object_permission means "no object permission", so it must unlink.
The merge-based upsert cannot express this: merging an empty grant set over the existing row
leaves every grant in place, and the empty-value filter drops the field before the upsert runs,
so without the explicit clear path the documented operation silently returns success unchanged.
A clear also leaves no incoming permission id, so invalidation keyed off one would skip it and
the gateway would keep enforcing the cleared grants until the cache expired.
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mock_prisma_client = _object_permission_mocks(mocker, "perm-existing")
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
await _update_single_user_helper(
user_request=UpdateUserRequest(user_id="target-user", object_permission={}),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
written = mock_prisma_client.update_data.call_args.kwargs["data"]
assert written["object_permission_id"] is None
assert "object_permission" not in written
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called()
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
assert deleted == {
"object_permission_id:perm-existing",
"user_object_permission_id:target-user",
"target-user",
}
@pytest.mark.asyncio
async def test_user_update_invalidates_both_the_old_and_new_permission_rows(mocker):
"""An upsert can mint a new permission row, which leaves the outgoing one cached under its id.
Only the link cache knows the user moved; the old row's own entry still holds the pre-update
grants, so anything still resolving that id keeps reading them.
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
_object_permission_mocks(mocker, "perm-existing")
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
await _update_single_user_helper(
user_request=UpdateUserRequest(
user_id="target-user",
object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}},
),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
deleted = {call.kwargs["key"] for call in cache.async_delete_cache.call_args_list}
assert deleted == {
"object_permission_id:perm-existing",
"object_permission_id:perm-new",
"user_object_permission_id:target-user",
"target-user",
}
@pytest.mark.asyncio
async def test_non_admin_cannot_clear_their_own_mcp_entitlement(mocker):
"""The empty-value filter drops `object_permission: {}` before the guard saw it, so a non-admin
could clear the very ceiling an admin placed on them. The guard reads the fields the caller SENT.
"""
from fastapi import HTTPException
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mock_prisma_client = _object_permission_mocks(mocker, "perm-existing")
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
with pytest.raises(HTTPException) as exc:
await _update_single_user_helper(
user_request=UpdateUserRequest(user_id="target-user", object_permission={}),
user_api_key_dict=UserAPIKeyAuth(
user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER
),
)
assert exc.value.status_code == 403
mock_prisma_client.update_data.assert_not_called()
@pytest.mark.asyncio
async def test_non_admin_cannot_rewrite_their_own_mcp_entitlement(mocker):
"""The entitlement bounds the human, so a self-write is an escalation path: an empty grant list
means "no restriction" and would lift a ceiling the admin placed on them."""
from fastapi import HTTPException
from litellm.proxy.management_endpoints.internal_user_endpoints import (
_update_single_user_helper,
)
mock_prisma_client = _object_permission_mocks(mocker, "perm-existing")
cache = mocker.MagicMock()
cache.async_delete_cache = mocker.AsyncMock()
mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache)
with pytest.raises(HTTPException) as exc:
await _update_single_user_helper(
user_request=UpdateUserRequest(
user_id="target-user",
object_permission={"mcp_servers": [], "mcp_tool_permissions": {}},
),
user_api_key_dict=UserAPIKeyAuth(
user_id="target-user", user_role=LitellmUserRoles.INTERNAL_USER
),
)
assert exc.value.status_code == 403
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_not_called()
mock_prisma_client.update_data.assert_not_called()
@pytest.mark.asyncio
async def test_new_user_persists_the_requested_mcp_entitlement(mocker):
"""generate_key_helper_fn only forwards object_permission_id, so /user/new has to create the
grants row itself; otherwise the entitlement the admin sent is silently dropped."""
mock_prisma_client = mocker.MagicMock()
mock_prisma_client.db.litellm_objectpermissiontable.create = mocker.AsyncMock(
return_value=SimpleNamespace(object_permission_id="perm-created")
)
mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock(
return_value=None
)
mock_prisma_client.db.litellm_usertable.count = mocker.AsyncMock(return_value=0)
mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.check_if_default_team_set",
return_value=None,
)
mock_generate = mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints.generate_key_helper_fn",
new=mocker.AsyncMock(
return_value={"user_id": "new-human", "token": "sk-x", "expires": None}
),
)
mocker.patch(
"litellm.proxy.hooks.user_management_event_hooks.UserManagementEventHooks.async_user_created_hook",
new=mocker.AsyncMock(),
)
await new_user(
data=NewUserRequest(
user_id="new-human",
object_permission={"mcp_tool_permissions": {"github": ["list_issues"]}},
),
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
created = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"]
assert json.loads(created["mcp_tool_permissions"]) == {"github": ["list_issues"]}
forwarded = mock_generate.call_args.kwargs
assert forwarded["object_permission_id"] == "perm-created"
assert "object_permission" not in forwarded
@pytest.mark.asyncio
async def test_user_info_v2_returns_the_mcp_entitlement(mocker):
"""The admin UI reads the current entitlement off this endpoint, so the grants have to come back
with the user row rather than only their id."""
from litellm.proxy.management_endpoints.internal_user_endpoints import user_info_v2
user_row = SimpleNamespace(
object_permission=SimpleNamespace(
object_permission_id="perm-1",
mcp_servers=["github"],
mcp_access_groups=[],
mcp_tool_permissions={"github": ["list_issues"]},
),
)
user_row.model_dump = lambda: {
"user_id": "human-1",
"object_permission": {
"object_permission_id": "perm-1",
"mcp_servers": ["github"],
"mcp_access_groups": [],
"mcp_tool_permissions": {"github": ["list_issues"]},
},
}
mocker.patch("litellm.proxy.proxy_server.prisma_client", mocker.MagicMock())
mocker.patch(
"litellm.proxy.management_endpoints.internal_user_endpoints._check_user_info_v2_access",
new=mocker.AsyncMock(return_value=user_row),
)
response = await user_info_v2(
request=SimpleNamespace(query_params={}),
user_id="human-1",
user_api_key_dict=UserAPIKeyAuth(
user_id="admin-1", user_role=LitellmUserRoles.PROXY_ADMIN
),
)
assert response.object_permission is not None
assert response.object_permission.mcp_servers == ["github"]
assert response.object_permission.mcp_tool_permissions == {
"github": ["list_issues"]
}

38
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
View file

@ -590,22 +590,20 @@ class TestVertexAIBatchCostCalculation:
assert usage.completion_tokens == 0
assert usage.total_tokens == 0
def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
@pytest.mark.asyncio
async def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
"""
Regression test for the bug where Vertex batch cost/usage was always 0.
After PR #25627 (transform_file_content_response), the GCS predictions.jsonl
is rewritten into OpenAI batch shape before the cost-tracking path sees it.
With disable_vertex_batch_output_transformation=False (default), the content
is OpenAI-shaped, so _batch_cost_calculator must fall through to the generic
path rather than calling calculate_vertex_ai_batch_cost_and_usage (which only
reads raw usageMetadata fields).
With disable_vertex_batch_output_transformation=False (default), the cost
dispatch must fall through to the generic aggregation path rather than
calling calculate_vertex_ai_batch_cost_and_usage (which only reads raw
usageMetadata fields).
"""
import litellm
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_total_usage_from_file_content,
)
from litellm.batches.batch_utils import calculate_batch_cost_and_usage
openai_shaped_responses = [
{
@ -668,12 +666,7 @@ class TestVertexAIBatchCostCalculation:
try:
litellm.disable_vertex_batch_output_transformation = False
cost = _batch_cost_calculator(
file_content_dictionary=openai_shaped_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
usage = _get_batch_job_total_usage_from_file_content(
cost, usage, _ = await calculate_batch_cost_and_usage(
file_content_dictionary=openai_shaped_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
@ -694,16 +687,14 @@ class TestVertexAIBatchCostCalculation:
cost > 0
), f"expected non-zero cost for completed Vertex batch, got {cost}"
def test_raw_vertex_output_still_works_when_transformation_disabled(self):
@pytest.mark.asyncio
async def test_raw_vertex_output_still_works_when_transformation_disabled(self):
"""
When disable_vertex_batch_output_transformation=True the GCS file is returned
as raw Vertex predictions.jsonl; the specialized reader must be used.
"""
import litellm
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_total_usage_from_file_content,
)
from litellm.batches.batch_utils import calculate_batch_cost_and_usage
raw_vertex_responses = [
{
@ -727,12 +718,7 @@ class TestVertexAIBatchCostCalculation:
try:
litellm.disable_vertex_batch_output_transformation = True
cost = _batch_cost_calculator(
file_content_dictionary=raw_vertex_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",
)
usage = _get_batch_job_total_usage_from_file_content(
cost, usage, _ = await calculate_batch_cost_and_usage(
file_content_dictionary=raw_vertex_responses,
custom_llm_provider="vertex_ai",
model_name="gemini-2.0-flash-001",

146
tests/test_litellm/proxy/proxy_server/test_routes_model_info.py
View file

@ -292,3 +292,149 @@ def test_model_group_info_invalid_method(client, auth_as, null_router):
response = client.post("/model_group/info", json={})
assert response.status_code == 405
assert len(response.content) > 0
# ---------------------------------------------------------------------------
# GET /v2/model/info?exclude_auto_routers
# ---------------------------------------------------------------------------
@pytest.fixture
def mixed_auto_router_router(monkeypatch):
"""Router carrying one ordinary deployment per auto-router strategy plus two plain ones."""
model_list = [
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
"model_info": {"id": "plain-1", "db_model": False},
},
{
"model_name": "tri-tier-router",
"litellm_params": {"model": "auto_router/complexity_router"},
"model_info": {"id": "auto-complexity", "db_model": True},
},
{
"model_name": "support-router",
"litellm_params": {"model": "auto_router/support-router"},
"model_info": {"id": "auto-semantic", "db_model": True},
},
{
"model_name": "adaptive-router",
"litellm_params": {"model": "auto_router/adaptive_router"},
"model_info": {"id": "auto-adaptive", "db_model": True},
},
{
"model_name": "claude-opus",
"litellm_params": {"model": "anthropic/claude-opus-4-6"},
"model_info": {"id": "plain-2", "db_model": False},
},
]
from unittest.mock import AsyncMock
router = MagicMock()
router.model_list = model_list
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "llm_model_list", model_list)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "user_model", None)
monkeypatch.setattr(proxy_server.proxy_config, "get_config", AsyncMock(return_value={}))
monkeypatch.setattr(
proxy_server,
"_apply_search_filter_to_models",
AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))),
)
monkeypatch.setattr(proxy_server, "_enrich_model_info_with_litellm_data", lambda model, **kw: model)
import litellm.proxy.agent_endpoints.model_list_helpers as mlh
monkeypatch.setattr(mlh, "append_agents_to_model_info", AsyncMock(side_effect=lambda models, **kw: models))
yield router
def _model_names(payload) -> list:
return [m["model_name"] for m in payload["data"]]
def test_v2_model_info_includes_auto_routers_by_default(client, auth_as, mixed_auto_router_router):
"""The new param is opt-in; omitting it must not change what any existing caller sees."""
with auth_as():
response = client.get("/v2/model/info")
assert response.status_code == 200
payload = response.json()
assert "tri-tier-router" in _model_names(payload)
assert payload["total_count"] == 5
def test_v2_model_info_excludes_every_auto_router_strategy(client, auth_as, mixed_auto_router_router):
"""All four `auto_router/*` strategies go, not just the semantic one that
Router._is_auto_router_deployment recognises."""
with auth_as():
response = client.get("/v2/model/info", params={"exclude_auto_routers": "true"})
assert response.status_code == 200
payload = response.json()
assert _model_names(payload) == ["gpt-4o-mini", "claude-opus"]
def test_v2_model_info_exclude_auto_routers_shrinks_total_count(client, auth_as, mixed_auto_router_router):
"""The filter must run before the count, or the table pages off a total that
includes rows it never renders (49 shown, 50 claimed)."""
with auth_as():
response = client.get("/v2/model/info", params={"exclude_auto_routers": "true"})
payload = response.json()
assert payload["total_count"] == 2
assert len(payload["data"]) == payload["total_count"]
def test_v2_model_info_exclude_auto_routers_paginates_over_the_filtered_set(
client, auth_as, mixed_auto_router_router
):
"""Page size applies to the filtered list, so no page silently comes back short."""
with auth_as():
response = client.get(
"/v2/model/info", params={"exclude_auto_routers": "true", "page": 1, "size": 1}
)
payload = response.json()
assert payload["total_count"] == 2
assert payload["total_pages"] == 2
assert len(payload["data"]) == 1
@pytest.mark.asyncio
async def test_model_info_v2_query_sentinel_does_not_filter(monkeypatch, mixed_auto_router_router):
"""Called directly (not through FastAPI) the default arrives as a truthy Query object.
Guarding on `is True` is what stops every direct-call test from silently filtering."""
from unittest.mock import AsyncMock
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server.proxy_config, "get_config", AsyncMock(return_value={}))
monkeypatch.setattr(
proxy_server,
"_apply_search_filter_to_models",
AsyncMock(side_effect=lambda all_models, **kw: (all_models, len(all_models))),
)
monkeypatch.setattr(proxy_server, "_enrich_model_info_with_litellm_data", lambda model, **kw: model)
import litellm.proxy.agent_endpoints.model_list_helpers as mlh
monkeypatch.setattr(mlh, "append_agents_to_model_info", AsyncMock(side_effect=lambda models, **kw: models))
admin = UserAPIKeyAuth(user_id="u", user_role=LitellmUserRoles.PROXY_ADMIN)
# Deliberately omit exclude_auto_routers, exactly as the pre-existing direct-call tests do.
resp = await proxy_server.model_info_v2(
user_api_key_dict=admin,
model=None,
user_models_only=False,
include_team_models=False,
debug=False,
page=1,
size=50,
search=None,
modelId=None,
teamId=None,
sortBy=None,
sortOrder="asc",
)
assert "tri-tier-router" in [m["model_name"] for m in resp["data"]]

6
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -2396,7 +2396,7 @@ class TestSpendLogsPayload:
"model": "gpt-4o",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}',
"cache_key": "Cache OFF",
"spend": 0.00022500000000000002,
"total_tokens": 30,
@ -2492,7 +2492,7 @@ class TestSpendLogsPayload:
"model": "claude-4-sonnet-20250514",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"cache_key": "Cache OFF",
"spend": 0.01383,
"total_tokens": 2598,
@ -2586,7 +2586,7 @@ class TestSpendLogsPayload:
"model": "claude-4-sonnet-20250514",
"user": "",
"team_id": "",
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}',
"cache_key": "Cache OFF",
"spend": 0.01383,
"total_tokens": 2598,

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