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

# Conflicts:
#	tests/integration/_support/upstream.py
This commit is contained in:
mateo-berri 2026-10-02 18:31:43 -07:00
commit eee6988b44
762 changed files with 75154 additions and 22372 deletions

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=$!
}

View file

@ -146,6 +146,7 @@ legacy_paths() {
echo tests/unit/proxy/test_proxy_token_counter.py
echo tests/unit/proxy/test_server_root_path.py ;;
proxy-db-proxy-server-core)
echo tests/unit/proxy/test__lazy_features.py
echo tests/unit/proxy/test_aproxy_startup.py
echo tests/unit/proxy/test_proxy_server.py ;;
proxy-db-proxy-utils) echo tests/unit/proxy/test_proxy_utils.py ;;

View file

@ -176,6 +176,7 @@ jobs:
TESTS: ${{ needs.detect.outputs.tests }}
E2E_FIXTURE_MODE: live
E2E_PROVIDER_EDGE_HOST_REACHABLE: '1'
E2E_OWNED_GATEWAY: '1'
COLUMNS: '400'
run: |
umask 077

View file

@ -17,7 +17,7 @@ concurrency:
jobs:
resolve:
runs-on: ubuntu-latest
timeout-minutes: 15
timeout-minutes: 25
strategy:
fail-fast: false
matrix:

View file

@ -89,7 +89,7 @@ jobs:
rust-test:
runs-on: ubuntu-latest
timeout-minutes: 20
timeout-minutes: 30
defaults:
run:
working-directory: litellm-rust

1
.gitignore vendored
View file

