Merge remote-tracking branch 'origin/main' into litellm_otel_v2_team_capture_message_content
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> # Conflicts: # tests/unit/integrations/otel/test_otel_v2_destinations.py
BIN
.github/assets/roi-calculator-integrations/after-github.jpg
vendored
Normal file
|
After Width: | Height: | Size: 77 KiB |
BIN
.github/assets/roi-calculator-integrations/after-gitlab-detail-top.jpg
vendored
Normal file
|
After Width: | Height: | Size: 79 KiB |
BIN
.github/assets/roi-calculator-integrations/after-gitlab-detail.jpg
vendored
Normal file
|
After Width: | Height: | Size: 61 KiB |
BIN
.github/assets/roi-calculator-integrations/after-gitlab.jpg
vendored
Normal file
|
After Width: | Height: | Size: 86 KiB |
BIN
.github/assets/roi-calculator-integrations/before-github.jpg
vendored
Normal file
|
After Width: | Height: | Size: 75 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-exit-loading.jpg
vendored
Normal file
|
After Width: | Height: | Size: 40 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-fallback-live.jpg
vendored
Normal file
|
After Width: | Height: | Size: 88 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-overview.jpg
vendored
Normal file
|
After Width: | Height: | Size: 95 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-people.jpg
vendored
Normal file
|
After Width: | Height: | Size: 96 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-pr-costs.jpg
vendored
Normal file
|
After Width: | Height: | Size: 94 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-pr-detail.jpg
vendored
Normal file
|
After Width: | Height: | Size: 70 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-preview-link.jpg
vendored
Normal file
|
After Width: | Height: | Size: 82 KiB |
BIN
.github/assets/roi-calculator-integrations/demo-with-live-errors.jpg
vendored
Normal file
|
After Width: | Height: | Size: 82 KiB |
BIN
.github/assets/roi-calculator-integrations/source-race-after.jpg
vendored
Normal file
|
After Width: | Height: | Size: 64 KiB |
BIN
.github/assets/roi-calculator-integrations/source-race-before.jpg
vendored
Normal file
|
After Width: | Height: | Size: 90 KiB |
4
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -49,6 +49,10 @@ jobs:
|
|||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm ci
|
||||
|
||||
- name: Check UI production source types
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
run: npm run typecheck
|
||||
|
||||
- name: Run UI type tests (Vitest)
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
|
|
|
|||
12
.github/workflows/test-rust.yml
vendored
|
|
@ -5,6 +5,8 @@ on:
|
|||
paths:
|
||||
- "litellm-rust/**"
|
||||
- "litellm/rust_bridge/**"
|
||||
- "scripts/generate_trace_types.py"
|
||||
- "scripts/trace_codegen/**"
|
||||
- "tests/test_litellm_rust/**"
|
||||
- "litellm/integrations/custom_logger.py"
|
||||
- "litellm/litellm_core_utils/litellm_logging.py"
|
||||
|
|
@ -32,6 +34,8 @@ on:
|
|||
paths:
|
||||
- "litellm-rust/**"
|
||||
- "litellm/rust_bridge/**"
|
||||
- "scripts/generate_trace_types.py"
|
||||
- "scripts/trace_codegen/**"
|
||||
- "tests/test_litellm_rust/**"
|
||||
- "litellm/integrations/custom_logger.py"
|
||||
- "litellm/litellm_core_utils/litellm_logging.py"
|
||||
|
|
@ -85,7 +89,7 @@ jobs:
|
|||
cache-on-failure: true
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- run: cargo clippy --workspace --all-targets --locked -- -D warnings
|
||||
- run: cargo clippy --workspace --all-targets --locked --features litellm-traces/schema,litellm-traces-clickhouse/schema -- -D warnings
|
||||
|
||||
rust-test:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
@ -124,7 +128,11 @@ jobs:
|
|||
cache-on-failure: true
|
||||
save-if: ${{ github.ref == 'refs/heads/main' }}
|
||||
|
||||
- run: cargo nextest run --workspace --locked
|
||||
- name: Check generated trace contracts
|
||||
working-directory: .
|
||||
run: uv run scripts/generate_trace_types.py --check
|
||||
|
||||
- run: cargo nextest run --workspace --locked --features litellm-traces/schema,litellm-traces-clickhouse/schema
|
||||
|
||||
- run: cargo test --workspace --doc --locked
|
||||
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
|
|
@ -165,6 +165,7 @@ jobs:
|
|||
tests/unit/proxy/response_api_endpoints
|
||||
tests/unit/proxy/image_endpoints
|
||||
tests/unit/proxy/ocr_endpoints
|
||||
tests/unit/proxy/search_endpoints
|
||||
tests/unit/proxy/vector_store_endpoints
|
||||
tests/unit/proxy/agent_endpoints
|
||||
tests/unit/proxy/a2a
|
||||
|
|
|
|||
1
.gitignore
vendored
|
|
@ -58,6 +58,7 @@ litellm/proxy/tests/package-lock.json
|
|||
ui/litellm-dashboard/.next
|
||||
ui/litellm-dashboard/node_modules
|
||||
ui/litellm-dashboard/next-env.d.ts
|
||||
*.tsbuildinfo
|
||||
ui/litellm-dashboard/package.json
|
||||
ui/litellm-dashboard/package-lock.json
|
||||
helm/litellm-helm/*.tgz
|
||||
|
|
|
|||
|
|
@ -2,5 +2,6 @@ FROM python:3.12-slim
|
|||
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 litellm/proxy/lens/prompts/ /app/lens/prompts/
|
||||
USER 65532:65532
|
||||
CMD ["python", "-m", "lens.worker"]
|
||||
|
|
|
|||
|
|
@ -4,5 +4,8 @@
|
|||
!litellm/proxy/lens/
|
||||
!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/prompts/
|
||||
!litellm/proxy/lens/prompts/**
|
||||
|
|
|
|||
|
|
@ -19,11 +19,13 @@ The URL, database, and retention settings can also come from `CLICKHOUSE_URL`, `
|
|||
|
||||
Retention changes require a proxy restart. ClickHouse removes expired rows during background merges, not immediately at startup. Enable request/response logging to analyze LLM requests. Lens can only inspect content you actually retain
|
||||
|
||||
In Lens, click **Set up analysis**, choose an existing virtual key or **Create worker key**, then **Generate setup command**. The LiteLLM address is filled in for you; change it only if the server running Docker needs a different network address. Copy the command and run it on your server. The dialog changes to **Analyzer connected** when the container checks in
|
||||
In **Lens > Investigations**, click **Connect worker**, choose an analysis model and monthly limit, then **Get install command**. Use **Advanced options** to select an existing virtual key or change the proxy URL if the server running Docker needs a different network address. Copy the command and run it on your server. The dashboard shows **Worker connected** when the container checks in
|
||||
|
||||
The command already contains the compatible worker image and one worker token. The selected virtual key stays on the proxy; its secret is never sent to the worker. No source checkout, environment file, or second LiteLLM deployment is needed. Keep the command private because it includes the token. The LiteLLM release provides the dashboard and APIs; the container only runs background analysis
|
||||
|
||||
The dashboard and Compose file pin a verified worker image by digest. The image uses Linux amd64, and the generated command selects that platform. Worker image releases are independent of proxy releases: update the pinned image when changing their API contract. CI also publishes immutable commit tags for reproducible builds
|
||||
The dashboard and Compose file pin a verified worker image by digest. The image uses Linux amd64, and the generated command selects that platform. CI also publishes immutable `:sha-<commit>` tags for successful worker builds on `main`. Keep the worker image compatible with your gateway version
|
||||
|
||||
After upgrading the gateway, update the worker image and redeploy it while keeping its proxy URL and token. Existing containers do not update automatically. If an investigation reports a worker compatibility error, update the image before retrying
|
||||
|
||||
For deployments managed with Compose, download `compose.yaml` and provide `LITELLM_URL` and `LENS_WORKER_TOKEN` in an environment file. Its default image is already selected:
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
services:
|
||||
lens-worker:
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:67eba741c1b97c749975c5c38e2370a603e1105babc908d613c1b79d7b995393}
|
||||
image: ${LENS_WORKER_IMAGE:-ghcr.io/berriai/litellm-lens-worker@sha256:44f0597c7583dcfef999ece9a8bc02cfeb9f0f5167a1221cee3bd10b1b79271b}
|
||||
environment:
|
||||
LITELLM_URL: ${LITELLM_URL:?Set the URL reachable from this container}
|
||||
LENS_WORKER_TOKEN: ${LENS_WORKER_TOKEN:?Create a worker credential in the Lens UI}
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ services:
|
|||
target: runtime
|
||||
command: ["--config", "/app/tracing-config.yaml", "--port", "4000"]
|
||||
environment:
|
||||
LITELLM_MASTER_KEY: local-tracing-master-key
|
||||
LITELLM_MASTER_KEY: sk-1234
|
||||
LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY: "true"
|
||||
LITELLM_SALT_KEY: sk-local-tracing-salt-key
|
||||
DATABASE_URL: postgresql://litellm:litellm@db:5432/litellm
|
||||
STORE_MODEL_IN_DB: "True"
|
||||
|
|
|
|||
|
|
@ -7,15 +7,19 @@
|
|||
## This accepts a list of user id's for whom calls will be rejected
|
||||
|
||||
|
||||
from typing import Optional, Literal
|
||||
import litellm
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth, LiteLLM_EndUserTable
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from typing import Literal, Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import LiteLLM_EndUserTable, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import AUTH_OBJECTS_TARGET
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class _ENTERPRISE_BlockedUserList(CustomLogger):
|
||||
enforces_request_content: bool = True
|
||||
|
|
@ -54,6 +58,7 @@ class _ENTERPRISE_BlockedUserList(CustomLogger):
|
|||
if litellm.set_verbose is True:
|
||||
print(print_statement) # noqa
|
||||
|
||||
@with_service_target(AUTH_OBJECTS_TARGET)
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ Base class for sending emails to user after creating keys or invite links
|
|||
import html
|
||||
import json
|
||||
import os
|
||||
from typing import List, Literal, Optional
|
||||
from typing import Final, List, Literal, Optional
|
||||
|
||||
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
|
||||
EmailEvent,
|
||||
|
|
@ -15,6 +15,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
|
|||
SendKeyRotatedEmailEvent,
|
||||
)
|
||||
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import (
|
||||
|
|
@ -48,6 +49,8 @@ from litellm.proxy._types import (
|
|||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
|
||||
|
||||
_BUDGET_ALERT_CLAIMS_TARGET: Final = "budget_alert_claims"
|
||||
|
||||
|
||||
def _max_budget_alert_id(user_info: CallInfo) -> str:
|
||||
if user_info.event_group == Litellm_EntityType.TEAM_MEMBER:
|
||||
|
|
@ -437,6 +440,7 @@ class BaseEmailLogger(CustomLogger):
|
|||
html_body=email_html_content,
|
||||
)
|
||||
|
||||
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
|
||||
async def budget_alerts(
|
||||
self,
|
||||
type: Literal[
|
||||
|
|
@ -606,6 +610,7 @@ class BaseEmailLogger(CustomLogger):
|
|||
await self._release_budget_alert_claim(_cache, _cache_key)
|
||||
return
|
||||
|
||||
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
|
||||
async def _handle_multi_threshold_max_budget_alert(
|
||||
self,
|
||||
user_info: CallInfo,
|
||||
|
|
@ -691,6 +696,7 @@ class BaseEmailLogger(CustomLogger):
|
|||
)
|
||||
await self._release_budget_alert_claim(_cache, _cache_key)
|
||||
|
||||
@with_service_target(_BUDGET_ALERT_CLAIMS_TARGET)
|
||||
async def _release_budget_alert_claim(self, cache: DualCache, cache_key: str) -> None:
|
||||
try:
|
||||
await cache.async_delete_cache(key=cache_key)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.db.db_span import db_span
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
|
|
@ -94,16 +95,17 @@ async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
|
|||
json_settings = json.dumps(general_settings, default=str)
|
||||
|
||||
# Save updated general settings
|
||||
await prisma_client.db.litellm_config.upsert(
|
||||
where={"param_name": "general_settings"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "general_settings",
|
||||
"param_value": json_settings,
|
||||
async with db_span("save_email_settings", "LiteLLM_Config"):
|
||||
await prisma_client.db.litellm_config.upsert(
|
||||
where={"param_name": "general_settings"},
|
||||
data={
|
||||
"create": {
|
||||
"param_name": "general_settings",
|
||||
"param_value": json_settings,
|
||||
},
|
||||
"update": {"param_value": json_settings},
|
||||
},
|
||||
"update": {"param_value": json_settings},
|
||||
},
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -26,6 +26,7 @@ from pydantic import ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm import Router, verbose_logger
|
||||
from litellm._internal_context import with_service_target
|
||||
from litellm._uuid import uuid
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import MAX_FILE_LIST_LIMIT
|
||||
|
|
@ -144,6 +145,7 @@ def _parse_managed_file_object(raw_file_object: object, unified_file_id: str) ->
|
|||
class _ManagedFileRow(Protocol):
|
||||
unified_file_id: str
|
||||
file_object: OpenAIFileObject
|
||||
flat_model_file_ids: Sequence[str]
|
||||
storage_backend: Optional[str]
|
||||
storage_url: Optional[str]
|
||||
created_by: Optional[str]
|
||||
|
|
@ -201,6 +203,16 @@ def _managed_file_table(prisma_client: PrismaClient) -> _ManagedFileTableActions
|
|||
return prisma_client.db.litellm_managedfiletable
|
||||
|
||||
|
||||
def _iter_provider_file_id_pairs(
|
||||
rows: Sequence[_ManagedFileRow],
|
||||
requested_provider_file_ids: frozenset[str],
|
||||
) -> Iterator[tuple[str, str]]:
|
||||
for row in rows:
|
||||
for provider_file_id in row.flat_model_file_ids:
|
||||
if provider_file_id in requested_provider_file_ids:
|
||||
yield provider_file_id, row.unified_file_id
|
||||
|
||||
|
||||
def _managed_object_table(prisma_client: PrismaClient) -> _ManagedObjectTableActions:
|
||||
return prisma_client.db.litellm_managedobjecttable
|
||||
|
||||
|
|
@ -218,6 +230,9 @@ def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, s
|
|||
)
|
||||
|
||||
|
||||
_MANAGED_FILES_TARGET: Final = "managed_files"
|
||||
|
||||
|
||||
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
||||
# Class variables or attributes
|
||||
def __init__(self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient):
|
||||
|
|
@ -231,6 +246,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
|
||||
return PrometheusLogger.get_instance()
|
||||
|
||||
@with_service_target(_MANAGED_FILES_TARGET)
|
||||
async def store_unified_file_id(
|
||||
self,
|
||||
file_id: str,
|
||||
|
|
@ -314,6 +330,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
verbose_logger.warning(f"could not resolve org for managed object attribution: {e}")
|
||||
return None
|
||||
|
||||
@with_service_target(_MANAGED_FILES_TARGET)
|
||||
async def store_unified_object_id(
|
||||
self,
|
||||
unified_object_id: str,
|
||||
|
|
@ -401,6 +418,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
},
|
||||
)
|
||||
|
||||
@with_service_target(_MANAGED_FILES_TARGET)
|
||||
async def get_unified_file_id(
|
||||
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
|
||||
) -> Optional[LiteLLM_ManagedFileTable]:
|
||||
|
|
@ -423,6 +441,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return LiteLLM_ManagedFileTable.model_validate(db_object.model_dump())
|
||||
return None
|
||||
|
||||
@with_service_target(_MANAGED_FILES_TARGET)
|
||||
async def delete_unified_file_id(
|
||||
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
|
||||
) -> OpenAIFileObject:
|
||||
|
|
@ -710,6 +729,39 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
return None
|
||||
return batch_obj
|
||||
|
||||
async def get_unified_file_ids_for_provider_file_ids(
|
||||
self,
|
||||
provider_file_ids: Sequence[str],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> Mapping[str, str]:
|
||||
if not provider_file_ids:
|
||||
return MappingProxyType({})
|
||||
|
||||
unique_provider_file_ids: Final = tuple(dict.fromkeys(provider_file_ids))
|
||||
owner_filter: Final = build_owner_filter(user_api_key_dict)
|
||||
if owner_filter is None:
|
||||
return MappingProxyType({})
|
||||
|
||||
provider_file_ids_list: Final = [ # mutable-ok: Prisma hasSome requires a list
|
||||
provider_file_id for provider_file_id in unique_provider_file_ids
|
||||
]
|
||||
rows: Final = await _managed_file_table(self.prisma_client).find_many(
|
||||
where={ # mutable-ok: Prisma requires a plain dictionary for where
|
||||
**owner_filter,
|
||||
"flat_model_file_ids": { # mutable-ok: Prisma requires a plain filter dictionary
|
||||
"hasSome": provider_file_ids_list,
|
||||
},
|
||||
}
|
||||
)
|
||||
return MappingProxyType(
|
||||
dict(
|
||||
_iter_provider_file_id_pairs(
|
||||
rows,
|
||||
frozenset(unique_provider_file_ids),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
async def get_user_created_file_ids(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str]
|
||||
) -> List[OpenAIFileObject]:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,16 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundInteractionSettlement" (
|
||||
"interaction_id" TEXT NOT NULL,
|
||||
"custom_llm_provider" TEXT NOT NULL,
|
||||
"create_context" JSONB NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"claimed_at" TIMESTAMP(3),
|
||||
"claimed_by" TEXT,
|
||||
"settled_at" TIMESTAMP(3),
|
||||
"outcome" TEXT,
|
||||
|
||||
CONSTRAINT "LiteLLM_BackgroundInteractionSettlement_pkey" PRIMARY KEY ("interaction_id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "idx_background_interaction_settlement_claimed_at" ON "LiteLLM_BackgroundInteractionSettlement"("claimed_at");
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
-- CreateIndex (CONCURRENTLY)
|
||||
--
|
||||
-- Disclaimer:
|
||||
-- - CREATE INDEX CONCURRENTLY cannot run inside a transaction. This migration must stay a
|
||||
-- single statement so Prisma Migrate on PostgreSQL can apply it outside a transaction.
|
||||
-- - Builds are slower and use more I/O than a blocking CREATE INDEX; if the build is
|
||||
-- interrupted, Postgres may leave an INVALID index that must be dropped and recreated.
|
||||
-- - Do not edit this file after it has been applied to any database: Prisma checksums
|
||||
-- migrations; add a new migration instead.
|
||||
-- - Requires PostgreSQL that supports CONCURRENTLY with IF NOT EXISTS (use a new migration
|
||||
-- without IF NOT EXISTS if you must support older versions).
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS "LiteLLM_ManagedFileTable_flat_model_file_ids_idx" ON "LiteLLM_ManagedFileTable" USING GIN ("flat_model_file_ids");
|
||||
|
|
@ -1144,6 +1144,7 @@ model LiteLLM_ManagedFileTable {
|
|||
updated_by String?
|
||||
|
||||
@@index([unified_file_id])
|
||||
@@index([flat_model_file_ids], type: Gin)
|
||||
@@index([team_id, created_at(sort: Desc)])
|
||||
}
|
||||
|
||||
|
|
@ -1916,6 +1917,22 @@ model LiteLLM_WorkflowMessage {
|
|||
@@index([run_id])
|
||||
}
|
||||
|
||||
// Pending billing settlements for background interactions, keyed by the
|
||||
// interaction id so any replica can settle one that another replica created.
|
||||
// `claimed_at` is the exactly-once gate: the first conditional update wins.
|
||||
model LiteLLM_BackgroundInteractionSettlement {
|
||||
interaction_id String @id
|
||||
custom_llm_provider String
|
||||
create_context Json
|
||||
created_at DateTime @default(now())
|
||||
claimed_at DateTime?
|
||||
claimed_by String?
|
||||
settled_at DateTime?
|
||||
outcome String?
|
||||
|
||||
@@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at")
|
||||
}
|
||||
|
||||
model LiteLLM_Lens {
|
||||
id String @id
|
||||
version Int @default(0)
|
||||
|
|
|
|||
|
|
@ -353,9 +353,8 @@ class ProxyExtrasDBManager:
|
|||
pass
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str) -> Optional[str]:
|
||||
"""Return failed migration logs, or None if the ledger is unavailable."""
|
||||
database_url = os.getenv("DATABASE_URL")
|
||||
def _read_migration_ledger(query: str, params: tuple[str, ...]) -> "tuple[object, ...] | None":
|
||||
database_url: Final = os.getenv("DATABASE_URL")
|
||||
if not database_url:
|
||||
return None
|
||||
|
||||
|
|
@ -364,28 +363,37 @@ class ProxyExtrasDBManager:
|
|||
except ImportError:
|
||||
return None
|
||||
|
||||
cleaned_url = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
ledger_table = psycopg.sql.SQL("{}.{}").format(
|
||||
psycopg.sql.Identifier(
|
||||
ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"
|
||||
),
|
||||
cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url)
|
||||
ledger_table: Final = psycopg.sql.SQL("{}.{}").format(
|
||||
psycopg.sql.Identifier(ProxyExtrasDBManager._prisma_schema_param(database_url) or "public"),
|
||||
psycopg.sql.Identifier("_prisma_migrations"),
|
||||
)
|
||||
try:
|
||||
with psycopg.connect(
|
||||
cleaned_url, connect_timeout=10, autocommit=True
|
||||
) as conn:
|
||||
row = conn.execute(
|
||||
psycopg.sql.SQL(
|
||||
"SELECT logs FROM {} "
|
||||
"WHERE migration_name = %s AND finished_at IS NULL "
|
||||
"AND rolled_back_at IS NULL"
|
||||
).format(ledger_table),
|
||||
(migration_name,),
|
||||
).fetchone()
|
||||
with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn:
|
||||
row: Final = conn.execute(psycopg.sql.SQL(query).format(ledger_table), params).fetchone()
|
||||
except (psycopg.OperationalError, psycopg.DatabaseError):
|
||||
return None
|
||||
return (row[0] or "") if row else ""
|
||||
return tuple(row) if row is not None else ()
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_logs(migration_name: str, started_at: str) -> Optional[str]:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT logs FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND finished_at IS NULL AND rolled_back_at IS NULL",
|
||||
(migration_name, started_at),
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
return row[0] if row and isinstance(row[0], str) else ""
|
||||
|
||||
@staticmethod
|
||||
def _failed_migration_recovered(migration_name: str, started_at: str) -> bool:
|
||||
row: Final = ProxyExtrasDBManager._read_migration_ledger(
|
||||
"SELECT 1 FROM {} WHERE migration_name = %s AND started_at = %s::timestamptz "
|
||||
"AND (finished_at IS NOT NULL OR rolled_back_at IS NOT NULL)",
|
||||
(migration_name, started_at),
|
||||
)
|
||||
return bool(row)
|
||||
|
||||
@staticmethod
|
||||
def _resolve_specific_migration(migration_name: str):
|
||||
|
|
@ -1102,6 +1110,11 @@ class ProxyExtrasDBManager:
|
|||
return match.group(1) if match else None
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _v2_failed_migration_started_at(stderr: str, migration_name: str) -> "str | None":
|
||||
match: Final = re.search(rf"`{re.escape(migration_name)}` migration started at ([^\r\n]+?) failed", stderr)
|
||||
return match.group(1) if match else None
|
||||
|
||||
@staticmethod
|
||||
def _v2_roll_back_migration_best_effort(migration_name: str) -> None:
|
||||
from litellm_proxy_extras.migration_lock import migration_environment
|
||||
|
|
@ -1130,8 +1143,11 @@ class ProxyExtrasDBManager:
|
|||
|
||||
if "P3009" in stderr:
|
||||
migration_name = ProxyExtrasDBManager._v2_failed_migration_name(stderr)
|
||||
if migration_name:
|
||||
ledger_logs = ProxyExtrasDBManager._failed_migration_logs(migration_name)
|
||||
started_at: Final = (
|
||||
ProxyExtrasDBManager._v2_failed_migration_started_at(stderr, migration_name) if migration_name else None
|
||||
)
|
||||
if migration_name and started_at:
|
||||
ledger_logs: Final = ProxyExtrasDBManager._failed_migration_logs(migration_name, started_at)
|
||||
if ledger_logs and _MIGRATION_DEADLOCK_MARKER in ledger_logs:
|
||||
logger.info(
|
||||
"Migration %s failed in a concurrent migrate deploy "
|
||||
|
|
@ -1140,6 +1156,14 @@ class ProxyExtrasDBManager:
|
|||
)
|
||||
ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name)
|
||||
return budget.spend()
|
||||
if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at):
|
||||
logger.info(
|
||||
"Migration %s started at %s was already rolled back or completed by a concurrent "
|
||||
"migrate deploy, retrying",
|
||||
migration_name,
|
||||
started_at,
|
||||
)
|
||||
return budget.spend()
|
||||
raise RuntimeError(
|
||||
"Migration completion could not be verified. LiteLLM startup has stopped.\n\n"
|
||||
f"Prisma migration history (migration name and start time):\n{stderr}\n\n"
|
||||
|
|
|
|||
10
litellm-rust/Cargo.lock
generated
|
|
@ -4453,15 +4453,20 @@ dependencies = [
|
|||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-llms-types",
|
||||
"macro_rules_attribute",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4469,15 +4474,19 @@ name = "litellm-traces-clickhouse"
|
|||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"hmac 0.12.1",
|
||||
"jsonschema",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
"litellm-storage-clickhouse",
|
||||
"litellm-traces",
|
||||
"macro_rules_attribute",
|
||||
"moka",
|
||||
"rstest",
|
||||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2 0.10.9",
|
||||
|
|
@ -4486,6 +4495,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -467,6 +467,10 @@ pub struct ModelInfo {
|
|||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub supports_audio_output: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub supports_bedrock_runtime_chat_completions_response_format: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub supports_bedrock_runtime_chat_completions_tools_with_reasoning: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub supports_computer_use: Option<bool>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub supports_embedding_image_input: Option<bool>,
|
||||
|
|
|
|||
|
|
@ -45,8 +45,7 @@ mod _native {
|
|||
use crate::routes::token_counter::TokenCounter;
|
||||
#[pymodule_export]
|
||||
use crate::routes::traces::{
|
||||
NativeTraceConfig, NativeTraceStorage, trace_decode_otlp, trace_encode_error,
|
||||
trace_normalized_field_definitions,
|
||||
NativeTraceConfig, NativeTraceStorage, trace_encode_error, trace_span_rows,
|
||||
};
|
||||
#[cfg(feature = "huggingface")]
|
||||
#[pymodule_export]
|
||||
|
|
@ -114,9 +113,8 @@ mod tests {
|
|||
"NativeDiagnosticProcessor",
|
||||
"NativeTraceConfig",
|
||||
"NativeTraceStorage",
|
||||
"trace_decode_otlp",
|
||||
"trace_encode_error",
|
||||
"trace_normalized_field_definitions",
|
||||
"trace_span_rows",
|
||||
"TokenCounter",
|
||||
"Tokenizer",
|
||||
"gil_stats",
|
||||
|
|
|
|||
|
|
@ -1,14 +1,13 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_host_python::{FromPythonCache, ToPythonCache};
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_traces::{QueryScope, ReadQuery, Shared};
|
||||
use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessParams};
|
||||
use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders};
|
||||
use prost::Message;
|
||||
use pyo3::{
|
||||
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
|
||||
prelude::*,
|
||||
types::{PyBytes, PyDict, PyList, PyMapping, PyString},
|
||||
types::PyBytes,
|
||||
};
|
||||
|
||||
#[derive(Message)]
|
||||
|
|
@ -36,14 +35,20 @@ fn map_error_ref(error: &Error) -> PyErr {
|
|||
use litellm_storage_clickhouse::Error as StorageError;
|
||||
|
||||
match error {
|
||||
Error::Decode(litellm_traces::Error::TooLarge) | Error::InsertTooLarge => {
|
||||
PyOverflowError::new_err(error.to_string())
|
||||
}
|
||||
Error::InvalidRow
|
||||
| Error::InvalidTable
|
||||
| Error::InvalidCursor(_)
|
||||
| Error::AmbiguousTrace
|
||||
| Error::Decode(_)
|
||||
| Error::InvalidSchema
|
||||
| Error::InvalidQuery
|
||||
| Error::InvalidParameters
|
||||
| Error::InvalidScope => PyValueError::new_err(error.to_string()),
|
||||
Error::InsertTooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
Error::SchemaFailed(_)
|
||||
Error::Task
|
||||
| Error::SchemaFailed(_)
|
||||
| Error::SchemaTransport
|
||||
| Error::MissingSecret
|
||||
| Error::Busy
|
||||
|
|
@ -87,9 +92,15 @@ pub struct NativeTraceConfig {
|
|||
#[pymethods]
|
||||
impl NativeTraceConfig {
|
||||
#[new]
|
||||
fn new(database: String, url: &str, retention_days: u32) -> PyResult<Self> {
|
||||
fn new(
|
||||
database: String,
|
||||
url: &str,
|
||||
retention_days: u32,
|
||||
max_attribute_value_bytes: usize,
|
||||
) -> PyResult<Self> {
|
||||
Ok(Self {
|
||||
inner: Config::new(database, url, retention_days).map_err(map_error)?,
|
||||
inner: Config::new(database, url, retention_days, max_attribute_value_bytes)
|
||||
.map_err(map_error)?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -137,7 +148,9 @@ impl NativeTraceStorage {
|
|||
&self,
|
||||
py: Python<'py>,
|
||||
table: &str,
|
||||
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces_clickhouse::InsertRow>,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] rows: Vec<
|
||||
BTreeMap<String, serde_json::Value>,
|
||||
>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let table = InsertTable::parse(table).map_err(map_error)?;
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
|
|
@ -146,13 +159,156 @@ impl NativeTraceStorage {
|
|||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::insert_rows(&client, &connection, &database, table, rows)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn ingest<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
payload: &[u8],
|
||||
content_type: Option<String>,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] tenant: Tenant,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let payload = payload.to_vec();
|
||||
let max_value_bytes = self.config.max_attribute_value_bytes();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().writer().clone();
|
||||
let database = self.config.storage().database().to_owned();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
let rows = tokio::task::spawn_blocking(move || {
|
||||
litellm_traces::decode_otlp(&payload, content_type.as_deref()).map(|spans| {
|
||||
litellm_traces_clickhouse::span_rows(spans, &tenant, max_value_bytes)
|
||||
})
|
||||
})
|
||||
.await
|
||||
.map_err(|_| Error::Task)??;
|
||||
let count = rows.len();
|
||||
litellm_traces_clickhouse::insert_shared_rows(
|
||||
&client,
|
||||
&connection,
|
||||
&database,
|
||||
table,
|
||||
InsertTable::OtelTraces,
|
||||
rows,
|
||||
)
|
||||
.await?;
|
||||
Ok(count)
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (scope, start_ms, end_ms, cursor, limit))]
|
||||
fn list_traces<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams,
|
||||
start_ms: i64,
|
||||
end_ms: i64,
|
||||
cursor: Option<String>,
|
||||
limit: u32,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::list_traces(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
start_ms,
|
||||
end_ms,
|
||||
cursor.as_deref(),
|
||||
limit,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn get_trace<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
trace_id: String,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams,
|
||||
trace_ref: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::get_trace(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&trace_ref,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn get_span<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
trace_id: String,
|
||||
span_id: String,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams,
|
||||
trace_ref: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces_clickhouse::get_span(
|
||||
&client,
|
||||
&connection,
|
||||
&scope,
|
||||
&trace_id,
|
||||
&span_id,
|
||||
&trace_ref,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
}
|
||||
|
||||
#[pyo3(signature = (trace_id, span_id, scope, trace_ref, cursor))]
|
||||
fn get_span_error<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
trace_id: String,
|
||||
span_id: String,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams,
|
||||
trace_ref: String,
|
||||
cursor: Option<String>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
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
|
||||
},
|
||||
map_error,
|
||||
|
|
@ -232,101 +388,23 @@ impl NativeTraceStorage {
|
|||
}
|
||||
}
|
||||
|
||||
/// The `otel_traces` rows an export would be stored as, without writing them.
|
||||
#[pyfunction]
|
||||
pub fn trace_decode_otlp<'py>(
|
||||
pub fn trace_span_rows<'py>(
|
||||
py: Python<'py>,
|
||||
body: &[u8],
|
||||
content_type: Option<&str>,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] tenant: Tenant,
|
||||
max_attribute_value_bytes: usize,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let spans = py
|
||||
.detach(|| litellm_traces::decode_otlp(body, content_type))
|
||||
.map_err(|error| match error {
|
||||
litellm_traces::Error::TooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
_ => PyValueError::new_err(error.to_string()),
|
||||
})?;
|
||||
spans_to_py(py, &spans).map(Bound::into_any)
|
||||
}
|
||||
|
||||
fn insert_rows_from_py(
|
||||
value: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Vec<litellm_traces_clickhouse::InsertRow>> {
|
||||
let mut resources = FromPythonCache::default();
|
||||
value
|
||||
.try_iter()?
|
||||
.map(|row| {
|
||||
let row = row?;
|
||||
let mut fields = BTreeMap::new();
|
||||
for item in row.cast::<PyMapping>()?.items()?.iter() {
|
||||
let (key, value): (String, Bound<'_, PyAny>) = item.extract()?;
|
||||
let converted = if matches!(
|
||||
key.as_str(),
|
||||
"ResourceAttributes" | "ScopeName" | "ScopeVersion"
|
||||
) {
|
||||
resources
|
||||
.get_or_try_insert_with(&value, |value| {
|
||||
litellm_host_python::from_py_argument::<serde_json::Value>(value)
|
||||
.map(Shared::new)
|
||||
})?
|
||||
.clone()
|
||||
} else {
|
||||
Shared::new(litellm_host_python::from_py_argument(&value)?)
|
||||
};
|
||||
fields.insert(key, converted);
|
||||
}
|
||||
Ok(fields)
|
||||
let rows = py
|
||||
.detach(|| {
|
||||
litellm_traces::decode_otlp(body, content_type).map(|spans| {
|
||||
litellm_traces_clickhouse::span_rows(spans, &tenant, max_attribute_value_bytes)
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn spans_to_py<'py>(
|
||||
py: Python<'py>,
|
||||
spans: &[litellm_traces::DecodedSpan],
|
||||
) -> PyResult<Bound<'py, PyList>> {
|
||||
let mut resources = ToPythonCache::default();
|
||||
let mut scopes = ToPythonCache::default();
|
||||
let result = PyList::empty(py);
|
||||
for span in spans {
|
||||
let resource = resources
|
||||
.get_or_try_insert_with(span.resource_attributes.as_ref(), |value| {
|
||||
litellm_host_python::Pythonized(value).into_pyobject(py)
|
||||
})?;
|
||||
let row = PyDict::new(py);
|
||||
row.set_item("trace_id", &span.trace_id)?;
|
||||
row.set_item("span_id", &span.span_id)?;
|
||||
row.set_item("parent_span_id", &span.parent_span_id)?;
|
||||
row.set_item("trace_state", &span.trace_state)?;
|
||||
row.set_item("name", &span.name)?;
|
||||
row.set_item("kind", &span.kind)?;
|
||||
row.set_item("resource_attributes", resource)?;
|
||||
for (key, value) in [
|
||||
("scope_name", &span.scope_name),
|
||||
("scope_version", &span.scope_version),
|
||||
] {
|
||||
let value = scopes.get_or_try_insert_with(value.as_ref(), |value| {
|
||||
Ok(PyString::new(py, value).into_any())
|
||||
})?;
|
||||
row.set_item(key, value)?;
|
||||
}
|
||||
row.set_item("attributes", &span.attributes)?;
|
||||
row.set_item("start_ns", span.start_ns)?;
|
||||
row.set_item("end_ns", span.end_ns)?;
|
||||
row.set_item("status_code", &span.status_code)?;
|
||||
row.set_item("status_message", &span.status_message)?;
|
||||
row.set_item(
|
||||
"events",
|
||||
litellm_host_python::Pythonized(&span.events).into_pyobject(py)?,
|
||||
)?;
|
||||
row.set_item(
|
||||
"normalized",
|
||||
litellm_host_python::Pythonized(&span.normalized).into_pyobject(py)?,
|
||||
)?;
|
||||
row.set_item(
|
||||
"consumed_attributes",
|
||||
litellm_host_python::Pythonized(&span.consumed_attributes).into_pyobject(py)?,
|
||||
)?;
|
||||
result.append(row)?;
|
||||
}
|
||||
Ok(result)
|
||||
.map_err(|error| map_error(error.into()))?;
|
||||
litellm_host_python::Pythonized(rows).into_pyobject(py)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -383,50 +461,20 @@ mod tests {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
fn insert_projection_preserves_identity_without_merging_equal_resources() {
|
||||
#[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")]
|
||||
fn trace_read_and_ingest_failures_preserve_public_exception_types(
|
||||
#[case] error: Error,
|
||||
#[case] exception_name: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let resource = PyDict::new(py);
|
||||
resource.set_item("service.name", "shared").unwrap();
|
||||
let equal_resource = resource.copy().unwrap();
|
||||
let rows = PyList::empty(py);
|
||||
for value in [&resource, &resource, &equal_resource] {
|
||||
let row = PyDict::new(py);
|
||||
row.set_item("ResourceAttributes", value).unwrap();
|
||||
rows.append(row).unwrap();
|
||||
}
|
||||
let projected = insert_rows_from_py(rows.as_any()).unwrap();
|
||||
assert!(Shared::shares_storage_with(
|
||||
&projected[0]["ResourceAttributes"],
|
||||
&projected[1]["ResourceAttributes"]
|
||||
));
|
||||
assert!(!Shared::shares_storage_with(
|
||||
&projected[0]["ResourceAttributes"],
|
||||
&projected[2]["ResourceAttributes"]
|
||||
));
|
||||
assert_eq!(projected[0], projected[2]);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn shared_conversion_preserves_every_decoded_field() {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let spans = litellm_traces::decode_otlp(
|
||||
include_bytes!("../../../../../tests/test_litellm/tracing/fixtures/langsmith_deep_agent_export.json"),
|
||||
Some("application/json"),
|
||||
).unwrap();
|
||||
let expected = litellm_host_python::Pythonized(&spans)
|
||||
.into_pyobject(py)
|
||||
.unwrap();
|
||||
let actual = spans_to_py(py, &spans).unwrap();
|
||||
assert!(actual.eq(expected).unwrap());
|
||||
assert_eq!(
|
||||
map_error(error).get_type(py).name().unwrap(),
|
||||
exception_name
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
litellm_host_python::Pythonized(litellm_traces_clickhouse::NORMALIZED_FIELD_DEFINITIONS)
|
||||
.into_pyobject(py)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ mod read;
|
|||
|
||||
pub use error::Error;
|
||||
pub use insert::{insert_compressed_rows, insert_encoded_rows};
|
||||
pub use read::{Parameter, Query, execute_read, fetch, fetch_json};
|
||||
pub use read::{Parameter, Query, READ_LIMITS, ReadLimits, execute_read, fetch, fetch_json};
|
||||
use url::Url;
|
||||
|
||||
#[derive(Clone)]
|
||||
|
|
|
|||
|
|
@ -5,7 +5,18 @@ use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
|||
|
||||
use crate::{Connection, Error};
|
||||
|
||||
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub struct ReadLimits {
|
||||
pub result_rows: u64,
|
||||
pub response_bytes: usize,
|
||||
pub execution_seconds: u64,
|
||||
}
|
||||
|
||||
pub const READ_LIMITS: ReadLimits = ReadLimits {
|
||||
result_rows: 1000,
|
||||
response_bytes: 4 * 1024 * 1024,
|
||||
execution_seconds: 10,
|
||||
};
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[serde(untagged)]
|
||||
|
|
@ -78,9 +89,12 @@ pub async fn execute_read(
|
|||
.clear()
|
||||
.extend_pairs(existing_pairs)
|
||||
.append_pair("readonly", "1")
|
||||
.append_pair("max_result_rows", "1000")
|
||||
.append_pair("max_result_rows", &READ_LIMITS.result_rows.to_string())
|
||||
.append_pair("result_overflow_mode", "throw")
|
||||
.append_pair("max_execution_time", "10")
|
||||
.append_pair(
|
||||
"max_execution_time",
|
||||
&READ_LIMITS.execution_seconds.to_string(),
|
||||
)
|
||||
.append_pair("wait_end_of_query", "1")
|
||||
.append_pair("default_format", "JSON");
|
||||
|
||||
|
|
@ -101,7 +115,7 @@ pub async fn execute_read(
|
|||
|
||||
let mut body = Vec::new();
|
||||
while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? {
|
||||
if body.len() + chunk.len() > MAX_RESPONSE_BYTES {
|
||||
if body.len() + chunk.len() > READ_LIMITS.response_bytes {
|
||||
return Err(Error::ResponseTooLarge);
|
||||
}
|
||||
body.extend_from_slice(&chunk);
|
||||
|
|
|
|||
|
|
@ -5,8 +5,14 @@ edition.workspace = true
|
|||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[features]
|
||||
schema = ["dep:schemars", "litellm-traces/schema"]
|
||||
|
||||
[dependencies]
|
||||
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"
|
||||
|
|
@ -22,10 +28,17 @@ strum.workspace = true
|
|||
thiserror.workspace = true
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
jsonschema = { version = "0.55.1", default-features = false }
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
wiremock.workspace = true
|
||||
|
||||
[[bin]]
|
||||
name = "export-traces-clickhouse-schema"
|
||||
path = "src/bin/export_schema.rs"
|
||||
required-features = ["schema"]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,5 @@
|
|||
ALTER TABLE {database}.otel_traces
|
||||
ADD COLUMN IF NOT EXISTS WrapperCandidate Bool DEFAULT false AFTER ObservationType,
|
||||
ADD COLUMN IF NOT EXISTS CallKeys Array(String) DEFAULT [] AFTER LiteLLMRequestId,
|
||||
ADD COLUMN IF NOT EXISTS CallEvidence LowCardinality(String) DEFAULT '' AFTER CallKeys,
|
||||
ADD COLUMN IF NOT EXISTS ToolCallId String DEFAULT '' AFTER Output
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
ALTER TABLE {database}.otel_traces
|
||||
ADD COLUMN IF NOT EXISTS AgentMetadata String DEFAULT '{}' CODEC(ZSTD(3))
|
||||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE {database}.spend_logs MODIFY COLUMN spend Nullable(Float64) DEFAULT NULL
|
||||
|
|
@ -18,8 +18,7 @@ SELECT TraceId AS trace_id,
|
|||
FROM agent_traces_by_key
|
||||
WHERE ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserIds = [{user_id:String}])
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, TeamId))
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
|
|
|
|||
|
|
@ -9,8 +9,7 @@ LEFT JOIN (
|
|||
AND ObservationType = 'llm' AND Output != ''
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, TeamId))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
|
||||
GROUP BY TeamId, ApiKeyHash, ParentSpanId
|
||||
|
|
@ -19,8 +18,7 @@ LEFT JOIN (
|
|||
WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, o.TeamId)
|
||||
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, o.TeamId))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
|
||||
LIMIT 1
|
||||
|
|
|
|||
|
|
@ -6,8 +6,7 @@ FROM otel_traces
|
|||
WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, TeamId))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
|
||||
AND ({error_version:String} = '' OR hex(SHA256(StatusMessage)) = {error_version:String})
|
||||
|
|
|
|||
|
|
@ -1,11 +1,21 @@
|
|||
SELECT request_id, response_id, team_id, api_key, user, spend,
|
||||
SELECT request_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend,
|
||||
toUnixTimestamp64Milli(start_time) AS start_ms
|
||||
FROM spend_logs FINAL
|
||||
FROM (
|
||||
SELECT *,
|
||||
-- A chat request served through the Responses API returns the upstream `resp_` id to the
|
||||
-- client but logs LiteLLM's managed `resp_<base64>` id, which embeds it.
|
||||
if(startsWith(response_id, 'resp_'),
|
||||
extract(tryBase64Decode(substring(response_id, 6)), 'response_id:([^;]+)'),
|
||||
'') AS upstream_response_id
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND user = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, team_id))
|
||||
)
|
||||
WHERE response_id IN {response_ids:Array(String)}
|
||||
AND start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND user = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, team_id)
|
||||
OR ({api_key_hash:String} != '' AND api_key = {api_key_hash:String}))
|
||||
OR upstream_response_id IN {response_ids:Array(String)}
|
||||
OR request_id IN {request_ids:Array(String)}
|
||||
OR (trace_id != '' AND trace_id IN {trace_ids:Array(String)})
|
||||
ORDER BY start_time DESC
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ FROM otel_traces
|
|||
WHERE TraceId = {trace_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, TeamId))
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
LIMIT 2
|
||||
|
|
|
|||
|
|
@ -0,0 +1,24 @@
|
|||
SELECT o.TraceId AS trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
|
||||
o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent,
|
||||
o.Framework AS framework, o.StatusCode AS status,
|
||||
substringUTF8(o.StatusMessage, 1, 128) AS status_message,
|
||||
lengthUTF8(o.StatusMessage) > 128 AS error_truncated,
|
||||
toUnixTimestamp64Nano(o.Timestamp) AS start_ns, o.Duration AS duration_ns,
|
||||
o.ServiceName AS service, o.InputPreview AS input_preview, o.Model AS model,
|
||||
o.InputTokens AS input_tokens, o.OutputTokens AS output_tokens,
|
||||
o.LiteLLMRequestId AS litellm_request_id,
|
||||
o.CallKeys AS call_keys, o.CallEvidence AS call_evidence,
|
||||
-- Rows written before ToolCallId keep the call id only in their attributes.
|
||||
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
|
||||
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
|
||||
AS tool_call_id,
|
||||
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
|
||||
FROM otel_traces AS o
|
||||
WHERE o.Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND o.Timestamp < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, o.TeamId))
|
||||
AND hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) IN {trace_refs:Array(String)}
|
||||
ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage
|
||||
LIMIT 1 BY o.TeamId, o.ApiKeyHash, o.TraceId, o.SpanId
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
|
||||
o.ObservationType AS type, o.AgentName AS agent,
|
||||
SELECT o.TraceId AS trace_id, o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name,
|
||||
o.ObservationType AS type, toUInt8(o.WrapperCandidate) AS wrapper_candidate, o.AgentName AS agent,
|
||||
o.Framework AS framework, o.StatusCode AS status,
|
||||
substringUTF8(o.StatusMessage, 1, 128) AS status_message,
|
||||
lengthUTF8(o.StatusMessage) > 128 AS error_truncated,
|
||||
|
|
@ -7,13 +7,17 @@ SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name
|
|||
o.ServiceName AS service, o.InputPreview AS input_preview, o.Model AS model,
|
||||
o.InputTokens AS input_tokens, o.OutputTokens AS output_tokens,
|
||||
o.LiteLLMRequestId AS litellm_request_id,
|
||||
o.CallKeys AS call_keys, o.CallEvidence AS call_evidence,
|
||||
-- Rows written before ToolCallId keep the call id only in their attributes.
|
||||
if(o.ToolCallId != '' OR o.ObservationType != 'tool', o.ToolCallId,
|
||||
coalesce(nullIf(o.SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(o.SpanAttributes['tool.id'], ''), ''))
|
||||
AS tool_call_id,
|
||||
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
|
||||
FROM otel_traces AS o
|
||||
WHERE o.TraceId = {trace_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, o.TeamId)
|
||||
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
|
||||
OR has({team_ids:Array(String)}, o.TeamId))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
|
||||
ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage
|
||||
|
|
|
|||
|
|
@ -0,0 +1,6 @@
|
|||
fn main() {
|
||||
println!(
|
||||
"{}",
|
||||
serde_json::to_string_pretty(&litellm_traces_clickhouse::wire_schema::schemas()).unwrap()
|
||||
);
|
||||
}
|
||||
|
|
@ -5,14 +5,21 @@ use litellm_storage_clickhouse::Storage;
|
|||
pub struct Config {
|
||||
storage: Storage,
|
||||
retention_days: u32,
|
||||
max_attribute_value_bytes: usize,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn new(database: String, url: &str, retention_days: u32) -> Result<Self, Error> {
|
||||
pub fn new(
|
||||
database: String,
|
||||
url: &str,
|
||||
retention_days: u32,
|
||||
max_attribute_value_bytes: usize,
|
||||
) -> Result<Self, Error> {
|
||||
super::schema_statements(&database, retention_days)?;
|
||||
Ok(Self {
|
||||
storage: Storage::new(database, url)?,
|
||||
retention_days,
|
||||
max_attribute_value_bytes,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -23,4 +30,9 @@ impl Config {
|
|||
pub fn retention_days(&self) -> u32 {
|
||||
self.retention_days
|
||||
}
|
||||
|
||||
/// Stored span attribute and payload values longer than this are truncated with a marker.
|
||||
pub fn max_attribute_value_bytes(&self) -> usize {
|
||||
self.max_attribute_value_bytes
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,14 @@ 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(transparent)]
|
||||
Decode(#[from] litellm_traces::Error),
|
||||
#[error("trace ingestion task failed")]
|
||||
Task,
|
||||
#[error(transparent)]
|
||||
Storage(#[from] litellm_storage_clickhouse::Error),
|
||||
#[error(transparent)]
|
||||
|
|
|
|||
|
|
@ -1,21 +1,39 @@
|
|||
macro_rules_attribute::attribute_alias! {
|
||||
#[apply(wire_type)] =
|
||||
#[derive(serde::Serialize, serde::Deserialize)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))];
|
||||
#[apply(response_type)] =
|
||||
#[derive(serde::Serialize)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))];
|
||||
#[apply(request_type)] =
|
||||
#[derive(serde::Deserialize)]
|
||||
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))];
|
||||
}
|
||||
|
||||
mod config;
|
||||
mod error;
|
||||
mod insert;
|
||||
pub mod query;
|
||||
mod query_access;
|
||||
mod reads;
|
||||
mod schema;
|
||||
mod span_row;
|
||||
mod sql;
|
||||
mod table;
|
||||
#[cfg(feature = "schema")]
|
||||
pub mod wire_schema;
|
||||
|
||||
pub use config::Config;
|
||||
pub use error::Error;
|
||||
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
|
||||
pub use litellm_storage_clickhouse::{Connection, Parameter};
|
||||
pub use litellm_traces::{QueryScope, ReadQuery};
|
||||
pub use query::{execute_read, query_help, query_sql};
|
||||
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, list_traces};
|
||||
pub use schema::{
|
||||
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements,
|
||||
};
|
||||
pub use span_row::span_rows;
|
||||
pub use sql::execute_named_read;
|
||||
pub use table::TraceTable;
|
||||
|
|
|
|||
|
|
@ -6,11 +6,15 @@ use futures_util::{
|
|||
stream::{self, TryStreamExt},
|
||||
};
|
||||
use litellm_http::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use litellm_traces::query::guide::{Example, QueryGuide, Section};
|
||||
use serde::{Deserialize, Serialize, Serializer};
|
||||
use serde_json::Value;
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use super::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, Parameter};
|
||||
use super::{
|
||||
Connection, Error, NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, Parameter,
|
||||
query_access::READER_LIMITS,
|
||||
};
|
||||
|
||||
mod guide;
|
||||
pub mod lens;
|
||||
|
|
@ -36,26 +40,65 @@ struct MetadataRow {
|
|||
metadata: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[macro_rules_attribute::apply(request_type)]
|
||||
struct AttributeRow {
|
||||
key: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
|
||||
#[serde(untagged)]
|
||||
enum PathPart {
|
||||
Key(String),
|
||||
Index(usize),
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, strum::Display)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
#[strum(serialize_all = "lowercase")]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "MetadataValueType"))]
|
||||
enum JsonKind {
|
||||
Array,
|
||||
Boolean,
|
||||
Integer,
|
||||
Null,
|
||||
Number,
|
||||
Object,
|
||||
String,
|
||||
}
|
||||
|
||||
impl JsonKind {
|
||||
fn of(value: &Value) -> Self {
|
||||
match value {
|
||||
Value::Null => Self::Null,
|
||||
Value::Bool(_) => Self::Boolean,
|
||||
Value::Number(number) if number.is_i64() || number.is_u64() => Self::Integer,
|
||||
Value::Number(_) => Self::Number,
|
||||
Value::String(_) => Self::String,
|
||||
Value::Array(_) => Self::Array,
|
||||
Value::Object(_) => Self::Object,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[derive(Clone, Copy, Debug, strum::Display)]
|
||||
enum MapValueType {
|
||||
String,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryMetadataField"))]
|
||||
struct MetadataField {
|
||||
path: Vec<PathPart>,
|
||||
types: BTreeSet<&'static str>,
|
||||
types: BTreeSet<JsonKind>,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryColumn"))]
|
||||
struct ColumnSchema {
|
||||
name: String,
|
||||
#[serde(rename = "type")]
|
||||
|
|
@ -64,44 +107,200 @@ struct ColumnSchema {
|
|||
details: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryTable"))]
|
||||
struct TableSchema {
|
||||
name: &'static str,
|
||||
name: TraceTable,
|
||||
columns: Vec<ColumnSchema>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MetadataCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
trait Unobserved {
|
||||
fn unobserved() -> Self;
|
||||
}
|
||||
|
||||
enum Discovery<T> {
|
||||
Observed(T),
|
||||
Unavailable(String),
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
impl<T: schemars::JsonSchema> schemars::JsonSchema for Discovery<T> {
|
||||
fn schema_name() -> std::borrow::Cow<'static, str> {
|
||||
format!("Discovery{}", T::schema_name()).into()
|
||||
}
|
||||
|
||||
fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
|
||||
let mut schema = T::json_schema(generator);
|
||||
schema
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.get_mut("properties")
|
||||
.unwrap()
|
||||
.as_object_mut()
|
||||
.unwrap()
|
||||
.insert(
|
||||
"error".into(),
|
||||
serde_json::json!({"type": ["string", "null"], "default": null}),
|
||||
);
|
||||
schema
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
pub(crate) fn help_schema() -> schemars::Schema {
|
||||
schemars::generate::SchemaSettings::draft2020_12()
|
||||
.for_serialize()
|
||||
.with_transform(litellm_traces::schema::integer_bounds)
|
||||
.into_generator()
|
||||
.into_root_schema_for::<QueryHelp>()
|
||||
}
|
||||
|
||||
impl<T: Serialize + Unobserved> Serialize for Discovery<T> {
|
||||
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
|
||||
#[derive(Serialize)]
|
||||
struct Unavailable<'a, T> {
|
||||
#[serde(flatten)]
|
||||
sample: T,
|
||||
error: &'a str,
|
||||
}
|
||||
match self {
|
||||
Self::Observed(sample) => sample.serialize(serializer),
|
||||
Self::Unavailable(error) => Unavailable {
|
||||
sample: T::unobserved(),
|
||||
error,
|
||||
}
|
||||
.serialize(serializer),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
struct MetadataSample {
|
||||
fields: Vec<MetadataField>,
|
||||
sampled_rows: usize,
|
||||
invalid_json_rows: usize,
|
||||
truncated: bool,
|
||||
sample_sql: &'static str,
|
||||
scope: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
impl Unobserved for MetadataSample {
|
||||
fn unobserved() -> Self {
|
||||
Self {
|
||||
fields: Vec::new(),
|
||||
sampled_rows: 0,
|
||||
invalid_json_rows: 0,
|
||||
truncated: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryMetadata"))]
|
||||
struct MetadataCatalog {
|
||||
table: TraceTable,
|
||||
column: &'static str,
|
||||
#[serde(flatten)]
|
||||
discovery: Discovery<MetadataSample>,
|
||||
sample_sql: &'static str,
|
||||
scope: &'static str,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryAttributeField"))]
|
||||
struct AttributeField {
|
||||
key: String,
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
kind: MapValueType,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AttributeCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
struct AttributeSample {
|
||||
fields: Vec<AttributeField>,
|
||||
truncated: bool,
|
||||
}
|
||||
|
||||
impl Unobserved for AttributeSample {
|
||||
fn unobserved() -> Self {
|
||||
Self {
|
||||
fields: Vec::new(),
|
||||
truncated: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryAttributes"))]
|
||||
struct AttributeCatalog {
|
||||
table: TraceTable,
|
||||
column: &'static str,
|
||||
#[serde(flatten)]
|
||||
discovery: Discovery<AttributeSample>,
|
||||
discovery_sql: String,
|
||||
scope: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryNormalizedField"))]
|
||||
struct NormalizedField {
|
||||
table: TraceTable,
|
||||
name: &'static str,
|
||||
column: &'static str,
|
||||
#[serde(rename = "type")]
|
||||
kind: &'static str,
|
||||
meaning: &'static str,
|
||||
}
|
||||
|
||||
impl From<&NormalizedFieldDefinition> for NormalizedField {
|
||||
fn from(field: &NormalizedFieldDefinition) -> Self {
|
||||
Self {
|
||||
table: TraceTable::OtelTraces,
|
||||
name: field.name,
|
||||
column: field.clickhouse_column,
|
||||
kind: field.clickhouse_type,
|
||||
meaning: field.meaning,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryRelationship"))]
|
||||
struct Relationship {
|
||||
left: &'static str,
|
||||
right: &'static str,
|
||||
additional_predicates: &'static str,
|
||||
meaning: &'static str,
|
||||
}
|
||||
|
||||
const RELATIONSHIPS: [Relationship; 1] = [Relationship {
|
||||
left: "otel_traces.LiteLLMRequestId",
|
||||
right: "spend_logs.response_id",
|
||||
additional_predicates: "otel_traces.TeamId = spend_logs.team_id AND ((otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' AND otel_traces.ApiKeyHash = spend_logs.api_key))",
|
||||
meaning: "LiteLLMRequestId contains the first normalized request or provider response ID. This relationship matches response IDs only; CallKeys retains all typed identifiers. Cached requests can share response_id; joins may return multiple spend rows",
|
||||
}];
|
||||
|
||||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceQueryHelp"))]
|
||||
pub struct QueryHelp {
|
||||
dialect: &'static str,
|
||||
access: &'static str,
|
||||
response: &'static str,
|
||||
tables: Vec<TableSchema>,
|
||||
normalized_fields: Vec<NormalizedField>,
|
||||
metadata: MetadataCatalog,
|
||||
attributes: Vec<AttributeCatalog>,
|
||||
relationships: &'static [Relationship],
|
||||
#[cfg_attr(feature = "schema", schemars(with = "Vec<Example>"))]
|
||||
examples: [Example; 9],
|
||||
#[cfg_attr(feature = "schema", schemars(with = "Vec<String>"))]
|
||||
gotchas: [String; 13],
|
||||
guide: String,
|
||||
}
|
||||
|
||||
pub async fn execute_read(
|
||||
|
|
@ -153,22 +352,16 @@ fn metadata_expression(path: &[PathPart]) -> String {
|
|||
fn discover(
|
||||
value: &Value,
|
||||
path: Vec<PathPart>,
|
||||
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
|
||||
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<JsonKind>>,
|
||||
) -> bool {
|
||||
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
|
||||
return true;
|
||||
}
|
||||
if !path.is_empty() {
|
||||
let kind = match value {
|
||||
Value::Null => "null",
|
||||
Value::Bool(_) => "boolean",
|
||||
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
|
||||
Value::Number(_) => "number",
|
||||
Value::String(_) => "string",
|
||||
Value::Array(_) => "array",
|
||||
Value::Object(_) => "object",
|
||||
};
|
||||
fields.entry(path.clone()).or_default().insert(kind);
|
||||
fields
|
||||
.entry(path.clone())
|
||||
.or_default()
|
||||
.insert(JsonKind::of(value));
|
||||
}
|
||||
match value {
|
||||
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
|
||||
|
|
@ -194,7 +387,7 @@ fn discover(
|
|||
}
|
||||
}
|
||||
|
||||
fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
|
||||
fn metadata_sample(sample: &[MetadataRow]) -> MetadataSample {
|
||||
let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold(
|
||||
(BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0),
|
||||
|(fields, limited, invalid_rows), row| match serde_json::from_str::<Value>(&row.metadata) {
|
||||
|
|
@ -214,24 +407,19 @@ fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
|
|||
types,
|
||||
})
|
||||
.collect();
|
||||
MetadataCatalog {
|
||||
table: "spend_logs",
|
||||
column: "metadata",
|
||||
MetadataSample {
|
||||
fields,
|
||||
sampled_rows: sample.len().min(SAMPLE_ROWS),
|
||||
invalid_json_rows: invalid_rows,
|
||||
truncated: limited,
|
||||
sample_sql: METADATA_SQL,
|
||||
error: None,
|
||||
scope: METADATA_SCOPE,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn query_help(client: &Client, connection: &Connection) -> Result<String, Error> {
|
||||
pub async fn query_help(client: &Client, connection: &Connection) -> Result<QueryHelp, Error> {
|
||||
let tables = stream::iter(TraceTable::iter())
|
||||
.then(|table| async move {
|
||||
Ok::<_, Error>(TableSchema {
|
||||
name: table.into(),
|
||||
name: table,
|
||||
columns: rows::<ColumnSchema>(
|
||||
client,
|
||||
connection,
|
||||
|
|
@ -242,13 +430,15 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
|
|||
})
|
||||
.try_collect::<Vec<_>>()
|
||||
.await?;
|
||||
let metadata = match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
|
||||
Ok(sample) => metadata_catalog(&sample),
|
||||
Err(error) => MetadataCatalog {
|
||||
error: Some(error.to_string()),
|
||||
truncated: true,
|
||||
..metadata_catalog(&[])
|
||||
let metadata = MetadataCatalog {
|
||||
table: TraceTable::SpendLogs,
|
||||
column: "metadata",
|
||||
discovery: match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
|
||||
Ok(sample) => Discovery::Observed(metadata_sample(&sample)),
|
||||
Err(error) => Discovery::Unavailable(error.to_string()),
|
||||
},
|
||||
sample_sql: METADATA_SQL,
|
||||
scope: METADATA_SCOPE,
|
||||
};
|
||||
let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
|
||||
.then(|column| async move {
|
||||
|
|
@ -257,27 +447,27 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
|
|||
(SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \
|
||||
LIMIT 200) ORDER BY key LIMIT 201"
|
||||
);
|
||||
let (keys, error) = match rows::<AttributeRow>(client, connection, &sql).await {
|
||||
Ok(keys) => (keys, None),
|
||||
Err(error) => (Vec::new(), Some(error.to_string())),
|
||||
let discovery = match rows::<AttributeRow>(client, connection, &sql).await {
|
||||
Ok(keys) => Discovery::Observed(AttributeSample {
|
||||
truncated: keys.len() > MAX_FIELDS,
|
||||
fields: keys
|
||||
.into_iter()
|
||||
.take(MAX_FIELDS)
|
||||
.map(|row| AttributeField {
|
||||
expression: format!("{column}[{}]", literal(&row.key)),
|
||||
key: row.key,
|
||||
kind: MapValueType::String,
|
||||
})
|
||||
.collect(),
|
||||
}),
|
||||
Err(error) => Discovery::Unavailable(error.to_string()),
|
||||
};
|
||||
let fields = keys
|
||||
.iter()
|
||||
.take(MAX_FIELDS)
|
||||
.map(|row| AttributeField {
|
||||
key: row.key.clone(),
|
||||
kind: "String",
|
||||
expression: format!("{column}[{}]", literal(&row.key)),
|
||||
})
|
||||
.collect();
|
||||
AttributeCatalog {
|
||||
table: "otel_traces",
|
||||
table: TraceTable::OtelTraces,
|
||||
column,
|
||||
fields,
|
||||
truncated: error.is_some() || keys.len() > MAX_FIELDS,
|
||||
discovery,
|
||||
discovery_sql: sql,
|
||||
scope: ATTRIBUTE_SCOPE,
|
||||
error,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
|
|
@ -287,33 +477,78 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
|
|||
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
|
||||
metadata: &metadata,
|
||||
attributes: &attributes,
|
||||
limits: &READER_LIMITS,
|
||||
};
|
||||
Ok(json!({
|
||||
"dialect": "ClickHouse SQL",
|
||||
"access": "Request-log visibility enforced by ClickHouse row policies; proxy admins see all rows, users see their own rows and permitted teams, and callers without user identity see their own key rows",
|
||||
"response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
|
||||
"tables": tables,
|
||||
"normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({
|
||||
"table": "otel_traces", "name": field.name, "column": field.clickhouse_column,
|
||||
"type": field.clickhouse_type, "meaning": field.meaning
|
||||
})).collect::<Vec<_>>(),
|
||||
"metadata": metadata,
|
||||
"attributes": attributes,
|
||||
"relationships": [{
|
||||
"left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id",
|
||||
"additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND (otel_traces.TeamId != '' OR (otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' AND otel_traces.ApiKeyHash = spend_logs.api_key))",
|
||||
"meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows"
|
||||
}],
|
||||
"examples": guide.examples()?,
|
||||
"gotchas": guide.gotchas()?,
|
||||
"guide": guide::render(&guide)?,
|
||||
}).to_string())
|
||||
let bodies = guide.sections()?;
|
||||
let sections = [
|
||||
"Live ClickHouse schema",
|
||||
"Normalized span fields",
|
||||
"Observed LLM call metadata",
|
||||
"Observed span and resource attributes",
|
||||
]
|
||||
.into_iter()
|
||||
.zip(&bodies)
|
||||
.map(|(title, body)| Section { title, body })
|
||||
.collect::<Vec<_>>();
|
||||
let examples = guide.examples()?;
|
||||
let gotchas = guide.gotchas()?;
|
||||
let rendered = QueryGuide {
|
||||
sections: §ions,
|
||||
examples: &examples,
|
||||
gotchas: &gotchas,
|
||||
}
|
||||
.render()
|
||||
.map_err(|_| Error::InvalidResponse)?;
|
||||
Ok(QueryHelp {
|
||||
dialect: "ClickHouse SQL",
|
||||
access: "Request-log visibility enforced by ClickHouse row policies; proxy admins see all rows, users see their own rows and permitted teams",
|
||||
response: "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
|
||||
examples,
|
||||
gotchas,
|
||||
guide: rendered,
|
||||
normalized_fields: NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(NormalizedField::from)
|
||||
.collect(),
|
||||
relationships: &RELATIONSHIPS,
|
||||
tables,
|
||||
metadata,
|
||||
attributes,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
#[cfg(feature = "schema")]
|
||||
#[rstest]
|
||||
#[case::observed(false)]
|
||||
#[case::unavailable(true)]
|
||||
fn discovery_serialization_matches_its_schema(#[case] unavailable: bool) {
|
||||
let discovery = if unavailable {
|
||||
Discovery::Unavailable("discovery failed".into())
|
||||
} else {
|
||||
Discovery::Observed(MetadataSample::unobserved())
|
||||
};
|
||||
let catalog = MetadataCatalog {
|
||||
table: TraceTable::SpendLogs,
|
||||
column: "metadata",
|
||||
discovery,
|
||||
sample_sql: METADATA_SQL,
|
||||
scope: METADATA_SCOPE,
|
||||
};
|
||||
let schema = schemars::generate::SchemaSettings::draft2020_12()
|
||||
.for_serialize()
|
||||
.into_generator()
|
||||
.into_root_schema_for::<MetadataCatalog>();
|
||||
let serialized = serde_json::to_value(&catalog).unwrap();
|
||||
assert!(jsonschema::is_valid(schema.as_value(), &serialized));
|
||||
assert_eq!(serialized.get("error").is_some(), unavailable);
|
||||
assert!(serialized["fields"].is_array());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() {
|
||||
|
|
@ -328,7 +563,7 @@ mod tests {
|
|||
metadata: "invalid".into(),
|
||||
},
|
||||
];
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
let catalog = json!(metadata_sample(&sample));
|
||||
assert_eq!(
|
||||
catalog["fields"],
|
||||
json!([{
|
||||
|
|
@ -351,7 +586,7 @@ mod tests {
|
|||
metadata: json!(metadata).to_string(),
|
||||
})
|
||||
.collect();
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
let catalog = json!(metadata_sample(&sample));
|
||||
assert_eq!(catalog["truncated"], true);
|
||||
assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS));
|
||||
assert_eq!(
|
||||
|
|
|
|||
|
|
@ -1,11 +1,15 @@
|
|||
use askama::Template;
|
||||
use serde::Serialize;
|
||||
use litellm_traces::query::guide::Example;
|
||||
|
||||
use super::{AttributeCatalog, MetadataCatalog, TableSchema};
|
||||
use crate::{Error, NormalizedFieldDefinition};
|
||||
use super::{AttributeCatalog, Discovery, MetadataCatalog, TableSchema};
|
||||
use crate::{Error, NormalizedFieldDefinition, query_access::ReaderLimits};
|
||||
|
||||
#[derive(Template)]
|
||||
#[template(path = "query_help.jinja", escape = "none", blocks = [
|
||||
"live_schema",
|
||||
"normalized_fields",
|
||||
"metadata",
|
||||
"attributes",
|
||||
"recent_spans_name",
|
||||
"recent_spans_sql",
|
||||
"custom_metadata_name",
|
||||
|
|
@ -16,6 +20,16 @@ use crate::{Error, NormalizedFieldDefinition};
|
|||
"correlated_calls_sql",
|
||||
"discover_keys_name",
|
||||
"discover_keys_sql",
|
||||
"recent_spend_name",
|
||||
"recent_spend_sql",
|
||||
"model_spend_name",
|
||||
"model_spend_sql",
|
||||
"trace_spend_name",
|
||||
"trace_spend_sql",
|
||||
"unmatched_spans_name",
|
||||
"unmatched_spans_sql",
|
||||
"missing_spend",
|
||||
"partial_spend",
|
||||
"time_window",
|
||||
"reader_limits",
|
||||
"reader_profile",
|
||||
|
|
@ -33,16 +47,20 @@ pub(super) struct QueryGuide<'a> {
|
|||
pub normalized_fields: &'a [NormalizedFieldDefinition],
|
||||
pub metadata: &'a MetadataCatalog,
|
||||
pub attributes: &'a [AttributeCatalog],
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct Example {
|
||||
name: String,
|
||||
sql: String,
|
||||
pub limits: &'a ReaderLimits,
|
||||
}
|
||||
|
||||
impl QueryGuide<'_> {
|
||||
pub fn examples(&self) -> Result<[Example; 5], Error> {
|
||||
pub fn sections(&self) -> Result<[String; 4], Error> {
|
||||
Ok([
|
||||
render(&self.as_live_schema())?,
|
||||
render(&self.as_normalized_fields())?,
|
||||
render(&self.as_metadata())?,
|
||||
render(&self.as_attributes())?,
|
||||
])
|
||||
}
|
||||
|
||||
pub fn examples(&self) -> Result<[Example; 9], Error> {
|
||||
Ok([
|
||||
Example {
|
||||
name: render(&self.as_recent_spans_name())?,
|
||||
|
|
@ -64,10 +82,26 @@ impl QueryGuide<'_> {
|
|||
name: render(&self.as_discover_keys_name())?,
|
||||
sql: render(&self.as_discover_keys_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_recent_spend_name())?,
|
||||
sql: render(&self.as_recent_spend_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_model_spend_name())?,
|
||||
sql: render(&self.as_model_spend_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_trace_spend_name())?,
|
||||
sql: render(&self.as_trace_spend_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_unmatched_spans_name())?,
|
||||
sql: render(&self.as_unmatched_spans_sql())?,
|
||||
},
|
||||
])
|
||||
}
|
||||
|
||||
pub fn gotchas(&self) -> Result<[String; 11], Error> {
|
||||
pub fn gotchas(&self) -> Result<[String; 13], Error> {
|
||||
Ok([
|
||||
render(&self.as_time_window())?,
|
||||
render(&self.as_reader_limits())?,
|
||||
|
|
@ -78,6 +112,8 @@ impl QueryGuide<'_> {
|
|||
render(&self.as_literal_keys())?,
|
||||
render(&self.as_time_units())?,
|
||||
render(&self.as_spend_totals())?,
|
||||
render(&self.as_missing_spend())?,
|
||||
render(&self.as_partial_spend())?,
|
||||
render(&self.as_trace_rollups())?,
|
||||
render(&self.as_sampling())?,
|
||||
])
|
||||
|
|
|
|||
|
|
@ -1,27 +1,72 @@
|
|||
use litellm_storage_clickhouse::Query;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub const LENS_QUERIES: [litellm_traces::ReadQuery; 5] = [
|
||||
litellm_traces::ReadQuery::Availability,
|
||||
litellm_traces::ReadQuery::Agents,
|
||||
litellm_traces::ReadQuery::Sample,
|
||||
litellm_traces::ReadQuery::Content,
|
||||
litellm_traces::ReadQuery::Evidence,
|
||||
];
|
||||
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ExecutionSource {
|
||||
Traces,
|
||||
Requests,
|
||||
Both,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum ContentSource {
|
||||
Traces,
|
||||
Requests,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(deny_unknown_fields))]
|
||||
pub struct LensAccessParams {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub all_teams: u8,
|
||||
#[serde(
|
||||
deserialize_with = "super::number::boolean",
|
||||
serialize_with = "litellm_traces::wire::serialize_flag"
|
||||
)]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "litellm_traces::schema::flag")
|
||||
)]
|
||||
pub all_teams: bool,
|
||||
pub team: String,
|
||||
pub key_hash: String,
|
||||
}
|
||||
|
||||
pub struct LensAvailability;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LensAvailabilityParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "ActivityAvailability"))]
|
||||
pub struct LensAvailabilityRow {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(default, deserialize_with = "super::number::flag")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::boolean_flag")
|
||||
)]
|
||||
pub traces: u8,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(default, deserialize_with = "super::number::flag")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::boolean_flag")
|
||||
)]
|
||||
pub requests: u8,
|
||||
}
|
||||
|
||||
|
|
@ -34,13 +79,17 @@ impl Query for LensAvailability {
|
|||
|
||||
pub struct LensAgents;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LensAgentsParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "AgentRow"))]
|
||||
pub struct LensAgentsRow {
|
||||
pub agent_name: String,
|
||||
}
|
||||
|
|
@ -54,11 +103,13 @@ impl Query for LensAgents {
|
|||
|
||||
pub struct LensSample;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LensSampleParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
pub source: ExecutionSource,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
|
|
@ -71,9 +122,14 @@ pub struct LensSampleParams {
|
|||
pub execution_ids: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub sample_cap: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(deserialize_with = "super::number::percent")]
|
||||
#[cfg_attr(feature = "schema", schemars(range(min = 0, max = 100)))]
|
||||
pub sample_percent: f64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(deserialize_with = "super::number::flag")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "litellm_traces::schema::flag")
|
||||
)]
|
||||
pub preview: u8,
|
||||
pub after: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
|
|
@ -82,26 +138,49 @@ pub struct LensSampleParams {
|
|||
pub offset: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "ExecutionRow"))]
|
||||
pub struct LensSampleRow {
|
||||
pub source: String,
|
||||
pub source: ContentSource,
|
||||
pub trace_id: String,
|
||||
pub team_id: String,
|
||||
#[serde(default)]
|
||||
pub trace_ref: String,
|
||||
pub name: String,
|
||||
pub start_time: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::u64_number")
|
||||
)]
|
||||
pub span_count: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(deserialize_with = "super::number::flag")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::flag_number")
|
||||
)]
|
||||
pub root_seen: u8,
|
||||
#[serde(default)]
|
||||
pub service: String,
|
||||
#[serde(default)]
|
||||
pub attributes: Vec<(String, String)>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::u64_number")
|
||||
)]
|
||||
pub eligible: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[cfg_attr(feature = "schema", schemars(skip))]
|
||||
pub position: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(default, deserialize_with = "super::number::deserialize")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::selected")
|
||||
)]
|
||||
pub selected: f64,
|
||||
#[serde(default)]
|
||||
pub selection_key: String,
|
||||
}
|
||||
|
||||
|
|
@ -114,11 +193,13 @@ impl Query for LensSample {
|
|||
|
||||
pub struct LensContent;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LensContentParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
pub source: ContentSource,
|
||||
pub id: String,
|
||||
pub record_team: String,
|
||||
pub trace_ref: String,
|
||||
|
|
@ -127,14 +208,20 @@ pub struct LensContentParams {
|
|||
pub offset: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "PartRow"))]
|
||||
pub struct LensContentRow {
|
||||
pub span_id: String,
|
||||
pub parent_span_id: String,
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub content: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[serde(deserialize_with = "super::number::flag")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::flag_number")
|
||||
)]
|
||||
pub truncated: u8,
|
||||
}
|
||||
|
||||
|
|
@ -147,11 +234,13 @@ impl Query for LensContent {
|
|||
|
||||
pub struct LensEvidence;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[serde(deny_unknown_fields)]
|
||||
pub struct LensEvidenceParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
pub source: ContentSource,
|
||||
pub id: String,
|
||||
pub record_team: String,
|
||||
pub trace_ref: String,
|
||||
|
|
@ -159,9 +248,15 @@ pub struct LensEvidenceParams {
|
|||
pub quote: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[macro_rules_attribute::apply(wire_type)]
|
||||
#[derive(Debug)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "CountRow"))]
|
||||
pub struct LensEvidenceRow {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
#[cfg_attr(
|
||||
feature = "schema",
|
||||
schemars(schema_with = "crate::wire_schema::u64_number")
|
||||
)]
|
||||
pub count: u64,
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -42,7 +42,8 @@ struct ListTracesRowEncoding {
|
|||
pub name: String,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub status: String,
|
||||
#[serde(serialize_with = "litellm_traces::wire::serialize_status")]
|
||||
pub status: litellm_traces::SpanStatus,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
|
|
@ -79,18 +80,30 @@ pub use contracts::TraceSpansParams;
|
|||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::TraceSpansRow")]
|
||||
struct TraceSpansRowEncoding {
|
||||
#[serde(default)]
|
||||
pub trace_id: String,
|
||||
pub span_id: String,
|
||||
pub parent_span_id: String,
|
||||
pub name: String,
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub kind: litellm_traces::ObservationType,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "super::number::boolean",
|
||||
serialize_with = "litellm_traces::wire::serialize_flag"
|
||||
)]
|
||||
pub wrapper_candidate: bool,
|
||||
pub agent: String,
|
||||
#[serde(default)]
|
||||
pub framework: String,
|
||||
pub status: String,
|
||||
#[serde(serialize_with = "litellm_traces::wire::serialize_status")]
|
||||
pub status: litellm_traces::SpanStatus,
|
||||
pub status_message: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub error_truncated: u8,
|
||||
#[serde(
|
||||
deserialize_with = "super::number::boolean",
|
||||
serialize_with = "litellm_traces::wire::serialize_flag"
|
||||
)]
|
||||
pub error_truncated: bool,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ns: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
|
|
@ -103,6 +116,16 @@ struct TraceSpansRowEncoding {
|
|||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub output_tokens: u32,
|
||||
pub litellm_request_id: String,
|
||||
#[serde(default)]
|
||||
pub call_keys: Vec<litellm_traces::CallKey>,
|
||||
#[serde(
|
||||
default,
|
||||
deserialize_with = "litellm_traces::wire::evidence",
|
||||
serialize_with = "litellm_traces::wire::serialize_evidence"
|
||||
)]
|
||||
pub call_evidence: Option<litellm_traces::CallEvidenceKind>,
|
||||
#[serde(default)]
|
||||
pub tool_call_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key_hash: String,
|
||||
pub user_id: String,
|
||||
|
|
@ -158,6 +181,8 @@ struct SpendByResponseIdsParamsEncoding {
|
|||
#[serde(flatten)]
|
||||
pub access: contracts::ReadAccessParams,
|
||||
pub response_ids: Vec<String>,
|
||||
pub request_ids: Vec<String>,
|
||||
pub trace_ids: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
|
|
@ -180,11 +205,14 @@ impl From<contracts::SpendByResponseIdsParams> for SpendByResponseIdsParams {
|
|||
struct SpendByResponseIdsRowEncoding {
|
||||
pub request_id: String,
|
||||
pub response_id: String,
|
||||
pub upstream_response_id: String,
|
||||
pub trace_id: String,
|
||||
pub span_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key: String,
|
||||
pub user: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub spend: f64,
|
||||
#[serde(deserialize_with = "super::number::optional_finite")]
|
||||
pub spend: Option<f64>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
}
|
||||
|
|
@ -203,6 +231,38 @@ impl Query for ListTraces {
|
|||
const SQL: &'static str = include_str!("../../query/list_traces.sql");
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::TracePageSpansParams")]
|
||||
struct TracePageSpansParamsEncoding {
|
||||
#[serde(flatten)]
|
||||
pub access: contracts::ReadAccessParams,
|
||||
pub trace_refs: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub end_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TracePageSpansParams(
|
||||
#[serde(with = "TracePageSpansParamsEncoding")] pub contracts::TracePageSpansParams,
|
||||
);
|
||||
|
||||
impl From<contracts::TracePageSpansParams> for TracePageSpansParams {
|
||||
fn from(value: contracts::TracePageSpansParams) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct TracePageSpans;
|
||||
|
||||
impl Query for TracePageSpans {
|
||||
type Params = TracePageSpansParams;
|
||||
type Row = TraceSpansRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/trace_page_spans.sql");
|
||||
}
|
||||
|
||||
pub struct TraceSpans;
|
||||
|
||||
impl Query for TraceSpans {
|
||||
|
|
@ -279,11 +339,11 @@ mod tests {
|
|||
#[case::quoted(true)]
|
||||
fn rows_decode_into_neutral_contracts(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesRow>(
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<TraceSpansRow>(
|
||||
json!({"span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "agent": "agent", "framework": "claude-agent-sdk", "status": "error", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
|
||||
json!({"trace_id": "trace", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpanDetailRow>(
|
||||
|
|
@ -295,7 +355,7 @@ mod tests {
|
|||
quoted,
|
||||
);
|
||||
round_trip::<SpendByResponseIdsRow>(
|
||||
json!({"request_id": "request", "response_id": "response", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
|
||||
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}),
|
||||
quoted,
|
||||
);
|
||||
}
|
||||
|
|
@ -305,16 +365,44 @@ mod tests {
|
|||
#[case::quoted(true)]
|
||||
fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesParams>(
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "key", "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}),
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpanErrorParams>(
|
||||
json!({"all_teams": 0, "user_id": "", "team_ids": [], "api_key_hash": "key", "trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
|
||||
json!({"all_teams": 0, "user_id": "", "team_ids": [], "trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpendByResponseIdsParams>(
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "", "response_ids": ["response"], "start_ms": -1, "end_ms": 10}),
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "response_ids": ["response"], "request_ids": ["request"], "trace_ids": ["trace"], "start_ms": -1, "end_ms": 10}),
|
||||
quoted,
|
||||
);
|
||||
}
|
||||
#[rstest]
|
||||
#[case::unknown(json!(null), None)]
|
||||
#[case::free(json!(0), Some(0.0))]
|
||||
#[case::paid(json!("0.125"), Some(0.125))]
|
||||
fn spend_rows_preserve_unknown_and_known_cost(
|
||||
#[case] cost: serde_json::Value,
|
||||
#[case] expected: Option<f64>,
|
||||
) {
|
||||
let row: SpendByResponseIdsRow = serde_json::from_value(json!({
|
||||
"request_id": "request", "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
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(row.0.spend, expected);
|
||||
}
|
||||
#[rstest]
|
||||
#[case::nan(json!("NaN"))]
|
||||
#[case::infinity(json!("1e999"))]
|
||||
#[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": "",
|
||||
"trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key",
|
||||
"user": "user", "spend": cost, "start_ms": 0
|
||||
}));
|
||||
assert!(row.is_err());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,11 +18,95 @@ where
|
|||
.map_err(serde::de::Error::custom)
|
||||
}
|
||||
|
||||
pub(super) fn optional_finite<'de, D: Deserializer<'de>>(
|
||||
deserializer: D,
|
||||
) -> Result<Option<f64>, D::Error> {
|
||||
let value = Option::<serde_json::Value>::deserialize(deserializer)?;
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let number: f64 = deserialize(value).map_err(serde::de::Error::custom)?;
|
||||
if number.is_finite() {
|
||||
Ok(Some(number))
|
||||
} else {
|
||||
Err(serde::de::Error::custom("expected finite spend"))
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn flag<'de, D: Deserializer<'de>>(deserializer: D) -> Result<u8, D::Error> {
|
||||
match deserialize(deserializer)? {
|
||||
value @ 0..=1 => Ok(value),
|
||||
_ => Err(serde::de::Error::custom("expected 0 or 1")),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn percent<'de, D: Deserializer<'de>>(deserializer: D) -> Result<f64, D::Error> {
|
||||
let value: f64 = deserialize(deserializer)?;
|
||||
if value.is_finite() && (0.0..=100.0).contains(&value) {
|
||||
Ok(value)
|
||||
} else {
|
||||
Err(serde::de::Error::custom(
|
||||
"expected a finite percentage between 0 and 100",
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn boolean<'de, D: Deserializer<'de>>(deserializer: D) -> Result<bool, D::Error> {
|
||||
flag(deserializer).map(|value| value == 1)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::query::named::SpanErrorRow;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::flag_zero(serde_json::json!(0), true)]
|
||||
#[case::flag_one(serde_json::json!("1"), true)]
|
||||
#[case::invalid_flag(serde_json::json!(2), false)]
|
||||
fn access_rejects_non_boolean_flags(#[case] value: serde_json::Value, #[case] valid: bool) {
|
||||
let parameters = serde_json::json!({"all_teams": value, "team": "team", "key_hash": ""});
|
||||
assert_eq!(
|
||||
serde_json::from_value::<crate::query::lens::LensAccessParams>(parameters).is_ok(),
|
||||
valid
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::zero(serde_json::json!(0), true)]
|
||||
#[case::hundred(serde_json::json!("100"), true)]
|
||||
#[case::negative(serde_json::json!(-0.1), false)]
|
||||
#[case::too_large(serde_json::json!(100.1), false)]
|
||||
#[case::nan(serde_json::json!("NaN"), false)]
|
||||
fn sampling_rejects_invalid_percentages(#[case] value: serde_json::Value, #[case] valid: bool) {
|
||||
let parameters = serde_json::json!({
|
||||
"all_teams": 0, "team": "team", "key_hash": "", "source": "both", "start": 0, "end": 1,
|
||||
"agent_name": "", "service": "", "filter_keys": [], "filter_values": [], "selected_team": "",
|
||||
"execution_ids": [], "sample_cap": 0, "sample_percent": value, "preview": 0, "after": "",
|
||||
"limit": 10, "offset": 0
|
||||
});
|
||||
assert_eq!(
|
||||
serde_json::from_value::<crate::query::lens::LensSampleParams>(parameters).is_ok(),
|
||||
valid
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::trace("traces", true)]
|
||||
#[case::request("requests", true)]
|
||||
#[case::both("both", false)]
|
||||
#[case::unknown("unknown", false)]
|
||||
fn content_rejects_unsupported_sources(#[case] source: &str, #[case] valid: bool) {
|
||||
let parameters = serde_json::json!({
|
||||
"all_teams": 0, "team": "team", "key_hash": "", "source": source, "id": "id",
|
||||
"record_team": "team", "trace_ref": "", "cursor": "", "offset": 0
|
||||
});
|
||||
assert_eq!(
|
||||
serde_json::from_value::<crate::query::lens::LensContentParams>(parameters).is_ok(),
|
||||
valid
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::quoted_max(serde_json::json!(u64::MAX.to_string()), Some(u64::MAX))]
|
||||
#[case::unquoted_max(serde_json::json!(u64::MAX), Some(u64::MAX))]
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ use std::{sync::Arc, time::Duration};
|
|||
|
||||
use hmac::{Hmac, Mac};
|
||||
use litellm_http::Client;
|
||||
use litellm_storage_clickhouse::READ_LIMITS;
|
||||
use litellm_traces::QueryScope;
|
||||
use moka::future::Cache;
|
||||
use strum::IntoEnumIterator;
|
||||
|
|
@ -11,6 +12,33 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
|||
|
||||
use super::{Connection, Error, TraceTable};
|
||||
|
||||
const MIB: u64 = 1024 * 1024;
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub(crate) struct ReaderLimits {
|
||||
pub result_rows: u64,
|
||||
pub result_bytes: u64,
|
||||
pub memory_bytes: u64,
|
||||
pub execution_seconds: u64,
|
||||
}
|
||||
|
||||
impl ReaderLimits {
|
||||
pub fn result_mib(&self) -> u64 {
|
||||
self.result_bytes / MIB
|
||||
}
|
||||
|
||||
pub fn memory_mib(&self) -> u64 {
|
||||
self.memory_bytes / MIB
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) const READER_LIMITS: ReaderLimits = ReaderLimits {
|
||||
result_rows: READ_LIMITS.result_rows,
|
||||
result_bytes: READ_LIMITS.response_bytes as u64,
|
||||
memory_bytes: 256 * MIB,
|
||||
execution_seconds: READ_LIMITS.execution_seconds,
|
||||
};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueryReaders {
|
||||
writer: Connection,
|
||||
|
|
@ -75,13 +103,19 @@ impl QueryReaders {
|
|||
return Err(Error::InvalidScope);
|
||||
}
|
||||
let password_hash = format!("{:x}", Sha256::digest(password));
|
||||
let ReaderLimits {
|
||||
result_rows,
|
||||
result_bytes,
|
||||
memory_bytes,
|
||||
execution_seconds,
|
||||
} = READER_LIMITS;
|
||||
self.execute(
|
||||
client,
|
||||
format!(
|
||||
"CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \
|
||||
SETTINGS readonly = 1 CONST, max_execution_time = 10 CONST, \
|
||||
max_result_rows = 1000 CONST, max_result_bytes = 4194304 CONST, \
|
||||
result_overflow_mode = 'throw' CONST, max_memory_usage = 268435456 CONST, \
|
||||
SETTINGS readonly = 1 CONST, max_execution_time = {execution_seconds} CONST, \
|
||||
max_result_rows = {result_rows} CONST, max_result_bytes = {result_bytes} CONST, \
|
||||
result_overflow_mode = 'throw' CONST, max_memory_usage = {memory_bytes} CONST, \
|
||||
max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST"
|
||||
),
|
||||
)
|
||||
|
|
@ -142,17 +176,13 @@ impl QueryReaders {
|
|||
}
|
||||
|
||||
fn predicate(scope: &QueryScope, table: TraceTable) -> String {
|
||||
let (team, key) = match table {
|
||||
TraceTable::OtelTraces | TraceTable::AgentTracesByKey => ("TeamId", "ApiKeyHash"),
|
||||
TraceTable::SpendLogs => ("team_id", "api_key"),
|
||||
let team = match table {
|
||||
TraceTable::OtelTraces | TraceTable::AgentTracesByKey => "TeamId",
|
||||
TraceTable::SpendLogs => "team_id",
|
||||
};
|
||||
match scope {
|
||||
QueryScope::Admin => "1".to_owned(),
|
||||
QueryScope::Logs {
|
||||
user_id,
|
||||
team_ids,
|
||||
api_key_hash,
|
||||
} => {
|
||||
QueryScope::All => "1".to_owned(),
|
||||
QueryScope::Owned { user_id, team_ids } => {
|
||||
let owner = literal(user_id);
|
||||
let user_clause = match table {
|
||||
TraceTable::OtelTraces => format!("UserId = {owner}"),
|
||||
|
|
@ -169,21 +199,8 @@ fn predicate(scope: &QueryScope, table: TraceTable) -> String {
|
|||
} else {
|
||||
format!("{team} IN ({teams})")
|
||||
};
|
||||
format!(
|
||||
"({owner} != '' AND {user_clause}) OR ({team_clause}) OR ({hash} != '' AND {key} = {hash})",
|
||||
owner = owner,
|
||||
hash = literal(api_key_hash),
|
||||
)
|
||||
format!("({owner} != '' AND {user_clause}) OR ({team_clause})")
|
||||
}
|
||||
QueryScope::Team { team_id } => format!("{team} = {}", literal(team_id)),
|
||||
QueryScope::Key {
|
||||
team_id,
|
||||
api_key_hash,
|
||||
} => format!(
|
||||
"{team} = {} AND {key} = {}",
|
||||
literal(team_id),
|
||||
literal(api_key_hash)
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -205,33 +222,24 @@ mod tests {
|
|||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::otel(TraceTable::OtelTraces, "TeamId", "ApiKeyHash")]
|
||||
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "ApiKeyHash")]
|
||||
#[case::spend(TraceTable::SpendLogs, "team_id", "api_key")]
|
||||
#[case::otel(TraceTable::OtelTraces, "TeamId", "UserId = ''")]
|
||||
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "UserIds = ['']")]
|
||||
#[case::spend(TraceTable::SpendLogs, "team_id", "user = ''")]
|
||||
fn predicates_preserve_scope_and_escape_values(
|
||||
#[case] table: TraceTable,
|
||||
#[case] team: &str,
|
||||
#[case] key: &str,
|
||||
#[case] user: &str,
|
||||
) {
|
||||
assert_eq!(predicate(&QueryScope::Admin, table), "1");
|
||||
assert_eq!(predicate(&QueryScope::All, table), "1");
|
||||
assert_eq!(
|
||||
predicate(
|
||||
&QueryScope::Team {
|
||||
team_id: "team'\\".into()
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team'\\".into()]
|
||||
},
|
||||
table
|
||||
),
|
||||
format!("{team} = 'team\\'\\\\'")
|
||||
);
|
||||
assert_eq!(
|
||||
predicate(
|
||||
&QueryScope::Key {
|
||||
team_id: "".into(),
|
||||
api_key_hash: "key'\\".into()
|
||||
},
|
||||
table
|
||||
),
|
||||
format!("{team} = '' AND {key} = 'key\\'\\\\'")
|
||||
format!("('' != '' AND {user}) OR ({team} IN ('team\\'\\\\'))")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
368
litellm-rust/crates/traces-clickhouse/src/reads.rs
Normal file
|
|
@ -0,0 +1,368 @@
|
|||
//! Scoped trace reads: the trace list, one trace resolved with its spend, and span payloads.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use base64::{Engine, engine::general_purpose::URL_SAFE};
|
||||
use litellm_http::Client;
|
||||
use litellm_storage_clickhouse::fetch;
|
||||
use litellm_traces::{
|
||||
SpanDetail, SpanErrorPage, SpendLookup, Trace, TracePage, listed_summary,
|
||||
query::named as contracts, resolve_trace, to_ui_content,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::{
|
||||
Connection, Error,
|
||||
query::named::{
|
||||
ListTraces, ListTracesParams, ReadAccessParams, SpanDetail as SpanDetailQuery,
|
||||
SpanDetailParams, SpanError, SpanErrorParams, SpendByResponseIds, SpendByResponseIdsParams,
|
||||
TraceIdentity, TraceIdentityParams, TracePageSpans, TracePageSpansParams, TraceSpans,
|
||||
TraceSpansParams,
|
||||
},
|
||||
};
|
||||
|
||||
const NANOS_PER_MS: i64 = 1_000_000;
|
||||
const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000;
|
||||
|
||||
fn encode_cursor<T: Serialize>(position: &T) -> String {
|
||||
URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default())
|
||||
}
|
||||
|
||||
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")),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct ErrorPosition {
|
||||
offset: u64,
|
||||
version: String,
|
||||
}
|
||||
|
||||
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"));
|
||||
}
|
||||
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()));
|
||||
}
|
||||
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 fetch::<SpendByResponseIds>(client, connection, ¶ms).await {
|
||||
Ok(rows) => rows.into_iter().map(|row| row.0).collect(),
|
||||
Err(error) => {
|
||||
tracing::warn!(%error, "trace spend lookup unavailable");
|
||||
Vec::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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> {
|
||||
let (cursor_ms, cursor_trace_id) = trace_position(cursor)?;
|
||||
let params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: access.clone(),
|
||||
start_ms,
|
||||
end_ms,
|
||||
cursor_ms,
|
||||
cursor_trace_id,
|
||||
limit,
|
||||
});
|
||||
let page: Vec<contracts::ListTracesRow> = fetch::<ListTraces>(client, connection, ¶ms)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| row.0)
|
||||
.collect();
|
||||
let next_cursor = page
|
||||
.last()
|
||||
.filter(|_| page.len() == limit as usize)
|
||||
.map(|last| encode_cursor(&(last.start_ms, &last.trace_ref)));
|
||||
let (Some(page_start), Some(page_end)) = (
|
||||
page.iter().map(|row| row.start_ms).min(),
|
||||
page.iter().map(|row| row.start_ms + row.duration_ms).max(),
|
||||
) else {
|
||||
return Ok(TracePage {
|
||||
data: Vec::new(),
|
||||
next_cursor,
|
||||
});
|
||||
};
|
||||
let span_params = TracePageSpansParams::from(contracts::TracePageSpansParams {
|
||||
access: access.clone(),
|
||||
trace_refs: page.iter().map(|row| row.trace_ref.clone()).collect(),
|
||||
start_ms: page_start,
|
||||
end_ms: page_end + 1,
|
||||
});
|
||||
let span_rows: Vec<contracts::TraceSpansRow> =
|
||||
fetch::<TracePageSpans>(client, connection, &span_params)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| row.0)
|
||||
.collect();
|
||||
let spend_rows = spend(client, connection, access, &span_rows).await;
|
||||
let mut by_trace: HashMap<(String, String, String), Vec<contracts::TraceSpansRow>> =
|
||||
HashMap::new();
|
||||
for span in span_rows {
|
||||
let key = (
|
||||
span.team_id.clone(),
|
||||
span.api_key_hash.clone(),
|
||||
span.trace_id.clone(),
|
||||
);
|
||||
by_trace.entry(key).or_default().push(span);
|
||||
}
|
||||
let data = page
|
||||
.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();
|
||||
resolve_trace(&row.trace_id, &row.trace_ref, spans, &spend_rows)
|
||||
.map_or_else(|| listed_summary(row), |trace| trace.summary)
|
||||
})
|
||||
.collect();
|
||||
Ok(TracePage { data, next_cursor })
|
||||
}
|
||||
|
||||
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: Vec<contracts::TraceSpansRow> = fetch::<TraceSpans>(client, connection, ¶ms)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| row.0)
|
||||
.collect();
|
||||
if rows.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
let spend_rows = spend(client, connection, access, &rows).await;
|
||||
Ok(resolve_trace(trace_id, &trace_ref, &rows, &spend_rows))
|
||||
}
|
||||
|
||||
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"))
|
||||
));
|
||||
}
|
||||
}
|
||||
|
|
@ -78,12 +78,18 @@ pub struct NormalizedFieldDefinition {
|
|||
pub meaning: &'static str,
|
||||
}
|
||||
|
||||
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
|
||||
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 15] = [
|
||||
NormalizedFieldDefinition {
|
||||
name: "observation_type",
|
||||
clickhouse_column: "ObservationType",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent, LLM, tool, chain, or framework span",
|
||||
meaning: "Operation recorded by the span, including agent, model, tool, retrieval and evaluation steps",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "wrapper_candidate",
|
||||
clickhouse_column: "WrapperCandidate",
|
||||
clickhouse_type: "Bool",
|
||||
meaning: "Span may only wrap the operation it names; the trace graph decides",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "agent_name",
|
||||
|
|
@ -97,12 +103,30 @@ pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
|
|||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent framework or SDK that emitted this span, e.g. claude-agent-sdk",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "agent_metadata",
|
||||
clickhouse_column: "AgentMetadata",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Typed agent metadata as JSON, including thread, subagent, runtime and repository identity",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "litellm_request_id",
|
||||
clickhouse_column: "LiteLLMRequestId",
|
||||
clickhouse_type: "String",
|
||||
meaning: "LiteLLM response ID used to link a span to a spend log",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "call_keys",
|
||||
clickhouse_column: "CallKeys",
|
||||
clickhouse_type: "Array(String)",
|
||||
meaning: "Model requests the span accounts for, as kind:id (litellm_request, provider_response, transport)",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "call_evidence",
|
||||
clickhouse_column: "CallEvidence",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Whether CallKeys are all of the span's requests: complete, partial or unknown",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "model",
|
||||
clickhouse_column: "Model",
|
||||
|
|
@ -127,10 +151,22 @@ pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
|
|||
clickhouse_type: "String",
|
||||
meaning: "Normalized input payload",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input_preview",
|
||||
clickhouse_column: "InputPreview",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Latest user message of the input, else the input's first characters",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output",
|
||||
clickhouse_column: "Output",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized output payload",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "tool_call_id",
|
||||
clickhouse_column: "ToolCallId",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Tool call the span executes, shared by instrumentations recording the same call",
|
||||
},
|
||||
];
|
||||
|
|
|
|||
217
litellm-rust/crates/traces-clickhouse/src/span_row.rs
Normal file
|
|
@ -0,0 +1,217 @@
|
|||
//! Decoded spans as `otel_traces` rows: payloads capped, the sending tenant stamped over whatever
|
||||
//! the export claimed, and resource maps shared across the rows that came from one resource.
|
||||
|
||||
use std::collections::{BTreeMap, HashMap};
|
||||
|
||||
use litellm_traces::{
|
||||
CallEvidence, CallKey, DecodedEvent, DecodedSpan, Shared, SharedIdentity, Tenant,
|
||||
truncate_messages, truncate_value,
|
||||
};
|
||||
use serde::Serialize;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::InsertRow;
|
||||
|
||||
/// Converts each distinct shared source once; keeping the source pins its identity.
|
||||
struct SharedValues<T>(HashMap<SharedIdentity, (Shared<T>, Shared<Value>)>);
|
||||
|
||||
impl<T: Clone> SharedValues<T> {
|
||||
fn new() -> Self {
|
||||
Self(HashMap::new())
|
||||
}
|
||||
|
||||
fn get(&mut self, source: &Shared<T>, convert: impl FnOnce(&T) -> Value) -> Shared<Value> {
|
||||
self.0
|
||||
.entry(source.identity())
|
||||
.or_insert_with(|| (source.clone(), Shared::new(convert(source))))
|
||||
.1
|
||||
.clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn stamped(attributes: &BTreeMap<String, String>, tenant: &Tenant) -> Value {
|
||||
let mut stamped: Map<String, Value> = attributes
|
||||
.iter()
|
||||
.map(|(key, value)| (key.clone(), Value::from(value.as_str())))
|
||||
.collect();
|
||||
for (key, value) in [
|
||||
("litellm.team_id", &tenant.team_id),
|
||||
("litellm.api_key_hash", &tenant.api_key_hash),
|
||||
("litellm.org_id", &tenant.org_id),
|
||||
("litellm.user_id", &tenant.user_id),
|
||||
] {
|
||||
stamped.insert(key.to_owned(), Value::from(value.as_str()));
|
||||
}
|
||||
Value::Object(stamped)
|
||||
}
|
||||
|
||||
fn exception_message(events: &[DecodedEvent]) -> String {
|
||||
events
|
||||
.iter()
|
||||
.find(|event| event.name == "exception")
|
||||
.and_then(|event| {
|
||||
event
|
||||
.attributes
|
||||
.get("exception.message")
|
||||
.filter(|message| !message.is_empty())
|
||||
.or_else(|| event.attributes.get("exception.type"))
|
||||
})
|
||||
.cloned()
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn json<T: Serialize>(value: T) -> Value {
|
||||
serde_json::to_value(value).unwrap_or(Value::Null)
|
||||
}
|
||||
|
||||
fn present_fields<T: Serialize>(value: &T) -> String {
|
||||
match json(value) {
|
||||
Value::Object(fields) => Value::Object(
|
||||
fields
|
||||
.into_iter()
|
||||
.filter(|(_, value)| !value.is_null())
|
||||
.collect(),
|
||||
)
|
||||
.to_string(),
|
||||
other => other.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn span_rows(
|
||||
spans: Vec<DecodedSpan>,
|
||||
tenant: &Tenant,
|
||||
max_value_bytes: usize,
|
||||
) -> Vec<InsertRow> {
|
||||
let mut resources = SharedValues::new();
|
||||
let mut scopes = SharedValues::new();
|
||||
spans
|
||||
.into_iter()
|
||||
.map(|span| {
|
||||
let normalized = span.normalized;
|
||||
let service = span
|
||||
.resource_attributes
|
||||
.get("service.name")
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let status_message = if span.status_message.is_empty() {
|
||||
exception_message(&span.events)
|
||||
} else {
|
||||
span.status_message
|
||||
};
|
||||
let attributes: Map<String, Value> = span
|
||||
.attributes
|
||||
.into_iter()
|
||||
.filter(|(key, _)| !span.consumed_attributes.contains(&key.as_str()))
|
||||
.map(|(key, value)| (key, Value::String(truncate_value(value, max_value_bytes))))
|
||||
.collect();
|
||||
let shared = [
|
||||
(
|
||||
"ResourceAttributes",
|
||||
resources.get(&span.resource_attributes, |attributes| {
|
||||
stamped(attributes, tenant)
|
||||
}),
|
||||
),
|
||||
(
|
||||
"ScopeName",
|
||||
scopes.get(&span.scope_name, |name| Value::from(name.as_str())),
|
||||
),
|
||||
(
|
||||
"ScopeVersion",
|
||||
scopes.get(&span.scope_version, |version| Value::from(version.as_str())),
|
||||
),
|
||||
];
|
||||
let owned = [
|
||||
("Timestamp", json(span.start_ns)),
|
||||
("TraceId", Value::String(span.trace_id)),
|
||||
("SpanId", Value::String(span.span_id)),
|
||||
("ParentSpanId", Value::String(span.parent_span_id)),
|
||||
("TraceState", Value::String(span.trace_state)),
|
||||
("SpanName", Value::String(span.name)),
|
||||
("SpanKind", Value::String(span.kind)),
|
||||
("ServiceName", Value::String(service)),
|
||||
("SpanAttributes", Value::Object(attributes)),
|
||||
("Duration", json(span.end_ns - span.start_ns)),
|
||||
("StatusCode", Value::String(span.status_code)),
|
||||
("StatusMessage", Value::String(status_message)),
|
||||
("TeamId", Value::from(tenant.team_id.as_str())),
|
||||
("ApiKeyHash", Value::from(tenant.api_key_hash.as_str())),
|
||||
("UserId", Value::from(tenant.user_id.as_str())),
|
||||
("ObservationType", json(normalized.observation_type)),
|
||||
(
|
||||
"WrapperCandidate",
|
||||
Value::Bool(normalized.wrapper_candidate),
|
||||
),
|
||||
(
|
||||
"AgentName",
|
||||
Value::String(normalized.agent_name.unwrap_or_default()),
|
||||
),
|
||||
(
|
||||
"Framework",
|
||||
Value::String(
|
||||
normalized
|
||||
.framework
|
||||
.map(|integration| integration.to_string())
|
||||
.unwrap_or_default(),
|
||||
),
|
||||
),
|
||||
(
|
||||
"AgentMetadata",
|
||||
Value::String(present_fields(&normalized.agent_metadata)),
|
||||
),
|
||||
(
|
||||
"LiteLLMRequestId",
|
||||
Value::String(request_id(&normalized.calls).to_owned()),
|
||||
),
|
||||
(
|
||||
"CallKeys",
|
||||
json(
|
||||
normalized
|
||||
.calls
|
||||
.key_set()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect::<Vec<_>>(),
|
||||
),
|
||||
),
|
||||
("CallEvidence", json(normalized.calls.kind())),
|
||||
("Model", Value::String(normalized.model.unwrap_or_default())),
|
||||
("InputTokens", Value::from(normalized.input_tokens)),
|
||||
("OutputTokens", Value::from(normalized.output_tokens)),
|
||||
(
|
||||
"Input",
|
||||
Value::String(truncate_messages(normalized.input, max_value_bytes)),
|
||||
),
|
||||
("InputPreview", Value::String(normalized.input_preview)),
|
||||
(
|
||||
"Output",
|
||||
Value::String(truncate_value(normalized.output, max_value_bytes)),
|
||||
),
|
||||
(
|
||||
"ToolCallId",
|
||||
Value::String(normalized.tool_call_id.unwrap_or_default()),
|
||||
),
|
||||
];
|
||||
shared
|
||||
.into_iter()
|
||||
.chain(
|
||||
owned
|
||||
.into_iter()
|
||||
.map(|(column, value)| (column, Shared::new(value))),
|
||||
)
|
||||
.map(|(column, value)| (column.to_owned(), value))
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn request_id(evidence: &CallEvidence) -> &str {
|
||||
evidence
|
||||
.key_set()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find_map(|key| match key {
|
||||
CallKey::LiteLlmRequest(id) | CallKey::ProviderResponse(id) => Some(id.as_str()),
|
||||
CallKey::Transport => None,
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
|
@ -19,6 +19,9 @@ pub async fn execute_named_read(
|
|||
named_json::<TraceIdentity>(client, connection, parameters).await
|
||||
}
|
||||
ReadQuery::TraceSpans => named_json::<TraceSpans>(client, connection, parameters).await,
|
||||
ReadQuery::TracePageSpans => {
|
||||
named_json::<TracePageSpans>(client, connection, parameters).await
|
||||
}
|
||||
ReadQuery::SpanDetail => named_json::<SpanDetail>(client, connection, parameters).await,
|
||||
ReadQuery::SpanError => named_json::<SpanError>(client, connection, parameters).await,
|
||||
ReadQuery::SpendByResponseIds => {
|
||||
|
|
@ -64,7 +67,7 @@ mod tests {
|
|||
#[case] specific: serde_json::Value,
|
||||
) {
|
||||
let common = serde_json::json!({
|
||||
"all_teams": 1, "user_id": "", "team_ids": [], "api_key_hash": "", "trace_id": "trace", "trace_ref": ""
|
||||
"all_teams": 1, "user_id": "", "team_ids": [], "trace_id": "trace", "trace_ref": ""
|
||||
});
|
||||
let parameters: BTreeMap<String, Parameter> = common
|
||||
.as_object()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
#[macro_rules_attribute::apply(response_type)]
|
||||
#[cfg_attr(feature = "schema", schemars(rename = "TraceTableName"))]
|
||||
#[derive(
|
||||
Clone, Copy, Debug, strum::Display, strum::AsRefStr, strum::EnumIter, strum::IntoStaticStr,
|
||||
)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum TraceTable {
|
||||
OtelTraces,
|
||||
|
|
|
|||
119
litellm-rust/crates/traces-clickhouse/src/wire_schema.rs
Normal file
|
|
@ -0,0 +1,119 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use schemars::{JsonSchema, Schema, SchemaGenerator, generate::SchemaSettings};
|
||||
use serde_json::json;
|
||||
|
||||
use crate::query::lens;
|
||||
|
||||
fn quoted_u64() -> Schema {
|
||||
let upper = u64::MAX.to_string();
|
||||
let alternatives = upper
|
||||
.char_indices()
|
||||
.filter_map(|(index, digit)| {
|
||||
let lower = if index == 0 { '1' } else { '0' };
|
||||
if digit <= lower {
|
||||
return None;
|
||||
}
|
||||
Some(format!(
|
||||
"{}[{}-{}][0-9]{{{}}}",
|
||||
&upper[..index],
|
||||
lower,
|
||||
char::from(digit as u8 - 1),
|
||||
upper.len() - index - 1
|
||||
))
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("|");
|
||||
json!({
|
||||
"type": "string",
|
||||
"pattern": format!("^(?:0|[1-9][0-9]{{0,{}}}|{alternatives}|{upper})$", upper.len() - 2),
|
||||
})
|
||||
.try_into()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn numeric_wire(normalized: Schema, python_type: String) -> Schema {
|
||||
json!({
|
||||
"anyOf": [normalized, quoted_u64()],
|
||||
"x-python-normalized": {"type": python_type, "minimum": 0, "maximum": u64::MAX},
|
||||
})
|
||||
.try_into()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn u64_number(generator: &mut SchemaGenerator) -> Schema {
|
||||
numeric_wire(u64::json_schema(generator), "int".to_owned())
|
||||
}
|
||||
|
||||
pub(crate) fn flag_number(_: &mut SchemaGenerator) -> Schema {
|
||||
json!({
|
||||
"anyOf": [{"type": "integer", "enum": [0, 1]}, {"type": "string", "enum": ["0", "1"]}],
|
||||
"x-python-normalized": {"type": "int", "minimum": 0, "maximum": 1}
|
||||
})
|
||||
.try_into()
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn boolean_flag(_: &mut SchemaGenerator) -> Schema {
|
||||
json!({
|
||||
"anyOf": [{"type": "boolean"}, {"type": "integer", "enum": [0, 1]}, {"type": "string", "enum": ["0", "1"]}],
|
||||
"default": false,
|
||||
"x-python-normalized": {"type": "bool"}
|
||||
}).try_into().unwrap()
|
||||
}
|
||||
|
||||
pub(crate) fn selected(generator: &mut SchemaGenerator) -> Schema {
|
||||
u64_number(generator)
|
||||
}
|
||||
|
||||
fn received<T: JsonSchema>() -> Schema {
|
||||
SchemaSettings::draft2020_12()
|
||||
.for_deserialize()
|
||||
.with_transform(litellm_traces::schema::integer_bounds)
|
||||
.into_generator()
|
||||
.into_root_schema_for::<T>()
|
||||
}
|
||||
|
||||
pub fn schemas() -> BTreeMap<&'static str, Schema> {
|
||||
BTreeMap::from([
|
||||
("ReadQueryName", json!({"$schema": "https://json-schema.org/draft/2020-12/schema", "title": "ReadQueryName", "type": "string", "enum": lens::LENS_QUERIES.map(|query| query.to_string())}).try_into().unwrap()),
|
||||
("LensAccessParams", received::<lens::LensAccessParams>()),
|
||||
("LensSampleParams", received::<lens::LensSampleParams>()),
|
||||
("LensContentParams", received::<lens::LensContentParams>()),
|
||||
("LensEvidenceParams", received::<lens::LensEvidenceParams>()),
|
||||
(
|
||||
"ActivityAvailability",
|
||||
received::<lens::LensAvailabilityRow>(),
|
||||
),
|
||||
("ExecutionRow", received::<lens::LensSampleRow>()),
|
||||
("PartRow", received::<lens::LensContentRow>()),
|
||||
("CountRow", received::<lens::LensEvidenceRow>()),
|
||||
("AgentRow", received::<lens::LensAgentsRow>()),
|
||||
("TraceQueryHelp", crate::query::help_schema()),
|
||||
])
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::zero(json!(0), true)]
|
||||
#[case::quoted_zero(json!("0"), true)]
|
||||
#[case::maximum(json!(u64::MAX), true)]
|
||||
#[case::quoted_maximum(json!(u64::MAX.to_string()), true)]
|
||||
#[case::negative(json!(-1), false)]
|
||||
#[case::overflow(json!((u128::from(u64::MAX) + 1).to_string()), false)]
|
||||
#[case::fraction(json!(1.5), false)]
|
||||
fn count_schema_enforces_the_native_range(
|
||||
#[case] value: serde_json::Value,
|
||||
#[case] valid: bool,
|
||||
) {
|
||||
let schema = received::<lens::LensEvidenceRow>();
|
||||
assert_eq!(
|
||||
jsonschema::is_valid(schema.as_value(), &json!({"count": value})),
|
||||
valid
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,65 +1,268 @@
|
|||
Trace SQL query guide
|
||||
|
||||
Live ClickHouse schema
|
||||
{% for table in tables %}
|
||||
{% block live_schema -%}
|
||||
{% for table in tables -%}
|
||||
{{ table.name }}
|
||||
{% for column in table.columns %}{{ column.name }}: {{ column.kind }}
|
||||
{% endfor %}{% endfor %}
|
||||
Normalized span fields
|
||||
{% for field in normalized_fields %}{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }})
|
||||
{{ field.meaning }}
|
||||
{% for column in table.columns -%}
|
||||
{{ column.name }}: {{ column.kind }}
|
||||
{% endfor %}
|
||||
Observed LLM call metadata
|
||||
{% endfor -%}
|
||||
{%- endblock %}
|
||||
|
||||
{% block normalized_fields -%}
|
||||
{% for field in normalized_fields -%}
|
||||
{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }})
|
||||
{{ field.meaning }}
|
||||
{% endfor -%}
|
||||
{%- endblock %}
|
||||
|
||||
{% block metadata -%}
|
||||
{{ metadata.scope }}
|
||||
Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }}
|
||||
{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }}
|
||||
{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows
|
||||
{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %}
|
||||
{% endfor %}{% endif %}
|
||||
Observed span and resource attributes
|
||||
{% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }}
|
||||
Sampling SQL:
|
||||
{{ metadata.sample_sql }}
|
||||
{% match metadata.discovery -%}
|
||||
{% when Discovery::Unavailable(error) -%}
|
||||
Metadata discovery unavailable: {{ error }}
|
||||
{% when Discovery::Observed(sample) -%}
|
||||
Sampled rows: {{ sample.sampled_rows }}; invalid JSON rows: {{ sample.invalid_json_rows }}; truncated: {{ sample.truncated }}
|
||||
{% if sample.fields.is_empty() -%}
|
||||
No metadata paths found in the sampled rows
|
||||
{% else -%}
|
||||
{% for field in sample.fields -%}
|
||||
{{ field.expression }}: {{ field.types|join(", ") }}
|
||||
{% endfor -%}
|
||||
{% endif -%}
|
||||
{% endmatch -%}
|
||||
{%- endblock %}
|
||||
|
||||
{% block attributes -%}
|
||||
{% for catalog in attributes -%}
|
||||
{{ catalog.table }}.{{ catalog.column }}
|
||||
{{ catalog.scope }}
|
||||
{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }}
|
||||
{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans
|
||||
{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }}
|
||||
{% endfor %}{% endif %}{% endfor %}
|
||||
Examples
|
||||
Discovery SQL:
|
||||
{{ catalog.discovery_sql }}
|
||||
{% match catalog.discovery -%}
|
||||
{% when Discovery::Unavailable(error) -%}
|
||||
Attribute discovery unavailable: {{ error }}
|
||||
{% when Discovery::Observed(sample) -%}
|
||||
Truncated: {{ sample.truncated }}
|
||||
{% if sample.fields.is_empty() -%}
|
||||
No attribute keys found in the sampled spans
|
||||
{% else -%}
|
||||
{% for field in sample.fields -%}
|
||||
{{ field.expression }}: {{ field.kind }}
|
||||
{% endfor -%}
|
||||
{% endif -%}
|
||||
{% endmatch %}
|
||||
{% endfor -%}
|
||||
{%- endblock %}
|
||||
|
||||
{% block recent_spans_name %}Recent normalized LLM spans{% endblock %}
|
||||
{% block recent_spans_sql %}SELECT TraceId, SpanId, Model, InputTokens, OutputTokens, Duration / 1000000 AS duration_ms FROM otel_traces WHERE Timestamp >= now() - INTERVAL 1 DAY AND ObservationType = 'llm' ORDER BY Timestamp DESC LIMIT 100{% endblock %}
|
||||
{% block recent_spans_name -%}
|
||||
Recent normalized LLM spans
|
||||
{%- endblock %}
|
||||
|
||||
{% block custom_metadata_name %}Find calls by custom metadata{% endblock %}
|
||||
{% block custom_metadata_sql %}SELECT request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'project') AND JSONExtractString(metadata, 'project') = 'example' ORDER BY start_time DESC LIMIT 100{% endblock %}
|
||||
{% block recent_spans_sql -%}
|
||||
SELECT
|
||||
TraceId, SpanId, Model, InputTokens, OutputTokens,
|
||||
Duration / 1000000 AS duration_ms
|
||||
FROM otel_traces
|
||||
WHERE Timestamp >= now() - INTERVAL 1 DAY
|
||||
AND ObservationType = 'llm'
|
||||
ORDER BY Timestamp DESC
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block nested_metadata_name %}Nested metadata with unknown types{% endblock %}
|
||||
{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %}
|
||||
{% block custom_metadata_name -%}
|
||||
Find calls by custom metadata
|
||||
{%- endblock %}
|
||||
|
||||
{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %}
|
||||
{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND (t.TeamId != '' OR (t.UserId != '' AND t.UserId = s.user) OR (t.ApiKeyHash != '' AND t.ApiKeyHash = s.api_key)) WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %}
|
||||
{% block custom_metadata_sql -%}
|
||||
SELECT
|
||||
request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
AND JSONHas(metadata, 'project')
|
||||
AND JSONExtractString(metadata, 'project') = 'example'
|
||||
ORDER BY start_time DESC
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block discover_keys_name %}Discover metadata keys over a different window{% endblock %}
|
||||
{% block discover_keys_sql %}SELECT DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 30 DAY ORDER BY key LIMIT 200{% endblock %}
|
||||
{% block nested_metadata_name -%}
|
||||
Nested metadata with unknown types
|
||||
{%- endblock %}
|
||||
|
||||
Gotchas
|
||||
{% block nested_metadata_sql -%}
|
||||
SELECT
|
||||
request_id,
|
||||
JSONType(metadata, 'labels', 'priority') AS type,
|
||||
JSONExtractRaw(metadata, 'labels', 'priority') AS value
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
AND JSONHas(metadata, 'labels', 'priority')
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %}
|
||||
{% block correlated_calls_name -%}
|
||||
Traces correlated with LLM call metadata
|
||||
{%- endblock %}
|
||||
|
||||
{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %}
|
||||
{% block correlated_calls_sql -%}
|
||||
SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata
|
||||
FROM otel_traces AS t
|
||||
INNER JOIN (
|
||||
SELECT *
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
) AS s
|
||||
ON t.LiteLLMRequestId = s.response_id
|
||||
AND t.TeamId = s.team_id
|
||||
AND ((t.UserId != '' AND t.UserId = s.user)
|
||||
OR (t.ApiKeyHash != '' AND t.ApiKeyHash = s.api_key))
|
||||
WHERE t.Timestamp >= now() - INTERVAL 1 DAY
|
||||
AND t.LiteLLMRequestId != ''
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block reader_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces request-log visibility through row policies. Callers see their own user rows and permitted teams, or their own key rows when no user identity is available. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %}
|
||||
{% block discover_keys_name -%}
|
||||
Discover metadata keys over a different window
|
||||
{%- endblock %}
|
||||
|
||||
{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %}
|
||||
{% block discover_keys_sql -%}
|
||||
SELECT
|
||||
DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 30 DAY
|
||||
ORDER BY key
|
||||
LIMIT 200
|
||||
{%- endblock %}
|
||||
|
||||
{% block json_values %}metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false{% endblock %}
|
||||
{% block recent_spend_name -%}
|
||||
Recent spend records
|
||||
{%- endblock %}
|
||||
|
||||
{% block map_values %}SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks{% endblock %}
|
||||
{% block recent_spend_sql -%}
|
||||
SELECT
|
||||
request_id, response_id, trace_id, span_id, model, spend,
|
||||
prompt_tokens, completion_tokens, status,
|
||||
JSONExtractBool(metadata, 'synthetic_spend') AS synthetic_spend
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
ORDER BY start_time DESC, request_id
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block literal_keys %}Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator{% endblock %}
|
||||
{% block model_spend_name -%}
|
||||
Spend and tokens by model
|
||||
{%- endblock %}
|
||||
|
||||
{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %}
|
||||
{% block model_spend_sql -%}
|
||||
SELECT
|
||||
team_id, model, requests, unknown_cost_requests,
|
||||
if(unknown_cost_requests = 0, recorded_spend, NULL) AS spend,
|
||||
input_tokens, output_tokens
|
||||
FROM (
|
||||
SELECT
|
||||
team_id, model, count() AS requests,
|
||||
countIf(isNull(spend) OR NOT isFinite(spend)) AS unknown_cost_requests,
|
||||
sum(spend) AS recorded_spend,
|
||||
sum(prompt_tokens) AS input_tokens,
|
||||
sum(completion_tokens) AS output_tokens
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
GROUP BY team_id, model
|
||||
)
|
||||
ORDER BY team_id, model
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; require one spend match per response ID and ownership before aggregating. Missing IDs or costs leave totals unknown{% endblock %}
|
||||
{% block trace_spend_name -%}
|
||||
Recorded spend by trace
|
||||
{%- endblock %}
|
||||
|
||||
{% block trace_rollups %}agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators{% endblock %}
|
||||
{% block trace_spend_sql -%}
|
||||
SELECT
|
||||
team_id, api_key, trace_id, count() AS requests,
|
||||
countIf(isNull(spend) OR NOT isFinite(spend)) AS unknown_cost_requests,
|
||||
if(unknown_cost_requests = 0, sum(spend), NULL) AS recorded_spend
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
AND trace_id != ''
|
||||
GROUP BY team_id, api_key, trace_id
|
||||
ORDER BY team_id, api_key, trace_id
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block sampling %}Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent'){% endblock %}
|
||||
{% block unmatched_spans_name -%}
|
||||
LLM spans without a direct spend match
|
||||
{%- endblock %}
|
||||
|
||||
{% block unmatched_spans_sql -%}
|
||||
SELECT
|
||||
t.TraceId, t.SpanId, t.Model, t.LiteLLMRequestId,
|
||||
t.InputTokens, t.OutputTokens
|
||||
FROM otel_traces AS t
|
||||
LEFT ANTI JOIN (
|
||||
SELECT *
|
||||
FROM spend_logs FINAL
|
||||
WHERE start_time >= now() - INTERVAL 1 DAY
|
||||
) AS s
|
||||
ON t.TeamId = s.team_id
|
||||
AND ((t.UserId != '' AND t.UserId = s.user)
|
||||
OR (t.ApiKeyHash != '' AND t.ApiKeyHash = s.api_key))
|
||||
AND t.LiteLLMRequestId != ''
|
||||
AND (t.LiteLLMRequestId = s.response_id OR t.LiteLLMRequestId = s.request_id)
|
||||
WHERE t.Timestamp >= now() - INTERVAL 1 DAY
|
||||
AND t.ObservationType = 'llm'
|
||||
ORDER BY t.Timestamp DESC, t.SpanId
|
||||
LIMIT 100
|
||||
{%- endblock %}
|
||||
|
||||
{% block time_window -%}
|
||||
Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant
|
||||
{%- endblock %}
|
||||
|
||||
{% block reader_limits -%}
|
||||
The reader enforces {{ limits.result_rows }} result rows, {{ limits.result_mib() }} MiB response bytes, {{ limits.memory_mib() }} MiB memory and a {{ limits.execution_seconds }} second query limit; exceeding limits fails instead of returning partial results
|
||||
{%- endblock %}
|
||||
|
||||
{% block reader_profile -%}
|
||||
LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces request-log visibility through row policies. Callers see their own user rows and permitted teams. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions
|
||||
{%- endblock %}
|
||||
|
||||
{% block output_format -%}
|
||||
Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output
|
||||
{%- endblock %}
|
||||
|
||||
{% block json_values -%}
|
||||
metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false
|
||||
{%- endblock %}
|
||||
|
||||
{% block map_values -%}
|
||||
SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks
|
||||
{%- endblock %}
|
||||
|
||||
{% block literal_keys -%}
|
||||
Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator
|
||||
{%- endblock %}
|
||||
|
||||
{% block time_units -%}
|
||||
Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision
|
||||
{%- endblock %}
|
||||
|
||||
{% block 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
|
||||
{%- endblock %}
|
||||
|
||||
{% block partial_spend -%}
|
||||
Recorded spend by trace totals only requests whose spend_logs.trace_id is populated. Direct ID joins do not resolve every CallKeys entry, managed Responses IDs, or transport correlation. Use the trace detail API for resolved totals; unmatched spans are a starting point for investigation
|
||||
{%- endblock %}
|
||||
|
||||
{% block spend_totals -%}
|
||||
Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; require one spend match per response ID and ownership before aggregating. Missing IDs or costs leave totals unknown
|
||||
{%- endblock %}
|
||||
|
||||
{% block trace_rollups -%}
|
||||
agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators
|
||||
{%- endblock %}
|
||||
|
||||
{% block sampling -%}
|
||||
Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent')
|
||||
{%- endblock %}
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
|||
}
|
||||
let readers = QueryReaders::new(Connection::writer(&admin_url)?, "litellm".into());
|
||||
let connection = readers
|
||||
.connection(&client, &QueryScope::Admin, "test-secret")
|
||||
.connection(&client, &QueryScope::All, "test-secret")
|
||||
.await?;
|
||||
let url = connection.url().to_string();
|
||||
Ok(Database {
|
||||
|
|
|
|||
|
|
@ -8,6 +8,18 @@ Raw OTLP exports live in `crates/traces/tests/fixtures/query_*.json`. The seeded
|
|||
|
||||
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 `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
|
||||
|
||||
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
|
||||
|
||||
The LlamaIndex captures contain provider IDs inside `output.value.raw.id`. Regression tests require normalization to retain those call keys and trace cost resolution to count nested model spans once. The Claude captures use the SDK example's local gateway adapter, which supplies the actual Anthropic message ID in the `request-id` response header. The two `claude_agent_sdk_missing_request_id_*` exports retain the earlier behavior: real spend rows exist, but model spans contain no matching call IDs, so trace spend remains unknown
|
||||
|
||||
`scripts/seed_tracing_fixtures.py` replays every JSON export in `crates/traces/tests/fixtures` through `POST /v1/traces`, then inserts all companion spend rows into ClickHouse through the production storage API and into Postgres through Prisma. The Requests table reads Postgres, while trace costs and Lens read ClickHouse. It shifts each capture into the current time window, keeping span, event, and paired spend timestamps aligned. The split `query_*.json` exports share a time shift and ID namespace to preserve cross-file parent links. Other captures get separate ID namespaces to avoid collisions between fixtures. It assigns fresh linked IDs for each run, including provider IDs inside managed response IDs, and reads the authenticated tenant from the ingested spans before stamping spend rows. The command exits unsuccessfully if any trace detail API result differs from its captured spend total or expected unknown cost. Exports without companion spend rows retain missing costs and do not create Requests entries
|
||||
|
||||
`tests/test_litellm_rust/test_traces.py` ingests these exports and spend rows into an isolated ClickHouse container, then checks trace detail costs and spend queries through the real FastAPI endpoints. Every SQL example returned by `/v1/traces/query/help` is executed through `/v1/traces/query`, including missing costs, free requests, replacement rows, and tenant ownership cases
|
||||
|
||||
`spend_logs.jsonl` contains spend insert rows with millisecond timestamps, including two versions of one request. Replace this small placeholder dataset when the actual data is available. The query fixture applies production migrations, then removes TTL from its isolated database so fixed timestamps do not expire. Background merges are stopped so rollup aggregation and `FINAL` deduplication are exercised on unmerged data. Retention behavior stays covered by the migration tests
|
||||
|
||||
Curated SQL lives in `tests/queries/*.sql`. Each query has a matching `.expected.json` containing ordered result rows for `admin`, `team`, `key`, and `other_team` readers. Update the exports and expected results together. Add a named case in `tests/queries.rs` for each new query. Assertions compare only result data, excluding server statistics and execution timing
|
||||
|
|
|
|||
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/claude_agent_sdk_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/crewai_simple_spend_logs.jsonl
vendored
Normal file
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/crewai_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/deepagents_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/deepagents_swarm_spend_logs.jsonl
vendored
Normal file
7
litellm-rust/crates/traces-clickhouse/tests/fixtures/deeplite_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/google_adk_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/langchain_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/langchain_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/langgraph_simple_spend_logs.jsonl
vendored
Normal file
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/langgraph_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/llamaindex_simple_spend_logs.jsonl
vendored
Normal file
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/llamaindex_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/openai_agents_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/openai_agents_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/opentelemetry_simple_spend_logs.jsonl
vendored
Normal file
2
litellm-rust/crates/traces-clickhouse/tests/fixtures/opentelemetry_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/pydantic_ai_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/strands_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/strands_swarm_spend_logs.jsonl
vendored
Normal file
1
litellm-rust/crates/traces-clickhouse/tests/fixtures/vercel_ai_sdk_simple_spend_logs.jsonl
vendored
Normal file
5
litellm-rust/crates/traces-clickhouse/tests/fixtures/vercel_ai_sdk_swarm_spend_logs.jsonl
vendored
Normal file
|
|
@ -101,10 +101,9 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
&reader,
|
||||
&litellm_traces_clickhouse::query::named::SpanDetailParams {
|
||||
access: litellm_traces_clickhouse::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
all_teams: false,
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-1".into()],
|
||||
api_key_hash: String::new(),
|
||||
},
|
||||
trace_id: "trace-1".into(),
|
||||
trace_ref: String::new(),
|
||||
|
|
@ -119,7 +118,6 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
|
|
@ -150,10 +148,11 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
"response_ids".into(),
|
||||
Parameter::Strings(vec!["response-1".into()]),
|
||||
),
|
||||
("request_ids".into(), Parameter::Strings(Vec::new())),
|
||||
("trace_ids".into(), Parameter::Strings(Vec::new())),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
|
|
@ -239,6 +238,44 @@ async fn normalized_fields_match_clickhouse_catalog(
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn agent_metadata_is_stored_and_queryable(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let ready = database?;
|
||||
ensure_schema(
|
||||
&ready.client,
|
||||
&Connection::writer(&ready.url)?,
|
||||
"trace_test",
|
||||
7,
|
||||
)
|
||||
.await?;
|
||||
let metadata = serde_json::json!({"thread_id": "thread-1", "ls_subagent_id": "agent-1"});
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
insert_rows(
|
||||
&ready,
|
||||
"otel_traces",
|
||||
vec![BTreeMap::from([
|
||||
("Timestamp".into(), timestamp.into()),
|
||||
("TraceId".into(), "trace-1".into()),
|
||||
("SpanId".into(), "span-1".into()),
|
||||
("AgentMetadata".into(), metadata.to_string().into()),
|
||||
])],
|
||||
)
|
||||
.await?;
|
||||
let response = read_json(
|
||||
&ready,
|
||||
"SELECT JSONExtractString(AgentMetadata, 'thread_id') AS thread_id, JSONExtractString(AgentMetadata, 'ls_subagent_id') AS subagent_id FROM trace_test.otel_traces WHERE TraceId = 'trace-1'",
|
||||
).await?;
|
||||
assert_eq!(response["data"][0]["thread_id"], metadata["thread_id"]);
|
||||
assert_eq!(
|
||||
response["data"][0]["subagent_id"],
|
||||
metadata["ls_subagent_id"]
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
|
||||
|
|
@ -421,6 +458,7 @@ async fn listed_agent_names_preserve_scope_and_cursor(
|
|||
vec![serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": trace, "SpanId": span, "ParentSpanId": parent,
|
||||
"ServiceName": "shared-app", "SpanName": span, "AgentName": agent,
|
||||
"UserId": if key == "one" { "owner" } else { "other" },
|
||||
"Framework": framework, "ObservationType": "agent",
|
||||
"ResourceAttributes": {"litellm.team_id": team, "litellm.api_key_hash": key}
|
||||
}))?],
|
||||
|
|
@ -442,9 +480,8 @@ async fn listed_agent_names_preserve_scope_and_cursor(
|
|||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("user_id".into(), Parameter::Text("owner".into())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text("one".into())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
|
|
@ -596,7 +633,6 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
|||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
Parameter::Integer(day_start / 1_000_000 - 2000),
|
||||
|
|
@ -795,12 +831,12 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
for (key, text) in [("one", "timeout"), ("two", "success")] {
|
||||
insert_rows(&database, "otel_traces", vec![serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "shared", "SpanId": "root", "ParentSpanId": "",
|
||||
"ServiceName": "review", "SpanName": "release", "Input": text,
|
||||
"ServiceName": "review", "SpanName": "release", "Input": text, "UserId": key,
|
||||
"ResourceAttributes": {"litellm.team_id": "team", "litellm.api_key_hash": key, "swarm": "release"}
|
||||
}))?]).await?;
|
||||
}
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
let sample_parameters = BTreeMap::from([
|
||||
("source".into(), Parameter::Text("traces".into())),
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("team".into(), Parameter::Text(String::new())),
|
||||
|
|
@ -837,7 +873,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
&sample_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
|
|
@ -849,7 +885,6 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
]);
|
||||
let identities: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
|
|
@ -861,19 +896,18 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
.await?,
|
||||
)?;
|
||||
assert_eq!(identities["data"].as_array().map(Vec::len), Some(2));
|
||||
let key_params = identity_params
|
||||
.into_iter()
|
||||
.chain([
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text("one".into())),
|
||||
])
|
||||
.collect();
|
||||
let user_params = BTreeMap::from([
|
||||
("trace_id".into(), Parameter::Text("shared".into())),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text("one".into())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let identity: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::TraceIdentity,
|
||||
&key_params,
|
||||
&user_params,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
|
|
@ -883,23 +917,23 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
.any(|row| row["trace_ref"] == identity["data"][0]["trace_ref"])
|
||||
);
|
||||
let first_ref = rows[0]["trace_ref"].as_str().expect("reference");
|
||||
let read_parameters: BTreeMap<_, _> = parameters
|
||||
.into_iter()
|
||||
.chain([
|
||||
("id".into(), Parameter::Text("shared".into())),
|
||||
("record_team".into(), Parameter::Text("team".into())),
|
||||
("trace_ref".into(), Parameter::Text(first_ref.into())),
|
||||
("cursor".into(), Parameter::Text(String::new())),
|
||||
("offset".into(), Parameter::Integer(1)),
|
||||
("span".into(), Parameter::Text("root".into())),
|
||||
])
|
||||
.collect();
|
||||
let content_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("team".into(), Parameter::Text(String::new())),
|
||||
("key_hash".into(), Parameter::Text(String::new())),
|
||||
("source".into(), Parameter::Text("traces".into())),
|
||||
("id".into(), Parameter::Text("shared".into())),
|
||||
("record_team".into(), Parameter::Text("team".into())),
|
||||
("trace_ref".into(), Parameter::Text(first_ref.into())),
|
||||
("cursor".into(), Parameter::Text(String::new())),
|
||||
("offset".into(), Parameter::Integer(1)),
|
||||
]);
|
||||
let content: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::Content,
|
||||
&read_parameters,
|
||||
&content_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
|
|
@ -910,10 +944,17 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
} else {
|
||||
"timeout"
|
||||
};
|
||||
let evidence_parameters = read_parameters
|
||||
.into_iter()
|
||||
.chain([("quote".into(), Parameter::Text(opposite.into()))])
|
||||
.collect();
|
||||
let evidence_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("team".into(), Parameter::Text(String::new())),
|
||||
("key_hash".into(), Parameter::Text(String::new())),
|
||||
("source".into(), Parameter::Text("traces".into())),
|
||||
("id".into(), Parameter::Text("shared".into())),
|
||||
("record_team".into(), Parameter::Text("team".into())),
|
||||
("trace_ref".into(), Parameter::Text(first_ref.into())),
|
||||
("span".into(), Parameter::Text("root".into())),
|
||||
("quote".into(), Parameter::Text(opposite.into())),
|
||||
]);
|
||||
let evidence: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
|
|
@ -1170,7 +1211,6 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
]);
|
||||
let body = execute_named_read(
|
||||
|
|
@ -1217,10 +1257,7 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
}
|
||||
assert_eq!(recovered, message);
|
||||
parameters.insert("all_teams".into(), Parameter::Integer(0));
|
||||
parameters.insert(
|
||||
"api_key_hash".into(),
|
||||
Parameter::Text("unrelated-key".into()),
|
||||
);
|
||||
parameters.insert("user_id".into(), Parameter::Text("unrelated-user".into()));
|
||||
let denied =
|
||||
execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters).await?;
|
||||
assert_eq!(
|
||||
|
|
@ -1265,7 +1302,6 @@ async fn duplicate_span_preview_matches_diagnostic(
|
|||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
("error_version".into(), Parameter::Text(String::new())),
|
||||
("error_offset".into(), Parameter::Integer(0)),
|
||||
|
|
@ -1327,7 +1363,7 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
.await?;
|
||||
}
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let scope_parameters = BTreeMap::from([
|
||||
let agent_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("team".into(), Parameter::Text("alpha".into())),
|
||||
("key_hash".into(), Parameter::Text("one".into())),
|
||||
|
|
@ -1337,7 +1373,7 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::Agents,
|
||||
&scope_parameters,
|
||||
&agent_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
|
|
@ -1347,53 +1383,58 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
{"agent_name": "research_agent"}, {"agent_name": "support_agent"}
|
||||
])
|
||||
);
|
||||
let parameters = scope_parameters
|
||||
.into_iter()
|
||||
.chain([
|
||||
("source".into(), Parameter::Text("traces".into())),
|
||||
(
|
||||
"start".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
),
|
||||
(
|
||||
"end".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 + 1000),
|
||||
),
|
||||
("service".into(), Parameter::Text("shared-app".into())),
|
||||
(
|
||||
"agent_name".into(),
|
||||
Parameter::Text("research_agent".into()),
|
||||
),
|
||||
("filter_keys".into(), Parameter::Strings(vec![])),
|
||||
("filter_values".into(), Parameter::Strings(vec![])),
|
||||
("limit".into(), Parameter::Integer(100)),
|
||||
("offset".into(), Parameter::Integer(0)),
|
||||
("after".into(), Parameter::Text(String::new())),
|
||||
("sample_percent".into(), Parameter::Text("100".into())),
|
||||
("sample_cap".into(), Parameter::Integer(0)),
|
||||
("preview".into(), Parameter::Integer(1)),
|
||||
("selected_team".into(), Parameter::Text(String::new())),
|
||||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
])
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let sample_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("team".into(), Parameter::Text("alpha".into())),
|
||||
("key_hash".into(), Parameter::Text("one".into())),
|
||||
("source".into(), Parameter::Text("traces".into())),
|
||||
(
|
||||
"start".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 - 1000),
|
||||
),
|
||||
(
|
||||
"end".into(),
|
||||
Parameter::Integer(timestamp / 1_000_000 + 1000),
|
||||
),
|
||||
("service".into(), Parameter::Text("shared-app".into())),
|
||||
(
|
||||
"agent_name".into(),
|
||||
Parameter::Text("research_agent".into()),
|
||||
),
|
||||
("filter_keys".into(), Parameter::Strings(vec![])),
|
||||
("filter_values".into(), Parameter::Strings(vec![])),
|
||||
("limit".into(), Parameter::Integer(100)),
|
||||
("offset".into(), Parameter::Integer(0)),
|
||||
("after".into(), Parameter::Text(String::new())),
|
||||
("sample_percent".into(), Parameter::Text("100".into())),
|
||||
("sample_cap".into(), Parameter::Integer(0)),
|
||||
("preview".into(), Parameter::Integer(1)),
|
||||
("selected_team".into(), Parameter::Text(String::new())),
|
||||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
&sample_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(sample["data"].as_array().expect("rows").len(), 1);
|
||||
assert_eq!(sample["data"][0]["trace_id"], "research");
|
||||
assert_eq!(sample["data"][0]["span_count"], 2);
|
||||
let availability_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("team".into(), Parameter::Text("alpha".into())),
|
||||
("key_hash".into(), Parameter::Text("one".into())),
|
||||
]);
|
||||
let available: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::Availability,
|
||||
¶meters,
|
||||
&availability_parameters,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
|
|
@ -1437,7 +1478,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
.await?;
|
||||
let metadata = serde_json::json!({
|
||||
"project": "example", "labels": {"priority": 3, "enabled": true},
|
||||
"dotted.key": "literal", "quote'\\key": null, "items": [{"name": "first"}],
|
||||
"dotted.key": "private-metadata-value", "quote'\\key": null, "items": [{"name": "first"}],
|
||||
"<custom>&{{key}}": {"nested.key": true}
|
||||
});
|
||||
insert_rows(
|
||||
|
|
@ -1445,7 +1486,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
"spend_logs",
|
||||
vec![serde_json::from_value(serde_json::json!({
|
||||
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
|
||||
"api_key": "key-1", "metadata": metadata.to_string(), "spend": 0.25,
|
||||
"api_key": "key-1", "trace_id": "trace-1", "metadata": metadata.to_string(), "spend": 0.25,
|
||||
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100
|
||||
}))?],
|
||||
)
|
||||
|
|
@ -1467,8 +1508,8 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
)
|
||||
.await?;
|
||||
}
|
||||
let help: serde_json::Value = serde_json::from_str(
|
||||
&litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
let help = serde_json::to_value(
|
||||
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
let keys: std::collections::BTreeSet<_> = help
|
||||
.as_object()
|
||||
|
|
@ -1511,9 +1552,31 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
)));
|
||||
}
|
||||
}
|
||||
for gotcha in help["gotchas"].as_array().ok_or("missing gotchas")? {
|
||||
assert!(guide.contains(gotcha.as_str().ok_or("gotcha text")?));
|
||||
let gotchas = help["gotchas"].as_array().ok_or("missing gotchas")?;
|
||||
let gotcha_positions = gotchas
|
||||
.iter()
|
||||
.map(|gotcha| {
|
||||
guide
|
||||
.find(gotcha.as_str().expect("gotcha text"))
|
||||
.expect("rendered gotcha")
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert!(gotcha_positions.windows(2).all(|pair| pair[0] < pair[1]));
|
||||
assert!(
|
||||
guide.contains(
|
||||
help["metadata"]["sample_sql"]
|
||||
.as_str()
|
||||
.ok_or("sampling SQL")?
|
||||
)
|
||||
);
|
||||
for catalog in help["attributes"].as_array().ok_or("attributes")? {
|
||||
assert!(guide.contains(catalog["discovery_sql"].as_str().ok_or("discovery SQL")?));
|
||||
assert!(guide.contains(&format!("Truncated: {}", catalog["truncated"])));
|
||||
}
|
||||
assert_eq!(
|
||||
guide.contains("No attribute keys found in the sampled spans"),
|
||||
!populated
|
||||
);
|
||||
let tables = help["tables"].as_array().ok_or("missing tables")?;
|
||||
assert_eq!(tables.len(), 3);
|
||||
let columns = tables[0]["columns"].as_array().ok_or("missing columns")?;
|
||||
|
|
@ -1567,6 +1630,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
.any(|field| field["path"] == serde_json::json!(["items", 1, "name"]))
|
||||
);
|
||||
assert!(guide.contains("CustomColumn: String"));
|
||||
assert!(!guide.contains("private-metadata-value"));
|
||||
assert!(guide.contains("JSONExtractRaw(metadata, '<custom>&{{key}}', 'nested.key')"));
|
||||
assert!(guide.contains("SpanAttributes['custom.tag']"));
|
||||
assert!(guide.contains("ResourceAttributes['custom.resource']"));
|
||||
|
|
@ -1585,7 +1649,24 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
assert_ne!(values["data"][0]["value"], "");
|
||||
}
|
||||
}
|
||||
for example in help["examples"].as_array().ok_or("missing examples")? {
|
||||
let examples = help["examples"].as_array().ok_or("missing examples")?;
|
||||
let example_positions = examples
|
||||
.iter()
|
||||
.map(|example| {
|
||||
let rendered = format!(
|
||||
"{}\n{}",
|
||||
example["name"].as_str().expect("name"),
|
||||
example["sql"].as_str().expect("SQL")
|
||||
);
|
||||
guide.find(&rendered).expect("rendered example")
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
assert!(example_positions.windows(2).all(|pair| pair[0] < pair[1]));
|
||||
assert!(
|
||||
example_positions.last().ok_or("last example")?
|
||||
< gotcha_positions.first().ok_or("first gotcha")?
|
||||
);
|
||||
for example in examples {
|
||||
let sql = example["sql"].as_str().ok_or("missing example SQL")?;
|
||||
assert!(guide.contains(example["name"].as_str().ok_or("missing example name")?));
|
||||
assert!(guide.contains(sql));
|
||||
|
|
@ -1602,7 +1683,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(
|
||||
values["data"].as_array().ok_or("missing data")?.is_empty(),
|
||||
!populated,
|
||||
!populated || example["name"] == "LLM spans without a direct spend match",
|
||||
"{sql}"
|
||||
);
|
||||
if populated && example["name"] == "Traces correlated with LLM call metadata" {
|
||||
|
|
@ -1654,8 +1735,8 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
}))).collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", spans).await?;
|
||||
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
|
||||
let help: serde_json::Value = serde_json::from_str(
|
||||
&litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
let help = serde_json::to_value(
|
||||
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3);
|
||||
assert!(!help["examples"].as_array().ok_or("examples")?.is_empty());
|
||||
|
|
@ -1676,6 +1757,18 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
guide.contains("Attribute discovery unavailable:"),
|
||||
span_rows > 1
|
||||
);
|
||||
assert!(!guide.contains("No metadata paths found in the sampled rows"));
|
||||
assert!(!guide.contains("No attribute keys found in the sampled spans"));
|
||||
assert!(
|
||||
guide.contains(
|
||||
help["metadata"]["sample_sql"]
|
||||
.as_str()
|
||||
.ok_or("sampling SQL")?
|
||||
)
|
||||
);
|
||||
for catalog in help["attributes"].as_array().ok_or("attributes")? {
|
||||
assert!(guide.contains(catalog["discovery_sql"].as_str().ok_or("discovery SQL")?));
|
||||
}
|
||||
for (catalog, unavailable) in [
|
||||
(&help["metadata"], spend_rows > 1),
|
||||
(&help["attributes"][0], span_rows > 1),
|
||||
|
|
@ -1691,41 +1784,99 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn query_help_displays_discovery_truncation(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
execute_write(
|
||||
&database,
|
||||
"INSERT INTO trace_test.otel_traces (Timestamp, TraceId, SpanId, SpanAttributes, ResourceAttributes) \
|
||||
SELECT now64(9), 'trace', 'span', \
|
||||
mapFromArrays(arrayMap(x -> concat('key-', toString(x)), range(1000)), arrayMap(x -> 'value', range(1000))) AS attributes, \
|
||||
attributes FROM numbers(1)",
|
||||
)
|
||||
.await?;
|
||||
execute_write(
|
||||
&database,
|
||||
"INSERT INTO trace_test.spend_logs (request_id, start_time, end_time, metadata) \
|
||||
SELECT toString(number), now64(3), now64(3), '{\"key\":true}' FROM numbers(1000)",
|
||||
)
|
||||
.await?;
|
||||
let reader = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let help = serde_json::to_value(
|
||||
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
let guide = help["guide"].as_str().ok_or("guide")?;
|
||||
assert_eq!(help["metadata"]["truncated"], true);
|
||||
assert!(guide.contains("truncated: true"));
|
||||
for catalog in help["attributes"].as_array().ok_or("attributes")? {
|
||||
assert_eq!(catalog["truncated"], true);
|
||||
let displayed = format!(
|
||||
"{}.{}",
|
||||
catalog["table"].as_str().ok_or("table")?,
|
||||
catalog["column"].as_str().ok_or("column")?
|
||||
);
|
||||
let section = guide.split(&displayed).nth(1).ok_or("attribute section")?;
|
||||
assert!(
|
||||
section
|
||||
.split("\n\n")
|
||||
.next()
|
||||
.ok_or("catalog body")?
|
||||
.contains("Truncated: true")
|
||||
);
|
||||
for field in catalog["fields"].as_array().ok_or("fields")? {
|
||||
assert!(section.contains(field["expression"].as_str().ok_or("expression")?));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn field_definitions_match_serialized_normalized_span() {
|
||||
use litellm_traces::decode_otlp;
|
||||
use litellm_traces::{Tenant, decode_otlp};
|
||||
use litellm_traces_clickhouse::span_rows;
|
||||
use std::collections::BTreeSet;
|
||||
let spans = decode_otlp(
|
||||
br#"{"resourceSpans":[{"scopeSpans":[{"spans":[{"traceId":"11111111111111111111111111111111","spanId":"2222222222222222","name":"root"}]}]}]}"#,
|
||||
Some("application/json"),
|
||||
)
|
||||
.expect("valid OTLP");
|
||||
let fields = &spans[0].normalized;
|
||||
let serialized = serde_json::to_value(fields).expect("serializable fields");
|
||||
let keys: BTreeSet<_> = serialized
|
||||
let tenant = Tenant {
|
||||
team_id: "team".into(),
|
||||
api_key_hash: "key".into(),
|
||||
..Tenant::default()
|
||||
};
|
||||
let rows = span_rows(spans, &tenant, 64 * 1024);
|
||||
let row =
|
||||
serde_json::to_value(rows.first().expect("storage row")).expect("serializable storage row");
|
||||
let keys: BTreeSet<_> = row
|
||||
.as_object()
|
||||
.expect("field object")
|
||||
.expect("storage row object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(|field| field.name)
|
||||
.map(|field| field.clickhouse_column)
|
||||
.collect();
|
||||
assert_eq!(keys, mapped);
|
||||
assert!(mapped.is_subset(&keys));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user("owner", vec![], "", vec!["own"])]
|
||||
#[case::own_user_and_permitted_team("owner", vec!["permitted"], "", vec!["own", "team"])]
|
||||
#[case::key_only("", vec![], "request-key", vec!["own"])]
|
||||
#[case::no_identity("", vec![], "", vec![])]
|
||||
#[case::own_user("owner", vec![], None, vec!["own"])]
|
||||
#[case::own_user_and_permitted_team("owner", vec!["permitted"], None, vec!["own", "team"])]
|
||||
#[case::no_identity("", vec![], None, vec![])]
|
||||
#[case::legacy_key_without_identity("", vec![], Some("request-key"), vec![])]
|
||||
#[tokio::test]
|
||||
async fn named_and_sql_readers_share_request_log_visibility(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] key: &str,
|
||||
#[case] legacy_key: Option<&str>,
|
||||
#[case] expected: Vec<&str>,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{
|
||||
|
|
@ -1748,13 +1899,13 @@ async fn named_and_sql_readers_share_request_log_visibility(
|
|||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let params =
|
||||
SpendByResponseIdsParams::from(litellm_traces::query::named::SpendByResponseIdsParams {
|
||||
access: ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: user.into(),
|
||||
team_ids: teams.iter().map(|team| (*team).into()).collect(),
|
||||
api_key_hash: key.into(),
|
||||
},
|
||||
access: serde_json::from_value::<ReadAccessParams>(serde_json::json!({
|
||||
"all_teams": 0, "user_id": user, "team_ids": teams,
|
||||
"api_key_hash": legacy_key.unwrap_or_default(),
|
||||
}))?,
|
||||
response_ids: vec!["shared-response".into()],
|
||||
request_ids: Vec::new(),
|
||||
trace_ids: Vec::new(),
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
});
|
||||
|
|
@ -1765,10 +1916,9 @@ async fn named_and_sql_readers_share_request_log_visibility(
|
|||
spend.iter().map(|row| row.0.request_id.as_str()).collect();
|
||||
let expected: std::collections::BTreeSet<_> = expected.into_iter().collect();
|
||||
assert_eq!(actual, expected);
|
||||
let scope = QueryScope::Logs {
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
api_key_hash: key.into(),
|
||||
};
|
||||
if user.is_empty() && scope.validate().is_err() {
|
||||
assert!(
|
||||
|
|
@ -1826,10 +1976,9 @@ async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_his
|
|||
let params = litellm_traces_clickhouse::query::named::ListTracesParams::from(
|
||||
litellm_traces::query::named::ListTracesParams {
|
||||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
all_teams: false,
|
||||
user_id: "".into(),
|
||||
team_ids: vec!["team".into()],
|
||||
api_key_hash: "".into(),
|
||||
},
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
|
|
@ -1849,8 +1998,7 @@ async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_his
|
|||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
user_id: "owner".into(),
|
||||
team_ids: vec![],
|
||||
all_teams: 0,
|
||||
api_key_hash: String::new(),
|
||||
all_teams: false,
|
||||
},
|
||||
..params.0
|
||||
},
|
||||
|
|
@ -1882,18 +2030,16 @@ async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_his
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::admin(1, "", vec![], "", "own answer")]
|
||||
#[case::user(0, "owner", vec![], "", "own answer")]
|
||||
#[case::team(0, "", vec!["alpha"], "", "own answer")]
|
||||
#[case::key(0, "", vec![], "one", "own answer")]
|
||||
#[case::no_identity(0, "", vec![], "", "")]
|
||||
#[case::admin(1, "", vec![], "own answer")]
|
||||
#[case::user(0, "owner", vec![], "own answer")]
|
||||
#[case::team(0, "", vec!["alpha"], "own answer")]
|
||||
#[case::no_identity(0, "", vec![], "")]
|
||||
#[tokio::test]
|
||||
async fn agent_final_answer_preserves_visibility_and_trace_ownership(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] all_teams: u8,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] key: &str,
|
||||
#[case] expected: &str,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{ReadAccessParams, SpanDetail, SpanDetailParams};
|
||||
|
|
@ -1949,10 +2095,9 @@ async fn agent_final_answer_preserves_visibility_and_trace_ownership(
|
|||
&reader,
|
||||
&SpanDetailParams {
|
||||
access: ReadAccessParams {
|
||||
all_teams,
|
||||
all_teams: all_teams == 1,
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
api_key_hash: key.into(),
|
||||
},
|
||||
trace_id: "shared".into(),
|
||||
trace_ref: String::new(),
|
||||
|
|
@ -1969,3 +2114,55 @@ async fn agent_final_answer_preserves_visibility_and_trace_ownership(
|
|||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn nullable_spend_upgrade_preserves_existing_costs_and_unknown_new_costs(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
let timestamp = (time::OffsetDateTime::now_utc().unix_timestamp_nanos() / 1_000_000) as i64;
|
||||
let statements = schema_statements("trace_test", 7)?;
|
||||
for statement in &statements[..statements.len() - 1] {
|
||||
execute_write(&database, statement).await?;
|
||||
}
|
||||
let legacy = serde_json::from_value(serde_json::json!({
|
||||
"request_id": "legacy", "response_id": "legacy-response", "spend": 0.25,
|
||||
"start_time": timestamp, "end_time": timestamp + 100
|
||||
}))?;
|
||||
insert_rows(&database, "spend_logs", vec![legacy]).await?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let unknown = serde_json::from_value(serde_json::json!({
|
||||
"request_id": "unknown", "response_id": "unknown-response", "spend": null,
|
||||
"start_time": timestamp, "end_time": timestamp + 100
|
||||
}))?;
|
||||
let free = serde_json::from_value(serde_json::json!({
|
||||
"request_id": "free", "response_id": "free-response", "spend": 0.0,
|
||||
"start_time": timestamp, "end_time": timestamp + 100
|
||||
}))?;
|
||||
insert_rows(&database, "spend_logs", vec![unknown, free]).await?;
|
||||
let result = read_json(
|
||||
&database,
|
||||
"SELECT request_id, spend FROM trace_test.spend_logs FINAL ORDER BY request_id",
|
||||
)
|
||||
.await?;
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct CostRow {
|
||||
request_id: String,
|
||||
spend: Option<f64>,
|
||||
}
|
||||
let rows: Vec<CostRow> = serde_json::from_value(result["data"].clone())?;
|
||||
assert_eq!(
|
||||
rows.iter()
|
||||
.map(|row| (row.request_id.as_str(), row.spend))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![
|
||||
("free", Some(0.0)),
|
||||
("legacy", Some(0.25)),
|
||||
("unknown", None)
|
||||
]
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -23,23 +23,20 @@ use support::TestResult;
|
|||
enum ScopeCase {
|
||||
Admin,
|
||||
Team,
|
||||
Key,
|
||||
OtherTeam,
|
||||
}
|
||||
|
||||
impl ScopeCase {
|
||||
fn scope(self) -> QueryScope {
|
||||
match self {
|
||||
Self::Admin => QueryScope::Admin,
|
||||
Self::Team => QueryScope::Team {
|
||||
team_id: "team-a".into(),
|
||||
Self::Admin => QueryScope::All,
|
||||
Self::Team => QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".into()],
|
||||
},
|
||||
Self::Key => QueryScope::Key {
|
||||
team_id: "team-a".into(),
|
||||
api_key_hash: "key-a".into(),
|
||||
},
|
||||
Self::OtherTeam => QueryScope::Team {
|
||||
team_id: "team-b".into(),
|
||||
Self::OtherTeam => QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-b".into()],
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
@ -60,13 +57,7 @@ async fn curated_queries_return_expected_rows(
|
|||
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
|
||||
#[case] sql: &str,
|
||||
#[case] expected_json: &str,
|
||||
#[values(
|
||||
ScopeCase::Admin,
|
||||
ScopeCase::Team,
|
||||
ScopeCase::Key,
|
||||
ScopeCase::OtherTeam
|
||||
)]
|
||||
scope: ScopeCase,
|
||||
#[values(ScopeCase::Admin, ScopeCase::Team, ScopeCase::OtherTeam)] scope: ScopeCase,
|
||||
) -> TestResult {
|
||||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
|
|
@ -103,11 +94,7 @@ async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate(
|
|||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.connection(&fixture.database.client, &QueryScope::All, "fixture-secret")
|
||||
.await?;
|
||||
let params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: admin_access?,
|
||||
|
|
@ -157,11 +144,11 @@ async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate(
|
|||
);
|
||||
assert_eq!(
|
||||
(
|
||||
spans[1].0.kind.as_str(),
|
||||
spans[1].0.kind,
|
||||
spans[1].0.input_tokens,
|
||||
spans[1].0.output_tokens
|
||||
),
|
||||
("llm", 12, 6)
|
||||
(litellm_traces::ObservationType::Llm, 12, 6)
|
||||
);
|
||||
assert_eq!(spans[2].0.status_message, "lookup timed out");
|
||||
Ok(())
|
||||
|
|
@ -176,11 +163,7 @@ async fn typed_trace_cursor_returns_the_next_fixture_trace(
|
|||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.connection(&fixture.database.client, &QueryScope::All, "fixture-secret")
|
||||
.await?;
|
||||
let params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: admin_access?,
|
||||
|
|
@ -218,11 +201,7 @@ async fn captured_deeplite_exports_round_trip_through_clickhouse(
|
|||
let decoded = insert_export(&fixture, export, "team-a", "key-a").await?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.connection(&fixture.database.client, &QueryScope::All, "fixture-secret")
|
||||
.await?;
|
||||
let params = TraceSpansParams {
|
||||
access: admin_access?,
|
||||
|
|
@ -246,7 +225,13 @@ async fn captured_deeplite_exports_round_trip_through_clickhouse(
|
|||
.filter(|span| span.parent_span_id.is_empty())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(roots.len(), 1);
|
||||
assert_eq!(traces[0].0.status, roots[0].status_code);
|
||||
assert_eq!(
|
||||
traces[0].0.status,
|
||||
serde_json::from_value::<litellm_traces::SpanStatus>(serde_json::json!(
|
||||
roots[0].status_code
|
||||
))
|
||||
.unwrap()
|
||||
);
|
||||
assert_eq!(
|
||||
traces[0].0.error_count,
|
||||
decoded
|
||||
|
|
@ -267,7 +252,13 @@ async fn captured_deeplite_exports_round_trip_through_clickhouse(
|
|||
assert_eq!(row.duration_ns, span.end_ns - span.start_ns);
|
||||
assert_eq!(row.input_tokens, span.normalized.input_tokens);
|
||||
assert_eq!(row.output_tokens, span.normalized.output_tokens);
|
||||
assert_eq!(row.status, span.status_code);
|
||||
assert_eq!(
|
||||
row.status,
|
||||
serde_json::from_value::<litellm_traces::SpanStatus>(serde_json::json!(
|
||||
span.status_code
|
||||
))
|
||||
.unwrap()
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
{
|
||||
"all_teams": 1,
|
||||
"user_id": "",
|
||||
"team_ids": ["team-a", "team-b"],
|
||||
"api_key_hash": ""
|
||||
"team_ids": [
|
||||
"team-a",
|
||||
"team-b"
|
||||
]
|
||||
}
|
||||
|
|
|
|||
|
|
@ -139,11 +139,29 @@ fn span_row(span: &DecodedSpan, team: &str, key: &str) -> BTreeMap<String, Value
|
|||
"ObservationType".into(),
|
||||
json!(span.normalized.observation_type),
|
||||
),
|
||||
("AgentName".into(), json!(span.normalized.agent_name)),
|
||||
("Model".into(), json!(span.normalized.model)),
|
||||
(
|
||||
"AgentName".into(),
|
||||
json!(span.normalized.agent_name.as_deref().unwrap_or_default()),
|
||||
),
|
||||
(
|
||||
"Model".into(),
|
||||
json!(span.normalized.model.as_deref().unwrap_or_default()),
|
||||
),
|
||||
(
|
||||
"LiteLLMRequestId".into(),
|
||||
json!(span.normalized.litellm_request_id),
|
||||
json!(
|
||||
span.normalized
|
||||
.calls
|
||||
.key_set()
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find_map(|key| match key {
|
||||
litellm_traces::CallKey::LiteLlmRequest(id)
|
||||
| litellm_traces::CallKey::ProviderResponse(id) => Some(id.as_str()),
|
||||
litellm_traces::CallKey::Transport => None,
|
||||
})
|
||||
.unwrap_or_default()
|
||||
),
|
||||
),
|
||||
("InputTokens".into(), json!(span.normalized.input_tokens)),
|
||||
("OutputTokens".into(), json!(span.normalized.output_tokens)),
|
||||
|
|
|
|||
|
|
@ -25,8 +25,8 @@ async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
|||
let writer = Connection::parse(&url)?;
|
||||
ensure_schema(&client, &writer, "trace_test", 7).await?;
|
||||
for sql in [
|
||||
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes, UserId) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a'), 'owner'), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a'), 'other'), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'), 'owner'), ('', 'key-teamless', 'shared-trace', 'teamless', now(), map('visible', 'teamless'), ''), ('', 'key-other', 'shared-trace', 'other-teamless', now(), map('visible', 'other'), '')",
|
||||
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata, user) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}', 'owner'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}', 'other'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}', 'owner'), ('', 'key-teamless', 'teamless', now(), now(), '{}', ''), ('', 'key-other', 'other-teamless', now(), now(), '{}', '')",
|
||||
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes, UserId) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a'), 'owner'), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a'), 'other'), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'), 'owner'), ('team-c', 'key-a1', 'shared-trace', 'same-key-foreign', now(), map('visible', 'foreign'), 'other'), ('', 'key-teamless', 'shared-trace', 'teamless', now(), map('visible', 'teamless'), ''), ('', 'key-other', 'shared-trace', 'other-teamless', now(), map('visible', 'other'), '')",
|
||||
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata, user) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}', 'owner'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}', 'other'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}', 'owner'), ('team-c', 'key-a1', 'same-key-foreign', now(), now(), '{}', 'other'), ('', 'key-teamless', 'teamless', now(), now(), '{}', ''), ('', 'key-other', 'other-teamless', now(), now(), '{}', '')",
|
||||
"CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory",
|
||||
"INSERT INTO trace_test.private_data VALUES ('hidden')",
|
||||
] {
|
||||
|
|
@ -43,15 +43,12 @@ async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
|||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user(QueryScope::Logs { user_id: "owner".into(), team_ids: vec![], api_key_hash: "".into() }, vec!["a1", "b"])]
|
||||
#[case::own_user_and_permitted_team(QueryScope::Logs { user_id: "owner".into(), team_ids: vec!["team-a".into()], api_key_hash: "".into() }, vec!["a1", "a2", "b"])]
|
||||
#[case::key_only_logs(QueryScope::Logs { user_id: "".into(), team_ids: vec![], api_key_hash: "key-teamless".into() }, vec!["teamless"])]
|
||||
#[case::quoted_user(QueryScope::Logs { user_id: "owner' OR 1=1 --".into(), team_ids: vec![], api_key_hash: "".into() }, vec![])]
|
||||
#[case::team(QueryScope::Team { team_id: "team-a".to_owned() }, vec!["a1", "a2"])]
|
||||
#[case::project_key(QueryScope::Key { team_id: "team-a".to_owned(), api_key_hash: "key-a1".to_owned() }, vec!["a1"])]
|
||||
#[case::teamless_key(QueryScope::Key { team_id: "".to_owned(), api_key_hash: "key-teamless".to_owned() }, vec!["teamless"])]
|
||||
#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b", "other-teamless", "teamless"])]
|
||||
#[case::quoted_team(QueryScope::Team { team_id: "team-a' OR 1=1 --\\".to_owned() }, vec![])]
|
||||
#[case::own_user(QueryScope::Owned { user_id: "owner".into(), team_ids: vec![] }, vec!["a1", "b"])]
|
||||
#[case::own_user_and_permitted_team(QueryScope::Owned { user_id: "owner".into(), team_ids: vec!["team-a".into()] }, vec!["a1", "a2", "b"])]
|
||||
#[case::quoted_user(QueryScope::Owned { user_id: "owner' OR 1=1 --".into(), team_ids: vec![] }, vec![])]
|
||||
#[case::team(QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a".to_owned() ] }, vec!["a1", "a2"])]
|
||||
#[case::admin(QueryScope::All, vec!["a1", "a2", "b", "other-teamless", "same-key-foreign", "teamless"])]
|
||||
#[case::quoted_team(QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a' OR 1=1 --\\".to_owned() ] }, vec![])]
|
||||
#[tokio::test]
|
||||
async fn queries_and_help_are_scoped_by_the_database(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
|
|
@ -94,7 +91,7 @@ async fn queries_and_help_are_scoped_by_the_database(
|
|||
.await?,
|
||||
)?;
|
||||
assert_eq!(summary["data"][0]["count"], json!(expected.len()));
|
||||
let help = query_help(&database.client, &reader).await?;
|
||||
let help = serde_json::to_string(&query_help(&database.client, &reader).await?)?;
|
||||
assert_eq!(help.contains("secret_b"), expected.contains(&"b"));
|
||||
assert_eq!(help.contains("secret-b"), expected.contains(&"b"));
|
||||
let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned());
|
||||
|
|
@ -111,8 +108,9 @@ async fn rotating_master_secret_revokes_previous_reader_credentials(
|
|||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
};
|
||||
let old_reader = database
|
||||
.readers
|
||||
|
|
@ -159,8 +157,9 @@ async fn managed_reader_rejects_privilege_and_scope_bypasses(
|
|||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
let scope = QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
};
|
||||
let reader = database
|
||||
.readers
|
||||
|
|
@ -213,14 +212,15 @@ async fn provisioning_failure_never_returns_a_writer_connection(
|
|||
let database = database?;
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::Admin, "test-master-secret")
|
||||
.connection(&database.client, &QueryScope::All, "test-master-secret")
|
||||
.await?;
|
||||
let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned());
|
||||
let result = no_provision_privileges
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-a".to_owned()],
|
||||
},
|
||||
"other-secret",
|
||||
)
|
||||
|
|
@ -232,7 +232,7 @@ async fn provisioning_failure_never_returns_a_writer_connection(
|
|||
assert!(matches!(
|
||||
database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::Admin, "")
|
||||
.connection(&database.client, &QueryScope::All, "")
|
||||
.await,
|
||||
Err(Error::MissingSecret)
|
||||
));
|
||||
|
|
@ -241,8 +241,9 @@ async fn provisioning_failure_never_returns_a_writer_connection(
|
|||
.readers
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Team {
|
||||
team_id: String::new()
|
||||
&QueryScope::Owned {
|
||||
user_id: String::new(),
|
||||
team_ids: vec![String::new()]
|
||||
},
|
||||
"test-master-secret"
|
||||
)
|
||||
|
|
@ -263,6 +264,6 @@ async fn provisioning_failure_never_returns_a_writer_connection(
|
|||
)
|
||||
.await?;
|
||||
let rows: Value = serde_json::from_str(&rows)?;
|
||||
assert_eq!(rows["data"][0]["count"], 5);
|
||||
assert_eq!(rows["data"][0]["count"], 6);
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||