scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 71

refactor(traces): type the ClickHouse query help response (#44285)

Browse source 
* 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 <yujong@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-02 18:01:27 -07:00 • committed by GitHub
parent 9e31afbc6d
commit dd86ca5175
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 392 additions and 134 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-rust/crates/storage-clickhouse/src/lib.rs
View file

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

22
litellm-rust/crates/storage-clickhouse/src/read.rs
View file

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

2
litellm-rust/crates/traces-clickhouse/src/lib.rs
View file

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

292
litellm-rust/crates/traces-clickhouse/src/query.rs
View file

@ -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<PathPart>,
types: BTreeSet<&'static str>,
types: BTreeSet<JsonKind>,
expression: String,
}
@ -66,42 +101,150 @@ struct ColumnSchema {
#[derive(Serialize)]
struct TableSchema {
name: &'static str,
name: TraceTable,
columns: Vec<ColumnSchema>,
}
trait Unobserved {
fn unobserved() -> Self;
}
enum Discovery<T> {
Observed(T),
Unavailable(String),
}
impl<T: Serialize + Unobserved> Serialize for Discovery<T> {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
#[derive(Serialize)]
struct Unavailable<'a, T> {
#[serde(flatten)]
sample: T,
error: &'a str,
}
match self {
Self::Observed(sample) => sample.serialize(serializer),
Self::Unavailable(error) => Unavailable {
sample: T::unobserved(),
error,
}
.serialize(serializer),
}
}
}
#[derive(Serialize)]
struct MetadataCatalog {
table: &'static str,
column: &'static str,
struct MetadataSample {
fields: Vec<MetadataField>,
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<MetadataSample>,
sample_sql: &'static str,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
#[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<AttributeField>,
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<AttributeSample>,
discovery_sql: String,
scope: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
error: Option<String>,
}
#[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<TableSchema>,
normalized_fields: Vec<NormalizedField>,
metadata: MetadataCatalog,
attributes: Vec<AttributeCatalog>,
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<PathPart>,
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<JsonKind>>,
) -> bool {
if path.len() > MAX_DEPTH || (fields.len() >= MAX_FIELDS && !fields.contains_key(&path)) {
return true;
}
if !path.is_empty() {
let kind = match value {
Value::Null => "null",
Value::Bool(_) => "boolean",
Value::Number(number) if number.is_i64() || number.is_u64() => "integer",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
};
fields.entry(path.clone()).or_default().insert(kind);
fields
.entry(path.clone())
.or_default()
.insert(JsonKind::of(value));
}
match value {
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
@ -194,7 +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::<Value>(&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<String, Error> {
pub async fn query_help(client: &Client, connection: &Connection) -> Result<QueryHelp, Error> {
let tables = stream::iter(TraceTable::iter())
.then(|table| async move {
Ok::<_, Error>(TableSchema {
name: table.into(),
name: table,
columns: rows::<ColumnSchema>(
client,
connection,
@ -242,13 +374,15 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
})
.try_collect::<Vec<_>>()
.await?;
let metadata = match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
Ok(sample) => metadata_catalog(&sample),
Err(error) => MetadataCatalog {
error: Some(error.to_string()),
truncated: true,
..metadata_catalog(&[])
let metadata = MetadataCatalog {
table: TraceTable::SpendLogs,
column: "metadata",
discovery: match rows::<MetadataRow>(client, connection, METADATA_SQL).await {
Ok(sample) => Discovery::Observed(metadata_sample(&sample)),
Err(error) => Discovery::Unavailable(error.to_string()),
},
sample_sql: METADATA_SQL,
scope: METADATA_SCOPE,
};
let attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
.then(|column| async move {
@ -257,27 +391,27 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
(SELECT {column} FROM otel_traces WHERE Timestamp >= now() - INTERVAL 7 DAY \
LIMIT 200) ORDER BY key LIMIT 201"
);
let (keys, error) = match rows::<AttributeRow>(client, connection, &sql).await {
Ok(keys) => (keys, None),
Err(error) => (Vec::new(), Some(error.to_string())),
let discovery = match rows::<AttributeRow>(client, connection, &sql).await {
Ok(keys) => Discovery::Observed(AttributeSample {
truncated: keys.len() > MAX_FIELDS,
fields: keys
.into_iter()
.take(MAX_FIELDS)
.map(|row| AttributeField {
expression: format!("{column}[{}]", literal(&row.key)),
key: row.key,
kind: MapValueType::String,
})
.collect(),
}),
Err(error) => Discovery::Unavailable(error.to_string()),
};
let fields = keys
.iter()
.take(MAX_FIELDS)
.map(|row| AttributeField {
key: row.key.clone(),
kind: "String",
expression: format!("{column}[{}]", literal(&row.key)),
})
.collect();
AttributeCatalog {
table: "otel_traces",
table: TraceTable::OtelTraces,
column,
fields,
truncated: error.is_some() || keys.len() > MAX_FIELDS,
discovery,
discovery_sql: sql,
scope: ATTRIBUTE_SCOPE,
error,
}
})
.collect::<Vec<_>>()
@ -287,33 +421,31 @@ pub async fn query_help(client: &Client, connection: &Connection) -> Result<Stri
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
metadata: &metadata,
attributes: &attributes,
limits: &READER_LIMITS,
};
Ok(json!({
"dialect": "ClickHouse SQL",
"access": "Request-log visibility enforced by ClickHouse row policies; proxy admins see all rows, users see their own rows and permitted teams, and callers without user identity see their own key rows",
"response": "ClickHouse JSON envelope: meta, data, rows, statistics; 64-bit integers may be strings",
"tables": tables,
"normalized_fields": NORMALIZED_FIELD_DEFINITIONS.iter().map(|field| json!({
"table": "otel_traces", "name": field.name, "column": field.clickhouse_column,
"type": field.clickhouse_type, "meaning": field.meaning
})).collect::<Vec<_>>(),
"metadata": metadata,
"attributes": attributes,
"relationships": [{
"left": "otel_traces.LiteLLMRequestId", "right": "spend_logs.response_id",
"additional_predicates": "otel_traces.TeamId = spend_logs.team_id AND (otel_traces.TeamId != '' OR (otel_traces.UserId != '' AND otel_traces.UserId = spend_logs.user) OR (otel_traces.ApiKeyHash != '' AND otel_traces.ApiKeyHash = spend_logs.api_key))",
"meaning": "The normalized ID is the response ID, not request_id. Cached requests can share response_id; joins may return multiple spend rows"
}],
"examples": guide.examples()?,
"gotchas": guide.gotchas()?,
"guide": guide::render(&guide)?,
}).to_string())
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!(

5
litellm-rust/crates/traces-clickhouse/src/query/guide.rs
View file

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

40
litellm-rust/crates/traces-clickhouse/src/query_access.rs
View file

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

10
litellm-rust/crates/traces-clickhouse/src/table.rs
View file

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

20
litellm-rust/crates/traces-clickhouse/templates/query_help.jinja
View file

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

8
litellm-rust/crates/traces-clickhouse/tests/migrations.rs
View file

@ -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::<Result<Vec<_>, _>>()?;
insert_rows(&database, "otel_traces", spans).await?;
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
let help: serde_json::Value = serde_json::from_str(
&litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
let help = serde_json::to_value(
litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
)?;
assert_eq!(help["tables"].as_array().ok_or("tables")?.len(), 3);
assert!(!help["examples"].as_array().ok_or("examples")?.is_empty());

2
litellm-rust/crates/traces-clickhouse/tests/query_access.rs
View file

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

4
litellm-rust/crates/traces/tests/query/named.rs
View file

@ -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::<ListTracesRow>(
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "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::<TraceSpansRow>(
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::<SpanDetailRow>(
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),

2
litellm/rust_bridge/_native.pyi
View file

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

35
litellm/rust_bridge/trace_query_responses.py
View file

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

13
litellm/rust_bridge/traces.py
View file

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

69
tests/unit/proxy/test_tracing_endpoints.py
View file

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