@ -58,6 +58,7 @@ litellm/proxy/tests/package-lock.json
ui/litellm-dashboard/.next
ui/litellm-dashboard/node_modules
ui/litellm-dashboard/next-env.d.ts
*.tsbuildinfo
ui/litellm-dashboard/package.json
ui/litellm-dashboard/package-lock.json
helm/litellm-helm/*.tgz

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

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

View file

@ -4,13 +4,28 @@ 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:
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
```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 > Investigations**, click **Connect worker**, choose an analysis model and monthly limit, then **Get install command**. Use **Advanced options** to select an existing virtual key or change the proxy URL if the server running Docker needs a different network address. Copy the command and run it on your server. The dashboard shows **Worker connected** when the container checks in
The command already contains the compatible worker image and one worker token. The selected virtual key stays on the proxy; its secret is never sent to the worker. No source checkout, environment file, or second LiteLLM deployment is needed. Keep the command private because it includes the token. The LiteLLM release provides the dashboard and APIs; the container only runs background analysis
The dashboard and Compose file pin a verified worker image by digest. The image uses Linux amd64, and the generated command selects that platform. Worker image releases are independent of proxy releases: update the pinned image when changing their API contract. CI also publishes immutable commit tags for reproducible builds
The dashboard and Compose file pin a verified worker image by digest. The image uses Linux amd64, and the generated command selects that platform. CI also publishes immutable `:sha-<commit>` tags for successful worker builds on `main`. Keep the worker image compatible with your gateway version
After upgrading the gateway, update the worker image and redeploy it while keeping its proxy URL and token. Existing containers do not update automatically. If an investigation reports a worker compatibility error, update the image before retrying
For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL` and `LENS_WORKER_TOKEN` in an environment file. Its default image is already selected:

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:44f0597c7583dcfef999ece9a8bc02cfeb9f0f5167a1221cee3bd10b1b79271b}
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}

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"

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

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"]

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:

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

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==",

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

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:

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".

View file

@ -0,0 +1,17 @@
CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterDailySpend" (
"date" TEXT NOT NULL,
"api_key" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"router_name" TEXT NOT NULL,
"router_type" TEXT NOT NULL,
"turns" INTEGER NOT NULL DEFAULT 0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_turns" INTEGER NOT NULL DEFAULT 0,
"savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0,
"classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
"classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0,
CONSTRAINT "LiteLLM_AutoRouterDailySpend_pkey" PRIMARY KEY ("date", "api_key", "user_id", "router_name", "router_type")
);

View file

@ -0,0 +1,3 @@
CREATE INDEX IF NOT EXISTS "LiteLLM_LensWorker_active_scope_idx"
ON "LiteLLM_LensWorker" USING GIN ((data->'scope') jsonb_path_ops)
WHERE data @> '{"revoked": false}'::jsonb;

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}")

View file

@ -1744,6 +1744,27 @@ model LiteLLM_AutoRouterUserSession {
@@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn")
}
// Auto-routed requests per UTC request day and router: the selected-day money behind the
// auto-router usage view. Written in the same statement as the session rollup, so a day row
// and its session row never disagree; corrected in the same transaction as late baselines.
model LiteLLM_AutoRouterDailySpend {
date String
api_key String
user_id String
router_name String
router_type String
turns Int @default(0)
spend Float @default(0)
saved_spend Float @default(0)
savings_estimated_turns Int @default(0)
savings_estimated_actual_spend Float @default(0)
savings_estimated_saved_spend Float @default(0)
classifier_cost Float @default(0)
classifier_cost_recorded_turns Int @default(0)
@@id([date, api_key, user_id, router_name, router_type])
}
// Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in
// either direction. forward duplicates the requests the keys did not route through the
// router through it, answering whether they should adopt it; reverse duplicates the

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:

View file

@ -1,6 +1,6 @@
[project]
name = "litellm-proxy-extras"
version = "0.4.103"
version = "0.4.105"
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.105"
version_files = [
"pyproject.toml:^version",
"../pyproject.toml:litellm-proxy-extras==",

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"
@ -4089,6 +4156,7 @@ dependencies = [
"litellm-storage-clickhouse",
"litellm-token-counter",
"litellm-traces",
"litellm-traces-clickhouse",
"litellm-tracing",
"prost",
"pyo3",
@ -4302,6 +4370,7 @@ dependencies = [
"thiserror 2.0.19",
"tokio",
"url",
"wiremock",
]
[[package]]
@ -4384,22 +4453,40 @@ dependencies = [
name = "litellm-traces"
version = "0.1.0"
dependencies = [
"base64 0.22.1",
"criterion",
"flate2",
"litellm-http",
"litellm-storage-clickhouse",
"indexmap 2.14.0",
"opentelemetry-proto",
"prost",
"rstest",
"serde",
"serde_json",
"strum",
"thiserror 2.0.19",
]
[[package]]
name = "litellm-traces-clickhouse"
version = "0.1.0"
dependencies = [
"askama",
"flate2",
"futures-util",
"hmac 0.12.1",
"litellm-http",
"litellm-migrate",
"litellm-storage-clickhouse",
"litellm-traces",
"moka",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
"strum",
"testcontainers-modules",
"thiserror 2.0.19",
"time",
"tokio",
"url",
"wiremock",
]

View file

@ -13,7 +13,10 @@ litellm-config = { path = "crates/config" }
litellm-router = { path = "crates/router" }
litellm-tracing = { path = "crates/tracing" }
litellm-traces = { path = "crates/traces" }
litellm-traces-clickhouse = { path = "crates/traces-clickhouse" }
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 +66,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 +97,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"] }

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)]

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()

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();

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

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 },
}

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 { .. })));
}
}

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

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

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,
}

View file

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

View file

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

View file

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

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);
}

View file

@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
fancy-regex.workspace = true
litellm-tracing.workspace = true
litellm-traces.workspace = true
litellm-traces-clickhouse.workspace = true
litellm-storage-clickhouse.workspace = true
litellm-host.workspace = true
bytes.workspace = true

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",

View file

@ -2,8 +2,8 @@ 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::{QueryScope, ReadQuery, Shared};
use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders};
use prost::Message;
use pyo3::{
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
@ -29,57 +29,103 @@ pub fn trace_encode_error<'py>(py: Python<'py>, message: &str) -> Bound<'py, PyB
}
fn map_error(error: Error) -> PyErr {
map_error_ref(&error)
}
fn map_error_ref(error: &Error) -> PyErr {
use litellm_storage_clickhouse::Error as StorageError;
match error {
Error::InvalidRow
| Error::InvalidTable
| Error::InvalidSchema
| Error::EmptySql
| Error::InvalidQuery => PyValueError::new_err(error.to_string()),
| Error::InvalidQuery
| Error::InvalidParameters
| Error::InvalidScope => PyValueError::new_err(error.to_string()),
Error::InsertTooLarge => PyOverflowError::new_err(error.to_string()),
Error::InvalidUrl
| Error::QueryFailed(_)
| Error::InsertFailed(_)
| Error::SchemaFailed(_)
| Error::ResponseTooLarge
| Error::InvalidResponse
| Error::Transport => PyRuntimeError::new_err(error.to_string()),
Error::SchemaFailed(_)
| Error::SchemaTransport
| Error::MissingSecret
| Error::Busy
| Error::ProvisionFailed(_)
| Error::ProvisionTransport
| Error::InvalidResponse => PyRuntimeError::new_err(error.to_string()),
Error::Cached(source) => map_error_ref(source),
Error::Storage(source) => match source {
StorageError::InvalidRow
| StorageError::InvalidTable
| StorageError::InvalidSchema
| StorageError::EmptySql
| StorageError::InvalidParameters
| StorageError::InvalidQuery => PyValueError::new_err(error.to_string()),
StorageError::InsertTooLarge => PyOverflowError::new_err(error.to_string()),
StorageError::InvalidUrl
| StorageError::QueryFailed(_)
| StorageError::InsertFailed(_)
| StorageError::SchemaFailed(_)
| StorageError::ResponseTooLarge
| StorageError::InvalidResponse
| StorageError::Transport => PyRuntimeError::new_err(error.to_string()),
},
}
}
fn map_sql_error(error: Error) -> PyErr {
match error {
Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => {
PyValueError::new_err(error.to_string())
}
error => map_error(error),
}
}
#[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(
litellm_traces_clickhouse::ensure_schema(
&client,
&connection,
&database,
trace_retention_days,
spend_log_retention_days,
retention_days,
)
.await
},
@ -91,42 +137,69 @@ impl NativeTraceStorage {
&self,
py: Python<'py>,
table: &str,
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces::InsertRow>,
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces_clickhouse::InsertRow>,
) -> 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 {
litellm_traces::insert_shared_rows(&client, &connection, &database, table, rows)
.await
litellm_traces_clickhouse::insert_shared_rows(
&client,
&connection,
&database,
table,
rows,
)
.await
},
map_error,
)
}
fn lens_query<'py>(
fn query_sql<'py>(
&self,
py: Python<'py>,
name: &str,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap<
String,
Parameter,
>,
sql: String,
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
secret: String,
) -> 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")
})?;
if sql.trim().is_empty() {
return Err(map_error(
litellm_storage_clickhouse::Error::EmptySql.into(),
));
}
let readers = self.query_readers.clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
litellm_traces::execute_read(&client, &connection, query.sql(), &parameters).await
let _permit = readers.acquire()?;
let connection = readers.connection(&client, &scope, &secret).await?;
litellm_traces_clickhouse::query_sql(&client, &connection, &sql).await
},
map_error,
map_sql_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_clickhouse::query_help(&client, &connection).await
},
map_sql_error,
)
}
@ -139,15 +212,20 @@ impl NativeTraceStorage {
Parameter,
>,
) -> 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 query =
ReadQuery::parse(query).map_err(|error| PyValueError::new_err(error.to_string()))?;
let connection = self.config.storage().reader().clone();
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
crate::execution::run_async(
py,
async move {
litellm_traces::execute_named_read(&client, &connection, query, &parameters).await
litellm_traces_clickhouse::execute_named_read(
&client,
&connection,
query,
&parameters,
)
.await
},
map_error,
)
@ -163,13 +241,15 @@ pub fn trace_decode_otlp<'py>(
let spans = py
.detach(|| litellm_traces::decode_otlp(body, content_type))
.map_err(|error| match error {
litellm_traces::DecodeError::TooLarge => PyOverflowError::new_err(error.to_string()),
litellm_traces::Error::TooLarge => PyOverflowError::new_err(error.to_string()),
_ => PyValueError::new_err(error.to_string()),
})?;
spans_to_py(py, &spans).map(Bound::into_any)
}
fn insert_rows_from_py(value: &Bound<'_, PyAny>) -> PyResult<Vec<litellm_traces::InsertRow>> {
fn insert_rows_from_py(
value: &Bound<'_, PyAny>,
) -> PyResult<Vec<litellm_traces_clickhouse::InsertRow>> {
let mut resources = FromPythonCache::default();
value
.try_iter()?
@ -236,6 +316,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)
@ -246,6 +334,54 @@ mod tests {
use super::*;
use rstest::rstest;
#[rstest]
#[case::row(Error::InvalidRow, "ValueError")]
#[case::insert_budget(Error::InsertTooLarge, "OverflowError")]
#[case::scope(Error::InvalidScope, "ValueError")]
#[case::schema(Error::SchemaFailed(503), "RuntimeError")]
#[case::reader(Error::MissingSecret, "RuntimeError")]
#[case::storage(
Error::Storage(litellm_storage_clickhouse::Error::InvalidUrl),
"RuntimeError"
)]
#[case::cached_scope(Error::Cached(std::sync::Arc::new(Error::InvalidScope)), "ValueError")]
fn trace_failures_preserve_public_exception_types(
#[case] error: Error,
#[case] exception_name: &str,
) {
Python::initialize();
Python::attach(|py| {
let message = error.to_string();
let exception = map_error(error);
assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
assert_eq!(
exception.value(py).str().unwrap().to_str().unwrap(),
message
);
});
}
#[rstest]
#[case::invalid_sql(400, "ValueError")]
#[case::missing_table(404, "ValueError")]
#[case::unavailable(503, "RuntimeError")]
fn wrapped_query_status_preserves_public_exception_type(
#[case] status: u16,
#[case] exception_name: &str,
) {
Python::initialize();
Python::attach(|py| {
let error = Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(status));
let message = error.to_string();
let exception = map_sql_error(error);
assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
assert_eq!(
exception.value(py).str().unwrap().to_str().unwrap(),
message
);
});
}
#[rstest]
fn insert_projection_preserves_identity_without_merging_equal_resources() {
Python::initialize();
@ -288,3 +424,9 @@ mod tests {
});
}
}
#[pyfunction]
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
litellm_host_python::Pythonized(litellm_traces_clickhouse::NORMALIZED_FIELD_DEFINITIONS)
.into_pyobject(py)
}

View file

@ -0,0 +1,5 @@
# ClickHouse storage
`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-clickhouse` supplies those rules and uses this storage for both trace rows and spend rows

View file

@ -18,3 +18,4 @@ url.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
rstest.workspace = true
tokio.workspace = true
wiremock.workspace = true

View file

@ -1,5 +0,0 @@
# 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
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

View file

@ -10,6 +10,8 @@ pub enum Error {
InvalidSchema,
#[error("SQL query must not be empty")]
EmptySql,
#[error("invalid ClickHouse query parameters")]
InvalidParameters,
#[error("unknown ClickHouse read query")]
InvalidQuery,
#[error("ClickHouse query failed with HTTP status {0}")]

View file

@ -4,7 +4,7 @@ mod read;
pub use error::Error;
pub use insert::{insert_compressed_rows, insert_encoded_rows};
pub use read::{Parameter, execute_read};
pub use read::{Parameter, Query, READ_LIMITS, ReadLimits, execute_read, fetch, fetch_json};
use url::Url;
#[derive(Clone)]
@ -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
}
}

View file

@ -1,17 +1,30 @@
use std::{collections::BTreeMap, time::Duration};
use litellm_http::Client;
use serde::Deserialize;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use crate::{Connection, Error};
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ReadLimits {
pub result_rows: u64,
pub response_bytes: usize,
pub execution_seconds: u64,
}
#[derive(Debug, Deserialize)]
pub const READ_LIMITS: ReadLimits = ReadLimits {
result_rows: 1000,
response_bytes: 4 * 1024 * 1024,
execution_seconds: 10,
};
#[derive(Debug, Deserialize, Serialize)]
#[serde(untagged)]
pub enum Parameter {
Text(String),
Integer(i64),
Unsigned(u64),
Float(f64),
Strings(Vec<String>),
}
@ -20,6 +33,8 @@ impl Parameter {
match self {
Self::Text(value) => escaped(value),
Self::Integer(value) => value.to_string(),
Self::Unsigned(value) => value.to_string(),
Self::Float(value) => value.to_string(),
Self::Strings(values) => format!(
"[{}]",
values
@ -74,9 +89,12 @@ pub async fn execute_read(
.clear()
.extend_pairs(existing_pairs)
.append_pair("readonly", "1")
.append_pair("max_result_rows", "1000")
.append_pair("max_result_rows", &READ_LIMITS.result_rows.to_string())
.append_pair("result_overflow_mode", "throw")
.append_pair("max_execution_time", "10")
.append_pair(
"max_execution_time",
&READ_LIMITS.execution_seconds.to_string(),
)
.append_pair("wait_end_of_query", "1")
.append_pair("default_format", "JSON");
@ -97,7 +115,7 @@ pub async fn execute_read(
let mut body = Vec::new();
while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? {
if body.len() + chunk.len() > MAX_RESPONSE_BYTES {
if body.len() + chunk.len() > READ_LIMITS.response_bytes {
return Err(Error::ResponseTooLarge);
}
body.extend_from_slice(&chunk);
@ -111,3 +129,45 @@ pub async fn execute_read(
}
String::from_utf8(body).map_err(|_| Error::InvalidResponse)
}
pub trait Query {
type Params: Serialize;
type Row: DeserializeOwned;
const SQL: &'static str;
}
#[derive(Deserialize)]
struct Rows<T> {
data: Vec<T>,
}
fn parameters<T: Serialize>(params: &T) -> Result<BTreeMap<String, Parameter>, Error> {
let value = serde_json::to_value(params).map_err(|_| Error::InvalidParameters)?;
serde_json::from_value(value).map_err(|_| Error::InvalidParameters)
}
pub async fn fetch<Q: Query>(
client: &Client,
connection: &Connection,
params: &Q::Params,
) -> Result<Vec<Q::Row>, Error> {
let body = execute_read(client, connection, Q::SQL, &parameters(params)?).await?;
decode_rows::<Q::Row>(&body)
}
pub async fn fetch_json<Q: Query>(
client: &Client,
connection: &Connection,
params: &Q::Params,
) -> Result<String, Error> {
let body = execute_read(client, connection, Q::SQL, &parameters(params)?).await?;
decode_rows::<Q::Row>(&body)?;
Ok(body)
}
fn decode_rows<T: DeserializeOwned>(body: &str) -> Result<Vec<T>, Error> {
serde_json::from_str::<Rows<T>>(body)
.map(|rows| rows.data)
.map_err(|_| Error::InvalidResponse)
}

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());
}

View file

@ -1,7 +1,7 @@
use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_storage_clickhouse::{Connection, Error, execute_read, insert_encoded_rows};
use litellm_storage_clickhouse::{Connection, Error, Query, execute_read, insert_encoded_rows};
use rstest::rstest;
#[rstest]
@ -32,3 +32,82 @@ async fn read_rejects_empty_sql() {
Err(Error::EmptySql)
));
}
#[derive(serde::Serialize)]
struct QueryParams {
signed: i64,
unsigned: u64,
float: f64,
text: String,
strings: Vec<String>,
}
#[derive(Debug, serde::Deserialize, PartialEq)]
struct QueryRow {
answer: String,
}
struct TypedQuery;
impl litellm_storage_clickhouse::Query for TypedQuery {
type Params = QueryParams;
type Row = QueryRow;
const SQL: &'static str = "SELECT typed_parameters";
}
#[rstest]
#[case::valid(
r#"{"meta":[],"data":[{"answer":"ok"}],"rows":1,"statistics":{"elapsed":0.1}}"#,
true
)]
#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, false)]
#[case::missing_column(r#"{"data":[{}]}"#, false)]
#[case::exception(r#"{"data":[],"exception":"failed"}"#, false)]
#[tokio::test]
async fn typed_fetch_encodes_parameters_and_validates_rows(
#[case] body: &str,
#[case] valid: bool,
) {
use litellm_storage_clickhouse::{fetch, fetch_json};
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_string, query_param},
};
let server = MockServer::start().await;
Mock::given(body_string(TypedQuery::SQL))
.and(query_param("param_signed", i64::MIN.to_string()))
.and(query_param("param_unsigned", u64::MAX.to_string()))
.and(query_param("param_float", "12.5"))
.and(query_param("param_text", "line\\nbreak"))
.and(query_param("param_strings", "['a\\'b','雪']"))
.and(query_param("readonly", "1"))
.and(query_param("max_result_rows", "1000"))
.respond_with(ResponseTemplate::new(200).set_body_string(body))
.expect(2)
.mount(&server)
.await;
let client = Client::no_redirect_for_test();
let connection = Connection::parse(&server.uri()).unwrap();
let params = QueryParams {
signed: i64::MIN,
unsigned: u64::MAX,
float: 12.5,
text: "line\nbreak".into(),
strings: vec!["a'b".into(), "雪".into()],
};
let rows = fetch::<TypedQuery>(&client, &connection, &params).await;
let envelope = fetch_json::<TypedQuery>(&client, &connection, &params).await;
if valid {
assert_eq!(
rows.unwrap(),
vec![QueryRow {
answer: "ok".into()
}]
);
assert_eq!(envelope.unwrap(), body);
} else {
assert!(matches!(rows, Err(Error::InvalidResponse)));
assert!(matches!(envelope, Err(Error::InvalidResponse)));
}
}

View file

@ -0,0 +1,7 @@
- Own trace schema, row encoding, SQL query adapters and reader provisioning; consume domain types from `litellm-traces`
- Keep generic ClickHouse connections and HTTP execution in `litellm-storage-clickhouse`; keep PyO3 conversion in `python-bridge`
- Keep schema definitions only in `migrations/NNNN_description.sql`, embedded by `litellm_migrate::migrate!`
- Require typed query parameters and SELECT-only readers with server-side limits and tenant isolation
- Bound insert time and encoded bytes; preserve shared values and explicit retry deduplication
- Test storage behavior through the public API against ClickHouse
- Expose one top-level `Error` enum in `src/error.rs`; own trace failures and wrap storage errors with `#[from]` or `#[source]`

View file

@ -0,0 +1,31 @@
[package]
name = "litellm-traces-clickhouse"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
askama.workspace = true
flate2.workspace = true
futures-util.workspace = true
hmac = "0.12.1"
litellm-http.workspace = true
litellm-migrate.workspace = true
litellm-storage-clickhouse.workspace = true
litellm-traces.workspace = true
moka.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
strum.workspace = true
thiserror.workspace = true
time = { workspace = true, features = ["formatting"] }
tokio.workspace = true
url.workspace = true
[dev-dependencies]
litellm-http = { workspace = true, features = ["test-support"] }
rstest.workspace = true
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
wiremock.workspace = true

View file

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

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

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

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

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

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

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

View file

@ -0,0 +1,2 @@
ALTER TABLE {database}.otel_traces
ADD COLUMN IF NOT EXISTS UserId String DEFAULT ''

View file

@ -0,0 +1,3 @@
ALTER TABLE {database}.agent_traces_by_key
ADD COLUMN IF NOT EXISTS UserIds SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [],
ADD COLUMN IF NOT EXISTS IdentifiedLlmCount SimpleAggregateFunction(sum, UInt64) DEFAULT 0

View file

@ -0,0 +1,22 @@
ALTER TABLE {database}.agent_traces_by_key_mv MODIFY QUERY
SELECT
TeamId, ApiKeyHash, TraceId, groupUniqArray(UserId) AS UserIds,
min(Timestamp) AS StartTs,
max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs,
any(ServiceName) AS ServiceName,
anyLastIf(toNullable(SpanName), ParentSpanId = '') AS RootName,
anyLastIf(toNullable(InputPreview), ParentSpanId = '') AS RootInput,
anyLastIf(toNullable(StatusCode), ParentSpanId = '') AS RootStatus,
count() AS SpanCount,
countIf(ObservationType = 'agent') AS AgentCount,
countIf(ObservationType = 'llm') AS LlmCount,
countIf(ObservationType = 'llm' AND LiteLLMRequestId != '') AS IdentifiedLlmCount,
countIf(ObservationType = 'tool') AS ToolCount,
countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount,
sum(InputTokens) AS InputTokens,
sum(OutputTokens) AS OutputTokens,
groupUniqArrayIf(toString(Model), Model != '') AS Models,
groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames,
groupArrayIf(LiteLLMRequestId, ObservationType = 'llm' OR LiteLLMRequestId != '') AS RequestIds
FROM {database}.otel_traces
GROUP BY TeamId, ApiKeyHash, TraceId

View file

@ -0,0 +1 @@
ALTER TABLE {database}.otel_traces ADD COLUMN IF NOT EXISTS Framework LowCardinality(String) AFTER AgentName

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

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

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

View file

@ -0,0 +1,49 @@
WITH page AS (
SELECT TraceId AS trace_id,
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
if(length(groupUniqArrayArray(UserIds)) = 1, arrayElement(groupUniqArrayArray(UserIds), 1), '') AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash,
ifNull(any(RootName), '') AS name, any(ServiceName) AS service,
ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status,
toUnixTimestamp64Milli(min(StartTs)) AS start_ms,
min(StartTs) AS trace_start, max(EndTs) AS trace_end,
dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms,
sum(SpanCount) AS span_count,
sum(AgentCount) AS agent_invocations,
sum(LlmCount) AS llm_calls, sum(ToolCount) AS tool_calls,
sum(InputTokens) AS input_tokens, sum(OutputTokens) AS output_tokens,
groupUniqArrayArray(Models) AS models, sum(ErrorCount) AS error_count,
arrayDistinct(if(sum(IdentifiedLlmCount) != sum(LlmCount),
arrayConcat(groupArrayArray(RequestIds), ['']),
groupArrayArray(RequestIds))) AS request_ids
FROM agent_traces_by_key
WHERE ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND UserIds = [{user_id:String}])
OR has({team_ids:Array(String)}, TeamId)
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
GROUP BY TeamId, ApiKeyHash, TraceId
HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64})
AND ({cursor_ms:Int64} = 0 OR (toUnixTimestamp64Milli(min(StartTs)), trace_ref)
< ({cursor_ms:Int64}, {cursor_trace_id:String}))
ORDER BY start_ms DESC, trace_ref DESC
LIMIT {limit:UInt32}
)
SELECT page.* EXCEPT (trace_start, trace_end),
identities.agent_names AS agent_names, identities.agent_count AS agent_count,
identities.frameworks AS frameworks
FROM page
LEFT JOIN (
SELECT TeamId, ApiKeyHash, TraceId,
arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names,
arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks,
uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count
FROM otel_traces
WHERE Timestamp >= (SELECT min(trace_start) FROM page)
AND Timestamp <= (SELECT max(trace_end) FROM page)
AND TraceId IN (SELECT trace_id FROM page)
AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page)
GROUP BY TeamId, ApiKeyHash, TraceId
) AS identities
ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash
AND page.trace_id = identities.TraceId
ORDER BY page.start_ms DESC, page.trace_ref DESC

View file

@ -0,0 +1,26 @@
SELECT o.SpanId AS span_id, o.Input AS input,
if(o.Output = '' AND o.ObservationType = 'agent', answer.output, o.Output) AS output,
o.SpanAttributes AS attributes
FROM otel_traces AS o
LEFT JOIN (
SELECT TeamId, ApiKeyHash, ParentSpanId AS parent_span_id, argMax(Output, Timestamp) AS output
FROM otel_traces
WHERE TraceId = {trace_id:String} AND ParentSpanId = {span_id:String}
AND ObservationType = 'llm' AND Output != ''
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND UserId = {user_id:String})
OR has({team_ids:Array(String)}, TeamId)
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
GROUP BY TeamId, ApiKeyHash, ParentSpanId
) AS answer ON answer.parent_span_id = o.SpanId
AND answer.TeamId = o.TeamId AND answer.ApiKeyHash = o.ApiKeyHash
WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String}
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
OR has({team_ids:Array(String)}, o.TeamId)
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
LIMIT 1

View file

@ -4,8 +4,10 @@ SELECT SpanId AS span_id,
hex(SHA256(StatusMessage)) AS version
FROM otel_traces
WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String}
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND UserId = {user_id:String})
OR has({team_ids:Array(String)}, TeamId)
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
AND ({error_version:String} = '' OR hex(SHA256(StatusMessage)) = {error_version:String})

View file

@ -0,0 +1,11 @@
SELECT request_id, response_id, team_id, api_key, user, spend,
toUnixTimestamp64Milli(start_time) AS start_ms
FROM spend_logs FINAL
WHERE response_id IN {response_ids:Array(String)}
AND start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND user = {user_id:String})
OR has({team_ids:Array(String)}, team_id)
OR ({api_key_hash:String} != '' AND api_key = {api_key_hash:String}))
ORDER BY start_time DESC

View file

@ -0,0 +1,9 @@
SELECT hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref
FROM otel_traces
WHERE TraceId = {trace_id:String}
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND UserId = {user_id:String})
OR has({team_ids:Array(String)}, TeamId)
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
GROUP BY TeamId, ApiKeyHash, TraceId
LIMIT 2

View file

@ -1,16 +1,19 @@
SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
o.ObservationType AS type, o.AgentName AS agent, o.StatusCode AS status,
o.ObservationType AS type, o.AgentName AS agent,
o.Framework AS framework, o.StatusCode AS status,
substringUTF8(o.StatusMessage, 1, 128) AS status_message,
lengthUTF8(o.StatusMessage) > 128 AS error_truncated,
toUnixTimestamp64Nano(o.Timestamp) AS start_ns, o.Duration AS duration_ns,
o.ServiceName AS service, o.InputPreview AS input_preview, o.Model AS model,
o.InputTokens AS input_tokens, o.OutputTokens AS output_tokens,
o.LiteLLMRequestId AS litellm_request_id,
o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
FROM otel_traces AS o
WHERE o.TraceId = {trace_id:String}
AND (empty({team_ids:Array(String)}) OR o.TeamId IN {team_ids:Array(String)})
AND ({api_key_hash:String} = '' OR o.ApiKeyHash = {api_key_hash:String})
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
OR has({team_ids:Array(String)}, o.TeamId)
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
AND ({trace_ref:String} = '' OR
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage

View file

@ -0,0 +1,26 @@
use crate::Error;
use litellm_storage_clickhouse::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> {
super::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
}
}

View file

@ -0,0 +1,37 @@
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("invalid ClickHouse insert row")]
InvalidRow,
#[error("invalid ClickHouse insert table")]
InvalidTable,
#[error("database must be a nonempty SQL identifier and retention must be positive")]
InvalidSchema,
#[error("unknown ClickHouse read query")]
InvalidQuery,
#[error("invalid ClickHouse query parameters")]
InvalidParameters,
#[error("ClickHouse returned an invalid or failed JSON query response")]
InvalidResponse,
#[error("ClickHouse insert exceeds the encoded size limit")]
InsertTooLarge,
#[error("ClickHouse schema setup failed with HTTP status {0}")]
SchemaFailed(u16),
#[error("ClickHouse schema setup transport failed")]
SchemaTransport,
#[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<Error>),
}

View file

@ -12,7 +12,8 @@ use serde_json::Value;
use sha2::{Digest, Sha256};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
use crate::{Connection, Error, Shared};
use super::{Connection, Error};
use litellm_traces::Shared;
const MAX_INSERT_BYTES: usize = 64 * 1024 * 1024;
@ -71,6 +72,7 @@ pub async fn insert_shared_rows(
body,
)
.await
.map_err(Error::from)
}
fn shared_rows(rows: Vec<BTreeMap<String, Value>>) -> Vec<InsertRow> {
@ -226,8 +228,8 @@ mod tests {
use rstest::rstest;
use serde_json::json;
use super::Error;
use super::{shared_rows, write_rows};
use crate::Error;
#[rstest]
fn encoded_limit_counts_utf8_bytes_across_rows() {

View file

@ -0,0 +1,21 @@
mod config;
mod error;
mod insert;
pub mod query;
mod query_access;
mod schema;
mod sql;
mod table;
pub use config::Config;
pub use error::Error;
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
pub use litellm_storage_clickhouse::{Connection, Parameter};
pub use litellm_traces::{QueryScope, ReadQuery};
pub use query::{QueryHelp, execute_read, query_help, query_sql};
pub use query_access::QueryReaders;
pub use schema::{
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements,
};
pub use sql::execute_named_read;
pub use table::TraceTable;

View file

@ -0,0 +1,494 @@
use std::collections::{BTreeMap, BTreeSet};
use crate::TraceTable;
use futures_util::{
StreamExt,
stream::{self, TryStreamExt},
};
use litellm_http::Client;
use serde::{Deserialize, Serialize, Serializer};
use serde_json::Value;
use strum::IntoEnumIterator;
use super::{
Connection, Error, NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, Parameter,
query_access::READER_LIMITS,
};
mod guide;
pub mod lens;
pub mod named;
mod number;
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(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, strum::Display)]
#[serde(rename_all = "lowercase")]
#[strum(serialize_all = "lowercase")]
enum JsonKind {
Array,
Boolean,
Integer,
Null,
Number,
Object,
String,
}
impl JsonKind {
fn of(value: &Value) -> Self {
match value {
Value::Null => Self::Null,
Value::Bool(_) => Self::Boolean,
Value::Number(number) if number.is_i64() || number.is_u64() => Self::Integer,
Value::Number(_) => Self::Number,
Value::String(_) => Self::String,
Value::Array(_) => Self::Array,
Value::Object(_) => Self::Object,
}
}
}
#[derive(Clone, Copy, Debug, Serialize, strum::Display)]
enum MapValueType {
String,
}
#[derive(Serialize)]
struct MetadataField {
path: Vec<PathPart>,
types: BTreeSet<JsonKind>,
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: TraceTable,
columns: Vec<ColumnSchema>,
}
trait Unobserved {
fn unobserved() -> Self;
}
enum Discovery<T> {
Observed(T),
Unavailable(String),
}
impl<T: Serialize + Unobserved> Serialize for Discovery<T> {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
#[derive(Serialize)]
struct Unavailable<'a, T> {
#[serde(flatten)]
sample: T,
error: &'a str,
}
match self {
Self::Observed(sample) => sample.serialize(serializer),
Self::Unavailable(error) => Unavailable {
sample: T::unobserved(),
error,
}
.serialize(serializer),
}
}
}
#[derive(Serialize)]
struct MetadataSample {
fields: Vec<MetadataField>,
sampled_rows: usize,
invalid_json_rows: usize,
truncated: bool,
}
impl Unobserved for MetadataSample {
fn unobserved() -> Self {
Self {
fields: Vec::new(),
sampled_rows: 0,
invalid_json_rows: 0,
truncated: true,
}
}
}
#[derive(Serialize)]
struct MetadataCatalog {
table: TraceTable,
column: &'static str,
#[serde(flatten)]
discovery: Discovery<MetadataSample>,
sample_sql: &'static str,
scope: &'static str,
}
#[derive(Serialize)]
struct AttributeField {
key: String,
#[serde(rename = "type")]
kind: MapValueType,
expression: String,
}
#[derive(Serialize)]
struct AttributeSample {
fields: Vec<AttributeField>,
truncated: bool,
}
impl Unobserved for AttributeSample {
fn unobserved() -> Self {
Self {
fields: Vec::new(),
truncated: true,
}
}
}
#[derive(Serialize)]
struct AttributeCatalog {
table: TraceTable,
column: &'static str,
#[serde(flatten)]
discovery: Discovery<AttributeSample>,
discovery_sql: String,
scope: &'static str,
}
#[derive(Serialize)]
struct NormalizedField {
table: TraceTable,
name: &'static str,
column: &'static str,
#[serde(rename = "type")]
kind: &'static str,
meaning: &'static str,
}
impl From<&NormalizedFieldDefinition> for NormalizedField {
fn from(field: &NormalizedFieldDefinition) -> Self {
Self {
table: TraceTable::OtelTraces,
name: field.name,
column: field.clickhouse_column,
kind: field.clickhouse_type,
meaning: field.meaning,
}
}
}
#[derive(Serialize)]
struct Relationship {
left: &'static str,
right: &'static str,
additional_predicates: &'static str,
meaning: &'static str,
}
const RELATIONSHIPS: [Relationship; 1] = [Relationship {
left: "otel_traces.LiteLLMRequestId",
right: "spend_logs.response_id",
additional_predicates: "otel_traces.TeamId = spend_logs.team_id AND (otel_traces.TeamId != '' OR (otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' 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",
}];
#[derive(Serialize)]
pub struct QueryHelp {
dialect: &'static str,
access: &'static str,
response: &'static str,
tables: Vec<TableSchema>,
normalized_fields: Vec<NormalizedField>,
metadata: MetadataCatalog,
attributes: Vec<AttributeCatalog>,
relationships: &'static [Relationship],
examples: [guide::Example; 5],
gotchas: [String; 11],
guide: String,
}
pub async fn execute_read(
client: &Client,
connection: &Connection,
sql: &str,
parameters: &BTreeMap<String, Parameter>,
) -> Result<String, Error> {
litellm_storage_clickhouse::execute_read(client, connection, sql, parameters)
.await
.map_err(Error::from)
}
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<JsonKind>>,
) -> bool {
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
return true;
}
if !path.is_empty() {
fields
.entry(path.clone())
.or_default()
.insert(JsonKind::of(value));
}
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_sample(sample: &[MetadataRow]) -> MetadataSample {
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();
MetadataSample {
fields,
sampled_rows: sample.len().min(SAMPLE_ROWS),
invalid_json_rows: invalid_rows,
truncated: limited,
}
}
pub async fn query_help(client: &Client, connection: &Connection) -> Result<QueryHelp, Error> {
let tables = stream::iter(TraceTable::iter())
.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 = MetadataCatalog {
table: TraceTable::SpendLogs,
column: "metadata",
discovery: match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
Ok(sample) => Discovery::Observed(metadata_sample(&sample)),
Err(error) => Discovery::Unavailable(error.to_string()),
},
sample_sql: METADATA_SQL,
scope: METADATA_SCOPE,
};
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 discovery = match rows::<AttributeRow>(client, connection, &sql).await {
Ok(keys) => Discovery::Observed(AttributeSample {
truncated: keys.len() > MAX_FIELDS,
fields: keys
.into_iter()
.take(MAX_FIELDS)
.map(|row| AttributeField {
expression: format!("{column}[{}]", literal(&row.key)),
key: row.key,
kind: MapValueType::String,
})
.collect(),
}),
Err(error) => Discovery::Unavailable(error.to_string()),
};
AttributeCatalog {
table: TraceTable::OtelTraces,
column,
discovery,
discovery_sql: sql,
scope: ATTRIBUTE_SCOPE,
}
})
.collect::<Vec<_>>()
.await;
let guide = guide::QueryGuide {
tables: &tables,
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
metadata: &metadata,
attributes: &attributes,
limits: &READER_LIMITS,
};
Ok(QueryHelp {
dialect: "ClickHouse SQL",
access: "Request-log visibility enforced by ClickHouse row policies; proxy admins see all rows, users see their own rows and permitted teams, and callers without user identity see their own key rows",
response: "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
examples: guide.examples()?,
gotchas: guide.gotchas()?,
guide: guide::render(&guide)?,
normalized_fields: NORMALIZED_FIELD_DEFINITIONS
.iter()
.map(NormalizedField::from)
.collect(),
relationships: &RELATIONSHIPS,
tables,
metadata,
attributes,
})
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use serde_json::json;
#[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_sample(&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_sample(&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)
);
}
}

View file

@ -0,0 +1,90 @@
use askama::Template;
use serde::Serialize;
use super::{AttributeCatalog, Discovery, MetadataCatalog, TableSchema};
use crate::{Error, NormalizedFieldDefinition, query_access::ReaderLimits};
#[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],
pub limits: &'a ReaderLimits,
}
#[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)
}

View file

@ -0,0 +1,173 @@
use litellm_storage_clickhouse::Query;
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize, Serialize)]
pub struct LensAccessParams {
#[serde(deserialize_with = "super::number::deserialize")]
pub all_teams: u8,
pub team: String,
pub key_hash: String,
}
pub struct LensAvailability;
#[derive(Debug, Deserialize, Serialize)]
pub struct LensAvailabilityParams {
#[serde(flatten)]
pub access: LensAccessParams,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct LensAvailabilityRow {
#[serde(deserialize_with = "super::number::deserialize")]
pub traces: u8,
#[serde(deserialize_with = "super::number::deserialize")]
pub requests: u8,
}
impl Query for LensAvailability {
type Params = LensAvailabilityParams;
type Row = LensAvailabilityRow;
const SQL: &'static str = include_str!("../../query/lens_availability.sql");
}
pub struct LensAgents;
#[derive(Debug, Deserialize, Serialize)]
pub struct LensAgentsParams {
#[serde(flatten)]
pub access: LensAccessParams,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct LensAgentsRow {
pub agent_name: String,
}
impl Query for LensAgents {
type Params = LensAgentsParams;
type Row = LensAgentsRow;
const SQL: &'static str = include_str!("../../query/lens_agents.sql");
}
pub struct LensSample;
#[derive(Debug, Deserialize, Serialize)]
pub struct LensSampleParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub source: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub start: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end: u64,
pub agent_name: String,
pub service: String,
pub filter_keys: Vec<String>,
pub filter_values: Vec<String>,
pub selected_team: String,
pub execution_ids: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub sample_cap: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub sample_percent: f64,
#[serde(deserialize_with = "super::number::deserialize")]
pub preview: u8,
pub after: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub limit: u32,
#[serde(deserialize_with = "super::number::deserialize")]
pub offset: u64,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct LensSampleRow {
pub source: String,
pub trace_id: String,
pub team_id: String,
pub trace_ref: String,
pub name: String,
pub start_time: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub span_count: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub root_seen: u8,
pub service: String,
pub attributes: Vec<(String, String)>,
#[serde(deserialize_with = "super::number::deserialize")]
pub eligible: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub position: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub selected: f64,
pub selection_key: String,
}
impl Query for LensSample {
type Params = LensSampleParams;
type Row = LensSampleRow;
const SQL: &'static str = include_str!("../../query/lens_sample.sql");
}
pub struct LensContent;
#[derive(Debug, Deserialize, Serialize)]
pub struct LensContentParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub source: String,
pub id: String,
pub record_team: String,
pub trace_ref: String,
pub cursor: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub offset: u32,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct LensContentRow {
pub span_id: String,
pub parent_span_id: String,
pub name: String,
pub kind: String,
pub content: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub truncated: u8,
}
impl Query for LensContent {
type Params = LensContentParams;
type Row = LensContentRow;
const SQL: &'static str = include_str!("../../query/lens_content.sql");
}
pub struct LensEvidence;
#[derive(Debug, Deserialize, Serialize)]
pub struct LensEvidenceParams {
#[serde(flatten)]
pub access: LensAccessParams,
pub source: String,
pub id: String,
pub record_team: String,
pub trace_ref: String,
pub span: String,
pub quote: String,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct LensEvidenceRow {
#[serde(deserialize_with = "super::number::deserialize")]
pub count: u64,
}
impl Query for LensEvidence {
type Params = LensEvidenceParams;
type Row = LensEvidenceRow;
const SQL: &'static str = include_str!("../../query/lens_evidence.sql");
}

View file

@ -0,0 +1,320 @@
use litellm_storage_clickhouse::Query;
use litellm_traces::query::named as contracts;
use serde::{Deserialize, Serialize};
pub use contracts::ReadAccessParams;
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::ListTracesParams")]
struct ListTracesParamsEncoding {
#[serde(flatten)]
pub access: contracts::ReadAccessParams,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub cursor_ms: i64,
pub cursor_trace_id: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub limit: u32,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ListTracesParams(
#[serde(with = "ListTracesParamsEncoding")] pub contracts::ListTracesParams,
);
impl From<contracts::ListTracesParams> for ListTracesParams {
fn from(value: contracts::ListTracesParams) -> Self {
Self(value)
}
}
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::ListTracesRow")]
struct ListTracesRowEncoding {
pub trace_id: String,
pub trace_ref: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,
pub name: String,
pub service: String,
pub input_preview: String,
pub status: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub duration_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub span_count: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub agent_count: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub agent_invocations: u64,
#[serde(default)]
pub agent_names: Vec<String>,
#[serde(default)]
pub frameworks: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub llm_calls: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub tool_calls: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub input_tokens: u64,
#[serde(deserialize_with = "super::number::deserialize")]
pub output_tokens: u64,
pub models: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub error_count: u64,
pub request_ids: Vec<String>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ListTracesRow(#[serde(with = "ListTracesRowEncoding")] pub contracts::ListTracesRow);
pub use contracts::TraceSpansParams;
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::TraceSpansRow")]
struct TraceSpansRowEncoding {
pub span_id: String,
pub parent_span_id: String,
pub name: String,
#[serde(rename = "type")]
pub kind: String,
pub agent: String,
#[serde(default)]
pub framework: String,
pub status: String,
pub status_message: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub error_truncated: u8,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ns: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub duration_ns: u64,
pub service: String,
pub input_preview: String,
pub model: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub input_tokens: u32,
#[serde(deserialize_with = "super::number::deserialize")]
pub output_tokens: u32,
pub litellm_request_id: String,
pub team_id: String,
pub api_key_hash: String,
pub user_id: String,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct TraceSpansRow(#[serde(with = "TraceSpansRowEncoding")] pub contracts::TraceSpansRow);
pub use contracts::SpanDetailParams;
pub use contracts::SpanDetailRow;
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::SpanErrorParams")]
struct SpanErrorParamsEncoding {
#[serde(flatten)]
pub access: contracts::ReadAccessParams,
pub trace_id: String,
pub trace_ref: String,
pub span_id: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub error_offset: u64,
pub error_version: String,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct SpanErrorParams(
#[serde(with = "SpanErrorParamsEncoding")] pub contracts::SpanErrorParams,
);
impl From<contracts::SpanErrorParams> for SpanErrorParams {
fn from(value: contracts::SpanErrorParams) -> Self {
Self(value)
}
}
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::SpanErrorRow")]
struct SpanErrorRowEncoding {
pub span_id: String,
pub message: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub total_chars: u64,
pub version: String,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct SpanErrorRow(#[serde(with = "SpanErrorRowEncoding")] pub contracts::SpanErrorRow);
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::SpendByResponseIdsParams")]
struct SpendByResponseIdsParamsEncoding {
#[serde(flatten)]
pub access: contracts::ReadAccessParams,
pub response_ids: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end_ms: i64,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct SpendByResponseIdsParams(
#[serde(with = "SpendByResponseIdsParamsEncoding")] pub contracts::SpendByResponseIdsParams,
);
impl From<contracts::SpendByResponseIdsParams> for SpendByResponseIdsParams {
fn from(value: contracts::SpendByResponseIdsParams) -> Self {
Self(value)
}
}
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::SpendByResponseIdsRow")]
struct SpendByResponseIdsRowEncoding {
pub request_id: String,
pub response_id: String,
pub team_id: String,
pub api_key: String,
pub user: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub spend: f64,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct SpendByResponseIdsRow(
#[serde(with = "SpendByResponseIdsRowEncoding")] pub contracts::SpendByResponseIdsRow,
);
pub struct ListTraces;
impl Query for ListTraces {
type Params = ListTracesParams;
type Row = ListTracesRow;
const SQL: &'static str = include_str!("../../query/list_traces.sql");
}
pub struct TraceSpans;
impl Query for TraceSpans {
type Params = TraceSpansParams;
type Row = TraceSpansRow;
const SQL: &'static str = include_str!("../../query/trace_spans.sql");
}
pub struct SpanDetail;
impl Query for SpanDetail {
type Params = SpanDetailParams;
type Row = SpanDetailRow;
const SQL: &'static str = include_str!("../../query/span_detail.sql");
}
pub struct SpanError;
impl Query for SpanError {
type Params = SpanErrorParams;
type Row = SpanErrorRow;
const SQL: &'static str = include_str!("../../query/span_error.sql");
}
pub struct SpendByResponseIds;
impl Query for SpendByResponseIds {
type Params = SpendByResponseIdsParams;
type Row = SpendByResponseIdsRow;
const SQL: &'static str = include_str!("../../query/spend_by_response_ids.sql");
}
pub use contracts::{TraceIdentityParams, TraceIdentityRow};
pub struct TraceIdentity;
impl Query for TraceIdentity {
type Params = TraceIdentityParams;
type Row = TraceIdentityRow;
const SQL: &'static str = include_str!("../../query/trace_identity.sql");
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use serde_json::{Value, json};
fn round_trip<T: serde::de::DeserializeOwned + Serialize>(wire: Value, quoted: bool) {
let encoded = Value::Object(
wire.as_object()
.unwrap()
.iter()
.map(|(name, value)| {
let encoded = if quoted && value.is_number() && name != "all_teams" {
json!(value.to_string())
} else {
value.clone()
};
(name.clone(), encoded)
})
.collect(),
);
let decoded: T = serde_json::from_value(encoded).unwrap();
assert_eq!(serde_json::to_value(decoded).unwrap(), wire);
}
#[rstest]
#[case::unquoted(false)]
#[case::quoted(true)]
fn rows_decode_into_neutral_contracts(#[case] quoted: bool) {
round_trip::<ListTracesRow>(
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
quoted,
);
round_trip::<TraceSpansRow>(
json!({"span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "agent": "agent", "framework": "claude-agent-sdk", "status": "error", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
quoted,
);
round_trip::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),
quoted,
);
round_trip::<SpanErrorRow>(
json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}),
quoted,
);
round_trip::<SpendByResponseIdsRow>(
json!({"request_id": "request", "response_id": "response", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
quoted,
);
}
#[rstest]
#[case::unquoted(false)]
#[case::quoted(true)]
fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) {
round_trip::<ListTracesParams>(
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "key", "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}),
quoted,
);
round_trip::<SpanErrorParams>(
json!({"all_teams": 0, "user_id": "", "team_ids": [], "api_key_hash": "key", "trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
quoted,
);
round_trip::<SpendByResponseIdsParams>(
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "", "response_ids": ["response"], "start_ms": -1, "end_ms": 10}),
quoted,
);
}
}

View file

@ -0,0 +1,44 @@
use serde::{Deserialize, Deserializer, de::DeserializeOwned};
pub(super) fn deserialize<'de, D, T>(deserializer: D) -> Result<T, D::Error>
where
D: Deserializer<'de>,
T: DeserializeOwned,
{
#[derive(Deserialize)]
#[serde(untagged)]
enum Number {
Quoted(String),
Unquoted(serde_json::Number),
}
match Number::deserialize(deserializer)? {
Number::Quoted(value) => serde_json::from_str(&value),
Number::Unquoted(value) => serde_json::from_value(serde_json::Value::Number(value)),
}
.map_err(serde::de::Error::custom)
}
#[cfg(test)]
mod tests {
use crate::query::named::SpanErrorRow;
use rstest::rstest;
#[rstest]
#[case::quoted_max(serde_json::json!(u64::MAX.to_string()), Some(u64::MAX))]
#[case::unquoted_max(serde_json::json!(u64::MAX), Some(u64::MAX))]
#[case::overflow(serde_json::json!("18446744073709551616"), None)]
#[case::negative(serde_json::json!(-1), None)]
#[case::fraction(serde_json::json!(1.5), None)]
fn numeric_rows_enforce_integer_range(
#[case] value: serde_json::Value,
#[case] expected: Option<u64>,
) {
let row = serde_json::from_value::<SpanErrorRow>(serde_json::json!({
"span_id": "span", "message": "error", "total_chars": value, "version": "hash"
}));
match expected {
Some(value) => assert_eq!(row.unwrap().0.total_chars, value),
None => assert!(row.is_err()),
}
}
}

View file

@ -0,0 +1,271 @@
use std::{sync::Arc, time::Duration};
use hmac::{Hmac, Mac};
use litellm_http::Client;
use litellm_storage_clickhouse::READ_LIMITS;
use litellm_traces::QueryScope;
use moka::future::Cache;
use strum::IntoEnumIterator;
use sha2::{Digest, Sha256};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use super::{Connection, Error, TraceTable};
const MIB: u64 = 1024 * 1024;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct ReaderLimits {
pub result_rows: u64,
pub result_bytes: u64,
pub memory_bytes: u64,
pub execution_seconds: u64,
}
impl ReaderLimits {
pub fn result_mib(&self) -> u64 {
self.result_bytes / MIB
}
pub fn memory_mib(&self) -> u64 {
self.memory_bytes / MIB
}
}
pub(crate) const READER_LIMITS: ReaderLimits = ReaderLimits {
result_rows: READ_LIMITS.result_rows,
result_bytes: READ_LIMITS.response_bytes as u64,
memory_bytes: 256 * MIB,
execution_seconds: READ_LIMITS.execution_seconds,
};
#[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, Error> {
self.slots
.clone()
.try_acquire_owned()
.map_err(|_| Error::Busy)
}
pub async fn connection(
&self,
client: &Client,
scope: &QueryScope,
secret: &str,
) -> Result<Connection, Error> {
scope.validate().map_err(|_| Error::InvalidScope)?;
if secret.is_empty() {
return Err(Error::MissingSecret);
}
let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope))
.map_err(|_| Error::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(Error::Cached)
}
async fn provision(
&self,
client: &Client,
scope: &QueryScope,
user: &str,
password: &str,
) -> Result<Connection, Error> {
let database = &self.database;
if database.is_empty()
|| !database
.bytes()
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
{
return Err(Error::InvalidScope);
}
let password_hash = format!("{:x}", Sha256::digest(password));
let ReaderLimits {
result_rows,
result_bytes,
memory_bytes,
execution_seconds,
} = READER_LIMITS;
self.execute(
client,
format!(
"CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \
SETTINGS readonly = 1 CONST, max_execution_time = {execution_seconds} CONST, \
max_result_rows = {result_rows} CONST, max_result_bytes = {result_bytes} CONST, \
result_overflow_mode = 'throw' CONST, max_memory_usage = {memory_bytes} 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 TraceTable::iter() {
let predicate = predicate(scope, 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 TraceTable::iter() {
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(Error::Storage)
}
async fn execute(&self, client: &Client, sql: String) -> Result<(), Error> {
let response = client
.post(self.writer.url().clone())
.timeout(Duration::from_secs(15))
.body(sql)
.send()
.await
.map_err(|_| Error::ProvisionTransport)?;
if !response.status().is_success() {
return Err(Error::ProvisionFailed(response.status().as_u16()));
}
Ok(())
}
}
fn predicate(scope: &QueryScope, table: TraceTable) -> String {
let (team, key) = match table {
TraceTable::OtelTraces | TraceTable::AgentTracesByKey => ("TeamId", "ApiKeyHash"),
TraceTable::SpendLogs => ("team_id", "api_key"),
};
match scope {
QueryScope::Admin => "1".to_owned(),
QueryScope::Logs {
user_id,
team_ids,
api_key_hash,
} => {
let owner = literal(user_id);
let user_clause = match table {
TraceTable::OtelTraces => format!("UserId = {owner}"),
TraceTable::AgentTracesByKey => format!("UserIds = [{owner}]"),
TraceTable::SpendLogs => format!("user = {owner}"),
};
let teams = team_ids
.iter()
.map(|value| literal(value))
.collect::<Vec<_>>()
.join(", ");
let team_clause = if team_ids.is_empty() {
"0".to_owned()
} else {
format!("{team} IN ({teams})")
};
format!(
"({owner} != '' AND {user_clause}) OR ({team_clause}) OR ({hash} != '' AND {key} = {hash})",
owner = owner,
hash = literal(api_key_hash),
)
}
QueryScope::Team { team_id } => format!("{team} = {}", literal(team_id)),
QueryScope::Key {
team_id,
api_key_hash,
} => format!(
"{team} = {} AND {key} = {}",
literal(team_id),
literal(api_key_hash)
),
}
}
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, Error> {
let mut mac =
Hmac::<Sha256>::new_from_slice(secret.as_bytes()).map_err(|_| Error::MissingSecret)?;
mac.update(purpose);
mac.update(identity);
Ok(format!("{:x}", mac.finalize().into_bytes()))
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
#[case::otel(TraceTable::OtelTraces, "TeamId", "ApiKeyHash")]
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "ApiKeyHash")]
#[case::spend(TraceTable::SpendLogs, "team_id", "api_key")]
fn predicates_preserve_scope_and_escape_values(
#[case] table: TraceTable,
#[case] team: &str,
#[case] key: &str,
) {
assert_eq!(predicate(&QueryScope::Admin, table), "1");
assert_eq!(
predicate(
&QueryScope::Team {
team_id: "team'\\".into()
},
table
),
format!("{team} = 'team\\'\\\\'")
);
assert_eq!(
predicate(
&QueryScope::Key {
team_id: "".into(),
api_key_hash: "key'\\".into()
},
table
),
format!("{team} = '' AND {key} = 'key\\'\\\\'")
);
}
}

View file

@ -0,0 +1,136 @@
use litellm_http::Client;
use litellm_migrate::Migration;
use serde::Serialize;
use std::time::Duration;
use super::Connection;
use super::Error;
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations");
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'_')
|| 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(|migration| {
migration
.sql
.replace("{database}", &database)
.replace("{retention_days}", &retention_days.to_string())
}))
.collect(),
)
}
pub async fn ensure_schema(
client: &Client,
connection: &Connection,
database: &str,
retention_days: u32,
) -> Result<(), Error> {
ensure_schema_with_timeout(
client,
connection,
database,
retention_days,
SCHEMA_REQUEST_TIMEOUT,
)
.await
}
async fn ensure_schema_with_timeout(
client: &Client,
connection: &Connection,
database: &str,
retention_days: u32,
request_timeout: Duration,
) -> Result<(), Error> {
for statement in schema_statements(database, retention_days)? {
let response = client
.post(connection.url().clone())
.timeout(request_timeout)
.body(statement)
.send()
.await
.map_err(|_| Error::SchemaTransport)?;
if !response.status().is_success() {
return Err(Error::SchemaFailed(response.status().as_u16()));
}
}
Ok(())
}
#[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; 9] = [
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: "framework",
clickhouse_column: "Framework",
clickhouse_type: "LowCardinality(String)",
meaning: "Agent framework or SDK that emitted this span, e.g. claude-agent-sdk",
},
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",
},
];

View file

@ -0,0 +1,83 @@
use std::collections::BTreeMap;
use litellm_http::Client;
use litellm_traces::ReadQuery;
use super::query::{lens::*, named::*};
use super::{Connection, Error, Parameter};
use litellm_storage_clickhouse::{Query, fetch_json};
pub async fn execute_named_read(
client: &Client,
connection: &Connection,
query: ReadQuery,
parameters: &BTreeMap<String, Parameter>,
) -> Result<String, Error> {
match query {
ReadQuery::ListTraces => named_json::<ListTraces>(client, connection, parameters).await,
ReadQuery::TraceIdentity => {
named_json::<TraceIdentity>(client, connection, parameters).await
}
ReadQuery::TraceSpans => named_json::<TraceSpans>(client, connection, parameters).await,
ReadQuery::SpanDetail => named_json::<SpanDetail>(client, connection, parameters).await,
ReadQuery::SpanError => named_json::<SpanError>(client, connection, parameters).await,
ReadQuery::SpendByResponseIds => {
named_json::<SpendByResponseIds>(client, connection, parameters).await
}
ReadQuery::Availability => {
named_json::<LensAvailability>(client, connection, parameters).await
}
ReadQuery::Agents => named_json::<LensAgents>(client, connection, parameters).await,
ReadQuery::Sample => named_json::<LensSample>(client, connection, parameters).await,
ReadQuery::Content => named_json::<LensContent>(client, connection, parameters).await,
ReadQuery::Evidence => named_json::<LensEvidence>(client, connection, parameters).await,
}
}
async fn named_json<Q: Query>(
client: &Client,
connection: &Connection,
parameters: &BTreeMap<String, Parameter>,
) -> Result<String, Error>
where
Q::Params: serde::de::DeserializeOwned,
{
let value = serde_json::to_value(parameters).map_err(|_| Error::InvalidParameters)?;
let params =
serde_json::from_value::<Q::Params>(value).map_err(|_| Error::InvalidParameters)?;
fetch_json::<Q>(client, connection, &params)
.await
.map_err(Error::from)
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
#[case::missing_span(serde_json::json!({}))]
#[case::negative_offset(serde_json::json!({"span_id": "span", "error_offset": -1, "error_version": ""}))]
#[case::overflow(serde_json::json!({"span_id": "span", "error_offset": "18446744073709551616", "error_version": ""}))]
#[tokio::test]
async fn named_read_rejects_invalid_parameters_before_transport(
#[case] specific: serde_json::Value,
) {
let common = serde_json::json!({
"all_teams": 1, "user_id": "", "team_ids": [], "api_key_hash": "", "trace_id": "trace", "trace_ref": ""
});
let parameters: BTreeMap<String, Parameter> = common
.as_object()
.unwrap()
.iter()
.chain(specific.as_object().unwrap().iter())
.map(|(name, value)| (name.clone(), serde_json::from_value(value.clone()).unwrap()))
.collect();
let client = Client::no_redirect_for_test();
let connection = Connection::parse("http://127.0.0.1:1").unwrap();
assert!(matches!(
execute_named_read(&client, &connection, ReadQuery::SpanError, &parameters).await,
Err(Error::InvalidParameters)
));
}
}

View file

@ -0,0 +1,17 @@
#[derive(
Clone,
Copy,
Debug,
serde::Serialize,
strum::Display,
strum::AsRefStr,
strum::EnumIter,
strum::IntoStaticStr,
)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum TraceTable {
OtelTraces,
AgentTracesByKey,
SpendLogs,
}

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 }}
{% match metadata.discovery %}{% when Discovery::Unavailable(error) %}Metadata discovery unavailable: {{ error }}
{% when Discovery::Observed(sample) %}Sampled rows: {{ sample.sampled_rows }}; invalid JSON rows: {{ sample.invalid_json_rows }}; truncated: {{ sample.truncated }}
{% if sample.fields.is_empty() %}No metadata paths found in the sampled rows
{% else %}{% for field in sample.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %}
{% endfor %}{% endif %}{% endmatch %}
Observed span and resource attributes
{% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }}
{{ catalog.scope }}
{% match catalog.discovery %}{% when Discovery::Unavailable(error) %}Attribute discovery unavailable: {{ error }}
{% when Discovery::Observed(sample) %}{% if sample.fields.is_empty() %}No attribute keys found in the sampled spans
{% else %}{% for field in sample.fields %}{{ field.expression }}: {{ field.kind }}
{% endfor %}{% endif %}{% endmatch %}{% 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.TeamId != '' OR (t.UserId != '' AND t.UserId = s.user) OR (t.ApiKeyHash != '' 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 {{ limits.result_rows }} result rows, {{ limits.result_mib() }} MiB response bytes, {{ limits.memory_mib() }} MiB memory and a {{ limits.execution_seconds }} 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 request-log visibility through row policies. Callers see their own user rows and permitted teams, or their own key rows when no user identity is available. 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; require one spend match per response ID and ownership before aggregating. Missing IDs or costs leave totals unknown{% 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 %}

View file

@ -1,18 +1,16 @@
use litellm_http::Client;
use litellm_traces::{Connection, Error, Parameter, execute_read};
use litellm_traces_clickhouse::{
Connection, Error, Parameter, QueryReaders, QueryScope, execute_read,
};
use rstest::{fixture, rstest};
use serde_json::Value;
use std::collections::BTreeMap;
use testcontainers_modules::{
clickhouse::ClickHouse,
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
};
mod support;
const CLICKHOUSE_TAG: &str =
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
use support::{ClickHouseDatabase, database as start_database};
struct Database {
_container: ContainerAsync<ClickHouse>,
_database: ClickHouseDatabase,
url: String,
admin_url: String,
client: Client,
@ -20,22 +18,9 @@ struct Database {
#[fixture]
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
let container = ClickHouse::default()
.with_tag(CLICKHOUSE_TAG)
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
.with_env_var("LITELLM_TRACES_READER_PASSWORD", "test_password")
.with_copy_to(
"/etc/clickhouse-server/users.d/litellm-traces-reader.xml",
include_bytes!("../config/reader.xml").to_vec(),
)
.start()
.await?;
let admin_url = format!(
"http://{}:{}",
container.get_host().await?,
container.get_host_port_ipv4(8123).await?,
);
let client = Client::no_redirect_for_test();
let instance = start_database().await?;
let admin_url = instance.url.clone();
let client = instance.client.clone();
for sql in [
"CREATE DATABASE litellm",
"CREATE TABLE litellm.otel_traces (n UInt8) ENGINE = Memory",
@ -54,12 +39,13 @@ async fn database() -> Result<Database, Box<dyn std::error::Error>> {
.await?
.error_for_status()?;
}
let url = format!(
"{}?database=litellm",
admin_url.replacen("http://", "http://litellm_traces_reader:test_password@", 1)
);
let readers = QueryReaders::new(Connection::writer(&admin_url)?, "litellm".into());
let connection = readers
.connection(&client, &QueryScope::Admin, "test-secret")
.await?;
let url = connection.url().to_string();
Ok(Database {
_container: container,
_database: instance,
url,
admin_url,
client,
@ -120,7 +106,15 @@ async fn reader_rejects_writes_and_privilege_escalation(
let result = read(&database.client, &connection, sql).await;
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
assert!(
matches!(
result,
Err(Error::Storage(
litellm_storage_clickhouse::Error::QueryFailed(_)
))
),
"{result:?}"
);
let rows = read(&database.client, &connection, "SELECT n FROM otel_traces").await?;
let json: Value = serde_json::from_str(&rows)?;
assert_eq!(json["data"], serde_json::json!([{ "n": 1 }]));
@ -147,7 +141,12 @@ async fn admin_sql_rejects_errors_after_output_starts(
.await;
assert!(
matches!(result, Err(Error::InvalidResponse)),
matches!(
result,
Err(Error::Storage(
litellm_storage_clickhouse::Error::InvalidResponse
))
),
"expected an error embedded in a successful HTTP response: {result:?}"
);
Ok(())
@ -171,7 +170,15 @@ async fn admin_sql_enforces_result_row_limit(
)
.await;
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
assert!(
matches!(
result,
Err(Error::Storage(
litellm_storage_clickhouse::Error::QueryFailed(_)
))
),
"{result:?}"
);
Ok(())
}
@ -190,7 +197,15 @@ async fn admin_sql_enforces_response_byte_limit(
)
.await;
assert!(matches!(result, Err(Error::ResponseTooLarge)), "{result:?}");
assert!(
matches!(
result,
Err(Error::Storage(
litellm_storage_clickhouse::Error::ResponseTooLarge
))
),
"{result:?}"
);
Ok(())
}

View file

@ -0,0 +1,15 @@
# ClickHouse query fixtures
Run `cargo test -p litellm-traces-clickhouse --test queries --locked -- --test-threads=2` from `litellm-rust` with Docker running
Raw OTLP exports live in `crates/traces/tests/fixtures/query_*.json`. The seeded fixture decodes and normalizes them through `litellm_traces::decode_otlp` at test startup, then projects the decoded fields into ClickHouse columns. Team and key identities come from fixture setup rather than exporter claims. Root and child exports are inserted separately through the public insert API so materialized views process multiple blocks
`crates/traces/tests/fixtures/deeplite_auth_error.json` and `deeplite_swarm.json` were captured from Deeplite runs against the local proxy on 2026-10-02. The first contains a failed model call. The second contains successful model calls, searches, handoff attempts, and virtual filesystem writes. Credentials, workspace identifiers, and local user paths were redacted, and the protobuf exports were converted to OTLP JSON. Their round-trip tests check span identities, parent links, timestamps, durations, token counts, and statuses without pinning the provider's error wording
The swarm capture has handoff spans marked ERROR with `ParentCommand` exception events and a root with UNSET status. These are exported diagnostic statuses, which do not establish a failed execution. The tests preserve incoming statuses and check root status separately from the count of error spans, deriving both from the decoded export. They do not infer an execution outcome from exception text, framework names, successful model calls, or output presence. Framework-specific interpretation of control-flow exceptions belongs in the instrumentation integration
`spend_logs.jsonl` contains spend insert rows with millisecond timestamps, including two versions of one request. Replace this small placeholder dataset when the actual data is available. The query fixture applies production migrations, then removes TTL from its isolated database so fixed timestamps do not expire. Background merges are stopped so rollup aggregation and `FINAL` deduplication are exercised on unmerged data. Retention behavior stays covered by the migration tests
Curated SQL lives in `tests/queries/*.sql`. Each query has a matching `.expected.json` containing ordered result rows for `admin`, `team`, `key`, and `other_team` readers. Update the exports and expected results together. Add a named case in `tests/queries.rs` for each new query. Assertions compare only result data, excluding server statistics and execution timing
Typed query tests execute the production SQL through `litellm_storage_clickhouse::fetch` using contracts from `litellm-traces`. The fixture projection is test setup, so this suite covers the Rust decoder, normalization, inserts, schema, readers, and queries. Python ingress transformations, including payload truncation and exception-event fallback, remain covered by the Python tests

View file

@ -0,0 +1,3 @@
{"request_id":"request-a","response_id":"response-shared","team_id":"team-a","api_key":"key-a","user":"user-a","spend":0.125,"start_time":1735689600100,"end_time":1735689600500,"metadata":"{\"labels\":{\"priority\":\"obsolete\"}}"}
{"request_id":"request-a","response_id":"response-shared","team_id":"team-a","api_key":"key-a","user":"user-a","spend":0.5,"start_time":1735689600100,"end_time":1735689600600,"metadata":"{\"labels\":{\"priority\":\"high\"}}"}
{"request_id":"request-b","response_id":"response-shared","team_id":"team-b","api_key":"key-b","user":"user-b","spend":0.25,"start_time":1735689602000,"end_time":1735689602500,"metadata":"{\"labels\":{\"priority\":\"low\"}}"}

View file

@ -5,8 +5,9 @@ use std::{
use flate2::read::GzDecoder;
use litellm_http::Client;
use litellm_traces::{
Connection, Error, InsertRow, InsertTable, Shared, encode_rows, insert_shared_rows,
use litellm_traces::Shared;
use litellm_traces_clickhouse::{
Connection, Error, InsertRow, InsertTable, encode_rows, insert_shared_rows,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,273 @@
use std::collections::BTreeMap;
use litellm_storage_clickhouse::fetch;
use litellm_traces::query::named as contracts;
use litellm_traces_clickhouse::{
QueryScope,
query::named::{ListTraces, ListTracesParams, TraceSpans, TraceSpansParams},
query_sql,
};
use rstest::{fixture, rstest};
use serde::Deserialize;
use serde_json::Value;
#[path = "queries/support.rs"]
mod fixtures;
mod support;
use fixtures::{SeededDatabase, insert_export, migrated_database, seeded_database};
use support::TestResult;
#[derive(Clone, Copy, strum::AsRefStr)]
#[strum(serialize_all = "snake_case")]
enum ScopeCase {
Admin,
Team,
Key,
OtherTeam,
}
impl ScopeCase {
fn scope(self) -> QueryScope {
match self {
Self::Admin => QueryScope::Admin,
Self::Team => QueryScope::Team {
team_id: "team-a".into(),
},
Self::Key => QueryScope::Key {
team_id: "team-a".into(),
api_key_hash: "key-a".into(),
},
Self::OtherTeam => QueryScope::Team {
team_id: "team-b".into(),
},
}
}
}
#[derive(Deserialize)]
struct QueryResult {
data: Vec<Value>,
}
#[rstest]
#[case::rollups(include_str!("queries/rollups.sql"), include_str!("queries/rollups.expected.json"))]
#[case::costs(include_str!("queries/trace_costs.sql"), include_str!("queries/trace_costs.expected.json"))]
#[case::errors(include_str!("queries/failed_spans.sql"), include_str!("queries/failed_spans.expected.json"))]
#[case::metadata(include_str!("queries/metadata_filters.sql"), include_str!("queries/metadata_filters.expected.json"))]
#[tokio::test]
async fn curated_queries_return_expected_rows(
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
#[case] sql: &str,
#[case] expected_json: &str,
#[values(
ScopeCase::Admin,
ScopeCase::Team,
ScopeCase::Key,
ScopeCase::OtherTeam
)]
scope: ScopeCase,
) -> TestResult {
let fixture = seeded_database?;
let reader = fixture
.readers
.connection(&fixture.database.client, &scope.scope(), "fixture-secret")
.await?;
let result: QueryResult =
serde_json::from_str(&query_sql(&fixture.database.client, &reader, sql).await?)?;
let expected: BTreeMap<String, Vec<Value>> = serde_json::from_str(expected_json)?;
assert_eq!(
&result.data,
expected
.get(scope.as_ref())
.ok_or("missing expected scope")?,
"{}: {sql}",
scope.as_ref()
);
Ok(())
}
#[fixture]
fn admin_access() -> TestResult<contracts::ReadAccessParams> {
Ok(serde_json::from_str(include_str!(
"queries/read_access.json"
))?)
}
#[rstest]
#[tokio::test]
async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate(
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
) -> TestResult {
let fixture = seeded_database?;
let reader = fixture
.readers
.connection(
&fixture.database.client,
&QueryScope::Admin,
"fixture-secret",
)
.await?;
let params = ListTracesParams::from(contracts::ListTracesParams {
access: admin_access?,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 10,
});
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, &params).await?;
assert_eq!(
traces
.iter()
.map(|row| row.0.api_key_hash.as_str())
.collect::<Vec<_>>(),
["key-b", "key-alt", "key-a"]
);
let trace = &traces[2].0;
assert_eq!(
(
trace.span_count,
trace.llm_calls,
trace.tool_calls,
trace.error_count
),
(3, 1, 1, 1)
);
assert_eq!((trace.input_tokens, trace.output_tokens), (12, 6));
assert_eq!(trace.input_preview, "Review the change");
let span_params = TraceSpansParams {
access: params.0.access,
trace_id: trace.trace_id.clone(),
trace_ref: trace.trace_ref.clone(),
};
let spans = fetch::<TraceSpans>(&fixture.database.client, &reader, &span_params).await?;
assert_eq!(
spans
.iter()
.map(|row| row.0.name.as_str())
.collect::<Vec<_>>(),
["review", "completion", "lookup"]
);
assert!(
spans
.iter()
.all(|row| row.0.api_key_hash == trace.api_key_hash)
);
assert_eq!(
(
spans[1].0.kind.as_str(),
spans[1].0.input_tokens,
spans[1].0.output_tokens
),
("llm", 12, 6)
);
assert_eq!(spans[2].0.status_message, "lookup timed out");
Ok(())
}
#[rstest]
#[tokio::test]
async fn typed_trace_cursor_returns_the_next_fixture_trace(
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
) -> TestResult {
let fixture = seeded_database?;
let reader = fixture
.readers
.connection(
&fixture.database.client,
&QueryScope::Admin,
"fixture-secret",
)
.await?;
let params = ListTracesParams::from(contracts::ListTracesParams {
access: admin_access?,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 1,
});
let first = fetch::<ListTraces>(&fixture.database.client, &reader, &params).await?;
assert_eq!(first.len(), 1);
assert_eq!(first[0].0.api_key_hash, "key-b");
let next_params = ListTracesParams::from(contracts::ListTracesParams {
cursor_ms: first[0].0.start_ms,
cursor_trace_id: first[0].0.trace_ref.clone(),
..params.0
});
let next = fetch::<ListTraces>(&fixture.database.client, &reader, &next_params).await?;
assert_eq!(next.len(), 1);
assert_eq!(next[0].0.api_key_hash, "key-alt");
assert_ne!(first[0].0.trace_ref, next[0].0.trace_ref);
Ok(())
}
#[rstest]
#[case::authentication_error(include_bytes!("../../traces/tests/fixtures/deeplite_auth_error.json"))]
#[case::swarm(include_bytes!("../../traces/tests/fixtures/deeplite_swarm.json"))]
#[tokio::test]
async fn captured_deeplite_exports_round_trip_through_clickhouse(
#[future(awt)] migrated_database: TestResult<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
#[case] export: &[u8],
) -> TestResult {
let fixture = migrated_database?;
let decoded = insert_export(&fixture, export, "team-a", "key-a").await?;
let reader = fixture
.readers
.connection(
&fixture.database.client,
&QueryScope::Admin,
"fixture-secret",
)
.await?;
let params = TraceSpansParams {
access: admin_access?,
trace_id: decoded[0].trace_id.clone(),
trace_ref: String::new(),
};
let stored = fetch::<TraceSpans>(&fixture.database.client, &reader, &params).await?;
assert_eq!(stored.len(), decoded.len());
let list_params = ListTracesParams::from(contracts::ListTracesParams {
access: params.access,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 10,
});
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, &list_params).await?;
assert_eq!(traces.len(), 1);
let roots = decoded
.iter()
.filter(|span| span.parent_span_id.is_empty())
.collect::<Vec<_>>();
assert_eq!(roots.len(), 1);
assert_eq!(traces[0].0.status, roots[0].status_code);
assert_eq!(
traces[0].0.error_count,
decoded
.iter()
.filter(|span| span.status_code == "STATUS_CODE_ERROR")
.count() as u64
);
let by_id: BTreeMap<_, _> = stored
.iter()
.map(|row| (row.0.span_id.as_str(), &row.0))
.collect();
for span in &decoded {
let row = by_id
.get(span.span_id.as_str())
.ok_or("missing captured span")?;
assert_eq!(row.parent_span_id, span.parent_span_id);
assert_eq!(row.start_ns as u64, span.start_ns);
assert_eq!(row.duration_ns, span.end_ns - span.start_ns);
assert_eq!(row.input_tokens, span.normalized.input_tokens);
assert_eq!(row.output_tokens, span.normalized.output_tokens);
assert_eq!(row.status, span.status_code);
}
Ok(())
}

View file

@ -0,0 +1,30 @@
{
"admin": [
{
"team": "team-a",
"api_key": "key-a",
"trace_id": "01010101010101010101010101010101",
"span_id": "0303030303030303",
"message": "lookup timed out"
}
],
"team": [
{
"team": "team-a",
"api_key": "key-a",
"trace_id": "01010101010101010101010101010101",
"span_id": "0303030303030303",
"message": "lookup timed out"
}
],
"key": [
{
"team": "team-a",
"api_key": "key-a",
"trace_id": "01010101010101010101010101010101",
"span_id": "0303030303030303",
"message": "lookup timed out"
}
],
"other_team": []
}

View file

@ -0,0 +1,5 @@
SELECT TeamId AS team, ApiKeyHash AS api_key, TraceId AS trace_id,
SpanId AS span_id, StatusMessage AS message
FROM otel_traces
WHERE StatusCode = 'STATUS_CODE_ERROR'
ORDER BY team, api_key, trace_id, span_id

View file

@ -0,0 +1,30 @@
{
"admin": [
{
"team": "team-a",
"api_key": "key-a",
"request_id": "request-a",
"spend": 0.5,
"priority": "high"
}
],
"team": [
{
"team": "team-a",
"api_key": "key-a",
"request_id": "request-a",
"spend": 0.5,
"priority": "high"
}
],
"key": [
{
"team": "team-a",
"api_key": "key-a",
"request_id": "request-a",
"spend": 0.5,
"priority": "high"
}
],
"other_team": []
}

View file

@ -0,0 +1,5 @@
SELECT team_id AS team, api_key, request_id, spend,
JSONExtractString(metadata, 'labels', 'priority') AS priority
FROM spend_logs FINAL
WHERE JSONExtractString(metadata, 'labels', 'priority') = 'high'
ORDER BY team, api_key, request_id

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