From dd86ca517584a6662859436aa3ecb0b13c276a0f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 2 Oct 2026 18:01:27 -0700 Subject: [PATCH] refactor(traces): type the ClickHouse query help response (#44285) * refactor(traces): type the ClickHouse query help response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(traces): include agent names and frameworks in named contract round trips Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(traces): cover native query help validation in the storage adapter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Yujong Lee Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/storage-clickhouse/src/lib.rs | 2 +- .../crates/storage-clickhouse/src/read.rs | 22 +- .../crates/traces-clickhouse/src/lib.rs | 2 +- .../crates/traces-clickhouse/src/query.rs | 292 +++++++++++++----- .../traces-clickhouse/src/query/guide.rs | 5 +- .../traces-clickhouse/src/query_access.rs | 40 ++- .../crates/traces-clickhouse/src/table.rs | 10 +- .../templates/query_help.jinja | 20 +- .../traces-clickhouse/tests/migrations.rs | 8 +- .../traces-clickhouse/tests/query_access.rs | 2 +- .../crates/traces/tests/query/named.rs | 4 +- litellm/rust_bridge/_native.pyi | 2 +- litellm/rust_bridge/trace_query_responses.py | 35 ++- litellm/rust_bridge/traces.py | 13 +- tests/unit/proxy/test_tracing_endpoints.py | 69 ++++- 15 files changed, 392 insertions(+), 134 deletions(-) diff --git a/litellm-rust/crates/storage-clickhouse/src/lib.rs b/litellm-rust/crates/storage-clickhouse/src/lib.rs index 7cba0a11653..7ab2aa9bc0a 100644 --- a/litellm-rust/crates/storage-clickhouse/src/lib.rs +++ b/litellm-rust/crates/storage-clickhouse/src/lib.rs @@ -4,7 +4,7 @@ mod read; pub use error::Error; pub use insert::{insert_compressed_rows, insert_encoded_rows}; -pub use read::{Parameter, Query, execute_read, fetch, fetch_json}; +pub use read::{Parameter, Query, READ_LIMITS, ReadLimits, execute_read, fetch, fetch_json}; use url::Url; #[derive(Clone)] diff --git a/litellm-rust/crates/storage-clickhouse/src/read.rs b/litellm-rust/crates/storage-clickhouse/src/read.rs index 0dae94109bc..59d5b0de558 100644 --- a/litellm-rust/crates/storage-clickhouse/src/read.rs +++ b/litellm-rust/crates/storage-clickhouse/src/read.rs @@ -5,7 +5,18 @@ use serde::{Deserialize, Serialize, de::DeserializeOwned}; use crate::{Connection, Error}; -const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024; +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct ReadLimits { + pub result_rows: u64, + pub response_bytes: usize, + pub execution_seconds: u64, +} + +pub const READ_LIMITS: ReadLimits = ReadLimits { + result_rows: 1000, + response_bytes: 4 * 1024 * 1024, + execution_seconds: 10, +}; #[derive(Debug, Deserialize, Serialize)] #[serde(untagged)] @@ -78,9 +89,12 @@ pub async fn execute_read( .clear() .extend_pairs(existing_pairs) .append_pair("readonly", "1") - .append_pair("max_result_rows", "1000") + .append_pair("max_result_rows", &READ_LIMITS.result_rows.to_string()) .append_pair("result_overflow_mode", "throw") - .append_pair("max_execution_time", "10") + .append_pair( + "max_execution_time", + &READ_LIMITS.execution_seconds.to_string(), + ) .append_pair("wait_end_of_query", "1") .append_pair("default_format", "JSON"); @@ -101,7 +115,7 @@ pub async fn execute_read( let mut body = Vec::new(); while let Some(chunk) = response.chunk().await.map_err(|_| Error::Transport)? { - if body.len() + chunk.len() > MAX_RESPONSE_BYTES { + if body.len() + chunk.len() > READ_LIMITS.response_bytes { return Err(Error::ResponseTooLarge); } body.extend_from_slice(&chunk); diff --git a/litellm-rust/crates/traces-clickhouse/src/lib.rs b/litellm-rust/crates/traces-clickhouse/src/lib.rs index cd5877e4993..ba2f776b51a 100644 --- a/litellm-rust/crates/traces-clickhouse/src/lib.rs +++ b/litellm-rust/crates/traces-clickhouse/src/lib.rs @@ -12,7 +12,7 @@ pub use error::Error; pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows}; pub use litellm_storage_clickhouse::{Connection, Parameter}; pub use litellm_traces::{QueryScope, ReadQuery}; -pub use query::{execute_read, query_help, query_sql}; +pub use query::{QueryHelp, execute_read, query_help, query_sql}; pub use query_access::QueryReaders; pub use schema::{ NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements, diff --git a/litellm-rust/crates/traces-clickhouse/src/query.rs b/litellm-rust/crates/traces-clickhouse/src/query.rs index c862cebf556..459f0f95fe7 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query.rs @@ -6,11 +6,14 @@ use futures_util::{ stream::{self, TryStreamExt}, }; use litellm_http::Client; -use serde::{Deserialize, Serialize}; -use serde_json::{Value, json}; +use serde::{Deserialize, Serialize, Serializer}; +use serde_json::Value; use strum::IntoEnumIterator; -use super::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, Parameter}; +use super::{ + Connection, Error, NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, Parameter, + query_access::READER_LIMITS, +}; mod guide; pub mod lens; @@ -48,10 +51,42 @@ enum PathPart { Index(usize), } +#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize, strum::Display)] +#[serde(rename_all = "lowercase")] +#[strum(serialize_all = "lowercase")] +enum JsonKind { + Array, + Boolean, + Integer, + Null, + Number, + Object, + String, +} + +impl JsonKind { + fn of(value: &Value) -> Self { + match value { + Value::Null => Self::Null, + Value::Bool(_) => Self::Boolean, + Value::Number(number) if number.is_i64() || number.is_u64() => Self::Integer, + Value::Number(_) => Self::Number, + Value::String(_) => Self::String, + Value::Array(_) => Self::Array, + Value::Object(_) => Self::Object, + } + } +} + +#[derive(Clone, Copy, Debug, Serialize, strum::Display)] +enum MapValueType { + String, +} + #[derive(Serialize)] struct MetadataField { path: Vec, - types: BTreeSet<&'static str>, + types: BTreeSet, expression: String, } @@ -66,42 +101,150 @@ struct ColumnSchema { #[derive(Serialize)] struct TableSchema { - name: &'static str, + name: TraceTable, columns: Vec, } +trait Unobserved { + fn unobserved() -> Self; +} + +enum Discovery { + Observed(T), + Unavailable(String), +} + +impl Serialize for Discovery { + fn serialize(&self, serializer: S) -> Result { + #[derive(Serialize)] + struct Unavailable<'a, T> { + #[serde(flatten)] + sample: T, + error: &'a str, + } + match self { + Self::Observed(sample) => sample.serialize(serializer), + Self::Unavailable(error) => Unavailable { + sample: T::unobserved(), + error, + } + .serialize(serializer), + } + } +} + #[derive(Serialize)] -struct MetadataCatalog { - table: &'static str, - column: &'static str, +struct MetadataSample { fields: Vec, sampled_rows: usize, invalid_json_rows: usize, truncated: bool, +} + +impl Unobserved for MetadataSample { + fn unobserved() -> Self { + Self { + fields: Vec::new(), + sampled_rows: 0, + invalid_json_rows: 0, + truncated: true, + } + } +} + +#[derive(Serialize)] +struct MetadataCatalog { + table: TraceTable, + column: &'static str, + #[serde(flatten)] + discovery: Discovery, sample_sql: &'static str, scope: &'static str, - #[serde(skip_serializing_if = "Option::is_none")] - error: Option, } #[derive(Serialize)] struct AttributeField { key: String, #[serde(rename = "type")] - kind: &'static str, + kind: MapValueType, expression: String, } #[derive(Serialize)] -struct AttributeCatalog { - table: &'static str, - column: &'static str, +struct AttributeSample { fields: Vec, truncated: bool, +} + +impl Unobserved for AttributeSample { + fn unobserved() -> Self { + Self { + fields: Vec::new(), + truncated: true, + } + } +} + +#[derive(Serialize)] +struct AttributeCatalog { + table: TraceTable, + column: &'static str, + #[serde(flatten)] + discovery: Discovery, discovery_sql: String, scope: &'static str, - #[serde(skip_serializing_if = "Option::is_none")] - error: Option, +} + +#[derive(Serialize)] +struct NormalizedField { + table: TraceTable, + name: &'static str, + column: &'static str, + #[serde(rename = "type")] + kind: &'static str, + meaning: &'static str, +} + +impl From<&NormalizedFieldDefinition> for NormalizedField { + fn from(field: &NormalizedFieldDefinition) -> Self { + Self { + table: TraceTable::OtelTraces, + name: field.name, + column: field.clickhouse_column, + kind: field.clickhouse_type, + meaning: field.meaning, + } + } +} + +#[derive(Serialize)] +struct Relationship { + left: &'static str, + right: &'static str, + additional_predicates: &'static str, + meaning: &'static str, +} + +const RELATIONSHIPS: [Relationship; 1] = [Relationship { + left: "otel_traces.LiteLLMRequestId", + right: "spend_logs.response_id", + additional_predicates: "otel_traces.TeamId = spend_logs.team_id AND (otel_traces.TeamId != '' OR (otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' AND otel_traces.ApiKeyHash = spend_logs.api_key))", + meaning: "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows", +}]; + +#[derive(Serialize)] +pub struct QueryHelp { + dialect: &'static str, + access: &'static str, + response: &'static str, + tables: Vec, + normalized_fields: Vec, + metadata: MetadataCatalog, + attributes: Vec, + relationships: &'static [Relationship], + examples: [guide::Example; 5], + gotchas: [String; 11], + guide: String, } pub async fn execute_read( @@ -153,22 +296,16 @@ fn metadata_expression(path: &[PathPart]) -> String { fn discover( value: &Value, path: Vec, - fields: &mut BTreeMap, BTreeSet<&'static str>>, + fields: &mut BTreeMap, BTreeSet>, ) -> bool { if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) { return true; } if !path.is_empty() { - let kind = match value { - Value::Null => "null", - Value::Bool(_) => "boolean", - Value::Number(number) if number.is_i64() || number.is_u64() => "integer", - Value::Number(_) => "number", - Value::String(_) => "string", - Value::Array(_) => "array", - Value::Object(_) => "object", - }; - fields.entry(path.clone()).or_default().insert(kind); + fields + .entry(path.clone()) + .or_default() + .insert(JsonKind::of(value)); } match value { Value::Object(object) => object.iter().fold(false, |limited, (key, value)| { @@ -194,7 +331,7 @@ fn discover( } } -fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog { +fn metadata_sample(sample: &[MetadataRow]) -> MetadataSample { let (fields, limited, invalid_rows) = sample.iter().take(SAMPLE_ROWS).fold( (BTreeMap::new(), sample.len() > SAMPLE_ROWS, 0), |(fields, limited, invalid_rows), row| match serde_json::from_str::(&row.metadata) { @@ -214,24 +351,19 @@ fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog { types, }) .collect(); - MetadataCatalog { - table: "spend_logs", - column: "metadata", + MetadataSample { fields, sampled_rows: sample.len().min(SAMPLE_ROWS), invalid_json_rows: invalid_rows, truncated: limited, - sample_sql: METADATA_SQL, - error: None, - scope: METADATA_SCOPE, } } -pub async fn query_help(client: &Client, connection: &Connection) -> Result { +pub async fn query_help(client: &Client, connection: &Connection) -> Result { let tables = stream::iter(TraceTable::iter()) .then(|table| async move { Ok::<_, Error>(TableSchema { - name: table.into(), + name: table, columns: rows::( client, connection, @@ -242,13 +374,15 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result>() .await?; - let metadata = match rows::(client, connection, METADATA_SQL).await { - Ok(sample) => metadata_catalog(&sample), - Err(error) => MetadataCatalog { - error: Some(error.to_string()), - truncated: true, - ..metadata_catalog(&[]) + let metadata = MetadataCatalog { + table: TraceTable::SpendLogs, + column: "metadata", + discovery: match rows::(client, connection, METADATA_SQL).await { + Ok(sample) => Discovery::Observed(metadata_sample(&sample)), + Err(error) => Discovery::Unavailable(error.to_string()), }, + sample_sql: METADATA_SQL, + scope: METADATA_SCOPE, }; let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"]) .then(|column| async move { @@ -257,27 +391,27 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result= now() - INTERVAL 7 DAY \ LIMIT 200) ORDER BY key LIMIT 201" ); - let (keys, error) = match rows::(client, connection, &sql).await { - Ok(keys) => (keys, None), - Err(error) => (Vec::new(), Some(error.to_string())), + let discovery = match rows::(client, connection, &sql).await { + Ok(keys) => Discovery::Observed(AttributeSample { + truncated: keys.len() > MAX_FIELDS, + fields: keys + .into_iter() + .take(MAX_FIELDS) + .map(|row| AttributeField { + expression: format!("{column}[{}]", literal(&row.key)), + key: row.key, + kind: MapValueType::String, + }) + .collect(), + }), + Err(error) => Discovery::Unavailable(error.to_string()), }; - let fields = keys - .iter() - .take(MAX_FIELDS) - .map(|row| AttributeField { - key: row.key.clone(), - kind: "String", - expression: format!("{column}[{}]", literal(&row.key)), - }) - .collect(); AttributeCatalog { - table: "otel_traces", + table: TraceTable::OtelTraces, column, - fields, - truncated: error.is_some() || keys.len() > MAX_FIELDS, + discovery, discovery_sql: sql, scope: ATTRIBUTE_SCOPE, - error, } }) .collect::>() @@ -287,33 +421,31 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result>(), - "metadata": metadata, - "attributes": attributes, - "relationships": [{ - "left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id", - "additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND (otel_traces.TeamId != '' OR (otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' AND otel_traces.ApiKeyHash = spend_logs.api_key))", - "meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows" - }], - "examples": guide.examples()?, - "gotchas": guide.gotchas()?, - "guide": guide::render(&guide)?, - }).to_string()) + Ok(QueryHelp { + dialect: "ClickHouse SQL", + access: "Request-log visibility enforced by ClickHouse row policies; proxy admins see all rows, users see their own rows and permitted teams, and callers without user identity see their own key rows", + response: "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings", + examples: guide.examples()?, + gotchas: guide.gotchas()?, + guide: guide::render(&guide)?, + normalized_fields: NORMALIZED_FIELD_DEFINITIONS + .iter() + .map(NormalizedField::from) + .collect(), + relationships: &RELATIONSHIPS, + tables, + metadata, + attributes, + }) } #[cfg(test)] mod tests { use super::*; use rstest::rstest; + use serde_json::json; #[rstest] fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() { @@ -328,7 +460,7 @@ mod tests { metadata: "invalid".into(), }, ]; - let catalog = json!(metadata_catalog(&sample)); + let catalog = json!(metadata_sample(&sample)); assert_eq!( catalog["fields"], json!([{ @@ -351,7 +483,7 @@ mod tests { metadata: json!(metadata).to_string(), }) .collect(); - let catalog = json!(metadata_catalog(&sample)); + let catalog = json!(metadata_sample(&sample)); assert_eq!(catalog["truncated"], true); assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS)); assert_eq!( diff --git a/litellm-rust/crates/traces-clickhouse/src/query/guide.rs b/litellm-rust/crates/traces-clickhouse/src/query/guide.rs index 3bf7336648d..a0b60f9666f 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query/guide.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query/guide.rs @@ -1,8 +1,8 @@ use askama::Template; use serde::Serialize; -use super::{AttributeCatalog, MetadataCatalog, TableSchema}; -use crate::{Error, NormalizedFieldDefinition}; +use super::{AttributeCatalog, Discovery, MetadataCatalog, TableSchema}; +use crate::{Error, NormalizedFieldDefinition, query_access::ReaderLimits}; #[derive(Template)] #[template(path = "query_help.jinja", escape = "none", blocks = [ @@ -33,6 +33,7 @@ pub(super) struct QueryGuide<'a> { pub normalized_fields: &'a [NormalizedFieldDefinition], pub metadata: &'a MetadataCatalog, pub attributes: &'a [AttributeCatalog], + pub limits: &'a ReaderLimits, } #[derive(Serialize)] diff --git a/litellm-rust/crates/traces-clickhouse/src/query_access.rs b/litellm-rust/crates/traces-clickhouse/src/query_access.rs index 94ead18dd94..5d82a5e2e93 100644 --- a/litellm-rust/crates/traces-clickhouse/src/query_access.rs +++ b/litellm-rust/crates/traces-clickhouse/src/query_access.rs @@ -2,6 +2,7 @@ use std::{sync::Arc, time::Duration}; use hmac::{Hmac, Mac}; use litellm_http::Client; +use litellm_storage_clickhouse::READ_LIMITS; use litellm_traces::QueryScope; use moka::future::Cache; use strum::IntoEnumIterator; @@ -11,6 +12,33 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use super::{Connection, Error, TraceTable}; +const MIB: u64 = 1024 * 1024; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct ReaderLimits { + pub result_rows: u64, + pub result_bytes: u64, + pub memory_bytes: u64, + pub execution_seconds: u64, +} + +impl ReaderLimits { + pub fn result_mib(&self) -> u64 { + self.result_bytes / MIB + } + + pub fn memory_mib(&self) -> u64 { + self.memory_bytes / MIB + } +} + +pub(crate) const READER_LIMITS: ReaderLimits = ReaderLimits { + result_rows: READ_LIMITS.result_rows, + result_bytes: READ_LIMITS.response_bytes as u64, + memory_bytes: 256 * MIB, + execution_seconds: READ_LIMITS.execution_seconds, +}; + #[derive(Clone)] pub struct QueryReaders { writer: Connection, @@ -75,13 +103,19 @@ impl QueryReaders { return Err(Error::InvalidScope); } let password_hash = format!("{:x}", Sha256::digest(password)); + let ReaderLimits { + result_rows, + result_bytes, + memory_bytes, + execution_seconds, + } = READER_LIMITS; self.execute( client, format!( "CREATE USER IF NOT EXISTS {user} IDENTIFIED WITH sha256_hash BY '{password_hash}' \ - SETTINGS readonly = 1 CONST, max_execution_time = 10 CONST, \ - max_result_rows = 1000 CONST, max_result_bytes = 4194304 CONST, \ - result_overflow_mode = 'throw' CONST, max_memory_usage = 268435456 CONST, \ + SETTINGS readonly = 1 CONST, max_execution_time = {execution_seconds} CONST, \ + max_result_rows = {result_rows} CONST, max_result_bytes = {result_bytes} CONST, \ + result_overflow_mode = 'throw' CONST, max_memory_usage = {memory_bytes} CONST, \ max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST" ), ) diff --git a/litellm-rust/crates/traces-clickhouse/src/table.rs b/litellm-rust/crates/traces-clickhouse/src/table.rs index 5e3fe83fa82..4cbbf3cee89 100644 --- a/litellm-rust/crates/traces-clickhouse/src/table.rs +++ b/litellm-rust/crates/traces-clickhouse/src/table.rs @@ -1,6 +1,14 @@ #[derive( - Clone, Copy, Debug, strum::Display, strum::AsRefStr, strum::EnumIter, strum::IntoStaticStr, + Clone, + Copy, + Debug, + serde::Serialize, + strum::Display, + strum::AsRefStr, + strum::EnumIter, + strum::IntoStaticStr, )] +#[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] pub enum TraceTable { OtelTraces, diff --git a/litellm-rust/crates/traces-clickhouse/templates/query_help.jinja b/litellm-rust/crates/traces-clickhouse/templates/query_help.jinja index a79342c8343..6e746221022 100644 --- a/litellm-rust/crates/traces-clickhouse/templates/query_help.jinja +++ b/litellm-rust/crates/traces-clickhouse/templates/query_help.jinja @@ -11,18 +11,18 @@ Normalized span fields {% endfor %} Observed LLM call metadata {{ metadata.scope }} -Sampled rows: {{ metadata.sampled_rows }}; invalid JSON rows: {{ metadata.invalid_json_rows }}; truncated: {{ metadata.truncated }} -{% if let Some(error) = metadata.error %}Metadata discovery unavailable: {{ error }} -{% else if metadata.fields.is_empty() %}No metadata paths found in the sampled rows -{% else %}{% for field in metadata.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %} -{% endfor %}{% endif %} +{% match metadata.discovery %}{% when Discovery::Unavailable(error) %}Metadata discovery unavailable: {{ error }} +{% when Discovery::Observed(sample) %}Sampled rows: {{ sample.sampled_rows }}; invalid JSON rows: {{ sample.invalid_json_rows }}; truncated: {{ sample.truncated }} +{% if sample.fields.is_empty() %}No metadata paths found in the sampled rows +{% else %}{% for field in sample.fields %}{{ field.expression }}: {% for kind in field.types %}{{ kind }} {% endfor %} +{% endfor %}{% endif %}{% endmatch %} Observed span and resource attributes {% for catalog in attributes %}{{ catalog.table }}.{{ catalog.column }} {{ catalog.scope }} -{% if let Some(error) = catalog.error %}Attribute discovery unavailable: {{ error }} -{% else if catalog.fields.is_empty() %}No attribute keys found in the sampled spans -{% else %}{% for field in catalog.fields %}{{ field.expression }}: {{ field.kind }} -{% endfor %}{% endif %}{% endfor %} +{% match catalog.discovery %}{% when Discovery::Unavailable(error) %}Attribute discovery unavailable: {{ error }} +{% when Discovery::Observed(sample) %}{% if sample.fields.is_empty() %}No attribute keys found in the sampled spans +{% else %}{% for field in sample.fields %}{{ field.expression }}: {{ field.kind }} +{% endfor %}{% endif %}{% endmatch %}{% endfor %} Examples {% block recent_spans_name %}Recent normalized LLM spans{% endblock %} @@ -44,7 +44,7 @@ Gotchas {% block time_window %}Always bound Timestamp or start_time and use LIMIT; add TeamId/ApiKeyHash or team_id/api_key filters when investigating one tenant{% endblock %} -{% block reader_limits %}The reader enforces 1000 result rows, 4 MiB response bytes, 256 MiB memory and a 10 second query limit; exceeding limits fails instead of returning partial results{% endblock %} +{% block reader_limits %}The reader enforces {{ limits.result_rows }} result rows, {{ limits.result_mib() }} MiB response bytes, {{ limits.memory_mib() }} MiB memory and a {{ limits.execution_seconds }} second query limit; exceeding limits fails instead of returning partial results{% endblock %} {% block reader_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces request-log visibility through row policies. Callers see their own user rows and permitted teams, or their own key rows when no user identity is available. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %} diff --git a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs index dff7d8d5940..9fdeb450c70 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/migrations.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/migrations.rs @@ -1467,8 +1467,8 @@ async fn query_help_discovers_live_schema_and_runs_its_examples( ) .await?; } - let help: serde_json::Value = serde_json::from_str( - &litellm_traces_clickhouse::query_help(&database.client, &reader).await?, + let help = serde_json::to_value( + litellm_traces_clickhouse::query_help(&database.client, &reader).await?, )?; let keys: std::collections::BTreeSet<_> = help .as_object() @@ -1654,8 +1654,8 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits }))).collect::, _>>()?; insert_rows(&database, "otel_traces", spans).await?; let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?; - let help: serde_json::Value = serde_json::from_str( - &litellm_traces_clickhouse::query_help(&database.client, &reader).await?, + let help = serde_json::to_value( + litellm_traces_clickhouse::query_help(&database.client, &reader).await?, )?; assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3); assert!(!help["examples"].as_array().ok_or("examples")?.is_empty()); diff --git a/litellm-rust/crates/traces-clickhouse/tests/query_access.rs b/litellm-rust/crates/traces-clickhouse/tests/query_access.rs index d4c7c886bd0..be749351a72 100644 --- a/litellm-rust/crates/traces-clickhouse/tests/query_access.rs +++ b/litellm-rust/crates/traces-clickhouse/tests/query_access.rs @@ -94,7 +94,7 @@ async fn queries_and_help_are_scoped_by_the_database( .await?, )?; assert_eq!(summary["data"][0]["count"], json!(expected.len())); - let help = query_help(&database.client, &reader).await?; + let help = serde_json::to_string(&query_help(&database.client, &reader).await?)?; assert_eq!(help.contains("secret_b"), expected.contains(&"b")); assert_eq!(help.contains("secret-b"), expected.contains(&"b")); let recreated = QueryReaders::new(database.writer.clone(), "trace_test".to_owned()); diff --git a/litellm-rust/crates/traces/tests/query/named.rs b/litellm-rust/crates/traces/tests/query/named.rs index b20da7bd800..7db77df8a29 100644 --- a/litellm-rust/crates/traces/tests/query/named.rs +++ b/litellm-rust/crates/traces/tests/query/named.rs @@ -46,10 +46,10 @@ fn named_requests_preserve_all_access_cases(#[case] teams: Vec<&str>, #[case] ke #[rstest] fn result_contracts_preserve_public_field_names() { round_trip::( - json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), + json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["framework"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}), ); round_trip::( - json!({"span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "agent": "agent", "status": "error", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), + json!({"span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "agent": "agent", "framework": "framework", "status": "error", "status_message": "error", "error_truncated": 1, "start_ns": -1, "duration_ns": u64::MAX, "service": "service", "input_preview": "input", "model": "model", "input_tokens": u32::MAX, "output_tokens": 6, "litellm_request_id": "request", "team_id": "team", "api_key_hash": "key", "user_id": "user"}), ); round_trip::( json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}), diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index a86daaa35ad..45bccad6179 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -41,7 +41,7 @@ class NativeTraceStorage: def ensure_schema(self) -> Future[None]: ... def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Future[None]: ... def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Future[str]: ... - def query_help(self, scope: QueryScope, secret: str) -> Future[str]: ... + def query_help(self, scope: QueryScope, secret: str) -> Future[JsonValue]: ... def query(self, query: ReadQueryName, parameters: Mapping[str, str | int | float | Sequence[str]]) -> Future[str]: ... @final diff --git a/litellm/rust_bridge/trace_query_responses.py b/litellm/rust_bridge/trace_query_responses.py index 914d9b6c8d4..ad7678fcd9b 100644 --- a/litellm/rust_bridge/trace_query_responses.py +++ b/litellm/rust_bridge/trace_query_responses.py @@ -1,9 +1,12 @@ from collections.abc import Mapping -from typing import Final +from typing import Final, Literal from pydantic import BaseModel, ConfigDict, JsonValue _RESPONSE_CONFIG: Final = ConfigDict(frozen=True, extra="allow") +_HELP_CONFIG: Final = ConfigDict(frozen=True, extra="forbid") +TraceTableName = Literal["otel_traces", "agent_traces_by_key", "spend_logs"] +MetadataValueType = Literal["array", "boolean", "integer", "null", "number", "object", "string"] class TraceQueryColumn(BaseModel): @@ -28,14 +31,14 @@ class TraceSQLResponse(BaseModel): class TraceQueryTable(BaseModel): - model_config = ConfigDict(frozen=True) - name: str + model_config = _HELP_CONFIG + name: TraceTableName columns: tuple[TraceQueryColumn, ...] class TraceQueryNormalizedField(BaseModel): - model_config = ConfigDict(frozen=True) - table: str + model_config = _HELP_CONFIG + table: TraceTableName name: str column: str type: str @@ -43,15 +46,15 @@ class TraceQueryNormalizedField(BaseModel): class TraceQueryMetadataField(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = _HELP_CONFIG path: tuple[str | int, ...] - types: tuple[str, ...] + types: tuple[MetadataValueType, ...] expression: str class TraceQueryMetadata(BaseModel): - model_config = ConfigDict(frozen=True) - table: str + model_config = _HELP_CONFIG + table: TraceTableName column: str fields: tuple[TraceQueryMetadataField, ...] sampled_rows: int @@ -63,15 +66,15 @@ class TraceQueryMetadata(BaseModel): class TraceQueryAttributeField(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = _HELP_CONFIG key: str - type: str + type: Literal["String"] expression: str class TraceQueryAttributes(BaseModel): - model_config = ConfigDict(frozen=True) - table: str + model_config = _HELP_CONFIG + table: TraceTableName column: str fields: tuple[TraceQueryAttributeField, ...] truncated: bool @@ -81,7 +84,7 @@ class TraceQueryAttributes(BaseModel): class TraceQueryRelationship(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = _HELP_CONFIG left: str right: str additional_predicates: str @@ -89,13 +92,13 @@ class TraceQueryRelationship(BaseModel): class TraceQueryExample(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = _HELP_CONFIG name: str sql: str class TraceQueryHelp(BaseModel): - model_config = ConfigDict(frozen=True) + model_config = _HELP_CONFIG dialect: str access: str response: str diff --git a/litellm/rust_bridge/traces.py b/litellm/rust_bridge/traces.py index 2ed2df4e55e..23778983d75 100644 --- a/litellm/rust_bridge/traces.py +++ b/litellm/rust_bridge/traces.py @@ -2,7 +2,7 @@ from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass from typing import Final, Literal, Protocol, TypedDict, TypeVar, runtime_checkable -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly from litellm.rust_bridge.loader import get_native_bridge @@ -111,7 +111,7 @@ class NativeStore(Protocol): def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Awaitable[str]: ... - def query_help(self, scope: QueryScope, secret: str) -> Awaitable[str]: ... + def query_help(self, scope: QueryScope, secret: str) -> Awaitable[JsonValue]: ... def query( self, name: ReadQueryName, parameters: Mapping[str, str | int | float | Sequence[str]] @@ -189,6 +189,13 @@ def _decode_query_response(adapter: TypeAdapter[_ResponseT], body: str) -> _Resp raise RuntimeError("Native trace query returned an invalid response") from error +def _validate_query_response(adapter: TypeAdapter[_ResponseT], value: JsonValue) -> _ResponseT: + try: + return adapter.validate_python(value) + except ValidationError as error: + raise RuntimeError("Native trace query returned an invalid response") from error + + class ClickHouseStorage: def __init__(self, config: TraceStorageConfig) -> None: native: Final = _native() @@ -216,7 +223,7 @@ class ClickHouseStorage: async def query_help(self, scope: QueryScope, secret: str) -> TraceQueryHelp: result: Final = await self._native.query_help(scope, secret) - return _decode_query_response(_HELP_RESPONSE, result) + return _validate_query_response(_HELP_RESPONSE, result) async def lens_sample(self, parameters: LensSampleParams) -> tuple[ExecutionRow, ...]: return await self.query(LENS_SAMPLE, parameters) diff --git a/tests/unit/proxy/test_tracing_endpoints.py b/tests/unit/proxy/test_tracing_endpoints.py index 7c69796ddfa..fd05d54dcb6 100644 --- a/tests/unit/proxy/test_tracing_endpoints.py +++ b/tests/unit/proxy/test_tracing_endpoints.py @@ -2,8 +2,9 @@ Tests for the agent tracing endpoints (litellm/proxy/tracing_endpoints.py). """ -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from contextlib import asynccontextmanager +from types import ModuleType from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -15,9 +16,10 @@ from litellm.proxy import tracing_endpoints from litellm.proxy._types import LitellmUserRoles, ProxyLifespanState, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.tracing_runtime import manage_tracing, provide_storage +from litellm.rust_bridge import loader from litellm.rust_bridge.trace_queries import SPAN_DETAIL, SpanDetailParams from litellm.rust_bridge.trace_query_responses import TraceQueryHelp, TraceSQLResponse -from litellm.rust_bridge.traces import ClickHouseStorage +from litellm.rust_bridge.traces import AdminQueryScope, ClickHouseStorage, TraceStorageConfig from litellm.tracing import TraceReceiver, TracingPayloadTooLargeError from litellm.tracing.store import TraceStore from litellm.tracing.types import TraceScope @@ -29,14 +31,14 @@ SQL_ENVELOPE: Final = { "statistics": {"elapsed": 0.01, "rows_read": 1, "bytes_read": 8}, "rows_before_limit_at_least": 1, } -QUERY_HELP: Final = { +QUERY_HELP: Final[Mapping[str, object]] = { "dialect": "test SQL", "access": "authenticated scope", "response": "JSON envelope", - "tables": [{"name": "traces", "columns": [{"name": "value", "type": "String", "comment": "label"}]}], + "tables": [{"name": "otel_traces", "columns": [{"name": "value", "type": "String", "comment": "label"}]}], "normalized_fields": [], "metadata": { - "table": "traces", + "table": "spend_logs", "column": "metadata", "fields": [], "sampled_rows": 0, @@ -668,3 +670,60 @@ def test_queries_require_a_proxy_secret( receiver.store.storage.query_sql.assert_awaited_once_with( "SELECT 1", {"kind": "logs", "user_id": "", "team_ids": (), "api_key_hash": "hashed-key"}, secret ) + + +class _NativeConfig: + def __init__(self, database: str, url: str, retention_days: int) -> None: + pass + + +class _NativeReturningHelp(ModuleType): + def __init__(self, help_payload: Mapping[str, object]) -> None: + super().__init__("native_traces") + + class Storage: + def __init__(self, config: _NativeConfig) -> None: + pass + + async def query_help(self, scope: AdminQueryScope, secret: str) -> Mapping[str, object]: + return help_payload + + self.NativeTraceConfig: Final = _NativeConfig + self.NativeTraceStorage: Final = Storage + self.trace_decode_otlp: Final = list + self.trace_encode_error: Final = bytes + self.trace_normalized_field_definitions: Final = list + + +async def test_storage_validates_the_native_query_help_value(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(loader, "_cached_bridge", _NativeReturningHelp(QUERY_HELP)) + storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) + assert await storage.query_help({"kind": "admin"}, "secret") == TraceQueryHelp.model_validate(QUERY_HELP) + + +@pytest.mark.parametrize( + "drift", + ( + { + "metadata": { + "table": "spend_logs", + "column": "metadata", + "fields": [{"path": ["a"], "types": ["boolen"], "expression": "a"}], + "sampled_rows": 1, + "invalid_json_rows": 0, + "truncated": False, + "sample_sql": "SELECT metadata FROM spend_logs", + "scope": "bounded sample", + } + }, + {"tables": [{"name": "traces", "columns": [{"name": "value", "type": "String"}]}]}, + {"unexpected": True}, + ), +) +async def test_storage_rejects_native_query_help_that_drifts_from_the_contract( + monkeypatch: pytest.MonkeyPatch, drift: Mapping[str, object] +) -> None: + monkeypatch.setattr(loader, "_cached_bridge", _NativeReturningHelp({**QUERY_HELP, **drift})) + storage: Final = ClickHouseStorage(TraceStorageConfig("http://clickhouse:8123")) + with pytest.raises(RuntimeError, match="invalid response"): + await storage.query_help({"kind": "admin"}, "secret")