mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
Merge branch 'main' of https://github.com/BerriAI/litellm into responses-tool-search-lowering
# Conflicts: # tests/unit/llms/openai/responses/test_openai_responses_transformation.py
This commit is contained in:
commit
c3ce477b3d
480 changed files with 43083 additions and 6145 deletions
|
|
@ -169,7 +169,7 @@ 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 \
|
||||
|
|
|
|||
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
|
||||
|
|
|
|||
|
|
@ -19,11 +19,13 @@ The URL, database, and retention settings can also come from `CLICKHOUSE_URL`, `
|
|||
|
||||
Retention changes require a proxy restart. ClickHouse removes expired rows during background merges, not immediately at startup. Enable request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
|
||||
|
||||
In Lens, click **Set up analysis**, choose an existing virtual key or **Create worker key**, then **Generate setup command**. The LiteLLM address is filled in for you; change it only if the server running Docker needs a different network address. Copy the command and run it on your server. The dialog changes to **Analyzer connected** when the container checks in
|
||||
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:67eba741c1b97c749975c5c38e2370a603e1105babc908d613c1b79d7b995393}
|
||||
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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -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)
|
||||
|
|
@ -336,9 +353,8 @@ class ProxyExtrasDBManager:
|
|||
pass
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str) -> Optional[str]:
|
||||
"""Return failed migration logs, or None if the ledger is unavailable."""
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
def _read_migration_ledger(query: str, params: tuple[str, ...]) -> "tuple[object, ...] | None":
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
return None
|
||||
|
||||
|
|
@ -347,28 +363,37 @@ class ProxyExtrasDBManager:
|
|||
except ImportError:
|
||||
return None
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
ledger_table = psycopg.sql.SQL("{}.{}").format(
|
||||
psycopg.sql.Identifier(
|
||||
ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
|
||||
),
|
||||
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
ledger_table: Final = psycopg.sql.SQL("{}.{}").format(
|
||||
psycopg.sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
|
||||
psycopg.sql.Identifier("_prisma_migrations"),
|
||||
)
|
||||
try:
|
||||
with psycopg.connect(
|
||||
cleaned_url, connect_timeout=10, autocommit=True
|
||||
) as conn:
|
||||
row = conn.execute(
|
||||
psycopg.sql.SQL(
|
||||
"SELECT logs FROM {} "
|
||||
"WHERE migration_name = %s AND finished_at IS NULL "
|
||||
"AND rolled_back_at IS NULL"
|
||||
).format(ledger_table),
|
||||
(migration_name,),
|
||||
).fetchone()
|
||||
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
|
||||
row: Final = conn.execute(psycopg.sql.SQL(query).format(ledger_table), params).fetchone()
|
||||
except (psycopg.OperationalError, psycopg.DatabaseError):
|
||||
return None
|
||||
return (row[0] or "") if row else ""
|
||||
return tuple(row) if row is not None else ()
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND finished_at IS NULL AND rolled_back_at IS NULL",
|
||||
(migration_name, started_at),
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
return row[0] if row and isinstance(row[0], str) else ""
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_recovered(migration_name: str, started_at: str) -> bool:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT 1 FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND (finished_at IS NOT NULL OR rolled_back_at IS NOT NULL)",
|
||||
(migration_name, started_at),
|
||||
)
|
||||
return bool(row)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_specific_migration(migration_name: str):
|
||||
|
|
@ -689,30 +714,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:
|
||||
|
|
@ -1072,6 +1110,11 @@ class ProxyExtrasDBManager:
|
|||
return match.group(1) if match else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _v2_failed_migration_started_at(stderr: str, migration_name: str) -> "str | None":
|
||||
match: Final = re.search(rf"`{re.escape(migration_name)}` migration started at ([^\r\n]+?) failed", stderr)
|
||||
return match.group(1) if match else None
|
||||
|
||||
@staticmethod
|
||||
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
|
|
@ -1100,8 +1143,11 @@ class ProxyExtrasDBManager:
|
|||
|
||||
if "P3009" in stderr:
|
||||
migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
|
||||
if migration_name:
|
||||
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
|
||||
started_at: Final = (
|
||||
ProxyExtrasDBManager._v2_failed_migration_started_at(stderr, migration_name) if migration_name else None
|
||||
)
|
||||
if migration_name and started_at:
|
||||
ledger_logs: Final = ProxyExtrasDBManager._failed_migration_logs(migration_name, started_at)
|
||||
if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs:
|
||||
logger.info(
|
||||
"Migration %s failed in a concurrent migrate deploy "
|
||||
|
|
@ -1110,6 +1156,14 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
return budget.spend()
|
||||
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
|
||||
logger.info(
|
||||
"Migration %s started at %s was already rolled back or completed by a concurrent "
|
||||
"migrate deploy, retrying",
|
||||
migration_name,
|
||||
started_at,
|
||||
)
|
||||
return budget.spend()
|
||||
raise RuntimeError(
|
||||
"Migration completion could not be verified. LiteLLM startup has stopped.\n\n"
|
||||
f"Prisma migration history (migration name and start time):\n{stderr}\n\n"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.104"
|
||||
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.104"
|
||||
version = "0.4.105"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
23
litellm-rust/Cargo.lock
generated
23
litellm-rust/Cargo.lock
generated
|
|
@ -4156,6 +4156,7 @@ dependencies = [
|
|||
"litellm-storage-clickhouse",
|
||||
"litellm-token-counter",
|
||||
"litellm-traces",
|
||||
"litellm-traces-clickhouse",
|
||||
"litellm-tracing",
|
||||
"prost",
|
||||
"pyo3",
|
||||
|
|
@ -4369,6 +4370,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4451,19 +4453,30 @@ dependencies = [
|
|||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"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",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
"litellm-storage-clickhouse",
|
||||
"litellm-traces",
|
||||
"moka",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ 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" }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -2,10 +2,8 @@ use std::collections::BTreeMap;
|
|||
|
||||
use litellm_host_python::{FromPythonCache, ToPythonCache};
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_traces::{
|
||||
Config, Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, 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},
|
||||
|
|
@ -31,38 +29,56 @@ 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::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
|
||||
Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => {
|
||||
PyValueError::new_err(error.to_string())
|
||||
}
|
||||
error => map_error(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_query_access_error(error: QueryAccessError) -> PyErr {
|
||||
match error {
|
||||
QueryAccessError::Storage(error) => map_sql_error(error),
|
||||
QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()),
|
||||
error => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass(frozen)]
|
||||
pub struct NativeTraceConfig {
|
||||
inner: Config,
|
||||
|
|
@ -105,7 +121,13 @@ impl NativeTraceStorage {
|
|||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces::ensure_schema(&client, &connection, &database, retention_days).await
|
||||
litellm_traces_clickhouse::ensure_schema(
|
||||
&client,
|
||||
&connection,
|
||||
&database,
|
||||
retention_days,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
|
|
@ -115,7 +137,7 @@ 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)?;
|
||||
|
|
@ -124,8 +146,14 @@ impl NativeTraceStorage {
|
|||
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,
|
||||
)
|
||||
|
|
@ -139,7 +167,9 @@ impl NativeTraceStorage {
|
|||
secret: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
if sql.trim().is_empty() {
|
||||
return Err(map_error(Error::EmptySql));
|
||||
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)?;
|
||||
|
|
@ -148,11 +178,9 @@ impl NativeTraceStorage {
|
|||
async move {
|
||||
let _permit = readers.acquire()?;
|
||||
let connection = readers.connection(&client, &scope, &secret).await?;
|
||||
litellm_traces::query_sql(&client, &connection, &sql)
|
||||
.await
|
||||
.map_err(QueryAccessError::Storage)
|
||||
litellm_traces_clickhouse::query_sql(&client, &connection, &sql).await
|
||||
},
|
||||
map_query_access_error,
|
||||
map_sql_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -169,32 +197,9 @@ impl NativeTraceStorage {
|
|||
async move {
|
||||
let _permit = readers.acquire()?;
|
||||
let connection = readers.connection(&client, &scope, &secret).await?;
|
||||
litellm_traces::query_help(&client, &connection)
|
||||
.await
|
||||
.map_err(QueryAccessError::Storage)
|
||||
litellm_traces_clickhouse::query_help(&client, &connection).await
|
||||
},
|
||||
map_query_access_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn lens_query<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
name: &str,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap<
|
||||
String,
|
||||
Parameter,
|
||||
>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
|
||||
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_read(&client, &connection, query.sql(), ¶meters).await
|
||||
},
|
||||
map_error,
|
||||
map_sql_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -207,13 +212,20 @@ impl NativeTraceStorage {
|
|||
Parameter,
|
||||
>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let query = ReadQuery::parse(query).map_err(map_error)?;
|
||||
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,
|
||||
)
|
||||
|
|
@ -229,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()?
|
||||
|
|
@ -320,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();
|
||||
|
|
@ -365,5 +427,6 @@ mod tests {
|
|||
|
||||
#[pyfunction]
|
||||
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
litellm_host_python::Pythonized(litellm_traces::NORMALIZED_FIELD_DEFINITIONS).into_pyobject(py)
|
||||
litellm_host_python::Pythonized(litellm_traces_clickhouse::NORMALIZED_FIELD_DEFINITIONS)
|
||||
.into_pyobject(py)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,4 +2,4 @@
|
|||
|
||||
`litellm-storage-clickhouse` exports `Storage`, a writer and bounded reader derived from one ClickHouse URL and database. It also exports bounded HTTP read and insert execution
|
||||
|
||||
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces` supplies those rules and uses this storage for both trace rows and spend rows
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
48
litellm-rust/crates/traces-clickhouse/query/list_traces.sql
Normal file
48
litellm-rust/crates/traces-clickhouse/query/list_traces.sql
Normal file
|
|
@ -0,0 +1,48 @@
|
|||
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))
|
||||
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
|
||||
24
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal file
24
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
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))
|
||||
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))
|
||||
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,9 @@ 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))
|
||||
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})
|
||||
|
|
@ -1,9 +1,10 @@
|
|||
SELECT request_id, response_id, team_id, api_key, spend,
|
||||
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 (empty({team_ids:Array(String)}) OR team_id IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR api_key = {api_key_hash:String})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND user = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, team_id))
|
||||
ORDER BY start_time DESC
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
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))
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
LIMIT 2
|
||||
|
|
@ -1,16 +1,18 @@
|
|||
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))
|
||||
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
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
use litellm_storage_clickhouse::{Error, Storage};
|
||||
use crate::Error;
|
||||
use litellm_storage_clickhouse::Storage;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Config {
|
||||
|
|
@ -8,7 +9,7 @@ pub struct Config {
|
|||
|
||||
impl Config {
|
||||
pub fn new(database: String, url: &str, retention_days: u32) -> Result<Self, Error> {
|
||||
crate::schema_statements(&database, retention_days)?;
|
||||
super::schema_statements(&database, retention_days)?;
|
||||
Ok(Self {
|
||||
storage: Storage::new(database, url)?,
|
||||
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",
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,8 +1,8 @@
|
|||
use askama::Template;
|
||||
use serde::Serialize;
|
||||
|
||||
use super::{AttributeCatalog, MetadataCatalog, TableSchema};
|
||||
use crate::{Error, NormalizedFieldDefinition};
|
||||
use super::{AttributeCatalog, Discovery, MetadataCatalog, TableSchema};
|
||||
use crate::{Error, NormalizedFieldDefinition, query_access::ReaderLimits};
|
||||
|
||||
#[derive(Template)]
|
||||
#[template(path = "query_help.jinja", escape = "none", blocks = [
|
||||
|
|
@ -33,6 +33,7 @@ pub(super) struct QueryGuide<'a> {
|
|||
pub normalized_fields: &'a [NormalizedFieldDefinition],
|
||||
pub metadata: &'a MetadataCatalog,
|
||||
pub attributes: &'a [AttributeCatalog],
|
||||
pub limits: &'a ReaderLimits,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
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"], "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": [], "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"], "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()),
|
||||
}
|
||||
}
|
||||
}
|
||||
245
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal file
245
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal file
|
|
@ -0,0 +1,245 @@
|
|||
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 = match table {
|
||||
TraceTable::OtelTraces | TraceTable::AgentTracesByKey => "TeamId",
|
||||
TraceTable::SpendLogs => "team_id",
|
||||
};
|
||||
match scope {
|
||||
QueryScope::All => "1".to_owned(),
|
||||
QueryScope::Owned { user_id, team_ids } => {
|
||||
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})")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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", "UserId = ''")]
|
||||
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "UserIds = ['']")]
|
||||
#[case::spend(TraceTable::SpendLogs, "team_id", "user = ''")]
|
||||
fn predicates_preserve_scope_and_escape_values(
|
||||
#[case] table: TraceTable,
|
||||
#[case] team: &str,
|
||||
#[case] user: &str,
|
||||
) {
|
||||
assert_eq!(predicate(&QueryScope::All, table), "1");
|
||||
assert_eq!(
|
||||
predicate(
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team'\\".into()]
|
||||
},
|
||||
table
|
||||
),
|
||||
format!("('' != '' AND {user}) OR ({team} IN ('team\\'\\\\'))")
|
||||
);
|
||||
}
|
||||
}
|
||||
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": [], "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,
|
||||
}
|
||||
|
|
@ -11,18 +11,18 @@ Normalized span fields
|
|||
{% endfor %}
|
||||
Observed LLM call metadata
|
||||
{{ metadata.scope }}
|
||||
Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }}
|
||||
{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }}
|
||||
{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows
|
||||
{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %}
|
||||
{% endfor %}{% endif %}
|
||||
{% 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 }}
|
||||
{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }}
|
||||
{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans
|
||||
{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }}
|
||||
{% endfor %}{% endif %}{% endfor %}
|
||||
{% 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 %}
|
||||
|
|
@ -35,7 +35,7 @@ Examples
|
|||
{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %}
|
||||
|
||||
{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %}
|
||||
{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND t.ApiKeyHash = s.api_key WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %}
|
||||
{% block 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 %}
|
||||
|
|
@ -44,9 +44,9 @@ Gotchas
|
|||
|
||||
{% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %}
|
||||
|
||||
{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %}
|
||||
{% block reader_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 authenticated team scope through row policies. Project-bound and teamless keys see only their own rows. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %}
|
||||
{% block 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. 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 %}
|
||||
|
||||
|
|
@ -58,7 +58,7 @@ Gotchas
|
|||
|
||||
{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %}
|
||||
|
||||
{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; aggregate spend separately{% endblock %}
|
||||
{% block 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 %}
|
||||
|
||||
|
|
@ -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::All, "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};
|
||||
|
|
@ -1,45 +1,14 @@
|
|||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::{
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery,
|
||||
encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
use rstest::rstest;
|
||||
mod support;
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
|
||||
|
||||
struct ClickHouseDatabase {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
url: String,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> TestResult<ClickHouseDatabase> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
);
|
||||
Ok(ClickHouseDatabase {
|
||||
_container: container,
|
||||
url,
|
||||
client: Client::no_redirect_for_test(),
|
||||
})
|
||||
}
|
||||
use support::{ClickHouseDatabase, TestResult, database};
|
||||
|
||||
async fn insert_rows(
|
||||
database: &ClickHouseDatabase,
|
||||
|
|
@ -115,7 +84,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
let span = serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "",
|
||||
"ServiceName": "proxy", "SpanName": "request", "Input": "hello world",
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1"},
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1", "litellm.user_id": "exporter-claim"},
|
||||
"SpanAttributes": {"gen_ai.response.id": "response-1", "gen_ai.usage.input_tokens": "12"}
|
||||
}))?;
|
||||
let spend = serde_json::from_value(serde_json::json!({
|
||||
|
|
@ -126,9 +95,29 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
insert_rows(&database, "otel_traces", vec![span]).await?;
|
||||
insert_rows(&database, "spend_logs", vec![spend]).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let detail =
|
||||
litellm_storage_clickhouse::fetch::<litellm_traces_clickhouse::query::named::SpanDetail>(
|
||||
&database.client,
|
||||
&reader,
|
||||
&litellm_traces_clickhouse::query::named::SpanDetailParams {
|
||||
access: litellm_traces_clickhouse::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-1".into()],
|
||||
},
|
||||
trace_id: "trace-1".into(),
|
||||
trace_ref: String::new(),
|
||||
span_id: "span-1".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(detail.len(), 1);
|
||||
assert_eq!(detail[0].input, "hello world");
|
||||
assert_eq!(detail[0].attributes["gen_ai.response.id"], "response-1");
|
||||
let list_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
|
|
@ -159,8 +148,9 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
"response_ids".into(),
|
||||
Parameter::Strings(vec!["response-1".into()]),
|
||||
),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
|
|
@ -182,7 +172,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
assert_eq!(matched["data"][0]["spend"], 0.125);
|
||||
let body = read_json(
|
||||
&database,
|
||||
"SELECT o.TeamId, o.ApiKeyHash, o.ObservationType, o.InputPreview, s.spend, \
|
||||
"SELECT o.TeamId, o.ApiKeyHash, o.UserId, o.ObservationType, o.InputPreview, s.spend, \
|
||||
toString(toUnixTimestamp64Nano(o.Timestamp)) AS timestamp_ns, \
|
||||
toString(toUnixTimestamp64Milli(s.start_time)) AS start_ms \
|
||||
FROM trace_test.otel_traces o JOIN trace_test.spend_logs s \
|
||||
|
|
@ -192,7 +182,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
assert_eq!(
|
||||
body["data"],
|
||||
serde_json::json!([{
|
||||
"TeamId": "team-1", "ApiKeyHash": "hash-1", "ObservationType": "agent",
|
||||
"TeamId": "team-1", "ApiKeyHash": "hash-1", "UserId": "", "ObservationType": "agent",
|
||||
"InputPreview": "hello world", "spend": 0.125,
|
||||
"timestamp_ns": timestamp.to_string(), "start_ms": (timestamp / 1_000_000).to_string()
|
||||
}])
|
||||
|
|
@ -269,7 +259,7 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
|
|||
]);
|
||||
|
||||
assert!(matches!(
|
||||
litellm_traces::insert_rows(
|
||||
litellm_traces_clickhouse::insert_rows(
|
||||
&database.client,
|
||||
&writer,
|
||||
"trace_test",
|
||||
|
|
@ -277,7 +267,9 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
|
|||
vec![row]
|
||||
)
|
||||
.await,
|
||||
Err(Error::InsertFailed(_))
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::InsertFailed(_)
|
||||
))
|
||||
));
|
||||
assert_eq!(table_rows(&database, "otel_traces").await?, 0);
|
||||
Ok(())
|
||||
|
|
@ -297,7 +289,7 @@ async fn retried_trace_insert_does_not_inflate_rollup(
|
|||
"TeamId": "team-1", "ApiKeyHash": "key-1", "SpanName": "root", "InputTokens": 7
|
||||
}))?;
|
||||
for _ in 0..2 {
|
||||
litellm_traces::insert_rows(
|
||||
litellm_traces_clickhouse::insert_rows(
|
||||
&database.client,
|
||||
&writer,
|
||||
"trace_test",
|
||||
|
|
@ -362,6 +354,194 @@ async fn keyed_rollup_keeps_same_trace_ids_separate_by_api_key(
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn listed_agent_names_preserve_scope_and_cursor(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
for (team, key, trace, agent, span, parent, framework) in [
|
||||
(
|
||||
"alpha",
|
||||
"one",
|
||||
"shared",
|
||||
"research_agent",
|
||||
"root",
|
||||
"",
|
||||
"claude-code",
|
||||
),
|
||||
(
|
||||
"alpha",
|
||||
"one",
|
||||
"shared",
|
||||
"reviewer",
|
||||
"child",
|
||||
"root",
|
||||
"claude-agent-sdk",
|
||||
),
|
||||
(
|
||||
"alpha",
|
||||
"one",
|
||||
"shared",
|
||||
"reviewer",
|
||||
"repeated",
|
||||
"root",
|
||||
"claude-agent-sdk",
|
||||
),
|
||||
("alpha", "one", "shared", "", "unnamed", "root", ""),
|
||||
("alpha", "one", "second", "support_agent", "root", "", ""),
|
||||
(
|
||||
"alpha",
|
||||
"two",
|
||||
"shared",
|
||||
"private_agent",
|
||||
"root",
|
||||
"",
|
||||
"private-sdk",
|
||||
),
|
||||
(
|
||||
"beta",
|
||||
"other",
|
||||
"shared",
|
||||
"other_agent",
|
||||
"root",
|
||||
"",
|
||||
"other-sdk",
|
||||
),
|
||||
] {
|
||||
insert_rows(
|
||||
&database,
|
||||
"otel_traces",
|
||||
vec![serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
|
||||
"ServiceName": "shared-app", "SpanName": span, "AgentName": agent,
|
||||
"UserId": if key == "one" { "owner" } else { "other" },
|
||||
"Framework": framework, "ObservationType": "agent",
|
||||
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key}
|
||||
}))?],
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let historical_rows = (0..5000)
|
||||
.map(|index| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp - 86_400_000_000_000_i64,
|
||||
"TraceId": "shared", "SpanId": format!("historical-{index}"),
|
||||
"ParentSpanId": "", "SpanName": "historical", "AgentName": "private_agent",
|
||||
"ObservationType": "agent", "ServiceName": "shared-app",
|
||||
"ResourceAttributes": {"litellm.team_id": "alpha", "litellm.api_key_hash": "history"}
|
||||
}))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", historical_rows).await?;
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text("owner".into())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
),
|
||||
(
|
||||
"end_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 + 1000),
|
||||
),
|
||||
("cursor_ms".into(), Parameter::Integer(0)),
|
||||
("cursor_trace_id".into(), Parameter::Text(String::new())),
|
||||
("limit".into(), Parameter::Integer(1)),
|
||||
]);
|
||||
let first: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::ListTraces,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
let cursor = first["data"][0]["trace_ref"]
|
||||
.as_str()
|
||||
.ok_or("missing cursor")?;
|
||||
let next_parameters = parameters
|
||||
.into_iter()
|
||||
.chain([
|
||||
(
|
||||
"cursor_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000),
|
||||
),
|
||||
("cursor_trace_id".into(), Parameter::Text(cursor.into())),
|
||||
])
|
||||
.collect();
|
||||
let second: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::ListTraces,
|
||||
&next_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(
|
||||
first["data"].as_array().ok_or("missing first page")?.len(),
|
||||
1
|
||||
);
|
||||
assert_eq!(
|
||||
second["data"]
|
||||
.as_array()
|
||||
.ok_or("missing second page")?
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
assert_ne!(first["data"][0]["trace_id"], second["data"][0]["trace_id"]);
|
||||
let names = [&first["data"][0], &second["data"][0]]
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row["trace_id"].as_str().unwrap(),
|
||||
row["agent_names"].clone(),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
assert_eq!(
|
||||
names["shared"],
|
||||
serde_json::json!(["research_agent", "reviewer"])
|
||||
);
|
||||
assert_eq!(names["second"], serde_json::json!(["support_agent"]));
|
||||
let frameworks = [&first["data"][0], &second["data"][0]]
|
||||
.into_iter()
|
||||
.map(|row| (row["trace_id"].as_str().unwrap(), row["frameworks"].clone()))
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
assert_eq!(
|
||||
frameworks["shared"],
|
||||
serde_json::json!(["claude-agent-sdk", "claude-code"])
|
||||
);
|
||||
assert_eq!(frameworks["second"], serde_json::json!([]));
|
||||
let counts = [&first["data"][0], &second["data"][0]]
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
(
|
||||
row["trace_id"].as_str().unwrap(),
|
||||
row["agent_count"].as_u64(),
|
||||
)
|
||||
})
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
assert_eq!(counts["shared"], Some(3));
|
||||
assert_eq!(counts["second"], Some(1));
|
||||
for page in [&first, &second] {
|
||||
assert!(
|
||||
page["statistics"]["rows_read"]
|
||||
.as_u64()
|
||||
.ok_or("missing read statistics")?
|
||||
< 5000
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
||||
|
|
@ -376,6 +556,7 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
|||
let root = serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": day_start - 1_000_000_000, "TraceId": "cross-day", "SpanId": "span-root",
|
||||
"ParentSpanId": "", "ServiceName": "proxy", "SpanName": "root", "Input": "root input",
|
||||
"AgentName": "lead", "ObservationType": "agent",
|
||||
"StatusCode": "STATUS_CODE_ERROR",
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1"}
|
||||
}))?;
|
||||
|
|
@ -383,6 +564,7 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
|||
let child = serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": day_start + 1_000_000_000, "TraceId": "cross-day", "SpanId": "span-child",
|
||||
"ParentSpanId": "span-root", "ServiceName": "proxy", "SpanName": "child",
|
||||
"AgentName": "researcher", "ObservationType": "agent",
|
||||
"StatusCode": "STATUS_CODE_UNSET",
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1"}
|
||||
}))?;
|
||||
|
|
@ -406,6 +588,34 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
|||
"RootStatus": "STATUS_CODE_ERROR", "SpanCount": 2
|
||||
}])
|
||||
);
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(day_start / 1_000_000 - 2000),
|
||||
),
|
||||
("end_ms".into(), Parameter::Integer(day_start / 1_000_000)),
|
||||
("cursor_ms".into(), Parameter::Integer(0)),
|
||||
("cursor_trace_id".into(), Parameter::Text(String::new())),
|
||||
("limit".into(), Parameter::Integer(10)),
|
||||
]);
|
||||
let listed: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::ListTraces,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(
|
||||
listed["data"][0]["agent_names"],
|
||||
serde_json::json!(["lead", "researcher"])
|
||||
);
|
||||
assert_eq!(listed["data"][0]["agent_count"], 2);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -553,7 +763,10 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
|
|||
)
|
||||
.await;
|
||||
server.abort();
|
||||
assert!(matches!(result, Ok(Err(Error::Transport))), "{result:?}");
|
||||
assert!(
|
||||
matches!(result, Ok(Err(Error::SchemaTransport))),
|
||||
"{result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -570,7 +783,7 @@ fn schema_rejects_invalid_configuration(#[case] database: &str, #[case] retentio
|
|||
async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::{LensQuery, Parameter};
|
||||
use litellm_traces_clickhouse::{Parameter, ReadQuery};
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -578,7 +791,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
for (key, text) in [("one", "timeout"), ("two", "success")] {
|
||||
insert_rows(&database, "otel_traces", vec![serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "shared", "SpanId": "root", "ParentSpanId": "",
|
||||
"ServiceName": "review", "SpanName": "release", "Input": text,
|
||||
"ServiceName": "review", "SpanName": "release", "Input": text, "UserId": key,
|
||||
"ResourceAttributes": {"litellm.team_id": "team", "litellm.api_key_hash": key, "swarm": "release"}
|
||||
}))?]).await?;
|
||||
}
|
||||
|
|
@ -616,10 +829,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -627,6 +840,43 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
let rows = sample["data"].as_array().expect("sample rows");
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_ne!(rows[0]["trace_ref"], rows[1]["trace_ref"]);
|
||||
let identity_params = BTreeMap::from([
|
||||
("trace_id".into(), Parameter::Text("shared".into())),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team".into()])),
|
||||
]);
|
||||
let identities: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::TraceIdentity,
|
||||
&identity_params,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(identities["data"].as_array().map(Vec::len), Some(2));
|
||||
let user_params = identity_params
|
||||
.into_iter()
|
||||
.chain([
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("user_id".into(), Parameter::Text("one".into())),
|
||||
])
|
||||
.collect();
|
||||
let identity: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::TraceIdentity,
|
||||
&user_params,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(identity["data"].as_array().map(Vec::len), Some(1));
|
||||
assert!(
|
||||
rows.iter()
|
||||
.any(|row| row["trace_ref"] == identity["data"][0]["trace_ref"])
|
||||
);
|
||||
let first_ref = rows[0]["trace_ref"].as_str().expect("reference");
|
||||
let read_parameters: BTreeMap<_, _> = parameters
|
||||
.into_iter()
|
||||
|
|
@ -640,10 +890,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
])
|
||||
.collect();
|
||||
let content: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
&read_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -660,10 +910,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
.chain([("quote".into(), Parameter::Text(opposite.into()))])
|
||||
.collect();
|
||||
let evidence: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Evidence.sql(),
|
||||
ReadQuery::Evidence,
|
||||
&evidence_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -677,7 +927,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
async fn lens_request_sample_does_not_trust_caller_tags(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::{LensQuery, Parameter};
|
||||
use litellm_traces_clickhouse::{Parameter, ReadQuery};
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -712,10 +962,10 @@ async fn lens_request_sample_does_not_trust_caller_tags(
|
|||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -742,7 +992,7 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
|
|||
#[case] page_size: usize,
|
||||
#[case] changing: bool,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database?;
|
||||
ensure_schema(
|
||||
&database.client,
|
||||
|
|
@ -781,10 +1031,10 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
|
|||
("selected_team".into(), Parameter::Text(String::new())),
|
||||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -822,7 +1072,7 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] input_length: usize,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database?;
|
||||
ensure_schema(
|
||||
&database.client,
|
||||
|
|
@ -846,10 +1096,10 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
("cursor".into(), Parameter::Text(String::new())),
|
||||
("offset".into(), Parameter::Integer(1)),
|
||||
]);
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -868,10 +1118,10 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
let mut recovered = String::new();
|
||||
for offset in (2..original.len() + 2).step_by(8000) {
|
||||
parameters.insert("offset".into(), Parameter::Integer(offset as i64));
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -912,8 +1162,9 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
"trace_id".into(),
|
||||
Parameter::Text("diagnostic-trace".into()),
|
||||
),
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
]);
|
||||
let body = execute_named_read(
|
||||
|
|
@ -959,10 +1210,8 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
}
|
||||
}
|
||||
assert_eq!(recovered, message);
|
||||
parameters.insert(
|
||||
"api_key_hash".into(),
|
||||
Parameter::Text("unrelated-key".into()),
|
||||
);
|
||||
parameters.insert("all_teams".into(), Parameter::Integer(0));
|
||||
parameters.insert("user_id".into(), Parameter::Text("unrelated-user".into()));
|
||||
let denied =
|
||||
execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters).await?;
|
||||
assert_eq!(
|
||||
|
|
@ -1004,8 +1253,9 @@ async fn duplicate_span_preview_matches_diagnostic(
|
|||
let parameters = BTreeMap::from([
|
||||
("trace_id".into(), Parameter::Text("duplicate-trace".into())),
|
||||
("span_id".into(), Parameter::Text("duplicate-span".into())),
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
("error_version".into(), Parameter::Text(String::new())),
|
||||
("error_offset".into(), Parameter::Integer(0)),
|
||||
|
|
@ -1042,7 +1292,7 @@ fn schema_includes_every_migration_file() -> TestResult {
|
|||
async fn lens_agent_discovery_and_selection_preserve_scope(
|
||||
#[future] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database.await?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -1073,10 +1323,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
("key_hash".into(), Parameter::Text("one".into())),
|
||||
]);
|
||||
let agents: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Agents.sql(),
|
||||
ReadQuery::Agents,
|
||||
&scope_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1117,10 +1367,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
])
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1129,10 +1379,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
assert_eq!(sample["data"][0]["trace_id"], "research");
|
||||
assert_eq!(sample["data"][0]["span_count"], 2);
|
||||
let available: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Availability.sql(),
|
||||
ReadQuery::Availability,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1207,8 +1457,9 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
)
|
||||
.await?;
|
||||
}
|
||||
let help: serde_json::Value =
|
||||
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
|
||||
let help = serde_json::to_value(
|
||||
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
let keys: std::collections::BTreeSet<_> = help
|
||||
.as_object()
|
||||
.ok_or("missing help object")?
|
||||
|
|
@ -1318,7 +1569,8 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
"missing plain-text expression: {expression}"
|
||||
);
|
||||
let sql = format!("SELECT {expression} AS value FROM spend_logs FINAL");
|
||||
let body = litellm_traces::query_sql(&database.client, &reader, &sql).await?;
|
||||
let body =
|
||||
litellm_traces_clickhouse::query_sql(&database.client, &reader, &sql).await?;
|
||||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_ne!(values["data"][0]["value"], "");
|
||||
}
|
||||
|
|
@ -1336,7 +1588,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
.collect::<std::collections::BTreeSet<_>>(),
|
||||
std::collections::BTreeSet::from(["name", "sql"])
|
||||
);
|
||||
let body = litellm_traces::query_sql(&database.client, &reader, sql).await?;
|
||||
let body = litellm_traces_clickhouse::query_sql(&database.client, &reader, sql).await?;
|
||||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(
|
||||
values["data"].as_array().ok_or("missing data")?.is_empty(),
|
||||
|
|
@ -1392,8 +1644,9 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
}))).collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", spans).await?;
|
||||
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
|
||||
let help: serde_json::Value =
|
||||
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
|
||||
let help = serde_json::to_value(
|
||||
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3);
|
||||
assert!(!help["examples"].as_array().ok_or("examples")?.is_empty());
|
||||
assert_eq!(
|
||||
|
|
@ -1427,3 +1680,274 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn field_definitions_match_serialized_normalized_span() {
|
||||
use litellm_traces::decode_otlp;
|
||||
use std::collections::BTreeSet;
|
||||
let spans = decode_otlp(
|
||||
br#"{"resourceSpans":[{"scopeSpans":[{"spans":[{"traceId":"11111111111111111111111111111111","spanId":"2222222222222222","name":"root"}]}]}]}"#,
|
||||
Some("application/json"),
|
||||
)
|
||||
.expect("valid OTLP");
|
||||
let fields = &spans[0].normalized;
|
||||
let serialized = serde_json::to_value(fields).expect("serializable fields");
|
||||
let keys: BTreeSet<_> = serialized
|
||||
.as_object()
|
||||
.expect("field object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(|field| field.name)
|
||||
.collect();
|
||||
assert_eq!(keys, mapped);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user("owner", vec![], None, vec!["own"])]
|
||||
#[case::own_user_and_permitted_team("owner", vec!["permitted"], None, vec!["own", "team"])]
|
||||
#[case::no_identity("", vec![], None, vec![])]
|
||||
#[case::legacy_key_without_identity("", vec![], Some("request-key"), vec![])]
|
||||
#[tokio::test]
|
||||
async fn named_and_sql_readers_share_request_log_visibility(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] legacy_key: Option<&str>,
|
||||
#[case] expected: Vec<&str>,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{
|
||||
ReadAccessParams, SpendByResponseIds, SpendByResponseIdsParams,
|
||||
};
|
||||
use litellm_traces_clickhouse::{QueryReaders, QueryScope};
|
||||
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [("own", "unpermitted", "owner", "request-key"), ("team", "permitted", "other", "other-key"), ("foreign", "foreign", "other", "foreign-key")]
|
||||
.into_iter()
|
||||
.map(|(id, team, owner, api_key)| serde_json::from_value(serde_json::json!({
|
||||
"request_id": id, "response_id": "shared-response", "team_id": team, "user": owner,
|
||||
"api_key": api_key, "spend": 0.25, "start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000,
|
||||
})))
|
||||
.collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
|
||||
insert_rows(&database, "spend_logs", rows).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let params =
|
||||
SpendByResponseIdsParams::from(litellm_traces::query::named::SpendByResponseIdsParams {
|
||||
access: serde_json::from_value::<ReadAccessParams>(serde_json::json!({
|
||||
"all_teams": 0, "user_id": user, "team_ids": teams,
|
||||
"api_key_hash": legacy_key.unwrap_or_default(),
|
||||
}))?,
|
||||
response_ids: vec!["shared-response".into()],
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
});
|
||||
let spend =
|
||||
litellm_storage_clickhouse::fetch::<SpendByResponseIds>(&database.client, &reader, ¶ms)
|
||||
.await?;
|
||||
let actual: std::collections::BTreeSet<_> =
|
||||
spend.iter().map(|row| row.0.request_id.as_str()).collect();
|
||||
let expected: std::collections::BTreeSet<_> = expected.into_iter().collect();
|
||||
assert_eq!(actual, expected);
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
};
|
||||
if user.is_empty() && scope.validate().is_err() {
|
||||
assert!(
|
||||
QueryReaders::new(writer, "trace_test".into())
|
||||
.connection(&database.client, &scope, "secret")
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
let scoped = QueryReaders::new(writer, "trace_test".into())
|
||||
.connection(&database.client, &scope, "secret")
|
||||
.await?;
|
||||
let result: serde_json::Value = serde_json::from_str(
|
||||
&litellm_traces_clickhouse::query_sql(
|
||||
&database.client,
|
||||
&scoped,
|
||||
"SELECT request_id FROM spend_logs FINAL ORDER BY request_id",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(
|
||||
result["data"],
|
||||
serde_json::json!(
|
||||
expected
|
||||
.into_iter()
|
||||
.map(|id| serde_json::json!({"request_id": id}))
|
||||
.collect::<Vec<_>>()
|
||||
)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_historical_rows(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let initial_mutations = mutation_rows(&database).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [("complete", "llm", "response"), ("complete", "llm", "response"), ("complete", "agent", ""), ("missing", "llm", "response"), ("missing", "llm", ""), ("missing", "agent", "extra-id"), ("mixed", "llm", "mine"), ("mixed", "llm", "other")]
|
||||
.into_iter().enumerate().map(|(index, (trace, kind, id))| serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": trace, "SpanId": index.to_string(), "TeamId": "team", "ApiKeyHash": "export",
|
||||
"UserId": if id == "other" { "other" } else { "owner" }, "ObservationType": kind, "LiteLLMRequestId": id,
|
||||
}))).collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", rows).await?;
|
||||
execute_write(&database, &format!(
|
||||
"INSERT INTO trace_test.agent_traces_by_key (TeamId, ApiKeyHash, TraceId, StartTs, EndTs, LlmCount, RequestIds) \
|
||||
VALUES ('team', 'export', 'historical', fromUnixTimestamp64Nano({timestamp}), fromUnixTimestamp64Nano({timestamp}), 2, ['response', 'non-llm-id'])"
|
||||
)).await?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let params = litellm_traces_clickhouse::query::named::ListTracesParams::from(
|
||||
litellm_traces::query::named::ListTracesParams {
|
||||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: "".into(),
|
||||
team_ids: vec!["team".into()],
|
||||
},
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
cursor_ms: 0,
|
||||
cursor_trace_id: "".into(),
|
||||
limit: 10,
|
||||
},
|
||||
);
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let listed = litellm_storage_clickhouse::fetch::<
|
||||
litellm_traces_clickhouse::query::named::ListTraces,
|
||||
>(&database.client, &reader, ¶ms)
|
||||
.await?;
|
||||
assert_eq!(listed.len(), 4);
|
||||
let owned_params = litellm_traces_clickhouse::query::named::ListTracesParams::from(
|
||||
litellm_traces::query::named::ListTracesParams {
|
||||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
user_id: "owner".into(),
|
||||
team_ids: vec![],
|
||||
all_teams: 0,
|
||||
},
|
||||
..params.0
|
||||
},
|
||||
);
|
||||
let owned = litellm_storage_clickhouse::fetch::<
|
||||
litellm_traces_clickhouse::query::named::ListTraces,
|
||||
>(&database.client, &reader, &owned_params)
|
||||
.await?;
|
||||
assert_eq!(owned.len(), 2);
|
||||
assert!(
|
||||
owned
|
||||
.iter()
|
||||
.all(|row| ["complete", "missing"].contains(&row.0.trace_id.as_str()))
|
||||
);
|
||||
for row in listed {
|
||||
match row.0.trace_id.as_str() {
|
||||
"complete" => {
|
||||
assert_eq!(row.0.user_id, "owner");
|
||||
assert_eq!(row.0.request_ids, ["response"]);
|
||||
assert_eq!(row.0.llm_calls, 2);
|
||||
}
|
||||
"missing" | "historical" => assert!(row.0.request_ids.iter().any(String::is_empty)),
|
||||
"mixed" => assert!(row.0.user_id.is_empty()),
|
||||
id => panic!("unexpected trace {id}"),
|
||||
}
|
||||
}
|
||||
assert_eq!(mutation_rows(&database).await?, initial_mutations);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::admin(1, "", vec![], "own answer")]
|
||||
#[case::user(0, "owner", vec![], "own answer")]
|
||||
#[case::team(0, "", vec!["alpha"], "own answer")]
|
||||
#[case::no_identity(0, "", vec![], "")]
|
||||
#[tokio::test]
|
||||
async fn agent_final_answer_preserves_visibility_and_trace_ownership(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] all_teams: u8,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] expected: &str,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{ReadAccessParams, SpanDetail, SpanDetailParams};
|
||||
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [
|
||||
("alpha", "one", "owner", "root", "", "agent", ""),
|
||||
(
|
||||
"alpha",
|
||||
"one",
|
||||
"owner",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"own answer",
|
||||
),
|
||||
(
|
||||
"alpha",
|
||||
"two",
|
||||
"other",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"other key answer",
|
||||
),
|
||||
(
|
||||
"beta",
|
||||
"one",
|
||||
"other",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"other team answer",
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, (team, key, user, span, parent, kind, output))| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp + index as i64, "TraceId": "shared", "SpanId": span,
|
||||
"ParentSpanId": parent, "TeamId": team, "ApiKeyHash": key, "UserId": user,
|
||||
"ObservationType": kind, "Input": "prompt", "Output": output,
|
||||
}))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", rows).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let details = litellm_storage_clickhouse::fetch::<SpanDetail>(
|
||||
&database.client,
|
||||
&reader,
|
||||
&SpanDetailParams {
|
||||
access: ReadAccessParams {
|
||||
all_teams,
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
},
|
||||
trace_id: "shared".into(),
|
||||
trace_ref: String::new(),
|
||||
span_id: "root".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if expected.is_empty() {
|
||||
assert!(details.is_empty());
|
||||
} else {
|
||||
assert_eq!(details.len(), 1);
|
||||
assert_eq!(details[0].input, "prompt");
|
||||
assert_eq!(details[0].output, expected);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
252
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal file
252
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
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,
|
||||
OtherTeam,
|
||||
}
|
||||
|
||||
impl ScopeCase {
|
||||
fn scope(self) -> QueryScope {
|
||||
match self {
|
||||
Self::Admin => QueryScope::All,
|
||||
Self::Team => QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".into()],
|
||||
},
|
||||
Self::OtherTeam => QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["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::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::All, "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::All, "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::All, "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
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
{
|
||||
"all_teams": 1,
|
||||
"user_id": "",
|
||||
"team_ids": [
|
||||
"team-a",
|
||||
"team-b"
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
},
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-alt",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "alternate",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
},
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "other-team",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
},
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-alt",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "alternate",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
}
|
||||
],
|
||||
"other_team": [
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "other-team",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
SELECT TeamId AS team, ApiKeyHash AS api_key, TraceId AS trace_id,
|
||||
ifNull(any(RootName), '') AS name,
|
||||
toUInt32(sum(SpanCount)) AS spans,
|
||||
toUInt32(sum(LlmCount)) AS llm_calls,
|
||||
toUInt32(sum(ErrorCount)) AS errors,
|
||||
toUInt32(sum(InputTokens)) AS input_tokens,
|
||||
toUInt32(sum(OutputTokens)) AS output_tokens
|
||||
FROM agent_traces_by_key
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
ORDER BY team, api_key, trace_id
|
||||
153
litellm-rust/crates/traces-clickhouse/tests/queries/support.rs
Normal file
153
litellm-rust/crates/traces-clickhouse/tests/queries/support.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_traces::{DecodedSpan, decode_otlp};
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, InsertTable, QueryReaders, ensure_schema, insert_rows,
|
||||
};
|
||||
use rstest::fixture;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::support::{ClickHouseDatabase, TestResult, database};
|
||||
|
||||
pub const DATABASE: &str = "trace_test";
|
||||
|
||||
pub struct SeededDatabase {
|
||||
pub database: ClickHouseDatabase,
|
||||
pub readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn migrated_database(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult<SeededDatabase> {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, DATABASE, 7).await?;
|
||||
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
|
||||
database
|
||||
.client
|
||||
.post(writer.url().clone())
|
||||
.body(format!("ALTER TABLE {DATABASE}.{table} REMOVE TTL"))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
database
|
||||
.client
|
||||
.post(writer.url().clone())
|
||||
.body(format!("SYSTEM STOP MERGES {DATABASE}.{table}"))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
}
|
||||
let readers = QueryReaders::new(writer, DATABASE.to_owned());
|
||||
Ok(SeededDatabase { database, readers })
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn seeded_database(
|
||||
#[future(awt)] migrated_database: TestResult<SeededDatabase>,
|
||||
) -> TestResult<SeededDatabase> {
|
||||
let fixture = migrated_database?;
|
||||
let writer = Connection::writer(&fixture.database.url)?;
|
||||
for (contents, team, key) in [
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_root.json").as_slice(),
|
||||
"team-a",
|
||||
"key-a",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_children.json").as_slice(),
|
||||
"team-a",
|
||||
"key-a",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_alternate.json").as_slice(),
|
||||
"team-a",
|
||||
"key-alt",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_other_team.json").as_slice(),
|
||||
"team-b",
|
||||
"key-b",
|
||||
),
|
||||
] {
|
||||
insert_export(&fixture, contents, team, key).await?;
|
||||
}
|
||||
let spend_rows = include_str!("../fixtures/spend_logs.jsonl")
|
||||
.lines()
|
||||
.map(serde_json::from_str::<BTreeMap<String, Value>>)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(
|
||||
&fixture.database.client,
|
||||
&writer,
|
||||
DATABASE,
|
||||
InsertTable::SpendLogs,
|
||||
spend_rows,
|
||||
)
|
||||
.await?;
|
||||
Ok(fixture)
|
||||
}
|
||||
|
||||
pub async fn insert_export(
|
||||
fixture: &SeededDatabase,
|
||||
contents: &[u8],
|
||||
team: &str,
|
||||
key: &str,
|
||||
) -> TestResult<Vec<DecodedSpan>> {
|
||||
let spans = decode_otlp(contents, Some("application/json"))?;
|
||||
let writer = Connection::writer(&fixture.database.url)?;
|
||||
let rows = spans.iter().map(|span| span_row(span, team, key)).collect();
|
||||
insert_rows(
|
||||
&fixture.database.client,
|
||||
&writer,
|
||||
DATABASE,
|
||||
InsertTable::OtelTraces,
|
||||
rows,
|
||||
)
|
||||
.await?;
|
||||
Ok(spans)
|
||||
}
|
||||
|
||||
fn span_row(span: &DecodedSpan, team: &str, key: &str) -> BTreeMap<String, Value> {
|
||||
BTreeMap::from([
|
||||
("Timestamp".into(), json!(span.start_ns)),
|
||||
("TraceId".into(), json!(span.trace_id)),
|
||||
("SpanId".into(), json!(span.span_id)),
|
||||
("ParentSpanId".into(), json!(span.parent_span_id)),
|
||||
("TraceState".into(), json!(span.trace_state)),
|
||||
("SpanName".into(), json!(span.name)),
|
||||
("SpanKind".into(), json!(span.kind)),
|
||||
(
|
||||
"ServiceName".into(),
|
||||
json!(
|
||||
span.resource_attributes
|
||||
.get("service.name")
|
||||
.map(String::as_str)
|
||||
.unwrap_or_default()
|
||||
),
|
||||
),
|
||||
("ResourceAttributes".into(), json!(span.resource_attributes)),
|
||||
("ScopeName".into(), json!(span.scope_name)),
|
||||
("ScopeVersion".into(), json!(span.scope_version)),
|
||||
("SpanAttributes".into(), json!(span.attributes)),
|
||||
("Duration".into(), json!(span.end_ns - span.start_ns)),
|
||||
("StatusCode".into(), json!(span.status_code)),
|
||||
("StatusMessage".into(), json!(span.status_message)),
|
||||
("TeamId".into(), json!(team)),
|
||||
("ApiKeyHash".into(), json!(key)),
|
||||
(
|
||||
"ObservationType".into(),
|
||||
json!(span.normalized.observation_type),
|
||||
),
|
||||
("AgentName".into(), json!(span.normalized.agent_name)),
|
||||
("Model".into(), json!(span.normalized.model)),
|
||||
(
|
||||
"LiteLLMRequestId".into(),
|
||||
json!(span.normalized.litellm_request_id),
|
||||
),
|
||||
("InputTokens".into(), json!(span.normalized.input_tokens)),
|
||||
("OutputTokens".into(), json!(span.normalized.output_tokens)),
|
||||
("Input".into(), json!(span.normalized.input)),
|
||||
("Output".into(), json!(span.normalized.output)),
|
||||
])
|
||||
}
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
},
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.25
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
}
|
||||
],
|
||||
"other_team": [
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.25
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
SELECT o.TeamId AS team, o.ApiKeyHash AS api_key, o.TraceId AS trace_id,
|
||||
sum(s.spend) AS spend
|
||||
FROM otel_traces AS o
|
||||
INNER JOIN (SELECT * FROM spend_logs FINAL) AS s
|
||||
ON o.TeamId = s.team_id
|
||||
AND o.ApiKeyHash = s.api_key
|
||||
AND o.LiteLLMRequestId = s.response_id
|
||||
GROUP BY o.TeamId, o.ApiKeyHash, o.TraceId
|
||||
ORDER BY team, api_key, trace_id
|
||||
269
litellm-rust/crates/traces-clickhouse/tests/query_access.rs
Normal file
269
litellm-rust/crates/traces-clickhouse/tests/query_access.rs
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
mod support;
|
||||
|
||||
use support::{ClickHouseDatabase, database as start_database};
|
||||
|
||||
struct Database {
|
||||
_database: ClickHouseDatabase,
|
||||
client: Client,
|
||||
writer: Connection,
|
||||
readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
||||
let instance = start_database().await?;
|
||||
let url = instance.url.clone();
|
||||
let client = instance.client.clone();
|
||||
let writer = Connection::parse(&url)?;
|
||||
ensure_schema(&client, &writer, "trace_test", 7).await?;
|
||||
for sql in [
|
||||
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes, UserId) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a'), 'owner'), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a'), 'other'), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'), 'owner'), ('team-c', 'key-a1', 'shared-trace', 'same-key-foreign', now(), map('visible', 'foreign'), 'other'), ('', 'key-teamless', 'shared-trace', 'teamless', now(), map('visible', 'teamless'), ''), ('', 'key-other', 'shared-trace', 'other-teamless', now(), map('visible', 'other'), '')",
|
||||
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata, user) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}', 'owner'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}', 'other'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}', 'owner'), ('team-c', 'key-a1', 'same-key-foreign', now(), now(), '{}', 'other'), ('', 'key-teamless', 'teamless', now(), now(), '{}', ''), ('', 'key-other', 'other-teamless', now(), now(), '{}', '')",
|
||||
"CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory",
|
||||
"INSERT INTO trace_test.private_data VALUES ('hidden')",
|
||||
] {
|
||||
let response = client.post(writer.url().clone()).body(sql).send().await?;
|
||||
assert!(response.status().is_success(), "{}", response.text().await?);
|
||||
}
|
||||
let readers = QueryReaders::new(writer.clone(), "trace_test".to_owned());
|
||||
Ok(Database {
|
||||
_database: instance,
|
||||
client,
|
||||
writer,
|
||||
readers,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user(QueryScope::Owned { user_id: "owner".into(), team_ids: vec![] }, vec!["a1", "b"])]
|
||||
#[case::own_user_and_permitted_team(QueryScope::Owned { user_id: "owner".into(), team_ids: vec!["team-a".into()] }, vec!["a1", "a2", "b"])]
|
||||
#[case::quoted_user(QueryScope::Owned { user_id: "owner' OR 1=1 --".into(), team_ids: vec![] }, vec![])]
|
||||
#[case::team(QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a".to_owned() ] }, vec!["a1", "a2"])]
|
||||
#[case::admin(QueryScope::All, vec!["a1", "a2", "b", "other-teamless", "same-key-foreign", "teamless"])]
|
||||
#[case::quoted_team(QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a' OR 1=1 --\\".to_owned() ] }, vec![])]
|
||||
#[tokio::test]
|
||||
async fn queries_and_help_are_scoped_by_the_database(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
#[case] scope: QueryScope,
|
||||
#[case] expected: Vec<&str>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
let queries = [
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
"SELECT SpanId AS id FROM trace_test.otel_traces WHERE 1 = 1 ORDER BY id",
|
||||
"SELECT SpanId AS id FROM merge('trace_test', '^otel_traces$') ORDER BY id",
|
||||
"WITH source AS (SELECT * FROM trace_test.otel_traces) SELECT SpanId AS id FROM source ORDER BY id",
|
||||
"SELECT id FROM (SELECT SpanId AS id FROM otel_traces UNION DISTINCT SELECT SpanId AS id FROM trace_test.otel_traces) ORDER BY id",
|
||||
"SELECT t.SpanId AS id FROM otel_traces t INNER JOIN spend_logs s ON t.SpanId = s.request_id ORDER BY id",
|
||||
"SELECT request_id AS id FROM spend_logs FINAL ORDER BY id",
|
||||
];
|
||||
for sql in queries {
|
||||
let body: Value = serde_json::from_str(&query_sql(&database.client, &reader, sql).await?)?;
|
||||
assert_eq!(
|
||||
body["data"],
|
||||
json!(
|
||||
expected
|
||||
.iter()
|
||||
.map(|id| json!({"id": id}))
|
||||
.collect::<Vec<_>>()
|
||||
),
|
||||
"{sql}"
|
||||
);
|
||||
}
|
||||
let summary: Value = serde_json::from_str(
|
||||
&query_sql(
|
||||
&database.client,
|
||||
&reader,
|
||||
"SELECT sum(SpanCount) AS count FROM agent_traces_by_key",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(summary["data"][0]["count"], json!(expected.len()));
|
||||
let help = serde_json::to_string(&query_help(&database.client, &reader).await?)?;
|
||||
assert_eq!(help.contains("secret_b"), expected.contains(&"b"));
|
||||
assert_eq!(help.contains("secret-b"), expected.contains(&"b"));
|
||||
let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned());
|
||||
let repeated = recreated
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
assert_eq!(reader.url(), repeated.url());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rotating_master_secret_revokes_previous_reader_credentials(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
};
|
||||
let old_reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "old-master-secret")
|
||||
.await?;
|
||||
let old_result = query_sql(
|
||||
&database.client,
|
||||
&old_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await?;
|
||||
let old_rows: Value = serde_json::from_str(&old_result)?;
|
||||
assert_eq!(old_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
|
||||
|
||||
let rotated_readers = QueryReaders::new(database.writer.clone(), "trace_test".into());
|
||||
let new_reader = rotated_readers
|
||||
.connection(&database.client, &scope, "new-master-secret")
|
||||
.await?;
|
||||
assert!(
|
||||
query_sql(
|
||||
&database.client,
|
||||
&old_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let new_result = query_sql(
|
||||
&database.client,
|
||||
&new_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await?;
|
||||
let new_rows: Value = serde_json::from_str(&new_result)?;
|
||||
assert_eq!(new_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
|
||||
assert_eq!(old_reader.url().username(), new_reader.url().username());
|
||||
assert_ne!(old_reader.url().password(), new_reader.url().password());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn managed_reader_rejects_privilege_and_scope_bypasses(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
};
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
for sql in [
|
||||
"INSERT INTO otel_traces (TraceId) VALUES ('injected')",
|
||||
"DROP TABLE otel_traces",
|
||||
"SELECT * FROM private_data",
|
||||
"SELECT * FROM otel_traces SETTINGS readonly = 0",
|
||||
"SELECT * FROM otel_traces SETTINGS max_memory_usage = 0",
|
||||
"SELECT * FROM otel_traces SETTINGS max_execution_time = 0",
|
||||
"CREATE USER scope_bypass",
|
||||
"CREATE NAMED COLLECTION scope_bypass AS host = 'localhost'",
|
||||
"BACKUP TABLE otel_traces TO Disk('default', 'scope-bypass')",
|
||||
"SELECT * FROM url('http://127.0.0.1:1/', 'LineAsString', 'line String')",
|
||||
"SELECT * FROM remote('127.0.0.1', 'trace_test', 'otel_traces')",
|
||||
] {
|
||||
assert!(
|
||||
matches!(
|
||||
query_sql(&database.client, &reader, sql).await,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::QueryFailed(_)
|
||||
))
|
||||
),
|
||||
"{sql}"
|
||||
);
|
||||
}
|
||||
let roles: Value = serde_json::from_str(
|
||||
&query_sql(&database.client, &reader, "SELECT enabledRoles() AS roles").await?,
|
||||
)?;
|
||||
assert_eq!(roles["data"], json!([{ "roles": [] }]));
|
||||
let rows: Value = serde_json::from_str(
|
||||
&query_sql(
|
||||
&database.client,
|
||||
&reader,
|
||||
"SELECT DISTINCT TeamId FROM otel_traces",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(rows["data"], json!([{ "TeamId": "team-a" }]));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn provisioning_failure_never_returns_a_writer_connection(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::All, "test-master-secret")
|
||||
.await?;
|
||||
let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned());
|
||||
let result = no_provision_privileges
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
},
|
||||
"other-secret",
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(Error::Cached(source)) if matches!(source.as_ref(), Error::ProvisionFailed(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::All, "")
|
||||
.await,
|
||||
Err(Error::MissingSecret)
|
||||
));
|
||||
assert!(matches!(
|
||||
database
|
||||
.readers
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec![String::new()]
|
||||
},
|
||||
"test-master-secret"
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidScope)
|
||||
));
|
||||
let permits = (0..8)
|
||||
.map(|_| database.readers.acquire())
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
assert!(matches!(database.readers.acquire(), Err(Error::Busy)));
|
||||
drop(permits);
|
||||
assert!(database.readers.acquire().is_ok());
|
||||
let rows = litellm_traces_clickhouse::execute_read(
|
||||
&database.client,
|
||||
&database.writer,
|
||||
"SELECT count() AS count FROM trace_test.otel_traces",
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.await?;
|
||||
let rows: Value = serde_json::from_str(&rows)?;
|
||||
assert_eq!(rows["data"][0]["count"], 6);
|
||||
Ok(())
|
||||
}
|
||||
36
litellm-rust/crates/traces-clickhouse/tests/support/mod.rs
Normal file
36
litellm-rust/crates/traces-clickhouse/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
use litellm_http::Client;
|
||||
use rstest::fixture;
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
pub type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
|
||||
|
||||
pub struct ClickHouseDatabase {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
pub url: String,
|
||||
pub client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn database() -> TestResult<ClickHouseDatabase> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
);
|
||||
Ok(ClickHouseDatabase {
|
||||
_container: container,
|
||||
url,
|
||||
client: Client::no_redirect_for_test(),
|
||||
})
|
||||
}
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
- Keep OTLP decoding, trace schema, row encoding and named query selection here. Generic ClickHouse connections and HTTP execution belong in `litellm-storage-clickhouse`
|
||||
- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge`
|
||||
- Keep the SQL migrations here as the only ClickHouse schema definition, as `migrations/NNNN_description.sql` files embedded by `litellm_migrate::migrate!`; adding a file is the only step
|
||||
- Use typed query parameters and a dedicated SELECT-only reader with server-side limits
|
||||
- Keep `config/reader.xml` grants on the database the schema is created in (CLICKHOUSE_DATABASE, default `litellm`)
|
||||
- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions
|
||||
- Test storage behavior through the crate's public API against ClickHouse
|
||||
- Own OTLP decoding, normalization, shared authorization and named query contracts; remain independent of storage and Python
|
||||
- Never depend on `litellm-traces-clickhouse` or `litellm-storage-clickhouse`
|
||||
- Preserve decoding limits, normalization precedence and shared resource identity
|
||||
- Keep ClickHouse schema, row encoding and queries in `litellm-traces-clickhouse`; keep PyO3 conversion in `python-bridge`
|
||||
- Test decoding and normalization through the public API
|
||||
- Expose one top-level `Error` enum in `src/error.rs` for decoding and normalization failures
|
||||
|
|
|
|||
|
|
@ -6,34 +6,17 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
askama.workspace = true
|
||||
base64.workspace = true
|
||||
flate2.workspace = true
|
||||
futures-util.workspace = true
|
||||
hmac = "0.12.1"
|
||||
indexmap = { version = "2", features = ["serde"] }
|
||||
moka.workspace = true
|
||||
opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
|
||||
prost.workspace = true
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
litellm-http.workspace = true
|
||||
litellm-migrate.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
sha2.workspace = true
|
||||
serde = { workspace = true, features = ["rc"] }
|
||||
serde_json.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
criterion.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
||||
[[bench]]
|
||||
name = "resource-fanout"
|
||||
|
|
|
|||
|
|
@ -1,32 +0,0 @@
|
|||
<clickhouse>
|
||||
<profiles>
|
||||
<litellm_traces_reader>
|
||||
<readonly>1</readonly>
|
||||
<max_execution_time>10</max_execution_time>
|
||||
<max_result_rows>1000</max_result_rows>
|
||||
<max_result_bytes>4194304</max_result_bytes>
|
||||
<result_overflow_mode>throw</result_overflow_mode>
|
||||
<max_memory_usage>268435456</max_memory_usage>
|
||||
<constraints>
|
||||
<readonly><readonly/></readonly>
|
||||
<max_execution_time><readonly/></max_execution_time>
|
||||
<max_result_rows><readonly/></max_result_rows>
|
||||
<max_result_bytes><readonly/></max_result_bytes>
|
||||
<result_overflow_mode><readonly/></result_overflow_mode>
|
||||
<max_memory_usage><readonly/></max_memory_usage>
|
||||
</constraints>
|
||||
</litellm_traces_reader>
|
||||
</profiles>
|
||||
<users>
|
||||
<litellm_traces_reader>
|
||||
<password from_env="LITELLM_TRACES_READER_PASSWORD"/>
|
||||
<networks><ip>::/0</ip></networks>
|
||||
<profile>litellm_traces_reader</profile>
|
||||
<grants>
|
||||
<query>GRANT SELECT ON litellm.otel_traces</query>
|
||||
<query>GRANT SELECT ON litellm.agent_traces_by_key</query>
|
||||
<query>GRANT SELECT ON litellm.spend_logs</query>
|
||||
</grants>
|
||||
</litellm_traces_reader>
|
||||
</users>
|
||||
</clickhouse>
|
||||
|
|
@ -1,23 +0,0 @@
|
|||
SELECT TraceId AS trace_id,
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
|
||||
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,
|
||||
dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms,
|
||||
sum(SpanCount) AS span_count, length(groupUniqArrayArray(AgentNames)) AS agent_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(groupArrayArray(RequestIds)) AS request_ids
|
||||
FROM agent_traces_by_key
|
||||
WHERE (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR 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}
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
SELECT SpanId AS span_id, Input AS input, Output AS output, SpanAttributes AS attributes
|
||||
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 ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
|
||||
LIMIT 1
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DecodeError {
|
||||
pub enum Error {
|
||||
#[error("invalid OTLP trace payload")]
|
||||
InvalidPayload,
|
||||
#[error("OTLP trace payload exceeds the decoding budget")]
|
||||
|
|
@ -9,21 +9,9 @@ pub enum DecodeError {
|
|||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum QueryAccessError {
|
||||
#[error("trace SQL queries require a configured proxy master key")]
|
||||
MissingSecret,
|
||||
#[error("invalid trace query scope")]
|
||||
InvalidScope,
|
||||
#[error("trace SQL query concurrency limit exceeded")]
|
||||
Busy,
|
||||
#[error(
|
||||
"ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables"
|
||||
)]
|
||||
ProvisionFailed(u16),
|
||||
#[error("ClickHouse reader provisioning transport failed")]
|
||||
ProvisionTransport,
|
||||
#[error(transparent)]
|
||||
Storage(#[from] litellm_storage_clickhouse::Error),
|
||||
#[error(transparent)]
|
||||
Cached(#[from] std::sync::Arc<QueryAccessError>),
|
||||
}
|
||||
#[error("invalid trace query scope")]
|
||||
pub struct InvalidScope;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("unknown ClickHouse read query")]
|
||||
pub struct InvalidQuery;
|
||||
|
|
|
|||
|
|
@ -1,25 +1,13 @@
|
|||
mod config;
|
||||
mod error;
|
||||
mod insert;
|
||||
mod normalize;
|
||||
mod otlp;
|
||||
mod query;
|
||||
pub mod query;
|
||||
mod query_access;
|
||||
mod schema;
|
||||
mod shared;
|
||||
mod sql;
|
||||
|
||||
pub use config::Config;
|
||||
pub use error::{DecodeError, QueryAccessError};
|
||||
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
|
||||
pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
|
||||
pub use normalize::{
|
||||
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType,
|
||||
};
|
||||
pub use error::{Error, InvalidQuery, InvalidScope};
|
||||
pub use normalize::{NormalizedSpan, ObservationType};
|
||||
pub use otlp::{DecodedSpan, decode_otlp};
|
||||
pub use query_access::{QueryReaders, QueryScope};
|
||||
pub use schema::{ensure_schema, schema_statements};
|
||||
pub use query::ReadQuery;
|
||||
pub use query_access::QueryScope;
|
||||
pub use shared::{Shared, SharedIdentity};
|
||||
pub use sql::{LensQuery, ReadQuery, execute_named_read};
|
||||
|
||||
pub use query::{query_help, query_sql};
|
||||
|
|
|
|||
367
litellm-rust/crates/traces/src/normalize/claude_code.rs
Normal file
367
litellm-rust/crates/traces/src/normalize/claude_code.rs
Normal file
|
|
@ -0,0 +1,367 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, tokens};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(crate) const CLAUDE_CODE_SCOPE: &str = "com.anthropic.claude_code.tracing";
|
||||
pub(crate) const CLAUDE_CODE_AGENT: &str = "claude-code";
|
||||
const AGENT_SDK_FRAMEWORK: &str = "claude-agent-sdk";
|
||||
|
||||
pub(super) struct ClaudeCodeNormalizer;
|
||||
|
||||
enum SpanType {
|
||||
Interaction,
|
||||
LlmRequest,
|
||||
Tool,
|
||||
Other,
|
||||
}
|
||||
|
||||
fn span_type(name: &str, attributes: &BTreeMap<String, String>) -> SpanType {
|
||||
let kind = attr(attributes, "span.type");
|
||||
let kind = if kind.is_empty() {
|
||||
name.strip_prefix("claude_code.").unwrap_or(name)
|
||||
} else {
|
||||
kind
|
||||
};
|
||||
match kind {
|
||||
"interaction" => SpanType::Interaction,
|
||||
"llm_request" => SpanType::LlmRequest,
|
||||
"tool" => SpanType::Tool,
|
||||
_ => SpanType::Other,
|
||||
}
|
||||
}
|
||||
|
||||
fn framework(attributes: &BTreeMap<String, String>) -> &'static str {
|
||||
if attr(attributes, "query_source_safe") == "sdk"
|
||||
|| attr(attributes, "system_prompt_preview").contains("cc_entrypoint=sdk")
|
||||
{
|
||||
AGENT_SDK_FRAMEWORK
|
||||
} else {
|
||||
CLAUDE_CODE_AGENT
|
||||
}
|
||||
}
|
||||
|
||||
fn split_header(text: &str) -> Option<(&str, &str)> {
|
||||
let (header, body) = text.strip_prefix('[')?.split_once("]\n")?;
|
||||
Some((header, body))
|
||||
}
|
||||
|
||||
fn without_header<'a>(text: &'a str, prefix: &str) -> &'a str {
|
||||
split_header(text)
|
||||
.filter(|(header, _)| header.starts_with(prefix))
|
||||
.map_or(text, |(_, body)| body)
|
||||
}
|
||||
|
||||
fn tool_arguments(attributes: &BTreeMap<String, String>) -> Option<&str> {
|
||||
let arguments = without_header(attr(attributes, "tool_input"), "TOOL INPUT");
|
||||
serde_json::from_str::<Map<String, Value>>(arguments)
|
||||
.is_ok()
|
||||
.then_some(arguments)
|
||||
}
|
||||
|
||||
fn tool_input(attributes: &BTreeMap<String, String>) -> String {
|
||||
if let Some(arguments) = tool_arguments(attributes) {
|
||||
return arguments.to_owned();
|
||||
}
|
||||
let fields: Map<String, Value> = [
|
||||
("command", "full_command"),
|
||||
("file_path", "file_path"),
|
||||
("bash_argv0", "bash_argv0"),
|
||||
]
|
||||
.into_iter()
|
||||
.filter_map(|(key, source)| {
|
||||
let value = attr(attributes, source);
|
||||
(!value.is_empty()).then(|| (key.to_owned(), Value::String(value.to_owned())))
|
||||
})
|
||||
.collect();
|
||||
if fields.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
Value::Object(fields).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn tool_output(attributes: &BTreeMap<String, String>, events: &[DecodedEvent]) -> String {
|
||||
events
|
||||
.iter()
|
||||
.filter(|event| event.name == "tool.output")
|
||||
.flat_map(|event| {
|
||||
["output", "content", "diff"]
|
||||
.into_iter()
|
||||
.map(|key| attr(&event.attributes, key))
|
||||
})
|
||||
.find(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| without_header(attr(attributes, "new_context"), "TOOL RESULT"))
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
fn context_message(context: &str) -> Value {
|
||||
let (role, content) = match split_header(context) {
|
||||
Some(("USER" | "USER PROMPT", body)) => ("user", body),
|
||||
Some(("ASSISTANT", body)) => ("assistant", body),
|
||||
Some((header, body)) if header.starts_with("TOOL RESULT") => ("tool", body),
|
||||
_ => ("user", context),
|
||||
};
|
||||
json!({"role": role, "content": content})
|
||||
}
|
||||
|
||||
fn user_prompt(attributes: &BTreeMap<String, String>) -> String {
|
||||
let prompt = attr(attributes, "user_prompt");
|
||||
if prompt.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
json!([{"role": "user", "content": prompt}]).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn llm_input(attributes: &BTreeMap<String, String>) -> String {
|
||||
let messages: Vec<Value> = [
|
||||
Some(attr(attributes, "system_prompt_preview"))
|
||||
.filter(|system| !system.is_empty())
|
||||
.map(|system| json!({"role": "system", "content": system})),
|
||||
Some(attr(attributes, "new_context"))
|
||||
.filter(|context| !context.is_empty())
|
||||
.map(context_message),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
if messages.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
Value::Array(messages).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn llm_output(attributes: &BTreeMap<String, String>) -> String {
|
||||
let output = attr(attributes, "response.model_output");
|
||||
if output.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
json!({"role": "assistant", "content": output}).to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn input_tokens(attributes: &BTreeMap<String, String>) -> Result<u32, Error> {
|
||||
["input_tokens", "cache_read_tokens", "cache_creation_tokens"]
|
||||
.into_iter()
|
||||
.try_fold(0u32, |total, key| {
|
||||
total
|
||||
.checked_add(tokens(attributes, key)?)
|
||||
.ok_or(Error::TokenCountOutOfRange)
|
||||
})
|
||||
}
|
||||
|
||||
impl SpanNormalizer for ClaudeCodeNormalizer {
|
||||
fn matches(&self, scope_name: &str, _attributes: &BTreeMap<String, String>) -> bool {
|
||||
scope_name == CLAUDE_CODE_SCOPE
|
||||
}
|
||||
|
||||
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2] {
|
||||
match span_type("", attributes) {
|
||||
SpanType::Interaction => ["user_prompt", ""],
|
||||
SpanType::LlmRequest => ["new_context", "response.model_output"],
|
||||
SpanType::Tool if tool_arguments(attributes).is_some() => ["tool_input", ""],
|
||||
SpanType::Tool | SpanType::Other => ["", ""],
|
||||
}
|
||||
}
|
||||
|
||||
fn display_name(&self, attributes: &BTreeMap<String, String>) -> Option<String> {
|
||||
let tool_name = attr(attributes, "tool_name");
|
||||
(matches!(span_type("", attributes), SpanType::Tool) && !tool_name.is_empty())
|
||||
.then(|| tool_name.to_owned())
|
||||
}
|
||||
|
||||
fn normalize(
|
||||
&self,
|
||||
name: &str,
|
||||
_parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let base = NormalizedSpan {
|
||||
observation_type: ObservationType::Framework,
|
||||
agent_name: CLAUDE_CODE_AGENT.to_owned(),
|
||||
framework: framework(attributes).to_owned(),
|
||||
litellm_request_id: String::new(),
|
||||
model: String::new(),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
input: String::new(),
|
||||
output: String::new(),
|
||||
};
|
||||
Ok(match span_type(name, attributes) {
|
||||
SpanType::Interaction => NormalizedSpan {
|
||||
observation_type: ObservationType::Agent,
|
||||
input: user_prompt(attributes),
|
||||
..base
|
||||
},
|
||||
SpanType::LlmRequest => NormalizedSpan {
|
||||
observation_type: ObservationType::Llm,
|
||||
litellm_request_id: first(attributes, "gen_ai.response.id", "request_id")
|
||||
.to_owned(),
|
||||
model: first(attributes, "model", "gen_ai.request.model").to_owned(),
|
||||
input_tokens: input_tokens(attributes)?,
|
||||
output_tokens: tokens(attributes, "output_tokens")?,
|
||||
input: llm_input(attributes),
|
||||
output: llm_output(attributes),
|
||||
..base
|
||||
},
|
||||
SpanType::Tool => NormalizedSpan {
|
||||
observation_type: ObservationType::Tool,
|
||||
input: tool_input(attributes),
|
||||
output: tool_output(attributes, events),
|
||||
..base
|
||||
},
|
||||
SpanType::Other => base,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use rstest::rstest;
|
||||
use serde_json::Value;
|
||||
|
||||
use super::{CLAUDE_CODE_SCOPE, ClaudeCodeNormalizer, SpanNormalizer};
|
||||
use crate::{Error, normalize::ObservationType, otlp::DecodedEvent};
|
||||
|
||||
fn attributes(pairs: &[(&str, &str)]) -> BTreeMap<String, String> {
|
||||
pairs
|
||||
.iter()
|
||||
.map(|(key, value)| ((*key).to_owned(), (*value).to_owned()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn tool_without_detailed_input_lists_known_arguments() {
|
||||
let span = ClaudeCodeNormalizer
|
||||
.normalize(
|
||||
"claude_code.tool",
|
||||
"parent",
|
||||
&attributes(&[
|
||||
("span.type", "tool"),
|
||||
("tool_name", "Bash"),
|
||||
("full_command", "git status"),
|
||||
("bash_argv0", "git"),
|
||||
]),
|
||||
&[],
|
||||
)
|
||||
.expect("valid span");
|
||||
let input: Value = serde_json::from_str(&span.input).expect("argument object");
|
||||
assert_eq!(input["command"], "git status");
|
||||
assert_eq!(input["bash_argv0"], "git");
|
||||
assert!(input.get("file_path").is_none());
|
||||
assert!(input.get("role").is_none());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn malformed_tool_input_falls_back_and_stays_in_attributes() {
|
||||
let attrs = attributes(&[
|
||||
("span.type", "tool"),
|
||||
("tool_input", "[TOOL INPUT: Read]\nnot json"),
|
||||
("file_path", "/workspace/a.py"),
|
||||
]);
|
||||
let span = ClaudeCodeNormalizer
|
||||
.normalize("claude_code.tool", "parent", &attrs, &[])
|
||||
.expect("valid span");
|
||||
let input: Value = serde_json::from_str(&span.input).expect("argument object");
|
||||
assert_eq!(input["file_path"], "/workspace/a.py");
|
||||
assert!(
|
||||
!ClaudeCodeNormalizer
|
||||
.consumed_attributes(&attrs)
|
||||
.contains(&"tool_input")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::event_output(
|
||||
vec![DecodedEvent { name: "tool.output".to_owned(), attributes: attributes(&[("output", "stdout text")]) }],
|
||||
"stdout text"
|
||||
)]
|
||||
#[case::event_diff(
|
||||
vec![DecodedEvent { name: "tool.output".to_owned(), attributes: attributes(&[("diff", "+line")]) }],
|
||||
"+line"
|
||||
)]
|
||||
#[case::other_event_ignored(
|
||||
vec![DecodedEvent { name: "other".to_owned(), attributes: attributes(&[("output", "nope")]) }],
|
||||
"{\"stdout\":\"ctx\"}"
|
||||
)]
|
||||
fn tool_output_prefers_event_then_context(
|
||||
#[case] events: Vec<DecodedEvent>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let span = ClaudeCodeNormalizer
|
||||
.normalize(
|
||||
"claude_code.tool",
|
||||
"parent",
|
||||
&attributes(&[
|
||||
("span.type", "tool"),
|
||||
("new_context", "[TOOL RESULT: Bash]\n{\"stdout\":\"ctx\"}"),
|
||||
]),
|
||||
&events,
|
||||
)
|
||||
.expect("valid span");
|
||||
assert_eq!(span.output, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn llm_tool_result_context_becomes_tool_message() {
|
||||
let span = ClaudeCodeNormalizer
|
||||
.normalize(
|
||||
"claude_code.llm_request",
|
||||
"parent",
|
||||
&attributes(&[
|
||||
("span.type", "llm_request"),
|
||||
("new_context", "[TOOL RESULT: toolu_1]\n1\timport os"),
|
||||
]),
|
||||
&[],
|
||||
)
|
||||
.expect("valid span");
|
||||
let input: Value = serde_json::from_str(&span.input).expect("messages");
|
||||
assert_eq!(input[0]["role"], "tool");
|
||||
assert_eq!(input[0]["content"], "1\timport os");
|
||||
assert_eq!(span.output, "");
|
||||
assert_eq!(span.framework, "claude-code");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn llm_token_sum_overflow_is_rejected() {
|
||||
let result = ClaudeCodeNormalizer.normalize(
|
||||
"claude_code.llm_request",
|
||||
"parent",
|
||||
&attributes(&[
|
||||
("span.type", "llm_request"),
|
||||
("input_tokens", "4294967295"),
|
||||
("cache_read_tokens", "1"),
|
||||
]),
|
||||
&[],
|
||||
);
|
||||
assert!(matches!(result, Err(Error::TokenCountOutOfRange)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::span_type_wins("claude_code.tool", "hook", ObservationType::Framework)]
|
||||
#[case::name_fallback("claude_code.interaction", "", ObservationType::Agent)]
|
||||
#[case::unknown("claude_code.something_new", "", ObservationType::Framework)]
|
||||
fn span_type_attribute_then_name_select_the_observation(
|
||||
#[case] name: &str,
|
||||
#[case] kind: &str,
|
||||
#[case] expected: ObservationType,
|
||||
) {
|
||||
let attrs = if kind.is_empty() {
|
||||
BTreeMap::new()
|
||||
} else {
|
||||
attributes(&[("span.type", kind)])
|
||||
};
|
||||
let span = ClaudeCodeNormalizer
|
||||
.normalize(name, "parent", &attrs, &[])
|
||||
.expect("valid span");
|
||||
assert_eq!(span.observation_type, expected);
|
||||
assert!(ClaudeCodeNormalizer.matches(CLAUDE_CODE_SCOPE, &attrs));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, usage_tokens};
|
||||
use crate::DecodeError;
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct GenAiNormalizer;
|
||||
|
||||
|
|
@ -30,7 +30,8 @@ impl SpanNormalizer for GenAiNormalizer {
|
|||
_name: &str,
|
||||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
|
||||
let observation_type = match attr(attributes, "gen_ai.operation.name") {
|
||||
"invoke_agent" => ObservationType::Agent,
|
||||
|
|
@ -42,6 +43,7 @@ impl SpanNormalizer for GenAiNormalizer {
|
|||
Ok(NormalizedSpan {
|
||||
observation_type,
|
||||
agent_name: attr(attributes, "gen_ai.agent.name").to_owned(),
|
||||
framework: String::new(),
|
||||
litellm_request_id: attr(attributes, "gen_ai.response.id").to_owned(),
|
||||
model: first(attributes, "gen_ai.request.model", "gen_ai.response.model").to_owned(),
|
||||
input_tokens,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::DeserializeOwned};
|
|||
use serde_json::{Value, ser::Formatter};
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, usage_tokens};
|
||||
use crate::DecodeError;
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct LangSmithNormalizer;
|
||||
|
||||
|
|
@ -404,13 +404,15 @@ impl SpanNormalizer for LangSmithNormalizer {
|
|||
name: &str,
|
||||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
|
||||
let observation_type = span_type(name, parent_span_id, attributes);
|
||||
let io = span_io(observation_type, attributes);
|
||||
Ok(NormalizedSpan {
|
||||
observation_type,
|
||||
agent_name: attr(attributes, "langsmith.metadata.lc_agent_name").to_owned(),
|
||||
framework: String::new(),
|
||||
litellm_request_id: io.request_id,
|
||||
model: attr(attributes, "gen_ai.request.model").to_owned(),
|
||||
input_tokens,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::DecodeError;
|
||||
use serde::Serialize;
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
|
|
@ -17,6 +17,7 @@ pub enum ObservationType {
|
|||
pub struct NormalizedSpan {
|
||||
pub observation_type: ObservationType,
|
||||
pub agent_name: String,
|
||||
pub framework: String,
|
||||
pub litellm_request_id: String,
|
||||
pub model: String,
|
||||
pub input_tokens: u32,
|
||||
|
|
@ -27,68 +28,10 @@ pub struct NormalizedSpan {
|
|||
|
||||
pub(crate) struct Normalization {
|
||||
pub span: NormalizedSpan,
|
||||
pub display_name: Option<String>,
|
||||
pub consumed_attributes: [&'static str; 2],
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
pub struct NormalizedFieldDefinition {
|
||||
pub name: &'static str,
|
||||
pub clickhouse_column: &'static str,
|
||||
pub clickhouse_type: &'static str,
|
||||
pub meaning: &'static str,
|
||||
}
|
||||
|
||||
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 8] = [
|
||||
NormalizedFieldDefinition {
|
||||
name: "observation_type",
|
||||
clickhouse_column: "ObservationType",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent, LLM, tool, chain, or framework span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "agent_name",
|
||||
clickhouse_column: "AgentName",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent associated with this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "litellm_request_id",
|
||||
clickhouse_column: "LiteLLMRequestId",
|
||||
clickhouse_type: "String",
|
||||
meaning: "LiteLLM response ID used to link a span to a spend log",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "model",
|
||||
clickhouse_column: "Model",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Model used by this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input_tokens",
|
||||
clickhouse_column: "InputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Input token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output_tokens",
|
||||
clickhouse_column: "OutputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Output token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input",
|
||||
clickhouse_column: "Input",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized input payload",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output",
|
||||
clickhouse_column: "Output",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized output payload",
|
||||
},
|
||||
];
|
||||
|
||||
trait SpanNormalizer {
|
||||
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool;
|
||||
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2];
|
||||
|
|
@ -97,13 +40,20 @@ trait SpanNormalizer {
|
|||
name: &str,
|
||||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
) -> Result<NormalizedSpan, DecodeError>;
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, Error>;
|
||||
fn display_name(&self, _attributes: &BTreeMap<String, String>) -> Option<String> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
mod claude_code;
|
||||
mod genai;
|
||||
mod langsmith;
|
||||
mod openinference;
|
||||
|
||||
use claude_code::ClaudeCodeNormalizer;
|
||||
pub(crate) use claude_code::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE};
|
||||
use genai::GenAiNormalizer;
|
||||
use langsmith::LangSmithNormalizer;
|
||||
use openinference::OpenInferenceNormalizer;
|
||||
|
|
@ -121,40 +71,96 @@ fn first<'a>(attributes: &'a BTreeMap<String, String>, left: &str, right: &str)
|
|||
}
|
||||
}
|
||||
|
||||
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, DecodeError> {
|
||||
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, Error> {
|
||||
let value = attr(attributes, key).trim();
|
||||
if value.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
match value.parse::<i128>() {
|
||||
Ok(number) if (0..=u32::MAX as i128).contains(&number) => Ok(number as u32),
|
||||
Ok(_) => Err(DecodeError::TokenCountOutOfRange),
|
||||
Ok(_) => Err(Error::TokenCountOutOfRange),
|
||||
Err(_)
|
||||
if value
|
||||
.trim_start_matches(['+', '-'])
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit()) =>
|
||||
{
|
||||
Err(DecodeError::TokenCountOutOfRange)
|
||||
Err(Error::TokenCountOutOfRange)
|
||||
}
|
||||
Err(_) => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), DecodeError> {
|
||||
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), Error> {
|
||||
Ok((
|
||||
tokens(attributes, "gen_ai.usage.input_tokens")?,
|
||||
tokens(attributes, "gen_ai.usage.output_tokens")?,
|
||||
))
|
||||
}
|
||||
|
||||
#[derive(Default, Deserialize)]
|
||||
struct AgentMetadata {
|
||||
#[serde(default)]
|
||||
lc_agent_name: String,
|
||||
#[serde(default)]
|
||||
ls_integration: String,
|
||||
}
|
||||
|
||||
fn recorded_agent_name(
|
||||
name: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
span: &NormalizedSpan,
|
||||
) -> String {
|
||||
let explicit = [
|
||||
span.agent_name.as_str(),
|
||||
attr(attributes, "gen_ai.agent.name"),
|
||||
attr(attributes, "agent.name"),
|
||||
attr(attributes, "openclaw.agent"),
|
||||
]
|
||||
.into_iter()
|
||||
.find(|value| !value.is_empty());
|
||||
if let Some(value) = explicit {
|
||||
return value.to_owned();
|
||||
}
|
||||
let metadata =
|
||||
serde_json::from_str::<AgentMetadata>(attr(attributes, "metadata")).unwrap_or_default();
|
||||
if !metadata.lc_agent_name.is_empty() {
|
||||
return metadata.lc_agent_name;
|
||||
}
|
||||
if span.observation_type == ObservationType::Agent {
|
||||
let node = attr(attributes, "graph.node.id");
|
||||
if !node.is_empty() {
|
||||
return node.to_owned();
|
||||
}
|
||||
if metadata.ls_integration == "langgraph" && name != "LangGraph" && !is_middleware(name) {
|
||||
return name.to_owned();
|
||||
}
|
||||
}
|
||||
String::new()
|
||||
}
|
||||
|
||||
fn is_middleware(name: &str) -> bool {
|
||||
[
|
||||
".wrap_model_call",
|
||||
".wrap_tool_call",
|
||||
".before_agent",
|
||||
".after_agent",
|
||||
".before_model",
|
||||
".after_model",
|
||||
]
|
||||
.iter()
|
||||
.any(|suffix| name.ends_with(suffix))
|
||||
}
|
||||
|
||||
pub fn normalize(
|
||||
scope_name: &str,
|
||||
name: &str,
|
||||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
) -> Result<Normalization, DecodeError> {
|
||||
let normalizers: [&dyn SpanNormalizer; 3] = [
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<Normalization, Error> {
|
||||
let normalizers: [&dyn SpanNormalizer; 4] = [
|
||||
&ClaudeCodeNormalizer,
|
||||
&LangSmithNormalizer,
|
||||
&OpenInferenceNormalizer,
|
||||
&GenAiNormalizer,
|
||||
|
|
@ -163,24 +169,40 @@ pub fn normalize(
|
|||
.into_iter()
|
||||
.find(|normalizer| normalizer.matches(scope_name, attributes))
|
||||
.expect("GenAI fallback always matches");
|
||||
let span = normalizer.normalize(name, parent_span_id, attributes, events)?;
|
||||
let agent_name = recorded_agent_name(name, attributes, &span);
|
||||
let observation_type = if !parent_span_id.is_empty()
|
||||
&& scope_name == "openinference.instrumentation.langchain"
|
||||
&& is_middleware(name)
|
||||
{
|
||||
ObservationType::Framework
|
||||
} else {
|
||||
span.observation_type
|
||||
};
|
||||
Ok(Normalization {
|
||||
span: normalizer.normalize(name, parent_span_id, attributes)?,
|
||||
span: NormalizedSpan {
|
||||
agent_name,
|
||||
observation_type,
|
||||
..span
|
||||
},
|
||||
display_name: normalizer.display_name(attributes),
|
||||
consumed_attributes: normalizer.consumed_attributes(attributes),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{NORMALIZED_FIELD_DEFINITIONS, ObservationType, normalize};
|
||||
use super::{ObservationType, normalize};
|
||||
|
||||
#[rstest]
|
||||
#[case::langsmith("langsmith", [("langsmith.span.kind", "llm"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
|
||||
#[case::openinference("other", [("openinference.span.kind", "LLM"), ("gen_ai.operation.name", "execute_tool")], ObservationType::Llm)]
|
||||
#[case::genai("other", [("gen_ai.operation.name", "execute_tool"), ("gen_ai.usage.input_tokens", "7")], ObservationType::Tool)]
|
||||
#[case::claude_code("com.anthropic.claude_code.tracing", [("span.type", "llm_request"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
|
||||
fn convention_dispatch_preserves_precedence(
|
||||
#[case] scope: &str,
|
||||
#[case] attributes: [(&str, &str); 2],
|
||||
|
|
@ -190,7 +212,7 @@ mod tests {
|
|||
.into_iter()
|
||||
.map(|(key, value)| (key.to_owned(), value.to_owned()))
|
||||
.collect();
|
||||
let fields = normalize(scope, "step", "parent", &attributes)
|
||||
let fields = normalize(scope, "step", "parent", &attributes, &[])
|
||||
.expect("valid tokens")
|
||||
.span;
|
||||
assert_eq!(fields.observation_type, expected);
|
||||
|
|
@ -199,30 +221,11 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn field_definitions_match_serialized_normalized_span() {
|
||||
let fields = normalize("", "root", "", &BTreeMap::new())
|
||||
.expect("valid tokens")
|
||||
.span;
|
||||
let serialized = serde_json::to_value(fields).expect("serializable fields");
|
||||
let keys: BTreeSet<_> = serialized
|
||||
.as_object()
|
||||
.expect("field object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(|field| field.name)
|
||||
.collect();
|
||||
assert_eq!(keys, mapped);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn token_counts_accept_surrounding_whitespace() {
|
||||
let attributes =
|
||||
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), " 7 ".to_owned())]);
|
||||
let fields = normalize("", "root", "", &attributes)
|
||||
let fields = normalize("", "root", "", &attributes, &[])
|
||||
.expect("valid tokens")
|
||||
.span;
|
||||
assert_eq!(fields.input_tokens, 7);
|
||||
|
|
@ -234,6 +237,6 @@ mod tests {
|
|||
fn token_counts_outside_storage_range_are_rejected(#[case] value: &str) {
|
||||
let attributes =
|
||||
BTreeMap::from([("gen_ai.usage.input_tokens".to_owned(), value.to_owned())]);
|
||||
assert!(normalize("", "root", "", &attributes).is_err());
|
||||
assert!(normalize("", "root", "", &attributes, &[]).is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, tokens, usage_tokens};
|
||||
use crate::DecodeError;
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct OpenInferenceNormalizer;
|
||||
|
||||
|
|
@ -19,7 +19,8 @@ impl SpanNormalizer for OpenInferenceNormalizer {
|
|||
_name: &str,
|
||||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (usage_input, usage_output) = usage_tokens(attributes)?;
|
||||
let observation_type = match attr(attributes, "openinference.span.kind")
|
||||
.to_ascii_uppercase()
|
||||
|
|
@ -34,6 +35,7 @@ impl SpanNormalizer for OpenInferenceNormalizer {
|
|||
Ok(NormalizedSpan {
|
||||
observation_type,
|
||||
agent_name: attr(attributes, "agent.name").to_owned(),
|
||||
framework: String::new(),
|
||||
litellm_request_id: String::new(),
|
||||
model: attr(attributes, "llm.model_name").to_owned(),
|
||||
input_tokens: if attributes.contains_key("llm.token_count.prompt") {
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ use serde::{
|
|||
};
|
||||
|
||||
use super::limits::{Budget, MAX_ATTRIBUTES};
|
||||
use crate::DecodeError;
|
||||
use crate::Error;
|
||||
|
||||
struct AttributeWriter<'a> {
|
||||
body: Vec<u8>,
|
||||
|
|
@ -32,9 +32,9 @@ impl Write for AttributeWriter<'_> {
|
|||
pub(super) fn attributes(
|
||||
values: Vec<KeyValue>,
|
||||
budget: &mut Budget,
|
||||
) -> Result<BTreeMap<String, String>, DecodeError> {
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
if values.len() > MAX_ATTRIBUTES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
values
|
||||
.into_iter()
|
||||
|
|
@ -59,8 +59,8 @@ pub(super) fn attributes(
|
|||
budget,
|
||||
};
|
||||
serde_json::to_writer(&mut writer, &AttributeJson(value.as_ref()))
|
||||
.map_err(|_| DecodeError::TooLarge)?;
|
||||
String::from_utf8(writer.body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
.map_err(|_| Error::TooLarge)?;
|
||||
String::from_utf8(writer.body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
};
|
||||
Ok((entry.key, text))
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use std::fmt;
|
|||
use prost::encoding::{DecodeContext, WireType, decode_key, decode_varint, skip_field};
|
||||
use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor};
|
||||
|
||||
use crate::{DecodeError, Shared};
|
||||
use crate::{Error, Shared};
|
||||
|
||||
pub(super) const MAX_DEPTH: usize = 32;
|
||||
pub(super) const MAX_NODES: usize = 65_536;
|
||||
|
|
@ -12,7 +12,7 @@ pub(super) const MAX_ATTRIBUTES: usize = 256;
|
|||
pub(super) const MAX_EVENTS: usize = 256;
|
||||
pub(super) const MAX_DECODED_SPAN_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
pub(super) fn json_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
||||
pub(super) fn json_preflight(payload: &[u8]) -> Result<(), Error> {
|
||||
let mut nodes = 0;
|
||||
let mut exceeded = false;
|
||||
let mut decoder = serde_json::Deserializer::from_slice(payload);
|
||||
|
|
@ -24,9 +24,9 @@ pub(super) fn json_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
|||
.deserialize(&mut decoder)
|
||||
.and_then(|()| decoder.end());
|
||||
if exceeded {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
result.map_err(|_| DecodeError::InvalidPayload)
|
||||
result.map_err(|_| Error::InvalidPayload)
|
||||
}
|
||||
|
||||
struct JsonBudget<'a> {
|
||||
|
|
@ -146,7 +146,7 @@ impl MessageKind {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) fn protobuf_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
||||
pub(super) fn protobuf_preflight(payload: &[u8]) -> Result<(), Error> {
|
||||
scan_message(payload, MessageKind::Export, 0, &mut 0)
|
||||
}
|
||||
|
||||
|
|
@ -155,27 +155,27 @@ fn scan_message(
|
|||
kind: MessageKind,
|
||||
depth: usize,
|
||||
nodes: &mut usize,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
if depth > MAX_DEPTH {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
while !payload.is_empty() {
|
||||
*nodes += 1;
|
||||
if *nodes > MAX_NODES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
let (tag, wire) = decode_key(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let (tag, wire) = decode_key(&mut payload).map_err(|_| Error::InvalidPayload)?;
|
||||
if let (WireType::LengthDelimited, Some(child)) = (wire, kind.child(tag)) {
|
||||
let length = decode_varint(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let length = usize::try_from(length).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let length = decode_varint(&mut payload).map_err(|_| Error::InvalidPayload)?;
|
||||
let length = usize::try_from(length).map_err(|_| Error::InvalidPayload)?;
|
||||
let (message, rest) = payload
|
||||
.split_at_checked(length)
|
||||
.ok_or(DecodeError::InvalidPayload)?;
|
||||
.ok_or(Error::InvalidPayload)?;
|
||||
scan_message(message, child, depth + 1, nodes)?;
|
||||
payload = rest;
|
||||
} else {
|
||||
skip_field(wire, tag, &mut payload, DecodeContext::default())
|
||||
.map_err(|_| DecodeError::InvalidPayload)?;
|
||||
.map_err(|_| Error::InvalidPayload)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
|
|
@ -194,7 +194,7 @@ impl Budget {
|
|||
&mut self,
|
||||
value: &Shared<T>,
|
||||
allocated_bytes: impl FnOnce(&T) -> usize,
|
||||
) -> Result<Shared<T>, DecodeError> {
|
||||
) -> Result<Shared<T>, Error> {
|
||||
let cloned = value.clone();
|
||||
if !value.shares_storage_with(&cloned) {
|
||||
self.consume(allocated_bytes(value))?;
|
||||
|
|
@ -202,11 +202,8 @@ impl Budget {
|
|||
Ok(cloned)
|
||||
}
|
||||
|
||||
pub(super) fn consume(&mut self, bytes: usize) -> Result<(), DecodeError> {
|
||||
self.remaining = self
|
||||
.remaining
|
||||
.checked_sub(bytes)
|
||||
.ok_or(DecodeError::TooLarge)?;
|
||||
pub(super) fn consume(&mut self, bytes: usize) -> Result<(), Error> {
|
||||
self.remaining = self.remaining.checked_sub(bytes).ok_or(Error::TooLarge)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ mod wire;
|
|||
use serde::Serialize;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::{DecodeError, NormalizedSpan, Shared};
|
||||
use crate::{Error, NormalizedSpan, Shared};
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct DecodedEvent {
|
||||
|
|
@ -35,10 +35,7 @@ pub struct DecodedSpan {
|
|||
pub consumed_attributes: [&'static str; 2],
|
||||
}
|
||||
|
||||
pub fn decode_otlp(
|
||||
body: &[u8],
|
||||
content_type: Option<&str>,
|
||||
) -> Result<Vec<DecodedSpan>, DecodeError> {
|
||||
pub fn decode_otlp(body: &[u8], content_type: Option<&str>) -> Result<Vec<DecodedSpan>, Error> {
|
||||
let request = wire::decode(body, content_type)?;
|
||||
span::flatten(request)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -10,9 +10,12 @@ use super::{
|
|||
attributes::attributes,
|
||||
limits::{Budget, MAX_ATTRIBUTES, MAX_DECODED_SPAN_BYTES, MAX_EVENTS, MAX_SPANS},
|
||||
};
|
||||
use crate::{DecodeError, Shared, normalize::normalize};
|
||||
use crate::{
|
||||
Error, Shared,
|
||||
normalize::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE, normalize},
|
||||
};
|
||||
|
||||
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, DecodeError> {
|
||||
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, Error> {
|
||||
let mut budget = Budget::new(MAX_DECODED_SPAN_BYTES);
|
||||
let mut spans = Vec::new();
|
||||
for resource in request.resource_spans {
|
||||
|
|
@ -25,7 +28,7 @@ fn append_resource(
|
|||
resource: ResourceSpans,
|
||||
budget: &mut Budget,
|
||||
spans: &mut Vec<DecodedSpan>,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
let attributes = Shared::new(attributes(
|
||||
resource
|
||||
.resource
|
||||
|
|
@ -44,17 +47,17 @@ fn append_scope(
|
|||
resource: &Shared<BTreeMap<String, String>>,
|
||||
budget: &mut Budget,
|
||||
spans: &mut Vec<DecodedSpan>,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
let scope = scope_spans.scope.unwrap_or_default();
|
||||
if scope.attributes.len() > MAX_ATTRIBUTES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
budget.consume(scope.name.len() + scope.version.len())?;
|
||||
let scope_name: Shared<String> = scope.name.into();
|
||||
let scope_version: Shared<String> = scope.version.into();
|
||||
for span in scope_spans.spans {
|
||||
if spans.len() >= MAX_SPANS {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
validate_span(&span)?;
|
||||
budget.consume(
|
||||
|
|
@ -82,7 +85,7 @@ fn valid_id(value: &[u8], length: usize) -> bool {
|
|||
value.len() == length && value.iter().any(|byte| *byte != 0)
|
||||
}
|
||||
|
||||
fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
||||
fn validate_span(span: &Span) -> Result<(), Error> {
|
||||
if !valid_id(&span.trace_id, 16)
|
||||
|| !valid_id(&span.span_id, 8)
|
||||
|| (!span.parent_span_id.is_empty() && !valid_id(&span.parent_span_id, 8))
|
||||
|
|
@ -94,7 +97,7 @@ fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
|||
.iter()
|
||||
.any(|link| !valid_id(&link.trace_id, 16) || !valid_id(&link.span_id, 8))
|
||||
{
|
||||
return Err(DecodeError::InvalidPayload);
|
||||
return Err(Error::InvalidPayload);
|
||||
}
|
||||
if span.events.len() > MAX_EVENTS
|
||||
|| span.links.len() > MAX_EVENTS
|
||||
|
|
@ -108,7 +111,7 @@ fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
|||
.iter()
|
||||
.any(|event| event.attributes.len() > MAX_ATTRIBUTES)
|
||||
{
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -123,30 +126,62 @@ fn decoded_span(
|
|||
scope_name: &Shared<String>,
|
||||
scope_version: &Shared<String>,
|
||||
budget: &mut Budget,
|
||||
) -> Result<DecodedSpan, DecodeError> {
|
||||
) -> Result<DecodedSpan, Error> {
|
||||
let status = span.status.unwrap_or_default();
|
||||
let parent_span_id = hex_bytes(&span.parent_span_id);
|
||||
let span_attributes = attributes(span.attributes, budget)?;
|
||||
let events = span
|
||||
.events
|
||||
.into_iter()
|
||||
.map(|event| {
|
||||
budget.consume(event.name.len() + 96)?;
|
||||
Ok(DecodedEvent {
|
||||
name: event.name,
|
||||
attributes: attributes(event.attributes, budget)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, Error>>()?;
|
||||
let normalization = normalize(
|
||||
scope_name.as_ref(),
|
||||
&span.name,
|
||||
&parent_span_id,
|
||||
&span_attributes,
|
||||
&events,
|
||||
)?;
|
||||
let normalized = normalization.span;
|
||||
let resource_agent_name = resource_attributes
|
||||
.get("gen_ai.agent.name")
|
||||
.filter(|name| !name.is_empty());
|
||||
let agent_name = match (resource_agent_name, normalization.span.agent_name.as_str()) {
|
||||
(Some(name), "") => name.clone(),
|
||||
(Some(name), "hermes-agent") if scope_name.as_ref() == "hermes-otel-plugin" => name.clone(),
|
||||
(Some(name), CLAUDE_CODE_AGENT) if scope_name.as_ref() == CLAUDE_CODE_SCOPE => name.clone(),
|
||||
(None, CLAUDE_CODE_AGENT) if scope_name.as_ref() == CLAUDE_CODE_SCOPE => {
|
||||
resource_attributes
|
||||
.get("service.name")
|
||||
.filter(|name| !name.is_empty())
|
||||
.map_or_else(|| CLAUDE_CODE_AGENT.to_owned(), Clone::clone)
|
||||
}
|
||||
(_, name) => name.to_owned(),
|
||||
};
|
||||
let normalized = crate::normalize::NormalizedSpan {
|
||||
agent_name,
|
||||
..normalization.span
|
||||
};
|
||||
budget.consume(
|
||||
normalized.input.len()
|
||||
+ normalized.output.len()
|
||||
+ normalized.agent_name.len()
|
||||
+ normalized.framework.len()
|
||||
+ normalized.litellm_request_id.len()
|
||||
+ normalized.model.len(),
|
||||
+ normalized.model.len()
|
||||
+ normalization.display_name.as_ref().map_or(0, String::len),
|
||||
)?;
|
||||
Ok(DecodedSpan {
|
||||
trace_id: hex_bytes(&span.trace_id),
|
||||
span_id: hex_bytes(&span.span_id),
|
||||
parent_span_id,
|
||||
trace_state: span.trace_state,
|
||||
name: span.name,
|
||||
name: normalization.display_name.unwrap_or(span.name),
|
||||
kind: SpanKind::try_from(span.kind)
|
||||
.unwrap_or(SpanKind::Unspecified)
|
||||
.as_str_name()
|
||||
|
|
@ -167,17 +202,7 @@ fn decoded_span(
|
|||
.as_str_name()
|
||||
.to_owned(),
|
||||
status_message: status.message,
|
||||
events: span
|
||||
.events
|
||||
.into_iter()
|
||||
.map(|event| {
|
||||
budget.consume(event.name.len() + 96)?;
|
||||
Ok(DecodedEvent {
|
||||
name: event.name,
|
||||
attributes: attributes(event.attributes, budget)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DecodeError>>()?,
|
||||
events,
|
||||
normalized,
|
||||
consumed_attributes: normalization.consumed_attributes,
|
||||
})
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
|
|||
use prost::Message;
|
||||
|
||||
use super::limits::{json_preflight, protobuf_preflight};
|
||||
use crate::DecodeError;
|
||||
use crate::Error;
|
||||
|
||||
#[derive(strum::EnumString)]
|
||||
#[strum(ascii_case_insensitive)]
|
||||
|
|
@ -19,7 +19,7 @@ enum OtlpMediaType {
|
|||
pub(super) fn decode(
|
||||
body: &[u8],
|
||||
content_type: Option<&str>,
|
||||
) -> Result<ExportTraceServiceRequest, DecodeError> {
|
||||
) -> Result<ExportTraceServiceRequest, Error> {
|
||||
let media_type = content_type
|
||||
.unwrap_or("application/x-protobuf")
|
||||
.split(';')
|
||||
|
|
@ -27,16 +27,16 @@ pub(super) fn decode(
|
|||
.unwrap_or_default()
|
||||
.trim()
|
||||
.parse::<OtlpMediaType>()
|
||||
.map_err(|_| DecodeError::InvalidPayload)?;
|
||||
.map_err(|_| Error::InvalidPayload)?;
|
||||
|
||||
let request = match media_type {
|
||||
OtlpMediaType::Json => {
|
||||
json_preflight(body)?;
|
||||
serde_json::from_slice(body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
serde_json::from_slice(body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
OtlpMediaType::Protobuf => {
|
||||
protobuf_preflight(body)?;
|
||||
ExportTraceServiceRequest::decode(body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
ExportTraceServiceRequest::decode(body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
};
|
||||
Ok(request)
|
||||
|
|
|
|||
|
|
@ -1,346 +1,23 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
pub mod named;
|
||||
|
||||
use futures_util::{
|
||||
StreamExt,
|
||||
stream::{self, TryStreamExt},
|
||||
};
|
||||
use litellm_http::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, execute_read};
|
||||
|
||||
mod guide;
|
||||
|
||||
const SAMPLE_ROWS: usize = 200;
|
||||
const MAX_FIELDS: usize = 200;
|
||||
const MAX_DEPTH: usize = 16;
|
||||
const METADATA_SQL: &str = "SELECT metadata FROM spend_logs FINAL \
|
||||
WHERE start_time >= now() - INTERVAL 7 DAY AND length(metadata) <= 8192 \
|
||||
LIMIT 201";
|
||||
const METADATA_SCOPE: &str = "Up to 200 unordered rows from the last 7 days, excluding metadata larger than 8192 bytes; up to 200 paths and 16 levels. Missing paths may exist outside this sample. Array indexes are 1-based and describe sampled positions, not a fixed schema";
|
||||
const ATTRIBUTE_SCOPE: &str = "Distinct keys from up to 200 unordered spans in the last 7 days; up to 200 keys per map. Missing keys may exist outside this sample";
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Rows<T> {
|
||||
data: Vec<T>,
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum ReadQuery {
|
||||
ListTraces,
|
||||
TraceSpans,
|
||||
TraceIdentity,
|
||||
SpanDetail,
|
||||
SpanError,
|
||||
SpendByResponseIds,
|
||||
Availability,
|
||||
Agents,
|
||||
Sample,
|
||||
Content,
|
||||
Evidence,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct MetadataRow {
|
||||
metadata: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AttributeRow {
|
||||
key: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
|
||||
#[serde(untagged)]
|
||||
enum PathPart {
|
||||
Key(String),
|
||||
Index(usize),
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MetadataField {
|
||||
path: Vec<PathPart>,
|
||||
types: BTreeSet<&'static str>,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct ColumnSchema {
|
||||
name: String,
|
||||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
#[serde(flatten)]
|
||||
details: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct TableSchema {
|
||||
name: &'static str,
|
||||
columns: Vec<ColumnSchema>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MetadataCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
fields: Vec<MetadataField>,
|
||||
sampled_rows: usize,
|
||||
invalid_json_rows: usize,
|
||||
truncated: bool,
|
||||
sample_sql: &'static str,
|
||||
scope: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AttributeField {
|
||||
key: String,
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AttributeCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
fields: Vec<AttributeField>,
|
||||
truncated: bool,
|
||||
discovery_sql: String,
|
||||
scope: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn query_sql(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
) -> Result<String, Error> {
|
||||
execute_read(client, connection, sql, &BTreeMap::new()).await
|
||||
}
|
||||
|
||||
async fn rows<T: serde::de::DeserializeOwned>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
) -> Result<Vec<T>, Error> {
|
||||
let body = query_sql(client, connection, sql).await?;
|
||||
serde_json::from_str::<Rows<T>>(&body)
|
||||
.map(|result| result.data)
|
||||
.map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
|
||||
fn literal(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
|
||||
}
|
||||
|
||||
fn metadata_expression(path: &[PathPart]) -> String {
|
||||
let arguments = path
|
||||
.iter()
|
||||
.map(|part| match part {
|
||||
PathPart::Key(key) => literal(key),
|
||||
PathPart::Index(index) => index.to_string(),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
format!("JSONExtractRaw(metadata, {arguments})")
|
||||
}
|
||||
|
||||
fn discover(
|
||||
value: &Value,
|
||||
path: Vec<PathPart>,
|
||||
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
|
||||
) -> bool {
|
||||
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
|
||||
return true;
|
||||
}
|
||||
if !path.is_empty() {
|
||||
let kind = match value {
|
||||
Value::Null => "null",
|
||||
Value::Bool(_) => "boolean",
|
||||
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
|
||||
Value::Number(_) => "number",
|
||||
Value::String(_) => "string",
|
||||
Value::Array(_) => "array",
|
||||
Value::Object(_) => "object",
|
||||
};
|
||||
fields.entry(path.clone()).or_default().insert(kind);
|
||||
}
|
||||
match value {
|
||||
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
|
||||
let child = path
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([PathPart::Key(key.clone())])
|
||||
.collect();
|
||||
discover(value, child, fields) | limited
|
||||
}),
|
||||
Value::Array(array) => array
|
||||
.iter()
|
||||
.enumerate()
|
||||
.fold(false, |limited, (index, value)| {
|
||||
let child = path
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([PathPart::Index(index + 1)])
|
||||
.collect();
|
||||
discover(value, child, fields) | limited
|
||||
}),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
|
||||
let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold(
|
||||
(BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0),
|
||||
|(fields, limited, invalid_rows), row| match serde_json::from_str::<Value>(&row.metadata) {
|
||||
Ok(value) => {
|
||||
let mut fields = fields;
|
||||
let limited = limited | discover(&value, Vec::new(), &mut fields);
|
||||
(fields, limited, invalid_rows)
|
||||
}
|
||||
Err(_) => (fields, limited, invalid_rows + 1),
|
||||
},
|
||||
);
|
||||
let fields: Vec<_> = fields
|
||||
.into_iter()
|
||||
.map(|(path, types)| MetadataField {
|
||||
expression: metadata_expression(&path),
|
||||
path,
|
||||
types,
|
||||
})
|
||||
.collect();
|
||||
MetadataCatalog {
|
||||
table: "spend_logs",
|
||||
column: "metadata",
|
||||
fields,
|
||||
sampled_rows: sample.len().min(SAMPLE_ROWS),
|
||||
invalid_json_rows: invalid_rows,
|
||||
truncated: limited,
|
||||
sample_sql: METADATA_SQL,
|
||||
error: None,
|
||||
scope: METADATA_SCOPE,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_help(client: &Client, connection: &Connection) -> Result<String, Error> {
|
||||
let tables = stream::iter(["otel_traces", "agent_traces_by_key", "spend_logs"])
|
||||
.then(|table| async move {
|
||||
Ok::<_, Error>(TableSchema {
|
||||
name: table,
|
||||
columns: rows::<ColumnSchema>(
|
||||
client,
|
||||
connection,
|
||||
&format!("DESCRIBE TABLE {table}"),
|
||||
)
|
||||
.await?,
|
||||
})
|
||||
})
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
let metadata = match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
|
||||
Ok(sample) => metadata_catalog(&sample),
|
||||
Err(error) => MetadataCatalog {
|
||||
error: Some(error.to_string()),
|
||||
truncated: true,
|
||||
..metadata_catalog(&[])
|
||||
},
|
||||
};
|
||||
let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
|
||||
.then(|column| async move {
|
||||
let sql = format!(
|
||||
"SELECT DISTINCT arrayJoin(mapKeys({column})) AS key FROM \
|
||||
(SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \
|
||||
LIMIT 200) ORDER BY key LIMIT 201"
|
||||
);
|
||||
let (keys, error) = match rows::<AttributeRow>(client, connection, &sql).await {
|
||||
Ok(keys) => (keys, None),
|
||||
Err(error) => (Vec::new(), Some(error.to_string())),
|
||||
};
|
||||
let fields = keys
|
||||
.iter()
|
||||
.take(MAX_FIELDS)
|
||||
.map(|row| AttributeField {
|
||||
key: row.key.clone(),
|
||||
kind: "String",
|
||||
expression: format!("{column}[{}]", literal(&row.key)),
|
||||
})
|
||||
.collect();
|
||||
AttributeCatalog {
|
||||
table: "otel_traces",
|
||||
column,
|
||||
fields,
|
||||
truncated: error.is_some() || keys.len() > MAX_FIELDS,
|
||||
discovery_sql: sql,
|
||||
scope: ATTRIBUTE_SCOPE,
|
||||
error,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
let guide = guide::QueryGuide {
|
||||
tables: &tables,
|
||||
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
|
||||
metadata: &metadata,
|
||||
attributes: &attributes,
|
||||
};
|
||||
Ok(json!({
|
||||
"dialect": "ClickHouse SQL",
|
||||
"access": "Authenticated team scope enforced by ClickHouse row policies; proxy admins can read all teams, while project-bound and teamless keys can read only their own rows",
|
||||
"response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
|
||||
"tables": tables,
|
||||
"normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({
|
||||
"table": "otel_traces", "name": field.name, "column": field.clickhouse_column,
|
||||
"type": field.clickhouse_type, "meaning": field.meaning
|
||||
})).collect::<Vec<_>>(),
|
||||
"metadata": metadata,
|
||||
"attributes": attributes,
|
||||
"relationships": [{
|
||||
"left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id",
|
||||
"additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND otel_traces.ApiKeyHash = spend_logs.api_key",
|
||||
"meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows"
|
||||
}],
|
||||
"examples": guide.examples()?,
|
||||
"gotchas": guide.gotchas()?,
|
||||
"guide": guide::render(&guide)?,
|
||||
}).to_string())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() {
|
||||
let sample = [
|
||||
MetadataRow {
|
||||
metadata: r#"{"x": 1}"#.into(),
|
||||
},
|
||||
MetadataRow {
|
||||
metadata: r#"{"x": "one"}"#.into(),
|
||||
},
|
||||
MetadataRow {
|
||||
metadata: "invalid".into(),
|
||||
},
|
||||
];
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
assert_eq!(
|
||||
catalog["fields"],
|
||||
json!([{
|
||||
"path": ["x"], "types": ["integer", "string"], "expression": "JSONExtractRaw(metadata, 'x')"
|
||||
}])
|
||||
);
|
||||
assert_eq!(catalog["invalid_json_rows"], 1);
|
||||
assert_eq!(catalog["sampled_rows"], sample.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::rows(SAMPLE_ROWS + 1, 1)]
|
||||
#[case::paths(1, MAX_FIELDS + 1)]
|
||||
fn metadata_discovery_reports_truncation(#[case] row_count: usize, #[case] field_count: usize) {
|
||||
let metadata: BTreeMap<_, _> = (0..field_count)
|
||||
.map(|index| (format!("field{index}"), index))
|
||||
.collect();
|
||||
let sample: Vec<_> = (0..row_count)
|
||||
.map(|_| MetadataRow {
|
||||
metadata: json!(metadata).to_string(),
|
||||
})
|
||||
.collect();
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
assert_eq!(catalog["truncated"], true);
|
||||
assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS));
|
||||
assert_eq!(
|
||||
catalog["fields"].as_array().unwrap().len(),
|
||||
field_count.min(MAX_FIELDS)
|
||||
);
|
||||
impl ReadQuery {
|
||||
pub fn parse(value: &str) -> Result<Self, crate::InvalidQuery> {
|
||||
value.parse().map_err(|_| crate::InvalidQuery)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
151
litellm-rust/crates/traces/src/query/named.rs
Normal file
151
litellm-rust/crates/traces/src/query/named.rs
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ReadAccessParams {
|
||||
pub all_teams: u8,
|
||||
pub user_id: String,
|
||||
pub team_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub start_ms: i64,
|
||||
pub end_ms: i64,
|
||||
pub cursor_ms: i64,
|
||||
pub cursor_trace_id: String,
|
||||
pub limit: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesRow {
|
||||
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,
|
||||
pub start_ms: i64,
|
||||
pub duration_ms: i64,
|
||||
pub span_count: u64,
|
||||
pub agent_count: u64,
|
||||
pub agent_invocations: u64,
|
||||
#[serde(default)]
|
||||
pub agent_names: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub frameworks: Vec<String>,
|
||||
pub llm_calls: u64,
|
||||
pub tool_calls: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub models: Vec<String>,
|
||||
pub error_count: u64,
|
||||
pub request_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceSpansParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceSpansRow {
|
||||
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,
|
||||
pub error_truncated: u8,
|
||||
pub start_ns: i64,
|
||||
pub duration_ns: u64,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub model: String,
|
||||
pub input_tokens: u32,
|
||||
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 SpanDetailParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub span_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanDetailRow {
|
||||
pub span_id: String,
|
||||
pub input: String,
|
||||
pub output: String,
|
||||
pub attributes: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub span_id: String,
|
||||
pub error_offset: u64,
|
||||
pub error_version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorRow {
|
||||
pub span_id: String,
|
||||
pub message: String,
|
||||
pub total_chars: u64,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub response_ids: Vec<String>,
|
||||
pub start_ms: i64,
|
||||
pub end_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsRow {
|
||||
pub request_id: String,
|
||||
pub response_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key: String,
|
||||
pub user: String,
|
||||
pub spend: f64,
|
||||
pub start_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceIdentityParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceIdentityRow {
|
||||
pub trace_ref: String,
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue