scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

chore(mcp): reconcile listed metadata branch with main

Browse source
This commit is contained in:
yucheng 2026-10-02 18:35:36 +00:00
commit ad827e0025
467 changed files with 35399 additions and 17276 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
.circleci/scripts/run_integration.sh
View file

@ -169,11 +169,11 @@ start_proxy() {
INTEGRATION_UPSTREAM_URL="$INTEGRATION_UPSTREAM_URL" \
LITELLM_MASTER_KEY="$LITELLM_MASTER_KEY" LITELLM_SALT_KEY="$LITELLM_SALT_KEY" LITELLM_UI_PATH="$LITELLM_UI_PATH" PROXY_BASE_URL="http://127.0.0.1:$port" \
LITELLM_LICENSE="${LITELLM_LICENSE:-}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True "${cost_map_env[@]}" \
LITELLM_MODE=PRODUCTION STORE_MODEL_IN_DB=True LITELLM_ENABLE_MCP_STDIO=true "${cost_map_env[@]}" \
AWS_EC2_METADATA_DISABLED=true DO_NOT_TRACK=1 COVERAGE_FILE="$coverage_data" \
"${proxy_command[@]}" --config tests/integration/proxy_config.yaml \
--host 127.0.0.1 --port "$port" --num_workers 1 --telemetry False \
--use_prisma_db_push --enforce_prisma_migration_check \
--use_prisma_db_push \
> "$results/$log_name" 2>&1 &
launched_pid=$!
}

16
backend/main.py
View file

@ -8,9 +8,13 @@ Run with:
uvicorn backend.main:app --host 0.0.0.0 --port 4001
"""
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from typing import Final
from fastapi.routing import Mount
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Lifespan
# See gateway/main.py for why we assemble DATABASE_URL(s) here before
# importing proxy_server.
@ -43,14 +47,16 @@ def _is_backend_route(route) -> bool:
# See gateway/main.py for why the trim runs inside the lifespan instead of at
# module scope.
_proxy_lifespan = app.router.lifespan_context
_proxy_lifespan: Final = app.router.lifespan_context
@asynccontextmanager
async def _backend_lifespan(app_):
async with _proxy_lifespan(app_):
async def _backend_lifespan(
app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan
) -> AsyncGenerator[Mapping[str, object], None]:
async with lifespan(app_) as state:
app_.router.routes = [r for r in app_.router.routes if _is_backend_route(r)]
yield
yield state if state is not None else {}
app.router.lifespan_context = _backend_lifespan

1
backend/routes/allowlist.py
View file

@ -60,6 +60,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
# Tools / agents (registry & policy admin)
"/v1/tool/",
"/v1/agents",
"/agent/daily/activity/",
# Guardrails admin
"/v2/guardrails/",
# MCP server admin + BYOK OAuth flow (UI-initiated) + dynamic per-server endpoints

15
deploy/lens/README.md
View file

@ -4,7 +4,20 @@ Lens reviews recorded activity and saves evidence-linked findings in the LiteLLM
## Start a worker
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL, agent tracing (`general_settings.tracing: {store: clickhouse}`), and ClickHouse configured through `CLICKHOUSE_URL` and a separate SELECT-only `CLICKHOUSE_READER_URL`. Enable the ClickHouse callback and request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL and agent tracing. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries:
```yaml
general_settings:
tracing:
store:
type: clickhouse
url: os.environ/CLICKHOUSE_URL
retention_days: 14
```
The URL, database, and retention settings can also come from `CLICKHOUSE_URL`, `CLICKHOUSE_DATABASE`, and `AGENT_TRACING_RETENTION_DAYS` when omitted from YAML. A YAML value wins when both are set. The database defaults to `litellm`. `retention_days` defaults to 14 and applies to both traces and spend logs
Retention changes require a proxy restart. ClickHouse removes expired rows during background merges, not immediately at startup. Enable request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
In Lens, click **Set up analysis**, choose an existing virtual key or **Create worker key**, then **Generate setup command**. The LiteLLM address is filled in for you; change it only if the server running Docker needs a different network address. Copy the command and run it on your server. The dialog changes to **Analyzer connected** when the container checks in

2
deploy/lens/compose.yaml
View file

@ -1,6 +1,6 @@
services:
lens-worker:
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:a8e8731d954916594eea462969946b9292fb771681ff515a9fd296b53f856c77}
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:67eba741c1b97c749975c5c38e2370a603e1105babc908d613c1b79d7b995393}
environment:
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}

2
docker-compose.hardened.yml
View file

@ -6,8 +6,6 @@ services:
context: .
dockerfile: docker/Dockerfile.non_root
target: runtime
args:
PROXY_EXTRAS_SOURCE: "local"
depends_on:
- squid
user: "101:101"

31
docker/Dockerfile.non_root
View file

@ -3,7 +3,6 @@
# Base images
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
@ -44,7 +43,6 @@ COPY ui/litellm-dashboard/ ./
RUN npm run build
FROM $LITELLM_BUILD_IMAGE AS builder
ARG PROXY_EXTRAS_SOURCE
WORKDIR /app
USER root
@ -107,26 +105,14 @@ RUN mkdir -p /var/lib/litellm/ui /var/lib/litellm/assets && \
touch /var/lib/litellm/ui/.litellm_ui_ready
RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \
if [ "$PROXY_EXTRAS_SOURCE" = "published" ]; then \
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13 \
--no-sources-package litellm-proxy-extras; \
else \
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13; \
fi
uv sync --frozen --no-default-groups --no-editable \
--extra proxy \
--extra proxy-runtime \
--extra extra_proxy \
--extra semantic-router \
--extra saml \
--extra bedrock-realtime \
--python python3.13
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
@ -136,7 +122,6 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
FROM $LITELLM_RUNTIME_IMAGE AS runtime
ARG PROXY_EXTRAS_SOURCE
WORKDIR /app
USER root

8
docker/docker-compose.quickstart.yml
View file

@ -13,11 +13,13 @@ services:
litellm:
image: docker.litellm.ai/berriai/litellm:main-stable
ports:
- "4000:4000"
# LITELLM_BIND is empty by default, so this stays "4000:4000". The quickstart
# script sets it to "127.0.0.1:" so new installs listen on this machine only.
- "${LITELLM_BIND:-}${LITELLM_PORT:-4000}:4000"
environment:
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?set it in .env - see the header of this file}
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?set it in .env - see the header of this file}
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
DATABASE_URL: postgresql://litellm:${POSTGRES_PASSWORD:-litellm}@db:5432/litellm
STORE_MODEL_IN_DB: "True"
depends_on:
db:
@ -27,7 +29,7 @@ services:
image: postgres:16
environment:
POSTGRES_USER: litellm
POSTGRES_PASSWORD: litellm
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-litellm}
POSTGRES_DB: litellm
healthcheck:
test: ["CMD-SHELL", "pg_isready -U litellm"]

1
docker/docker-compose.tracing.yml
View file

@ -12,7 +12,6 @@ services:
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
STORE_MODEL_IN_DB: "True"
CLICKHOUSE_URL: http://default:local-tracing@clickhouse:8123
CLICKHOUSE_READER_URL: http://default:local-tracing@clickhouse:8123
CLICKHOUSE_DATABASE: litellm
OPENAI_API_KEY: ${OPENAI_API_KEY:-}
volumes:

5
docker/tracing-config.yaml
View file

@ -7,4 +7,7 @@ model_list:
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
tracing:
store: clickhouse
store:
type: clickhouse
url: os.environ/CLICKHOUSE_URL
retention_days: 14

4
enterprise/pyproject.toml
View file

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

16
gateway/main.py
View file

@ -9,9 +9,13 @@ Run with:
uvicorn gateway.main:app --host 0.0.0.0 --port 4000
"""
from collections.abc import AsyncGenerator, Mapping
from contextlib import asynccontextmanager
from typing import Final
from fastapi.routing import Mount
from starlette.applications import Starlette
from starlette.routing import Mount
from starlette.types import Lifespan
# Assemble DATABASE_URL (+ DATABASE_URL_READ_REPLICA) from the discrete
# DATABASE_* env vars before proxy_server imports spin up Prisma. Handles
@ -54,14 +58,16 @@ def _is_gateway_route(route) -> bool:
# register routes. A module-load filter would miss routes added during
# startup; running inside the lifespan, after the inner __aenter__, catches
# them while still completing before uvicorn opens the listener.
_proxy_lifespan = app.router.lifespan_context
_proxy_lifespan: Final = app.router.lifespan_context
@asynccontextmanager
async def _gateway_lifespan(app_):
async with _proxy_lifespan(app_):
async def _gateway_lifespan(
app_: Starlette, lifespan: Lifespan[Starlette] = _proxy_lifespan
) -> AsyncGenerator[Mapping[str, object], None]:
async with lifespan(app_) as state:
app_.router.routes = [r for r in app_.router.routes if _is_gateway_route(r)]
yield
yield state if state is not None else {}
app.router.lifespan_context = _gateway_lifespan

18
helm/litellm-helm/tests/migrations-job_tests.yaml
View file

@ -112,6 +112,24 @@ tests:
name: CUSTOM_VAR
value: "custom_value"
- it: should override a user-supplied DISABLE_SCHEMA_UPDATE so the Job always migrates
template: migrations-job.yaml
set:
envVars:
DISABLE_SCHEMA_UPDATE: "true"
migrationJob:
enabled: true
asserts:
# The Job is what owns the schema, so it renders its own
# DISABLE_SCHEMA_UPDATE=false after envVars and extraEnvVars. Kubernetes
# takes the last value for a duplicated name, so the user's "true" cannot
# leave the schema unmigrated. Skipping migrations is migrationJob.enabled.
- equal:
path: spec.template.spec.containers[0].env[-1]
value:
name: DISABLE_SCHEMA_UPDATE
value: "false"
- it: should not include DATABASE_URL when deployStandalone is false
template: migrations-job.yaml
set:

2
helm/litellm-helm/values.yaml
View file

@ -545,7 +545,6 @@ redis:
# Prisma migration job settings
migrationJob:
enabled: true # Enable or disable the schema migration Job
retries: 3 # Number of retries for the Job in case of failure
backoffLimit: 4 # Backoff limit for Job restarts
# Wall-clock budget for the whole Job, shared across every `backoffLimit`
# retry rather than granted per attempt. Without it a migration that blocks
@ -554,7 +553,6 @@ migrationJob:
# stop reconciling the whole chart until someone deletes the Job by hand.
# Set to null to opt out and restore the unbounded behaviour.
activeDeadlineSeconds: 1800
disableSchemaUpdate: false # Skip schema migrations for specific environments. When True, the job will exit with code 0.
# Optional service account for the migration job.
# Only used when migrationJob.hooks.helm.enabled=true and serviceAccount.create=true.
# In that case, pre-install/pre-upgrade hooks run before normal resources, so this defaults to "default".

57
litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py
View file

@ -56,8 +56,10 @@ REQUEST_LOG_INDEXES: Final = (
)
_IDENTIFIER_MAX_BYTES: Final = 63
_PARENT_LOCK_TIMEOUT: Final = "2s"
_PARENT_LOCK_ATTEMPTS: Final = 30
_DDL_LOCK_TIMEOUT: Final = "200ms"
_DDL_LOCK_ATTEMPTS: Final = 10
_DDL_RETRY_BASE_SECONDS: Final = 0.25
_DDL_RETRY_MAX_SECONDS: Final = 8.0
_LOCK_HANDOVER_SECONDS: Final = 2.0
_DIGEST_LENGTH: Final = 8
_CREATE_INDEX_STATEMENT: Final = re.compile(
@ -181,6 +183,32 @@ def _under_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]",
return step()
def _with_bounded_lock(
connection: "psycopg.Connection[tuple[object, ...]]", step: Callable[[], bool], what: str
) -> bool:
"""Run `step` under the migration lock with a short lock_timeout, so a DDL statement that has to wait for open
transactions holds new writes back for at most that long; retry with capped exponential backoff, holding the
migration lock per attempt only and releasing it while sleeping. False when another process holds the migration
lock or every attempt timed out."""
import psycopg
from psycopg import sql
for attempt in range(_DDL_LOCK_ATTEMPTS):
if attempt:
time.sleep(min(_DDL_RETRY_MAX_SECONDS, _DDL_RETRY_BASE_SECONDS * 2.0**attempt) * random.uniform(0.5, 1.0))
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_DDL_LOCK_TIMEOUT)))
try:
return _under_migration_lock(connection, step)
except psycopg.errors.LockNotAvailable:
logger.info("Waiting for open transactions before %s", what)
finally:
connection.execute("SET lock_timeout = 0")
logger.warning(
"Could not get the lock for %s without holding writes back, leaving it for the next index build", what
)
return False
def _ensure_index(connection: "psycopg.Connection[tuple[object, ...]]", schema: str, index: RequestLogIndex) -> bool:
from psycopg.rows import class_row
@ -368,12 +396,13 @@ def build_index_on_partitioned_table(
"Index %s already exists on %s rather than %s, leaving it alone", parent_index, existing.table, parent_table
)
return False
if existing is None and not _under_migration_lock(
if existing is None and not _with_bounded_lock(
connection,
lambda: (
_adopt_equivalent_index(connection, schema, parent_table, parent_index, index)
or _create_parent_index(connection, schema, parent_index, parent_table, index)
),
f"creating the parent index {parent_index}",
):
return False
children: Final = _children_without_the_index(connection, schema, parent_table, parent_index)
@ -393,29 +422,15 @@ def _create_parent_index(
table: str,
index: RequestLogIndex,
) -> bool:
"""Create the metadata-only parent index. Postgres takes a SHARE lock on the
parent for that statement, so it waits for in-flight writes and queues new ones
behind it; a short lock_timeout with retries keeps every such pause bounded."""
import psycopg
"""Create the metadata-only parent index. The caller bounds Postgres's SHARE lock wait on the parent."""
from psycopg import sql
prefix: Final = sql.SQL("CREATE INDEX IF NOT EXISTS {} ON ONLY {} ").format(
sql.Identifier(name), sql.Identifier(schema, table)
)
statement: Final = _create_index_statement(connection, prefix, index.definition)
connection.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(_PARENT_LOCK_TIMEOUT)))
try:
for _ in range(_PARENT_LOCK_ATTEMPTS):
try:
connection.execute(statement)
return True
except psycopg.errors.LockNotAvailable:
logger.info("Waiting for in-flight writes to %s before creating the parent index %s", table, name)
time.sleep(random.uniform(0.1, 0.5))
finally:
connection.execute("SET lock_timeout = 0")
logger.warning("Could not get the parent lock on %s to create %s, leaving it for the next index build", table, name)
return False
connection.execute(statement)
return True
def _attach_child_index(
@ -445,4 +460,4 @@ def _attach_child_index(
logger.info("Attached index %s on partition %s to %s", child_index, child.name, parent_index)
return True
return _under_migration_lock(connection, attach)
return _with_bounded_lock(connection, attach, f"attaching {child_index}")

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

@ -78,6 +78,23 @@ class _InvalidIndex:
table_size: str
MAX_MIGRATE_DEPLOY_ATTEMPTS = 4
LIBPQ_URL_PARAMS: Final = frozenset(
{
"sslmode",
"sslcert",
"sslkey",
"sslrootcert",
"sslpassword",
"application_name",
"connect_timeout",
"client_encoding",
"options",
"service",
"gssencmode",
"krbsrvname",
"target_session_attrs",
}
)
@dataclass(frozen=True)
@ -689,30 +706,43 @@ class ProxyExtrasDBManager:
@staticmethod
def _strip_prisma_query_params(url: str) -> str:
"""Remove Prisma-specific query params (connection_limit, pool_timeout,
schema, etc.) from DATABASE_URL so psycopg can parse it."""
"""Rewrite a Prisma-dialect URL for libpq: drop the Prisma-only params
(connection_limit, pool_timeout, schema, pgbouncer, sslaccept, ...) and
translate Prisma's TLS params back, since libpq reads ``sslcert`` as a
client certificate where Prisma reads it as the CA."""
from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse
parsed = urlparse(url)
parsed: Final = urlparse(url)
if not parsed.query:
return url
libpq_params = {
"sslmode",
"sslcert",
"sslkey",
"sslrootcert",
"sslpassword",
"application_name",
"connect_timeout",
"client_encoding",
"options",
"service",
"gssencmode",
"krbsrvname",
"target_session_attrs",
}
kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params]
return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote)))
pairs: Final = tuple(parse_qsl(parsed.query))
kept: Final = tuple((k, v) for k, v in pairs if k in LIBPQ_URL_PARAMS)
sslaccept: Final = next((v for k, v in pairs if k == "sslaccept"), None)
libpq_pairs: Final = ProxyExtrasDBManager._libpq_tls_params(kept, sslaccept)
return urlunparse(parsed._replace(query=urlencode(libpq_pairs, quote_via=quote)))
@staticmethod
def _libpq_tls_params(
pairs: "tuple[tuple[str, str], ...]", sslaccept: "str | None"
) -> "tuple[tuple[str, str], ...]":
"""Undo ``translate_libpq_ssl_params``. Prisma's ``sslcert`` is the CA and
``sslaccept=strict`` checks chain and hostname, which libpq only does in
``sslmode=verify-full``, so strict becomes ``sslrootcert`` plus
``verify-full`` whatever ``sslmode`` said (``disable`` stays off). Prisma
defaults an absent ``sslaccept`` to ``accept_invalid_certs`` and anything
else to strict. Without strict it checks nothing, so the CA is dropped and
``sslmode`` is kept as is: libpq only verifies when a root cert is present.
A URL that also carries ``sslkey`` is libpq's own client-certificate form
and is kept."""
keys: Final = frozenset(k for k, _ in pairs)
if "sslcert" not in keys or "sslkey" in keys:
return pairs
sslmode: Final = next((v for k, v in pairs if k == "sslmode"), None)
rest: Final = tuple((k, v) for k, v in pairs if k not in ("sslcert", "sslmode"))
if sslaccept in (None, "accept_invalid_certs") or sslmode == "disable":
return rest if sslmode is None else rest + (("sslmode", sslmode),)
root_cert: Final = tuple(("sslrootcert", v) for k, v in pairs if k == "sslcert" and "sslrootcert" not in keys)
return rest + root_cert + (("sslmode", "verify-full"),)
@staticmethod
def _warn_if_db_ahead_of_head(migrations_dir: str) -> None:

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

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

74
litellm-rust/Cargo.lock generated
View file

@ -97,6 +97,53 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
[[package]]
name = "askama"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6024d73179f43f15ccd2b881bfea6fee7f3a46ec53f33b52210dea749ebebaa4"
dependencies = [
"askama_macros",
"itoa",
"percent-encoding",
"serde",
"serde_json",
]
[[package]]
name = "askama_derive"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "071ee5ebf2138e3ad180e0aacf6940c2cab5e6d8333741d9925c7bee2b153f39"
dependencies = [
"askama_parser",
"memchr",
"proc-macro2",
"quote",
"rustc-hash",
"syn 3.0.6",
]
[[package]]
name = "askama_macros"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "643e1c7cbb6aec1d920332fe51a7c0d8219e273dcb8602db03f5263e4d16487b"
dependencies = [
"askama_derive",
]
[[package]]
name = "askama_parser"
version = "0.16.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2c5ae75772275d268b03ab8bdccdd12117b6169ee23256942b34e46c9f476583"
dependencies = [
"rustc-hash",
"unicode-ident",
"winnow 1.0.4",
]
[[package]]
name = "asn1-rs"
version = "0.7.2"
@ -4038,6 +4085,26 @@ dependencies = [
"strum",
]
[[package]]
name = "litellm-migrate"
version = "0.1.0"
dependencies = [
"litellm-migrate-macros",
"rstest",
]
[[package]]
name = "litellm-migrate-macros"
version = "0.1.0"
dependencies = [
"proc-macro2",
"quote",
"rstest",
"syn 2.0.119",
"tempfile",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-model-catalog"
version = "0.1.0"
@ -4384,11 +4451,17 @@ dependencies = [
name = "litellm-traces"
version = "0.1.0"
dependencies = [
"askama",
"base64 0.22.1",
"criterion",
"flate2",
"futures-util",
"hmac 0.12.1",
"indexmap 2.14.0",
"litellm-http",
"litellm-migrate",
"litellm-storage-clickhouse",
"moka",
"opentelemetry-proto",
"prost",
"rstest",
@ -4400,6 +4473,7 @@ dependencies = [
"thiserror 2.0.19",
"time",
"tokio",
"url",
"wiremock",
]

6
litellm-rust/Cargo.toml
View file

@ -14,6 +14,8 @@ litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
litellm-migrate = { path = "crates/migrate" }
litellm-migrate-macros = { path = "crates/migrate-macros" }
litellm-core = { path = "crates/core" }
litellm-gateway-mcp = { path = "crates/gateway-mcp" }
litellm-gateway = { path = "crates/gateway" }
@ -63,6 +65,7 @@ litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
litellm-host-python = { path = "crates/host-python" }
litellm-python-compat = { path = "crates/python-compat" }
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
tracing = "0.1"
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
axum-login = "0.18.0"
@ -93,7 +96,10 @@ serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["float_roundtrip"] }
serde_with = { version = "=3.16.1", default-features = false, features = ["std", "macros"] }
sha2 = "0.10"
syn = { version = "2", default-features = false }
sqlx = { version = "0.9.0", default-features = false, features = ["json", "macros", "postgres", "runtime-tokio", "chrono", "tls-rustls-ring-native-roots"] }
proc-macro2 = "1"
quote = "1"
subtle = "2"
thiserror = "2.0"
tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] }

5
litellm-rust/crates/config/src/lib.rs
View file

@ -12,7 +12,10 @@ use serde::Deserialize;
pub use error::Error;
pub use mcp::{McpAuth, McpServer, McpTransport};
pub use model::{LiteLlmParams, Model};
pub use settings::{GeneralSettings, LiteLlmSettings, RouterSettings};
pub use settings::{
ClickHouseStoreSettings, GeneralSettings, LiteLlmSettings, RouterSettings, TracingSettings,
TracingStoreSettings,
};
pub use value::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value};
#[derive(Clone, Default, Deserialize)]

44
litellm-rust/crates/config/src/settings.rs
View file

@ -5,6 +5,47 @@ use serde::Deserialize;
use crate::{AdditionalFields, Flag, NumberOrString, Object, OneOrMany, Value};
#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum TracingStoreKind {
Clickhouse,
}
#[derive(Clone, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ClickHouseStoreSettings {
#[serde(rename = "type")]
pub kind: TracingStoreKind,
pub url: Option<SecretValue>,
pub database: Option<String>,
pub retention_days: Option<NumberOrString>,
}
impl fmt::Debug for ClickHouseStoreSettings {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ClickHouseStoreSettings")
.field("kind", &self.kind)
.field("database", &self.database)
.field("retention_days", &self.retention_days)
.finish()
}
}
#[derive(Clone, Debug, Deserialize)]
#[serde(untagged)]
pub enum TracingStoreSettings {
ClickHouse(ClickHouseStoreSettings),
}
#[derive(Clone, Default, Debug, Deserialize)]
#[serde(default)]
pub struct TracingSettings {
pub store: Option<TracingStoreSettings>,
#[serde(flatten)]
pub additional_fields: AdditionalFields,
}
#[derive(Clone, Deserialize)]
#[serde(default)]
pub struct GeneralSettings {
@ -14,6 +55,7 @@ pub struct GeneralSettings {
pub admission_queue_timeout_seconds: f64,
pub master_key: Option<SecretValue>,
pub database_url: Option<SecretValue>,
pub tracing: Option<TracingSettings>,
pub database_connection_pool_limit: Option<u64>,
pub database_connection_timeout: Option<f64>,
pub database_connect_timeout: Option<f64>,
@ -50,6 +92,7 @@ impl Default for GeneralSettings {
admission_queue_timeout_seconds: 1.0,
master_key: None,
database_url: None,
tracing: None,
database_connection_pool_limit: Some(10),
database_connection_timeout: Some(60.0),
database_connect_timeout: None,
@ -97,6 +140,7 @@ impl fmt::Debug for GeneralSettings {
)
.field("master_key", &self.master_key)
.field("database_url", &self.database_url)
.field("tracing", &self.tracing)
.field("store_model_in_db", &self.store_model_in_db)
.field("additional_fields", &self.additional_fields.keys())
.finish_non_exhaustive()

50
litellm-rust/crates/config/tests/config.rs
View file

@ -1,4 +1,4 @@
use litellm_config::{Config, Error, Flag, NumberOrString};
use litellm_config::{Config, Error, Flag, NumberOrString, TracingStoreSettings};
use rstest::{fixture, rstest};
use tempfile::TempDir;
@ -113,6 +113,54 @@ fn missing_general_settings_has_no_master_key() {
assert!(config.general_settings.master_key.is_none());
}
#[test]
fn tracing_settings_are_typed_and_redact_the_url() {
let config = Config::from_yaml(
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: https://writer:password@example.com\n database: analytics\n retention_days: 7\n",
)
.unwrap();
let tracing = config.general_settings.tracing.as_ref().unwrap();
let Some(TracingStoreSettings::ClickHouse(store)) = tracing.store.as_ref() else {
panic!("expected ClickHouse tracing store")
};
assert_eq!(
store.url.as_ref().unwrap().expose(),
"https://writer:password@example.com"
);
assert_eq!(store.database.as_deref(), Some("analytics"));
assert_eq!(store.retention_days, Some(NumberOrString::Number(7.0)));
assert!(!format!("{config:?}").contains("password"));
}
#[test]
fn tracing_settings_accept_environment_references() {
let config = Config::from_yaml(
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: os.environ/CLICKHOUSE_URL\n retention_days: os.environ/RETENTION_DAYS\n",
)
.unwrap();
let Some(TracingStoreSettings::ClickHouse(store)) =
config.general_settings.tracing.unwrap().store
else {
panic!("expected ClickHouse tracing store")
};
assert_eq!(
store.retention_days,
Some(NumberOrString::String(
"os.environ/RETENTION_DAYS".to_owned()
))
);
}
#[test]
fn tracing_settings_reject_string_store() {
assert!(Config::from_yaml("general_settings:\n tracing:\n store: clickhouse\n").is_err());
}
#[test]
fn tracing_settings_reject_removed_reader_configuration() {
assert!(Config::from_yaml("general_settings:\n tracing:\n store:\n type: clickhouse\n reader_url: http://localhost:8123\n").is_err());
}
#[rstest]
fn empty_config_matches_python_defaults() {
let config = Config::from_yaml("{}").unwrap();

19
litellm-rust/crates/migrate-macros/Cargo.toml Normal file
View file

@ -0,0 +1,19 @@
[package]
name = "litellm-migrate-macros"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[lib]
proc-macro = true
[dependencies]
proc-macro2.workspace = true
quote.workspace = true
syn = { workspace = true, features = ["parsing", "printing", "proc-macro"] }
thiserror.workspace = true
[dev-dependencies]
rstest.workspace = true
tempfile.workspace = true

21
litellm-rust/crates/migrate-macros/src/error.rs Normal file
View file

@ -0,0 +1,21 @@
use std::io;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("could not read migrations directory `{path}`")]
ReadDirectory {
path: String,
#[source]
source: io::Error,
},
#[error(
"migration name `{name}` must be `<digits>_<description>.sql` with a `[a-z0-9_]` description"
)]
InvalidName { name: String },
#[error("migration version `{version}` is declared more than once")]
DuplicateVersion { version: u64 },
#[error("migrations directory `{path}` contains no migrations")]
Empty { path: String },
#[error("migration path `{path}` is not valid UTF-8")]
NonUtf8Path { path: String },
}

