mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(tracing): add scoped SQL queries and schema-aware help (#44085)
* feat(tracing): add SQL queries and schema-aware query help * test(tracing): verify help requests and sync API types * refactor(tracing): render query help with Askama * refactor(tracing): use jinja extension for query guide * fix(tracing): preserve query help when discovery fails * feat(tracing): enforce team SQL scope with managed ClickHouse readers * test(tracing): verify reads with one ClickHouse URL * fix(tracing): revoke rotated trace reader credentials Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(tracing): streamline query help catalog assembly Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tracing): run query help discovery sequentially Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(tracing): update reader setup request expectations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee
GitHub
parent
e6639d15b5
commit
8d28e8d776
20 changed files with 1827 additions and 6 deletions
52
litellm-rust/Cargo.lock
generated
52
litellm-rust/Cargo.lock
generated
|
|
@ -97,6 +97,53 @@ version = "1.2.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d"
|
||||
|
||||
[[package]]
|
||||
name = "askama"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "6024d73179f43f15ccd2b881bfea6fee7f3a46ec53f33b52210dea749ebebaa4"
|
||||
dependencies = [
|
||||
"askama_macros",
|
||||
"itoa",
|
||||
"percent-encoding",
|
||||
"serde",
|
||||
"serde_json",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "askama_derive"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "071ee5ebf2138e3ad180e0aacf6940c2cab5e6d8333741d9925c7bee2b153f39"
|
||||
dependencies = [
|
||||
"askama_parser",
|
||||
"memchr",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"rustc-hash",
|
||||
"syn 3.0.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "askama_macros"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "643e1c7cbb6aec1d920332fe51a7c0d8219e273dcb8602db03f5263e4d16487b"
|
||||
dependencies = [
|
||||
"askama_derive",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "askama_parser"
|
||||
version = "0.16.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2c5ae75772275d268b03ab8bdccdd12117b6169ee23256942b34e46c9f476583"
|
||||
dependencies = [
|
||||
"rustc-hash",
|
||||
"unicode-ident",
|
||||
"winnow 1.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "asn1-rs"
|
||||
version = "0.7.2"
|
||||
|
|
@ -4404,13 +4451,17 @@ dependencies = [
|
|||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"hmac 0.12.1",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
"litellm-storage-clickhouse",
|
||||
"moka",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
|
|
@ -4422,6 +4473,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -65,6 +65,7 @@ litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
|
|||
litellm-host-python = { path = "crates/host-python" }
|
||||
litellm-python-compat = { path = "crates/python-compat" }
|
||||
|
||||
askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
|
||||
tracing = "0.1"
|
||||
axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
|
||||
axum-login = "0.18.0"
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ use std::collections::BTreeMap;
|
|||
use litellm_host_python::{FromPythonCache, ToPythonCache};
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_storage_clickhouse::Storage;
|
||||
use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared};
|
||||
use litellm_traces::{
|
||||
Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery, Shared,
|
||||
};
|
||||
use prost::Message;
|
||||
use pyo3::{
|
||||
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
|
||||
|
|
@ -46,9 +48,25 @@ fn map_error(error: Error) -> PyErr {
|
|||
}
|
||||
}
|
||||
|
||||
fn map_sql_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
|
||||
error => map_error(error),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_query_access_error(error: QueryAccessError) -> PyErr {
|
||||
match error {
|
||||
QueryAccessError::Storage(error) => map_sql_error(error),
|
||||
QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()),
|
||||
error => PyRuntimeError::new_err(error.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[pyclass]
|
||||
pub struct NativeTraceStorage {
|
||||
storage: Storage,
|
||||
query_readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[pymethods]
|
||||
|
|
@ -57,8 +75,13 @@ impl NativeTraceStorage {
|
|||
#[pyo3(signature = (database, url, reader_url = None))]
|
||||
fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
|
||||
litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
|
||||
let storage = Storage::new(database, url, reader_url).map_err(map_error)?;
|
||||
Ok(Self {
|
||||
storage: Storage::new(database, url, reader_url).map_err(map_error)?,
|
||||
query_readers: QueryReaders::new(
|
||||
storage.writer().clone(),
|
||||
storage.database().to_owned(),
|
||||
),
|
||||
storage,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -107,6 +130,52 @@ impl NativeTraceStorage {
|
|||
)
|
||||
}
|
||||
|
||||
fn query_sql<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
sql: String,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
|
||||
secret: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
if sql.trim().is_empty() {
|
||||
return Err(map_error(Error::EmptySql));
|
||||
}
|
||||
let readers = self.query_readers.clone();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
let _permit = readers.acquire()?;
|
||||
let connection = readers.connection(&client, &scope, &secret).await?;
|
||||
litellm_traces::query_sql(&client, &connection, &sql)
|
||||
.await
|
||||
.map_err(QueryAccessError::Storage)
|
||||
},
|
||||
map_query_access_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn query_help<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
|
||||
secret: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let readers = self.query_readers.clone();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
let _permit = readers.acquire()?;
|
||||
let connection = readers.connection(&client, &scope, &secret).await?;
|
||||
litellm_traces::query_help(&client, &connection)
|
||||
.await
|
||||
.map_err(QueryAccessError::Storage)
|
||||
},
|
||||
map_query_access_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn lens_query<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
|
|
|
|||
|
|
@ -6,9 +6,13 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
askama.workspace = true
|
||||
base64.workspace = true
|
||||
flate2.workspace = true
|
||||
futures-util.workspace = true
|
||||
hmac = "0.12.1"
|
||||
indexmap = { version = "2", features = ["serde"] }
|
||||
moka.workspace = true
|
||||
opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
|
||||
prost.workspace = true
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
|
|
@ -20,6 +24,8 @@ serde = { workspace = true, features = ["rc"] }
|
|||
serde_json.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
criterion.workspace = true
|
||||
|
|
|
|||
|
|
@ -7,3 +7,23 @@ pub enum DecodeError {
|
|||
#[error("OTLP token count is outside the storage range")]
|
||||
TokenCountOutOfRange,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum QueryAccessError {
|
||||
#[error("trace SQL queries require a configured proxy master key")]
|
||||
MissingSecret,
|
||||
#[error("invalid trace query scope")]
|
||||
InvalidScope,
|
||||
#[error("trace SQL query concurrency limit exceeded")]
|
||||
Busy,
|
||||
#[error(
|
||||
"ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables"
|
||||
)]
|
||||
ProvisionFailed(u16),
|
||||
#[error("ClickHouse reader provisioning transport failed")]
|
||||
ProvisionTransport,
|
||||
#[error(transparent)]
|
||||
Storage(#[from] litellm_storage_clickhouse::Error),
|
||||
#[error(transparent)]
|
||||
Cached(#[from] std::sync::Arc<QueryAccessError>),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,17 +2,22 @@ mod error;
|
|||
mod insert;
|
||||
mod normalize;
|
||||
mod otlp;
|
||||
mod query;
|
||||
mod query_access;
|
||||
mod schema;
|
||||
mod shared;
|
||||
mod sql;
|
||||
|
||||
pub use error::DecodeError;
|
||||
pub use error::{DecodeError, QueryAccessError};
|
||||
pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
|
||||
pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
|
||||
pub use normalize::{
|
||||
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType,
|
||||
};
|
||||
pub use otlp::{DecodedSpan, decode_otlp};
|
||||
pub use query_access::{QueryReaders, QueryScope};
|
||||
pub use schema::{ensure_schema, schema_statements};
|
||||
pub use shared::{Shared, SharedIdentity};
|
||||
pub use sql::{LensQuery, ReadQuery, execute_named_read};
|
||||
|
||||
pub use query::{query_help, query_sql};
|
||||
|
|
|
|||
346
litellm-rust/crates/traces/src/query.rs
Normal file
346
litellm-rust/crates/traces/src/query.rs
Normal file
|
|
@ -0,0 +1,346 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use futures_util::{
|
||||
StreamExt,
|
||||
stream::{self, TryStreamExt},
|
||||
};
|
||||
use litellm_http::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, execute_read};
|
||||
|
||||
mod guide;
|
||||
|
||||
const SAMPLE_ROWS: usize = 200;
|
||||
const MAX_FIELDS: usize = 200;
|
||||
const MAX_DEPTH: usize = 16;
|
||||
const METADATA_SQL: &str = "SELECT metadata FROM spend_logs FINAL \
|
||||
WHERE start_time >= now() - INTERVAL 7 DAY AND length(metadata) <= 8192 \
|
||||
LIMIT 201";
|
||||
const METADATA_SCOPE: &str = "Up to 200 unordered rows from the last 7 days, excluding metadata larger than 8192 bytes; up to 200 paths and 16 levels. Missing paths may exist outside this sample. Array indexes are 1-based and describe sampled positions, not a fixed schema";
|
||||
const ATTRIBUTE_SCOPE: &str = "Distinct keys from up to 200 unordered spans in the last 7 days; up to 200 keys per map. Missing keys may exist outside this sample";
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Rows<T> {
|
||||
data: Vec<T>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct MetadataRow {
|
||||
metadata: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AttributeRow {
|
||||
key: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Serialize)]
|
||||
#[serde(untagged)]
|
||||
enum PathPart {
|
||||
Key(String),
|
||||
Index(usize),
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MetadataField {
|
||||
path: Vec<PathPart>,
|
||||
types: BTreeSet<&'static str>,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct ColumnSchema {
|
||||
name: String,
|
||||
#[serde(rename = "type")]
|
||||
kind: String,
|
||||
#[serde(flatten)]
|
||||
details: BTreeMap<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct TableSchema {
|
||||
name: &'static str,
|
||||
columns: Vec<ColumnSchema>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MetadataCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
fields: Vec<MetadataField>,
|
||||
sampled_rows: usize,
|
||||
invalid_json_rows: usize,
|
||||
truncated: bool,
|
||||
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,
|
||||
expression: String,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AttributeCatalog {
|
||||
table: &'static str,
|
||||
column: &'static str,
|
||||
fields: Vec<AttributeField>,
|
||||
truncated: bool,
|
||||
discovery_sql: String,
|
||||
scope: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
error: Option<String>,
|
||||
}
|
||||
|
||||
pub async fn query_sql(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
) -> Result<String, Error> {
|
||||
execute_read(client, connection, sql, &BTreeMap::new()).await
|
||||
}
|
||||
|
||||
async fn rows<T: serde::de::DeserializeOwned>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
) -> Result<Vec<T>, Error> {
|
||||
let body = query_sql(client, connection, sql).await?;
|
||||
serde_json::from_str::<Rows<T>>(&body)
|
||||
.map(|result| result.data)
|
||||
.map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
|
||||
fn literal(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
|
||||
}
|
||||
|
||||
fn metadata_expression(path: &[PathPart]) -> String {
|
||||
let arguments = path
|
||||
.iter()
|
||||
.map(|part| match part {
|
||||
PathPart::Key(key) => literal(key),
|
||||
PathPart::Index(index) => index.to_string(),
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
format!("JSONExtractRaw(metadata, {arguments})")
|
||||
}
|
||||
|
||||
fn discover(
|
||||
value: &Value,
|
||||
path: Vec<PathPart>,
|
||||
fields: &mut BTreeMap<Vec<PathPart>, BTreeSet<&'static str>>,
|
||||
) -> 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);
|
||||
}
|
||||
match value {
|
||||
Value::Object(object) => object.iter().fold(false, |limited, (key, value)| {
|
||||
let child = path
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([PathPart::Key(key.clone())])
|
||||
.collect();
|
||||
discover(value, child, fields) | limited
|
||||
}),
|
||||
Value::Array(array) => array
|
||||
.iter()
|
||||
.enumerate()
|
||||
.fold(false, |limited, (index, value)| {
|
||||
let child = path
|
||||
.iter()
|
||||
.cloned()
|
||||
.chain([PathPart::Index(index + 1)])
|
||||
.collect();
|
||||
discover(value, child, fields) | limited
|
||||
}),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn metadata_catalog(sample: &[MetadataRow]) -> MetadataCatalog {
|
||||
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) {
|
||||
Ok(value) => {
|
||||
let mut fields = fields;
|
||||
let limited = limited | discover(&value, Vec::new(), &mut fields);
|
||||
(fields, limited, invalid_rows)
|
||||
}
|
||||
Err(_) => (fields, limited, invalid_rows + 1),
|
||||
},
|
||||
);
|
||||
let fields: Vec<_> = fields
|
||||
.into_iter()
|
||||
.map(|(path, types)| MetadataField {
|
||||
expression: metadata_expression(&path),
|
||||
path,
|
||||
types,
|
||||
})
|
||||
.collect();
|
||||
MetadataCatalog {
|
||||
table: "spend_logs",
|
||||
column: "metadata",
|
||||
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> {
|
||||
let tables = stream::iter(["otel_traces", "agent_traces_by_key", "spend_logs"])
|
||||
.then(|table| async move {
|
||||
Ok::<_, Error>(TableSchema {
|
||||
name: table,
|
||||
columns: rows::<ColumnSchema>(
|
||||
client,
|
||||
connection,
|
||||
&format!("DESCRIBE TABLE {table}"),
|
||||
)
|
||||
.await?,
|
||||
})
|
||||
})
|
||||
.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 attributes = stream::iter(["SpanAttributes", "ResourceAttributes"])
|
||||
.then(|column| async move {
|
||||
let sql = format!(
|
||||
"SELECT DISTINCT arrayJoin(mapKeys({column})) AS key FROM \
|
||||
(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 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",
|
||||
column,
|
||||
fields,
|
||||
truncated: error.is_some() || keys.len() > MAX_FIELDS,
|
||||
discovery_sql: sql,
|
||||
scope: ATTRIBUTE_SCOPE,
|
||||
error,
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.await;
|
||||
let guide = guide::QueryGuide {
|
||||
tables: &tables,
|
||||
normalized_fields: &NORMALIZED_FIELD_DEFINITIONS,
|
||||
metadata: &metadata,
|
||||
attributes: &attributes,
|
||||
};
|
||||
Ok(json!({
|
||||
"dialect": "ClickHouse SQL",
|
||||
"access": "Authenticated team scope enforced by ClickHouse row policies; proxy admins can read all teams, while project-bound and teamless keys can read only their own 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.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())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
fn metadata_discovery_preserves_mixed_types_and_reports_invalid_rows() {
|
||||
let sample = [
|
||||
MetadataRow {
|
||||
metadata: r#"{"x": 1}"#.into(),
|
||||
},
|
||||
MetadataRow {
|
||||
metadata: r#"{"x": "one"}"#.into(),
|
||||
},
|
||||
MetadataRow {
|
||||
metadata: "invalid".into(),
|
||||
},
|
||||
];
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
assert_eq!(
|
||||
catalog["fields"],
|
||||
json!([{
|
||||
"path": ["x"], "types": ["integer", "string"], "expression": "JSONExtractRaw(metadata, 'x')"
|
||||
}])
|
||||
);
|
||||
assert_eq!(catalog["invalid_json_rows"], 1);
|
||||
assert_eq!(catalog["sampled_rows"], sample.len());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::rows(SAMPLE_ROWS + 1, 1)]
|
||||
#[case::paths(1, MAX_FIELDS + 1)]
|
||||
fn metadata_discovery_reports_truncation(#[case] row_count: usize, #[case] field_count: usize) {
|
||||
let metadata: BTreeMap<_, _> = (0..field_count)
|
||||
.map(|index| (format!("field{index}"), index))
|
||||
.collect();
|
||||
let sample: Vec<_> = (0..row_count)
|
||||
.map(|_| MetadataRow {
|
||||
metadata: json!(metadata).to_string(),
|
||||
})
|
||||
.collect();
|
||||
let catalog = json!(metadata_catalog(&sample));
|
||||
assert_eq!(catalog["truncated"], true);
|
||||
assert_eq!(catalog["sampled_rows"], row_count.min(SAMPLE_ROWS));
|
||||
assert_eq!(
|
||||
catalog["fields"].as_array().unwrap().len(),
|
||||
field_count.min(MAX_FIELDS)
|
||||
);
|
||||
}
|
||||
}
|
||||
89
litellm-rust/crates/traces/src/query/guide.rs
Normal file
89
litellm-rust/crates/traces/src/query/guide.rs
Normal file
|
|
@ -0,0 +1,89 @@
|
|||
use askama::Template;
|
||||
use serde::Serialize;
|
||||
|
||||
use super::{AttributeCatalog, MetadataCatalog, TableSchema};
|
||||
use crate::{Error, NormalizedFieldDefinition};
|
||||
|
||||
#[derive(Template)]
|
||||
#[template(path = "query_help.jinja", escape = "none", blocks = [
|
||||
"recent_spans_name",
|
||||
"recent_spans_sql",
|
||||
"custom_metadata_name",
|
||||
"custom_metadata_sql",
|
||||
"nested_metadata_name",
|
||||
"nested_metadata_sql",
|
||||
"correlated_calls_name",
|
||||
"correlated_calls_sql",
|
||||
"discover_keys_name",
|
||||
"discover_keys_sql",
|
||||
"time_window",
|
||||
"reader_limits",
|
||||
"reader_profile",
|
||||
"output_format",
|
||||
"json_values",
|
||||
"map_values",
|
||||
"literal_keys",
|
||||
"time_units",
|
||||
"spend_totals",
|
||||
"trace_rollups",
|
||||
"sampling",
|
||||
])]
|
||||
pub(super) struct QueryGuide<'a> {
|
||||
pub tables: &'a [TableSchema],
|
||||
pub normalized_fields: &'a [NormalizedFieldDefinition],
|
||||
pub metadata: &'a MetadataCatalog,
|
||||
pub attributes: &'a [AttributeCatalog],
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub(super) struct Example {
|
||||
name: String,
|
||||
sql: String,
|
||||
}
|
||||
|
||||
impl QueryGuide<'_> {
|
||||
pub fn examples(&self) -> Result<[Example; 5], Error> {
|
||||
Ok([
|
||||
Example {
|
||||
name: render(&self.as_recent_spans_name())?,
|
||||
sql: render(&self.as_recent_spans_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_custom_metadata_name())?,
|
||||
sql: render(&self.as_custom_metadata_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_nested_metadata_name())?,
|
||||
sql: render(&self.as_nested_metadata_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_correlated_calls_name())?,
|
||||
sql: render(&self.as_correlated_calls_sql())?,
|
||||
},
|
||||
Example {
|
||||
name: render(&self.as_discover_keys_name())?,
|
||||
sql: render(&self.as_discover_keys_sql())?,
|
||||
},
|
||||
])
|
||||
}
|
||||
|
||||
pub fn gotchas(&self) -> Result<[String; 11], Error> {
|
||||
Ok([
|
||||
render(&self.as_time_window())?,
|
||||
render(&self.as_reader_limits())?,
|
||||
render(&self.as_reader_profile())?,
|
||||
render(&self.as_output_format())?,
|
||||
render(&self.as_json_values())?,
|
||||
render(&self.as_map_values())?,
|
||||
render(&self.as_literal_keys())?,
|
||||
render(&self.as_time_units())?,
|
||||
render(&self.as_spend_totals())?,
|
||||
render(&self.as_trace_rollups())?,
|
||||
render(&self.as_sampling())?,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn render(template: &impl Template) -> Result<String, Error> {
|
||||
template.render().map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
200
litellm-rust/crates/traces/src/query_access.rs
Normal file
200
litellm-rust/crates/traces/src/query_access.rs
Normal file
|
|
@ -0,0 +1,200 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use litellm_http::Client;
|
||||
use moka::future::Cache;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use crate::{Connection, QueryAccessError};
|
||||
|
||||
const TABLES: [&str; 3] = ["otel_traces", "agent_traces_by_key", "spend_logs"];
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
|
||||
pub enum QueryScope {
|
||||
Admin,
|
||||
Team {
|
||||
team_id: String,
|
||||
},
|
||||
Key {
|
||||
team_id: String,
|
||||
api_key_hash: String,
|
||||
},
|
||||
}
|
||||
|
||||
impl QueryScope {
|
||||
fn validate(&self) -> Result<(), QueryAccessError> {
|
||||
match self {
|
||||
Self::Admin => Ok(()),
|
||||
Self::Team { team_id } if !team_id.is_empty() => Ok(()),
|
||||
Self::Key { api_key_hash, .. } if !api_key_hash.is_empty() => Ok(()),
|
||||
_ => Err(QueryAccessError::InvalidScope),
|
||||
}
|
||||
}
|
||||
|
||||
fn predicate(&self, table: &str) -> String {
|
||||
let (team, key) = if table == "spend_logs" {
|
||||
("team_id", "api_key")
|
||||
} else {
|
||||
("TeamId", "ApiKeyHash")
|
||||
};
|
||||
match self {
|
||||
Self::Admin => "1".to_owned(),
|
||||
Self::Team { team_id } => format!("{team} = {}", literal(team_id)),
|
||||
Self::Key {
|
||||
team_id,
|
||||
api_key_hash,
|
||||
} => format!(
|
||||
"{team} = {} AND {key} = {}",
|
||||
literal(team_id),
|
||||
literal(api_key_hash)
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct QueryReaders {
|
||||
writer: Connection,
|
||||
database: String,
|
||||
readers: Cache<String, Connection>,
|
||||
slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
impl QueryReaders {
|
||||
pub fn new(writer: Connection, database: String) -> Self {
|
||||
Self {
|
||||
writer,
|
||||
database,
|
||||
readers: Cache::builder().max_capacity(1024).build(),
|
||||
slots: Arc::new(Semaphore::new(8)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn acquire(&self) -> Result<OwnedSemaphorePermit, QueryAccessError> {
|
||||
self.slots
|
||||
.clone()
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| QueryAccessError::Busy)
|
||||
}
|
||||
|
||||
pub async fn connection(
|
||||
&self,
|
||||
client: &Client,
|
||||
scope: &QueryScope,
|
||||
secret: &str,
|
||||
) -> Result<Connection, QueryAccessError> {
|
||||
scope.validate()?;
|
||||
if secret.is_empty() {
|
||||
return Err(QueryAccessError::MissingSecret);
|
||||
}
|
||||
let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope))
|
||||
.map_err(|_| QueryAccessError::InvalidScope)?;
|
||||
let user = format!("litellm_traces_{:x}", Sha256::digest(&identity));
|
||||
let password = credential(secret, b"password", &identity)?;
|
||||
self.readers
|
||||
.try_get_with(
|
||||
user.clone(),
|
||||
self.provision(client, scope, &user, &password),
|
||||
)
|
||||
.await
|
||||
.map_err(QueryAccessError::Cached)
|
||||
}
|
||||
|
||||
async fn provision(
|
||||
&self,
|
||||
client: &Client,
|
||||
scope: &QueryScope,
|
||||
user: &str,
|
||||
password: &str,
|
||||
) -> Result<Connection, QueryAccessError> {
|
||||
let database = &self.database;
|
||||
if database.is_empty()
|
||||
|| !database
|
||||
.bytes()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|
||||
{
|
||||
return Err(QueryAccessError::InvalidScope);
|
||||
}
|
||||
let password_hash = format!("{:x}", Sha256::digest(password));
|
||||
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, \
|
||||
max_threads = 2 CONST, max_concurrent_queries_for_user = 8 CONST"
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
self.execute(
|
||||
client,
|
||||
format!("ALTER USER {user} IDENTIFIED WITH sha256_hash BY '{password_hash}'"),
|
||||
)
|
||||
.await?;
|
||||
for table in TABLES {
|
||||
let predicate = scope.predicate(table);
|
||||
self.execute(
|
||||
client,
|
||||
format!(
|
||||
"CREATE ROW POLICY IF NOT EXISTS {user}_allow ON `{database}`.{table} \
|
||||
USING 1 TO {user}"
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
self.execute(
|
||||
client,
|
||||
format!(
|
||||
"CREATE ROW POLICY IF NOT EXISTS {user}_scope ON `{database}`.{table} \
|
||||
AS RESTRICTIVE USING {predicate} TO {user}"
|
||||
),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
for table in TABLES {
|
||||
self.execute(
|
||||
client,
|
||||
format!("GRANT SELECT ON `{database}`.{table} TO {user}"),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
Connection::configured(
|
||||
&self.writer.url()[..url::Position::AfterPath],
|
||||
database,
|
||||
user,
|
||||
password,
|
||||
)
|
||||
.map_err(QueryAccessError::Storage)
|
||||
}
|
||||
|
||||
async fn execute(&self, client: &Client, sql: String) -> Result<(), QueryAccessError> {
|
||||
let response = client
|
||||
.post(self.writer.url().clone())
|
||||
.timeout(Duration::from_secs(15))
|
||||
.body(sql)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| QueryAccessError::ProvisionTransport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(QueryAccessError::ProvisionFailed(
|
||||
response.status().as_u16(),
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, QueryAccessError> {
|
||||
let mut mac = Hmac::<Sha256>::new_from_slice(secret.as_bytes())
|
||||
.map_err(|_| QueryAccessError::MissingSecret)?;
|
||||
mac.update(purpose);
|
||||
mac.update(identity);
|
||||
Ok(format!("{:x}", mac.finalize().into_bytes()))
|
||||
}
|
||||
|
||||
fn literal(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
|
||||
}
|
||||
65
litellm-rust/crates/traces/templates/query_help.jinja
Normal file
65
litellm-rust/crates/traces/templates/query_help.jinja
Normal file
|
|
@ -0,0 +1,65 @@
|
|||
Trace SQL query guide
|
||||
|
||||
Live ClickHouse schema
|
||||
{% for table in tables %}
|
||||
{{ table.name }}
|
||||
{% for column in table.columns %}{{ column.name }}: {{ column.kind }}
|
||||
{% endfor %}{% endfor %}
|
||||
Normalized span fields
|
||||
{% for field in normalized_fields %}{{ field.name }}: otel_traces.{{ field.clickhouse_column }} ({{ field.clickhouse_type }})
|
||||
{{ field.meaning }}
|
||||
{% 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 %}
|
||||
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 %}
|
||||
Examples
|
||||
|
||||
{% block recent_spans_name %}Recent normalized LLM spans{% endblock %}
|
||||
{% block recent_spans_sql %}SELECT TraceId, SpanId, Model, InputTokens, OutputTokens, Duration / 1000000 AS duration_ms FROM otel_traces WHERE Timestamp >= now() - INTERVAL 1 DAY AND ObservationType = 'llm' ORDER BY Timestamp DESC LIMIT 100{% endblock %}
|
||||
|
||||
{% block custom_metadata_name %}Find calls by custom metadata{% endblock %}
|
||||
{% block custom_metadata_sql %}SELECT request_id, response_id, model, spend, JSONExtractString(metadata, 'project') AS project FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'project') AND JSONExtractString(metadata, 'project') = 'example' ORDER BY start_time DESC LIMIT 100{% endblock %}
|
||||
|
||||
{% block nested_metadata_name %}Nested metadata with unknown types{% endblock %}
|
||||
{% block nested_metadata_sql %}SELECT request_id, JSONType(metadata, 'labels', 'priority') AS type, JSONExtractRaw(metadata, 'labels', 'priority') AS value FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY AND JSONHas(metadata, 'labels', 'priority') LIMIT 100{% endblock %}
|
||||
|
||||
{% block correlated_calls_name %}Traces correlated with LLM call metadata{% endblock %}
|
||||
{% block correlated_calls_sql %}SELECT t.TraceId, t.SpanId, s.request_id, s.spend, s.metadata FROM otel_traces AS t INNER JOIN (SELECT * FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 1 DAY) AS s ON t.LiteLLMRequestId = s.response_id AND t.TeamId = s.team_id AND t.ApiKeyHash = s.api_key WHERE t.Timestamp >= now() - INTERVAL 1 DAY AND t.LiteLLMRequestId != '' AND JSONExtractString(s.metadata, 'project') = 'example' LIMIT 100{% endblock %}
|
||||
|
||||
{% block discover_keys_name %}Discover metadata keys over a different window{% endblock %}
|
||||
{% block discover_keys_sql %}SELECT DISTINCT arrayJoin(JSONExtractKeys(metadata)) AS key FROM spend_logs FINAL WHERE start_time >= now() - INTERVAL 30 DAY ORDER BY key LIMIT 200{% endblock %}
|
||||
|
||||
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_profile %}LiteLLM provisions SELECT-only readers from the configured ClickHouse connection and enforces authenticated team scope through row policies. Project-bound and teamless keys see only their own rows. Provisioning requires CREATE USER, ALTER USER, CREATE ROW POLICY, and GRANT SELECT permissions{% endblock %}
|
||||
|
||||
{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %}
|
||||
|
||||
{% block json_values %}metadata is a JSON-encoded String; use JSONHas before typed extraction to distinguish missing values from empty strings, zero and false{% endblock %}
|
||||
|
||||
{% block map_values %}SpanAttributes and ResourceAttributes are Map(String, String); missing map keys return an empty string, so use mapContains for existence checks{% endblock %}
|
||||
|
||||
{% block literal_keys %}Use the discovered path components as separate JSONExtract arguments; a dot inside a key is literal, not a path separator{% endblock %}
|
||||
|
||||
{% block time_units %}Duration is nanoseconds; Timestamp has nanosecond precision, spend start_time has millisecond precision{% endblock %}
|
||||
|
||||
{% block spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; aggregate spend separately{% endblock %}
|
||||
|
||||
{% block trace_rollups %}agent_traces_by_key uses SimpleAggregateFunction columns; group by TeamId, ApiKeyHash and TraceId, using min(StartTs), max(EndTs), sum(SpanCount) and groupUniqArrayArray(Models). Do not use Merge combinators{% endblock %}
|
||||
|
||||
{% block sampling %}Discovery is sampled, contains no metadata values, and is not an exhaustive schema. Edit the supplied discovery SQL for older data or nested JSONExtractKeys(metadata, 'parent'){% endblock %}
|
||||
|
|
@ -1149,3 +1149,289 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
assert_eq!(available["data"][0]["requests"], 0);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(false)]
|
||||
#[case::custom_metadata(true)]
|
||||
#[tokio::test]
|
||||
async fn query_help_discovers_live_schema_and_runs_its_examples(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] populated: bool,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
|
||||
execute_write(&database, "CREATE USER help_reader").await?;
|
||||
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
|
||||
execute_write(
|
||||
&database,
|
||||
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let reader = Connection::configured(&database.url, "trace_test", "help_reader", "")?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
if populated {
|
||||
execute_write(&database, "SYSTEM STOP MERGES trace_test.spend_logs").await?;
|
||||
insert_rows(
|
||||
&database,
|
||||
"spend_logs",
|
||||
vec![serde_json::from_value(serde_json::json!({
|
||||
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
|
||||
"api_key": "key-1", "metadata": r#"{"obsolete":true,"labels":{"priority":"old"}}"#,
|
||||
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000
|
||||
}))?],
|
||||
)
|
||||
.await?;
|
||||
let metadata = serde_json::json!({
|
||||
"project": "example", "labels": {"priority": 3, "enabled": true},
|
||||
"dotted.key": "literal", "quote'\\key": null, "items": [{"name": "first"}],
|
||||
"<custom>&{{key}}": {"nested.key": true}
|
||||
});
|
||||
insert_rows(
|
||||
&database,
|
||||
"spend_logs",
|
||||
vec![serde_json::from_value(serde_json::json!({
|
||||
"request_id": "request-1", "response_id": "response-1", "team_id": "team-1",
|
||||
"api_key": "key-1", "metadata": metadata.to_string(), "spend": 0.25,
|
||||
"start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000 + 100
|
||||
}))?],
|
||||
)
|
||||
.await?;
|
||||
insert_rows(
|
||||
&database,
|
||||
"otel_traces",
|
||||
vec![serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1",
|
||||
"TeamId": "team-1", "ApiKeyHash": "key-1", "ObservationType": "llm",
|
||||
"LiteLLMRequestId": "response-1", "SpanAttributes": {"custom.tag": "value"},
|
||||
"ResourceAttributes": {"custom.resource": "value"}
|
||||
}))?],
|
||||
)
|
||||
.await?;
|
||||
execute_write(
|
||||
&database,
|
||||
"ALTER TABLE trace_test.otel_traces ADD COLUMN CustomColumn String",
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let help: serde_json::Value =
|
||||
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
|
||||
let keys: std::collections::BTreeSet<_> = help
|
||||
.as_object()
|
||||
.ok_or("missing help object")?
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
assert_eq!(
|
||||
keys,
|
||||
std::collections::BTreeSet::from([
|
||||
"access",
|
||||
"attributes",
|
||||
"dialect",
|
||||
"examples",
|
||||
"gotchas",
|
||||
"guide",
|
||||
"metadata",
|
||||
"normalized_fields",
|
||||
"relationships",
|
||||
"response",
|
||||
"tables",
|
||||
])
|
||||
);
|
||||
let guide = help["guide"].as_str().ok_or("missing rendered guide")?;
|
||||
assert!(guide.starts_with("Trace SQL query guide"));
|
||||
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
|
||||
let described = read_json(&database, &format!("DESCRIBE TABLE {table}")).await?;
|
||||
let schema = help["tables"]
|
||||
.as_array()
|
||||
.ok_or("missing tables")?
|
||||
.iter()
|
||||
.find(|schema| schema["name"] == table)
|
||||
.ok_or("missing table")?;
|
||||
assert_eq!(schema["columns"], described["data"]);
|
||||
for column in described["data"].as_array().ok_or("missing live columns")? {
|
||||
assert!(guide.contains(&format!(
|
||||
"{}: {}",
|
||||
column["name"].as_str().ok_or("column name")?,
|
||||
column["type"].as_str().ok_or("column type")?
|
||||
)));
|
||||
}
|
||||
}
|
||||
for gotcha in help["gotchas"].as_array().ok_or("missing gotchas")? {
|
||||
assert!(guide.contains(gotcha.as_str().ok_or("gotcha text")?));
|
||||
}
|
||||
let tables = help["tables"].as_array().ok_or("missing tables")?;
|
||||
assert_eq!(tables.len(), 3);
|
||||
let columns = tables[0]["columns"].as_array().ok_or("missing columns")?;
|
||||
for field in NORMALIZED_FIELD_DEFINITIONS {
|
||||
assert!(
|
||||
columns
|
||||
.iter()
|
||||
.any(|column| column["name"] == field.clickhouse_column
|
||||
&& column["type"] == field.clickhouse_type)
|
||||
);
|
||||
assert!(
|
||||
help["normalized_fields"]
|
||||
.as_array()
|
||||
.ok_or("missing mappings")?
|
||||
.iter()
|
||||
.any(|mapped| {
|
||||
mapped["name"] == field.name && mapped["column"] == field.clickhouse_column
|
||||
})
|
||||
);
|
||||
}
|
||||
let fields = help["metadata"]["fields"]
|
||||
.as_array()
|
||||
.ok_or("missing metadata fields")?;
|
||||
assert_eq!(fields.is_empty(), !populated);
|
||||
assert_eq!(help["metadata"]["truncated"], false);
|
||||
assert!(guide.contains(help["metadata"]["scope"].as_str().ok_or("missing scope")?));
|
||||
assert_eq!(
|
||||
guide.contains("No metadata paths found in the sampled rows"),
|
||||
!populated
|
||||
);
|
||||
if populated {
|
||||
let versions = read_json(&database, "SELECT count() AS count FROM spend_logs").await?;
|
||||
assert_eq!(versions["data"][0]["count"], 2);
|
||||
assert_eq!(help["metadata"]["sampled_rows"], 1);
|
||||
assert!(
|
||||
!fields
|
||||
.iter()
|
||||
.any(|field| field["path"] == serde_json::json!(["obsolete"]))
|
||||
);
|
||||
assert!(
|
||||
columns
|
||||
.iter()
|
||||
.any(|column| column["name"] == "CustomColumn")
|
||||
);
|
||||
assert!(fields.iter().any(|field| field["path"]
|
||||
== serde_json::json!(["labels", "priority"])
|
||||
&& field["types"] == serde_json::json!(["integer"])));
|
||||
assert!(
|
||||
fields
|
||||
.iter()
|
||||
.any(|field| field["path"] == serde_json::json!(["items", 1, "name"]))
|
||||
);
|
||||
assert!(guide.contains("CustomColumn: String"));
|
||||
assert!(guide.contains("JSONExtractRaw(metadata, '<custom>&{{key}}', 'nested.key')"));
|
||||
assert!(guide.contains("SpanAttributes['custom.tag']"));
|
||||
assert!(guide.contains("ResourceAttributes['custom.resource']"));
|
||||
assert_eq!(help["attributes"][0]["fields"][0]["key"], "custom.tag");
|
||||
assert_eq!(help["attributes"][1]["fields"][0]["key"], "custom.resource");
|
||||
for field in fields {
|
||||
let expression = field["expression"].as_str().ok_or("missing expression")?;
|
||||
assert!(
|
||||
guide.contains(expression),
|
||||
"missing plain-text expression: {expression}"
|
||||
);
|
||||
let sql = format!("SELECT {expression} AS value FROM spend_logs FINAL");
|
||||
let body = litellm_traces::query_sql(&database.client, &reader, &sql).await?;
|
||||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_ne!(values["data"][0]["value"], "");
|
||||
}
|
||||
}
|
||||
for example in help["examples"].as_array().ok_or("missing examples")? {
|
||||
let sql = example["sql"].as_str().ok_or("missing example SQL")?;
|
||||
assert!(guide.contains(example["name"].as_str().ok_or("missing example name")?));
|
||||
assert!(guide.contains(sql));
|
||||
assert_eq!(
|
||||
example
|
||||
.as_object()
|
||||
.ok_or("example object")?
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect::<std::collections::BTreeSet<_>>(),
|
||||
std::collections::BTreeSet::from(["name", "sql"])
|
||||
);
|
||||
let body = litellm_traces::query_sql(&database.client, &reader, sql).await?;
|
||||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_eq!(
|
||||
values["data"].as_array().ok_or("missing data")?.is_empty(),
|
||||
!populated,
|
||||
"{sql}"
|
||||
);
|
||||
if populated && example["name"] == "Traces correlated with LLM call metadata" {
|
||||
assert_eq!(values["data"][0]["TraceId"], "trace-1");
|
||||
assert_eq!(values["data"][0]["spend"], 0.25);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::metadata(2, 1)]
|
||||
#[case::attributes(1, 2)]
|
||||
#[case::all(2, 2)]
|
||||
#[tokio::test]
|
||||
async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] spend_rows: usize,
|
||||
#[case] span_rows: usize,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7, 14).await?;
|
||||
execute_write(
|
||||
&database,
|
||||
"CREATE USER help_reader SETTINGS max_rows_to_read = 1",
|
||||
)
|
||||
.await?;
|
||||
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
|
||||
execute_write(
|
||||
&database,
|
||||
&format!("GRANT SELECT ON trace_test.{table} TO help_reader"),
|
||||
)
|
||||
.await?;
|
||||
}
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let spend = (0..spend_rows)
|
||||
.map(|index| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"request_id": format!("request-{index}"), "start_time": timestamp / 1_000_000,
|
||||
"end_time": timestamp / 1_000_000, "metadata": r#"{"custom":{"enabled":true}}"#
|
||||
}))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "spend_logs", spend).await?;
|
||||
let spans = (0..span_rows).map(|index| serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "trace", "SpanId": format!("span-{index}"),
|
||||
"SpanAttributes": {"custom.span": "value"}, "ResourceAttributes": {"custom.resource": "value"}
|
||||
}))).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::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());
|
||||
assert_eq!(
|
||||
help["normalized_fields"]
|
||||
.as_array()
|
||||
.ok_or("normalized fields")?
|
||||
.len(),
|
||||
NORMALIZED_FIELD_DEFINITIONS.len()
|
||||
);
|
||||
let guide = help["guide"].as_str().ok_or("guide")?;
|
||||
assert!(guide.contains("TraceId: String"));
|
||||
assert_eq!(
|
||||
guide.contains("Metadata discovery unavailable:"),
|
||||
spend_rows > 1
|
||||
);
|
||||
assert_eq!(
|
||||
guide.contains("Attribute discovery unavailable:"),
|
||||
span_rows > 1
|
||||
);
|
||||
for (catalog, unavailable) in [
|
||||
(&help["metadata"], spend_rows > 1),
|
||||
(&help["attributes"][0], span_rows > 1),
|
||||
(&help["attributes"][1], span_rows > 1),
|
||||
] {
|
||||
assert_eq!(catalog.get("error").is_some(), unavailable);
|
||||
assert_eq!(catalog["truncated"], unavailable);
|
||||
assert_eq!(
|
||||
catalog["fields"].as_array().ok_or("fields")?.is_empty(),
|
||||
unavailable
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
268
litellm-rust/crates/traces/tests/query_access.rs
Normal file
268
litellm-rust/crates/traces/tests/query_access.rs
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::{
|
||||
Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
struct Database {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
client: Client,
|
||||
writer: Connection,
|
||||
readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e",
|
||||
)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let writer = Connection::parse(&format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
))?;
|
||||
let client = Client::no_redirect_for_test();
|
||||
ensure_schema(&client, &writer, "trace_test", 7, 7).await?;
|
||||
for sql in [
|
||||
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a')), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a')), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'))",
|
||||
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}')",
|
||||
"CREATE TABLE trace_test.private_data (secret String) ENGINE = Memory",
|
||||
"INSERT INTO trace_test.private_data VALUES ('hidden')",
|
||||
] {
|
||||
let response = client.post(writer.url().clone()).body(sql).send().await?;
|
||||
assert!(response.status().is_success(), "{}", response.text().await?);
|
||||
}
|
||||
let readers = QueryReaders::new(writer.clone(), "trace_test".to_owned());
|
||||
Ok(Database {
|
||||
_container: container,
|
||||
client,
|
||||
writer,
|
||||
readers,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::team(QueryScope::Team { team_id: "team-a".to_owned() }, vec!["a1", "a2"])]
|
||||
#[case::project_key(QueryScope::Key { team_id: "team-a".to_owned(), api_key_hash: "key-a1".to_owned() }, vec!["a1"])]
|
||||
#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b"])]
|
||||
#[case::quoted_team(QueryScope::Team { team_id: "team-a' OR 1=1 --\\".to_owned() }, vec![])]
|
||||
#[tokio::test]
|
||||
async fn queries_and_help_are_scoped_by_the_database(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
#[case] scope: QueryScope,
|
||||
#[case] expected: Vec<&str>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
let queries = [
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
"SELECT SpanId AS id FROM trace_test.otel_traces WHERE 1 = 1 ORDER BY id",
|
||||
"SELECT SpanId AS id FROM merge('trace_test', '^otel_traces$') ORDER BY id",
|
||||
"WITH source AS (SELECT * FROM trace_test.otel_traces) SELECT SpanId AS id FROM source ORDER BY id",
|
||||
"SELECT SpanId AS id FROM otel_traces UNION DISTINCT SELECT SpanId AS id FROM trace_test.otel_traces ORDER BY id",
|
||||
"SELECT t.SpanId AS id FROM otel_traces t INNER JOIN spend_logs s ON t.SpanId = s.request_id ORDER BY id",
|
||||
"SELECT request_id AS id FROM spend_logs FINAL ORDER BY id",
|
||||
];
|
||||
for sql in queries {
|
||||
let body: Value = serde_json::from_str(&query_sql(&database.client, &reader, sql).await?)?;
|
||||
assert_eq!(
|
||||
body["data"],
|
||||
json!(
|
||||
expected
|
||||
.iter()
|
||||
.map(|id| json!({"id": id}))
|
||||
.collect::<Vec<_>>()
|
||||
),
|
||||
"{sql}"
|
||||
);
|
||||
}
|
||||
let summary: Value = serde_json::from_str(
|
||||
&query_sql(
|
||||
&database.client,
|
||||
&reader,
|
||||
"SELECT sum(SpanCount) AS count FROM agent_traces_by_key",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(summary["data"][0]["count"], json!(expected.len()));
|
||||
let help = 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());
|
||||
let repeated = recreated
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
assert_eq!(reader.url(), repeated.url());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rotating_master_secret_revokes_previous_reader_credentials(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
};
|
||||
let old_reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "old-master-secret")
|
||||
.await?;
|
||||
let old_result = query_sql(
|
||||
&database.client,
|
||||
&old_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await?;
|
||||
let old_rows: Value = serde_json::from_str(&old_result)?;
|
||||
assert_eq!(old_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
|
||||
|
||||
let rotated_readers = QueryReaders::new(database.writer.clone(), "trace_test".into());
|
||||
let new_reader = rotated_readers
|
||||
.connection(&database.client, &scope, "new-master-secret")
|
||||
.await?;
|
||||
assert!(
|
||||
query_sql(
|
||||
&database.client,
|
||||
&old_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let new_result = query_sql(
|
||||
&database.client,
|
||||
&new_reader,
|
||||
"SELECT SpanId AS id FROM otel_traces ORDER BY id",
|
||||
)
|
||||
.await?;
|
||||
let new_rows: Value = serde_json::from_str(&new_result)?;
|
||||
assert_eq!(new_rows["data"], json!([{ "id": "a1" }, { "id": "a2" }]));
|
||||
assert_eq!(old_reader.url().username(), new_reader.url().username());
|
||||
assert_ne!(old_reader.url().password(), new_reader.url().password());
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn managed_reader_rejects_privilege_and_scope_bypasses(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let scope = QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
};
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &scope, "test-master-secret")
|
||||
.await?;
|
||||
for sql in [
|
||||
"INSERT INTO otel_traces (TraceId) VALUES ('injected')",
|
||||
"DROP TABLE otel_traces",
|
||||
"SELECT * FROM private_data",
|
||||
"SELECT * FROM otel_traces SETTINGS readonly = 0",
|
||||
"SELECT * FROM otel_traces SETTINGS max_memory_usage = 0",
|
||||
"SELECT * FROM otel_traces SETTINGS max_execution_time = 0",
|
||||
"CREATE USER scope_bypass",
|
||||
"CREATE NAMED COLLECTION scope_bypass AS host = 'localhost'",
|
||||
"BACKUP TABLE otel_traces TO Disk('default', 'scope-bypass')",
|
||||
"SELECT * FROM url('http://127.0.0.1:1/', 'LineAsString', 'line String')",
|
||||
"SELECT * FROM remote('127.0.0.1', 'trace_test', 'otel_traces')",
|
||||
] {
|
||||
assert!(
|
||||
matches!(
|
||||
query_sql(&database.client, &reader, sql).await,
|
||||
Err(Error::QueryFailed(_))
|
||||
),
|
||||
"{sql}"
|
||||
);
|
||||
}
|
||||
let roles: Value = serde_json::from_str(
|
||||
&query_sql(&database.client, &reader, "SELECT enabledRoles() AS roles").await?,
|
||||
)?;
|
||||
assert_eq!(roles["data"], json!([{ "roles": [] }]));
|
||||
let rows: Value = serde_json::from_str(
|
||||
&query_sql(
|
||||
&database.client,
|
||||
&reader,
|
||||
"SELECT DISTINCT TeamId FROM otel_traces",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(rows["data"], json!([{ "TeamId": "team-a" }]));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn provisioning_failure_never_returns_a_writer_connection(
|
||||
#[future(awt)] database: Result<Database, Box<dyn std::error::Error>>,
|
||||
) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let database = database?;
|
||||
let reader = database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::Admin, "test-master-secret")
|
||||
.await?;
|
||||
let no_provision_privileges = QueryReaders::new(reader, "trace_test".to_owned());
|
||||
let result = no_provision_privileges
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Team {
|
||||
team_id: "team-a".to_owned(),
|
||||
},
|
||||
"other-secret",
|
||||
)
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
assert!(
|
||||
database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::Admin, "")
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
database
|
||||
.readers
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Team {
|
||||
team_id: String::new()
|
||||
},
|
||||
"test-master-secret"
|
||||
)
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
let permits = (0..8)
|
||||
.map(|_| database.readers.acquire())
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
assert!(database.readers.acquire().is_err());
|
||||
drop(permits);
|
||||
assert!(database.readers.acquire().is_ok());
|
||||
let rows = litellm_traces::execute_read(
|
||||
&database.client,
|
||||
&database.writer,
|
||||
"SELECT count() AS count FROM trace_test.otel_traces",
|
||||
&BTreeMap::new(),
|
||||
)
|
||||
.await?;
|
||||
let rows: Value = serde_json::from_str(&rows)?;
|
||||
assert_eq!(rows["data"][0]["count"], 3);
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -537,6 +537,8 @@ class LiteLLMRoutes(enum.Enum):
|
|||
"/lens/workers/register",
|
||||
"/lens/workers/{worker_id}",
|
||||
"/v1/traces",
|
||||
"/v1/traces/query",
|
||||
"/v1/traces/query/help",
|
||||
"/v1/traces/{trace_id}",
|
||||
"/v1/traces/{trace_id}/spans/{span_id}",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -15,12 +15,15 @@ from types import MappingProxyType
|
|||
from typing import Annotated, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import OTLP_RETRY_AFTER_SECONDS
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.http_parsing_utils import is_otlp_trace_request
|
||||
from litellm.proxy.tracing_runtime import provide_receiver, require_receiver
|
||||
from litellm.rust_bridge.traces import ClickHouseStorage, QueryScope
|
||||
from litellm.tracing import (
|
||||
Tenant,
|
||||
TraceReceiver,
|
||||
|
|
@ -137,6 +140,77 @@ async def list_agent_traces(
|
|||
raise HTTPException(status_code=400, detail=str(error)) from error
|
||||
|
||||
|
||||
class TraceQueryRequest(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
sql: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TraceQueryAccess:
|
||||
storage: ClickHouseStorage
|
||||
scope: QueryScope
|
||||
secret: str
|
||||
|
||||
|
||||
def provide_trace_query_secret() -> str:
|
||||
from litellm.proxy.proxy_server import master_key
|
||||
|
||||
if not master_key:
|
||||
raise HTTPException(status_code=503, detail="Trace SQL queries require a configured proxy master key")
|
||||
return master_key
|
||||
|
||||
|
||||
def trace_query_scope(auth: UserAPIKeyAuth) -> QueryScope:
|
||||
if auth.user_role in (LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY):
|
||||
return {"kind": "admin"}
|
||||
if auth.project_id and auth.token:
|
||||
return {"kind": "key", "team_id": auth.team_id or "", "api_key_hash": auth.token}
|
||||
if auth.project_id:
|
||||
raise HTTPException(status_code=403, detail="Project trace SQL queries require a project key")
|
||||
if auth.team_id:
|
||||
return {"kind": "team", "team_id": auth.team_id}
|
||||
if auth.token:
|
||||
return {"kind": "key", "team_id": "", "api_key_hash": auth.token}
|
||||
raise HTTPException(status_code=403, detail="Trace SQL queries require an authenticated trace scope")
|
||||
|
||||
|
||||
async def provide_trace_query_access(
|
||||
auth: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
tracing: Annotated[TraceReceiver | None, Depends(provide_receiver)],
|
||||
secret: Annotated[str, Depends(provide_trace_query_secret)],
|
||||
) -> TraceQueryAccess:
|
||||
return TraceQueryAccess(require_receiver(tracing).store.storage, trace_query_scope(auth), secret)
|
||||
|
||||
|
||||
@router.post("/v1/traces/query")
|
||||
async def query_agent_traces(
|
||||
body: TraceQueryRequest,
|
||||
access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)],
|
||||
) -> Response:
|
||||
try:
|
||||
return Response(
|
||||
content=await access.storage.query_sql(body.sql, access.scope, access.secret), media_type="application/json"
|
||||
)
|
||||
except ValueError as error:
|
||||
raise HTTPException(status_code=400, detail=str(error)) from error
|
||||
except RuntimeError as error:
|
||||
verbose_proxy_logger.warning("Trace SQL query unavailable: %s", error)
|
||||
raise HTTPException(status_code=503, detail="Trace SQL query failed or exceeded reader limits") from error
|
||||
|
||||
|
||||
@router.get("/v1/traces/query/help")
|
||||
async def help_agent_trace_queries(
|
||||
access: Annotated[TraceQueryAccess, Depends(provide_trace_query_access)],
|
||||
) -> Response:
|
||||
try:
|
||||
return Response(
|
||||
content=await access.storage.query_help(access.scope, access.secret), media_type="application/json"
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@router.get("/v1/traces/{trace_id}", response_model=Trace)
|
||||
async def get_agent_trace(
|
||||
trace_id: str,
|
||||
|
|
|
|||
|
|
@ -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.traces import DecodedSpan
|
||||
from litellm.rust_bridge.traces import DecodedSpan, QueryScope
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.utils import EmbeddingResponse, ModelResponse
|
||||
|
|
@ -30,6 +30,8 @@ class NativeTraceStorage:
|
|||
def __new__(cls, database: str, url: str, reader_url: str | None = None) -> NativeTraceStorage: ...
|
||||
def ensure_schema(self, trace_retention_days: int, spend_log_retention_days: int) -> 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 lens_query(self, name: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Future[str]: ...
|
||||
def query(self, query: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Future[str]: ...
|
||||
|
||||
|
|
|
|||
|
|
@ -58,6 +58,24 @@ class DecodedSpan(TypedDict):
|
|||
ReadQueryName = Literal["list_traces", "trace_spans", "span_detail", "span_error", "spend_by_response_ids"]
|
||||
|
||||
|
||||
class AdminQueryScope(TypedDict):
|
||||
kind: ReadOnly[Literal["admin"]]
|
||||
|
||||
|
||||
class TeamQueryScope(TypedDict):
|
||||
kind: ReadOnly[Literal["team"]]
|
||||
team_id: ReadOnly[str]
|
||||
|
||||
|
||||
class KeyQueryScope(TypedDict):
|
||||
kind: ReadOnly[Literal["key"]]
|
||||
team_id: ReadOnly[str]
|
||||
api_key_hash: ReadOnly[str]
|
||||
|
||||
|
||||
QueryScope = AdminQueryScope | TeamQueryScope | KeyQueryScope
|
||||
|
||||
|
||||
class NativeStore(Protocol):
|
||||
def __init__(self, database: str, url: str, reader_url: str | None = None) -> None: ...
|
||||
|
||||
|
|
@ -65,6 +83,10 @@ class NativeStore(Protocol):
|
|||
|
||||
def insert_rows(self, table: str, rows: Sequence[Mapping[str, object]]) -> Awaitable[None]: ...
|
||||
|
||||
def query_sql(self, sql: str, scope: QueryScope, secret: str) -> Awaitable[str]: ...
|
||||
|
||||
def query_help(self, scope: QueryScope, secret: str) -> Awaitable[str]: ...
|
||||
|
||||
def lens_query(self, name: str, parameters: Mapping[str, str | int | Sequence[str]]) -> Awaitable[str]: ...
|
||||
|
||||
def query(self, name: ReadQueryName, parameters: Mapping[str, str | int | Sequence[str]]) -> Awaitable[str]: ...
|
||||
|
|
@ -138,6 +160,12 @@ class ClickHouseStorage:
|
|||
)
|
||||
return QueryResponse.model_validate_json(result).data
|
||||
|
||||
async def query_sql(self, sql: str, scope: QueryScope, secret: str) -> str:
|
||||
return await self._native.query_sql(sql, scope, secret)
|
||||
|
||||
async def query_help(self, scope: QueryScope, secret: str) -> str:
|
||||
return await self._native.query_help(scope, secret)
|
||||
|
||||
async def _lens_query(self, name: str, parameters: Mapping[str, object]) -> list[dict[str, JsonValue]]:
|
||||
result: Final = await self._native.lens_query(name, QUERY_PARAMETERS.validate_python(parameters))
|
||||
return QueryResponse.model_validate_json(result).data
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ class TraceReceiver:
|
|||
ClickHouseStorage(
|
||||
database=os.getenv("CLICKHOUSE_DATABASE", "litellm"),
|
||||
url=os.environ["CLICKHOUSE_URL"],
|
||||
reader_url=os.environ["CLICKHOUSE_READER_URL"],
|
||||
reader_url=os.getenv("CLICKHOUSE_READER_URL", os.environ["CLICKHOUSE_URL"]),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -207,3 +207,124 @@ async def test_insert_validates_values_without_pydantic_copy(recording_server: R
|
|||
assert stored["Timestamp"] == "1970-01-01T00:00:00.000000001Z"
|
||||
assert stored["ResourceAttributes"] == attributes
|
||||
assert stored["SpanAttributes"] == attributes
|
||||
|
||||
|
||||
@pytest.mark.parametrize("role", ["proxy_admin", "proxy_admin_viewer", "internal_user"])
|
||||
def test_trace_sql_endpoint_executes_for_admin_and_preserves_clickhouse_envelope(
|
||||
recording_server: RecordingServer, role: str
|
||||
) -> None:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router
|
||||
|
||||
envelope: Final = {"meta": [{"name": "answer", "type": "UInt8"}], "data": [{"answer": 42}], "rows": 1}
|
||||
recording_server.expected_requests = 12
|
||||
for _ in range(11):
|
||||
recording_server.enqueue(ResponseSpec(body=""))
|
||||
recording_server.enqueue(ResponseSpec(body=envelope))
|
||||
storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url)
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret"
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=role, token="test")
|
||||
app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage))
|
||||
with TestClient(app) as client:
|
||||
result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 42 AS answer"})
|
||||
assert result.status_code == 200, result.text
|
||||
assert result.json() == envelope
|
||||
assert recording_server.requests[-1].raw_body == b"SELECT 42 AS answer"
|
||||
assert client.post("/v1/traces/query", json={"sql": " "}).status_code == 400
|
||||
assert client.post("/v1/traces/query", json={}).status_code == 422
|
||||
|
||||
|
||||
def test_trace_help_endpoint_runs_native_schema_and_metadata_discovery(recording_server: RecordingServer) -> None:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router
|
||||
|
||||
recording_server.expected_requests = 17
|
||||
for _ in range(11):
|
||||
recording_server.enqueue(ResponseSpec(body=""))
|
||||
for response in (
|
||||
{"data": [{"name": "Model", "type": "String"}]},
|
||||
{"data": []},
|
||||
{"data": []},
|
||||
{"data": [{"metadata": '{"custom": {"label": "hello"}}'}]},
|
||||
{"data": [{"key": "custom.span"}]},
|
||||
{"data": [{"key": "custom.resource"}]},
|
||||
):
|
||||
recording_server.enqueue(ResponseSpec(body=response))
|
||||
storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url)
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret"
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin", token="test")
|
||||
app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage))
|
||||
with TestClient(app) as client:
|
||||
result: Final = client.get("/v1/traces/query/help")
|
||||
assert result.status_code == 200, result.text
|
||||
body: Final = result.json()
|
||||
assert body["guide"].startswith("Trace SQL query guide")
|
||||
assert "JSONExtractRaw(metadata, 'custom', 'label')" in body["guide"]
|
||||
assert body["tables"][0]["columns"] == [{"name": "Model", "type": "String"}]
|
||||
assert body["metadata"]["fields"][1] == {
|
||||
"path": ["custom", "label"],
|
||||
"types": ["string"],
|
||||
"expression": "JSONExtractRaw(metadata, 'custom', 'label')",
|
||||
}
|
||||
assert body["attributes"][0]["fields"][0]["expression"] == "SpanAttributes['custom.span']"
|
||||
assert body["attributes"][1]["fields"][0]["expression"] == "ResourceAttributes['custom.resource']"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("clickhouse_status, expected_status", [(400, 400), (404, 400), (500, 503), (503, 503)])
|
||||
def test_trace_sql_endpoint_distinguishes_query_errors_from_reader_failures(
|
||||
recording_server: RecordingServer, clickhouse_status: int, expected_status: int
|
||||
) -> None:
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.tracing_endpoints import provide_receiver, provide_trace_query_secret, router
|
||||
|
||||
recording_server.expected_requests = 13
|
||||
for _ in range(11):
|
||||
recording_server.enqueue(ResponseSpec(body=""))
|
||||
recording_server.enqueue(ResponseSpec(status=clickhouse_status, body=b"ClickHouse rejected the query"))
|
||||
envelope: Final = {"meta": [{"name": "answer", "type": "UInt8"}], "data": [{"answer": 42}], "rows": 1}
|
||||
recording_server.enqueue(ResponseSpec(body=envelope))
|
||||
storage: Final = ClickHouseStorage("trace_test", recording_server.base_url, recording_server.base_url)
|
||||
app: Final = FastAPI()
|
||||
app.include_router(router)
|
||||
app.dependency_overrides[provide_trace_query_secret] = lambda: "test-master-secret"
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin", token="test")
|
||||
app.dependency_overrides[provide_receiver] = lambda: TraceReceiver(TraceStore(storage))
|
||||
with TestClient(app) as client:
|
||||
failed: Final = client.post("/v1/traces/query", json={"sql": "SELEC 42"})
|
||||
assert failed.status_code == expected_status, failed.text
|
||||
recovered: Final = client.post("/v1/traces/query", json={"sql": "SELECT 42 AS answer"})
|
||||
assert recovered.status_code == 200, recovered.text
|
||||
assert recovered.json() == envelope
|
||||
assert recording_server.requests[-2].raw_body == b"SELEC 42"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_trace_receiver_reads_with_only_one_clickhouse_url(
|
||||
recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("CLICKHOUSE_URL", recording_server.base_url)
|
||||
monkeypatch.setenv("CLICKHOUSE_DATABASE", "trace_test")
|
||||
monkeypatch.delenv("CLICKHOUSE_READER_URL", raising=False)
|
||||
recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]}))
|
||||
receiver: Final = TraceReceiver.from_env()
|
||||
rows: Final = await receiver.store.storage.query("trace_spans", {"trace_id": "trace-1"})
|
||||
assert rows == [{"trace_id": "trace-1"}]
|
||||
parameters: Final = parse_qs(urlsplit(recording_server.requests[0].path).query)
|
||||
assert parameters["database"] == ["trace_test"]
|
||||
assert parameters["readonly"] == ["1"]
|
||||
|
|
|
|||
|
|
@ -141,7 +141,9 @@ def test_501_when_tracing_not_enabled(
|
|||
assert response.status_code == 501
|
||||
assert response.headers["content-type"] == "application/x-protobuf"
|
||||
assert Status.FromString(response.content).message == (
|
||||
"Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL." if native_available else ""
|
||||
"Agent tracing is not enabled. Set `tracing:` in general_settings and CLICKHOUSE_URL."
|
||||
if native_available
|
||||
else ""
|
||||
)
|
||||
assert client.get("/v1/traces").status_code == 501
|
||||
|
||||
|
|
@ -525,3 +527,96 @@ def test_lens_reads_from_injected_storage_without_receiver() -> None:
|
|||
assert response.status_code == 200, response.text
|
||||
assert response.json()["executions"] == []
|
||||
storage.lens_sample.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("auth", "expected_scope"),
|
||||
(
|
||||
(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), {"kind": "admin"}),
|
||||
(UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY), {"kind": "admin"}),
|
||||
(TEAM_KEY, {"kind": "team", "team_id": "team-research"}),
|
||||
(
|
||||
UserAPIKeyAuth(token="project-key", team_id="team-a", project_id="project-a"),
|
||||
{"kind": "key", "team_id": "team-a", "api_key_hash": "project-key"},
|
||||
),
|
||||
(UserAPIKeyAuth(token="solo-key"), {"kind": "key", "team_id": "", "api_key_hash": "solo-key"}),
|
||||
),
|
||||
)
|
||||
def test_sql_and_help_use_authenticated_scope(
|
||||
client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth, expected_scope: dict[str, str]
|
||||
) -> None:
|
||||
client.app.dependency_overrides[user_api_key_auth] = lambda: auth
|
||||
client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
|
||||
receiver.store.storage.query_sql = AsyncMock(return_value='{"data":[{"value":1}]}')
|
||||
receiver.store.storage.query_help = AsyncMock(return_value='{"guide":"scoped"}')
|
||||
result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
|
||||
assert result.status_code == 200, result.text
|
||||
assert result.json() == {"data": [{"value": 1}]}
|
||||
receiver.store.storage.query_sql.assert_awaited_once_with(
|
||||
"SELECT * FROM otel_traces", expected_scope, "test-secret"
|
||||
)
|
||||
help_result: Final = client.get("/v1/traces/query/help")
|
||||
assert help_result.status_code == 200, help_result.text
|
||||
assert help_result.json() == {"guide": "scoped"}
|
||||
receiver.store.storage.query_help.assert_awaited_once_with(expected_scope, "test-secret")
|
||||
forged: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1", "scope": {"kind": "admin"}})
|
||||
assert forged.status_code == 422, forged.text
|
||||
assert receiver.store.storage.query_sql.await_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("auth", (UserAPIKeyAuth(), UserAPIKeyAuth(team_id="a", project_id="p")))
|
||||
def test_sql_rejects_missing_identity_without_querying(
|
||||
client: TestClient, receiver: MagicMock, auth: UserAPIKeyAuth
|
||||
) -> None:
|
||||
client.app.dependency_overrides[user_api_key_auth] = lambda: auth
|
||||
client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
|
||||
result: Final = client.post("/v1/traces/query", json={"sql": "SELECT * FROM otel_traces"})
|
||||
assert result.status_code == 403, result.text
|
||||
assert client.get("/v1/traces/query/help").status_code == 403
|
||||
receiver.store.storage.query_sql.assert_not_called()
|
||||
receiver.store.storage.query_help.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "status"), ((ValueError("invalid SQL"), 400), (RuntimeError("reader unavailable"), 503))
|
||||
)
|
||||
def test_sql_reports_rejected_queries_and_unavailable_readers(
|
||||
client: TestClient, receiver: MagicMock, error: Exception, status: int
|
||||
) -> None:
|
||||
client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
|
||||
receiver.store.storage.query_sql = AsyncMock(side_effect=error)
|
||||
result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"})
|
||||
assert result.status_code == status, result.text
|
||||
receiver.store.storage.query_sql.assert_awaited_once_with(
|
||||
"SELECT 1", {"kind": "team", "team_id": "team-research"}, "test-secret"
|
||||
)
|
||||
|
||||
|
||||
def test_query_help_does_not_fall_back_when_reader_provisioning_fails(client: TestClient, receiver: MagicMock) -> None:
|
||||
client.app.dependency_overrides[tracing_endpoints.provide_trace_query_secret] = lambda: "test-secret"
|
||||
receiver.store.storage.query_help = AsyncMock(side_effect=RuntimeError("reader provisioning failed"))
|
||||
result: Final = client.get("/v1/traces/query/help")
|
||||
assert result.status_code == 503, result.text
|
||||
receiver.store.storage.query_help.assert_awaited_once_with(
|
||||
{"kind": "team", "team_id": "team-research"}, "test-secret"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("secret", (None, "configured-master-key"))
|
||||
def test_queries_require_a_proxy_secret(
|
||||
client: TestClient, receiver: MagicMock, monkeypatch: pytest.MonkeyPatch, secret: str | None
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "master_key", secret)
|
||||
receiver.store.storage.query_sql = AsyncMock(return_value='{"data":[]}')
|
||||
result: Final = client.post("/v1/traces/query", json={"sql": "SELECT 1"})
|
||||
if secret is None:
|
||||
assert result.status_code == 503, result.text
|
||||
assert "master key" in result.json()["detail"]
|
||||
receiver.store.storage.query_sql.assert_not_awaited()
|
||||
return
|
||||
assert result.status_code == 200, result.text
|
||||
receiver.store.storage.query_sql.assert_awaited_once_with(
|
||||
"SELECT 1", {"kind": "team", "team_id": "team-research"}, secret
|
||||
)
|
||||
|
|
|
|||
92
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
92
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -22357,6 +22357,40 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/traces/query": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
get?: never;
|
||||
put?: never;
|
||||
/** Query Agent Traces */
|
||||
post: operations["query_agent_traces_v1_traces_query_post"];
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/traces/query/help": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/** Help Agent Trace Queries */
|
||||
get: operations["help_agent_trace_queries_v1_traces_query_help_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/traces/{trace_id}": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -46804,6 +46838,11 @@ export interface components {
|
|||
*/
|
||||
truncated: boolean;
|
||||
};
|
||||
/** TraceQueryRequest */
|
||||
TraceQueryRequest: {
|
||||
/** Sql */
|
||||
sql: string;
|
||||
};
|
||||
/** TraceSummary */
|
||||
TraceSummary: {
|
||||
/** Agent Count */
|
||||
|
|
@ -79639,6 +79678,59 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
query_agent_traces_v1_traces_query_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody: {
|
||||
content: {
|
||||
"application/json": components["schemas"]["TraceQueryRequest"];
|
||||
};
|
||||
};
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": unknown;
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
help_agent_trace_queries_v1_traces_query_help_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": unknown;
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
get_agent_trace_v1_traces__trace_id__get: {
|
||||
parameters: {
|
||||
query?: {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue