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"