199
litellm-rust/crates/migrate-macros/src/lib.rs Normal file
View file

@ -0,0 +1,199 @@
mod error;
use std::path::{Path, PathBuf};
use error::Error;
use proc_macro::TokenStream;
use quote::quote;
use syn::LitStr;
struct Entry {
version: u64,
description: String,
path: PathBuf,
}
fn resolve(dir: &Path) -> Result<Vec<Entry>, Error> {
let mut entries = Vec::new();
let files = std::fs::read_dir(dir).map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
for file in files {
let file = file.map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
let path = file.path();
let name = path
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| Error::NonUtf8Path {
path: path.display().to_string(),
})?
.to_owned();
let invalid = || Error::InvalidName { name: name.clone() };
let stem = name
.strip_suffix(".sql")
.filter(|_| file.file_type().is_ok_and(|kind| kind.is_file()))
.and_then(|stem| stem.split_once('_'))
.filter(|(version, description)| {
!version.is_empty()
&& version.bytes().all(|b| b.is_ascii_digit())
&& !description.is_empty()
&& description
.bytes()
.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_')
})
.ok_or_else(invalid)?;
let version = stem.0.parse::<u64>().map_err(|_| invalid())?;
entries.push(Entry {
version,
description: stem.1.to_owned(),
path,
});
}
if entries.is_empty() {
return Err(Error::Empty {
path: dir.display().to_string(),
});
}
entries.sort_by_key(|entry| entry.version);
for pair in entries.windows(2) {
if pair[0].version == pair[1].version {
return Err(Error::DuplicateVersion {
version: pair[0].version,
});
}
}
Ok(entries)
}
fn resolve_input(lit: &LitStr) -> Result<Vec<Entry>, Error> {
let root = std::env::var("CARGO_MANIFEST_DIR")
.map(PathBuf::from)
.unwrap_or_default();
let dir = root.join(lit.value());
let dir = dir.canonicalize().map_err(|source| Error::ReadDirectory {
path: dir.display().to_string(),
source,
})?;
if dir.to_str().is_none() {
return Err(Error::NonUtf8Path {
path: dir.display().to_string(),
});
}
resolve(&dir)
}
#[proc_macro]
pub fn migrate(input: TokenStream) -> TokenStream {
let lit = syn::parse_macro_input!(input as LitStr);
match resolve_input(&lit) {
Ok(entries) => {
let migrations = entries.iter().map(|entry| {
let version = entry.version;
let description = &entry.description;
let path = entry
.path
.to_str()
.expect("canonical migration path is UTF-8");
quote! {
::litellm_migrate::Migration {
version: #version,
description: #description,
sql: ::core::include_str!(#path),
}
}
});
quote! { &[#(#migrations),*] }.into()
}
Err(err) => syn::Error::new(lit.span(), err).to_compile_error().into(),
}
}
#[cfg(test)]
mod tests {
use std::fs;
use rstest::rstest;
use tempfile::TempDir;
use super::{Error, resolve};
fn migrations_dir(files: &[&str]) -> TempDir {
let dir = TempDir::new().expect("tempdir");
for file in files {
fs::write(dir.path().join(file), "SELECT 1").expect("write fixture");
}
dir
}
#[rstest]
fn orders_versions_numerically() {
let dir = migrations_dir(&["10_tenth.sql", "2_second.sql", "1_first.sql"]);
let entries = resolve(dir.path()).expect("resolves");
let versions: Vec<u64> = entries.iter().map(|entry| entry.version).collect();
let descriptions: Vec<&str> = entries
.iter()
.map(|entry| entry.description.as_str())
.collect();
assert_eq!(versions, [1, 2, 10]);
assert_eq!(descriptions, ["first", "second", "tenth"]);
}
#[rstest]
#[case::dash_in_version(&["0001-dash.sql"])]
#[case::not_sql(&["notes.txt"])]
#[case::empty_description(&["0001_.sql"])]
#[case::non_digit_version(&["x_name.sql"])]
#[case::uppercase_description(&["0001_Upper.sql"])]
#[case::no_underscore(&["0001.sql"])]
#[case::plus_sign_version(&["+10_add.sql"])]
fn rejects_invalid_names(#[case] files: &[&str]) {
let dir = migrations_dir(files);
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[rstest]
fn rejects_subdirectories() {
let dir = migrations_dir(&["0001_a.sql"]);
fs::create_dir(dir.path().join("0002_b.sql")).expect("subdir");
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[cfg(unix)]
#[rstest]
fn rejects_symlinks() {
let dir = migrations_dir(&["0001_a.sql"]);
let target = TempDir::new().expect("tempdir");
let target_file = target.path().join("real.sql");
fs::write(&target_file, "SELECT 2").expect("write fixture");
std::os::unix::fs::symlink(&target_file, dir.path().join("0002_b.sql")).expect("symlink");
assert!(matches!(
resolve(dir.path()),
Err(Error::InvalidName { .. })
));
}
#[rstest]
fn rejects_duplicate_versions() {
let dir = migrations_dir(&["0001_a.sql", "1_b.sql"]);
assert!(matches!(
resolve(dir.path()),
Err(Error::DuplicateVersion { version: 1 })
));
}
#[rstest]
fn rejects_empty_directory() {
let dir = migrations_dir(&[]);
assert!(matches!(resolve(dir.path()), Err(Error::Empty { .. })));
}
}

12
litellm-rust/crates/migrate/Cargo.toml Normal file
View file

@ -0,0 +1,12 @@
[package]
name = "litellm-migrate"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-migrate-macros.workspace = true
[dev-dependencies]
rstest.workspace = true

5
litellm-rust/crates/migrate/README.md Normal file
View file

@ -0,0 +1,5 @@
# Migrations
`litellm-migrate` exports the `Migration` struct and the `migrate!` macro that embeds a directory of `<digits>_<description>.sql` files at compile time, sorted by numeric version
The crate does not apply or track migrations; callers decide how and when the embedded SQL runs

8
litellm-rust/crates/migrate/src/lib.rs Normal file
View file

@ -0,0 +1,8 @@
pub use litellm_migrate_macros::migrate;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Migration {
pub version: u64,
pub description: &'static str,
pub sql: &'static str,
}

1
litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 10;

1
litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 1;

1
litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql vendored Normal file
View file

@ -0,0 +1 @@
SELECT 2;

21
litellm-rust/crates/migrate/tests/migrate.rs Normal file
View file

@ -0,0 +1,21 @@
use litellm_migrate::Migration;
use rstest::rstest;
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("tests/fixtures/migrations");
#[rstest]
#[case::first(0, 1, "first", include_str!("fixtures/migrations/1_first.sql"))]
#[case::second(1, 2, "second", include_str!("fixtures/migrations/2_second.sql"))]
#[case::tenth(2, 10, "tenth", include_str!("fixtures/migrations/10_tenth.sql"))]
fn embeds_every_file_sorted_by_numeric_version(
#[case] index: usize,
#[case] version: u64,
#[case] description: &str,
#[case] sql: &str,
) {
assert_eq!(MIGRATIONS.len(), 3);
let migration = &MIGRATIONS[index];
assert_eq!(migration.version, version);
assert_eq!(migration.description, description);
assert_eq!(migration.sql, sql);
}

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

@ -44,7 +44,10 @@ mod _native {
#[pymodule_export]
use crate::routes::token_counter::TokenCounter;
#[pymodule_export]
use crate::routes::traces::{NativeTraceStorage, trace_decode_otlp, trace_encode_error};
use crate::routes::traces::{
NativeTraceConfig, NativeTraceStorage, trace_decode_otlp, trace_encode_error,
trace_normalized_field_definitions,
};
#[cfg(feature = "huggingface")]
#[pymodule_export]
use crate::tokenizer::HuggingFaceEncoding;
@ -109,9 +112,11 @@ mod tests {
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"NativeTraceConfig",
"NativeTraceStorage",
"trace_decode_otlp",
"trace_encode_error",
"trace_normalized_field_definitions",
"TokenCounter",
"Tokenizer",
"gil_stats",

141
litellm-rust/crates/python-bridge/src/routes/traces.rs
View file

@ -2,8 +2,10 @@ use std::collections::BTreeMap;
use litellm_host_python::{FromPythonCache, ToPythonCache};
use litellm_http::ClientVariant;
use litellm_storage_clickhouse::Storage;
use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared};
use litellm_traces::{
Config, Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery,
Shared,
};
use prost::Message;
use pyo3::{
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
@ -46,42 +48,64 @@ fn map_error(error: Error) -> PyErr {
}
}
fn map_sql_error(error: Error) -> PyErr {
match error {
Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
error => map_error(error),
}
}
fn map_query_access_error(error: QueryAccessError) -> PyErr {
match error {
QueryAccessError::Storage(error) => map_sql_error(error),
QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()),
error => PyRuntimeError::new_err(error.to_string()),
}
}
#[pyclass(frozen)]
pub struct NativeTraceConfig {
inner: Config,
}
#[pymethods]
impl NativeTraceConfig {
#[new]
fn new(database: String, url: &str, retention_days: u32) -> PyResult<Self> {
Ok(Self {
inner: Config::new(database, url, retention_days).map_err(map_error)?,
})
}
}
#[pyclass]
pub struct NativeTraceStorage {
storage: Storage,
config: Config,
query_readers: QueryReaders,
}
#[pymethods]
impl NativeTraceStorage {
#[new]
#[pyo3(signature = (database, url, reader_url = None))]
fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
fn new(config: PyRef<'_, NativeTraceConfig>) -> PyResult<Self> {
Ok(Self {
storage: Storage::new(database, url, reader_url).map_err(map_error)?,
query_readers: QueryReaders::new(
config.inner.storage().writer().clone(),
config.inner.storage().database().to_owned(),
),
config: config.inner.clone(),
})
}
fn ensure_schema<'py>(
&self,
py: Python<'py>,
trace_retention_days: u32,
spend_log_retention_days: u32,
) -> PyResult<Bound<'py, PyAny>> {
fn ensure_schema<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
let connection = self.config.storage().writer().clone();
let database = self.config.storage().database().to_owned();
let retention_days = self.config.retention_days();
crate::execution::run_async(
py,
async move {
litellm_traces::ensure_schema(
&client,
&connection,
&database,
trace_retention_days,
spend_log_retention_days,
)
.await
litellm_traces::ensure_schema(&client, &connection, &database, retention_days).await
},
map_error,
)
@ -95,8 +119,8 @@ impl NativeTraceStorage {
) -> PyResult<Bound<'py, PyAny>> {
let table = InsertTable::parse(table).map_err(map_error)?;
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
let connection = self.storage.writer().clone();
let database = self.storage.database().to_owned();
let connection = self.config.storage().writer().clone();
let database = self.config.storage().database().to_owned();
crate::execution::run_async(
py,
async move {
@ -107,6 +131,52 @@ impl NativeTraceStorage {
)
}
fn query_sql<'py>(
&self,
py: Python<'py>,
sql: String,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
secret: String,
) -> PyResult<Bound<'py, PyAny>> {
if sql.trim().is_empty() {
return Err(map_error(Error::EmptySql));
}
let readers = self.query_readers.clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
let _permit = readers.acquire()?;
let connection = readers.connection(&client, &scope, &secret).await?;
litellm_traces::query_sql(&client, &connection, &sql)
.await
.map_err(QueryAccessError::Storage)
},
map_query_access_error,
)
}
fn query_help<'py>(
&self,
py: Python<'py>,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
secret: String,
) -> PyResult<Bound<'py, PyAny>> {
let readers = self.query_readers.clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
let _permit = readers.acquire()?;
let connection = readers.connection(&client, &scope, &secret).await?;
litellm_traces::query_help(&client, &connection)
.await
.map_err(QueryAccessError::Storage)
},
map_query_access_error,
)
}
fn lens_query<'py>(
&self,
py: Python<'py>,
@ -117,9 +187,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let connection = self.config.storage().reader().clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
@ -140,9 +208,7 @@ impl NativeTraceStorage {
>,
) -> PyResult<Bound<'py, PyAny>> {
let query = ReadQuery::parse(query).map_err(map_error)?;
let connection = self.storage.reader().cloned().ok_or_else(|| {
PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
})?;
let connection = self.config.storage().reader().clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
@ -236,6 +302,14 @@ fn spans_to_py<'py>(
"events",
litellm_host_python::Pythonized(&span.events).into_pyobject(py)?,
)?;
row.set_item(
"normalized",
litellm_host_python::Pythonized(&span.normalized).into_pyobject(py)?,
)?;
row.set_item(
"consumed_attributes",
litellm_host_python::Pythonized(&span.consumed_attributes).into_pyobject(py)?,
)?;
result.append(row)?;
}
Ok(result)
@ -288,3 +362,8 @@ mod tests {
});
}
}
#[pyfunction]
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
litellm_host_python::Pythonized(litellm_traces::NORMALIZED_FIELD_DEFINITIONS).into_pyobject(py)
}

2
litellm-rust/crates/storage-clickhouse/README.md
View file

@ -1,5 +1,5 @@
# ClickHouse storage
`litellm-storage-clickhouse` exports `Storage`, a shared writer connection and optional reader connection for one ClickHouse database. It also exports bounded HTTP read and insert execution
`litellm-storage-clickhouse` exports `Storage`, a writer and bounded reader derived from one ClickHouse URL and database. It also exports bounded HTTP read and insert execution
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces` supplies those rules and uses this storage for both trace rows and spend rows

12
litellm-rust/crates/storage-clickhouse/src/lib.rs
View file

@ -89,19 +89,17 @@ impl Connection {
pub struct Storage {
database: String,
writer: Connection,
reader: Option<Connection>,
reader: Connection,
}
impl Storage {
pub fn new(database: String, url: &str, reader_url: Option<&str>) -> Result<Self, Error> {
pub fn new(database: String, url: &str) -> Result<Self, Error> {
if !valid_identifier(&database) {
return Err(Error::InvalidSchema);
}
Ok(Self {
writer: Connection::writer(url)?,
reader: reader_url
.map(|value| Connection::reader(value, &database))
.transpose()?,
reader: Connection::reader(url, &database)?,
database,
})
}
@ -114,8 +112,8 @@ impl Storage {
&self.writer
}
pub fn reader(&self) -> Option<&Connection> {
self.reader.as_ref()
pub fn reader(&self) -> &Connection {
&self.reader
}
}

27
litellm-rust/crates/storage-clickhouse/tests/connection.rs
View file

@ -10,25 +10,30 @@ fn accepts_only_clickhouse_http_urls(#[case] value: &str, #[case] expected: bool
assert_eq!(Connection::parse(value).is_ok(), expected);
}
#[rstest]
#[case::writer_only(None, false)]
#[case::separate_reader(Some("http://localhost:8124"), true)]
fn storage_exports_writer_and_optional_reader(
#[case] reader_url: Option<&str>,
#[case] has_reader: bool,
) {
let storage = Storage::new("litellm".to_owned(), "http://localhost:8123", reader_url)
.expect("valid ClickHouse URLs");
#[test]
fn storage_uses_one_url_for_writes_and_bounded_reads() {
let storage =
Storage::new("litellm".to_owned(), "http://localhost:8123").expect("valid ClickHouse URLs");
assert_eq!(storage.database(), "litellm");
assert_eq!(storage.writer().url().host_str(), Some("localhost"));
assert_eq!(storage.writer().url().port(), Some(8123));
assert_eq!(storage.reader().is_some(), has_reader);
assert_eq!(storage.reader().url().port(), Some(8123));
assert_eq!(
storage
.reader()
.url()
.query_pairs()
.find(|(key, _)| key == "database")
.unwrap()
.1,
"litellm"
);
}
#[rstest]
#[case::empty("")]
#[case::injection("db; DROP DATABASE default")]
fn storage_rejects_invalid_database(#[case] database: &str) {
assert!(Storage::new(database.to_owned(), "http://localhost:8123", None).is_err());
assert!(Storage::new(database.to_owned(), "http://localhost:8123").is_err());
}

2
litellm-rust/crates/traces/AGENTS.md
View file

@ -1,6 +1,6 @@
- Keep OTLP decoding, trace schema, row encoding and named query selection here. Generic ClickHouse connections and HTTP execution belong in `litellm-storage-clickhouse`
- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge`
- Keep the SQL migrations here as the only ClickHouse schema definition
- Keep the SQL migrations here as the only ClickHouse schema definition, as `migrations/NNNN_description.sql` files embedded by `litellm_migrate::migrate!`; adding a file is the only step
- Use typed query parameters and a dedicated SELECT-only reader with server-side limits
- Keep `config/reader.xml` grants on the database the schema is created in (CLICKHOUSE_DATABASE, default `litellm`)
- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions

8
litellm-rust/crates/traces/Cargo.toml
View file

@ -6,18 +6,26 @@ license.workspace = true
repository.workspace = true
[dependencies]
askama.workspace = true
base64.workspace = true
flate2.workspace = true
futures-util.workspace = true
hmac = "0.12.1"
indexmap = { version = "2", features = ["serde"] }
moka.workspace = true
opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
prost.workspace = true
time = { workspace = true, features = ["formatting"] }
litellm-http.workspace = true
litellm-migrate.workspace = true
litellm-storage-clickhouse.workspace = true
sha2.workspace = true
serde = { workspace = true, features = ["rc"] }
serde_json.workspace = true
strum.workspace = true
thiserror.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
criterion.workspace = true

3
litellm-rust/crates/traces/build.rs Normal file
View file

@ -0,0 +1,3 @@
fn main() {
println!("cargo:rerun-if-changed=migrations");
}

3
litellm-rust/crates/traces/migrations/0001_otel_traces.sql
View file

@ -38,10 +38,11 @@ CREATE TABLE IF NOT EXISTS {database}.otel_traces
Input String CODEC(ZSTD(3)),
Output String CODEC(ZSTD(3)),
InputPreview String DEFAULT substring(Input, 1, 240),
EngineReceivedMs UInt64 DEFAULT 0,
INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1,
INDEX idx_req_id LiteLLMRequestId TYPE bloom_filter(0.01) GRANULARITY 1
)
ENGINE = MergeTree
PARTITION BY toDate(Timestamp)
ORDER BY (TeamId, ServiceName, toDateTime(Timestamp), TraceId)
SETTINGS ttl_only_drop_parts = 1, non_replicated_deduplication_window = 1000
SETTINGS ttl_only_drop_parts = 1, materialize_ttl_recalculate_only = 1, non_replicated_deduplication_window = 1000

2
litellm-rust/crates/traces/migrations/0005_otel_traces_ttl.sql → litellm-rust/crates/traces/migrations/0002_otel_traces_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.otel_traces MODIFY TTL toDateTime(Timestamp) + INTERVAL {trace_retention_days} DAY
ALTER TABLE {database}.otel_traces MODIFY TTL toDateTime(Timestamp) + INTERVAL {retention_days} DAY

2
litellm-rust/crates/traces/migrations/0002_agent_traces.sql → litellm-rust/crates/traces/migrations/0003_agent_traces.sql
View file

@ -22,4 +22,4 @@ CREATE TABLE IF NOT EXISTS {database}.agent_traces_by_key
)
ENGINE = AggregatingMergeTree
ORDER BY (TeamId, ApiKeyHash, TraceId)
SETTINGS non_replicated_deduplication_window = 1000
SETTINGS materialize_ttl_recalculate_only = 1, non_replicated_deduplication_window = 1000

2
litellm-rust/crates/traces/migrations/0006_agent_traces_ttl.sql → litellm-rust/crates/traces/migrations/0004_agent_traces_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.agent_traces_by_key MODIFY TTL toDateTime(StartTs) + INTERVAL {trace_retention_days} DAY
ALTER TABLE {database}.agent_traces_by_key MODIFY TTL toDateTime(StartTs) + INTERVAL {retention_days} DAY

0
litellm-rust/crates/traces/migrations/0003_agent_traces_mv.sql → litellm-rust/crates/traces/migrations/0005_agent_traces_mv.sql
View file

2
litellm-rust/crates/traces/migrations/0004_spend_logs.sql → litellm-rust/crates/traces/migrations/0006_spend_logs.sql
View file

@ -34,9 +34,11 @@ CREATE TABLE IF NOT EXISTS {database}.spend_logs
metadata String CODEC(ZSTD(3)),
messages String CODEC(ZSTD(3)),
response String CODEC(ZSTD(3)),
EngineReceivedMs UInt64 DEFAULT 0,
INDEX idx_response_id response_id TYPE bloom_filter(0.001) GRANULARITY 1,
INDEX idx_trace_id trace_id TYPE bloom_filter(0.001) GRANULARITY 1
)
ENGINE = ReplacingMergeTree(end_time)
PARTITION BY toYYYYMM(start_time)
ORDER BY (team_id, start_time, request_id)
SETTINGS materialize_ttl_recalculate_only = 1

2
litellm-rust/crates/traces/migrations/0007_spend_logs_ttl.sql
View file

@ -1 +1 @@
ALTER TABLE {database}.spend_logs MODIFY TTL toDateTime(start_time) + INTERVAL {spend_log_retention_days} DAY
ALTER TABLE {database}.spend_logs MODIFY TTL toDateTime(start_time) + INTERVAL {retention_days} DAY

1
litellm-rust/crates/traces/migrations/0008_trace_received.sql
View file

@ -1 +0,0 @@
ALTER TABLE {database}.otel_traces ADD COLUMN IF NOT EXISTS EngineReceivedMs UInt64 DEFAULT 0

1
litellm-rust/crates/traces/migrations/0009_spend_received.sql
View file

@ -1 +0,0 @@
ALTER TABLE {database}.spend_logs ADD COLUMN IF NOT EXISTS EngineReceivedMs UInt64 DEFAULT 0

6
litellm-rust/crates/traces/query/lens_agents.sql Normal file
View file

@ -0,0 +1,6 @@
SELECT DISTINCT AgentName AS agent_name
FROM otel_traces
WHERE AgentName != ''
AND ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})
ORDER BY agent_name

8
litellm-rust/crates/traces/query/lens_availability.sql Normal file
View file

@ -0,0 +1,8 @@
SELECT
EXISTS(SELECT 1 FROM otel_traces
WHERE ({all_teams:UInt8}=1 OR TeamId={team:String})
AND ({key_hash:String}='' OR ApiKeyHash={key_hash:String})) AS traces,
EXISTS(SELECT 1 FROM spend_logs
WHERE ({all_teams:UInt8}=1 OR team_id={team:String})
AND ({key_hash:String}='' OR api_key={key_hash:String})
AND NOT JSONExtractBool(metadata,'litellm_lens_internal')) AS requests

2
litellm-rust/crates/traces/query/lens_sample.sql
View file

