mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
chore: merge main into the prompt-cache token count seam branch
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
Some checks failed
LiteLLM Rust / rust-lint (push) Has been cancelled
LiteLLM Rust / rust-test (push) Has been cancelled
LiteLLM Rust / rust-wheel (push) Has been cancelled
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
main already patches _get_messages_reach_token_count in this test, so the branch keeps main's copy and no longer carries a change of its own.
This commit is contained in:
commit
0fb0fa1055
945 changed files with 90572 additions and 31931 deletions
|
|
@ -41,6 +41,7 @@ legacy_paths() {
|
|||
echo tests/unit/enterprise/proxy/hooks
|
||||
echo tests/unit/enterprise/proxy/management_endpoints
|
||||
echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py
|
||||
echo tests/unit/enterprise/proxy/test_liteadmin.py
|
||||
echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;;
|
||||
enterprise-routing)
|
||||
echo tests/unit/google_genai
|
||||
|
|
|
|||
65
.github/scripts/assert_ci_coverage.py
vendored
65
.github/scripts/assert_ci_coverage.py
vendored
|
|
@ -130,37 +130,24 @@ def _unit_selection_arms(repo_root: pathlib.Path = REPO_ROOT) -> Mapping[str, fr
|
|||
text: Final = _uncommented(script.read_text())
|
||||
return MappingProxyType(
|
||||
{
|
||||
label: frozenset(
|
||||
match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(body)
|
||||
)
|
||||
label: frozenset(match.group(0).rstrip("/") for match in TEST_TOKEN_RE.finditer(body))
|
||||
for label, body in SELECTION_ARM_RE.findall(text)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _unit_selection_tokens(repo_root: pathlib.Path = REPO_ROOT) -> frozenset[str]:
|
||||
return frozenset(
|
||||
token for tokens in _unit_selection_arms(repo_root).values() for token in tokens
|
||||
)
|
||||
return frozenset(token for tokens in _unit_selection_arms(repo_root).values() for token in tokens)
|
||||
|
||||
|
||||
def _wired_unit_flags(scalars: Iterable[Scalar]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
scalar.value
|
||||
for scalar in scalars
|
||||
if scalar.key == "unit-flag" and "${{" not in scalar.value
|
||||
)
|
||||
return frozenset(scalar.value for scalar in scalars if scalar.key == "unit-flag" and "${{" not in scalar.value)
|
||||
|
||||
|
||||
def _shard_tokens(
|
||||
scalars: Iterable[Scalar], arms: Mapping[str, frozenset[str]]
|
||||
) -> frozenset[str]:
|
||||
def _shard_tokens(scalars: Iterable[Scalar], arms: Mapping[str, frozenset[str]]) -> frozenset[str]:
|
||||
wired: Final = _wired_unit_flags(scalars)
|
||||
return _invoked_test_tokens(scalars) | frozenset(
|
||||
token
|
||||
for label, tokens in arms.items()
|
||||
if label in wired
|
||||
for token in tokens
|
||||
token for label, tokens in arms.items() if label in wired for token in tokens
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -544,17 +531,37 @@ def _integration_groups(runner: pathlib.Path) -> dict[str, tuple[str, ...]]:
|
|||
return {group: tuple(folders) for group, folders in ast.literal_eval(mapping).items()}
|
||||
|
||||
|
||||
def _integration_github_files(runner: pathlib.Path) -> frozenset[str]:
|
||||
module: Final = ast.parse(runner.read_text())
|
||||
literal: Final = next(
|
||||
(
|
||||
node.value
|
||||
for node in module.body
|
||||
if isinstance(node, ast.AnnAssign)
|
||||
and isinstance(node.target, ast.Name)
|
||||
and node.target.id == "GITHUB_FILES"
|
||||
),
|
||||
None,
|
||||
)
|
||||
if literal is None:
|
||||
return frozenset()
|
||||
values: Final = literal.args[0] if isinstance(literal, ast.Call) else literal
|
||||
return frozenset(ast.literal_eval(values))
|
||||
|
||||
|
||||
def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozenset[str], tuple[Finding, ...]]:
|
||||
runner: Final = repo_root / "tests/integration/run.py"
|
||||
if not runner.exists():
|
||||
return frozenset(), ()
|
||||
groups: Final = _integration_groups(runner)
|
||||
github_files: Final = _integration_github_files(runner)
|
||||
integration_root: Final = repo_root / "tests/integration"
|
||||
paths: Final = frozenset(
|
||||
str(path.relative_to(repo_root))
|
||||
for folders in groups.values()
|
||||
for folder in folders
|
||||
for path in (integration_root / folder).rglob("test_*.py")
|
||||
if str(path.relative_to(repo_root)) not in github_files
|
||||
)
|
||||
browser_manifest: Final = repo_root / "tests/e2e/ui/tests/integrationCritical/expected.json"
|
||||
browser_nodes: Final = json.loads(browser_manifest.read_text()) if browser_manifest.exists() else ()
|
||||
|
|
@ -595,10 +602,22 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
for path in (repo_root / ".github/workflows").glob("*.y*ml")
|
||||
for scalar in _scalars(yaml.safe_load(path.read_text()), path.name)
|
||||
)
|
||||
findings: Final = tuple(
|
||||
Finding(path, "integration contract is also selected by GitHub Actions")
|
||||
for path in paths
|
||||
if any(_token_covers(token, path) for token in gha_tokens)
|
||||
findings: Final = (
|
||||
tuple(
|
||||
Finding(path, "integration contract is also selected by GitHub Actions")
|
||||
for path in paths
|
||||
if any(_token_covers(token, path) for token in gha_tokens)
|
||||
)
|
||||
+ tuple(
|
||||
Finding(path, "GitHub-owned integration contract has no invoking workflow")
|
||||
for path in sorted(github_files)
|
||||
if not any(_token_covers(token, path) for token in gha_tokens)
|
||||
)
|
||||
+ tuple(
|
||||
Finding(path, "GitHub-owned integration file is missing")
|
||||
for path in sorted(github_files)
|
||||
if not (repo_root / path).is_file()
|
||||
)
|
||||
)
|
||||
browser_commands: Final = tuple(
|
||||
scalar.value
|
||||
|
|
@ -642,7 +661,7 @@ def _integration_ownership(repo_root: pathlib.Path = REPO_ROOT) -> tuple[frozens
|
|||
return frozenset(), findings + (
|
||||
Finding(str(runner.relative_to(repo_root)), "dedicated CircleCI runner is missing"),
|
||||
)
|
||||
return paths | browser_paths, findings + group_findings + browser_findings + exclusion_findings
|
||||
return paths | browser_paths | github_files, findings + group_findings + browser_findings + exclusion_findings
|
||||
|
||||
|
||||
def main() -> int:
|
||||
|
|
|
|||
84
.github/workflows/image-scan.yml
vendored
84
.github/workflows/image-scan.yml
vendored
|
|
@ -15,6 +15,9 @@ on:
|
|||
- gateway/main.py
|
||||
- backend/Dockerfile
|
||||
- backend/main.py
|
||||
- deploy/lens/**
|
||||
- litellm/proxy/lens/**
|
||||
- tests/e2e/migrations/lens_compose_smoke.sh
|
||||
- docker/component_entrypoint.sh
|
||||
- docker/entrypoint.sh
|
||||
- litellm/proxy/prisma_migration.py
|
||||
|
|
@ -37,6 +40,80 @@ concurrency:
|
|||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
lens-worker-image:
|
||||
name: lens-worker-image (${{ matrix.arch }})
|
||||
runs-on: ${{ matrix.runner }}
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 15
|
||||
permissions:
|
||||
contents: read
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- arch: amd64
|
||||
runner: ubuntu-latest
|
||||
grype_sha256: edda0968d8827daab01d32b3cd7de192ae0915005e7bbfcfef9e68e79bc43343
|
||||
- arch: arm64
|
||||
runner: ubuntu-24.04-arm
|
||||
grype_sha256: 553e4c36d9d61349830ba6034d43b8700a7f10576d3e2f4981c0fd2b96086465
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
- name: Build the release worker
|
||||
env:
|
||||
RELEASE_TAG: sha-${{ github.sha }}
|
||||
run: docker build --build-arg LITELLM_RELEASE_TAG="${RELEASE_TAG}" -f deploy/lens/Dockerfile -t lens-worker-scan .
|
||||
- name: Verify the standalone worker on a read-only filesystem
|
||||
env:
|
||||
RELEASE_TAG: sha-${{ github.sha }}
|
||||
run: |
|
||||
docker run --rm --network none --read-only --cap-drop ALL \
|
||||
--tmpfs /tmp:rw,noexec,nosuid,size=1g --security-opt no-new-privileges \
|
||||
-e EXPECTED_RELEASE_TAG="${RELEASE_TAG}" --entrypoint python lens-worker-scan -c '
|
||||
import os
|
||||
import lens.worker
|
||||
from lens.release import release_tag
|
||||
from lens.trace_store import trace_store
|
||||
assert os.getuid() == 65532
|
||||
assert release_tag() == os.environ["EXPECTED_RELEASE_TAG"]
|
||||
with trace_store() as store:
|
||||
assert store.count() == 0
|
||||
'
|
||||
- name: Reject a dependency whose hash has changed
|
||||
run: |
|
||||
docker build --target builder -f deploy/lens/Dockerfile -t lens-worker-deps .
|
||||
sed -E 's/sha256:[0-9a-f]{64}/sha256:0000000000000000000000000000000000000000000000000000000000000000/g' \
|
||||
deploy/lens/requirements.lock > "$RUNNER_TEMP/tampered.lock"
|
||||
if docker run --rm -v "$RUNNER_TEMP/tampered.lock:/tmp/tampered.lock:ro" \
|
||||
--entrypoint uv lens-worker-deps pip sync --python /app/.venv/bin/python \
|
||||
--require-hashes --only-binary :all: --reinstall --no-cache /tmp/tampered.lock \
|
||||
> "$RUNNER_TEMP/hash-check.log" 2>&1; then
|
||||
echo "::error::Dependency hash mismatch was accepted"
|
||||
exit 1
|
||||
fi
|
||||
cat "$RUNNER_TEMP/hash-check.log"
|
||||
grep -qi 'hash mismatch' "$RUNNER_TEMP/hash-check.log"
|
||||
- name: Download Grype v0.114.0
|
||||
env:
|
||||
ARCH: ${{ matrix.arch }}
|
||||
GRYPE_SHA256: ${{ matrix.grype_sha256 }}
|
||||
run: |
|
||||
curl -fsSL --retry 3 -o "$RUNNER_TEMP/grype.tar.gz" \
|
||||
"https://github.com/anchore/grype/releases/download/v0.114.0/grype_0.114.0_linux_${ARCH}.tar.gz"
|
||||
echo "${GRYPE_SHA256} $RUNNER_TEMP/grype.tar.gz" | sha256sum -c -
|
||||
tar xzf "$RUNNER_TEMP/grype.tar.gz" -C "$RUNNER_TEMP" grype
|
||||
chmod +x "$RUNNER_TEMP/grype"
|
||||
- name: Scan the worker for fixable HIGH/CRITICAL CVEs
|
||||
env:
|
||||
GRYPE_MATCH_PYTHON_USING_CPES: "true"
|
||||
run: |
|
||||
"$RUNNER_TEMP/grype" lens-worker-scan \
|
||||
--config .grype.yaml --only-fixed --fail-on high --output table
|
||||
|
||||
image-scan:
|
||||
name: image-scan
|
||||
runs-on: ubuntu-latest
|
||||
|
|
@ -113,7 +190,7 @@ jobs:
|
|||
persist-credentials: false
|
||||
|
||||
- name: Build runtime image
|
||||
run: docker build -f Dockerfile -t litellm-runtime-scan:${{ github.sha }} .
|
||||
run: docker build --build-arg LITELLM_RELEASE_TAG=v0.0.0-lens-ci -f Dockerfile -t litellm-runtime-scan:${{ github.sha }} .
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
|
|
@ -127,6 +204,11 @@ jobs:
|
|||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py tests/proxy_migration_tests/test_image_bedrock_realtime_extra.py -v
|
||||
|
||||
- name: Verify the bundled Lens Compose installation and restart
|
||||
env:
|
||||
LITELLM_IMAGE: litellm-runtime-scan:${{ github.sha }}
|
||||
run: bash tests/e2e/migrations/lens_compose_smoke.sh
|
||||
|
||||
migrations-image:
|
||||
name: migrations-image
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
13
.github/workflows/lens-worker.yml
vendored
13
.github/workflows/lens-worker.yml
vendored
|
|
@ -34,7 +34,14 @@ jobs:
|
|||
with:
|
||||
persist-credentials: false
|
||||
- name: Build Lens worker
|
||||
run: docker build -f deploy/lens/Dockerfile -t lens-worker:${{ github.sha }} .
|
||||
run: docker build --build-arg LITELLM_RELEASE_TAG=sha-${{ github.sha }} -f deploy/lens/Dockerfile -t lens-worker:${{ github.sha }} .
|
||||
- name: Reject custom builds without a matching release tag
|
||||
run: |
|
||||
if docker build --progress plain -f deploy/lens/Dockerfile -t lens-worker:unversioned . > missing-tag.log 2>&1; then
|
||||
echo "::error::An unversioned worker build unexpectedly succeeded"
|
||||
exit 1
|
||||
fi
|
||||
grep -F 'LITELLM_RELEASE_TAG: Pass --build-arg LITELLM_RELEASE_TAG matching the gateway' missing-tag.log
|
||||
- name: Verify standalone imports with a read-only filesystem
|
||||
run: |
|
||||
docker run --rm --network none --read-only --cap-drop ALL --tmpfs /tmp:rw,noexec,nosuid,size=1g \
|
||||
|
|
@ -54,11 +61,11 @@ jobs:
|
|||
-v "$PWD/tests/proxy_behavior/lens/worker_storage_smoke.py:/app/storage_smoke.py:ro" \
|
||||
--entrypoint python lens-worker:${{ github.sha }} /app/storage_smoke.py
|
||||
- name: Publish versioned Lens worker
|
||||
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm'
|
||||
if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm' && github.ref == 'refs/heads/main'
|
||||
env:
|
||||
REGISTRY_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
REGISTRY_USER: ${{ github.actor }}
|
||||
IMAGE: ghcr.io/berriai/litellm-lens-worker:sha-${{ github.sha }}
|
||||
IMAGE: ghcr.io/berriai/litellm-lens-worker-dev:sha-${{ github.sha }}
|
||||
run: |
|
||||
printf '%s' "$REGISTRY_TOKEN" | docker login ghcr.io -u "$REGISTRY_USER" --password-stdin
|
||||
docker tag lens-worker:${{ github.sha }} "$IMAGE"
|
||||
|
|
|
|||
20
.github/workflows/test-postgres.yml
vendored
20
.github/workflows/test-postgres.yml
vendored
|
|
@ -45,6 +45,13 @@ jobs:
|
|||
fail-fast: false
|
||||
matrix:
|
||||
include:
|
||||
- shard: roi-database
|
||||
test-path: "tests/integration/database/test_roi_observed.py"
|
||||
seed: none
|
||||
workers: 0
|
||||
timeout-minutes: 10
|
||||
job-timeout-minutes: 35
|
||||
|
||||
- shard: proxy-behavior
|
||||
test-path: "tests/proxy_behavior"
|
||||
seed: db-push
|
||||
|
|
@ -147,7 +154,7 @@ jobs:
|
|||
env:
|
||||
TEST_PATH: ${{ matrix.test-path }}
|
||||
WORKERS: ${{ matrix.workers }}
|
||||
PYTEST_ADDOPTS: ${{ matrix.shard == 'proxy-behavior' && '--cov=./litellm --cov-report=xml:coverage-lens-postgres.xml' || '' }}
|
||||
PYTEST_ADDOPTS: ${{ matrix.shard == 'proxy-behavior' && '--cov=./litellm --cov-report=xml:coverage-lens-postgres.xml' || matrix.shard == 'roi-database' && '--cov=./litellm --cov-report=xml:coverage-roi-postgres.xml' || '' }}
|
||||
run: |
|
||||
if [ "${WORKERS}" = "0" ]; then
|
||||
uv run --no-sync pytest ${TEST_PATH:?} -vv --tb=short --durations=10
|
||||
|
|
@ -165,3 +172,14 @@ jobs:
|
|||
files: coverage-lens-postgres.xml
|
||||
flags: lens-postgres
|
||||
fail_ci_if_error: true
|
||||
|
||||
- name: Upload ROI database coverage
|
||||
if: steps.changes.outputs.decision != 'skip' && matrix.shard == 'roi-database' && !cancelled()
|
||||
uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1
|
||||
with:
|
||||
use_oidc: true
|
||||
version: v11.3.1
|
||||
root_dir: ${{ github.workspace }}
|
||||
files: coverage-roi-postgres.xml
|
||||
flags: roi-postgres
|
||||
fail_ci_if_error: true
|
||||
|
|
|
|||
13
Dockerfile
13
Dockerfile
|
|
@ -114,8 +114,20 @@ RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/
|
|||
RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
||||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
||||
FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder
|
||||
COPY --from=uvbin /uv /usr/local/bin/uv
|
||||
RUN apk add --no-cache python-3.13
|
||||
ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \
|
||||
https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz
|
||||
RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \
|
||||
uv venv /opt/liteadmin --python python3.13 && \
|
||||
uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \
|
||||
uv pip install --python /opt/liteadmin/bin/python --no-deps /tmp/liteadmin
|
||||
|
||||
# Runtime stage
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
ARG LITELLM_RELEASE_TAG=""
|
||||
ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
|
||||
|
||||
USER root
|
||||
|
||||
|
|
@ -141,6 +153,7 @@ ENV PATH="/app/.venv/bin:${PATH}" \
|
|||
# ship (manifest-scanning tools attribute everything in it to this image).
|
||||
# entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path.
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY --from=liteadmin-builder /opt/liteadmin /opt/liteadmin
|
||||
COPY --from=builder /app/docker /app/docker
|
||||
COPY --from=builder /app/schema.prisma /app/schema.prisma
|
||||
COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py
|
||||
|
|
|
|||
4
Makefile
4
Makefile
|
|
@ -58,7 +58,7 @@ help:
|
|||
@echo " make test-unit-helm - Run helm unit tests"
|
||||
@echo " make test-rust-extension - Build the Rust extension and run its public Python tests"
|
||||
@echo " make rust-sqlx-prepare - Refresh litellm-rust/crates/db/.sqlx against a migrated Postgres container"
|
||||
@echo " make lens-dev - Run proxy + Lens worker + hot-reload dashboard (LENS_DEV_PROXY_PORT, LENS_DEV_UI_PORT)"
|
||||
@echo " make lens-dev - Run proxy + Lens worker + hot-reload dashboard (ARGS=\"--seed large --seed-logs\", LENS_DEV_PROXY_PORT, LENS_DEV_UI_PORT)"
|
||||
@echo ""
|
||||
@echo "Heavy targets (check, lint) queue for LITELLM_GATE_SLOTS machine-wide"
|
||||
@echo "slots (default 2; 0 disables) so parallel sessions don't thrash one machine."
|
||||
|
|
@ -313,7 +313,7 @@ rust-sqlx-prepare:
|
|||
cd litellm-rust && cargo run -p litellm-db-testing --bin sqlx-prepare
|
||||
|
||||
lens-dev:
|
||||
./scripts/lens_dev.sh
|
||||
./scripts/lens_dev.sh $(ARGS)
|
||||
|
||||
test: install-test-deps
|
||||
$(UV_RUN) pytest tests/
|
||||
|
|
|
|||
|
|
@ -71,6 +71,8 @@ RUN sed -i 's/\r$//' docker/component_entrypoint.sh && chmod +x docker/component
|
|||
|
||||
# ---------- Runtime ----------
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
ARG LITELLM_RELEASE_TAG=""
|
||||
ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
|
||||
|
||||
USER root
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = (
|
|||
"/customer/",
|
||||
"/end_user/",
|
||||
"/sso/",
|
||||
"/liteadmin/slack/connect/",
|
||||
"/login",
|
||||
"/v2/login",
|
||||
"/v3/login",
|
||||
|
|
|
|||
|
|
@ -5710,6 +5710,17 @@
|
|||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "histogram_quantile(0.95, sum(rate(litellm_anthropic_wif_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "anthropic_wif",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
|
|
@ -5719,7 +5730,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_auth_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "auth",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
"refId": "B"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5730,7 +5741,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_batch_write_to_db_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "batch_write_to_db",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
"refId": "C"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5741,7 +5752,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_postgres_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "postgres",
|
||||
"range": true,
|
||||
"refId": "C"
|
||||
"refId": "D"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5752,7 +5763,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_proxy_pre_call_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "proxy_pre_call",
|
||||
"range": true,
|
||||
"refId": "D"
|
||||
"refId": "E"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5763,7 +5774,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "redis",
|
||||
"range": true,
|
||||
"refId": "E"
|
||||
"refId": "F"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5774,7 +5785,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_org_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "redis_daily_org_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "F"
|
||||
"refId": "G"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5785,7 +5796,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_tag_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "redis_daily_tag_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "G"
|
||||
"refId": "H"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5796,7 +5807,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_daily_team_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "redis_daily_team_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "H"
|
||||
"refId": "I"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5807,7 +5818,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_redis_window_spend_update_queue_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "redis_window_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "I"
|
||||
"refId": "J"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5818,7 +5829,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_reset_budget_job_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "reset_budget_job",
|
||||
"range": true,
|
||||
"refId": "J"
|
||||
"refId": "K"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5829,7 +5840,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_router_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "router",
|
||||
"range": true,
|
||||
"refId": "K"
|
||||
"refId": "L"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5840,7 +5851,7 @@
|
|||
"expr": "histogram_quantile(0.95, sum(rate(litellm_self_latency_bucket[$__rate_interval])) by (le))",
|
||||
"legendFormat": "self",
|
||||
"range": true,
|
||||
"refId": "L"
|
||||
"refId": "M"
|
||||
}
|
||||
],
|
||||
"title": "Service latency p95 (litellm_<service>_latency)",
|
||||
|
|
@ -5888,6 +5899,28 @@
|
|||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_anthropic_wif_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "anthropic_wif",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_anthropic_wif_cache_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "anthropic_wif_cache",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
|
|
@ -5897,7 +5930,7 @@
|
|||
"expr": "sum(rate(litellm_auth_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "auth",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
"refId": "C"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5908,7 +5941,7 @@
|
|||
"expr": "sum(rate(litellm_batch_write_to_db_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "batch_write_to_db",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
"refId": "D"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5919,7 +5952,7 @@
|
|||
"expr": "sum(rate(litellm_postgres_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "postgres",
|
||||
"range": true,
|
||||
"refId": "C"
|
||||
"refId": "E"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5930,7 +5963,7 @@
|
|||
"expr": "sum(rate(litellm_proxy_pre_call_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "proxy_pre_call",
|
||||
"range": true,
|
||||
"refId": "D"
|
||||
"refId": "F"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5941,7 +5974,7 @@
|
|||
"expr": "sum(rate(litellm_redis_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "redis",
|
||||
"range": true,
|
||||
"refId": "E"
|
||||
"refId": "G"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5952,7 +5985,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_org_spend_update_queue_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "redis_daily_org_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "F"
|
||||
"refId": "H"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5963,7 +5996,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_tag_spend_update_queue_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "redis_daily_tag_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "G"
|
||||
"refId": "I"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5974,7 +6007,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_team_spend_update_queue_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "redis_daily_team_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "H"
|
||||
"refId": "J"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5985,7 +6018,7 @@
|
|||
"expr": "sum(rate(litellm_redis_window_spend_update_queue_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "redis_window_spend_update_queue",
|
||||
"range": true,
|
||||
"refId": "I"
|
||||
"refId": "K"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -5996,7 +6029,7 @@
|
|||
"expr": "sum(rate(litellm_reset_budget_job_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "reset_budget_job",
|
||||
"range": true,
|
||||
"refId": "J"
|
||||
"refId": "L"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6007,7 +6040,7 @@
|
|||
"expr": "sum(rate(litellm_router_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "router",
|
||||
"range": true,
|
||||
"refId": "K"
|
||||
"refId": "M"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6018,7 +6051,7 @@
|
|||
"expr": "sum(rate(litellm_self_total_requests_total[$__rate_interval]))",
|
||||
"legendFormat": "self",
|
||||
"range": true,
|
||||
"refId": "L"
|
||||
"refId": "N"
|
||||
}
|
||||
],
|
||||
"title": "Service request rate (litellm_<service>_total_requests)",
|
||||
|
|
@ -6066,6 +6099,28 @@
|
|||
}
|
||||
},
|
||||
"targets": [
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_anthropic_wif_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "anthropic_wif / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
"uid": "${DS_PROMETHEUS}"
|
||||
},
|
||||
"editorMode": "code",
|
||||
"expr": "sum(rate(litellm_anthropic_wif_cache_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "anthropic_wif_cache / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
"type": "prometheus",
|
||||
|
|
@ -6075,7 +6130,7 @@
|
|||
"expr": "sum(rate(litellm_auth_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "auth / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "A"
|
||||
"refId": "C"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6086,7 +6141,7 @@
|
|||
"expr": "sum(rate(litellm_batch_write_to_db_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "batch_write_to_db / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "B"
|
||||
"refId": "D"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6097,7 +6152,7 @@
|
|||
"expr": "sum(rate(litellm_postgres_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "postgres / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "C"
|
||||
"refId": "E"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6108,7 +6163,7 @@
|
|||
"expr": "sum(rate(litellm_proxy_pre_call_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "proxy_pre_call / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "D"
|
||||
"refId": "F"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6119,7 +6174,7 @@
|
|||
"expr": "sum(rate(litellm_redis_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "redis / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "E"
|
||||
"refId": "G"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6130,7 +6185,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_org_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "redis_daily_org_spend_update_queue / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "F"
|
||||
"refId": "H"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6141,7 +6196,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_tag_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "redis_daily_tag_spend_update_queue / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "G"
|
||||
"refId": "I"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6152,7 +6207,7 @@
|
|||
"expr": "sum(rate(litellm_redis_daily_team_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "redis_daily_team_spend_update_queue / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "H"
|
||||
"refId": "J"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6163,7 +6218,7 @@
|
|||
"expr": "sum(rate(litellm_redis_window_spend_update_queue_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "redis_window_spend_update_queue / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "I"
|
||||
"refId": "K"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6174,7 +6229,7 @@
|
|||
"expr": "sum(rate(litellm_reset_budget_job_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "reset_budget_job / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "J"
|
||||
"refId": "L"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6185,7 +6240,7 @@
|
|||
"expr": "sum(rate(litellm_router_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "router / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "K"
|
||||
"refId": "M"
|
||||
},
|
||||
{
|
||||
"datasource": {
|
||||
|
|
@ -6196,7 +6251,7 @@
|
|||
"expr": "sum(rate(litellm_self_failed_requests_total[$__rate_interval])) by (error_class)",
|
||||
"legendFormat": "self / {{error_class}}",
|
||||
"range": true,
|
||||
"refId": "L"
|
||||
"refId": "N"
|
||||
}
|
||||
],
|
||||
"title": "Service failure rate (litellm_<service>_failed_requests)",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
# LiteLLM All Prometheus Metrics dashboard
|
||||
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (136 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about
|
||||
|
||||
Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,28 @@
|
|||
FROM python:3.12-slim
|
||||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
||||
FROM $LITELLM_BUILD_IMAGE AS builder
|
||||
COPY --from=uvbin /uv /usr/local/bin/uv
|
||||
RUN apk add --no-cache python-3.13
|
||||
ENV UV_PYTHON_DOWNLOADS=0 UV_LINK_MODE=copy
|
||||
WORKDIR /app
|
||||
RUN pip install --no-cache-dir httpx==0.28.1 pydantic==2.11.7
|
||||
COPY litellm/proxy/lens/__init__.py litellm/proxy/lens/models.py litellm/proxy/lens/trace_store.py litellm/proxy/lens/analysis.py litellm/proxy/lens/worker.py /app/lens/
|
||||
COPY deploy/lens/requirements.lock /tmp/requirements.lock
|
||||
RUN uv venv --python python3.13 /app/.venv && \
|
||||
uv pip sync --python /app/.venv/bin/python --require-hashes --only-binary :all: /tmp/requirements.lock
|
||||
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
ARG LITELLM_RELEASE_TAG=""
|
||||
RUN : "${LITELLM_RELEASE_TAG:?Pass --build-arg LITELLM_RELEASE_TAG matching the gateway}"
|
||||
RUN apk add --no-cache python-3.13
|
||||
ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} \
|
||||
PATH="/app/.venv/bin:${PATH}" \
|
||||
PYTHONDONTWRITEBYTECODE=1
|
||||
WORKDIR /app
|
||||
COPY --from=builder /app/.venv /app/.venv
|
||||
COPY litellm/proxy/lens/__init__.py litellm/proxy/lens/models.py litellm/proxy/lens/trace_store.py litellm/proxy/lens/analysis.py litellm/proxy/lens/worker.py litellm/proxy/lens/release.py /app/lens/
|
||||
COPY litellm/proxy/lens/prompts/ /app/lens/prompts/
|
||||
USER 65532:65532
|
||||
CMD ["python", "-m", "lens.worker"]
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
**
|
||||
!deploy/
|
||||
!deploy/lens/
|
||||
!deploy/lens/requirements.lock
|
||||
!litellm/
|
||||
!litellm/proxy/
|
||||
!litellm/proxy/lens/
|
||||
|
|
@ -7,5 +10,6 @@
|
|||
!litellm/proxy/lens/trace_store.py
|
||||
!litellm/proxy/lens/analysis.py
|
||||
!litellm/proxy/lens/worker.py
|
||||
!litellm/proxy/lens/release.py
|
||||
!litellm/proxy/lens/prompts/
|
||||
!litellm/proxy/lens/prompts/**
|
||||
|
|
|
|||
|
|
@ -2,9 +2,74 @@
|
|||
|
||||
Lens reviews recorded activity and saves evidence-linked findings in the LiteLLM dashboard under Observability, Lens (`/ui/lens/`)
|
||||
|
||||
## Start a worker
|
||||
## Install
|
||||
|
||||
Upgrade your existing LiteLLM proxy to a release that includes Lens with PostgreSQL and agent tracing. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries:
|
||||
Build LiteLLM and its worker from the same source commit with the same release identity. The worker runs separately and connects to your gateway using a limited worker token
|
||||
|
||||
### New local installation
|
||||
|
||||
Install Docker with Compose and Git. This builds LiteLLM and its worker from the same checkout and starts the existing local tracing stack:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/BerriAI/litellm.git
|
||||
cd litellm
|
||||
export LITELLM_RELEASE_TAG="sha-$(git rev-parse HEAD)"
|
||||
export LENS_WORKER_IMAGE="litellm-lens-worker:${LITELLM_RELEASE_TAG}"
|
||||
export OPENAI_API_KEY='sk-...'
|
||||
docker build --build-arg LITELLM_RELEASE_TAG="$LITELLM_RELEASE_TAG" \
|
||||
-f deploy/lens/Dockerfile -t "$LENS_WORKER_IMAGE" .
|
||||
docker compose -f docker/docker-compose.tracing.yml up -d --build
|
||||
```
|
||||
|
||||
Open `http://localhost:4002/ui/` and sign in as `admin` with password `sk-1234`. Go to **Lens > Investigations > Connect worker**, choose a model and monthly budget, then **Get install command**. Expand **Using Docker Compose or Helm?** and copy the worker token. In the same terminal, run:
|
||||
|
||||
```bash
|
||||
export LITELLM_URL=http://litellm:4000
|
||||
export LENS_WORKER_TOKEN='<paste-your-worker-token>'
|
||||
docker compose -f docker/docker-compose.tracing.yml -f deploy/lens/compose.yaml up -d
|
||||
```
|
||||
|
||||
The worker joins the gateway's Docker network, and the dashboard shows **Worker connected**. Save the token privately for restarts and upgrades
|
||||
|
||||
This stack is for local evaluation: it binds to localhost and uses development database credentials. For a hosted deployment, keep your normal database, keys, networking, and deployment process. Build both images from one source revision with the same `LITELLM_RELEASE_TAG`, publish the worker to your registry, and set `LENS_WORKER_IMAGE` on LiteLLM to that image
|
||||
|
||||
### Existing LiteLLM installation
|
||||
|
||||
Keep your deployment and PostgreSQL database. A working gateway/worker pair can stay as it is until you upgrade both. For a gateway built from source, use its exact commit and `LITELLM_RELEASE_TAG`; a release version or the latest commit on `main` is not a substitute for that source identity
|
||||
|
||||
The public development package is `ghcr.io/berriai/litellm-lens-worker-dev:sha-<full-commit>`. It publishes amd64 images on Lens-related changes, so an arbitrary source commit may have no image. Check the exact image exists before using it. If it is unavailable, your gateway uses a different release identity, or you need native arm64, build the worker from the gateway's checkout:
|
||||
|
||||
```bash
|
||||
export LITELLM_RELEASE_TAG='<gateway-release-identity>'
|
||||
export LENS_WORKER_IMAGE='<your-registry>/litellm-lens-worker:<your-image-tag>'
|
||||
docker build --build-arg LITELLM_RELEASE_TAG="$LITELLM_RELEASE_TAG" \
|
||||
-f deploy/lens/Dockerfile -t "$LENS_WORKER_IMAGE" .
|
||||
```
|
||||
|
||||
For a remote worker host, publish that image to a registry the host can pull from. Set the gateway's `LENS_WORKER_IMAGE` to the resulting image reference, restart the gateway using its normal deployment process, then copy its install command. Prefer the published image digest for hosted installations. Do not change the gateway's release identity just to accept another worker
|
||||
|
||||
For Kubernetes or Render, run the standalone worker using `LITELLM_URL` and `LENS_WORKER_TOKEN` from setup. Keep existing databases and secrets. The worker needs no inbound port.
|
||||
|
||||
## Helm
|
||||
|
||||
The componentized source chart at `helm/litellm` includes an optional Lens worker. Use the chart from the same checkout as your gateway and keep your component image overrides in your values. Configure PostgreSQL and ClickHouse as usual, install the chart, then obtain a limited worker token from Lens setup. Store it in a Kubernetes Secret and enable the worker in your values:
|
||||
|
||||
```yaml
|
||||
lensWorker:
|
||||
enabled: true
|
||||
image:
|
||||
repository: <your-worker-image-repository>
|
||||
digest: sha256:<matching-worker-image-digest>
|
||||
tokenSecret:
|
||||
name: litellm-lens-worker
|
||||
key: token
|
||||
```
|
||||
|
||||
Set the worker repository and digest explicitly to an image built from the gateway's source commit and release identity. The chart connects the worker to the backend service. Keep these values and the Secret when upgrading the chart and update the gateway and worker image overrides together. `lensWorker.replicaCount` controls simultaneous investigations. To use a private registry or external proxy, set `lensWorker.image.repository`, `lensWorker.image.digest` (or `tag` for a source build), and `lensWorker.url`. A digest takes precedence over the tag. The dashboard uses the chart's worker image for standalone install commands too
|
||||
|
||||
## Standalone worker
|
||||
|
||||
Start with a source deployment that includes Lens, PostgreSQL, and agent tracing, and prepare its matching worker as described above. Configure one ClickHouse URL for trace writes, bounded reads, and Lens queries:
|
||||
|
||||
```yaml
|
||||
general_settings:
|
||||
|
|
@ -21,19 +86,19 @@ Retention changes require a proxy restart. ClickHouse removes expired rows durin
|
|||
|
||||
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 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. Once the matching image is available on the worker host, no 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. CI also publishes immutable `:sha-<commit>` tags for successful worker builds on `main`. Keep the worker image compatible with your gateway version
|
||||
The dashboard uses the gateway's `LENS_WORKER_IMAGE` override when set. Public `:sha-<commit>` development images must match both the gateway commit and release identity. Build from source for the worker host's native architecture
|
||||
|
||||
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:
|
||||
For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL`, `LENS_WORKER_TOKEN`, and an explicit `LENS_WORKER_IMAGE` in a private environment file:
|
||||
|
||||
```bash
|
||||
docker compose --env-file /path/to/lens.env -f compose.yaml up -d
|
||||
```
|
||||
|
||||
Developers can build locally with `LENS_WORKER_IMAGE=litellm-lens-worker:local docker compose -f deploy/lens/compose.yaml -f deploy/lens/compose.build.yaml up -d --build`. To work on Lens itself, `make lens-dev` runs the proxy, a worker from source and the hot-reload dashboard together; set `LENS_DEV_PROXY_PORT` / `LENS_DEV_UI_PORT` to move them off 4000/3000
|
||||
To work on Lens itself, `make lens-dev` runs the proxy, a worker from source and the hot-reload dashboard together; set `LENS_DEV_PROXY_PORT` / `LENS_DEV_UI_PORT` to move them off 4000/3000. For a local container build, set `LENS_WORKER_IMAGE=litellm-lens-worker:local` and `LITELLM_RELEASE_TAG` to the gateway's release tag, then use `docker compose -f deploy/lens/compose.yaml -f deploy/lens/compose.build.yaml up -d --build`
|
||||
|
||||
The generated command gives the worker 1 GiB of temporary memory-backed storage, shared across parallel reviews. Change `size=1g` in the Docker command or set `LENS_WORKER_TMP_SIZE` with Compose to fit your server and workload. A storage failure marks the scan as failed, cleans up temporary traces, and leaves the worker available for other scans; it does not silently truncate the review. Existing workers must be recreated with the new image and mount options
|
||||
|
||||
|
|
@ -107,6 +172,37 @@ curl "$LITELLM_URL/lens/$LENS_ID/runs/$BATCH_ID" -H "Authorization: Bearer $LITE
|
|||
|
||||
Creation queues the first batch. Posting to `/lens/{id}/runs` queues another, or returns the existing active batch. The run response contains its ID under `jobs[0].id`. Poll the batch URL for status, findings and assessments. List responses omit large result payloads; request a batch to retrieve them. Supply an optional complete `settings` object on the runs POST for a one-off override; the saved lens stays unchanged. Selection accepts `team_id`, exact `filters`, and opaque `execution_ids` returned by `/lens/preview/sample`. Preview accepts `offset` and `as_of` to keep the time window fixed while paging. Feedback uses `PATCH /lens/{id}/findings/{finding_id}` with `status` and `reason`
|
||||
|
||||
## Local development
|
||||
|
||||
`make lens-dev ARGS=--seed` starts the full dev stack. The live dashboard is at `http://localhost:3000/ui/lens/`, with login at `http://localhost:3000/ui/login/`. Next.js forwards API requests to the proxy on port 4000, so login and navigation stay in the live UI and edits hot-reload
|
||||
|
||||
The default is Next.js dev with no production build (`LENS_DEV_BUILD_UI=0`). Set `LENS_DEV_BUILD_UI=1` when you also want a fresh static dashboard at `http://localhost:4000/ui/`. Build output goes to `.lens-dev/logs/ui-build.log`; a failed build stops startup. Both modes keep the live dashboard on port 3000. Startup checks the live login route before seeding and fails with the UI log path if Next.js exits. `LENS_DEV_STARTUP_TIMEOUT_SECONDS` controls startup readiness retries (default 300; `LENS_DEV_READINESS_REQUEST_TIMEOUT_SECONDS` caps each HTTP probe, default 5)
|
||||
|
||||
For local fixture data, run `make lens-dev ARGS=--seed`. Use `make lens-dev ARGS="--seed large"` for 2,000 fixture copies spread over the last 24 hours, about 860,000 spans with linked request logs, plus three long sessions of roughly 1,150, 9,200 and 92,000 spans in a single trace for drawer paging and the oversized read path. Their trace IDs are printed at the end. To seed a running stack without restarting it, use `make lens-dev ARGS="--seed-only --seed large --copies 100"`. Every profile replays one copy of every checked-in capture through authenticated `/v1/traces`, including failures, retries, streaming and multiple agent frameworks, and verifies linked spend totals through the proxy. Large seeds then copy that first copy inside ClickHouse and PostgreSQL with `INSERT ... SELECT`, rewriting trace, span and call IDs so each copy keeps its own spend, and verify the last copy through the proxy
|
||||
|
||||
Seeds append fresh IDs on every invocation and spread copies over recent timestamps. Restarts without `SEED` do not add data. Lens excludes activity received in the last two minutes, so wait two minutes after seeding before checking investigation previews. `LENS_DEV_SEED_COPIES` overrides total copies. Large seeds test data volume and pagination, rather than concurrent ingestion throughput or review accuracy. They can use substantial disk space; adjust `--copies` for your machine. Seeding expects the generated local tracing configuration. The old `run_tracing_proxy_local.sh --seed` command forwards to Lens dev, using its ports and saved master key
|
||||
|
||||
Local ingestion limits are explicit and configurable. Set OTLP and ClickHouse variables before starting the proxy and seeder so both processes use the same settings. Invalid, zero and negative values fail instead of silently falling back. Changing these limits does not require rebuilding Rust
|
||||
|
||||
| Environment variable | Default | Controls |
|
||||
| --- | --- | --- |
|
||||
| `LENS_DEV_SEED_COPIES` | 1 default, 2000 large | Total fixture copies |
|
||||
| `LENS_DEV_SEED_TIMEOUT_SECONDS` | 120 | Seeder HTTP timeout |
|
||||
| `OTLP_MAX_BODY_BYTES` | 16777216 | HTTP body and decompressed payload bytes |
|
||||
| `OTLP_MAX_CONCURRENT_INGESTS` | 2 | Concurrent proxy ingestion requests |
|
||||
| `OTLP_MAX_ATTRIBUTE_VALUE_BYTES` | 65536 | Stored attribute/content bytes |
|
||||
| `OTLP_MAX_DECODE_DEPTH` | 32 | Nested decode depth |
|
||||
| `OTLP_MAX_DECODE_NODES` | 65536 | JSON values or protobuf fields per export |
|
||||
| `OTLP_MAX_SPANS` | 4096 | Spans per export |
|
||||
| `OTLP_MAX_ATTRIBUTES` | 256 | Attributes per resource, scope, span, event or link |
|
||||
| `OTLP_MAX_EVENTS` | 256 | Events per span |
|
||||
| `OTLP_MAX_LINKS` | 256 | Links per span |
|
||||
| `OTLP_MAX_DECODED_SPAN_BYTES` | 16777216 | Decoded span allocation budget |
|
||||
| `CLICKHOUSE_TRACE_MAX_INSERT_BYTES` | 67108864 | Encoded trace or spend insert bytes |
|
||||
| `CLICKHOUSE_INSERT_TIMEOUT_SECONDS` | 30 | ClickHouse insert HTTP timeout |
|
||||
|
||||
The wire parsers also enforce their library recursion limits (128 levels for JSON, 100 for protobuf). Raising the configured depth does not remove those parser limits.
|
||||
|
||||
## Quality evaluation
|
||||
|
||||
Run the checked-in cases against a configured real model. Expected labels are used only for scoring, never passed to the model. Dev and held-out cases include missing outcomes, failed tools, recovery, handoffs, unsupported claims, repeated work, long evidence and prompt injection. The background option adds clean arithmetic traces to test rare-issue discovery at scale; those repeated synthetic cases do not establish accuracy on every production workload
|
||||
|
|
@ -130,3 +226,19 @@ The Lens API now uses `/lens` instead of `/engine`, list responses use `lenses`,
|
|||
Stop workers and let active scans finish before upgrading. Deploy proxy instances together: older proxies cannot use the renamed database tables. The schema migration renames the three Lens tables and the run-history identifier column in place, preserving saved investigations, findings, history, worker credentials, and billing assignments. Existing migration files retain their original names and checksums
|
||||
|
||||
Upgrades using `--use_prisma_db_push` stop before schema changes if any legacy Lens table exists, preventing Prisma from dropping saved data. Apply `litellm-proxy-extras/litellm_proxy_extras/migrations/20261001100000_rename_lens/migration.sql` to the configured database schema before retrying. Deployments already using migration history can instead start without `--use_prisma_db_push` to apply the shipped migration normally. Fresh databases and databases already using the renamed tables can continue using database push
|
||||
|
||||
|
||||
## Release compatibility
|
||||
|
||||
Gateway and worker builds carry the same `LITELLM_RELEASE_TAG`. A worker announces its release and protocol before claiming an investigation. A mismatch returns HTTP 409 with the required image, leaving queued investigations untouched. During a rolling upgrade, workers wait for a gateway from their release
|
||||
|
||||
The dashboard reads its image from the running gateway. `LENS_WORKER_IMAGE` overrides the registry/image for private deployments. Set an explicit `LENS_WORKER_IMAGE` for worker-only Compose. Verify that the image exists and matches the gateway before deploying it
|
||||
|
||||
For source development, use `make lens-dev`, which gives the proxy and source worker the same commit identity. For custom containers, build both from the same checkout with `--build-arg LITELLM_RELEASE_TAG=sha-$(git rev-parse HEAD)` and set the proxy's `LENS_WORKER_IMAGE` to the worker image you built. An unlabelled custom build refuses worker setup and claims instead of guessing from the Python package version. Normal package-index installations use their installed release version
|
||||
|
||||
The hourly development pipeline pins all component images to the same selected commit and publishes its chart only after every build and worker smoke test succeeds. The public commit-tagged worker workflow publishes to `ghcr.io/berriai/litellm-lens-worker-dev` on Lens-related changes, so an arbitrary `main` commit may require building your own pair; do not substitute the newest available worker
|
||||
|
||||
|
||||
## Worker dependencies
|
||||
|
||||
The worker uses the same digest-pinned Wolfi base and Python version as the component images. Python dependencies and their hashes are locked in `deploy/lens/requirements.lock`. To update them, edit `deploy/lens/requirements.in`, then run `uv pip compile --universal --python-version 3.13 --generate-hashes --no-emit-index-url deploy/lens/requirements.in -o deploy/lens/requirements.lock`. The image installs only the locked wheels with hash verification. CI builds and scans both native architectures
|
||||
|
|
|
|||
|
|
@ -3,4 +3,6 @@ services:
|
|||
build:
|
||||
context: ../..
|
||||
dockerfile: deploy/lens/Dockerfile
|
||||
args:
|
||||
LITELLM_RELEASE_TAG: ${LITELLM_RELEASE_TAG:?Set the release tag used by the gateway}
|
||||
image: litellm-lens-worker:local
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
services:
|
||||
lens-worker:
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:44f0597c7583dcfef999ece9a8bc02cfeb9f0f5167a1221cee3bd10b1b79271b}
|
||||
image: ${LENS_WORKER_IMAGE:-${LITELLM_VERSION:+ghcr.io/berriai/litellm-lens-worker:v}${LITELLM_VERSION:-}}
|
||||
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}
|
||||
|
|
|
|||
7
deploy/lens/config.yaml
Normal file
7
deploy/lens/config.yaml
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
general_settings:
|
||||
master_key: os.environ/LITELLM_MASTER_KEY
|
||||
tracing:
|
||||
store:
|
||||
type: clickhouse
|
||||
url: os.environ/CLICKHOUSE_URL
|
||||
retention_days: 14
|
||||
2
deploy/lens/requirements.in
Normal file
2
deploy/lens/requirements.in
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
httpx==0.28.1
|
||||
pydantic==2.13.4
|
||||
172
deploy/lens/requirements.lock
Normal file
172
deploy/lens/requirements.lock
Normal file
|
|
@ -0,0 +1,172 @@
|
|||
# This file was autogenerated by uv via the following command:
|
||||
# uv pip compile --universal --python-version 3.13 --generate-hashes --no-emit-index-url deploy/lens/requirements.in -o deploy/lens/requirements.lock
|
||||
annotated-types==0.8.0 \
|
||||
--hash=sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7 \
|
||||
--hash=sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0
|
||||
# via pydantic
|
||||
anyio==4.15.1 \
|
||||
--hash=sha256:6152fdbbf9a77fdec97731721bebf7c4c44f7c29b424b0065826173efc7ed101 \
|
||||
--hash=sha256:9f28306018cbd6d329e64a36d58256edff76dd996fe423bc957326e578b82a94
|
||||
# via httpx
|
||||
certifi==2026.7.22 \
|
||||
--hash=sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775 \
|
||||
--hash=sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55
|
||||
# via
|
||||
# httpcore
|
||||
# httpx
|
||||
h11==0.16.0 \
|
||||
--hash=sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1 \
|
||||
--hash=sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86
|
||||
# via httpcore
|
||||
httpcore==1.0.9 \
|
||||
--hash=sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55 \
|
||||
--hash=sha256:6e34463af53fd2ab5d807f399a9b45ea31c3dfa2276f15a2c3f00afff6e176e8
|
||||
# via httpx
|
||||
httpx==0.28.1 \
|
||||
--hash=sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc \
|
||||
--hash=sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad
|
||||
# via -r deploy/lens/requirements.in
|
||||
idna==3.20 \
|
||||
--hash=sha256:a7db850025b95ded1eae8a46181a1a6c56c92c96f0e2b005d9ff8dc0210cab44 \
|
||||
--hash=sha256:ab7ae7122974553370f0bdb919e1a960b2cd1bc1ef0276416d896db81c14582c
|
||||
# via
|
||||
# anyio
|
||||
# httpx
|
||||
pydantic==2.13.4 \
|
||||
--hash=sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba \
|
||||
--hash=sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6
|
||||
# via -r deploy/lens/requirements.in
|
||||
pydantic-core==2.46.4 \
|
||||
--hash=sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0 \
|
||||
--hash=sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262 \
|
||||
--hash=sha256:041bde0a48fd37cf71cab1c9d56d3e8625a3793fef1f7dd232b3ff37e978ecda \
|
||||
--hash=sha256:0c563b08bca408dc7f65f700633d8442fffb2421fc47b8101377e9fd65051ff0 \
|
||||
--hash=sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e \
|
||||
--hash=sha256:0ce40cd7b21210e99342afafbd4d0f76d784eb5b1d60f3bdc566be4983c6c73b \
|
||||
--hash=sha256:0e96592440881c74a213e5ad528e2b24d3d4f940de2766bed9010ab1d9e51594 \
|
||||
--hash=sha256:10e17cbb10a330363733efc4d7c4d0dd827ac0909b8f6a6542298fed1ea62f29 \
|
||||
--hash=sha256:133878133d271ade3d41d1bfb2a45ec38dbdbda40bc065921c6b04e4630127e2 \
|
||||
--hash=sha256:14d4edf427bdcf950a8a02d7cb44a08614388dd6e1bdcbf4f67504fa7887da9c \
|
||||
--hash=sha256:14f4c5d6db102bd796a627bbb3a17b4cf4574b9ae861d8b7c9a9661c6dd3362d \
|
||||
--hash=sha256:17299feefe090f2caa5b8e37222bb5f663e4935a8bfa6931d4102e5df1a9f398 \
|
||||
--hash=sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d \
|
||||
--hash=sha256:18e5ceec2ab67e6d5f1a9085e5a24c9c4e2ac4545730bfe668680bca05e555f3 \
|
||||
--hash=sha256:19e51f073cd3df251856a8a4189fbdf1de4012c3ebacfb1884f94f1eb406079f \
|
||||
--hash=sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb \
|
||||
--hash=sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7 \
|
||||
--hash=sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5 \
|
||||
--hash=sha256:228ee9bae8bef5b1e97ec58302f80357c37199e0d0a99174e138d28e6957b9d9 \
|
||||
--hash=sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462 \
|
||||
--hash=sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4 \
|
||||
--hash=sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b \
|
||||
--hash=sha256:2f84c03c8607173d16b5a854ec68a2f9079ae03237a54fb506d13af47e1d018d \
|
||||
--hash=sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df \
|
||||
--hash=sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2 \
|
||||
--hash=sha256:3447661d99f75a3683a4cf5c87da72f2161964611864dbbeac7fbb118bb4bfc0 \
|
||||
--hash=sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519 \
|
||||
--hash=sha256:395aebd9183f9d112f569aeb5b2214d1a10a33bec8456447f7fbdfa51d38d4cd \
|
||||
--hash=sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7 \
|
||||
--hash=sha256:3be77f45df024d789a672ae34f8b06fb346c4f9f46ea714956660ea4862e89ac \
|
||||
--hash=sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6 \
|
||||
--hash=sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565 \
|
||||
--hash=sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898 \
|
||||
--hash=sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb \
|
||||
--hash=sha256:432c179df7874eeb73307aad2df0755e1ae0efa61ff0ea89b93e194411ae3928 \
|
||||
--hash=sha256:4a05d69cba51d852c5c3e92758653245a50c0b646ced0cf05bd793ed592839d6 \
|
||||
--hash=sha256:4c63ebc82684aa89d9a3bcbd13d515b3be44250dc68dd3bd81526c1cb31286c3 \
|
||||
--hash=sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a \
|
||||
--hash=sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596 \
|
||||
--hash=sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987 \
|
||||
--hash=sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e \
|
||||
--hash=sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d \
|
||||
--hash=sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712 \
|
||||
--hash=sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008 \
|
||||
--hash=sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd \
|
||||
--hash=sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1 \
|
||||
--hash=sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be \
|
||||
--hash=sha256:66ce7632c22d837c95301830e111ad0128a32b8207533b60896a96c4915192ea \
|
||||
--hash=sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292 \
|
||||
--hash=sha256:6f2eeda33a839975441c86a4119e1383c50b47faf0cbb5176985565c6bb02c33 \
|
||||
--hash=sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3 \
|
||||
--hash=sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4 \
|
||||
--hash=sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b \
|
||||
--hash=sha256:7bfb192b3f4b9e8a89b6277b6ce787564f62cfd272055f6e685726b111dc7826 \
|
||||
--hash=sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac \
|
||||
--hash=sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7 \
|
||||
--hash=sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d \
|
||||
--hash=sha256:8358a950c8909158e3df31538a7e4edc2d7265a7c54b47f0864d9e5bae9dcebf \
|
||||
--hash=sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4 \
|
||||
--hash=sha256:86e1a4418c6cd97d60c95c71164158eaf7324fae7b0923264016baa993eba6fc \
|
||||
--hash=sha256:8b9bab013d1c7a79d3501ff86d0bc9c31bf587db4551677b96bec07df78c6b15 \
|
||||
--hash=sha256:8c5dac79fa1614d1e06ca695109c6105923bd9c7d1d6c918d4e637b7e6b32fd3 \
|
||||
--hash=sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b \
|
||||
--hash=sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914 \
|
||||
--hash=sha256:9037063db01f09b09e237c282b6792bd4da634b5402c4e7f0c61effed7701a04 \
|
||||
--hash=sha256:905a0ed8ea6f2d61c1738835f99b699348d7857379083e5fc497fa0c967a407c \
|
||||
--hash=sha256:90884113d8b48f760e9587002789ddd741e76ab9f89518cd1e43b1f1a52ec44b \
|
||||
--hash=sha256:91a06d2e259ecfbd8c901d70c3c507900458498142b3026a296b7de4d1322cc9 \
|
||||
--hash=sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce \
|
||||
--hash=sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4 \
|
||||
--hash=sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a \
|
||||
--hash=sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f \
|
||||
--hash=sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424 \
|
||||
--hash=sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894 \
|
||||
--hash=sha256:9bc519fbf2b7578398853d815009ae5e4d4603d12f4e3f91da8c06852d3da3e9 \
|
||||
--hash=sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76 \
|
||||
--hash=sha256:9f444c499b3eefd3a92e348059471ea0c3a6e303d9c1cec09fa748fd9f895201 \
|
||||
--hash=sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb \
|
||||
--hash=sha256:a0f62d0a58f4e7da165457e995725421e0064f2255d8eccebc49f41bbc23b109 \
|
||||
--hash=sha256:a396dcc17e5a0b164dbe026896245a4fa9ff402edca1dff0be3d53a517f74de4 \
|
||||
--hash=sha256:aaa2a54443eff1950ba5ddc6b6ccda0d9c84a364276a62f969bdf2a390650848 \
|
||||
--hash=sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526 \
|
||||
--hash=sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0 \
|
||||
--hash=sha256:b078afbc25f3a1436c7a1d2cd3e322497ee99615ba97c563566fdf46aff1ee01 \
|
||||
--hash=sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458 \
|
||||
--hash=sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e \
|
||||
--hash=sha256:bb63e0198ca18aad131c089b9204c23079c3afa95487e561f4c522d519e55aba \
|
||||
--hash=sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a \
|
||||
--hash=sha256:c1747f85cee84c26985853c6f3d9bd3e75da5212912443fa111c113b9c246f39 \
|
||||
--hash=sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c \
|
||||
--hash=sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000 \
|
||||
--hash=sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b \
|
||||
--hash=sha256:c7a7bd4e39e8e4c12c39cd480356842b6a8a06e41b23a55a5e3e191718838ddf \
|
||||
--hash=sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4 \
|
||||
--hash=sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd \
|
||||
--hash=sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28 \
|
||||
--hash=sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9 \
|
||||
--hash=sha256:d396ec2b979760aaf3218e76c24e65bd0aca24983298653b3a9d7a45f9e47b30 \
|
||||
--hash=sha256:d51026d73fcfd93610abc7b27789c26b313920fcfb20e27462d74a7f8b06e983 \
|
||||
--hash=sha256:d80ee3d731373b24cebbc10d689ca4ee1875caf0d5703a245db18efd4dd37fc1 \
|
||||
--hash=sha256:d995260fdf4e1db774581b4900e0f832abe3c7c84996726bbc161b19c8f29e76 \
|
||||
--hash=sha256:da4b951fe36dc7c3a1ccb4e3cd1747c3542b8c9ceede8fc86cae054e764485f5 \
|
||||
--hash=sha256:daa27d92c36f24388fe3ad306b174781c747627f134452e4f128ea00ce1fe8c4 \
|
||||
--hash=sha256:db06ffe51636ffe9ca531fe9023dd64bdd794be8754cb5df57c5498ae5b518a7 \
|
||||
--hash=sha256:e0d65b8c354be7fb5f720c3caa8bc940bc2d20ce749c8e06135f07f8ed95dd7c \
|
||||
--hash=sha256:e68b7a074f65a2fd746c52a7ce6142ab7006074ac269ace0c25cd8ba171f8066 \
|
||||
--hash=sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3 \
|
||||
--hash=sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02 \
|
||||
--hash=sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89 \
|
||||
--hash=sha256:ea793e075b70290d89d8142074262885d3f7da19634845135751bd6344f73b50 \
|
||||
--hash=sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76 \
|
||||
--hash=sha256:f13a646d65d09fbf1bc6b3a9635d30095c8e7e5cc419ff35ecc563c5fd04cd49 \
|
||||
--hash=sha256:f47286a97f0bc9b8859519809077b91b2cefe4ae47fcbf5e466a009c1c5d742b \
|
||||
--hash=sha256:f747929cf940cddb5b3668a390056ddd5ba2e5010615ea2dcf4f9c4f3ab8791d \
|
||||
--hash=sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7 \
|
||||
--hash=sha256:f9fa868638bf362d3d138ea55829cefb3d5f4b0d7f142234382a15e2485dbec4 \
|
||||
--hash=sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c \
|
||||
--hash=sha256:fc010ab034c8c7452522748bf937df58020d256ccae0874463d1f4d01758af8e \
|
||||
--hash=sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff \
|
||||
--hash=sha256:fd8b3d9fd264be37976686c7f65cd52a83f5e84f4bfd2adf9c1d469676bbb6ae
|
||||
# via pydantic
|
||||
typing-extensions==4.16.0 \
|
||||
--hash=sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8 \
|
||||
--hash=sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5
|
||||
# via
|
||||
# anyio
|
||||
# pydantic
|
||||
# pydantic-core
|
||||
# typing-inspection
|
||||
typing-inspection==0.4.4 \
|
||||
--hash=sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47 \
|
||||
--hash=sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147
|
||||
# via pydantic
|
||||
91
deploy/lens/stack.yaml
Normal file
91
deploy/lens/stack.yaml
Normal file
|
|
@ -0,0 +1,91 @@
|
|||
name: litellm-lens
|
||||
|
||||
services:
|
||||
litellm:
|
||||
image: ghcr.io/berriai/litellm:${LITELLM_VERSION:?Set LITELLM_VERSION to a published release, without the v prefix}
|
||||
entrypoint:
|
||||
- python3
|
||||
- -c
|
||||
- |
|
||||
import os, sys
|
||||
from urllib.parse import quote
|
||||
postgres_password = quote(os.environ["POSTGRES_PASSWORD"], safe="")
|
||||
clickhouse_password = quote(os.environ["CLICKHOUSE_PASSWORD"], safe="")
|
||||
os.environ["DATABASE_URL"] = f"postgresql://litellm:{postgres_password}@db:5432/litellm"
|
||||
os.environ["CLICKHOUSE_URL"] = f"http://default:{clickhouse_password}@clickhouse:8123"
|
||||
os.execv("docker/prod_entrypoint.sh", ["docker/prod_entrypoint.sh", *sys.argv[1:]])
|
||||
command: ["--config", "/app/lens-config.yaml", "--port", "4000"]
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: ${LITELLM_MASTER_KEY:?Set a strong master key}
|
||||
LITELLM_SALT_KEY: ${LITELLM_SALT_KEY:?Set a permanent encryption key and keep it across upgrades}
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?Set a permanent database password}
|
||||
STORE_MODEL_IN_DB: "True"
|
||||
CLICKHOUSE_PASSWORD: ${CLICKHOUSE_PASSWORD:?Set a permanent ClickHouse password}
|
||||
LENS_WORKER_IMAGE: ghcr.io/berriai/litellm-lens-worker:v${LITELLM_VERSION}
|
||||
volumes:
|
||||
- ./config.yaml:/app/lens-config.yaml:ro
|
||||
ports:
|
||||
- "127.0.0.1:${LITELLM_PORT:-4000}:4000"
|
||||
networks: [proxy, storage]
|
||||
depends_on:
|
||||
db:
|
||||
condition: service_healthy
|
||||
clickhouse:
|
||||
condition: service_healthy
|
||||
restart: unless-stopped
|
||||
|
||||
lens-worker:
|
||||
profiles: [lens]
|
||||
image: ghcr.io/berriai/litellm-lens-worker:v${LITELLM_VERSION}
|
||||
environment:
|
||||
LITELLM_URL: http://litellm:4000
|
||||
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:-}
|
||||
depends_on: [litellm]
|
||||
networks: [proxy]
|
||||
restart: unless-stopped
|
||||
read_only: true
|
||||
tmpfs:
|
||||
- /tmp:rw,noexec,nosuid,size=${LENS_WORKER_TMP_SIZE:-1g}
|
||||
cap_drop: [ALL]
|
||||
security_opt: [no-new-privileges:true]
|
||||
|
||||
db:
|
||||
image: postgres:16
|
||||
environment:
|
||||
POSTGRES_DB: litellm
|
||||
POSTGRES_USER: litellm
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
|
||||
networks: [storage]
|
||||
volumes:
|
||||
- postgres_data:/var/lib/postgresql/data
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U litellm -d litellm"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 20
|
||||
restart: unless-stopped
|
||||
|
||||
clickhouse:
|
||||
image: clickhouse/clickhouse-server:26.9.6.6
|
||||
environment:
|
||||
CLICKHOUSE_USER: default
|
||||
CLICKHOUSE_PASSWORD: ${CLICKHOUSE_PASSWORD}
|
||||
CLICKHOUSE_DEFAULT_ACCESS_MANAGEMENT: "1"
|
||||
volumes:
|
||||
- clickhouse_data:/var/lib/clickhouse
|
||||
healthcheck:
|
||||
test: ["CMD", "clickhouse-client", "--user", "default", "--password", "${CLICKHOUSE_PASSWORD}", "--query", "SELECT 1"]
|
||||
interval: 5s
|
||||
timeout: 5s
|
||||
retries: 20
|
||||
restart: unless-stopped
|
||||
networks: [storage]
|
||||
|
||||
networks:
|
||||
proxy:
|
||||
storage:
|
||||
internal: true
|
||||
|
||||
volumes:
|
||||
postgres_data:
|
||||
clickhouse_data:
|
||||
44
docker-compose.liteadmin.yml
Normal file
44
docker-compose.liteadmin.yml
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
services:
|
||||
litellm:
|
||||
image: ${LITELLM_IMAGE:?Set the native-enabled gateway image}
|
||||
environment:
|
||||
LITELLM_ADMIN_AGENT_URL: http://liteadmin:10000
|
||||
ADMIN_AGENT_SERVICE_TOKEN: ${ADMIN_AGENT_SERVICE_TOKEN:?Set a shared worker token}
|
||||
PROXY_BASE_URL: ${LITELLM_PUBLIC_URL:?Set the existing HTTPS gateway URL}
|
||||
|
||||
liteadmin:
|
||||
image: ${LITELLM_IMAGE:?Set the same native-enabled image used by the gateway}
|
||||
command: ["--admin-agent"]
|
||||
restart: unless-stopped
|
||||
init: true
|
||||
read_only: true
|
||||
cap_drop: [ALL]
|
||||
security_opt: [no-new-privileges:true]
|
||||
stop_grace_period: 75s
|
||||
environment:
|
||||
CONNECTION_AUTH_MODE: native
|
||||
LITELLM_BASE_URL: ${LITELLM_PUBLIC_URL:?Set the existing HTTPS gateway URL}
|
||||
LITELLM_MODEL: ${LITELLM_ADMIN_MODEL:?Set a gateway model with tool support}
|
||||
SLACK_BOT_TOKEN: ${SLACK_BOT_TOKEN:?Install the Slack app}
|
||||
SLACK_APP_TOKEN: ${SLACK_APP_TOKEN:?Enable Socket Mode}
|
||||
SLACK_WORKSPACE_ID: ${SLACK_WORKSPACE_ID:?Set the Slack workspace ID}
|
||||
ADMIN_AGENT_SERVICE_TOKEN: ${ADMIN_AGENT_SERVICE_TOKEN:?Set a shared worker token}
|
||||
CREDENTIAL_ENCRYPTION_KEY: ${CREDENTIAL_ENCRYPTION_KEY:?Set a persistent Fernet key}
|
||||
STATE_DB: /var/data/events.sqlite3
|
||||
ADMIN_READ_ONLY: ${ADMIN_READ_ONLY:-false}
|
||||
OPENAI_AGENTS_DISABLE_TRACING: "1"
|
||||
volumes:
|
||||
- liteadmin_state:/var/data
|
||||
tmpfs:
|
||||
- /tmp:rw,noexec,nosuid,size=64m
|
||||
healthcheck:
|
||||
test: ["CMD", "/opt/liteadmin/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:10000/readyz', timeout=3)"]
|
||||
interval: 30s
|
||||
timeout: 5s
|
||||
start_period: 30s
|
||||
depends_on:
|
||||
litellm:
|
||||
condition: service_healthy
|
||||
|
||||
volumes:
|
||||
liteadmin_state:
|
||||
|
|
@ -113,6 +113,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
|||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
ARG LITELLM_RELEASE_TAG=""
|
||||
ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
|
||||
|
||||
USER root
|
||||
|
||||
|
|
|
|||
|
|
@ -122,6 +122,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \
|
|||
sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh
|
||||
|
||||
FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
||||
ARG LITELLM_RELEASE_TAG=""
|
||||
ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG}
|
||||
WORKDIR /app
|
||||
USER root
|
||||
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ services:
|
|||
build:
|
||||
context: ..
|
||||
target: runtime
|
||||
args:
|
||||
LITELLM_RELEASE_TAG: ${LITELLM_RELEASE_TAG:-}
|
||||
command: ["--config", "/app/tracing-config.yaml", "--port", "4000"]
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: sk-1234
|
||||
|
|
@ -15,6 +17,7 @@ services:
|
|||
CLICKHOUSE_URL: http://default:local-tracing@clickhouse:8123
|
||||
CLICKHOUSE_DATABASE: litellm
|
||||
OPENAI_API_KEY: ${OPENAI_API_KEY:-}
|
||||
LENS_WORKER_IMAGE: ${LENS_WORKER_IMAGE:-}
|
||||
volumes:
|
||||
- ./tracing-config.yaml:/app/tracing-config.yaml:ro
|
||||
ports:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,11 @@
|
|||
#!/bin/sh
|
||||
|
||||
if [ "$1" = "--admin-agent" ]; then
|
||||
shift
|
||||
export CONNECTION_AUTH_MODE=native
|
||||
exec /opt/liteadmin/bin/litellm-admin-agent --web "$@"
|
||||
fi
|
||||
|
||||
case "$USE_DDTRACE" in
|
||||
[Tt][Rr][Uu][Ee])
|
||||
export DD_TRACE_OPENAI_ENABLED="False"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from litellm_enterprise.enterprise_callbacks.send_emails.endpoints import (
|
|||
|
||||
from . import ui_crud_endpoints # side-effect: registers extra UI settings
|
||||
from .audit_logging_endpoints import router as audit_logging_router
|
||||
from .liteadmin import router as liteadmin_router
|
||||
from .management_endpoints import management_endpoints_router
|
||||
from .utils import _should_block_robots
|
||||
|
||||
|
|
@ -14,6 +15,7 @@ __all__ = ["router", "ui_crud_endpoints"]
|
|||
router = APIRouter()
|
||||
router.include_router(email_events_router)
|
||||
router.include_router(audit_logging_router)
|
||||
router.include_router(liteadmin_router)
|
||||
router.include_router(management_endpoints_router)
|
||||
|
||||
|
||||
|
|
|
|||
283
enterprise/litellm_enterprise/proxy/liteadmin.py
Normal file
283
enterprise/litellm_enterprise/proxy/liteadmin.py
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import html
|
||||
import os
|
||||
import re
|
||||
import secrets
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Annotated, Final
|
||||
from urllib.parse import urlencode, urlsplit
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, RedirectResponse, Response
|
||||
from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.proxy._experimental.mcp_server.oauth_utils import get_request_base_url
|
||||
from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
|
||||
router: Final = APIRouter()
|
||||
_PREFIX: Final = "/liteadmin/slack/connect/"
|
||||
_COOKIE: Final = "__Host-litellm-slack-connect-"
|
||||
_HEADERS: Final = {
|
||||
"Cache-Control": "no-store",
|
||||
"Referrer-Policy": "same-origin",
|
||||
"X-Frame-Options": "DENY",
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
"Content-Security-Policy": "default-src 'none'; style-src 'unsafe-inline'; form-action 'self'; frame-ancestors 'none'; base-uri 'none'",
|
||||
}
|
||||
|
||||
|
||||
class LinkDetails(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, strict=True, extra="forbid")
|
||||
workspace_id: str = Field(min_length=1, max_length=64)
|
||||
slack_user_id: str = Field(min_length=1, max_length=64)
|
||||
email: str = Field(min_length=1, max_length=320)
|
||||
|
||||
|
||||
class AdminSession(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
user_id: str
|
||||
credential: SecretStr
|
||||
expires_at: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NativeAdminContext:
|
||||
worker_url: str
|
||||
service_token: SecretStr
|
||||
client: httpx.AsyncClient
|
||||
session_user: Callable[[Request], Awaitable[str | None]]
|
||||
load_user: Callable[[str], Awaitable[LiteLLM_UserTable | None]]
|
||||
mint_session: Callable[[LiteLLM_UserTable], AdminSession]
|
||||
|
||||
async def worker_request(self, token: str, session: AdminSession | None = None) -> httpx.Response:
|
||||
if re.fullmatch(r"[A-Za-z0-9_-]{43}", token) is None:
|
||||
raise HTTPException(410, "Connection link expired. Send connect in Slack for a new link")
|
||||
try:
|
||||
response: Final = await self.client.request(
|
||||
"GET" if session is None else "POST",
|
||||
f"{self.worker_url}/internal/liteadmin/links/{token}",
|
||||
headers={"X-LiteLLM-Admin-Agent-Token": self.service_token.get_secret_value()},
|
||||
json=None
|
||||
if session is None
|
||||
else {
|
||||
"user_id": session.user_id,
|
||||
"credential": session.credential.get_secret_value(),
|
||||
"expires_at": session.expires_at,
|
||||
},
|
||||
timeout=15,
|
||||
follow_redirects=False,
|
||||
)
|
||||
except httpx.HTTPError:
|
||||
raise HTTPException(503, "LiteAdmin is temporarily unavailable") from None
|
||||
if response.status_code == 410:
|
||||
raise HTTPException(410, "Connection link expired. Send connect in Slack for a new link")
|
||||
if response.status_code == 403:
|
||||
raise HTTPException(403, "Connect your own active LiteLLM proxy-admin account with the same email as Slack")
|
||||
if response.status_code != 200:
|
||||
raise HTTPException(503, "LiteAdmin could not verify this connection")
|
||||
return response
|
||||
|
||||
async def details(self, token: str) -> LinkDetails:
|
||||
response: Final = await self.worker_request(token)
|
||||
try:
|
||||
return LinkDetails.model_validate_json(response.content)
|
||||
except ValidationError:
|
||||
raise HTTPException(503, "LiteAdmin could not verify this connection") from None
|
||||
|
||||
async def admin(self, user_id: str, details: LinkDetails) -> LiteLLM_UserTable:
|
||||
user: Final = await self.load_user(user_id)
|
||||
if (
|
||||
user is None
|
||||
or user.user_role != LitellmUserRoles.PROXY_ADMIN.value
|
||||
or not user.user_email
|
||||
or user.user_email.strip().casefold() != details.email.strip().casefold()
|
||||
):
|
||||
raise HTTPException(403, "Connect your own active LiteLLM proxy-admin account with the same email as Slack")
|
||||
return user
|
||||
|
||||
|
||||
def _page(title: str, body: str) -> HTMLResponse:
|
||||
return HTMLResponse(
|
||||
f'<!doctype html><html lang="en"><meta charset="utf-8">'
|
||||
f'<meta name="viewport" content="width=device-width,initial-scale=1"><title>{html.escape(title)}</title>'
|
||||
"<style>body{font:17px system-ui;color:#18252f;max-width:560px;margin:10vh auto;padding:24px}"
|
||||
"p{line-height:1.6}button{font:inherit;border:0;border-radius:8px;padding:14px 20px;background:#5b3fd1;"
|
||||
"color:white;cursor:pointer}small{color:#556}</style>"
|
||||
f"<main><h1>{html.escape(title)}</h1>{body}</main></html>",
|
||||
headers=_HEADERS,
|
||||
)
|
||||
|
||||
|
||||
def _cookie_name(token: str) -> str:
|
||||
return _COOKIE + hashlib.sha256(token.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
async def _session_user(request: Request) -> str | None:
|
||||
from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import (
|
||||
get_authenticated_browser_user_id,
|
||||
)
|
||||
|
||||
return await get_authenticated_browser_user_id(request)
|
||||
|
||||
|
||||
async def _load_user(user_id: str) -> LiteLLM_UserTable | None:
|
||||
from litellm.proxy.auth.auth_checks import get_user_object
|
||||
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
|
||||
|
||||
if prisma_client is None:
|
||||
raise HTTPException(503, "LiteAdmin requires a database")
|
||||
try:
|
||||
return await get_user_object(
|
||||
user_id=user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
user_id_upsert=False,
|
||||
check_db_only=True,
|
||||
)
|
||||
except UserNotFoundError:
|
||||
return None
|
||||
except Exception:
|
||||
raise HTTPException(503, "LiteAdmin could not verify your current permissions") from None
|
||||
|
||||
|
||||
def mint_admin_session(user: LiteLLM_UserTable) -> AdminSession:
|
||||
from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token
|
||||
|
||||
expires: Final = datetime.now(timezone.utc) + timedelta(hours=24)
|
||||
auth: Final = UserAPIKeyAuth(
|
||||
token="liteadmin-" + secrets.token_urlsafe(24),
|
||||
key_name="LiteAdmin Slack",
|
||||
key_alias="LiteAdmin Slack",
|
||||
user_id=user.user_id,
|
||||
user_role=LitellmUserRoles.PROXY_ADMIN,
|
||||
models=TypeAdapter(list[str]).validate_python(user.model_dump().get("models", [])),
|
||||
expires=expires,
|
||||
is_session_token=True,
|
||||
)
|
||||
return AdminSession(
|
||||
user_id=user.user_id,
|
||||
credential=SecretStr(
|
||||
encrypt_bearer_token(auth.model_dump_json(exclude_none=True), LITELLM_SESSION_TOKEN_PREFIX)
|
||||
),
|
||||
expires_at=expires.timestamp(),
|
||||
)
|
||||
|
||||
|
||||
def validate_native_configuration(
|
||||
worker_url: str, service_token: str, enterprise: bool, database_available: bool
|
||||
) -> None:
|
||||
if not worker_url:
|
||||
raise HTTPException(404, "LiteAdmin Slack is not enabled")
|
||||
if not enterprise:
|
||||
raise HTTPException(403, "LiteAdmin Slack requires LiteLLM Enterprise")
|
||||
if not database_available:
|
||||
raise HTTPException(503, "LiteAdmin requires a database")
|
||||
try:
|
||||
parsed: Final = urlsplit(worker_url)
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
raise HTTPException(503, "LiteAdmin worker configuration is invalid") from None
|
||||
if (
|
||||
parsed.scheme not in {"http", "https"}
|
||||
or not parsed.hostname
|
||||
or port == 0
|
||||
or parsed.username
|
||||
or parsed.password
|
||||
or parsed.path
|
||||
or parsed.query
|
||||
or parsed.fragment
|
||||
or len(service_token) < 32
|
||||
or any(character.isspace() for character in service_token)
|
||||
):
|
||||
raise HTTPException(503, "LiteAdmin worker configuration is invalid")
|
||||
|
||||
|
||||
async def native_admin_context() -> NativeAdminContext:
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
|
||||
worker_url: Final = os.getenv("LITELLM_ADMIN_AGENT_URL", "").rstrip("/")
|
||||
service_token: Final = os.getenv("ADMIN_AGENT_SERVICE_TOKEN", "")
|
||||
validate_native_configuration(worker_url, service_token, premium_user is True, prisma_client is not None)
|
||||
client: Final = get_async_httpx_client(
|
||||
llm_provider="liteadmin_native", params={"timeout": 15.0, "follow_redirects": False}
|
||||
).client
|
||||
return NativeAdminContext(
|
||||
worker_url, SecretStr(service_token), client, _session_user, _load_user, mint_admin_session
|
||||
)
|
||||
|
||||
|
||||
@router.get(_PREFIX + "{token}", include_in_schema=False, response_class=HTMLResponse)
|
||||
async def connect_page(
|
||||
request: Request,
|
||||
token: str,
|
||||
context: Annotated[NativeAdminContext, Depends(native_admin_context)],
|
||||
) -> Response:
|
||||
details: Final = await context.details(token)
|
||||
base_url: Final = get_request_base_url(request)
|
||||
parsed_base: Final = urlsplit(base_url)
|
||||
if parsed_base.scheme != "https":
|
||||
raise HTTPException(400, "LiteAdmin account connections require HTTPS")
|
||||
user_id: Final = await context.session_user(request)
|
||||
if user_id is None:
|
||||
return RedirectResponse(
|
||||
base_url + "/sso/key/generate?" + urlencode({"return_to": parsed_base.path + _PREFIX + token}),
|
||||
status_code=303,
|
||||
headers=_HEADERS,
|
||||
)
|
||||
await context.admin(user_id, details)
|
||||
csrf: Final = secrets.token_urlsafe(32)
|
||||
page: Final = _page(
|
||||
"Connect LiteAdmin to Slack",
|
||||
f"<p>Connect <strong>{html.escape(details.email)}</strong> to LiteAdmin in your Slack workspace?</p>"
|
||||
"<p>Model requests and administrative actions will use your own LiteLLM account and current permissions</p>"
|
||||
f'<form method="post"><input type="hidden" name="csrf" value="{csrf}">'
|
||||
'<button type="submit">Connect account</button></form>'
|
||||
"<p><small>This connection lasts 24 hours. Send disconnect in Slack to remove the saved session</small></p>",
|
||||
)
|
||||
page.set_cookie(_cookie_name(token), csrf, max_age=600, secure=True, httponly=True, samesite="strict", path="/")
|
||||
return page
|
||||
|
||||
|
||||
@router.post(_PREFIX + "{token}", include_in_schema=False, response_class=HTMLResponse)
|
||||
async def connect_account(
|
||||
request: Request,
|
||||
token: str,
|
||||
context: Annotated[NativeAdminContext, Depends(native_admin_context)],
|
||||
) -> Response:
|
||||
base_url: Final = get_request_base_url(request)
|
||||
parsed_base: Final = urlsplit(base_url)
|
||||
origin: Final = f"{parsed_base.scheme}://{parsed_base.netloc}"
|
||||
if parsed_base.scheme != "https" or request.headers.get("Origin") != origin:
|
||||
raise HTTPException(403, "Reopen your private Slack connection link")
|
||||
if request.headers.get("Content-Type", "").split(";", 1)[0] != "application/x-www-form-urlencoded":
|
||||
raise HTTPException(400, "Expected a connection form")
|
||||
form: Final = await request.form(max_fields=1, max_files=0, max_part_size=1024)
|
||||
supplied: Final = form.get("csrf")
|
||||
expected: Final = request.cookies.get(_cookie_name(token), "")
|
||||
if (
|
||||
not isinstance(supplied, str)
|
||||
or len(expected) != 43
|
||||
or len(supplied) != 43
|
||||
or not hmac.compare_digest(supplied.encode(), expected.encode())
|
||||
):
|
||||
raise HTTPException(403, "Reopen your private Slack connection link")
|
||||
user_id: Final = await context.session_user(request)
|
||||
if user_id is None:
|
||||
raise HTTPException(401, "Your login expired. Reopen your private Slack connection link")
|
||||
details: Final = await context.details(token)
|
||||
user: Final = await context.admin(user_id, details)
|
||||
await context.worker_request(token, context.mint_session(user))
|
||||
page: Final = _page(
|
||||
"Account connected", "<p>Return to Slack and ask LiteAdmin to list your teams or check a budget</p>"
|
||||
)
|
||||
page.delete_cookie(_cookie_name(token), path="/", secure=True, httponly=True, samesite="strict")
|
||||
return page
|
||||
|
|
@ -57,6 +57,19 @@ spec:
|
|||
imagePullPolicy: {{ .Values.image.pullPolicy }}
|
||||
env:
|
||||
{{- include "litellm.proxyEnv" . | nindent 12 }}
|
||||
{{- if .Values.liteadmin.enabled }}
|
||||
- name: LITELLM_ADMIN_AGENT_URL
|
||||
value: {{ printf "http://%s-liteadmin:10000" (include "litellm.fullname" . | trunc 53 | trimSuffix "-") | quote }}
|
||||
- name: ADMIN_AGENT_SERVICE_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "liteadmin.existingSecret is required" .Values.liteadmin.existingSecret }}
|
||||
key: ADMIN_AGENT_SERVICE_TOKEN
|
||||
{{- if not (hasKey (default dict .Values.envVars) "PROXY_BASE_URL") }}
|
||||
- name: PROXY_BASE_URL
|
||||
value: {{ required "liteadmin.gatewayUrl is required" .Values.liteadmin.gatewayUrl | quote }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- include "litellm.proxyMetricsEnv" . | nindent 12 }}
|
||||
{{- if .Values.collector.enabled }}
|
||||
{{- include "litellm.collectorEnv" . | nindent 12 }}
|
||||
|
|
|
|||
112
helm/litellm-helm/templates/liteadmin.yaml
Normal file
112
helm/litellm-helm/templates/liteadmin.yaml
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
{{- if .Values.liteadmin.enabled }}
|
||||
{{- $name := printf "%s-liteadmin" (include "litellm.fullname" . | trunc 53 | trimSuffix "-") }}
|
||||
{{- $secret := required "liteadmin.existingSecret is required" .Values.liteadmin.existingSecret }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ $name }}
|
||||
spec:
|
||||
replicas: 1
|
||||
strategy:
|
||||
type: Recreate
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/name: {{ $name }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
app.kubernetes.io/name: {{ $name }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
spec:
|
||||
automountServiceAccountToken: false
|
||||
terminationGracePeriodSeconds: 75
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
runAsUser: 10001
|
||||
runAsGroup: 10001
|
||||
fsGroup: 10001
|
||||
runAsNonRoot: true
|
||||
containers:
|
||||
- name: liteadmin
|
||||
image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}"
|
||||
imagePullPolicy: {{ .Values.image.pullPolicy }}
|
||||
args: ["--admin-agent"]
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop: [ALL]
|
||||
envFrom:
|
||||
- secretRef:
|
||||
name: {{ $secret }}
|
||||
env:
|
||||
- name: CONNECTION_AUTH_MODE
|
||||
value: native
|
||||
- name: LITELLM_BASE_URL
|
||||
value: {{ required "liteadmin.gatewayUrl is required" .Values.liteadmin.gatewayUrl | quote }}
|
||||
- name: LITELLM_MODEL
|
||||
value: {{ required "liteadmin.model is required" .Values.liteadmin.model | quote }}
|
||||
- name: STATE_DB
|
||||
value: /var/data/events.sqlite3
|
||||
- name: ADMIN_READ_ONLY
|
||||
value: {{ .Values.liteadmin.readOnly | quote }}
|
||||
- name: OPENAI_AGENTS_DISABLE_TRACING
|
||||
value: "1"
|
||||
ports:
|
||||
- name: health
|
||||
containerPort: 10000
|
||||
readinessProbe:
|
||||
httpGet:
|
||||
path: /readyz
|
||||
port: health
|
||||
periodSeconds: 15
|
||||
livenessProbe:
|
||||
httpGet:
|
||||
path: /healthz
|
||||
port: health
|
||||
periodSeconds: 30
|
||||
resources:
|
||||
{{- toYaml .Values.liteadmin.resources | nindent 12 }}
|
||||
volumeMounts:
|
||||
- name: state
|
||||
mountPath: /var/data
|
||||
- name: tmp
|
||||
mountPath: /tmp
|
||||
volumes:
|
||||
- name: state
|
||||
persistentVolumeClaim:
|
||||
claimName: {{ $name }}
|
||||
- name: tmp
|
||||
emptyDir:
|
||||
sizeLimit: 64Mi
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: {{ $name }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
selector:
|
||||
app.kubernetes.io/name: {{ $name }}
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
ports:
|
||||
- port: 10000
|
||||
targetPort: health
|
||||
---
|
||||
apiVersion: v1
|
||||
kind: PersistentVolumeClaim
|
||||
metadata:
|
||||
name: {{ $name }}
|
||||
spec:
|
||||
accessModes: [ReadWriteOnce]
|
||||
{{- with .Values.liteadmin.storageClassName }}
|
||||
storageClassName: {{ . | quote }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.liteadmin.storageSize }}
|
||||
{{- end }}
|
||||
|
|
@ -3,6 +3,20 @@
|
|||
# Declare variables to be passed into your templates.
|
||||
|
||||
replicaCount: 1
|
||||
liteadmin:
|
||||
enabled: false
|
||||
existingSecret: ""
|
||||
gatewayUrl: ""
|
||||
model: ""
|
||||
readOnly: false
|
||||
storageSize: 1Gi
|
||||
storageClassName: ""
|
||||
resources:
|
||||
requests:
|
||||
cpu: 100m
|
||||
memory: 256Mi
|
||||
limits:
|
||||
memory: 1Gi
|
||||
# numWorkers: 2
|
||||
|
||||
image:
|
||||
|
|
|
|||
|
|
@ -471,6 +471,24 @@ Directory of the collector's unix socket, shared by the gateway and
|
|||
collector containers through an emptyDir. Empty when the sidecar is off
|
||||
or gateway.collector.address is a tcp://127.0.0.1:<port> address.
|
||||
*/}}
|
||||
{{- define "litellm.lensWorker.image" -}}
|
||||
{{- if .Values.lensWorker.image.digest -}}
|
||||
{{- if not (regexMatch "^sha256:[0-9a-f]{64}$" .Values.lensWorker.image.digest) -}}
|
||||
{{- fail "lensWorker.image.digest must be sha256 followed by 64 lowercase hex characters" -}}
|
||||
{{- end -}}
|
||||
{{- printf "%s@%s" .Values.lensWorker.image.repository .Values.lensWorker.image.digest -}}
|
||||
{{- else -}}
|
||||
{{- $backendTag := .Values.backend.image.tag | default .Chart.AppVersion -}}
|
||||
{{- $releaseTag := ternary (printf "v%s" $backendTag) $backendTag (regexMatch "^[0-9]" $backendTag) -}}
|
||||
{{- $tag := .Values.lensWorker.image.tag | default $releaseTag -}}
|
||||
{{- $repository := .Values.lensWorker.image.repository -}}
|
||||
{{- if and (hasPrefix "sha-" $tag) (eq $repository "ghcr.io/berriai/litellm-lens-worker") -}}
|
||||
{{- $repository = "ghcr.io/berriai/litellm-lens-worker-dev" -}}
|
||||
{{- end -}}
|
||||
{{- printf "%s:%s" $repository $tag -}}
|
||||
{{- end -}}
|
||||
{{- end -}}
|
||||
|
||||
{{- define "litellm.gateway.collectorSocketDir" -}}
|
||||
{{- if and .Values.gateway.collector.enabled (hasPrefix "unix://" .Values.gateway.collector.address) -}}
|
||||
{{- dir (trimPrefix "unix://" .Values.gateway.collector.address) -}}
|
||||
|
|
|
|||
|
|
@ -57,6 +57,8 @@ spec:
|
|||
containerPort: 4001
|
||||
protocol: TCP
|
||||
env:
|
||||
- name: LENS_WORKER_IMAGE
|
||||
value: {{ include "litellm.lensWorker.image" . | quote }}
|
||||
{{- include "litellm.serverEnv" (dict "root" $ "component" .Values.backend) | nindent 12 }}
|
||||
{{- if .Values.gateway.config.create }}
|
||||
- name: CONFIG_FILE_PATH
|
||||
|
|
|
|||
72
helm/litellm/templates/lens/deployment.yaml
Normal file
72
helm/litellm/templates/lens/deployment.yaml
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
{{- if .Values.lensWorker.enabled }}
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: {{ include "litellm.fullname" . }}-lens-worker
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 4 }}
|
||||
app.kubernetes.io/component: lens-worker
|
||||
spec:
|
||||
replicas: {{ .Values.lensWorker.replicaCount }}
|
||||
selector:
|
||||
matchLabels:
|
||||
app.kubernetes.io/instance: {{ .Release.Name }}
|
||||
app.kubernetes.io/component: lens-worker
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
{{- include "litellm.commonLabels" . | nindent 8 }}
|
||||
app.kubernetes.io/component: lens-worker
|
||||
spec:
|
||||
automountServiceAccountToken: false
|
||||
{{- with .Values.imagePullSecrets }}
|
||||
imagePullSecrets:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
securityContext:
|
||||
runAsNonRoot: true
|
||||
runAsUser: 65532
|
||||
runAsGroup: 65532
|
||||
fsGroup: 65532
|
||||
seccompProfile:
|
||||
type: RuntimeDefault
|
||||
containers:
|
||||
- name: lens-worker
|
||||
image: {{ include "litellm.lensWorker.image" . | quote }}
|
||||
imagePullPolicy: {{ .Values.lensWorker.image.pullPolicy }}
|
||||
securityContext:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
capabilities:
|
||||
drop: [ALL]
|
||||
env:
|
||||
- name: LITELLM_URL
|
||||
value: {{ .Values.lensWorker.url | default (printf "http://%s:%v" (include "litellm.backend.fullname" .) .Values.backend.service.port) | quote }}
|
||||
- name: LENS_WORKER_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ required "lensWorker.tokenSecret.name must reference a Lens worker token" .Values.lensWorker.tokenSecret.name | quote }}
|
||||
key: {{ .Values.lensWorker.tokenSecret.key | quote }}
|
||||
resources:
|
||||
{{- toYaml .Values.lensWorker.resources | nindent 12 }}
|
||||
volumeMounts:
|
||||
- name: tmp
|
||||
mountPath: /tmp
|
||||
volumes:
|
||||
- name: tmp
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
sizeLimit: {{ .Values.lensWorker.tmpSizeLimit }}
|
||||
{{- with .Values.lensWorker.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.lensWorker.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.lensWorker.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
174
helm/litellm/tests/lens_worker_tests.yaml
Normal file
174
helm/litellm/tests/lens_worker_tests.yaml
Normal file
|
|
@ -0,0 +1,174 @@
|
|||
suite: Lens worker release and credentials
|
||||
templates:
|
||||
- lens/deployment.yaml
|
||||
- backend/deployment.yaml
|
||||
- gateway/configmap.yaml
|
||||
values:
|
||||
- ./values/required.yaml
|
||||
tests:
|
||||
- it: installs the development package for a source commit
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: sha-0123456789abcdef
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker-dev:sha-0123456789abcdef
|
||||
- it: advertises the development package for standalone source workers
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: sha-0123456789abcdef
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LENS_WORKER_IMAGE
|
||||
value: ghcr.io/berriai/litellm-lens-worker-dev:sha-0123456789abcdef
|
||||
- it: preserves an explicit private source image repository
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: sha-0123456789abcdef
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
lensWorker.image.repository: registry.example/lens-worker
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: registry.example/lens-worker:sha-0123456789abcdef
|
||||
- it: pins the worker to its approved digest even when its tag changes
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
lensWorker.image.tag: replaced-release
|
||||
lensWorker.image.digest: sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
|
||||
- it: advertises the approved digest to standalone installers
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
lensWorker.image.tag: replaced-release
|
||||
lensWorker.image.digest: sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LENS_WORKER_IMAGE
|
||||
value: ghcr.io/berriai/litellm-lens-worker@sha256:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa
|
||||
- it: refuses a malformed digest instead of falling back to the tag
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
lensWorker.image.digest: sha256:invalid
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: lensWorker.image.digest must be sha256 followed by 64 lowercase hex characters
|
||||
- it: keeps the worker opt in
|
||||
template: lens/deployment.yaml
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
- it: requires a limited worker credential when enabled
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
lensWorker.enabled: true
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: lensWorker.tokenSecret.name must reference a Lens worker token
|
||||
- it: uses the chart release and a secret without granting Kubernetes access
|
||||
template: lens/deployment.yaml
|
||||
chart:
|
||||
appVersion: v1.2.3
|
||||
set:
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker:v1.2.3
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[1].valueFrom.secretKeyRef
|
||||
value:
|
||||
name: lens-credential
|
||||
key: token
|
||||
- equal:
|
||||
path: spec.template.spec.automountServiceAccountToken
|
||||
value: false
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].securityContext.readOnlyRootFilesystem
|
||||
value: true
|
||||
- equal:
|
||||
path: spec.template.spec.volumes[0].emptyDir
|
||||
value:
|
||||
medium: Memory
|
||||
sizeLimit: 1Gi
|
||||
- it: advertises the same private dev image to standalone installers
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
lensWorker.image.repository: registry.example/lens-worker
|
||||
lensWorker.image.tag: branch-main-1234567
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LENS_WORKER_IMAGE
|
||||
value: registry.example/lens-worker:branch-main-1234567
|
||||
- it: supports an external gateway and a registry override
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
lensWorker.url: https://gateway.example/proxy
|
||||
lensWorker.image.repository: registry.example/lens-worker
|
||||
lensWorker.image.tag: branch-main-1234567
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: registry.example/lens-worker:branch-main-1234567
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].env[0].value
|
||||
value: https://gateway.example/proxy
|
||||
- it: prefixes a numeric chart release with v
|
||||
template: lens/deployment.yaml
|
||||
chart:
|
||||
appVersion: 1.2.3-rc.4
|
||||
set:
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker:v1.2.3-rc.4
|
||||
- it: follows a backend image override when no worker tag is set
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: branch-main-1234567
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker:branch-main-1234567
|
||||
- it: recommends the overridden backend release for standalone installers
|
||||
template: backend/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: v1.2.3-dev.4
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: LENS_WORKER_IMAGE
|
||||
value: ghcr.io/berriai/litellm-lens-worker:v1.2.3-dev.4
|
||||
- it: normalizes a numeric backend tag to the published worker tag
|
||||
template: lens/deployment.yaml
|
||||
set:
|
||||
backend.image.tag: 1.2.3-dev.4
|
||||
lensWorker.enabled: true
|
||||
lensWorker.tokenSecret.name: lens-credential
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].image
|
||||
value: ghcr.io/berriai/litellm-lens-worker:v1.2.3-dev.4
|
||||
|
|
@ -629,3 +629,26 @@ ui:
|
|||
affinity: {}
|
||||
# Same shape as gateway.topologySpreadConstraints.
|
||||
topologySpreadConstraints: []
|
||||
|
||||
lensWorker:
|
||||
enabled: false
|
||||
replicaCount: 1
|
||||
image:
|
||||
repository: ghcr.io/berriai/litellm-lens-worker
|
||||
tag: ""
|
||||
digest: ""
|
||||
pullPolicy: IfNotPresent
|
||||
tokenSecret:
|
||||
name: ""
|
||||
key: token
|
||||
url: ""
|
||||
tmpSizeLimit: 1Gi
|
||||
resources:
|
||||
requests:
|
||||
cpu: 100m
|
||||
memory: 256Mi
|
||||
limits:
|
||||
memory: 2Gi
|
||||
nodeSelector: {}
|
||||
tolerations: []
|
||||
affinity: {}
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import functools
|
||||
import glob
|
||||
import os
|
||||
import random
|
||||
|
|
@ -10,7 +11,8 @@ import time
|
|||
from collections.abc import Callable
|
||||
from dataclasses import dataclass, replace
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Final, Optional
|
||||
from typing import TYPE_CHECKING, Final, Optional, Union
|
||||
from urllib.parse import unquote, urlsplit
|
||||
|
||||
from litellm_proxy_extras import prisma_toolchain
|
||||
from litellm_proxy_extras._logging import logger
|
||||
|
|
@ -202,6 +204,66 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
_REDACTED: Final = "REDACTED"
|
||||
_PASSWORD_QUERY_KEYS: Final = frozenset(("password", "sslpassword"))
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _secret_shape_redactor() -> Callable[[str], str]:
|
||||
try:
|
||||
from litellm._logging import redact_secrets
|
||||
except ImportError:
|
||||
return lambda text: text
|
||||
return redact_secrets
|
||||
|
||||
|
||||
def _url_passwords(url: str) -> frozenset[str]:
|
||||
try:
|
||||
parts: Final = urlsplit(url)
|
||||
except ValueError:
|
||||
return frozenset()
|
||||
query_pairs: Final = tuple(pair.partition("=") for pair in parts.query.split("&"))
|
||||
raw_query_passwords: Final = tuple(
|
||||
value for key, separator, value in query_pairs if separator and key.lower() in _PASSWORD_QUERY_KEYS
|
||||
)
|
||||
raw_passwords: Final = ((parts.password,) if parts.password else ()) + raw_query_passwords
|
||||
return frozenset(password for password in raw_passwords + tuple(map(unquote, raw_passwords)) if password)
|
||||
|
||||
|
||||
def _configured_database_passwords() -> frozenset[str]:
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
direct_url: Final = os.getenv("DIRECT_URL")
|
||||
database_passwords: Final = _url_passwords(database_url) if database_url else frozenset()
|
||||
direct_passwords: Final = _url_passwords(direct_url) if direct_url else frozenset()
|
||||
return database_passwords | direct_passwords
|
||||
|
||||
|
||||
def _redact_credentials(text: str) -> str:
|
||||
"""Mask configured database passwords before passing the text to LiteLLM redaction."""
|
||||
passwords: Final = sorted(_configured_database_passwords(), key=len, reverse=True)
|
||||
alternation: Final = "|".join(re.escape(password) for password in passwords)
|
||||
password_pattern: Final = (
|
||||
re.compile(rf"(?P<lead>:|password=)(?:{alternation})(?=@|&|$|[\s'\"\]),])", re.IGNORECASE)
|
||||
if passwords
|
||||
else None
|
||||
)
|
||||
result: Final = password_pattern.sub(rf"\g<lead>{_REDACTED}", text) if password_pattern is not None else text
|
||||
return _secret_shape_redactor()(result)
|
||||
|
||||
|
||||
def _redacted_command(command: object) -> Union[str, tuple[str, ...], list[str]]:
|
||||
if isinstance(command, tuple):
|
||||
return tuple(_redact_credentials(str(argument)) for argument in command)
|
||||
if isinstance(command, list):
|
||||
return [_redact_credentials(str(argument)) for argument in command]
|
||||
return _redact_credentials(str(command))
|
||||
|
||||
|
||||
def _redact_command_error(error: subprocess.CalledProcessError) -> str:
|
||||
redacted_command: Final = _redacted_command(error.cmd)
|
||||
return str(subprocess.CalledProcessError(error.returncode, redacted_command))
|
||||
|
||||
|
||||
def _get_prisma_command() -> str:
|
||||
"""Get the Prisma command to use, bypassing Python wrapper in offline mode."""
|
||||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
|
|
@ -315,7 +377,8 @@ class ProxyExtrasDBManager:
|
|||
return False
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(
|
||||
f"Error creating baseline migration: {e}, {e.stderr}, {e.stdout}"
|
||||
f"Error creating baseline migration: {_redact_command_error(e)}, "
|
||||
f"{_redact_credentials(str(e.stderr))}, {_redact_credentials(str(e.stdout))}"
|
||||
)
|
||||
raise e
|
||||
|
||||
|
|
@ -1572,6 +1635,11 @@ class ProxyExtrasDBManager:
|
|||
f"Error: {stderr}"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
logger.error(
|
||||
"prisma migrate deploy failed with an error the resolver does not handle: "
|
||||
f"{_redact_credentials(stderr)}"
|
||||
)
|
||||
else:
|
||||
if ProxyExtrasDBManager.spend_logs_is_partitioned():
|
||||
raise RuntimeError(PARTITIONED_SPEND_LOGS_PUSH_ERROR)
|
||||
|
|
@ -1586,7 +1654,7 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
logger.error(
|
||||
"Attempt %s timed out. Raise %s if this database needs longer to apply its schema.",
|
||||
attempt + 1,
|
||||
PRISMA_MIGRATE_DEPLOY_TIMEOUT_ENV_VAR if use_migrate else PRISMA_COMMAND_TIMEOUT_ENV_VAR,
|
||||
|
|
@ -1599,7 +1667,12 @@ class ProxyExtrasDBManager:
|
|||
if attempts_left > 0
|
||||
else ""
|
||||
)
|
||||
logger.info(f"The process failed to execute. Details: {e}.{retry_msg}")
|
||||
stderr_detail: Final = (
|
||||
f" stderr: {_redact_credentials(str(e.stderr))}" if e.stderr else ""
|
||||
)
|
||||
logger.error(
|
||||
f"The process failed to execute. Details: {_redact_command_error(e)}.{stderr_detail}{retry_msg}"
|
||||
)
|
||||
time.sleep(random.randrange(5, 15))
|
||||
finally:
|
||||
os.chdir(original_dir)
|
||||
|
|
|
|||
9
litellm-rust/Cargo.lock
generated
9
litellm-rust/Cargo.lock
generated
|
|
@ -4156,6 +4156,7 @@ dependencies = [
|
|||
"litellm-storage-clickhouse",
|
||||
"litellm-token-counter",
|
||||
"litellm-traces",
|
||||
"litellm-traces-cache",
|
||||
"litellm-traces-clickhouse",
|
||||
"litellm-tracing",
|
||||
"prost",
|
||||
|
|
@ -4454,6 +4455,7 @@ name = "litellm-traces"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-llms-types",
|
||||
|
|
@ -4473,13 +4475,17 @@ dependencies = [
|
|||
name = "litellm-traces-cache"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"litellm-traces",
|
||||
"moka",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4487,11 +4493,9 @@ name = "litellm-traces-clickhouse"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"hmac 0.12.1",
|
||||
"itertools 0.14.0",
|
||||
"jsonschema",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
|
|
@ -4510,7 +4514,6 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
|
|||
fancy-regex.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
litellm-traces.workspace = true
|
||||
litellm-traces-cache.workspace = true
|
||||
litellm-traces-clickhouse.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
litellm-host.workspace = true
|
||||
|
|
|
|||
|
|
@ -1,8 +1,11 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::{collections::BTreeMap, sync::Arc};
|
||||
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessParams};
|
||||
use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders};
|
||||
use litellm_traces_cache::{ReadError, TraceReader};
|
||||
use litellm_traces_clickhouse::{
|
||||
ClickHouseTraces, Config, Error, InsertTable, Parameter, QueryReaders,
|
||||
};
|
||||
use prost::Message;
|
||||
use pyo3::{
|
||||
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
|
||||
|
|
@ -10,6 +13,8 @@ use pyo3::{
|
|||
types::PyBytes,
|
||||
};
|
||||
|
||||
pyo3::import_exception!(litellm.rust_bridge.trace.errors, TraceChanged);
|
||||
|
||||
#[derive(Message)]
|
||||
struct OtlpErrorStatus {
|
||||
#[prost(int32, tag = "1")]
|
||||
|
|
@ -35,14 +40,12 @@ fn map_error_ref(error: &Error) -> PyErr {
|
|||
use litellm_storage_clickhouse::Error as StorageError;
|
||||
|
||||
match error {
|
||||
Error::Decode(litellm_traces::Error::TooLarge)
|
||||
| Error::InsertTooLarge
|
||||
| Error::ReadTooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
Error::Decode(litellm_traces::Error::TooLarge) | Error::InsertTooLarge => {
|
||||
PyOverflowError::new_err(error.to_string())
|
||||
}
|
||||
Error::InvalidRow
|
||||
| Error::InvalidLimit(_)
|
||||
| Error::InvalidTable
|
||||
| Error::InvalidCursor(_)
|
||||
| Error::AmbiguousTrace
|
||||
| Error::TraceChanged
|
||||
| Error::Decode(_)
|
||||
| Error::InvalidSchema
|
||||
| Error::InvalidQuery
|
||||
|
|
@ -59,6 +62,7 @@ fn map_error_ref(error: &Error) -> PyErr {
|
|||
Error::Cached(source) => map_error_ref(source),
|
||||
Error::Storage(source) => match source {
|
||||
StorageError::InvalidRow
|
||||
| StorageError::InvalidLimit(_)
|
||||
| StorageError::InvalidTable
|
||||
| StorageError::InvalidSchema
|
||||
| StorageError::EmptySql
|
||||
|
|
@ -76,6 +80,18 @@ fn map_error_ref(error: &Error) -> PyErr {
|
|||
}
|
||||
}
|
||||
|
||||
fn map_read_error(error: ReadError<Error>) -> PyErr {
|
||||
match error {
|
||||
error @ (ReadError::InvalidParameters
|
||||
| ReadError::InvalidCursor(_)
|
||||
| ReadError::AmbiguousTrace) => PyValueError::new_err(error.to_string()),
|
||||
error @ ReadError::TraceChanged => TraceChanged::new_err(error.to_string()),
|
||||
error @ ReadError::TooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
error @ ReadError::Encode(_) => PyRuntimeError::new_err(error.to_string()),
|
||||
ReadError::Store(error) => map_error_ref(&error),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_sql_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => {
|
||||
|
|
@ -110,6 +126,7 @@ impl NativeTraceConfig {
|
|||
pub struct NativeTraceStorage {
|
||||
config: Config,
|
||||
query_readers: QueryReaders,
|
||||
reader: Arc<TraceReader>,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
|
|
@ -121,6 +138,9 @@ impl NativeTraceStorage {
|
|||
config.inner.storage().writer().clone(),
|
||||
config.inner.storage().database().to_owned(),
|
||||
),
|
||||
reader: Arc::new(TraceReader::new(
|
||||
litellm_storage_clickhouse::READ_LIMITS.response_bytes,
|
||||
)),
|
||||
config: config.inner.clone(),
|
||||
})
|
||||
}
|
||||
|
|
@ -216,21 +236,16 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let reader = Arc::clone(&self.reader);
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::list_traces(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
start_ms,
|
||||
end_ms,
|
||||
cursor.as_deref(),
|
||||
limit,
|
||||
)
|
||||
.await
|
||||
let store = ClickHouseTraces::new(client, connection);
|
||||
reader
|
||||
.list_traces(&store, &scope, start_ms, end_ms, cursor.as_deref(), limit)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -246,34 +261,31 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let reader = Arc::clone(&self.reader);
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
let store = ClickHouseTraces::new(client, connection);
|
||||
if let Some(page_size) = page_size {
|
||||
litellm_traces_clickhouse::get_trace_page(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&trace_ref,
|
||||
cursor.as_deref(),
|
||||
page_size,
|
||||
)
|
||||
.await
|
||||
reader
|
||||
.get_trace_page(
|
||||
&store,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&trace_ref,
|
||||
cursor.as_deref(),
|
||||
page_size,
|
||||
)
|
||||
.await
|
||||
} else if cursor.is_some() {
|
||||
Err(Error::InvalidParameters)
|
||||
Err(ReadError::InvalidParameters)
|
||||
} else {
|
||||
litellm_traces_clickhouse::get_trace(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&trace_ref,
|
||||
)
|
||||
.await
|
||||
reader
|
||||
.get_trace(&store, &scope, &trace_id, &trace_ref)
|
||||
.await
|
||||
}
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -287,20 +299,16 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let reader = Arc::clone(&self.reader);
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::get_span(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&span_id,
|
||||
&trace_ref,
|
||||
)
|
||||
.await
|
||||
let store = ClickHouseTraces::new(client, connection);
|
||||
reader
|
||||
.get_span(&store, &scope, &trace_id, &span_id, &trace_ref)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -316,21 +324,23 @@ impl NativeTraceStorage {
|
|||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let reader = Arc::clone(&self.reader);
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::get_span_error(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&span_id,
|
||||
&trace_ref,
|
||||
cursor.as_deref(),
|
||||
)
|
||||
.await
|
||||
let store = ClickHouseTraces::new(client, connection);
|
||||
reader
|
||||
.get_span_error(
|
||||
&store,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&span_id,
|
||||
&trace_ref,
|
||||
cursor.as_deref(),
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
map_read_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -433,6 +443,13 @@ mod tests {
|
|||
|
||||
#[rstest]
|
||||
#[case::row(Error::InvalidRow, "ValueError")]
|
||||
#[case::insert_limit(Error::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES"), "ValueError")]
|
||||
#[case::insert_timeout(
|
||||
Error::Storage(litellm_storage_clickhouse::Error::InvalidLimit(
|
||||
"CLICKHOUSE_INSERT_TIMEOUT_SECONDS"
|
||||
)),
|
||||
"ValueError"
|
||||
)]
|
||||
#[case::insert_budget(Error::InsertTooLarge, "OverflowError")]
|
||||
#[case::scope(Error::InvalidScope, "ValueError")]
|
||||
#[case::schema(Error::SchemaFailed(503), "RuntimeError")]
|
||||
|
|
@ -482,11 +499,11 @@ mod tests {
|
|||
#[rstest]
|
||||
#[case::decode_budget(Error::Decode(litellm_traces::Error::TooLarge), "OverflowError")]
|
||||
#[case::invalid_export(Error::Decode(litellm_traces::Error::InvalidPayload), "ValueError")]
|
||||
#[case::cursor(Error::InvalidCursor("trace"), "ValueError")]
|
||||
#[case::ambiguous(Error::AmbiguousTrace, "ValueError")]
|
||||
#[case::changed_snapshot(Error::TraceChanged, "ValueError")]
|
||||
#[case::read_budget(Error::ReadTooLarge, "OverflowError")]
|
||||
fn trace_read_and_ingest_failures_preserve_public_exception_types(
|
||||
#[case::invalid_decode_limit(
|
||||
Error::Decode(litellm_traces::Error::InvalidLimit("OTLP_MAX_SPANS")),
|
||||
"ValueError"
|
||||
)]
|
||||
fn trace_ingest_failures_preserve_public_exception_types(
|
||||
#[case] error: Error,
|
||||
#[case] exception_name: &str,
|
||||
) {
|
||||
|
|
@ -498,4 +515,43 @@ mod tests {
|
|||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_parameters(ReadError::InvalidParameters, "ValueError")]
|
||||
#[case::invalid_cursor(ReadError::InvalidCursor("trace"), "ValueError")]
|
||||
#[case::ambiguous(ReadError::AmbiguousTrace, "ValueError")]
|
||||
#[case::changed_snapshot(ReadError::TraceChanged, "TraceChanged")]
|
||||
#[case::read_budget(ReadError::TooLarge, "OverflowError")]
|
||||
#[case::encode(
|
||||
ReadError::Encode(Arc::new(serde_json::Error::io(std::io::Error::other("invalid")))),
|
||||
"RuntimeError"
|
||||
)]
|
||||
#[case::store(ReadError::Store(Arc::new(Error::InvalidScope)), "ValueError")]
|
||||
fn trace_read_failures_preserve_public_exception_types(
|
||||
#[case] error: ReadError<Error>,
|
||||
#[case] exception_name: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let repository = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.ancestors()
|
||||
.nth(3)
|
||||
.unwrap()
|
||||
.to_str()
|
||||
.unwrap();
|
||||
pyo3::types::PyModule::import(py, "sys")
|
||||
.unwrap()
|
||||
.getattr("path")
|
||||
.unwrap()
|
||||
.call_method1("insert", (0, repository))
|
||||
.unwrap();
|
||||
let message = error.to_string();
|
||||
let exception = map_read_error(error);
|
||||
assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
|
||||
assert_eq!(
|
||||
exception.value(py).str().unwrap().to_str().unwrap(),
|
||||
message
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
pub enum Error {
|
||||
#[error("invalid ClickHouse insert row")]
|
||||
InvalidRow,
|
||||
#[error("{0} must be a positive integer")]
|
||||
InvalidLimit(&'static str),
|
||||
#[error("invalid ClickHouse insert table")]
|
||||
InvalidTable,
|
||||
#[error("invalid ClickHouse HTTP URL")]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,19 @@ use litellm_http::Client;
|
|||
|
||||
use crate::{Connection, Error, valid_identifier};
|
||||
|
||||
const INSERT_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
fn insert_timeout() -> Result<Duration, Error> {
|
||||
let name = "CLICKHOUSE_INSERT_TIMEOUT_SECONDS";
|
||||
match std::env::var(name) {
|
||||
Ok(value) => value
|
||||
.parse::<u64>()
|
||||
.ok()
|
||||
.filter(|value| *value > 0)
|
||||
.map(Duration::from_secs)
|
||||
.ok_or(Error::InvalidLimit(name)),
|
||||
Err(std::env::VarError::NotPresent) => Ok(Duration::from_secs(30)),
|
||||
Err(_) => Err(Error::InvalidLimit(name)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn insert_encoded_rows(
|
||||
client: &Client,
|
||||
|
|
@ -74,7 +86,7 @@ pub async fn insert_compressed_rows(
|
|||
.append_pair("date_time_input_format", "best_effort");
|
||||
let response = client
|
||||
.post(url)
|
||||
.timeout(INSERT_TIMEOUT)
|
||||
.timeout(insert_timeout()?)
|
||||
.header("Content-Encoding", "gzip")
|
||||
.body(body)
|
||||
.send()
|
||||
|
|
|
|||
|
|
@ -144,3 +144,50 @@ async fn server_result_limits_allow_smaller_pages_without_retrying_other_failure
|
|||
assert!(matches!(error, Error::QueryFailed(500)));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_timeout_environment_controls_transport() {
|
||||
for value in ["1", "3", "0", "invalid"] {
|
||||
let result = std::process::Command::new(std::env::current_exe().unwrap())
|
||||
.args(["--exact", "insert_timeout_environment_child"])
|
||||
.env("LITELLM_TEST_INSERT_TIMEOUT", value)
|
||||
.env("CLICKHOUSE_INSERT_TIMEOUT_SECONDS", value)
|
||||
.output()
|
||||
.unwrap();
|
||||
assert!(
|
||||
result.status.success(),
|
||||
"{}",
|
||||
String::from_utf8_lossy(&result.stdout)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn insert_timeout_environment_child() {
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
|
||||
let Ok(value) = std::env::var("LITELLM_TEST_INSERT_TIMEOUT") else {
|
||||
return;
|
||||
};
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200).set_delay(std::time::Duration::from_millis(1500)))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let result = insert_encoded_rows(
|
||||
&Client::no_redirect_for_test(),
|
||||
&Connection::parse(&server.uri()).unwrap(),
|
||||
"traces",
|
||||
"otel_traces",
|
||||
"token",
|
||||
"{}",
|
||||
)
|
||||
.await;
|
||||
match value.as_str() {
|
||||
"1" => assert!(matches!(result, Err(Error::Transport))),
|
||||
"3" => assert!(result.is_ok()),
|
||||
_ => assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidLimit("CLICKHOUSE_INSERT_TIMEOUT_SECONDS"))
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
Own resolved-trace snapshot storage, cache identity, weighting, and expiry
|
||||
Depend on trace domain types, never storage, HTTP, or Python
|
||||
Own storage-independent trace reads over `TraceStore`: the in-process read cache (identity, single-flight, freshness expiry, weighting), cursor formats, paging, response splitting, spend windows and run batching
|
||||
Depend on trace domain types, never storage, HTTP or Python
|
||||
Preserve the full source and authorization scope in every cache key
|
||||
Keep snapshots immutable and expose borrowed data
|
||||
Keep cursor formats and database reads in their existing owners
|
||||
Storage adapters implement `TraceStore`; keep SQL and row encoding there
|
||||
|
|
|
|||
|
|
@ -6,11 +6,15 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
litellm-traces.workspace = true
|
||||
moka.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
time.workspace = true
|
||||
thiserror.workspace = true
|
||||
tracing.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
|
|
|
|||
395
litellm-rust/crates/traces-cache/src/cache.rs
Normal file
395
litellm-rust/crates/traces-cache/src/cache.rs
Normal file
|
|
@ -0,0 +1,395 @@
|
|||
use std::{future::Future, sync::Arc, time::Duration};
|
||||
|
||||
use litellm_traces::{
|
||||
Trace, TraceSummary,
|
||||
query::named::{ReadAccessParams, TraceSpansRow},
|
||||
};
|
||||
use moka::{Expiry, future::Cache};
|
||||
use serde::Serialize;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
use crate::Error;
|
||||
|
||||
pub const LIVE_TTL: Duration = Duration::from_secs(5);
|
||||
pub const SETTLED_TTL: Duration = Duration::from_secs(10 * 60);
|
||||
const SETTLED_AFTER_MS: u64 = 5 * 60 * 1000;
|
||||
const MAX_INDEX_ENTRIES: u64 = 100_000;
|
||||
|
||||
#[derive(Clone, Eq, Hash, PartialEq)]
|
||||
pub struct SnapshotKey(String);
|
||||
|
||||
impl SnapshotKey {
|
||||
fn digest(fields: &impl Serialize) -> Result<Self, Error> {
|
||||
Ok(Self(format!(
|
||||
"{:x}",
|
||||
Sha256::digest(serde_json::to_vec(fields)?)
|
||||
)))
|
||||
}
|
||||
|
||||
pub fn new(
|
||||
source: &str,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Self, Error> {
|
||||
Self::digest(&(source, access, trace_id, trace_ref, snapshot_ms))
|
||||
}
|
||||
|
||||
pub fn latest(
|
||||
source: &str,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Self, Error> {
|
||||
Self::digest(&(source, access, trace_id, trace_ref))
|
||||
}
|
||||
|
||||
pub(crate) fn run(
|
||||
source: &str,
|
||||
access: &ReadAccessParams,
|
||||
run: (&str, &str, &str, &str),
|
||||
) -> Result<Self, Error> {
|
||||
Self::digest(&("run", source, access, run))
|
||||
}
|
||||
|
||||
pub(crate) fn scope(source: &str, access: &ReadAccessParams) -> Result<Self, Error> {
|
||||
Self::digest(&("scope", source, access))
|
||||
}
|
||||
}
|
||||
|
||||
/// How long a read result stays reusable: traces still receiving spans, or read with spend
|
||||
/// unavailable, are re-read after `LIVE_TTL`; traces quiet for `SETTLED_AFTER_MS` are kept for
|
||||
/// `SETTLED_TTL`.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum Freshness {
|
||||
Live,
|
||||
Settled,
|
||||
}
|
||||
|
||||
impl Freshness {
|
||||
pub fn of(rows: &[TraceSpansRow], spend_known: bool, snapshot_ms: u64) -> Self {
|
||||
let last_end_ms = rows
|
||||
.iter()
|
||||
.map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns) / 1_000_000)
|
||||
.max()
|
||||
.unwrap_or(i64::MAX);
|
||||
let quiet_ms = i64::try_from(snapshot_ms)
|
||||
.unwrap_or(i64::MAX)
|
||||
.saturating_sub(last_end_ms);
|
||||
if spend_known && quiet_ms >= SETTLED_AFTER_MS as i64 {
|
||||
Self::Settled
|
||||
} else {
|
||||
Self::Live
|
||||
}
|
||||
}
|
||||
|
||||
fn ttl(self) -> Duration {
|
||||
match self {
|
||||
Self::Live => LIVE_TTL,
|
||||
Self::Settled => SETTLED_TTL,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trait Fresh {
|
||||
fn freshness(&self) -> Freshness;
|
||||
}
|
||||
|
||||
struct ByFreshness;
|
||||
|
||||
impl<K, V: Fresh> Expiry<K, V> for ByFreshness {
|
||||
fn expire_after_create(&self, _: &K, value: &V, _: std::time::Instant) -> Option<Duration> {
|
||||
Some(value.freshness().ttl())
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Snapshot {
|
||||
trace: Trace,
|
||||
version: String,
|
||||
snapshot_ms: u64,
|
||||
freshness: Freshness,
|
||||
weight: u32,
|
||||
}
|
||||
|
||||
impl Snapshot {
|
||||
pub fn trace(&self) -> &Trace {
|
||||
&self.trace
|
||||
}
|
||||
|
||||
pub fn version(&self) -> &str {
|
||||
&self.version
|
||||
}
|
||||
|
||||
pub fn snapshot_ms(&self) -> u64 {
|
||||
self.snapshot_ms
|
||||
}
|
||||
|
||||
pub fn freshness(&self) -> Freshness {
|
||||
self.freshness
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
struct Latest {
|
||||
snapshot_ms: u64,
|
||||
freshness: Freshness,
|
||||
}
|
||||
|
||||
impl Fresh for Latest {
|
||||
fn freshness(&self) -> Freshness {
|
||||
self.freshness
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolved trace snapshots pinned by `snapshot_ms` for paging, plus which snapshot each trace
|
||||
/// currently serves so repeated opens reuse one read until its freshness expires.
|
||||
pub struct SnapshotCache {
|
||||
pinned: Cache<SnapshotKey, Arc<Snapshot>>,
|
||||
latest: Cache<SnapshotKey, Latest>,
|
||||
max_graph_bytes: usize,
|
||||
}
|
||||
|
||||
impl SnapshotCache {
|
||||
pub fn new(max_graph_bytes: usize, idle: Duration) -> Self {
|
||||
Self {
|
||||
pinned: Cache::builder()
|
||||
.max_capacity((max_graph_bytes as u64).saturating_mul(2))
|
||||
.weigher(|_: &SnapshotKey, snapshot: &Arc<Snapshot>| snapshot.weight)
|
||||
.time_to_idle(idle)
|
||||
.build(),
|
||||
latest: Cache::builder()
|
||||
.max_capacity(MAX_INDEX_ENTRIES)
|
||||
.expire_after(ByFreshness)
|
||||
.build(),
|
||||
max_graph_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get(&self, key: &SnapshotKey) -> Option<Arc<Snapshot>> {
|
||||
self.pinned.get(key).await
|
||||
}
|
||||
|
||||
/// Returns the snapshot pinned at `key`, running `load` once for all concurrent callers on a
|
||||
/// miss. A failed load is not cached.
|
||||
pub async fn pinned_or_load<E, F>(
|
||||
&self,
|
||||
key: SnapshotKey,
|
||||
snapshot_ms: u64,
|
||||
load: F,
|
||||
) -> Result<Arc<Snapshot>, Arc<E>>
|
||||
where
|
||||
E: From<Error> + Send + Sync + 'static,
|
||||
F: Future<Output = Result<(Trace, Freshness), E>>,
|
||||
{
|
||||
self.pinned
|
||||
.try_get_with(key, async {
|
||||
let (trace, freshness) = load.await?;
|
||||
Ok(Arc::new(self.snapshot(trace, snapshot_ms, freshness)?))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Returns the snapshot `latest` currently serves. On a miss, `load_at(now_ms)` runs once for
|
||||
/// all concurrent callers and its snapshot is served until its freshness expires.
|
||||
pub async fn latest_or_load<E, F, Fut>(
|
||||
&self,
|
||||
latest: SnapshotKey,
|
||||
now_ms: u64,
|
||||
load_at: F,
|
||||
) -> Result<Arc<Snapshot>, Arc<E>>
|
||||
where
|
||||
E: Send + Sync + 'static,
|
||||
F: Fn(u64) -> Fut,
|
||||
Fut: Future<Output = Result<Arc<Snapshot>, Arc<E>>>,
|
||||
{
|
||||
let entry = self
|
||||
.latest
|
||||
.try_get_with(latest, async {
|
||||
let snapshot = load_at(now_ms).await?;
|
||||
Ok::<_, Arc<E>>(Latest {
|
||||
snapshot_ms: snapshot.snapshot_ms,
|
||||
freshness: snapshot.freshness,
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|error| Arc::clone(&*error))?;
|
||||
load_at(entry.snapshot_ms).await
|
||||
}
|
||||
|
||||
fn snapshot(
|
||||
&self,
|
||||
trace: Trace,
|
||||
snapshot_ms: u64,
|
||||
freshness: Freshness,
|
||||
) -> Result<Snapshot, Error> {
|
||||
let encoded = serde_json::to_vec(&trace)?;
|
||||
if encoded.len() > self.max_graph_bytes {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
let span_ids: Vec<&str> = trace
|
||||
.spans
|
||||
.iter()
|
||||
.map(|span| span.span_id.as_str())
|
||||
.collect();
|
||||
let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?));
|
||||
Ok(Snapshot {
|
||||
trace,
|
||||
version,
|
||||
snapshot_ms,
|
||||
freshness,
|
||||
weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX),
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
async fn weighted_size(&self) -> u64 {
|
||||
self.pinned.run_pending_tasks().await;
|
||||
self.pinned.weighted_size()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) enum ListedRun {
|
||||
Resolved(Box<TraceSummary>, Freshness),
|
||||
Limited,
|
||||
}
|
||||
|
||||
impl Fresh for ListedRun {
|
||||
fn freshness(&self) -> Freshness {
|
||||
match self {
|
||||
Self::Resolved(_, freshness) => *freshness,
|
||||
Self::Limited => Freshness::Settled,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) struct ListCache {
|
||||
pub(crate) runs: Cache<SnapshotKey, ListedRun>,
|
||||
pub(crate) limits: Cache<SnapshotKey, u32>,
|
||||
}
|
||||
|
||||
impl ListCache {
|
||||
pub(crate) fn new() -> Self {
|
||||
Self {
|
||||
runs: Cache::builder()
|
||||
.max_capacity(MAX_INDEX_ENTRIES)
|
||||
.expire_after(ByFreshness)
|
||||
.build(),
|
||||
limits: Cache::builder()
|
||||
.max_capacity(MAX_INDEX_ENTRIES)
|
||||
.time_to_live(SETTLED_TTL)
|
||||
.build(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_traces::{
|
||||
SpanStatus,
|
||||
query::named::{SpendByResponseIdsRow, TraceSpansRow},
|
||||
resolve_trace,
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn row(span_id: &str) -> TraceSpansRow {
|
||||
TraceSpansRow {
|
||||
trace_id: String::new(),
|
||||
span_id: span_id.into(),
|
||||
parent_span_id: String::new(),
|
||||
name: "run".into(),
|
||||
kind: litellm_traces::ObservationType::Agent,
|
||||
wrapper_candidate: false,
|
||||
agent: "agent".into(),
|
||||
framework: String::new(),
|
||||
status: SpanStatus::Ok,
|
||||
status_message: String::new(),
|
||||
error_truncated: false,
|
||||
start_ns: 1_790_742_989_000_000_000,
|
||||
duration_ns: 10_000_000,
|
||||
service: "agent-demo".into(),
|
||||
input_preview: format!("input of {span_id}"),
|
||||
model: String::new(),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
litellm_request_id: String::new(),
|
||||
call_keys: Vec::new(),
|
||||
call_evidence: None,
|
||||
tool_call_id: String::new(),
|
||||
team_id: String::new(),
|
||||
api_key_hash: String::new(),
|
||||
user_id: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn trace(span_id: &str) -> Trace {
|
||||
resolve_trace(
|
||||
"trace",
|
||||
"ref",
|
||||
&[row(span_id)],
|
||||
&[] as &[SpendByResponseIdsRow],
|
||||
)
|
||||
.expect("fixture should resolve")
|
||||
}
|
||||
|
||||
fn key(suffix: &str) -> SnapshotKey {
|
||||
SnapshotKey::new(
|
||||
"source",
|
||||
&ReadAccessParams {
|
||||
all_teams: false,
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team".into()],
|
||||
},
|
||||
suffix,
|
||||
"ref",
|
||||
100,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn weighted_capacity_bounds_retained_snapshots() {
|
||||
let limit = ["first", "second", "third"]
|
||||
.iter()
|
||||
.map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len())
|
||||
.max()
|
||||
.unwrap();
|
||||
let cache = SnapshotCache::new(limit, Duration::from_secs(120));
|
||||
|
||||
for (key, span_id) in [
|
||||
(key("a"), "first"),
|
||||
(key("b"), "second"),
|
||||
(key("c"), "third"),
|
||||
] {
|
||||
cache
|
||||
.pinned_or_load(key, 100, async {
|
||||
Ok::<_, Error>((trace(span_id), Freshness::Settled))
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
assert!(cache.weighted_size().await <= (limit as u64) * 2);
|
||||
}
|
||||
|
||||
const LAST_END_MS: u64 = 1_790_742_989_010;
|
||||
|
||||
#[rstest]
|
||||
#[case::just_ended(LAST_END_MS, true, Freshness::Live)]
|
||||
#[case::quiet_just_under(LAST_END_MS + SETTLED_AFTER_MS - 1, true, Freshness::Live)]
|
||||
#[case::quiet_long_enough(LAST_END_MS + SETTLED_AFTER_MS, true, Freshness::Settled)]
|
||||
#[case::spend_unknown(LAST_END_MS + SETTLED_AFTER_MS, false, Freshness::Live)]
|
||||
fn freshness_settles_once_spans_stop_and_spend_is_known(
|
||||
#[case] snapshot_ms: u64,
|
||||
#[case] spend_known: bool,
|
||||
#[case] expected: Freshness,
|
||||
) {
|
||||
assert_eq!(
|
||||
Freshness::of(&[row("root")], spend_known, snapshot_ms),
|
||||
expected
|
||||
);
|
||||
}
|
||||
}
|
||||
124
litellm-rust/crates/traces-cache/src/cursor.rs
Normal file
124
litellm-rust/crates/traces-cache/src/cursor.rs
Normal file
|
|
@ -0,0 +1,124 @@
|
|||
use base64::{Engine, engine::general_purpose::URL_SAFE};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::ReadError;
|
||||
|
||||
pub(super) fn encode_cursor<T: Serialize>(position: &T) -> String {
|
||||
URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default())
|
||||
}
|
||||
|
||||
pub(super) fn decode_cursor<T: for<'de> Deserialize<'de>, E>(
|
||||
cursor: &str,
|
||||
kind: &'static str,
|
||||
) -> Result<T, ReadError<E>> {
|
||||
URL_SAFE
|
||||
.decode(cursor)
|
||||
.ok()
|
||||
.and_then(|json| serde_json::from_slice(&json).ok())
|
||||
.ok_or(ReadError::InvalidCursor(kind))
|
||||
}
|
||||
|
||||
pub(super) fn trace_position<E>(cursor: Option<&str>) -> Result<(i64, String), ReadError<E>> {
|
||||
let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else {
|
||||
return Ok((0, String::new()));
|
||||
};
|
||||
match decode_cursor::<(i64, String), E>(cursor, "trace")? {
|
||||
(start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)),
|
||||
_ => Err(ReadError::InvalidCursor("trace")),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
pub(super) struct ErrorPosition {
|
||||
pub(super) offset: u64,
|
||||
pub(super) version: String,
|
||||
}
|
||||
|
||||
pub(super) fn error_position<E>(
|
||||
cursor: Option<&str>,
|
||||
) -> Result<Option<ErrorPosition>, ReadError<E>> {
|
||||
let Some(cursor) = cursor else {
|
||||
return Ok(None);
|
||||
};
|
||||
let position: ErrorPosition = decode_cursor(cursor, "diagnostic")?;
|
||||
let valid_version = position.version.len() == 64
|
||||
&& position
|
||||
.version
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte));
|
||||
if i64::try_from(position.offset).is_err() || !valid_version {
|
||||
return Err(ReadError::InvalidCursor("diagnostic"));
|
||||
}
|
||||
Ok(Some(position))
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
pub(super) struct SpanPosition {
|
||||
pub(super) trace_ref: String,
|
||||
pub(super) snapshot_ms: u64,
|
||||
pub(super) offset: usize,
|
||||
pub(super) version: String,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
fn trace_cursor_round_trips_the_last_listed_run() {
|
||||
let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023"));
|
||||
assert_eq!(
|
||||
trace_position::<std::io::Error>(Some(&cursor)).unwrap(),
|
||||
(
|
||||
1_790_742_989_377,
|
||||
"4bad42b84e9de3ba46fc870185f8f023".to_owned()
|
||||
)
|
||||
);
|
||||
assert_eq!(
|
||||
trace_position::<std::io::Error>(None).unwrap(),
|
||||
(0, String::new())
|
||||
);
|
||||
assert_eq!(
|
||||
trace_position::<std::io::Error>(Some("")).unwrap(),
|
||||
(0, String::new())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_base64("abc")]
|
||||
#[case::not_json("bm90LWpzb24=")]
|
||||
#[case::numeric_reference("WzEsIDJd")]
|
||||
#[case::zero_start("WzAsICJ0Il0=")]
|
||||
fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) {
|
||||
let result: Result<(i64, String), ReadError<std::io::Error>> = trace_position(Some(cursor));
|
||||
assert!(matches!(result, Err(ReadError::InvalidCursor("trace"))));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_base64("garbage")]
|
||||
#[case::missing_fields("e30=")]
|
||||
#[case::not_an_object("WzEsMl0=")]
|
||||
fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) {
|
||||
let result: Result<Option<ErrorPosition>, ReadError<std::io::Error>> =
|
||||
error_position(Some(cursor));
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ReadError::InvalidCursor("diagnostic"))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lowercase_version("a".repeat(64))]
|
||||
#[case::short_version("A".repeat(63))]
|
||||
fn diagnostic_cursor_requires_a_content_version(#[case] version: String) {
|
||||
let cursor = encode_cursor(&ErrorPosition { offset: 1, version });
|
||||
let result: Result<Option<ErrorPosition>, ReadError<std::io::Error>> =
|
||||
error_position(Some(&cursor));
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(ReadError::InvalidCursor("diagnostic"))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,5 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("trace snapshot serialization failed")]
|
||||
|
|
@ -5,3 +7,45 @@ pub enum Error {
|
|||
#[error("trace snapshot exceeds the size limit")]
|
||||
ReadTooLarge,
|
||||
}
|
||||
|
||||
/// Cheap to clone so one failed single-flight read can be returned to every waiting caller.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ReadError<E> {
|
||||
#[error("invalid trace read parameters")]
|
||||
InvalidParameters,
|
||||
#[error("Invalid {0} cursor")]
|
||||
InvalidCursor(&'static str),
|
||||
#[error("Multiple traces have this ID; provide trace_ref")]
|
||||
AmbiguousTrace,
|
||||
#[error("Trace changed while paging; refresh the trace to continue")]
|
||||
TraceChanged,
|
||||
#[error("Trace exceeds the interactive read budget; use a filtered trace query")]
|
||||
TooLarge,
|
||||
#[error("trace could not be encoded")]
|
||||
Encode(#[source] Arc<serde_json::Error>),
|
||||
#[error(transparent)]
|
||||
Store(Arc<E>),
|
||||
}
|
||||
|
||||
impl<E> Clone for ReadError<E> {
|
||||
fn clone(&self) -> Self {
|
||||
match self {
|
||||
Self::InvalidParameters => Self::InvalidParameters,
|
||||
Self::InvalidCursor(kind) => Self::InvalidCursor(kind),
|
||||
Self::AmbiguousTrace => Self::AmbiguousTrace,
|
||||
Self::TraceChanged => Self::TraceChanged,
|
||||
Self::TooLarge => Self::TooLarge,
|
||||
Self::Encode(error) => Self::Encode(Arc::clone(error)),
|
||||
Self::Store(error) => Self::Store(Arc::clone(error)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<E> From<Error> for ReadError<E> {
|
||||
fn from(error: Error) -> Self {
|
||||
match error {
|
||||
Error::ReadTooLarge => Self::TooLarge,
|
||||
Error::Serialization(error) => Self::Encode(Arc::new(error)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,166 +1,12 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_traces::{Trace, query::named::ReadAccessParams};
|
||||
use moka::future::Cache;
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
mod cache;
|
||||
mod cursor;
|
||||
mod error;
|
||||
mod list;
|
||||
mod reader;
|
||||
mod spend;
|
||||
mod store;
|
||||
|
||||
pub use error::Error;
|
||||
|
||||
#[derive(Clone, Eq, Hash, PartialEq)]
|
||||
pub struct SnapshotKey(String);
|
||||
|
||||
impl SnapshotKey {
|
||||
pub fn new(
|
||||
source: &str,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Self, Error> {
|
||||
let encoded = serde_json::to_vec(&(source, access, trace_id, trace_ref, snapshot_ms))?;
|
||||
Ok(Self(format!("{:x}", Sha256::digest(encoded))))
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Snapshot {
|
||||
trace: Trace,
|
||||
version: String,
|
||||
weight: u32,
|
||||
}
|
||||
|
||||
impl Snapshot {
|
||||
pub fn trace(&self) -> &Trace {
|
||||
&self.trace
|
||||
}
|
||||
|
||||
pub fn version(&self) -> &str {
|
||||
&self.version
|
||||
}
|
||||
}
|
||||
|
||||
pub struct SnapshotCache {
|
||||
entries: Cache<SnapshotKey, Arc<Snapshot>>,
|
||||
max_graph_bytes: usize,
|
||||
}
|
||||
|
||||
impl SnapshotCache {
|
||||
pub fn new(max_graph_bytes: usize, ttl: Duration) -> Self {
|
||||
Self {
|
||||
entries: Cache::builder()
|
||||
.max_capacity((max_graph_bytes as u64).saturating_mul(2))
|
||||
.weigher(|_: &SnapshotKey, snapshot: &Arc<Snapshot>| snapshot.weight)
|
||||
.time_to_live(ttl)
|
||||
.build(),
|
||||
max_graph_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get(&self, key: &SnapshotKey) -> Option<Arc<Snapshot>> {
|
||||
self.entries.get(key).await
|
||||
}
|
||||
|
||||
pub async fn insert(&self, key: SnapshotKey, trace: Trace) -> Result<Arc<Snapshot>, Error> {
|
||||
let encoded = serde_json::to_vec(&trace)?;
|
||||
if encoded.len() > self.max_graph_bytes {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
|
||||
let span_ids: Vec<&str> = trace
|
||||
.spans
|
||||
.iter()
|
||||
.map(|span| span.span_id.as_str())
|
||||
.collect();
|
||||
|
||||
let version = format!("{:x}", Sha256::digest(serde_json::to_vec(&span_ids)?));
|
||||
|
||||
let snapshot = Arc::new(Snapshot {
|
||||
trace,
|
||||
version,
|
||||
weight: u32::try_from(encoded.len().saturating_mul(2)).unwrap_or(u32::MAX),
|
||||
});
|
||||
|
||||
self.entries.insert(key, Arc::clone(&snapshot)).await;
|
||||
Ok(snapshot)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_traces::{
|
||||
SpanStatus,
|
||||
query::named::{SpendByResponseIdsRow, TraceSpansRow},
|
||||
resolve_trace,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
|
||||
fn trace(span_id: &str) -> Trace {
|
||||
let rows = [TraceSpansRow {
|
||||
trace_id: String::new(),
|
||||
span_id: span_id.into(),
|
||||
parent_span_id: String::new(),
|
||||
name: "run".into(),
|
||||
kind: litellm_traces::ObservationType::Agent,
|
||||
wrapper_candidate: false,
|
||||
agent: "agent".into(),
|
||||
framework: String::new(),
|
||||
status: SpanStatus::Ok,
|
||||
status_message: String::new(),
|
||||
error_truncated: false,
|
||||
start_ns: 1_790_742_989_000_000_000,
|
||||
duration_ns: 10_000_000,
|
||||
service: "agent-demo".into(),
|
||||
input_preview: format!("input of {span_id}"),
|
||||
model: String::new(),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
litellm_request_id: String::new(),
|
||||
call_keys: Vec::new(),
|
||||
call_evidence: None,
|
||||
tool_call_id: String::new(),
|
||||
team_id: String::new(),
|
||||
api_key_hash: String::new(),
|
||||
user_id: String::new(),
|
||||
}];
|
||||
resolve_trace("trace", "ref", &rows, &[] as &[SpendByResponseIdsRow])
|
||||
.expect("fixture should resolve")
|
||||
}
|
||||
|
||||
fn key(suffix: &str) -> SnapshotKey {
|
||||
SnapshotKey::new(
|
||||
"source",
|
||||
&ReadAccessParams {
|
||||
all_teams: false,
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team".into()],
|
||||
},
|
||||
suffix,
|
||||
"ref",
|
||||
100,
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn weighted_capacity_bounds_retained_snapshots() {
|
||||
let limit = ["first", "second", "third"]
|
||||
.iter()
|
||||
.map(|span_id| serde_json::to_vec(&trace(span_id)).unwrap().len())
|
||||
.max()
|
||||
.unwrap();
|
||||
let cache = SnapshotCache::new(limit, Duration::from_secs(120));
|
||||
|
||||
for (key, span_id) in [
|
||||
(key("a"), "first"),
|
||||
(key("b"), "second"),
|
||||
(key("c"), "third"),
|
||||
] {
|
||||
cache.insert(key, trace(span_id)).await.unwrap();
|
||||
}
|
||||
|
||||
cache.entries.run_pending_tasks().await;
|
||||
assert!(cache.entries.weighted_size() <= (limit as u64) * 2);
|
||||
}
|
||||
}
|
||||
pub use cache::{Freshness, LIVE_TTL, SETTLED_TTL, Snapshot, SnapshotCache, SnapshotKey};
|
||||
pub use error::{Error, ReadError};
|
||||
pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, TraceReader};
|
||||
pub use store::{StoreError, TraceStore};
|
||||
|
|
|
|||
197
litellm-rust/crates/traces-cache/src/list.rs
Normal file
197
litellm-rust/crates/traces-cache/src/list.rs
Normal file
|
|
@ -0,0 +1,197 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use crate::{
|
||||
ReadError, SnapshotKey, TraceReader, TraceStore,
|
||||
cache::{Freshness, ListedRun},
|
||||
reader::{map_store_error, now_ms},
|
||||
spend::{spend, spend_window, spend_within},
|
||||
store::StoreError,
|
||||
};
|
||||
use litellm_traces::{
|
||||
TraceSummary, listed_summary,
|
||||
query::named::{ListTracesRow, ReadAccessParams, TracePageSpansParams, TraceSpansRow},
|
||||
resolve_trace,
|
||||
};
|
||||
|
||||
const RUNS_PER_SPAN_READ: usize = 16;
|
||||
|
||||
fn run_key(team_id: &str, api_key_hash: &str, trace_id: &str) -> (String, String, String) {
|
||||
(
|
||||
team_id.to_owned(),
|
||||
api_key_hash.to_owned(),
|
||||
trace_id.to_owned(),
|
||||
)
|
||||
}
|
||||
|
||||
fn cache_key<E>(
|
||||
source: &str,
|
||||
access: &ReadAccessParams,
|
||||
row: &ListTracesRow,
|
||||
) -> Result<SnapshotKey, ReadError<E>> {
|
||||
Ok(SnapshotKey::run(
|
||||
source,
|
||||
access,
|
||||
(
|
||||
&row.team_id,
|
||||
&row.api_key_hash,
|
||||
&row.trace_id,
|
||||
&row.trace_ref,
|
||||
),
|
||||
)?)
|
||||
}
|
||||
|
||||
fn summary(row: &ListTracesRow, listed: Option<&ListedRun>) -> TraceSummary {
|
||||
match listed {
|
||||
Some(ListedRun::Resolved(summary, _)) => (**summary).clone(),
|
||||
Some(ListedRun::Limited) | None => listed_summary(row),
|
||||
}
|
||||
}
|
||||
|
||||
/// Summaries for one batch of listed runs. Runs resolved within their freshness window come
|
||||
/// from the cache; only the rest are read from storage, with one span and one spend read.
|
||||
pub(super) async fn list_summaries<S: TraceStore>(
|
||||
reader: &TraceReader,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
runs: &[ListTracesRow],
|
||||
) -> Result<Vec<TraceSummary>, ReadError<S::Error>> {
|
||||
let mut keys = Vec::with_capacity(runs.len());
|
||||
let mut listed = Vec::with_capacity(runs.len());
|
||||
for row in runs {
|
||||
let key = cache_key(store.source(), access, row)?;
|
||||
listed.push(reader.lists.runs.get(&key).await);
|
||||
keys.push(key);
|
||||
}
|
||||
let misses: Vec<&ListTracesRow> = runs
|
||||
.iter()
|
||||
.zip(&listed)
|
||||
.filter_map(|(row, listed)| listed.is_none().then_some(row))
|
||||
.collect();
|
||||
let mut resolved = resolve_runs(reader, store, access, &misses)
|
||||
.await?
|
||||
.into_iter();
|
||||
let mut summaries = Vec::with_capacity(runs.len());
|
||||
for ((row, key), cached) in runs.iter().zip(keys).zip(listed) {
|
||||
let listed = match cached {
|
||||
Some(listed) => Some(listed),
|
||||
None => {
|
||||
let listed = resolved.next().flatten();
|
||||
if let Some(listed) = &listed {
|
||||
reader.lists.runs.insert(key, listed.clone()).await;
|
||||
}
|
||||
listed
|
||||
}
|
||||
};
|
||||
summaries.push(summary(row, listed.as_ref()));
|
||||
}
|
||||
Ok(summaries)
|
||||
}
|
||||
|
||||
/// One entry per run; `None` means the run could not be resolved and keeps its listed summary
|
||||
/// without being cached.
|
||||
async fn resolve_runs<S: TraceStore>(
|
||||
reader: &TraceReader,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
runs: &[&ListTracesRow],
|
||||
) -> Result<Vec<Option<ListedRun>>, ReadError<S::Error>> {
|
||||
let (Some(start_ms), Some(end_ms)) = (
|
||||
runs.iter().map(|row| row.start_ms).min(),
|
||||
runs.iter()
|
||||
.map(|row| row.start_ms.saturating_add(row.duration_ms))
|
||||
.max(),
|
||||
) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let params = TracePageSpansParams {
|
||||
access: access.clone(),
|
||||
trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(),
|
||||
start_ms,
|
||||
end_ms: end_ms.saturating_add(1),
|
||||
};
|
||||
let snapshot_ms = now_ms();
|
||||
let spans = match store.run_spans(¶ms, snapshot_ms).await {
|
||||
Ok(spans) => spans,
|
||||
Err(StoreError::TooLarge) => {
|
||||
let mut resolved = Vec::with_capacity(runs.len());
|
||||
for row in runs {
|
||||
resolved.push(resolve_run(reader, store, access, row).await?);
|
||||
}
|
||||
return Ok(resolved);
|
||||
}
|
||||
Err(error) => return Err(map_store_error(error)),
|
||||
};
|
||||
let Some(spend_rows) = spend(store, access, &spans).await else {
|
||||
// The batch's combined spend read failed; a run's own narrower window may still
|
||||
// resolve, so fall back per run instead of leaving every run in the batch costless.
|
||||
let mut resolved = Vec::with_capacity(runs.len());
|
||||
for row in runs {
|
||||
resolved.push(resolve_run(reader, store, access, row).await?);
|
||||
}
|
||||
return Ok(resolved);
|
||||
};
|
||||
let mut spans = spans;
|
||||
spans.sort_by(|left, right| {
|
||||
run_key(&left.team_id, &left.api_key_hash, &left.trace_id)
|
||||
.cmp(&run_key(
|
||||
&right.team_id,
|
||||
&right.api_key_hash,
|
||||
&right.trace_id,
|
||||
))
|
||||
.then(left.start_ns.cmp(&right.start_ns))
|
||||
});
|
||||
let by_run: HashMap<_, &[TraceSpansRow]> = spans
|
||||
.chunk_by(|left, right| {
|
||||
(&left.team_id, &left.api_key_hash, &left.trace_id)
|
||||
== (&right.team_id, &right.api_key_hash, &right.trace_id)
|
||||
})
|
||||
.map(|run| {
|
||||
(
|
||||
run_key(&run[0].team_id, &run[0].api_key_hash, &run[0].trace_id),
|
||||
run,
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
Ok(runs
|
||||
.iter()
|
||||
.map(|row| {
|
||||
let spans = by_run
|
||||
.get(&run_key(&row.team_id, &row.api_key_hash, &row.trace_id))
|
||||
.copied()
|
||||
.unwrap_or_default();
|
||||
let spend =
|
||||
spend_window(spans).map_or(&[][..], |window| spend_within(&spend_rows, window));
|
||||
resolve_trace(&row.trace_id, &row.trace_ref, spans, spend).map(|trace| {
|
||||
ListedRun::Resolved(
|
||||
Box::new(trace.summary),
|
||||
Freshness::of(spans, true, snapshot_ms),
|
||||
)
|
||||
})
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn resolve_run<S: TraceStore>(
|
||||
reader: &TraceReader,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
row: &ListTracesRow,
|
||||
) -> Result<Option<ListedRun>, ReadError<S::Error>> {
|
||||
match reader
|
||||
.current(store, access, &row.trace_id, &row.trace_ref)
|
||||
.await
|
||||
{
|
||||
Ok(snapshot) => Ok(snapshot.map(|snapshot| {
|
||||
ListedRun::Resolved(
|
||||
Box::new(snapshot.trace().summary.clone()),
|
||||
snapshot.freshness(),
|
||||
)
|
||||
})),
|
||||
Err(ReadError::TooLarge) => Ok(Some(ListedRun::Limited)),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn run_batches<T>(runs: &[T]) -> impl Iterator<Item = &[T]> + '_ {
|
||||
runs.chunks(RUNS_PER_SPAN_READ)
|
||||
}
|
||||
365
litellm-rust/crates/traces-cache/src/reader.rs
Normal file
365
litellm-rust/crates/traces-cache/src/reader.rs
Normal file
|
|
@ -0,0 +1,365 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use crate::{
|
||||
ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore,
|
||||
cache::{Freshness, ListCache},
|
||||
cursor::{
|
||||
ErrorPosition, SpanPosition, decode_cursor, encode_cursor, error_position, trace_position,
|
||||
},
|
||||
list::{list_summaries, run_batches},
|
||||
spend::spend,
|
||||
};
|
||||
use litellm_traces::{
|
||||
SpanDetail, SpanErrorPage, Trace, TracePage,
|
||||
query::named::{
|
||||
ListTracesParams, ReadAccessParams, SpanDetailParams, SpanErrorParams, TraceIdentityParams,
|
||||
TraceSpansParams,
|
||||
},
|
||||
resolve_trace, to_ui_content,
|
||||
};
|
||||
|
||||
pub const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024;
|
||||
pub const MAX_GRAPH_SPANS: usize = 100_000;
|
||||
|
||||
const SNAPSHOT_IDLE: Duration = Duration::from_secs(120);
|
||||
|
||||
/// A read that found no trace, kept apart from failures so single-flight waiters share it
|
||||
/// without it being cached.
|
||||
pub(super) enum Miss<E> {
|
||||
Absent,
|
||||
Read(ReadError<E>),
|
||||
}
|
||||
|
||||
impl<E> From<crate::Error> for Miss<E> {
|
||||
fn from(error: crate::Error) -> Self {
|
||||
Self::Read(error.into())
|
||||
}
|
||||
}
|
||||
|
||||
fn settle<T, E>(result: Result<T, Arc<Miss<E>>>) -> Result<Option<T>, ReadError<E>> {
|
||||
match result {
|
||||
Ok(value) => Ok(Some(value)),
|
||||
Err(miss) => match &*miss {
|
||||
Miss::Absent => Ok(None),
|
||||
Miss::Read(error) => Err(error.clone()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TraceReader {
|
||||
snapshots: SnapshotCache,
|
||||
pub(super) lists: ListCache,
|
||||
response_bytes: usize,
|
||||
}
|
||||
|
||||
impl TraceReader {
|
||||
pub fn new(response_bytes: usize) -> Self {
|
||||
Self {
|
||||
snapshots: SnapshotCache::new(MAX_GRAPH_BYTES, SNAPSHOT_IDLE),
|
||||
lists: ListCache::new(),
|
||||
response_bytes,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_traces<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
cursor: Option<&str>,
|
||||
limit: u32,
|
||||
) -> Result<TracePage, ReadError<S::Error>> {
|
||||
if limit == 0 {
|
||||
return Err(ReadError::InvalidParameters);
|
||||
}
|
||||
let (cursor_ms, cursor_trace_id) = trace_position(cursor)?;
|
||||
let scope = SnapshotKey::scope(store.source(), access)?;
|
||||
let accepted = self.lists.limits.get(&scope).await.unwrap_or(u32::MAX);
|
||||
let mut params = ListTracesParams {
|
||||
access: access.clone(),
|
||||
start_ms,
|
||||
end_ms,
|
||||
cursor_ms,
|
||||
cursor_trace_id,
|
||||
limit: limit.min(500).min(accepted),
|
||||
};
|
||||
let page = loop {
|
||||
match store.list_runs(¶ms).await {
|
||||
Err(StoreError::TooLarge) if params.limit > 1 => {
|
||||
params.limit /= 2;
|
||||
self.lists.limits.insert(scope.clone(), params.limit).await;
|
||||
}
|
||||
Err(StoreError::TooLarge) => return Err(ReadError::TooLarge),
|
||||
result => break result.map_err(map_store_error)?,
|
||||
}
|
||||
};
|
||||
let next_cursor = page
|
||||
.last()
|
||||
.filter(|_| page.len() == params.limit as usize)
|
||||
.map(|last| encode_cursor(&(last.start_ms, &last.trace_ref)));
|
||||
let data = {
|
||||
let mut summaries = Vec::with_capacity(page.len());
|
||||
for batch in run_batches(&page) {
|
||||
summaries.extend(list_summaries(self, store, access, batch).await?);
|
||||
}
|
||||
summaries
|
||||
};
|
||||
Ok(TracePage { data, next_cursor })
|
||||
}
|
||||
|
||||
pub async fn get_trace<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<Trace>, ReadError<S::Error>> {
|
||||
let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
Ok(self
|
||||
.current(store, access, trace_id, &trace_ref)
|
||||
.await?
|
||||
.map(|snapshot| snapshot.trace().clone()))
|
||||
}
|
||||
|
||||
pub async fn get_trace_page<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
cursor: Option<&str>,
|
||||
page_size: u32,
|
||||
) -> Result<Option<Trace>, ReadError<S::Error>> {
|
||||
if !(1..=500).contains(&page_size) {
|
||||
return Err(ReadError::InvalidParameters);
|
||||
}
|
||||
let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let Some(cursor) = cursor else {
|
||||
let Some(snapshot) = self.current(store, access, trace_id, &trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let position = SpanPosition {
|
||||
trace_ref,
|
||||
snapshot_ms: snapshot.snapshot_ms(),
|
||||
offset: 0,
|
||||
version: snapshot.version().to_owned(),
|
||||
};
|
||||
return page(&snapshot, &position, page_size, self.response_bytes).map(Some);
|
||||
};
|
||||
let position: SpanPosition = decode_cursor(cursor, "span")?;
|
||||
if position.trace_ref != trace_ref || position.snapshot_ms == 0 {
|
||||
return Err(ReadError::InvalidCursor("span"));
|
||||
}
|
||||
let Some(snapshot) = settle(
|
||||
self.pinned(store, access, trace_id, &trace_ref, position.snapshot_ms)
|
||||
.await,
|
||||
)?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if position.version != snapshot.version() {
|
||||
return Err(ReadError::TraceChanged);
|
||||
}
|
||||
if position.offset > snapshot.trace().spans.len() {
|
||||
return Err(ReadError::InvalidCursor("span"));
|
||||
}
|
||||
page(&snapshot, &position, page_size, self.response_bytes).map(Some)
|
||||
}
|
||||
|
||||
pub(super) async fn current<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<Arc<Snapshot>>, ReadError<S::Error>> {
|
||||
let latest = SnapshotKey::latest(store.source(), access, trace_id, trace_ref)?;
|
||||
settle(
|
||||
self.snapshots
|
||||
.latest_or_load(latest, now_ms(), |snapshot_ms| {
|
||||
self.pinned(store, access, trace_id, trace_ref, snapshot_ms)
|
||||
})
|
||||
.await,
|
||||
)
|
||||
}
|
||||
|
||||
async fn pinned<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Arc<Snapshot>, Arc<Miss<S::Error>>> {
|
||||
let key = SnapshotKey::new(store.source(), access, trace_id, trace_ref, snapshot_ms)
|
||||
.map_err(|error| Arc::new(error.into()))?;
|
||||
self.snapshots
|
||||
.pinned_or_load(key, snapshot_ms, async {
|
||||
let params = TraceSpansParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref: trace_ref.to_owned(),
|
||||
};
|
||||
let rows = store
|
||||
.trace_spans(¶ms, snapshot_ms)
|
||||
.await
|
||||
.map_err(|error| Miss::Read(map_store_error(error)))?;
|
||||
let spend_rows = spend(store, access, &rows).await;
|
||||
let freshness = Freshness::of(&rows, spend_rows.is_some(), snapshot_ms);
|
||||
resolve_trace(
|
||||
trace_id,
|
||||
trace_ref,
|
||||
&rows,
|
||||
spend_rows.as_deref().unwrap_or_default(),
|
||||
)
|
||||
.map(|trace| (trace, freshness))
|
||||
.ok_or(Miss::Absent)
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn get_span<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
span_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<SpanDetail>, ReadError<S::Error>> {
|
||||
let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let params = SpanDetailParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref,
|
||||
span_id: span_id.to_owned(),
|
||||
};
|
||||
let row = store.span_detail(¶ms).await.map_err(map_store_error)?;
|
||||
Ok(row.map(|row| SpanDetail {
|
||||
input_ui: to_ui_content(&row.input),
|
||||
output_ui: to_ui_content(&row.output),
|
||||
span_id: row.span_id,
|
||||
input: row.input,
|
||||
output: row.output,
|
||||
attributes: row.attributes,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn get_span_error<S: TraceStore>(
|
||||
&self,
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
span_id: &str,
|
||||
trace_ref: &str,
|
||||
cursor: Option<&str>,
|
||||
) -> Result<Option<SpanErrorPage>, ReadError<S::Error>> {
|
||||
let position = error_position(cursor)?;
|
||||
let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let offset = position.as_ref().map_or(0, |position| position.offset);
|
||||
let params = SpanErrorParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref,
|
||||
span_id: span_id.to_owned(),
|
||||
error_offset: offset,
|
||||
error_version: position
|
||||
.map(|position| position.version)
|
||||
.unwrap_or_default(),
|
||||
};
|
||||
let Some(row) = store.span_error(¶ms).await.map_err(map_store_error)? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let next_offset = offset + row.message.chars().count() as u64;
|
||||
let next_cursor = (next_offset < row.total_chars).then(|| {
|
||||
encode_cursor(&ErrorPosition {
|
||||
offset: next_offset,
|
||||
version: row.version,
|
||||
})
|
||||
});
|
||||
Ok(Some(SpanErrorPage {
|
||||
span_id: row.span_id,
|
||||
message: row.message,
|
||||
total_chars: row.total_chars,
|
||||
next_cursor,
|
||||
}))
|
||||
}
|
||||
}
|
||||
|
||||
async fn reference<S: TraceStore>(
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<String>, ReadError<S::Error>> {
|
||||
if !trace_ref.is_empty() {
|
||||
return Ok(Some(trace_ref.to_owned()));
|
||||
}
|
||||
let params = TraceIdentityParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
};
|
||||
let identities = store.trace_refs(¶ms).await.map_err(map_store_error)?;
|
||||
if identities.len() > 1 {
|
||||
return Err(ReadError::AmbiguousTrace);
|
||||
}
|
||||
Ok(identities.into_iter().next())
|
||||
}
|
||||
|
||||
fn page<E>(
|
||||
snapshot: &Snapshot,
|
||||
position: &SpanPosition,
|
||||
page_size: u32,
|
||||
response_bytes: usize,
|
||||
) -> Result<Trace, ReadError<E>> {
|
||||
let spans = &snapshot.trace().spans;
|
||||
let create_page = |count: usize| {
|
||||
let end = position.offset.saturating_add(count).min(spans.len());
|
||||
Trace {
|
||||
summary: snapshot.trace().summary.clone(),
|
||||
agents: snapshot.trace().agents.clone(),
|
||||
spans: spans[position.offset..end].to_vec(),
|
||||
next_cursor: (end < spans.len()).then(|| {
|
||||
encode_cursor(&SpanPosition {
|
||||
trace_ref: position.trace_ref.clone(),
|
||||
snapshot_ms: position.snapshot_ms,
|
||||
offset: end,
|
||||
version: snapshot.version().to_owned(),
|
||||
})
|
||||
}),
|
||||
}
|
||||
};
|
||||
let mut trace = create_page(page_size as usize);
|
||||
loop {
|
||||
if serde_json::to_vec(&trace)
|
||||
.map_err(|error| ReadError::Encode(Arc::new(error)))?
|
||||
.len()
|
||||
<= response_bytes
|
||||
{
|
||||
return Ok(trace);
|
||||
}
|
||||
if trace.spans.len() <= 1 {
|
||||
return Err(ReadError::TooLarge);
|
||||
}
|
||||
trace = create_page(trace.spans.len() / 2);
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn now_ms() -> u64 {
|
||||
(time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64
|
||||
}
|
||||
|
||||
pub(super) fn map_store_error<E>(error: StoreError<E>) -> ReadError<E> {
|
||||
match error {
|
||||
StoreError::TooLarge => ReadError::TooLarge,
|
||||
StoreError::Failed(error) => ReadError::Store(Arc::new(error)),
|
||||
}
|
||||
}
|
||||
67
litellm-rust/crates/traces-cache/src/spend.rs
Normal file
67
litellm-rust/crates/traces-cache/src/spend.rs
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
use std::ops::Range;
|
||||
|
||||
use crate::TraceStore;
|
||||
use litellm_traces::{
|
||||
SpendLookup,
|
||||
query::named::{
|
||||
ReadAccessParams, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow,
|
||||
},
|
||||
};
|
||||
|
||||
const NANOS_PER_MS: i64 = 1_000_000;
|
||||
const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000;
|
||||
|
||||
pub(super) fn spend_window(rows: &[TraceSpansRow]) -> Option<Range<i64>> {
|
||||
let start_ns = rows.iter().map(|row| row.start_ns).min()?;
|
||||
let end_ns = rows
|
||||
.iter()
|
||||
.map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns))
|
||||
.max()?;
|
||||
Some(
|
||||
start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS
|
||||
..end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS,
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn spend_within(
|
||||
spend: &[SpendByResponseIdsRow],
|
||||
window: Range<i64>,
|
||||
) -> &[SpendByResponseIdsRow] {
|
||||
let first = spend.partition_point(|row| row.start_ms < window.start);
|
||||
let end = spend.partition_point(|row| row.start_ms < window.end);
|
||||
&spend[first..end.max(first)]
|
||||
}
|
||||
|
||||
/// Spend rows sorted by `start_ms`, or `None` when the lookup failed and spend is unknown.
|
||||
pub(super) async fn spend<S: TraceStore>(
|
||||
store: &S,
|
||||
access: &ReadAccessParams,
|
||||
rows: &[TraceSpansRow],
|
||||
) -> Option<Vec<SpendByResponseIdsRow>> {
|
||||
let lookup = SpendLookup::new(rows);
|
||||
let Some(window) = spend_window(rows) else {
|
||||
return Some(Vec::new());
|
||||
};
|
||||
if lookup.is_empty() {
|
||||
return Some(Vec::new());
|
||||
}
|
||||
let params = SpendByResponseIdsParams {
|
||||
access: access.clone(),
|
||||
response_ids: lookup.response_ids,
|
||||
request_ids: lookup.request_ids,
|
||||
trace_ids: lookup.trace_ids,
|
||||
start_ms: window.start,
|
||||
end_ms: window.end,
|
||||
};
|
||||
match store.spend(¶ms).await {
|
||||
Ok(rows) => {
|
||||
let mut rows = rows;
|
||||
rows.sort_by_key(|row| row.start_ms);
|
||||
Some(rows)
|
||||
}
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "trace spend lookup unavailable");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
62
litellm-rust/crates/traces-cache/src/store.rs
Normal file
62
litellm-rust/crates/traces-cache/src/store.rs
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
use std::future::Future;
|
||||
|
||||
use litellm_traces::query::named::{
|
||||
ListTracesParams, ListTracesRow, SpanDetailParams, SpanDetailRow, SpanErrorParams,
|
||||
SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceIdentityParams,
|
||||
TracePageSpansParams, TraceSpansParams, TraceSpansRow,
|
||||
};
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum StoreError<E> {
|
||||
#[error("trace read exceeds the storage read budget")]
|
||||
TooLarge,
|
||||
#[error(transparent)]
|
||||
Failed(E),
|
||||
}
|
||||
|
||||
pub trait TraceStore: Sync {
|
||||
type Error: std::error::Error + Send + Sync + 'static;
|
||||
|
||||
/// Identifies the backing storage for snapshot cache keys.
|
||||
fn source(&self) -> &str;
|
||||
|
||||
fn trace_refs(
|
||||
&self,
|
||||
params: &TraceIdentityParams,
|
||||
) -> impl Future<Output = Result<Vec<String>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
/// Returns `TooLarge` when the response exceeds the storage limit so the reader can halve `limit`.
|
||||
fn list_runs(
|
||||
&self,
|
||||
params: &ListTracesParams,
|
||||
) -> impl Future<Output = Result<Vec<ListTracesRow>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
/// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`.
|
||||
fn trace_spans(
|
||||
&self,
|
||||
params: &TraceSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> impl Future<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
/// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`.
|
||||
fn run_spans(
|
||||
&self,
|
||||
params: &TracePageSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> impl Future<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
fn spend(
|
||||
&self,
|
||||
params: &SpendByResponseIdsParams,
|
||||
) -> impl Future<Output = Result<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
fn span_detail(
|
||||
&self,
|
||||
params: &SpanDetailParams,
|
||||
) -> impl Future<Output = Result<Option<SpanDetailRow>, StoreError<Self::Error>>> + Send;
|
||||
|
||||
fn span_error(
|
||||
&self,
|
||||
params: &SpanErrorParams,
|
||||
) -> impl Future<Output = Result<Option<SpanErrorRow>, StoreError<Self::Error>>> + Send;
|
||||
}
|
||||
768
litellm-rust/crates/traces-cache/tests/read.rs
Normal file
768
litellm-rust/crates/traces-cache/tests/read.rs
Normal file
|
|
@ -0,0 +1,768 @@
|
|||
use std::{
|
||||
collections::{HashMap, HashSet},
|
||||
sync::{
|
||||
Mutex,
|
||||
atomic::{AtomicUsize, Ordering},
|
||||
},
|
||||
time::Duration,
|
||||
};
|
||||
|
||||
use litellm_traces::{
|
||||
CallEvidenceKind, CallKey, ObservationType, SpanStatus,
|
||||
query::named::{
|
||||
ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetailParams, SpanDetailRow,
|
||||
SpanErrorParams, SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow,
|
||||
TraceIdentityParams, TracePageSpansParams, TraceSpansParams, TraceSpansRow,
|
||||
},
|
||||
};
|
||||
use litellm_traces_cache::{LIVE_TTL, ReadError, StoreError, TraceReader, TraceStore};
|
||||
use rstest::rstest;
|
||||
|
||||
const START_NS: i64 = 1_790_742_989_000_000_000;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
|
||||
enum Operation {
|
||||
TraceRefs,
|
||||
ListRuns,
|
||||
TraceSpans,
|
||||
RunSpans,
|
||||
Spend,
|
||||
SpanDetail,
|
||||
SpanError,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Failure {
|
||||
TooLarge,
|
||||
Failed,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("fake trace store failed")]
|
||||
struct FakeError;
|
||||
|
||||
#[derive(Default)]
|
||||
struct State {
|
||||
failures: HashMap<Operation, Failure>,
|
||||
trace_refs: Vec<String>,
|
||||
list_runs: Vec<ListTracesRow>,
|
||||
trace_spans: HashMap<String, Vec<TraceSpansRow>>,
|
||||
run_spans: Vec<TraceSpansRow>,
|
||||
spend: Vec<SpendByResponseIdsRow>,
|
||||
span_detail: Option<SpanDetailRow>,
|
||||
span_error: Option<SpanErrorRow>,
|
||||
list_runs_too_large_above: Option<u32>,
|
||||
trace_too_large_refs: HashSet<String>,
|
||||
spend_fails_above_response_ids: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Calls {
|
||||
trace_refs: AtomicUsize,
|
||||
list_runs: AtomicUsize,
|
||||
trace_spans: AtomicUsize,
|
||||
run_spans: AtomicUsize,
|
||||
spend: AtomicUsize,
|
||||
span_detail: AtomicUsize,
|
||||
span_error: AtomicUsize,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct FakeStore {
|
||||
state: Mutex<State>,
|
||||
calls: Calls,
|
||||
}
|
||||
|
||||
impl FakeStore {
|
||||
fn with_spans(trace_ref: &str, spans: Vec<TraceSpansRow>) -> Self {
|
||||
Self {
|
||||
state: Mutex::new(State {
|
||||
trace_spans: HashMap::from([(trace_ref.to_owned(), spans)]),
|
||||
..State::default()
|
||||
}),
|
||||
calls: Calls::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_failure(&self, operation: Operation, failure: Failure) {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap()
|
||||
.failures
|
||||
.insert(operation, failure);
|
||||
}
|
||||
|
||||
fn set_trace_refs(&self, trace_refs: Vec<String>) {
|
||||
self.state.lock().unwrap().trace_refs = trace_refs;
|
||||
}
|
||||
|
||||
fn set_list_runs(&self, rows: Vec<ListTracesRow>) {
|
||||
self.state.lock().unwrap().list_runs = rows;
|
||||
}
|
||||
|
||||
fn set_list_runs_too_large_above(&self, limit: u32) {
|
||||
self.state.lock().unwrap().list_runs_too_large_above = Some(limit);
|
||||
}
|
||||
|
||||
fn set_run_spans(&self, rows: Vec<TraceSpansRow>) {
|
||||
self.state.lock().unwrap().run_spans = rows;
|
||||
}
|
||||
|
||||
fn set_trace_spans_too_large(&self, trace_ref: &str) {
|
||||
self.state
|
||||
.lock()
|
||||
.unwrap()
|
||||
.trace_too_large_refs
|
||||
.insert(trace_ref.to_owned());
|
||||
}
|
||||
|
||||
/// Fails `spend` only when the lookup covers more than `limit` response ids, so a batch
|
||||
/// covering several runs fails while each run's own narrower lookup still succeeds.
|
||||
fn set_spend_fails_above_response_ids(&self, limit: usize) {
|
||||
self.state.lock().unwrap().spend_fails_above_response_ids = Some(limit);
|
||||
}
|
||||
|
||||
fn calls(&self, operation: Operation) -> usize {
|
||||
match operation {
|
||||
Operation::TraceRefs => self.calls.trace_refs.load(Ordering::SeqCst),
|
||||
Operation::ListRuns => self.calls.list_runs.load(Ordering::SeqCst),
|
||||
Operation::TraceSpans => self.calls.trace_spans.load(Ordering::SeqCst),
|
||||
Operation::RunSpans => self.calls.run_spans.load(Ordering::SeqCst),
|
||||
Operation::Spend => self.calls.spend.load(Ordering::SeqCst),
|
||||
Operation::SpanDetail => self.calls.span_detail.load(Ordering::SeqCst),
|
||||
Operation::SpanError => self.calls.span_error.load(Ordering::SeqCst),
|
||||
}
|
||||
}
|
||||
|
||||
fn failure(state: &State, operation: Operation) -> Result<(), StoreError<FakeError>> {
|
||||
match state.failures.get(&operation) {
|
||||
Some(Failure::TooLarge) => Err(StoreError::TooLarge),
|
||||
Some(Failure::Failed) => Err(StoreError::Failed(FakeError)),
|
||||
None => Ok(()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl TraceStore for FakeStore {
|
||||
type Error = FakeError;
|
||||
|
||||
fn source(&self) -> &str {
|
||||
"fake"
|
||||
}
|
||||
|
||||
async fn trace_refs(
|
||||
&self,
|
||||
_: &TraceIdentityParams,
|
||||
) -> Result<Vec<String>, StoreError<Self::Error>> {
|
||||
self.calls.trace_refs.fetch_add(1, Ordering::SeqCst);
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::TraceRefs)?;
|
||||
Ok(state.trace_refs.clone())
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
&self,
|
||||
params: &ListTracesParams,
|
||||
) -> Result<Vec<ListTracesRow>, StoreError<Self::Error>> {
|
||||
self.calls.list_runs.fetch_add(1, Ordering::SeqCst);
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::ListRuns)?;
|
||||
if state
|
||||
.list_runs_too_large_above
|
||||
.is_some_and(|limit| params.limit > limit)
|
||||
{
|
||||
return Err(StoreError::TooLarge);
|
||||
}
|
||||
Ok(state
|
||||
.list_runs
|
||||
.iter()
|
||||
.take(params.limit as usize)
|
||||
.cloned()
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn trace_spans(
|
||||
&self,
|
||||
params: &TraceSpansParams,
|
||||
_: u64,
|
||||
) -> Result<Vec<TraceSpansRow>, StoreError<Self::Error>> {
|
||||
self.calls.trace_spans.fetch_add(1, Ordering::SeqCst);
|
||||
tokio::task::yield_now().await;
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::TraceSpans)?;
|
||||
if state.trace_too_large_refs.contains(¶ms.trace_ref) {
|
||||
return Err(StoreError::TooLarge);
|
||||
}
|
||||
Ok(state
|
||||
.trace_spans
|
||||
.get(¶ms.trace_ref)
|
||||
.cloned()
|
||||
.unwrap_or_default())
|
||||
}
|
||||
|
||||
async fn run_spans(
|
||||
&self,
|
||||
_: &TracePageSpansParams,
|
||||
_: u64,
|
||||
) -> Result<Vec<TraceSpansRow>, StoreError<Self::Error>> {
|
||||
self.calls.run_spans.fetch_add(1, Ordering::SeqCst);
|
||||
tokio::task::yield_now().await;
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::RunSpans)?;
|
||||
Ok(state.run_spans.clone())
|
||||
}
|
||||
|
||||
async fn spend(
|
||||
&self,
|
||||
params: &SpendByResponseIdsParams,
|
||||
) -> Result<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>> {
|
||||
self.calls.spend.fetch_add(1, Ordering::SeqCst);
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::Spend)?;
|
||||
if state
|
||||
.spend_fails_above_response_ids
|
||||
.is_some_and(|limit| params.response_ids.len() > limit)
|
||||
{
|
||||
return Err(StoreError::Failed(FakeError));
|
||||
}
|
||||
Ok(state.spend.clone())
|
||||
}
|
||||
|
||||
async fn span_detail(
|
||||
&self,
|
||||
_: &SpanDetailParams,
|
||||
) -> Result<Option<SpanDetailRow>, StoreError<Self::Error>> {
|
||||
self.calls.span_detail.fetch_add(1, Ordering::SeqCst);
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::SpanDetail)?;
|
||||
Ok(state.span_detail.clone())
|
||||
}
|
||||
|
||||
async fn span_error(
|
||||
&self,
|
||||
_: &SpanErrorParams,
|
||||
) -> Result<Option<SpanErrorRow>, StoreError<Self::Error>> {
|
||||
self.calls.span_error.fetch_add(1, Ordering::SeqCst);
|
||||
let state = self.state.lock().unwrap();
|
||||
Self::failure(&state, Operation::SpanError)?;
|
||||
Ok(state.span_error.clone())
|
||||
}
|
||||
}
|
||||
|
||||
fn access() -> ReadAccessParams {
|
||||
ReadAccessParams {
|
||||
all_teams: true,
|
||||
user_id: String::new(),
|
||||
team_ids: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn span(index: usize) -> TraceSpansRow {
|
||||
TraceSpansRow {
|
||||
trace_id: "trace".into(),
|
||||
span_id: format!("span-{index}"),
|
||||
parent_span_id: if index == 0 {
|
||||
String::new()
|
||||
} else {
|
||||
"span-0".into()
|
||||
},
|
||||
name: "agent".into(),
|
||||
kind: ObservationType::Agent,
|
||||
wrapper_candidate: false,
|
||||
agent: "agent".into(),
|
||||
framework: String::new(),
|
||||
status: SpanStatus::Ok,
|
||||
status_message: String::new(),
|
||||
error_truncated: false,
|
||||
start_ns: START_NS + index as i64 * 1_000_000,
|
||||
duration_ns: 10_000_000,
|
||||
service: "test".into(),
|
||||
input_preview: format!("span input {index}"),
|
||||
model: String::new(),
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
litellm_request_id: String::new(),
|
||||
call_keys: Vec::new(),
|
||||
call_evidence: None,
|
||||
tool_call_id: String::new(),
|
||||
team_id: "team".into(),
|
||||
api_key_hash: "key".into(),
|
||||
user_id: "user".into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn run(trace_id: &str, trace_ref: &str) -> ListTracesRow {
|
||||
ListTracesRow {
|
||||
trace_id: trace_id.into(),
|
||||
trace_ref: trace_ref.into(),
|
||||
team_id: "team".into(),
|
||||
api_key_hash: "key".into(),
|
||||
user_id: "user".into(),
|
||||
name: "listed".into(),
|
||||
service: "test".into(),
|
||||
input_preview: String::new(),
|
||||
status: SpanStatus::Ok,
|
||||
start_ms: 1_790_742_989_000,
|
||||
duration_ms: 10,
|
||||
span_count: 1,
|
||||
agent_count: 1,
|
||||
agent_invocations: 1,
|
||||
agent_names: vec!["agent".into()],
|
||||
frameworks: Vec::new(),
|
||||
llm_calls: 0,
|
||||
tool_calls: 0,
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
models: Vec::new(),
|
||||
error_count: 0,
|
||||
request_ids: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn pages_reuse_one_trace_snapshot_and_concatenate_in_order() {
|
||||
let store = FakeStore::with_spans("ref", (0..5).map(span).collect());
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
let first = reader
|
||||
.get_trace_page(&store, &access, "trace", "ref", None, 2)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let second = reader
|
||||
.get_trace_page(
|
||||
&store,
|
||||
&access,
|
||||
"trace",
|
||||
"ref",
|
||||
first.next_cursor.as_deref(),
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let third = reader
|
||||
.get_trace_page(
|
||||
&store,
|
||||
&access,
|
||||
"trace",
|
||||
"ref",
|
||||
second.next_cursor.as_deref(),
|
||||
2,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let ids: Vec<_> = first
|
||||
.spans
|
||||
.iter()
|
||||
.chain(&second.spans)
|
||||
.chain(&third.spans)
|
||||
.map(|span| span.span_id.as_str())
|
||||
.collect();
|
||||
assert_eq!(ids, ["span-0", "span-1", "span-2", "span-3", "span-4"]);
|
||||
assert!(third.next_cursor.is_none());
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn snapshot_versions_are_stable_across_readers_and_detect_changes() {
|
||||
let access = access();
|
||||
let original = FakeStore::with_spans("ref", vec![span(0), span(1)]);
|
||||
let reader_a = TraceReader::new(usize::MAX);
|
||||
let first = reader_a
|
||||
.get_trace_page(&original, &access, "trace", "ref", None, 1)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let cursor = first.next_cursor.unwrap();
|
||||
|
||||
let changed = FakeStore::with_spans("ref", vec![span(0), span(1), span(2)]);
|
||||
let reader_b = TraceReader::new(usize::MAX);
|
||||
let result = reader_b
|
||||
.get_trace_page(&changed, &access, "trace", "ref", Some(&cursor), 1)
|
||||
.await;
|
||||
assert!(matches!(result, Err(ReadError::TraceChanged)));
|
||||
|
||||
let unchanged = FakeStore::with_spans("ref", vec![span(0), span(1)]);
|
||||
let reader_c = TraceReader::new(usize::MAX);
|
||||
let next = reader_c
|
||||
.get_trace_page(&unchanged, &access, "trace", "ref", Some(&cursor), 1)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(next.spans[0].span_id, "span-1");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn response_size_splits_pages_and_rejects_a_single_oversized_span() {
|
||||
let spans: Vec<_> = (0..4)
|
||||
.map(|index| {
|
||||
let mut row = span(index);
|
||||
row.input_preview = "x".repeat(256);
|
||||
row
|
||||
})
|
||||
.collect();
|
||||
let access = access();
|
||||
let full_budget_reader = TraceReader::new(usize::MAX);
|
||||
let one_span = full_budget_reader
|
||||
.get_trace_page(
|
||||
&FakeStore::with_spans("ref", spans.clone()),
|
||||
&access,
|
||||
"trace",
|
||||
"ref",
|
||||
None,
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let response_bytes = serde_json::to_vec(&one_span).unwrap().len() + 128;
|
||||
let reader = TraceReader::new(response_bytes);
|
||||
let store = FakeStore::with_spans("ref", spans.clone());
|
||||
let page = reader
|
||||
.get_trace_page(&store, &access, "trace", "ref", None, 4)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(!page.spans.is_empty());
|
||||
assert!(page.spans.len() < 4);
|
||||
assert!(page.next_cursor.is_some());
|
||||
let continued = reader
|
||||
.get_trace_page(
|
||||
&store,
|
||||
&access,
|
||||
"trace",
|
||||
"ref",
|
||||
page.next_cursor.as_deref(),
|
||||
4,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert!(!continued.spans.is_empty());
|
||||
assert!(matches!(
|
||||
TraceReader::new(1)
|
||||
.get_trace_page(&store, &access, "trace", "ref", None, 1)
|
||||
.await,
|
||||
Err(ReadError::TooLarge)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() {
|
||||
let store = FakeStore::default();
|
||||
store.set_list_runs(
|
||||
(0..3)
|
||||
.map(|index| run(&format!("trace-{index}"), &format!("ref-{index}")))
|
||||
.collect(),
|
||||
);
|
||||
store.set_list_runs_too_large_above(2);
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
let page = reader
|
||||
.list_traces(&store, &access, 0, i64::MAX, None, 8)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.data.len(), 2);
|
||||
assert!(page.next_cursor.is_some());
|
||||
assert_eq!(store.calls(Operation::ListRuns), 3);
|
||||
|
||||
let shorter = FakeStore::default();
|
||||
shorter.set_list_runs(vec![run("only", "ref-only")]);
|
||||
shorter.set_list_runs_too_large_above(2);
|
||||
let page = reader
|
||||
.list_traces(&shorter, &access, 0, i64::MAX, None, 8)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.data.len(), 1);
|
||||
assert!(page.next_cursor.is_none());
|
||||
assert_eq!(shorter.calls(Operation::ListRuns), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() {
|
||||
let store = FakeStore::with_spans("ref-good", vec![span(0)]);
|
||||
store.set_list_runs(vec![
|
||||
run("trace-large", "ref-large"),
|
||||
run("trace-good", "ref-good"),
|
||||
]);
|
||||
store.set_trace_spans_too_large("ref-large");
|
||||
store.set_run_spans(Vec::new());
|
||||
store.set_failure(Operation::RunSpans, Failure::TooLarge);
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let page = reader
|
||||
.list_traces(&store, &access(), 0, i64::MAX, None, 2)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(page.data.len(), 2);
|
||||
assert!(page.data[0].resolution_limited);
|
||||
assert_eq!(page.data[0].trace_ref, "ref-large");
|
||||
assert!(!page.data[1].resolution_limited);
|
||||
assert_eq!(page.data[1].trace_ref, "ref-good");
|
||||
assert_eq!(store.calls(Operation::RunSpans), 1);
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 2);
|
||||
|
||||
let again = reader
|
||||
.list_traces(&store, &access(), 0, i64::MAX, None, 2)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(again.data, page.data);
|
||||
assert_eq!(store.calls(Operation::RunSpans), 1);
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 2);
|
||||
}
|
||||
|
||||
fn spend_row(response_id: &str, cost: f64) -> SpendByResponseIdsRow {
|
||||
SpendByResponseIdsRow {
|
||||
request_id: response_id.into(),
|
||||
litellm_call_id: String::new(),
|
||||
response_id: response_id.into(),
|
||||
upstream_response_id: String::new(),
|
||||
trace_id: String::new(),
|
||||
span_id: String::new(),
|
||||
team_id: "team".into(),
|
||||
api_key: "key".into(),
|
||||
user: "user".into(),
|
||||
spend: Some(cost),
|
||||
start_ms: START_NS / 1_000_000,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn failed_batch_spend_lookup_falls_back_to_each_run_instead_of_losing_every_cost() {
|
||||
let mut first = span(0);
|
||||
first.trace_id = "trace-a".into();
|
||||
first.kind = ObservationType::Llm;
|
||||
first.litellm_request_id = "response-a".into();
|
||||
first.call_keys = vec![CallKey::ProviderResponse("response-a".into())];
|
||||
first.call_evidence = Some(CallEvidenceKind::Complete);
|
||||
let mut second = span(0);
|
||||
second.trace_id = "trace-b".into();
|
||||
second.kind = ObservationType::Llm;
|
||||
second.litellm_request_id = "response-b".into();
|
||||
second.call_keys = vec![CallKey::ProviderResponse("response-b".into())];
|
||||
second.call_evidence = Some(CallEvidenceKind::Complete);
|
||||
|
||||
let store = FakeStore::default();
|
||||
store.set_list_runs(vec![run("trace-a", "ref-a"), run("trace-b", "ref-b")]);
|
||||
store.set_run_spans(vec![first.clone(), second.clone()]);
|
||||
{
|
||||
let mut state = store.state.lock().unwrap();
|
||||
state.trace_spans.insert("ref-a".to_owned(), vec![first]);
|
||||
state.trace_spans.insert("ref-b".to_owned(), vec![second]);
|
||||
state.spend = vec![spend_row("response-a", 1.5), spend_row("response-b", 2.5)];
|
||||
}
|
||||
// The batch covers both runs' response ids (2); each run resolved on its own only ever
|
||||
// asks for its own (1), so this fails only the combined read, not the per-run fallback.
|
||||
store.set_spend_fails_above_response_ids(1);
|
||||
|
||||
let page = TraceReader::new(usize::MAX)
|
||||
.list_traces(&store, &access(), 0, i64::MAX, None, 8)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(page.data.len(), 2);
|
||||
let by_ref: HashMap<&str, f64> = page
|
||||
.data
|
||||
.iter()
|
||||
.map(|run| {
|
||||
(
|
||||
run.trace_ref.as_str(),
|
||||
run.spend
|
||||
.expect("run's own spend read should have succeeded"),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(by_ref["ref-a"], 1.5);
|
||||
assert_eq!(by_ref["ref-b"], 2.5);
|
||||
assert_eq!(store.calls(Operation::RunSpans), 1);
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn failed_spend_lookup_preserves_the_trace_with_unknown_spend() {
|
||||
let mut row = span(0);
|
||||
row.litellm_request_id = "response".into();
|
||||
row.call_keys = vec![CallKey::ProviderResponse("response".into())];
|
||||
row.call_evidence = Some(CallEvidenceKind::Complete);
|
||||
let store = FakeStore::with_spans("ref", vec![row]);
|
||||
store.set_failure(Operation::Spend, Failure::Failed);
|
||||
let trace = TraceReader::new(usize::MAX)
|
||||
.get_trace(&store, &access(), "trace", "ref")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(trace.summary.spend, None);
|
||||
assert_eq!(trace.spans[0].spend, None);
|
||||
assert_eq!(store.calls(Operation::Spend), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn ambiguous_trace_references_fail_and_a_single_reference_is_resolved() {
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
let ambiguous = FakeStore::default();
|
||||
ambiguous.set_trace_refs(vec!["ref-a".into(), "ref-b".into()]);
|
||||
assert!(matches!(
|
||||
reader.get_trace(&ambiguous, &access, "trace", "").await,
|
||||
Err(ReadError::AmbiguousTrace)
|
||||
));
|
||||
|
||||
let unique = FakeStore::with_spans("ref-only", vec![span(0)]);
|
||||
unique.set_trace_refs(vec!["ref-only".into()]);
|
||||
let trace = reader
|
||||
.get_trace(&unique, &access, "trace", "")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(trace.summary.trace_ref, "ref-only");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::zero(0)]
|
||||
#[case::above_max(501)]
|
||||
#[tokio::test]
|
||||
async fn invalid_page_sizes_are_rejected(#[case] page_size: u32) {
|
||||
let store = FakeStore::default();
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
assert!(matches!(
|
||||
reader
|
||||
.get_trace_page(&store, &access, "trace", "ref", None, page_size)
|
||||
.await,
|
||||
Err(ReadError::InvalidParameters)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn zero_list_limit_is_rejected() {
|
||||
let store = FakeStore::default();
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
assert!(matches!(
|
||||
reader
|
||||
.list_traces(&store, &access, 0, i64::MAX, None, 0)
|
||||
.await,
|
||||
Err(ReadError::InvalidParameters)
|
||||
));
|
||||
}
|
||||
|
||||
fn now_ns() -> i64 {
|
||||
time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn concurrent_and_repeated_opens_share_one_storage_read() {
|
||||
let store = FakeStore::with_spans("ref", (0..3).map(span).collect());
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
let (first, second) = tokio::join!(
|
||||
reader.get_trace(&store, &access, "trace", "ref"),
|
||||
reader.get_trace_page(&store, &access, "trace", "ref", None, 2),
|
||||
);
|
||||
let first = first.unwrap().unwrap();
|
||||
let second = second.unwrap().unwrap();
|
||||
let reopened = reader
|
||||
.get_trace_page(&store, &access, "trace", "ref", None, 2)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(first.spans.len(), 3);
|
||||
assert_eq!(second.spans, first.spans[..2]);
|
||||
assert_eq!(reopened.next_cursor, second.next_cursor);
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn failed_reads_are_not_cached() {
|
||||
let store = FakeStore::with_spans("ref", vec![span(0)]);
|
||||
store.set_failure(Operation::TraceSpans, Failure::Failed);
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
assert!(matches!(
|
||||
reader.get_trace(&store, &access, "trace", "ref").await,
|
||||
Err(ReadError::Store(_))
|
||||
));
|
||||
store.state.lock().unwrap().failures.clear();
|
||||
let trace = reader
|
||||
.get_trace(&store, &access, "trace", "ref")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(trace.spans.len(), 1);
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn listed_runs_are_read_once_until_a_live_run_expires() {
|
||||
let live = TraceSpansRow {
|
||||
trace_id: "trace-live".into(),
|
||||
start_ns: now_ns(),
|
||||
..span(0)
|
||||
};
|
||||
let settled = TraceSpansRow {
|
||||
trace_id: "trace-settled".into(),
|
||||
..span(0)
|
||||
};
|
||||
let store = FakeStore::default();
|
||||
store.set_list_runs(vec![
|
||||
run("trace-live", "ref-live"),
|
||||
run("trace-settled", "ref-settled"),
|
||||
]);
|
||||
store.set_run_spans(vec![live, settled]);
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let access = access();
|
||||
let list = || reader.list_traces(&store, &access, 0, i64::MAX, None, 2);
|
||||
|
||||
let first = list().await.unwrap();
|
||||
assert!(first.data.iter().all(|summary| summary.name == "agent"));
|
||||
list().await.unwrap();
|
||||
assert_eq!(store.calls(Operation::RunSpans), 1);
|
||||
|
||||
store.set_run_spans(Vec::new());
|
||||
tokio::time::sleep(LIVE_TTL + Duration::from_millis(200)).await;
|
||||
let after = list().await.unwrap();
|
||||
assert_eq!(store.calls(Operation::RunSpans), 2);
|
||||
assert_eq!(after.data[0].name, "listed");
|
||||
assert_eq!(after.data[1], first.data[1]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn concurrent_pages_of_an_evicted_snapshot_share_one_storage_read() {
|
||||
let access = access();
|
||||
let first = TraceReader::new(usize::MAX)
|
||||
.get_trace_page(
|
||||
&FakeStore::with_spans("ref", (0..3).map(span).collect()),
|
||||
&access,
|
||||
"trace",
|
||||
"ref",
|
||||
None,
|
||||
1,
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let cursor = first.next_cursor.as_deref();
|
||||
let store = FakeStore::with_spans("ref", (0..3).map(span).collect());
|
||||
let reader = TraceReader::new(usize::MAX);
|
||||
let (left, right) = tokio::join!(
|
||||
reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1),
|
||||
reader.get_trace_page(&store, &access, "trace", "ref", cursor, 1),
|
||||
);
|
||||
assert_eq!(left.unwrap().unwrap().spans[0].span_id, "span-1");
|
||||
assert_eq!(right.unwrap().unwrap().spans[0].span_id, "span-1");
|
||||
assert_eq!(store.calls(Operation::TraceSpans), 1);
|
||||
}
|
||||
|
|
@ -5,7 +5,9 @@ use litellm_traces::{
|
|||
query::named::{ReadAccessParams, SpendByResponseIdsRow, TraceSpansRow},
|
||||
resolve_trace,
|
||||
};
|
||||
use litellm_traces_cache::{Error, SnapshotCache, SnapshotKey};
|
||||
use std::sync::Arc;
|
||||
|
||||
use litellm_traces_cache::{Error, Freshness, Snapshot, SnapshotCache, SnapshotKey};
|
||||
use rstest::{fixture, rstest};
|
||||
|
||||
const T0: i64 = 1_790_742_989_000_000_000;
|
||||
|
|
@ -60,6 +62,18 @@ fn key(
|
|||
SnapshotKey::new(source, access, trace_id, trace_ref, ms).unwrap()
|
||||
}
|
||||
|
||||
async fn insert(
|
||||
cache: &SnapshotCache,
|
||||
key: SnapshotKey,
|
||||
trace: Trace,
|
||||
) -> Result<Arc<Snapshot>, Arc<Error>> {
|
||||
cache
|
||||
.pinned_or_load(key, 100, async {
|
||||
Ok::<_, Error>((trace, Freshness::Settled))
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn trace() -> Trace {
|
||||
resolve_trace(
|
||||
|
|
@ -85,7 +99,7 @@ async fn cached_trace_is_isolated_by_access_scope(
|
|||
let cache = SnapshotCache::new(1024 * 1024, TTL);
|
||||
let stored = key("source", &access(), "trace", "ref", 100);
|
||||
|
||||
cache.insert(stored.clone(), trace.clone()).await.unwrap();
|
||||
insert(&cache, stored.clone(), trace.clone()).await.unwrap();
|
||||
|
||||
let other_access = ReadAccessParams {
|
||||
all_teams,
|
||||
|
|
@ -115,7 +129,7 @@ async fn cached_trace_is_isolated_by_key_fields(
|
|||
let cache = SnapshotCache::new(1024 * 1024, TTL);
|
||||
let stored = key("source", &access(), "trace", "ref", 100);
|
||||
|
||||
cache.insert(stored.clone(), trace.clone()).await.unwrap();
|
||||
insert(&cache, stored.clone(), trace.clone()).await.unwrap();
|
||||
|
||||
let other = key(source, &access(), trace_id, trace_ref, snapshot_ms);
|
||||
assert!(cache.get(&other).await.is_none());
|
||||
|
|
@ -129,7 +143,7 @@ async fn snapshot_at_the_size_limit_is_accepted(trace: Trace) {
|
|||
let cache = SnapshotCache::new(size, TTL);
|
||||
let stored = key("source", &access(), "trace", "ref", 100);
|
||||
|
||||
cache.insert(stored.clone(), trace).await.unwrap();
|
||||
insert(&cache, stored.clone(), trace).await.unwrap();
|
||||
assert!(cache.get(&stored).await.is_some());
|
||||
}
|
||||
|
||||
|
|
@ -141,8 +155,8 @@ async fn snapshot_one_byte_over_the_size_limit_is_rejected(trace: Trace) {
|
|||
let stored = key("source", &access(), "trace", "ref", 100);
|
||||
|
||||
assert!(matches!(
|
||||
cache.insert(stored.clone(), trace).await,
|
||||
Err(Error::ReadTooLarge)
|
||||
insert(&cache, stored.clone(), trace).await,
|
||||
Err(error) if matches!(*error, Error::ReadTooLarge)
|
||||
));
|
||||
assert!(cache.get(&stored).await.is_none());
|
||||
}
|
||||
|
|
@ -166,25 +180,31 @@ async fn snapshot_version_tracks_the_ordered_span_ids(
|
|||
};
|
||||
let cache = SnapshotCache::new(1024 * 1024, TTL);
|
||||
|
||||
let first = cache
|
||||
.insert(key("source", &access(), "a", "ref", 100), build(first_ids))
|
||||
.await
|
||||
.unwrap();
|
||||
let second = cache
|
||||
.insert(key("source", &access(), "b", "ref", 100), build(second_ids))
|
||||
.await
|
||||
.unwrap();
|
||||
let first = insert(
|
||||
&cache,
|
||||
key("source", &access(), "a", "ref", 100),
|
||||
build(first_ids),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
let second = insert(
|
||||
&cache,
|
||||
key("source", &access(), "b", "ref", 100),
|
||||
build(second_ids),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(first.version() == second.version(), equal);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn snapshots_expire_after_the_ttl(trace: Trace) {
|
||||
async fn snapshots_expire_when_idle(trace: Trace) {
|
||||
let cache = SnapshotCache::new(1024 * 1024, Duration::from_millis(50));
|
||||
let stored = key("source", &access(), "trace", "ref", 100);
|
||||
|
||||
cache.insert(stored.clone(), trace).await.unwrap();
|
||||
insert(&cache, stored.clone(), trace).await.unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(200)).await;
|
||||
|
||||
assert!(cache.get(&stored).await.is_none());
|
||||
|
|
|
|||
|
|
@ -12,11 +12,9 @@ schema = ["dep:schemars", "litellm-traces/schema"]
|
|||
macro_rules_attribute.workspace = true
|
||||
schemars = { workspace = true, optional = true }
|
||||
askama.workspace = true
|
||||
base64.workspace = true
|
||||
flate2.workspace = true
|
||||
futures-util.workspace = true
|
||||
hmac = "0.12.1"
|
||||
itertools = "0.14.0"
|
||||
litellm-http.workspace = true
|
||||
litellm-migrate.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
|
|
@ -30,7 +28,6 @@ strum.workspace = true
|
|||
thiserror.workspace = true
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,6 @@
|
|||
ALTER TABLE {database}.spend_logs
|
||||
ADD COLUMN IF NOT EXISTS litellm_call_id String DEFAULT '' AFTER response_id,
|
||||
ADD INDEX IF NOT EXISTS idx_litellm_call_id litellm_call_id
|
||||
TYPE bloom_filter(0.001) GRANULARITY 1,
|
||||
ADD INDEX IF NOT EXISTS idx_request_id request_id
|
||||
TYPE bloom_filter(0.001) GRANULARITY 1
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
SELECT
|
||||
request_id, response_id, trace_id, span_id, model, spend,
|
||||
prompt_tokens, completion_tokens, status,
|
||||
JSONExtractBool(metadata, 'synthetic_spend') AS synthetic_spend
|
||||
prompt_tokens, completion_tokens, status
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
ORDER BY start_time DESC, request_id
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
SELECT * FROM (
|
||||
SELECT request_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend,
|
||||
SELECT request_id, litellm_call_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend,
|
||||
toUnixTimestamp64Milli(start_time) AS start_ms
|
||||
FROM (
|
||||
SELECT *,
|
||||
|
|
@ -17,7 +17,8 @@ FROM (
|
|||
)
|
||||
WHERE response_id IN {response_ids:Array(String)}
|
||||
OR upstream_response_id IN {response_ids:Array(String)}
|
||||
OR request_id IN {request_ids:Array(String)}
|
||||
OR litellm_call_id IN {request_ids:Array(String)}
|
||||
OR (litellm_call_id = '' AND request_id IN {request_ids:Array(String)})
|
||||
OR (trace_id != '' AND trace_id IN {trace_ids:Array(String)})
|
||||
ORDER BY start_time DESC
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
SELECT request_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend,
|
||||
SELECT request_id, litellm_call_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend,
|
||||
toUnixTimestamp64Milli(start_time) AS start_ms
|
||||
FROM (
|
||||
SELECT *,
|
||||
|
|
@ -16,6 +16,7 @@ FROM (
|
|||
)
|
||||
WHERE response_id IN {response_ids:Array(String)}
|
||||
OR upstream_response_id IN {response_ids:Array(String)}
|
||||
OR request_id IN {request_ids:Array(String)}
|
||||
OR litellm_call_id IN {request_ids:Array(String)}
|
||||
OR (litellm_call_id = '' AND request_id IN {request_ids:Array(String)})
|
||||
OR (trace_id != '' AND trace_id IN {trace_ids:Array(String)})
|
||||
ORDER BY start_time DESC
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
pub enum Error {
|
||||
#[error("invalid ClickHouse insert row")]
|
||||
InvalidRow,
|
||||
#[error("{0} must be a positive integer")]
|
||||
InvalidLimit(&'static str),
|
||||
#[error("invalid ClickHouse insert table")]
|
||||
InvalidTable,
|
||||
#[error("database must be a nonempty SQL identifier and retention must be positive")]
|
||||
|
|
@ -14,8 +16,6 @@ pub enum Error {
|
|||
InvalidResponse,
|
||||
#[error("ClickHouse insert exceeds the encoded size limit")]
|
||||
InsertTooLarge,
|
||||
#[error("Trace exceeds the interactive read budget; use a filtered trace query")]
|
||||
ReadTooLarge,
|
||||
#[error("ClickHouse schema setup failed with HTTP status {0}")]
|
||||
SchemaFailed(u16),
|
||||
#[error("ClickHouse schema setup transport failed")]
|
||||
|
|
@ -32,12 +32,6 @@ pub enum Error {
|
|||
ProvisionFailed(u16),
|
||||
#[error("ClickHouse reader provisioning transport failed")]
|
||||
ProvisionTransport,
|
||||
#[error("Invalid {0} cursor")]
|
||||
InvalidCursor(&'static str),
|
||||
#[error("Multiple traces have this ID; provide trace_ref")]
|
||||
AmbiguousTrace,
|
||||
#[error("Trace changed while paging; refresh the trace to continue")]
|
||||
TraceChanged,
|
||||
#[error(transparent)]
|
||||
Decode(#[from] litellm_traces::Error),
|
||||
#[error("trace ingestion task failed")]
|
||||
|
|
@ -47,12 +41,3 @@ pub enum Error {
|
|||
#[error(transparent)]
|
||||
Cached(#[from] std::sync::Arc<Error>),
|
||||
}
|
||||
|
||||
impl From<litellm_traces_cache::Error> for Error {
|
||||
fn from(error: litellm_traces_cache::Error) -> Self {
|
||||
match error {
|
||||
litellm_traces_cache::Error::Serialization(_) => Self::InvalidResponse,
|
||||
litellm_traces_cache::Error::ReadTooLarge => Self::ReadTooLarge,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -15,7 +15,18 @@ use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
|||
use super::{Connection, Error};
|
||||
use litellm_traces::Shared;
|
||||
|
||||
const MAX_INSERT_BYTES: usize = 64 * 1024 * 1024;
|
||||
fn max_insert_bytes() -> Result<usize, Error> {
|
||||
let name = "CLICKHOUSE_TRACE_MAX_INSERT_BYTES";
|
||||
match std::env::var(name) {
|
||||
Ok(value) => value
|
||||
.parse::<usize>()
|
||||
.ok()
|
||||
.filter(|value| *value > 0)
|
||||
.ok_or(Error::InvalidLimit(name)),
|
||||
Err(std::env::VarError::NotPresent) => Ok(64 * 1024 * 1024),
|
||||
Err(_) => Err(Error::InvalidLimit(name)),
|
||||
}
|
||||
}
|
||||
|
||||
pub type InsertRow = BTreeMap<String, Shared<Value>>;
|
||||
|
||||
|
|
@ -62,7 +73,7 @@ pub async fn insert_shared_rows(
|
|||
return Ok(());
|
||||
}
|
||||
let received_ms = (OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64;
|
||||
let (token, body) = prepare_insert(&rows, received_ms, MAX_INSERT_BYTES)?;
|
||||
let (token, body) = prepare_insert(&rows, received_ms, max_insert_bytes()?)?;
|
||||
litellm_storage_clickhouse::insert_compressed_rows(
|
||||
client,
|
||||
connection,
|
||||
|
|
@ -228,8 +239,7 @@ mod tests {
|
|||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::Error;
|
||||
use super::{shared_rows, write_rows};
|
||||
use super::{Error, shared_rows, write_rows};
|
||||
|
||||
#[rstest]
|
||||
fn encoded_limit_counts_utf8_bytes_across_rows() {
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ 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 reads::{get_span, get_span_error, get_trace, get_trace_page, list_traces};
|
||||
pub use reads::ClickHouseTraces;
|
||||
pub use schema::{
|
||||
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -204,6 +204,7 @@ impl From<contracts::SpendByResponseIdsParams> for SpendByResponseIdsParams {
|
|||
#[serde(remote = "contracts::SpendByResponseIdsRow")]
|
||||
struct SpendByResponseIdsRowEncoding {
|
||||
pub request_id: String,
|
||||
pub litellm_call_id: String,
|
||||
pub response_id: String,
|
||||
pub upstream_response_id: String,
|
||||
pub trace_id: String,
|
||||
|
|
@ -355,7 +356,7 @@ mod tests {
|
|||
quoted,
|
||||
);
|
||||
round_trip::<SpendByResponseIdsRow>(
|
||||
json!({"request_id": "request", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
|
||||
json!({"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
|
||||
quoted,
|
||||
);
|
||||
}
|
||||
|
|
@ -386,7 +387,7 @@ mod tests {
|
|||
#[case] expected: Option<f64>,
|
||||
) {
|
||||
let row: SpendByResponseIdsRow = serde_json::from_value(json!({
|
||||
"request_id": "request", "response_id": "response", "upstream_response_id": "",
|
||||
"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "",
|
||||
"trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key",
|
||||
"user": "user", "spend": cost, "start_ms": 0
|
||||
}))
|
||||
|
|
@ -399,7 +400,7 @@ mod tests {
|
|||
#[case::boolean(json!(true))]
|
||||
fn spend_rows_reject_invalid_cost(#[case] cost: serde_json::Value) {
|
||||
let row = serde_json::from_value::<SpendByResponseIdsRow>(json!({
|
||||
"request_id": "request", "response_id": "response", "upstream_response_id": "",
|
||||
"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "",
|
||||
"trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key",
|
||||
"user": "user", "spend": cost, "start_ms": 0
|
||||
}));
|
||||
|
|
|
|||
|
|
@ -1,26 +1,13 @@
|
|||
//! Scoped trace reads: the trace list, one trace resolved with its spend, and span payloads.
|
||||
|
||||
use std::sync::LazyLock;
|
||||
use std::time::Duration;
|
||||
|
||||
use base64::{Engine, engine::general_purpose::URL_SAFE};
|
||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||
use itertools::Itertools;
|
||||
use litellm_http::Client;
|
||||
use litellm_storage_clickhouse::{Query, fetch};
|
||||
use litellm_traces::{
|
||||
SpanDetail, SpanErrorPage, SpendLookup, Trace, TracePage, listed_summary,
|
||||
query::named as contracts, resolve_trace, to_ui_content,
|
||||
};
|
||||
use litellm_traces_cache::{SnapshotCache, SnapshotKey};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use litellm_storage_clickhouse::{Error as StorageError, Query, fetch};
|
||||
use litellm_traces::query::named as contracts;
|
||||
use litellm_traces_cache::{StoreError, TraceStore};
|
||||
|
||||
use crate::{
|
||||
Connection, Error,
|
||||
query::named::{
|
||||
ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetail as SpanDetailQuery,
|
||||
SpanDetailParams, SpanError, SpanErrorParams, SpendByResponseIdsParams, TraceIdentity,
|
||||
TraceIdentityParams, TracePageSpansParams, TraceSpansParams,
|
||||
ListTracesParams, ListTracesRow, SpanDetail as SpanDetailQuery, SpanError, SpanErrorParams,
|
||||
SpendByResponseIdsParams, TraceIdentity, TracePageSpansParams,
|
||||
},
|
||||
};
|
||||
|
||||
|
|
@ -36,501 +23,103 @@ impl Query for RunCandidates {
|
|||
);
|
||||
}
|
||||
|
||||
// Cursor pages share a bounded snapshot so advancing does not resolve the whole graph again.
|
||||
static TRACE_SNAPSHOTS: LazyLock<SnapshotCache> = LazyLock::new(|| {
|
||||
SnapshotCache::new(
|
||||
crate::span_batches::MAX_GRAPH_BYTES,
|
||||
Duration::from_secs(120),
|
||||
)
|
||||
});
|
||||
|
||||
const NANOS_PER_MS: i64 = 1_000_000;
|
||||
const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000;
|
||||
const SPEND_CONCURRENCY: usize = 4;
|
||||
|
||||
fn encode_cursor<T: Serialize>(position: &T) -> String {
|
||||
URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default())
|
||||
pub struct ClickHouseTraces {
|
||||
client: Client,
|
||||
connection: Connection,
|
||||
}
|
||||
|
||||
fn decode_cursor<T: for<'de> Deserialize<'de>>(
|
||||
cursor: &str,
|
||||
kind: &'static str,
|
||||
) -> Result<T, Error> {
|
||||
URL_SAFE
|
||||
.decode(cursor)
|
||||
.ok()
|
||||
.and_then(|json| serde_json::from_slice(&json).ok())
|
||||
.ok_or(Error::InvalidCursor(kind))
|
||||
}
|
||||
|
||||
fn trace_position(cursor: Option<&str>) -> Result<(i64, String), Error> {
|
||||
let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else {
|
||||
return Ok((0, String::new()));
|
||||
};
|
||||
match decode_cursor::<(i64, String)>(cursor, "trace")? {
|
||||
(start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)),
|
||||
_ => Err(Error::InvalidCursor("trace")),
|
||||
impl ClickHouseTraces {
|
||||
pub fn new(client: Client, connection: Connection) -> Self {
|
||||
Self { client, connection }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct ErrorPosition {
|
||||
offset: u64,
|
||||
version: String,
|
||||
}
|
||||
impl TraceStore for ClickHouseTraces {
|
||||
type Error = Error;
|
||||
|
||||
fn error_position(cursor: Option<&str>) -> Result<Option<ErrorPosition>, Error> {
|
||||
let Some(cursor) = cursor else {
|
||||
return Ok(None);
|
||||
};
|
||||
let position = decode_cursor::<ErrorPosition>(cursor, "diagnostic")?;
|
||||
let valid_version = position.version.len() == 64
|
||||
&& position
|
||||
.version
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte));
|
||||
if i64::try_from(position.offset).is_err() || !valid_version {
|
||||
return Err(Error::InvalidCursor("diagnostic"));
|
||||
fn source(&self) -> &str {
|
||||
self.connection.url().as_str()
|
||||
}
|
||||
Ok(Some(position))
|
||||
}
|
||||
|
||||
/// The stored run a trace id names for this caller; ids can repeat across tenants and runs.
|
||||
async fn reference(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<String>, Error> {
|
||||
if !trace_ref.is_empty() {
|
||||
return Ok(Some(trace_ref.to_owned()));
|
||||
async fn trace_refs(
|
||||
&self,
|
||||
params: &contracts::TraceIdentityParams,
|
||||
) -> Result<Vec<String>, StoreError<Self::Error>> {
|
||||
fetch::<TraceIdentity>(&self.client, &self.connection, params)
|
||||
.await
|
||||
.map(|rows| rows.into_iter().map(|row| row.trace_ref).collect())
|
||||
.map_err(failed)
|
||||
}
|
||||
let params = TraceIdentityParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
};
|
||||
let mut identities = fetch::<TraceIdentity>(client, connection, ¶ms).await?;
|
||||
if identities.len() > 1 {
|
||||
return Err(Error::AmbiguousTrace);
|
||||
}
|
||||
Ok(identities.pop().map(|identity| identity.trace_ref))
|
||||
}
|
||||
|
||||
/// Spend records behind the spans' calls. A failed lookup leaves cost unknown instead of failing
|
||||
/// the read.
|
||||
async fn spend(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
rows: &[contracts::TraceSpansRow],
|
||||
) -> Vec<contracts::SpendByResponseIdsRow> {
|
||||
let lookup = SpendLookup::new(rows);
|
||||
let (Some(start_ns), Some(end_ns)) = (
|
||||
rows.iter().map(|row| row.start_ns).min(),
|
||||
rows.iter()
|
||||
.map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns))
|
||||
.max(),
|
||||
) else {
|
||||
return Vec::new();
|
||||
};
|
||||
if lookup.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let params = SpendByResponseIdsParams::from(contracts::SpendByResponseIdsParams {
|
||||
access: access.clone(),
|
||||
response_ids: lookup.response_ids,
|
||||
request_ids: lookup.request_ids,
|
||||
trace_ids: lookup.trace_ids,
|
||||
start_ms: start_ns.div_euclid(NANOS_PER_MS) - SPEND_WINDOW_MS,
|
||||
end_ms: end_ns.div_euclid(NANOS_PER_MS) + SPEND_WINDOW_MS,
|
||||
});
|
||||
match crate::span_batches::read_spend(client, connection, params).await {
|
||||
Ok(rows) => rows,
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "trace spend lookup unavailable");
|
||||
Vec::new()
|
||||
async fn list_runs(
|
||||
&self,
|
||||
params: &contracts::ListTracesParams,
|
||||
) -> Result<Vec<contracts::ListTracesRow>, StoreError<Self::Error>> {
|
||||
let storage_params = ListTracesParams::from(params.clone());
|
||||
match fetch::<RunCandidates>(&self.client, &self.connection, &storage_params).await {
|
||||
Ok(rows) => Ok(rows.into_iter().map(|row| row.0).collect()),
|
||||
Err(StorageError::ResponseTooLarge) => Err(StoreError::TooLarge),
|
||||
Err(error) => Err(failed(error)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_traces(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
cursor: Option<&str>,
|
||||
limit: u32,
|
||||
) -> Result<TracePage, Error> {
|
||||
if limit == 0 {
|
||||
return Err(Error::InvalidParameters);
|
||||
async fn trace_spans(
|
||||
&self,
|
||||
params: &contracts::TraceSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
|
||||
crate::span_batches::read_spans(&self.client, &self.connection, params.clone(), snapshot_ms)
|
||||
.await
|
||||
}
|
||||
let (cursor_ms, cursor_trace_id) = trace_position(cursor)?;
|
||||
let mut params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: access.clone(),
|
||||
start_ms,
|
||||
end_ms,
|
||||
cursor_ms,
|
||||
cursor_trace_id,
|
||||
limit: limit.min(500),
|
||||
});
|
||||
let page: Vec<contracts::ListTracesRow> = loop {
|
||||
match fetch::<RunCandidates>(client, connection, ¶ms).await {
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if params.0.limit > 1 => {
|
||||
params.0.limit /= 2;
|
||||
}
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
result => break result?.into_iter().map(|row| row.0).collect(),
|
||||
}
|
||||
};
|
||||
let next_cursor = page
|
||||
.last()
|
||||
.filter(|_| page.len() == params.0.limit as usize)
|
||||
.map(|last| encode_cursor(&(last.start_ms, &last.trace_ref)));
|
||||
let data = stream::iter(page.chunks(16))
|
||||
.then(|batch| list_summaries(client, connection, access, batch))
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect();
|
||||
Ok(TracePage { data, next_cursor })
|
||||
}
|
||||
|
||||
async fn list_summaries(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
runs: &[contracts::ListTracesRow],
|
||||
) -> Result<Vec<litellm_traces::TraceSummary>, Error> {
|
||||
let (Some(start_ms), Some(end_ms)) = (
|
||||
runs.iter().map(|row| row.start_ms).min(),
|
||||
runs.iter()
|
||||
.map(|row| row.start_ms.saturating_add(row.duration_ms))
|
||||
.max(),
|
||||
) else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let params = TracePageSpansParams::from(contracts::TracePageSpansParams {
|
||||
access: access.clone(),
|
||||
trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(),
|
||||
start_ms,
|
||||
end_ms: end_ms.saturating_add(1),
|
||||
});
|
||||
let spans = match crate::span_batches::read_list_spans(client, connection, params).await {
|
||||
Ok(spans) => spans,
|
||||
Err(Error::ReadTooLarge) => {
|
||||
return stream::iter(runs)
|
||||
.then(|row| async move {
|
||||
match get_trace(client, connection, access, &row.trace_id, &row.trace_ref).await
|
||||
{
|
||||
Ok(trace) => {
|
||||
Ok(trace.map_or_else(|| listed_summary(row), |trace| trace.summary))
|
||||
}
|
||||
Err(Error::ReadTooLarge) => Ok(listed_summary(row)),
|
||||
Err(error) => Err(error),
|
||||
}
|
||||
})
|
||||
.try_collect()
|
||||
.await;
|
||||
}
|
||||
Err(error) => return Err(error),
|
||||
};
|
||||
let by_trace = spans.into_iter().into_group_map_by(|span| {
|
||||
(
|
||||
span.team_id.clone(),
|
||||
span.api_key_hash.clone(),
|
||||
span.trace_id.clone(),
|
||||
async fn run_spans(
|
||||
&self,
|
||||
params: &contracts::TracePageSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
|
||||
crate::span_batches::read_list_spans(
|
||||
&self.client,
|
||||
&self.connection,
|
||||
TracePageSpansParams::from(params.clone()),
|
||||
snapshot_ms,
|
||||
)
|
||||
});
|
||||
let summaries = runs
|
||||
.iter()
|
||||
.map(|row| {
|
||||
let spans = by_trace
|
||||
.get(&(
|
||||
row.team_id.clone(),
|
||||
row.api_key_hash.clone(),
|
||||
row.trace_id.clone(),
|
||||
))
|
||||
.map(Vec::as_slice)
|
||||
.unwrap_or_default();
|
||||
async move {
|
||||
let spend_rows = spend(client, connection, access, spans).await;
|
||||
resolve_trace(&row.trace_id, &row.trace_ref, spans, &spend_rows)
|
||||
.map_or_else(|| listed_summary(row), |trace| trace.summary)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(stream::iter(summaries)
|
||||
.buffered(SPEND_CONCURRENCY)
|
||||
.collect()
|
||||
.await)
|
||||
}
|
||||
|
||||
pub async fn get_trace(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<Trace>, Error> {
|
||||
let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let params = TraceSpansParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref: trace_ref.clone(),
|
||||
};
|
||||
let rows = crate::span_batches::read_spans(client, connection, params, u64::MAX).await?;
|
||||
if rows.is_empty() {
|
||||
return Ok(None);
|
||||
.await
|
||||
}
|
||||
let spend_rows = spend(client, connection, access, &rows).await;
|
||||
Ok(resolve_trace(trace_id, &trace_ref, &rows, &spend_rows))
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct SpanPosition {
|
||||
trace_ref: String,
|
||||
snapshot_ms: u64,
|
||||
offset: usize,
|
||||
version: String,
|
||||
}
|
||||
|
||||
pub async fn get_trace_page(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
trace_ref: &str,
|
||||
cursor: Option<&str>,
|
||||
page_size: u32,
|
||||
) -> Result<Option<Trace>, Error> {
|
||||
if !(1..=500).contains(&page_size) {
|
||||
return Err(Error::InvalidParameters);
|
||||
async fn spend(
|
||||
&self,
|
||||
params: &contracts::SpendByResponseIdsParams,
|
||||
) -> Result<Vec<contracts::SpendByResponseIdsRow>, StoreError<Self::Error>> {
|
||||
crate::span_batches::read_spend(
|
||||
&self.client,
|
||||
&self.connection,
|
||||
SpendByResponseIdsParams::from(params.clone()),
|
||||
)
|
||||
.await
|
||||
}
|
||||
let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let position = match cursor {
|
||||
Some(cursor) => {
|
||||
let position: SpanPosition = decode_cursor(cursor, "span")?;
|
||||
if position.trace_ref != trace_ref || position.snapshot_ms == 0 {
|
||||
return Err(Error::InvalidCursor("span"));
|
||||
}
|
||||
position
|
||||
|
||||
async fn span_detail(
|
||||
&self,
|
||||
params: &contracts::SpanDetailParams,
|
||||
) -> Result<Option<contracts::SpanDetailRow>, StoreError<Self::Error>> {
|
||||
match fetch::<SpanDetailQuery>(&self.client, &self.connection, params).await {
|
||||
Ok(rows) => Ok(rows.into_iter().next()),
|
||||
Err(error) => Err(failed(error)),
|
||||
}
|
||||
None => SpanPosition {
|
||||
trace_ref: trace_ref.clone(),
|
||||
snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000)
|
||||
as u64,
|
||||
offset: 0,
|
||||
version: String::new(),
|
||||
},
|
||||
};
|
||||
let key = SnapshotKey::new(
|
||||
connection.url().as_str(),
|
||||
access,
|
||||
trace_id,
|
||||
&trace_ref,
|
||||
position.snapshot_ms,
|
||||
)
|
||||
.map_err(|_| Error::InvalidParameters)?;
|
||||
let snapshot = match TRACE_SNAPSHOTS.get(&key).await {
|
||||
Some(snapshot) => snapshot,
|
||||
None => {
|
||||
let params = TraceSpansParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref: trace_ref.clone(),
|
||||
};
|
||||
let rows =
|
||||
crate::span_batches::read_spans(client, connection, params, position.snapshot_ms)
|
||||
.await?;
|
||||
let spend_rows = spend(client, connection, access, &rows).await;
|
||||
let Some(trace) = resolve_trace(trace_id, &trace_ref, &rows, &spend_rows) else {
|
||||
return Ok(None);
|
||||
};
|
||||
TRACE_SNAPSHOTS.insert(key, trace).await?
|
||||
}
|
||||
|
||||
async fn span_error(
|
||||
&self,
|
||||
params: &contracts::SpanErrorParams,
|
||||
) -> Result<Option<contracts::SpanErrorRow>, StoreError<Self::Error>> {
|
||||
let storage_params = SpanErrorParams::from(params.clone());
|
||||
match fetch::<SpanError>(&self.client, &self.connection, &storage_params).await {
|
||||
Ok(rows) => Ok(rows.into_iter().next().map(|row| row.0)),
|
||||
Err(error) => Err(failed(error)),
|
||||
}
|
||||
};
|
||||
let spans = &snapshot.trace().spans;
|
||||
if cursor.is_some() && position.version != snapshot.version() {
|
||||
return Err(Error::TraceChanged);
|
||||
}
|
||||
let mut trace = Trace {
|
||||
summary: snapshot.trace().summary.clone(),
|
||||
agents: snapshot.trace().agents.clone(),
|
||||
spans: Vec::new(),
|
||||
next_cursor: None,
|
||||
};
|
||||
if position.offset > spans.len() {
|
||||
return Err(Error::InvalidCursor("span"));
|
||||
}
|
||||
let end = position
|
||||
.offset
|
||||
.saturating_add(page_size as usize)
|
||||
.min(spans.len());
|
||||
trace.next_cursor = (end < spans.len()).then(|| {
|
||||
encode_cursor(&SpanPosition {
|
||||
offset: end,
|
||||
version: snapshot.version().to_owned(),
|
||||
..position
|
||||
})
|
||||
});
|
||||
trace.spans = spans[position.offset..end].to_vec();
|
||||
while serde_json::to_vec(&trace)
|
||||
.map_err(|_| Error::InvalidResponse)?
|
||||
.len()
|
||||
> litellm_storage_clickhouse::READ_LIMITS.response_bytes
|
||||
{
|
||||
if trace.spans.len() <= 1 {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
trace.spans.truncate(trace.spans.len() / 2);
|
||||
trace.next_cursor = Some(encode_cursor(&SpanPosition {
|
||||
trace_ref: trace_ref.clone(),
|
||||
snapshot_ms: position.snapshot_ms,
|
||||
offset: position.offset + trace.spans.len(),
|
||||
version: snapshot.version().to_owned(),
|
||||
}));
|
||||
}
|
||||
Ok(Some(trace))
|
||||
}
|
||||
|
||||
pub async fn get_span(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
span_id: &str,
|
||||
trace_ref: &str,
|
||||
) -> Result<Option<SpanDetail>, Error> {
|
||||
let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let params = SpanDetailParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref,
|
||||
span_id: span_id.to_owned(),
|
||||
};
|
||||
let row = fetch::<SpanDetailQuery>(client, connection, ¶ms)
|
||||
.await?
|
||||
.into_iter()
|
||||
.next();
|
||||
Ok(row.map(|row| SpanDetail {
|
||||
input_ui: to_ui_content(&row.input),
|
||||
output_ui: to_ui_content(&row.output),
|
||||
span_id: row.span_id,
|
||||
input: row.input,
|
||||
output: row.output,
|
||||
attributes: row.attributes,
|
||||
}))
|
||||
}
|
||||
|
||||
pub async fn get_span_error(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
access: &ReadAccessParams,
|
||||
trace_id: &str,
|
||||
span_id: &str,
|
||||
trace_ref: &str,
|
||||
cursor: Option<&str>,
|
||||
) -> Result<Option<SpanErrorPage>, Error> {
|
||||
let position = error_position(cursor)?;
|
||||
let Some(trace_ref) = reference(client, connection, access, trace_id, trace_ref).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
let offset = position.as_ref().map_or(0, |position| position.offset);
|
||||
let params = SpanErrorParams::from(contracts::SpanErrorParams {
|
||||
access: access.clone(),
|
||||
trace_id: trace_id.to_owned(),
|
||||
trace_ref,
|
||||
span_id: span_id.to_owned(),
|
||||
error_offset: offset,
|
||||
error_version: position
|
||||
.map(|position| position.version)
|
||||
.unwrap_or_default(),
|
||||
});
|
||||
let Some(row) = fetch::<SpanError>(client, connection, ¶ms)
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let row = row.0;
|
||||
let next_offset = offset + row.message.chars().count() as u64;
|
||||
let next_cursor = (next_offset < row.total_chars).then(|| {
|
||||
encode_cursor(&ErrorPosition {
|
||||
offset: next_offset,
|
||||
version: row.version,
|
||||
})
|
||||
});
|
||||
Ok(Some(SpanErrorPage {
|
||||
span_id: row.span_id,
|
||||
message: row.message,
|
||||
total_chars: row.total_chars,
|
||||
next_cursor,
|
||||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
fn trace_cursor_round_trips_the_last_listed_run() {
|
||||
let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023"));
|
||||
assert_eq!(
|
||||
trace_position(Some(&cursor)).unwrap(),
|
||||
(
|
||||
1_790_742_989_377,
|
||||
"4bad42b84e9de3ba46fc870185f8f023".to_owned()
|
||||
)
|
||||
);
|
||||
assert_eq!(trace_position(None).unwrap(), (0, String::new()));
|
||||
assert_eq!(trace_position(Some("")).unwrap(), (0, String::new()));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_base64("abc")]
|
||||
#[case::not_json("bm90LWpzb24=")]
|
||||
#[case::numeric_reference("WzEsIDJd")]
|
||||
#[case::zero_start("WzAsICJ0Il0=")]
|
||||
fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) {
|
||||
assert!(matches!(
|
||||
trace_position(Some(cursor)),
|
||||
Err(Error::InvalidCursor("trace"))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_base64("garbage")]
|
||||
#[case::missing_fields("e30=")]
|
||||
#[case::not_an_object("WzEsMl0=")]
|
||||
fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) {
|
||||
assert!(matches!(
|
||||
error_position(Some(cursor)),
|
||||
Err(Error::InvalidCursor("diagnostic"))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::lowercase_version("a".repeat(64))]
|
||||
#[case::short_version("A".repeat(63))]
|
||||
fn diagnostic_cursor_requires_a_content_version(#[case] version: String) {
|
||||
let cursor = encode_cursor(&ErrorPosition { offset: 1, version });
|
||||
assert!(matches!(
|
||||
error_position(Some(&cursor)),
|
||||
Err(Error::InvalidCursor("diagnostic"))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
fn failed(error: StorageError) -> StoreError<Error> {
|
||||
StoreError::Failed(Error::Storage(error))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,8 +3,7 @@ use litellm_migrate::Migration;
|
|||
use serde::Serialize;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::Connection;
|
||||
use super::Error;
|
||||
use super::{Connection, Error};
|
||||
|
||||
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,20 @@
|
|||
use futures_util::{TryStreamExt, stream};
|
||||
use itertools::Itertools;
|
||||
//! Keyset-paged reads that shrink their page when ClickHouse rejects a response as too large and
|
||||
//! stop accumulating once a graph exceeds the interactive budget.
|
||||
|
||||
use std::{future::Future, marker::PhantomData};
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_storage_clickhouse::{Query, fetch};
|
||||
use litellm_traces::query::named as contracts;
|
||||
use serde::Serialize;
|
||||
use litellm_traces_cache::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, StoreError};
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::{Connection, Error, query::named::TraceSpansRow};
|
||||
use crate::{
|
||||
Connection, Error,
|
||||
query::named::{SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow},
|
||||
};
|
||||
|
||||
const PAGE_SIZE: u32 = 256;
|
||||
pub(crate) const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024;
|
||||
const MAX_GRAPH_SPANS: usize = 100_000;
|
||||
|
||||
#[derive(Default)]
|
||||
struct ReadBudget {
|
||||
|
|
@ -18,40 +23,137 @@ struct ReadBudget {
|
|||
}
|
||||
|
||||
impl ReadBudget {
|
||||
fn checked_add(&self, bytes: usize, rows: usize) -> Result<Self, Error> {
|
||||
let next = Self {
|
||||
bytes: self.bytes.saturating_add(bytes),
|
||||
rows: self.rows.saturating_add(rows),
|
||||
};
|
||||
if next.bytes > MAX_GRAPH_BYTES || next.rows > MAX_GRAPH_SPANS {
|
||||
return Err(Error::ReadTooLarge);
|
||||
fn reserve(&mut self, bytes: usize) -> Result<(), StoreError<Error>> {
|
||||
self.bytes = self.bytes.saturating_add(bytes);
|
||||
if self.bytes > MAX_GRAPH_BYTES || self.rows == MAX_GRAPH_SPANS {
|
||||
return Err(StoreError::TooLarge);
|
||||
}
|
||||
Ok(next)
|
||||
}
|
||||
|
||||
fn record(&mut self, row: &impl Serialize) -> Result<(), Error> {
|
||||
let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?;
|
||||
*self = self.checked_add(bytes.len(), 1)?;
|
||||
self.rows += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn record(&mut self, row: &impl Serialize) -> Result<(), StoreError<Error>> {
|
||||
let bytes =
|
||||
serde_json::to_vec(row).map_err(|_| StoreError::Failed(Error::InvalidResponse))?;
|
||||
self.reserve(bytes.len())
|
||||
}
|
||||
}
|
||||
|
||||
/// One keyset position in a paged query: the SQL reads the cursor fields of `Self` plus the
|
||||
/// `page_size` that [`Batch`] adds.
|
||||
trait Keyset: Serialize + Sized + Send + Sync {
|
||||
type Row: Serialize + DeserializeOwned + Send;
|
||||
const SQL: &'static str;
|
||||
|
||||
fn after(self, last: &Self::Row) -> Self;
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Parameters {
|
||||
struct Batch<K> {
|
||||
#[serde(flatten)]
|
||||
keyset: K,
|
||||
page_size: u32,
|
||||
}
|
||||
|
||||
trait PageSource<K: Keyset> {
|
||||
fn page(
|
||||
&self,
|
||||
batch: &Batch<K>,
|
||||
) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send;
|
||||
}
|
||||
|
||||
struct Paged<K>(PhantomData<K>);
|
||||
|
||||
impl<K: Keyset> Query for Paged<K> {
|
||||
type Params = Batch<K>;
|
||||
type Row = K::Row;
|
||||
const SQL: &'static str = K::SQL;
|
||||
}
|
||||
|
||||
/// Reads every row after `keyset`. A page ClickHouse rejects as too large is retried at half the
|
||||
/// size, and the smaller page is kept for the rest of the read because row sizes within one graph
|
||||
/// rarely shrink again. Halving a one-row page means a single row exceeds the response limit.
|
||||
async fn read_all<K: Keyset, S: PageSource<K>>(
|
||||
source: &S,
|
||||
keyset: K,
|
||||
) -> Result<Vec<K::Row>, StoreError<Error>> {
|
||||
let mut batch = Batch {
|
||||
keyset,
|
||||
page_size: PAGE_SIZE,
|
||||
};
|
||||
let mut rows = Vec::new();
|
||||
let mut budget = ReadBudget::default();
|
||||
loop {
|
||||
let page = match source.page(&batch).await {
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if batch.page_size > 1 => {
|
||||
batch.page_size /= 2;
|
||||
continue;
|
||||
}
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
|
||||
return Err(StoreError::TooLarge);
|
||||
}
|
||||
result => result.map_err(|error| StoreError::Failed(Error::Storage(error)))?,
|
||||
};
|
||||
let complete = page.len() < batch.page_size as usize;
|
||||
for row in &page {
|
||||
budget.record(row)?;
|
||||
}
|
||||
if let Some(last) = page.last() {
|
||||
batch.keyset = batch.keyset.after(last);
|
||||
}
|
||||
rows.extend(page);
|
||||
if complete {
|
||||
return Ok(rows);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ClickHouse<'a> {
|
||||
client: &'a Client,
|
||||
connection: &'a Connection,
|
||||
}
|
||||
|
||||
impl<K: Keyset> PageSource<K> for ClickHouse<'_> {
|
||||
fn page(
|
||||
&self,
|
||||
batch: &Batch<K>,
|
||||
) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send {
|
||||
fetch::<Paged<K>>(self.client, self.connection, batch)
|
||||
}
|
||||
}
|
||||
|
||||
async fn read_paged<K: Keyset>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
keyset: K,
|
||||
) -> Result<Vec<K::Row>, StoreError<Error>> {
|
||||
let source = ClickHouse { client, connection };
|
||||
read_all(&source, keyset).await
|
||||
}
|
||||
|
||||
fn by_start(mut rows: Vec<contracts::TraceSpansRow>) -> Vec<contracts::TraceSpansRow> {
|
||||
rows.sort_by_key(|row| row.start_ns);
|
||||
rows
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SpanKeyset {
|
||||
#[serde(flatten)]
|
||||
trace: contracts::TraceSpansParams,
|
||||
after_span_id: String,
|
||||
page_size: u32,
|
||||
snapshot_ms: u64,
|
||||
}
|
||||
|
||||
struct SpanBatch;
|
||||
|
||||
impl Query for SpanBatch {
|
||||
type Params = Parameters;
|
||||
impl Keyset for SpanKeyset {
|
||||
type Row = TraceSpansRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../query/trace_span_batch.sql");
|
||||
|
||||
fn after(self, last: &TraceSpansRow) -> Self {
|
||||
Self {
|
||||
after_span_id: last.0.span_id.clone(),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_spans(
|
||||
|
|
@ -59,199 +161,109 @@ pub(crate) async fn read_spans(
|
|||
connection: &Connection,
|
||||
trace: contracts::TraceSpansParams,
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, Error> {
|
||||
let mut parameters = Parameters {
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
|
||||
let keyset = SpanKeyset {
|
||||
trace,
|
||||
after_span_id: String::new(),
|
||||
page_size: PAGE_SIZE,
|
||||
snapshot_ms,
|
||||
};
|
||||
let mut spans = Vec::new();
|
||||
let mut budget = ReadBudget::default();
|
||||
loop {
|
||||
let page = match fetch::<SpanBatch>(client, connection, ¶meters).await {
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
|
||||
if parameters.page_size > 1 =>
|
||||
{
|
||||
parameters.page_size /= 2;
|
||||
continue;
|
||||
}
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
result => result?,
|
||||
};
|
||||
let complete = page.len() < parameters.page_size as usize;
|
||||
if let Some(last) = page.last() {
|
||||
parameters.after_span_id.clone_from(&last.0.span_id);
|
||||
}
|
||||
for row in page {
|
||||
budget.record(&row)?;
|
||||
spans.push(row.0);
|
||||
}
|
||||
if complete {
|
||||
spans.sort_by_key(|row| row.start_ns);
|
||||
return Ok(spans);
|
||||
}
|
||||
parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE);
|
||||
}
|
||||
let rows = read_paged(client, connection, keyset).await?;
|
||||
Ok(by_start(rows.into_iter().map(|row| row.0).collect()))
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ListParameters {
|
||||
struct ListSpanKeyset {
|
||||
#[serde(flatten)]
|
||||
runs: crate::query::named::TracePageSpansParams,
|
||||
after_team: String,
|
||||
after_key: String,
|
||||
after_trace: String,
|
||||
after_span: String,
|
||||
page_size: u32,
|
||||
snapshot_ms: u64,
|
||||
}
|
||||
|
||||
struct ListSpanBatch;
|
||||
|
||||
impl Query for ListSpanBatch {
|
||||
type Params = ListParameters;
|
||||
impl Keyset for ListSpanKeyset {
|
||||
type Row = TraceSpansRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../query/trace_list_span_batch.sql");
|
||||
|
||||
fn after(self, last: &TraceSpansRow) -> Self {
|
||||
Self {
|
||||
after_team: last.0.team_id.clone(),
|
||||
after_key: last.0.api_key_hash.clone(),
|
||||
after_trace: last.0.trace_id.clone(),
|
||||
after_span: last.0.span_id.clone(),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_list_spans(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
runs: crate::query::named::TracePageSpansParams,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, Error> {
|
||||
let parameters = ListParameters {
|
||||
snapshot_ms: u64,
|
||||
) -> Result<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
|
||||
let keyset = ListSpanKeyset {
|
||||
runs,
|
||||
after_team: String::new(),
|
||||
after_key: String::new(),
|
||||
after_trace: String::new(),
|
||||
after_span: String::new(),
|
||||
page_size: PAGE_SIZE,
|
||||
snapshot_ms: (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as u64,
|
||||
snapshot_ms,
|
||||
};
|
||||
let pages = stream::try_unfold(
|
||||
(Some(parameters), ReadBudget::default()),
|
||||
|(parameters, budget)| async move {
|
||||
let Some(parameters) = parameters else {
|
||||
return Ok(None);
|
||||
};
|
||||
let page = match fetch::<ListSpanBatch>(client, connection, ¶meters).await {
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
|
||||
if parameters.page_size > 1 =>
|
||||
{
|
||||
let retry = ListParameters {
|
||||
page_size: parameters.page_size / 2,
|
||||
..parameters
|
||||
};
|
||||
return Ok(Some((Vec::new(), (Some(retry), budget))));
|
||||
}
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
result => result?,
|
||||
};
|
||||
let next = page
|
||||
.last()
|
||||
.filter(|_| page.len() == parameters.page_size as usize)
|
||||
.map(|last| ListParameters {
|
||||
after_team: last.0.team_id.clone(),
|
||||
after_key: last.0.api_key_hash.clone(),
|
||||
after_trace: last.0.trace_id.clone(),
|
||||
after_span: last.0.span_id.clone(),
|
||||
page_size: (parameters.page_size * 2).min(PAGE_SIZE),
|
||||
..parameters
|
||||
});
|
||||
let next_budget = page.iter().try_fold(budget, |budget, row| {
|
||||
let bytes = serde_json::to_vec(row).map_err(|_| Error::InvalidResponse)?;
|
||||
budget.checked_add(bytes.len(), 1)
|
||||
})?;
|
||||
Ok(Some((page, (next, next_budget))))
|
||||
},
|
||||
)
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
Ok(pages
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.map(|row| row.0)
|
||||
.sorted_by_key(|row| row.start_ns)
|
||||
.collect())
|
||||
let rows = read_paged(client, connection, keyset).await?;
|
||||
Ok(by_start(rows.into_iter().map(|row| row.0).collect()))
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct SpendParameters {
|
||||
struct SpendKeyset {
|
||||
#[serde(flatten)]
|
||||
lookup: crate::query::named::SpendByResponseIdsParams,
|
||||
lookup: SpendByResponseIdsParams,
|
||||
has_cursor: u8,
|
||||
after_team: String,
|
||||
after_ms: i64,
|
||||
after_id: String,
|
||||
page_size: u32,
|
||||
}
|
||||
|
||||
struct SpendBatch;
|
||||
|
||||
impl Query for SpendBatch {
|
||||
type Params = SpendParameters;
|
||||
type Row = crate::query::named::SpendByResponseIdsRow;
|
||||
|
||||
impl Keyset for SpendKeyset {
|
||||
type Row = SpendByResponseIdsRow;
|
||||
const SQL: &'static str = include_str!("../query/spend_batch.sql");
|
||||
|
||||
fn after(self, last: &SpendByResponseIdsRow) -> Self {
|
||||
Self {
|
||||
has_cursor: 1,
|
||||
after_team: last.0.team_id.clone(),
|
||||
after_ms: last.0.start_ms,
|
||||
after_id: last.0.request_id.clone(),
|
||||
..self
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn read_spend(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
lookup: crate::query::named::SpendByResponseIdsParams,
|
||||
) -> Result<Vec<contracts::SpendByResponseIdsRow>, Error> {
|
||||
let mut parameters = SpendParameters {
|
||||
lookup: SpendByResponseIdsParams,
|
||||
) -> Result<Vec<contracts::SpendByResponseIdsRow>, StoreError<Error>> {
|
||||
let keyset = SpendKeyset {
|
||||
lookup,
|
||||
has_cursor: 0,
|
||||
after_team: String::new(),
|
||||
after_ms: 0,
|
||||
after_id: String::new(),
|
||||
page_size: PAGE_SIZE,
|
||||
};
|
||||
let mut rows = Vec::new();
|
||||
let mut budget = ReadBudget::default();
|
||||
loop {
|
||||
let page = match fetch::<SpendBatch>(client, connection, ¶meters).await {
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge)
|
||||
if parameters.page_size > 1 =>
|
||||
{
|
||||
parameters.page_size /= 2;
|
||||
continue;
|
||||
}
|
||||
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
|
||||
return Err(Error::ReadTooLarge);
|
||||
}
|
||||
result => result?,
|
||||
};
|
||||
let complete = page.len() < parameters.page_size as usize;
|
||||
if let Some(last) = page.last() {
|
||||
parameters.has_cursor = 1;
|
||||
parameters.after_team.clone_from(&last.0.team_id);
|
||||
parameters.after_ms = last.0.start_ms;
|
||||
parameters.after_id.clone_from(&last.0.request_id);
|
||||
}
|
||||
for row in page {
|
||||
budget.record(&row)?;
|
||||
rows.push(row.0);
|
||||
}
|
||||
if complete {
|
||||
return Ok(rows);
|
||||
}
|
||||
parameters.page_size = (parameters.page_size * 2).min(PAGE_SIZE);
|
||||
}
|
||||
let rows = read_paged(client, connection, keyset).await?;
|
||||
Ok(rows.into_iter().map(|row| row.0).collect())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::byte_boundary(MAX_GRAPH_BYTES - 1, 0, 1, false)]
|
||||
#[case::byte_overflow(MAX_GRAPH_BYTES - 1, 0, 2, true)]
|
||||
|
|
@ -264,7 +276,77 @@ mod tests {
|
|||
#[case] next: usize,
|
||||
#[case] rejected: bool,
|
||||
) {
|
||||
let budget = ReadBudget { bytes, rows };
|
||||
assert_eq!(budget.checked_add(next, 1).is_err(), rejected);
|
||||
let mut budget = ReadBudget { bytes, rows };
|
||||
assert_eq!(budget.reserve(next).is_err(), rejected);
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct Numbers {
|
||||
after: u32,
|
||||
}
|
||||
|
||||
impl Keyset for Numbers {
|
||||
type Row = u32;
|
||||
const SQL: &'static str = "";
|
||||
|
||||
fn after(self, last: &u32) -> Self {
|
||||
Self { after: *last }
|
||||
}
|
||||
}
|
||||
|
||||
/// A table of `total` rows whose transport rejects any page larger than `largest_page`.
|
||||
struct Table {
|
||||
total: u32,
|
||||
largest_page: u32,
|
||||
requests: Mutex<Vec<u32>>,
|
||||
}
|
||||
|
||||
impl PageSource<Numbers> for Table {
|
||||
async fn page(
|
||||
&self,
|
||||
batch: &Batch<Numbers>,
|
||||
) -> Result<Vec<u32>, litellm_storage_clickhouse::Error> {
|
||||
self.requests.lock().unwrap().push(batch.page_size);
|
||||
if batch.page_size > self.largest_page {
|
||||
return Err(litellm_storage_clickhouse::Error::ResponseTooLarge);
|
||||
}
|
||||
let end = (batch.keyset.after + batch.page_size).min(self.total);
|
||||
Ok((batch.keyset.after + 1..=end).collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])]
|
||||
#[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])]
|
||||
#[tokio::test]
|
||||
async fn a_rejected_page_size_is_not_retried(
|
||||
#[case] total: u32,
|
||||
#[case] largest_page: u32,
|
||||
#[case] requests: &[u32],
|
||||
) {
|
||||
let table = Table {
|
||||
total,
|
||||
largest_page,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
};
|
||||
let rows = read_all(&table, Numbers { after: 0 }).await.unwrap();
|
||||
assert_eq!(rows, (1..=total).collect::<Vec<_>>());
|
||||
assert_eq!(table.requests.lock().unwrap().as_slice(), requests);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_single_oversized_row_fails_the_read() {
|
||||
let table = Table {
|
||||
total: 10,
|
||||
largest_page: 0,
|
||||
requests: Mutex::new(Vec::new()),
|
||||
};
|
||||
let result = read_all(&table, Numbers { after: 0 }).await;
|
||||
assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}");
|
||||
assert_eq!(
|
||||
table.requests.lock().unwrap().as_slice(),
|
||||
&[256, 128, 64, 32, 16, 8, 4, 2, 1]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -210,8 +210,8 @@ fn request_id(evidence: &CallEvidence) -> &str {
|
|||
.into_iter()
|
||||
.flatten()
|
||||
.find_map(|key| match key {
|
||||
CallKey::LiteLlmRequest(id) | CallKey::ProviderResponse(id) => Some(id.as_str()),
|
||||
CallKey::Transport => None,
|
||||
CallKey::ProviderResponse(id) => Some(id.as_str()),
|
||||
CallKey::LiteLlmRequest(_) | CallKey::Transport | CallKey::GatewayAttempt => None,
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,8 +3,10 @@ use std::collections::BTreeMap;
|
|||
use litellm_http::Client;
|
||||
use litellm_traces::ReadQuery;
|
||||
|
||||
use super::query::{lens::*, named::*};
|
||||
use super::{Connection, Error, Parameter};
|
||||
use super::{
|
||||
Connection, Error, Parameter,
|
||||
query::{lens::*, named::*},
|
||||
};
|
||||
use litellm_storage_clickhouse::{Query, fetch_json};
|
||||
|
||||
pub async fn execute_named_read(
|
||||
|
|
|
|||
|
|
@ -184,7 +184,7 @@ Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time ha
|
|||
{%- endblock %}
|
||||
|
||||
{% block missing_spend -%}
|
||||
Token usage does not establish billed spend. OTLP exports without companion spend_logs rows have unknown cost; synthetic fixture spend is marked by metadata.synthetic_spend
|
||||
Token usage does not establish billed spend. OTLP exports without companion spend_logs rows have unknown cost
|
||||
{%- endblock %}
|
||||
|
||||
{% block partial_spend -%}
|
||||
|
|
|
|||
|
|
@ -4,13 +4,11 @@ Run `cargo test -p litellm-traces-clickhouse --test queries --locked -- --test-t
|
|||
|
||||
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 ClickHouse round-trip test replays the `google_adk_billed_failure`, `pydantic_ai_retry`, and `deepagents_swarm` exports and checks span identities, parent links, timestamps, durations, token counts, and statuses without pinning the provider's error wording. Exported ERROR and UNSET statuses are diagnostic and do not establish a failed execution, so the test preserves incoming statuses and checks root status separately from the count of error spans, deriving both from the decoded export. Framework-specific interpretation of control-flow exceptions belongs in the instrumentation integration
|
||||
|
||||
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
|
||||
For a local dashboard with linked requests and traces, run `make lens-dev ARGS=--seed` from the repository root and open `http://localhost:3000/ui/lens/`. Log in as `admin` with the master key saved in `.lens-dev/master_key`. The launcher keeps the stack running until Ctrl-C and leaves the database volumes intact
|
||||
|
||||
For a local dashboard with linked requests and traces, run `bash scripts/run_tracing_proxy_local.sh --seed` from the repository root and open `http://127.0.0.1:4002/ui/`. Log in as `admin` with password `sk-1234`, matching the UI E2E harness. The launcher keeps the proxy running until Ctrl-C and leaves the database volumes intact
|
||||
|
||||
`deeplite_swarm_spend_logs.jsonl` pairs every LLM span in the swarm export with a ClickHouse spend row. Response IDs, trace and span IDs, token counts, input, output, and timestamps come from the export. Messages and responses use the chat completion format supported by the request viewer. Spend is synthetic, set to $0.01 per request and marked in metadata, because the export does not include actual billed costs. These rows are stored here because `traces-clickhouse` owns the spend row schema
|
||||
Spend rows are stored here because `traces-clickhouse` owns the spend row schema
|
||||
|
||||
The simple and swarm exports for all twelve SDK examples were captured on 2026-10-03 against port 4002 using `openai/gpt-6-luna`. Each export has a matching `<name>_spend_logs.jsonl` with actual proxy spend, usage, request and response IDs, messages, and timestamps. Authorization headers, provider cookies, organization and project identifiers, and local paths were redacted. OTLP identifiers and enums use their canonical JSON encodings. `metadata.fixture_capture` identifies the associated export and whether model spans contain sufficient identity to join spend
|
||||
|
||||
|
|
@ -25,3 +23,5 @@ The LlamaIndex captures contain provider IDs inside `output.value.raw.id`. Regre
|
|||
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
|
||||
|
||||
`crates/traces/tests/captures.rs` resolves every capture against its spend rows without ClickHouse and checks the unrelated-transport and redundant-response-ID invariants
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1 @@
|
|||
{"request_id":"chatcmpl-EV1CtycizaTWvf3rjCEenOd7gXLRl","response_id":"chatcmpl-EV1CtycizaTWvf3rjCEenOd7gXLRl","litellm_call_id":"5034129d-560a-4a8b-85da-4f0f3628e70b","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":2.37e-05,"prompt_tokens":32,"completion_tokens":41,"total_tokens":73,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"8110744290a0840bc6029401723da262","trace_id":"8110744290a0840bc6029401723da262","span_id":"43984ccbf7e93ffd","request_tags":["User-Agent: AsyncOpenAI","User-Agent: AsyncOpenAI/Python 2.54.0"],"metadata":"{\"fixture_capture\":{\"name\":\"google_adk_billed_failure\",\"trace_id\":\"8110744290a0840bc6029401723da262\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"system\", \"content\": \"You are an agent. Your internal name is \\\"research_agent\\\".\"}, {\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1CtycizaTWvf3rjCEenOd7gXLRl\", \"created\": 1791061515, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces record an agent\\u2019s actions, decisions, and tool calls.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 41, \"prompt_tokens\": 32, \"total_tokens\": 73, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 18, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061515632,"end_time":1791061516529,"completion_start_time":1791061516529}
|
||||
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_retry_spend_logs.jsonl
vendored
Normal file
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_retry_spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
{"request_id":"chatcmpl-EV1DqRSm66p9hLCDdlrh9loQ2slG2","response_id":"chatcmpl-EV1DqRSm66p9hLCDdlrh9loQ2slG2","litellm_call_id":"4d1e094d-24e7-49f6-8d33-656a7d7712ae","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":3.7199999999999996e-05,"prompt_tokens":32,"completion_tokens":68,"total_tokens":100,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"cdd888fb2417b10d57e2b2cee4dd428c","trace_id":"cdd888fb2417b10d57e2b2cee4dd428c","span_id":"977a62497f2dda9d","request_tags":["User-Agent: AsyncOpenAI","User-Agent: AsyncOpenAI/Python 2.54.0"],"metadata":"{\"fixture_capture\":{\"name\":\"google_adk_retry\",\"trace_id\":\"cdd888fb2417b10d57e2b2cee4dd428c\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"system\", \"content\": \"You are an agent. Your internal name is \\\"research_agent\\\".\"}, {\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1DqRSm66p9hLCDdlrh9loQ2slG2\", \"created\": 1791061574, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces record an agent\\u2019s actions, tool calls, and outcomes.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 68, \"prompt_tokens\": 32, \"total_tokens\": 100, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 45, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061574518,"end_time":1791061576007,"completion_start_time":1791061576007}
|
||||
{"request_id":"chatcmpl-EV1Dsf8Gh5U6VZykhfciNU1lmsTgS","response_id":"chatcmpl-EV1Dsf8Gh5U6VZykhfciNU1lmsTgS","litellm_call_id":"38acfc26-d8f3-4adc-9528-0196c28b8640","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":2.67e-05,"prompt_tokens":32,"completion_tokens":47,"total_tokens":79,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"cdd888fb2417b10d57e2b2cee4dd428c","trace_id":"cdd888fb2417b10d57e2b2cee4dd428c","span_id":"2856caaa423beba1","request_tags":["User-Agent: AsyncOpenAI","User-Agent: AsyncOpenAI/Python 2.54.0"],"metadata":"{\"fixture_capture\":{\"name\":\"google_adk_retry\",\"trace_id\":\"cdd888fb2417b10d57e2b2cee4dd428c\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"system\", \"content\": \"You are an agent. Your internal name is \\\"research_agent\\\".\"}, {\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1Dsf8Gh5U6VZykhfciNU1lmsTgS\", \"created\": 1791061576, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces show the steps an agent takes to complete a task.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 47, \"prompt_tokens\": 32, \"total_tokens\": 79, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 25, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061576510,"end_time":1791061577760,"completion_start_time":1791061577760}
|
||||
File diff suppressed because one or more lines are too long
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_stream_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_stream_spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
{"request_id":"chatcmpl-EV1C3uQ69CnjfRcZy2MDlreaXkRVd","response_id":"chatcmpl-EV1C3uQ69CnjfRcZy2MDlreaXkRVd","litellm_call_id":"73d0b177-86e5-4e11-9704-bff2667a0a75","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":9.489999999999999e-05,"prompt_tokens":29,"completion_tokens":184,"total_tokens":213,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"59fb70ffc358dfe30a5e0d4615e4ca7f","trace_id":"59fb70ffc358dfe30a5e0d4615e4ca7f","span_id":"fc3bcda50363d82c","request_tags":["User-Agent: AsyncOpenAI","User-Agent: AsyncOpenAI/Python 2.54.0"],"metadata":"{\"fixture_capture\":{\"name\":\"google_adk_stream\",\"trace_id\":\"59fb70ffc358dfe30a5e0d4615e4ca7f\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"system\", \"content\": \"You are an agent. Your internal name is \\\"research_agent\\\".\"}, {\"role\": \"user\", \"content\": \"What is an agent trace?\"}]","response":"{\"id\": \"chatcmpl-EV1C3uQ69CnjfRcZy2MDlreaXkRVd\", \"created\": 1791061464, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"An **agent trace** is a record of an AI agent\\u2019s execution: what it received, the actions it took (such as tool calls), the results it got back, and how the task ended.\\n\\nTraces help developers debug failures, understand behavior, and evaluate performance. They may include timestamps, inputs and outputs, tool errors, or state changes. A trace doesn\\u2019t have to include the model\\u2019s private reasoning; it can record only observable steps.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null}}], \"usage\": {\"completion_tokens\": 184, \"prompt_tokens\": 29, \"total_tokens\": 213, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 84, \"rejected_prediction_tokens\": 0, \"text_tokens\": null, \"image_tokens\": null, \"video_tokens\": null}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"text_tokens\": null, \"image_tokens\": null, \"video_tokens\": null, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}, \"cost\": 9.489999999999999e-05}, \"service_tier\": \"default\"}","start_time":1791061463025,"end_time":1791061465613,"completion_start_time":1791061464712}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/mastra_simple_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/mastra_simple_spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
{"request_id":"chatcmpl-EV19kkSsaD6PMG5TGol2OVhYek5te","response_id":"chatcmpl-EV19kkSsaD6PMG5TGol2OVhYek5te","litellm_call_id":"afc4951d-7141-4476-b4d5-529b294e56b7","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":7.329999999999999e-05,"prompt_tokens":23,"completion_tokens":142,"total_tokens":165,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"6573bbd66c74e1ceb09fd173fe4cdc45","trace_id":"6573bbd66c74e1ceb09fd173fe4cdc45","span_id":"453659faa8fb4313","request_tags":["User-Agent: ai-sdk-openai-compatible","User-Agent: ai-sdk-openai-compatible/3.0.62 ai-sdk-provider-utils/5.0.53 node.js/25"],"metadata":"{\"fixture_capture\":{\"name\":\"mastra_simple\",\"trace_id\":\"6573bbd66c74e1ceb09fd173fe4cdc45\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"system\", \"content\": \"Answer the question concisely.\"}, {\"role\": \"user\", \"content\": \"What is an agent trace?\"}]","response":"{\"id\": \"chatcmpl-EV19kkSsaD6PMG5TGol2OVhYek5te\", \"created\": 1791061320, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"An **agent trace** is a record of an AI agent\\u2019s steps during a task\\u2014such as the inputs it received, actions or tool calls it made, results it got back, and its final response. Traces help people debug, evaluate, and audit agent behavior.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 142, \"prompt_tokens\": 23, \"total_tokens\": 165, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 78, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061319992,"end_time":1791061322400,"completion_start_time":1791061322400}
|
||||
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/mastra_swarm_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/mastra_swarm_spend_logs.jsonl
vendored
Normal file
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -0,0 +1 @@
|
|||
{"request_id":"chatcmpl-EV1DWkeu7A7n33K6iCaSxCDK1G71w","response_id":"chatcmpl-EV1DWkeu7A7n33K6iCaSxCDK1G71w","litellm_call_id":"27fa6242-6416-4749-9c1f-bc3ab3f83f35","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":2.45e-05,"prompt_tokens":15,"completion_tokens":46,"total_tokens":61,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"478c93b0465e5e11cc19e599738eadeb","trace_id":"478c93b0465e5e11cc19e599738eadeb","span_id":"877363da53522020","request_tags":["User-Agent: pydantic-ai","User-Agent: pydantic-ai/2.53.0"],"metadata":"{\"fixture_capture\":{\"name\":\"pydantic_ai_billed_failure\",\"trace_id\":\"478c93b0465e5e11cc19e599738eadeb\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1DWkeu7A7n33K6iCaSxCDK1G71w\", \"created\": 1791061554, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces show the steps an AI agent takes to complete a task.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 46, \"prompt_tokens\": 15, \"total_tokens\": 61, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 23, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061554409,"end_time":1791061555799,"completion_start_time":1791061555799}
|
||||
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_retry_spend_logs.jsonl
vendored
Normal file
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_retry_spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1,2 @@
|
|||
{"request_id":"chatcmpl-EV1EcoG4Hwa0tTYV7gLus0fwAJ7wY","response_id":"chatcmpl-EV1EcoG4Hwa0tTYV7gLus0fwAJ7wY","litellm_call_id":"bbff9fcc-930d-47fe-b0eb-0ab1eba0eff0","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":2.8e-05,"prompt_tokens":15,"completion_tokens":53,"total_tokens":68,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"213790ddf7f01594d52d54818775b98d","trace_id":"213790ddf7f01594d52d54818775b98d","span_id":"10cd8cbd5f72b220","request_tags":["User-Agent: pydantic-ai","User-Agent: pydantic-ai/2.53.0"],"metadata":"{\"fixture_capture\":{\"name\":\"pydantic_ai_retry\",\"trace_id\":\"213790ddf7f01594d52d54818775b98d\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1EcoG4Hwa0tTYV7gLus0fwAJ7wY\", \"created\": 1791061622, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces reveal the steps an agent takes to complete a task.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 53, \"prompt_tokens\": 15, \"total_tokens\": 68, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 31, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061622557,"end_time":1791061623931,"completion_start_time":1791061623931}
|
||||
{"request_id":"chatcmpl-EV1Ee6ZsIx2AYZLIIcwgTLCt8rYAp","response_id":"chatcmpl-EV1Ee6ZsIx2AYZLIIcwgTLCt8rYAp","litellm_call_id":"b1b0392c-9022-4369-88ce-7d8bdc57c90c","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":2.55e-05,"prompt_tokens":15,"completion_tokens":48,"total_tokens":63,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"213790ddf7f01594d52d54818775b98d","trace_id":"213790ddf7f01594d52d54818775b98d","span_id":"c80b675f9912b72c","request_tags":["User-Agent: pydantic-ai","User-Agent: pydantic-ai/2.53.0"],"metadata":"{\"fixture_capture\":{\"name\":\"pydantic_ai_retry\",\"trace_id\":\"213790ddf7f01594d52d54818775b98d\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"user\", \"content\": \"Reply with one short sentence about agent traces.\"}]","response":"{\"id\": \"chatcmpl-EV1Ee6ZsIx2AYZLIIcwgTLCt8rYAp\", \"created\": 1791061624, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"Agent traces record an agent\\u2019s actions and decisions over time.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null, \"provider_specific_fields\": {\"refusal\": null}, \"annotations\": []}, \"provider_specific_fields\": {}}], \"usage\": {\"completion_tokens\": 48, \"prompt_tokens\": 15, \"total_tokens\": 63, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 27, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}}, \"service_tier\": \"default\"}","start_time":1791061624398,"end_time":1791061625775,"completion_start_time":1791061625775}
|
||||
File diff suppressed because one or more lines are too long
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_stream_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_stream_spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1 @@
|
|||
{"request_id":"chatcmpl-EV1CVbDkWDEsratBuf9bhvEG3Av51","response_id":"chatcmpl-EV1CVbDkWDEsratBuf9bhvEG3Av51","litellm_call_id":"94605f85-b0da-4596-8538-9072fce7119e","call_type":"acompletion","api_key":"fixture-key","key_alias":"","team_id":"fixture-team","team_alias":"","organization_id":"","user":"fixture-user","end_user":"","model":"openai/gpt-6-luna","model_group":"openai/gpt-6-luna","model_id":"e68fbd1ce26aa7a89b059c9a6df12007e624c6a2106cd3e6af492ab7a259f117","custom_llm_provider":"openai","api_base":"https://api.openai.com/v1","spend":0.0001117,"prompt_tokens":12,"completion_tokens":221,"total_tokens":233,"cache_read_tokens":0,"cache_write_tokens":0,"status":"success","error_str":"","cache_hit":false,"session_id":"f3f8852c4f360b03ce2a36e6d3ed6f51","trace_id":"f3f8852c4f360b03ce2a36e6d3ed6f51","span_id":"a383ecf3f8beb7a6","request_tags":["User-Agent: pydantic-ai","User-Agent: pydantic-ai/2.53.0"],"metadata":"{\"fixture_capture\":{\"name\":\"pydantic_ai_stream\",\"trace_id\":\"f3f8852c4f360b03ce2a36e6d3ed6f51\",\"spend_linked\":true,\"spend_complete\":true}}","messages":"[{\"role\": \"user\", \"content\": \"What is an agent trace?\"}]","response":"{\"id\": \"chatcmpl-EV1CVbDkWDEsratBuf9bhvEG3Av51\", \"created\": 1791061493, \"model\": \"openai/gpt-6-luna\", \"object\": \"chat.completion\", \"system_fingerprint\": null, \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"message\": {\"content\": \"An **agent trace** is a chronological record of an AI agent\\u2019s activity while completing a task. It may include:\\n\\n- The task or input it received\\n- Actions it took, such as tool calls\\n- Results or observations it received\\n- Changes in state or decisions\\n- The final response or outcome\\n\\nFor example: *User asks for the weather \\u2192 agent calls a weather API \\u2192 API returns the forecast \\u2192 agent summarizes it.*\\n\\nA trace helps people debug or evaluate an agent. It doesn\\u2019t necessarily include the agent\\u2019s private reasoning; often it records only observable steps and tool interactions.\", \"role\": \"assistant\", \"tool_calls\": null, \"function_call\": null}}], \"usage\": {\"completion_tokens\": 221, \"prompt_tokens\": 12, \"total_tokens\": 233, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 93, \"rejected_prediction_tokens\": 0, \"text_tokens\": null, \"image_tokens\": null, \"video_tokens\": null}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0, \"text_tokens\": null, \"image_tokens\": null, \"video_tokens\": null, \"cache_write_tokens\": 0, \"cache_creation_tokens\": 0}, \"cost\": 0.0001117}, \"service_tier\": \"default\"}","start_time":1791061491444,"end_time":1791061494417,"completion_start_time":1791061493102}
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue