From 180b5ee4820eac2c12d72cbbd80a5c8ae18f85fc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 4 Oct 2026 15:39:45 -0700 Subject: [PATCH] refactor(traces): read traces through a five-method storage port Trace storage now answers five storage-neutral reads: runs, run_counts, spans, span_text and calls. Query and row types live in litellm_traces::store with no ClickHouse encodings, and traces-clickhouse only adapts them, so another backend implements five queries instead of one SQL file per UI widget The runs list, histogram, field values, trace graph, span detail, error paging and cost lookup all go through the port. Paging, the graph budget, excerpting and the agent output fallback moved out of SQL into TraceReader. QueryScope is the only access wire, and every SQL file reads owned_spans, owned_runs and owned_calls CTEs built from one ownership rule per table that the console row policy also renders Cursors are tagged with their kind, and run search terms reach storage as raw text and globs. This also carries the run histogram and values endpoints, the gateway-traces router and the agent label rollup --- litellm-rust/Cargo.lock | 16 + litellm-rust/crates/gateway-traces/Cargo.toml | 20 + .../crates/gateway-traces/src/error.rs | 87 +++ litellm-rust/crates/gateway-traces/src/lib.rs | 24 + .../crates/gateway-traces/src/runs.rs | 126 ++++ .../crates/gateway-traces/tests/routes.rs | 283 ++++++++ .../crates/python-bridge/src/routes/traces.rs | 93 ++- litellm-rust/crates/traces-cache/src/cache.rs | 36 +- .../crates/traces-cache/src/cursor.rs | 254 ++++--- litellm-rust/crates/traces-cache/src/lib.rs | 5 +- litellm-rust/crates/traces-cache/src/list.rs | 42 +- litellm-rust/crates/traces-cache/src/pages.rs | 184 +++++ .../crates/traces-cache/src/reader.rs | 327 ++++++--- litellm-rust/crates/traces-cache/src/spend.rs | 44 +- litellm-rust/crates/traces-cache/src/store.rs | 66 +- .../crates/traces-cache/tests/read.rs | 593 ++++++++++++---- .../crates/traces-cache/tests/snapshots.rs | 51 +- .../0016_trace_rollup_agent_labels.sql | 2 + .../0017_trace_rollup_agent_labels_mv.sql | 23 + .../query/{spend_batch.sql => calls.sql} | 8 +- .../traces-clickhouse/query/list_traces.sql | 48 -- .../traces-clickhouse/query/matching_runs.sql | 33 + .../traces-clickhouse/query/run_counts.sql | 20 + .../traces-clickhouse/query/run_spans.sql | 3 + .../crates/traces-clickhouse/query/runs.sql | 26 + .../traces-clickhouse/query/span_columns.sql | 16 + .../traces-clickhouse/query/span_detail.sql | 24 - .../traces-clickhouse/query/span_error.sql | 14 - .../traces-clickhouse/query/span_page.sql | 7 + .../traces-clickhouse/query/span_text.sql | 15 + .../query/spend_by_response_ids.sql | 22 - .../query/trace_identity.sql | 8 - .../query/trace_list_span_batch.sql | 31 - .../query/trace_page_spans.sql | 24 - .../query/trace_span_batch.sql | 30 - .../traces-clickhouse/query/trace_spans.sql | 26 +- .../crates/traces-clickhouse/src/access.rs | 93 +++ .../crates/traces-clickhouse/src/config.rs | 3 +- .../crates/traces-clickhouse/src/insert.rs | 8 +- .../crates/traces-clickhouse/src/lib.rs | 2 +- .../crates/traces-clickhouse/src/query.rs | 5 +- .../traces-clickhouse/src/query/named.rs | 636 +++++++++++------ .../traces-clickhouse/src/query/number.rs | 7 +- .../traces-clickhouse/src/query_access.rs | 66 +- .../crates/traces-clickhouse/src/reads.rs | 139 ++-- .../crates/traces-clickhouse/src/schema.rs | 3 +- .../traces-clickhouse/src/span_batches.rs | 352 ---------- .../crates/traces-clickhouse/src/sql.rs | 34 +- .../traces-clickhouse/src/wire_schema.rs | 3 +- .../traces-clickhouse/tests/admin_sql.rs | 3 +- .../traces-clickhouse/tests/migrations.rs | 640 ++++++++---------- .../crates/traces-clickhouse/tests/queries.rs | 134 ++-- .../tests/queries/read_access.json | 8 - .../crates/traces-clickhouse/tests/reads.rs | 123 +++- .../crates/traces-clickhouse/tests/search.rs | 421 ++++++++++++ litellm-rust/crates/traces/src/lib.rs | 2 + .../normalize/instrumentation/claude_code.rs | 3 +- .../normalize/instrumentation/http_client.rs | 5 +- .../normalize/instrumentation/langchain.rs | 3 +- .../normalize/instrumentation/llama_index.rs | 3 +- .../normalize/instrumentation/pydantic_ai.rs | 3 +- .../crates/traces/src/normalize/messages.rs | 8 +- .../crates/traces/src/normalize/mod.rs | 7 +- litellm-rust/crates/traces/src/otlp/mod.rs | 6 +- litellm-rust/crates/traces/src/query.rs | 8 - litellm-rust/crates/traces/src/query/named.rs | 203 ------ .../crates/traces/src/query_access.rs | 2 +- .../crates/traces/src/resolve/graph.rs | 14 +- .../crates/traces/src/resolve/resolution.rs | 15 +- .../crates/traces/src/resolve/spend.rs | 8 +- .../crates/traces/src/resolve/view.rs | 17 +- litellm-rust/crates/traces/src/schema.rs | 7 +- litellm-rust/crates/traces/src/search.rs | 217 ++++++ litellm-rust/crates/traces/src/store.rs | 282 ++++++++ litellm-rust/crates/traces/tests/captures.rs | 31 +- litellm-rust/crates/traces/tests/otlp.rs | 3 +- litellm-rust/crates/traces/tests/query.rs | 13 +- .../crates/traces/tests/query/named.rs | 67 -- litellm-rust/crates/traces/tests/resolve.rs | 102 ++- litellm-rust/crates/traces/tests/search.rs | 92 +++ litellm/proxy/lens/endpoints.py | 27 +- litellm/proxy/lens/repository.py | 23 + litellm/proxy/lens/search.py | 104 +++ litellm/proxy/tracing_endpoints.py | 92 ++- litellm/rust_bridge/_native.pyi | 16 +- litellm/rust_bridge/trace/generated/types.py | 30 +- litellm/rust_bridge/trace/storage.py | 44 +- litellm/tracing/receiver.py | 35 +- .../schemas/traces/RunField.json | 13 + .../schemas/traces/RunValues.json | 17 + .../schemas/traces/TraceHistogram.json | 80 +++ .../schemas/traces/TraceScope.json | 28 - .../database/test_lens_repository.py | 106 ++- tests/unit/proxy/lens/test_search.py | 28 + tests/unit/proxy/test_tracing_endpoints.py | 118 +++- 95 files changed, 4770 insertions(+), 2423 deletions(-) create mode 100644 litellm-rust/crates/gateway-traces/Cargo.toml create mode 100644 litellm-rust/crates/gateway-traces/src/error.rs create mode 100644 litellm-rust/crates/gateway-traces/src/lib.rs create mode 100644 litellm-rust/crates/gateway-traces/src/runs.rs create mode 100644 litellm-rust/crates/gateway-traces/tests/routes.rs create mode 100644 litellm-rust/crates/traces-cache/src/pages.rs create mode 100644 litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql create mode 100644 litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql rename litellm-rust/crates/traces-clickhouse/query/{spend_batch.sql => calls.sql} (83%) delete mode 100644 litellm-rust/crates/traces-clickhouse/query/list_traces.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/matching_runs.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/run_counts.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/run_spans.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/runs.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/span_columns.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/span_detail.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/span_error.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/span_page.sql create mode 100644 litellm-rust/crates/traces-clickhouse/query/span_text.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/spend_by_response_ids.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/trace_identity.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql delete mode 100644 litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql create mode 100644 litellm-rust/crates/traces-clickhouse/src/access.rs delete mode 100644 litellm-rust/crates/traces-clickhouse/src/span_batches.rs delete mode 100644 litellm-rust/crates/traces-clickhouse/tests/queries/read_access.json create mode 100644 litellm-rust/crates/traces-clickhouse/tests/search.rs delete mode 100644 litellm-rust/crates/traces/src/query/named.rs create mode 100644 litellm-rust/crates/traces/src/search.rs create mode 100644 litellm-rust/crates/traces/src/store.rs delete mode 100644 litellm-rust/crates/traces/tests/query/named.rs create mode 100644 litellm-rust/crates/traces/tests/search.rs create mode 100644 litellm/proxy/lens/search.py create mode 100644 scripts/trace_codegen/schemas/traces/RunField.json create mode 100644 scripts/trace_codegen/schemas/traces/RunValues.json create mode 100644 scripts/trace_codegen/schemas/traces/TraceHistogram.json delete mode 100644 scripts/trace_codegen/schemas/traces/TraceScope.json create mode 100644 tests/unit/proxy/lens/test_search.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f7b667c8ab2..daa0cf631bd 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3932,6 +3932,22 @@ dependencies = [ "uuid", ] +[[package]] +name = "litellm-gateway-traces" +version = "0.1.0" +dependencies = [ + "axum", + "litellm-traces", + "litellm-traces-cache", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tower", + "tracing", +] + [[package]] name = "litellm-gateway-ui" version = "0.1.0" diff --git a/litellm-rust/crates/gateway-traces/Cargo.toml b/litellm-rust/crates/gateway-traces/Cargo.toml new file mode 100644 index 00000000000..a8abedd3675 --- /dev/null +++ b/litellm-rust/crates/gateway-traces/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-gateway-traces" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +axum = { workspace = true, features = ["json", "query"] } +litellm-traces.workspace = true +litellm-traces-cache.workspace = true +serde.workspace = true +tracing.workspace = true + +[dev-dependencies] +rstest.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +tower = { version = "0.5", features = ["util"] } diff --git a/litellm-rust/crates/gateway-traces/src/error.rs b/litellm-rust/crates/gateway-traces/src/error.rs new file mode 100644 index 00000000000..b557c120652 --- /dev/null +++ b/litellm-rust/crates/gateway-traces/src/error.rs @@ -0,0 +1,87 @@ +use axum::{ + Json, + http::{HeaderValue, StatusCode, header}, + response::{IntoResponse, Response}, +}; +use litellm_traces_cache::ReadError; +use serde::Serialize; + +const RETRY_AFTER_SECONDS: &str = "2"; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ReadFailureCode { + InvalidRequest, + TraceChanged, + TooLarge, + Unavailable, +} + +#[derive(Debug)] +pub struct ReadFailure { + code: ReadFailureCode, + message: String, +} + +impl From> for ReadFailure { + fn from(error: ReadError) -> Self { + let (code, message) = match &error { + ReadError::InvalidParameters + | ReadError::InvalidCursor(_) + | ReadError::AmbiguousTrace => (ReadFailureCode::InvalidRequest, error.to_string()), + ReadError::TraceChanged => (ReadFailureCode::TraceChanged, error.to_string()), + ReadError::TooLarge => ( + ReadFailureCode::TooLarge, + "Trace is too large for this view. Use a filtered trace query.".to_owned(), + ), + ReadError::Encode(_) | ReadError::Store(_) => { + tracing::warn!(%error, "trace read unavailable"); + ( + ReadFailureCode::Unavailable, + "Traces are temporarily unavailable. Please try again.".to_owned(), + ) + } + }; + Self { code, message } + } +} + +#[derive(Serialize)] +struct Body<'a> { + detail: Detail<'a>, +} + +#[derive(Serialize)] +struct Detail<'a> { + code: ReadFailureCode, + message: &'a str, +} + +impl IntoResponse for ReadFailure { + fn into_response(self) -> Response { + let status = match self.code { + ReadFailureCode::InvalidRequest => StatusCode::BAD_REQUEST, + ReadFailureCode::TraceChanged => StatusCode::CONFLICT, + ReadFailureCode::TooLarge => StatusCode::PAYLOAD_TOO_LARGE, + ReadFailureCode::Unavailable => StatusCode::SERVICE_UNAVAILABLE, + }; + let body = Json(Body { + detail: Detail { + code: self.code, + message: &self.message, + }, + }); + match self.code { + ReadFailureCode::Unavailable => ( + status, + [( + header::RETRY_AFTER, + HeaderValue::from_static(RETRY_AFTER_SECONDS), + )], + body, + ) + .into_response(), + _ => (status, body).into_response(), + } + } +} diff --git a/litellm-rust/crates/gateway-traces/src/lib.rs b/litellm-rust/crates/gateway-traces/src/lib.rs new file mode 100644 index 00000000000..2b31f22c0a5 --- /dev/null +++ b/litellm-rust/crates/gateway-traces/src/lib.rs @@ -0,0 +1,24 @@ +mod error; +mod runs; + +use std::sync::Arc; + +use axum::{Router, routing::get}; +pub use error::ReadFailure; +use litellm_traces_cache::{TraceReader, TraceStore}; + +pub struct Traces { + pub reader: TraceReader, + pub store: S, +} + +pub fn router(traces: Arc>) -> Router +where + S: TraceStore + Send + 'static, +{ + Router::new() + .route("/v1/traces", get(runs::list::)) + .route("/v1/traces/histogram", get(runs::histogram::)) + .route("/v1/traces/values/{field}", get(runs::values::)) + .with_state(traces) +} diff --git a/litellm-rust/crates/gateway-traces/src/runs.rs b/litellm-rust/crates/gateway-traces/src/runs.rs new file mode 100644 index 00000000000..b5ee4b92641 --- /dev/null +++ b/litellm-rust/crates/gateway-traces/src/runs.rs @@ -0,0 +1,126 @@ +use std::{ + sync::Arc, + time::{SystemTime, UNIX_EPOCH}, +}; + +use axum::{ + Extension, Json, + extract::{Path, Query, State}, +}; +use litellm_traces::{ + QueryScope, TracePage, + search::{RunField, RunFilter, RunSearch, RunValues, TraceHistogram}, +}; +use litellm_traces_cache::{PageRequest, TraceStore}; +use serde::Deserialize; + +use crate::{ReadFailure, Traces}; + +const DAY_MS: i64 = 24 * 60 * 60 * 1000; +const PAGE_SIZE: u32 = 50; + +#[derive(Deserialize)] +pub(crate) struct Runs { + start_ms: Option, + end_ms: Option, + #[serde(default)] + q: String, +} + +impl Runs { + fn filter(&self) -> RunFilter { + let now_ms = now_ms(); + RunFilter { + start_ms: self.start_ms.unwrap_or(now_ms - DAY_MS), + end_ms: self.end_ms.unwrap_or(now_ms), + search: RunSearch::parse(&self.q), + } + } +} + +#[derive(Deserialize)] +pub(crate) struct Page { + cursor: Option, +} + +pub(crate) async fn list( + State(traces): State>>, + Extension(access): Extension, + Query(runs): Query, + Query(Page { cursor }): Query, +) -> Result, ReadFailure> { + let page = PageRequest { + cursor, + limit: PAGE_SIZE, + }; + Ok(Json( + traces + .reader + .list_traces(&traces.store, &access, &runs.filter(), &page) + .await?, + )) +} + +#[derive(Deserialize)] +pub(crate) struct Buckets { + #[serde(default = "default_buckets")] + buckets: u32, +} + +fn default_buckets() -> u32 { + 60 +} + +pub(crate) async fn histogram( + State(traces): State>>, + Extension(access): Extension, + Query(runs): Query, + Query(Buckets { buckets }): Query, +) -> Result, ReadFailure> { + Ok(Json( + traces + .reader + .histogram(&traces.store, &access, &runs.filter(), buckets) + .await?, + )) +} + +#[derive(Deserialize)] +pub(crate) struct Values { + #[serde(default)] + contains: String, + #[serde(default = "default_values")] + limit: u32, +} + +fn default_values() -> u32 { + 20 +} + +pub(crate) async fn values( + State(traces): State>>, + Extension(access): Extension, + Path(field): Path, + Query(runs): Query, + Query(values): Query, +) -> Result, ReadFailure> { + Ok(Json( + traces + .reader + .values( + &traces.store, + &access, + &runs.filter(), + field, + &values.contains, + values.limit, + ) + .await?, + )) +} + +fn now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |elapsed| elapsed.as_millis() as i64) +} diff --git a/litellm-rust/crates/gateway-traces/tests/routes.rs b/litellm-rust/crates/gateway-traces/tests/routes.rs new file mode 100644 index 00000000000..444c5bf4bb8 --- /dev/null +++ b/litellm-rust/crates/gateway-traces/tests/routes.rs @@ -0,0 +1,283 @@ +use std::{ + sync::{Arc, Mutex}, + time::{SystemTime, UNIX_EPOCH}, +}; + +use axum::{ + Extension, Router, + body::{Body, to_bytes}, + http::{Request, StatusCode, header}, + response::Response, +}; +use litellm_gateway_traces::{Traces, router}; +use litellm_traces::{ + QueryScope, + search::{RunField, RunFilter, RunSearch}, + store::{ + CallQuery, CallRow, CountValue, RunCount, RunCountQuery, RunQuery, RunRow, RunSelection, + SpanQuery, SpanRow, SpanText, SpanTextQuery, + }, +}; +use litellm_traces_cache::{StoreError, StoreResult, TraceReader, TraceStore}; +use rstest::rstest; +use serde_json::{Value, json}; +use tower::ServiceExt; + +#[derive(Debug, thiserror::Error)] +#[error("fake trace store failed")] +struct FakeError; + +#[derive(Clone, Copy)] +enum Outcome { + Empty, + TooLarge, + Failed, +} + +struct FakeStore { + outcome: Outcome, + lists: Mutex>, + counts: Mutex>, +} + +impl FakeStore { + fn listed_filter(&self) -> RunFilter { + let lists = self.lists.lock().unwrap(); + let [(_, query)] = lists.as_slice() else { + panic!("expected one list read, got {}", lists.len()); + }; + let RunSelection::Matching(filter) = &query.selection else { + panic!("expected a run search"); + }; + filter.clone() + } + + fn counted(&self) -> RunCountQuery { + self.counts.lock().unwrap()[0].clone() + } +} + +impl TraceStore for FakeStore { + type Error = FakeError; + + fn source(&self) -> &str { + "fake" + } + + async fn runs( + &self, + access: &QueryScope, + query: &RunQuery, + ) -> StoreResult, FakeError> { + self.lists + .lock() + .unwrap() + .push((access.clone(), query.clone())); + match self.outcome { + Outcome::Empty => Ok(Vec::new()), + Outcome::TooLarge => Err(StoreError::TooLarge), + Outcome::Failed => Err(StoreError::Failed(FakeError)), + } + } + + async fn run_counts( + &self, + _: &QueryScope, + query: &RunCountQuery, + ) -> StoreResult, FakeError> { + self.counts.lock().unwrap().push(query.clone()); + Ok(vec![RunCount { + bucket: 0, + failed: false, + value: "researcher".into(), + runs: 2, + }]) + } + + async fn spans(&self, _: &QueryScope, _: &SpanQuery) -> StoreResult, FakeError> { + Ok(Vec::new()) + } + + async fn span_text( + &self, + _: &QueryScope, + _: &SpanTextQuery, + ) -> StoreResult, FakeError> { + Ok(None) + } + + async fn calls(&self, _: &QueryScope, _: &CallQuery) -> StoreResult, FakeError> { + Ok(Vec::new()) + } +} + +fn access() -> QueryScope { + QueryScope::Owned { + user_id: "user".into(), + team_ids: vec!["team".into()], + } +} + +fn app(outcome: Outcome) -> (Router, Arc>) { + let traces = Arc::new(Traces { + reader: TraceReader::new(usize::MAX), + store: FakeStore { + outcome, + lists: Mutex::new(Vec::new()), + counts: Mutex::new(Vec::new()), + }, + }); + let app = router(Arc::clone(&traces)).layer(Extension(access())); + (app, traces) +} + +async fn get(app: Router, uri: &str) -> Response { + app.oneshot(Request::get(uri).body(Body::empty()).unwrap()) + .await + .unwrap() +} + +async fn json_body(response: Response) -> Value { + serde_json::from_slice(&to_bytes(response.into_body(), 65536).await.unwrap()).unwrap() +} + +#[tokio::test] +async fn list_reads_the_window_and_parsed_search_in_the_callers_scope() { + let (app, traces) = app(Outcome::Empty); + let response = get( + app, + "/v1/traces?start_ms=10&end_ms=20&q=plan%20-agent:res*%20model:%22gpt%20x%22", + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + json_body(response).await, + json!({"data": [], "next_cursor": null}) + ); + let filter = traces.store.listed_filter(); + assert_eq!((filter.start_ms, filter.end_ms), (10, 20)); + let lists = traces.store.lists.lock().unwrap(); + assert_eq!((&lists[0].0, lists[0].1.limit), (&access(), 50)); + assert_eq!( + filter.search, + RunSearch::parse(r#"plan -agent:res* model:"gpt x""#) + ); +} + +#[tokio::test] +async fn list_defaults_to_the_last_day() { + let (app, traces) = app(Outcome::Empty); + let before = now_ms(); + assert_eq!(get(app, "/v1/traces").await.status(), StatusCode::OK); + let after = now_ms(); + + let filter = traces.store.listed_filter(); + assert!((before..=after).contains(&filter.end_ms)); + assert_eq!(filter.end_ms - filter.start_ms, 24 * 60 * 60 * 1000); + assert_eq!(filter.search, RunSearch::default()); +} + +#[tokio::test] +async fn histogram_buckets_the_matching_runs_of_the_window() { + let (app, traces) = app(Outcome::Empty); + let response = get( + app, + "/v1/traces/histogram?start_ms=0&end_ms=40&q=status:ok&buckets=4", + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + let body = json_body(response).await; + assert_eq!(body["buckets"].as_array().map(Vec::len), Some(4)); + assert_eq!( + body["buckets"][0], + json!({"start_ms": 0, "end_ms": 10, "total": 2, "failed": 0, "agents": [{"agent": "researcher", "runs": 2}]}) + ); + let query = traces.store.counted(); + assert_eq!(query.by.buckets, Some(4)); + assert_eq!(query.filter.search, RunSearch::parse("status:ok")); +} + +#[tokio::test] +async fn values_suggest_a_field_narrowed_by_the_search() { + let (app, traces) = app(Outcome::Empty); + let response = get( + app, + "/v1/traces/values/agent?start_ms=0&end_ms=40&q=model:gpt*&contains=res&limit=5", + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(json_body(response).await, json!({"values": ["researcher"]})); + let query = traces.store.counted(); + assert_eq!(query.by.value, Some(CountValue::Field(RunField::Agent))); + assert_eq!((query.contains.as_str(), query.limit), ("res", Some(5))); + assert_eq!(query.filter.search, RunSearch::parse("model:gpt*")); +} + +#[rstest] +#[case::bad_cursor( + Outcome::Empty, + "/v1/traces?cursor=not-a-cursor", + StatusCode::BAD_REQUEST, + "invalid_request" +)] +#[case::bad_window(Outcome::Empty, "/v1/traces?start_ms=x", StatusCode::BAD_REQUEST, "")] +#[case::reversed_window( + Outcome::Empty, + "/v1/traces?start_ms=20&end_ms=10", + StatusCode::BAD_REQUEST, + "invalid_request" +)] +#[case::too_many_buckets( + Outcome::Empty, + "/v1/traces/histogram?buckets=241", + StatusCode::BAD_REQUEST, + "invalid_request" +)] +#[case::too_many_values( + Outcome::Empty, + "/v1/traces/values/agent?limit=101", + StatusCode::BAD_REQUEST, + "invalid_request" +)] +#[case::unknown_field(Outcome::Empty, "/v1/traces/values/color", StatusCode::BAD_REQUEST, "")] +#[case::too_large( + Outcome::TooLarge, + "/v1/traces", + StatusCode::PAYLOAD_TOO_LARGE, + "too_large" +)] +#[case::store_down( + Outcome::Failed, + "/v1/traces", + StatusCode::SERVICE_UNAVAILABLE, + "unavailable" +)] +#[tokio::test] +async fn failures_keep_the_python_status_and_code( + #[case] outcome: Outcome, + #[case] uri: &str, + #[case] status: StatusCode, + #[case] code: &str, +) { + let (app, _) = app(outcome); + let response = get(app, uri).await; + + assert_eq!(response.status(), status); + assert_eq!( + response.headers().get(header::RETRY_AFTER).is_some(), + status == StatusCode::SERVICE_UNAVAILABLE + ); + if !code.is_empty() { + assert_eq!(json_body(response).await["detail"]["code"], code); + } +} + +fn now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as i64 +} diff --git a/litellm-rust/crates/python-bridge/src/routes/traces.rs b/litellm-rust/crates/python-bridge/src/routes/traces.rs index f442d6c31de..cf8d9c46fa8 100644 --- a/litellm-rust/crates/python-bridge/src/routes/traces.rs +++ b/litellm-rust/crates/python-bridge/src/routes/traces.rs @@ -1,8 +1,11 @@ use std::{collections::BTreeMap, sync::Arc}; use litellm_http::ClientVariant; -use litellm_traces::{QueryScope, ReadQuery, Tenant, query::named::ReadAccessParams}; -use litellm_traces_cache::{ReadError, TraceReader}; +use litellm_traces::{ + QueryScope, ReadQuery, Tenant, + search::{RunField, RunFilter, RunSearch}, +}; +use litellm_traces_cache::{PageRequest, ReadError, TraceReader}; use litellm_traces_clickhouse::{ ClickHouseTraces, Config, Error, InsertTable, Parameter, QueryReaders, }; @@ -92,6 +95,14 @@ fn map_read_error(error: ReadError) -> PyErr { } } +fn run_filter(start_ms: i64, end_ms: i64, q: &str) -> RunFilter { + RunFilter { + start_ms, + end_ms, + search: RunSearch::parse(q), + } +} + fn map_sql_error(error: Error) -> PyErr { match error { Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(400 | 404)) => { @@ -224,16 +235,79 @@ impl NativeTraceStorage { ) } - #[pyo3(signature = (scope, start_ms, end_ms, cursor, limit))] + #[pyo3(signature = (scope, start_ms, end_ms, q, cursor, limit))] + #[expect( + clippy::too_many_arguments, + reason = "one parameter per Python argument" + )] fn list_traces<'py>( &self, py: Python<'py>, - #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, start_ms: i64, end_ms: i64, + q: &str, cursor: Option, limit: u32, ) -> PyResult> { + let filter = run_filter(start_ms, end_ms, q); + let page = PageRequest { cursor, limit }; + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); + crate::execution::run_async( + py, + async move { + let store = ClickHouseTraces::new(client, connection); + reader.list_traces(&store, &scope, &filter, &page).await + }, + map_read_error, + ) + } + + fn trace_histogram<'py>( + &self, + py: Python<'py>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + start_ms: i64, + end_ms: i64, + q: &str, + buckets: u32, + ) -> PyResult> { + let filter = run_filter(start_ms, end_ms, q); + let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; + let connection = self.config.storage().reader().clone(); + let reader = Arc::clone(&self.reader); + crate::execution::run_async( + py, + async move { + let store = ClickHouseTraces::new(client, connection); + reader.histogram(&store, &scope, &filter, buckets).await + }, + map_read_error, + ) + } + + #[expect( + clippy::too_many_arguments, + reason = "one parameter per Python argument" + )] + fn run_values<'py>( + &self, + py: Python<'py>, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, + start_ms: i64, + end_ms: i64, + q: &str, + field: &str, + contains: &str, + limit: u32, + ) -> PyResult> { + let field = field + .parse::() + .map_err(|_| PyValueError::new_err(format!("unknown run field {field}")))?; + let filter = run_filter(start_ms, end_ms, q); + let contains = contains.to_owned(); let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; let connection = self.config.storage().reader().clone(); let reader = Arc::clone(&self.reader); @@ -242,7 +316,7 @@ impl NativeTraceStorage { async move { let store = ClickHouseTraces::new(client, connection); reader - .list_traces(&store, &scope, start_ms, end_ms, cursor.as_deref(), limit) + .values(&store, &scope, &filter, field, &contains, limit) .await }, map_read_error, @@ -254,7 +328,7 @@ impl NativeTraceStorage { &self, py: Python<'py>, trace_id: String, - #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, trace_ref: String, cursor: Option, page_size: Option, @@ -294,7 +368,7 @@ impl NativeTraceStorage { py: Python<'py>, trace_id: String, span_id: String, - #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, trace_ref: String, ) -> PyResult> { let client = crate::http::host_client(py, ClientVariant::NoRedirect)?; @@ -318,7 +392,7 @@ impl NativeTraceStorage { py: Python<'py>, trace_id: String, span_id: String, - #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: ReadAccessParams, + #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope, trace_ref: String, cursor: Option, ) -> PyResult> { @@ -438,9 +512,10 @@ pub fn trace_span_rows<'py>( #[cfg(test)] mod tests { - use super::*; use rstest::rstest; + use super::*; + #[rstest] #[case::row(Error::InvalidRow, "ValueError")] #[case::insert_limit(Error::InvalidLimit("CLICKHOUSE_TRACE_MAX_INSERT_BYTES"), "ValueError")] diff --git a/litellm-rust/crates/traces-cache/src/cache.rs b/litellm-rust/crates/traces-cache/src/cache.rs index f308714146e..367526ea613 100644 --- a/litellm-rust/crates/traces-cache/src/cache.rs +++ b/litellm-rust/crates/traces-cache/src/cache.rs @@ -1,9 +1,6 @@ use std::{future::Future, sync::Arc, time::Duration}; -use litellm_traces::{ - Trace, TraceSummary, - query::named::{ReadAccessParams, TraceSpansRow}, -}; +use litellm_traces::{QueryScope, Trace, TraceSummary, store::SpanRow}; use moka::{Expiry, future::Cache}; use serde::Serialize; use sha2::{Digest, Sha256}; @@ -28,7 +25,7 @@ impl SnapshotKey { pub fn new( source: &str, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, snapshot_ms: u64, @@ -38,7 +35,7 @@ impl SnapshotKey { pub fn latest( source: &str, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, ) -> Result { @@ -47,13 +44,13 @@ impl SnapshotKey { pub(crate) fn run( source: &str, - access: &ReadAccessParams, + access: &QueryScope, run: (&str, &str, &str, &str), ) -> Result { Self::digest(&("run", source, access, run)) } - pub(crate) fn scope(source: &str, access: &ReadAccessParams) -> Result { + pub(crate) fn scope(source: &str, access: &QueryScope) -> Result { Self::digest(&("scope", source, access)) } } @@ -68,7 +65,7 @@ pub enum Freshness { } impl Freshness { - pub fn of(rows: &[TraceSpansRow], spend_known: bool, snapshot_ms: u64) -> Self { + pub fn of(rows: &[SpanRow], spend_known: bool, snapshot_ms: u64) -> Self { let last_end_ms = rows .iter() .map(|row| row.start_ns.saturating_add_unsigned(row.duration_ns) / 1_000_000) @@ -287,16 +284,15 @@ impl ListCache { #[cfg(test)] mod tests { use litellm_traces::{ - SpanStatus, - query::named::{SpendByResponseIdsRow, TraceSpansRow}, - resolve_trace, + SpanStatus, resolve_trace, + store::{CallRow, SpanRow}, }; use rstest::rstest; use super::*; - fn row(span_id: &str) -> TraceSpansRow { - TraceSpansRow { + fn row(span_id: &str) -> SpanRow { + SpanRow { trace_id: String::new(), span_id: span_id.into(), parent_span_id: String::new(), @@ -326,20 +322,14 @@ mod tests { } fn trace(span_id: &str) -> Trace { - resolve_trace( - "trace", - "ref", - &[row(span_id)], - &[] as &[SpendByResponseIdsRow], - ) - .expect("fixture should resolve") + resolve_trace("trace", "ref", &[row(span_id)], &[] as &[CallRow]) + .expect("fixture should resolve") } fn key(suffix: &str) -> SnapshotKey { SnapshotKey::new( "source", - &ReadAccessParams { - all_teams: false, + &QueryScope::Owned { user_id: String::new(), team_ids: vec!["team".into()], }, diff --git a/litellm-rust/crates/traces-cache/src/cursor.rs b/litellm-rust/crates/traces-cache/src/cursor.rs index 8858f243f09..436fbc3049d 100644 --- a/litellm-rust/crates/traces-cache/src/cursor.rs +++ b/litellm-rust/crates/traces-cache/src/cursor.rs @@ -1,58 +1,47 @@ use base64::{Engine, engine::general_purpose::URL_SAFE}; +use litellm_traces::store::{RunCursor, SpanPart}; use serde::{Deserialize, Serialize}; use crate::ReadError; -pub(super) fn encode_cursor(position: &T) -> String { - URL_SAFE.encode(serde_json::to_vec(position).unwrap_or_default()) +#[derive(Deserialize, Serialize)] +#[serde( + tag = "kind", + content = "position", + rename_all = "snake_case", + deny_unknown_fields +)] +pub(super) enum Cursor { + Run(RunCursor), + Span(SpanPosition), + Text(TextPosition), } -pub(super) fn decode_cursor Deserialize<'de>, E>( - cursor: &str, - kind: &'static str, -) -> Result> { - URL_SAFE - .decode(cursor) - .ok() - .and_then(|json| serde_json::from_slice(&json).ok()) - .ok_or(ReadError::InvalidCursor(kind)) -} +impl Cursor { + fn kind(&self) -> &'static str { + match self { + Self::Run(_) => "trace", + Self::Span(_) => "span", + Self::Text(_) => "diagnostic", + } + } -pub(super) fn trace_position(cursor: Option<&str>) -> Result<(i64, String), ReadError> { - let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else { - return Ok((0, String::new())); - }; - match decode_cursor::<(i64, String), E>(cursor, "trace")? { - (start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)), - _ => Err(ReadError::InvalidCursor("trace")), + pub(super) fn encode(&self) -> String { + URL_SAFE.encode(serde_json::to_vec(self).unwrap_or_default()) + } + + fn decode(cursor: &str, kind: &'static str) -> Result> { + URL_SAFE + .decode(cursor) + .ok() + .and_then(|json| serde_json::from_slice::(&json).ok()) + .filter(|decoded| decoded.kind() == kind) + .ok_or(ReadError::InvalidCursor(kind)) } } #[derive(Deserialize, Serialize)] -pub(super) struct ErrorPosition { - pub(super) offset: u64, - pub(super) version: String, -} - -pub(super) fn error_position( - cursor: Option<&str>, -) -> Result, ReadError> { - let Some(cursor) = cursor else { - return Ok(None); - }; - let position: ErrorPosition = decode_cursor(cursor, "diagnostic")?; - let valid_version = position.version.len() == 64 - && position - .version - .bytes() - .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)); - if i64::try_from(position.offset).is_err() || !valid_version { - return Err(ReadError::InvalidCursor("diagnostic")); - } - Ok(Some(position)) -} - -#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] pub(super) struct SpanPosition { pub(super) trace_ref: String, pub(super) snapshot_ms: u64, @@ -60,64 +49,169 @@ pub(super) struct SpanPosition { pub(super) version: String, } +#[derive(Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub(super) struct TextPosition { + pub(super) part: SpanPart, + pub(super) offset: u64, + pub(super) version: String, +} + +pub(super) fn run_position(cursor: Option<&str>) -> Result, ReadError> { + let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else { + return Ok(None); + }; + match Cursor::decode(cursor, "trace")? { + Cursor::Run(position) if position.start_ms > 0 && !position.trace_ref.is_empty() => { + Ok(Some(position)) + } + _ => Err(ReadError::InvalidCursor("trace")), + } +} + +pub(super) fn span_position(cursor: &str) -> Result> { + match Cursor::decode(cursor, "span")? { + Cursor::Span(position) => Ok(position), + _ => Err(ReadError::InvalidCursor("span")), + } +} + +pub(super) fn text_position( + cursor: Option<&str>, + part: SpanPart, +) -> Result, ReadError> { + let Some(cursor) = cursor else { + return Ok(None); + }; + let Cursor::Text(position) = Cursor::decode(cursor, "diagnostic")? else { + return Err(ReadError::InvalidCursor("diagnostic")); + }; + let valid_version = position.version.len() == 64 + && position + .version + .bytes() + .all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte)); + if position.part != part || i64::try_from(position.offset).is_err() || !valid_version { + return Err(ReadError::InvalidCursor("diagnostic")); + } + Ok(Some(position)) +} + #[cfg(test)] mod tests { use rstest::rstest; use super::*; + fn run(start_ms: i64, trace_ref: &str) -> String { + Cursor::Run(RunCursor { + start_ms, + trace_ref: trace_ref.into(), + }) + .encode() + } + + fn text(part: SpanPart, offset: u64, version: String) -> String { + Cursor::Text(TextPosition { + part, + offset, + version, + }) + .encode() + } + + fn span() -> String { + Cursor::Span(SpanPosition { + trace_ref: "ref".into(), + snapshot_ms: 1, + offset: 2, + version: "A".repeat(64), + }) + .encode() + } + + fn json(value: serde_json::Value) -> String { + URL_SAFE.encode(value.to_string()) + } + #[rstest] - fn trace_cursor_round_trips_the_last_listed_run() { - let cursor = encode_cursor(&(1_790_742_989_377_i64, "4bad42b84e9de3ba46fc870185f8f023")); + fn run_cursor_round_trips_the_last_listed_run() { + let position = run_position::(Some(&run(1_790_742_989_377, "4BAD"))) + .unwrap() + .unwrap(); 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()) + (position.start_ms, position.trace_ref.as_str()), + (1_790_742_989_377, "4BAD") ); } #[rstest] - #[case::not_base64("abc")] - #[case::not_json("bm90LWpzb24=")] - #[case::numeric_reference("WzEsIDJd")] - #[case::zero_start("WzAsICJ0Il0=")] - fn malformed_trace_cursors_are_rejected(#[case] cursor: &str) { - let result: Result<(i64, String), ReadError> = trace_position(Some(cursor)); - assert!(matches!(result, Err(ReadError::InvalidCursor("trace")))); + #[case::absent(None)] + #[case::empty(Some(""))] + fn missing_run_cursor_starts_from_the_newest(#[case] cursor: Option<&str>) { + assert!(run_position::(cursor).unwrap().is_none()); } #[rstest] - #[case::not_base64("garbage")] - #[case::missing_fields("e30=")] - #[case::not_an_object("WzEsMl0=")] - fn malformed_diagnostic_cursors_are_rejected(#[case] cursor: &str) { - let result: Result, ReadError> = - error_position(Some(cursor)); + #[case::not_base64("abc".into())] + #[case::not_json(URL_SAFE.encode("not-json"))] + #[case::untagged_tuple(json(serde_json::json!([1, "ref"])))] + #[case::zero_start(run(0, "ref"))] + #[case::empty_ref(run(1, ""))] + #[case::span_cursor(span())] + #[case::text_cursor(text(SpanPart::Error, 0, "A".repeat(64)))] + #[case::unknown_field(json(serde_json::json!({"kind": "run", "position": {"start_ms": 1, "trace_ref": "r", "extra": 1}})))] + fn malformed_run_cursors_are_rejected(#[case] cursor: String) { assert!(matches!( - result, - Err(ReadError::InvalidCursor("diagnostic")) + run_position::(Some(&cursor)), + Err(ReadError::InvalidCursor("trace")) )); } #[rstest] - #[case::lowercase_version("a".repeat(64))] - #[case::short_version("A".repeat(63))] - fn diagnostic_cursor_requires_a_content_version(#[case] version: String) { - let cursor = encode_cursor(&ErrorPosition { offset: 1, version }); - let result: Result, ReadError> = - error_position(Some(&cursor)); + fn span_cursor_round_trips() { + let position = span_position::(&span()).unwrap(); + assert_eq!( + ( + position.trace_ref.as_str(), + position.snapshot_ms, + position.offset + ), + ("ref", 1, 2) + ); + } + + #[rstest] + #[case::run_cursor(run(1, "ref"))] + #[case::text_cursor(text(SpanPart::Error, 0, "a".repeat(64)))] + fn other_kinds_are_not_span_cursors(#[case] cursor: String) { assert!(matches!( - result, + span_position::(&cursor), + Err(ReadError::InvalidCursor("span")) + )); + } + + #[rstest] + fn text_cursor_round_trips() { + let position = text_position::( + Some(&text(SpanPart::Error, 7, "A".repeat(64))), + SpanPart::Error, + ) + .unwrap() + .unwrap(); + assert_eq!((position.offset, position.version), (7, "A".repeat(64))); + } + + #[rstest] + #[case::lowercase_version(text(SpanPart::Error, 1, "a".repeat(64)))] + #[case::short_version(text(SpanPart::Error, 1, "A".repeat(63)))] + #[case::offset_past_i64(text(SpanPart::Error, u64::MAX, "A".repeat(64)))] + #[case::other_part(text(SpanPart::Output, 1, "A".repeat(64)))] + #[case::span_cursor(span())] + #[case::missing_fields(json(serde_json::json!({"kind": "text", "position": {}})))] + fn malformed_text_cursors_are_rejected(#[case] cursor: String) { + assert!(matches!( + text_position::(Some(&cursor), SpanPart::Error), Err(ReadError::InvalidCursor("diagnostic")) )); } diff --git a/litellm-rust/crates/traces-cache/src/lib.rs b/litellm-rust/crates/traces-cache/src/lib.rs index 8ed79c1a547..8e8863cc232 100644 --- a/litellm-rust/crates/traces-cache/src/lib.rs +++ b/litellm-rust/crates/traces-cache/src/lib.rs @@ -2,11 +2,12 @@ mod cache; mod cursor; mod error; mod list; +mod pages; mod reader; mod spend; mod store; pub use cache::{Freshness, LIVE_TTL, SETTLED_TTL, Snapshot, SnapshotCache, SnapshotKey}; pub use error::{Error, ReadError}; -pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, TraceReader}; -pub use store::{StoreError, TraceStore}; +pub use reader::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, PageRequest, TraceReader}; +pub use store::{StoreError, StoreResult, TraceStore}; diff --git a/litellm-rust/crates/traces-cache/src/list.rs b/litellm-rust/crates/traces-cache/src/list.rs index 1b1dcbbefc3..c00ece2d1bd 100644 --- a/litellm-rust/crates/traces-cache/src/list.rs +++ b/litellm-rust/crates/traces-cache/src/list.rs @@ -1,17 +1,17 @@ use std::collections::HashMap; +use litellm_traces::{ + QueryScope, TraceSummary, listed_summary, resolve_trace, + store::{RunRow, SpanRow, SpanSelection}, +}; + use crate::{ ReadError, SnapshotKey, TraceReader, TraceStore, cache::{Freshness, ListedRun}, - reader::{map_store_error, now_ms}, + reader::{map_store_error, now_ms, spans}, spend::{spend, spend_window, spend_within}, store::StoreError, }; -use litellm_traces::{ - TraceSummary, listed_summary, - query::named::{ListTracesRow, ReadAccessParams, TracePageSpansParams, TraceSpansRow}, - resolve_trace, -}; const RUNS_PER_SPAN_READ: usize = 16; @@ -25,8 +25,8 @@ fn run_key(team_id: &str, api_key_hash: &str, trace_id: &str) -> (String, String fn cache_key( source: &str, - access: &ReadAccessParams, - row: &ListTracesRow, + access: &QueryScope, + row: &RunRow, ) -> Result> { Ok(SnapshotKey::run( source, @@ -40,7 +40,7 @@ fn cache_key( )?) } -fn summary(row: &ListTracesRow, listed: Option<&ListedRun>) -> TraceSummary { +fn summary(row: &RunRow, listed: Option<&ListedRun>) -> TraceSummary { match listed { Some(ListedRun::Resolved(summary, _)) => (**summary).clone(), Some(ListedRun::Limited) | None => listed_summary(row), @@ -52,8 +52,8 @@ fn summary(row: &ListTracesRow, listed: Option<&ListedRun>) -> TraceSummary { pub(super) async fn list_summaries( reader: &TraceReader, store: &S, - access: &ReadAccessParams, - runs: &[ListTracesRow], + access: &QueryScope, + runs: &[RunRow], ) -> Result, ReadError> { let mut keys = Vec::with_capacity(runs.len()); let mut listed = Vec::with_capacity(runs.len()); @@ -62,7 +62,7 @@ pub(super) async fn list_summaries( listed.push(reader.lists.runs.get(&key).await); keys.push(key); } - let misses: Vec<&ListTracesRow> = runs + let misses: Vec<&RunRow> = runs .iter() .zip(&listed) .filter_map(|(row, listed)| listed.is_none().then_some(row)) @@ -92,8 +92,8 @@ pub(super) async fn list_summaries( async fn resolve_runs( reader: &TraceReader, store: &S, - access: &ReadAccessParams, - runs: &[&ListTracesRow], + access: &QueryScope, + runs: &[&RunRow], ) -> Result>, ReadError> { let (Some(start_ms), Some(end_ms)) = ( runs.iter().map(|row| row.start_ms).min(), @@ -103,14 +103,12 @@ async fn resolve_runs( ) else { return Ok(Vec::new()); }; - let params = TracePageSpansParams { - access: access.clone(), + let selection = SpanSelection::Runs { trace_refs: runs.iter().map(|row| row.trace_ref.clone()).collect(), - start_ms, - end_ms: end_ms.saturating_add(1), + window: start_ms..end_ms.saturating_add(1), }; let snapshot_ms = now_ms(); - let spans = match store.run_spans(¶ms, snapshot_ms).await { + let spans = match spans(store, access, selection, snapshot_ms).await { Ok(spans) => spans, Err(StoreError::TooLarge) => { let mut resolved = Vec::with_capacity(runs.len()); @@ -140,7 +138,7 @@ async fn resolve_runs( )) .then(left.start_ns.cmp(&right.start_ns)) }); - let by_run: HashMap<_, &[TraceSpansRow]> = spans + let by_run: HashMap<_, &[SpanRow]> = spans .chunk_by(|left, right| { (&left.team_id, &left.api_key_hash, &left.trace_id) == (&right.team_id, &right.api_key_hash, &right.trace_id) @@ -174,8 +172,8 @@ async fn resolve_runs( async fn resolve_run( reader: &TraceReader, store: &S, - access: &ReadAccessParams, - row: &ListTracesRow, + access: &QueryScope, + row: &RunRow, ) -> Result, ReadError> { match reader .current(store, access, &row.trace_id, &row.trace_ref) diff --git a/litellm-rust/crates/traces-cache/src/pages.rs b/litellm-rust/crates/traces-cache/src/pages.rs new file mode 100644 index 00000000000..4b0b7e959d1 --- /dev/null +++ b/litellm-rust/crates/traces-cache/src/pages.rs @@ -0,0 +1,184 @@ +use std::future::Future; + +use litellm_traces::store::{CallCursor, CallRow, SpanCursor, SpanRow}; +use serde::Serialize; + +use crate::{ + MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, + store::{StoreError, StoreResult}, +}; + +pub(crate) const PAGE_SIZE: u32 = 256; + +pub(crate) trait Keyed: Serialize { + type Cursor: Clone; + + fn cursor(&self) -> Self::Cursor; +} + +impl Keyed for SpanRow { + type Cursor = SpanCursor; + + fn cursor(&self) -> SpanCursor { + SpanRow::cursor(self) + } +} + +impl Keyed for CallRow { + type Cursor = CallCursor; + + fn cursor(&self) -> CallCursor { + CallRow::cursor(self) + } +} + +#[derive(Default)] +struct ReadBudget { + bytes: usize, + rows: usize, +} + +impl ReadBudget { + fn reserve(&mut self, bytes: usize) -> bool { + self.bytes = self.bytes.saturating_add(bytes); + if self.bytes > MAX_GRAPH_BYTES || self.rows == MAX_GRAPH_SPANS { + return false; + } + self.rows += 1; + true + } +} + +/// Reads every row by following the keyset. A page storage rejects as too large is retried at +/// half the size, and the smaller page is kept for the rest of the read because row sizes within +/// one graph rarely shrink again. The read stops with `TooLarge` once it passes the graph budget +/// or a single row exceeds the response limit. +pub(crate) async fn read_all(read: F) -> StoreResult, E> +where + R: Keyed, + F: Fn(Option, u32) -> Fut, + Fut: Future, E>>, +{ + let mut after = None; + let mut limit = PAGE_SIZE; + let mut rows = Vec::new(); + let mut budget = ReadBudget::default(); + loop { + let page = match read(after.clone(), limit).await { + Err(StoreError::TooLarge) if limit > 1 => { + limit /= 2; + continue; + } + result => result?, + }; + let complete = page.len() < limit as usize; + for row in &page { + let bytes = serde_json::to_vec(row).map_or(usize::MAX, |json| json.len()); + if !budget.reserve(bytes) { + return Err(StoreError::TooLarge); + } + } + if let Some(last) = page.last() { + after = Some(last.cursor()); + } + rows.extend(page); + if complete { + return Ok(rows); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::byte_boundary(MAX_GRAPH_BYTES - 1, 0, 1, true)] + #[case::byte_overflow(MAX_GRAPH_BYTES - 1, 0, 2, false)] + #[case::integer_overflow(MAX_GRAPH_BYTES, 0, usize::MAX, false)] + #[case::row_boundary(0, MAX_GRAPH_SPANS - 1, 1, true)] + #[case::row_overflow(0, MAX_GRAPH_SPANS, 1, false)] + fn accumulation_stops_at_the_graph_budget( + #[case] bytes: usize, + #[case] rows: usize, + #[case] next: usize, + #[case] accepted: bool, + ) { + let mut budget = ReadBudget { bytes, rows }; + assert_eq!(budget.reserve(next), accepted); + } + + #[derive(Serialize)] + struct Number(u32); + + impl Keyed for Number { + type Cursor = u32; + + fn cursor(&self) -> u32 { + self.0 + } + } + + /// A table of `total` rows whose transport rejects any page larger than `largest_page`. + struct Table { + total: u32, + largest_page: u32, + requests: Mutex>, + } + + impl Table { + async fn page(&self, after: Option, limit: u32) -> StoreResult, String> { + self.requests.lock().unwrap().push(limit); + if limit > self.largest_page { + return Err(StoreError::TooLarge); + } + let after = after.unwrap_or(0); + let end = (after + limit).min(self.total); + Ok((after + 1..=end).map(Number).collect()) + } + } + + #[rstest] + #[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])] + #[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])] + #[tokio::test] + async fn a_rejected_page_size_is_not_retried( + #[case] total: u32, + #[case] largest_page: u32, + #[case] requests: &[u32], + ) { + let table = Table { + total, + largest_page, + requests: Mutex::new(Vec::new()), + }; + let rows = read_all(|after, limit| table.page(after, limit)) + .await + .unwrap(); + assert_eq!( + rows.iter().map(|row| row.0).collect::>(), + (1..=total).collect::>() + ); + assert_eq!(table.requests.lock().unwrap().as_slice(), requests); + } + + #[rstest] + #[tokio::test] + async fn a_single_oversized_row_fails_the_read() { + let table = Table { + total: 10, + largest_page: 0, + requests: Mutex::new(Vec::new()), + }; + let result = read_all(|after, limit| table.page(after, limit)).await; + assert!(matches!(result, Err(StoreError::TooLarge))); + assert_eq!( + table.requests.lock().unwrap().as_slice(), + &[256, 128, 64, 32, 16, 8, 4, 2, 1] + ); + } +} diff --git a/litellm-rust/crates/traces-cache/src/reader.rs b/litellm-rust/crates/traces-cache/src/reader.rs index 5c68d6e58bd..78d58a80f1b 100644 --- a/litellm-rust/crates/traces-cache/src/reader.rs +++ b/litellm-rust/crates/traces-cache/src/reader.rs @@ -1,27 +1,38 @@ -use std::{sync::Arc, time::Duration}; +use std::{collections::BTreeMap, sync::Arc, time::Duration}; + +use litellm_traces::{ + ObservationType, QueryScope, SpanDetail, SpanErrorPage, Trace, TracePage, resolve_trace, + search::{ + MAX_HISTOGRAM_BUCKETS, MAX_RUN_VALUES, RunField, RunFilter, RunValues, TraceHistogram, + histogram, + }, + store::{ + CountBy, CountValue, RunCountQuery, RunQuery, RunSelection, SpanPart, SpanQuery, SpanRow, + SpanSelection, SpanText, SpanTextQuery, + }, + to_ui_content, +}; use crate::{ ReadError, Snapshot, SnapshotCache, SnapshotKey, StoreError, TraceStore, cache::{Freshness, ListCache}, - cursor::{ - ErrorPosition, SpanPosition, decode_cursor, encode_cursor, error_position, trace_position, - }, + cursor::{Cursor, SpanPosition, TextPosition, run_position, span_position, text_position}, list::{list_summaries, run_batches}, + pages::read_all, spend::spend, }; -use litellm_traces::{ - SpanDetail, SpanErrorPage, Trace, TracePage, - query::named::{ - ListTracesParams, ReadAccessParams, SpanDetailParams, SpanErrorParams, TraceIdentityParams, - TraceSpansParams, - }, - resolve_trace, to_ui_content, -}; pub const MAX_GRAPH_BYTES: usize = 64 * 1024 * 1024; pub const MAX_GRAPH_SPANS: usize = 100_000; const SNAPSHOT_IDLE: Duration = Duration::from_secs(120); +const ERROR_PAGE_CHARS: u64 = 16_384; + +#[derive(Clone, Debug, Default)] +pub struct PageRequest { + pub cursor: Option, + pub limit: u32, +} /// A read that found no trace, kept apart from failures so single-flight waiters share it /// without it being cached. @@ -64,43 +75,38 @@ impl TraceReader { pub async fn list_traces( &self, store: &S, - access: &ReadAccessParams, - start_ms: i64, - end_ms: i64, - cursor: Option<&str>, - limit: u32, + access: &QueryScope, + filter: &RunFilter, + page: &PageRequest, ) -> Result> { - if limit == 0 { + if page.limit == 0 || filter.start_ms >= filter.end_ms { return Err(ReadError::InvalidParameters); } - let (cursor_ms, cursor_trace_id) = trace_position(cursor)?; + let after = run_position(page.cursor.as_deref())?; let scope = SnapshotKey::scope(store.source(), access)?; let accepted = self.lists.limits.get(&scope).await.unwrap_or(u32::MAX); - let mut params = ListTracesParams { - access: access.clone(), - start_ms, - end_ms, - cursor_ms, - cursor_trace_id, - limit: limit.min(500).min(accepted), + let mut query = RunQuery { + selection: RunSelection::Matching(filter.clone()), + after, + limit: page.limit.min(500).min(accepted), }; - let page = loop { - match store.list_runs(¶ms).await { - Err(StoreError::TooLarge) if params.limit > 1 => { - params.limit /= 2; - self.lists.limits.insert(scope.clone(), params.limit).await; + let rows = loop { + match store.runs(access, &query).await { + Err(StoreError::TooLarge) if query.limit > 1 => { + query.limit /= 2; + self.lists.limits.insert(scope.clone(), query.limit).await; } Err(StoreError::TooLarge) => return Err(ReadError::TooLarge), result => break result.map_err(map_store_error)?, } }; - let next_cursor = page + let next_cursor = rows .last() - .filter(|_| page.len() == params.limit as usize) - .map(|last| encode_cursor(&(last.start_ms, &last.trace_ref))); + .filter(|_| rows.len() == query.limit as usize) + .map(|last| Cursor::Run(last.cursor()).encode()); let data = { - let mut summaries = Vec::with_capacity(page.len()); - for batch in run_batches(&page) { + let mut summaries = Vec::with_capacity(rows.len()); + for batch in run_batches(&rows) { summaries.extend(list_summaries(self, store, access, batch).await?); } summaries @@ -108,10 +114,67 @@ impl TraceReader { Ok(TracePage { data, next_cursor }) } + pub async fn histogram( + &self, + store: &S, + access: &QueryScope, + filter: &RunFilter, + buckets: u32, + ) -> Result> { + if filter.start_ms >= filter.end_ms || !(1..=MAX_HISTOGRAM_BUCKETS).contains(&buckets) { + return Err(ReadError::InvalidParameters); + } + let query = RunCountQuery { + filter: filter.clone(), + by: CountBy { + buckets: Some(buckets), + failed: true, + value: Some(CountValue::PrimaryAgent), + }, + contains: String::new(), + limit: None, + }; + let rows = store + .run_counts(access, &query) + .await + .map_err(map_store_error)?; + Ok(histogram(&rows, filter.start_ms, filter.end_ms, buckets)) + } + + pub async fn values( + &self, + store: &S, + access: &QueryScope, + filter: &RunFilter, + field: RunField, + contains: &str, + limit: u32, + ) -> Result> { + if filter.start_ms >= filter.end_ms || !(1..=MAX_RUN_VALUES).contains(&limit) { + return Err(ReadError::InvalidParameters); + } + let query = RunCountQuery { + filter: filter.clone(), + by: CountBy { + value: Some(CountValue::Field(field)), + ..CountBy::default() + }, + contains: contains.to_owned(), + limit: Some(limit), + }; + let rows = store + .run_counts(access, &query) + .await + .map_err(map_store_error)?; + Ok(RunValues { + values: rows.into_iter().map(|row| row.value).collect(), + }) + } + pub async fn get_trace( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, ) -> Result, ReadError> { @@ -127,7 +190,7 @@ impl TraceReader { pub async fn get_trace_page( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, cursor: Option<&str>, @@ -151,7 +214,7 @@ impl TraceReader { }; return page(&snapshot, &position, page_size, self.response_bytes).map(Some); }; - let position: SpanPosition = decode_cursor(cursor, "span")?; + let position = span_position(cursor)?; if position.trace_ref != trace_ref || position.snapshot_ms == 0 { return Err(ReadError::InvalidCursor("span")); } @@ -174,7 +237,7 @@ impl TraceReader { pub(super) async fn current( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, ) -> Result>, ReadError> { @@ -191,7 +254,7 @@ impl TraceReader { async fn pinned( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, snapshot_ms: u64, @@ -200,13 +263,11 @@ impl TraceReader { .map_err(|error| Arc::new(error.into()))?; self.snapshots .pinned_or_load(key, snapshot_ms, async { - let params = TraceSpansParams { - access: access.clone(), + let selection = SpanSelection::Trace { trace_id: trace_id.to_owned(), trace_ref: trace_ref.to_owned(), }; - let rows = store - .trace_spans(¶ms, snapshot_ms) + let rows = spans(store, access, selection, snapshot_ms) .await .map_err(|error| Miss::Read(map_store_error(error)))?; let spend_rows = spend(store, access, &rows).await; @@ -226,7 +287,7 @@ impl TraceReader { pub async fn get_span( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, span_id: &str, trace_ref: &str, @@ -234,84 +295,183 @@ impl TraceReader { let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { return Ok(None); }; - let params = SpanDetailParams { - access: access.clone(), - trace_id: trace_id.to_owned(), - trace_ref, - span_id: span_id.to_owned(), + let read = |part| { + let query = SpanTextQuery { + trace_id: trace_id.to_owned(), + trace_ref: trace_ref.clone(), + span_id: span_id.to_owned(), + part, + offset: 0, + max_chars: None, + }; + async move { store.span_text(access, &query).await } }; - let row = store.span_detail(¶ms).await.map_err(map_store_error)?; - Ok(row.map(|row| SpanDetail { - input_ui: to_ui_content(&row.input), - output_ui: to_ui_content(&row.output), - span_id: row.span_id, - input: row.input, - output: row.output, - attributes: row.attributes, + let Some(input) = read(SpanPart::Input).await.map_err(map_store_error)? else { + return Ok(None); + }; + let output = text_of(read(SpanPart::Output).await)?; + let attributes = text_of(read(SpanPart::Attributes).await)?; + let output = if output.is_empty() { + self.agent_answer(store, access, trace_id, &trace_ref, span_id) + .await? + } else { + output + }; + Ok(Some(SpanDetail { + input_ui: to_ui_content(&input.text), + output_ui: to_ui_content(&output), + span_id: span_id.to_owned(), + input: input.text, + output, + attributes: parse_attributes(&attributes)?, })) } + async fn agent_answer( + &self, + store: &S, + access: &QueryScope, + trace_id: &str, + trace_ref: &str, + span_id: &str, + ) -> Result> { + let Some(snapshot) = self.current(store, access, trace_id, trace_ref).await? else { + return Ok(String::new()); + }; + let spans = &snapshot.trace().spans; + if !spans + .iter() + .any(|span| span.span_id == span_id && span.kind == ObservationType::Agent) + { + return Ok(String::new()); + } + let mut calls: Vec<_> = spans + .iter() + .filter(|span| { + span.kind == ObservationType::Llm && span.parent_span_id.as_deref() == Some(span_id) + }) + .collect(); + calls.sort_by(|left, right| right.start_offset_ms.total_cmp(&left.start_offset_ms)); + for call in calls { + let query = SpanTextQuery { + trace_id: trace_id.to_owned(), + trace_ref: trace_ref.to_owned(), + span_id: call.span_id.clone(), + part: SpanPart::Output, + offset: 0, + max_chars: None, + }; + let output = text_of(store.span_text(access, &query).await)?; + if !output.is_empty() { + return Ok(output); + } + } + Ok(String::new()) + } + pub async fn get_span_error( &self, store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, span_id: &str, trace_ref: &str, cursor: Option<&str>, ) -> Result, ReadError> { - let position = error_position(cursor)?; + let position = text_position(cursor, SpanPart::Error)?; let Some(trace_ref) = reference(store, access, trace_id, trace_ref).await? else { return Ok(None); }; let offset = position.as_ref().map_or(0, |position| position.offset); - let params = SpanErrorParams { - access: access.clone(), + let query = SpanTextQuery { 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(), + part: SpanPart::Error, + offset, + max_chars: Some(ERROR_PAGE_CHARS), }; - let Some(row) = store.span_error(¶ms).await.map_err(map_store_error)? else { + let Some(text) = store + .span_text(access, &query) + .await + .map_err(map_store_error)? + else { return Ok(None); }; - let next_offset = offset + row.message.chars().count() as u64; - let next_cursor = (next_offset < row.total_chars).then(|| { - encode_cursor(&ErrorPosition { + if position.is_some_and(|position| position.version != text.version) { + return Ok(None); + } + let next_offset = offset + text.text.chars().count() as u64; + let next_cursor = (next_offset < text.total_chars).then(|| { + Cursor::Text(TextPosition { + part: SpanPart::Error, offset: next_offset, - version: row.version, + version: text.version, }) + .encode() }); Ok(Some(SpanErrorPage { - span_id: row.span_id, - message: row.message, - total_chars: row.total_chars, + span_id: span_id.to_owned(), + message: text.text, + total_chars: text.total_chars, next_cursor, })) } } +fn text_of(result: Result, StoreError>) -> Result> { + Ok(result + .map_err(map_store_error)? + .map(|text| text.text) + .unwrap_or_default()) +} + +fn parse_attributes(json: &str) -> Result, ReadError> { + if json.is_empty() { + return Ok(BTreeMap::new()); + } + serde_json::from_str(json).map_err(|error| ReadError::Encode(Arc::new(error))) +} + +pub(super) async fn spans( + store: &S, + access: &QueryScope, + selection: SpanSelection, + as_of_ms: u64, +) -> Result, StoreError> { + let mut rows = read_all(|after, limit| { + let query = SpanQuery { + selection: selection.clone(), + as_of_ms, + after, + limit, + }; + async move { store.spans(access, &query).await } + }) + .await?; + rows.sort_by_key(|row| row.start_ns); + Ok(rows) +} + async fn reference( store: &S, - access: &ReadAccessParams, + access: &QueryScope, trace_id: &str, trace_ref: &str, ) -> Result, ReadError> { 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 query = RunQuery { + selection: RunSelection::TraceId(trace_id.to_owned()), + after: None, + limit: 2, }; - let identities = store.trace_refs(¶ms).await.map_err(map_store_error)?; - if identities.len() > 1 { + let runs = store.runs(access, &query).await.map_err(map_store_error)?; + if runs.len() > 1 { return Err(ReadError::AmbiguousTrace); } - Ok(identities.into_iter().next()) + Ok(runs.into_iter().next().map(|run| run.trace_ref)) } fn page( @@ -328,12 +488,13 @@ fn page( agents: snapshot.trace().agents.clone(), spans: spans[position.offset..end].to_vec(), next_cursor: (end < spans.len()).then(|| { - encode_cursor(&SpanPosition { + Cursor::Span(SpanPosition { trace_ref: position.trace_ref.clone(), snapshot_ms: position.snapshot_ms, offset: end, version: snapshot.version().to_owned(), }) + .encode() }), } }; diff --git a/litellm-rust/crates/traces-cache/src/spend.rs b/litellm-rust/crates/traces-cache/src/spend.rs index 1cfcde841d0..ffcac16f34e 100644 --- a/litellm-rust/crates/traces-cache/src/spend.rs +++ b/litellm-rust/crates/traces-cache/src/spend.rs @@ -1,17 +1,16 @@ use std::ops::Range; -use crate::TraceStore; use litellm_traces::{ - SpendLookup, - query::named::{ - ReadAccessParams, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow, - }, + QueryScope, SpendLookup, + store::{CallQuery, CallRow, SpanRow}, }; +use crate::{TraceStore, pages::read_all}; + const NANOS_PER_MS: i64 = 1_000_000; const SPEND_WINDOW_MS: i64 = 30 * 60 * 1000; -pub(super) fn spend_window(rows: &[TraceSpansRow]) -> Option> { +pub(super) fn spend_window(rows: &[SpanRow]) -> Option> { let start_ns = rows.iter().map(|row| row.start_ns).min()?; let end_ns = rows .iter() @@ -23,10 +22,7 @@ pub(super) fn spend_window(rows: &[TraceSpansRow]) -> Option> { ) } -pub(super) fn spend_within( - spend: &[SpendByResponseIdsRow], - window: Range, -) -> &[SpendByResponseIdsRow] { +pub(super) fn spend_within(spend: &[CallRow], window: Range) -> &[CallRow] { let first = spend.partition_point(|row| row.start_ms < window.start); let end = spend.partition_point(|row| row.start_ms < window.end); &spend[first..end.max(first)] @@ -35,9 +31,9 @@ pub(super) fn spend_within( /// Spend rows sorted by `start_ms`, or `None` when the lookup failed and spend is unknown. pub(super) async fn spend( store: &S, - access: &ReadAccessParams, - rows: &[TraceSpansRow], -) -> Option> { + access: &QueryScope, + rows: &[SpanRow], +) -> Option> { let lookup = SpendLookup::new(rows); let Some(window) = spend_window(rows) else { return Some(Vec::new()); @@ -45,15 +41,19 @@ pub(super) async fn spend( if lookup.is_empty() { return Some(Vec::new()); } - let params = SpendByResponseIdsParams { - access: access.clone(), - response_ids: lookup.response_ids, - request_ids: lookup.request_ids, - trace_ids: lookup.trace_ids, - start_ms: window.start, - end_ms: window.end, - }; - match store.spend(¶ms).await { + let calls = read_all(|after, limit| { + let query = CallQuery { + window: window.clone(), + response_ids: lookup.response_ids.clone(), + request_ids: lookup.request_ids.clone(), + trace_ids: lookup.trace_ids.clone(), + after, + limit, + }; + async move { store.calls(access, &query).await } + }) + .await; + match calls { Ok(rows) => { let mut rows = rows; rows.sort_by_key(|row| row.start_ms); diff --git a/litellm-rust/crates/traces-cache/src/store.rs b/litellm-rust/crates/traces-cache/src/store.rs index 3f2a536d4cf..9fb7e96bf14 100644 --- a/litellm-rust/crates/traces-cache/src/store.rs +++ b/litellm-rust/crates/traces-cache/src/store.rs @@ -1,62 +1,60 @@ use std::future::Future; -use litellm_traces::query::named::{ - ListTracesParams, ListTracesRow, SpanDetailParams, SpanDetailRow, SpanErrorParams, - SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, TraceIdentityParams, - TracePageSpansParams, TraceSpansParams, TraceSpansRow, +use litellm_traces::{ + QueryScope, + store::{ + CallQuery, CallRow, RunCount, RunCountQuery, RunQuery, RunRow, SpanQuery, SpanRow, + SpanText, SpanTextQuery, + }, }; #[derive(Debug, thiserror::Error)] pub enum StoreError { + /// The answer would exceed what storage returns in one response; a smaller `limit` may fit. #[error("trace read exceeds the storage read budget")] TooLarge, #[error(transparent)] Failed(E), } +pub type StoreResult = Result>; + +/// Every read trace storage serves. Each call returns at most `limit` rows from one round trip; +/// paging, budgets and caching live above it. pub trait TraceStore: Sync { type Error: std::error::Error + Send + Sync + 'static; /// Identifies the backing storage for snapshot cache keys. fn source(&self) -> &str; - fn trace_refs( + fn runs( &self, - params: &TraceIdentityParams, - ) -> impl Future, StoreError>> + Send; + access: &QueryScope, + query: &RunQuery, + ) -> impl Future, Self::Error>> + Send; - /// Returns `TooLarge` when the response exceeds the storage limit so the reader can halve `limit`. - fn list_runs( + fn run_counts( &self, - params: &ListTracesParams, - ) -> impl Future, StoreError>> + Send; + access: &QueryScope, + query: &RunCountQuery, + ) -> impl Future, Self::Error>> + Send; - /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`. - fn trace_spans( + fn spans( &self, - params: &TraceSpansParams, - snapshot_ms: u64, - ) -> impl Future, StoreError>> + Send; + access: &QueryScope, + query: &SpanQuery, + ) -> impl Future, Self::Error>> + Send; - /// Returns spans visible at `snapshot_ms`, sorted by `start_ns`, or `TooLarge` past `MAX_GRAPH_BYTES`/`MAX_GRAPH_SPANS`. - fn run_spans( + /// `None` when the span is not visible to `access`. + fn span_text( &self, - params: &TracePageSpansParams, - snapshot_ms: u64, - ) -> impl Future, StoreError>> + Send; + access: &QueryScope, + query: &SpanTextQuery, + ) -> impl Future, Self::Error>> + Send; - fn spend( + fn calls( &self, - params: &SpendByResponseIdsParams, - ) -> impl Future, StoreError>> + Send; - - fn span_detail( - &self, - params: &SpanDetailParams, - ) -> impl Future, StoreError>> + Send; - - fn span_error( - &self, - params: &SpanErrorParams, - ) -> impl Future, StoreError>> + Send; + access: &QueryScope, + query: &CallQuery, + ) -> impl Future, Self::Error>> + Send; } diff --git a/litellm-rust/crates/traces-cache/tests/read.rs b/litellm-rust/crates/traces-cache/tests/read.rs index d714d7c9ed4..02c54e45049 100644 --- a/litellm-rust/crates/traces-cache/tests/read.rs +++ b/litellm-rust/crates/traces-cache/tests/read.rs @@ -8,14 +8,16 @@ use std::{ }; use litellm_traces::{ - CallEvidenceKind, CallKey, ObservationType, SpanStatus, - query::named::{ - ListTracesParams, ListTracesRow, ReadAccessParams, SpanDetailParams, SpanDetailRow, - SpanErrorParams, SpanErrorRow, SpendByResponseIdsParams, SpendByResponseIdsRow, - TraceIdentityParams, TracePageSpansParams, TraceSpansParams, TraceSpansRow, + CallEvidenceKind, CallKey, ObservationType, QueryScope, SpanStatus, + search::{AgentRuns, HistogramBucket, RunField, RunFilter, RunSearch}, + store::{ + CallQuery, CallRow, CountBy, CountValue, RunCount, RunCountQuery, RunQuery, RunRow, + RunSelection, SpanPart, SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery, }, }; -use litellm_traces_cache::{LIVE_TTL, ReadError, StoreError, TraceReader, TraceStore}; +use litellm_traces_cache::{ + LIVE_TTL, PageRequest, ReadError, StoreError, StoreResult, TraceReader, TraceStore, +}; use rstest::rstest; const START_NS: i64 = 1_790_742_989_000_000_000; @@ -27,8 +29,7 @@ enum Operation { TraceSpans, RunSpans, Spend, - SpanDetail, - SpanError, + SpanText, } #[derive(Clone, Copy)] @@ -45,42 +46,32 @@ struct FakeError; struct State { failures: HashMap, trace_refs: Vec, - list_runs: Vec, - trace_spans: HashMap>, - run_spans: Vec, - spend: Vec, - span_detail: Option, - span_error: Option, + list_runs: Vec, + trace_spans: HashMap>, + run_spans: Vec, + spend: Vec, + texts: HashMap<(String, SpanPart), String>, list_runs_too_large_above: Option, trace_too_large_refs: HashSet, spend_fails_above_response_ids: Option, -} - -#[derive(Default)] -struct Calls { - trace_refs: AtomicUsize, - list_runs: AtomicUsize, - trace_spans: AtomicUsize, - run_spans: AtomicUsize, - spend: AtomicUsize, - span_detail: AtomicUsize, - span_error: AtomicUsize, + count_rows: Vec, + count_reads: Vec, } #[derive(Default)] struct FakeStore { state: Mutex, - calls: Calls, + calls: Mutex>, } impl FakeStore { - fn with_spans(trace_ref: &str, spans: Vec) -> Self { + fn with_spans(trace_ref: &str, spans: Vec) -> Self { Self { state: Mutex::new(State { trace_spans: HashMap::from([(trace_ref.to_owned(), spans)]), ..State::default() }), - calls: Calls::default(), + calls: Mutex::default(), } } @@ -96,7 +87,7 @@ impl FakeStore { self.state.lock().unwrap().trace_refs = trace_refs; } - fn set_list_runs(&self, rows: Vec) { + fn set_list_runs(&self, rows: Vec) { self.state.lock().unwrap().list_runs = rows; } @@ -104,7 +95,7 @@ impl FakeStore { self.state.lock().unwrap().list_runs_too_large_above = Some(limit); } - fn set_run_spans(&self, rows: Vec) { + fn set_run_spans(&self, rows: Vec) { self.state.lock().unwrap().run_spans = rows; } @@ -116,25 +107,38 @@ impl FakeStore { .insert(trace_ref.to_owned()); } - /// Fails `spend` only when the lookup covers more than `limit` response ids, so a batch + fn set_text(&self, span_id: &str, part: SpanPart, text: &str) { + self.state + .lock() + .unwrap() + .texts + .insert((span_id.to_owned(), part), text.to_owned()); + } + + /// Fails `calls` only when the lookup covers more than `limit` response ids, so a batch /// covering several runs fails while each run's own narrower lookup still succeeds. fn set_spend_fails_above_response_ids(&self, limit: usize) { self.state.lock().unwrap().spend_fails_above_response_ids = Some(limit); } fn calls(&self, operation: Operation) -> usize { - match operation { - Operation::TraceRefs => self.calls.trace_refs.load(Ordering::SeqCst), - Operation::ListRuns => self.calls.list_runs.load(Ordering::SeqCst), - Operation::TraceSpans => self.calls.trace_spans.load(Ordering::SeqCst), - Operation::RunSpans => self.calls.run_spans.load(Ordering::SeqCst), - Operation::Spend => self.calls.spend.load(Ordering::SeqCst), - Operation::SpanDetail => self.calls.span_detail.load(Ordering::SeqCst), - Operation::SpanError => self.calls.span_error.load(Ordering::SeqCst), - } + self.calls + .lock() + .unwrap() + .get(&operation) + .map_or(0, |count| count.load(Ordering::SeqCst)) } - fn failure(state: &State, operation: Operation) -> Result<(), StoreError> { + fn record(&self, operation: Operation) { + self.calls + .lock() + .unwrap() + .entry(operation) + .or_default() + .fetch_add(1, Ordering::SeqCst); + } + + fn failure(state: &State, operation: Operation) -> StoreResult<(), FakeError> { match state.failures.get(&operation) { Some(Failure::TooLarge) => Err(StoreError::TooLarge), Some(Failure::Failed) => Err(StoreError::Failed(FakeError)), @@ -143,6 +147,20 @@ impl FakeStore { } } +fn keyset(rows: &[R], cursor: impl Fn(&R) -> C, after: Option, limit: u32) -> Vec +where + R: Clone, +{ + let mut rows: Vec = rows + .iter() + .filter(|row| after.as_ref().is_none_or(|after| cursor(row) > *after)) + .cloned() + .collect(); + rows.sort_by_key(|row| cursor(row)); + rows.truncate(limit as usize); + rows +} + impl TraceStore for FakeStore { type Error = FakeError; @@ -150,115 +168,152 @@ impl TraceStore for FakeStore { "fake" } - async fn trace_refs( - &self, - _: &TraceIdentityParams, - ) -> Result, StoreError> { - self.calls.trace_refs.fetch_add(1, Ordering::SeqCst); - let state = self.state.lock().unwrap(); - Self::failure(&state, Operation::TraceRefs)?; - Ok(state.trace_refs.clone()) - } - - async fn list_runs( - &self, - params: &ListTracesParams, - ) -> Result, StoreError> { - self.calls.list_runs.fetch_add(1, Ordering::SeqCst); + async fn runs(&self, _: &QueryScope, query: &RunQuery) -> StoreResult, FakeError> { let state = self.state.lock().unwrap(); + if let RunSelection::TraceId(trace_id) = &query.selection { + self.record(Operation::TraceRefs); + Self::failure(&state, Operation::TraceRefs)?; + return Ok(state + .trace_refs + .iter() + .take(query.limit as usize) + .map(|trace_ref| run(trace_id, trace_ref)) + .collect()); + } + self.record(Operation::ListRuns); Self::failure(&state, Operation::ListRuns)?; if state .list_runs_too_large_above - .is_some_and(|limit| params.limit > limit) + .is_some_and(|limit| query.limit > limit) { return Err(StoreError::TooLarge); } Ok(state .list_runs .iter() - .take(params.limit as usize) + .take(query.limit as usize) .cloned() .collect()) } - async fn trace_spans( + async fn run_counts( &self, - params: &TraceSpansParams, - _: u64, - ) -> Result, StoreError> { - self.calls.trace_spans.fetch_add(1, Ordering::SeqCst); - tokio::task::yield_now().await; - let state = self.state.lock().unwrap(); - Self::failure(&state, Operation::TraceSpans)?; - if state.trace_too_large_refs.contains(¶ms.trace_ref) { - return Err(StoreError::TooLarge); - } - Ok(state - .trace_spans - .get(¶ms.trace_ref) - .cloned() - .unwrap_or_default()) + _: &QueryScope, + query: &RunCountQuery, + ) -> StoreResult, FakeError> { + let mut state = self.state.lock().unwrap(); + state.count_reads.push(query.clone()); + Ok(state.count_rows.clone()) } - async fn run_spans( + async fn spans( &self, - _: &TracePageSpansParams, - _: u64, - ) -> Result, StoreError> { - self.calls.run_spans.fetch_add(1, Ordering::SeqCst); + _: &QueryScope, + query: &SpanQuery, + ) -> StoreResult, FakeError> { tokio::task::yield_now().await; let state = self.state.lock().unwrap(); - Self::failure(&state, Operation::RunSpans)?; - Ok(state.run_spans.clone()) + let rows = match &query.selection { + SpanSelection::Trace { trace_ref, .. } => { + self.record(Operation::TraceSpans); + Self::failure(&state, Operation::TraceSpans)?; + if state.trace_too_large_refs.contains(trace_ref) { + return Err(StoreError::TooLarge); + } + state + .trace_spans + .get(trace_ref) + .cloned() + .unwrap_or_default() + } + SpanSelection::Runs { .. } => { + self.record(Operation::RunSpans); + Self::failure(&state, Operation::RunSpans)?; + state.run_spans.clone() + } + }; + Ok(keyset( + &rows, + SpanRow::cursor, + query.after.clone(), + query.limit, + )) } - async fn spend( + async fn span_text( &self, - params: &SpendByResponseIdsParams, - ) -> Result, StoreError> { - self.calls.spend.fetch_add(1, Ordering::SeqCst); + _: &QueryScope, + query: &SpanTextQuery, + ) -> StoreResult, FakeError> { + self.record(Operation::SpanText); + let state = self.state.lock().unwrap(); + Self::failure(&state, Operation::SpanText)?; + let Some(text) = state.texts.get(&(query.span_id.clone(), query.part)) else { + return Ok(None); + }; + let rest = text.chars().skip(query.offset as usize); + Ok(Some(SpanText { + text: match query.max_chars { + Some(max) => rest.take(max as usize).collect(), + None => rest.collect(), + }, + total_chars: text.chars().count() as u64, + version: format!("{:0>64}", text.len()), + })) + } + + async fn calls( + &self, + _: &QueryScope, + query: &CallQuery, + ) -> StoreResult, FakeError> { + self.record(Operation::Spend); let state = self.state.lock().unwrap(); Self::failure(&state, Operation::Spend)?; if state .spend_fails_above_response_ids - .is_some_and(|limit| params.response_ids.len() > limit) + .is_some_and(|limit| query.response_ids.len() > limit) { return Err(StoreError::Failed(FakeError)); } - Ok(state.spend.clone()) - } - - async fn span_detail( - &self, - _: &SpanDetailParams, - ) -> Result, StoreError> { - self.calls.span_detail.fetch_add(1, Ordering::SeqCst); - let state = self.state.lock().unwrap(); - Self::failure(&state, Operation::SpanDetail)?; - Ok(state.span_detail.clone()) - } - - async fn span_error( - &self, - _: &SpanErrorParams, - ) -> Result, StoreError> { - self.calls.span_error.fetch_add(1, Ordering::SeqCst); - let state = self.state.lock().unwrap(); - Self::failure(&state, Operation::SpanError)?; - Ok(state.span_error.clone()) + Ok(keyset( + &state.spend, + CallRow::cursor, + query.after.clone(), + query.limit, + )) } } -fn access() -> ReadAccessParams { - ReadAccessParams { - all_teams: true, - user_id: String::new(), - team_ids: Vec::new(), +fn access() -> QueryScope { + QueryScope::All +} + +fn everything() -> RunFilter { + RunFilter { + start_ms: 0, + end_ms: i64::MAX, + search: RunSearch::default(), } } -fn span(index: usize) -> TraceSpansRow { - TraceSpansRow { +fn window(start_ms: i64, end_ms: i64, q: &str) -> RunFilter { + RunFilter { + start_ms, + end_ms, + search: RunSearch::parse(q), + } +} + +fn newest(limit: u32) -> PageRequest { + PageRequest { + cursor: None, + limit, + } +} + +fn span(index: usize) -> SpanRow { + SpanRow { trace_id: "trace".into(), span_id: format!("span-{index}"), parent_span_id: if index == 0 { @@ -291,8 +346,8 @@ fn span(index: usize) -> TraceSpansRow { } } -fn run(trace_id: &str, trace_ref: &str) -> ListTracesRow { - ListTracesRow { +fn run(trace_id: &str, trace_ref: &str) -> RunRow { + RunRow { trace_id: trace_id.into(), trace_ref: trace_ref.into(), team_id: "team".into(), @@ -315,7 +370,6 @@ fn run(trace_id: &str, trace_ref: &str) -> ListTracesRow { output_tokens: 0, models: Vec::new(), error_count: 0, - request_ids: Vec::new(), } } @@ -465,7 +519,7 @@ async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() { let reader = TraceReader::new(usize::MAX); let access = access(); let page = reader - .list_traces(&store, &access, 0, i64::MAX, None, 8) + .list_traces(&store, &access, &everything(), &newest(8)) .await .unwrap(); assert_eq!(page.data.len(), 2); @@ -476,7 +530,7 @@ async fn list_run_budget_halves_the_limit_and_cursor_requires_a_full_page() { shorter.set_list_runs(vec![run("only", "ref-only")]); shorter.set_list_runs_too_large_above(2); let page = reader - .list_traces(&shorter, &access, 0, i64::MAX, None, 8) + .list_traces(&shorter, &access, &everything(), &newest(8)) .await .unwrap(); assert_eq!(page.data.len(), 1); @@ -497,7 +551,7 @@ async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() store.set_failure(Operation::RunSpans, Failure::TooLarge); let reader = TraceReader::new(usize::MAX); let page = reader - .list_traces(&store, &access(), 0, i64::MAX, None, 2) + .list_traces(&store, &access(), &everything(), &newest(2)) .await .unwrap(); assert_eq!(page.data.len(), 2); @@ -505,20 +559,27 @@ async fn oversized_run_batch_falls_back_to_each_run_and_keeps_listed_summaries() assert_eq!(page.data[0].trace_ref, "ref-large"); assert!(!page.data[1].resolution_limited); assert_eq!(page.data[1].trace_ref, "ref-good"); - assert_eq!(store.calls(Operation::RunSpans), 1); - assert_eq!(store.calls(Operation::TraceSpans), 2); + let reads = ( + store.calls(Operation::RunSpans), + store.calls(Operation::TraceSpans), + ); let again = reader - .list_traces(&store, &access(), 0, i64::MAX, None, 2) + .list_traces(&store, &access(), &everything(), &newest(2)) .await .unwrap(); assert_eq!(again.data, page.data); - assert_eq!(store.calls(Operation::RunSpans), 1); - assert_eq!(store.calls(Operation::TraceSpans), 2); + assert_eq!( + ( + store.calls(Operation::RunSpans), + store.calls(Operation::TraceSpans) + ), + reads + ); } -fn spend_row(response_id: &str, cost: f64) -> SpendByResponseIdsRow { - SpendByResponseIdsRow { +fn spend_row(response_id: &str, cost: f64) -> CallRow { + CallRow { request_id: response_id.into(), litellm_call_id: String::new(), response_id: response_id.into(), @@ -563,7 +624,7 @@ async fn failed_batch_spend_lookup_falls_back_to_each_run_instead_of_losing_ever store.set_spend_fails_above_response_ids(1); let page = TraceReader::new(usize::MAX) - .list_traces(&store, &access(), 0, i64::MAX, None, 8) + .list_traces(&store, &access(), &everything(), &newest(8)) .await .unwrap(); @@ -643,17 +704,166 @@ async fn invalid_page_sizes_are_rejected(#[case] page_size: u32) { } #[rstest] +#[case::zero_limit(window(0, 10, ""), 0)] +#[case::empty_window(window(10, 10, ""), 8)] +#[case::reversed_window(window(10, 0, ""), 8)] #[tokio::test] -async fn zero_list_limit_is_rejected() { +async fn invalid_list_reads_are_rejected_before_storage( + #[case] filter: RunFilter, + #[case] limit: u32, +) { let store = FakeStore::default(); - let reader = TraceReader::new(usize::MAX); - let access = access(); assert!(matches!( - reader - .list_traces(&store, &access, 0, i64::MAX, None, 0) + TraceReader::new(usize::MAX) + .list_traces(&store, &access(), &filter, &newest(limit)) .await, Err(ReadError::InvalidParameters) )); + assert_eq!(store.calls(Operation::ListRuns), 0); +} + +#[rstest] +#[case::empty_window(window(10, 10, ""), 4)] +#[case::reversed_window(window(10, 0, ""), 4)] +#[case::zero_buckets(window(0, 10, ""), 0)] +#[case::too_many_buckets(window(0, 10, ""), 241)] +#[tokio::test] +async fn invalid_histogram_reads_are_rejected_before_storage( + #[case] filter: RunFilter, + #[case] buckets: u32, +) { + let store = FakeStore::default(); + assert!(matches!( + TraceReader::new(usize::MAX) + .histogram(&store, &access(), &filter, buckets) + .await, + Err(ReadError::InvalidParameters) + )); + assert!(store.state.lock().unwrap().count_reads.is_empty()); +} + +#[rstest] +#[tokio::test] +async fn histogram_reads_the_search_and_fills_every_bucket() { + let store = FakeStore::default(); + store.state.lock().unwrap().count_rows = vec![ + RunCount { + bucket: 1, + failed: false, + value: "writer".into(), + runs: 2, + }, + RunCount { + bucket: 1, + failed: true, + value: "reviewer".into(), + runs: 1, + }, + ]; + let filter = window(0, 30, "model:gpt*"); + let histogram = TraceReader::new(usize::MAX) + .histogram(&store, &access(), &filter, 3) + .await + .unwrap(); + + let empty = |start_ms, end_ms| HistogramBucket { + start_ms, + end_ms, + total: 0, + failed: 0, + agents: Vec::new(), + }; + assert_eq!( + histogram.buckets, + vec![ + empty(0, 10), + HistogramBucket { + start_ms: 10, + end_ms: 20, + total: 3, + failed: 1, + agents: vec![AgentRuns { + agent: "writer".into(), + runs: 2, + }], + }, + empty(20, 30), + ] + ); + let state = store.state.lock().unwrap(); + let [query] = state.count_reads.as_slice() else { + panic!("expected one histogram read"); + }; + assert_eq!( + query.by, + CountBy { + buckets: Some(3), + failed: true, + value: Some(CountValue::PrimaryAgent), + } + ); + assert_eq!((query.filter.clone(), query.limit), (filter, None)); +} + +#[rstest] +#[case::empty_window(window(10, 10, ""), 20)] +#[case::zero_limit(window(0, 10, ""), 0)] +#[case::too_many_values(window(0, 10, ""), 101)] +#[tokio::test] +async fn invalid_value_reads_are_rejected_before_storage( + #[case] filter: RunFilter, + #[case] limit: u32, +) { + let store = FakeStore::default(); + assert!(matches!( + TraceReader::new(usize::MAX) + .values(&store, &access(), &filter, RunField::Agent, "", limit) + .await, + Err(ReadError::InvalidParameters) + )); + assert!(store.state.lock().unwrap().count_reads.is_empty()); +} + +#[rstest] +#[tokio::test] +async fn values_narrow_by_the_search_and_the_needle() { + let store = FakeStore::default(); + store.state.lock().unwrap().count_rows = vec![ + RunCount { + bucket: 0, + failed: false, + value: "researcher".into(), + runs: 3, + }, + RunCount { + bucket: 0, + failed: false, + value: "re_writer".into(), + runs: 1, + }, + ]; + let filter = window(0, 10, "-status:error"); + let values = TraceReader::new(usize::MAX) + .values(&store, &access(), &filter, RunField::Agent, "re_", 20) + .await + .unwrap(); + + assert_eq!(values.values, vec!["researcher", "re_writer"]); + let state = store.state.lock().unwrap(); + let [query] = state.count_reads.as_slice() else { + panic!("expected one values read"); + }; + assert_eq!( + query.by, + CountBy { + value: Some(CountValue::Field(RunField::Agent)), + ..CountBy::default() + } + ); + assert_eq!( + (query.contains.as_str(), query.limit, &query.filter), + ("re_", Some(20), &filter) + ); } fn now_ns() -> i64 { @@ -707,12 +917,12 @@ async fn failed_reads_are_not_cached() { #[rstest] #[tokio::test] async fn listed_runs_are_read_once_until_a_live_run_expires() { - let live = TraceSpansRow { + let live = SpanRow { trace_id: "trace-live".into(), start_ns: now_ns(), ..span(0) }; - let settled = TraceSpansRow { + let settled = SpanRow { trace_id: "trace-settled".into(), ..span(0) }; @@ -724,7 +934,9 @@ async fn listed_runs_are_read_once_until_a_live_run_expires() { store.set_run_spans(vec![live, settled]); let reader = TraceReader::new(usize::MAX); let access = access(); - let list = || reader.list_traces(&store, &access, 0, i64::MAX, None, 2); + let filter = everything(); + let page = newest(2); + let list = || reader.list_traces(&store, &access, &filter, &page); let first = list().await.unwrap(); assert!(first.data.iter().all(|summary| summary.name == "agent")); @@ -766,3 +978,112 @@ async fn concurrent_pages_of_an_evicted_snapshot_share_one_storage_read() { assert_eq!(right.unwrap().unwrap().spans[0].span_id, "span-1"); assert_eq!(store.calls(Operation::TraceSpans), 1); } + +#[rstest] +#[tokio::test] +async fn trace_reads_follow_the_keyset_past_one_storage_page() { + let spans: Vec<_> = (0..600).map(span).collect(); + let store = FakeStore::with_spans("ref", spans); + let trace = TraceReader::new(usize::MAX) + .get_trace(&store, &access(), "trace", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(trace.spans.len(), 600); + assert_eq!(store.calls(Operation::TraceSpans), 3); +} + +fn llm(index: usize, offset_ms: i64) -> SpanRow { + SpanRow { + kind: ObservationType::Llm, + start_ns: START_NS + offset_ms * 1_000_000, + ..span(index) + } +} + +#[rstest] +#[case::agent_with_output("answer", "answer")] +#[case::agent_without_output("", "latest")] +#[tokio::test] +async fn span_detail_answers_a_silent_agent_with_its_latest_llm_output( + #[case] recorded: &str, + #[case] expected: &str, +) { + let store = FakeStore::with_spans("ref", vec![span(0), llm(1, 1), llm(2, 5), llm(3, 9)]); + store.set_text("span-0", SpanPart::Input, "question"); + store.set_text("span-0", SpanPart::Output, recorded); + store.set_text("span-0", SpanPart::Attributes, r#"{"k":"v"}"#); + store.set_text("span-1", SpanPart::Output, "earliest"); + store.set_text("span-2", SpanPart::Output, "latest"); + store.set_text("span-3", SpanPart::Output, ""); + let detail = TraceReader::new(usize::MAX) + .get_span(&store, &access(), "trace", "span-0", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!( + (detail.input.as_str(), detail.output.as_str()), + ("question", expected) + ); + assert_eq!(detail.attributes["k"], "v"); +} + +#[rstest] +#[tokio::test] +async fn span_detail_keeps_an_empty_output_for_spans_that_are_not_agents() { + let store = FakeStore::with_spans("ref", vec![span(0), llm(1, 1), llm(2, 2)]); + store.set_text("span-1", SpanPart::Input, "question"); + store.set_text("span-2", SpanPart::Output, "child output"); + let detail = TraceReader::new(usize::MAX) + .get_span(&store, &access(), "trace", "span-1", "ref") + .await + .unwrap() + .unwrap(); + assert_eq!(detail.output, ""); +} + +#[rstest] +#[tokio::test] +async fn span_detail_of_an_invisible_span_is_absent() { + let store = FakeStore::with_spans("ref", vec![span(0)]); + assert!( + TraceReader::new(usize::MAX) + .get_span(&store, &access(), "trace", "span-0", "ref") + .await + .unwrap() + .is_none() + ); +} + +#[rstest] +#[tokio::test] +async fn span_errors_page_by_characters_until_the_message_changes() { + let store = FakeStore::default(); + let message = "é".repeat(20_000); + store.set_text("span", SpanPart::Error, &message); + let reader = TraceReader::new(usize::MAX); + let first = reader + .get_span_error(&store, &access(), "trace", "span", "ref", None) + .await + .unwrap() + .unwrap(); + assert_eq!(first.message.chars().count(), 16_384); + assert_eq!(first.total_chars, 20_000); + let cursor = first.next_cursor.unwrap(); + let second = reader + .get_span_error(&store, &access(), "trace", "span", "ref", Some(&cursor)) + .await + .unwrap() + .unwrap(); + assert_eq!(second.message.chars().count(), 20_000 - 16_384); + assert!(second.next_cursor.is_none()); + + store.set_text("span", SpanPart::Error, "changed"); + assert!( + reader + .get_span_error(&store, &access(), "trace", "span", "ref", Some(&cursor)) + .await + .unwrap() + .is_none() + ); +} diff --git a/litellm-rust/crates/traces-cache/tests/snapshots.rs b/litellm-rust/crates/traces-cache/tests/snapshots.rs index 42085d3e95b..c79f5d4cae2 100644 --- a/litellm-rust/crates/traces-cache/tests/snapshots.rs +++ b/litellm-rust/crates/traces-cache/tests/snapshots.rs @@ -1,12 +1,9 @@ -use std::time::Duration; +use std::{sync::Arc, time::Duration}; use litellm_traces::{ - SpanStatus, Trace, - query::named::{ReadAccessParams, SpendByResponseIdsRow, TraceSpansRow}, - resolve_trace, + QueryScope, SpanStatus, Trace, resolve_trace, + store::{CallRow, SpanRow}, }; -use std::sync::Arc; - use litellm_traces_cache::{Error, Freshness, Snapshot, SnapshotCache, SnapshotKey}; use rstest::{fixture, rstest}; @@ -14,8 +11,8 @@ const T0: i64 = 1_790_742_989_000_000_000; const MS: i64 = 1_000_000; const TTL: Duration = Duration::from_secs(120); -fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> TraceSpansRow { - TraceSpansRow { +fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> SpanRow { + SpanRow { trace_id: String::new(), span_id: span_id.into(), parent_span_id: parent.into(), @@ -44,21 +41,14 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac } } -fn access() -> ReadAccessParams { - ReadAccessParams { - all_teams: false, +fn access() -> QueryScope { + QueryScope::Owned { user_id: String::new(), team_ids: vec!["team".into()], } } -fn key( - source: &str, - access: &ReadAccessParams, - trace_id: &str, - trace_ref: &str, - ms: u64, -) -> SnapshotKey { +fn key(source: &str, access: &QueryScope, trace_id: &str, trace_ref: &str, ms: u64) -> SnapshotKey { SnapshotKey::new(source, access, trace_id, trace_ref, ms).unwrap() } @@ -80,32 +70,22 @@ fn trace() -> Trace { "trace", "ref", &[row("root", "", "run", "agent", "agent")], - &[] as &[SpendByResponseIdsRow], + &[] as &[CallRow], ) .expect("fixture should resolve") } #[rstest] -#[case::different_team(false, "", "other-team")] -#[case::different_user(false, "other-user", "team")] -#[case::different_scope(true, "", "team")] +#[case::different_team(QueryScope::Owned { user_id: String::new(), team_ids: vec!["other-team".into()] })] +#[case::different_user(QueryScope::Owned { user_id: "other-user".into(), team_ids: vec!["team".into()] })] +#[case::different_scope(QueryScope::All)] #[tokio::test] -async fn cached_trace_is_isolated_by_access_scope( - trace: Trace, - #[case] all_teams: bool, - #[case] user_id: &str, - #[case] team_id: &str, -) { +async fn cached_trace_is_isolated_by_access_scope(trace: Trace, #[case] other_access: QueryScope) { let cache = SnapshotCache::new(1024 * 1024, TTL); let stored = key("source", &access(), "trace", "ref", 100); insert(&cache, stored.clone(), trace.clone()).await.unwrap(); - let other_access = ReadAccessParams { - all_teams, - user_id: user_id.into(), - team_ids: vec![team_id.into()], - }; let other = key("source", &other_access, "trace", "ref", 100); assert!(cache.get(&other).await.is_none()); @@ -171,12 +151,11 @@ async fn snapshot_version_tracks_the_ordered_span_ids( #[case] equal: bool, ) { let build = |ids: &[&str]| -> Trace { - let rows: Vec = ids + let rows: Vec = ids .iter() .map(|span_id| row(span_id, "", "run", "agent", "agent")) .collect(); - resolve_trace("trace", "ref", &rows, &[] as &[SpendByResponseIdsRow]) - .expect("fixture should resolve") + resolve_trace("trace", "ref", &rows, &[] as &[CallRow]).expect("fixture should resolve") }; let cache = SnapshotCache::new(1024 * 1024, TTL); diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql b/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql new file mode 100644 index 00000000000..b59a21f2d4f --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/migrations/0016_trace_rollup_agent_labels.sql @@ -0,0 +1,2 @@ +ALTER TABLE {database}.agent_traces_by_key + ADD COLUMN IF NOT EXISTS AgentLabels SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [] diff --git a/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql b/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql new file mode 100644 index 00000000000..c20373ef706 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/migrations/0017_trace_rollup_agent_labels_mv.sql @@ -0,0 +1,23 @@ +ALTER TABLE {database}.agent_traces_by_key_mv MODIFY QUERY +SELECT + TeamId, ApiKeyHash, TraceId, groupUniqArray(UserId) AS UserIds, + min(Timestamp) AS StartTs, + max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs, + any(ServiceName) AS ServiceName, + anyLastIf(toNullable(SpanName), ParentSpanId = '') AS RootName, + anyLastIf(toNullable(InputPreview), ParentSpanId = '') AS RootInput, + anyLastIf(toNullable(StatusCode), ParentSpanId = '') AS RootStatus, + count() AS SpanCount, + countIf(ObservationType = 'agent') AS AgentCount, + countIf(ObservationType = 'llm') AS LlmCount, + countIf(ObservationType = 'llm' AND LiteLLMRequestId != '') AS IdentifiedLlmCount, + countIf(ObservationType = 'tool') AS ToolCount, + countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount, + sum(InputTokens) AS InputTokens, + sum(OutputTokens) AS OutputTokens, + groupUniqArrayIf(toString(Model), Model != '') AS Models, + groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames, + groupUniqArrayIf(if(AgentName != '', AgentName, SpanName), ObservationType = 'agent' OR AgentName != '') AS AgentLabels, + groupArrayIf(LiteLLMRequestId, ObservationType = 'llm' OR LiteLLMRequestId != '') AS RequestIds +FROM {database}.otel_traces +GROUP BY TeamId, ApiKeyHash, TraceId diff --git a/litellm-rust/crates/traces-clickhouse/query/spend_batch.sql b/litellm-rust/crates/traces-clickhouse/query/calls.sql similarity index 83% rename from litellm-rust/crates/traces-clickhouse/query/spend_batch.sql rename to litellm-rust/crates/traces-clickhouse/query/calls.sql index 3918286a61f..3ec431c9ad6 100644 --- a/litellm-rust/crates/traces-clickhouse/query/spend_batch.sql +++ b/litellm-rust/crates/traces-clickhouse/query/calls.sql @@ -8,21 +8,17 @@ FROM ( if(startsWith(response_id, 'resp_'), extract(tryBase64Decode(substring(response_id, 6)), 'response_id:([^;]+)'), '') AS upstream_response_id - FROM spend_logs FINAL + FROM owned_calls 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)} OR upstream_response_id IN {response_ids:Array(String)} OR litellm_call_id IN {request_ids:Array(String)} OR (litellm_call_id = '' AND request_id IN {request_ids:Array(String)}) OR (trace_id != '' AND trace_id IN {trace_ids:Array(String)}) -ORDER BY start_time DESC ) WHERE {has_cursor:UInt8} = 0 OR (team_id, start_ms, request_id) > ({after_team:String}, {after_ms:Int64}, {after_id:String}) ORDER BY team_id, start_ms, request_id -LIMIT {page_size:UInt32} +LIMIT {limit:UInt32} diff --git a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql b/litellm-rust/crates/traces-clickhouse/query/list_traces.sql deleted file mode 100644 index c52adf7ef49..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/list_traces.sql +++ /dev/null @@ -1,48 +0,0 @@ -WITH page AS ( -SELECT TraceId AS trace_id, - hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref, - if(length(groupUniqArrayArray(UserIds)) = 1, arrayElement(groupUniqArrayArray(UserIds), 1), '') AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash, - ifNull(any(RootName), '') AS name, any(ServiceName) AS service, - ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status, - toUnixTimestamp64Milli(min(StartTs)) AS start_ms, - min(StartTs) AS trace_start, max(EndTs) AS trace_end, - dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms, - sum(SpanCount) AS span_count, - sum(AgentCount) AS agent_invocations, - sum(LlmCount) AS llm_calls, sum(ToolCount) AS tool_calls, - sum(InputTokens) AS input_tokens, sum(OutputTokens) AS output_tokens, - groupUniqArrayArray(Models) AS models, sum(ErrorCount) AS error_count, - arrayDistinct(if(sum(IdentifiedLlmCount) != sum(LlmCount), - arrayConcat(groupArrayArray(RequestIds), ['']), - groupArrayArray(RequestIds))) AS request_ids -FROM agent_traces_by_key -WHERE ({all_teams:UInt8} = 1 - OR ({user_id:String} != '' AND UserIds = [{user_id:String}]) - OR has({team_ids:Array(String)}, TeamId)) -GROUP BY TeamId, ApiKeyHash, TraceId -HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64}) - AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64}) - AND ({cursor_ms:Int64} = 0 OR (toUnixTimestamp64Milli(min(StartTs)), trace_ref) - < ({cursor_ms:Int64}, {cursor_trace_id:String})) -ORDER BY start_ms DESC, trace_ref DESC -LIMIT {limit:UInt32} -) -SELECT page.* EXCEPT (trace_start, trace_end), - identities.agent_names AS agent_names, identities.agent_count AS agent_count, - identities.frameworks AS frameworks -FROM page -LEFT JOIN ( - SELECT TeamId, ApiKeyHash, TraceId, - arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names, - arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks, - uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count - FROM otel_traces - WHERE Timestamp >= (SELECT min(trace_start) FROM page) - AND Timestamp <= (SELECT max(trace_end) FROM page) - AND TraceId IN (SELECT trace_id FROM page) - AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) - GROUP BY TeamId, ApiKeyHash, TraceId -) AS identities -ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash - AND page.trace_id = identities.TraceId -ORDER BY page.start_ms DESC, page.trace_ref DESC diff --git a/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql b/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql new file mode 100644 index 00000000000..c033e7f1f60 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/matching_runs.sql @@ -0,0 +1,33 @@ +SELECT TraceId AS trace_id, + hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref, + if(length(groupUniqArrayArray(UserIds)) = 1, arrayElement(groupUniqArrayArray(UserIds), 1), '') AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash, + ifNull(any(RootName), '') AS name, any(ServiceName) AS service, + ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status, + toUnixTimestamp64Milli(min(StartTs)) AS start_ms, + min(StartTs) AS trace_start, max(EndTs) AS trace_end, + dateDiff('millisecond', min(StartTs), max(EndTs)) AS duration_ms, + sum(SpanCount) AS span_count, + sum(AgentCount) AS agent_invocations, + sum(LlmCount) AS llm_calls, sum(ToolCount) AS tool_calls, + sum(InputTokens) AS input_tokens, sum(OutputTokens) AS output_tokens, + groupUniqArrayArray(Models) AS models, sum(ErrorCount) AS error_count, + arraySort(if(empty(groupUniqArrayArray(AgentLabels)), + groupUniqArrayArray(AgentNames), + groupUniqArrayArray(AgentLabels))) AS search_agents, + if(error_count > 0, 'error', 'ok') AS search_status +FROM owned_runs +WHERE {trace_id:String} = '' OR TraceId = {trace_id:String} +GROUP BY TeamId, ApiKeyHash, TraceId +HAVING {trace_id:String} != '' + OR (min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64}) + AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64}) + AND arrayAll(t -> trace_id ILIKE t OR input_preview ILIKE t OR name ILIKE t, {text:Array(String)}) + AND arrayAll((f, p, m) -> (m = 'exclude') != multiIf( + f = 'name', name ILIKE p, + f = 'agent', arrayExists(a -> a ILIKE p, search_agents), + f = 'status', search_status ILIKE p, + f = 'model', arrayExists(x -> x ILIKE p, models), + f = 'input', input_preview ILIKE p, + f = 'trace_id', trace_id ILIKE p, + false), + {filter_fields:Array(String)}, {filter_patterns:Array(String)}, {filter_modes:Array(String)})) diff --git a/litellm-rust/crates/traces-clickhouse/query/run_counts.sql b/litellm-rust/crates/traces-clickhouse/query/run_counts.sql new file mode 100644 index 00000000000..91ff6974a53 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/run_counts.sql @@ -0,0 +1,20 @@ +SELECT if({buckets:UInt32} = 0, toUInt32(0), + toUInt32(intDiv((start_ms - {start_ms:Int64}) * {buckets:UInt32}, {end_ms:Int64} - {start_ms:Int64}))) AS bucket, + toUInt8({by_failed:UInt8} = 1 AND error_count > 0) AS failed, + value, + count() AS runs +FROM runs +ARRAY JOIN multiIf( + {value:String} = '', [''], + {value:String} = 'primary_agent', [ifNull(nullIf(arrayElement(search_agents, 1), ''), service)], + {value:String} = 'name', [name], + {value:String} = 'agent', search_agents, + {value:String} = 'status', [search_status], + {value:String} = 'model', models, + {value:String} = 'input', [input_preview], + {value:String} = 'trace_id', [trace_id], + []) AS value +WHERE ({value:String} = '' OR value != '') AND value ILIKE {contains:String} +GROUP BY bucket, failed, value +ORDER BY runs DESC, bucket, failed, value +LIMIT {limit:UInt64} diff --git a/litellm-rust/crates/traces-clickhouse/query/run_spans.sql b/litellm-rust/crates/traces-clickhouse/query/run_spans.sql new file mode 100644 index 00000000000..65b06d9e3b8 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/run_spans.sql @@ -0,0 +1,3 @@ +WHERE Timestamp >= fromUnixTimestamp64Milli({start_ms:Int64}) + AND Timestamp < fromUnixTimestamp64Milli({end_ms:Int64}) + AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) IN {trace_refs:Array(String)} diff --git a/litellm-rust/crates/traces-clickhouse/query/runs.sql b/litellm-rust/crates/traces-clickhouse/query/runs.sql new file mode 100644 index 00000000000..6cf5d5c9f70 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/runs.sql @@ -0,0 +1,26 @@ +page AS ( +SELECT * EXCEPT (search_agents, search_status) +FROM runs +WHERE {has_cursor:UInt8} = 0 OR (start_ms, trace_ref) < ({cursor_ms:Int64}, {cursor_ref:String}) +ORDER BY start_ms DESC, trace_ref DESC +LIMIT {limit:UInt32} +) +SELECT page.* EXCEPT (trace_start, trace_end), + identities.agent_names AS agent_names, identities.agent_count AS agent_count, + identities.frameworks AS frameworks +FROM page +LEFT JOIN ( + SELECT TeamId, ApiKeyHash, TraceId, + arraySort(groupUniqArrayIf(AgentName, AgentName != '')) AS agent_names, + arraySort(groupUniqArrayIf(toString(Framework), Framework != '')) AS frameworks, + uniqExactIf(if(AgentName = '', SpanName, AgentName), ObservationType = 'agent') AS agent_count + FROM owned_spans + WHERE Timestamp >= (SELECT min(trace_start) FROM page) + AND Timestamp <= (SELECT max(trace_end) FROM page) + AND TraceId IN (SELECT trace_id FROM page) + AND (TeamId, ApiKeyHash, TraceId) IN (SELECT team_id, api_key_hash, trace_id FROM page) + GROUP BY TeamId, ApiKeyHash, TraceId +) AS identities +ON page.team_id = identities.TeamId AND page.api_key_hash = identities.ApiKeyHash + AND page.trace_id = identities.TraceId +ORDER BY page.start_ms DESC, page.trace_ref DESC diff --git a/litellm-rust/crates/traces-clickhouse/query/span_columns.sql b/litellm-rust/crates/traces-clickhouse/query/span_columns.sql new file mode 100644 index 00000000000..8c1a3a9b151 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/span_columns.sql @@ -0,0 +1,16 @@ +SELECT TraceId AS trace_id, SpanId AS span_id, ParentSpanId AS parent_span_id, SpanName AS name, + ObservationType AS type, toUInt8(WrapperCandidate) AS wrapper_candidate, AgentName AS agent, + Framework AS framework, StatusCode AS status, + substringUTF8(StatusMessage, 1, 128) AS status_message, + lengthUTF8(StatusMessage) > 128 AS error_truncated, + toUnixTimestamp64Nano(Timestamp) AS start_ns, Duration AS duration_ns, + ServiceName AS service, InputPreview AS input_preview, Model AS model, + InputTokens AS input_tokens, OutputTokens AS output_tokens, + LiteLLMRequestId AS litellm_request_id, + CallKeys AS call_keys, CallEvidence AS call_evidence, + -- Rows written before ToolCallId keep the call id only in their attributes. + if(ToolCallId != '' OR ObservationType != 'tool', ToolCallId, + coalesce(nullIf(SpanAttributes['gen_ai.tool.call.id'], ''), nullIf(SpanAttributes['tool.id'], ''), '')) + AS tool_call_id, + UserId AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash +FROM owned_spans diff --git a/litellm-rust/crates/traces-clickhouse/query/span_detail.sql b/litellm-rust/crates/traces-clickhouse/query/span_detail.sql deleted file mode 100644 index 7db742ea3ee..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/span_detail.sql +++ /dev/null @@ -1,24 +0,0 @@ -SELECT o.SpanId AS span_id, o.Input AS input, - if(o.Output = '' AND o.ObservationType = 'agent', answer.output, o.Output) AS output, - o.SpanAttributes AS attributes -FROM otel_traces AS o -LEFT JOIN ( - SELECT TeamId, ApiKeyHash, ParentSpanId AS parent_span_id, argMax(Output, Timestamp) AS output - FROM otel_traces - WHERE TraceId = {trace_id:String} AND ParentSpanId = {span_id:String} - AND ObservationType = 'llm' AND Output != '' - AND ({all_teams:UInt8} = 1 - OR ({user_id:String} != '' AND UserId = {user_id:String}) - OR has({team_ids:Array(String)}, TeamId)) - AND ({trace_ref:String} = '' OR - hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String}) - GROUP BY TeamId, ApiKeyHash, ParentSpanId -) AS answer ON answer.parent_span_id = o.SpanId - AND answer.TeamId = o.TeamId AND answer.ApiKeyHash = o.ApiKeyHash -WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String} - AND ({all_teams:UInt8} = 1 - OR ({user_id:String} != '' AND o.UserId = {user_id:String}) - OR has({team_ids:Array(String)}, o.TeamId)) - AND ({trace_ref:String} = '' OR - hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String}) -LIMIT 1 diff --git a/litellm-rust/crates/traces-clickhouse/query/span_error.sql b/litellm-rust/crates/traces-clickhouse/query/span_error.sql deleted file mode 100644 index e1226c4d23c..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/span_error.sql +++ /dev/null @@ -1,14 +0,0 @@ -SELECT SpanId AS span_id, - substringUTF8(StatusMessage, {error_offset:UInt64} + 1, 16384) AS message, - lengthUTF8(StatusMessage) AS total_chars, - hex(SHA256(StatusMessage)) AS version -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)) - 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}) -ORDER BY Timestamp, EngineReceivedMs, StatusMessage -LIMIT 1 diff --git a/litellm-rust/crates/traces-clickhouse/query/span_page.sql b/litellm-rust/crates/traces-clickhouse/query/span_page.sql new file mode 100644 index 00000000000..d51c278adc0 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/span_page.sql @@ -0,0 +1,7 @@ + AND EngineReceivedMs <= {as_of_ms:UInt64} +ORDER BY Timestamp, EngineReceivedMs, StatusMessage +LIMIT 1 BY TeamId, ApiKeyHash, TraceId, SpanId +) +WHERE (team_id, api_key_hash, trace_id, span_id) > ({after_team:String}, {after_key:String}, {after_trace:String}, {after_span:String}) +ORDER BY team_id, api_key_hash, trace_id, span_id +LIMIT {limit:UInt32} diff --git a/litellm-rust/crates/traces-clickhouse/query/span_text.sql b/litellm-rust/crates/traces-clickhouse/query/span_text.sql new file mode 100644 index 00000000000..ebec7b82fab --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/query/span_text.sql @@ -0,0 +1,15 @@ +SELECT if({bounded:UInt8} = 1, substringUTF8(part, {offset:UInt64} + 1, {max_chars:UInt64}), + substringUTF8(part, {offset:UInt64} + 1)) AS text, + lengthUTF8(part) AS total_chars, + hex(SHA256(part)) AS version +FROM ( + SELECT multiIf({part:String} = 'input', Input, + {part:String} = 'output', Output, + {part:String} = 'error', StatusMessage, + toJSONString(SpanAttributes)) AS part + FROM owned_spans + WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String} + AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String} + ORDER BY Timestamp, EngineReceivedMs, StatusMessage + LIMIT 1 +) diff --git a/litellm-rust/crates/traces-clickhouse/query/spend_by_response_ids.sql b/litellm-rust/crates/traces-clickhouse/query/spend_by_response_ids.sql deleted file mode 100644 index da64dafbc39..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/spend_by_response_ids.sql +++ /dev/null @@ -1,22 +0,0 @@ -SELECT request_id, litellm_call_id, response_id, upstream_response_id, trace_id, span_id, team_id, api_key, user, spend, - toUnixTimestamp64Milli(start_time) AS start_ms -FROM ( - SELECT *, - -- A chat request served through the Responses API returns the upstream `resp_` id to the - -- client but logs LiteLLM's managed `resp_` 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)} - OR upstream_response_id IN {response_ids:Array(String)} - OR litellm_call_id IN {request_ids:Array(String)} - OR (litellm_call_id = '' AND request_id IN {request_ids:Array(String)}) - OR (trace_id != '' AND trace_id IN {trace_ids:Array(String)}) -ORDER BY start_time DESC diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_identity.sql b/litellm-rust/crates/traces-clickhouse/query/trace_identity.sql deleted file mode 100644 index e3881b150b7..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/trace_identity.sql +++ /dev/null @@ -1,8 +0,0 @@ -SELECT hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref -FROM otel_traces -WHERE TraceId = {trace_id:String} - AND ({all_teams:UInt8} = 1 - OR ({user_id:String} != '' AND UserId = {user_id:String}) - OR has({team_ids:Array(String)}, TeamId)) -GROUP BY TeamId, ApiKeyHash, TraceId -LIMIT 2 diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql b/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql deleted file mode 100644 index 679edfdec2e..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/trace_list_span_batch.sql +++ /dev/null @@ -1,31 +0,0 @@ -SELECT * FROM ( -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)} - AND o.EngineReceivedMs <= {snapshot_ms:UInt64} -ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage -LIMIT 1 BY o.TeamId, o.ApiKeyHash, o.TraceId, o.SpanId - -) -WHERE (team_id, api_key_hash, trace_id, span_id) > ({after_team:String}, {after_key:String}, {after_trace:String}, {after_span:String}) -ORDER BY team_id, api_key_hash, trace_id, span_id -LIMIT {page_size:UInt32} diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql b/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql deleted file mode 100644 index b81bba61e7c..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/trace_page_spans.sql +++ /dev/null @@ -1,24 +0,0 @@ -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 diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql b/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql deleted file mode 100644 index 967ef2fcf1f..00000000000 --- a/litellm-rust/crates/traces-clickhouse/query/trace_span_batch.sql +++ /dev/null @@ -1,30 +0,0 @@ -SELECT * FROM ( -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.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)) - AND ({trace_ref:String} = '' OR - hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String}) - AND o.EngineReceivedMs <= {snapshot_ms:UInt64} -ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage -LIMIT 1 BY o.SpanId -) -WHERE span_id > {after_span_id:String} -ORDER BY span_id -LIMIT {page_size:UInt32} diff --git a/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql b/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql index 2e0ac4f6dfb..e68dbbaf0d2 100644 --- a/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql +++ b/litellm-rust/crates/traces-clickhouse/query/trace_spans.sql @@ -1,24 +1,2 @@ -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.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)) - 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 -LIMIT 1 BY o.SpanId +WHERE TraceId = {trace_id:String} + AND hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String} diff --git a/litellm-rust/crates/traces-clickhouse/src/access.rs b/litellm-rust/crates/traces-clickhouse/src/access.rs new file mode 100644 index 00000000000..941b1358dd6 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/src/access.rs @@ -0,0 +1,93 @@ +use litellm_traces::QueryScope; +use serde::Serialize; + +use crate::TraceTable; + +macro_rules! owned_by { + (otel_traces) => { + "({access_all:UInt8} = 1 OR ({access_user:String} != '' AND UserId = {access_user:String}) OR has({access_teams:Array(String)}, TeamId))" + }; + (agent_traces_by_key) => { + "({access_all:UInt8} = 1 OR ({access_user:String} != '' AND UserIds = [{access_user:String}]) OR has({access_teams:Array(String)}, TeamId))" + }; + (spend_logs) => { + "({access_all:UInt8} = 1 OR ({access_user:String} != '' AND user = {access_user:String}) OR has({access_teams:Array(String)}, team_id))" + }; +} + +/// Prefixes a read with the rows its caller may see: `owned_spans`, `owned_runs` and +/// `owned_calls`. Trusted SQL reads only these, never the tables. +macro_rules! owned { + ($($sql:expr),+ $(,)?) => { + concat!( + "WITH owned_spans AS (SELECT * FROM otel_traces WHERE ", + $crate::access::owned_by!(otel_traces), + "),\nowned_runs AS (SELECT * FROM agent_traces_by_key WHERE ", + $crate::access::owned_by!(agent_traces_by_key), + "),\nowned_calls AS (SELECT * FROM spend_logs FINAL WHERE ", + $crate::access::owned_by!(spend_logs), + ")", + $($sql),+ + ) + }; +} + +pub(crate) use owned; +pub(crate) use owned_by; + +#[derive(Debug, Serialize)] +pub(crate) struct AccessParams { + access_all: u8, + access_user: String, + access_teams: Vec, +} + +impl From<&QueryScope> for AccessParams { + fn from(scope: &QueryScope) -> Self { + match scope { + QueryScope::All => Self { + access_all: 1, + access_user: String::new(), + access_teams: Vec::new(), + }, + QueryScope::Owned { user_id, team_ids } => Self { + access_all: 0, + access_user: user_id.clone(), + access_teams: team_ids.clone(), + }, + } + } +} + +/// The ownership rule of `table` with `scope` written in, for a row policy. +pub(crate) fn predicate(scope: &QueryScope, table: TraceTable) -> String { + if *scope == QueryScope::All { + return "1".to_owned(); + } + let template = match table { + TraceTable::OtelTraces => owned_by!(otel_traces), + TraceTable::AgentTracesByKey => owned_by!(agent_traces_by_key), + TraceTable::SpendLogs => owned_by!(spend_logs), + }; + let AccessParams { + access_all, + access_user, + access_teams, + } = scope.into(); + let teams = access_teams + .iter() + .map(|team| literal(team)) + .collect::>() + .join(", "); + template + .replace("{access_all:UInt8}", &access_all.to_string()) + .replace("{access_user:String}", &literal(&access_user)) + .replace( + "{access_teams:Array(String)}", + &format!("CAST([{teams}], 'Array(String)')"), + ) +} + +fn literal(value: &str) -> String { + format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) +} diff --git a/litellm-rust/crates/traces-clickhouse/src/config.rs b/litellm-rust/crates/traces-clickhouse/src/config.rs index dccd1c368f6..285021f113e 100644 --- a/litellm-rust/crates/traces-clickhouse/src/config.rs +++ b/litellm-rust/crates/traces-clickhouse/src/config.rs @@ -1,6 +1,7 @@ -use crate::Error; use litellm_storage_clickhouse::Storage; +use crate::Error; + #[derive(Clone)] pub struct Config { storage: Storage, diff --git a/litellm-rust/crates/traces-clickhouse/src/insert.rs b/litellm-rust/crates/traces-clickhouse/src/insert.rs index d45cc8d53b8..8e3689a34ba 100644 --- a/litellm-rust/crates/traces-clickhouse/src/insert.rs +++ b/litellm-rust/crates/traces-clickhouse/src/insert.rs @@ -4,16 +4,15 @@ use std::{ io::{BufWriter, Write}, }; -use serde::{Serialize, Serializer, ser::SerializeMap}; - use flate2::{Compression, write::GzEncoder}; use litellm_http::Client; +use litellm_traces::Shared; +use serde::{Serialize, Serializer, ser::SerializeMap}; use serde_json::Value; use sha2::{Digest, Sha256}; use time::{OffsetDateTime, format_description::well_known::Rfc3339}; use super::{Connection, Error}; -use litellm_traces::Shared; fn max_insert_bytes() -> Result { let name = "CLICKHOUSE_TRACE_MAX_INSERT_BYTES"; @@ -260,9 +259,10 @@ mod tests { #[case::absent(None)] #[case::submitted(Some(123))] fn streamed_insert_preserves_token_and_stamps_receive_time(#[case] submitted: Option) { + use std::io::Read; + use flate2::read::GzDecoder; use sha2::{Digest, Sha256}; - use std::io::Read; let mut row = BTreeMap::from([ ("ApiKeyHash".into(), json!("key")), ("ResourceAttributes".into(), json!({"message": "雪\n\""})), diff --git a/litellm-rust/crates/traces-clickhouse/src/lib.rs b/litellm-rust/crates/traces-clickhouse/src/lib.rs index 43ff0b8bf33..f44116b8c86 100644 --- a/litellm-rust/crates/traces-clickhouse/src/lib.rs +++ b/litellm-rust/crates/traces-clickhouse/src/lib.rs @@ -10,6 +10,7 @@ macro_rules_attribute::attribute_alias! { #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]; } +mod access; mod config; mod error; mod insert; @@ -17,7 +18,6 @@ pub mod query; mod query_access; mod reads; mod schema; -mod span_batches; mod span_row; mod sql; mod table; diff --git a/litellm-rust/crates/traces-clickhouse/src/query.rs b/litellm-rust/crates/traces-clickhouse/src/query.rs index d8f7041386c..98603ca142c 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query.rs @@ -1,6 +1,5 @@ use std::collections::{BTreeMap, BTreeSet}; -use crate::TraceTable; use futures_util::{ StreamExt, stream::{self, TryStreamExt}, @@ -15,6 +14,7 @@ use super::{ Connection, Error, NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, Parameter, query_access::READER_LIMITS, }; +use crate::TraceTable; mod guide; pub mod lens; @@ -519,10 +519,11 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result String { + value + .chars() + .flat_map(|char| match char { + '\\' | '%' | '_' => vec!['\\', char], + _ => vec![char], + }) + .collect() } -#[derive(Debug, Deserialize, Serialize)] -pub struct ListTracesParams( - #[serde(with = "ListTracesParamsEncoding")] pub contracts::ListTracesParams, -); +fn contains(value: &str) -> String { + format!("%{}%", like_literal(value)) +} -impl From for ListTracesParams { - fn from(value: contracts::ListTracesParams) -> Self { - Self(value) +/// Query parameters only carry string arrays, so filters travel as parallel columns. +#[derive(Debug, Default, Serialize)] +struct SearchColumns { + text: Vec, + filter_fields: Vec<&'static str>, + filter_patterns: Vec, + filter_modes: Vec<&'static str>, +} + +impl From<&RunSearch> for SearchColumns { + fn from(search: &RunSearch) -> Self { + Self { + text: search.text.iter().map(|term| contains(term)).collect(), + filter_fields: search + .filters + .iter() + .map(|filter| filter.field.into()) + .collect(), + filter_patterns: search + .filters + .iter() + .map(|filter| like_literal(&filter.pattern).replace('*', "%")) + .collect(), + filter_modes: search + .filters + .iter() + .map(|filter| if filter.exclude { "exclude" } else { "include" }) + .collect(), + } + } +} + +#[derive(Debug, Default, Serialize)] +struct RunsFilter { + trace_id: String, + start_ms: i64, + end_ms: i64, + #[serde(flatten)] + search: SearchColumns, +} + +impl From<&RunFilter> for RunsFilter { + fn from(filter: &RunFilter) -> Self { + Self { + trace_id: String::new(), + start_ms: filter.start_ms, + end_ms: filter.end_ms, + search: (&filter.search).into(), + } + } +} + +impl From<&RunSelection> for RunsFilter { + fn from(selection: &RunSelection) -> Self { + match selection { + RunSelection::Matching(filter) => filter.into(), + RunSelection::TraceId(trace_id) => Self { + trace_id: trace_id.clone(), + ..Self::default() + }, + } + } +} + +#[derive(Debug, Serialize)] +pub(crate) struct RunsParams { + #[serde(flatten)] + access: AccessParams, + #[serde(flatten)] + filter: RunsFilter, + has_cursor: u8, + cursor_ms: i64, + cursor_ref: String, + limit: u32, +} + +impl RunsParams { + pub(crate) fn new(access: &QueryScope, query: &RunQuery) -> Self { + let after = query.after.clone().unwrap_or_default(); + Self { + access: access.into(), + filter: (&query.selection).into(), + has_cursor: query.after.is_some().into(), + cursor_ms: after.start_ms, + cursor_ref: after.trace_ref, + limit: query.limit, + } } } #[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::ListTracesRow")] -struct ListTracesRowEncoding { +#[serde(remote = "RunRow")] +struct RunRowEncoding { pub trace_id: String, pub trace_ref: String, pub team_id: String, @@ -69,17 +154,166 @@ struct ListTracesRowEncoding { pub models: Vec, #[serde(deserialize_with = "super::number::deserialize")] pub error_count: u64, - pub request_ids: Vec, } #[derive(Debug, Deserialize, Serialize)] -pub struct ListTracesRow(#[serde(with = "ListTracesRowEncoding")] pub contracts::ListTracesRow); +pub(crate) struct RunRowWire(#[serde(with = "RunRowEncoding")] pub RunRow); -pub use contracts::TraceSpansParams; +macro_rules! over_matching_runs { + ($($tail:expr),+ $(,)?) => { + owned!( + ",\nruns AS (\n", + include_str!("../../query/matching_runs.sql"), + ")", + $($tail),+ + ) + }; +} + +pub(crate) struct Runs; + +impl Query for Runs { + type Params = RunsParams; + type Row = RunRowWire; + + const SQL: &'static str = over_matching_runs!(",\n", include_str!("../../query/runs.sql")); +} + +#[derive(Debug, Serialize)] +pub(crate) struct RunCountsParams { + #[serde(flatten)] + access: AccessParams, + #[serde(flatten)] + filter: RunsFilter, + buckets: u32, + by_failed: u8, + value: &'static str, + contains: String, + limit: u64, +} + +impl RunCountsParams { + pub(crate) fn new(access: &QueryScope, query: &RunCountQuery) -> Self { + Self { + access: access.into(), + filter: (&query.filter).into(), + buckets: query.by.buckets.unwrap_or(0), + by_failed: query.by.failed.into(), + value: match query.by.value { + None => "", + Some(CountValue::PrimaryAgent) => "primary_agent", + Some(CountValue::Field(field)) => field.into(), + }, + contains: contains(&query.contains), + limit: query.limit.map_or(u64::MAX, u64::from), + } + } +} #[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::TraceSpansRow")] -struct TraceSpansRowEncoding { +#[serde(remote = "RunCount")] +struct RunCountEncoding { + #[serde(deserialize_with = "super::number::deserialize")] + pub bucket: u32, + #[serde( + deserialize_with = "super::number::boolean", + serialize_with = "litellm_traces::wire::serialize_flag" + )] + pub failed: bool, + pub value: String, + #[serde(deserialize_with = "super::number::deserialize")] + pub runs: u64, +} + +#[derive(Debug, Deserialize, Serialize)] +pub(crate) struct RunCountRow(#[serde(with = "RunCountEncoding")] pub RunCount); + +pub(crate) struct RunCounts; + +impl Query for RunCounts { + type Params = RunCountsParams; + type Row = RunCountRow; + + const SQL: &'static str = over_matching_runs!("\n", include_str!("../../query/run_counts.sql")); +} + +#[derive(Debug, Default, Serialize)] +struct SpanKeyset { + as_of_ms: u64, + after_team: String, + after_key: String, + after_trace: String, + after_span: String, + limit: u32, +} + +impl From<&SpanQuery> for SpanKeyset { + fn from(query: &SpanQuery) -> Self { + let after = query.after.clone().unwrap_or_default(); + Self { + as_of_ms: query.as_of_ms, + after_team: after.team_id, + after_key: after.api_key_hash, + after_trace: after.trace_id, + after_span: after.span_id, + limit: query.limit, + } + } +} + +#[derive(Debug, Serialize)] +pub(crate) struct TraceSpansParams { + #[serde(flatten)] + access: AccessParams, + trace_id: String, + trace_ref: String, + #[serde(flatten)] + keyset: SpanKeyset, +} + +#[derive(Debug, Serialize)] +pub(crate) struct RunSpansParams { + #[serde(flatten)] + access: AccessParams, + trace_refs: Vec, + start_ms: i64, + end_ms: i64, + #[serde(flatten)] + keyset: SpanKeyset, +} + +pub(crate) enum SpansParams { + Trace(TraceSpansParams), + Runs(RunSpansParams), +} + +impl SpansParams { + pub(crate) fn new(access: &QueryScope, query: &SpanQuery) -> Self { + let keyset = query.into(); + match &query.selection { + SpanSelection::Trace { + trace_id, + trace_ref, + } => Self::Trace(TraceSpansParams { + access: access.into(), + trace_id: trace_id.clone(), + trace_ref: trace_ref.clone(), + keyset, + }), + SpanSelection::Runs { trace_refs, window } => Self::Runs(RunSpansParams { + access: access.into(), + trace_refs: trace_refs.clone(), + start_ms: window.start, + end_ms: window.end, + keyset, + }), + } + } +} + +#[derive(Deserialize, Serialize)] +#[serde(remote = "SpanRow")] +struct SpanRowEncoding { #[serde(default)] pub trace_id: String, pub span_id: String, @@ -132,77 +366,124 @@ struct TraceSpansRowEncoding { } #[derive(Debug, Deserialize, Serialize)] -pub struct TraceSpansRow(#[serde(with = "TraceSpansRowEncoding")] pub contracts::TraceSpansRow); +pub(crate) struct SpanRowWire(#[serde(with = "SpanRowEncoding")] pub SpanRow); -pub use contracts::SpanDetailParams; - -pub use contracts::SpanDetailRow; - -#[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::SpanErrorParams")] -struct SpanErrorParamsEncoding { - #[serde(flatten)] - pub access: contracts::ReadAccessParams, - pub trace_id: String, - pub trace_ref: String, - pub span_id: String, - #[serde(deserialize_with = "super::number::deserialize")] - pub error_offset: u64, - pub error_version: String, +macro_rules! span_page { + ($selection:literal) => { + owned!( + "\nSELECT * FROM (\n", + include_str!("../../query/span_columns.sql"), + include_str!(concat!("../../query/", $selection)), + include_str!("../../query/span_page.sql"), + ) + }; } -#[derive(Debug, Deserialize, Serialize)] -pub struct SpanErrorParams( - #[serde(with = "SpanErrorParamsEncoding")] pub contracts::SpanErrorParams, -); +pub(crate) struct TraceSpans; -impl From for SpanErrorParams { - fn from(value: contracts::SpanErrorParams) -> Self { - Self(value) +impl Query for TraceSpans { + type Params = TraceSpansParams; + type Row = SpanRowWire; + + const SQL: &'static str = span_page!("trace_spans.sql"); +} + +pub(crate) struct RunSpans; + +impl Query for RunSpans { + type Params = RunSpansParams; + type Row = SpanRowWire; + + const SQL: &'static str = span_page!("run_spans.sql"); +} + +#[derive(Debug, Serialize)] +pub(crate) struct SpanTextParams { + #[serde(flatten)] + access: AccessParams, + trace_id: String, + trace_ref: String, + span_id: String, + part: &'static str, + offset: u64, + bounded: u8, + max_chars: u64, +} + +impl SpanTextParams { + pub(crate) fn new(access: &QueryScope, query: &SpanTextQuery) -> Self { + Self { + access: access.into(), + trace_id: query.trace_id.clone(), + trace_ref: query.trace_ref.clone(), + span_id: query.span_id.clone(), + part: query.part.into(), + offset: query.offset, + bounded: query.max_chars.is_some().into(), + max_chars: query.max_chars.unwrap_or(0), + } } } #[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::SpanErrorRow")] -struct SpanErrorRowEncoding { - pub span_id: String, - pub message: String, +#[serde(remote = "SpanText")] +struct SpanTextEncoding { + pub text: String, #[serde(deserialize_with = "super::number::deserialize")] pub total_chars: u64, pub version: String, } #[derive(Debug, Deserialize, Serialize)] -pub struct SpanErrorRow(#[serde(with = "SpanErrorRowEncoding")] pub contracts::SpanErrorRow); +pub(crate) struct SpanTextRow(#[serde(with = "SpanTextEncoding")] pub SpanText); -#[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::SpendByResponseIdsParams")] -struct SpendByResponseIdsParamsEncoding { - #[serde(flatten)] - pub access: contracts::ReadAccessParams, - pub response_ids: Vec, - pub request_ids: Vec, - pub trace_ids: Vec, - #[serde(deserialize_with = "super::number::deserialize")] - pub start_ms: i64, - #[serde(deserialize_with = "super::number::deserialize")] - pub end_ms: i64, +pub(crate) struct SpanTexts; + +impl Query for SpanTexts { + type Params = SpanTextParams; + type Row = SpanTextRow; + + const SQL: &'static str = owned!("\n", include_str!("../../query/span_text.sql")); } -#[derive(Debug, Deserialize, Serialize)] -pub struct SpendByResponseIdsParams( - #[serde(with = "SpendByResponseIdsParamsEncoding")] pub contracts::SpendByResponseIdsParams, -); +#[derive(Debug, Serialize)] +pub(crate) struct CallsParams { + #[serde(flatten)] + access: AccessParams, + start_ms: i64, + end_ms: i64, + response_ids: Vec, + request_ids: Vec, + trace_ids: Vec, + has_cursor: u8, + after_team: String, + after_ms: i64, + after_id: String, + limit: u32, +} -impl From for SpendByResponseIdsParams { - fn from(value: contracts::SpendByResponseIdsParams) -> Self { - Self(value) +impl CallsParams { + pub(crate) fn new(access: &QueryScope, query: &CallQuery) -> Self { + let after = query.after.clone().unwrap_or_default(); + Self { + access: access.into(), + start_ms: query.window.start, + end_ms: query.window.end, + response_ids: query.response_ids.clone(), + request_ids: query.request_ids.clone(), + trace_ids: query.trace_ids.clone(), + has_cursor: query.after.is_some().into(), + after_team: after.team_id, + after_ms: after.start_ms, + after_id: after.request_id, + limit: query.limit, + } } } #[derive(Deserialize, Serialize)] -#[serde(remote = "contracts::SpendByResponseIdsRow")] -struct SpendByResponseIdsRowEncoding { +#[serde(remote = "CallRow")] +struct CallRowEncoding { pub request_id: String, pub litellm_call_id: String, pub response_id: String, @@ -219,110 +500,68 @@ struct SpendByResponseIdsRowEncoding { } #[derive(Debug, Deserialize, Serialize)] -pub struct SpendByResponseIdsRow( - #[serde(with = "SpendByResponseIdsRowEncoding")] pub contracts::SpendByResponseIdsRow, -); +pub(crate) struct CallRowWire(#[serde(with = "CallRowEncoding")] pub CallRow); -pub struct ListTraces; +pub(crate) struct Calls; -impl Query for ListTraces { - type Params = ListTracesParams; - type Row = ListTracesRow; +impl Query for Calls { + type Params = CallsParams; + type Row = CallRowWire; - 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, - #[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 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 { - type Params = TraceSpansParams; - type Row = TraceSpansRow; - - const SQL: &'static str = include_str!("../../query/trace_spans.sql"); -} - -pub struct SpanDetail; - -impl Query for SpanDetail { - type Params = SpanDetailParams; - type Row = SpanDetailRow; - - const SQL: &'static str = include_str!("../../query/span_detail.sql"); -} - -pub struct SpanError; - -impl Query for SpanError { - type Params = SpanErrorParams; - type Row = SpanErrorRow; - - const SQL: &'static str = include_str!("../../query/span_error.sql"); -} - -pub struct SpendByResponseIds; - -impl Query for SpendByResponseIds { - type Params = SpendByResponseIdsParams; - type Row = SpendByResponseIdsRow; - - const SQL: &'static str = include_str!("../../query/spend_by_response_ids.sql"); -} - -pub use contracts::{TraceIdentityParams, TraceIdentityRow}; - -pub struct TraceIdentity; - -impl Query for TraceIdentity { - type Params = TraceIdentityParams; - type Row = TraceIdentityRow; - const SQL: &'static str = include_str!("../../query/trace_identity.sql"); + const SQL: &'static str = owned!("\n", include_str!("../../query/calls.sql")); } #[cfg(test)] mod tests { - use super::*; + use litellm_traces::search::{FieldFilter, RunField}; use rstest::rstest; use serde_json::{Value, json}; - fn round_trip(wire: Value, quoted: bool) { + use super::*; + + #[rstest] + #[case::plain("plan", "%plan%")] + #[case::like_metacharacters("50%_off\\", "%50\\%\\_off\\\\%")] + fn text_terms_match_literally_anywhere(#[case] term: &str, #[case] pattern: &str) { + let columns = SearchColumns::from(&RunSearch { + text: vec![term.into()], + filters: Vec::new(), + }); + assert_eq!(columns.text, [pattern]); + } + + #[rstest] + #[case::glob("res*er", "res%er")] + #[case::glob_escapes_the_rest("gpt_4*", "gpt\\_4%")] + #[case::exact("plan trip", "plan trip")] + fn field_patterns_turn_globs_into_like(#[case] pattern: &str, #[case] like: &str) { + let columns = SearchColumns::from(&RunSearch { + text: Vec::new(), + filters: vec![ + FieldFilter { + field: RunField::Model, + pattern: pattern.into(), + exclude: true, + }, + FieldFilter { + field: RunField::Agent, + pattern: "x".into(), + exclude: false, + }, + ], + }); + assert_eq!(columns.filter_fields, ["model", "agent"]); + assert_eq!(columns.filter_patterns, [like, "x"]); + assert_eq!(columns.filter_modes, ["exclude", "include"]); + } + + fn decoded(wire: Value, quoted: bool) -> Value { let encoded = Value::Object( wire.as_object() .unwrap() .iter() .map(|(name, value)| { - let encoded = if quoted && value.is_number() && name != "all_teams" { + let encoded = if quoted && value.is_number() { json!(value.to_string()) } else { value.clone() @@ -331,62 +570,34 @@ mod tests { }) .collect(), ); - let decoded: T = serde_json::from_value(encoded).unwrap(); - assert_eq!(serde_json::to_value(decoded).unwrap(), wire); + serde_json::to_value(serde_json::from_value::(encoded).unwrap()).unwrap() } #[rstest] #[case::unquoted(false)] #[case::quoted(true)] - fn rows_decode_into_neutral_contracts(#[case] quoted: bool) { - round_trip::( - 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::( - 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::( - json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}), - quoted, - ); - round_trip::( - json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}), - quoted, - ); - round_trip::( - json!({"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}), - quoted, - ); + fn rows_decode_clickhouse_numbers(#[case] quoted: bool) { + let run = 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}); + assert_eq!(decoded::(run.clone(), quoted), run); + let span = 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"}); + assert_eq!(decoded::(span.clone(), quoted), span); + let count = json!({"bucket": 2, "failed": 1, "value": "v", "runs": u64::MAX}); + assert_eq!(decoded::(count.clone(), quoted), count); + let text = json!({"text": "error", "total_chars": u64::MAX, "version": "version"}); + assert_eq!(decoded::(text.clone(), quoted), text); + let call = json!({"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}); + assert_eq!(decoded::(call.clone(), quoted), call); } - #[rstest] - #[case::unquoted(false)] - #[case::quoted(true)] - fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) { - round_trip::( - 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::( - 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::( - 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] cost: Value, #[case] expected: Option, ) { - let row: SpendByResponseIdsRow = serde_json::from_value(json!({ + let row: CallRowWire = serde_json::from_value(json!({ "request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": cost, "start_ms": 0 @@ -394,12 +605,13 @@ mod tests { .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::(json!({ + fn spend_rows_reject_invalid_cost(#[case] cost: Value) { + let row = serde_json::from_value::(json!({ "request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": cost, "start_ms": 0 diff --git a/litellm-rust/crates/traces-clickhouse/src/query/number.rs b/litellm-rust/crates/traces-clickhouse/src/query/number.rs index 9283903fee1..27c88bd3768 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/number.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/number.rs @@ -57,9 +57,10 @@ pub(super) fn boolean<'de, D: Deserializer<'de>>(deserializer: D) -> Result, ) { - let row = serde_json::from_value::(serde_json::json!({ - "span_id": "span", "message": "error", "total_chars": value, "version": "hash" + let row = serde_json::from_value::(serde_json::json!({ + "text": "error", "total_chars": value, "version": "hash" })); match expected { Some(value) => assert_eq!(row.unwrap().0.total_chars, value), diff --git a/litellm-rust/crates/traces-clickhouse/src/query_access.rs b/litellm-rust/crates/traces-clickhouse/src/query_access.rs index e6ca322098e..cb6b1d25f7a 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query_access.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query_access.rs @@ -5,9 +5,8 @@ use litellm_http::Client; use litellm_storage_clickhouse::READ_LIMITS; use litellm_traces::QueryScope; use moka::future::Cache; -use strum::IntoEnumIterator; - use sha2::{Digest, Sha256}; +use strum::IntoEnumIterator; use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use super::{Connection, Error, TraceTable}; @@ -126,7 +125,7 @@ impl QueryReaders { ) .await?; for table in TraceTable::iter() { - let predicate = predicate(scope, table); + let predicate = crate::access::predicate(scope, table); self.execute( client, format!( @@ -175,35 +174,6 @@ impl QueryReaders { } } -fn predicate(scope: &QueryScope, table: TraceTable) -> String { - let team = match table { - TraceTable::OtelTraces | TraceTable::AgentTracesByKey => "TeamId", - TraceTable::SpendLogs => "team_id", - }; - match scope { - QueryScope::All => "1".to_owned(), - QueryScope::Owned { user_id, team_ids } => { - let owner = literal(user_id); - let user_clause = match table { - TraceTable::OtelTraces => format!("UserId = {owner}"), - TraceTable::AgentTracesByKey => format!("UserIds = [{owner}]"), - TraceTable::SpendLogs => format!("user = {owner}"), - }; - let teams = team_ids - .iter() - .map(|value| literal(value)) - .collect::>() - .join(", "); - let team_clause = if team_ids.is_empty() { - "0".to_owned() - } else { - format!("{team} IN ({teams})") - }; - format!("({owner} != '' AND {user_clause}) OR ({team_clause})") - } - } -} - fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result { let mut mac = Hmac::::new_from_slice(secret.as_bytes()).map_err(|_| Error::MissingSecret)?; @@ -211,35 +181,3 @@ fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result String { - format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'")) -} - -#[cfg(test)] -mod tests { - use super::*; - use rstest::rstest; - - #[rstest] - #[case::otel(TraceTable::OtelTraces, "TeamId", "UserId = ''")] - #[case::agent(TraceTable::AgentTracesByKey, "TeamId", "UserIds = ['']")] - #[case::spend(TraceTable::SpendLogs, "team_id", "user = ''")] - fn predicates_preserve_scope_and_escape_values( - #[case] table: TraceTable, - #[case] team: &str, - #[case] user: &str, - ) { - assert_eq!(predicate(&QueryScope::All, table), "1"); - assert_eq!( - predicate( - &QueryScope::Owned { - user_id: String::new(), - team_ids: vec!["team'\\".into()] - }, - table - ), - format!("('' != '' AND {user}) OR ({team} IN ('team\\'\\\\'))") - ); - } -} diff --git a/litellm-rust/crates/traces-clickhouse/src/reads.rs b/litellm-rust/crates/traces-clickhouse/src/reads.rs index eaaeb4fef82..32562972ee7 100644 --- a/litellm-rust/crates/traces-clickhouse/src/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/src/reads.rs @@ -1,28 +1,22 @@ use litellm_http::Client; use litellm_storage_clickhouse::{Error as StorageError, Query, fetch}; -use litellm_traces::query::named as contracts; -use litellm_traces_cache::{StoreError, TraceStore}; +use litellm_traces::{ + QueryScope, + store::{ + CallQuery, CallRow, RunCount, RunCountQuery, RunQuery, RunRow, SpanQuery, SpanRow, + SpanText, SpanTextQuery, + }, +}; +use litellm_traces_cache::{StoreError, StoreResult, TraceStore}; use crate::{ Connection, Error, query::named::{ - ListTracesParams, ListTracesRow, SpanDetail as SpanDetailQuery, SpanError, SpanErrorParams, - SpendByResponseIdsParams, TraceIdentity, TracePageSpansParams, + Calls, CallsParams, RunCounts, RunCountsParams, RunSpans, Runs, RunsParams, SpanTextParams, + SpanTexts, SpansParams, TraceSpans, }, }; -struct RunCandidates; - -impl Query for RunCandidates { - type Params = ListTracesParams; - type Row = ListTracesRow; - const SQL: &'static str = concat!( - "SELECT * EXCEPT (request_ids), [] AS request_ids FROM (", - include_str!("../query/list_traces.sql"), - ") ORDER BY start_ms DESC, trace_ref DESC" - ); -} - pub struct ClickHouseTraces { client: Client, connection: Connection, @@ -32,6 +26,15 @@ impl ClickHouseTraces { pub fn new(client: Client, connection: Connection) -> Self { Self { client, connection } } + + async fn fetch(&self, params: &Q::Params) -> StoreResult, Error> { + fetch::(&self.client, &self.connection, params) + .await + .map_err(|error| match error { + StorageError::ResponseTooLarge => StoreError::TooLarge, + error => StoreError::Failed(Error::Storage(error)), + }) + } } impl TraceStore for ClickHouseTraces { @@ -41,85 +44,53 @@ impl TraceStore for ClickHouseTraces { self.connection.url().as_str() } - async fn trace_refs( - &self, - params: &contracts::TraceIdentityParams, - ) -> Result, StoreError> { - fetch::(&self.client, &self.connection, params) - .await - .map(|rows| rows.into_iter().map(|row| row.trace_ref).collect()) - .map_err(failed) + async fn runs(&self, access: &QueryScope, query: &RunQuery) -> StoreResult, Error> { + let rows = self.fetch::(&RunsParams::new(access, query)).await?; + Ok(rows.into_iter().map(|row| row.0).collect()) } - async fn list_runs( + async fn run_counts( &self, - params: &contracts::ListTracesParams, - ) -> Result, StoreError> { - let storage_params = ListTracesParams::from(params.clone()); - match fetch::(&self.client, &self.connection, &storage_params).await { - Ok(rows) => Ok(rows.into_iter().map(|row| row.0).collect()), - Err(StorageError::ResponseTooLarge) => Err(StoreError::TooLarge), - Err(error) => Err(failed(error)), - } + access: &QueryScope, + query: &RunCountQuery, + ) -> StoreResult, Error> { + let rows = self + .fetch::(&RunCountsParams::new(access, query)) + .await?; + Ok(rows.into_iter().map(|row| row.0).collect()) } - async fn trace_spans( + async fn spans( &self, - params: &contracts::TraceSpansParams, - snapshot_ms: u64, - ) -> Result, StoreError> { - crate::span_batches::read_spans(&self.client, &self.connection, params.clone(), snapshot_ms) - .await + access: &QueryScope, + query: &SpanQuery, + ) -> StoreResult, Error> { + let rows = match SpansParams::new(access, query) { + SpansParams::Trace(params) => self.fetch::(¶ms).await?, + SpansParams::Runs(params) => self.fetch::(¶ms).await?, + }; + Ok(rows.into_iter().map(|row| row.0).collect()) } - async fn run_spans( + async fn span_text( &self, - params: &contracts::TracePageSpansParams, - snapshot_ms: u64, - ) -> Result, StoreError> { - crate::span_batches::read_list_spans( - &self.client, - &self.connection, - TracePageSpansParams::from(params.clone()), - snapshot_ms, - ) - .await + access: &QueryScope, + query: &SpanTextQuery, + ) -> StoreResult, Error> { + let rows = self + .fetch::(&SpanTextParams::new(access, query)) + .await?; + Ok(rows.into_iter().next().map(|row| row.0)) } - async fn spend( + async fn calls( &self, - params: &contracts::SpendByResponseIdsParams, - ) -> Result, StoreError> { - crate::span_batches::read_spend( - &self.client, - &self.connection, - SpendByResponseIdsParams::from(params.clone()), - ) - .await - } - - async fn span_detail( - &self, - params: &contracts::SpanDetailParams, - ) -> Result, StoreError> { - match fetch::(&self.client, &self.connection, params).await { - Ok(rows) => Ok(rows.into_iter().next()), - Err(error) => Err(failed(error)), - } - } - - async fn span_error( - &self, - params: &contracts::SpanErrorParams, - ) -> Result, StoreError> { - let storage_params = SpanErrorParams::from(params.clone()); - match fetch::(&self.client, &self.connection, &storage_params).await { - Ok(rows) => Ok(rows.into_iter().next().map(|row| row.0)), - Err(error) => Err(failed(error)), - } + access: &QueryScope, + query: &CallQuery, + ) -> StoreResult, Error> { + let rows = self + .fetch::(&CallsParams::new(access, query)) + .await?; + Ok(rows.into_iter().map(|row| row.0).collect()) } } - -fn failed(error: StorageError) -> StoreError { - StoreError::Failed(Error::Storage(error)) -} diff --git a/litellm-rust/crates/traces-clickhouse/src/schema.rs b/litellm-rust/crates/traces-clickhouse/src/schema.rs index 6a39bd24041..48df69123a4 100644 --- a/litellm-rust/crates/traces-clickhouse/src/schema.rs +++ b/litellm-rust/crates/traces-clickhouse/src/schema.rs @@ -1,7 +1,8 @@ +use std::time::Duration; + use litellm_http::Client; use litellm_migrate::Migration; use serde::Serialize; -use std::time::Duration; use super::{Connection, Error}; diff --git a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs b/litellm-rust/crates/traces-clickhouse/src/span_batches.rs deleted file mode 100644 index 5e406d67cfa..00000000000 --- a/litellm-rust/crates/traces-clickhouse/src/span_batches.rs +++ /dev/null @@ -1,352 +0,0 @@ -//! Keyset-paged reads that shrink their page when ClickHouse rejects a response as too large and -//! stop accumulating once a graph exceeds the interactive budget. - -use std::{future::Future, marker::PhantomData}; - -use litellm_http::Client; -use litellm_storage_clickhouse::{Query, fetch}; -use litellm_traces::query::named as contracts; -use litellm_traces_cache::{MAX_GRAPH_BYTES, MAX_GRAPH_SPANS, StoreError}; -use serde::{Serialize, de::DeserializeOwned}; - -use crate::{ - Connection, Error, - query::named::{SpendByResponseIdsParams, SpendByResponseIdsRow, TraceSpansRow}, -}; - -const PAGE_SIZE: u32 = 256; - -#[derive(Default)] -struct ReadBudget { - bytes: usize, - rows: usize, -} - -impl ReadBudget { - fn reserve(&mut self, bytes: usize) -> Result<(), StoreError> { - self.bytes = self.bytes.saturating_add(bytes); - if self.bytes > MAX_GRAPH_BYTES || self.rows == MAX_GRAPH_SPANS { - return Err(StoreError::TooLarge); - } - self.rows += 1; - Ok(()) - } - - fn record(&mut self, row: &impl Serialize) -> Result<(), StoreError> { - let bytes = - serde_json::to_vec(row).map_err(|_| StoreError::Failed(Error::InvalidResponse))?; - self.reserve(bytes.len()) - } -} - -/// One keyset position in a paged query: the SQL reads the cursor fields of `Self` plus the -/// `page_size` that [`Batch`] adds. -trait Keyset: Serialize + Sized + Send + Sync { - type Row: Serialize + DeserializeOwned + Send; - const SQL: &'static str; - - fn after(self, last: &Self::Row) -> Self; -} - -#[derive(Serialize)] -struct Batch { - #[serde(flatten)] - keyset: K, - page_size: u32, -} - -trait PageSource { - fn page( - &self, - batch: &Batch, - ) -> impl Future, litellm_storage_clickhouse::Error>> + Send; -} - -struct Paged(PhantomData); - -impl Query for Paged { - type Params = Batch; - type Row = K::Row; - const SQL: &'static str = K::SQL; -} - -/// Reads every row after `keyset`. A page ClickHouse rejects as too large is retried at half the -/// size, and the smaller page is kept for the rest of the read because row sizes within one graph -/// rarely shrink again. Halving a one-row page means a single row exceeds the response limit. -async fn read_all>( - source: &S, - keyset: K, -) -> Result, StoreError> { - let mut batch = Batch { - keyset, - page_size: PAGE_SIZE, - }; - let mut rows = Vec::new(); - let mut budget = ReadBudget::default(); - loop { - let page = match source.page(&batch).await { - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if batch.page_size > 1 => { - batch.page_size /= 2; - continue; - } - Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => { - return Err(StoreError::TooLarge); - } - result => result.map_err(|error| StoreError::Failed(Error::Storage(error)))?, - }; - let complete = page.len() < batch.page_size as usize; - for row in &page { - budget.record(row)?; - } - if let Some(last) = page.last() { - batch.keyset = batch.keyset.after(last); - } - rows.extend(page); - if complete { - return Ok(rows); - } - } -} - -struct ClickHouse<'a> { - client: &'a Client, - connection: &'a Connection, -} - -impl PageSource for ClickHouse<'_> { - fn page( - &self, - batch: &Batch, - ) -> impl Future, litellm_storage_clickhouse::Error>> + Send { - fetch::>(self.client, self.connection, batch) - } -} - -async fn read_paged( - client: &Client, - connection: &Connection, - keyset: K, -) -> Result, StoreError> { - let source = ClickHouse { client, connection }; - read_all(&source, keyset).await -} - -fn by_start(mut rows: Vec) -> Vec { - rows.sort_by_key(|row| row.start_ns); - rows -} - -#[derive(Serialize)] -struct SpanKeyset { - #[serde(flatten)] - trace: contracts::TraceSpansParams, - after_span_id: String, - snapshot_ms: u64, -} - -impl Keyset for SpanKeyset { - type Row = TraceSpansRow; - const SQL: &'static str = include_str!("../query/trace_span_batch.sql"); - - fn after(self, last: &TraceSpansRow) -> Self { - Self { - after_span_id: last.0.span_id.clone(), - ..self - } - } -} - -pub(crate) async fn read_spans( - client: &Client, - connection: &Connection, - trace: contracts::TraceSpansParams, - snapshot_ms: u64, -) -> Result, StoreError> { - let keyset = SpanKeyset { - trace, - after_span_id: String::new(), - snapshot_ms, - }; - let rows = read_paged(client, connection, keyset).await?; - Ok(by_start(rows.into_iter().map(|row| row.0).collect())) -} - -#[derive(Serialize)] -struct ListSpanKeyset { - #[serde(flatten)] - runs: crate::query::named::TracePageSpansParams, - after_team: String, - after_key: String, - after_trace: String, - after_span: String, - snapshot_ms: u64, -} - -impl Keyset for ListSpanKeyset { - type Row = TraceSpansRow; - const SQL: &'static str = include_str!("../query/trace_list_span_batch.sql"); - - fn after(self, last: &TraceSpansRow) -> Self { - Self { - after_team: last.0.team_id.clone(), - after_key: last.0.api_key_hash.clone(), - after_trace: last.0.trace_id.clone(), - after_span: last.0.span_id.clone(), - ..self - } - } -} - -pub(crate) async fn read_list_spans( - client: &Client, - connection: &Connection, - runs: crate::query::named::TracePageSpansParams, - snapshot_ms: u64, -) -> Result, StoreError> { - let keyset = ListSpanKeyset { - runs, - after_team: String::new(), - after_key: String::new(), - after_trace: String::new(), - after_span: String::new(), - snapshot_ms, - }; - let rows = read_paged(client, connection, keyset).await?; - Ok(by_start(rows.into_iter().map(|row| row.0).collect())) -} - -#[derive(Serialize)] -struct SpendKeyset { - #[serde(flatten)] - lookup: SpendByResponseIdsParams, - has_cursor: u8, - after_team: String, - after_ms: i64, - after_id: String, -} - -impl Keyset for SpendKeyset { - type Row = SpendByResponseIdsRow; - const SQL: &'static str = include_str!("../query/spend_batch.sql"); - - fn after(self, last: &SpendByResponseIdsRow) -> Self { - Self { - has_cursor: 1, - after_team: last.0.team_id.clone(), - after_ms: last.0.start_ms, - after_id: last.0.request_id.clone(), - ..self - } - } -} - -pub(crate) async fn read_spend( - client: &Client, - connection: &Connection, - lookup: SpendByResponseIdsParams, -) -> Result, StoreError> { - let keyset = SpendKeyset { - lookup, - has_cursor: 0, - after_team: String::new(), - after_ms: 0, - after_id: String::new(), - }; - let rows = read_paged(client, connection, keyset).await?; - Ok(rows.into_iter().map(|row| row.0).collect()) -} - -#[cfg(test)] -mod tests { - use std::sync::Mutex; - - use rstest::rstest; - - use super::*; - - #[rstest] - #[case::byte_boundary(MAX_GRAPH_BYTES - 1, 0, 1, false)] - #[case::byte_overflow(MAX_GRAPH_BYTES - 1, 0, 2, true)] - #[case::integer_overflow(MAX_GRAPH_BYTES, 0, usize::MAX, true)] - #[case::row_boundary(0, MAX_GRAPH_SPANS - 1, 1, false)] - #[case::row_overflow(0, MAX_GRAPH_SPANS, 1, true)] - fn accumulation_stops_at_the_graph_budget( - #[case] bytes: usize, - #[case] rows: usize, - #[case] next: usize, - #[case] rejected: bool, - ) { - let mut budget = ReadBudget { bytes, rows }; - assert_eq!(budget.reserve(next).is_err(), rejected); - } - - #[derive(Serialize)] - struct Numbers { - after: u32, - } - - impl Keyset for Numbers { - type Row = u32; - const SQL: &'static str = ""; - - fn after(self, last: &u32) -> Self { - Self { after: *last } - } - } - - /// A table of `total` rows whose transport rejects any page larger than `largest_page`. - struct Table { - total: u32, - largest_page: u32, - requests: Mutex>, - } - - impl PageSource for Table { - async fn page( - &self, - batch: &Batch, - ) -> Result, litellm_storage_clickhouse::Error> { - self.requests.lock().unwrap().push(batch.page_size); - if batch.page_size > self.largest_page { - return Err(litellm_storage_clickhouse::Error::ResponseTooLarge); - } - let end = (batch.keyset.after + batch.page_size).min(self.total); - Ok((batch.keyset.after + 1..=end).collect()) - } - } - - #[rstest] - #[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])] - #[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])] - #[tokio::test] - async fn a_rejected_page_size_is_not_retried( - #[case] total: u32, - #[case] largest_page: u32, - #[case] requests: &[u32], - ) { - let table = Table { - total, - largest_page, - requests: Mutex::new(Vec::new()), - }; - let rows = read_all(&table, Numbers { after: 0 }).await.unwrap(); - assert_eq!(rows, (1..=total).collect::>()); - assert_eq!(table.requests.lock().unwrap().as_slice(), requests); - } - - #[rstest] - #[tokio::test] - async fn a_single_oversized_row_fails_the_read() { - let table = Table { - total: 10, - largest_page: 0, - requests: Mutex::new(Vec::new()), - }; - let result = read_all(&table, Numbers { after: 0 }).await; - assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}"); - assert_eq!( - table.requests.lock().unwrap().as_slice(), - &[256, 128, 64, 32, 16, 8, 4, 2, 1] - ); - } -} diff --git a/litellm-rust/crates/traces-clickhouse/src/sql.rs b/litellm-rust/crates/traces-clickhouse/src/sql.rs index dfa0618773a..40eb278c76c 100644 --- a/litellm-rust/crates/traces-clickhouse/src/sql.rs +++ b/litellm-rust/crates/traces-clickhouse/src/sql.rs @@ -1,13 +1,10 @@ use std::collections::BTreeMap; use litellm_http::Client; +use litellm_storage_clickhouse::{Query, fetch_json}; use litellm_traces::ReadQuery; -use super::{ - Connection, Error, Parameter, - query::{lens::*, named::*}, -}; -use litellm_storage_clickhouse::{Query, fetch_json}; +use super::{Connection, Error, Parameter, query::lens::*}; pub async fn execute_named_read( client: &Client, @@ -16,19 +13,6 @@ pub async fn execute_named_read( parameters: &BTreeMap, ) -> Result { match query { - ReadQuery::ListTraces => named_json::(client, connection, parameters).await, - ReadQuery::TraceIdentity => { - named_json::(client, connection, parameters).await - } - ReadQuery::TraceSpans => named_json::(client, connection, parameters).await, - ReadQuery::TracePageSpans => { - named_json::(client, connection, parameters).await - } - ReadQuery::SpanDetail => named_json::(client, connection, parameters).await, - ReadQuery::SpanError => named_json::(client, connection, parameters).await, - ReadQuery::SpendByResponseIds => { - named_json::(client, connection, parameters).await - } ReadQuery::Availability => { named_json::(client, connection, parameters).await } @@ -57,19 +41,21 @@ where #[cfg(test)] mod tests { - use super::*; use rstest::rstest; + use super::*; + #[rstest] - #[case::missing_span(serde_json::json!({}))] - #[case::negative_offset(serde_json::json!({"span_id": "span", "error_offset": -1, "error_version": ""}))] - #[case::overflow(serde_json::json!({"span_id": "span", "error_offset": "18446744073709551616", "error_version": ""}))] + #[case::missing_cursor(serde_json::json!({"offset": 1}))] + #[case::negative_offset(serde_json::json!({"cursor": "", "offset": -1}))] + #[case::overflow(serde_json::json!({"cursor": "", "offset": "4294967296"}))] #[tokio::test] async fn named_read_rejects_invalid_parameters_before_transport( #[case] specific: serde_json::Value, ) { let common = serde_json::json!({ - "all_teams": 1, "user_id": "", "team_ids": [], "trace_id": "trace", "trace_ref": "" + "all_teams": 1, "team": "", "key_hash": "", "source": "traces", "id": "trace", + "record_team": "team", "trace_ref": "" }); let parameters: BTreeMap = common .as_object() @@ -81,7 +67,7 @@ mod tests { let client = Client::no_redirect_for_test(); let connection = Connection::parse("http://127.0.0.1:1").unwrap(); assert!(matches!( - execute_named_read(&client, &connection, ReadQuery::SpanError, ¶meters).await, + execute_named_read(&client, &connection, ReadQuery::Content, ¶meters).await, Err(Error::InvalidParameters) )); } diff --git a/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs b/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs index 6c89d401c85..82473f1727f 100644 --- a/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs +++ b/litellm-rust/crates/traces-clickhouse/src/wire_schema.rs @@ -95,9 +95,10 @@ pub fn schemas() -> BTreeMap<&'static str, Schema> { #[cfg(test)] mod tests { - use super::*; use rstest::rstest; + use super::*; + #[rstest] #[case::zero(json!(0), true)] #[case::quoted_zero(json!("0"), true)] diff --git a/litellm-rust/crates/traces-clickhouse/tests/admin_sql.rs b/litellm-rust/crates/traces-clickhouse/tests/admin_sql.rs index 73b5125929a..5c3303e9e8a 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/admin_sql.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/admin_sql.rs @@ -1,10 +1,11 @@ +use std::collections::BTreeMap; + use litellm_http::Client; use litellm_traces_clickhouse::{ Connection, Error, Parameter, QueryReaders, QueryScope, execute_read, }; use rstest::{fixture, rstest}; use serde_json::Value; -use std::collections::BTreeMap; mod support; use support::{ClickHouseDatabase, database as start_database}; diff --git a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs index 09fd3e78da8..dc806182c1e 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs @@ -1,15 +1,60 @@ use std::{collections::BTreeMap, time::Duration}; use litellm_http::Client; +use litellm_traces::{ + search::RunFilter, + store::{CallQuery, RunCursor, RunQuery, RunSelection, SpanQuery, SpanSelection}, +}; +use litellm_traces_cache::{TraceReader, TraceStore}; use litellm_traces_clickhouse::{ - Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery, - encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements, + ClickHouseTraces, Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, + QueryScope, encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements, }; use rstest::rstest; +use sha2::{Digest, Sha256}; mod support; use support::{ClickHouseDatabase, TestResult, database}; +fn owned(user_id: &str, team_ids: &[&str]) -> QueryScope { + QueryScope::Owned { + user_id: user_id.into(), + team_ids: team_ids.iter().map(|team| (*team).to_owned()).collect(), + } +} + +fn traces(database: &ClickHouseDatabase, connection: &Connection) -> ClickHouseTraces { + ClickHouseTraces::new(database.client.clone(), connection.clone()) +} + +fn trace_ref(team_id: &str, api_key_hash: &str, trace_id: &str) -> String { + format!( + "{:X}", + Sha256::digest(format!("{team_id}\0{api_key_hash}\0{trace_id}")) + ) +} + +async fn list_runs( + database: &ClickHouseDatabase, + connection: &Connection, + access: &QueryScope, + window: std::ops::Range, + after: Option, + limit: u32, +) -> TestResult { + let query = RunQuery { + selection: RunSelection::Matching(RunFilter { + start_ms: window.start, + end_ms: window.end, + search: Default::default(), + }), + after, + limit, + }; + let rows = traces(database, connection).runs(access, &query).await?; + Ok(serde_json::json!({ "data": rows })) +} + async fn insert_rows( database: &ClickHouseDatabase, table: &str, @@ -95,83 +140,37 @@ async fn schema_supports_span_rollups_and_spend_joins( insert_rows(&database, "otel_traces", vec![span]).await?; insert_rows(&database, "spend_logs", vec![spend]).await?; let reader = Connection::reader(&database.url, "trace_test")?; - let detail = - litellm_storage_clickhouse::fetch::( - &database.client, - &reader, - &litellm_traces_clickhouse::query::named::SpanDetailParams { - access: litellm_traces_clickhouse::query::named::ReadAccessParams { - all_teams: false, - user_id: String::new(), - team_ids: vec!["team-1".into()], - }, - trace_id: "trace-1".into(), - trace_ref: String::new(), - span_id: "span-1".into(), + let team = owned("", &["team-1"]); + let detail = TraceReader::new(usize::MAX) + .get_span(&traces(&database, &reader), &team, "trace-1", "span-1", "") + .await? + .ok_or("missing span detail")?; + assert_eq!(detail.input, "hello world"); + assert_eq!(detail.attributes["gen_ai.response.id"], "response-1"); + let request_ids = read_json( + &database, + "SELECT groupArrayArray(RequestIds) AS ids FROM trace_test.agent_traces_by_key \ + WHERE TraceId = 'trace-1'", + ) + .await?; + assert_eq!( + request_ids["data"][0]["ids"], + serde_json::json!(["response-1"]) + ); + let calls = traces(&database, &reader) + .calls( + &team, + &CallQuery { + window: timestamp / 1_000_000 - 1000..timestamp / 1_000_000 + 1000, + response_ids: vec!["response-1".into()], + request_ids: Vec::new(), + trace_ids: Vec::new(), + after: None, + limit: 10, }, ) .await?; - assert_eq!(detail.len(), 1); - assert_eq!(detail[0].input, "hello world"); - assert_eq!(detail[0].attributes["gen_ai.response.id"], "response-1"); - let list_parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("user_id".into(), Parameter::Text(String::new())), - ("team_ids".into(), Parameter::Strings(vec!["team-1".into()])), - ( - "start_ms".into(), - Parameter::Integer(timestamp / 1_000_000 - 1000), - ), - ( - "end_ms".into(), - Parameter::Integer(timestamp / 1_000_000 + 1000), - ), - ("cursor_ms".into(), Parameter::Integer(0)), - ("cursor_trace_id".into(), Parameter::Text(String::new())), - ("limit".into(), Parameter::Integer(10)), - ]); - let listed: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &reader, - ReadQuery::ListTraces, - &list_parameters, - ) - .await?, - )?; - assert_eq!( - listed["data"][0]["request_ids"], - serde_json::json!(["response-1"]) - ); - let spend_parameters = BTreeMap::from([ - ( - "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()])), - ( - "start_ms".into(), - Parameter::Integer(timestamp / 1_000_000 - 1000), - ), - ( - "end_ms".into(), - Parameter::Integer(timestamp / 1_000_000 + 1000), - ), - ]); - let matched: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &reader, - ReadQuery::SpendByResponseIds, - &spend_parameters, - ) - .await?, - )?; - assert_eq!(matched["data"][0]["spend"], 0.125); + assert_eq!(calls[0].spend, Some(0.125)); let body = read_json( &database, "SELECT o.TeamId, o.ApiKeyHash, o.UserId, o.ObservationType, o.InputPreview, s.spend, \ @@ -478,53 +477,19 @@ async fn listed_agent_names_preserve_scope_and_cursor( .collect::, _>>()?; insert_rows(&database, "otel_traces", historical_rows).await?; let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("user_id".into(), Parameter::Text("owner".into())), - ("team_ids".into(), Parameter::Strings(vec![])), - ( - "start_ms".into(), - Parameter::Integer(timestamp / 1_000_000 - 1000), - ), - ( - "end_ms".into(), - Parameter::Integer(timestamp / 1_000_000 + 1000), - ), - ("cursor_ms".into(), Parameter::Integer(0)), - ("cursor_trace_id".into(), Parameter::Text(String::new())), - ("limit".into(), Parameter::Integer(1)), - ]); - let first: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::ListTraces, - ¶meters, - ) - .await?, - )?; - let cursor = first["data"][0]["trace_ref"] - .as_str() - .ok_or("missing cursor")?; - let next_parameters = parameters - .into_iter() - .chain([ - ( - "cursor_ms".into(), - Parameter::Integer(timestamp / 1_000_000), - ), - ("cursor_trace_id".into(), Parameter::Text(cursor.into())), - ]) - .collect(); - let second: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::ListTraces, - &next_parameters, - ) - .await?, - )?; + let owner = owned("owner", &[]); + let window = timestamp / 1_000_000 - 1000..timestamp / 1_000_000 + 1000; + let first = list_runs(&database, &connection, &owner, window.clone(), None, 1).await?; + let after = RunCursor { + start_ms: first["data"][0]["start_ms"] + .as_i64() + .ok_or("missing start")?, + trace_ref: first["data"][0]["trace_ref"] + .as_str() + .ok_or("missing cursor")? + .into(), + }; + let second = list_runs(&database, &connection, &owner, window, Some(after), 1).await?; assert_eq!( first["data"].as_array().ok_or("missing first page")?.len(), 1 @@ -571,14 +536,21 @@ async fn listed_agent_names_preserve_scope_and_cursor( .collect::>(); assert_eq!(counts["shared"], Some(3)); assert_eq!(counts["second"], Some(1)); - for page in [&first, &second] { - assert!( - page["statistics"]["rows_read"] - .as_u64() - .ok_or("missing read statistics")? - < 5000 - ); - } + execute_write(&database, "SYSTEM FLUSH LOGS").await?; + let reads = read_json( + &database, + "SELECT count() AS pages, max(read_rows) AS rows FROM system.query_log \ + WHERE type = 'QueryFinish' AND current_database = 'trace_test' \ + AND position(query, 'page AS (') > 0 AND query NOT LIKE '%system.query_log%'", + ) + .await?; + assert_eq!(reads["data"][0]["pages"], 2); + assert!( + reads["data"][0]["rows"] + .as_u64() + .ok_or("missing read statistics")? + < 5000 + ); Ok(()) } @@ -629,28 +601,15 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields( }]) ); let connection = Connection::configured(&database.url, "trace_test", "default", "")?; - let parameters = BTreeMap::from([ - ("all_teams".into(), Parameter::Integer(0)), - ("user_id".into(), Parameter::Text(String::new())), - ("team_ids".into(), Parameter::Strings(vec!["team-1".into()])), - ( - "start_ms".into(), - Parameter::Integer(day_start / 1_000_000 - 2000), - ), - ("end_ms".into(), Parameter::Integer(day_start / 1_000_000)), - ("cursor_ms".into(), Parameter::Integer(0)), - ("cursor_trace_id".into(), Parameter::Text(String::new())), - ("limit".into(), Parameter::Integer(10)), - ]); - let listed: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::ListTraces, - ¶meters, - ) - .await?, - )?; + let listed = list_runs( + &database, + &connection, + &owned("", &["team-1"]), + day_start / 1_000_000 - 2000..day_start / 1_000_000, + None, + 10, + ) + .await?; assert_eq!( listed["data"][0]["agent_names"], serde_json::json!(["lead", "researcher"]) @@ -880,37 +839,17 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate( let rows = sample["data"].as_array().expect("sample rows"); assert_eq!(rows.len(), 2); assert_ne!(rows[0]["trace_ref"], rows[1]["trace_ref"]); - let identity_params = BTreeMap::from([ - ("trace_id".into(), Parameter::Text("shared".into())), - ("all_teams".into(), Parameter::Integer(0)), - ("user_id".into(), Parameter::Text(String::new())), - ("team_ids".into(), Parameter::Strings(vec!["team".into()])), - ]); - let identities: serde_json::Value = serde_json::from_str( - &execute_named_read( - &database.client, - &connection, - ReadQuery::TraceIdentity, - &identity_params, - ) - .await?, - )?; - assert_eq!(identities["data"].as_array().map(Vec::len), Some(2)); - let user_params = 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, - &user_params, - ) - .await?, - )?; + let by_trace_id = RunQuery { + selection: RunSelection::TraceId("shared".into()), + after: None, + limit: 10, + }; + let store = traces(&database, &connection); + let identities = store.runs(&owned("", &["team"]), &by_trace_id).await?; + assert_eq!(identities.len(), 2); + let identity = serde_json::json!({ + "data": store.runs(&owned("one", &[]), &by_trace_id).await?, + }); assert_eq!(identity["data"].as_array().map(Vec::len), Some(1)); assert!( rows.iter() @@ -1203,67 +1142,65 @@ async fn trace_error_previews_preserve_paginated_diagnostics( .collect::>, _>>()?; insert_rows(&database, "otel_traces", rows).await?; let reader = Connection::reader(&database.url, "trace_test")?; - let mut parameters = BTreeMap::from([ - ( - "trace_id".into(), - Parameter::Text("diagnostic-trace".into()), - ), - ("all_teams".into(), Parameter::Integer(1)), - ("user_id".into(), Parameter::Text(String::new())), - ("team_ids".into(), Parameter::Strings(vec![])), - ("trace_ref".into(), Parameter::Text(String::new())), - ]); - let body = execute_named_read( - &database.client, - &reader, - ReadQuery::TraceSpans, - ¶meters, - ) - .await?; - let response: serde_json::Value = serde_json::from_str(&body)?; - let spans = response["data"].as_array().expect("trace spans"); + let store = traces(&database, &reader); + let reference = trace_ref("", "", "diagnostic-trace"); + let spans = store + .spans( + &QueryScope::All, + &SpanQuery { + selection: SpanSelection::Trace { + trace_id: "diagnostic-trace".into(), + trace_ref: reference.clone(), + }, + as_of_ms: u64::MAX, + after: None, + limit: 1000, + }, + ) + .await?; assert_eq!(spans.len(), span_count); let prefix: String = message.chars().take(128).collect(); assert!(!prefix.is_empty()); assert!( spans .iter() - .all(|span| span["status_message"] == prefix && span["error_truncated"] == 1) + .all(|span| span.status_message == prefix && span.error_truncated) ); - parameters.insert("span_id".into(), Parameter::Text("span-0".into())); - parameters.insert("error_version".into(), Parameter::Text(String::new())); + let reader = TraceReader::new(usize::MAX); let mut recovered = String::new(); + let mut cursor = None; loop { - parameters.insert( - "error_offset".into(), - Parameter::Integer(recovered.chars().count() as i64), - ); - let body = execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters) - .await?; - assert!(body.len() < 128 * 1024); - let response: serde_json::Value = serde_json::from_str(&body)?; - let chunk = response["data"][0]["message"] - .as_str() - .expect("diagnostic chunk"); - assert!(!chunk.is_empty()); - recovered.push_str(chunk); - let version = response["data"][0]["version"] - .as_str() - .expect("diagnostic version"); - parameters.insert("error_version".into(), Parameter::Text(version.into())); - if recovered.chars().count() >= message.chars().count() { + let page = reader + .get_span_error( + &store, + &QueryScope::All, + "diagnostic-trace", + "span-0", + &reference, + cursor.as_deref(), + ) + .await? + .ok_or("missing diagnostic page")?; + assert!(!page.message.is_empty()); + assert!(page.message.chars().count() <= 16_384); + recovered.push_str(&page.message); + let Some(next) = page.next_cursor else { break; - } + }; + cursor = Some(next); } assert_eq!(recovered, message); - parameters.insert("all_teams".into(), Parameter::Integer(0)); - parameters.insert("user_id".into(), Parameter::Text("unrelated-user".into())); - let denied = - execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters).await?; - assert_eq!( - serde_json::from_str::(&denied)?["data"], - serde_json::json!([]) - ); + let denied = reader + .get_span_error( + &store, + &owned("unrelated-user", &[]), + "diagnostic-trace", + "span-0", + &reference, + None, + ) + .await?; + assert!(denied.is_none()); Ok(()) } @@ -1296,30 +1233,36 @@ async fn duplicate_span_preview_matches_diagnostic( .collect::>, _>>()?; insert_rows(&database, "otel_traces", rows).await?; let reader = Connection::reader(&database.url, "trace_test")?; - let parameters = BTreeMap::from([ - ("trace_id".into(), Parameter::Text("duplicate-trace".into())), - ("span_id".into(), Parameter::Text("duplicate-span".into())), - ("all_teams".into(), Parameter::Integer(1)), - ("user_id".into(), Parameter::Text(String::new())), - ("team_ids".into(), Parameter::Strings(vec![])), - ("trace_ref".into(), Parameter::Text(String::new())), - ("error_version".into(), Parameter::Text(String::new())), - ("error_offset".into(), Parameter::Integer(0)), - ]); - let preview = execute_named_read( - &database.client, - &reader, - ReadQuery::TraceSpans, - ¶meters, - ) - .await?; - let diagnostic = - execute_named_read(&database.client, &reader, ReadQuery::SpanError, ¶meters).await?; - let preview: serde_json::Value = serde_json::from_str(&preview)?; - let diagnostic: serde_json::Value = serde_json::from_str(&diagnostic)?; - assert_eq!(preview["data"].as_array().unwrap().len(), 1); - assert_eq!(preview["data"][0]["status_message"], message[..128]); - assert_eq!(diagnostic["data"][0]["message"], message); + let store = traces(&database, &reader); + let reference = trace_ref("", "", "duplicate-trace"); + let preview = store + .spans( + &QueryScope::All, + &SpanQuery { + selection: SpanSelection::Trace { + trace_id: "duplicate-trace".into(), + trace_ref: reference.clone(), + }, + as_of_ms: u64::MAX, + after: None, + limit: 10, + }, + ) + .await?; + let diagnostic = TraceReader::new(usize::MAX) + .get_span_error( + &store, + &QueryScope::All, + "duplicate-trace", + "duplicate-span", + &reference, + None, + ) + .await? + .ok_or("missing diagnostic")?; + assert_eq!(preview.len(), 1); + assert_eq!(preview[0].status_message, message[..128]); + assert_eq!(diagnostic.message, message); Ok(()) } @@ -1845,9 +1788,10 @@ async fn query_help_displays_discovery_truncation( #[rstest] fn field_definitions_match_serialized_normalized_span() { + use std::collections::BTreeSet; + 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"), @@ -1875,22 +1819,17 @@ fn field_definitions_match_serialized_normalized_span() { } #[rstest] -#[case::own_user("owner", vec![], None, vec!["own"])] -#[case::own_user_and_permitted_team("owner", vec!["permitted"], None, vec!["own", "team"])] -#[case::no_identity("", vec![], None, vec![])] -#[case::legacy_key_without_identity("", vec![], Some("request-key"), vec![])] +#[case::own_user("owner", vec![], vec!["own"])] +#[case::own_user_and_permitted_team("owner", vec!["permitted"], vec!["own", "team"])] +#[case::no_identity("", vec![], vec![])] #[tokio::test] -async fn named_and_sql_readers_share_request_log_visibility( +async fn trusted_and_sql_readers_share_request_log_visibility( #[future(awt)] database: TestResult, #[case] user: &str, #[case] teams: Vec<&str>, - #[case] legacy_key: Option<&str>, #[case] expected: Vec<&str>, ) -> TestResult { - use litellm_traces_clickhouse::{ - QueryReaders, QueryScope, - query::named::{ReadAccessParams, SpendByResponseIds, SpendByResponseIdsParams}, - }; + use litellm_traces_clickhouse::QueryReaders; let database = database?; let writer = Connection::writer(&database.url)?; @@ -1905,23 +1844,21 @@ async fn named_and_sql_readers_share_request_log_visibility( .collect::>, _>>()?; insert_rows(&database, "spend_logs", rows).await?; let reader = Connection::reader(&database.url, "trace_test")?; - let params = - SpendByResponseIdsParams::from(litellm_traces::query::named::SpendByResponseIdsParams { - access: serde_json::from_value::(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, - }); - let spend = - litellm_storage_clickhouse::fetch::(&database.client, &reader, ¶ms) - .await?; + let spend = traces(&database, &reader) + .calls( + &owned(user, &teams), + &CallQuery { + window: timestamp / 1_000_000 - 1..timestamp / 1_000_000 + 1, + response_ids: vec!["shared-response".into()], + request_ids: Vec::new(), + trace_ids: Vec::new(), + after: None, + limit: 10, + }, + ) + .await?; let actual: std::collections::BTreeSet<_> = - spend.iter().map(|row| row.0.request_id.as_str()).collect(); + spend.iter().map(|row| row.request_id.as_str()).collect(); let expected: std::collections::BTreeSet<_> = expected.into_iter().collect(); assert_eq!(actual, expected); let scope = QueryScope::Owned { @@ -1981,77 +1918,69 @@ async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_his VALUES ('team', 'export', 'historical', fromUnixTimestamp64Nano({timestamp}), fromUnixTimestamp64Nano({timestamp}), 2, ['response', 'non-llm-id'])" )).await?; ensure_schema(&database.client, &writer, "trace_test", 7).await?; - let params = litellm_traces_clickhouse::query::named::ListTracesParams::from( - litellm_traces::query::named::ListTracesParams { - access: litellm_traces::query::named::ReadAccessParams { - all_teams: false, - user_id: "".into(), - team_ids: vec!["team".into()], - }, - start_ms: timestamp / 1_000_000 - 1, - end_ms: timestamp / 1_000_000 + 1, - cursor_ms: 0, - cursor_trace_id: "".into(), - limit: 10, - }, - ); let reader = Connection::reader(&database.url, "trace_test")?; - let listed = litellm_storage_clickhouse::fetch::< - litellm_traces_clickhouse::query::named::ListTraces, - >(&database.client, &reader, ¶ms) + let window = timestamp / 1_000_000 - 1..timestamp / 1_000_000 + 1; + let listed = list_runs( + &database, + &reader, + &owned("", &["team"]), + window.clone(), + None, + 10, + ) .await?; + let listed = listed["data"].as_array().ok_or("missing runs")?; assert_eq!(listed.len(), 4); - let owned_params = litellm_traces_clickhouse::query::named::ListTracesParams::from( - litellm_traces::query::named::ListTracesParams { - access: litellm_traces::query::named::ReadAccessParams { - user_id: "owner".into(), - team_ids: vec![], - all_teams: false, - }, - ..params.0 - }, - ); - let owned = litellm_storage_clickhouse::fetch::< - litellm_traces_clickhouse::query::named::ListTraces, - >(&database.client, &reader, &owned_params) - .await?; - assert_eq!(owned.len(), 2); - assert!( - owned - .iter() - .all(|row| ["complete", "missing"].contains(&row.0.trace_id.as_str())) - ); + let owner = list_runs(&database, &reader, &owned("owner", &[]), window, None, 10).await?; + let owner: std::collections::BTreeSet<_> = owner["data"] + .as_array() + .ok_or("missing owned runs")? + .iter() + .filter_map(|row| row["trace_id"].as_str()) + .collect(); + assert_eq!(owner, ["complete", "missing"].into()); for row in listed { - match row.0.trace_id.as_str() { + match row["trace_id"].as_str().ok_or("missing trace id")? { "complete" => { - assert_eq!(row.0.user_id, "owner"); - assert_eq!(row.0.request_ids, ["response"]); - assert_eq!(row.0.llm_calls, 2); + assert_eq!(row["user_id"], "owner"); + assert_eq!(row["llm_calls"], 2); } - "missing" | "historical" => assert!(row.0.request_ids.iter().any(String::is_empty)), - "mixed" => assert!(row.0.user_id.is_empty()), + "missing" | "historical" => {} + "mixed" => assert_eq!(row["user_id"], ""), id => panic!("unexpected trace {id}"), } } + let completeness = read_json( + &database, + "SELECT TraceId AS trace, sum(IdentifiedLlmCount) = sum(LlmCount) AS complete, \ + arraySort(groupArrayArray(RequestIds)) AS ids \ + FROM trace_test.agent_traces_by_key GROUP BY TraceId ORDER BY TraceId", + ) + .await?; + assert_eq!( + completeness["data"], + serde_json::json!([ + {"trace": "complete", "complete": 1, "ids": ["response", "response"]}, + {"trace": "historical", "complete": 0, "ids": ["non-llm-id", "response"]}, + {"trace": "missing", "complete": 0, "ids": ["", "extra-id", "response"]}, + {"trace": "mixed", "complete": 1, "ids": ["mine", "other"]}, + ]) + ); assert_eq!(mutation_rows(&database).await?, initial_mutations); Ok(()) } #[rstest] -#[case::admin(1, "", vec![], "own answer")] -#[case::user(0, "owner", vec![], "own answer")] -#[case::team(0, "", vec!["alpha"], "own answer")] -#[case::no_identity(0, "", vec![], "")] +#[case::admin(QueryScope::All, Some("own answer"))] +#[case::user(owned("owner", &[]), Some("own answer"))] +#[case::team(owned("", &["alpha"]), Some("own answer"))] +#[case::no_identity(owned("", &[]), None)] #[tokio::test] async fn agent_final_answer_preserves_visibility_and_trace_ownership( #[future(awt)] database: TestResult, - #[case] all_teams: u8, - #[case] user: &str, - #[case] teams: Vec<&str>, - #[case] expected: &str, + #[case] access: QueryScope, + #[case] expected: Option<&str>, ) -> TestResult { - use litellm_traces_clickhouse::query::named::{ReadAccessParams, SpanDetail, SpanDetailParams}; - let database = database?; let writer = Connection::writer(&database.url)?; ensure_schema(&database.client, &writer, "trace_test", 7).await?; @@ -2098,28 +2027,21 @@ async fn agent_final_answer_preserves_visibility_and_trace_ownership( .collect::, _>>()?; insert_rows(&database, "otel_traces", rows).await?; let reader = Connection::reader(&database.url, "trace_test")?; - let details = litellm_storage_clickhouse::fetch::( - &database.client, - &reader, - &SpanDetailParams { - access: ReadAccessParams { - all_teams: all_teams == 1, - user_id: user.into(), - team_ids: teams.into_iter().map(str::to_owned).collect(), - }, - trace_id: "shared".into(), - trace_ref: String::new(), - span_id: "root".into(), - }, - ) - .await?; - if expected.is_empty() { - assert!(details.is_empty()); - } else { - assert_eq!(details.len(), 1); - assert_eq!(details[0].input, "prompt"); - assert_eq!(details[0].output, expected); - } + let detail = TraceReader::new(usize::MAX) + .get_span( + &traces(&database, &reader), + &access, + "shared", + "root", + &trace_ref("alpha", "one", "shared"), + ) + .await?; + assert_eq!( + detail + .as_ref() + .map(|detail| (detail.input.as_str(), detail.output.as_str())), + expected.map(|output| ("prompt", output)) + ); Ok(()) } diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries.rs b/litellm-rust/crates/traces-clickhouse/tests/queries.rs index 5cb0ee4bfcd..b9235b5c775 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/queries.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/queries.rs @@ -1,12 +1,11 @@ use std::collections::BTreeMap; -use litellm_storage_clickhouse::fetch; -use litellm_traces::query::named as contracts; -use litellm_traces_clickhouse::{ - QueryScope, - query::named::{ListTraces, ListTracesParams, TraceSpans, TraceSpansParams}, - query_help, query_sql, +use litellm_traces::{ + search::RunFilter, + store::{RunQuery, RunRow, RunSelection, SpanQuery, SpanRow, SpanSelection}, }; +use litellm_traces_cache::{StoreResult, TraceStore}; +use litellm_traces_clickhouse::{ClickHouseTraces, Error, QueryScope, query_help, query_sql}; use rstest::{fixture, rstest}; use serde::Deserialize; use serde_json::Value; @@ -133,41 +132,57 @@ fn fixture_clock() -> TestResult { Ok(spans.first().ok_or("missing fixture root")?.start_ns / 1_000_000_000) } -#[fixture] -fn admin_access() -> TestResult { - Ok(serde_json::from_str(include_str!( - "queries/read_access.json" - ))?) +fn newest(limit: u32, after: Option<&RunRow>) -> RunQuery { + RunQuery { + selection: RunSelection::Matching(RunFilter { + start_ms: 0, + end_ms: i64::MAX / 1_000_000, + search: Default::default(), + }), + after: after.map(RunRow::cursor), + limit, + } +} + +async fn trace_spans( + store: &ClickHouseTraces, + trace_id: &str, + trace_ref: &str, +) -> StoreResult, Error> { + let query = SpanQuery { + selection: SpanSelection::Trace { + trace_id: trace_id.into(), + trace_ref: trace_ref.into(), + }, + as_of_ms: u64::MAX, + after: None, + limit: 1000, + }; + let mut spans = store.spans(&QueryScope::All, &query).await?; + spans.sort_by_key(|span| span.start_ns); + Ok(spans) } #[rstest] #[tokio::test] async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate( #[future(awt)] seeded_database: TestResult, - admin_access: TestResult, ) -> TestResult { let fixture = seeded_database?; let reader = fixture .readers .connection(&fixture.database.client, &QueryScope::All, "fixture-secret") .await?; - let params = ListTracesParams::from(contracts::ListTracesParams { - access: admin_access?, - start_ms: 0, - end_ms: i64::MAX / 1_000_000, - cursor_ms: 0, - cursor_trace_id: String::new(), - limit: 10, - }); - let traces = fetch::(&fixture.database.client, &reader, ¶ms).await?; + let store = ClickHouseTraces::new(fixture.database.client.clone(), reader); + let traces = store.runs(&QueryScope::All, &newest(10, None)).await?; assert_eq!( traces .iter() - .map(|row| row.0.api_key_hash.as_str()) + .map(|row| row.api_key_hash.as_str()) .collect::>(), ["key-b", "key-alt", "key-a"] ); - let trace = &traces[2].0; + let trace = &traces[2]; assert_eq!( ( trace.span_count, @@ -179,33 +194,24 @@ async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate( ); assert_eq!((trace.input_tokens, trace.output_tokens), (12, 6)); assert_eq!(trace.input_preview, "Review the change"); - let span_params = TraceSpansParams { - access: params.0.access, - trace_id: trace.trace_id.clone(), - trace_ref: trace.trace_ref.clone(), - }; - let spans = fetch::(&fixture.database.client, &reader, &span_params).await?; + let spans = trace_spans(&store, &trace.trace_id, &trace.trace_ref).await?; assert_eq!( spans .iter() - .map(|row| row.0.name.as_str()) + .map(|row| row.name.as_str()) .collect::>(), ["review", "completion", "lookup"] ); assert!( spans .iter() - .all(|row| row.0.api_key_hash == trace.api_key_hash) + .all(|row| row.api_key_hash == trace.api_key_hash) ); assert_eq!( - ( - spans[1].0.kind, - spans[1].0.input_tokens, - spans[1].0.output_tokens - ), + (spans[1].kind, spans[1].input_tokens, spans[1].output_tokens), (litellm_traces::ObservationType::Llm, 12, 6) ); - assert_eq!(spans[2].0.status_message, "lookup timed out"); + assert_eq!(spans[2].status_message, "lookup timed out"); Ok(()) } @@ -213,33 +219,22 @@ async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate( #[tokio::test] async fn typed_trace_cursor_returns_the_next_fixture_trace( #[future(awt)] seeded_database: TestResult, - admin_access: TestResult, ) -> TestResult { let fixture = seeded_database?; let reader = fixture .readers .connection(&fixture.database.client, &QueryScope::All, "fixture-secret") .await?; - let params = ListTracesParams::from(contracts::ListTracesParams { - access: admin_access?, - start_ms: 0, - end_ms: i64::MAX / 1_000_000, - cursor_ms: 0, - cursor_trace_id: String::new(), - limit: 1, - }); - let first = fetch::(&fixture.database.client, &reader, ¶ms).await?; + let store = ClickHouseTraces::new(fixture.database.client.clone(), reader); + let first = store.runs(&QueryScope::All, &newest(1, None)).await?; assert_eq!(first.len(), 1); - assert_eq!(first[0].0.api_key_hash, "key-b"); - let next_params = ListTracesParams::from(contracts::ListTracesParams { - cursor_ms: first[0].0.start_ms, - cursor_trace_id: first[0].0.trace_ref.clone(), - ..params.0 - }); - let next = fetch::(&fixture.database.client, &reader, &next_params).await?; + assert_eq!(first[0].api_key_hash, "key-b"); + let next = store + .runs(&QueryScope::All, &newest(1, Some(&first[0]))) + .await?; assert_eq!(next.len(), 1); - assert_eq!(next[0].0.api_key_hash, "key-alt"); - assert_ne!(first[0].0.trace_ref, next[0].0.trace_ref); + assert_eq!(next[0].api_key_hash, "key-alt"); + assert_ne!(first[0].trace_ref, next[0].trace_ref); Ok(()) } @@ -250,7 +245,6 @@ async fn typed_trace_cursor_returns_the_next_fixture_trace( #[tokio::test] async fn captured_sdk_exports_round_trip_through_clickhouse( #[future(awt)] migrated_database: TestResult, - admin_access: TestResult, #[case] export: &[u8], ) -> TestResult { let fixture = migrated_database?; @@ -259,37 +253,25 @@ async fn captured_sdk_exports_round_trip_through_clickhouse( .readers .connection(&fixture.database.client, &QueryScope::All, "fixture-secret") .await?; - let params = TraceSpansParams { - access: admin_access?, - trace_id: decoded[0].trace_id.clone(), - trace_ref: String::new(), - }; - let stored = fetch::(&fixture.database.client, &reader, ¶ms).await?; - assert_eq!(stored.len(), decoded.len()); - let list_params = ListTracesParams::from(contracts::ListTracesParams { - access: params.access, - start_ms: 0, - end_ms: i64::MAX / 1_000_000, - cursor_ms: 0, - cursor_trace_id: String::new(), - limit: 10, - }); - let traces = fetch::(&fixture.database.client, &reader, &list_params).await?; + let store = ClickHouseTraces::new(fixture.database.client.clone(), reader); + let traces = store.runs(&QueryScope::All, &newest(10, None)).await?; assert_eq!(traces.len(), 1); + let stored = trace_spans(&store, &decoded[0].trace_id, &traces[0].trace_ref).await?; + assert_eq!(stored.len(), decoded.len()); let roots = decoded .iter() .filter(|span| span.parent_span_id.is_empty()) .collect::>(); assert_eq!(roots.len(), 1); assert_eq!( - traces[0].0.status, + traces[0].status, serde_json::from_value::(serde_json::json!( roots[0].status_code )) .unwrap() ); assert_eq!( - traces[0].0.error_count, + traces[0].error_count, decoded .iter() .filter(|span| span.status_code == "STATUS_CODE_ERROR") @@ -297,7 +279,7 @@ async fn captured_sdk_exports_round_trip_through_clickhouse( ); let by_id: BTreeMap<_, _> = stored .iter() - .map(|row| (row.0.span_id.as_str(), &row.0)) + .map(|row| (row.span_id.as_str(), row)) .collect(); for span in &decoded { let row = by_id diff --git a/litellm-rust/crates/traces-clickhouse/tests/queries/read_access.json b/litellm-rust/crates/traces-clickhouse/tests/queries/read_access.json deleted file mode 100644 index f0af446092e..00000000000 --- a/litellm-rust/crates/traces-clickhouse/tests/queries/read_access.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "all_teams": 1, - "user_id": "", - "team_ids": [ - "team-a", - "team-b" - ] -} diff --git a/litellm-rust/crates/traces-clickhouse/tests/reads.rs b/litellm-rust/crates/traces-clickhouse/tests/reads.rs index e1996d6af04..45c4b34745c 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/reads.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/reads.rs @@ -1,14 +1,22 @@ use std::collections::BTreeMap; use litellm_http::Client; -use litellm_traces::query::named::ReadAccessParams; -use litellm_traces_cache::{ReadError, TraceReader}; +use litellm_traces::search::RunFilter; +use litellm_traces_cache::{PageRequest, ReadError, TraceReader}; use litellm_traces_clickhouse::{ ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows, }; use rstest::rstest; use serde_json::json; +fn all_runs() -> RunFilter { + RunFilter { + start_ms: 0, + end_ms: 2_000_000_000_000, + search: Default::default(), + } +} + #[path = "queries/support.rs"] mod fixtures; mod support; @@ -88,13 +96,20 @@ async fn list_costs_match_each_run_when_response_ids_are_reused( .connection(client, &QueryScope::All, "fixture-secret") .await?; let (reader, store) = make_reader(client, connection); - let access = ReadAccessParams { - all_teams: false, + let access = QueryScope::Owned { user_id: user_id.into(), team_ids: vec!["team-a".into()], }; let page = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 50, + }, + ) .await?; assert_eq!(page.data.len(), runs.len()); for (trace_id, _, cost) in runs { @@ -217,13 +232,20 @@ async fn large_runs_remain_complete_under_default_reader_limits( .connection(client, &QueryScope::All, "fixture-secret") .await?; let (reader, store) = make_reader(client, connection); - let access = ReadAccessParams { - all_teams: false, + let access = QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a".into()], }; let page = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 500) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 500, + }, + ) .await?; assert_eq!(page.data.len(), runs); assert!( @@ -238,17 +260,20 @@ async fn large_runs_remain_complete_under_default_reader_limits( .send() .await? .error_for_status()?; - for table in ["otel_traces AS o", "spend_logs FINAL"] - .into_iter() - .take(if costed { 2 } else { 1 }) + for read in [ + "LIMIT 1 BY TeamId, ApiKeyHash, TraceId, SpanId", + "FROM owned_calls", + ] + .into_iter() + .take(if costed { 2 } else { 1 }) { let read_queries = client.post(writer.url().clone()).body(format!( - "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND query LIKE '%FROM {table}%' AND query NOT LIKE '%system.query_log%'" + "SELECT count() FROM system.query_log WHERE type = 'QueryFinish' AND current_database = '{DATABASE}' AND position(query, '{read}') > 0 AND query NOT LIKE '%system.query_log%'" )).send().await?.error_for_status()?.text().await?; let read_queries = read_queries.trim().parse::()?; assert!( read_queries > 0 && read_queries < runs, - "{read_queries} {table} queries for {runs} runs" + "{read_queries} `{read}` queries for {runs} runs" ); } } @@ -290,9 +315,9 @@ async fn large_runs_remain_complete_under_default_reader_limits( }, (steps - 1) as u64 ); - let denied = ReadAccessParams { + let denied = QueryScope::Owned { + user_id: String::new(), team_ids: vec!["other-team".into()], - ..access.clone() }; assert!( reader @@ -370,13 +395,17 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( .connection(client, &QueryScope::All, "fixture-secret") .await?; let (reader, store) = make_reader(client, connection.clone()); - let access = ReadAccessParams { - all_teams: true, - user_id: String::new(), - team_ids: Vec::new(), - }; + let access = QueryScope::All; let listed = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 10) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 10, + }, + ) .await?; let summary = listed .data @@ -419,8 +448,7 @@ async fn cursor_pages_keep_a_tenant_scoped_snapshot_when_more_spans_arrive( ])], ) .await?; - let denied = ReadAccessParams { - all_teams: false, + let denied = QueryScope::Owned { user_id: String::new(), team_ids: vec!["not-this-team".into()], }; @@ -533,13 +561,17 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( .connection(client, &QueryScope::All, "fixture-secret") .await?; let (reader, store) = make_reader(client, connection.clone()); - let access = ReadAccessParams { - all_teams: true, - user_id: String::new(), - team_ids: Vec::new(), - }; + let access = QueryScope::All; let before = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 50, + }, + ) .await?; let run = before .data @@ -568,12 +600,28 @@ async fn an_oversized_span_keeps_the_run_list_available_with_partial_totals( ) .await?; let cached = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 50, + }, + ) .await?; assert_eq!(cached.data, before.data); let (reader, store) = make_reader(client, connection); let after = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 50, + }, + ) .await?; assert_eq!(after.data.len(), before.data.len()); let limited = after @@ -686,13 +734,20 @@ async fn gateway_ids_resolve_through_detail_and_batch_reads_with_legacy_fallback .connection(client, &QueryScope::All, "fixture-secret") .await?; let (reader, store) = make_reader(client, connection); - let access = ReadAccessParams { - all_teams: false, + let access = QueryScope::Owned { user_id: String::new(), team_ids: vec!["team-a".into()], }; let page = reader - .list_traces(&store, &access, 0, 2_000_000_000_000, None, 50) + .list_traces( + &store, + &access, + &all_runs(), + &PageRequest { + cursor: None, + limit: 50, + }, + ) .await?; assert_eq!(page.data.len(), cases.len()); for (id, _, _, _, _, expected) in cases { diff --git a/litellm-rust/crates/traces-clickhouse/tests/search.rs b/litellm-rust/crates/traces-clickhouse/tests/search.rs new file mode 100644 index 00000000000..defbb87f5b1 --- /dev/null +++ b/litellm-rust/crates/traces-clickhouse/tests/search.rs @@ -0,0 +1,421 @@ +use std::collections::{BTreeMap, BTreeSet}; + +use litellm_traces::search::{AgentRuns, HistogramBucket, RunField, RunFilter, RunSearch}; +use litellm_traces_cache::{PageRequest, TraceReader}; +use litellm_traces_clickhouse::{ + ClickHouseTraces, Connection, InsertTable, QueryScope, insert_rows, +}; +use rstest::rstest; +use serde_json::{Value, json}; +use strum::IntoEnumIterator; + +#[path = "queries/support.rs"] +mod fixtures; +mod support; + +use fixtures::{DATABASE, SeededDatabase, migrated_database}; +use support::TestResult; + +const HOUR_MS: i64 = 3_600_000; +const T0_MS: i64 = 1_790_000_000_000; +const WINDOW_END_MS: i64 = T0_MS + 3 * HOUR_MS; + +type ExpectedBucket = (u64, u64, Vec<(&'static str, u64)>); + +fn filter(start_ms: i64, q: &str) -> RunFilter { + RunFilter { + start_ms, + end_ms: WINDOW_END_MS, + search: RunSearch::parse(q), + } +} + +fn page(cursor: Option, limit: u32) -> PageRequest { + PageRequest { cursor, limit } +} + +fn reader() -> TraceReader { + TraceReader::new(litellm_storage_clickhouse::READ_LIMITS.response_bytes) +} + +struct Step { + span_id: &'static str, + name: &'static str, + kind: &'static str, + agent: &'static str, + model: &'static str, + error: bool, +} + +const fn step(span_id: &'static str, name: &'static str, kind: &'static str) -> Step { + Step { + span_id, + name, + kind, + agent: "", + model: "", + error: false, + } +} + +struct Run { + trace_id: &'static str, + team: &'static str, + start_ms: i64, + root: &'static str, + input: &'static str, + steps: Vec, +} + +fn runs() -> Vec { + vec![ + Run { + trace_id: "alpha", + team: "team-a", + start_ms: T0_MS, + root: "plan trip", + input: "book a flight to Paris", + steps: vec![ + Step { + agent: "researcher", + ..step("alpha-agent", "invoke_agent", "agent") + }, + Step { + model: "gpt-x", + ..step("alpha-llm", "chat", "llm") + }, + ], + }, + Run { + trace_id: "beta", + team: "team-a", + start_ms: T0_MS + HOUR_MS, + root: "write report", + input: "summarize 100% of Q3_sales", + steps: vec![ + step("beta-agent", "writer", "agent"), + Step { + model: "claude-y", + error: true, + ..step("beta-llm", "chat", "llm") + }, + ], + }, + Run { + trace_id: "gamma", + team: "team-a", + start_ms: T0_MS + 2 * HOUR_MS, + root: "plan trip", + input: "hello", + steps: vec![Step { + model: "gpt-x", + ..step("gamma-llm", "chat", "llm") + }], + }, + Run { + trace_id: "foreign", + team: "team-b", + start_ms: T0_MS + 30 * 60_000, + root: "plan trip", + input: "book a flight to Paris", + steps: vec![Step { + agent: "spy", + model: "gpt-x", + ..step("foreign-agent", "invoke_agent", "agent") + }], + }, + ] +} + +fn span_row(run: &Run, offset_ms: i64, span: &Step, parent: &str) -> BTreeMap { + BTreeMap::from([ + ( + "Timestamp".into(), + json!((run.start_ms + offset_ms) * 1_000_000), + ), + ("TraceId".into(), json!(run.trace_id)), + ("SpanId".into(), json!(span.span_id)), + ("ParentSpanId".into(), json!(parent)), + ("SpanName".into(), json!(span.name)), + ("ServiceName".into(), json!("svc")), + ("ObservationType".into(), json!(span.kind)), + ("AgentName".into(), json!(span.agent)), + ("Model".into(), json!(span.model)), + ( + "StatusCode".into(), + json!(if span.error { + "STATUS_CODE_ERROR" + } else { + "STATUS_CODE_OK" + }), + ), + ("TeamId".into(), json!(run.team)), + ("ApiKeyHash".into(), json!("key-a")), + ("Duration".into(), json!(1_000_000)), + ]) +} + +fn rows(run: &Run) -> Vec> { + let root = step("root", run.root, "chain"); + let root_id = format!("{}-root", run.trace_id); + let root_row = BTreeMap::from_iter(span_row(run, 0, &root, "").into_iter().chain([ + ("SpanId".into(), json!(root_id)), + ("Input".into(), json!(run.input)), + ])); + std::iter::once(root_row) + .chain( + run.steps + .iter() + .zip(1..) + .map(|(span, offset)| span_row(run, offset, span, &root_id)), + ) + .collect() +} + +async fn seed(fixture: &SeededDatabase) -> TestResult { + let writer = Connection::writer(&fixture.database.url)?; + insert_rows( + &fixture.database.client, + &writer, + DATABASE, + InsertTable::OtelTraces, + runs().iter().flat_map(rows).collect(), + ) + .await?; + let connection = fixture + .readers + .connection(&fixture.database.client, &QueryScope::All, "fixture-secret") + .await?; + Ok(ClickHouseTraces::new( + fixture.database.client.clone(), + connection, + )) +} + +fn team_a() -> QueryScope { + QueryScope::Owned { + user_id: String::new(), + team_ids: vec!["team-a".into()], + } +} + +#[rstest] +#[tokio::test] +async fn list_q_selects_matching_runs_before_paging( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let reader = reader(); + let cases: &[(&str, &[&str])] = &[ + ("", &["gamma", "beta", "alpha"]), + ("status:", &["gamma", "beta", "alpha"]), + (r#"name:"plan trip""#, &["gamma", "alpha"]), + (r#"-name:"plan trip""#, &["beta"]), + ("NAME:PLAN*", &["gamma", "alpha"]), + ("agent:researcher", &["alpha"]), + ("agent:writer", &["beta"]), + ("agent:RESEARCH*", &["alpha"]), + ("agent:research", &[]), + ("-agent:*", &["gamma"]), + ("status:error", &["beta"]), + ("status:ok", &["gamma", "alpha"]), + ("status:*r*", &["beta"]), + ("model:gpt-x", &["gamma", "alpha"]), + ("model:gpt", &[]), + ("-model:gpt-x", &["beta"]), + ("input:*paris", &["alpha"]), + ("trace_id:gam*", &["gamma"]), + ("paris", &["alpha"]), + ("PLAN", &["gamma", "alpha"]), + ("100%", &["beta"]), + ("summarize%sales", &[]), + (r#""100% of""#, &["beta"]), + (r#""100% of""#, &[]), + ("q3_", &["beta"]), + ("flight_to", &[]), + ("plan hello", &["gamma"]), + ("plan -trace_id:gamma", &["alpha"]), + ("unknown:x", &[]), + ]; + for (q, expected) in cases { + let page = reader + .list_traces(&store, &team_a(), &filter(0, q), &page(None, 50)) + .await?; + let listed: Vec<&str> = page.data.iter().map(|run| run.trace_id.as_str()).collect(); + assert_eq!(&listed, expected, "q = {q:?}"); + } + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn list_q_pages_through_matches_only( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let reader = reader(); + let filter = filter(0, "model:gpt-x"); + let first = reader + .list_traces(&store, &team_a(), &filter, &page(None, 1)) + .await?; + let second = reader + .list_traces( + &store, + &team_a(), + &filter, + &page(first.next_cursor.clone(), 1), + ) + .await?; + let ids = |page: &litellm_traces::TracePage| -> Vec { + page.data.iter().map(|run| run.trace_id.clone()).collect() + }; + assert_eq!(ids(&first), ["gamma"]); + assert_eq!(ids(&second), ["alpha"]); + Ok(()) +} + +#[rstest] +#[case::all("", [ + (1, 0, vec![("researcher", 1)]), + (1, 1, vec![]), + (1, 0, vec![("svc", 1)]), +])] +#[case::filtered("model:gpt-x", [ + (1, 0, vec![("researcher", 1)]), + (0, 0, vec![]), + (1, 0, vec![("svc", 1)]), +])] +#[tokio::test] +async fn histogram_counts_matching_runs_per_bucket( + #[future(awt)] migrated_database: TestResult, + #[case] q: &str, + #[case] expected: [ExpectedBucket; 3], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let histogram = reader() + .histogram(&store, &team_a(), &filter(T0_MS, q), 3) + .await?; + let expected: Vec = expected + .into_iter() + .zip(0..) + .map(|((total, failed, agents), index)| HistogramBucket { + start_ms: T0_MS + index * HOUR_MS, + end_ms: T0_MS + (index + 1) * HOUR_MS, + total, + failed, + agents: agents + .into_iter() + .map(|(agent, runs)| AgentRuns { + agent: agent.into(), + runs, + }) + .collect(), + }) + .collect(); + assert_eq!(histogram.buckets, expected); + Ok(()) +} + +#[rstest] +#[case::agents_in_scope(RunField::Agent, "", &["researcher", "writer"])] +#[case::names_by_frequency(RunField::Name, "", &["plan trip", "write report"])] +#[case::statuses_by_frequency(RunField::Status, "", &["ok", "error"])] +#[case::models_by_frequency(RunField::Model, "", &["gpt-x", "claude-y"])] +#[case::needle_ignores_case(RunField::Name, "WR", &["write report"])] +#[case::needle_matches_inside(RunField::Name, "trip", &["plan trip"])] +#[case::needle_underscore_is_literal(RunField::Name, "n_t", &[])] +#[case::needle_percent_is_literal(RunField::Name, "n%t", &[])] +#[tokio::test] +async fn values_list_distinct_field_values_in_scope( + #[future(awt)] migrated_database: TestResult, + #[case] field: RunField, + #[case] needle: &str, + #[case] expected: &[&str], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let values = reader() + .values(&store, &team_a(), &filter(0, ""), field, needle, 10) + .await?; + assert_eq!(values.values, expected); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn admin_scope_sees_every_team( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let admin = QueryScope::All; + let values = reader() + .values(&store, &admin, &filter(0, ""), RunField::Agent, "", 10) + .await?; + assert_eq!( + values.values.into_iter().collect::>(), + BTreeSet::from(["researcher".into(), "spy".into(), "writer".into()]) + ); + Ok(()) +} + +#[rstest] +#[case::by_model(RunField::Agent, "model:claude-y", &["writer"])] +#[case::by_status(RunField::Name, "status:error", &["write report"])] +#[case::excluding(RunField::Model, "-trace_id:alpha", &["claude-y", "gpt-x"])] +#[tokio::test] +async fn values_narrow_to_runs_matching_the_search( + #[future(awt)] migrated_database: TestResult, + #[case] field: RunField, + #[case] q: &str, + #[case] expected: &[&str], +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let values = reader() + .values(&store, &team_a(), &filter(0, q), field, "", 10) + .await?; + assert_eq!(values.values, expected); + Ok(()) +} + +#[rstest] +#[tokio::test] +async fn every_field_suggests_values_that_filter_back_to_their_runs( + #[future(awt)] migrated_database: TestResult, +) -> TestResult { + let fixture = migrated_database?; + let store = seed(&fixture).await?; + let reader = reader(); + for field in RunField::iter() { + let key: &str = field.into(); + let values = reader + .values(&store, &team_a(), &filter(0, ""), field, "", 10) + .await?; + assert!(!values.values.is_empty(), "{key} suggests nothing"); + for value in &values.values { + let q = format!(r#"{key}:"{value}""#); + let included = reader + .list_traces(&store, &team_a(), &filter(0, &q), &page(None, 50)) + .await?; + let excluded = reader + .list_traces( + &store, + &team_a(), + &filter(0, &format!("-{q}")), + &page(None, 50), + ) + .await?; + assert!(!included.data.is_empty(), "{q} lists nothing"); + assert_eq!( + included.data.len() + excluded.data.len(), + 3, + "{q} does not partition the runs" + ); + } + } + Ok(()) +} diff --git a/litellm-rust/crates/traces/src/lib.rs b/litellm-rust/crates/traces/src/lib.rs index a44679a59f3..72c4baa2e67 100644 --- a/litellm-rust/crates/traces/src/lib.rs +++ b/litellm-rust/crates/traces/src/lib.rs @@ -18,7 +18,9 @@ mod query_access; mod resolve; #[cfg(feature = "schema")] pub mod schema; +pub mod search; mod shared; +pub mod store; mod tenant; mod truncate; mod ui; diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/claude_code.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/claude_code.rs index 27c722dac78..9f469beafd2 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/claude_code.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/claude_code.rs @@ -1,8 +1,9 @@ +use std::collections::BTreeMap; + use super::{ Integration, ObservationType, RoleEvidence, Rule, SpanContext, SpanFacts, attr, present, }; use crate::normalize::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE}; -use std::collections::BTreeMap; pub(super) const SCOPE: &str = CLAUDE_CODE_SCOPE; diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs index 1ee6c24b9d4..38a16ee76de 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/http_client.rs @@ -1,5 +1,6 @@ -use super::{CallEvidence, CallKey, ObservationType, RoleEvidence, SpanContext, SpanFacts}; -use super::{Integration, Rule}; +use super::{ + CallEvidence, CallKey, Integration, ObservationType, RoleEvidence, Rule, SpanContext, SpanFacts, +}; const SCOPES: [&str; 7] = [ "opentelemetry.instrumentation.httpx", diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/langchain.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/langchain.rs index 0e311169ad3..26d60a2d5aa 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/langchain.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/langchain.rs @@ -1,9 +1,10 @@ +use std::collections::BTreeMap; + use super::{ AgentMetadata, Integration, ObservationType, RoleEvidence, SpanContext, SpanFacts, attr, messages, }; use crate::normalize::present; -use std::collections::BTreeMap; pub(super) fn adjust(context: &SpanContext<'_>, facts: SpanFacts) -> SpanFacts { let middleware = !context.parent_span_id.is_empty() && is_langchain_middleware(context.name); diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/llama_index.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/llama_index.rs index 093d27955ac..745730cc575 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/llama_index.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/llama_index.rs @@ -1,6 +1,7 @@ -use super::{ObservationType, RoleEvidence, SpanContext, SpanFacts, attr}; use serde_json::Value; +use super::{ObservationType, RoleEvidence, SpanContext, SpanFacts, attr}; + pub(super) fn adjust(context: &SpanContext<'_>, facts: SpanFacts) -> SpanFacts { let agent = context .name diff --git a/litellm-rust/crates/traces/src/normalize/instrumentation/pydantic_ai.rs b/litellm-rust/crates/traces/src/normalize/instrumentation/pydantic_ai.rs index 4c6c81052f6..60d41bbcb1c 100644 --- a/litellm-rust/crates/traces/src/normalize/instrumentation/pydantic_ai.rs +++ b/litellm-rust/crates/traces/src/normalize/instrumentation/pydantic_ai.rs @@ -1,5 +1,4 @@ -use super::{Extraction, SpanContext, SpanFacts, messages, select_attribute}; -use super::{Integration, Rule}; +use super::{Extraction, Integration, Rule, SpanContext, SpanFacts, messages, select_attribute}; use crate::normalize::format::genai::Operation; pub(super) const SCOPE: &str = "pydantic-ai"; diff --git a/litellm-rust/crates/traces/src/normalize/messages.rs b/litellm-rust/crates/traces/src/normalize/messages.rs index 0aa5fde2754..d18cbdb946f 100644 --- a/litellm-rust/crates/traces/src/normalize/messages.rs +++ b/litellm-rust/crates/traces/src/normalize/messages.rs @@ -1,15 +1,15 @@ //! The common message format normalizers emit for span input and output: a JSON array of //! `{role, content, tool_calls?, name?}` that the UI renders as a conversation. -use indexmap::IndexMap; -use serde::{Deserialize, Deserializer, Serialize}; -use serde_json::{Value, ser::Formatter}; use std::{ collections::{BTreeMap, BTreeSet}, io, }; +use indexmap::IndexMap; use litellm_llms_types::{formats::chat_completions::ChatMessageContent, recognized::Recognized}; +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Value, ser::Formatter}; use super::{CallEvidence, CallKey, attr}; @@ -592,9 +592,9 @@ pub(super) fn state_conversation(input: &str) -> Option> { #[cfg(test)] mod tests { use rstest::rstest; + use serde_json::Value; use super::{state_conversation, state_preview}; - use serde_json::Value; #[rstest] #[case::latest_user(r#"{"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"reply"},{"role":"user","content":"last"}]}"#, Some("last"))] diff --git a/litellm-rust/crates/traces/src/normalize/mod.rs b/litellm-rust/crates/traces/src/normalize/mod.rs index c3ad76fde02..1cde89341ad 100644 --- a/litellm-rust/crates/traces/src/normalize/mod.rs +++ b/litellm-rust/crates/traces/src/normalize/mod.rs @@ -9,9 +9,10 @@ use std::{ str::FromStr, }; -use crate::{Error, otlp::DecodedEvent}; use serde::{Deserialize, Serialize, Serializer}; +use crate::{Error, otlp::DecodedEvent}; + mod format; mod instrumentation; mod messages; @@ -118,7 +119,7 @@ pub enum CallEvidence { } impl CallEvidence { - pub(crate) fn row_keys(row: &crate::query::named::TraceSpansRow) -> BTreeSet { + pub(crate) fn row_keys(row: &crate::store::SpanRow) -> BTreeSet { if row.call_keys.is_empty() && !row.litellm_request_id.is_empty() { BTreeSet::from([CallKey::ProviderResponse(row.litellm_request_id.clone())]) } else { @@ -126,7 +127,7 @@ impl CallEvidence { } } - pub(crate) fn from_row(row: &crate::query::named::TraceSpansRow) -> Self { + pub(crate) fn from_row(row: &crate::store::SpanRow) -> Self { let kind = row .call_evidence .unwrap_or(if Self::row_keys(row).is_empty() { diff --git a/litellm-rust/crates/traces/src/otlp/mod.rs b/litellm-rust/crates/traces/src/otlp/mod.rs index 2e17d868090..929163e0291 100644 --- a/litellm-rust/crates/traces/src/otlp/mod.rs +++ b/litellm-rust/crates/traces/src/otlp/mod.rs @@ -3,11 +3,11 @@ mod limits; mod span; mod wire; -pub use limits::DecodeLimits; - -use serde::Serialize; use std::collections::BTreeMap; +pub use limits::DecodeLimits; +use serde::Serialize; + use crate::{Error, NormalizedSpan, Shared}; #[derive(Serialize)] diff --git a/litellm-rust/crates/traces/src/query.rs b/litellm-rust/crates/traces/src/query.rs index c39b1f26a52..a0524dbb22b 100644 --- a/litellm-rust/crates/traces/src/query.rs +++ b/litellm-rust/crates/traces/src/query.rs @@ -1,16 +1,8 @@ pub mod guide; -pub mod named; #[derive(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)] #[strum(serialize_all = "snake_case")] pub enum ReadQuery { - ListTraces, - TraceSpans, - TracePageSpans, - TraceIdentity, - SpanDetail, - SpanError, - SpendByResponseIds, Availability, Agents, Sample, diff --git a/litellm-rust/crates/traces/src/query/named.rs b/litellm-rust/crates/traces/src/query/named.rs deleted file mode 100644 index a981bc46471..00000000000 --- a/litellm-rust/crates/traces/src/query/named.rs +++ /dev/null @@ -1,203 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; - -#[macro_rules_attribute::apply(wire_type)] -#[derive(Clone, Debug)] -#[cfg_attr(feature = "schema", schemars(rename = "TraceScope"))] -pub struct ReadAccessParams { - #[serde( - deserialize_with = "crate::wire::flag", - serialize_with = "crate::wire::serialize_flag" - )] - #[cfg_attr(feature = "schema", schemars(schema_with = "crate::schema::flag"))] - pub all_teams: bool, - pub user_id: String, - pub team_ids: Vec, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct ListTracesParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub start_ms: i64, - pub end_ms: i64, - pub cursor_ms: i64, - pub cursor_trace_id: String, - pub limit: u32, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct ListTracesRow { - pub trace_id: String, - pub trace_ref: String, - pub team_id: String, - pub api_key_hash: String, - pub user_id: String, - pub name: String, - pub service: String, - pub input_preview: String, - #[serde(serialize_with = "crate::wire::serialize_status")] - pub status: crate::SpanStatus, - pub start_ms: i64, - pub duration_ms: i64, - pub span_count: u64, - pub agent_count: u64, - pub agent_invocations: u64, - #[serde(default)] - pub agent_names: Vec, - #[serde(default)] - pub frameworks: Vec, - pub llm_calls: u64, - pub tool_calls: u64, - pub input_tokens: u64, - pub output_tokens: u64, - pub models: Vec, - pub error_count: u64, - pub request_ids: Vec, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct TraceSpansParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub trace_id: String, - pub trace_ref: String, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct TraceSpansRow { - #[serde(default)] - pub trace_id: String, - pub span_id: String, - pub parent_span_id: String, - pub name: String, - #[serde(rename = "type")] - pub kind: crate::ObservationType, - #[serde( - default, - deserialize_with = "crate::wire::flag", - serialize_with = "crate::wire::serialize_flag" - )] - pub wrapper_candidate: bool, - pub agent: String, - #[serde(default)] - pub framework: String, - #[serde(serialize_with = "crate::wire::serialize_status")] - pub status: crate::SpanStatus, - pub status_message: String, - #[serde( - deserialize_with = "crate::wire::flag", - serialize_with = "crate::wire::serialize_flag" - )] - pub error_truncated: bool, - pub start_ns: i64, - pub duration_ns: u64, - pub service: String, - pub input_preview: String, - pub model: String, - pub input_tokens: u32, - pub output_tokens: u32, - pub litellm_request_id: String, - #[serde(default)] - pub call_keys: Vec, - #[serde( - default, - deserialize_with = "crate::wire::evidence", - serialize_with = "crate::wire::serialize_evidence" - )] - pub call_evidence: Option, - #[serde(default)] - pub tool_call_id: String, - pub team_id: String, - pub api_key_hash: String, - pub user_id: String, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct TracePageSpansParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub trace_refs: Vec, - pub start_ms: i64, - pub end_ms: i64, -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct SpanDetailParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub trace_id: String, - pub trace_ref: String, - pub span_id: String, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct SpanDetailRow { - pub span_id: String, - pub input: String, - pub output: String, - pub attributes: BTreeMap, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct SpanErrorParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub trace_id: String, - pub trace_ref: String, - pub span_id: String, - pub error_offset: u64, - pub error_version: String, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct SpanErrorRow { - pub span_id: String, - pub message: String, - pub total_chars: u64, - pub version: String, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct SpendByResponseIdsParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub response_ids: Vec, - pub request_ids: Vec, - pub trace_ids: Vec, - pub start_ms: i64, - pub end_ms: i64, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct SpendByResponseIdsRow { - pub request_id: String, - pub litellm_call_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, - pub spend: Option, - pub start_ms: i64, -} - -impl SpendByResponseIdsRow { - pub(crate) fn identity(&self) -> (&str, i64, &str) { - (&self.team_id, self.start_ms, &self.request_id) - } -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct TraceIdentityParams { - #[serde(flatten)] - pub access: ReadAccessParams, - pub trace_id: String, -} - -#[derive(Debug, Deserialize, Serialize)] -pub struct TraceIdentityRow { - pub trace_ref: String, -} diff --git a/litellm-rust/crates/traces/src/query_access.rs b/litellm-rust/crates/traces/src/query_access.rs index c543ddb5808..f2d6f338e5a 100644 --- a/litellm-rust/crates/traces/src/query_access.rs +++ b/litellm-rust/crates/traces/src/query_access.rs @@ -1,7 +1,7 @@ use crate::InvalidScope; #[macro_rules_attribute::apply(wire_type)] -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Eq, PartialEq)] #[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] pub enum QueryScope { #[cfg_attr(feature = "schema", schemars(title = "AllQueryScope"))] diff --git a/litellm-rust/crates/traces/src/resolve/graph.rs b/litellm-rust/crates/traces/src/resolve/graph.rs index 733f5e87213..b891d1ee1bb 100644 --- a/litellm-rust/crates/traces/src/resolve/graph.rs +++ b/litellm-rust/crates/traces/src/resolve/graph.rs @@ -1,15 +1,15 @@ use std::collections::{HashMap, HashSet}; -use crate::query::named::TraceSpansRow; +use crate::store::SpanRow; pub(super) struct Graph<'a> { - pub(super) rows: &'a [TraceSpansRow], + pub(super) rows: &'a [SpanRow], by_id: HashMap<&'a str, usize>, children: HashMap<&'a str, Vec>, } impl<'a> Graph<'a> { - pub(super) fn new(rows: &'a [TraceSpansRow]) -> Self { + pub(super) fn new(rows: &'a [SpanRow]) -> Self { let by_id: HashMap<&str, usize> = rows .iter() .enumerate() @@ -90,9 +90,9 @@ mod tests { use rstest::{fixture, rstest}; use super::Graph; - use crate::query::named::TraceSpansRow; + use crate::store::SpanRow; - fn row(id: &str, parent: &str) -> TraceSpansRow { + fn row(id: &str, parent: &str) -> SpanRow { serde_json::from_value(serde_json::json!({ "span_id": id, "parent_span_id": parent, @@ -118,7 +118,7 @@ mod tests { } #[fixture] - fn unordered_rows() -> Vec { + fn unordered_rows() -> Vec { vec![ row("leaf", "middle"), row("sibling", "root"), @@ -128,7 +128,7 @@ mod tests { } #[rstest] - fn traversal_follows_links_instead_of_export_order(unordered_rows: Vec) { + fn traversal_follows_links_instead_of_export_order(unordered_rows: Vec) { let graph = Graph::new(&unordered_rows); assert_eq!(graph.ancestors(0), [2, 3]); let descendants: std::collections::BTreeSet<&str> = graph diff --git a/litellm-rust/crates/traces/src/resolve/resolution.rs b/litellm-rust/crates/traces/src/resolve/resolution.rs index b469fb95019..6d961bf173b 100644 --- a/litellm-rust/crates/traces/src/resolve/resolution.rs +++ b/litellm-rust/crates/traces/src/resolve/resolution.rs @@ -2,17 +2,16 @@ use std::collections::HashMap; use indexmap::IndexMap; -use crate::{ - normalize::{CallKey, ObservationType}, - query::named::{SpendByResponseIdsRow as SpendRow, TraceSpansRow}, -}; - use super::{ graph::Graph, spend::{self, Ownership, Requests, SpendEvidence}, }; +use crate::{ + normalize::{CallKey, ObservationType}, + store::{CallRow as SpendRow, SpanRow}, +}; -pub(super) fn agent_label(row: &TraceSpansRow) -> &str { +pub(super) fn agent_label(row: &SpanRow) -> &str { if row.agent.is_empty() { &row.name } else { @@ -29,7 +28,7 @@ pub(super) struct Resolution<'a> { } impl<'a> Resolution<'a> { - pub(super) fn new(rows: &'a [TraceSpansRow], spend: &'a [SpendRow]) -> Self { + pub(super) fn new(rows: &'a [SpanRow], spend: &'a [SpendRow]) -> Self { let graph = Graph::new(rows); let named_agents = rows.iter().any(|row| !row.agent.is_empty()); let types: HashMap<&str, ObservationType> = (0..rows.len()) @@ -57,7 +56,7 @@ impl<'a> Resolution<'a> { } } - pub(super) fn row(&self, index: usize) -> &'a TraceSpansRow { + pub(super) fn row(&self, index: usize) -> &'a SpanRow { &self.graph.rows[index] } diff --git a/litellm-rust/crates/traces/src/resolve/spend.rs b/litellm-rust/crates/traces/src/resolve/spend.rs index ce23362b7f8..e74c5c79129 100644 --- a/litellm-rust/crates/traces/src/resolve/spend.rs +++ b/litellm-rust/crates/traces/src/resolve/spend.rs @@ -4,7 +4,7 @@ use indexmap::IndexMap; use crate::{ CallEvidence, CallEvidenceKind, CallKey, - query::named::{SpendByResponseIdsRow as SpendRow, TraceSpansRow}, + store::{CallRow as SpendRow, SpanRow}, }; /// The spend records to fetch for a set of spans. @@ -17,7 +17,7 @@ pub struct SpendLookup { } impl SpendLookup { - pub fn new(rows: &[TraceSpansRow]) -> Self { + pub fn new(rows: &[SpanRow]) -> Self { let evidence: Vec<_> = rows .iter() .map(|row| (row, CallEvidence::row_keys(row))) @@ -180,7 +180,7 @@ fn matches<'a>( ownership: &Ownership<'_>, spend_rows: &'a [SpendRow], key: &CallKey, - row: &TraceSpansRow, + row: &SpanRow, ) -> IndexMap<(&'a str, i64, &'a str), &'a SpendRow> { let matches = |spend: &SpendRow| match key { CallKey::ProviderResponse(id) => { @@ -206,7 +206,7 @@ fn matches<'a>( } pub(super) fn requests<'a>( - row: &TraceSpansRow, + row: &SpanRow, ownership: &Ownership<'_>, spend_rows: &'a [SpendRow], ) -> SpendEvidence<'a> { diff --git a/litellm-rust/crates/traces/src/resolve/view.rs b/litellm-rust/crates/traces/src/resolve/view.rs index 9c1c3756856..2b3a49d4e12 100644 --- a/litellm-rust/crates/traces/src/resolve/view.rs +++ b/litellm-rust/crates/traces/src/resolve/view.rs @@ -3,16 +3,15 @@ use std::collections::{BTreeSet, HashSet}; use indexmap::IndexMap; use time::OffsetDateTime; -use crate::{ - normalize::ObservationType, - query::named::{ListTracesRow, SpendByResponseIdsRow as SpendRow, TraceSpansRow}, - view::{AgentNode, Span, SpanStatus, Trace, TraceSummary}, -}; - use super::{ resolution::{Resolution, agent_label}, spend::{Requests, request_cost, total}, }; +use crate::{ + normalize::ObservationType, + store::{CallRow as SpendRow, RunRow, SpanRow}, + view::{AgentNode, Span, SpanStatus, Trace, TraceSummary}, +}; const NANOS_PER_MS: f64 = 1_000_000.0; @@ -136,7 +135,7 @@ fn sorted_unique<'a>(values: impl Iterator) -> Vec { pub fn resolve_trace( trace_id: &str, trace_ref: &str, - rows: &[TraceSpansRow], + rows: &[SpanRow], spend: &[SpendRow], ) -> Option { let first = rows.first()?; @@ -154,7 +153,7 @@ pub fn resolve_trace( .unwrap_or_default(); let agents = agents(&resolution); let calls = &resolution.model_calls; - let counted: Vec<&TraceSpansRow> = if calls.is_empty() { + let counted: Vec<&SpanRow> = if calls.is_empty() { rows.iter().collect() } else { calls.iter().map(|call| &rows[*call]).collect() @@ -214,7 +213,7 @@ pub fn resolve_trace( }) } -pub fn listed_summary(row: &ListTracesRow) -> TraceSummary { +pub fn listed_summary(row: &RunRow) -> TraceSummary { TraceSummary { resolution_limited: true, trace_id: row.trace_id.clone(), diff --git a/litellm-rust/crates/traces/src/schema.rs b/litellm-rust/crates/traces/src/schema.rs index cfa1d8e201e..1f9fa8f491c 100644 --- a/litellm-rust/crates/traces/src/schema.rs +++ b/litellm-rust/crates/traces/src/schema.rs @@ -46,15 +46,14 @@ fn emitted() -> Schema { pub fn schemas() -> BTreeMap<&'static str, Schema> { BTreeMap::from([ - ( - "TraceScope", - received::(), - ), ("QueryScope", received::()), ("Tenant", received::()), ("TracePage", emitted::()), ("Trace", emitted::()), ("SpanDetail", emitted::()), ("SpanErrorPage", emitted::()), + ("TraceHistogram", emitted::()), + ("RunValues", emitted::()), + ("RunField", received::()), ]) } diff --git a/litellm-rust/crates/traces/src/search.rs b/litellm-rust/crates/traces/src/search.rs new file mode 100644 index 00000000000..b2b5ceef126 --- /dev/null +++ b/litellm-rust/crates/traces/src/search.rs @@ -0,0 +1,217 @@ +use serde::{Deserialize, Serialize}; + +use crate::store::RunCount; + +#[derive( + Clone, + Copy, + Debug, + Eq, + PartialEq, + Deserialize, + Serialize, + strum::EnumIter, + strum::EnumString, + strum::IntoStaticStr, +)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case", ascii_case_insensitive)] +pub enum RunField { + Name, + Agent, + Status, + Model, + Input, + TraceId, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct RunFilter { + pub start_ms: i64, + pub end_ms: i64, + pub search: RunSearch, +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct FieldFilter { + pub field: RunField, + /// Matched against the whole value, ignoring case; `*` matches any run of characters. + pub pattern: String, + pub exclude: bool, +} + +/// The parsed `q` of the runs list. Every text term and every filter must hold. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct RunSearch { + /// Each must appear in the trace id, input or name, ignoring case. + pub text: Vec, + pub filters: Vec, +} + +impl RunSearch { + /// Mirrors the dashboard's search box: `key:value` filters on known keys, `-key:value` negates, + /// `*` globs, double quotes keep spaces, and anything else is a free-text term. + /// A key typed without a value yet narrows nothing. + pub fn parse(q: &str) -> Self { + let (text, filters): (Vec<_>, Vec<_>) = tokens(q) + .map(clause) + .filter(|clause| !clause.value().is_empty()) + .partition(|clause| matches!(clause, Clause::Text(_))); + Self { + text: text + .into_iter() + .map(|clause| clause.value().to_owned()) + .collect(), + filters: filters + .into_iter() + .filter_map(|clause| match clause { + Clause::Field { + field, + exclude, + value, + } => Some(FieldFilter { + field, + pattern: value, + exclude, + }), + Clause::Text(_) => None, + }) + .collect(), + } + } +} + +enum Clause { + Text(String), + Field { + field: RunField, + exclude: bool, + value: String, + }, +} + +impl Clause { + fn value(&self) -> &str { + match self { + Self::Text(value) | Self::Field { value, .. } => value, + } + } +} + +/// Whitespace-separated tokens; a double-quoted stretch keeps its spaces, and an unclosed quote runs to the end. +fn tokens(q: &str) -> impl Iterator { + let mut rest = q; + std::iter::from_fn(move || { + rest = rest.trim_start(); + if rest.is_empty() { + return None; + } + let mut quoted = false; + let end = rest + .char_indices() + .find(|&(_, char)| { + if char == '"' { + quoted = !quoted; + } + !quoted && char.is_whitespace() + }) + .map_or(rest.len(), |(index, _)| index); + let (token, tail) = rest.split_at(end); + rest = tail; + Some(token) + }) +} + +fn unquote(raw: &str) -> String { + raw.strip_prefix('"') + .map(|inner| inner.strip_suffix('"').unwrap_or(inner)) + .filter(|inner| !inner.contains('"')) + .unwrap_or(raw) + .to_owned() +} + +fn clause(raw: &str) -> Clause { + let (exclude, body) = raw + .strip_prefix('-') + .map_or((false, raw), |body| (true, body)); + let field = body.split_once(':').and_then(|(key, value)| { + let named = !key.is_empty() && key.chars().all(|c| c.is_ascii_alphabetic() || c == '_'); + named + .then(|| key.parse::().ok()) + .flatten() + .map(|field| (field, value)) + }); + match field { + Some((field, value)) => Clause::Field { + field, + exclude, + value: unquote(value), + }, + None => Clause::Text(unquote(raw)), + } +} + +pub const MAX_HISTOGRAM_BUCKETS: u32 = 240; +pub const MAX_RUN_VALUES: u32 = 100; + +/// Matching runs per equal-width slice of the window. +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, PartialEq)] +pub struct TraceHistogram { + pub buckets: Vec, +} + +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, PartialEq)] +pub struct HistogramBucket { + pub start_ms: i64, + pub end_ms: i64, + pub total: u64, + pub failed: u64, + /// Runs that did not fail, by their alphabetically first agent label, or service when unlabelled. + pub agents: Vec, +} + +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, PartialEq)] +pub struct AgentRuns { + pub agent: String, + pub runs: u64, +} + +/// Distinct values of one run field, most common first. +#[macro_rules_attribute::apply(response_type)] +#[derive(Clone, Debug, PartialEq)] +pub struct RunValues { + pub values: Vec, +} + +/// Bucket `i` covers `[start + span * i / buckets, start + span * (i + 1) / buckets)`. +pub fn histogram(rows: &[RunCount], start_ms: i64, end_ms: i64, buckets: u32) -> TraceHistogram { + let span = i128::from(end_ms - start_ms); + let edge = |index: u32| start_ms + (span * i128::from(index) / i128::from(buckets)) as i64; + TraceHistogram { + buckets: (0..buckets) + .map(|index| { + let hits = rows.iter().filter(|row| row.bucket == index); + let mut agents: Vec = hits + .clone() + .filter(|row| !row.failed) + .map(|row| AgentRuns { + agent: row.value.clone(), + runs: row.runs, + }) + .collect(); + agents.sort_by(|left, right| left.agent.cmp(&right.agent)); + HistogramBucket { + start_ms: edge(index), + end_ms: edge(index + 1), + total: hits.clone().map(|row| row.runs).sum(), + failed: hits.filter(|row| row.failed).map(|row| row.runs).sum(), + agents, + } + }) + .collect(), + } +} diff --git a/litellm-rust/crates/traces/src/store.rs b/litellm-rust/crates/traces/src/store.rs new file mode 100644 index 00000000000..f9a67ab8af6 --- /dev/null +++ b/litellm-rust/crates/traces/src/store.rs @@ -0,0 +1,282 @@ +//! What trace storage must answer, independent of the engine behind it. + +use std::ops::Range; + +use serde::{Deserialize, Serialize}; + +use crate::search::{RunField, RunFilter}; + +#[derive(Clone, Debug, PartialEq)] +pub enum RunSelection { + Matching(RunFilter), + /// Every run with this trace id, whenever it happened. + TraceId(String), +} + +#[derive(Clone, Debug, Default, Eq, PartialEq, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct RunCursor { + pub start_ms: i64, + pub trace_ref: String, +} + +/// Runs newest first, by `(start_ms, trace_ref)` descending. +#[derive(Clone, Debug, PartialEq)] +pub struct RunQuery { + pub selection: RunSelection, + pub after: Option, + pub limit: u32, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct RunRow { + pub trace_id: String, + pub trace_ref: String, + pub team_id: String, + pub api_key_hash: String, + pub user_id: String, + pub name: String, + pub service: String, + pub input_preview: String, + #[serde(serialize_with = "crate::wire::serialize_status")] + pub status: crate::SpanStatus, + pub start_ms: i64, + pub duration_ms: i64, + pub span_count: u64, + pub agent_count: u64, + pub agent_invocations: u64, + #[serde(default)] + pub agent_names: Vec, + #[serde(default)] + pub frameworks: Vec, + pub llm_calls: u64, + pub tool_calls: u64, + pub input_tokens: u64, + pub output_tokens: u64, + pub models: Vec, + pub error_count: u64, +} + +impl RunRow { + pub fn cursor(&self) -> RunCursor { + RunCursor { + start_ms: self.start_ms, + trace_ref: self.trace_ref.clone(), + } + } +} + +/// What a run is counted under. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum CountValue { + Field(RunField), + /// The run's alphabetically first agent label, or its service when it has none. + PrimaryAgent, +} + +/// Each dimension left unset collapses to one group: bucket 0, not failed, or an empty value. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct CountBy { + /// Equal-width slices of the filter window; run `i` lands in + /// `(start_ms - window.start) * buckets / window.len()`. + pub buckets: Option, + pub failed: bool, + pub value: Option, +} + +/// Matching runs per group, most runs first. A run with several values for a field, such as +/// several models, counts once under each; empty values are not counted. +#[derive(Clone, Debug, PartialEq)] +pub struct RunCountQuery { + pub filter: RunFilter, + pub by: CountBy, + /// Keeps values containing this text, ignoring case. + pub contains: String, + /// `None` returns every group. + pub limit: Option, +} + +#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)] +pub struct RunCount { + pub bucket: u32, + #[serde( + deserialize_with = "crate::wire::flag", + serialize_with = "crate::wire::serialize_flag" + )] + pub failed: bool, + pub value: String, + pub runs: u64, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum SpanSelection { + Trace { + trace_id: String, + trace_ref: String, + }, + /// Spans of several runs that started within `window`. + Runs { + trace_refs: Vec, + window: Range, + }, +} + +#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd)] +pub struct SpanCursor { + pub team_id: String, + pub api_key_hash: String, + pub trace_id: String, + pub span_id: String, +} + +/// One row per span, the earliest received copy when a span was exported twice, ordered by +/// [`SpanCursor`] ascending. Spans received after `as_of_ms` are left out. +#[derive(Clone, Debug, PartialEq)] +pub struct SpanQuery { + pub selection: SpanSelection, + pub as_of_ms: u64, + pub after: Option, + pub limit: u32, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct SpanRow { + #[serde(default)] + pub trace_id: String, + pub span_id: String, + pub parent_span_id: String, + pub name: String, + #[serde(rename = "type")] + pub kind: crate::ObservationType, + #[serde( + default, + deserialize_with = "crate::wire::flag", + serialize_with = "crate::wire::serialize_flag" + )] + pub wrapper_candidate: bool, + pub agent: String, + #[serde(default)] + pub framework: String, + #[serde(serialize_with = "crate::wire::serialize_status")] + pub status: crate::SpanStatus, + /// The first 128 characters; [`SpanPart::Error`] reads the rest. + pub status_message: String, + #[serde( + deserialize_with = "crate::wire::flag", + serialize_with = "crate::wire::serialize_flag" + )] + pub error_truncated: bool, + pub start_ns: i64, + pub duration_ns: u64, + pub service: String, + pub input_preview: String, + pub model: String, + pub input_tokens: u32, + pub output_tokens: u32, + pub litellm_request_id: String, + #[serde(default)] + pub call_keys: Vec, + #[serde( + default, + deserialize_with = "crate::wire::evidence", + serialize_with = "crate::wire::serialize_evidence" + )] + pub call_evidence: Option, + #[serde(default)] + pub tool_call_id: String, + pub team_id: String, + pub api_key_hash: String, + pub user_id: String, +} + +impl SpanRow { + pub fn cursor(&self) -> SpanCursor { + SpanCursor { + team_id: self.team_id.clone(), + api_key_hash: self.api_key_hash.clone(), + trace_id: self.trace_id.clone(), + span_id: self.span_id.clone(), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, Deserialize, Serialize, strum::IntoStaticStr)] +#[serde(rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum SpanPart { + Input, + Output, + Error, + /// The span's attributes as one JSON object of strings. + Attributes, +} + +/// A character range of one part of one span, read from the copy [`SpanQuery`] would return. +#[derive(Clone, Debug, PartialEq)] +pub struct SpanTextQuery { + pub trace_id: String, + pub trace_ref: String, + pub span_id: String, + pub part: SpanPart, + pub offset: u64, + /// `None` reads to the end. + pub max_chars: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct SpanText { + pub text: String, + pub total_chars: u64, + /// Uppercase hex SHA-256 of the whole part, so a reader can tell when it changed. + pub version: String, +} + +/// Gateway calls that can be priced against spans: those whose response id, call id or trace id +/// is listed, oldest first by `(team_id, start_ms, request_id)`. +#[derive(Clone, Debug, PartialEq)] +pub struct CallQuery { + pub window: Range, + /// Also matches the upstream id a managed `resp_` id wraps. + pub response_ids: Vec, + /// Matches the gateway call id, or the request id when a call has none. + pub request_ids: Vec, + pub trace_ids: Vec, + pub after: Option, + pub limit: u32, +} + +#[derive(Clone, Debug, Default, Eq, Ord, PartialEq, PartialOrd)] +pub struct CallCursor { + pub team_id: String, + pub start_ms: i64, + pub request_id: String, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct CallRow { + pub request_id: String, + pub litellm_call_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, + pub spend: Option, + pub start_ms: i64, +} + +impl CallRow { + pub fn cursor(&self) -> CallCursor { + CallCursor { + team_id: self.team_id.clone(), + start_ms: self.start_ms, + request_id: self.request_id.clone(), + } + } + + pub(crate) fn identity(&self) -> (&str, i64, &str) { + (&self.team_id, self.start_ms, &self.request_id) + } +} diff --git a/litellm-rust/crates/traces/tests/captures.rs b/litellm-rust/crates/traces/tests/captures.rs index ca498fc914a..bddf111b02d 100644 --- a/litellm-rust/crates/traces/tests/captures.rs +++ b/litellm-rust/crates/traces/tests/captures.rs @@ -1,11 +1,13 @@ -use std::collections::{BTreeMap, BTreeSet}; -use std::path::{Path, PathBuf}; +use std::{ + collections::{BTreeMap, BTreeSet}, + path::{Path, PathBuf}, +}; use base64::{Engine as _, engine::general_purpose::STANDARD}; use litellm_traces::{ CallEvidence, CallEvidenceKind, CallKey, DecodedSpan, ObservationType, SpanStatus, decode_otlp, - query::named::{SpendByResponseIdsRow, TraceSpansRow}, resolve_trace, + store::{CallRow, SpanRow}, }; use rstest::rstest; use serde::Deserialize; @@ -99,7 +101,7 @@ fn upstream_response_id(response_id: &str) -> String { .unwrap_or_default() } -fn captured_spend_rows(spend_logs: &str) -> (FixtureCapture, Vec) { +fn captured_spend_rows(spend_logs: &str) -> (FixtureCapture, Vec) { let records: Vec = spend_logs .lines() .filter(|line| !line.trim().is_empty()) @@ -112,7 +114,7 @@ fn captured_spend_rows(spend_logs: &str) -> (FixtureCapture, Vec &str { .unwrap_or_default() } -fn trace_span(span: DecodedSpan) -> TraceSpansRow { +fn trace_span(span: DecodedSpan) -> SpanRow { let service = span .resource_attributes .get("service.name") @@ -175,7 +177,7 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow { CallKey::LiteLlmRequest(_) | CallKey::Transport | CallKey::GatewayAttempt => None, }) .unwrap_or_default(); - TraceSpansRow { + SpanRow { trace_id: span.trace_id, span_id: span.span_id, parent_span_id: span.parent_span_id, @@ -211,7 +213,7 @@ fn trace_span(span: DecodedSpan) -> TraceSpansRow { } } -fn trace_rows(otlp: &[u8]) -> Vec { +fn trace_rows(otlp: &[u8]) -> Vec { let mut seen = BTreeSet::new(); decode_otlp(otlp, Some("application/json")) .expect("valid OTLP fixture") @@ -220,14 +222,7 @@ fn trace_rows(otlp: &[u8]) -> Vec { .collect() } -fn fixture( - spend_log_path: &Path, -) -> ( - CaptureData, - FixtureCapture, - Vec, - Vec, -) { +fn fixture(spend_log_path: &Path) -> (CaptureData, FixtureCapture, Vec, Vec) { let spend_log_path = manifest_path(spend_log_path); let name = capture_name(&spend_log_path); let spend_log_contents = std::fs::read_to_string(&spend_log_path).unwrap_or_else(|error| { @@ -261,11 +256,11 @@ fn agent_spends(trace: &litellm_traces::Trace) -> BTreeMap> .collect() } -fn unrelated_transport(call: &TraceSpansRow) -> TraceSpansRow { +fn unrelated_transport(call: &SpanRow) -> SpanRow { let start_ns = i64::try_from(i128::from(call.start_ns) + i128::from(call.duration_ns) + 1_000_000) .expect("valid unrelated transport timestamp"); - TraceSpansRow { + SpanRow { trace_id: call.trace_id.clone(), span_id: format!("unrelated-transport-{}", call.span_id), parent_span_id: call.parent_span_id.clone(), diff --git a/litellm-rust/crates/traces/tests/otlp.rs b/litellm-rust/crates/traces/tests/otlp.rs index eb17bdc374b..c2d37ec3c75 100644 --- a/litellm-rust/crates/traces/tests/otlp.rs +++ b/litellm-rust/crates/traces/tests/otlp.rs @@ -1,5 +1,4 @@ -use litellm_traces::decode_otlp; -use litellm_traces::{AgentType, Integration, ObservationType, Shared}; +use litellm_traces::{AgentType, Integration, ObservationType, Shared, decode_otlp}; use opentelemetry_proto::tonic::trace::v1::Span; use rstest::rstest; diff --git a/litellm-rust/crates/traces/tests/query.rs b/litellm-rust/crates/traces/tests/query.rs index 79f9886dae6..f0ad89bc54e 100644 --- a/litellm-rust/crates/traces/tests/query.rs +++ b/litellm-rust/crates/traces/tests/query.rs @@ -2,12 +2,6 @@ use litellm_traces::{InvalidQuery, ReadQuery}; use rstest::rstest; #[rstest] -#[case::list_traces("list_traces", ReadQuery::ListTraces)] -#[case::trace_spans("trace_spans", ReadQuery::TraceSpans)] -#[case::span_detail("span_detail", ReadQuery::SpanDetail)] -#[case::span_error("span_error", ReadQuery::SpanError)] -#[case::identity("trace_identity", ReadQuery::TraceIdentity)] -#[case::spend("spend_by_response_ids", ReadQuery::SpendByResponseIds)] #[case::availability("availability", ReadQuery::Availability)] #[case::agents("agents", ReadQuery::Agents)] #[case::sample("sample", ReadQuery::Sample)] @@ -21,14 +15,11 @@ fn names_select_the_public_query(#[case] name: &str, #[case] query: ReadQuery) { #[rstest] #[case::unknown("unknown")] -#[case::case_sensitive("List_Traces")] -#[case::whitespace(" list_traces")] +#[case::case_sensitive("Sample")] +#[case::whitespace(" sample")] #[case::empty("")] fn invalid_names_preserve_the_public_error(#[case] name: &str) { let error = ReadQuery::parse(name).unwrap_err(); assert!(matches!(error, InvalidQuery)); assert_eq!(error.to_string(), "unknown ClickHouse read query"); } - -#[path = "query/named.rs"] -mod named; diff --git a/litellm-rust/crates/traces/tests/query/named.rs b/litellm-rust/crates/traces/tests/query/named.rs deleted file mode 100644 index 6b853ebf151..00000000000 --- a/litellm-rust/crates/traces/tests/query/named.rs +++ /dev/null @@ -1,67 +0,0 @@ -use litellm_traces::query::named::*; -use rstest::rstest; -use serde::{Serialize, de::DeserializeOwned}; -use serde_json::{Value, json}; - -fn round_trip(wire: Value) { - let contract: T = serde_json::from_value(wire.clone()).unwrap(); - assert_eq!(serde_json::to_value(contract).unwrap(), wire); -} - -#[rstest] -#[case::admin(1, "", vec![])] -#[case::own_user(0, "user", vec![])] -#[case::multiple_teams(0, "user", vec!["team-a", "team-b"])] -#[case::no_identity(0, "", vec![])] -fn named_requests_preserve_all_access_cases( - #[case] all_teams: u8, - #[case] user: &str, - #[case] teams: Vec<&str>, -) { - let access = json!({"all_teams": all_teams, "user_id": user, "team_ids": teams}); - round_trip::(access.clone()); - let request = |specific: Value| { - Value::Object( - access - .as_object() - .unwrap() - .iter() - .chain(specific.as_object().unwrap()) - .map(|(key, value)| (key.clone(), value.clone())) - .collect(), - ) - }; - round_trip::(request( - json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100}), - )); - round_trip::(request(json!({"trace_id": "trace"}))); - round_trip::(request(json!({"trace_id": "trace", "trace_ref": "ref"}))); - round_trip::(request( - json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span"}), - )); - round_trip::(request( - json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}), - )); - round_trip::(request( - json!({"response_ids": ["response"], "request_ids": ["request"], "trace_ids": ["trace"], "start_ms": -1, "end_ms": 10}), - )); -} - -#[rstest] -fn result_contracts_preserve_public_field_names() { - round_trip::( - 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": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), - ); - round_trip::( - json!({"trace_id": "trace", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "framework", "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"}), - ); - round_trip::( - json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}), - ); - round_trip::( - json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}), - ); - round_trip::( - json!({"request_id": "request", "litellm_call_id": "gateway", "response_id": "response", "upstream_response_id": "upstream", "trace_id": "trace", "span_id": "span", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}), - ); -} diff --git a/litellm-rust/crates/traces/tests/resolve.rs b/litellm-rust/crates/traces/tests/resolve.rs index 6e987cd4970..2c62dce27a1 100644 --- a/litellm-rust/crates/traces/tests/resolve.rs +++ b/litellm-rust/crates/traces/tests/resolve.rs @@ -1,15 +1,14 @@ use litellm_traces::{ - AgentNode, SpanStatus, iso_time, listed_summary, - query::named::{ListTracesRow, SpendByResponseIdsRow, TraceSpansRow}, - resolve_trace, + AgentNode, SpanStatus, iso_time, listed_summary, resolve_trace, + store::{CallRow, RunRow, SpanRow}, }; use rstest::rstest; const T0: i64 = 1_790_742_989_000_000_000; const MS: i64 = 1_000_000; -fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> TraceSpansRow { - TraceSpansRow { +fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> SpanRow { + SpanRow { trace_id: String::new(), span_id: span_id.into(), parent_span_id: parent.into(), @@ -38,14 +37,14 @@ fn row(span_id: &str, parent: &str, name: &str, kind: &str, agent: &str) -> Trac } } -fn at(mut span: TraceSpansRow, start_ms: i64, duration_ms: u64) -> TraceSpansRow { +fn at(mut span: SpanRow, start_ms: i64, duration_ms: u64) -> SpanRow { span.start_ns = T0 + start_ms * MS; span.duration_ns = duration_ms * MS as u64; span } -fn llm(span_id: &str, parent: &str, agent: &str, response_id: &str) -> TraceSpansRow { - TraceSpansRow { +fn llm(span_id: &str, parent: &str, agent: &str, response_id: &str) -> SpanRow { + SpanRow { model: "claude-sonnet-4-5".into(), input_tokens: 100, output_tokens: 20, @@ -55,7 +54,7 @@ fn llm(span_id: &str, parent: &str, agent: &str, response_id: &str) -> TraceSpan } } -fn owned(mut span: TraceSpansRow, team: &str, user: &str, key: &str) -> TraceSpansRow { +fn owned(mut span: SpanRow, team: &str, user: &str, key: &str) -> SpanRow { span.team_id = team.into(); span.user_id = user.into(); span.api_key_hash = key.into(); @@ -69,8 +68,8 @@ fn spend( user: &str, key: &str, cost: f64, -) -> SpendByResponseIdsRow { - SpendByResponseIdsRow { +) -> CallRow { + CallRow { request_id: request_id.into(), response_id: response_id.into(), litellm_call_id: String::new(), @@ -86,7 +85,7 @@ fn spend( } /// root agent -> llm, task tool -> researcher subagent (N times) -> llm + search tool + middleware. -fn deep_agent(researchers: usize) -> Vec { +fn deep_agent(researchers: usize) -> Vec { let mut rows = vec![ at( row( @@ -147,7 +146,7 @@ fn deep_agent(researchers: usize) -> Vec { rows } -fn agents(rows: &[TraceSpansRow]) -> Vec { +fn agents(rows: &[SpanRow]) -> Vec { resolve_trace("t", "", rows, &[]) .map(|trace| trace.agents) .unwrap_or_default() @@ -259,7 +258,7 @@ fn parent_agent_skips_same_name_ancestors_and_stops_at_cycles() { #[rstest] fn unnamed_calls_belong_to_the_nearest_agent_and_wrappers_are_not_agents() { - let crew = TraceSpansRow { + let crew = SpanRow { wrapper_candidate: true, ..row("crew", "", "crew.kickoff", "agent", "") }; @@ -285,7 +284,7 @@ fn unnamed_calls_belong_to_the_nearest_agent_and_wrappers_are_not_agents() { #[rstest] fn named_wrapper_inside_the_same_agent_is_a_chain() { - let wrapper = TraceSpansRow { + let wrapper = SpanRow { wrapper_candidate: true, ..row("w", "a", "researcher.run", "agent", "researcher") }; @@ -315,7 +314,7 @@ fn agents_named_only_by_their_tools_are_agents() { #[rstest] fn overlapping_tool_spans_count_one_call() { - let tool = |span_id: &str| TraceSpansRow { + let tool = |span_id: &str| SpanRow { tool_call_id: "call-1".into(), ..row(span_id, "a", "search", "tool", "") }; @@ -336,7 +335,7 @@ fn overlapping_tool_spans_count_one_call() { #[rstest] fn names_and_frameworks_are_sorted_and_distinct() { - let framed = |span: TraceSpansRow, framework: &str| TraceSpansRow { + let framed = |span: SpanRow, framework: &str| SpanRow { framework: framework.into(), ..span }; @@ -617,7 +616,7 @@ fn sibling_transport_does_not_lose_model_call_spend( #[case] expected: Option, ) { let call = owned( - TraceSpansRow { + SpanRow { trace_id: "trace".into(), call_keys: vec![litellm_traces::CallKey::ProviderResponse( "chatcmpl-1".into(), @@ -654,7 +653,7 @@ fn sibling_transport_does_not_lose_model_call_spend( })) .collect(); let logs: Vec<_> = std::iter::once(spend("chatcmpl-1", "chatcmpl-1", "team", "", "key", 0.5)) - .chain(unrelated_spend.then(|| SpendByResponseIdsRow { + .chain(unrelated_spend.then(|| CallRow { trace_id: "trace".into(), span_id: "tool-http".into(), ..spend("unrelated", "unrelated", "team", "", "key", 0.75) @@ -714,7 +713,7 @@ fn gateway_attempt_identifiers_must_match_one_spend_row( ), owned(llm("call", "agent", "agent", ""), "team", "", "key"), owned( - TraceSpansRow { + SpanRow { trace_id: "trace".into(), call_keys: keys, call_evidence: Some(litellm_traces::CallEvidenceKind::Complete), @@ -726,13 +725,13 @@ fn gateway_attempt_identifiers_must_match_one_spend_row( ), ]; let logs = [ - SpendByResponseIdsRow { + CallRow { litellm_call_id: "call-a".into(), trace_id: "trace".into(), span_id: "attempt".into(), ..spend("request-a", "response-a", "team", "", "key", 0.25) }, - SpendByResponseIdsRow { + CallRow { litellm_call_id: "call-b".into(), trace_id: "trace".into(), span_id: "other-attempt".into(), @@ -781,7 +780,7 @@ fn gateway_id_miss_only_vetoes_rows_that_carry_a_call_id( #[rstest] fn listed_summary_keeps_rollup_counts_with_unknown_cost() { - let summary = listed_summary(&ListTracesRow { + let summary = listed_summary(&RunRow { trace_id: "t1".into(), trace_ref: "ref".into(), team_id: "team".into(), @@ -804,7 +803,6 @@ fn listed_summary_keeps_rollup_counts_with_unknown_cost() { output_tokens: 2_620, models: vec!["claude-sonnet-4-5".into()], error_count: 1, - request_ids: Vec::new(), }); assert_eq!(summary.spend, None); assert_eq!(summary.status, SpanStatus::Ok); @@ -834,7 +832,7 @@ fn complete_wrapper_reconciles_ambiguous_response( #[case] exact_id: &str, #[case] expected: Option, ) { - let wrapper = TraceSpansRow { + let wrapper = SpanRow { call_keys: vec![litellm_traces::CallKey::LiteLlmRequest(exact_id.to_owned())], call_evidence: Some(litellm_traces::CallEvidenceKind::Complete), ..owned(llm("wrapper", "", "agent", ""), "team", "", "key") @@ -874,7 +872,7 @@ fn complete_correlation_requires_known_finite_cost( "", "key", )]; - let logged = SpendByResponseIdsRow { + let logged = CallRow { spend: cost, ..spend("request", "response", "team", "", "key", 0.25) }; @@ -896,7 +894,7 @@ fn transports_preserve_retry_spend_without_counting_unrelated_cached_rows( ) { let transport = |id: &str| { owned( - TraceSpansRow { + SpanRow { trace_id: "trace".into(), call_keys: vec!["transport:".parse().unwrap()], call_evidence: Some(if id == "first" { @@ -922,12 +920,12 @@ fn transports_preserve_retry_spend_without_counting_unrelated_cached_rows( transport("second"), ]; let logs = [ - SpendByResponseIdsRow { + CallRow { trace_id: "trace".into(), span_id: "first".into(), ..spend("retry", "retry-response", "team", "", "key", 0.25) }, - SpendByResponseIdsRow { + CallRow { trace_id: "trace".into(), span_id: "second".into(), ..spend("final", "final-response", "team", "", "key", 0.5) @@ -945,7 +943,7 @@ fn transports_preserve_retry_spend_without_counting_unrelated_cached_rows( #[case::ambiguous_response(true)] fn multiple_identifiers_for_one_request_count_its_spend_once(#[case] cached_row: bool) { let rows = [owned( - TraceSpansRow { + SpanRow { call_keys: vec![ "provider_response:response".parse().unwrap(), "litellm_request:request".parse().unwrap(), @@ -989,7 +987,7 @@ fn trace_cost_requires_a_finite_total(#[case] cost: f64, #[case] expected: Optio fn complete_wrapper_accounts_for_retries_missing_from_the_call_span() { let rows = [ owned( - TraceSpansRow { + SpanRow { call_keys: vec![ "litellm_request:retry".parse().unwrap(), "litellm_request:final".parse().unwrap(), @@ -1027,7 +1025,7 @@ fn legacy_request_id_fallback_respects_recorded_evidence( #[case] expected: Option, ) { let span = owned( - TraceSpansRow { + SpanRow { call_evidence: evidence, ..llm("call", "", "agent", "response") }, @@ -1036,7 +1034,7 @@ fn legacy_request_id_fallback_respects_recorded_evidence( "key", ); let stored = serde_json::to_value(span).unwrap(); - let decoded: TraceSpansRow = serde_json::from_value(stored).unwrap(); + let decoded: SpanRow = serde_json::from_value(stored).unwrap(); let logs = [spend("request", "response", "team", "", "key", 0.25)]; let trace = resolve_trace("trace", "ref", &[decoded], &logs).unwrap(); assert_eq!(trace.summary.spend, expected); @@ -1055,7 +1053,7 @@ fn malformed_stored_span_fields_are_rejected( ) { let mut encoded = serde_json::to_value(row("span", "", "agent", "agent", "agent")).unwrap(); encoded[field] = value; - assert!(serde_json::from_value::(encoded).is_err()); + assert!(serde_json::from_value::(encoded).is_err()); } #[rstest] @@ -1064,7 +1062,7 @@ fn malformed_stored_span_fields_are_rejected( fn spend_lookup_fetches_recorded_keys_before_resolving_completeness( #[case] evidence: litellm_traces::CallEvidenceKind, ) { - let recorded = TraceSpansRow { + let recorded = SpanRow { trace_id: "trace".to_owned(), call_keys: vec![ litellm_traces::CallKey::ProviderResponse("response".to_owned()), @@ -1084,12 +1082,12 @@ fn spend_lookup_fetches_recorded_keys_before_resolving_completeness( #[case::parent_first(false)] #[case::child_first(true)] fn overlapping_model_spans_count_leaf_usage_and_keep_agent_ownership(#[case] reverse: bool) { - let root = TraceSpansRow { + let root = SpanRow { input_tokens: 900, output_tokens: 800, ..row("root", "", "planner", "agent", "planner") }; - let wrapper = TraceSpansRow { + let wrapper = SpanRow { input_tokens: 700, output_tokens: 600, ..llm("wrapper", "root", "", "") @@ -1116,20 +1114,20 @@ fn overlapping_model_spans_count_leaf_usage_and_keep_agent_ownership(#[case] rev fn empty_root_preview_uses_the_earliest_agent_or_model_input() { let rows = [ at( - TraceSpansRow { + SpanRow { input_preview: "later input".into(), ..llm("later", "root", "", "") }, 20, 1, ), - TraceSpansRow { + SpanRow { input_preview: String::new(), ..row("root", "", "planner", "agent", "planner") }, at(row("tool", "root", "search", "tool", ""), 1, 1), at( - TraceSpansRow { + SpanRow { input_preview: "earlier input".into(), ..llm("earlier", "root", "", "") }, @@ -1148,11 +1146,11 @@ fn empty_root_preview_uses_the_earliest_agent_or_model_input() { #[case::oldest_first(false)] #[case::newest_first(true)] fn repeated_request_ids_preserve_storage_identity(#[case] reverse: bool) { - let first = SpendByResponseIdsRow { + let first = CallRow { start_ms: 100, ..spend("same", "response", "team", "", "key", 0.25) }; - let second = SpendByResponseIdsRow { + let second = CallRow { start_ms: 200, ..spend("same", "response", "team", "", "key", 0.5) }; @@ -1177,7 +1175,7 @@ fn repeated_request_ids_preserve_storage_identity(#[case] reverse: bool) { #[case::transport(true)] fn independent_key_disambiguates_repeated_request_ids(#[case] transport: bool) { let rows = [owned( - TraceSpansRow { + SpanRow { trace_id: "trace".into(), call_keys: vec![ litellm_traces::CallKey::ProviderResponse("response".into()), @@ -1194,14 +1192,14 @@ fn independent_key_disambiguates_repeated_request_ids(#[case] transport: bool) { "key", )]; let logs = [ - SpendByResponseIdsRow { + CallRow { start_ms: 100, litellm_call_id: "gateway".into(), trace_id: "trace".into(), span_id: "call".into(), ..spend("same", "response", "team", "", "key", 0.25) }, - SpendByResponseIdsRow { + CallRow { start_ms: 200, litellm_call_id: "other".into(), ..spend("same", "response", "team", "", "key", 0.5) @@ -1229,11 +1227,11 @@ fn totals_deduplicate_only_equal_storage_identities( ), ]; let logs = [ - SpendByResponseIdsRow { + CallRow { start_ms: 100, ..spend("same", "a", "team", "", "key", 0.25) }, - SpendByResponseIdsRow { + CallRow { start_ms: if duplicate { 100 } else { 200 }, ..spend( "same", @@ -1258,7 +1256,7 @@ fn totals_deduplicate_only_equal_storage_identities( fn conflicting_keys_cannot_agree_on_request_id_alone() { let rows = [ owned( - TraceSpansRow { + SpanRow { call_keys: vec![litellm_traces::CallKey::LiteLlmRequest("gateway".into())], ..llm("wrapper", "", "agent", "") }, @@ -1274,11 +1272,11 @@ fn conflicting_keys_cannot_agree_on_request_id_alone() { ), ]; let logs = [ - SpendByResponseIdsRow { + CallRow { start_ms: 100, ..spend("same", "response", "team", "", "key", 0.25) }, - SpendByResponseIdsRow { + CallRow { start_ms: 200, litellm_call_id: "gateway".into(), ..spend("same", "other", "team", "", "key", 0.5) @@ -1307,7 +1305,7 @@ fn gateway_lookup_respects_legacy_fallback_and_ownership( #[case] expected: Option, ) { let rows = [owned( - TraceSpansRow { + SpanRow { call_keys: vec![litellm_traces::CallKey::LiteLlmRequest("gateway".into())], ..llm("call", "", "agent", "") }, @@ -1315,7 +1313,7 @@ fn gateway_lookup_respects_legacy_fallback_and_ownership( "", "key", )]; - let logs = [SpendByResponseIdsRow { + let logs = [CallRow { litellm_call_id: call_id.into(), ..spend(request_id, "provider", team, "", key, 0.25) }]; diff --git a/litellm-rust/crates/traces/tests/search.rs b/litellm-rust/crates/traces/tests/search.rs new file mode 100644 index 00000000000..d2c7f4a8cd4 --- /dev/null +++ b/litellm-rust/crates/traces/tests/search.rs @@ -0,0 +1,92 @@ +use litellm_traces::{ + search::{AgentRuns, FieldFilter, RunField, RunSearch, histogram}, + store::RunCount, +}; +use rstest::rstest; + +fn filter(field: RunField, pattern: &str, exclude: bool) -> FieldFilter { + FieldFilter { + field, + pattern: pattern.into(), + exclude, + } +} + +#[rstest] +#[case::empty("", &[], vec![])] +#[case::words("foo bar", &["foo", "bar"], vec![])] +#[case::quoted_phrase(r#""foo bar""#, &["foo bar"], vec![])] +#[case::unclosed_quote(r#""foo bar"#, &["foo bar"], vec![])] +#[case::inner_quote_kept(r#"a"b"c"#, &[r#"a"b"c"#], vec![])] +#[case::like_metacharacters_stay_literal("50%_off\\", &["50%_off\\"], vec![])] +#[case::exact(r#"name:"plan trip""#, &[], vec![filter(RunField::Name, "plan trip", false)])] +#[case::glob("agent:res*er", &[], vec![filter(RunField::Agent, "res*er", false)])] +#[case::glob_keeps_the_rest("model:gpt_4*", &[], vec![filter(RunField::Model, "gpt_4*", false)])] +#[case::negated("-status:error", &[], vec![filter(RunField::Status, "error", true)])] +#[case::key_ignores_case("Trace_ID:abc", &[], vec![filter(RunField::TraceId, "abc", false)])] +#[case::value_keeps_colons("input:a:b", &[], vec![filter(RunField::Input, "a:b", false)])] +#[case::missing_value_narrows_nothing("status: foo", &["foo"], vec![])] +#[case::unknown_key_is_text("color:red", &["color:red"], vec![])] +#[case::non_word_key_is_text("k1:v", &["k1:v"], vec![])] +#[case::negated_text_stays_text("-foo", &["-foo"], vec![])] +fn parse_matches_the_dashboard_search_grammar( + #[case] q: &str, + #[case] text: &[&str], + #[case] filters: Vec, +) { + assert_eq!( + RunSearch::parse(q), + RunSearch { + text: text.iter().map(|term| (*term).to_owned()).collect(), + filters, + } + ); +} + +fn row(bucket: u32, failed: bool, agent: &str, runs: u64) -> RunCount { + RunCount { + bucket, + failed, + value: agent.into(), + runs, + } +} + +#[test] +fn histogram_fills_every_bucket_and_splits_failures_from_agents() { + let shaped = histogram( + &[ + row(0, true, "a", 3), + row(0, false, "b", 1), + row(0, false, "a", 2), + row(0, true, "c", 1), + row(2, true, "a", 1), + ], + 100, + 110, + 3, + ); + let summary: Vec<_> = shaped + .buckets + .iter() + .map(|bucket| (bucket.start_ms, bucket.end_ms, bucket.total, bucket.failed)) + .collect(); + assert_eq!( + summary, + [(100, 103, 7, 4), (103, 106, 0, 0), (106, 110, 1, 1)] + ); + assert_eq!( + shaped.buckets[0].agents, + [ + AgentRuns { + agent: "a".into(), + runs: 2 + }, + AgentRuns { + agent: "b".into(), + runs: 1 + }, + ] + ); + assert!(shaped.buckets[2].agents.is_empty()); +} diff --git a/litellm/proxy/lens/endpoints.py b/litellm/proxy/lens/endpoints.py index 12119f07962..32339e7befe 100644 --- a/litellm/proxy/lens/endpoints.py +++ b/litellm/proxy/lens/endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.lens.models import ( ) from litellm.proxy.lens.release import PROTOCOL_VERSION, release_tag, worker_image from litellm.proxy.lens.repository import LensRepository, WriterDatabase +from litellm.proxy.lens.search import LensField, parse_search from litellm.proxy.lens.sources import ActivityAvailability, SourceReader, Storage, parse_execution from litellm.proxy.lens.state import ( add_step, @@ -197,10 +198,21 @@ async def validate_workers(settings: LensSettings, scope: Scope) -> None: @router.get("", response_model=LensList) -async def list_lenses(auth: Auth, storage: StorageDep) -> LensList: +async def list_lenses( + auth: Auth, + storage: StorageDep, + q: Annotated[ + str, + Query( + max_length=1000, + description="Free text and key:value filters, e.g. `status:failed -schedule:paused`. " + "Keys: name, agent, status, schedule. `*` globs and a leading `-` negates", + ), + ] = "", +) -> LensList: scope: Final = user_scope(auth) return LensList( - lenses=tuple(e for e in await repository().lenses() if can_access(scope, e.scope)), + lenses=tuple(e for e in await repository().search(parse_search(q)) if can_access(scope, e.scope)), workers=tuple(w for w in await repository().workers() if can_access(scope, w.scope)), tracing_enabled=storage is not None, ) @@ -235,6 +247,17 @@ async def list_agents(auth: Auth, storage: StorageDep) -> tuple[str, ...]: return await source_reader(storage).agents(scope) if storage is not None else () +@router.get("/values/{field}", response_model=tuple[str, ...]) +async def list_lens_values( + field: LensField, + auth: Auth, + contains: Annotated[str, Query(max_length=200)] = "", + limit: Annotated[int, Query(ge=1, le=100)] = 20, +) -> tuple[str, ...]: + user_scope(auth) + return await repository().values(field, contains, limit) + + def watching(lens: Lens) -> Lens: if lens.settings.enabled: return lens diff --git a/litellm/proxy/lens/repository.py b/litellm/proxy/lens/repository.py index 6e1e2da112a..62a2f9d76ee 100644 --- a/litellm/proxy/lens/repository.py +++ b/litellm/proxy/lens/repository.py @@ -7,6 +7,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter from litellm.proxy.db.prisma_client import PrismaWrapper from litellm.proxy.lens.models import Job, Lens, Scope, Worker +from litellm.proxy.lens.search import FIELD_VALUES, LensField, LensSearch, like_literal, search_predicate class Database(Protocol): @@ -29,6 +30,28 @@ class LensRepository: rows: Final = _ROWS.validate_python(await self.db.query_raw('SELECT data FROM "LiteLLM_Lens" ORDER BY id')) return tuple(Lens.model_validate(row.data) for row in rows) + async def search(self, search: LensSearch) -> tuple[Lens, ...]: + where, args = search_predicate(search) + rows: Final = _ROWS.validate_python( + await self.db.query_raw( + f"""SELECT data FROM "LiteLLM_Lens" WHERE {where} + ORDER BY (data->>'created_at')::timestamptz DESC, id""", + *args, + ) + ) + return tuple(Lens.model_validate(row.data) for row in rows) + + async def values(self, field: LensField, contains: str, limit: int) -> tuple[str, ...]: + rows: Final = _ROWS.validate_python( + await self.db.query_raw( + f"""SELECT v AS data FROM "LiteLLM_Lens", unnest(array_remove({FIELD_VALUES[field]}, '')) AS v + WHERE v ILIKE $1 GROUP BY v ORDER BY count(*) DESC, v LIMIT $2""", + f"%{like_literal(contains)}%", + limit, + ) + ) + return tuple(TypeAdapter(str).validate_python(row.data) for row in rows) + async def get(self, lens_id: str) -> Lens | None: rows: Final = _ROWS.validate_python( await self.db.query_raw( diff --git a/litellm/proxy/lens/search.py b/litellm/proxy/lens/search.py new file mode 100644 index 00000000000..78701cc1517 --- /dev/null +++ b/litellm/proxy/lens/search.py @@ -0,0 +1,104 @@ +import re +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal, TypeAlias + +LensField: TypeAlias = Literal["name", "agent", "status", "schedule"] + +_SETTINGS: Final = "data->'settings'" +FIELD_VALUES: Final[MappingProxyType[LensField, str]] = MappingProxyType( + { + "name": f"ARRAY[{_SETTINGS}->>'name']", + "agent": f"ARRAY[{_SETTINGS}->>'agent_name', {_SETTINGS}->>'service']", + "status": "ARRAY[COALESCE(data->'jobs'->0->>'status', 'never')]", + "schedule": f"ARRAY[CASE WHEN ({_SETTINGS}->>'enabled')::boolean THEN 'watching' ELSE 'paused' END]", + } +) + +_SCOPE_LABEL: Final = f"""COALESCE(NULLIF(concat_ws(' · ', + NULLIF({_SETTINGS}->>'agent_name', ''), + NULLIF({_SETTINGS}->>'service', ''), + (SELECT string_agg((f->>'key') || ': ' || (f->>'value'), ' · ') + FROM jsonb_array_elements({_SETTINGS}->'filters') AS f)), ''), 'All activity')""" +FREE_TEXT: Final = f"ARRAY[{_SETTINGS}->>'name', {_SCOPE_LABEL}]" + +_TOKEN: Final = re.compile(r'(?:"[^"]*"?|\S)+') +_FIELD_TOKEN: Final = re.compile(r"^(-?)([A-Za-z_]+):(.*)$", re.DOTALL) +_QUOTED: Final = re.compile(r'^"([^"]*)"?$') + + +@dataclass(frozen=True, slots=True) +class LensFilter: + field: LensField + pattern: str + exclude: bool + + +@dataclass(frozen=True, slots=True) +class LensSearch: + text: tuple[str, ...] + filters: tuple[LensFilter, ...] + + +@dataclass(frozen=True, slots=True) +class _Text: + value: str + + +@dataclass(frozen=True, slots=True) +class _Field: + field: LensField + exclude: bool + value: str + + +def like_literal(value: str) -> str: + return re.sub(r"([\\%_])", r"\\\1", value) + + +def _unquote(raw: str) -> str: + match: Final = _QUOTED.match(raw) + return match.group(1) if match else raw + + +def _as_field(key: str) -> LensField | None: + match key: + case "name" | "agent" | "status" | "schedule": + return key + case _: + return None + + +def _clause(raw: str) -> _Text | _Field: + match: Final = _FIELD_TOKEN.match(raw) + field: Final = _as_field(match.group(2).lower()) if match else None + if match is None or field is None: + return _Text(_unquote(raw)) + return _Field(field=field, exclude=match.group(1) == "-", value=_unquote(match.group(3))) + + +def parse_search(q: str) -> LensSearch: + clauses: Final = tuple(clause for clause in map(_clause, _TOKEN.findall(q)) if clause.value) + return LensSearch( + text=tuple(f"%{like_literal(clause.value)}%" for clause in clauses if isinstance(clause, _Text)), + filters=tuple( + LensFilter(field=clause.field, pattern=like_literal(clause.value).replace("*", "%"), exclude=clause.exclude) + for clause in clauses + if isinstance(clause, _Field) + ), + ) + + +def _any_matches(values: str, placeholder: str) -> str: + return f"EXISTS (SELECT 1 FROM unnest(array_remove({values}, '')) AS v WHERE v ILIKE {placeholder})" + + +def search_predicate(search: LensSearch) -> tuple[str, tuple[str, ...]]: + terms: Final = tuple((FREE_TEXT, pattern, False) for pattern in search.text) + tuple( + (FIELD_VALUES[f.field], f.pattern, f.exclude) for f in search.filters + ) + predicates: Final = tuple( + f"{'NOT ' if exclude else ''}{_any_matches(values, f'${index}')}" + for index, (values, _, exclude) in enumerate(terms, start=1) + ) + return " AND ".join(predicates) or "TRUE", tuple(pattern for _, pattern, _ in terms) diff --git a/litellm/proxy/tracing_endpoints.py b/litellm/proxy/tracing_endpoints.py index 0d3723c45b8..1d91615da6e 100644 --- a/litellm/proxy/tracing_endpoints.py +++ b/litellm/proxy/tracing_endpoints.py @@ -3,6 +3,8 @@ Agent tracing endpoints. Thin wrappers over `TraceReceiver`: auth -> tenant/scop POST /v1/traces OTLP/HTTP trace export (protobuf or JSON) GET /v1/traces TracePage +GET /v1/traces/histogram TraceHistogram +GET /v1/traces/values/{field} RunValues GET /v1/traces/{trace_id} Trace GET /v1/traces/{trace_id}/spans/{span_id} SpanDetail """ @@ -33,11 +35,13 @@ from litellm.rust_bridge.trace.generated.types import ( AllQueryScope, OwnedQueryScope, QueryScope, + RunField, + RunValues, SpanDetail, SpanErrorPage, Trace, + TraceHistogram, TracePage, - TraceScope, ) from litellm.rust_bridge.trace.queries import TraceSQLResponse from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant @@ -55,11 +59,11 @@ class TraceAccessContext: read_scope: ReadScope | None write_tenant: Tenant | None - def reader(self) -> tuple[TraceReceiver, TraceScope]: + def reader(self) -> tuple[TraceReceiver, QueryScope]: tracing: Final = require_receiver(self.receiver) if self.read_scope is None: raise HTTPException(status_code=403, detail="Not allowed to view agent traces") - return tracing, _trace_scope(self.read_scope) + return tracing, read_access(self.read_scope) def writer(self) -> tuple[TraceReceiver, Tenant]: if self.write_tenant is None: @@ -80,16 +84,6 @@ async def provide_trace_access( return TraceAccessContext(tracing, read_scope, write_tenant) -def _trace_scope(scope: ReadScope) -> TraceScope: - if isinstance(scope, AllRows): - return TraceScope(all_teams=1, user_id="", team_ids=()) - return TraceScope( - all_teams=0, - user_id=scope.user_id or "", - team_ids=scope.team_ids, - ) - - def otlp_error_response( request: Request, status_code: int, headers: Mapping[str, str] | None = None ) -> Response | None: @@ -180,26 +174,78 @@ def read_failure(error: TraceChanged | ValueError | OverflowError | RuntimeError return assert_never(error) +StartMs = Annotated[int | None, Query(description="Window start, unix ms. Default: 24h ago")] +EndMs = Annotated[int | None, Query(description="Window end, unix ms. Default: now")] +RunQuery = Annotated[ + str, + Query( + max_length=1000, + description='Free text and key:value filters, e.g. `agent:research* -status:ok "book a flight"`. ' + "Keys: name, agent, status, model, input, trace_id. `*` globs and a leading `-` negates", + ), +] + + +@dataclass(frozen=True, slots=True) +class TraceWindow: + start_ms: int + end_ms: int + + +def trace_window(start_ms: StartMs = None, end_ms: EndMs = None) -> TraceWindow: + now_ms: Final = int(time.time() * 1000) + return TraceWindow( + start_ms=start_ms if start_ms is not None else now_ms - MS_PER_DAY, + end_ms=end_ms if end_ms is not None else now_ms, + ) + + @router.get("/v1/traces", response_model=TracePage) async def list_agent_traces( context: Annotated[TraceAccessContext, Depends(provide_trace_access)], - start_ms: Annotated[int | None, Query(description="Window start, unix ms. Default: 24h ago")] = None, - end_ms: Annotated[int | None, Query(description="Window end, unix ms. Default: now")] = None, + window: Annotated[TraceWindow, Depends(trace_window)], + q: RunQuery = "", cursor: Annotated[str | None, Query(max_length=512)] = None, ) -> TracePage: - now_ms: Final = int(time.time() * 1000) try: tracing, scope = context.reader() return await tracing.list_traces( - scope=scope, - start_ms=start_ms if start_ms is not None else now_ms - MS_PER_DAY, - end_ms=end_ms if end_ms is not None else now_ms, - cursor=cursor, + scope=scope, start_ms=window.start_ms, end_ms=window.end_ms, q=q, cursor=cursor ) except (TraceChanged, ValueError, OverflowError, RuntimeError) as error: raise read_failure(error) from error +@router.get("/v1/traces/histogram", response_model=TraceHistogram) +async def agent_trace_histogram( + context: Annotated[TraceAccessContext, Depends(provide_trace_access)], + window: Annotated[TraceWindow, Depends(trace_window)], + q: RunQuery = "", + buckets: Annotated[int, Query(ge=1, le=240)] = 60, +) -> TraceHistogram: + try: + tracing, scope = context.reader() + return await tracing.trace_histogram(scope, window.start_ms, window.end_ms, q, buckets) + except (ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error + + +@router.get("/v1/traces/values/{field}", response_model=RunValues) +async def agent_trace_values( + field: RunField, + context: Annotated[TraceAccessContext, Depends(provide_trace_access)], + window: Annotated[TraceWindow, Depends(trace_window)], + q: RunQuery = "", + contains: Annotated[str, Query(max_length=200)] = "", + limit: Annotated[int, Query(ge=1, le=100)] = 20, +) -> RunValues: + try: + tracing, scope = context.reader() + return await tracing.run_values(scope, window.start_ms, window.end_ms, q, field, contains, limit) + except (ValueError, OverflowError, RuntimeError) as error: + raise read_failure(error) from error + + class TraceQueryRequest(BaseModel): model_config = ConfigDict(frozen=True, extra="forbid") sql: str @@ -220,7 +266,7 @@ def provide_trace_query_secret() -> str: return master_key -def trace_query_scope(scope: ReadScope) -> QueryScope: +def read_access(scope: ReadScope) -> QueryScope: if isinstance(scope, AllRows): return AllQueryScope(kind="all") return OwnedQueryScope( @@ -249,7 +295,7 @@ async def query_agent_traces( access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)], ) -> TraceSQLResponse: try: - return await access.storage.query_sql(body.sql, trace_query_scope(access.scope), access.secret) + return await access.storage.query_sql(body.sql, read_access(access.scope), access.secret) except ValueError as error: raise HTTPException(status_code=400, detail=str(error)) from error except RuntimeError as error: @@ -262,7 +308,7 @@ async def help_agent_trace_queries( access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)], ) -> TraceQueryHelp: try: - return await access.storage.query_help(trace_query_scope(access.scope), access.secret) + return await access.storage.query_help(read_access(access.scope), access.secret) except RuntimeError as error: verbose_proxy_logger.warning("Trace query help unavailable: %s", error) raise HTTPException(status_code=503, detail="Trace query help is temporarily unavailable") from error diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c146a6eac92..f3f7b4114d8 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -11,7 +11,7 @@ from litellm.rust_bridge.embeddings.entrypoints import LiteLLMEmbeddingRequest from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest -from litellm.rust_bridge.trace.generated.types import QueryScope, ReadQueryName, TraceScope +from litellm.rust_bridge.trace.generated.types import QueryScope, ReadQueryName from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import EmbeddingResponse, ModelResponse @@ -43,14 +43,20 @@ class NativeTraceStorage: def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... def ingest(self, payload: bytes, content_type: str | None, tenant: Mapping[str, str]) -> Future[int]: ... def list_traces( - self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None, limit: int + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, cursor: str | None, limit: int + ) -> Future[JsonValue]: ... + def trace_histogram( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int + ) -> Future[JsonValue]: ... + def run_values( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, field: str, contains: str, limit: int ) -> Future[JsonValue]: ... def get_trace( - self, trace_id: str, scope: TraceScope, trace_ref: str, cursor: str | None = None, page_size: int | None = None + self, trace_id: str, scope: QueryScope, trace_ref: str, cursor: str | None = None, page_size: int | None = None ) -> Future[JsonValue]: ... - def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str) -> Future[JsonValue]: ... + def get_span(self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str) -> Future[JsonValue]: ... def get_span_error( - self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str, cursor: str | None + self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str, cursor: str | None ) -> Future[JsonValue]: ... def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Future[str]: ... def query_help(self, scope: QueryScope, secret: str) -> Future[JsonValue]: ... diff --git a/litellm/rust_bridge/trace/generated/types.py b/litellm/rust_bridge/trace/generated/types.py index 127e86e9160..b7499d38f07 100644 --- a/litellm/rust_bridge/trace/generated/types.py +++ b/litellm/rust_bridge/trace/generated/types.py @@ -23,6 +23,13 @@ class OwnedQueryScope(typing_extensions.TypedDict): QueryScope: TypeAlias = AllQueryScope | OwnedQueryScope +RunField: TypeAlias = Literal["name", "agent", "status", "model", "input", "trace_id"] + + +class RunValues(typing_extensions.TypedDict): + values: ReadOnly[tuple[str, ...]] + + class UIText(typing_extensions.TypedDict): text: ReadOnly[str] kind: ReadOnly[Literal["text"]] @@ -77,10 +84,9 @@ SpanType: TypeAlias = Literal[ ] -class TraceScope(typing_extensions.TypedDict): - all_teams: ReadOnly[Literal[0, 1]] - user_id: ReadOnly[str] - team_ids: ReadOnly[tuple[str, ...]] +class AgentRuns(typing_extensions.TypedDict): + agent: ReadOnly[str] + runs: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] ReadQueryName: TypeAlias = Literal["availability", "agents", "sample", "content", "evidence"] @@ -149,6 +155,18 @@ class Trace(typing_extensions.TypedDict): next_cursor: ReadOnly[NotRequired[str | None]] +class HistogramBucket(typing_extensions.TypedDict): + start_ms: ReadOnly[Annotated[int, Field(ge=-9223372036854775808, le=9223372036854775807)]] + end_ms: ReadOnly[Annotated[int, Field(ge=-9223372036854775808, le=9223372036854775807)]] + total: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + failed: ReadOnly[Annotated[int, Field(ge=0, le=18446744073709551615)]] + agents: ReadOnly[tuple[AgentRuns, ...]] + + +class TraceHistogram(typing_extensions.TypedDict): + buckets: ReadOnly[tuple[HistogramBucket, ...]] + + class TracePage(typing_extensions.TypedDict): data: ReadOnly[tuple[TraceSummary, ...]] next_cursor: ReadOnly[str | None] @@ -171,4 +189,6 @@ class SpanDetail(typing_extensions.TypedDict): attributes: ReadOnly[Mapping[str, str]] -TraceWireTypes: TypeAlias = QueryScope | SpanDetail | SpanErrorPage | Trace | TracePage | TraceScope | ReadQueryName +TraceWireTypes: TypeAlias = ( + QueryScope | RunField | RunValues | SpanDetail | SpanErrorPage | Trace | TraceHistogram | TracePage | ReadQueryName +) diff --git a/litellm/rust_bridge/trace/storage.py b/litellm/rust_bridge/trace/storage.py index a80e41a5aa3..600ce4cdea3 100644 --- a/litellm/rust_bridge/trace/storage.py +++ b/litellm/rust_bridge/trace/storage.py @@ -32,11 +32,12 @@ from litellm.rust_bridge.trace.queries import ( from .generated.models import TraceQueryHelp from .generated.types import ( QueryScope, + RunValues, SpanDetail, SpanErrorPage, Trace, + TraceHistogram, TracePage, - TraceScope, ) from .queries import TraceSQLResponse @@ -64,17 +65,25 @@ class NativeStore(Protocol): def ingest(self, payload: bytes, content_type: str | None, tenant: Mapping[str, str]) -> Awaitable[int]: ... def list_traces( - self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None, limit: int + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, cursor: str | None, limit: int + ) -> Awaitable[JsonValue]: ... + + def trace_histogram( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int + ) -> Awaitable[JsonValue]: ... + + def run_values( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, field: str, contains: str, limit: int ) -> Awaitable[JsonValue]: ... def get_trace( - self, trace_id: str, scope: TraceScope, trace_ref: str, cursor: str | None = None, page_size: int | None = None + self, trace_id: str, scope: QueryScope, trace_ref: str, cursor: str | None = None, page_size: int | None = None ) -> Awaitable[JsonValue]: ... - def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str) -> Awaitable[JsonValue]: ... + def get_span(self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str) -> Awaitable[JsonValue]: ... def get_span_error( - self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str, cursor: str | None + self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str, cursor: str | None ) -> Awaitable[JsonValue]: ... def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Awaitable[str]: ... @@ -102,6 +111,8 @@ QUERY_PARAMETERS: Final = TypeAdapter(dict[str, str | int | float | list[str]]) _SQL_RESPONSE: Final = TypeAdapter(TraceSQLResponse) _HELP_RESPONSE: Final = TypeAdapter(TraceQueryHelp) _TRACE_PAGE: Final = TypeAdapter(TracePage) +_TRACE_HISTOGRAM: Final = TypeAdapter(TraceHistogram) +_RUN_VALUES: Final = TypeAdapter(RunValues) _TRACE: Final[TypeAdapter[Trace | None]] = TypeAdapter(Trace | None) _SPAN_DETAIL: Final[TypeAdapter[SpanDetail | None]] = TypeAdapter(SpanDetail | None) _SPAN_ERROR_PAGE: Final[TypeAdapter[SpanErrorPage | None]] = TypeAdapter(SpanErrorPage | None) @@ -182,19 +193,32 @@ class ClickHouseStorage: async def list_traces( self, - scope: TraceScope, + scope: QueryScope, start_ms: int, end_ms: int, + q: str = "", cursor: str | None = None, limit: int = AGENT_TRACING_LIST_PAGE_SIZE, ) -> TracePage: - result: Final = await self._native.list_traces(scope, start_ms, end_ms, cursor, limit) + result: Final = await self._native.list_traces(scope, start_ms, end_ms, q, cursor, limit) return _validate_query_response(_TRACE_PAGE, result) + async def trace_histogram( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int + ) -> TraceHistogram: + result: Final = await self._native.trace_histogram(scope, start_ms, end_ms, q, buckets) + return _validate_query_response(_TRACE_HISTOGRAM, result) + + async def run_values( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, field: str, contains: str, limit: int + ) -> RunValues: + result: Final = await self._native.run_values(scope, start_ms, end_ms, q, field, contains, limit) + return _validate_query_response(_RUN_VALUES, result) + async def get_trace( self, trace_id: str, - scope: TraceScope, + scope: QueryScope, trace_ref: str = "", cursor: str | None = None, page_size: int | None = None, @@ -202,12 +226,12 @@ class ClickHouseStorage: result: Final = await self._native.get_trace(trace_id, scope, trace_ref, cursor, page_size) return _validate_query_response(_TRACE, result) - async def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "") -> SpanDetail | None: + async def get_span(self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str = "") -> SpanDetail | None: result: Final = await self._native.get_span(trace_id, span_id, scope, trace_ref) return _validate_query_response(_SPAN_DETAIL, result) async def get_span_error( - self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "", cursor: str | None = None + self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str = "", cursor: str | None = None ) -> SpanErrorPage | None: result: Final = await self._native.get_span_error(trace_id, span_id, scope, trace_ref, cursor) return _validate_query_response(_SPAN_ERROR_PAGE, result) diff --git a/litellm/tracing/receiver.py b/litellm/tracing/receiver.py index cf6bf8feb01..7e6303f0acd 100644 --- a/litellm/tracing/receiver.py +++ b/litellm/tracing/receiver.py @@ -5,7 +5,7 @@ await tracing.start() # create tables if missing tracing.ingest(otlp_body, content_type, content_encoding, tenant) # POST /v1/traces - await tracing.list_traces(scope, start_ms, end_ms, cursor) # GET /v1/traces + await tracing.list_traces(scope, start_ms, end_ms, q, cursor) # GET /v1/traces await tracing.get_trace(trace_id, scope) # GET /v1/traces/{id} await tracing.get_span(trace_id, span_id, scope) # GET /v1/traces/{id}/spans/{span_id} @@ -19,7 +19,16 @@ from threading import BoundedSemaphore from typing import Final from litellm.constants import AGENT_TRACING_LIST_PAGE_SIZE, OTLP_MAX_BODY_BYTES, OTLP_MAX_CONCURRENT_INGESTS -from litellm.rust_bridge.trace.generated.types import SpanDetail, SpanErrorPage, Trace, TracePage, TraceScope +from litellm.rust_bridge.trace.generated.types import ( + QueryScope, + RunField, + RunValues, + SpanDetail, + SpanErrorPage, + Trace, + TraceHistogram, + TracePage, +) from litellm.rust_bridge.trace.storage import ClickHouseStorage, Tenant from litellm.tracing.config import trace_storage_config from litellm.tracing.otlp_http import InvalidOTLPPayloadError, TracingPayloadTooLargeError, decompress @@ -96,24 +105,36 @@ class TraceReceiver: except ValueError as error: raise InvalidOTLPPayloadError(str(error)) from error - async def list_traces(self, scope: TraceScope, start_ms: int, end_ms: int, cursor: str | None = None) -> TracePage: - return await self.storage.list_traces(scope, start_ms, end_ms, cursor, AGENT_TRACING_LIST_PAGE_SIZE) + async def list_traces( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str = "", cursor: str | None = None + ) -> TracePage: + return await self.storage.list_traces(scope, start_ms, end_ms, q, cursor, AGENT_TRACING_LIST_PAGE_SIZE) + + async def trace_histogram( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, buckets: int + ) -> TraceHistogram: + return await self.storage.trace_histogram(scope, start_ms, end_ms, q, buckets) + + async def run_values( + self, scope: QueryScope, start_ms: int, end_ms: int, q: str, field: RunField, contains: str, limit: int + ) -> RunValues: + return await self.storage.run_values(scope, start_ms, end_ms, q, field, contains, limit) async def get_trace( self, trace_id: str, - scope: TraceScope, + scope: QueryScope, trace_ref: str = "", cursor: str | None = None, page_size: int | None = None, ) -> Trace | None: return await self.storage.get_trace(trace_id, scope, trace_ref, cursor, page_size) - async def get_span(self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "") -> SpanDetail | None: + async def get_span(self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str = "") -> SpanDetail | None: return await self.storage.get_span(trace_id, span_id, scope, trace_ref) async def get_span_error( - self, trace_id: str, span_id: str, scope: TraceScope, trace_ref: str = "", cursor: str | None = None + self, trace_id: str, span_id: str, scope: QueryScope, trace_ref: str = "", cursor: str | None = None ) -> SpanErrorPage | None: return await self.storage.get_span_error(trace_id, span_id, scope, trace_ref, cursor) diff --git a/scripts/trace_codegen/schemas/traces/RunField.json b/scripts/trace_codegen/schemas/traces/RunField.json new file mode 100644 index 00000000000..3879382fa35 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/RunField.json @@ -0,0 +1,13 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "enum": [ + "name", + "agent", + "status", + "model", + "input", + "trace_id" + ], + "title": "RunField", + "type": "string" +} diff --git a/scripts/trace_codegen/schemas/traces/RunValues.json b/scripts/trace_codegen/schemas/traces/RunValues.json new file mode 100644 index 00000000000..9105ffb28c1 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/RunValues.json @@ -0,0 +1,17 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "description": "Distinct values of one run field, most common first.", + "properties": { + "values": { + "items": { + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "values" + ], + "title": "RunValues", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceHistogram.json b/scripts/trace_codegen/schemas/traces/TraceHistogram.json new file mode 100644 index 00000000000..9871d6d5683 --- /dev/null +++ b/scripts/trace_codegen/schemas/traces/TraceHistogram.json @@ -0,0 +1,80 @@ +{ + "$defs": { + "AgentRuns": { + "properties": { + "agent": { + "type": "string" + }, + "runs": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "agent", + "runs" + ], + "type": "object" + }, + "HistogramBucket": { + "properties": { + "agents": { + "description": "Runs that did not fail, by their alphabetically first agent label, or service when unlabelled.", + "items": { + "$ref": "#/$defs/AgentRuns" + }, + "type": "array" + }, + "end_ms": { + "format": "int64", + "maximum": 9223372036854775807, + "minimum": -9223372036854775808, + "type": "integer" + }, + "failed": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + }, + "start_ms": { + "format": "int64", + "maximum": 9223372036854775807, + "minimum": -9223372036854775808, + "type": "integer" + }, + "total": { + "format": "uint64", + "maximum": 18446744073709551615, + "minimum": 0, + "type": "integer" + } + }, + "required": [ + "start_ms", + "end_ms", + "total", + "failed", + "agents" + ], + "type": "object" + } + }, + "$schema": "https://json-schema.org/draft/2020-12/schema", + "description": "Matching runs per equal-width slice of the window.", + "properties": { + "buckets": { + "items": { + "$ref": "#/$defs/HistogramBucket" + }, + "type": "array" + } + }, + "required": [ + "buckets" + ], + "title": "TraceHistogram", + "type": "object" +} diff --git a/scripts/trace_codegen/schemas/traces/TraceScope.json b/scripts/trace_codegen/schemas/traces/TraceScope.json deleted file mode 100644 index c5c2159e646..00000000000 --- a/scripts/trace_codegen/schemas/traces/TraceScope.json +++ /dev/null @@ -1,28 +0,0 @@ -{ - "$schema": "https://json-schema.org/draft/2020-12/schema", - "properties": { - "all_teams": { - "enum": [ - 0, - 1 - ], - "type": "integer" - }, - "team_ids": { - "items": { - "type": "string" - }, - "type": "array" - }, - "user_id": { - "type": "string" - } - }, - "required": [ - "all_teams", - "user_id", - "team_ids" - ], - "title": "TraceScope", - "type": "object" -} diff --git a/tests/integration/database/test_lens_repository.py b/tests/integration/database/test_lens_repository.py index 29c6ad13825..3073e5895a3 100644 --- a/tests/integration/database/test_lens_repository.py +++ b/tests/integration/database/test_lens_repository.py @@ -1,7 +1,7 @@ import asyncio import os from collections.abc import AsyncIterator -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit @@ -14,8 +14,9 @@ from prisma import Prisma from psycopg import sql from litellm.proxy.db.prisma_client import PrismaWrapper -from litellm.proxy.lens.models import Check, Lens, LensSettings, Scope, Worker +from litellm.proxy.lens.models import Check, Lens, LensSettings, MetadataFilter, Scope, Worker from litellm.proxy.lens.repository import LensRepository, WriterDatabase +from litellm.proxy.lens.search import LensField, parse_search from litellm.proxy.lens.state import claim_job, queue_job @@ -53,6 +54,107 @@ async def test_concurrent_workers_cannot_both_acquire_the_same_job(lens_db: Pris await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id=$1', lens.id) +@pytest_asyncio.fixture(loop_scope="function") +async def searchable_lenses(lens_db: Prisma) -> AsyncIterator[tuple[LensRepository, str]]: + tag: Final = uuid4().hex[:12] + start: Final = datetime.now(timezone.utc) + repo: Final = LensRepository(WriterDatabase(PrismaWrapper(lens_db))) + + def lens(lens_id: str, minutes: int, status: str | None, **settings: object) -> Lens: + created: Final = start + timedelta(minutes=minutes) + base: Final = Lens( + id=f"{tag}-{lens_id}", + scope=Scope(all_teams=True), + settings=LensSettings.model_validate({"model": "test", "context": "Find failures", **settings}), + created_at=created, + next_run_at=created, + budget_month=created.strftime("%Y-%m"), + ) + if status is None: + return base + queued: Final = queue_job(base, created, uuid4().hex) + return queued.model_copy(update={"jobs": (queued.jobs[0].model_copy(update={"status": status}),)}) + + lenses: Final = ( + lens( + "a", + 0, + "failed", + name=f"{tag} nightly", + agent_name=f"{tag}-researcher", + filters=(MetadataFilter(key="env", value="prod"),), + ), + lens("b", 1, None, name=f"{tag} weekly 100%_x", service=f"{tag}-billing", enabled=False), + lens("c", 2, "completed", name=f"{tag} adhoc"), + ) + for created in lenses: + await repo.create(created) + try: + yield repo, tag + finally: + await lens_db.execute_raw('DELETE FROM "LiteLLM_Lens" WHERE id LIKE $1', f"{tag}-%") + + +@pytest.mark.parametrize( + "q, expected", + ( + ("", ("c", "b", "a")), + ("status:failed", ("a",)), + ("status:never", ("b",)), + ("-status:never", ("c", "a")), + ("status:FAIL*", ("a",)), + ("schedule:paused", ("b",)), + ("-schedule:paused", ("c", "a")), + ("agent:{tag}-researcher", ("a",)), + ("agent:{tag}-bill*", ("b",)), + ("agent:researcher", ()), + ("-agent:*", ("c",)), + ("name:*nightly", ("a",)), + ("name:nightly", ()), + ('"env: prod"', ("a",)), + ('"all activity"', ("c",)), + ("100%_x", ("b",)), + ("100%x", ()), + ("100__x", ()), + ("weekly status:never -schedule:watching", ("b",)), + ), +) +@pytest.mark.asyncio +async def test_search_selects_lenses_newest_first( + searchable_lenses: tuple[LensRepository, str], q: str, expected: tuple[str, ...] +) -> None: + repo, tag = searchable_lenses + found: Final = await repo.search(parse_search(f"{tag} {q.format(tag=tag)}")) + assert tuple(lens.id.removeprefix(f"{tag}-") for lens in found if lens.id.startswith(tag)) == expected + + +@pytest.mark.parametrize( + "field, needle, expected", + ( + ("name", "{tag}", ("{tag} adhoc", "{tag} nightly", "{tag} weekly 100%_x")), + ("name", "{tag} W", ("{tag} weekly 100%_x",)), + ("name", "{tag}_", ()), + ("name", "{tag} weekly 100%", ("{tag} weekly 100%_x",)), + ("name", "{tag} weekly 1%", ()), + ("agent", "{tag}", ("{tag}-billing", "{tag}-researcher")), + ("agent", "{tag}-RESEARCH", ("{tag}-researcher",)), + ), +) +@pytest.mark.asyncio +async def test_values_list_distinct_field_values_containing_the_needle( + searchable_lenses: tuple[LensRepository, str], field: LensField, needle: str, expected: tuple[str, ...] +) -> None: + repo, tag = searchable_lenses + assert await repo.values(field, needle.format(tag=tag), 100) == tuple(e.format(tag=tag) for e in expected) + + +@pytest.mark.asyncio +async def test_values_include_derived_status_and_schedule(searchable_lenses: tuple[LensRepository, str]) -> None: + repo, _ = searchable_lenses + assert {"failed", "completed", "never"} <= set(await repo.values("status", "", 100)) + assert {"watching", "paused"} <= set(await repo.values("schedule", "", 100)) + + @pytest.mark.asyncio async def test_heartbeat_never_restores_revoked_access(lens_db: Prisma) -> None: now: Final = datetime.now(timezone.utc) diff --git a/tests/unit/proxy/lens/test_search.py b/tests/unit/proxy/lens/test_search.py new file mode 100644 index 00000000000..26236c07628 --- /dev/null +++ b/tests/unit/proxy/lens/test_search.py @@ -0,0 +1,28 @@ +import pytest + +from litellm.proxy.lens.search import LensFilter, LensSearch, parse_search + + +@pytest.mark.parametrize( + "q, text, filters", + ( + pytest.param("", (), (), id="empty"), + pytest.param("foo bar", ("%foo%", "%bar%"), (), id="words"), + pytest.param('"foo bar"', ("%foo bar%",), (), id="quoted-phrase"), + pytest.param('"foo bar', ("%foo bar%",), (), id="unclosed-quote"), + pytest.param("50%_off\\", ("%50\\%\\_off\\\\%",), (), id="like-metacharacters"), + pytest.param('name:"nightly scan"', (), (LensFilter("name", "nightly scan", False),), id="exact"), + pytest.param("agent:res*er", (), (LensFilter("agent", "res%er", False),), id="glob"), + pytest.param("name:a_b*", (), (LensFilter("name", "a\\_b%", False),), id="glob-escapes-rest"), + pytest.param("-schedule:paused", (), (LensFilter("schedule", "paused", True),), id="negated"), + pytest.param("STATUS:failed", (), (LensFilter("status", "failed", False),), id="key-ignores-case"), + pytest.param("name:a:b", (), (LensFilter("name", "a:b", False),), id="value-keeps-colons"), + pytest.param("status: foo", ("%foo%",), (), id="missing-value-narrows-nothing"), + pytest.param("model:x", ("%model:x%",), (), id="unknown-key-is-text"), + pytest.param("-foo", ("%-foo%",), (), id="negated-text-stays-text"), + ), +) +def test_parse_matches_the_dashboard_search_grammar( + q: str, text: tuple[str, ...], filters: tuple[LensFilter, ...] +) -> None: + assert parse_search(q) == LensSearch(text=text, filters=filters) diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 913b1aec618..87ba4d1f61b 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -22,7 +22,7 @@ from litellm.proxy.tracing_runtime import manage_tracing, provide_storage from litellm.rust_bridge import loader from litellm.rust_bridge.trace.errors import TraceChanged from litellm.rust_bridge.trace.generated.models import TraceQueryHelp -from litellm.rust_bridge.trace.generated.types import AllQueryScope, TraceScope +from litellm.rust_bridge.trace.generated.types import AllQueryScope, OwnedQueryScope, QueryScope from litellm.rust_bridge.trace.queries import TraceSQLResponse from litellm.rust_bridge.trace.storage import ClickHouseStorage, TraceStorageConfig from litellm.tracing import Tenant, TraceReceiver, TracingPayloadTooLargeError @@ -104,25 +104,25 @@ SPAN_DETAIL_RESPONSE: Final = { ( pytest.param( UserAPIKeyAuth(token="admin-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN), - TraceScope(all_teams=1, user_id="", team_ids=()), + AllQueryScope(kind="all"), True, id="admin", ), pytest.param( UserAPIKeyAuth(token="view-key", team_id="team-a", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), - TraceScope(all_teams=1, user_id="", team_ids=()), + AllQueryScope(kind="all"), False, id="view-only-admin", ), pytest.param( TEAM_KEY, - TraceScope(all_teams=0, user_id="user", team_ids=()), + OwnedQueryScope(kind="owned", user_id="user", team_ids=()), True, id="team-key", ), pytest.param( UserAPIKeyAuth(user_id="user", token="hashed-key", user_role=LitellmUserRoles.INTERNAL_USER), - TraceScope(all_teams=0, user_id="user", team_ids=()), + OwnedQueryScope(kind="owned", user_id="user", team_ids=()), True, id="teamless-key", ), @@ -135,7 +135,7 @@ SPAN_DETAIL_RESPONSE: Final = { ), ) def test_trace_read_and_write_permissions( - client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: TraceScope | None, can_write: bool + client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, scope: QueryScope | None, can_write: bool ) -> None: client.app.dependency_overrides[user_api_key_auth] = lambda: auth @@ -144,7 +144,7 @@ def test_trace_read_and_write_permissions( if scope is None: receiver.list_traces.assert_not_awaited() else: - receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, cursor=None) + receiver.list_traces.assert_awaited_once_with(scope=scope, start_ms=1, end_ms=2, q="", cursor=None) write: Final = client.post("/v1/traces", json={}) assert write.status_code == (200 if can_write else 403), write.text @@ -165,6 +165,8 @@ def receiver(client) -> MagicMock: fake = MagicMock() fake.ingest = AsyncMock(return_value=1) fake.list_traces = AsyncMock(return_value={"data": [], "next_cursor": None}) + fake.trace_histogram = AsyncMock(return_value={"buckets": []}) + fake.run_values = AsyncMock(return_value={"values": ["researcher"]}) fake.get_trace = AsyncMock(return_value=None) fake.get_span = AsyncMock(return_value=None) client.app.dependency_overrides[tracing_endpoints.provide_receiver] = lambda: fake @@ -243,23 +245,83 @@ def test_post_too_large_is_413(client, receiver): assert "exceeds" in Status.FromString(response.content).message -def test_list_traces_passes_scope_window_and_cursor(client, receiver): - response = client.get("/v1/traces", params={"start_ms": 1, "end_ms": 2, "cursor": "abc"}) +def test_list_traces_passes_scope_window_query_and_cursor(client, receiver): + response = client.get( + "/v1/traces", params={"start_ms": 1, "end_ms": 2, "q": "agent:research* -status:ok", "cursor": "abc"} + ) assert response.status_code == 200 assert response.json() == {"data": [], "next_cursor": None} receiver.list_traces.assert_awaited_once_with( - scope={"all_teams": 0, "user_id": "user", "team_ids": ()}, + scope={"kind": "owned", "user_id": "user", "team_ids": ()}, start_ms=1, end_ms=2, + q="agent:research* -status:ok", cursor="abc", ) +def test_histogram_passes_scope_window_query_and_buckets(client, receiver): + response = client.get("/v1/traces/histogram", params={"start_ms": 1, "end_ms": 2, "q": "x", "buckets": 12}) + assert response.status_code == 200, response.text + assert response.json() == {"buckets": []} + receiver.trace_histogram.assert_awaited_once_with( + {"kind": "owned", "user_id": "user", "team_ids": ()}, 1, 2, "x", 12 + ) + + +def test_values_pass_scope_window_query_field_and_needle(client, receiver): + response = client.get( + "/v1/traces/values/agent", params={"start_ms": 1, "end_ms": 2, "q": "model:gpt*", "contains": "res"} + ) + assert response.status_code == 200, response.text + assert response.json() == {"values": ["researcher"]} + receiver.run_values.assert_awaited_once_with( + {"kind": "owned", "user_id": "user", "team_ids": ()}, 1, 2, "model:gpt*", "agent", "res", 20 + ) + + +@pytest.mark.parametrize( + "path", + ( + "/v1/traces?q=" + "x" * 1001, + "/v1/traces/histogram?buckets=0", + "/v1/traces/histogram?buckets=241", + "/v1/traces/values/color", + "/v1/traces/values/agent?limit=101", + "/v1/traces/values/agent?contains=" + "x" * 201, + "/v1/traces/values/agent?q=" + "x" * 1001, + ), +) +def test_search_reads_reject_unbounded_parameters(client: TestClient, receiver: MagicMock, path: str) -> None: + assert client.get(path).status_code == 422 + receiver.list_traces.assert_not_awaited() + receiver.trace_histogram.assert_not_awaited() + receiver.run_values.assert_not_awaited() + + +@pytest.mark.parametrize( + "error, status", + ( + (ValueError("invalid trace read parameters"), 400), + (OverflowError("too large"), 413), + (RuntimeError("down"), 503), + ), +) +def test_histogram_and_values_failures_map_to_read_statuses( + client: TestClient, receiver: MagicMock, error: Exception, status: int +) -> None: + receiver.trace_histogram.side_effect = error + receiver.run_values.side_effect = error + assert client.get("/v1/traces/histogram").status_code == status + assert client.get("/v1/traces/values/name").status_code == status + + def test_list_traces_defaults_to_last_24h(client, receiver): client.get("/v1/traces") kwargs = receiver.list_traces.call_args.kwargs assert kwargs["end_ms"] - kwargs["start_ms"] == tracing_endpoints.MS_PER_DAY assert kwargs["cursor"] is None + assert kwargs["q"] == "" def test_get_trace_404_and_200(client, receiver): @@ -268,7 +330,7 @@ def test_get_trace_404_and_200(client, receiver): response = client.get("/v1/traces/t1") assert response.status_code == 200 assert response.json() == TRACE_RESPONSE - receiver.get_trace.assert_awaited_with("t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "", None, None) + receiver.get_trace.assert_awaited_with("t1", {"kind": "owned", "user_id": "user", "team_ids": ()}, "", None, None) def test_get_span_404_and_200(client, receiver): @@ -277,7 +339,7 @@ def test_get_span_404_and_200(client, receiver): response = client.get("/v1/traces/t1/spans/s1") assert response.status_code == 200 assert response.json()["span_id"] == "s1" - receiver.get_span.assert_awaited_with("t1", "s1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "") + receiver.get_span.assert_awaited_with("t1", "s1", {"kind": "owned", "user_id": "user", "team_ids": ()}, "") @pytest.mark.parametrize("suffix,cursor,page_size", [("", None, None), ("&cursor=next&page_size=200", "next", 200)]) @@ -285,7 +347,7 @@ def test_trace_detail_passes_scoped_reference(client, receiver, suffix, cursor, receiver.get_trace.return_value = TRACE_RESPONSE assert client.get(f"/v1/traces/t1?trace_ref=run-one{suffix}").status_code == 200 receiver.get_trace.assert_awaited_with( - "t1", {"all_teams": 0, "user_id": "user", "team_ids": ()}, "run-one", cursor, page_size + "t1", {"kind": "owned", "user_id": "user", "team_ids": ()}, "run-one", cursor, page_size ) @@ -372,6 +434,8 @@ def test_key_without_user_cannot_read_traces(client: TestClient, auth: UserAPIKe client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret" for path in ( "/v1/traces", + "/v1/traces/histogram", + "/v1/traces/values/agent", "/v1/traces/t1", "/v1/traces/t1/spans/s1", "/v1/traces/t1/spans/s1/error", @@ -381,7 +445,14 @@ def test_key_without_user_cannot_read_traces(client: TestClient, auth: UserAPIKe assert response.status_code == 403, response.text query: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) assert query.status_code == 403, query.text - for read in (storage.list_traces, storage.get_trace, storage.get_span, storage.get_span_error): + for read in ( + storage.list_traces, + storage.trace_histogram, + storage.run_values, + storage.get_trace, + storage.get_span, + storage.get_span_error, + ): read.assert_not_called() storage.query_sql.assert_not_called() storage.query_help.assert_not_called() @@ -491,7 +562,7 @@ def test_lifespan_receivers_are_app_local() -> None: assert first_response.status_code == second_response.status_code == 200 assert first_response.json()["span_id"] == "first-span" assert second_response.json()["span_id"] == "second-span" - scope: Final = TraceScope(all_teams=0, user_id=TEAM_KEY.user_id or "", team_ids=()) + scope: Final = OwnedQueryScope(kind="owned", user_id=TEAM_KEY.user_id or "", team_ids=()) assert first_storage.get_span.await_count == 2 first_storage.get_span.assert_awaited_with("t1", "first-span", scope, "first-run") second_storage.get_span.assert_awaited_once_with("t1", "second-span", scope, "second-run") @@ -704,13 +775,6 @@ def test_shared_trace_permissions_reach_read_and_sql_boundaries( response: Final = client.get("/v1/traces/t1/spans/s1?trace_ref=run-one") assert response.status_code == 200, response.text assert response.json()["span_id"] == "s1" - storage.get_span.assert_awaited_once_with( - "t1", "s1", TraceScope(all_teams=expected[0], user_id=expected[1], team_ids=expected[2]), "run-one" - ) - sql_response: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) - assert sql_response.status_code == 200, sql_response.text - assert sql_response.json() == SQL_ENVELOPE - assert client.get("/v1/traces/query/help").json() == QUERY_HELP query_scope: Final = ( {"kind": "all"} if expected[0] @@ -720,6 +784,11 @@ def test_shared_trace_permissions_reach_read_and_sql_boundaries( "team_ids": expected[2], } ) + storage.get_span.assert_awaited_once_with("t1", "s1", query_scope, "run-one") + sql_response: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"}) + assert sql_response.status_code == 200, sql_response.text + assert sql_response.json() == SQL_ENVELOPE + assert client.get("/v1/traces/query/help").json() == QUERY_HELP storage.query_sql.assert_awaited_once_with("SELECT * FROM otel_traces", query_scope, "test-secret") storage.query_help.assert_awaited_once_with(query_scope, "test-secret") assert team_lookup.await_count == ( @@ -741,8 +810,7 @@ def test_trace_storage_permissions_map_owned_rows( scope: ReadScope, expected: tuple[str, tuple[str, ...]], ) -> None: - assert tracing_endpoints._trace_scope(scope) == TraceScope(all_teams=0, user_id=expected[0], team_ids=expected[1]) - assert tracing_endpoints.trace_query_scope(scope) == { + assert tracing_endpoints.read_access(scope) == { "kind": "owned", "user_id": expected[0], "team_ids": expected[1], @@ -781,7 +849,7 @@ async def test_storage_preserves_page_cursor_and_normalizes_native_trace_data( native: Final = _NativeReturningHelp(QUERY_HELP, {**TRACE_RESPONSE, "next_cursor": "more"}) monkeypatch.setattr(loader, "_cached_bridge", native) storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) - scope: Final[TraceScope] = {"all_teams": 0, "user_id": "owner", "team_ids": ()} + scope: Final[QueryScope] = {"kind": "owned", "user_id": "owner", "team_ids": ()} trace: Final = await storage.get_trace("t1", scope, "run", cursor, page_size) assert trace is not None assert trace["next_cursor"] == "more"