@ -28,6 +28,7 @@ SELECT *, selection_key FROM (
GROUP BY TeamId,ApiKeyHash,TraceId
HAVING max(EngineReceivedMs) < {end:UInt64}
AND max(toUnixTimestamp64Milli(Timestamp)+toInt64(intDiv(Duration,1000000))) < {end:UInt64}
AND ({agent_name:String}='' OR countIf(AgentName={agent_name:String}) > 0)
AND countIf(arrayAll((k,v) -> ResourceAttributes[k]=v OR SpanAttributes[k]=v,
{filter_keys:Array(String)},{filter_values:Array(String)})
AND ({service:String}='' OR ServiceName={service:String})) > 0
@ -48,6 +49,7 @@ SELECT *, selection_key FROM (
OR JSONExtractString(metadata,'requester_metadata',k)=v OR (k='tag' AND has(request_tags,v)),
{filter_keys:Array(String)},{filter_values:Array(String)})
AND ({service:String}='' OR model_group={service:String})
AND {agent_name:String}=''
AND NOT JSONExtractBool(metadata,'litellm_lens_internal')
AND ({source:String}!='both' OR (team_id,api_key,response_id) NOT IN (
SELECT TeamId,ApiKeyHash,LiteLLMRequestId FROM otel_traces

25
litellm-rust/crates/traces/src/config.rs Normal file
View file

@ -0,0 +1,25 @@
use litellm_storage_clickhouse::{Error, Storage};
#[derive(Clone)]
pub struct Config {
storage: Storage,
retention_days: u32,
}
impl Config {
pub fn new(database: String, url: &str, retention_days: u32) -> Result<Self, Error> {
crate::schema_statements(&database, retention_days)?;
Ok(Self {
storage: Storage::new(database, url)?,
retention_days,
})
}
pub fn storage(&self) -> &Storage {
&self.storage
}
pub fn retention_days(&self) -> u32 {
self.retention_days
}
}

22
litellm-rust/crates/traces/src/error.rs
View file

@ -4,4 +4,26 @@ pub enum DecodeError {
InvalidPayload,
#[error("OTLP trace payload exceeds the decoding budget")]
TooLarge,
#[error("OTLP token count is outside the storage range")]
TokenCountOutOfRange,
}
#[derive(Debug, thiserror::Error)]
pub enum QueryAccessError {
#[error("trace SQL queries require a configured proxy master key")]
MissingSecret,
#[error("invalid trace query scope")]
InvalidScope,
#[error("trace SQL query concurrency limit exceeded")]
Busy,
#[error(
"ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables"
)]
ProvisionFailed(u16),
#[error("ClickHouse reader provisioning transport failed")]
ProvisionTransport,
#[error(transparent)]
Storage(#[from] litellm_storage_clickhouse::Error),
#[error(transparent)]
Cached(#[from] std::sync::Arc<QueryAccessError>),
}

13
litellm-rust/crates/traces/src/lib.rs
View file

@ -1,14 +1,25 @@
mod config;
mod error;
mod insert;
mod normalize;
mod otlp;
mod query;
mod query_access;
mod schema;
mod shared;
mod sql;
pub use error::DecodeError;
pub use config::Config;
pub use error::{DecodeError, QueryAccessError};
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
pub use normalize::{
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType,
};
pub use otlp::{DecodedSpan, decode_otlp};
pub use query_access::{QueryReaders, QueryScope};
pub use schema::{ensure_schema, schema_statements};
pub use shared::{Shared, SharedIdentity};
pub use sql::{LensQuery, ReadQuery, execute_named_read};
pub use query::{query_help, query_sql};

63
litellm-rust/crates/traces/src/normalize/genai.rs Normal file
View file

@ -0,0 +1,63 @@
use std::collections::BTreeMap;
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, usage_tokens};
use crate::DecodeError;
pub(super) struct GenAiNormalizer;
impl SpanNormalizer for GenAiNormalizer {
fn matches(&self, _scope_name: &str, _attributes: &BTreeMap<String, String>) -> bool {
true
}
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
[
if attr(attributes, "gen_ai.input.messages").is_empty() {
"gen_ai.tool.call.arguments"
} else {
"gen_ai.input.messages"
},
if attr(attributes, "gen_ai.output.messages").is_empty() {
"gen_ai.tool.call.result"
} else {
"gen_ai.output.messages"
},
]
}
fn normalize(
&self,
_name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> Result<NormalizedSpan, DecodeError> {
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
let observation_type = match attr(attributes, "gen_ai.operation.name") {
"invoke_agent" => ObservationType::Agent,
"chat" | "text_completion" | "generate_content" => ObservationType::Llm,
"execute_tool" => ObservationType::Tool,
_ if parent_span_id.is_empty() => ObservationType::Agent,
_ => ObservationType::Chain,
};
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "gen_ai.agent.name").to_owned(),
litellm_request_id: attr(attributes, "gen_ai.response.id").to_owned(),
model: first(attributes, "gen_ai.request.model", "gen_ai.response.model").to_owned(),
input_tokens,
output_tokens,
input: first(
attributes,
"gen_ai.input.messages",
"gen_ai.tool.call.arguments",
)
.to_owned(),
output: first(
attributes,
"gen_ai.output.messages",
"gen_ai.tool.call.result",
)
.to_owned(),
})
}
}

468
litellm-rust/crates/traces/src/normalize/langsmith.rs Normal file
View file

@ -0,0 +1,468 @@
use std::{collections::BTreeMap, io};
use indexmap::IndexMap;
use serde::{Deserialize, Deserializer, Serialize, de::DeserializeOwned};
use serde_json::{Value, ser::Formatter};
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, usage_tokens};
use crate::DecodeError;
pub(super) struct LangSmithNormalizer;
#[derive(Deserialize)]
#[serde(untagged)]
enum MessageContent {
Text(String),
Blocks(Vec<ContentBlock>),
Other(Value),
}
impl MessageContent {
fn display_text(&self) -> String {
match self {
Self::Text(text) => text.clone(),
Self::Blocks(blocks) => blocks
.iter()
.filter_map(|block| match block {
ContentBlock::Text { text } => Some(text.as_str()),
ContentBlock::Hidden(kind) => match kind {
HiddenBlock::Reasoning
| HiddenBlock::Thinking
| HiddenBlock::RedactedThinking
| HiddenBlock::FunctionCall
| HiddenBlock::ToolUse
| HiddenBlock::ToolCall => None,
},
})
.collect::<Vec<_>>()
.join("\n\n"),
Self::Other(value) => encode(value),
}
}
}
#[derive(Deserialize)]
#[serde(untagged)]
enum ContentBlock {
Text { text: String },
Hidden(HiddenBlock),
}
#[derive(Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
enum HiddenBlock {
Reasoning,
Thinking,
RedactedThinking,
FunctionCall,
ToolUse,
ToolCall,
}
#[derive(Deserialize, Serialize)]
#[serde(transparent)]
struct RawToolCall(IndexMap<String, Value>);
#[derive(Deserialize)]
struct ResponseMetadata {
id: Option<String>,
}
#[derive(Deserialize)]
struct RawMessage {
kwargs: Option<Box<RawMessage>>,
#[serde(rename = "type")]
kind: Option<String>,
role: Option<String>,
content: Option<MessageContent>,
tool_calls: Option<Vec<RawToolCall>>,
name: Option<Value>,
response_metadata: Option<ResponseMetadata>,
}
impl RawMessage {
fn unwrapped(&self) -> &Self {
self.kwargs.as_deref().unwrap_or(self)
}
fn normalized(&self) -> NormalizedMessage<'_> {
let fields = self.unwrapped();
let raw_role = fields
.kind
.as_deref()
.filter(|role| !role.is_empty())
.or_else(|| fields.role.as_deref().filter(|role| !role.is_empty()))
.unwrap_or_default();
let role = match raw_role {
"human" => "user",
"ai" => "assistant",
other => other,
};
NormalizedMessage {
role,
content: fields
.content
.as_ref()
.map_or_else(String::new, MessageContent::display_text),
tool_calls: fields
.tool_calls
.as_deref()
.filter(|calls| !calls.is_empty()),
name: (role == "tool")
.then_some(fields.name.as_ref())
.flatten()
.filter(|name| !name.is_null() && name != &&Value::String(String::new())),
}
}
}
#[derive(Serialize)]
struct NormalizedMessage<'a> {
role: &'a str,
content: String,
#[serde(skip_serializing_if = "Option::is_none")]
tool_calls: Option<&'a [RawToolCall]>,
#[serde(skip_serializing_if = "Option::is_none")]
name: Option<&'a Value>,
}
enum MessageBatch {
Flat(Vec<RawMessage>),
Nested(Vec<Vec<RawMessage>>),
}
impl<'de> Deserialize<'de> for MessageBatch {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let value = Value::deserialize(deserializer)?;
let Value::Array(items) = value else {
return Err(serde::de::Error::custom("messages must be an array"));
};
let parse = |items: Vec<Value>| {
items
.into_iter()
.filter_map(|item| serde_json::from_value(item).ok())
.collect()
};
Ok(if items.first().is_some_and(Value::is_array) {
Self::Nested(
items
.into_iter()
.filter_map(|item| item.as_array().cloned())
.map(parse)
.collect(),
)
} else {
Self::Flat(parse(items))
})
}
}
fn lenient<'de, D: Deserializer<'de>, T: DeserializeOwned>(
deserializer: D,
) -> Result<Option<T>, D::Error> {
let value = Value::deserialize(deserializer)?;
Ok(serde_json::from_value(value).ok())
}
impl MessageBatch {
fn first_batch(&self) -> &[RawMessage] {
match self {
Self::Flat(messages) => messages,
Self::Nested(batches) => batches.first().map(Vec::as_slice).unwrap_or_default(),
}
}
fn agent_messages(&self) -> &[RawMessage] {
match self {
Self::Flat(messages) => messages,
Self::Nested(_) => &[],
}
}
}
#[derive(Deserialize)]
struct GenerationMessage {
kwargs: Option<RawMessage>,
}
#[derive(Deserialize)]
struct Generation {
message: Option<GenerationMessage>,
}
#[derive(Default, Deserialize)]
struct Payload {
#[serde(default, deserialize_with = "lenient")]
messages: Option<MessageBatch>,
#[serde(default, deserialize_with = "lenient")]
generations: Option<Vec<Vec<Generation>>>,
}
#[derive(Deserialize)]
struct Command {
update: CommandUpdate,
}
#[derive(Deserialize)]
struct CommandUpdate {
messages: Vec<Value>,
}
#[derive(Deserialize)]
struct ContentValue {
content: Value,
}
struct SpanIo {
input: String,
output: String,
request_id: String,
}
struct PythonJsonFormatter;
impl Formatter for PythonJsonFormatter {
fn begin_array_value<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_key<W: ?Sized + io::Write>(
&mut self,
writer: &mut W,
first: bool,
) -> io::Result<()> {
if first {
Ok(())
} else {
writer.write_all(b", ")
}
}
fn begin_object_value<W: ?Sized + io::Write>(&mut self, writer: &mut W) -> io::Result<()> {
writer.write_all(b": ")
}
}
fn encode<T: Serialize>(value: &T) -> String {
let mut output = Vec::new();
let mut serializer = serde_json::Serializer::with_formatter(&mut output, PythonJsonFormatter);
if value.serialize(&mut serializer).is_err() {
return String::new();
}
String::from_utf8(output).unwrap_or_default()
}
fn normalized_messages(messages: &[RawMessage]) -> String {
encode(
&messages
.iter()
.map(RawMessage::normalized)
.collect::<Vec<_>>(),
)
}
fn span_type(
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> ObservationType {
match attr(attributes, "langsmith.span.kind") {
"llm" => ObservationType::Llm,
"tool" => ObservationType::Tool,
_ if parent_span_id.is_empty()
|| name == attr(attributes, "langsmith.metadata.lc_agent_name") =>
{
ObservationType::Agent
}
_ if [
".wrap_model_call",
".wrap_tool_call",
".before_agent",
".after_agent",
".before_model",
".after_model",
]
.iter()
.any(|suffix| name.ends_with(suffix)) =>
{
ObservationType::Framework
}
_ => ObservationType::Chain,
}
}
fn tool_output(raw_completion: &str) -> String {
let completion = serde_json::from_str::<Value>(raw_completion).unwrap_or(Value::Null);
let raw = completion.get("output").cloned().unwrap_or(completion);
let selected = serde_json::from_value::<Command>(raw.clone())
.ok()
.and_then(|command| command.update.messages.into_iter().last())
.unwrap_or(raw);
let output = serde_json::from_value::<ContentValue>(selected.clone())
.map(|message| message.content)
.unwrap_or(selected);
output
.as_str()
.map(str::to_owned)
.unwrap_or_else(|| encode(&output))
}
fn span_io(kind: ObservationType, attributes: &BTreeMap<String, String>) -> SpanIo {
let raw_prompt = attr(attributes, "gen_ai.prompt");
let raw_completion = attr(attributes, "gen_ai.completion");
let prompt = serde_json::from_str::<Payload>(raw_prompt).unwrap_or_default();
let completion = serde_json::from_str::<Payload>(raw_completion).unwrap_or_default();
if kind == ObservationType::Llm
&& serde_json::from_str::<Value>(raw_completion).is_ok_and(|value| value.is_object())
{
let input = prompt.messages.as_ref().map_or_else(
|| "[]".to_owned(),
|messages| normalized_messages(messages.first_batch()),
);
let generation = completion
.generations
.as_ref()
.and_then(|batches| batches.first())
.and_then(|batch| batch.first())
.and_then(|generation| generation.message.as_ref())
.and_then(|message| message.kwargs.as_ref());
if let Some(generation) = generation {
let id = generation
.response_metadata
.as_ref()
.and_then(|metadata| metadata.id.as_deref())
.unwrap_or_default()
.to_owned();
return SpanIo {
input,
output: encode(&generation.normalized()),
request_id: id,
};
}
return SpanIo {
input,
output: raw_completion.to_owned(),
request_id: String::new(),
};
}
if kind == ObservationType::Tool {
return SpanIo {
input: raw_prompt.to_owned(),
output: tool_output(raw_completion),
request_id: String::new(),
};
}
if kind == ObservationType::Agent {
let input = prompt
.messages
.as_ref()
.filter(|messages| !messages.agent_messages().is_empty())
.map_or_else(
|| raw_prompt.to_owned(),
|messages| normalized_messages(messages.agent_messages()),
);
let output = completion
.messages
.as_ref()
.and_then(|messages| messages.agent_messages().last())
.map_or_else(
|| raw_completion.to_owned(),
|message| encode(&message.normalized()),
);
return SpanIo {
input,
output,
request_id: String::new(),
};
}
SpanIo {
input: raw_prompt.to_owned(),
output: raw_completion.to_owned(),
request_id: String::new(),
}
}
impl SpanNormalizer for LangSmithNormalizer {
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool {
scope_name == "langsmith" || attributes.contains_key("langsmith.span.kind")
}
fn consumed_attributes(&self, _attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
["gen_ai.prompt", "gen_ai.completion"]
}
fn normalize(
&self,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> Result<NormalizedSpan, DecodeError> {
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
let observation_type = span_type(name, parent_span_id, attributes);
let io = span_io(observation_type, attributes);
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "langsmith.metadata.lc_agent_name").to_owned(),
litellm_request_id: io.request_id,
model: attr(attributes, "gen_ai.request.model").to_owned(),
input_tokens,
output_tokens,
input: io.input,
output: io.output,
})
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use rstest::rstest;
use serde_json::Value;
use super::{ObservationType, span_io};
#[rstest]
fn malformed_messages_preserve_valid_input_and_response_id() {
let attributes = BTreeMap::from([
(
"gen_ai.prompt".to_owned(),
r#"{"messages":[[{"kwargs":{"type":"human","content":"hello"}},null]]}"#.to_owned(),
),
(
"gen_ai.completion".to_owned(),
r#"{"messages":"unexpected","generations":[[{"message":{"kwargs":{"type":"ai","content":"hi","response_metadata":{"id":"response-1"}}}}]]}"#.to_owned(),
),
]);
let io = span_io(ObservationType::Llm, &attributes);
let input: Value = serde_json::from_str(&io.input).expect("normalized input");
assert_eq!(input.as_array().expect("messages").len(), 1);
assert_eq!(input[0]["content"], "hello");
assert_eq!(io.request_id, "response-1");
}
#[rstest]
fn explicit_null_tool_output_is_preserved() {
let attributes = BTreeMap::from([(
"gen_ai.completion".to_owned(),
r#"{"output":null}"#.to_owned(),
)]);
let io = span_io(ObservationType::Tool, &attributes);
assert_eq!(io.output, "null");
}
#[rstest]
fn absent_llm_messages_render_as_an_empty_list() {
let attributes = BTreeMap::from([("gen_ai.completion".to_owned(), "{}".to_owned())]);
let io = span_io(ObservationType::Llm, &attributes);
assert_eq!(io.input, "[]");
}
}

239
litellm-rust/crates/traces/src/normalize/mod.rs Normal file
View file

@ -0,0 +1,239 @@
use std::collections::BTreeMap;
use crate::DecodeError;
use serde::Serialize;
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum ObservationType {
Agent,
Llm,
Tool,
Chain,
Framework,
}
#[derive(Debug, Serialize)]
pub struct NormalizedSpan {
pub observation_type: ObservationType,
pub agent_name: String,
pub litellm_request_id: String,
pub model: String,
pub input_tokens: u32,
pub output_tokens: u32,
pub input: String,
pub output: String,
}
pub(crate) struct Normalization {
pub span: NormalizedSpan,
pub consumed_attributes: [&'static str; 2],
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
pub struct NormalizedFieldDefinition {
pub name: &'static str,
pub clickhouse_column: &'static str,
pub clickhouse_type: &'static str,
pub meaning: &'static str,
}
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 8] = [
NormalizedFieldDefinition {
name: "observation_type",
clickhouse_column: "ObservationType",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent, LLM, tool, chain, or framework span",
},
NormalizedFieldDefinition {
name: "agent_name",
clickhouse_column: "AgentName",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent associated with this span",
},
NormalizedFieldDefinition {
name: "litellm_request_id",
clickhouse_column: "LiteLLMRequestId",
clickhouse_type: "String",
meaning: "LiteLLM response ID used to link a span to a spend log",
},
NormalizedFieldDefinition {
name: "model",
clickhouse_column: "Model",
clickhouse_type: "LowCardinality(String)",
meaning: "Model used by this span",
},
NormalizedFieldDefinition {
name: "input_tokens",
clickhouse_column: "InputTokens",
clickhouse_type: "UInt32",
meaning: "Input token count",
},
NormalizedFieldDefinition {
name: "output_tokens",
clickhouse_column: "OutputTokens",
clickhouse_type: "UInt32",
meaning: "Output token count",
},
NormalizedFieldDefinition {
name: "input",
clickhouse_column: "Input",
clickhouse_type: "String",
meaning: "Normalized input payload",
},
NormalizedFieldDefinition {
name: "output",
clickhouse_column: "Output",
clickhouse_type: "String",
meaning: "Normalized output payload",
},
];
trait SpanNormalizer {
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool;
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2];
fn normalize(
&self,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> Result<NormalizedSpan, DecodeError>;
}
mod genai;
mod langsmith;
mod openinference;
use genai::GenAiNormalizer;
use langsmith::LangSmithNormalizer;
use openinference::OpenInferenceNormalizer;
fn attr<'a>(attributes: &'a BTreeMap<String, String>, key: &str) -> &'a str {
attributes.get(key).map(String::as_str).unwrap_or_default()
}
fn first<'a>(attributes: &'a BTreeMap<String, String>, left: &str, right: &str) -> &'a str {
let value = attr(attributes, left);
if value.is_empty() {
attr(attributes, right)
} else {
value
}
}
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, DecodeError> {
let value = attr(attributes, key).trim();
if value.is_empty() {
return Ok(0);
}
match value.parse::<i128>() {
Ok(number) if (0..=u32::MAX as i128).contains(&number) => Ok(number as u32),
Ok(_) => Err(DecodeError::TokenCountOutOfRange),
Err(_)
if value
.trim_start_matches(['+', '-'])
.bytes()
.all(|byte| byte.is_ascii_digit()) =>
{
Err(DecodeError::TokenCountOutOfRange)
}
Err(_) => Ok(0),
}
}
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), DecodeError> {
Ok((
tokens(attributes, "gen_ai.usage.input_tokens")?,
tokens(attributes, "gen_ai.usage.output_tokens")?,
))
}
pub fn normalize(
scope_name: &str,
name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> Result<Normalization, DecodeError> {
let normalizers: [&dyn SpanNormalizer; 3] = [
&LangSmithNormalizer,
&OpenInferenceNormalizer,
&GenAiNormalizer,
];
let normalizer = normalizers
.into_iter()
.find(|normalizer| normalizer.matches(scope_name, attributes))
.expect("GenAI fallback always matches");
Ok(Normalization {
span: normalizer.normalize(name, parent_span_id, attributes)?,
consumed_attributes: normalizer.consumed_attributes(attributes),
})
}
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use rstest::rstest;
use super::{NORMALIZED_FIELD_DEFINITIONS, ObservationType, normalize};
#[rstest]
#[case::langsmith("langsmith", [("langsmith.span.kind", "llm"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
#[case::openinference("other", [("openinference.span.kind", "LLM"), ("gen_ai.operation.name", "execute_tool")], ObservationType::Llm)]
#[case::genai("other", [("gen_ai.operation.name", "execute_tool"), ("gen_ai.usage.input_tokens", "7")], ObservationType::Tool)]
fn convention_dispatch_preserves_precedence(
#[case] scope: &str,
#[case] attributes: [(&str, &str); 2],
#[case] expected: ObservationType,
) {
let attributes = attributes
.into_iter()
.map(|(key, value)| (key.to_owned(), value.to_owned()))
.collect();
let fields = normalize(scope, "step", "parent", &attributes)
.expect("valid tokens")
.span;
assert_eq!(fields.observation_type, expected);
if expected == ObservationType::Tool {
assert_eq!(fields.input_tokens, 7);
}
}
#[rstest]
fn field_definitions_match_serialized_normalized_span() {
let fields = normalize("", "root", "", &BTreeMap::new())
.expect("valid tokens")
.span;
let serialized = serde_json::to_value(fields).expect("serializable fields");
let keys: BTreeSet<_> = serialized
.as_object()
.expect("field object")
.keys()
.map(String::as_str)
.collect();
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
.iter()
.map(|field| field.name)
.collect();
assert_eq!(keys, mapped);
}
#[rstest]
fn token_counts_accept_surrounding_whitespace() {
let attributes =
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), " 7 ".to_owned())]);
let fields = normalize("", "root", "", &attributes)
.expect("valid tokens")
.span;
assert_eq!(fields.input_tokens, 7);
}
#[rstest]
#[case::negative("-1")]
#[case::overflow("4294967296")]
fn token_counts_outside_storage_range_are_rejected(#[case] value: &str) {
let attributes =
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), value.to_owned())]);
assert!(normalize("", "root", "", &attributes).is_err());
}
}

53
litellm-rust/crates/traces/src/normalize/openinference.rs Normal file
View file

@ -0,0 +1,53 @@
use std::collections::BTreeMap;
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, tokens, usage_tokens};
use crate::DecodeError;
pub(super) struct OpenInferenceNormalizer;
impl SpanNormalizer for OpenInferenceNormalizer {
fn matches(&self, _scope_name: &str, attributes: &BTreeMap<String, String>) -> bool {
attributes.contains_key("openinference.span.kind")
}
fn consumed_attributes(&self, _attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
["input.value", "output.value"]
}
fn normalize(
&self,
_name: &str,
parent_span_id: &str,
attributes: &BTreeMap<String, String>,
) -> Result<NormalizedSpan, DecodeError> {
let (usage_input, usage_output) = usage_tokens(attributes)?;
let observation_type = match attr(attributes, "openinference.span.kind")
.to_ascii_uppercase()
.as_str()
{
"AGENT" => ObservationType::Agent,
"LLM" => ObservationType::Llm,
"TOOL" => ObservationType::Tool,
_ if parent_span_id.is_empty() => ObservationType::Agent,
_ => ObservationType::Chain,
};
Ok(NormalizedSpan {
observation_type,
agent_name: attr(attributes, "agent.name").to_owned(),
litellm_request_id: String::new(),
model: attr(attributes, "llm.model_name").to_owned(),
input_tokens: if attributes.contains_key("llm.token_count.prompt") {
tokens(attributes, "llm.token_count.prompt")?
} else {
usage_input
},
output_tokens: if attributes.contains_key("llm.token_count.completion") {
tokens(attributes, "llm.token_count.completion")?
} else {
usage_output
},
input: attr(attributes, "input.value").to_owned(),
output: attr(attributes, "output.value").to_owned(),
})
}
}

4
litellm-rust/crates/traces/src/otlp/mod.rs
View file

