refactor(traces): read traces through a five-method storage port

Trace storage now answers five storage-neutral reads: runs, run_counts,
spans, span_text and calls. Query and row types live in
litellm_traces::store with no ClickHouse encodings, and traces-clickhouse
only adapts them, so another backend implements five queries instead of
one SQL file per UI widget

The runs list, histogram, field values, trace graph, span detail, error
paging and cost lookup all go through the port. Paging, the graph budget,
excerpting and the agent output fallback moved out of SQL into
TraceReader. QueryScope is the only access wire, and every SQL file reads
owned_spans, owned_runs and owned_calls CTEs built from one ownership rule
per table that the console row policy also renders

Cursors are tagged with their kind, and run search terms reach storage as
raw text and globs. This also carries the run histogram and values
endpoints, the gateway-traces router and the agent label rollup
This commit is contained in:
Yujong Lee 2026-10-04 15:39:45 -07:00
parent 60e7726116
commit 180b5ee482
95 changed files with 4770 additions and 2423 deletions

View file

@ -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"

View file

@ -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"] }

View file

@ -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<E: std::error::Error> From<ReadError<E>> for ReadFailure {
fn from(error: ReadError<E>) -> 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(),
}
}
}

View file

@ -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<S> {
pub reader: TraceReader,
pub store: S,
}
pub fn router<S>(traces: Arc<Traces<S>>) -> Router
where
S: TraceStore + Send + 'static,
{
Router::new()
.route("/v1/traces", get(runs::list::<S>))
.route("/v1/traces/histogram", get(runs::histogram::<S>))
.route("/v1/traces/values/{field}", get(runs::values::<S>))
.with_state(traces)
}

View file

@ -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<i64>,
end_ms: Option<i64>,
#[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<String>,
}
pub(crate) async fn list<S: TraceStore>(
State(traces): State<Arc<Traces<S>>>,
Extension(access): Extension<QueryScope>,
Query(runs): Query<Runs>,
Query(Page { cursor }): Query<Page>,
) -> Result<Json<TracePage>, 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<S: TraceStore>(
State(traces): State<Arc<Traces<S>>>,
Extension(access): Extension<QueryScope>,
Query(runs): Query<Runs>,
Query(Buckets { buckets }): Query<Buckets>,
) -> Result<Json<TraceHistogram>, 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<S: TraceStore>(
State(traces): State<Arc<Traces<S>>>,
Extension(access): Extension<QueryScope>,
Path(field): Path<RunField>,
Query(runs): Query<Runs>,
Query(values): Query<Values>,
) -> Result<Json<RunValues>, 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)
}

View file

@ -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<Vec<(QueryScope, RunQuery)>>,
counts: Mutex<Vec<RunCountQuery>>,
}
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<Vec<RunRow>, 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<Vec<RunCount>, 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<Vec<SpanRow>, FakeError> {
Ok(Vec::new())
}
async fn span_text(
&self,
_: &QueryScope,
_: &SpanTextQuery,
) -> StoreResult<Option<SpanText>, FakeError> {
Ok(None)
}
async fn calls(&self, _: &QueryScope, _: &CallQuery) -> StoreResult<Vec<CallRow>, FakeError> {
Ok(Vec::new())
}
}
fn access() -> QueryScope {
QueryScope::Owned {
user_id: "user".into(),
team_ids: vec!["team".into()],
}
}
fn app(outcome: Outcome) -> (Router, Arc<Traces<FakeStore>>) {
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
}

View file

@ -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<Error>) -> 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<String>,
limit: u32,
) -> PyResult<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
let field = field
.parse::<RunField>()
.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<String>,
page_size: Option<u32>,
@ -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<Bound<'py, PyAny>> {
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<String>,
) -> PyResult<Bound<'py, PyAny>> {
@ -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")]

View file

@ -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<Self, Error> {
@ -47,13 +44,13 @@ impl SnapshotKey {
pub(crate) fn run(
source: &str,
access: &ReadAccessParams,
access: &QueryScope,
run: (&str, &str, &str, &str),
) -> Result<Self, Error> {
Self::digest(&("run", source, access, run))
}
pub(crate) fn scope(source: &str, access: &ReadAccessParams) -> Result<Self, Error> {
pub(crate) fn scope(source: &str, access: &QueryScope) -> Result<Self, Error> {
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()],
},

View file

@ -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<T: Serialize>(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<T: for<'de> Deserialize<'de>, E>(
cursor: &str,
kind: &'static str,
) -> Result<T, ReadError<E>> {
URL_SAFE
.decode(cursor)
.ok()
.and_then(|json| serde_json::from_slice(&json).ok())
.ok_or(ReadError::InvalidCursor(kind))
}
impl Cursor {
fn kind(&self) -> &'static str {
match self {
Self::Run(_) => "trace",
Self::Span(_) => "span",
Self::Text(_) => "diagnostic",
}
}
pub(super) fn trace_position<E>(cursor: Option<&str>) -> Result<(i64, String), ReadError<E>> {
let Some(cursor) = cursor.filter(|cursor| !cursor.is_empty()) else {
return Ok((0, String::new()));
};
match decode_cursor::<(i64, String), E>(cursor, "trace")? {
(start_ms, trace_ref) if start_ms > 0 && !trace_ref.is_empty() => Ok((start_ms, trace_ref)),
_ => Err(ReadError::InvalidCursor("trace")),
pub(super) fn encode(&self) -> String {
URL_SAFE.encode(serde_json::to_vec(self).unwrap_or_default())
}
fn decode<E>(cursor: &str, kind: &'static str) -> Result<Self, ReadError<E>> {
URL_SAFE
.decode(cursor)
.ok()
.and_then(|json| serde_json::from_slice::<Self>(&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<E>(
cursor: Option<&str>,
) -> Result<Option<ErrorPosition>, ReadError<E>> {
let Some(cursor) = cursor else {
return Ok(None);
};
let position: ErrorPosition = decode_cursor(cursor, "diagnostic")?;
let valid_version = position.version.len() == 64
&& position
.version
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'A'..=b'F').contains(&byte));
if i64::try_from(position.offset).is_err() || !valid_version {
return Err(ReadError::InvalidCursor("diagnostic"));
}
Ok(Some(position))
}
#[derive(Deserialize, Serialize)]
#[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<E>(cursor: Option<&str>) -> Result<Option<RunCursor>, ReadError<E>> {
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<E>(cursor: &str) -> Result<SpanPosition, ReadError<E>> {
match Cursor::decode(cursor, "span")? {
Cursor::Span(position) => Ok(position),
_ => Err(ReadError::InvalidCursor("span")),
}
}
pub(super) fn text_position<E>(
cursor: Option<&str>,
part: SpanPart,
) -> Result<Option<TextPosition>, ReadError<E>> {
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::<std::io::Error>(Some(&run(1_790_742_989_377, "4BAD")))
.unwrap()
.unwrap();
assert_eq!(
trace_position::<std::io::Error>(Some(&cursor)).unwrap(),
(
1_790_742_989_377,
"4bad42b84e9de3ba46fc870185f8f023".to_owned()
)
);
assert_eq!(
trace_position::<std::io::Error>(None).unwrap(),
(0, String::new())
);
assert_eq!(
trace_position::<std::io::Error>(Some("")).unwrap(),
(0, String::new())
(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<std::io::Error>> = 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::<std::io::Error>(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<Option<ErrorPosition>, ReadError<std::io::Error>> =
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::<std::io::Error>(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<Option<ErrorPosition>, ReadError<std::io::Error>> =
error_position(Some(&cursor));
fn span_cursor_round_trips() {
let position = span_position::<std::io::Error>(&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::<std::io::Error>(&cursor),
Err(ReadError::InvalidCursor("span"))
));
}
#[rstest]
fn text_cursor_round_trips() {
let position = text_position::<std::io::Error>(
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::<std::io::Error>(Some(&cursor), SpanPart::Error),
Err(ReadError::InvalidCursor("diagnostic"))
));
}

View file

@ -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};

View file

@ -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<E>(
source: &str,
access: &ReadAccessParams,
row: &ListTracesRow,
access: &QueryScope,
row: &RunRow,
) -> Result<SnapshotKey, ReadError<E>> {
Ok(SnapshotKey::run(
source,
@ -40,7 +40,7 @@ fn cache_key<E>(
)?)
}
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<S: TraceStore>(
reader: &TraceReader,
store: &S,
access: &ReadAccessParams,
runs: &[ListTracesRow],
access: &QueryScope,
runs: &[RunRow],
) -> Result<Vec<TraceSummary>, ReadError<S::Error>> {
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<S: TraceStore>(
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<S: TraceStore>(
async fn resolve_runs<S: TraceStore>(
reader: &TraceReader,
store: &S,
access: &ReadAccessParams,
runs: &[&ListTracesRow],
access: &QueryScope,
runs: &[&RunRow],
) -> Result<Vec<Option<ListedRun>>, ReadError<S::Error>> {
let (Some(start_ms), Some(end_ms)) = (
runs.iter().map(|row| row.start_ms).min(),
@ -103,14 +103,12 @@ async fn resolve_runs<S: TraceStore>(
) 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(&params, 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<S: TraceStore>(
))
.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<S: TraceStore>(
async fn resolve_run<S: TraceStore>(
reader: &TraceReader,
store: &S,
access: &ReadAccessParams,
row: &ListTracesRow,
access: &QueryScope,
row: &RunRow,
) -> Result<Option<ListedRun>, ReadError<S::Error>> {
match reader
.current(store, access, &row.trace_id, &row.trace_ref)

View file

@ -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<R, E, F, Fut>(read: F) -> StoreResult<Vec<R>, E>
where
R: Keyed,
F: Fn(Option<R::Cursor>, u32) -> Fut,
Fut: Future<Output = StoreResult<Vec<R>, 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<Vec<u32>>,
}
impl Table {
async fn page(&self, after: Option<u32>, limit: u32) -> StoreResult<Vec<Number>, 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::<Vec<_>>(),
(1..=total).collect::<Vec<_>>()
);
assert_eq!(table.requests.lock().unwrap().as_slice(), requests);
}
#[rstest]
#[tokio::test]
async fn a_single_oversized_row_fails_the_read() {
let table = Table {
total: 10,
largest_page: 0,
requests: Mutex::new(Vec::new()),
};
let result = read_all(|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]
);
}
}

View file

@ -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<String>,
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<S: TraceStore>(
&self,
store: &S,
access: &ReadAccessParams,
start_ms: i64,
end_ms: i64,
cursor: Option<&str>,
limit: u32,
access: &QueryScope,
filter: &RunFilter,
page: &PageRequest,
) -> Result<TracePage, ReadError<S::Error>> {
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(&params).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<S: TraceStore>(
&self,
store: &S,
access: &QueryScope,
filter: &RunFilter,
buckets: u32,
) -> Result<TraceHistogram, ReadError<S::Error>> {
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<S: TraceStore>(
&self,
store: &S,
access: &QueryScope,
filter: &RunFilter,
field: RunField,
contains: &str,
limit: u32,
) -> Result<RunValues, ReadError<S::Error>> {
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<S: TraceStore>(
&self,
store: &S,
access: &ReadAccessParams,
access: &QueryScope,
trace_id: &str,
trace_ref: &str,
) -> Result<Option<Trace>, ReadError<S::Error>> {
@ -127,7 +190,7 @@ impl TraceReader {
pub async fn get_trace_page<S: TraceStore>(
&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<S: TraceStore>(
&self,
store: &S,
access: &ReadAccessParams,
access: &QueryScope,
trace_id: &str,
trace_ref: &str,
) -> Result<Option<Arc<Snapshot>>, ReadError<S::Error>> {
@ -191,7 +254,7 @@ impl TraceReader {
async fn pinned<S: TraceStore>(
&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(&params, 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<S: TraceStore>(
&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(&params).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<S: TraceStore>(
&self,
store: &S,
access: &QueryScope,
trace_id: &str,
trace_ref: &str,
span_id: &str,
) -> Result<String, ReadError<S::Error>> {
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<S: TraceStore>(
&self,
store: &S,
access: &ReadAccessParams,
access: &QueryScope,
trace_id: &str,
span_id: &str,
trace_ref: &str,
cursor: Option<&str>,
) -> Result<Option<SpanErrorPage>, ReadError<S::Error>> {
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(&params).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<E>(result: Result<Option<SpanText>, StoreError<E>>) -> Result<String, ReadError<E>> {
Ok(result
.map_err(map_store_error)?
.map(|text| text.text)
.unwrap_or_default())
}
fn parse_attributes<E>(json: &str) -> Result<BTreeMap<String, String>, ReadError<E>> {
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<S: TraceStore>(
store: &S,
access: &QueryScope,
selection: SpanSelection,
as_of_ms: u64,
) -> Result<Vec<SpanRow>, StoreError<S::Error>> {
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<S: TraceStore>(
store: &S,
access: &ReadAccessParams,
access: &QueryScope,
trace_id: &str,
trace_ref: &str,
) -> Result<Option<String>, ReadError<S::Error>> {
if !trace_ref.is_empty() {
return Ok(Some(trace_ref.to_owned()));
}
let params = TraceIdentityParams {
access: access.clone(),
trace_id: trace_id.to_owned(),
let query = RunQuery {
selection: RunSelection::TraceId(trace_id.to_owned()),
after: None,
limit: 2,
};
let identities = store.trace_refs(&params).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<E>(
@ -328,12 +488,13 @@ fn page<E>(
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()
}),
}
};

View file

@ -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<Range<i64>> {
pub(super) fn spend_window(rows: &[SpanRow]) -> Option<Range<i64>> {
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<Range<i64>> {
)
}
pub(super) fn spend_within(
spend: &[SpendByResponseIdsRow],
window: Range<i64>,
) -> &[SpendByResponseIdsRow] {
pub(super) fn spend_within(spend: &[CallRow], window: Range<i64>) -> &[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<S: TraceStore>(
store: &S,
access: &ReadAccessParams,
rows: &[TraceSpansRow],
) -> Option<Vec<SpendByResponseIdsRow>> {
access: &QueryScope,
rows: &[SpanRow],
) -> Option<Vec<CallRow>> {
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<S: TraceStore>(
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(&params).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);

View file

@ -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<E> {
/// 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<T, E> = Result<T, StoreError<E>>;
/// 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<Output = Result<Vec<String>, StoreError<Self::Error>>> + Send;
access: &QueryScope,
query: &RunQuery,
) -> impl Future<Output = StoreResult<Vec<RunRow>, 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<Output = Result<Vec<ListTracesRow>, StoreError<Self::Error>>> + Send;
access: &QueryScope,
query: &RunCountQuery,
) -> impl Future<Output = StoreResult<Vec<RunCount>, 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<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
access: &QueryScope,
query: &SpanQuery,
) -> impl Future<Output = StoreResult<Vec<SpanRow>, 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<Output = Result<Vec<TraceSpansRow>, StoreError<Self::Error>>> + Send;
access: &QueryScope,
query: &SpanTextQuery,
) -> impl Future<Output = StoreResult<Option<SpanText>, Self::Error>> + Send;
fn spend(
fn calls(
&self,
params: &SpendByResponseIdsParams,
) -> impl Future<Output = Result<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>>> + Send;
fn span_detail(
&self,
params: &SpanDetailParams,
) -> impl Future<Output = Result<Option<SpanDetailRow>, StoreError<Self::Error>>> + Send;
fn span_error(
&self,
params: &SpanErrorParams,
) -> impl Future<Output = Result<Option<SpanErrorRow>, StoreError<Self::Error>>> + Send;
access: &QueryScope,
query: &CallQuery,
) -> impl Future<Output = StoreResult<Vec<CallRow>, Self::Error>> + Send;
}

View file

@ -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<Operation, Failure>,
trace_refs: Vec<String>,
list_runs: Vec<ListTracesRow>,
trace_spans: HashMap<String, Vec<TraceSpansRow>>,
run_spans: Vec<TraceSpansRow>,
spend: Vec<SpendByResponseIdsRow>,
span_detail: Option<SpanDetailRow>,
span_error: Option<SpanErrorRow>,
list_runs: Vec<RunRow>,
trace_spans: HashMap<String, Vec<SpanRow>>,
run_spans: Vec<SpanRow>,
spend: Vec<CallRow>,
texts: HashMap<(String, SpanPart), String>,
list_runs_too_large_above: Option<u32>,
trace_too_large_refs: HashSet<String>,
spend_fails_above_response_ids: Option<usize>,
}
#[derive(Default)]
struct Calls {
trace_refs: AtomicUsize,
list_runs: AtomicUsize,
trace_spans: AtomicUsize,
run_spans: AtomicUsize,
spend: AtomicUsize,
span_detail: AtomicUsize,
span_error: AtomicUsize,
count_rows: Vec<RunCount>,
count_reads: Vec<RunCountQuery>,
}
#[derive(Default)]
struct FakeStore {
state: Mutex<State>,
calls: Calls,
calls: Mutex<HashMap<Operation, AtomicUsize>>,
}
impl FakeStore {
fn with_spans(trace_ref: &str, spans: Vec<TraceSpansRow>) -> Self {
fn with_spans(trace_ref: &str, spans: Vec<SpanRow>) -> 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<ListTracesRow>) {
fn set_list_runs(&self, rows: Vec<RunRow>) {
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<TraceSpansRow>) {
fn set_run_spans(&self, rows: Vec<SpanRow>) {
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<FakeError>> {
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<R, C: Ord>(rows: &[R], cursor: impl Fn(&R) -> C, after: Option<C>, limit: u32) -> Vec<R>
where
R: Clone,
{
let mut rows: Vec<R> = 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<Vec<String>, StoreError<Self::Error>> {
self.calls.trace_refs.fetch_add(1, Ordering::SeqCst);
let state = self.state.lock().unwrap();
Self::failure(&state, Operation::TraceRefs)?;
Ok(state.trace_refs.clone())
}
async fn list_runs(
&self,
params: &ListTracesParams,
) -> Result<Vec<ListTracesRow>, StoreError<Self::Error>> {
self.calls.list_runs.fetch_add(1, Ordering::SeqCst);
async fn runs(&self, _: &QueryScope, query: &RunQuery) -> StoreResult<Vec<RunRow>, 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<Vec<TraceSpansRow>, StoreError<Self::Error>> {
self.calls.trace_spans.fetch_add(1, Ordering::SeqCst);
tokio::task::yield_now().await;
let state = self.state.lock().unwrap();
Self::failure(&state, Operation::TraceSpans)?;
if state.trace_too_large_refs.contains(&params.trace_ref) {
return Err(StoreError::TooLarge);
}
Ok(state
.trace_spans
.get(&params.trace_ref)
.cloned()
.unwrap_or_default())
_: &QueryScope,
query: &RunCountQuery,
) -> StoreResult<Vec<RunCount>, 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<Vec<TraceSpansRow>, StoreError<Self::Error>> {
self.calls.run_spans.fetch_add(1, Ordering::SeqCst);
_: &QueryScope,
query: &SpanQuery,
) -> StoreResult<Vec<SpanRow>, 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<Vec<SpendByResponseIdsRow>, StoreError<Self::Error>> {
self.calls.spend.fetch_add(1, Ordering::SeqCst);
_: &QueryScope,
query: &SpanTextQuery,
) -> StoreResult<Option<SpanText>, 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<Vec<CallRow>, 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<Option<SpanDetailRow>, StoreError<Self::Error>> {
self.calls.span_detail.fetch_add(1, Ordering::SeqCst);
let state = self.state.lock().unwrap();
Self::failure(&state, Operation::SpanDetail)?;
Ok(state.span_detail.clone())
}
async fn span_error(
&self,
_: &SpanErrorParams,
) -> Result<Option<SpanErrorRow>, StoreError<Self::Error>> {
self.calls.span_error.fetch_add(1, Ordering::SeqCst);
let state = self.state.lock().unwrap();
Self::failure(&state, Operation::SpanError)?;
Ok(state.span_error.clone())
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()
);
}

View file

@ -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<TraceSpansRow> = ids
let rows: Vec<SpanRow> = 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);

View file

@ -0,0 +1,2 @@
ALTER TABLE {database}.agent_traces_by_key
ADD COLUMN IF NOT EXISTS AgentLabels SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT []

View file

@ -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

View file

@ -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}

View file

@ -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

View file

@ -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)}))

View file

@ -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}

View file

@ -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)}

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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}

View file

@ -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
)

View file

@ -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_<base64>` id, which embeds it.
if(startsWith(response_id, 'resp_'),
extract(tryBase64Decode(substring(response_id, 6)), 'response_id:([^;]+)'),
'') AS upstream_response_id
FROM spend_logs FINAL
WHERE start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
AND ({all_teams:UInt8} = 1
OR ({user_id:String} != '' AND user = {user_id:String})
OR has({team_ids:Array(String)}, team_id))
)
WHERE response_id IN {response_ids:Array(String)}
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

View file

@ -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

View file

@ -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}

View file

@ -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

View file

@ -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}

View file

@ -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}

View file

@ -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<String>,
}
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::<Vec<_>>()
.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('\'', "\\'"))
}

View file

@ -1,6 +1,7 @@
use crate::Error;
use litellm_storage_clickhouse::Storage;
use crate::Error;
#[derive(Clone)]
pub struct Config {
storage: Storage,

View file

@ -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<usize, Error> {
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<u64>) {
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\""})),

View file

@ -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;

View file

@ -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<Quer
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use serde_json::json;
use super::*;
#[cfg(feature = "schema")]
#[rstest]
#[case::observed(false)]

View file

@ -1,39 +1,124 @@
use litellm_storage_clickhouse::Query;
use litellm_traces::query::named as contracts;
use litellm_traces::{
QueryScope,
search::{RunFilter, RunSearch},
store::{
CallQuery, CallRow, CountValue, RunCount, RunCountQuery, RunQuery, RunRow, RunSelection,
SpanQuery, SpanRow, SpanSelection, SpanText, SpanTextQuery,
},
};
use serde::{Deserialize, Serialize};
pub use contracts::ReadAccessParams;
use crate::access::{AccessParams, owned};
#[derive(Deserialize, Serialize)]
#[serde(remote = "contracts::ListTracesParams")]
struct ListTracesParamsEncoding {
#[serde(flatten)]
pub access: contracts::ReadAccessParams,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub cursor_ms: i64,
pub cursor_trace_id: String,
#[serde(deserialize_with = "super::number::deserialize")]
pub limit: u32,
/// `LIKE`'s own metacharacters match literally.
fn like_literal(value: &str) -> 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<contracts::ListTracesParams> 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<String>,
filter_fields: Vec<&'static str>,
filter_patterns: Vec<String>,
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<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub error_count: u64,
pub request_ids: Vec<String>,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ListTracesRow(#[serde(with = "ListTracesRowEncoding")] pub contracts::ListTracesRow);
pub(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<String>,
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<contracts::SpanErrorParams> 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<String>,
pub request_ids: Vec<String>,
pub trace_ids: Vec<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
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<String>,
request_ids: Vec<String>,
trace_ids: Vec<String>,
has_cursor: u8,
after_team: String,
after_ms: i64,
after_id: String,
limit: u32,
}
impl From<contracts::SpendByResponseIdsParams> 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<String>,
#[serde(deserialize_with = "super::number::deserialize")]
pub start_ms: i64,
#[serde(deserialize_with = "super::number::deserialize")]
pub end_ms: i64,
}
#[derive(Debug, Deserialize, Serialize)]
pub struct TracePageSpansParams(
#[serde(with = "TracePageSpansParamsEncoding")] pub contracts::TracePageSpansParams,
);
impl From<contracts::TracePageSpansParams> for TracePageSpansParams {
fn from(value: contracts::TracePageSpansParams) -> Self {
Self(value)
}
}
pub struct TracePageSpans;
impl Query for TracePageSpans {
type Params = TracePageSpansParams;
type Row = TraceSpansRow;
const SQL: &'static str = include_str!("../../query/trace_page_spans.sql");
}
pub struct TraceSpans;
impl Query for TraceSpans {
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<T: serde::de::DeserializeOwned + Serialize>(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<T: serde::de::DeserializeOwned + Serialize>(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::<T>(encoded).unwrap()).unwrap()
}
#[rstest]
#[case::unquoted(false)]
#[case::quoted(true)]
fn rows_decode_into_neutral_contracts(#[case] quoted: bool) {
round_trip::<ListTracesRow>(
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "STATUS_CODE_OK", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
quoted,
);
round_trip::<TraceSpansRow>(
json!({"trace_id": "trace", "span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "wrapper_candidate": 1, "agent": "agent", "framework": "claude-agent-sdk", "status": "STATUS_CODE_ERROR", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "call_keys": ["provider_response:request"], "call_evidence": "complete", "tool_call_id": "call", "team_id": "team", "api_key_hash": "key", "user_id": "user"}),
quoted,
);
round_trip::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),
quoted,
);
round_trip::<SpanErrorRow>(
json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}),
quoted,
);
round_trip::<SpendByResponseIdsRow>(
json!({"request_id": "request", "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::<RunRowWire>(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::<SpanRowWire>(span.clone(), quoted), span);
let count = json!({"bucket": 2, "failed": 1, "value": "v", "runs": u64::MAX});
assert_eq!(decoded::<RunCountRow>(count.clone(), quoted), count);
let text = json!({"text": "error", "total_chars": u64::MAX, "version": "version"});
assert_eq!(decoded::<SpanTextRow>(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::<CallRowWire>(call.clone(), quoted), call);
}
#[rstest]
#[case::unquoted(false)]
#[case::quoted(true)]
fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) {
round_trip::<ListTracesParams>(
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}),
quoted,
);
round_trip::<SpanErrorParams>(
json!({"all_teams": 0, "user_id": "", "team_ids": [], "trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
quoted,
);
round_trip::<SpendByResponseIdsParams>(
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "response_ids": ["response"], "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<f64>,
) {
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::<SpendByResponseIdsRow>(json!({
fn spend_rows_reject_invalid_cost(#[case] cost: Value) {
let row = serde_json::from_value::<CallRowWire>(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

View file

@ -57,9 +57,10 @@ pub(super) fn boolean<'de, D: Deserializer<'de>>(deserializer: D) -> Result<bool
#[cfg(test)]
mod tests {
use crate::query::named::SpanErrorRow;
use rstest::rstest;
use crate::query::named::SpanTextRow;
#[rstest]
#[case::flag_zero(serde_json::json!(0), true)]
#[case::flag_one(serde_json::json!("1"), true)]
@ -117,8 +118,8 @@ mod tests {
#[case] value: serde_json::Value,
#[case] expected: Option<u64>,
) {
let row = serde_json::from_value::<SpanErrorRow>(serde_json::json!({
"span_id": "span", "message": "error", "total_chars": value, "version": "hash"
let row = serde_json::from_value::<SpanTextRow>(serde_json::json!({
"text": "error", "total_chars": value, "version": "hash"
}));
match expected {
Some(value) => assert_eq!(row.unwrap().0.total_chars, value),

View file

@ -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::<Vec<_>>()
.join(", ");
let team_clause = if team_ids.is_empty() {
"0".to_owned()
} else {
format!("{team} IN ({teams})")
};
format!("({owner} != '' AND {user_clause}) OR ({team_clause})")
}
}
}
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, Error> {
let mut mac =
Hmac::<Sha256>::new_from_slice(secret.as_bytes()).map_err(|_| Error::MissingSecret)?;
@ -211,35 +181,3 @@ fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, E
mac.update(identity);
Ok(format!("{:x}", mac.finalize().into_bytes()))
}
fn literal(value: &str) -> String {
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
#[case::otel(TraceTable::OtelTraces, "TeamId", "UserId = ''")]
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "UserIds = ['']")]
#[case::spend(TraceTable::SpendLogs, "team_id", "user = ''")]
fn predicates_preserve_scope_and_escape_values(
#[case] table: TraceTable,
#[case] team: &str,
#[case] user: &str,
) {
assert_eq!(predicate(&QueryScope::All, table), "1");
assert_eq!(
predicate(
&QueryScope::Owned {
user_id: String::new(),
team_ids: vec!["team'\\".into()]
},
table
),
format!("('' != '' AND {user}) OR ({team} IN ('team\\'\\\\'))")
);
}
}

View file

@ -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<Q: Query>(&self, params: &Q::Params) -> StoreResult<Vec<Q::Row>, Error> {
fetch::<Q>(&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<Vec<String>, StoreError<Self::Error>> {
fetch::<TraceIdentity>(&self.client, &self.connection, params)
.await
.map(|rows| rows.into_iter().map(|row| row.trace_ref).collect())
.map_err(failed)
async fn runs(&self, access: &QueryScope, query: &RunQuery) -> StoreResult<Vec<RunRow>, Error> {
let rows = self.fetch::<Runs>(&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<Vec<contracts::ListTracesRow>, StoreError<Self::Error>> {
let storage_params = ListTracesParams::from(params.clone());
match fetch::<RunCandidates>(&self.client, &self.connection, &storage_params).await {
Ok(rows) => Ok(rows.into_iter().map(|row| row.0).collect()),
Err(StorageError::ResponseTooLarge) => Err(StoreError::TooLarge),
Err(error) => Err(failed(error)),
}
access: &QueryScope,
query: &RunCountQuery,
) -> StoreResult<Vec<RunCount>, Error> {
let rows = self
.fetch::<RunCounts>(&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<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
crate::span_batches::read_spans(&self.client, &self.connection, params.clone(), snapshot_ms)
.await
access: &QueryScope,
query: &SpanQuery,
) -> StoreResult<Vec<SpanRow>, Error> {
let rows = match SpansParams::new(access, query) {
SpansParams::Trace(params) => self.fetch::<TraceSpans>(&params).await?,
SpansParams::Runs(params) => self.fetch::<RunSpans>(&params).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<Vec<contracts::TraceSpansRow>, StoreError<Self::Error>> {
crate::span_batches::read_list_spans(
&self.client,
&self.connection,
TracePageSpansParams::from(params.clone()),
snapshot_ms,
)
.await
access: &QueryScope,
query: &SpanTextQuery,
) -> StoreResult<Option<SpanText>, Error> {
let rows = self
.fetch::<SpanTexts>(&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<Vec<contracts::SpendByResponseIdsRow>, StoreError<Self::Error>> {
crate::span_batches::read_spend(
&self.client,
&self.connection,
SpendByResponseIdsParams::from(params.clone()),
)
.await
}
async fn span_detail(
&self,
params: &contracts::SpanDetailParams,
) -> Result<Option<contracts::SpanDetailRow>, StoreError<Self::Error>> {
match fetch::<SpanDetailQuery>(&self.client, &self.connection, params).await {
Ok(rows) => Ok(rows.into_iter().next()),
Err(error) => Err(failed(error)),
}
}
async fn span_error(
&self,
params: &contracts::SpanErrorParams,
) -> Result<Option<contracts::SpanErrorRow>, StoreError<Self::Error>> {
let storage_params = SpanErrorParams::from(params.clone());
match fetch::<SpanError>(&self.client, &self.connection, &storage_params).await {
Ok(rows) => Ok(rows.into_iter().next().map(|row| row.0)),
Err(error) => Err(failed(error)),
}
access: &QueryScope,
query: &CallQuery,
) -> StoreResult<Vec<CallRow>, Error> {
let rows = self
.fetch::<Calls>(&CallsParams::new(access, query))
.await?;
Ok(rows.into_iter().map(|row| row.0).collect())
}
}
fn failed(error: StorageError) -> StoreError<Error> {
StoreError::Failed(Error::Storage(error))
}

View file

@ -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};

View file

@ -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<Error>> {
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<Error>> {
let bytes =
serde_json::to_vec(row).map_err(|_| StoreError::Failed(Error::InvalidResponse))?;
self.reserve(bytes.len())
}
}
/// One keyset position in a paged query: the SQL reads the cursor fields of `Self` plus the
/// `page_size` that [`Batch`] adds.
trait Keyset: Serialize + Sized + Send + Sync {
type Row: Serialize + DeserializeOwned + Send;
const SQL: &'static str;
fn after(self, last: &Self::Row) -> Self;
}
#[derive(Serialize)]
struct Batch<K> {
#[serde(flatten)]
keyset: K,
page_size: u32,
}
trait PageSource<K: Keyset> {
fn page(
&self,
batch: &Batch<K>,
) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send;
}
struct Paged<K>(PhantomData<K>);
impl<K: Keyset> Query for Paged<K> {
type Params = Batch<K>;
type Row = K::Row;
const SQL: &'static str = K::SQL;
}
/// Reads every row after `keyset`. A page ClickHouse rejects as too large is retried at half the
/// size, and the smaller page is kept for the rest of the read because row sizes within one graph
/// rarely shrink again. Halving a one-row page means a single row exceeds the response limit.
async fn read_all<K: Keyset, S: PageSource<K>>(
source: &S,
keyset: K,
) -> Result<Vec<K::Row>, StoreError<Error>> {
let mut batch = Batch {
keyset,
page_size: PAGE_SIZE,
};
let mut rows = Vec::new();
let mut budget = ReadBudget::default();
loop {
let page = match source.page(&batch).await {
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) if batch.page_size > 1 => {
batch.page_size /= 2;
continue;
}
Err(litellm_storage_clickhouse::Error::ResponseTooLarge) => {
return Err(StoreError::TooLarge);
}
result => result.map_err(|error| StoreError::Failed(Error::Storage(error)))?,
};
let complete = page.len() < batch.page_size as usize;
for row in &page {
budget.record(row)?;
}
if let Some(last) = page.last() {
batch.keyset = batch.keyset.after(last);
}
rows.extend(page);
if complete {
return Ok(rows);
}
}
}
struct ClickHouse<'a> {
client: &'a Client,
connection: &'a Connection,
}
impl<K: Keyset> PageSource<K> for ClickHouse<'_> {
fn page(
&self,
batch: &Batch<K>,
) -> impl Future<Output = Result<Vec<K::Row>, litellm_storage_clickhouse::Error>> + Send {
fetch::<Paged<K>>(self.client, self.connection, batch)
}
}
async fn read_paged<K: Keyset>(
client: &Client,
connection: &Connection,
keyset: K,
) -> Result<Vec<K::Row>, StoreError<Error>> {
let source = ClickHouse { client, connection };
read_all(&source, keyset).await
}
fn by_start(mut rows: Vec<contracts::TraceSpansRow>) -> Vec<contracts::TraceSpansRow> {
rows.sort_by_key(|row| row.start_ns);
rows
}
#[derive(Serialize)]
struct SpanKeyset {
#[serde(flatten)]
trace: contracts::TraceSpansParams,
after_span_id: String,
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<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
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<Vec<contracts::TraceSpansRow>, StoreError<Error>> {
let keyset = ListSpanKeyset {
runs,
after_team: String::new(),
after_key: String::new(),
after_trace: String::new(),
after_span: String::new(),
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<Vec<contracts::SpendByResponseIdsRow>, StoreError<Error>> {
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<Vec<u32>>,
}
impl PageSource<Numbers> for Table {
async fn page(
&self,
batch: &Batch<Numbers>,
) -> Result<Vec<u32>, litellm_storage_clickhouse::Error> {
self.requests.lock().unwrap().push(batch.page_size);
if batch.page_size > self.largest_page {
return Err(litellm_storage_clickhouse::Error::ResponseTooLarge);
}
let end = (batch.keyset.after + batch.page_size).min(self.total);
Ok((batch.keyset.after + 1..=end).collect())
}
}
#[rstest]
#[case::fits(1000, PAGE_SIZE, &[256, 256, 256, 256])]
#[case::uniform_large_rows(1000, 100, &[256, 128, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64])]
#[tokio::test]
async fn a_rejected_page_size_is_not_retried(
#[case] total: u32,
#[case] largest_page: u32,
#[case] requests: &[u32],
) {
let table = Table {
total,
largest_page,
requests: Mutex::new(Vec::new()),
};
let rows = read_all(&table, Numbers { after: 0 }).await.unwrap();
assert_eq!(rows, (1..=total).collect::<Vec<_>>());
assert_eq!(table.requests.lock().unwrap().as_slice(), requests);
}
#[rstest]
#[tokio::test]
async fn a_single_oversized_row_fails_the_read() {
let table = Table {
total: 10,
largest_page: 0,
requests: Mutex::new(Vec::new()),
};
let result = read_all(&table, Numbers { after: 0 }).await;
assert!(matches!(result, Err(StoreError::TooLarge)), "{result:?}");
assert_eq!(
table.requests.lock().unwrap().as_slice(),
&[256, 128, 64, 32, 16, 8, 4, 2, 1]
);
}
}

View file

@ -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<String, Parameter>,
) -> Result<String, Error> {
match query {
ReadQuery::ListTraces => named_json::<ListTraces>(client, connection, parameters).await,
ReadQuery::TraceIdentity => {
named_json::<TraceIdentity>(client, connection, parameters).await
}
ReadQuery::TraceSpans => named_json::<TraceSpans>(client, connection, parameters).await,
ReadQuery::TracePageSpans => {
named_json::<TracePageSpans>(client, connection, parameters).await
}
ReadQuery::SpanDetail => named_json::<SpanDetail>(client, connection, parameters).await,
ReadQuery::SpanError => named_json::<SpanError>(client, connection, parameters).await,
ReadQuery::SpendByResponseIds => {
named_json::<SpendByResponseIds>(client, connection, parameters).await
}
ReadQuery::Availability => {
named_json::<LensAvailability>(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<String, Parameter> = 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, &parameters).await,
execute_named_read(&client, &connection, ReadQuery::Content, &parameters).await,
Err(Error::InvalidParameters)
));
}

View file

@ -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)]

View file

@ -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};

View file

@ -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<i64>,
after: Option<RunCursor>,
limit: u32,
) -> TestResult<serde_json::Value> {
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::<litellm_traces_clickhouse::query::named::SpanDetail>(
&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::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", historical_rows).await?;
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
let parameters = BTreeMap::from([
("all_teams".into(), Parameter::Integer(0)),
("user_id".into(), Parameter::Text("owner".into())),
("team_ids".into(), Parameter::Strings(vec![])),
(
"start_ms".into(),
Parameter::Integer(timestamp / 1_000_000 - 1000),
),
(
"end_ms".into(),
Parameter::Integer(timestamp / 1_000_000 + 1000),
),
("cursor_ms".into(), Parameter::Integer(0)),
("cursor_trace_id".into(), Parameter::Text(String::new())),
("limit".into(), Parameter::Integer(1)),
]);
let first: serde_json::Value = serde_json::from_str(
&execute_named_read(
&database.client,
&connection,
ReadQuery::ListTraces,
&parameters,
)
.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::<BTreeMap<_, _>>();
assert_eq!(counts["shared"], Some(3));
assert_eq!(counts["second"], Some(1));
for page in [&first, &second] {
assert!(
page["statistics"]["rows_read"]
.as_u64()
.ok_or("missing read statistics")?
< 5000
);
}
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,
&parameters,
)
.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::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
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,
&parameters,
)
.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, &parameters)
.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, &parameters).await?;
assert_eq!(
serde_json::from_str::<serde_json::Value>(&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::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
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,
&parameters,
)
.await?;
let diagnostic =
execute_named_read(&database.client, &reader, ReadQuery::SpanError, &parameters).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<ClickHouseDatabase>,
#[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::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
insert_rows(&database, "spend_logs", rows).await?;
let reader = Connection::reader(&database.url, "trace_test")?;
let params =
SpendByResponseIdsParams::from(litellm_traces::query::named::SpendByResponseIdsParams {
access: serde_json::from_value::<ReadAccessParams>(serde_json::json!({
"all_teams": 0, "user_id": user, "team_ids": teams,
"api_key_hash": legacy_key.unwrap_or_default(),
}))?,
response_ids: vec!["shared-response".into()],
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::<SpendByResponseIds>(&database.client, &reader, &params)
.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, &params)
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<ClickHouseDatabase>,
#[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::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", rows).await?;
let reader = Connection::reader(&database.url, "trace_test")?;
let details = litellm_storage_clickhouse::fetch::<SpanDetail>(
&database.client,
&reader,
&SpanDetailParams {
access: ReadAccessParams {
all_teams: 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(())
}

View file

@ -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<u64> {
Ok(spans.first().ok_or("missing fixture root")?.start_ns / 1_000_000_000)
}
#[fixture]
fn admin_access() -> TestResult<contracts::ReadAccessParams> {
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<Vec<SpanRow>, 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<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
) -> TestResult {
let fixture = seeded_database?;
let reader = fixture
.readers
.connection(&fixture.database.client, &QueryScope::All, "fixture-secret")
.await?;
let params = ListTracesParams::from(contracts::ListTracesParams {
access: admin_access?,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 10,
});
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, &params).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::<Vec<_>>(),
["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::<TraceSpans>(&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::<Vec<_>>(),
["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<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
) -> TestResult {
let fixture = seeded_database?;
let reader = fixture
.readers
.connection(&fixture.database.client, &QueryScope::All, "fixture-secret")
.await?;
let params = ListTracesParams::from(contracts::ListTracesParams {
access: admin_access?,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 1,
});
let first = fetch::<ListTraces>(&fixture.database.client, &reader, &params).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::<ListTraces>(&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<SeededDatabase>,
admin_access: TestResult<contracts::ReadAccessParams>,
#[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::<TraceSpans>(&fixture.database.client, &reader, &params).await?;
assert_eq!(stored.len(), decoded.len());
let list_params = ListTracesParams::from(contracts::ListTracesParams {
access: params.access,
start_ms: 0,
end_ms: i64::MAX / 1_000_000,
cursor_ms: 0,
cursor_trace_id: String::new(),
limit: 10,
});
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, &list_params).await?;
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::<Vec<_>>();
assert_eq!(roots.len(), 1);
assert_eq!(
traces[0].0.status,
traces[0].status,
serde_json::from_value::<litellm_traces::SpanStatus>(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

View file

@ -1,8 +0,0 @@
{
"all_teams": 1,
"user_id": "",
"team_ids": [
"team-a",
"team-b"
]
}

View file

@ -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::<usize>()?;
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 {

View file

@ -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<String>, 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<Step>,
}
fn runs() -> Vec<Run> {
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<String, Value> {
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<BTreeMap<String, Value>> {
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<ClickHouseTraces> {
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<SeededDatabase>,
) -> 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<SeededDatabase>,
) -> 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<String> {
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<SeededDatabase>,
#[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<HistogramBucket> = 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<SeededDatabase>,
#[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<SeededDatabase>,
) -> 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<_>>(),
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<SeededDatabase>,
#[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<SeededDatabase>,
) -> 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(())
}

View file

@ -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;

View file

@ -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;

View file

@ -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",

View file

@ -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);

View file

@ -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

View file

@ -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";

View file

@ -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<Vec<Message>> {
#[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"))]

View file

@ -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<CallKey> {
pub(crate) fn row_keys(row: &crate::store::SpanRow) -> BTreeSet<CallKey> {
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() {

View file

@ -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)]

View file

@ -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,

View file

@ -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<String>,
}
#[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<String>,
#[serde(default)]
pub frameworks: Vec<String>,
pub llm_calls: u64,
pub tool_calls: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub models: Vec<String>,
pub error_count: u64,
pub request_ids: Vec<String>,
}
#[derive(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<crate::CallKey>,
#[serde(
default,
deserialize_with = "crate::wire::evidence",
serialize_with = "crate::wire::serialize_evidence"
)]
pub call_evidence: Option<crate::CallEvidenceKind>,
#[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<String>,
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<String, String>,
}
#[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<String>,
pub request_ids: Vec<String>,
pub trace_ids: Vec<String>,
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<f64>,
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,
}

View file

@ -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"))]

View file

@ -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<usize>>,
}
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<TraceSpansRow> {
fn unordered_rows() -> Vec<SpanRow> {
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<TraceSpansRow>) {
fn traversal_follows_links_instead_of_export_order(unordered_rows: Vec<SpanRow>) {
let graph = Graph::new(&unordered_rows);
assert_eq!(graph.ancestors(0), [2, 3]);
let descendants: std::collections::BTreeSet<&str> = graph

View file

@ -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]
}

View file

@ -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> {

View file

@ -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<Item = &'a str>) -> Vec<String> {
pub fn resolve_trace(
trace_id: &str,
trace_ref: &str,
rows: &[TraceSpansRow],
rows: &[SpanRow],
spend: &[SpendRow],
) -> Option<Trace> {
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(),

View file

@ -46,15 +46,14 @@ fn emitted<T: JsonSchema>() -> Schema {
pub fn schemas() -> BTreeMap<&'static str, Schema> {
BTreeMap::from([
(
"TraceScope",
received::<crate::query::named::ReadAccessParams>(),
),
("QueryScope", received::<crate::QueryScope>()),
("Tenant", received::<crate::Tenant>()),
("TracePage", emitted::<crate::TracePage>()),
("Trace", emitted::<crate::Trace>()),
("SpanDetail", emitted::<crate::SpanDetail>()),
("SpanErrorPage", emitted::<crate::SpanErrorPage>()),
("TraceHistogram", emitted::<crate::search::TraceHistogram>()),
("RunValues", emitted::<crate::search::RunValues>()),
("RunField", received::<crate::search::RunField>()),
])
}

View file

@ -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<String>,
pub filters: Vec<FieldFilter>,
}
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<Item = &str> {
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::<RunField>().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<HistogramBucket>,
}
#[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<AgentRuns>,
}
#[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<String>,
}
/// 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<AgentRuns> = 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(),
}
}

View file

@ -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<RunCursor>,
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<String>,
#[serde(default)]
pub frameworks: Vec<String>,
pub llm_calls: u64,
pub tool_calls: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub models: Vec<String>,
pub error_count: u64,
}
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<u32>,
pub failed: bool,
pub value: Option<CountValue>,
}
/// 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<u32>,
}
#[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<String>,
window: Range<i64>,
},
}
#[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<SpanCursor>,
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<crate::CallKey>,
#[serde(
default,
deserialize_with = "crate::wire::evidence",
serialize_with = "crate::wire::serialize_evidence"
)]
pub call_evidence: Option<crate::CallEvidenceKind>,
#[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<u64>,
}
#[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<i64>,
/// Also matches the upstream id a managed `resp_` id wraps.
pub response_ids: Vec<String>,
/// Matches the gateway call id, or the request id when a call has none.
pub request_ids: Vec<String>,
pub trace_ids: Vec<String>,
pub after: Option<CallCursor>,
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<f64>,
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)
}
}

View file

@ -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<SpendByResponseIdsRow>) {
fn captured_spend_rows(spend_logs: &str) -> (FixtureCapture, Vec<CallRow>) {
let records: Vec<CapturedSpend> = spend_logs
.lines()
.filter(|line| !line.trim().is_empty())
@ -112,7 +114,7 @@ fn captured_spend_rows(spend_logs: &str) -> (FixtureCapture, Vec<SpendByResponse
.into_iter()
.map(|record| {
let upstream_response_id = upstream_response_id(&record.response_id);
SpendByResponseIdsRow {
CallRow {
request_id: record.request_id,
litellm_call_id: record.litellm_call_id,
response_id: record.response_id,
@ -148,7 +150,7 @@ fn status_message(span: &DecodedSpan) -> &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<TraceSpansRow> {
fn trace_rows(otlp: &[u8]) -> Vec<SpanRow> {
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<TraceSpansRow> {
.collect()
}
fn fixture(
spend_log_path: &Path,
) -> (
CaptureData,
FixtureCapture,
Vec<TraceSpansRow>,
Vec<SpendByResponseIdsRow>,
) {
fn fixture(spend_log_path: &Path) -> (CaptureData, FixtureCapture, Vec<SpanRow>, Vec<CallRow>) {
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<String, Option<f64>>
.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(),

View file

@ -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;

View file

@ -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;

View file

@ -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<T: DeserializeOwned + Serialize>(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::<ReadAccessParams>(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::<ListTracesParams>(request(
json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100}),
));
round_trip::<TraceIdentityParams>(request(json!({"trace_id": "trace"})));
round_trip::<TraceSpansParams>(request(json!({"trace_id": "trace", "trace_ref": "ref"})));
round_trip::<SpanDetailParams>(request(
json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span"}),
));
round_trip::<SpanErrorParams>(request(
json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
));
round_trip::<SpendByResponseIdsParams>(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::<ListTracesRow>(
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "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::<TraceSpansRow>(
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::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),
);
round_trip::<SpanErrorRow>(
json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}),
);
round_trip::<SpendByResponseIdsRow>(
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}),
);
}

View file

@ -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<TraceSpansRow> {
fn deep_agent(researchers: usize) -> Vec<SpanRow> {
let mut rows = vec![
at(
row(
@ -147,7 +146,7 @@ fn deep_agent(researchers: usize) -> Vec<TraceSpansRow> {
rows
}
fn agents(rows: &[TraceSpansRow]) -> Vec<AgentNode> {
fn agents(rows: &[SpanRow]) -> Vec<AgentNode> {
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<f64>,
) {
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<f64>,
) {
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<f64>,
) {
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::<TraceSpansRow>(encoded).is_err());
assert!(serde_json::from_value::<SpanRow>(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<f64>,
) {
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)
}];

View file

@ -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<FieldFilter>,
) {
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());
}

View file

@ -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

View file

@ -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(

View file

@ -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)

View file

@ -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

View file

@ -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]: ...

View file

@ -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
)

View file

@ -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)

View file

@ -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)

View file

@ -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"
}

View file

@ -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"
}

View file

@ -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"
}

View file

@ -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"
}

View file

@ -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)

View file

@ -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)

View file

@ -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"