mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
Merge remote-tracking branch 'origin/main' into litellm_durable_background_interaction_settlement
# Conflicts: # tests/integration/_support/upstream.py
This commit is contained in:
commit
eee6988b44
762 changed files with 75154 additions and 22372 deletions
|
|
@ -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=$!
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 ;;
|
||||
|
|
|
|||
1
.github/workflows/test-e2e-changed.yml
vendored
1
.github/workflows/test-e2e-changed.yml
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ concurrency:
|
|||
jobs:
|
||||
resolve:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 25
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
|
|
|
|||
2
.github/workflows/test-rust.yml
vendored
2
.github/workflows/test-rust.yml
vendored
|
|
@ -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
1
.gitignore
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -6,8 +6,6 @@ services:
|
|||
context: .
|
||||
dockerfile: docker/Dockerfile.non_root
|
||||
target: runtime
|
||||
args:
|
||||
PROXY_EXTRAS_SOURCE: "local"
|
||||
depends_on:
|
||||
- squid
|
||||
user: "101:101"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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".
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
);
|
||||
|
|
@ -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;
|
||||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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==",
|
||||
|
|
|
|||
95
litellm-rust/Cargo.lock
generated
95
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use litellm_config::{Config, Error, Flag, NumberOrString};
|
||||
use litellm_config::{Config, Error, Flag, NumberOrString, TracingStoreSettings};
|
||||
use rstest::{fixture, rstest};
|
||||
use tempfile::TempDir;
|
||||
|
||||
|
|
@ -113,6 +113,54 @@ fn missing_general_settings_has_no_master_key() {
|
|||
assert!(config.general_settings.master_key.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_settings_are_typed_and_redact_the_url() {
|
||||
let config = Config::from_yaml(
|
||||
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: https://writer:password@example.com\n database: analytics\n retention_days: 7\n",
|
||||
)
|
||||
.unwrap();
|
||||
let tracing = config.general_settings.tracing.as_ref().unwrap();
|
||||
let Some(TracingStoreSettings::ClickHouse(store)) = tracing.store.as_ref() else {
|
||||
panic!("expected ClickHouse tracing store")
|
||||
};
|
||||
assert_eq!(
|
||||
store.url.as_ref().unwrap().expose(),
|
||||
"https://writer:password@example.com"
|
||||
);
|
||||
assert_eq!(store.database.as_deref(), Some("analytics"));
|
||||
assert_eq!(store.retention_days, Some(NumberOrString::Number(7.0)));
|
||||
assert!(!format!("{config:?}").contains("password"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_settings_accept_environment_references() {
|
||||
let config = Config::from_yaml(
|
||||
"general_settings:\n tracing:\n store:\n type: clickhouse\n url: os.environ/CLICKHOUSE_URL\n retention_days: os.environ/RETENTION_DAYS\n",
|
||||
)
|
||||
.unwrap();
|
||||
let Some(TracingStoreSettings::ClickHouse(store)) =
|
||||
config.general_settings.tracing.unwrap().store
|
||||
else {
|
||||
panic!("expected ClickHouse tracing store")
|
||||
};
|
||||
assert_eq!(
|
||||
store.retention_days,
|
||||
Some(NumberOrString::String(
|
||||
"os.environ/RETENTION_DAYS".to_owned()
|
||||
))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_settings_reject_string_store() {
|
||||
assert!(Config::from_yaml("general_settings:\n tracing:\n store: clickhouse\n").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tracing_settings_reject_removed_reader_configuration() {
|
||||
assert!(Config::from_yaml("general_settings:\n tracing:\n store:\n type: clickhouse\n reader_url: http://localhost:8123\n").is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn empty_config_matches_python_defaults() {
|
||||
let config = Config::from_yaml("{}").unwrap();
|
||||
|
|
|
|||
19
litellm-rust/crates/migrate-macros/Cargo.toml
Normal file
19
litellm-rust/crates/migrate-macros/Cargo.toml
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
[package]
|
||||
name = "litellm-migrate-macros"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[lib]
|
||||
proc-macro = true
|
||||
|
||||
[dependencies]
|
||||
proc-macro2.workspace = true
|
||||
quote.workspace = true
|
||||
syn = { workspace = true, features = ["parsing", "printing", "proc-macro"] }
|
||||
thiserror.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tempfile.workspace = true
|
||||
21
litellm-rust/crates/migrate-macros/src/error.rs
Normal file
21
litellm-rust/crates/migrate-macros/src/error.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
use std::io;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("could not read migrations directory `{path}`")]
|
||||
ReadDirectory {
|
||||
path: String,
|
||||
#[source]
|
||||
source: io::Error,
|
||||
},
|
||||
#[error(
|
||||
"migration name `{name}` must be `<digits>_<description>.sql` with a `[a-z0-9_]` description"
|
||||
)]
|
||||
InvalidName { name: String },
|
||||
#[error("migration version `{version}` is declared more than once")]
|
||||
DuplicateVersion { version: u64 },
|
||||
#[error("migrations directory `{path}` contains no migrations")]
|
||||
Empty { path: String },
|
||||
#[error("migration path `{path}` is not valid UTF-8")]
|
||||
NonUtf8Path { path: String },
|
||||
}
|
||||
199
litellm-rust/crates/migrate-macros/src/lib.rs
Normal file
199
litellm-rust/crates/migrate-macros/src/lib.rs
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
mod error;
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use error::Error;
|
||||
use proc_macro::TokenStream;
|
||||
use quote::quote;
|
||||
use syn::LitStr;
|
||||
|
||||
struct Entry {
|
||||
version: u64,
|
||||
description: String,
|
||||
path: PathBuf,
|
||||
}
|
||||
|
||||
fn resolve(dir: &Path) -> Result<Vec<Entry>, Error> {
|
||||
let mut entries = Vec::new();
|
||||
let files = std::fs::read_dir(dir).map_err(|source| Error::ReadDirectory {
|
||||
path: dir.display().to_string(),
|
||||
source,
|
||||
})?;
|
||||
for file in files {
|
||||
let file = file.map_err(|source| Error::ReadDirectory {
|
||||
path: dir.display().to_string(),
|
||||
source,
|
||||
})?;
|
||||
let path = file.path();
|
||||
let name = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.ok_or_else(|| Error::NonUtf8Path {
|
||||
path: path.display().to_string(),
|
||||
})?
|
||||
.to_owned();
|
||||
let invalid = || Error::InvalidName { name: name.clone() };
|
||||
let stem = name
|
||||
.strip_suffix(".sql")
|
||||
.filter(|_| file.file_type().is_ok_and(|kind| kind.is_file()))
|
||||
.and_then(|stem| stem.split_once('_'))
|
||||
.filter(|(version, description)| {
|
||||
!version.is_empty()
|
||||
&& version.bytes().all(|b| b.is_ascii_digit())
|
||||
&& !description.is_empty()
|
||||
&& description
|
||||
.bytes()
|
||||
.all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'_')
|
||||
})
|
||||
.ok_or_else(invalid)?;
|
||||
let version = stem.0.parse::<u64>().map_err(|_| invalid())?;
|
||||
entries.push(Entry {
|
||||
version,
|
||||
description: stem.1.to_owned(),
|
||||
path,
|
||||
});
|
||||
}
|
||||
if entries.is_empty() {
|
||||
return Err(Error::Empty {
|
||||
path: dir.display().to_string(),
|
||||
});
|
||||
}
|
||||
entries.sort_by_key(|entry| entry.version);
|
||||
for pair in entries.windows(2) {
|
||||
if pair[0].version == pair[1].version {
|
||||
return Err(Error::DuplicateVersion {
|
||||
version: pair[0].version,
|
||||
});
|
||||
}
|
||||
}
|
||||
Ok(entries)
|
||||
}
|
||||
|
||||
fn resolve_input(lit: &LitStr) -> Result<Vec<Entry>, Error> {
|
||||
let root = std::env::var("CARGO_MANIFEST_DIR")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_default();
|
||||
let dir = root.join(lit.value());
|
||||
let dir = dir.canonicalize().map_err(|source| Error::ReadDirectory {
|
||||
path: dir.display().to_string(),
|
||||
source,
|
||||
})?;
|
||||
if dir.to_str().is_none() {
|
||||
return Err(Error::NonUtf8Path {
|
||||
path: dir.display().to_string(),
|
||||
});
|
||||
}
|
||||
resolve(&dir)
|
||||
}
|
||||
|
||||
#[proc_macro]
|
||||
pub fn migrate(input: TokenStream) -> TokenStream {
|
||||
let lit = syn::parse_macro_input!(input as LitStr);
|
||||
match resolve_input(&lit) {
|
||||
Ok(entries) => {
|
||||
let migrations = entries.iter().map(|entry| {
|
||||
let version = entry.version;
|
||||
let description = &entry.description;
|
||||
let path = entry
|
||||
.path
|
||||
.to_str()
|
||||
.expect("canonical migration path is UTF-8");
|
||||
quote! {
|
||||
::litellm_migrate::Migration {
|
||||
version: #version,
|
||||
description: #description,
|
||||
sql: ::core::include_str!(#path),
|
||||
}
|
||||
}
|
||||
});
|
||||
quote! { &[#(#migrations),*] }.into()
|
||||
}
|
||||
Err(err) => syn::Error::new(lit.span(), err).to_compile_error().into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::fs;
|
||||
|
||||
use rstest::rstest;
|
||||
use tempfile::TempDir;
|
||||
|
||||
use super::{Error, resolve};
|
||||
|
||||
fn migrations_dir(files: &[&str]) -> TempDir {
|
||||
let dir = TempDir::new().expect("tempdir");
|
||||
for file in files {
|
||||
fs::write(dir.path().join(file), "SELECT 1").expect("write fixture");
|
||||
}
|
||||
dir
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn orders_versions_numerically() {
|
||||
let dir = migrations_dir(&["10_tenth.sql", "2_second.sql", "1_first.sql"]);
|
||||
let entries = resolve(dir.path()).expect("resolves");
|
||||
let versions: Vec<u64> = entries.iter().map(|entry| entry.version).collect();
|
||||
let descriptions: Vec<&str> = entries
|
||||
.iter()
|
||||
.map(|entry| entry.description.as_str())
|
||||
.collect();
|
||||
assert_eq!(versions, [1, 2, 10]);
|
||||
assert_eq!(descriptions, ["first", "second", "tenth"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::dash_in_version(&["0001-dash.sql"])]
|
||||
#[case::not_sql(&["notes.txt"])]
|
||||
#[case::empty_description(&["0001_.sql"])]
|
||||
#[case::non_digit_version(&["x_name.sql"])]
|
||||
#[case::uppercase_description(&["0001_Upper.sql"])]
|
||||
#[case::no_underscore(&["0001.sql"])]
|
||||
#[case::plus_sign_version(&["+10_add.sql"])]
|
||||
fn rejects_invalid_names(#[case] files: &[&str]) {
|
||||
let dir = migrations_dir(files);
|
||||
assert!(matches!(
|
||||
resolve(dir.path()),
|
||||
Err(Error::InvalidName { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rejects_subdirectories() {
|
||||
let dir = migrations_dir(&["0001_a.sql"]);
|
||||
fs::create_dir(dir.path().join("0002_b.sql")).expect("subdir");
|
||||
assert!(matches!(
|
||||
resolve(dir.path()),
|
||||
Err(Error::InvalidName { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[rstest]
|
||||
fn rejects_symlinks() {
|
||||
let dir = migrations_dir(&["0001_a.sql"]);
|
||||
let target = TempDir::new().expect("tempdir");
|
||||
let target_file = target.path().join("real.sql");
|
||||
fs::write(&target_file, "SELECT 2").expect("write fixture");
|
||||
std::os::unix::fs::symlink(&target_file, dir.path().join("0002_b.sql")).expect("symlink");
|
||||
assert!(matches!(
|
||||
resolve(dir.path()),
|
||||
Err(Error::InvalidName { .. })
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rejects_duplicate_versions() {
|
||||
let dir = migrations_dir(&["0001_a.sql", "1_b.sql"]);
|
||||
assert!(matches!(
|
||||
resolve(dir.path()),
|
||||
Err(Error::DuplicateVersion { version: 1 })
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rejects_empty_directory() {
|
||||
let dir = migrations_dir(&[]);
|
||||
assert!(matches!(resolve(dir.path()), Err(Error::Empty { .. })));
|
||||
}
|
||||
}
|
||||
12
litellm-rust/crates/migrate/Cargo.toml
Normal file
12
litellm-rust/crates/migrate/Cargo.toml
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
[package]
|
||||
name = "litellm-migrate"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
litellm-migrate-macros.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
5
litellm-rust/crates/migrate/README.md
Normal file
5
litellm-rust/crates/migrate/README.md
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
# Migrations
|
||||
|
||||
`litellm-migrate` exports the `Migration` struct and the `migrate!` macro that embeds a directory of `<digits>_<description>.sql` files at compile time, sorted by numeric version
|
||||
|
||||
The crate does not apply or track migrations; callers decide how and when the embedded SQL runs
|
||||
8
litellm-rust/crates/migrate/src/lib.rs
Normal file
8
litellm-rust/crates/migrate/src/lib.rs
Normal file
|
|
@ -0,0 +1,8 @@
|
|||
pub use litellm_migrate_macros::migrate;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct Migration {
|
||||
pub version: u64,
|
||||
pub description: &'static str,
|
||||
pub sql: &'static str,
|
||||
}
|
||||
1
litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql
vendored
Normal file
1
litellm-rust/crates/migrate/tests/fixtures/migrations/10_tenth.sql
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
SELECT 10;
|
||||
1
litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql
vendored
Normal file
1
litellm-rust/crates/migrate/tests/fixtures/migrations/1_first.sql
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
SELECT 1;
|
||||
1
litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql
vendored
Normal file
1
litellm-rust/crates/migrate/tests/fixtures/migrations/2_second.sql
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
SELECT 2;
|
||||
21
litellm-rust/crates/migrate/tests/migrate.rs
Normal file
21
litellm-rust/crates/migrate/tests/migrate.rs
Normal 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);
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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(), ¶meters).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, ¶meters).await
|
||||
litellm_traces_clickhouse::execute_named_read(
|
||||
&client,
|
||||
&connection,
|
||||
query,
|
||||
¶meters,
|
||||
)
|
||||
.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)
|
||||
}
|
||||
|
|
|
|||
5
litellm-rust/crates/storage-clickhouse/AGENTS.md
Normal file
5
litellm-rust/crates/storage-clickhouse/AGENTS.md
Normal 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
|
||||
|
|
@ -18,3 +18,4 @@ url.workspace = true
|
|||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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}")]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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, ¶meters(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, ¶meters(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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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());
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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, ¶ms).await;
|
||||
let envelope = fetch_json::<TypedQuery>(&client, &connection, ¶ms).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)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
7
litellm-rust/crates/traces-clickhouse/AGENTS.md
Normal file
7
litellm-rust/crates/traces-clickhouse/AGENTS.md
Normal 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]`
|
||||
31
litellm-rust/crates/traces-clickhouse/Cargo.toml
Normal file
31
litellm-rust/crates/traces-clickhouse/Cargo.toml
Normal 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
|
||||
3
litellm-rust/crates/traces-clickhouse/build.rs
Normal file
3
litellm-rust/crates/traces-clickhouse/build.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
fn main() {
|
||||
println!("cargo:rerun-if-changed=migrations");
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
ALTER TABLE {database}.otel_traces
|
||||
ADD COLUMN IF NOT EXISTS UserId String DEFAULT ''
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE {database}.otel_traces ADD COLUMN IF NOT EXISTS Framework LowCardinality(String) AFTER AgentName
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
49
litellm-rust/crates/traces-clickhouse/query/list_traces.sql
Normal file
49
litellm-rust/crates/traces-clickhouse/query/list_traces.sql
Normal 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
|
||||
26
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal file
26
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal 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
|
||||
|
|
@ -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})
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
26
litellm-rust/crates/traces-clickhouse/src/config.rs
Normal file
26
litellm-rust/crates/traces-clickhouse/src/config.rs
Normal 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
|
||||
}
|
||||
}
|
||||
37
litellm-rust/crates/traces-clickhouse/src/error.rs
Normal file
37
litellm-rust/crates/traces-clickhouse/src/error.rs
Normal 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>),
|
||||
}
|
||||
|
|
@ -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() {
|
||||
21
litellm-rust/crates/traces-clickhouse/src/lib.rs
Normal file
21
litellm-rust/crates/traces-clickhouse/src/lib.rs
Normal 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;
|
||||
494
litellm-rust/crates/traces-clickhouse/src/query.rs
Normal file
494
litellm-rust/crates/traces-clickhouse/src/query.rs
Normal 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)
|
||||
);
|
||||
}
|
||||
}
|
||||
90
litellm-rust/crates/traces-clickhouse/src/query/guide.rs
Normal file
90
litellm-rust/crates/traces-clickhouse/src/query/guide.rs
Normal 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)
|
||||
}
|
||||
173
litellm-rust/crates/traces-clickhouse/src/query/lens.rs
Normal file
173
litellm-rust/crates/traces-clickhouse/src/query/lens.rs
Normal 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");
|
||||
}
|
||||
320
litellm-rust/crates/traces-clickhouse/src/query/named.rs
Normal file
320
litellm-rust/crates/traces-clickhouse/src/query/named.rs
Normal 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,
|
||||
);
|
||||
}
|
||||
}
|
||||
44
litellm-rust/crates/traces-clickhouse/src/query/number.rs
Normal file
44
litellm-rust/crates/traces-clickhouse/src/query/number.rs
Normal 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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
271
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal file
271
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal 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\\'\\\\'")
|
||||
);
|
||||
}
|
||||
}
|
||||
136
litellm-rust/crates/traces-clickhouse/src/schema.rs
Normal file
136
litellm-rust/crates/traces-clickhouse/src/schema.rs
Normal 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",
|
||||
},
|
||||
];
|
||||
83
litellm-rust/crates/traces-clickhouse/src/sql.rs
Normal file
83
litellm-rust/crates/traces-clickhouse/src/sql.rs
Normal 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, ¶ms)
|
||||
.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, ¶meters).await,
|
||||
Err(Error::InvalidParameters)
|
||||
));
|
||||
}
|
||||
}
|
||||
17
litellm-rust/crates/traces-clickhouse/src/table.rs
Normal file
17
litellm-rust/crates/traces-clickhouse/src/table.rs
Normal 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,
|
||||
}
|
||||
|
|
@ -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 %}
|
||||
|
|
@ -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(())
|
||||
}
|
||||
|
||||
15
litellm-rust/crates/traces-clickhouse/tests/fixtures/README.md
vendored
Normal file
15
litellm-rust/crates/traces-clickhouse/tests/fixtures/README.md
vendored
Normal 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
|
||||
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/spend_logs.jsonl
vendored
Normal file
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/spend_logs.jsonl
vendored
Normal 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\"}}"}
|
||||
|
|
@ -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};
|
||||
1971
litellm-rust/crates/traces-clickhouse/tests/migrations.rs
Normal file
1971
litellm-rust/crates/traces-clickhouse/tests/migrations.rs
Normal file
File diff suppressed because it is too large
Load diff
273
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal file
273
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal 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, ¶ms).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, ¶ms).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, ¶ms).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(())
|
||||
}
|
||||
|
|
@ -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": []
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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": []
|
||||
}
|
||||
|
|
@ -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
Loading…
Add table
Reference in a new issue