@ -6,7 +6,7 @@ mod wire;
use serde::Serialize;
use std::collections::BTreeMap;
use crate::{DecodeError, Shared};
use crate::{DecodeError, NormalizedSpan, Shared};
#[derive(Serialize)]
pub struct DecodedEvent {
@ -31,6 +31,8 @@ pub struct DecodedSpan {
pub status_code: String,
pub status_message: String,
pub events: Vec<DecodedEvent>,
pub normalized: NormalizedSpan,
pub consumed_attributes: [&'static str; 2],
}
pub fn decode_otlp(

24
litellm-rust/crates/traces/src/otlp/span.rs
View file

@ -10,7 +10,7 @@ use super::{
attributes::attributes,
limits::{Budget, MAX_ATTRIBUTES, MAX_DECODED_SPAN_BYTES, MAX_EVENTS, MAX_SPANS},
};
use crate::{DecodeError, Shared};
use crate::{DecodeError, Shared, normalize::normalize};
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, DecodeError> {
let mut budget = Budget::new(MAX_DECODED_SPAN_BYTES);
@ -125,10 +125,26 @@ fn decoded_span(
budget: &mut Budget,
) -> Result<DecodedSpan, DecodeError> {
let status = span.status.unwrap_or_default();
let parent_span_id = hex_bytes(&span.parent_span_id);
let span_attributes = attributes(span.attributes, budget)?;
let normalization = normalize(
scope_name.as_ref(),
&span.name,
&parent_span_id,
&span_attributes,
)?;
let normalized = normalization.span;
budget.consume(
normalized.input.len()
+ normalized.output.len()
+ normalized.agent_name.len()
+ normalized.litellm_request_id.len()
+ normalized.model.len(),
)?;
Ok(DecodedSpan {
trace_id: hex_bytes(&span.trace_id),
span_id: hex_bytes(&span.span_id),
parent_span_id: hex_bytes(&span.parent_span_id),
parent_span_id,
trace_state: span.trace_state,
name: span.name,
kind: SpanKind::try_from(span.kind)
@ -143,7 +159,7 @@ fn decoded_span(
})?,
scope_name: budget.clone_shared(scope_name, String::len)?,
scope_version: budget.clone_shared(scope_version, String::len)?,
attributes: attributes(span.attributes, budget)?,
attributes: span_attributes,
start_ns: span.start_time_unix_nano,
end_ns: span.end_time_unix_nano,
status_code: StatusCode::try_from(status.code)
@ -162,5 +178,7 @@ fn decoded_span(
})
})
.collect::<Result<Vec<_>, DecodeError>>()?,
normalized,
consumed_attributes: normalization.consumed_attributes,
})
}

346
litellm-rust/crates/traces/src/query.rs Normal file
View file

@ -0,0 +1,346 @@
use std::collections::{BTreeMap, BTreeSet};
use futures_util::{
StreamExt,
stream::{self, TryStreamExt},
};
use litellm_http::Client;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use crate::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, execute_read};
mod guide;
const SAMPLE_ROWS: usize = 200;
const MAX_FIELDS: usize = 200;
const MAX_DEPTH: usize = 16;
const METADATA_SQL: &str = "SELECT metadata FROM spend_logs FINAL \
WHERE start_time >= now() - INTERVAL 7 DAY AND length(metadata) <= 8192 \
LIMIT 201";
const METADATA_SCOPE: &str = "Up to 200 unordered rows from the last 7 days, excluding metadata larger than 8192 bytes; up to 200 paths and 16 levels. Missing paths may exist outside this sample. Array indexes are 1-based and describe sampled positions, not a fixed schema";
const ATTRIBUTE_SCOPE: &str = "Distinct keys from up to 200 unordered spans in the last 7 days; up to 200 keys per map. Missing keys may exist outside this sample";
#[derive(Deserialize)]
struct Rows<T> {
data: Vec<T>,
}
#[derive(Deserialize)]
struct MetadataRow {
metadata: String,
}
#[derive(Deserialize)]
struct AttributeRow {
key: String,
}
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
#[serde(untagged)]
enum PathPart {
Key(String),
Index(usize),
}
#[derive(Serialize)]
struct MetadataField {
path: Vec<PathPart>,
types: BTreeSet<&'static str>,
expression: String,
}
#[derive(Deserialize, Serialize)]
struct ColumnSchema {
name: String,
#[serde(rename = "type")]
kind: String,
#[serde(flatten)]
details: BTreeMap<String, Value>,
}
#[derive(Serialize)]
struct TableSchema {
name: &'static str,
columns: Vec<ColumnSchema>,
}
#[derive(Serialize)]
struct MetadataCatalog {
table: &'static str,
column: &'static str,
fields: Vec<MetadataField>,
sampled_rows: usize,
invalid_json_rows: usize,
truncated: bool,
sample_sql: &'static str,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
#[derive(Serialize)]
struct AttributeField {
key: String,
#[serde(rename = "type")]
kind: &'static str,
expression: String,
}
#[derive(Serialize)]
struct AttributeCatalog {
table: &'static str,
column: &'static str,
fields: Vec<AttributeField>,
truncated: bool,
discovery_sql: String,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
pub async fn query_sql(
client: &Client,
connection: &Connection,
sql: &str,
) -> Result<String, Error> {
execute_read(client, connection, sql, &BTreeMap::new()).await
}
async fn rows<T: serde::de::DeserializeOwned>(
client: &Client,
connection: &Connection,
sql: &str,
) -> Result<Vec<T>, Error> {
let body = query_sql(client, connection, sql).await?;
serde_json::from_str::<Rows<T>>(&body)
.map(|result| result.data)
.map_err(|_| Error::InvalidResponse)
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}
fn metadata_expression(path: &[PathPart]) -> String {
let arguments = path
.iter()
.map(|part| match part {
PathPart::Key(key) => literal(key),
PathPart::Index(index) => index.to_string(),
})
.collect::<Vec<_>>()
.join(", ");
format!("JSONExtractRaw(metadata, {arguments})")
}
fn discover(
value: &Value,
path: Vec<PathPart>,
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
) -> bool {
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
return true;
}
if !path.is_empty() {
let kind = match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
};
fields.entry(path.clone()).or_default().insert(kind);
}
match value {
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
let child = path
.iter()
.cloned()
.chain([PathPart::Key(key.clone())])
.collect();
discover(value, child, fields) | limited
}),
Value::Array(array) => array
.iter()
.enumerate()
.fold(false, |limited, (index, value)| {
let child = path
.iter()
.cloned()
.chain([PathPart::Index(index + 1)])
.collect();
discover(value, child, fields) | limited
}),
_ => false,
}
}
fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold(
(BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0),
|(fields, limited, invalid_rows), row| match serde_json::from_str::<Value>(&row.metadata) {
Ok(value) => {
let mut fields = fields;
let limited = limited | discover(&value, Vec::new(), &mut fields);
(fields, limited, invalid_rows)
}
Err(_) => (fields, limited, invalid_rows + 1),
},
);
let fields: Vec<_> = fields
.into_iter()
.map(|(path, types)| MetadataField {
expression: metadata_expression(&path),
path,
types,
})
.collect();
MetadataCatalog {
table: "spend_logs",
column: "metadata",
fields,
sampled_rows: sample.len().min(SAMPLE_ROWS),
invalid_json_rows: invalid_rows,
truncated: limited,
sample_sql: METADATA_SQL,
error: None,
scope: METADATA_SCOPE,
}
}
pub async fn query_help(client: &Client, connection: &Connection) -> Result<String, Error> {
let tables = stream::iter(["otel_traces", "agent_traces_by_key", "spend_logs"])
.then(|table| async move {
Ok::<_, Error>(TableSchema {
name: table,
columns: rows::<ColumnSchema>(
client,
connection,
&format!("DESCRIBE TABLE {table}"),
)
.await?,
})
})
.try_collect::<Vec<_>>()
.await?;
let metadata = match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
Ok(sample) => metadata_catalog(&sample),
Err(error) => MetadataCatalog {
error: Some(error.to_string()),
truncated: true,
..metadata_catalog(&[])
},
};
let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
.then(|column| async move {
let sql = format!(
"SELECT DISTINCT arrayJoin(mapKeys({column})) AS key FROM \
(SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \
LIMIT 200) ORDER BY key LIMIT 201"
);
let (keys, error) = match rows::<AttributeRow>(client, connection, &sql).await {
Ok(keys) => (keys, None),
Err(error) => (Vec::new(), Some(error.to_string())),
};
let fields = keys
.iter()
.take(MAX_FIELDS)
.map(|row| AttributeField {
key: row.key.clone(),
kind: "String",
expression: format!("{column}[{}]", literal(&row.key)),
})
.collect();
AttributeCatalog {
table: "otel_traces",
column,
fields,
truncated: error.is_some() || keys.len() > MAX_FIELDS,
discovery_sql: sql,
scope: ATTRIBUTE_SCOPE,
error,
}
})
.collect::<Vec<_>>()
.await;
let guide = guide::QueryGuide {
tables: &tables,
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
metadata: &metadata,
attributes: &attributes,
};
Ok(json!({
"dialect": "ClickHouse SQL",
"access": "Authenticated team scope enforced by ClickHouse row policies; proxy admins can read all teams, while project-bound and teamless keys can read only their own rows",
"response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
"tables": tables,
"normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({
"table": "otel_traces", "name": field.name, "column": field.clickhouse_column,
"type": field.clickhouse_type, "meaning": field.meaning
})).collect::<Vec<_>>(),
"metadata": metadata,
"attributes": attributes,
"relationships": [{
"left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id",
"additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND otel_traces.ApiKeyHash = spend_logs.api_key",
"meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows"
}],
"examples": guide.examples()?,
"gotchas": guide.gotchas()?,
"guide": guide::render(&guide)?,
}).to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() {
let sample = [
MetadataRow {
metadata: r#"{"x": 1}"#.into(),
},
MetadataRow {
metadata: r#"{"x": "one"}"#.into(),
},
MetadataRow {
metadata: "invalid".into(),
},
];
let catalog = json!(metadata_catalog(&sample));
assert_eq!(
catalog["fields"],
json!([{
"path": ["x"], "types": ["integer", "string"], "expression": "JSONExtractRaw(metadata, 'x')"
}])
);
assert_eq!(catalog["invalid_json_rows"], 1);
assert_eq!(catalog["sampled_rows"], sample.len());
}
#[rstest]
#[case::rows(SAMPLE_ROWS + 1, 1)]
#[case::paths(1, MAX_FIELDS + 1)]
fn metadata_discovery_reports_truncation(#[case] row_count: usize, #[case] field_count: usize) {
let metadata: BTreeMap<_, _> = (0..field_count)
.map(|index| (format!("field{index}"), index))
.collect();
let sample: Vec<_> = (0..row_count)
.map(|_| MetadataRow {
metadata: json!(metadata).to_string(),
})
.collect();
let catalog = json!(metadata_catalog(&sample));
assert_eq!(catalog["truncated"], true);
assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS));
assert_eq!(
catalog["fields"].as_array().unwrap().len(),
field_count.min(MAX_FIELDS)
);
}
}

89
litellm-rust/crates/traces/src/query/guide.rs Normal file
View file

@ -0,0 +1,89 @@
use askama::Template;
use serde::Serialize;
use super::{AttributeCatalog, MetadataCatalog, TableSchema};
use crate::{Error, NormalizedFieldDefinition};
#[derive(Template)]
#[template(path = "query_help.jinja", escape = "none", blocks = [
"recent_spans_name",
"recent_spans_sql",
"custom_metadata_name",
"custom_metadata_sql",
"nested_metadata_name",
"nested_metadata_sql",
"correlated_calls_name",
"correlated_calls_sql",
"discover_keys_name",
"discover_keys_sql",
"time_window",
"reader_limits",
"reader_profile",
"output_format",
"json_values",
"map_values",
"literal_keys",
"time_units",
"spend_totals",
"trace_rollups",
"sampling",
])]
pub(super) struct QueryGuide<'a> {
pub tables: &'a [TableSchema],
pub normalized_fields: &'a [NormalizedFieldDefinition],
pub metadata: &'a MetadataCatalog,
pub attributes: &'a [AttributeCatalog],
}
#[derive(Serialize)]
pub(super) struct Example {
name: String,
sql: String,
}
impl QueryGuide<'_> {
pub fn examples(&self) -> Result<[Example; 5], Error> {
Ok([
Example {
name: render(&self.as_recent_spans_name())?,
sql: render(&self.as_recent_spans_sql())?,
},
Example {
name: render(&self.as_custom_metadata_name())?,
sql: render(&self.as_custom_metadata_sql())?,
},
Example {
name: render(&self.as_nested_metadata_name())?,
sql: render(&self.as_nested_metadata_sql())?,
},
Example {
name: render(&self.as_correlated_calls_name())?,
sql: render(&self.as_correlated_calls_sql())?,
},
Example {
name: render(&self.as_discover_keys_name())?,
sql: render(&self.as_discover_keys_sql())?,
},
])
}
pub fn gotchas(&self) -> Result<[String; 11], Error> {
Ok([
render(&self.as_time_window())?,
render(&self.as_reader_limits())?,
render(&self.as_reader_profile())?,
render(&self.as_output_format())?,
render(&self.as_json_values())?,
render(&self.as_map_values())?,
render(&self.as_literal_keys())?,
render(&self.as_time_units())?,
render(&self.as_spend_totals())?,
render(&self.as_trace_rollups())?,
render(&self.as_sampling())?,
])
}
}
pub(super) fn render(template: &impl Template) -> Result<String, Error> {
template.render().map_err(|_| Error::InvalidResponse)
}

200
litellm-rust/crates/traces/src/query_access.rs Normal file
View file

@ -0,0 +1,200 @@
use std::{sync::Arc, time::Duration};
use hmac::{Hmac, Mac};
use litellm_http::Client;
use moka::future::Cache;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use crate::{Connection, QueryAccessError};
const TABLES: [&str; 3] = ["otel_traces", "agent_traces_by_key", "spend_logs"];
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum QueryScope {
Admin,
Team {
team_id: String,
},
Key {
team_id: String,
api_key_hash: String,
},
}
impl QueryScope {
fn validate(&self) -> Result<(), QueryAccessError> {
match self {
Self::Admin => Ok(()),
Self::Team { team_id } if !team_id.is_empty() => Ok(()),
Self::Key { api_key_hash, .. } if !api_key_hash.is_empty() => Ok(()),
_ => Err(QueryAccessError::InvalidScope),
}
}
fn predicate(&self, table: &str) -> String {
let (team, key) = if table == "spend_logs" {
("team_id", "api_key")
} else {
("TeamId", "ApiKeyHash")
};
match self {
Self::Admin => "1".to_owned(),
Self::Team { team_id } => format!("{team} = {}", literal(team_id)),
Self::Key {
team_id,
api_key_hash,
} => format!(
"{team} = {} AND {key} = {}",
literal(team_id),
literal(api_key_hash)
),
}
}
}
#[derive(Clone)]
pub struct QueryReaders {
writer: Connection,
database: String,
readers: Cache<String, Connection>,
slots: Arc<Semaphore>,
}
impl QueryReaders {
pub fn new(writer: Connection, database: String) -> Self {
Self {
writer,
database,
readers: Cache::builder().max_capacity(1024).build(),
slots: Arc::new(Semaphore::new(8)),
}
}
pub fn acquire(&self) -> Result<OwnedSemaphorePermit, QueryAccessError> {
self.slots
.clone()
.try_acquire_owned()
.map_err(|_| QueryAccessError::Busy)
}
pub async fn connection(
&self,
client: &Client,
scope: &QueryScope,
secret: &str,
) -> Result<Connection, QueryAccessError> {
scope.validate()?;
if secret.is_empty() {
return Err(QueryAccessError::MissingSecret);
}
let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope))
.map_err(|_| QueryAccessError::InvalidScope)?;
let user = format!("litellm_traces_{:x}", Sha256::digest(&identity));
let password = credential(secret, b"password", &identity)?;
self.readers
.try_get_with(
user.clone(),
self.provision(client, scope, &user, &password),
)
.await
.map_err(QueryAccessError::Cached)
}
async fn provision(
&self,
client: &Client,
scope: &QueryScope,
user: &str,
password: &str,
) -> Result<Connection, QueryAccessError> {
let database = &self.database;
if database.is_empty()
|| !database
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
{
return Err(QueryAccessError::InvalidScope);
}
let password_hash = format!("{:x}", Sha256::digest(password));
self.execute(
client,
format!(
"CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \
SETTINGS readonly = 1 CONST, max_execution_time = 10 CONST, \
max_result_rows = 1000 CONST, max_result_bytes = 4194304 CONST, \
result_overflow_mode = 'throw' CONST, max_memory_usage = 268435456 CONST, \
max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST"
),
)
.await?;
self.execute(
client,
format!("ALTER USER {user} IDENTIFIED WITH sha256_hash BY '{password_hash}'"),
)
.await?;
for table in TABLES {
let predicate = scope.predicate(table);
self.execute(
client,
format!(
"CREATE ROW POLICY IF NOT EXISTS {user}_allow ON `{database}`.{table} \
USING 1 TO {user}"
),
)
.await?;
self.execute(
client,
format!(
"CREATE ROW POLICY IF NOT EXISTS {user}_scope ON `{database}`.{table} \
AS RESTRICTIVE USING {predicate} TO {user}"
),
)
.await?;
}
for table in TABLES {
self.execute(
client,
format!("GRANT SELECT ON `{database}`.{table} TO {user}"),
)
.await?;
}
Connection::configured(
&self.writer.url()[..url::Position::AfterPath],
database,
user,
password,
)
.map_err(QueryAccessError::Storage)
}
async fn execute(&self, client: &Client, sql: String) -> Result<(), QueryAccessError> {
let response = client
.post(self.writer.url().clone())
.timeout(Duration::from_secs(15))
.body(sql)
.send()
.await
.map_err(|_| QueryAccessError::ProvisionTransport)?;
if !response.status().is_success() {
return Err(QueryAccessError::ProvisionFailed(
response.status().as_u16(),
));
}
Ok(())
}
}
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, QueryAccessError> {
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes())
.map_err(|_| QueryAccessError::MissingSecret)?;
mac.update(purpose);
mac.update(identity);
Ok(format!("{:x}", mac.finalize().into_bytes()))
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}

45
litellm-rust/crates/traces/src/schema.rs
View file

@ -1,4 +1,5 @@
use litellm_http::Client;
use litellm_migrate::Migration;
use std::time::Duration;
use crate::Connection;
@ -6,42 +7,25 @@ use crate::Error;
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MIGRATIONS: [&str; 9] = [
include_str!("../migrations/0001_otel_traces.sql"),
include_str!("../migrations/0002_agent_traces.sql"),
include_str!("../migrations/0003_agent_traces_mv.sql"),
include_str!("../migrations/0004_spend_logs.sql"),
include_str!("../migrations/0005_otel_traces_ttl.sql"),
include_str!("../migrations/0006_agent_traces_ttl.sql"),
include_str!("../migrations/0007_spend_logs_ttl.sql"),
include_str!("../migrations/0008_trace_received.sql"),
include_str!("../migrations/0009_spend_received.sql"),
];
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations");
pub fn schema_statements(
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
) -> Result<Vec<String>, Error> {
pub fn schema_statements(database: &str, retention_days: u32) -> Result<Vec<String>, Error> {
if database.is_empty()
|| !database
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|| trace_retention_days == 0
|| spend_log_retention_days == 0
|| retention_days == 0
{
return Err(Error::InvalidSchema);
}
let database = format!("`{database}`");
Ok(
std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}"))
.chain(MIGRATIONS.iter().map(|sql| {
sql.replace("{database}", &database)
.replace("{trace_retention_days}", &trace_retention_days.to_string())
.replace(
"{spend_log_retention_days}",
&spend_log_retention_days.to_string(),
)
.chain(MIGRATIONS.iter().map(|migration| {
migration
.sql
.replace("{database}", &database)
.replace("{retention_days}", &retention_days.to_string())
}))
.collect(),
)
@ -51,15 +35,13 @@ pub async fn ensure_schema(
client: &Client,
connection: &Connection,
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
retention_days: u32,
) -> Result<(), Error> {
ensure_schema_with_timeout(
client,
connection,
database,
trace_retention_days,
spend_log_retention_days,
retention_days,
SCHEMA_REQUEST_TIMEOUT,
)
.await
@ -69,11 +51,10 @@ async fn ensure_schema_with_timeout(
client: &Client,
connection: &Connection,
database: &str,
trace_retention_days: u32,
spend_log_retention_days: u32,
retention_days: u32,
request_timeout: Duration,
) -> Result<(), Error> {
for statement in schema_statements(database, trace_retention_days, spend_log_retention_days)? {
for statement in schema_statements(database, retention_days)? {
let response = client
.post(connection.url().clone())
.timeout(request_timeout)

6
litellm-rust/crates/traces/src/sql.rs
View file

@ -37,6 +37,8 @@ impl ReadQuery {
#[derive(Clone, Copy)]
pub enum LensQuery {
Availability,
Agents,
Sample,
Content,
Evidence,
@ -45,6 +47,8 @@ pub enum LensQuery {
impl LensQuery {
pub fn parse(name: &str) -> Result<Self, Error> {
match name {
"availability" => Ok(Self::Availability),
"agents" => Ok(Self::Agents),
"sample" => Ok(Self::Sample),
"content" => Ok(Self::Content),
"evidence" => Ok(Self::Evidence),
@ -53,6 +57,8 @@ impl LensQuery {
}
pub fn sql(self) -> &'static str {
match self {
Self::Availability => include_str!("../query/lens_availability.sql"),
Self::Agents => include_str!("../query/lens_agents.sql"),
Self::Sample => include_str!("../query/lens_sample.sql"),
Self::Content => include_str!("../query/lens_content.sql"),
Self::Evidence => include_str!("../query/lens_evidence.sql"),

65
litellm-rust/crates/traces/templates/query_help.jinja Normal file
View file

@ -0,0 +1,65 @@
Trace SQL query guide
Live ClickHouse schema
{% for table in tables %}
{{ table.name }}
{% for column in table.columns %}{{ column.name }}: {{ column.kind }}
{% endfor %}{% endfor %}
Normalized span fields
{% for field in normalized_fields %}{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }})
{{ field.meaning }}
{% endfor %}
Observed LLM call metadata
{{ metadata.scope }}
Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }}
{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }}
{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows
{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %}
{% endfor %}{% endif %}
Observed span and resource attributes
{% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }}
{{ catalog.scope }}
{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }}
{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans
{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }}
{% endfor %}{% endif %}{% endfor %}
Examples
{% block recent_spans_name %}Recent normalized LLM spans{% endblock %}
{% block recent_spans_sql %}SELECT TraceId, SpanId, Model, InputTokens, OutputTokens, Duration / 1000000 AS duration_ms FROM otel_traces WHERE Timestamp >= now() - INTERVAL 1 DAY AND ObservationType = 'llm' ORDER BY Timestamp DESC LIMIT 100{% endblock %}
{% block custom_metadata_name %}Find calls by custom metadata{% endblock %}
{% block custom_metadata_sql %}SELECT request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'project') AND JSONExtractString(metadata, 'project') = 'example' ORDER BY start_time DESC LIMIT 100{% endblock %}
{% block nested_metadata_name %}Nested metadata with unknown types{% endblock %}
{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %}
{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %}
{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND t.ApiKeyHash = s.api_key WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %}
{% block discover_keys_name %}Discover metadata keys over a different window{% endblock %}
{% block discover_keys_sql %}SELECT DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 30 DAY ORDER BY key LIMIT 200{% endblock %}
Gotchas
{% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %}
{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %}
{% block reader_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces authenticated team scope through row policies. Project-bound and teamless keys see only their own rows. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %}
{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %}
{% block json_values %}metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false{% endblock %}
{% block map_values %}SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks{% endblock %}
{% block literal_keys %}Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator{% endblock %}
{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %}
{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; aggregate spend separately{% endblock %}
{% block trace_rollups %}agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators{% endblock %}
{% block sampling %}Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent'){% endblock %}

505
litellm-rust/crates/traces/tests/migrations.rs
View file

@ -2,8 +2,8 @@ use std::{collections::BTreeMap, time::Duration};
use litellm_http::Client;
use litellm_traces::{
Connection, Error, InsertTable, Parameter, ReadQuery, encode_rows, ensure_schema,
execute_named_read, execute_read, schema_statements,
Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery,
encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements,
};
use rstest::{fixture, rstest};
use testcontainers_modules::{
@ -109,8 +109,8 @@ async fn schema_supports_span_rollups_and_spend_joins(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let span = serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "",
@ -210,6 +210,42 @@ async fn schema_supports_span_rollups_and_spend_joins(
Ok(())
}
#[rstest]
#[tokio::test]
async fn normalized_fields_match_clickhouse_catalog(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
let database = database?;
ensure_schema(
&database.client,
&Connection::writer(&database.url)?,
"trace_test",
7,
)
.await?;
let catalog = read_json(&database, "SELECT name, type FROM system.columns WHERE database = 'trace_test' AND table = 'otel_traces'").await?;
let columns: BTreeMap<&str, &str> = catalog["data"]
.as_array()
.expect("catalog rows")
.iter()
.map(|row| {
(
row["name"].as_str().expect("column name"),
row["type"].as_str().expect("column type"),
)
})
.collect();
for field in NORMALIZED_FIELD_DEFINITIONS {
assert_eq!(
columns.get(field.clickhouse_column).copied(),
Some(field.clickhouse_type),
"{}",
field.name
);
}
Ok(())
}
#[rstest]
#[tokio::test]
async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
@ -220,7 +256,7 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
"{}?input_format_skip_unknown_fields=1",
database.url
))?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let row = BTreeMap::from([
(
"Timestamp".to_owned(),
@ -254,7 +290,7 @@ async fn retried_trace_insert_does_not_inflate_rollup(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let row: BTreeMap<String, serde_json::Value> = serde_json::from_value(serde_json::json!({
"Timestamp": time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64,
"TraceId": "retried-trace", "SpanId": "span-1", "ParentSpanId": "",
@ -289,7 +325,7 @@ async fn keyed_rollup_keeps_same_trace_ids_separate_by_api_key(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let rows = vec![
serde_json::from_value(serde_json::json!({
@ -333,7 +369,7 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let day_start = time::OffsetDateTime::now_utc()
.replace_time(time::Time::MIDNIGHT)
.unix_timestamp_nanos() as i64;
@ -380,7 +416,7 @@ async fn spend_deduplication_preserves_subsecond_requests_and_retries(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let now_ms = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000_000;
let base_start_time = now_ms / 1000 * 1000;
let first_start_time = base_start_time + 100;
@ -432,7 +468,21 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 30, 30).await?;
ensure_schema(&database.client, &writer, "trace_test", 30).await?;
let tables = read_json(
&database,
"SELECT name FROM system.tables WHERE database = 'trace_test' \
AND match(engine_full, 'materialize_ttl_recalculate_only = 1') ORDER BY name",
)
.await?;
assert_eq!(
tables["data"],
serde_json::json!([
{"name": "agent_traces_by_key"},
{"name": "otel_traces"},
{"name": "spend_logs"}
])
);
let old_time = time::OffsetDateTime::now_utc() - time::Duration::days(20);
let old_timestamp_ns = old_time.unix_timestamp_nanos() as i64;
let old_timestamp_ms = old_timestamp_ns / 1_000_000;
@ -448,7 +498,7 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
insert_rows(&database, "otel_traces", vec![span]).await?;
insert_rows(&database, "spend_logs", vec![spend]).await?;
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 1);
ensure_schema(&database.client, &writer, "trace_test", 14, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 14).await?;
let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
loop {
let response = read_json(
@ -480,7 +530,7 @@ async fn retention_changes_materialize_existing_rows_and_remain_idempotent(
assert_eq!(table_rows(&database, "agent_traces_by_key").await?, 0);
assert_eq!(table_rows(&database, "spend_logs").await?, 0);
let mutation_count = mutation_rows(&database).await?;
ensure_schema(&database.client, &writer, "trace_test", 14, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 14).await?;
assert_eq!(mutation_rows(&database).await?, mutation_count);
Ok(())
}
@ -499,7 +549,7 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
let writer = Connection::writer(&url)?;
let result = tokio::time::timeout(
Duration::from_secs(35),
ensure_schema(&client, &writer, "trace_test", 7, 14),
ensure_schema(&client, &writer, "trace_test", 7),
)
.await;
server.abort();
@ -508,16 +558,11 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
}
#[rstest]
#[case::empty("", 7, 14)]
#[case::sql("db; DROP DATABASE default", 7, 14)]
#[case::trace_retention("traces", 0, 14)]
#[case::spend_retention("traces", 7, 0)]
fn schema_rejects_invalid_configuration(
#[case] database: &str,
#[case] traces: u32,
#[case] spend: u32,
) {
assert!(schema_statements(database, traces, spend).is_err());
#[case::empty("", 7)]
#[case::sql("db; DROP DATABASE default", 7)]
#[case::retention("traces", 0)]
fn schema_rejects_invalid_configuration(#[case] database: &str, #[case] retention_days: u32) {
assert!(schema_statements(database, retention_days).is_err());
}
#[rstest]
@ -528,7 +573,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
use litellm_traces::{LensQuery, Parameter};
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
for (key, text) in [("one", "timeout"), ("two", "success")] {
insert_rows(&database, "otel_traces", vec![serde_json::from_value(serde_json::json!({
@ -551,6 +596,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
"end".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text("review".into())),
(
"filter_keys".into(),
@ -634,7 +680,7 @@ async fn lens_request_sample_does_not_trust_caller_tags(
use litellm_traces::{LensQuery, Parameter};
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64 / 1_000_000;
for (id, internal) in [("external", false), ("internal", true)] {
let row = serde_json::from_value(serde_json::json!({
@ -652,6 +698,7 @@ async fn lens_request_sample_does_not_trust_caller_tags(
("key_hash".into(), Parameter::Text(String::new())),
("start".into(), Parameter::Integer(timestamp - 1000)),
("end".into(), Parameter::Integer(timestamp + 60000)),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text(String::new())),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
@ -702,7 +749,6 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
&Connection::writer(&database.url)?,
"trace_test",
7,
14,
)
.await?;
execute_write(&database, "INSERT INTO trace_test.spend_logs (request_id,team_id,start_time,end_time) SELECT toString(number),'team',now64(3)-INTERVAL 5 MINUTE,now64(3)-INTERVAL 5 MINUTE FROM numbers(1001)").await?;
@ -719,6 +765,7 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
("key_hash".into(), Parameter::Text(String::new())),
("start".into(), Parameter::Integer(0)),
("end".into(), Parameter::Integer(end)),
("agent_name".into(), Parameter::Text(String::new())),
("service".into(), Parameter::Text(String::new())),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
@ -782,7 +829,6 @@ async fn lens_content_keeps_output_visible_after_long_input(
&Connection::writer(&database.url)?,
"trace_test",
7,
14,
)
.await?;
insert_rows(&database, "spend_logs", vec![serde_json::from_value(serde_json::json!({
@ -848,7 +894,7 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let rows = (0..span_count)
.map(|index| {
@ -938,7 +984,7 @@ async fn duplicate_span_preview_matches_diagnostic(
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let message = "a".repeat(200);
let rows = [
@ -980,3 +1026,404 @@ async fn duplicate_span_preview_matches_diagnostic(
assert_eq!(diagnostic["data"][0]["message"], message);
Ok(())
}
#[rstest]
fn schema_includes_every_migration_file() -> TestResult {
let files = std::fs::read_dir(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"))?
.filter_map(|entry| entry.ok())
.filter(|entry| entry.path().extension().is_some_and(|ext| ext == "sql"))
.count();
assert_eq!(schema_statements("trace_test", 7)?.len(), 1 + files);
Ok(())
}
#[rstest]
#[tokio::test]
async fn lens_agent_discovery_and_selection_preserve_scope(
#[future] database: TestResult<ClickHouseDatabase>,
) -> TestResult {
use litellm_traces::LensQuery;
let database = database.await?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
for (team, key, trace, agent, span, parent) in [
("alpha", "one", "research", "research_agent", "root", ""),
("alpha", "one", "research", "", "tool", "root"),
("alpha", "one", "support", "support_agent", "root", ""),
("alpha", "two", "hidden-key", "private_agent", "root", ""),
("beta", "one", "hidden-team", "other_agent", "root", ""),
] {
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
"ServiceName": "shared-app", "SpanName": "run", "Input": "test",
"SpanAttributes": {"gen_ai.agent.name": agent},
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key}
}))?],
)
.await?;
}
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let scope_parameters = BTreeMap::from([
("all_teams".into(), Parameter::Integer(0)),
("team".into(), Parameter::Text("alpha".into())),
("key_hash".into(), Parameter::Text("one".into())),
]);
let agents: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Agents.sql(),
&scope_parameters,
)
.await?,
)?;
assert_eq!(
agents["data"],
serde_json::json!([
{"agent_name": "research_agent"}, {"agent_name": "support_agent"}
])
);
let parameters = scope_parameters
.into_iter()
.chain([
("source".into(), Parameter::Text("traces".into())),
(
"start".into(),
Parameter::Integer(timestamp / 1_000_000 - 1000),
),
(
"end".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("service".into(), Parameter::Text("shared-app".into())),
(
"agent_name".into(),
Parameter::Text("research_agent".into()),
),
("filter_keys".into(), Parameter::Strings(vec![])),
("filter_values".into(), Parameter::Strings(vec![])),
("limit".into(), Parameter::Integer(100)),
("offset".into(), Parameter::Integer(0)),
("after".into(), Parameter::Text(String::new())),
("sample_percent".into(), Parameter::Text("100".into())),
("sample_cap".into(), Parameter::Integer(0)),
("preview".into(), Parameter::Integer(1)),
("selected_team".into(), Parameter::Text(String::new())),
("execution_ids".into(), Parameter::Strings(vec![])),
])
.collect::<BTreeMap<_, _>>();
let sample: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Sample.sql(),
&parameters,
)
.await?,
)?;
assert_eq!(sample["data"].as_array().expect("rows").len(), 1);
assert_eq!(sample["data"][0]["trace_id"], "research");
assert_eq!(sample["data"][0]["span_count"], 2);
let available: serde_json::Value = serde_json::from_str(
&execute_read(
&database.client,
&connection,
LensQuery::Availability.sql(),
&parameters,
)
.await?,
)?;
assert_eq!(available["data"][0]["traces"], 1);
assert_eq!(available["data"][0]["requests"], 0);
Ok(())
}
#[rstest]
#[case::empty(false)]
#[case::custom_metadata(true)]
#[tokio::test]
async fn query_help_discovers_live_schema_and_runs_its_examples(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] populated: bool,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
execute_write(&database, "CREATE USER help_reader").await?;
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
execute_write(
&database,
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
)
.await?;
}
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
if populated {
execute_write(&database, "SYSTEM STOP MERGES trace_test.spend_logs").await?;
insert_rows(
&database,
"spend_logs",
vec![serde_json::from_value(serde_json::json!({
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
"api_key": "key-1", "metadata": r#"{"obsolete":true,"labels":{"priority":"old"}}"#,
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000
}))?],
)
.await?;
let metadata = serde_json::json!({
"project": "example", "labels": {"priority": 3, "enabled": true},
"dotted.key": "literal", "quote'\\key": null, "items": [{"name": "first"}],
"<custom>&{{key}}": {"nested.key": true}
});
insert_rows(
&database,
"spend_logs",
vec![serde_json::from_value(serde_json::json!({
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
"api_key": "key-1", "metadata": metadata.to_string(), "spend": 0.25,
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100
}))?],
)
.await?;
insert_rows(
&database,
"otel_traces",
vec![serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1",
"TeamId": "team-1", "ApiKeyHash": "key-1", "ObservationType": "llm",
"LiteLLMRequestId": "response-1", "SpanAttributes": {"custom.tag": "value"},
"ResourceAttributes": {"custom.resource": "value"}
}))?],
)
.await?;
execute_write(
&database,
"ALTER TABLE trace_test.otel_traces ADD COLUMN CustomColumn String",
)
.await?;
}
let help: serde_json::Value =
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
let keys: std::collections::BTreeSet<_> = help
.as_object()
.ok_or("missing help object")?
.keys()
.map(String::as_str)
.collect();
assert_eq!(
keys,
std::collections::BTreeSet::from([
"access",
"attributes",
"dialect",
"examples",
"gotchas",
"guide",
"metadata",
"normalized_fields",
"relationships",
"response",
"tables",
])
);
let guide = help["guide"].as_str().ok_or("missing rendered guide")?;
assert!(guide.starts_with("Trace SQL query guide"));
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
let described = read_json(&database, &format!("DESCRIBE TABLE {table}")).await?;
let schema = help["tables"]
.as_array()
.ok_or("missing tables")?
.iter()
.find(|schema| schema["name"] == table)
.ok_or("missing table")?;
assert_eq!(schema["columns"], described["data"]);
for column in described["data"].as_array().ok_or("missing live columns")? {
assert!(guide.contains(&format!(
"{}: {}",
column["name"].as_str().ok_or("column name")?,
column["type"].as_str().ok_or("column type")?
)));
}
}
for gotcha in help["gotchas"].as_array().ok_or("missing gotchas")? {
assert!(guide.contains(gotcha.as_str().ok_or("gotcha text")?));
}
let tables = help["tables"].as_array().ok_or("missing tables")?;
assert_eq!(tables.len(), 3);
let columns = tables[0]["columns"].as_array().ok_or("missing columns")?;
for field in NORMALIZED_FIELD_DEFINITIONS {
assert!(
columns
.iter()
.any(|column| column["name"] == field.clickhouse_column
&& column["type"] == field.clickhouse_type)
);
assert!(
help["normalized_fields"]
.as_array()
.ok_or("missing mappings")?
.iter()
.any(|mapped| {
mapped["name"] == field.name && mapped["column"] == field.clickhouse_column
})
);
}
let fields = help["metadata"]["fields"]
.as_array()
.ok_or("missing metadata fields")?;
assert_eq!(fields.is_empty(), !populated);
assert_eq!(help["metadata"]["truncated"], false);
assert!(guide.contains(help["metadata"]["scope"].as_str().ok_or("missing scope")?));
assert_eq!(
guide.contains("No metadata paths found in the sampled rows"),
!populated
);
if populated {
let versions = read_json(&database, "SELECT count() AS count FROM spend_logs").await?;
assert_eq!(versions["data"][0]["count"], 2);
assert_eq!(help["metadata"]["sampled_rows"], 1);
assert!(
!fields
.iter()
.any(|field| field["path"] == serde_json::json!(["obsolete"]))
);
assert!(
columns
.iter()
.any(|column| column["name"] == "CustomColumn")
);
assert!(fields.iter().any(|field| field["path"]
== serde_json::json!(["labels", "priority"])
&& field["types"] == serde_json::json!(["integer"])));
assert!(
fields
.iter()
.any(|field| field["path"] == serde_json::json!(["items", 1, "name"]))
);
assert!(guide.contains("CustomColumn: String"));
assert!(guide.contains("JSONExtractRaw(metadata, '<custom>&{{key}}', 'nested.key')"));
assert!(guide.contains("SpanAttributes['custom.tag']"));
assert!(guide.contains("ResourceAttributes['custom.resource']"));
assert_eq!(help["attributes"][0]["fields"][0]["key"], "custom.tag");
assert_eq!(help["attributes"][1]["fields"][0]["key"], "custom.resource");
for field in fields {
let expression = field["expression"].as_str().ok_or("missing expression")?;
assert!(
guide.contains(expression),
"missing plain-text expression: {expression}"
);
let sql = format!("SELECT {expression} AS value FROM spend_logs FINAL");
let body = litellm_traces::query_sql(&database.client, &reader, &sql).await?;
let values: serde_json::Value = serde_json::from_str(&body)?;
assert_ne!(values["data"][0]["value"], "");
}
}
for example in help["examples"].as_array().ok_or("missing examples")? {
let sql = example["sql"].as_str().ok_or("missing example SQL")?;
assert!(guide.contains(example["name"].as_str().ok_or("missing example name")?));
assert!(guide.contains(sql));
assert_eq!(
example
.as_object()
.ok_or("example object")?
.keys()
.map(String::as_str)
.collect::<std::collections::BTreeSet<_>>(),
std::collections::BTreeSet::from(["name", "sql"])
);
let body = litellm_traces::query_sql(&database.client, &reader, sql).await?;
let values: serde_json::Value = serde_json::from_str(&body)?;
assert_eq!(
values["data"].as_array().ok_or("missing data")?.is_empty(),
!populated,
"{sql}"
);
if populated && example["name"] == "Traces correlated with LLM call metadata" {
assert_eq!(values["data"][0]["TraceId"], "trace-1");
assert_eq!(values["data"][0]["spend"], 0.25);
}
}
Ok(())
}
#[rstest]
#[case::metadata(2, 1)]
#[case::attributes(1, 2)]
#[case::all(2, 2)]
#[tokio::test]
async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits(
#[future(awt)] database: TestResult<ClickHouseDatabase>,
#[case] spend_rows: usize,
#[case] span_rows: usize,
) -> TestResult {
let database = database?;
let writer = Connection::writer(&database.url)?;
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
execute_write(
&database,
"CREATE USER help_reader SETTINGS max_rows_to_read = 1",
)
.await?;
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
execute_write(
&database,
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
)
.await?;
}
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
let spend = (0..spend_rows)
.map(|index| {
serde_json::from_value(serde_json::json!({
"request_id": format!("request-{index}"), "start_time": timestamp / 1_000_000,
"end_time": timestamp / 1_000_000, "metadata": r#"{"custom":{"enabled":true}}"#
}))
})
.collect::<Result<Vec<_>, _>>()?;
insert_rows(&database, "spend_logs", spend).await?;
let spans = (0..span_rows).map(|index| serde_json::from_value(serde_json::json!({
"Timestamp": timestamp, "TraceId": "trace", "SpanId": format!("span-{index}"),
"SpanAttributes": {"custom.span": "value"}, "ResourceAttributes": {"custom.resource": "value"}
}))).collect::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", spans).await?;
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
let help: serde_json::Value =
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3);
assert!(!help["examples"].as_array().ok_or("examples")?.is_empty());
assert_eq!(
help["normalized_fields"]
.as_array()
.ok_or("normalized fields")?
.len(),
NORMALIZED_FIELD_DEFINITIONS.len()
);
let guide = help["guide"].as_str().ok_or("guide")?;
assert!(guide.contains("TraceId: String"));
assert_eq!(
guide.contains("Metadata discovery unavailable:"),
spend_rows > 1
);
assert_eq!(
guide.contains("Attribute discovery unavailable:"),
span_rows > 1
);
for (catalog, unavailable) in [
(&help["metadata"], spend_rows > 1),
(&help["attributes"][0], span_rows > 1),
(&help["attributes"][1], span_rows > 1),
] {
assert_eq!(catalog.get("error").is_some(), unavailable);
assert_eq!(catalog["truncated"], unavailable);
assert_eq!(
catalog["fields"].as_array().ok_or("fields")?.is_empty(),
unavailable
);
}
Ok(())
}

47
litellm-rust/crates/traces/tests/otlp.rs
View file

@ -1,5 +1,5 @@
use litellm_traces::Shared;
use litellm_traces::decode_otlp;
use litellm_traces::{ObservationType, Shared};
use rstest::rstest;
const FIXTURE: &[u8] = include_bytes!(
@ -341,3 +341,48 @@ fn escaped_attribute_expansion_is_bounded_below_four_mib(
Err(litellm_traces::DecodeError::TooLarge)
));
}
#[rstest]
fn normalizes_langsmith_fixture() {
let spans = decode_otlp(FIXTURE, Some("application/json")).expect("valid OTLP export");
let llm = spans
.iter()
.find(|span| span.name == "ChatOpenAI")
.expect("LLM span");
assert_eq!(llm.normalized.observation_type, ObservationType::Llm);
assert_eq!(llm.normalized.agent_name, "deep_research_agent");
assert_eq!(llm.normalized.model, "claude-sonnet-4-5");
assert_eq!(
(llm.normalized.input_tokens, llm.normalized.output_tokens),
(3332, 467)
);
assert_eq!(
llm.normalized.litellm_request_id,
"chatcmpl-4077bb36-9380-4a3b-9481-245700cef09a"
);
let input: serde_json::Value =
serde_json::from_str(&llm.normalized.input).expect("message input");
assert_eq!(input[0]["role"], "system");
assert_eq!(input[1]["role"], "user");
let output: serde_json::Value =
serde_json::from_str(&llm.normalized.output).expect("message output");
assert_eq!(output["role"], "assistant");
assert!(output["tool_calls"][0]["name"].is_string());
assert!(output["tool_calls"][0]["id"].is_string());
assert_eq!(output["tool_calls"][0]["type"], "tool_call");
let root = spans
.iter()
.find(|span| span.name == "deep_research_agent")
.expect("root span");
assert_eq!(root.normalized.observation_type, ObservationType::Agent);
assert_eq!(
root.normalized.input,
"[{\"role\": \"user\", \"content\": \"Should we store OTEL agent spans in ClickHouse or Postgres at 50k spans/sec?\"}]"
);
let tool = spans
.iter()
.find(|span| span.name == "task")
.expect("tool span");
assert_eq!(tool.normalized.observation_type, ObservationType::Tool);
assert!(tool.normalized.output.starts_with("Based on my research"));
}

268
litellm-rust/crates/traces/tests/query_access.rs Normal file
View file

@ -0,0 +1,268 @@
use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_traces::{
Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use testcontainers_modules::{
clickhouse::ClickHouse,
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
};
struct Database {
_container: ContainerAsync<ClickHouse>,
client: Client,
writer: Connection,
readers: QueryReaders,
}
#[fixture]
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
let container = ClickHouse::default()
.with_tag(
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e",
)
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
.start()
.await?;
let writer = Connection::parse(&format!(
"http://{}:{}",
container.get_host().await?,
container.get_host_port_ipv4(8123).await?
))?;
let client = Client::no_redirect_for_test();
ensure_schema(&client, &writer, "trace_test", 7).await?;
for sql in [
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a')), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a')), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'))",
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}')",
"CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory",
"INSERT INTO trace_test.private_data VALUES ('hidden')",
] {
let response = client.post(writer.url().clone()).body(sql).send().await?;
assert!(response.status().is_success(), "{}", response.text().await?);
}
let readers = QueryReaders::new(writer.clone(), "trace_test".to_owned());
Ok(Database {
_container: container,
client,
writer,
readers,
})
}
#[rstest]
#[case::team(QueryScope::Team { team_id: "team-a".to_owned() }, vec!["a1", "a2"])]
#[case::project_key(QueryScope::Key { team_id: "team-a".to_owned(), api_key_hash: "key-a1".to_owned() }, vec!["a1"])]
#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b"])]
#[case::quoted_team(QueryScope::Team { team_id: "team-a' OR 1=1 --\\".to_owned() }, vec![])]
#[tokio::test]
async fn queries_and_help_are_scoped_by_the_database(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
#[case] scope: QueryScope,
#[case] expected: Vec<&str>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let reader = database
.readers
.connection(&database.client, &scope, "test-master-secret")
.await?;
let queries = [
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
"SELECT SpanId AS id FROM trace_test.otel_traces WHERE 1 = 1 ORDER BY id",
"SELECT SpanId AS id FROM merge('trace_test', '^otel_traces$') ORDER BY id",
"WITH source AS (SELECT * FROM trace_test.otel_traces) SELECT SpanId AS id FROM source ORDER BY id",
"SELECT SpanId AS id FROM otel_traces UNION DISTINCT SELECT SpanId AS id FROM trace_test.otel_traces ORDER BY id",
"SELECT t.SpanId AS id FROM otel_traces t INNER JOIN spend_logs s ON t.SpanId = s.request_id ORDER BY id",
"SELECT request_id AS id FROM spend_logs FINAL ORDER BY id",
];
for sql in queries {
let body: Value = serde_json::from_str(&query_sql(&database.client, &reader, sql).await?)?;
assert_eq!(
body["data"],
json!(
expected
.iter()
.map(|id| json!({"id": id}))
.collect::<Vec<_>>()
),
"{sql}"
);
}
let summary: Value = serde_json::from_str(
&query_sql(
&database.client,
&reader,
"SELECT sum(SpanCount) AS count FROM agent_traces_by_key",
)
.await?,
)?;
assert_eq!(summary["data"][0]["count"], json!(expected.len()));
let help = query_help(&database.client, &reader).await?;
assert_eq!(help.contains("secret_b"), expected.contains(&"b"));
assert_eq!(help.contains("secret-b"), expected.contains(&"b"));
let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned());
let repeated = recreated
.connection(&database.client, &scope, "test-master-secret")
.await?;
assert_eq!(reader.url(), repeated.url());
Ok(())
}
#[rstest]
#[tokio::test]
async fn rotating_master_secret_revokes_previous_reader_credentials(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let scope = QueryScope::Team {
team_id: "team-a".to_owned(),
};
let old_reader = database
.readers
.connection(&database.client, &scope, "old-master-secret")
.await?;
let old_result = query_sql(
&database.client,
&old_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await?;
let old_rows: Value = serde_json::from_str(&old_result)?;
assert_eq!(old_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
let rotated_readers = QueryReaders::new(database.writer.clone(), "trace_test".into());
let new_reader = rotated_readers
.connection(&database.client, &scope, "new-master-secret")
.await?;
assert!(
query_sql(
&database.client,
&old_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await
.is_err()
);
let new_result = query_sql(
&database.client,
&new_reader,
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
)
.await?;
let new_rows: Value = serde_json::from_str(&new_result)?;
assert_eq!(new_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
assert_eq!(old_reader.url().username(), new_reader.url().username());
assert_ne!(old_reader.url().password(), new_reader.url().password());
Ok(())
}
#[rstest]
#[tokio::test]
async fn managed_reader_rejects_privilege_and_scope_bypasses(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let scope = QueryScope::Team {
team_id: "team-a".to_owned(),
};
let reader = database
.readers
.connection(&database.client, &scope, "test-master-secret")
.await?;
for sql in [
"INSERT INTO otel_traces (TraceId) VALUES ('injected')",
"DROP TABLE otel_traces",
"SELECT * FROM private_data",
"SELECT * FROM otel_traces SETTINGS readonly = 0",
"SELECT * FROM otel_traces SETTINGS max_memory_usage = 0",
"SELECT * FROM otel_traces SETTINGS max_execution_time = 0",
"CREATE USER scope_bypass",
"CREATE NAMED COLLECTION scope_bypass AS host = 'localhost'",
"BACKUP TABLE otel_traces TO Disk('default', 'scope-bypass')",
"SELECT * FROM url('http://127.0.0.1:1/', 'LineAsString', 'line String')",
"SELECT * FROM remote('127.0.0.1', 'trace_test', 'otel_traces')",
] {
assert!(
matches!(
query_sql(&database.client, &reader, sql).await,
Err(Error::QueryFailed(_))
),
"{sql}"
);
}
let roles: Value = serde_json::from_str(
&query_sql(&database.client, &reader, "SELECT enabledRoles() AS roles").await?,
)?;
assert_eq!(roles["data"], json!([{ "roles": [] }]));
let rows: Value = serde_json::from_str(
&query_sql(
&database.client,
&reader,
"SELECT DISTINCT TeamId FROM otel_traces",
)
.await?,
)?;
assert_eq!(rows["data"], json!([{ "TeamId": "team-a" }]));
Ok(())
}
#[rstest]
#[tokio::test]
async fn provisioning_failure_never_returns_a_writer_connection(
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
) -> Result<(), Box<dyn std::error::Error>> {
let database = database?;
let reader = database
.readers
.connection(&database.client, &QueryScope::Admin, "test-master-secret")
.await?;
let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned());
let result = no_provision_privileges
.connection(
&database.client,
&QueryScope::Team {
team_id: "team-a".to_owned(),
},
"other-secret",
)
.await;
assert!(result.is_err());
assert!(
database
.readers
.connection(&database.client, &QueryScope::Admin, "")
.await
.is_err()
);
assert!(
database
.readers
.connection(
&database.client,
&QueryScope::Team {
team_id: String::new()
},
"test-master-secret"
)
.await
.is_err()
);
let permits = (0..8)
.map(|_| database.readers.acquire())
.collect::<Result<Vec<_>, _>>()?;
assert!(database.readers.acquire().is_err());
drop(permits);
assert!(database.readers.acquire().is_ok());
let rows = litellm_traces::execute_read(
&database.client,
&database.writer,
"SELECT count() AS count FROM trace_test.otel_traces",
&BTreeMap::new(),
)
.await?;
let rows: Value = serde_json::from_str(&rows)?;
assert_eq!(rows["data"][0]["count"], 3);
Ok(())
}

2
litellm/__init__.py
View file

@ -2282,7 +2282,6 @@ if TYPE_CHECKING:
# Track if async client cleanup has been registered (for lazy loading)
_async_client_cleanup_registered = False
# litellm.agent() entrypoints, resolved lazily from litellm.harness by __getattr__.
_AGENT_EXPORTS: Final = frozenset(
{
"agent",
@ -2333,7 +2332,6 @@ def __getattr__(name: str) -> Any:
handler_func: Final = registry[name]
return handler_func(name)
# litellm.agent() and friends: imported on first access (not needed for completion calls)
if name == "harness" or name in _AGENT_EXPORTS:
import importlib

4
litellm/anthropic_beta_headers_config.json
View file

@ -67,6 +67,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": null,
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"web-fetch-2025-09-10": "web-fetch-2025-09-10",
"web-search-2025-03-05": "web-search-2025-03-05"
@ -100,6 +101,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": null,
"web-fetch-2025-09-10": null,
@ -172,6 +174,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": "token-efficient-tools-2025-02-19",
"tool-examples-2025-10-29": "tool-examples-2025-10-29",
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
@ -207,6 +210,7 @@
"text_editor_20241022": null,
"text_editor_20250124": null,
"thinking-binding-controls-2026-08-01": "thinking-binding-controls-2026-08-01",
"thinking-display-updates-2026-08-18": "thinking-display-updates-2026-08-18",
"token-efficient-tools-2025-02-19": null,
"tool-search-tool-2025-10-19": "tool-search-tool-2025-10-19",
"web-fetch-2025-09-10": null,

4
litellm/constants.py
View file

@ -50,8 +50,8 @@ CLICKHOUSE_BATCH_SIZE: Final = get_env_int("CLICKHOUSE_BATCH_SIZE", 10_000)
CLICKHOUSE_FLUSH_INTERVAL_SECONDS: Final = float(os.getenv("CLICKHOUSE_FLUSH_INTERVAL_SECONDS", "1.0"))
CLICKHOUSE_MAX_BUFFERED_ROWS: Final = get_env_int("CLICKHOUSE_MAX_BUFFERED_ROWS", 200_000)
CLICKHOUSE_MAX_RETRIES: Final = get_env_int("CLICKHOUSE_MAX_RETRIES", 3)
AGENT_TRACING_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_RETENTION_DAYS", 30)
AGENT_TRACING_SPEND_LOG_RETENTION_DAYS: Final = get_env_int("AGENT_TRACING_SPEND_LOG_RETENTION_DAYS", 90)
DEFAULT_CLICKHOUSE_DATABASE: Final = "litellm"
DEFAULT_AGENT_TRACING_RETENTION_DAYS: Final = 14
OTLP_MAX_BODY_BYTES: Final = get_env_int("OTLP_MAX_BODY_BYTES", 16 * 1024 * 1024)
OTLP_MAX_ATTRIBUTE_VALUE_BYTES: Final = get_env_int("OTLP_MAX_ATTRIBUTE_VALUE_BYTES", 64 * 1024)
OTLP_RETRY_AFTER_SECONDS: Final = get_env_int("OTLP_RETRY_AFTER_SECONDS", 2)

24
litellm/harness/endpoint.py
View file

@ -131,11 +131,6 @@ class UsageTracker:
)
# ---------------------------------------------------------------------------
# Usage parsing
# ---------------------------------------------------------------------------
def _as_int(value: object) -> int:
if isinstance(value, bool):
return 0
@ -230,11 +225,6 @@ class SSEUsageParser:
self.output_tokens = output_tokens
# ---------------------------------------------------------------------------
# Cost + helpers
# ---------------------------------------------------------------------------
def compute_cost(model: str | None, input_tokens: int, output_tokens: int) -> float:
"""Cost from LiteLLM's price map. Never raises; unknown models cost 0.0."""
if not model or not (input_tokens or output_tokens):
@ -373,11 +363,6 @@ def _noop() -> None:
return None
# ---------------------------------------------------------------------------
# ModelEndpoint
# ---------------------------------------------------------------------------
class ModelEndpoint:
"""Local HTTP endpoint for one harness session. Use as an async context manager."""
@ -401,7 +386,6 @@ class ModelEndpoint:
self.token = secrets.token_urlsafe(HARNESS_SESSION_TOKEN_BYTES)
self.usage = UsageTracker()
self.port = 0
# Injected client (tests); production uses LiteLLM's shared cached client.
self._injected_client = client
self._deps: _ServerDeps | None = None
self._client: httpx.AsyncClient | None = None
@ -412,8 +396,6 @@ class ModelEndpoint:
def url(self) -> str:
return f"http://{HARNESS_ENDPOINT_HOST}:{self.port}"
# -- lifecycle ----------------------------------------------------------
async def __aenter__(self) -> ModelEndpoint:
await self.start()
return self
@ -500,8 +482,6 @@ class ModelEndpoint:
routes=[*post_routes, *get_routes] # mutable-ok: Starlette takes a routes list
)
# -- request handling ---------------------------------------------------
@property
def _responses(self) -> ModuleType:
if self._deps is None:
@ -565,8 +545,6 @@ class ModelEndpoint:
cost = compute_cost(model, input_tokens, output_tokens)
self.usage.add(input_tokens, output_tokens, cost)
# -- gateway mode -------------------------------------------------------
async def _forward(self, request: Request, route: str, body: Mapping[str, Any]) -> Response:
if self._client is None or self.gateway is None:
raise HarnessError("gateway client is not started")
@ -622,8 +600,6 @@ class ModelEndpoint:
tokens = (0, 0)
self._record(model, tokens[0], tokens[1], header_cost(upstream.headers))
# -- SDK mode -----------------------------------------------------------
def _sdk_kwargs(
self, body: Mapping[str, Any]
) -> dict[str, Any]: # mutable-ok: SDK call kwargs, mutated by _invoke_sdk then splatted

42
litellm/harness/runtime.py
View file

@ -68,11 +68,6 @@ verbose_logger: Final = logging.getLogger("LiteLLM")
PROPAGATED_ERRORS: Final = (HarnessInstallFailed, CapabilityUnsupported)
# ---------------------------------------------------------------------------
# Configuration + validation
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class SessionConfig:
"""Every per-session parameter a caller can pass, already normalized."""
@ -249,11 +244,6 @@ def _context_for(config: SessionConfig) -> SessionContext:
)
# ---------------------------------------------------------------------------
# Structured output
# ---------------------------------------------------------------------------
def parse_output(output: type[BaseModel], output_json: str | None, text: str) -> tuple[BaseModel | None, str | None]:
"""Return (model, None) on success or (None, error message) on failure."""
raw = output_json or last_json_object(text)
@ -265,11 +255,6 @@ def parse_output(output: type[BaseModel], output_json: str | None, text: str) ->
return None, str(e)
# ---------------------------------------------------------------------------
# Turn machinery
# ---------------------------------------------------------------------------
@dataclass
class _End:
"""Sentinel the producer puts on the queue when the handler turn is over."""
@ -392,8 +377,6 @@ class _Turn:
self.usage_before: tuple[int, int, int, float] = (0, 0, 0, 0.0)
self.deadline: float | None = None
# -- setup / teardown ---------------------------------------------------
async def _begin(self) -> None:
sandbox = self.ctx.sandbox
self.before = await sandbox.snapshot()
@ -419,8 +402,6 @@ class _Turn:
if pending:
await asyncio.wait(pending)
# -- event loop ---------------------------------------------------------
async def _next_item(self) -> Event | _End:
if self.deadline is None:
return await self.queue.get()
@ -476,8 +457,6 @@ class _Turn:
# The consumer asked for the next event without answering.
item.deny("approval not answered")
# -- results ------------------------------------------------------------
async def _file_changes(self) -> list[FileChange]: # mutable-ok: becomes the public Result.files list
sandbox = self.ctx.sandbox
after = await sandbox.snapshot()
@ -535,8 +514,6 @@ class _Turn:
parsed, error = parse_output(output_type, self.ctx.output_json, text)
return parsed, self.ctx.output_json or text, error
# -- entry --------------------------------------------------------------
async def run(self) -> AsyncIterator[Event]:
await self._begin()
self._start_producer()
@ -564,11 +541,6 @@ class _Turn:
)
# ---------------------------------------------------------------------------
# Streams
# ---------------------------------------------------------------------------
class AsyncEventStream:
"""Async iterator of events for one turn. `.result` is set once Done is seen."""
@ -609,11 +581,6 @@ async def _one_shot(session: AsyncSession, prompt: str, control: TurnControl) ->
await session.aclose()
# ---------------------------------------------------------------------------
# Sessions
# ---------------------------------------------------------------------------
class AsyncSession:
"""A multi-turn conversation with one harness. Use `async with` or `await`."""
@ -641,8 +608,6 @@ class AsyncSession:
self._busy = False
self._restart_needed = False
# -- lifecycle ----------------------------------------------------------
def __await__(self) -> Generator[object, None, AsyncSession]:
return self.start().__await__()
@ -752,8 +717,6 @@ class AsyncSession:
detach = adetach
stop = astop
# -- turns --------------------------------------------------------------
def usage_counters(self) -> tuple[int, int, int, float]:
"""(input_tokens, output_tokens, calls, cost) so far, from endpoint or handler."""
endpoint = self.ctx.endpoint
@ -823,11 +786,6 @@ async def _collect(events: AsyncIterator[Event]) -> Result:
return result
# ---------------------------------------------------------------------------
# Public async API
# ---------------------------------------------------------------------------
def aagent_session(
harness: Harness,
*,

8
litellm/harness/sandbox/docker.py
View file

@ -77,8 +77,6 @@ class DockerSandbox:
def __repr__(self) -> str:
return f"DockerSandbox({self.image!r}, workdir={self.workdir!r})"
# -- docker CLI plumbing (tests monkeypatch these two) ---------------------
def _docker_binary(self) -> str:
binary = shutil.which("docker")
if binary is None:
@ -116,8 +114,6 @@ class DockerSandbox:
await handle.kill()
raise SandboxError(f"docker {args[0]} timed out after {timeout}s")
# -- command construction --------------------------------------------------
def run_args(
self,
) -> list[str]: # mutable-ok: argv is returned as a list, the shape callers and tests compare against
@ -169,8 +165,6 @@ class DockerSandbox:
joined = path if posixpath.isabs(path) else posixpath.join(self.workdir, path)
return posixpath.normpath(joined)
# -- lifecycle ---------------------------------------------------------------
async def start(self) -> str:
"""Start the container if needed and return its id."""
if self._closed:
@ -191,8 +185,6 @@ class DockerSandbox:
container_id = await self.start()
return await self._docker(self.exec_args(container_id, cmd), input=input)
# -- Sandbox protocol --------------------------------------------------------
async def exec(
self,
cmd: Sequence[str],

3
litellm/harness/sync.py
View file

@ -62,9 +62,6 @@ class _LoopThread:
self._thread.start()
return self._loop
def in_loop_thread(self) -> bool:
return self._thread is not None and threading.current_thread() is self._thread
def submit(self, coro: Coroutine[Any, Any, T]) -> Future[T]:
return asyncio.run_coroutine_threadsafe(coro, self.loop())

15
litellm/integrations/azure_storage/azure_storage.py
View file

@ -5,6 +5,7 @@ from collections.abc import Callable
from datetime import datetime, timedelta
from functools import cache
from typing import Final
from urllib.parse import unquote
from litellm._logging import verbose_logger
from litellm._uuid import uuid
@ -31,13 +32,19 @@ from litellm.types.utils import StandardLoggingPayload
AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
_ADLS_SAFE_NAME: Final = str.maketrans("/", "_", "=")
_DOT_OR_EMPTY_SEGMENTS: Final = frozenset(("", ".", ".."))
def adls_safe_file_name(payload_id: str | None) -> str:
"""`=` padding and `/` in a base64 payload id are what the Data Lake service rejects, so the name drops the
padding and maps `/` to `_`. Standard base64 has no `_` and its padding is fixed by the length, so ids from
that alphabet stay distinct; anything else is left as is."""
return f"{(payload_id or str(uuid.uuid4())).translate(_ADLS_SAFE_NAME)}.json"
"""A Responses API id is base64 behind `resp_`, and the Data Lake service rejects its `=` padding and `/`, so
that name drops the padding and maps `/` to `_`. Standard base64 has no `_` and its padding is fixed by the
length, so those ids stay distinct. Every other id, including a caller's `x-litellm-call-id`, is used as is
unless it has an empty, `.` or `..` path segment, which gets the same rewrite so the file keeps its own name in
the log directory"""
name: Final = payload_id or str(uuid.uuid4())
if not name.startswith("resp_") and _DOT_OR_EMPTY_SEGMENTS.isdisjoint(unquote(name).split("/")):
return f"{name}.json"
return f"{name.translate(_ADLS_SAFE_NAME)}.json"
@cache

7
litellm/integrations/clickhouse/clickhouse_batch_logger.py
View file

@ -9,7 +9,6 @@ gzip JSONEachRow insert, either every `CLICKHOUSE_FLUSH_INTERVAL_SECONDS` or as
"""
import asyncio
import os
from collections.abc import Mapping, Sequence
from contextlib import suppress
from typing import Any, ClassVar, Final
@ -23,13 +22,11 @@ from litellm.constants import (
)
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.rust_bridge.traces import ClickHouseStorage
from litellm.tracing.config import trace_storage_config
def clickhouse_storage_from_env() -> ClickHouseStorage:
return ClickHouseStorage(
database=os.getenv("CLICKHOUSE_DATABASE", "litellm"),
url=os.getenv("CLICKHOUSE_URL", ""),
)
return ClickHouseStorage(trace_storage_config({}))
class ClickHouseBatchLogger(CustomBatchLogger):

4
litellm/integrations/clickhouse/schema.py
View file

@ -7,5 +7,5 @@ AGENT_TRACES_BY_KEY_TABLE: Final = "agent_traces_by_key"
SPEND_LOGS_TABLE: Final = "spend_logs"
async def ensure_schema(storage: ClickHouseStorage, trace_retention_days: int, spend_log_retention_days: int) -> None:
await storage.ensure_schema(trace_retention_days, spend_log_retention_days)
async def ensure_schema(storage: ClickHouseStorage) -> None:
await storage.ensure_schema()

5
litellm/litellm_core_utils/litellm_logging.py
View file

@ -112,6 +112,7 @@ from litellm.litellm_core_utils.served_output_texts import (
SERVED_OUTPUT_TEXTS_KEY,
overlay_served_output_texts,
)
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.llms.base_llm.search.transformation import SearchResponse
from litellm.responses.utils import ResponseAPILoggingUtils
@ -175,7 +176,7 @@ from litellm.types.utils import (
Usage,
)
from litellm.types.videos.main import VideoObject
from litellm.utils import _get_base_model_from_metadata, executor, print_verbose
from litellm.utils import _get_base_model_from_metadata, print_verbose
from ..integrations.argilla import ArgillaLogger
from ..integrations.arize.arize_phoenix import ArizePhoenixLogger
@ -3970,7 +3971,7 @@ class Logging(LiteLLMLoggingBaseClass):
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
cache_hit: object | None = None,
cache_hit: bool | None = None,
) -> None:
"""
Handles calling success callbacks for Async calls.

2
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1104,7 +1104,7 @@ class CustomStreamWrapper:
self,
chunk: Any,
model_response: ModelResponseStream,
completion_obj: dict[str, Any],
completion_obj: dict[str, object],
) -> _ProviderChunkResult:
response_obj: dict[str, Any] = {}
if (

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

@ -35,6 +35,8 @@ from ..common_utils import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
_CHATGPT_SERVICE_TIERS: Final = {"default": "default", "priority": "priority", "fast": "priority"}
class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def __init__(self) -> None:
@ -108,7 +110,11 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
"truncation",
}
return {k: v for k, v in request.items() if k in allowed_keys}
filtered: Final = {k: v for k, v in request.items() if k in allowed_keys}
service_tier: Final = _CHATGPT_SERVICE_TIERS.get(request.get("service_tier"))
if service_tier is not None:
filtered["service_tier"] = service_tier
return filtered
def transform_response_api_response(
self,

12
litellm/llms/deepagents/harness/sandbox_backend.py
View file

@ -129,8 +129,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
def id(self) -> str:
return self._id
# -- paths --------------------------------------------------------------
def to_real(self, path: str) -> str:
"""Sandbox path for a virtual path (or an absolute path already under workdir)."""
normalized = posixpath.normpath("/" + path.lstrip("/"))
@ -186,8 +184,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
async def _run(self, cmd: Sequence[str], timeout: float | None = DEEPAGENTS_FS_TIMEOUT_SECONDS) -> CompletedRun:
return await self._sandbox.run(cmd, timeout=timeout)
# -- ls -----------------------------------------------------------------
async def als(self, path: str) -> LsResult:
try:
real = await self.to_confined(path)
@ -207,8 +203,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
def ls(self, path: str) -> LsResult:
return self._sync(self.als(path))
# -- read / write / edit ------------------------------------------------
async def _read_bytes(self, path: str) -> bytes:
return await self._sandbox.read(await self.to_confined(path))
@ -296,8 +290,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
def delete(self, file_path: str) -> DeleteResult:
return self._sync(self.adelete(file_path))
# -- glob / grep --------------------------------------------------------
def _find_cmd(self, root: str) -> tuple[str, ...]:
prune = tuple(
itertools.chain.from_iterable(
@ -378,8 +370,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
) -> GrepResult:
return self._sync(self.agrep(pattern, path, glob, max_count=max_count))
# -- upload / download --------------------------------------------------
async def _upload_one(self, path: str, data: bytes) -> FileUploadResponse:
if not self._writable:
return FileUploadResponse(path=path, error="permission_denied")
@ -425,8 +415,6 @@ class SandboxBackend(SandboxBackendProtocol): # pyright: ignore[reportUntypedBa
) -> list[FileDownloadResponse]: # mutable-ok: return type fixed by deepagents BackendProtocol
return self._sync(self.adownload_files(paths))
# -- execute ------------------------------------------------------------
async def aexecute(self, command: str, *, timeout: int | None = None) -> ExecuteResponse:
if not self._allow_execute:
return ExecuteResponse(output=_NO_EXECUTE_ERROR, exit_code=1)

5
litellm/llms/deepagents/harness/transformation.py
View file

@ -70,11 +70,6 @@ APPROVAL_TOOLS: Final = WRITE_TOOLS | EXECUTE_TOOLS
_APPROVAL_DECISIONS: Final = ("approve", "reject")
# ---------------------------------------------------------------------------
# Pure helpers (unit tested directly; kept module-level so they port cleanly)
# ---------------------------------------------------------------------------
def gateway_headers(
ctx: SessionContext,
) -> dict[str, str]: # mutable-ok: ChatLiteLLM.extra_headers is a pydantic dict field

1
litellm/llms/laya/__init__.py Normal file
View file

@ -0,0 +1 @@

60
litellm/llms/laya/common_utils.py Normal file
View file

@ -0,0 +1,60 @@
from collections.abc import Mapping
from dataclasses import dataclass, field
from typing import Final, Literal, TypeAlias
from pydantic import AnyHttpUrl, BaseModel, TypeAdapter, ValidationError
from litellm.secret_managers.main import get_secret_str
LayaCheckpoint: TypeAlias = Literal["english", "multilingual", "typed-decisions"]
def validate_laya_model(value: object) -> LayaCheckpoint:
try:
return TypeAdapter(LayaCheckpoint).validate_python(value)
except ValidationError as exc:
raise ValueError("Laya model must be 'english', 'multilingual', or 'typed-decisions'") from exc
def validate_laya_request(body: Mapping[str, object]) -> LayaCheckpoint:
if "custom_body" in body:
raise ValueError("custom_body is not supported for Laya requests")
if body.get("stream"):
raise ValueError("Streaming is not supported for Laya requests")
return validate_laya_model(body.get("model"))
@dataclass(frozen=True, slots=True)
class LayaConnection:
api_base: str
api_key: str | None = field(repr=False)
def validate_laya_api_base(value: str) -> str:
try:
url: Final = TypeAdapter(AnyHttpUrl).validate_python(value)
except ValidationError as exc:
raise ValueError("Laya api_base must be an HTTP or HTTPS server URL") from exc
if url.username or url.password or url.query or url.fragment:
raise ValueError("Laya api_base must not contain credentials, a query, or a fragment")
return str(url).rstrip("/")
def laya_connection(api_base: str | None = None, api_key: str | None = None) -> LayaConnection:
base: Final = api_base if api_base is not None else get_secret_str("LAYA_API_BASE")
if not base:
raise ValueError("Laya requires api_base or LAYA_API_BASE pointing to a self-hosted server")
key: Final = api_key if api_base is not None else api_key or get_secret_str("LAYA_API_KEY")
return LayaConnection(api_base=validate_laya_api_base(base), api_key=key)
class _LayaRouting(BaseModel):
model: str | None = None
def laya_response_model(response: Mapping[str, object], requested_model: str | None) -> str:
try:
routing: Final = TypeAdapter(_LayaRouting).validate_python(response.get("routing") or _LayaRouting())
except ValidationError:
return requested_model or "unknown"
return routing.model or requested_model or "unknown"

147
litellm/llms/oci/common_utils.py
View file

@ -1,16 +1,21 @@
import base64
import hashlib
import importlib
import json
import os
import re
from collections.abc import Mapping
from dataclasses import dataclass
from email.utils import formatdate
from typing import Final, Protocol
from pathlib import Path
from types import MappingProxyType
from typing import Final, Protocol, runtime_checkable
from urllib.parse import urlparse
import httpx
from pydantic import JsonValue
from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError, field_validator
from litellm._logging import verbose_logger
from litellm.llms.base_llm.chat.transformation import BaseLLMException
try:
@ -154,7 +159,7 @@ _OCI_KEY_ENV: Final = "OCI_KEY"
_OCI_COMPARTMENT_ID_ENV: Final = "OCI_COMPARTMENT_ID"
def resolve_oci_credentials(optional_params: dict) -> dict:
def resolve_oci_credentials(optional_params: Mapping[str, object]) -> dict:
"""
Merge OCI credentials from optional_params (explicit, always wins) and
environment variables (fallback).
@ -174,11 +179,140 @@ def resolve_oci_credentials(optional_params: dict) -> dict:
}
_OCI_REGION_RE: Final = re.compile(r"^[a-z][a-z0-9-]{0,30}[a-z0-9]$")
_OCI_REGION_PATTERN: Final = r"^[a-z][a-z0-9-]{0,30}[a-z0-9]$"
_OCI_REALM_DOMAIN_PATTERN: Final = r"^[a-z0-9]([a-z0-9-]*[a-z0-9])?(\.[a-z0-9]([a-z0-9-]*[a-z0-9])?)*$"
_OCI_REGION_RE: Final = re.compile(_OCI_REGION_PATTERN)
_OCI_ACTION_PATH_RE: Final = re.compile(rf"/{OCI_API_VERSION}/actions/[^/?#]+/?$")
_OCI_COMMERCIAL_REALM_DOMAIN: Final = "oraclecloud.com"
_OCI_INFERENCE_ENDPOINT_TEMPLATE: Final = "https://inference.generativeai.{region}.oci.{secondLevelDomain}"
_OCI_REGION_METADATA_ENV: Final = "OCI_REGION_METADATA"
_OCI_REGIONS_CONFIG_FILE: Final = "~/.oci/regions-config.json"
_OCID_REALM_RE: Final = re.compile(r"^ocid1\.[a-z0-9]+\.([a-z0-9]+)\.", re.IGNORECASE)
_OCI_REALM_DOMAINS: Final = MappingProxyType(
{
"oc1": "oraclecloud.com",
"oc2": "oraclegovcloud.com",
"oc3": "oraclegovcloud.com",
"oc4": "oraclegovcloud.uk",
"oc8": "oraclecloud8.com",
"oc9": "oraclecloud9.com",
"oc10": "oraclecloud10.com",
"oc14": "oraclecloud14.com",
"oc15": "oraclecloud15.com",
"oc19": "oraclecloud.eu",
"oc20": "oraclecloud20.com",
"oc21": "oraclecloud21.com",
"oc23": "oraclecloud23.com",
"oc24": "oraclecloud24.com",
"oc26": "oraclecloud26.com",
"oc29": "oraclecloud29.com",
"oc35": "oraclecloud35.com",
"oc42": "oraclecloud42.com",
"oc51": "oraclecloud51.com",
"oc52": "oraclecloud52.com",
}
)
def get_oci_base_url(optional_params: dict, api_base: str | None = None) -> str:
class OCIRegionMetadata(BaseModel):
"""One entry of the OCI SDK's region metadata schema, as found in
``~/.oci/regions-config.json`` (a JSON array) or ``OCI_REGION_METADATA`` (one object).
Values are lowercased before validation, as the SDK does."""
model_config = ConfigDict(frozen=True, extra="ignore")
region_identifier: str = Field(alias="regionIdentifier", pattern=_OCI_REGION_PATTERN)
realm_domain_component: str = Field(alias="realmDomainComponent", pattern=_OCI_REALM_DOMAIN_PATTERN)
@field_validator("region_identifier", "realm_domain_component", mode="before")
@classmethod
def _lowercase(cls, value: object) -> object:
return value.lower() if isinstance(value, str) else value
_JSON_ARRAY: Final = TypeAdapter(tuple[JsonValue, ...])
def _validated_region_metadata(raw: JsonValue, source: str) -> OCIRegionMetadata | None:
try:
return OCIRegionMetadata.model_validate(raw)
except ValidationError as e:
verbose_logger.warning("Ignoring OCI region metadata entry in %s: %s", source, e)
return None
def _region_metadata_from_file() -> tuple[OCIRegionMetadata, ...]:
path: Final = Path(os.path.expanduser(_OCI_REGIONS_CONFIG_FILE))
if not path.is_file():
return ()
try:
raw_entries: Final = _JSON_ARRAY.validate_json(path.read_bytes())
except (OSError, ValidationError) as e:
verbose_logger.warning("Ignoring OCI region metadata in %s: %s", path, e)
return ()
candidates: Final = (_validated_region_metadata(raw, str(path)) for raw in raw_entries)
return tuple(entry for entry in candidates if entry is not None)
def _region_metadata_from_env() -> tuple[OCIRegionMetadata, ...]:
raw: Final = os.environ.get(_OCI_REGION_METADATA_ENV)
if not raw:
return ()
try:
return (OCIRegionMetadata.model_validate_json(raw),)
except ValidationError as e:
verbose_logger.warning("Ignoring OCI region metadata in %s: %s", _OCI_REGION_METADATA_ENV, e)
return ()
def _realm_domain_from_ocid(ocid: str | None) -> str | None:
match: Final = _OCID_REALM_RE.match(ocid) if ocid else None
return _OCI_REALM_DOMAINS.get(match.group(1).lower()) if match else None
def _realm_domain_from_metadata(region: str) -> str | None:
entries: Final = (*_region_metadata_from_file(), *_region_metadata_from_env())
return next((entry.realm_domain_component for entry in entries if entry.region_identifier == region), None)
@runtime_checkable
class _OCIRegionRegistry(Protocol):
def endpoint_for(self, service: str, region: str, service_endpoint_template: str) -> str: ...
def _load_oci_region_registry() -> _OCIRegionRegistry | None:
try:
registry: Final = importlib.import_module("oci.regions")
except ImportError:
return None
return registry if isinstance(registry, _OCIRegionRegistry) else None
def resolve_oci_inference_endpoint(region: str, compartment_id: str | None = None) -> str:
"""Return the GenAI inference endpoint for ``region`` in whichever OCI realm hosts it.
The realm's second-level domain comes first from the realm key inside ``compartment_id``
(``ocid1.compartment.oc2..`` is the Government realm), then from the OCI SDK's region
registry when the SDK is installed, then from the per-region metadata sources the SDK
reads, ``~/.oci/regions-config.json`` and ``OCI_REGION_METADATA``, and otherwise defaults
to the commercial realm. Realm domains per ``oci/regions_definitions.py`` in oci 2.187.0.
A region that is not described anywhere therefore keeps its commercial endpoint, so one
government deployment never redirects the others.
"""
realm_domain: Final = _realm_domain_from_ocid(compartment_id)
if realm_domain is not None:
return _OCI_INFERENCE_ENDPOINT_TEMPLATE.format(region=region, secondLevelDomain=realm_domain)
registry: Final = _load_oci_region_registry()
if registry is not None:
return registry.endpoint_for(
"generative_ai_inference", region=region, service_endpoint_template=_OCI_INFERENCE_ENDPOINT_TEMPLATE
)
return _OCI_INFERENCE_ENDPOINT_TEMPLATE.format(
region=region, secondLevelDomain=_realm_domain_from_metadata(region) or _OCI_COMMERCIAL_REALM_DOMAIN
)
def get_oci_base_url(optional_params: Mapping[str, object], api_base: str | None = None) -> str:
"""Return the OCI inference base URL, respecting any explicit api_base override.
If ``api_base`` already ends with a fully-formed OCI action path
@ -196,7 +330,8 @@ def get_oci_base_url(optional_params: dict, api_base: str | None = None) -> str:
f"Invalid OCI region {region!r}: must match ^[a-z][a-z0-9-]{{0,30}}[a-z0-9]$ (e.g. 'us-ashburn-1')."
),
)
return f"https://inference.generativeai.{region}.oci.oraclecloud.com"
compartment_id: Final = creds["oci_compartment_id"]
return resolve_oci_inference_endpoint(region, compartment_id if isinstance(compartment_id, str) else None)
# ---------------------------------------------------------------------------

6
litellm/llms/oci/embed/transformation.py
View file

@ -77,7 +77,11 @@ class OCIEmbedConfig(BaseEmbeddingConfig):
Required call-time params (via optional_params or env vars):
- ``oci_compartment_id`` / ``OCI_COMPARTMENT_ID``
- ``oci_region`` / ``OCI_REGION`` (default: ``us-ashburn-1``)
- ``oci_region`` / ``OCI_REGION`` (default: ``us-ashburn-1``). The realm comes from the realm
key in ``oci_compartment_id`` (``ocid1.compartment.oc2..`` is the Government realm), so
non-commercial realms need no extra setting. A realm unknown to litellm can be described in
``OCI_REGION_METADATA`` or ``~/.oci/regions-config.json``, resolved through the OCI SDK when
it is installed, or given as ``api_base``.
Optional call-time params:
- ``oci_serving_mode``: ``"ON_DEMAND"`` (default) or ``"DEDICATED"``

199
litellm/model_prices_and_context_window_backup.json
View file

@ -3642,7 +3642,7 @@
"prompt_cache_min_tokens": 1024
},
"azure_ai/claude-sonnet-4-5": {
"deprecation_date": "2026-11-15",
"deprecation_date": "2026-11-30",
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_1hr": 6e-06,
"cache_read_input_token_cost": 3e-07,
@ -5516,7 +5516,7 @@
"supports_web_search": false
},
"azure/gpt-4.1-nano": {
"deprecation_date": "2026-10-14",
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_token": 1e-07,
"input_cost_per_token_batches": 5e-08,
@ -5550,7 +5550,7 @@
"supports_vision": true
},
"azure/gpt-4.1-nano-2025-04-14": {
"deprecation_date": "2026-10-14",
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 2.5e-08,
"input_cost_per_token": 1e-07,
"input_cost_per_token_batches": 5e-08,
@ -6133,7 +6133,7 @@
"cache_creation_input_audio_token_cost": 4e-07,
"cache_read_input_audio_token_cost": 4e-07,
"cache_read_input_token_cost": 4e-07,
"deprecation_date": "2027-06-25",
"deprecation_date": "2027-07-31",
"input_cost_per_audio_token": 3.2e-05,
"input_cost_per_image_token": 5e-06,
"input_cost_per_token": 4e-06,
@ -6168,7 +6168,7 @@
"cache_creation_input_audio_token_cost": 3e-07,
"cache_read_input_audio_token_cost": 3e-07,
"cache_read_input_token_cost": 6e-08,
"deprecation_date": "2027-06-25",
"deprecation_date": "2027-07-31",
"input_cost_per_audio_token": 1e-05,
"input_cost_per_image_token": 8e-07,
"input_cost_per_token": 6e-07,
@ -6346,7 +6346,7 @@
"supports_tool_choice": true
},
"azure/gpt-4o-transcribe": {
"deprecation_date": "2026-10-15",
"deprecation_date": "2026-12-31",
"input_cost_per_audio_token": 2.5e-06,
"input_cost_per_token": 2.5e-06,
"litellm_provider": "azure",
@ -6354,6 +6354,7 @@
"max_output_tokens": 2000,
"mode": "audio_transcription",
"output_cost_per_token": 1e-05,
"source": "https://management.azure.com/subscriptions/c873328e-b572-4770-8dff-aaeb6f1f0e79/providers/Microsoft.CognitiveServices/locations/eastus2/models?api-version=2024-10-01",
"supported_endpoints": [
"/v1/audio/transcriptions"
]
@ -8571,6 +8572,7 @@
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"deprecation_date": "2028-03-11",
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "azure",
@ -8619,6 +8621,7 @@
"cache_creation_input_token_cost_above_272k_tokens": 5e-06,
"cache_read_input_token_cost": 1e-07,
"cache_read_input_token_cost_above_272k_tokens": 2e-07,
"deprecation_date": "2028-03-11",
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_272k_tokens": 4e-06,
"litellm_provider": "azure",
@ -10972,7 +10975,7 @@
"supports_web_search": false
},
"azure/us/gpt-4.1-nano-2025-04-14": {
"deprecation_date": "2026-10-14",
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 2.8e-08,
"input_cost_per_token": 1.1e-07,
"input_cost_per_token_batches": 5.5e-08,
@ -11528,7 +11531,7 @@
"max_tokens": 5000,
"mode": "image_generation",
"output_cost_per_image": 0.04,
"source": "https://marketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice",
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/black-forest-labs/",
"supported_endpoints": [
"/v1/images/generations"
]
@ -11958,6 +11961,7 @@
"supports_vision": true
},
"azure_ai/Meta-Llama-3-70B-Instruct": {
"deprecation_date": "2025-06-30",
"input_cost_per_token": 1.1e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 8192,
@ -11965,9 +11969,11 @@
"max_tokens": 2048,
"mode": "chat",
"output_cost_per_token": 3.7e-07,
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true
},
"azure_ai/Meta-Llama-3.1-70B-Instruct": {
"deprecation_date": "2025-06-30",
"input_cost_per_token": 2.68e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -11975,10 +11981,11 @@
"max_tokens": 2048,
"mode": "chat",
"output_cost_per_token": 3.54e-06,
"source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-70b-instruct-offer?tab=PlansAndPrice",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true
},
"azure_ai/Phi-3-medium-128k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.7e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -11986,11 +11993,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6.8e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3-medium-4k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.7e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 4096,
@ -11998,11 +12006,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6.8e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3-mini-128k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.3e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12010,11 +12019,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 5.2e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3-mini-4k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.3e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 4096,
@ -12022,11 +12032,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 5.2e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3-small-128k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12034,11 +12045,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3-small-8k-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 8192,
@ -12046,11 +12058,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3.5-MoE-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.6e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12058,11 +12071,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6.4e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3.5-mini-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.3e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12070,11 +12084,12 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 5.2e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": false
},
"azure_ai/Phi-3.5-vision-instruct": {
"deprecation_date": "2025-08-30",
"input_cost_per_token": 1.3e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12082,7 +12097,7 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 5.2e-07,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true,
"supports_vision": true
},
@ -12217,6 +12232,7 @@
"source": "https://azure.microsoft.com/en-us/pricing/details/ai-document-intelligence/"
},
"azure_ai/MAI-DS-R1": {
"deprecation_date": "2026-02-27",
"input_cost_per_token": 1.35e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12224,11 +12240,12 @@
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 5.4e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_reasoning": true,
"supports_tool_choice": true
},
"azure_ai/cohere-rerank-v3-english": {
"deprecation_date": "2025-06-30",
"input_cost_per_query": 0.002,
"input_cost_per_token": 0.0,
"litellm_provider": "azure_ai",
@ -12236,9 +12253,11 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "rerank",
"output_cost_per_token": 0.0
"output_cost_per_token": 0.0,
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models"
},
"azure_ai/cohere-rerank-v3-multilingual": {
"deprecation_date": "2025-06-30",
"input_cost_per_query": 0.002,
"input_cost_per_token": 0.0,
"litellm_provider": "azure_ai",
@ -12246,7 +12265,8 @@
"max_output_tokens": 4096,
"max_tokens": 4096,
"mode": "rerank",
"output_cost_per_token": 0.0
"output_cost_per_token": 0.0,
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models"
},
"azure_ai/cohere-rerank-v4.0-pro": {
"input_cost_per_query": 0.0025,
@ -12301,6 +12321,7 @@
"supports_tool_choice": true
},
"azure_ai/deepseek-v3": {
"deprecation_date": "2025-08-31",
"input_cost_per_token": 1.14e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12308,7 +12329,7 @@
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 4.56e-06,
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_tool_choice": true
},
"azure_ai/deepseek-v4-pro": {
@ -12515,6 +12536,7 @@
"supports_web_search": true
},
"azure_ai/jais-30b-chat": {
"deprecation_date": "2026-01-30",
"input_cost_per_token": 0.0032,
"litellm_provider": "azure_ai",
"max_input_tokens": 8192,
@ -12522,7 +12544,7 @@
"max_tokens": 8192,
"mode": "chat",
"output_cost_per_token": 0.00971,
"source": "https://ai.azure.com/catalog/models/jais-30b-chat"
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models"
},
"azure_ai/jamba-instruct": {
"input_cost_per_token": 5e-07,
@ -12588,6 +12610,7 @@
"supports_tool_choice": true
},
"azure_ai/mistral-large": {
"deprecation_date": "2025-04-15",
"input_cost_per_token": 4e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 32000,
@ -12595,10 +12618,12 @@
"max_tokens": 8191,
"mode": "chat",
"output_cost_per_token": 1.2e-05,
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_function_calling": true,
"supports_tool_choice": true
},
"azure_ai/mistral-large-2407": {
"deprecation_date": "2025-05-13",
"input_cost_per_token": 2e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 128000,
@ -12606,7 +12631,7 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 6e-06,
"source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -12648,6 +12673,7 @@
"supports_tool_choice": true
},
"azure_ai/mistral-nemo": {
"deprecation_date": "2026-01-30",
"input_cost_per_token": 1.5e-07,
"litellm_provider": "azure_ai",
"max_input_tokens": 131072,
@ -12655,10 +12681,11 @@
"max_tokens": 4096,
"mode": "chat",
"output_cost_per_token": 1.5e-07,
"source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-nemo-12b-2407?tab=PlansAndPrice",
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_function_calling": true
},
"azure_ai/mistral-small": {
"deprecation_date": "2025-07-31",
"input_cost_per_token": 1e-06,
"litellm_provider": "azure_ai",
"max_input_tokens": 32000,
@ -12666,6 +12693,7 @@
"max_tokens": 8191,
"mode": "chat",
"output_cost_per_token": 3e-06,
"source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/retired-models",
"supports_function_calling": true,
"supports_tool_choice": true
},
@ -32787,6 +32815,7 @@
]
},
"gpt-4o-mini-tts-2025-03-20": {
"deprecation_date": "2027-01-06",
"input_cost_per_token": 6e-07,
"litellm_provider": "openai",
"mode": "audio_speech",
@ -33486,7 +33515,8 @@
"output_cost_per_token_flex": 5e-06,
"source": "https://developers.openai.com/api/docs/pricing",
"supports_xhigh_reasoning_effort": false,
"supports_minimal_reasoning_effort": false
"supports_minimal_reasoning_effort": false,
"deprecation_date": "2027-04-01"
},
"gpt-5.1-codex-mini": {
"cache_read_input_token_cost": 2.5e-08,
@ -35283,7 +35313,8 @@
"default_reasoning_effort": "none",
"source": "https://developers.openai.com/api/docs/pricing",
"supports_xhigh_reasoning_effort": true,
"supports_minimal_reasoning_effort": false
"supports_minimal_reasoning_effort": false,
"deprecation_date": "2027-04-01"
},
"gpt-5.4-nano-2026-03-17": {
"cache_read_input_token_cost": 2e-08,
@ -35582,7 +35613,8 @@
"supports_web_search": true,
"supports_none_reasoning_effort": false,
"supports_xhigh_reasoning_effort": false,
"supports_minimal_reasoning_effort": true
"supports_minimal_reasoning_effort": true,
"deprecation_date": "2027-04-01"
},
"gpt-5.3-chat-latest": {
"cache_read_input_token_cost": 1.75e-07,
@ -42206,14 +42238,14 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4-pro": {
"cache_read_input_token_cost": 6.525e-08,
"input_cost_per_token": 7.83e-07,
"cache_read_input_token_cost": 1.74e-08,
"input_cost_per_token": 2.088e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"mode": "chat",
"output_cost_per_token": 1.566e-06,
"output_cost_per_token": 4.176e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -42226,14 +42258,14 @@
"supports_web_search": false
},
"openrouter/deepseek/deepseek-v4.1-flash": {
"cache_read_input_token_cost": 2.91e-09,
"input_cost_per_token": 1.98e-08,
"cache_read_input_token_cost": 1e-08,
"input_cost_per_token": 3e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"output_cost_per_token": 3.96e-07,
"output_cost_per_token": 5e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -42820,14 +42852,14 @@
"supports_web_search": false
},
"openrouter/nvidia/nemotron-3.5-lightning": {
"cache_read_input_token_cost": 3e-08,
"input_cost_per_token": 6e-08,
"cache_read_input_token_cost": 2.975e-08,
"input_cost_per_token": 5.95e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 32768,
"max_tokens": 32768,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.6e-07,
"output_cost_per_token": 1.7e-07,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -43829,6 +43861,7 @@
"openrouter/z-ai/glm-4.7": {
"cache_creation_input_token_cost": 0.0,
"cache_read_input_token_cost": 1.1e-07,
"deprecation_date": "2026-12-31",
"input_cost_per_token": 6e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 204800,
@ -46840,6 +46873,7 @@
"supports_vision": true
},
"tts-1": {
"deprecation_date": "2027-01-06",
"input_cost_per_character": 1.5e-05,
"litellm_provider": "openai",
"mode": "audio_speech",
@ -46849,6 +46883,7 @@
]
},
"tts-1-hd": {
"deprecation_date": "2027-01-06",
"input_cost_per_character": 3e-05,
"litellm_provider": "openai",
"mode": "audio_speech",
@ -56731,6 +56766,7 @@
]
},
"gpt-4o-mini-tts-2025-12-15": {
"deprecation_date": "2027-01-06",
"input_cost_per_token": 6e-07,
"litellm_provider": "openai",
"mode": "audio_speech",
@ -67382,13 +67418,13 @@
"supports_web_search": false
},
"openrouter/z-ai/glm-5.3": {
"input_cost_per_token": 1.4e-06,
"output_cost_per_token": 4.4e-06,
"cache_read_input_token_cost": 2.6e-07,
"input_cost_per_token": 2.219e-07,
"output_cost_per_token": 3.39e-06,
"cache_read_input_token_cost": 1.775e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1310720,
"max_output_tokens": 943717,
"max_tokens": 943717,
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
"max_tokens": 943718,
"mode": "chat",
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
@ -67519,8 +67555,8 @@
"supports_prompt_caching": true
},
"openrouter/deepseek/deepseek-v4-flash-0731": {
"cache_read_input_token_cost": 8.9e-09,
"input_cost_per_token": 8.9e-09,
"cache_read_input_token_cost": 1.08e-08,
"input_cost_per_token": 1.08e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
@ -67608,8 +67644,8 @@
"supports_web_search": false
},
"openrouter/moonshotai/kimi-k3": {
"cache_read_input_token_cost": 2.7e-07,
"input_cost_per_token": 2.8e-07,
"cache_read_input_token_cost": 4.357e-07,
"input_cost_per_token": 4.357e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
@ -67732,7 +67768,7 @@
},
"openrouter/z-ai/glm-5.2": {
"cache_read_input_token_cost": 2.6e-07,
"input_cost_per_token": 3.249e-07,
"input_cost_per_token": 4.1e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 943718,
@ -68094,14 +68130,14 @@
"supports_web_search": true
},
"openrouter/deepseek/deepseek-v4-flash": {
"cache_read_input_token_cost": 1.5708e-08,
"input_cost_per_token": 7.854e-08,
"cache_read_input_token_cost": 8.372e-09,
"input_cost_per_token": 4.186e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 384000,
"max_tokens": 384000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5708e-07,
"output_cost_per_token": 8.372e-08,
"source": "https://openrouter.ai/api/v1/models",
"supports_audio_input": false,
"supports_function_calling": true,
@ -68114,9 +68150,9 @@
"supports_web_search": false
},
"openrouter/moonshotai/kimi-k2.6": {
"input_cost_per_token": 6.5e-07,
"output_cost_per_token": 3.41e-06,
"cache_read_input_token_cost": 1.5e-07,
"input_cost_per_token": 4.3415e-07,
"output_cost_per_token": 1.828e-06,
"cache_read_input_token_cost": 7.312e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 262144,
"max_output_tokens": 235929,
@ -68563,7 +68599,7 @@
"openrouter/z-ai/glm-4.6v": {
"input_cost_per_token": 3e-07,
"output_cost_per_token": 9e-07,
"cache_read_input_token_cost": 5.5e-08,
"cache_read_input_token_cost": 5e-08,
"litellm_provider": "openrouter",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
@ -69229,7 +69265,7 @@
"supports_web_search": false
},
"openrouter/minimax/minimax-m1": {
"input_cost_per_token": 4e-07,
"input_cost_per_token": 5.5e-07,
"output_cost_per_token": 2.2e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
@ -70653,7 +70689,7 @@
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/eu/gpt-4.1-nano": {
"deprecation_date": "2026-10-14",
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 2.8e-08,
"input_cost_per_token": 1.1e-07,
"input_cost_per_token_batches": 5.5e-08,
@ -71102,7 +71138,7 @@
"source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'"
},
"azure/us/gpt-4.1-nano": {
"deprecation_date": "2026-10-14",
"deprecation_date": "2027-04-14",
"cache_read_input_token_cost": 2.8e-08,
"input_cost_per_token": 1.1e-07,
"input_cost_per_token_batches": 5.5e-08,
@ -72586,6 +72622,45 @@
"supports_audio_input": true,
"supports_video_input": true
},
"laya/english": {
"input_cost_per_token": 0.0,
"litellm_provider": "laya",
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://github.com/NandhaKishorM/laya",
"supported_endpoints": [
"/v1/systemone"
],
"metadata": {
"notes": "Self-hosted decision model; infrastructure costs are paid separately"
}
},
"laya/multilingual": {
"input_cost_per_token": 0.0,
"litellm_provider": "laya",
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://github.com/NandhaKishorM/laya",
"supported_endpoints": [
"/v1/systemone"
],
"metadata": {
"notes": "Self-hosted decision model; infrastructure costs are paid separately"
}
},
"laya/typed-decisions": {
"input_cost_per_token": 0.0,
"litellm_provider": "laya",
"mode": "evaluation",
"output_cost_per_token": 0.0,
"source": "https://github.com/NandhaKishorM/laya",
"supported_endpoints": [
"/v1/systemone"
],
"metadata": {
"notes": "Self-hosted decision model; infrastructure costs are paid separately"
}
},
"typesafe/jev-1.13.0": {
"input_cost_per_token": 4.2e-08,
"litellm_provider": "typesafe",

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

@ -142,6 +142,12 @@ from litellm.proxy._experimental.mcp_server.result_conversion import (
from litellm.proxy._experimental.mcp_server.sampling_handler import (
MCP_SAMPLING_AVAILABLE,
)
from litellm.proxy._experimental.mcp_server.stdio_gate import (
MCP_STDIO_DISABLED_MESSAGE,
is_mcp_stdio_blocked,
is_mcp_stdio_enabled,
warn_if_mcp_stdio_blocked,
)
from litellm.proxy._experimental.mcp_server.tool_catalog_guard import (
CatalogAlert,
apply_description_overrides,
@ -1754,7 +1760,7 @@ def _create_elicitation_callback():
def _record_mcp_guardrail_evaluations(
synthetic_llm_data: dict[str, Any], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict
synthetic_llm_data: dict[str, object], # mutable-ok: `_sync_guardrail_info_to_logging_obj` takes a concrete dict
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Bridge guardrail decision records off an MCP synthetic request onto the request's logger.
@ -2541,6 +2547,7 @@ class MCPServerManager:
alias=alias,
server_name=server_name,
)
warn_if_mcp_stdio_blocked(server_name, server_config.get("transport"))
auth_type = server_config.get("auth_type", None)
manual_issuer = _blank_to_none(server_config.get("issuer"))
@ -3359,6 +3366,7 @@ class MCPServerManager:
# `credentials` field is the only one still encrypted here).
# Re-decrypting plaintext would zero the values, so build with
# env_vars_are_encrypted=False.
self._warn_if_newly_blocked_stdio(mcp_server, None)
new_server: Final = await self.build_mcp_server_from_table(mcp_server, env_vars_are_encrypted=False)
self._assign_unique_short_prefix(new_server)
self._invalidate_server_definition_caches(mcp_server.server_id)
@ -4350,6 +4358,8 @@ class MCPServerManager:
# Handle stdio transport
if transport == MCPTransport.stdio:
if not is_mcp_stdio_enabled():
raise HTTPException(status_code=403, detail=MCP_STDIO_DISABLED_MESSAGE)
resolved_env: Final = (
stdio_env
if stdio_env is not None
@ -4516,6 +4526,9 @@ class MCPServerManager:
global_mcp_tool_registry,
)
if self._skip_blocked_stdio_listing(server, "tool"):
return []
verbose_logger.debug("Connecting to url: %s", server.url)
verbose_logger.info("_get_tools_from_server for %s...", server.name)
@ -4802,6 +4815,19 @@ class MCPServerManager:
)
return server.server_id, hashlib.sha256(material.encode()).hexdigest()
@staticmethod
def _warn_if_newly_blocked_stdio(row: LiteLLM_MCPServerTable, previous: MCPServer | None) -> None:
if previous is None or previous.transport != row.transport:
warn_if_mcp_stdio_blocked(row.alias or row.server_name, row.transport)
def _skip_blocked_stdio_listing(self, server: MCPServer, listing: str) -> bool:
if not is_mcp_stdio_blocked(server.transport):
return False
verbose_logger.debug(
"Skipping %s listing for MCP server %s: %s", listing, server.name, MCP_STDIO_DISABLED_MESSAGE
)
return True
async def get_prompts_from_server(
self,
server: MCPServer,
@ -4812,6 +4838,8 @@ class MCPServerManager:
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[Prompt]:
if self._skip_blocked_stdio_listing(server, "prompt"):
return []
try:
headers: Final = (
dict(
@ -4858,6 +4886,8 @@ class MCPServerManager:
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[Resource]:
if self._skip_blocked_stdio_listing(server, "resource"):
return []
try:
headers: Final = (
dict(
@ -4904,6 +4934,8 @@ class MCPServerManager:
raw_headers: dict[str, str] | None = None,
client_ip: str | None = None,
) -> list[ResourceTemplate]:
if self._skip_blocked_stdio_listing(server, "resource template"):
return []
try:
headers: Final = (
dict(
@ -6492,6 +6524,8 @@ class MCPServerManager:
mcp_server = fallback
if mcp_server is None:
raise ValueError(f"Tool {name} not found")
if is_mcp_stdio_blocked(mcp_server.transport):
raise HTTPException(status_code=403, detail=MCP_STDIO_DISABLED_MESSAGE)
if resolved_by_server_name_only and not self.server_exposes_tool(mcp_server, name):
raise ValueError(f"Tool {name} not found")
@ -6886,7 +6920,10 @@ class MCPServerManager:
if matched is not None:
matched_prefix, original_tool_name = matched
matched_server: Final = prefix_to_server.get(matched_prefix)
if matched_server is not None and self.server_exposes_tool(matched_server, original_tool_name):
if matched_server is not None and (
self.server_exposes_tool(matched_server, original_tool_name)
or is_mcp_stdio_blocked(matched_server.transport)
):
return matched_server
return None
@ -6948,6 +6985,7 @@ class MCPServerManager:
alias=getattr(server, "alias", None),
server_name=getattr(server, "server_name", None),
)
self._warn_if_newly_blocked_stdio(server, existing_server)
verbose_logger.debug("Building server from DB: %s (%s)", server.server_id, server.server_name)
# raw_rows come straight from the DB, so their global env var
# values (like credentials) are still encrypted here, unlike the

28
litellm/proxy/_experimental/mcp_server/stdio_gate.py Normal file
View file

@ -0,0 +1,28 @@
import os
from typing import Final
from litellm._logging import verbose_logger
from litellm.types.mcp import MCPTransport
MCP_STDIO_ENABLED_ENV_VAR: Final = "LITELLM_ENABLE_MCP_STDIO"
MCP_STDIO_DISABLED_MESSAGE: Final = (
f"stdio MCP servers are disabled on this proxy. "
f"Set {MCP_STDIO_ENABLED_ENV_VAR}=true on the proxy and restart to enable them"
)
def is_mcp_stdio_enabled() -> bool:
return os.getenv(MCP_STDIO_ENABLED_ENV_VAR, "").strip().lower() == "true"
def is_mcp_stdio_flag_key(env_var_name: str) -> bool:
return env_var_name.upper() == MCP_STDIO_ENABLED_ENV_VAR
def is_mcp_stdio_blocked(transport: str | None) -> bool:
return transport == MCPTransport.stdio and not is_mcp_stdio_enabled()
def warn_if_mcp_stdio_blocked(server_name: str | None, transport: str | None) -> None:
if is_mcp_stdio_blocked(transport):
verbose_logger.warning("MCP server '%s' will not start: %s", server_name, MCP_STDIO_DISABLED_MESSAGE)

1
litellm/proxy/_experimental/out/assets/logos/microsoft_365.svg Normal file
View file

@ -0,0 +1 @@
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" role="img" aria-label="Microsoft 365"><rect x="1" y="1" width="10.5" height="10.5" fill="#F25022"/><rect x="12.5" y="1" width="10.5" height="10.5" fill="#7FBA00"/><rect x="1" y="12.5" width="10.5" height="10.5" fill="#00A4EF"/><rect x="12.5" y="12.5" width="10.5" height="10.5" fill="#FFB900"/></svg>
Side by side

After

Width:  |  Height:  |  Size: 361 B

7
litellm/proxy/_lazy_features.py
View file

@ -229,6 +229,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = (
"/tinyfish/",
"/transcribe",
"/typesafe/",
"/laya/",
"/openrouter/",
"/vertex-ai/",
"/vertex_ai/",
@ -530,11 +531,11 @@ def _register_all_on_startup(inner: "Lifespan[FastAPI]", features: tuple[LazyFea
(config pass-through endpoints), so the table is put back in lazy mode's order once it is up."""
@asynccontextmanager
async def lifespan(app: "FastAPI") -> AsyncGenerator[None]:
async def lifespan(app: "FastAPI") -> AsyncGenerator[Mapping[str, object]]:
register_all_features(app, features)
async with inner(app):
async with inner(app) as state:
_restore_registry_order(app, features)
yield
yield state if state is not None else {}
return lifespan

1070
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

116
litellm/proxy/_types.py
View file

@ -28,6 +28,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
validate_langfuse_span_scope_value,
validate_no_callback_env_reference,
)
from litellm.proxy._experimental.mcp_server.stdio_gate import MCP_STDIO_DISABLED_MESSAGE, is_mcp_stdio_enabled
from litellm.types.agents import AgentCaller, AgentResponse
from litellm.types.integrations.compression_interception import (
CompressionSavingsMetadata,
@ -506,6 +507,7 @@ class LiteLLMRoutes(enum.Enum):
"/vllm",
"/mistral",
"/typesafe",
"/laya",
"/openrouter",
"/milvus",
"/gigachat",
@ -537,6 +539,8 @@ class LiteLLMRoutes(enum.Enum):
"/lens/workers/register",
"/lens/workers/{worker_id}",
"/v1/traces",
"/v1/traces/query",
"/v1/traces/query/help",
"/v1/traces/{trace_id}",
"/v1/traces/{trace_id}/spans/{span_id}",
]
@ -699,6 +703,10 @@ class LiteLLMRoutes(enum.Enum):
KeyManagementRoutes.TEAM_KEY_BULK_UPDATE.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY.value,
KeyManagementRoutes.TEAM_DAILY_ACTIVITY_AGGREGATED.value,
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/team/daily/activity/export",
KeyManagementRoutes.SPEND_LOGS.value,
KeyManagementRoutes.SPEND_LOGS_V2.value,
KeyManagementRoutes.KEY_RESET_SPEND.value,
@ -725,6 +733,11 @@ class LiteLLMRoutes(enum.Enum):
"/user/list",
"/user/daily/activity",
"/user/daily/activity/aggregated",
"/user/daily/activity/aggregated/keys",
"/user/daily/activity/aggregated/search",
"/user/daily/activity/aggregated/model_top_keys",
"/user/daily/activity/export",
"/user/daily/activity/aggregated/cache_leakage_keys",
# team
"/team/new",
"/team/update",
@ -742,6 +755,10 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_bulk_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/team/daily/activity/export",
"/team/spend/by_user",
# gateway request counts (SGR); deployment-wide, admin-only
"/gateway/daily/activity",
@ -870,6 +887,11 @@ class LiteLLMRoutes(enum.Enum):
# Tag usage endpoints scope internal users to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/daily/activity/aggregated",
"/tag/daily/activity/aggregated/keys",
"/tag/daily/activity/aggregated/search",
"/tag/daily/activity/aggregated/model_top_keys",
"/tag/daily/activity/export",
"/tag/list",
"/v1/models/{model_id}",
"/models/{model_id}",
@ -894,6 +916,11 @@ class LiteLLMRoutes(enum.Enum):
# Tag usage endpoints scope internal viewers to tags produced by
# their own keys in tag_management_endpoints.py.
"/tag/daily/activity",
"/tag/daily/activity/aggregated",
"/tag/daily/activity/aggregated/keys",
"/tag/daily/activity/aggregated/search",
"/tag/daily/activity/aggregated/model_top_keys",
"/tag/daily/activity/export",
"/tag/list",
]
)
@ -913,6 +940,10 @@ class LiteLLMRoutes(enum.Enum):
"/team/permissions_update",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/team/daily/activity/export",
"/team/spend/by_user",
"/team/{team_id}/members/me",
# POST/GET the team's logging callbacks, and DELETE one of them. Every
@ -928,9 +959,19 @@ class LiteLLMRoutes(enum.Enum):
"/model/delete",
"/user/daily/activity",
"/user/daily/activity/aggregated",
"/user/daily/activity/aggregated/keys",
"/user/daily/activity/aggregated/search",
"/user/daily/activity/aggregated/model_top_keys",
"/user/daily/activity/export",
"/user/daily/activity/aggregated/cache_leakage_keys",
# Endpoint restricts results to organizations the caller is ORG_ADMIN
# of; a caller who administers none gets an empty result set.
"/organization/daily/activity",
"/organization/daily/activity/aggregated",
"/organization/daily/activity/aggregated/keys",
"/organization/daily/activity/aggregated/search",
"/organization/daily/activity/aggregated/model_top_keys",
"/organization/daily/activity/export",
"/user/available_roles", # read-only role metadata; any authenticated user may read
# Claude Code gateway: the signed-in CLI fetches its managed settings and posts its own telemetry
"/claude_code_gateway/managed/settings",
@ -1009,9 +1050,24 @@ class LiteLLMRoutes(enum.Enum):
"/user/available_users",
"/user/available_roles",
"/user/daily/activity",
"/user/daily/activity/aggregated",
"/user/daily/activity/aggregated/keys",
"/user/daily/activity/aggregated/search",
"/user/daily/activity/aggregated/model_top_keys",
"/user/daily/activity/export",
"/user/daily/activity/aggregated/cache_leakage_keys",
"/team/daily/activity",
"/team/daily/activity/aggregated",
"/team/daily/activity/aggregated/keys",
"/team/daily/activity/aggregated/search",
"/team/daily/activity/aggregated/model_top_keys",
"/team/daily/activity/export",
"/tag/daily/activity",
"/tag/daily/activity/aggregated",
"/tag/daily/activity/aggregated/keys",
"/tag/daily/activity/aggregated/search",
"/tag/daily/activity/aggregated/model_top_keys",
"/tag/daily/activity/export",
"/tag/list",
"/audit",
"/audit/{id}",
@ -1529,6 +1585,30 @@ def _reject_unsupported_per_server_oauth_discovery(values: object, require_auth_
raise _per_server_oauth_discovery_error()
def _validate_mcp_transport_fields(values: object) -> None:
if not isinstance(values, dict):
return
transport: Final = values.get("transport")
if transport in (MCPTransport.http, MCPTransport.sse):
if not values.get("url") and not values.get("spec_path"):
raise ValueError("url or spec_path is required for HTTP/SSE transport")
return
if transport != MCPTransport.stdio:
return
if not is_mcp_stdio_enabled():
raise ValueError(MCP_STDIO_DISABLED_MESSAGE)
command: Final = values.get("command")
if not command:
raise ValueError("command is required for stdio transport")
if not values.get("args"):
raise ValueError("args is required for stdio transport")
if os.path.basename(str(command)) not in MCP_STDIO_ALLOWED_COMMANDS:
raise ValueError(
f"Command '{command}' is not in the allowed commands list "
f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}"
)
class NewMCPServerRequest(LiteLLMPydanticObjectBase):
server_id: str | None = None
server_name: str | None = None
@ -1595,23 +1675,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
@model_validator(mode="before")
@classmethod
def validate_transport_fields(cls, values):
if isinstance(values, dict):
transport: Final = values.get("transport")
if transport == MCPTransport.stdio:
if not values.get("command"):
raise ValueError("command is required for stdio transport")
if not values.get("args"):
raise ValueError("args is required for stdio transport")
# Validate command against allowlist to prevent arbitrary execution
base_command: Final = os.path.basename(values["command"])
if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
raise ValueError(
f"Command '{values['command']}' is not in the allowed commands list "
f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}"
)
elif transport in [MCPTransport.http, MCPTransport.sse]:
if not values.get("url") and not values.get("spec_path"):
raise ValueError("url or spec_path is required for HTTP/SSE transport")
_validate_mcp_transport_fields(values)
return values
@model_validator(mode="before")
@ -1694,23 +1758,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
@model_validator(mode="before")
@classmethod
def validate_transport_fields(cls, values):
if isinstance(values, dict):
transport: Final = values.get("transport")
if transport == MCPTransport.stdio:
if not values.get("command"):
raise ValueError("command is required for stdio transport")
if not values.get("args"):
raise ValueError("args is required for stdio transport")
# Validate command against allowlist to prevent arbitrary execution
base_command: Final = os.path.basename(values["command"])
if base_command not in MCP_STDIO_ALLOWED_COMMANDS:
raise ValueError(
f"Command '{values['command']}' is not in the allowed commands list "
f"for stdio transport. Allowed commands: {sorted(MCP_STDIO_ALLOWED_COMMANDS)}"
)
elif transport in [MCPTransport.http, MCPTransport.sse]:
if not values.get("url") and not values.get("spec_path"):
raise ValueError("url or spec_path is required for HTTP/SSE transport")
_validate_mcp_transport_fields(values)
return values
@model_validator(mode="before")

175
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -13,7 +13,7 @@ import os
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final, TypedDict
from typing import Annotated, Final, NamedTuple, TypedDict
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from pydantic import ValidationError
@ -47,7 +47,11 @@ from litellm.proxy.agent_endpoints.agent_search import (
global_agent_search_index,
search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
UnrestrictedAgentAccess,
accessible_agents,
)
from litellm.proxy.agent_endpoints.identity import reject_legacy_identity
from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore
from litellm.proxy.agent_endpoints.kill_switch import (
@ -63,7 +67,8 @@ from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failur
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
from litellm.proxy.utils import get_custom_url
from litellm.proxy.utils import PrismaClient, get_custom_url
from litellm.repositories.chunked_in import find_many_in
from litellm.types.agents import (
AgentCard,
AgentConfig,
@ -303,6 +308,73 @@ async def _rank_agents_by_query(
assert_never(outcome)
class _AgentDailyActivityScope(NamedTuple):
agent_ids: tuple[str, ...] | None
agent_metadata: Mapping[str, dict[str, object]]
async def _owned_agent_ids(*, user_id: str | None, prisma_client: PrismaClient) -> frozenset[str]:
if user_id is None:
return frozenset()
owned_records: Final = await agents_table(prisma_client).find_many(where={"created_by": user_id})
return frozenset(agent.agent_id for agent in owned_records)
async def _permitted_daily_activity_agent_ids(
*, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient
) -> frozenset[str]:
access: Final = await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict)
if isinstance(access, UnrestrictedAgentAccess):
return await _owned_agent_ids(user_id=user_api_key_dict.user_id, prisma_client=prisma_client)
return access.agent_ids
async def _resolve_daily_activity_agent_ids(
*,
agent_ids: tuple[str, ...] | None,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> tuple[str, ...] | None:
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
if _user_has_admin_view(user_api_key_dict):
return agent_ids
permitted_agent_ids: Final = await _permitted_daily_activity_agent_ids(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
return (
tuple(agent_id for agent_id in agent_ids if agent_id in permitted_agent_ids)
if agent_ids
else tuple(permitted_agent_ids)
)
async def resolve_agent_daily_activity_scope(
*,
agent_ids: tuple[str, ...] | None,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
) -> _AgentDailyActivityScope:
await check_feature_access_for_user(user_api_key_dict, "agents")
resolved_agent_ids: Final = await _resolve_daily_activity_agent_ids(
agent_ids=agent_ids,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
agent_records: Final = (
await agents_table(prisma_client).find_many(where={})
if resolved_agent_ids is None
else await find_many_in(agents_table(prisma_client), "agent_id", resolved_agent_ids)
)
agent_metadata: Final[Mapping[str, dict[str, object]]] = MappingProxyType(
{agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records}
)
return _AgentDailyActivityScope(resolved_agent_ids, agent_metadata)
@router.get(
"/v1/agents",
tags=["[beta] A2A Agents"],
@ -1288,82 +1360,39 @@ async def get_agent_daily_activity(
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
agent_ids_list = agent_ids.split(",") if agent_ids else None
exclude_agent_ids_list: list[str] | None = None
if exclude_agent_ids:
exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None
# Without scoping, an empty `agent_ids` query returned every agent's
# spend/token rows on the proxy. Restrict non-admin callers to the
# agents they're permitted to invoke (or that they created), and
# intersect their explicit `agent_ids` filter with the same allowlist.
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
AgentRequestHandler,
RestrictedAgentAccess,
UnrestrictedAgentAccess,
requested_agent_ids: Final = tuple(agent_ids.split(",")) if agent_ids else None
exclude_agent_ids_list: Final[list[str] | None] = exclude_agent_ids.split(",") if exclude_agent_ids else None
agent_scope: Final = await resolve_agent_daily_activity_scope(
agent_ids=requested_agent_ids,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
where_condition: Final[dict[str, object]] = {}
if not _user_has_admin_view(user_api_key_dict):
permitted_agent_ids: list[str] = []
# An unrestricted caller is not "see everything" for activity scoping. Fall
# back to the agents the caller created so they cannot enumerate other
# tenants' agents.
# Guard against `user_id is None`: a literal None in Prisma
# `where={"created_by": None}` resolves to ``created_by IS NULL``
# and would expose every ownerless agent's rows.
match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
case RestrictedAgentAccess(allowed_agent_ids):
permitted_agent_ids = list(allowed_agent_ids)
case UnrestrictedAgentAccess():
if user_api_key_dict.user_id is not None:
owned_records: Final = await agents_table(prisma_client).find_many(
where={"created_by": user_api_key_dict.user_id}
)
permitted_agent_ids = [a.agent_id for a in owned_records]
if agent_ids_list:
permitted_agent_id_set: Final = set(permitted_agent_ids)
agent_ids_list = [aid for aid in agent_ids_list if aid in permitted_agent_id_set]
else:
agent_ids_list = list(permitted_agent_ids)
# No accessible agents → return an empty page without querying.
if not agent_ids_list:
return SpendAnalyticsPaginatedResponse(
results=[],
metadata=DailySpendMetadata(
total_spend=0.0,
total_prompt_tokens=0,
total_completion_tokens=0,
total_tokens=0,
total_api_requests=0,
total_successful_requests=0,
total_failed_requests=0,
total_cache_read_input_tokens=0,
total_cache_creation_input_tokens=0,
total_compression_saved_tokens=0,
page=page,
total_pages=0,
has_more=False,
),
)
if agent_ids_list:
where_condition["agent_id"] = {"in": list(agent_ids_list)}
agent_records: Final = await agents_table(prisma_client).find_many(where=where_condition)
agent_metadata: Final[Mapping[str, dict[str, object]]] = {
agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records
}
if agent_scope.agent_ids == ():
return SpendAnalyticsPaginatedResponse(
results=[],
metadata=DailySpendMetadata(
total_spend=0.0,
total_prompt_tokens=0,
total_completion_tokens=0,
total_tokens=0,
total_api_requests=0,
total_successful_requests=0,
total_failed_requests=0,
total_cache_read_input_tokens=0,
total_cache_creation_input_tokens=0,
total_compression_saved_tokens=0,
page=page,
total_pages=0,
has_more=False,
),
)
return await get_daily_activity(
prisma_client=prisma_client,
table_name="litellm_dailyagentspend",
entity_id_field="agent_id",
entity_id=agent_ids_list,
entity_metadata_field=agent_metadata,
entity_id=None if agent_scope.agent_ids is None else list(agent_scope.agent_ids),
entity_metadata_field=agent_scope.agent_metadata,
exclude_entity_ids=exclude_agent_ids_list,
start_date=start_date,
end_date=end_date,

9
litellm/proxy/auth/auth_utils.py
View file

@ -1883,6 +1883,15 @@ def _extract_model_candidates_from_request(
llm_router: Router | None = None,
team_id: str | None = None,
) -> list[str]:
if route.rstrip("/") == "/laya/v1/systemone":
from litellm.llms.laya.common_utils import validate_laya_model
try:
laya_request: Final = TypeAdapter(Mapping[str, object]).validate_python(request_data)
laya_model: Final = validate_laya_model(laya_request.get("model"))
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
return _dedupe_model_candidates((f"laya/{laya_model}",))
if route == "/cost/predict-cache":
prediction_models: Final = _cache_prediction_model_candidates(request_data, llm_router, team_id) # pyright: ignore[reportUnknownArgumentType] # the typed reader validates each deployment ID from this legacy payload
return _dedupe_model_candidates(prediction_models)

2
litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py
View file

@ -4,6 +4,7 @@ from typing import Final
from fastapi import APIRouter
from litellm.proxy._experimental.mcp_server.stdio_gate import is_mcp_stdio_enabled
from litellm.proxy.common_utils.html_forms.default_credentials_hint import should_hide_default_credentials_hint
from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import (
UiDiscoveryEndpoints,
@ -41,4 +42,5 @@ async def get_ui_config():
hide_default_credentials_hint=hide_default_credentials_hint,
is_control_plane=is_control_plane,
workers=proxy_config.worker_registry if is_control_plane else [],
mcp_stdio_enabled=is_mcp_stdio_enabled(),
)

18
litellm/proxy/guardrails/guardrail_hooks/azure/base.py
View file

@ -138,6 +138,20 @@ class AzureGuardrailBase:
return chunks
def get_user_prompt(self, messages: list[AllMessageValues]) -> str | None:
"""
Get the last consecutive block of messages from the user.
Example:
messages = [
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm good, thank you!"},
{"role": "user", "content": "What is the weather in Tokyo?"},
]
get_user_prompt(messages) -> "What is the weather in Tokyo?"
"""
return get_last_user_message(messages)
def get_user_prompt_from_request(self, data: Mapping[str, object], call_type: CallTypesLiteral) -> str | None:
if call_type in _RESPONSES_API_CALL_TYPES:
responses_input: Final = data.get("input")
@ -147,6 +161,6 @@ class AzureGuardrailBase:
return get_last_user_message(ResponsesAPIRequestUtils.responses_input_to_chat_messages(validated_input))
messages: Final = data.get("messages")
if not isinstance(messages, list):
if messages is None:
return None
return get_last_user_message(cast(list[AllMessageValues], messages)) # cast-ok: narrowed to list
return self.get_user_prompt(cast(list[AllMessageValues], messages)) # cast-ok: sequence of request messages

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