mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(traces): type queries and align read access with log visibility (#44228)
* wip * wip * test(traces): separate root status from diagnostic error counts * test(traces): cover normalization precedence and fallbacks * chore(cache): remove stray comments from trace PR * test(traces): name lens test for shared query path * fix(traces): place query implementation before test module * test(traces): use unified read scope in migration tests * ci(rust): allow feature checks to finish * ci(mcp): allow dependency resolution to finish * fix(traces): preserve key visibility and safe spend attribution
This commit is contained in:
parent
9b5562f89b
commit
688d791fa0
124 changed files with 17105 additions and 1820 deletions
|
|
@ -17,7 +17,7 @@ concurrency:
|
|||
jobs:
|
||||
resolve:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
timeout-minutes: 25
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
|
|
|
|||
2
.github/workflows/test-rust.yml
vendored
2
.github/workflows/test-rust.yml
vendored
|
|
@ -89,7 +89,7 @@ jobs:
|
|||
|
||||
rust-test:
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 20
|
||||
timeout-minutes: 30
|
||||
defaults:
|
||||
run:
|
||||
working-directory: litellm-rust
|
||||
|
|
|
|||
23
litellm-rust/Cargo.lock
generated
23
litellm-rust/Cargo.lock
generated
|
|
@ -4156,6 +4156,7 @@ dependencies = [
|
|||
"litellm-storage-clickhouse",
|
||||
"litellm-token-counter",
|
||||
"litellm-traces",
|
||||
"litellm-traces-clickhouse",
|
||||
"litellm-tracing",
|
||||
"prost",
|
||||
"pyo3",
|
||||
|
|
@ -4369,6 +4370,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"url",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4451,19 +4453,30 @@ dependencies = [
|
|||
name = "litellm-traces"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"base64 0.22.1",
|
||||
"criterion",
|
||||
"indexmap 2.14.0",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "litellm-traces-clickhouse"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"askama",
|
||||
"flate2",
|
||||
"futures-util",
|
||||
"hmac 0.12.1",
|
||||
"indexmap 2.14.0",
|
||||
"litellm-http",
|
||||
"litellm-migrate",
|
||||
"litellm-storage-clickhouse",
|
||||
"litellm-traces",
|
||||
"moka",
|
||||
"opentelemetry-proto",
|
||||
"prost",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ litellm-config = { path = "crates/config" }
|
|||
litellm-router = { path = "crates/router" }
|
||||
litellm-tracing = { path = "crates/tracing" }
|
||||
litellm-traces = { path = "crates/traces" }
|
||||
litellm-traces-clickhouse = { path = "crates/traces-clickhouse" }
|
||||
litellm-storage-clickhouse = { path = "crates/storage-clickhouse" }
|
||||
litellm-migrate = { path = "crates/migrate" }
|
||||
litellm-migrate-macros = { path = "crates/migrate-macros" }
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
|
|||
fancy-regex.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
litellm-traces.workspace = true
|
||||
litellm-traces-clickhouse.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
litellm-host.workspace = true
|
||||
bytes.workspace = true
|
||||
|
|
|
|||
|
|
@ -2,10 +2,8 @@ use std::collections::BTreeMap;
|
|||
|
||||
use litellm_host_python::{FromPythonCache, ToPythonCache};
|
||||
use litellm_http::ClientVariant;
|
||||
use litellm_traces::{
|
||||
Config, Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery,
|
||||
Shared,
|
||||
};
|
||||
use litellm_traces::{QueryScope, ReadQuery, Shared};
|
||||
use litellm_traces_clickhouse::{Config, Error, InsertTable, Parameter, QueryReaders};
|
||||
use prost::Message;
|
||||
use pyo3::{
|
||||
exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
|
||||
|
|
@ -31,38 +29,56 @@ pub fn trace_encode_error<'py>(py: Python<'py>, message: &str) -> Bound<'py, PyB
|
|||
}
|
||||
|
||||
fn map_error(error: Error) -> PyErr {
|
||||
map_error_ref(&error)
|
||||
}
|
||||
|
||||
fn map_error_ref(error: &Error) -> PyErr {
|
||||
use litellm_storage_clickhouse::Error as StorageError;
|
||||
|
||||
match error {
|
||||
Error::InvalidRow
|
||||
| Error::InvalidTable
|
||||
| Error::InvalidSchema
|
||||
| Error::EmptySql
|
||||
| Error::InvalidQuery => PyValueError::new_err(error.to_string()),
|
||||
| Error::InvalidQuery
|
||||
| Error::InvalidParameters
|
||||
| Error::InvalidScope => PyValueError::new_err(error.to_string()),
|
||||
Error::InsertTooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
Error::InvalidUrl
|
||||
| Error::QueryFailed(_)
|
||||
| Error::InsertFailed(_)
|
||||
| Error::SchemaFailed(_)
|
||||
| Error::ResponseTooLarge
|
||||
| Error::InvalidResponse
|
||||
| Error::Transport => PyRuntimeError::new_err(error.to_string()),
|
||||
Error::SchemaFailed(_)
|
||||
| Error::SchemaTransport
|
||||
| Error::MissingSecret
|
||||
| Error::Busy
|
||||
| Error::ProvisionFailed(_)
|
||||
| Error::ProvisionTransport
|
||||
| Error::InvalidResponse => PyRuntimeError::new_err(error.to_string()),
|
||||
Error::Cached(source) => map_error_ref(source),
|
||||
Error::Storage(source) => match source {
|
||||
StorageError::InvalidRow
|
||||
| StorageError::InvalidTable
|
||||
| StorageError::InvalidSchema
|
||||
| StorageError::EmptySql
|
||||
| StorageError::InvalidParameters
|
||||
| StorageError::InvalidQuery => PyValueError::new_err(error.to_string()),
|
||||
StorageError::InsertTooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
StorageError::InvalidUrl
|
||||
| StorageError::QueryFailed(_)
|
||||
| StorageError::InsertFailed(_)
|
||||
| StorageError::SchemaFailed(_)
|
||||
| StorageError::ResponseTooLarge
|
||||
| StorageError::InvalidResponse
|
||||
| StorageError::Transport => PyRuntimeError::new_err(error.to_string()),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn map_sql_error(error: Error) -> PyErr {
|
||||
match error {
|
||||
Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
|
||||
Error::Storage(litellm_storage_clickhouse::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(frozen)]
|
||||
pub struct NativeTraceConfig {
|
||||
inner: Config,
|
||||
|
|
@ -105,7 +121,13 @@ impl NativeTraceStorage {
|
|||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces::ensure_schema(&client, &connection, &database, retention_days).await
|
||||
litellm_traces_clickhouse::ensure_schema(
|
||||
&client,
|
||||
&connection,
|
||||
&database,
|
||||
retention_days,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
|
|
@ -115,7 +137,7 @@ impl NativeTraceStorage {
|
|||
&self,
|
||||
py: Python<'py>,
|
||||
table: &str,
|
||||
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces::InsertRow>,
|
||||
#[pyo3(from_py_with = insert_rows_from_py)] rows: Vec<litellm_traces_clickhouse::InsertRow>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let table = InsertTable::parse(table).map_err(map_error)?;
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
|
|
@ -124,8 +146,14 @@ impl NativeTraceStorage {
|
|||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces::insert_shared_rows(&client, &connection, &database, table, rows)
|
||||
.await
|
||||
litellm_traces_clickhouse::insert_shared_rows(
|
||||
&client,
|
||||
&connection,
|
||||
&database,
|
||||
table,
|
||||
rows,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
|
|
@ -139,7 +167,9 @@ impl NativeTraceStorage {
|
|||
secret: String,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
if sql.trim().is_empty() {
|
||||
return Err(map_error(Error::EmptySql));
|
||||
return Err(map_error(
|
||||
litellm_storage_clickhouse::Error::EmptySql.into(),
|
||||
));
|
||||
}
|
||||
let readers = self.query_readers.clone();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
|
|
@ -148,11 +178,9 @@ impl NativeTraceStorage {
|
|||
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)
|
||||
litellm_traces_clickhouse::query_sql(&client, &connection, &sql).await
|
||||
},
|
||||
map_query_access_error,
|
||||
map_sql_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -169,32 +197,9 @@ impl NativeTraceStorage {
|
|||
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)
|
||||
litellm_traces_clickhouse::query_help(&client, &connection).await
|
||||
},
|
||||
map_query_access_error,
|
||||
)
|
||||
}
|
||||
|
||||
fn lens_query<'py>(
|
||||
&self,
|
||||
py: Python<'py>,
|
||||
name: &str,
|
||||
#[pyo3(from_py_with = litellm_host_python::from_py_argument)] parameters: BTreeMap<
|
||||
String,
|
||||
Parameter,
|
||||
>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces::execute_read(&client, &connection, query.sql(), ¶meters).await
|
||||
},
|
||||
map_error,
|
||||
map_sql_error,
|
||||
)
|
||||
}
|
||||
|
||||
|
|
@ -207,13 +212,20 @@ impl NativeTraceStorage {
|
|||
Parameter,
|
||||
>,
|
||||
) -> PyResult<Bound<'py, PyAny>> {
|
||||
let query = ReadQuery::parse(query).map_err(map_error)?;
|
||||
let query =
|
||||
ReadQuery::parse(query).map_err(|error| PyValueError::new_err(error.to_string()))?;
|
||||
let connection = self.config.storage().reader().clone();
|
||||
let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
|
||||
crate::execution::run_async(
|
||||
py,
|
||||
async move {
|
||||
litellm_traces::execute_named_read(&client, &connection, query, ¶meters).await
|
||||
litellm_traces_clickhouse::execute_named_read(
|
||||
&client,
|
||||
&connection,
|
||||
query,
|
||||
¶meters,
|
||||
)
|
||||
.await
|
||||
},
|
||||
map_error,
|
||||
)
|
||||
|
|
@ -229,13 +241,15 @@ pub fn trace_decode_otlp<'py>(
|
|||
let spans = py
|
||||
.detach(|| litellm_traces::decode_otlp(body, content_type))
|
||||
.map_err(|error| match error {
|
||||
litellm_traces::DecodeError::TooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
litellm_traces::Error::TooLarge => PyOverflowError::new_err(error.to_string()),
|
||||
_ => PyValueError::new_err(error.to_string()),
|
||||
})?;
|
||||
spans_to_py(py, &spans).map(Bound::into_any)
|
||||
}
|
||||
|
||||
fn insert_rows_from_py(value: &Bound<'_, PyAny>) -> PyResult<Vec<litellm_traces::InsertRow>> {
|
||||
fn insert_rows_from_py(
|
||||
value: &Bound<'_, PyAny>,
|
||||
) -> PyResult<Vec<litellm_traces_clickhouse::InsertRow>> {
|
||||
let mut resources = FromPythonCache::default();
|
||||
value
|
||||
.try_iter()?
|
||||
|
|
@ -320,6 +334,54 @@ mod tests {
|
|||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::row(Error::InvalidRow, "ValueError")]
|
||||
#[case::insert_budget(Error::InsertTooLarge, "OverflowError")]
|
||||
#[case::scope(Error::InvalidScope, "ValueError")]
|
||||
#[case::schema(Error::SchemaFailed(503), "RuntimeError")]
|
||||
#[case::reader(Error::MissingSecret, "RuntimeError")]
|
||||
#[case::storage(
|
||||
Error::Storage(litellm_storage_clickhouse::Error::InvalidUrl),
|
||||
"RuntimeError"
|
||||
)]
|
||||
#[case::cached_scope(Error::Cached(std::sync::Arc::new(Error::InvalidScope)), "ValueError")]
|
||||
fn trace_failures_preserve_public_exception_types(
|
||||
#[case] error: Error,
|
||||
#[case] exception_name: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let message = error.to_string();
|
||||
let exception = map_error(error);
|
||||
assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
|
||||
assert_eq!(
|
||||
exception.value(py).str().unwrap().to_str().unwrap(),
|
||||
message
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::invalid_sql(400, "ValueError")]
|
||||
#[case::missing_table(404, "ValueError")]
|
||||
#[case::unavailable(503, "RuntimeError")]
|
||||
fn wrapped_query_status_preserves_public_exception_type(
|
||||
#[case] status: u16,
|
||||
#[case] exception_name: &str,
|
||||
) {
|
||||
Python::initialize();
|
||||
Python::attach(|py| {
|
||||
let error = Error::Storage(litellm_storage_clickhouse::Error::QueryFailed(status));
|
||||
let message = error.to_string();
|
||||
let exception = map_sql_error(error);
|
||||
assert_eq!(exception.get_type(py).name().unwrap(), exception_name);
|
||||
assert_eq!(
|
||||
exception.value(py).str().unwrap().to_str().unwrap(),
|
||||
message
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn insert_projection_preserves_identity_without_merging_equal_resources() {
|
||||
Python::initialize();
|
||||
|
|
@ -365,5 +427,6 @@ mod tests {
|
|||
|
||||
#[pyfunction]
|
||||
pub fn trace_normalized_field_definitions<'py>(py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
|
||||
litellm_host_python::Pythonized(litellm_traces::NORMALIZED_FIELD_DEFINITIONS).into_pyobject(py)
|
||||
litellm_host_python::Pythonized(litellm_traces_clickhouse::NORMALIZED_FIELD_DEFINITIONS)
|
||||
.into_pyobject(py)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,4 +2,4 @@
|
|||
|
||||
`litellm-storage-clickhouse` exports `Storage`, a writer and bounded reader derived from one ClickHouse URL and database. It also exports bounded HTTP read and insert execution
|
||||
|
||||
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces` supplies those rules and uses this storage for both trace rows and spend rows
|
||||
The crate has no trace tables, OTLP types, or named trace queries. `litellm-traces-clickhouse` supplies those rules and uses this storage for both trace rows and spend rows
|
||||
|
|
@ -18,3 +18,4 @@ url.workspace = true
|
|||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -10,6 +10,8 @@ pub enum Error {
|
|||
InvalidSchema,
|
||||
#[error("SQL query must not be empty")]
|
||||
EmptySql,
|
||||
#[error("invalid ClickHouse query parameters")]
|
||||
InvalidParameters,
|
||||
#[error("unknown ClickHouse read query")]
|
||||
InvalidQuery,
|
||||
#[error("ClickHouse query failed with HTTP status {0}")]
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ mod read;
|
|||
|
||||
pub use error::Error;
|
||||
pub use insert::{insert_compressed_rows, insert_encoded_rows};
|
||||
pub use read::{Parameter, execute_read};
|
||||
pub use read::{Parameter, Query, execute_read, fetch, fetch_json};
|
||||
use url::Url;
|
||||
|
||||
#[derive(Clone)]
|
||||
|
|
|
|||
|
|
@ -1,17 +1,19 @@
|
|||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use litellm_http::Client;
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize, de::DeserializeOwned};
|
||||
|
||||
use crate::{Connection, Error};
|
||||
|
||||
const MAX_RESPONSE_BYTES: usize = 4 * 1024 * 1024;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
#[serde(untagged)]
|
||||
pub enum Parameter {
|
||||
Text(String),
|
||||
Integer(i64),
|
||||
Unsigned(u64),
|
||||
Float(f64),
|
||||
Strings(Vec<String>),
|
||||
}
|
||||
|
||||
|
|
@ -20,6 +22,8 @@ impl Parameter {
|
|||
match self {
|
||||
Self::Text(value) => escaped(value),
|
||||
Self::Integer(value) => value.to_string(),
|
||||
Self::Unsigned(value) => value.to_string(),
|
||||
Self::Float(value) => value.to_string(),
|
||||
Self::Strings(values) => format!(
|
||||
"[{}]",
|
||||
values
|
||||
|
|
@ -111,3 +115,45 @@ pub async fn execute_read(
|
|||
}
|
||||
String::from_utf8(body).map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
|
||||
pub trait Query {
|
||||
type Params: Serialize;
|
||||
type Row: DeserializeOwned;
|
||||
|
||||
const SQL: &'static str;
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Rows<T> {
|
||||
data: Vec<T>,
|
||||
}
|
||||
|
||||
fn parameters<T: Serialize>(params: &T) -> Result<BTreeMap<String, Parameter>, Error> {
|
||||
let value = serde_json::to_value(params).map_err(|_| Error::InvalidParameters)?;
|
||||
serde_json::from_value(value).map_err(|_| Error::InvalidParameters)
|
||||
}
|
||||
|
||||
pub async fn fetch<Q: Query>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
params: &Q::Params,
|
||||
) -> Result<Vec<Q::Row>, Error> {
|
||||
let body = execute_read(client, connection, Q::SQL, ¶meters(params)?).await?;
|
||||
decode_rows::<Q::Row>(&body)
|
||||
}
|
||||
|
||||
pub async fn fetch_json<Q: Query>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
params: &Q::Params,
|
||||
) -> Result<String, Error> {
|
||||
let body = execute_read(client, connection, Q::SQL, ¶meters(params)?).await?;
|
||||
decode_rows::<Q::Row>(&body)?;
|
||||
Ok(body)
|
||||
}
|
||||
|
||||
fn decode_rows<T: DeserializeOwned>(body: &str) -> Result<Vec<T>, Error> {
|
||||
serde_json::from_str::<Rows<T>>(body)
|
||||
.map(|rows| rows.data)
|
||||
.map_err(|_| Error::InvalidResponse)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_storage_clickhouse::{Connection, Error, execute_read, insert_encoded_rows};
|
||||
use litellm_storage_clickhouse::{Connection, Error, Query, execute_read, insert_encoded_rows};
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
|
|
@ -32,3 +32,82 @@ async fn read_rejects_empty_sql() {
|
|||
Err(Error::EmptySql)
|
||||
));
|
||||
}
|
||||
|
||||
#[derive(serde::Serialize)]
|
||||
struct QueryParams {
|
||||
signed: i64,
|
||||
unsigned: u64,
|
||||
float: f64,
|
||||
text: String,
|
||||
strings: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize, PartialEq)]
|
||||
struct QueryRow {
|
||||
answer: String,
|
||||
}
|
||||
|
||||
struct TypedQuery;
|
||||
|
||||
impl litellm_storage_clickhouse::Query for TypedQuery {
|
||||
type Params = QueryParams;
|
||||
type Row = QueryRow;
|
||||
const SQL: &'static str = "SELECT typed_parameters";
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::valid(
|
||||
r#"{"meta":[],"data":[{"answer":"ok"}],"rows":1,"statistics":{"elapsed":0.1}}"#,
|
||||
true
|
||||
)]
|
||||
#[case::wrong_type(r#"{"data":[{"answer":1}]}"#, false)]
|
||||
#[case::missing_column(r#"{"data":[{}]}"#, false)]
|
||||
#[case::exception(r#"{"data":[],"exception":"failed"}"#, false)]
|
||||
#[tokio::test]
|
||||
async fn typed_fetch_encodes_parameters_and_validates_rows(
|
||||
#[case] body: &str,
|
||||
#[case] valid: bool,
|
||||
) {
|
||||
use litellm_storage_clickhouse::{fetch, fetch_json};
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{body_string, query_param},
|
||||
};
|
||||
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(body_string(TypedQuery::SQL))
|
||||
.and(query_param("param_signed", i64::MIN.to_string()))
|
||||
.and(query_param("param_unsigned", u64::MAX.to_string()))
|
||||
.and(query_param("param_float", "12.5"))
|
||||
.and(query_param("param_text", "line\\nbreak"))
|
||||
.and(query_param("param_strings", "['a\\'b','雪']"))
|
||||
.and(query_param("readonly", "1"))
|
||||
.and(query_param("max_result_rows", "1000"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_string(body))
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let client = Client::no_redirect_for_test();
|
||||
let connection = Connection::parse(&server.uri()).unwrap();
|
||||
let params = QueryParams {
|
||||
signed: i64::MIN,
|
||||
unsigned: u64::MAX,
|
||||
float: 12.5,
|
||||
text: "line\nbreak".into(),
|
||||
strings: vec!["a'b".into(), "雪".into()],
|
||||
};
|
||||
let rows = fetch::<TypedQuery>(&client, &connection, ¶ms).await;
|
||||
let envelope = fetch_json::<TypedQuery>(&client, &connection, ¶ms).await;
|
||||
if valid {
|
||||
assert_eq!(
|
||||
rows.unwrap(),
|
||||
vec![QueryRow {
|
||||
answer: "ok".into()
|
||||
}]
|
||||
);
|
||||
assert_eq!(envelope.unwrap(), body);
|
||||
} else {
|
||||
assert!(matches!(rows, Err(Error::InvalidResponse)));
|
||||
assert!(matches!(envelope, Err(Error::InvalidResponse)));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
7
litellm-rust/crates/traces-clickhouse/AGENTS.md
Normal file
7
litellm-rust/crates/traces-clickhouse/AGENTS.md
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
- Own trace schema, row encoding, SQL query adapters and reader provisioning; consume domain types from `litellm-traces`
|
||||
- Keep generic ClickHouse connections and HTTP execution in `litellm-storage-clickhouse`; keep PyO3 conversion in `python-bridge`
|
||||
- Keep schema definitions only in `migrations/NNNN_description.sql`, embedded by `litellm_migrate::migrate!`
|
||||
- Require typed query parameters and SELECT-only readers with server-side limits and tenant isolation
|
||||
- Bound insert time and encoded bytes; preserve shared values and explicit retry deduplication
|
||||
- Test storage behavior through the public API against ClickHouse
|
||||
- Expose one top-level `Error` enum in `src/error.rs`; own trace failures and wrap storage errors with `#[from]` or `#[source]`
|
||||
31
litellm-rust/crates/traces-clickhouse/Cargo.toml
Normal file
31
litellm-rust/crates/traces-clickhouse/Cargo.toml
Normal file
|
|
@ -0,0 +1,31 @@
|
|||
[package]
|
||||
name = "litellm-traces-clickhouse"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
askama.workspace = true
|
||||
flate2.workspace = true
|
||||
futures-util.workspace = true
|
||||
hmac = "0.12.1"
|
||||
litellm-http.workspace = true
|
||||
litellm-migrate.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
litellm-traces.workspace = true
|
||||
moka.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
sha2.workspace = true
|
||||
strum.workspace = true
|
||||
thiserror.workspace = true
|
||||
time = { workspace = true, features = ["formatting"] }
|
||||
tokio.workspace = true
|
||||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
wiremock.workspace = true
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
ALTER TABLE {database}.otel_traces
|
||||
ADD COLUMN IF NOT EXISTS UserId String DEFAULT ''
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE {database}.agent_traces_by_key
|
||||
ADD COLUMN IF NOT EXISTS UserIds SimpleAggregateFunction(groupUniqArrayArray, Array(String)) DEFAULT [],
|
||||
ADD COLUMN IF NOT EXISTS IdentifiedLlmCount SimpleAggregateFunction(sum, UInt64) DEFAULT 0
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
ALTER TABLE {database}.agent_traces_by_key_mv MODIFY QUERY
|
||||
SELECT
|
||||
TeamId, ApiKeyHash, TraceId, groupUniqArray(UserId) AS UserIds,
|
||||
min(Timestamp) AS StartTs,
|
||||
max(Timestamp + toIntervalNanosecond(Duration)) AS EndTs,
|
||||
any(ServiceName) AS ServiceName,
|
||||
anyLastIf(toNullable(SpanName), ParentSpanId = '') AS RootName,
|
||||
anyLastIf(toNullable(InputPreview), ParentSpanId = '') AS RootInput,
|
||||
anyLastIf(toNullable(StatusCode), ParentSpanId = '') AS RootStatus,
|
||||
count() AS SpanCount,
|
||||
countIf(ObservationType = 'agent') AS AgentCount,
|
||||
countIf(ObservationType = 'llm') AS LlmCount,
|
||||
countIf(ObservationType = 'llm' AND LiteLLMRequestId != '') AS IdentifiedLlmCount,
|
||||
countIf(ObservationType = 'tool') AS ToolCount,
|
||||
countIf(StatusCode = 'STATUS_CODE_ERROR') AS ErrorCount,
|
||||
sum(InputTokens) AS InputTokens,
|
||||
sum(OutputTokens) AS OutputTokens,
|
||||
groupUniqArrayIf(toString(Model), Model != '') AS Models,
|
||||
groupUniqArrayIf(SpanName, ObservationType = 'agent') AS AgentNames,
|
||||
groupArrayIf(LiteLLMRequestId, ObservationType = 'llm' OR LiteLLMRequestId != '') AS RequestIds
|
||||
FROM {database}.otel_traces
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
WITH page AS (
|
||||
SELECT TraceId AS trace_id,
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref,
|
||||
TeamId AS team_id, ApiKeyHash AS api_key_hash,
|
||||
if(length(groupUniqArrayArray(UserIds)) = 1, arrayElement(groupUniqArrayArray(UserIds), 1), '') AS user_id, TeamId AS team_id, ApiKeyHash AS api_key_hash,
|
||||
ifNull(any(RootName), '') AS name, any(ServiceName) AS service,
|
||||
ifNull(any(RootInput), '') AS input_preview, ifNull(any(RootStatus), '') AS status,
|
||||
toUnixTimestamp64Milli(min(StartTs)) AS start_ms,
|
||||
|
|
@ -12,10 +12,14 @@ SELECT TraceId AS trace_id,
|
|||
sum(LlmCount) AS llm_calls, sum(ToolCount) AS tool_calls,
|
||||
sum(InputTokens) AS input_tokens, sum(OutputTokens) AS output_tokens,
|
||||
groupUniqArrayArray(Models) AS models, sum(ErrorCount) AS error_count,
|
||||
arrayDistinct(groupArrayArray(RequestIds)) AS request_ids
|
||||
arrayDistinct(if(sum(IdentifiedLlmCount) != sum(LlmCount),
|
||||
arrayConcat(groupArrayArray(RequestIds), ['']),
|
||||
groupArrayArray(RequestIds))) AS request_ids
|
||||
FROM agent_traces_by_key
|
||||
WHERE (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
|
||||
WHERE ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserIds = [{user_id:String}])
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
HAVING min(StartTs) >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND min(StartTs) < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
26
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal file
26
litellm-rust/crates/traces-clickhouse/query/span_detail.sql
Normal file
|
|
@ -0,0 +1,26 @@
|
|||
SELECT o.SpanId AS span_id, o.Input AS input,
|
||||
if(o.Output = '' AND o.ObservationType = 'agent', answer.output, o.Output) AS output,
|
||||
o.SpanAttributes AS attributes
|
||||
FROM otel_traces AS o
|
||||
LEFT JOIN (
|
||||
SELECT TeamId, ApiKeyHash, ParentSpanId AS parent_span_id, argMax(Output, Timestamp) AS output
|
||||
FROM otel_traces
|
||||
WHERE TraceId = {trace_id:String} AND ParentSpanId = {span_id:String}
|
||||
AND ObservationType = 'llm' AND Output != ''
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
|
||||
GROUP BY TeamId, ApiKeyHash, ParentSpanId
|
||||
) AS answer ON answer.parent_span_id = o.SpanId
|
||||
AND answer.TeamId = o.TeamId AND answer.ApiKeyHash = o.ApiKeyHash
|
||||
WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, o.TeamId)
|
||||
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
|
||||
LIMIT 1
|
||||
|
|
@ -4,8 +4,10 @@ SELECT SpanId AS span_id,
|
|||
hex(SHA256(StatusMessage)) AS version
|
||||
FROM otel_traces
|
||||
WHERE TraceId = {trace_id:String} AND SpanId = {span_id:String}
|
||||
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) = {trace_ref:String})
|
||||
AND ({error_version:String} = '' OR hex(SHA256(StatusMessage)) = {error_version:String})
|
||||
|
|
@ -0,0 +1,11 @@
|
|||
SELECT request_id, response_id, team_id, api_key, user, spend,
|
||||
toUnixTimestamp64Milli(start_time) AS start_ms
|
||||
FROM spend_logs FINAL
|
||||
WHERE response_id IN {response_ids:Array(String)}
|
||||
AND start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND user = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, team_id)
|
||||
OR ({api_key_hash:String} != '' AND api_key = {api_key_hash:String}))
|
||||
ORDER BY start_time DESC
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
SELECT hex(SHA256(concat(TeamId, char(0), ApiKeyHash, char(0), TraceId))) AS trace_ref
|
||||
FROM otel_traces
|
||||
WHERE TraceId = {trace_id:String}
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, TeamId)
|
||||
OR ({api_key_hash:String} != '' AND ApiKeyHash = {api_key_hash:String}))
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
LIMIT 2
|
||||
|
|
@ -7,11 +7,13 @@ SELECT o.SpanId AS span_id, o.ParentSpanId AS parent_span_id, o.SpanName AS name
|
|||
o.ServiceName AS service, o.InputPreview AS input_preview, o.Model AS model,
|
||||
o.InputTokens AS input_tokens, o.OutputTokens AS output_tokens,
|
||||
o.LiteLLMRequestId AS litellm_request_id,
|
||||
o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
|
||||
o.UserId AS user_id, o.TeamId AS team_id, o.ApiKeyHash AS api_key_hash
|
||||
FROM otel_traces AS o
|
||||
WHERE o.TraceId = {trace_id:String}
|
||||
AND (empty({team_ids:Array(String)}) OR o.TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR o.ApiKeyHash = {api_key_hash:String})
|
||||
AND ({all_teams:UInt8} = 1
|
||||
OR ({user_id:String} != '' AND o.UserId = {user_id:String})
|
||||
OR has({team_ids:Array(String)}, o.TeamId)
|
||||
OR ({api_key_hash:String} != '' AND o.ApiKeyHash = {api_key_hash:String}))
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
|
||||
ORDER BY o.Timestamp, o.EngineReceivedMs, o.StatusMessage
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
use litellm_storage_clickhouse::{Error, Storage};
|
||||
use crate::Error;
|
||||
use litellm_storage_clickhouse::Storage;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Config {
|
||||
|
|
@ -8,7 +9,7 @@ pub struct Config {
|
|||
|
||||
impl Config {
|
||||
pub fn new(database: String, url: &str, retention_days: u32) -> Result<Self, Error> {
|
||||
crate::schema_statements(&database, retention_days)?;
|
||||
super::schema_statements(&database, retention_days)?;
|
||||
Ok(Self {
|
||||
storage: Storage::new(database, url)?,
|
||||
retention_days,
|
||||
37
litellm-rust/crates/traces-clickhouse/src/error.rs
Normal file
37
litellm-rust/crates/traces-clickhouse/src/error.rs
Normal file
|
|
@ -0,0 +1,37 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("invalid ClickHouse insert row")]
|
||||
InvalidRow,
|
||||
#[error("invalid ClickHouse insert table")]
|
||||
InvalidTable,
|
||||
#[error("database must be a nonempty SQL identifier and retention must be positive")]
|
||||
InvalidSchema,
|
||||
#[error("unknown ClickHouse read query")]
|
||||
InvalidQuery,
|
||||
#[error("invalid ClickHouse query parameters")]
|
||||
InvalidParameters,
|
||||
#[error("ClickHouse returned an invalid or failed JSON query response")]
|
||||
InvalidResponse,
|
||||
#[error("ClickHouse insert exceeds the encoded size limit")]
|
||||
InsertTooLarge,
|
||||
#[error("ClickHouse schema setup failed with HTTP status {0}")]
|
||||
SchemaFailed(u16),
|
||||
#[error("ClickHouse schema setup transport failed")]
|
||||
SchemaTransport,
|
||||
#[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<Error>),
|
||||
}
|
||||
|
|
@ -12,7 +12,8 @@ use serde_json::Value;
|
|||
use sha2::{Digest, Sha256};
|
||||
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
||||
|
||||
use crate::{Connection, Error, Shared};
|
||||
use super::{Connection, Error};
|
||||
use litellm_traces::Shared;
|
||||
|
||||
const MAX_INSERT_BYTES: usize = 64 * 1024 * 1024;
|
||||
|
||||
|
|
@ -71,6 +72,7 @@ pub async fn insert_shared_rows(
|
|||
body,
|
||||
)
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
fn shared_rows(rows: Vec<BTreeMap<String, Value>>) -> Vec<InsertRow> {
|
||||
|
|
@ -226,8 +228,8 @@ mod tests {
|
|||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::Error;
|
||||
use super::{shared_rows, write_rows};
|
||||
use crate::Error;
|
||||
|
||||
#[rstest]
|
||||
fn encoded_limit_counts_utf8_bytes_across_rows() {
|
||||
21
litellm-rust/crates/traces-clickhouse/src/lib.rs
Normal file
21
litellm-rust/crates/traces-clickhouse/src/lib.rs
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
mod config;
|
||||
mod error;
|
||||
mod insert;
|
||||
pub mod query;
|
||||
mod query_access;
|
||||
mod schema;
|
||||
mod sql;
|
||||
mod table;
|
||||
|
||||
pub use config::Config;
|
||||
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_access::QueryReaders;
|
||||
pub use schema::{
|
||||
NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, ensure_schema, schema_statements,
|
||||
};
|
||||
pub use sql::execute_named_read;
|
||||
pub use table::TraceTable;
|
||||
362
litellm-rust/crates/traces-clickhouse/src/query.rs
Normal file
362
litellm-rust/crates/traces-clickhouse/src/query.rs
Normal file
|
|
@ -0,0 +1,362 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
use crate::TraceTable;
|
||||
use futures_util::{
|
||||
StreamExt,
|
||||
stream::{self, TryStreamExt},
|
||||
};
|
||||
use litellm_http::Client;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Value, json};
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use super::{Connection, Error, NORMALIZED_FIELD_DEFINITIONS, Parameter};
|
||||
|
||||
mod guide;
|
||||
pub mod lens;
|
||||
pub mod named;
|
||||
mod number;
|
||||
|
||||
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 execute_read(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
sql: &str,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error> {
|
||||
litellm_storage_clickhouse::execute_read(client, connection, sql, parameters)
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
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(TraceTable::iter())
|
||||
.then(|table| async move {
|
||||
Ok::<_, Error>(TableSchema {
|
||||
name: table.into(),
|
||||
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": "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())
|
||||
}
|
||||
|
||||
#[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)
|
||||
);
|
||||
}
|
||||
}
|
||||
173
litellm-rust/crates/traces-clickhouse/src/query/lens.rs
Normal file
173
litellm-rust/crates/traces-clickhouse/src/query/lens.rs
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
use litellm_storage_clickhouse::Query;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensAccessParams {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub all_teams: u8,
|
||||
pub team: String,
|
||||
pub key_hash: String,
|
||||
}
|
||||
|
||||
pub struct LensAvailability;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensAvailabilityParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensAvailabilityRow {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub traces: u8,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub requests: u8,
|
||||
}
|
||||
|
||||
impl Query for LensAvailability {
|
||||
type Params = LensAvailabilityParams;
|
||||
type Row = LensAvailabilityRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/lens_availability.sql");
|
||||
}
|
||||
|
||||
pub struct LensAgents;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensAgentsParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensAgentsRow {
|
||||
pub agent_name: String,
|
||||
}
|
||||
|
||||
impl Query for LensAgents {
|
||||
type Params = LensAgentsParams;
|
||||
type Row = LensAgentsRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/lens_agents.sql");
|
||||
}
|
||||
|
||||
pub struct LensSample;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensSampleParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub end: u64,
|
||||
pub agent_name: String,
|
||||
pub service: String,
|
||||
pub filter_keys: Vec<String>,
|
||||
pub filter_values: Vec<String>,
|
||||
pub selected_team: String,
|
||||
pub execution_ids: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub sample_cap: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub sample_percent: f64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub preview: u8,
|
||||
pub after: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub limit: u32,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub offset: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensSampleRow {
|
||||
pub source: String,
|
||||
pub trace_id: String,
|
||||
pub team_id: String,
|
||||
pub trace_ref: String,
|
||||
pub name: String,
|
||||
pub start_time: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub span_count: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub root_seen: u8,
|
||||
pub service: String,
|
||||
pub attributes: Vec<(String, String)>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub eligible: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub position: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub selected: f64,
|
||||
pub selection_key: String,
|
||||
}
|
||||
|
||||
impl Query for LensSample {
|
||||
type Params = LensSampleParams;
|
||||
type Row = LensSampleRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/lens_sample.sql");
|
||||
}
|
||||
|
||||
pub struct LensContent;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensContentParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
pub id: String,
|
||||
pub record_team: String,
|
||||
pub trace_ref: String,
|
||||
pub cursor: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub offset: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensContentRow {
|
||||
pub span_id: String,
|
||||
pub parent_span_id: String,
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub content: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub truncated: u8,
|
||||
}
|
||||
|
||||
impl Query for LensContent {
|
||||
type Params = LensContentParams;
|
||||
type Row = LensContentRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/lens_content.sql");
|
||||
}
|
||||
|
||||
pub struct LensEvidence;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensEvidenceParams {
|
||||
#[serde(flatten)]
|
||||
pub access: LensAccessParams,
|
||||
pub source: String,
|
||||
pub id: String,
|
||||
pub record_team: String,
|
||||
pub trace_ref: String,
|
||||
pub span: String,
|
||||
pub quote: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct LensEvidenceRow {
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub count: u64,
|
||||
}
|
||||
|
||||
impl Query for LensEvidence {
|
||||
type Params = LensEvidenceParams;
|
||||
type Row = LensEvidenceRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/lens_evidence.sql");
|
||||
}
|
||||
320
litellm-rust/crates/traces-clickhouse/src/query/named.rs
Normal file
320
litellm-rust/crates/traces-clickhouse/src/query/named.rs
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
use litellm_storage_clickhouse::Query;
|
||||
use litellm_traces::query::named as contracts;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub use contracts::ReadAccessParams;
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::ListTracesParams")]
|
||||
struct ListTracesParamsEncoding {
|
||||
#[serde(flatten)]
|
||||
pub access: contracts::ReadAccessParams,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub end_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub cursor_ms: i64,
|
||||
pub cursor_trace_id: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub limit: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesParams(
|
||||
#[serde(with = "ListTracesParamsEncoding")] pub contracts::ListTracesParams,
|
||||
);
|
||||
|
||||
impl From<contracts::ListTracesParams> for ListTracesParams {
|
||||
fn from(value: contracts::ListTracesParams) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::ListTracesRow")]
|
||||
struct ListTracesRowEncoding {
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub team_id: String,
|
||||
pub api_key_hash: String,
|
||||
pub user_id: String,
|
||||
pub name: String,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub status: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub duration_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub span_count: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub agent_count: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub agent_invocations: u64,
|
||||
#[serde(default)]
|
||||
pub agent_names: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub frameworks: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub llm_calls: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub tool_calls: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub input_tokens: u64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub output_tokens: u64,
|
||||
pub models: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub error_count: u64,
|
||||
pub request_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesRow(#[serde(with = "ListTracesRowEncoding")] pub contracts::ListTracesRow);
|
||||
|
||||
pub use contracts::TraceSpansParams;
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::TraceSpansRow")]
|
||||
struct TraceSpansRowEncoding {
|
||||
pub span_id: String,
|
||||
pub parent_span_id: String,
|
||||
pub name: String,
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub agent: String,
|
||||
#[serde(default)]
|
||||
pub framework: String,
|
||||
pub status: String,
|
||||
pub status_message: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub error_truncated: u8,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ns: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub duration_ns: u64,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub model: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub input_tokens: u32,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub output_tokens: u32,
|
||||
pub litellm_request_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key_hash: String,
|
||||
pub user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceSpansRow(#[serde(with = "TraceSpansRowEncoding")] pub contracts::TraceSpansRow);
|
||||
|
||||
pub use contracts::SpanDetailParams;
|
||||
|
||||
pub use contracts::SpanDetailRow;
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::SpanErrorParams")]
|
||||
struct SpanErrorParamsEncoding {
|
||||
#[serde(flatten)]
|
||||
pub access: contracts::ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub span_id: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub error_offset: u64,
|
||||
pub error_version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorParams(
|
||||
#[serde(with = "SpanErrorParamsEncoding")] pub contracts::SpanErrorParams,
|
||||
);
|
||||
|
||||
impl From<contracts::SpanErrorParams> for SpanErrorParams {
|
||||
fn from(value: contracts::SpanErrorParams) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::SpanErrorRow")]
|
||||
struct SpanErrorRowEncoding {
|
||||
pub span_id: String,
|
||||
pub message: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub total_chars: u64,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorRow(#[serde(with = "SpanErrorRowEncoding")] pub contracts::SpanErrorRow);
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::SpendByResponseIdsParams")]
|
||||
struct SpendByResponseIdsParamsEncoding {
|
||||
#[serde(flatten)]
|
||||
pub access: contracts::ReadAccessParams,
|
||||
pub response_ids: Vec<String>,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub end_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsParams(
|
||||
#[serde(with = "SpendByResponseIdsParamsEncoding")] pub contracts::SpendByResponseIdsParams,
|
||||
);
|
||||
|
||||
impl From<contracts::SpendByResponseIdsParams> for SpendByResponseIdsParams {
|
||||
fn from(value: contracts::SpendByResponseIdsParams) -> Self {
|
||||
Self(value)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
#[serde(remote = "contracts::SpendByResponseIdsRow")]
|
||||
struct SpendByResponseIdsRowEncoding {
|
||||
pub request_id: String,
|
||||
pub response_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key: String,
|
||||
pub user: String,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub spend: f64,
|
||||
#[serde(deserialize_with = "super::number::deserialize")]
|
||||
pub start_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsRow(
|
||||
#[serde(with = "SpendByResponseIdsRowEncoding")] pub contracts::SpendByResponseIdsRow,
|
||||
);
|
||||
|
||||
pub struct ListTraces;
|
||||
|
||||
impl Query for ListTraces {
|
||||
type Params = ListTracesParams;
|
||||
type Row = ListTracesRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/list_traces.sql");
|
||||
}
|
||||
|
||||
pub struct TraceSpans;
|
||||
|
||||
impl Query for TraceSpans {
|
||||
type Params = TraceSpansParams;
|
||||
type Row = TraceSpansRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/trace_spans.sql");
|
||||
}
|
||||
|
||||
pub struct SpanDetail;
|
||||
|
||||
impl Query for SpanDetail {
|
||||
type Params = SpanDetailParams;
|
||||
type Row = SpanDetailRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/span_detail.sql");
|
||||
}
|
||||
|
||||
pub struct SpanError;
|
||||
|
||||
impl Query for SpanError {
|
||||
type Params = SpanErrorParams;
|
||||
type Row = SpanErrorRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/span_error.sql");
|
||||
}
|
||||
|
||||
pub struct SpendByResponseIds;
|
||||
|
||||
impl Query for SpendByResponseIds {
|
||||
type Params = SpendByResponseIdsParams;
|
||||
type Row = SpendByResponseIdsRow;
|
||||
|
||||
const SQL: &'static str = include_str!("../../query/spend_by_response_ids.sql");
|
||||
}
|
||||
|
||||
pub use contracts::{TraceIdentityParams, TraceIdentityRow};
|
||||
|
||||
pub struct TraceIdentity;
|
||||
|
||||
impl Query for TraceIdentity {
|
||||
type Params = TraceIdentityParams;
|
||||
type Row = TraceIdentityRow;
|
||||
const SQL: &'static str = include_str!("../../query/trace_identity.sql");
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
fn round_trip<T: serde::de::DeserializeOwned + Serialize>(wire: Value, quoted: bool) {
|
||||
let encoded = Value::Object(
|
||||
wire.as_object()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
let encoded = if quoted && value.is_number() && name != "all_teams" {
|
||||
json!(value.to_string())
|
||||
} else {
|
||||
value.clone()
|
||||
};
|
||||
(name.clone(), encoded)
|
||||
})
|
||||
.collect(),
|
||||
);
|
||||
let decoded: T = serde_json::from_value(encoded).unwrap();
|
||||
assert_eq!(serde_json::to_value(decoded).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unquoted(false)]
|
||||
#[case::quoted(true)]
|
||||
fn rows_decode_into_neutral_contracts(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesRow>(
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "ok", "start_ms": -1, "duration_ms": 20, "span_count": u64::MAX, "agent_count": 1, "agent_invocations": 2, "agent_names": ["agent"], "frameworks": ["claude-agent-sdk"], "llm_calls": 3, "tool_calls": 4, "input_tokens": 5, "output_tokens": 6, "models": ["model"], "error_count": 0, "request_ids": ["request"]}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<TraceSpansRow>(
|
||||
json!({"span_id": "span", "parent_span_id": "parent", "name": "agent", "type": "agent", "agent": "agent", "framework": "claude-agent-sdk", "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"}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpanDetailRow>(
|
||||
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpanErrorRow>(
|
||||
json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpendByResponseIdsRow>(
|
||||
json!({"request_id": "request", "response_id": "response", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
|
||||
quoted,
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unquoted(false)]
|
||||
#[case::quoted(true)]
|
||||
fn parameters_preserve_flattened_multi_team_access(#[case] quoted: bool) {
|
||||
round_trip::<ListTracesParams>(
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "key", "start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": u32::MAX}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpanErrorParams>(
|
||||
json!({"all_teams": 0, "user_id": "", "team_ids": [], "api_key_hash": "key", "trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
|
||||
quoted,
|
||||
);
|
||||
round_trip::<SpendByResponseIdsParams>(
|
||||
json!({"all_teams": 0, "user_id": "user", "team_ids": ["team-a", "team-b"], "api_key_hash": "", "response_ids": ["response"], "start_ms": -1, "end_ms": 10}),
|
||||
quoted,
|
||||
);
|
||||
}
|
||||
}
|
||||
44
litellm-rust/crates/traces-clickhouse/src/query/number.rs
Normal file
44
litellm-rust/crates/traces-clickhouse/src/query/number.rs
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
use serde::{Deserialize, Deserializer, de::DeserializeOwned};
|
||||
|
||||
pub(super) fn deserialize<'de, D, T>(deserializer: D) -> Result<T, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
T: DeserializeOwned,
|
||||
{
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum Number {
|
||||
Quoted(String),
|
||||
Unquoted(serde_json::Number),
|
||||
}
|
||||
match Number::deserialize(deserializer)? {
|
||||
Number::Quoted(value) => serde_json::from_str(&value),
|
||||
Number::Unquoted(value) => serde_json::from_value(serde_json::Value::Number(value)),
|
||||
}
|
||||
.map_err(serde::de::Error::custom)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::query::named::SpanErrorRow;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::quoted_max(serde_json::json!(u64::MAX.to_string()), Some(u64::MAX))]
|
||||
#[case::unquoted_max(serde_json::json!(u64::MAX), Some(u64::MAX))]
|
||||
#[case::overflow(serde_json::json!("18446744073709551616"), None)]
|
||||
#[case::negative(serde_json::json!(-1), None)]
|
||||
#[case::fraction(serde_json::json!(1.5), None)]
|
||||
fn numeric_rows_enforce_integer_range(
|
||||
#[case] value: serde_json::Value,
|
||||
#[case] expected: Option<u64>,
|
||||
) {
|
||||
let row = serde_json::from_value::<SpanErrorRow>(serde_json::json!({
|
||||
"span_id": "span", "message": "error", "total_chars": value, "version": "hash"
|
||||
}));
|
||||
match expected {
|
||||
Some(value) => assert_eq!(row.unwrap().0.total_chars, value),
|
||||
None => assert!(row.is_err()),
|
||||
}
|
||||
}
|
||||
}
|
||||
237
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal file
237
litellm-rust/crates/traces-clickhouse/src/query_access.rs
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use hmac::{Hmac, Mac};
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::QueryScope;
|
||||
use moka::future::Cache;
|
||||
use strum::IntoEnumIterator;
|
||||
|
||||
use sha2::{Digest, Sha256};
|
||||
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
|
||||
|
||||
use super::{Connection, Error, TraceTable};
|
||||
|
||||
#[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, Error> {
|
||||
self.slots
|
||||
.clone()
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| Error::Busy)
|
||||
}
|
||||
|
||||
pub async fn connection(
|
||||
&self,
|
||||
client: &Client,
|
||||
scope: &QueryScope,
|
||||
secret: &str,
|
||||
) -> Result<Connection, Error> {
|
||||
scope.validate().map_err(|_| Error::InvalidScope)?;
|
||||
if secret.is_empty() {
|
||||
return Err(Error::MissingSecret);
|
||||
}
|
||||
let identity = serde_json::to_vec(&("litellm_trace_reader_v1", &self.database, scope))
|
||||
.map_err(|_| Error::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(Error::Cached)
|
||||
}
|
||||
|
||||
async fn provision(
|
||||
&self,
|
||||
client: &Client,
|
||||
scope: &QueryScope,
|
||||
user: &str,
|
||||
password: &str,
|
||||
) -> Result<Connection, Error> {
|
||||
let database = &self.database;
|
||||
if database.is_empty()
|
||||
|| !database
|
||||
.bytes()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|
||||
{
|
||||
return Err(Error::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 TraceTable::iter() {
|
||||
let predicate = predicate(scope, 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 TraceTable::iter() {
|
||||
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(Error::Storage)
|
||||
}
|
||||
|
||||
async fn execute(&self, client: &Client, sql: String) -> Result<(), Error> {
|
||||
let response = client
|
||||
.post(self.writer.url().clone())
|
||||
.timeout(Duration::from_secs(15))
|
||||
.body(sql)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::ProvisionTransport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::ProvisionFailed(response.status().as_u16()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn predicate(scope: &QueryScope, table: TraceTable) -> String {
|
||||
let (team, key) = match table {
|
||||
TraceTable::OtelTraces | TraceTable::AgentTracesByKey => ("TeamId", "ApiKeyHash"),
|
||||
TraceTable::SpendLogs => ("team_id", "api_key"),
|
||||
};
|
||||
match scope {
|
||||
QueryScope::Admin => "1".to_owned(),
|
||||
QueryScope::Logs {
|
||||
user_id,
|
||||
team_ids,
|
||||
api_key_hash,
|
||||
} => {
|
||||
let owner = literal(user_id);
|
||||
let user_clause = match table {
|
||||
TraceTable::OtelTraces => format!("UserId = {owner}"),
|
||||
TraceTable::AgentTracesByKey => format!("UserIds = [{owner}]"),
|
||||
TraceTable::SpendLogs => format!("user = {owner}"),
|
||||
};
|
||||
let teams = team_ids
|
||||
.iter()
|
||||
.map(|value| literal(value))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let team_clause = if team_ids.is_empty() {
|
||||
"0".to_owned()
|
||||
} else {
|
||||
format!("{team} IN ({teams})")
|
||||
};
|
||||
format!(
|
||||
"({owner} != '' AND {user_clause}) OR ({team_clause}) OR ({hash} != '' AND {key} = {hash})",
|
||||
owner = owner,
|
||||
hash = literal(api_key_hash),
|
||||
)
|
||||
}
|
||||
QueryScope::Team { team_id } => format!("{team} = {}", literal(team_id)),
|
||||
QueryScope::Key {
|
||||
team_id,
|
||||
api_key_hash,
|
||||
} => format!(
|
||||
"{team} = {} AND {key} = {}",
|
||||
literal(team_id),
|
||||
literal(api_key_hash)
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn credential(secret: &str, purpose: &[u8], identity: &[u8]) -> Result<String, Error> {
|
||||
let mut mac =
|
||||
Hmac::<Sha256>::new_from_slice(secret.as_bytes()).map_err(|_| Error::MissingSecret)?;
|
||||
mac.update(purpose);
|
||||
mac.update(identity);
|
||||
Ok(format!("{:x}", mac.finalize().into_bytes()))
|
||||
}
|
||||
|
||||
fn literal(value: &str) -> String {
|
||||
format!("'{}'", value.replace('\\', "\\\\").replace('\'', "\\'"))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::otel(TraceTable::OtelTraces, "TeamId", "ApiKeyHash")]
|
||||
#[case::agent(TraceTable::AgentTracesByKey, "TeamId", "ApiKeyHash")]
|
||||
#[case::spend(TraceTable::SpendLogs, "team_id", "api_key")]
|
||||
fn predicates_preserve_scope_and_escape_values(
|
||||
#[case] table: TraceTable,
|
||||
#[case] team: &str,
|
||||
#[case] key: &str,
|
||||
) {
|
||||
assert_eq!(predicate(&QueryScope::Admin, table), "1");
|
||||
assert_eq!(
|
||||
predicate(
|
||||
&QueryScope::Team {
|
||||
team_id: "team'\\".into()
|
||||
},
|
||||
table
|
||||
),
|
||||
format!("{team} = 'team\\'\\\\'")
|
||||
);
|
||||
assert_eq!(
|
||||
predicate(
|
||||
&QueryScope::Key {
|
||||
team_id: "".into(),
|
||||
api_key_hash: "key'\\".into()
|
||||
},
|
||||
table
|
||||
),
|
||||
format!("{team} = '' AND {key} = 'key\\'\\\\'")
|
||||
);
|
||||
}
|
||||
}
|
||||
136
litellm-rust/crates/traces-clickhouse/src/schema.rs
Normal file
136
litellm-rust/crates/traces-clickhouse/src/schema.rs
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
use litellm_http::Client;
|
||||
use litellm_migrate::Migration;
|
||||
use serde::Serialize;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::Connection;
|
||||
use super::Error;
|
||||
|
||||
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations");
|
||||
|
||||
pub fn schema_statements(database: &str, retention_days: u32) -> Result<Vec<String>, Error> {
|
||||
if database.is_empty()
|
||||
|| !database
|
||||
.bytes()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|
||||
|| retention_days == 0
|
||||
{
|
||||
return Err(Error::InvalidSchema);
|
||||
}
|
||||
let database = format!("`{database}`");
|
||||
Ok(
|
||||
std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}"))
|
||||
.chain(MIGRATIONS.iter().map(|migration| {
|
||||
migration
|
||||
.sql
|
||||
.replace("{database}", &database)
|
||||
.replace("{retention_days}", &retention_days.to_string())
|
||||
}))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn ensure_schema(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
database: &str,
|
||||
retention_days: u32,
|
||||
) -> Result<(), Error> {
|
||||
ensure_schema_with_timeout(
|
||||
client,
|
||||
connection,
|
||||
database,
|
||||
retention_days,
|
||||
SCHEMA_REQUEST_TIMEOUT,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn ensure_schema_with_timeout(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
database: &str,
|
||||
retention_days: u32,
|
||||
request_timeout: Duration,
|
||||
) -> Result<(), Error> {
|
||||
for statement in schema_statements(database, retention_days)? {
|
||||
let response = client
|
||||
.post(connection.url().clone())
|
||||
.timeout(request_timeout)
|
||||
.body(statement)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::SchemaTransport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::SchemaFailed(response.status().as_u16()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
pub struct NormalizedFieldDefinition {
|
||||
pub name: &'static str,
|
||||
pub clickhouse_column: &'static str,
|
||||
pub clickhouse_type: &'static str,
|
||||
pub meaning: &'static str,
|
||||
}
|
||||
|
||||
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
|
||||
NormalizedFieldDefinition {
|
||||
name: "observation_type",
|
||||
clickhouse_column: "ObservationType",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent, LLM, tool, chain, or framework span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "agent_name",
|
||||
clickhouse_column: "AgentName",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent associated with this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "framework",
|
||||
clickhouse_column: "Framework",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent framework or SDK that emitted this span, e.g. claude-agent-sdk",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "litellm_request_id",
|
||||
clickhouse_column: "LiteLLMRequestId",
|
||||
clickhouse_type: "String",
|
||||
meaning: "LiteLLM response ID used to link a span to a spend log",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "model",
|
||||
clickhouse_column: "Model",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Model used by this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input_tokens",
|
||||
clickhouse_column: "InputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Input token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output_tokens",
|
||||
clickhouse_column: "OutputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Output token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input",
|
||||
clickhouse_column: "Input",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized input payload",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output",
|
||||
clickhouse_column: "Output",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized output payload",
|
||||
},
|
||||
];
|
||||
83
litellm-rust/crates/traces-clickhouse/src/sql.rs
Normal file
83
litellm-rust/crates/traces-clickhouse/src/sql.rs
Normal file
|
|
@ -0,0 +1,83 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::ReadQuery;
|
||||
|
||||
use super::query::{lens::*, named::*};
|
||||
use super::{Connection, Error, Parameter};
|
||||
use litellm_storage_clickhouse::{Query, fetch_json};
|
||||
|
||||
pub async fn execute_named_read(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
query: ReadQuery,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error> {
|
||||
match query {
|
||||
ReadQuery::ListTraces => named_json::<ListTraces>(client, connection, parameters).await,
|
||||
ReadQuery::TraceIdentity => {
|
||||
named_json::<TraceIdentity>(client, connection, parameters).await
|
||||
}
|
||||
ReadQuery::TraceSpans => named_json::<TraceSpans>(client, connection, parameters).await,
|
||||
ReadQuery::SpanDetail => named_json::<SpanDetail>(client, connection, parameters).await,
|
||||
ReadQuery::SpanError => named_json::<SpanError>(client, connection, parameters).await,
|
||||
ReadQuery::SpendByResponseIds => {
|
||||
named_json::<SpendByResponseIds>(client, connection, parameters).await
|
||||
}
|
||||
ReadQuery::Availability => {
|
||||
named_json::<LensAvailability>(client, connection, parameters).await
|
||||
}
|
||||
ReadQuery::Agents => named_json::<LensAgents>(client, connection, parameters).await,
|
||||
ReadQuery::Sample => named_json::<LensSample>(client, connection, parameters).await,
|
||||
ReadQuery::Content => named_json::<LensContent>(client, connection, parameters).await,
|
||||
ReadQuery::Evidence => named_json::<LensEvidence>(client, connection, parameters).await,
|
||||
}
|
||||
}
|
||||
|
||||
async fn named_json<Q: Query>(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error>
|
||||
where
|
||||
Q::Params: serde::de::DeserializeOwned,
|
||||
{
|
||||
let value = serde_json::to_value(parameters).map_err(|_| Error::InvalidParameters)?;
|
||||
let params =
|
||||
serde_json::from_value::<Q::Params>(value).map_err(|_| Error::InvalidParameters)?;
|
||||
fetch_json::<Q>(client, connection, ¶ms)
|
||||
.await
|
||||
.map_err(Error::from)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::missing_span(serde_json::json!({}))]
|
||||
#[case::negative_offset(serde_json::json!({"span_id": "span", "error_offset": -1, "error_version": ""}))]
|
||||
#[case::overflow(serde_json::json!({"span_id": "span", "error_offset": "18446744073709551616", "error_version": ""}))]
|
||||
#[tokio::test]
|
||||
async fn named_read_rejects_invalid_parameters_before_transport(
|
||||
#[case] specific: serde_json::Value,
|
||||
) {
|
||||
let common = serde_json::json!({
|
||||
"all_teams": 1, "user_id": "", "team_ids": [], "api_key_hash": "", "trace_id": "trace", "trace_ref": ""
|
||||
});
|
||||
let parameters: BTreeMap<String, Parameter> = common
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.chain(specific.as_object().unwrap().iter())
|
||||
.map(|(name, value)| (name.clone(), serde_json::from_value(value.clone()).unwrap()))
|
||||
.collect();
|
||||
let client = Client::no_redirect_for_test();
|
||||
let connection = Connection::parse("http://127.0.0.1:1").unwrap();
|
||||
assert!(matches!(
|
||||
execute_named_read(&client, &connection, ReadQuery::SpanError, ¶meters).await,
|
||||
Err(Error::InvalidParameters)
|
||||
));
|
||||
}
|
||||
}
|
||||
9
litellm-rust/crates/traces-clickhouse/src/table.rs
Normal file
9
litellm-rust/crates/traces-clickhouse/src/table.rs
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
#[derive(
|
||||
Clone, Copy, Debug, strum::Display, strum::AsRefStr, strum::EnumIter, strum::IntoStaticStr,
|
||||
)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum TraceTable {
|
||||
OtelTraces,
|
||||
AgentTracesByKey,
|
||||
SpendLogs,
|
||||
}
|
||||
|
|
@ -35,7 +35,7 @@ Examples
|
|||
{% 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 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.TeamId != '' OR (t.UserId != '' AND t.UserId = s.user) OR (t.ApiKeyHash != '' 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 %}
|
||||
|
|
@ -46,7 +46,7 @@ Gotchas
|
|||
|
||||
{% 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 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 %}
|
||||
|
||||
{% block output_format %}Do not add FORMAT clauses; the endpoint requires ClickHouse JSON output{% endblock %}
|
||||
|
||||
|
|
@ -58,7 +58,7 @@ Gotchas
|
|||
|
||||
{% 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 spend_totals %}Use spend_logs FINAL to collapse replacement rows before totals. Shared response IDs and multiple spans can multiply costs in joins; require one spend match per response ID and ownership before aggregating. Missing IDs or costs leave totals unknown{% 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 %}
|
||||
|
||||
|
|
@ -1,18 +1,16 @@
|
|||
use litellm_http::Client;
|
||||
use litellm_traces::{Connection, Error, Parameter, execute_read};
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, Parameter, QueryReaders, QueryScope, execute_read,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::Value;
|
||||
use std::collections::BTreeMap;
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
mod support;
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
use support::{ClickHouseDatabase, database as start_database};
|
||||
|
||||
struct Database {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
_database: ClickHouseDatabase,
|
||||
url: String,
|
||||
admin_url: String,
|
||||
client: Client,
|
||||
|
|
@ -20,22 +18,9 @@ struct Database {
|
|||
|
||||
#[fixture]
|
||||
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.with_env_var("LITELLM_TRACES_READER_PASSWORD", "test_password")
|
||||
.with_copy_to(
|
||||
"/etc/clickhouse-server/users.d/litellm-traces-reader.xml",
|
||||
include_bytes!("../config/reader.xml").to_vec(),
|
||||
)
|
||||
.start()
|
||||
.await?;
|
||||
let admin_url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?,
|
||||
);
|
||||
let client = Client::no_redirect_for_test();
|
||||
let instance = start_database().await?;
|
||||
let admin_url = instance.url.clone();
|
||||
let client = instance.client.clone();
|
||||
for sql in [
|
||||
"CREATE DATABASE litellm",
|
||||
"CREATE TABLE litellm.otel_traces (n UInt8) ENGINE = Memory",
|
||||
|
|
@ -54,12 +39,13 @@ async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
|||
.await?
|
||||
.error_for_status()?;
|
||||
}
|
||||
let url = format!(
|
||||
"{}?database=litellm",
|
||||
admin_url.replacen("http://", "http://litellm_traces_reader:test_password@", 1)
|
||||
);
|
||||
let readers = QueryReaders::new(Connection::writer(&admin_url)?, "litellm".into());
|
||||
let connection = readers
|
||||
.connection(&client, &QueryScope::Admin, "test-secret")
|
||||
.await?;
|
||||
let url = connection.url().to_string();
|
||||
Ok(Database {
|
||||
_container: container,
|
||||
_database: instance,
|
||||
url,
|
||||
admin_url,
|
||||
client,
|
||||
|
|
@ -120,7 +106,15 @@ async fn reader_rejects_writes_and_privilege_escalation(
|
|||
|
||||
let result = read(&database.client, &connection, sql).await;
|
||||
|
||||
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
|
||||
assert!(
|
||||
matches!(
|
||||
result,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::QueryFailed(_)
|
||||
))
|
||||
),
|
||||
"{result:?}"
|
||||
);
|
||||
let rows = read(&database.client, &connection, "SELECT n FROM otel_traces").await?;
|
||||
let json: Value = serde_json::from_str(&rows)?;
|
||||
assert_eq!(json["data"], serde_json::json!([{ "n": 1 }]));
|
||||
|
|
@ -147,7 +141,12 @@ async fn admin_sql_rejects_errors_after_output_starts(
|
|||
.await;
|
||||
|
||||
assert!(
|
||||
matches!(result, Err(Error::InvalidResponse)),
|
||||
matches!(
|
||||
result,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::InvalidResponse
|
||||
))
|
||||
),
|
||||
"expected an error embedded in a successful HTTP response: {result:?}"
|
||||
);
|
||||
Ok(())
|
||||
|
|
@ -171,7 +170,15 @@ async fn admin_sql_enforces_result_row_limit(
|
|||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::QueryFailed(_))), "{result:?}");
|
||||
assert!(
|
||||
matches!(
|
||||
result,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::QueryFailed(_)
|
||||
))
|
||||
),
|
||||
"{result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -190,7 +197,15 @@ async fn admin_sql_enforces_response_byte_limit(
|
|||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(Error::ResponseTooLarge)), "{result:?}");
|
||||
assert!(
|
||||
matches!(
|
||||
result,
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::ResponseTooLarge
|
||||
))
|
||||
),
|
||||
"{result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
15
litellm-rust/crates/traces-clickhouse/tests/fixtures/README.md
vendored
Normal file
15
litellm-rust/crates/traces-clickhouse/tests/fixtures/README.md
vendored
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
# ClickHouse query fixtures
|
||||
|
||||
Run `cargo test -p litellm-traces-clickhouse --test queries --locked -- --test-threads=2` from `litellm-rust` with Docker running
|
||||
|
||||
Raw OTLP exports live in `crates/traces/tests/fixtures/query_*.json`. The seeded fixture decodes and normalizes them through `litellm_traces::decode_otlp` at test startup, then projects the decoded fields into ClickHouse columns. Team and key identities come from fixture setup rather than exporter claims. Root and child exports are inserted separately through the public insert API so materialized views process multiple blocks
|
||||
|
||||
`crates/traces/tests/fixtures/deeplite_auth_error.json` and `deeplite_swarm.json` were captured from Deeplite runs against the local proxy on 2026-10-02. The first contains a failed model call. The second contains successful model calls, searches, handoff attempts, and virtual filesystem writes. Credentials, workspace identifiers, and local user paths were redacted, and the protobuf exports were converted to OTLP JSON. Their round-trip tests check span identities, parent links, timestamps, durations, token counts, and statuses without pinning the provider's error wording
|
||||
|
||||
The swarm capture has handoff spans marked ERROR with `ParentCommand` exception events and a root with UNSET status. These are exported diagnostic statuses, which do not establish a failed execution. The tests preserve incoming statuses and check root status separately from the count of error spans, deriving both from the decoded export. They do not infer an execution outcome from exception text, framework names, successful model calls, or output presence. Framework-specific interpretation of control-flow exceptions belongs in the instrumentation integration
|
||||
|
||||
`spend_logs.jsonl` contains spend insert rows with millisecond timestamps, including two versions of one request. Replace this small placeholder dataset when the actual data is available. The query fixture applies production migrations, then removes TTL from its isolated database so fixed timestamps do not expire. Background merges are stopped so rollup aggregation and `FINAL` deduplication are exercised on unmerged data. Retention behavior stays covered by the migration tests
|
||||
|
||||
Curated SQL lives in `tests/queries/*.sql`. Each query has a matching `.expected.json` containing ordered result rows for `admin`, `team`, `key`, and `other_team` readers. Update the exports and expected results together. Add a named case in `tests/queries.rs` for each new query. Assertions compare only result data, excluding server statistics and execution timing
|
||||
|
||||
Typed query tests execute the production SQL through `litellm_storage_clickhouse::fetch` using contracts from `litellm-traces`. The fixture projection is test setup, so this suite covers the Rust decoder, normalization, inserts, schema, readers, and queries. Python ingress transformations, including payload truncation and exception-event fallback, remain covered by the Python tests
|
||||
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/spend_logs.jsonl
vendored
Normal file
3
litellm-rust/crates/traces-clickhouse/tests/fixtures/spend_logs.jsonl
vendored
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
{"request_id":"request-a","response_id":"response-shared","team_id":"team-a","api_key":"key-a","user":"user-a","spend":0.125,"start_time":1735689600100,"end_time":1735689600500,"metadata":"{\"labels\":{\"priority\":\"obsolete\"}}"}
|
||||
{"request_id":"request-a","response_id":"response-shared","team_id":"team-a","api_key":"key-a","user":"user-a","spend":0.5,"start_time":1735689600100,"end_time":1735689600600,"metadata":"{\"labels\":{\"priority\":\"high\"}}"}
|
||||
{"request_id":"request-b","response_id":"response-shared","team_id":"team-b","api_key":"key-b","user":"user-b","spend":0.25,"start_time":1735689602000,"end_time":1735689602500,"metadata":"{\"labels\":{\"priority\":\"low\"}}"}
|
||||
|
|
@ -5,8 +5,9 @@ use std::{
|
|||
|
||||
use flate2::read::GzDecoder;
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::{
|
||||
Connection, Error, InsertRow, InsertTable, Shared, encode_rows, insert_shared_rows,
|
||||
use litellm_traces::Shared;
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, InsertRow, InsertTable, encode_rows, insert_shared_rows,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
|
|
@ -1,45 +1,14 @@
|
|||
use std::{collections::BTreeMap, time::Duration};
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces::{
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, InsertTable, NORMALIZED_FIELD_DEFINITIONS, Parameter, ReadQuery,
|
||||
encode_rows, ensure_schema, execute_named_read, execute_read, schema_statements,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
use rstest::rstest;
|
||||
mod support;
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
|
||||
|
||||
struct ClickHouseDatabase {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
url: String,
|
||||
client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> TestResult<ClickHouseDatabase> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
);
|
||||
Ok(ClickHouseDatabase {
|
||||
_container: container,
|
||||
url,
|
||||
client: Client::no_redirect_for_test(),
|
||||
})
|
||||
}
|
||||
use support::{ClickHouseDatabase, TestResult, database};
|
||||
|
||||
async fn insert_rows(
|
||||
database: &ClickHouseDatabase,
|
||||
|
|
@ -115,7 +84,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
let span = serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": "trace-1", "SpanId": "span-1", "ParentSpanId": "",
|
||||
"ServiceName": "proxy", "SpanName": "request", "Input": "hello world",
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1"},
|
||||
"ResourceAttributes": {"litellm.team_id": "team-1", "litellm.api_key_hash": "hash-1", "litellm.user_id": "exporter-claim"},
|
||||
"SpanAttributes": {"gen_ai.response.id": "response-1", "gen_ai.usage.input_tokens": "12"}
|
||||
}))?;
|
||||
let spend = serde_json::from_value(serde_json::json!({
|
||||
|
|
@ -126,7 +95,29 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
insert_rows(&database, "otel_traces", vec![span]).await?;
|
||||
insert_rows(&database, "spend_logs", vec![spend]).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let detail =
|
||||
litellm_storage_clickhouse::fetch::<litellm_traces_clickhouse::query::named::SpanDetail>(
|
||||
&database.client,
|
||||
&reader,
|
||||
&litellm_traces_clickhouse::query::named::SpanDetailParams {
|
||||
access: litellm_traces_clickhouse::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: String::new(),
|
||||
team_ids: vec!["team-1".into()],
|
||||
api_key_hash: String::new(),
|
||||
},
|
||||
trace_id: "trace-1".into(),
|
||||
trace_ref: String::new(),
|
||||
span_id: "span-1".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
assert_eq!(detail.len(), 1);
|
||||
assert_eq!(detail[0].input, "hello world");
|
||||
assert_eq!(detail[0].attributes["gen_ai.response.id"], "response-1");
|
||||
let list_parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
|
|
@ -159,6 +150,8 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
"response_ids".into(),
|
||||
Parameter::Strings(vec!["response-1".into()]),
|
||||
),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
|
|
@ -182,7 +175,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
assert_eq!(matched["data"][0]["spend"], 0.125);
|
||||
let body = read_json(
|
||||
&database,
|
||||
"SELECT o.TeamId, o.ApiKeyHash, o.ObservationType, o.InputPreview, s.spend, \
|
||||
"SELECT o.TeamId, o.ApiKeyHash, o.UserId, o.ObservationType, o.InputPreview, s.spend, \
|
||||
toString(toUnixTimestamp64Nano(o.Timestamp)) AS timestamp_ns, \
|
||||
toString(toUnixTimestamp64Milli(s.start_time)) AS start_ms \
|
||||
FROM trace_test.otel_traces o JOIN trace_test.spend_logs s \
|
||||
|
|
@ -192,7 +185,7 @@ async fn schema_supports_span_rollups_and_spend_joins(
|
|||
assert_eq!(
|
||||
body["data"],
|
||||
serde_json::json!([{
|
||||
"TeamId": "team-1", "ApiKeyHash": "hash-1", "ObservationType": "agent",
|
||||
"TeamId": "team-1", "ApiKeyHash": "hash-1", "UserId": "", "ObservationType": "agent",
|
||||
"InputPreview": "hello world", "spend": 0.125,
|
||||
"timestamp_ns": timestamp.to_string(), "start_ms": (timestamp / 1_000_000).to_string()
|
||||
}])
|
||||
|
|
@ -269,7 +262,7 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
|
|||
]);
|
||||
|
||||
assert!(matches!(
|
||||
litellm_traces::insert_rows(
|
||||
litellm_traces_clickhouse::insert_rows(
|
||||
&database.client,
|
||||
&writer,
|
||||
"trace_test",
|
||||
|
|
@ -277,7 +270,9 @@ async fn insert_rejects_unknown_columns_even_if_url_requests_skipping_them(
|
|||
vec![row]
|
||||
)
|
||||
.await,
|
||||
Err(Error::InsertFailed(_))
|
||||
Err(Error::Storage(
|
||||
litellm_storage_clickhouse::Error::InsertFailed(_)
|
||||
))
|
||||
));
|
||||
assert_eq!(table_rows(&database, "otel_traces").await?, 0);
|
||||
Ok(())
|
||||
|
|
@ -297,7 +292,7 @@ async fn retried_trace_insert_does_not_inflate_rollup(
|
|||
"TeamId": "team-1", "ApiKeyHash": "key-1", "SpanName": "root", "InputTokens": 7
|
||||
}))?;
|
||||
for _ in 0..2 {
|
||||
litellm_traces::insert_rows(
|
||||
litellm_traces_clickhouse::insert_rows(
|
||||
&database.client,
|
||||
&writer,
|
||||
"trace_test",
|
||||
|
|
@ -412,7 +407,7 @@ async fn listed_agent_names_preserve_scope_and_cursor(
|
|||
),
|
||||
(
|
||||
"beta",
|
||||
"one",
|
||||
"other",
|
||||
"shared",
|
||||
"other_agent",
|
||||
"root",
|
||||
|
|
@ -446,7 +441,9 @@ async fn listed_agent_names_preserve_scope_and_cursor(
|
|||
insert_rows(&database, "otel_traces", historical_rows).await?;
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
("team_ids".into(), Parameter::Strings(vec!["alpha".into()])),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text("one".into())),
|
||||
(
|
||||
"start_ms".into(),
|
||||
|
|
@ -596,6 +593,8 @@ async fn rollup_merges_spans_across_days_without_losing_root_fields(
|
|||
);
|
||||
let connection = Connection::configured(&database.url, "trace_test", "default", "")?;
|
||||
let parameters = BTreeMap::from([
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team-1".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
(
|
||||
|
|
@ -768,7 +767,10 @@ async fn schema_statement_timeout_maps_to_transport_error() -> TestResult {
|
|||
)
|
||||
.await;
|
||||
server.abort();
|
||||
assert!(matches!(result, Ok(Err(Error::Transport))), "{result:?}");
|
||||
assert!(
|
||||
matches!(result, Ok(Err(Error::SchemaTransport))),
|
||||
"{result:?}"
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
|
@ -785,7 +787,7 @@ fn schema_rejects_invalid_configuration(#[case] database: &str, #[case] retentio
|
|||
async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::{LensQuery, Parameter};
|
||||
use litellm_traces_clickhouse::{Parameter, ReadQuery};
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -831,10 +833,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -842,6 +844,44 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
let rows = sample["data"].as_array().expect("sample rows");
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_ne!(rows[0]["trace_ref"], rows[1]["trace_ref"]);
|
||||
let identity_params = BTreeMap::from([
|
||||
("trace_id".into(), Parameter::Text("shared".into())),
|
||||
("all_teams".into(), Parameter::Integer(0)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec!["team".into()])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
]);
|
||||
let identities: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::TraceIdentity,
|
||||
&identity_params,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(identities["data"].as_array().map(Vec::len), Some(2));
|
||||
let key_params = identity_params
|
||||
.into_iter()
|
||||
.chain([
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text("one".into())),
|
||||
])
|
||||
.collect();
|
||||
let identity: serde_json::Value = serde_json::from_str(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
ReadQuery::TraceIdentity,
|
||||
&key_params,
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(identity["data"].as_array().map(Vec::len), Some(1));
|
||||
assert!(
|
||||
rows.iter()
|
||||
.any(|row| row["trace_ref"] == identity["data"][0]["trace_ref"])
|
||||
);
|
||||
let first_ref = rows[0]["trace_ref"].as_str().expect("reference");
|
||||
let read_parameters: BTreeMap<_, _> = parameters
|
||||
.into_iter()
|
||||
|
|
@ -855,10 +895,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
])
|
||||
.collect();
|
||||
let content: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
&read_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -875,10 +915,10 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
.chain([("quote".into(), Parameter::Text(opposite.into()))])
|
||||
.collect();
|
||||
let evidence: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Evidence.sql(),
|
||||
ReadQuery::Evidence,
|
||||
&evidence_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -892,7 +932,7 @@ async fn lens_filters_reads_and_evidence_keep_reused_trace_ids_separate(
|
|||
async fn lens_request_sample_does_not_trust_caller_tags(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::{LensQuery, Parameter};
|
||||
use litellm_traces_clickhouse::{Parameter, ReadQuery};
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -927,10 +967,10 @@ async fn lens_request_sample_does_not_trust_caller_tags(
|
|||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -957,7 +997,7 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
|
|||
#[case] page_size: usize,
|
||||
#[case] changing: bool,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database?;
|
||||
ensure_schema(
|
||||
&database.client,
|
||||
|
|
@ -996,10 +1036,10 @@ async fn lens_selection_pages_without_losing_or_repeating_runs(
|
|||
("selected_team".into(), Parameter::Text(String::new())),
|
||||
("execution_ids".into(), Parameter::Strings(vec![])),
|
||||
]);
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -1037,7 +1077,7 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] input_length: usize,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database?;
|
||||
ensure_schema(
|
||||
&database.client,
|
||||
|
|
@ -1061,10 +1101,10 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
("cursor".into(), Parameter::Text(String::new())),
|
||||
("offset".into(), Parameter::Integer(1)),
|
||||
]);
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -1083,10 +1123,10 @@ async fn lens_content_keeps_output_visible_after_long_input(
|
|||
let mut recovered = String::new();
|
||||
for offset in (2..original.len() + 2).step_by(8000) {
|
||||
parameters.insert("offset".into(), Parameter::Integer(offset as i64));
|
||||
let body = execute_read(
|
||||
let body = execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Content.sql(),
|
||||
ReadQuery::Content,
|
||||
¶meters,
|
||||
)
|
||||
.await?;
|
||||
|
|
@ -1127,6 +1167,8 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
"trace_id".into(),
|
||||
Parameter::Text("diagnostic-trace".into()),
|
||||
),
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
|
|
@ -1174,6 +1216,7 @@ async fn trace_error_previews_preserve_paginated_diagnostics(
|
|||
}
|
||||
}
|
||||
assert_eq!(recovered, message);
|
||||
parameters.insert("all_teams".into(), Parameter::Integer(0));
|
||||
parameters.insert(
|
||||
"api_key_hash".into(),
|
||||
Parameter::Text("unrelated-key".into()),
|
||||
|
|
@ -1219,6 +1262,8 @@ async fn duplicate_span_preview_matches_diagnostic(
|
|||
let parameters = BTreeMap::from([
|
||||
("trace_id".into(), Parameter::Text("duplicate-trace".into())),
|
||||
("span_id".into(), Parameter::Text("duplicate-span".into())),
|
||||
("all_teams".into(), Parameter::Integer(1)),
|
||||
("user_id".into(), Parameter::Text(String::new())),
|
||||
("team_ids".into(), Parameter::Strings(vec![])),
|
||||
("api_key_hash".into(), Parameter::Text(String::new())),
|
||||
("trace_ref".into(), Parameter::Text(String::new())),
|
||||
|
|
@ -1257,7 +1302,7 @@ fn schema_includes_every_migration_file() -> TestResult {
|
|||
async fn lens_agent_discovery_and_selection_preserve_scope(
|
||||
#[future] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
use litellm_traces::LensQuery;
|
||||
use litellm_traces_clickhouse::ReadQuery;
|
||||
let database = database.await?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
|
|
@ -1288,10 +1333,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
("key_hash".into(), Parameter::Text("one".into())),
|
||||
]);
|
||||
let agents: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Agents.sql(),
|
||||
ReadQuery::Agents,
|
||||
&scope_parameters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1332,10 +1377,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
])
|
||||
.collect::<BTreeMap<_, _>>();
|
||||
let sample: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Sample.sql(),
|
||||
ReadQuery::Sample,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1344,10 +1389,10 @@ async fn lens_agent_discovery_and_selection_preserve_scope(
|
|||
assert_eq!(sample["data"][0]["trace_id"], "research");
|
||||
assert_eq!(sample["data"][0]["span_count"], 2);
|
||||
let available: serde_json::Value = serde_json::from_str(
|
||||
&execute_read(
|
||||
&execute_named_read(
|
||||
&database.client,
|
||||
&connection,
|
||||
LensQuery::Availability.sql(),
|
||||
ReadQuery::Availability,
|
||||
¶meters,
|
||||
)
|
||||
.await?,
|
||||
|
|
@ -1422,8 +1467,9 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
)
|
||||
.await?;
|
||||
}
|
||||
let help: serde_json::Value =
|
||||
serde_json::from_str(&litellm_traces::query_help(&database.client, &reader).await?)?;
|
||||
let help: serde_json::Value = serde_json::from_str(
|
||||
&litellm_traces_clickhouse::query_help(&database.client, &reader).await?,
|
||||
)?;
|
||||
let keys: std::collections::BTreeSet<_> = help
|
||||
.as_object()
|
||||
.ok_or("missing help object")?
|
||||
|
|
@ -1533,7 +1579,8 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
"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 body =
|
||||
litellm_traces_clickhouse::query_sql(&database.client, &reader, &sql).await?;
|
||||
let values: serde_json::Value = serde_json::from_str(&body)?;
|
||||
assert_ne!(values["data"][0]["value"], "");
|
||||
}
|
||||
|
|
@ -1551,7 +1598,7 @@ async fn query_help_discovers_live_schema_and_runs_its_examples(
|
|||
.collect::<std::collections::BTreeSet<_>>(),
|
||||
std::collections::BTreeSet::from(["name", "sql"])
|
||||
);
|
||||
let body = litellm_traces::query_sql(&database.client, &reader, sql).await?;
|
||||
let body = litellm_traces_clickhouse::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(),
|
||||
|
|
@ -1607,8 +1654,9 @@ 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::query_help(&database.client, &reader).await?)?;
|
||||
let help: serde_json::Value = serde_json::from_str(
|
||||
&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());
|
||||
assert_eq!(
|
||||
|
|
@ -1642,3 +1690,282 @@ async fn query_help_preserves_schema_and_guide_when_discovery_hits_reader_limits
|
|||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn field_definitions_match_serialized_normalized_span() {
|
||||
use litellm_traces::decode_otlp;
|
||||
use std::collections::BTreeSet;
|
||||
let spans = decode_otlp(
|
||||
br#"{"resourceSpans":[{"scopeSpans":[{"spans":[{"traceId":"11111111111111111111111111111111","spanId":"2222222222222222","name":"root"}]}]}]}"#,
|
||||
Some("application/json"),
|
||||
)
|
||||
.expect("valid OTLP");
|
||||
let fields = &spans[0].normalized;
|
||||
let serialized = serde_json::to_value(fields).expect("serializable fields");
|
||||
let keys: BTreeSet<_> = serialized
|
||||
.as_object()
|
||||
.expect("field object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(|field| field.name)
|
||||
.collect();
|
||||
assert_eq!(keys, mapped);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user("owner", vec![], "", vec!["own"])]
|
||||
#[case::own_user_and_permitted_team("owner", vec!["permitted"], "", vec!["own", "team"])]
|
||||
#[case::key_only("", vec![], "request-key", vec!["own"])]
|
||||
#[case::no_identity("", vec![], "", vec![])]
|
||||
#[tokio::test]
|
||||
async fn named_and_sql_readers_share_request_log_visibility(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] key: &str,
|
||||
#[case] expected: Vec<&str>,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{
|
||||
ReadAccessParams, SpendByResponseIds, SpendByResponseIdsParams,
|
||||
};
|
||||
use litellm_traces_clickhouse::{QueryReaders, QueryScope};
|
||||
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [("own", "unpermitted", "owner", "request-key"), ("team", "permitted", "other", "other-key"), ("foreign", "foreign", "other", "foreign-key")]
|
||||
.into_iter()
|
||||
.map(|(id, team, owner, api_key)| serde_json::from_value(serde_json::json!({
|
||||
"request_id": id, "response_id": "shared-response", "team_id": team, "user": owner,
|
||||
"api_key": api_key, "spend": 0.25, "start_time": timestamp / 1_000_000, "end_time": timestamp / 1_000_000,
|
||||
})))
|
||||
.collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
|
||||
insert_rows(&database, "spend_logs", rows).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let params =
|
||||
SpendByResponseIdsParams::from(litellm_traces::query::named::SpendByResponseIdsParams {
|
||||
access: ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: user.into(),
|
||||
team_ids: teams.iter().map(|team| (*team).into()).collect(),
|
||||
api_key_hash: key.into(),
|
||||
},
|
||||
response_ids: vec!["shared-response".into()],
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
});
|
||||
let spend =
|
||||
litellm_storage_clickhouse::fetch::<SpendByResponseIds>(&database.client, &reader, ¶ms)
|
||||
.await?;
|
||||
let actual: std::collections::BTreeSet<_> =
|
||||
spend.iter().map(|row| row.0.request_id.as_str()).collect();
|
||||
let expected: std::collections::BTreeSet<_> = expected.into_iter().collect();
|
||||
assert_eq!(actual, expected);
|
||||
let scope = QueryScope::Logs {
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
api_key_hash: key.into(),
|
||||
};
|
||||
if user.is_empty() && scope.validate().is_err() {
|
||||
assert!(
|
||||
QueryReaders::new(writer, "trace_test".into())
|
||||
.connection(&database.client, &scope, "secret")
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
let scoped = QueryReaders::new(writer, "trace_test".into())
|
||||
.connection(&database.client, &scope, "secret")
|
||||
.await?;
|
||||
let result: serde_json::Value = serde_json::from_str(
|
||||
&litellm_traces_clickhouse::query_sql(
|
||||
&database.client,
|
||||
&scoped,
|
||||
"SELECT request_id FROM spend_logs FINAL ORDER BY request_id",
|
||||
)
|
||||
.await?,
|
||||
)?;
|
||||
assert_eq!(
|
||||
result["data"],
|
||||
serde_json::json!(
|
||||
expected
|
||||
.into_iter()
|
||||
.map(|id| serde_json::json!({"request_id": id}))
|
||||
.collect::<Vec<_>>()
|
||||
)
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn rollup_cost_completeness_preserves_missing_ids_and_fails_closed_for_historical_rows(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let initial_mutations = mutation_rows(&database).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [("complete", "llm", "response"), ("complete", "llm", "response"), ("complete", "agent", ""), ("missing", "llm", "response"), ("missing", "llm", ""), ("missing", "agent", "extra-id"), ("mixed", "llm", "mine"), ("mixed", "llm", "other")]
|
||||
.into_iter().enumerate().map(|(index, (trace, kind, id))| serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp, "TraceId": trace, "SpanId": index.to_string(), "TeamId": "team", "ApiKeyHash": "export",
|
||||
"UserId": if id == "other" { "other" } else { "owner" }, "ObservationType": kind, "LiteLLMRequestId": id,
|
||||
}))).collect::<Result<Vec<BTreeMap<String, serde_json::Value>>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", rows).await?;
|
||||
execute_write(&database, &format!(
|
||||
"INSERT INTO trace_test.agent_traces_by_key (TeamId, ApiKeyHash, TraceId, StartTs, EndTs, LlmCount, RequestIds) \
|
||||
VALUES ('team', 'export', 'historical', fromUnixTimestamp64Nano({timestamp}), fromUnixTimestamp64Nano({timestamp}), 2, ['response', 'non-llm-id'])"
|
||||
)).await?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let params = litellm_traces_clickhouse::query::named::ListTracesParams::from(
|
||||
litellm_traces::query::named::ListTracesParams {
|
||||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
all_teams: 0,
|
||||
user_id: "".into(),
|
||||
team_ids: vec!["team".into()],
|
||||
api_key_hash: "".into(),
|
||||
},
|
||||
start_ms: timestamp / 1_000_000 - 1,
|
||||
end_ms: timestamp / 1_000_000 + 1,
|
||||
cursor_ms: 0,
|
||||
cursor_trace_id: "".into(),
|
||||
limit: 10,
|
||||
},
|
||||
);
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let listed = litellm_storage_clickhouse::fetch::<
|
||||
litellm_traces_clickhouse::query::named::ListTraces,
|
||||
>(&database.client, &reader, ¶ms)
|
||||
.await?;
|
||||
assert_eq!(listed.len(), 4);
|
||||
let owned_params = litellm_traces_clickhouse::query::named::ListTracesParams::from(
|
||||
litellm_traces::query::named::ListTracesParams {
|
||||
access: litellm_traces::query::named::ReadAccessParams {
|
||||
user_id: "owner".into(),
|
||||
team_ids: vec![],
|
||||
all_teams: 0,
|
||||
api_key_hash: String::new(),
|
||||
},
|
||||
..params.0
|
||||
},
|
||||
);
|
||||
let owned = litellm_storage_clickhouse::fetch::<
|
||||
litellm_traces_clickhouse::query::named::ListTraces,
|
||||
>(&database.client, &reader, &owned_params)
|
||||
.await?;
|
||||
assert_eq!(owned.len(), 2);
|
||||
assert!(
|
||||
owned
|
||||
.iter()
|
||||
.all(|row| ["complete", "missing"].contains(&row.0.trace_id.as_str()))
|
||||
);
|
||||
for row in listed {
|
||||
match row.0.trace_id.as_str() {
|
||||
"complete" => {
|
||||
assert_eq!(row.0.user_id, "owner");
|
||||
assert_eq!(row.0.request_ids, ["response"]);
|
||||
assert_eq!(row.0.llm_calls, 2);
|
||||
}
|
||||
"missing" | "historical" => assert!(row.0.request_ids.iter().any(String::is_empty)),
|
||||
"mixed" => assert!(row.0.user_id.is_empty()),
|
||||
id => panic!("unexpected trace {id}"),
|
||||
}
|
||||
}
|
||||
assert_eq!(mutation_rows(&database).await?, initial_mutations);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::admin(1, "", vec![], "", "own answer")]
|
||||
#[case::user(0, "owner", vec![], "", "own answer")]
|
||||
#[case::team(0, "", vec!["alpha"], "", "own answer")]
|
||||
#[case::key(0, "", vec![], "one", "own answer")]
|
||||
#[case::no_identity(0, "", vec![], "", "")]
|
||||
#[tokio::test]
|
||||
async fn agent_final_answer_preserves_visibility_and_trace_ownership(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
#[case] all_teams: u8,
|
||||
#[case] user: &str,
|
||||
#[case] teams: Vec<&str>,
|
||||
#[case] key: &str,
|
||||
#[case] expected: &str,
|
||||
) -> TestResult {
|
||||
use litellm_traces_clickhouse::query::named::{ReadAccessParams, SpanDetail, SpanDetailParams};
|
||||
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, "trace_test", 7).await?;
|
||||
let timestamp = time::OffsetDateTime::now_utc().unix_timestamp_nanos() as i64;
|
||||
let rows = [
|
||||
("alpha", "one", "owner", "root", "", "agent", ""),
|
||||
(
|
||||
"alpha",
|
||||
"one",
|
||||
"owner",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"own answer",
|
||||
),
|
||||
(
|
||||
"alpha",
|
||||
"two",
|
||||
"other",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"other key answer",
|
||||
),
|
||||
(
|
||||
"beta",
|
||||
"one",
|
||||
"other",
|
||||
"child",
|
||||
"root",
|
||||
"llm",
|
||||
"other team answer",
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, (team, key, user, span, parent, kind, output))| {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"Timestamp": timestamp + index as i64, "TraceId": "shared", "SpanId": span,
|
||||
"ParentSpanId": parent, "TeamId": team, "ApiKeyHash": key, "UserId": user,
|
||||
"ObservationType": kind, "Input": "prompt", "Output": output,
|
||||
}))
|
||||
})
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(&database, "otel_traces", rows).await?;
|
||||
let reader = Connection::reader(&database.url, "trace_test")?;
|
||||
let details = litellm_storage_clickhouse::fetch::<SpanDetail>(
|
||||
&database.client,
|
||||
&reader,
|
||||
&SpanDetailParams {
|
||||
access: ReadAccessParams {
|
||||
all_teams,
|
||||
user_id: user.into(),
|
||||
team_ids: teams.into_iter().map(str::to_owned).collect(),
|
||||
api_key_hash: key.into(),
|
||||
},
|
||||
trace_id: "shared".into(),
|
||||
trace_ref: String::new(),
|
||||
span_id: "root".into(),
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if expected.is_empty() {
|
||||
assert!(details.is_empty());
|
||||
} else {
|
||||
assert_eq!(details.len(), 1);
|
||||
assert_eq!(details[0].input, "prompt");
|
||||
assert_eq!(details[0].output, expected);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
273
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal file
273
litellm-rust/crates/traces-clickhouse/tests/queries.rs
Normal file
|
|
@ -0,0 +1,273 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_storage_clickhouse::fetch;
|
||||
use litellm_traces::query::named as contracts;
|
||||
use litellm_traces_clickhouse::{
|
||||
QueryScope,
|
||||
query::named::{ListTraces, ListTracesParams, TraceSpans, TraceSpansParams},
|
||||
query_sql,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde::Deserialize;
|
||||
use serde_json::Value;
|
||||
|
||||
#[path = "queries/support.rs"]
|
||||
mod fixtures;
|
||||
mod support;
|
||||
|
||||
use fixtures::{SeededDatabase, insert_export, migrated_database, seeded_database};
|
||||
use support::TestResult;
|
||||
|
||||
#[derive(Clone, Copy, strum::AsRefStr)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
enum ScopeCase {
|
||||
Admin,
|
||||
Team,
|
||||
Key,
|
||||
OtherTeam,
|
||||
}
|
||||
|
||||
impl ScopeCase {
|
||||
fn scope(self) -> QueryScope {
|
||||
match self {
|
||||
Self::Admin => QueryScope::Admin,
|
||||
Self::Team => QueryScope::Team {
|
||||
team_id: "team-a".into(),
|
||||
},
|
||||
Self::Key => QueryScope::Key {
|
||||
team_id: "team-a".into(),
|
||||
api_key_hash: "key-a".into(),
|
||||
},
|
||||
Self::OtherTeam => QueryScope::Team {
|
||||
team_id: "team-b".into(),
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct QueryResult {
|
||||
data: Vec<Value>,
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::rollups(include_str!("queries/rollups.sql"), include_str!("queries/rollups.expected.json"))]
|
||||
#[case::costs(include_str!("queries/trace_costs.sql"), include_str!("queries/trace_costs.expected.json"))]
|
||||
#[case::errors(include_str!("queries/failed_spans.sql"), include_str!("queries/failed_spans.expected.json"))]
|
||||
#[case::metadata(include_str!("queries/metadata_filters.sql"), include_str!("queries/metadata_filters.expected.json"))]
|
||||
#[tokio::test]
|
||||
async fn curated_queries_return_expected_rows(
|
||||
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
|
||||
#[case] sql: &str,
|
||||
#[case] expected_json: &str,
|
||||
#[values(
|
||||
ScopeCase::Admin,
|
||||
ScopeCase::Team,
|
||||
ScopeCase::Key,
|
||||
ScopeCase::OtherTeam
|
||||
)]
|
||||
scope: ScopeCase,
|
||||
) -> TestResult {
|
||||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(&fixture.database.client, &scope.scope(), "fixture-secret")
|
||||
.await?;
|
||||
let result: QueryResult =
|
||||
serde_json::from_str(&query_sql(&fixture.database.client, &reader, sql).await?)?;
|
||||
let expected: BTreeMap<String, Vec<Value>> = serde_json::from_str(expected_json)?;
|
||||
assert_eq!(
|
||||
&result.data,
|
||||
expected
|
||||
.get(scope.as_ref())
|
||||
.ok_or("missing expected scope")?,
|
||||
"{}: {sql}",
|
||||
scope.as_ref()
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn admin_access() -> TestResult<contracts::ReadAccessParams> {
|
||||
Ok(serde_json::from_str(include_str!(
|
||||
"queries/read_access.json"
|
||||
))?)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn typed_queries_read_normalized_spans_and_keep_trace_identities_separate(
|
||||
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
|
||||
admin_access: TestResult<contracts::ReadAccessParams>,
|
||||
) -> TestResult {
|
||||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.await?;
|
||||
let params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: admin_access?,
|
||||
start_ms: 0,
|
||||
end_ms: i64::MAX / 1_000_000,
|
||||
cursor_ms: 0,
|
||||
cursor_trace_id: String::new(),
|
||||
limit: 10,
|
||||
});
|
||||
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, ¶ms).await?;
|
||||
assert_eq!(
|
||||
traces
|
||||
.iter()
|
||||
.map(|row| row.0.api_key_hash.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["key-b", "key-alt", "key-a"]
|
||||
);
|
||||
let trace = &traces[2].0;
|
||||
assert_eq!(
|
||||
(
|
||||
trace.span_count,
|
||||
trace.llm_calls,
|
||||
trace.tool_calls,
|
||||
trace.error_count
|
||||
),
|
||||
(3, 1, 1, 1)
|
||||
);
|
||||
assert_eq!((trace.input_tokens, trace.output_tokens), (12, 6));
|
||||
assert_eq!(trace.input_preview, "Review the change");
|
||||
let span_params = TraceSpansParams {
|
||||
access: params.0.access,
|
||||
trace_id: trace.trace_id.clone(),
|
||||
trace_ref: trace.trace_ref.clone(),
|
||||
};
|
||||
let spans = fetch::<TraceSpans>(&fixture.database.client, &reader, &span_params).await?;
|
||||
assert_eq!(
|
||||
spans
|
||||
.iter()
|
||||
.map(|row| row.0.name.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["review", "completion", "lookup"]
|
||||
);
|
||||
assert!(
|
||||
spans
|
||||
.iter()
|
||||
.all(|row| row.0.api_key_hash == trace.api_key_hash)
|
||||
);
|
||||
assert_eq!(
|
||||
(
|
||||
spans[1].0.kind.as_str(),
|
||||
spans[1].0.input_tokens,
|
||||
spans[1].0.output_tokens
|
||||
),
|
||||
("llm", 12, 6)
|
||||
);
|
||||
assert_eq!(spans[2].0.status_message, "lookup timed out");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn typed_trace_cursor_returns_the_next_fixture_trace(
|
||||
#[future(awt)] seeded_database: TestResult<SeededDatabase>,
|
||||
admin_access: TestResult<contracts::ReadAccessParams>,
|
||||
) -> TestResult {
|
||||
let fixture = seeded_database?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.await?;
|
||||
let params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: admin_access?,
|
||||
start_ms: 0,
|
||||
end_ms: i64::MAX / 1_000_000,
|
||||
cursor_ms: 0,
|
||||
cursor_trace_id: String::new(),
|
||||
limit: 1,
|
||||
});
|
||||
let first = fetch::<ListTraces>(&fixture.database.client, &reader, ¶ms).await?;
|
||||
assert_eq!(first.len(), 1);
|
||||
assert_eq!(first[0].0.api_key_hash, "key-b");
|
||||
let next_params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
cursor_ms: first[0].0.start_ms,
|
||||
cursor_trace_id: first[0].0.trace_ref.clone(),
|
||||
..params.0
|
||||
});
|
||||
let next = fetch::<ListTraces>(&fixture.database.client, &reader, &next_params).await?;
|
||||
assert_eq!(next.len(), 1);
|
||||
assert_eq!(next[0].0.api_key_hash, "key-alt");
|
||||
assert_ne!(first[0].0.trace_ref, next[0].0.trace_ref);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::authentication_error(include_bytes!("../../traces/tests/fixtures/deeplite_auth_error.json"))]
|
||||
#[case::swarm(include_bytes!("../../traces/tests/fixtures/deeplite_swarm.json"))]
|
||||
#[tokio::test]
|
||||
async fn captured_deeplite_exports_round_trip_through_clickhouse(
|
||||
#[future(awt)] migrated_database: TestResult<SeededDatabase>,
|
||||
admin_access: TestResult<contracts::ReadAccessParams>,
|
||||
#[case] export: &[u8],
|
||||
) -> TestResult {
|
||||
let fixture = migrated_database?;
|
||||
let decoded = insert_export(&fixture, export, "team-a", "key-a").await?;
|
||||
let reader = fixture
|
||||
.readers
|
||||
.connection(
|
||||
&fixture.database.client,
|
||||
&QueryScope::Admin,
|
||||
"fixture-secret",
|
||||
)
|
||||
.await?;
|
||||
let params = TraceSpansParams {
|
||||
access: admin_access?,
|
||||
trace_id: decoded[0].trace_id.clone(),
|
||||
trace_ref: String::new(),
|
||||
};
|
||||
let stored = fetch::<TraceSpans>(&fixture.database.client, &reader, ¶ms).await?;
|
||||
assert_eq!(stored.len(), decoded.len());
|
||||
let list_params = ListTracesParams::from(contracts::ListTracesParams {
|
||||
access: params.access,
|
||||
start_ms: 0,
|
||||
end_ms: i64::MAX / 1_000_000,
|
||||
cursor_ms: 0,
|
||||
cursor_trace_id: String::new(),
|
||||
limit: 10,
|
||||
});
|
||||
let traces = fetch::<ListTraces>(&fixture.database.client, &reader, &list_params).await?;
|
||||
assert_eq!(traces.len(), 1);
|
||||
let roots = decoded
|
||||
.iter()
|
||||
.filter(|span| span.parent_span_id.is_empty())
|
||||
.collect::<Vec<_>>();
|
||||
assert_eq!(roots.len(), 1);
|
||||
assert_eq!(traces[0].0.status, roots[0].status_code);
|
||||
assert_eq!(
|
||||
traces[0].0.error_count,
|
||||
decoded
|
||||
.iter()
|
||||
.filter(|span| span.status_code == "STATUS_CODE_ERROR")
|
||||
.count() as u64
|
||||
);
|
||||
let by_id: BTreeMap<_, _> = stored
|
||||
.iter()
|
||||
.map(|row| (row.0.span_id.as_str(), &row.0))
|
||||
.collect();
|
||||
for span in &decoded {
|
||||
let row = by_id
|
||||
.get(span.span_id.as_str())
|
||||
.ok_or("missing captured span")?;
|
||||
assert_eq!(row.parent_span_id, span.parent_span_id);
|
||||
assert_eq!(row.start_ns as u64, span.start_ns);
|
||||
assert_eq!(row.duration_ns, span.end_ns - span.start_ns);
|
||||
assert_eq!(row.input_tokens, span.normalized.input_tokens);
|
||||
assert_eq!(row.output_tokens, span.normalized.output_tokens);
|
||||
assert_eq!(row.status, span.status_code);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"span_id": "0303030303030303",
|
||||
"message": "lookup timed out"
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"span_id": "0303030303030303",
|
||||
"message": "lookup timed out"
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"span_id": "0303030303030303",
|
||||
"message": "lookup timed out"
|
||||
}
|
||||
],
|
||||
"other_team": []
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
SELECT TeamId AS team, ApiKeyHash AS api_key, TraceId AS trace_id,
|
||||
SpanId AS span_id, StatusMessage AS message
|
||||
FROM otel_traces
|
||||
WHERE StatusCode = 'STATUS_CODE_ERROR'
|
||||
ORDER BY team, api_key, trace_id, span_id
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"request_id": "request-a",
|
||||
"spend": 0.5,
|
||||
"priority": "high"
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"request_id": "request-a",
|
||||
"spend": 0.5,
|
||||
"priority": "high"
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"request_id": "request-a",
|
||||
"spend": 0.5,
|
||||
"priority": "high"
|
||||
}
|
||||
],
|
||||
"other_team": []
|
||||
}
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
SELECT team_id AS team, api_key, request_id, spend,
|
||||
JSONExtractString(metadata, 'labels', 'priority') AS priority
|
||||
FROM spend_logs FINAL
|
||||
WHERE JSONExtractString(metadata, 'labels', 'priority') = 'high'
|
||||
ORDER BY team, api_key, request_id
|
||||
|
|
@ -0,0 +1,6 @@
|
|||
{
|
||||
"all_teams": 1,
|
||||
"user_id": "",
|
||||
"team_ids": ["team-a", "team-b"],
|
||||
"api_key_hash": ""
|
||||
}
|
||||
|
|
@ -0,0 +1,87 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
},
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-alt",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "alternate",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
},
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "other-team",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
},
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-alt",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "alternate",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "review",
|
||||
"spans": 3,
|
||||
"llm_calls": 1,
|
||||
"errors": 1,
|
||||
"input_tokens": 12,
|
||||
"output_tokens": 6
|
||||
}
|
||||
],
|
||||
"other_team": [
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"name": "other-team",
|
||||
"spans": 1,
|
||||
"llm_calls": 0,
|
||||
"errors": 0,
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
SELECT TeamId AS team, ApiKeyHash AS api_key, TraceId AS trace_id,
|
||||
ifNull(any(RootName), '') AS name,
|
||||
toUInt32(sum(SpanCount)) AS spans,
|
||||
toUInt32(sum(LlmCount)) AS llm_calls,
|
||||
toUInt32(sum(ErrorCount)) AS errors,
|
||||
toUInt32(sum(InputTokens)) AS input_tokens,
|
||||
toUInt32(sum(OutputTokens)) AS output_tokens
|
||||
FROM agent_traces_by_key
|
||||
GROUP BY TeamId, ApiKeyHash, TraceId
|
||||
ORDER BY team, api_key, trace_id
|
||||
153
litellm-rust/crates/traces-clickhouse/tests/queries/support.rs
Normal file
153
litellm-rust/crates/traces-clickhouse/tests/queries/support.rs
Normal file
|
|
@ -0,0 +1,153 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_traces::{DecodedSpan, decode_otlp};
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, InsertTable, QueryReaders, ensure_schema, insert_rows,
|
||||
};
|
||||
use rstest::fixture;
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use crate::support::{ClickHouseDatabase, TestResult, database};
|
||||
|
||||
pub const DATABASE: &str = "trace_test";
|
||||
|
||||
pub struct SeededDatabase {
|
||||
pub database: ClickHouseDatabase,
|
||||
pub readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn migrated_database(
|
||||
#[future(awt)] database: TestResult<ClickHouseDatabase>,
|
||||
) -> TestResult<SeededDatabase> {
|
||||
let database = database?;
|
||||
let writer = Connection::writer(&database.url)?;
|
||||
ensure_schema(&database.client, &writer, DATABASE, 7).await?;
|
||||
for table in ["otel_traces", "agent_traces_by_key", "spend_logs"] {
|
||||
database
|
||||
.client
|
||||
.post(writer.url().clone())
|
||||
.body(format!("ALTER TABLE {DATABASE}.{table} REMOVE TTL"))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
database
|
||||
.client
|
||||
.post(writer.url().clone())
|
||||
.body(format!("SYSTEM STOP MERGES {DATABASE}.{table}"))
|
||||
.send()
|
||||
.await?
|
||||
.error_for_status()?;
|
||||
}
|
||||
let readers = QueryReaders::new(writer, DATABASE.to_owned());
|
||||
Ok(SeededDatabase { database, readers })
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn seeded_database(
|
||||
#[future(awt)] migrated_database: TestResult<SeededDatabase>,
|
||||
) -> TestResult<SeededDatabase> {
|
||||
let fixture = migrated_database?;
|
||||
let writer = Connection::writer(&fixture.database.url)?;
|
||||
for (contents, team, key) in [
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_root.json").as_slice(),
|
||||
"team-a",
|
||||
"key-a",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_children.json").as_slice(),
|
||||
"team-a",
|
||||
"key-a",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_alternate.json").as_slice(),
|
||||
"team-a",
|
||||
"key-alt",
|
||||
),
|
||||
(
|
||||
include_bytes!("../../../traces/tests/fixtures/query_other_team.json").as_slice(),
|
||||
"team-b",
|
||||
"key-b",
|
||||
),
|
||||
] {
|
||||
insert_export(&fixture, contents, team, key).await?;
|
||||
}
|
||||
let spend_rows = include_str!("../fixtures/spend_logs.jsonl")
|
||||
.lines()
|
||||
.map(serde_json::from_str::<BTreeMap<String, Value>>)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
insert_rows(
|
||||
&fixture.database.client,
|
||||
&writer,
|
||||
DATABASE,
|
||||
InsertTable::SpendLogs,
|
||||
spend_rows,
|
||||
)
|
||||
.await?;
|
||||
Ok(fixture)
|
||||
}
|
||||
|
||||
pub async fn insert_export(
|
||||
fixture: &SeededDatabase,
|
||||
contents: &[u8],
|
||||
team: &str,
|
||||
key: &str,
|
||||
) -> TestResult<Vec<DecodedSpan>> {
|
||||
let spans = decode_otlp(contents, Some("application/json"))?;
|
||||
let writer = Connection::writer(&fixture.database.url)?;
|
||||
let rows = spans.iter().map(|span| span_row(span, team, key)).collect();
|
||||
insert_rows(
|
||||
&fixture.database.client,
|
||||
&writer,
|
||||
DATABASE,
|
||||
InsertTable::OtelTraces,
|
||||
rows,
|
||||
)
|
||||
.await?;
|
||||
Ok(spans)
|
||||
}
|
||||
|
||||
fn span_row(span: &DecodedSpan, team: &str, key: &str) -> BTreeMap<String, Value> {
|
||||
BTreeMap::from([
|
||||
("Timestamp".into(), json!(span.start_ns)),
|
||||
("TraceId".into(), json!(span.trace_id)),
|
||||
("SpanId".into(), json!(span.span_id)),
|
||||
("ParentSpanId".into(), json!(span.parent_span_id)),
|
||||
("TraceState".into(), json!(span.trace_state)),
|
||||
("SpanName".into(), json!(span.name)),
|
||||
("SpanKind".into(), json!(span.kind)),
|
||||
(
|
||||
"ServiceName".into(),
|
||||
json!(
|
||||
span.resource_attributes
|
||||
.get("service.name")
|
||||
.map(String::as_str)
|
||||
.unwrap_or_default()
|
||||
),
|
||||
),
|
||||
("ResourceAttributes".into(), json!(span.resource_attributes)),
|
||||
("ScopeName".into(), json!(span.scope_name)),
|
||||
("ScopeVersion".into(), json!(span.scope_version)),
|
||||
("SpanAttributes".into(), json!(span.attributes)),
|
||||
("Duration".into(), json!(span.end_ns - span.start_ns)),
|
||||
("StatusCode".into(), json!(span.status_code)),
|
||||
("StatusMessage".into(), json!(span.status_message)),
|
||||
("TeamId".into(), json!(team)),
|
||||
("ApiKeyHash".into(), json!(key)),
|
||||
(
|
||||
"ObservationType".into(),
|
||||
json!(span.normalized.observation_type),
|
||||
),
|
||||
("AgentName".into(), json!(span.normalized.agent_name)),
|
||||
("Model".into(), json!(span.normalized.model)),
|
||||
(
|
||||
"LiteLLMRequestId".into(),
|
||||
json!(span.normalized.litellm_request_id),
|
||||
),
|
||||
("InputTokens".into(), json!(span.normalized.input_tokens)),
|
||||
("OutputTokens".into(), json!(span.normalized.output_tokens)),
|
||||
("Input".into(), json!(span.normalized.input)),
|
||||
("Output".into(), json!(span.normalized.output)),
|
||||
])
|
||||
}
|
||||
|
|
@ -0,0 +1,40 @@
|
|||
{
|
||||
"admin": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
},
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.25
|
||||
}
|
||||
],
|
||||
"team": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
}
|
||||
],
|
||||
"key": [
|
||||
{
|
||||
"team": "team-a",
|
||||
"api_key": "key-a",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.5
|
||||
}
|
||||
],
|
||||
"other_team": [
|
||||
{
|
||||
"team": "team-b",
|
||||
"api_key": "key-b",
|
||||
"trace_id": "01010101010101010101010101010101",
|
||||
"spend": 0.25
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -0,0 +1,9 @@
|
|||
SELECT o.TeamId AS team, o.ApiKeyHash AS api_key, o.TraceId AS trace_id,
|
||||
sum(s.spend) AS spend
|
||||
FROM otel_traces AS o
|
||||
INNER JOIN (SELECT * FROM spend_logs FINAL) AS s
|
||||
ON o.TeamId = s.team_id
|
||||
AND o.ApiKeyHash = s.api_key
|
||||
AND o.LiteLLMRequestId = s.response_id
|
||||
GROUP BY o.TeamId, o.ApiKeyHash, o.TraceId
|
||||
ORDER BY team, api_key, trace_id
|
||||
268
litellm-rust/crates/traces-clickhouse/tests/query_access.rs
Normal file
268
litellm-rust/crates/traces-clickhouse/tests/query_access.rs
Normal file
|
|
@ -0,0 +1,268 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
use litellm_traces_clickhouse::{
|
||||
Connection, Error, QueryReaders, QueryScope, ensure_schema, query_help, query_sql,
|
||||
};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
mod support;
|
||||
|
||||
use support::{ClickHouseDatabase, database as start_database};
|
||||
|
||||
struct Database {
|
||||
_database: ClickHouseDatabase,
|
||||
client: Client,
|
||||
writer: Connection,
|
||||
readers: QueryReaders,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
async fn database() -> Result<Database, Box<dyn std::error::Error>> {
|
||||
let instance = start_database().await?;
|
||||
let url = instance.url.clone();
|
||||
let client = instance.client.clone();
|
||||
let writer = Connection::parse(&url)?;
|
||||
ensure_schema(&client, &writer, "trace_test", 7).await?;
|
||||
for sql in [
|
||||
"INSERT INTO trace_test.otel_traces (TeamId, ApiKeyHash, TraceId, SpanId, Timestamp, SpanAttributes, UserId) VALUES ('team-a', 'key-a1', 'shared-trace', 'a1', now(), map('visible', 'a'), 'owner'), ('team-a', 'key-a2', 'shared-trace', 'a2', now(), map('visible', 'a'), 'other'), ('team-b', 'key-b', 'shared-trace', 'b', now(), map('secret-b', 'b'), 'owner'), ('', 'key-teamless', 'shared-trace', 'teamless', now(), map('visible', 'teamless'), ''), ('', 'key-other', 'shared-trace', 'other-teamless', now(), map('visible', 'other'), '')",
|
||||
"INSERT INTO trace_test.spend_logs (team_id, api_key, request_id, start_time, end_time, metadata, user) VALUES ('team-a', 'key-a1', 'a1', now(), now(), '{\"visible\":1}', 'owner'), ('team-a', 'key-a2', 'a2', now(), now(), '{\"visible\":1}', 'other'), ('team-b', 'key-b', 'b', now(), now(), '{\"secret_b\":1}', 'owner'), ('', 'key-teamless', 'teamless', now(), now(), '{}', ''), ('', 'key-other', 'other-teamless', now(), now(), '{}', '')",
|
||||
"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 {
|
||||
_database: instance,
|
||||
client,
|
||||
writer,
|
||||
readers,
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::own_user(QueryScope::Logs { user_id: "owner".into(), team_ids: vec![], api_key_hash: "".into() }, vec!["a1", "b"])]
|
||||
#[case::own_user_and_permitted_team(QueryScope::Logs { user_id: "owner".into(), team_ids: vec!["team-a".into()], api_key_hash: "".into() }, vec!["a1", "a2", "b"])]
|
||||
#[case::key_only_logs(QueryScope::Logs { user_id: "".into(), team_ids: vec![], api_key_hash: "key-teamless".into() }, vec!["teamless"])]
|
||||
#[case::quoted_user(QueryScope::Logs { user_id: "owner' OR 1=1 --".into(), team_ids: vec![], api_key_hash: "".into() }, vec![])]
|
||||
#[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::teamless_key(QueryScope::Key { team_id: "".to_owned(), api_key_hash: "key-teamless".to_owned() }, vec!["teamless"])]
|
||||
#[case::admin(QueryScope::Admin, vec!["a1", "a2", "b", "other-teamless", "teamless"])]
|
||||
#[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 id FROM (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::Storage(
|
||||
litellm_storage_clickhouse::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!(matches!(
|
||||
result,
|
||||
Err(Error::Cached(source)) if matches!(source.as_ref(), Error::ProvisionFailed(_))
|
||||
));
|
||||
assert!(matches!(
|
||||
database
|
||||
.readers
|
||||
.connection(&database.client, &QueryScope::Admin, "")
|
||||
.await,
|
||||
Err(Error::MissingSecret)
|
||||
));
|
||||
assert!(matches!(
|
||||
database
|
||||
.readers
|
||||
.connection(
|
||||
&database.client,
|
||||
&QueryScope::Team {
|
||||
team_id: String::new()
|
||||
},
|
||||
"test-master-secret"
|
||||
)
|
||||
.await,
|
||||
Err(Error::InvalidScope)
|
||||
));
|
||||
let permits = (0..8)
|
||||
.map(|_| database.readers.acquire())
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
assert!(matches!(database.readers.acquire(), Err(Error::Busy)));
|
||||
drop(permits);
|
||||
assert!(database.readers.acquire().is_ok());
|
||||
let rows = litellm_traces_clickhouse::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"], 5);
|
||||
Ok(())
|
||||
}
|
||||
36
litellm-rust/crates/traces-clickhouse/tests/support/mod.rs
Normal file
36
litellm-rust/crates/traces-clickhouse/tests/support/mod.rs
Normal file
|
|
@ -0,0 +1,36 @@
|
|||
use litellm_http::Client;
|
||||
use rstest::fixture;
|
||||
use testcontainers_modules::{
|
||||
clickhouse::ClickHouse,
|
||||
testcontainers::{ContainerAsync, ImageExt, runners::AsyncRunner},
|
||||
};
|
||||
|
||||
const CLICKHOUSE_TAG: &str =
|
||||
"26.9.6.6@sha256:eb4870e7ca7ed70c259eebfcfbee6cf797017f6b5436c2926bbbfe3d4d28486e";
|
||||
|
||||
pub type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
|
||||
|
||||
pub struct ClickHouseDatabase {
|
||||
_container: ContainerAsync<ClickHouse>,
|
||||
pub url: String,
|
||||
pub client: Client,
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
pub async fn database() -> TestResult<ClickHouseDatabase> {
|
||||
let container = ClickHouse::default()
|
||||
.with_tag(CLICKHOUSE_TAG)
|
||||
.with_env_var("CLICKHOUSE_SKIP_USER_SETUP", "1")
|
||||
.start()
|
||||
.await?;
|
||||
let url = format!(
|
||||
"http://{}:{}",
|
||||
container.get_host().await?,
|
||||
container.get_host_port_ipv4(8123).await?
|
||||
);
|
||||
Ok(ClickHouseDatabase {
|
||||
_container: container,
|
||||
url,
|
||||
client: Client::no_redirect_for_test(),
|
||||
})
|
||||
}
|
||||
|
|
@ -1,7 +1,6 @@
|
|||
- Keep OTLP decoding, trace schema, row encoding and named query selection here. Generic ClickHouse connections and HTTP execution belong in `litellm-storage-clickhouse`
|
||||
- Keep this crate independent of Python; PyO3 conversion and public Python exceptions belong in `python-bridge`
|
||||
- Keep the SQL migrations here as the only ClickHouse schema definition, as `migrations/NNNN_description.sql` files embedded by `litellm_migrate::migrate!`; adding a file is the only step
|
||||
- Use typed query parameters and a dedicated SELECT-only reader with server-side limits
|
||||
- Keep `config/reader.xml` grants on the database the schema is created in (CLICKHOUSE_DATABASE, default `litellm`)
|
||||
- Bound insert time and encoded bytes; make retry deduplication behavior explicit for supported ClickHouse versions
|
||||
- Test storage behavior through the crate's public API against ClickHouse
|
||||
- Own OTLP decoding, normalization, shared authorization and named query contracts; remain independent of storage and Python
|
||||
- Never depend on `litellm-traces-clickhouse` or `litellm-storage-clickhouse`
|
||||
- Preserve decoding limits, normalization precedence and shared resource identity
|
||||
- Keep ClickHouse schema, row encoding and queries in `litellm-traces-clickhouse`; keep PyO3 conversion in `python-bridge`
|
||||
- Test decoding and normalization through the public API
|
||||
- Expose one top-level `Error` enum in `src/error.rs` for decoding and normalization failures
|
||||
|
|
|
|||
|
|
@ -6,34 +6,17 @@ 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"] }
|
||||
litellm-http.workspace = true
|
||||
litellm-migrate.workspace = true
|
||||
litellm-storage-clickhouse.workspace = true
|
||||
sha2.workspace = true
|
||||
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
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
rstest.workspace = true
|
||||
testcontainers-modules = { version = "0.15.0", features = ["clickhouse"] }
|
||||
tokio.workspace = true
|
||||
wiremock.workspace = true
|
||||
|
||||
[[bench]]
|
||||
name = "resource-fanout"
|
||||
|
|
|
|||
|
|
@ -1,32 +0,0 @@
|
|||
<clickhouse>
|
||||
<profiles>
|
||||
<litellm_traces_reader>
|
||||
<readonly>1</readonly>
|
||||
<max_execution_time>10</max_execution_time>
|
||||
<max_result_rows>1000</max_result_rows>
|
||||
<max_result_bytes>4194304</max_result_bytes>
|
||||
<result_overflow_mode>throw</result_overflow_mode>
|
||||
<max_memory_usage>268435456</max_memory_usage>
|
||||
<constraints>
|
||||
<readonly><readonly/></readonly>
|
||||
<max_execution_time><readonly/></max_execution_time>
|
||||
<max_result_rows><readonly/></max_result_rows>
|
||||
<max_result_bytes><readonly/></max_result_bytes>
|
||||
<result_overflow_mode><readonly/></result_overflow_mode>
|
||||
<max_memory_usage><readonly/></max_memory_usage>
|
||||
</constraints>
|
||||
</litellm_traces_reader>
|
||||
</profiles>
|
||||
<users>
|
||||
<litellm_traces_reader>
|
||||
<password from_env="LITELLM_TRACES_READER_PASSWORD"/>
|
||||
<networks><ip>::/0</ip></networks>
|
||||
<profile>litellm_traces_reader</profile>
|
||||
<grants>
|
||||
<query>GRANT SELECT ON litellm.otel_traces</query>
|
||||
<query>GRANT SELECT ON litellm.agent_traces_by_key</query>
|
||||
<query>GRANT SELECT ON litellm.spend_logs</query>
|
||||
</grants>
|
||||
</litellm_traces_reader>
|
||||
</users>
|
||||
</clickhouse>
|
||||
|
|
@ -1,19 +0,0 @@
|
|||
SELECT o.SpanId AS span_id, o.Input AS input,
|
||||
if(o.Output = '' AND o.ObservationType = 'agent', answer.output, o.Output) AS output,
|
||||
o.SpanAttributes AS attributes
|
||||
FROM otel_traces AS o
|
||||
LEFT JOIN (
|
||||
SELECT ParentSpanId AS parent_span_id, argMax(Output, Timestamp) AS output
|
||||
FROM otel_traces
|
||||
WHERE TraceId = {trace_id:String} AND ParentSpanId = {span_id:String}
|
||||
AND ObservationType = 'llm' AND Output != ''
|
||||
AND (empty({team_ids:Array(String)}) OR TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR ApiKeyHash = {api_key_hash:String})
|
||||
GROUP BY ParentSpanId
|
||||
) AS answer ON answer.parent_span_id = o.SpanId
|
||||
WHERE o.TraceId = {trace_id:String} AND o.SpanId = {span_id:String}
|
||||
AND (empty({team_ids:Array(String)}) OR o.TeamId IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR o.ApiKeyHash = {api_key_hash:String})
|
||||
AND ({trace_ref:String} = '' OR
|
||||
hex(SHA256(concat(o.TeamId, char(0), o.ApiKeyHash, char(0), o.TraceId))) = {trace_ref:String})
|
||||
LIMIT 1
|
||||
|
|
@ -1,9 +0,0 @@
|
|||
SELECT request_id, response_id, team_id, api_key, spend,
|
||||
toUnixTimestamp64Milli(start_time) AS start_ms
|
||||
FROM spend_logs FINAL
|
||||
WHERE response_id IN {response_ids:Array(String)}
|
||||
AND start_time >= fromUnixTimestamp64Milli({start_ms:Int64})
|
||||
AND start_time < fromUnixTimestamp64Milli({end_ms:Int64})
|
||||
AND (empty({team_ids:Array(String)}) OR team_id IN {team_ids:Array(String)})
|
||||
AND ({api_key_hash:String} = '' OR api_key = {api_key_hash:String})
|
||||
ORDER BY start_time DESC
|
||||
|
|
@ -1,5 +1,5 @@
|
|||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum DecodeError {
|
||||
pub enum Error {
|
||||
#[error("invalid OTLP trace payload")]
|
||||
InvalidPayload,
|
||||
#[error("OTLP trace payload exceeds the decoding budget")]
|
||||
|
|
@ -9,21 +9,9 @@ pub enum DecodeError {
|
|||
}
|
||||
|
||||
#[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>),
|
||||
}
|
||||
#[error("invalid trace query scope")]
|
||||
pub struct InvalidScope;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
#[error("unknown ClickHouse read query")]
|
||||
pub struct InvalidQuery;
|
||||
|
|
|
|||
|
|
@ -1,25 +1,13 @@
|
|||
mod config;
|
||||
mod error;
|
||||
mod insert;
|
||||
mod normalize;
|
||||
mod otlp;
|
||||
mod query;
|
||||
pub mod query;
|
||||
mod query_access;
|
||||
mod schema;
|
||||
mod shared;
|
||||
mod sql;
|
||||
|
||||
pub use config::Config;
|
||||
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 error::{Error, InvalidQuery, InvalidScope};
|
||||
pub use normalize::{NormalizedSpan, ObservationType};
|
||||
pub use otlp::{DecodedSpan, decode_otlp};
|
||||
pub use query_access::{QueryReaders, QueryScope};
|
||||
pub use schema::{ensure_schema, schema_statements};
|
||||
pub use query::ReadQuery;
|
||||
pub use query_access::QueryScope;
|
||||
pub use shared::{Shared, SharedIdentity};
|
||||
pub use sql::{LensQuery, ReadQuery, execute_named_read};
|
||||
|
||||
pub use query::{query_help, query_sql};
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use std::collections::BTreeMap;
|
|||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, tokens};
|
||||
use crate::{DecodeError, otlp::DecodedEvent};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(crate) const CLAUDE_CODE_SCOPE: &str = "com.anthropic.claude_code.tracing";
|
||||
pub(crate) const CLAUDE_CODE_AGENT: &str = "claude-code";
|
||||
|
|
@ -144,13 +144,13 @@ fn llm_output(attributes: &BTreeMap<String, String>) -> String {
|
|||
}
|
||||
}
|
||||
|
||||
fn input_tokens(attributes: &BTreeMap<String, String>) -> Result<u32, DecodeError> {
|
||||
fn input_tokens(attributes: &BTreeMap<String, String>) -> Result<u32, Error> {
|
||||
["input_tokens", "cache_read_tokens", "cache_creation_tokens"]
|
||||
.into_iter()
|
||||
.try_fold(0u32, |total, key| {
|
||||
total
|
||||
.checked_add(tokens(attributes, key)?)
|
||||
.ok_or(DecodeError::TokenCountOutOfRange)
|
||||
.ok_or(Error::TokenCountOutOfRange)
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -180,7 +180,7 @@ impl SpanNormalizer for ClaudeCodeNormalizer {
|
|||
_parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let base = NormalizedSpan {
|
||||
observation_type: ObservationType::Framework,
|
||||
agent_name: CLAUDE_CODE_AGENT.to_owned(),
|
||||
|
|
@ -228,7 +228,7 @@ mod tests {
|
|||
use serde_json::Value;
|
||||
|
||||
use super::{CLAUDE_CODE_SCOPE, ClaudeCodeNormalizer, SpanNormalizer};
|
||||
use crate::{DecodeError, normalize::ObservationType, otlp::DecodedEvent};
|
||||
use crate::{Error, normalize::ObservationType, otlp::DecodedEvent};
|
||||
|
||||
fn attributes(pairs: &[(&str, &str)]) -> BTreeMap<String, String> {
|
||||
pairs
|
||||
|
|
@ -341,7 +341,7 @@ mod tests {
|
|||
]),
|
||||
&[],
|
||||
);
|
||||
assert!(matches!(result, Err(DecodeError::TokenCountOutOfRange)));
|
||||
assert!(matches!(result, Err(Error::TokenCountOutOfRange)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, first, usage_tokens};
|
||||
use crate::{DecodeError, otlp::DecodedEvent};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct GenAiNormalizer;
|
||||
|
||||
|
|
@ -31,7 +31,7 @@ impl SpanNormalizer for GenAiNormalizer {
|
|||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
|
||||
let observation_type = match attr(attributes, "gen_ai.operation.name") {
|
||||
"invoke_agent" => ObservationType::Agent,
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ use serde::{Deserialize, Deserializer, Serialize, de::DeserializeOwned};
|
|||
use serde_json::{Value, ser::Formatter};
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, usage_tokens};
|
||||
use crate::{DecodeError, otlp::DecodedEvent};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct LangSmithNormalizer;
|
||||
|
||||
|
|
@ -405,7 +405,7 @@ impl SpanNormalizer for LangSmithNormalizer {
|
|||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (input_tokens, output_tokens) = usage_tokens(attributes)?;
|
||||
let observation_type = span_type(name, parent_span_id, attributes);
|
||||
let io = span_io(observation_type, attributes);
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::{DecodeError, otlp::DecodedEvent};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
|
|
@ -32,71 +32,6 @@ pub(crate) struct Normalization {
|
|||
pub consumed_attributes: [&'static str; 2],
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize)]
|
||||
pub struct NormalizedFieldDefinition {
|
||||
pub name: &'static str,
|
||||
pub clickhouse_column: &'static str,
|
||||
pub clickhouse_type: &'static str,
|
||||
pub meaning: &'static str,
|
||||
}
|
||||
|
||||
pub const NORMALIZED_FIELD_DEFINITIONS: [NormalizedFieldDefinition; 9] = [
|
||||
NormalizedFieldDefinition {
|
||||
name: "observation_type",
|
||||
clickhouse_column: "ObservationType",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent, LLM, tool, chain, or framework span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "agent_name",
|
||||
clickhouse_column: "AgentName",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent associated with this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "framework",
|
||||
clickhouse_column: "Framework",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Agent framework or SDK that emitted this span, e.g. claude-agent-sdk",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "litellm_request_id",
|
||||
clickhouse_column: "LiteLLMRequestId",
|
||||
clickhouse_type: "String",
|
||||
meaning: "LiteLLM response ID used to link a span to a spend log",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "model",
|
||||
clickhouse_column: "Model",
|
||||
clickhouse_type: "LowCardinality(String)",
|
||||
meaning: "Model used by this span",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input_tokens",
|
||||
clickhouse_column: "InputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Input token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output_tokens",
|
||||
clickhouse_column: "OutputTokens",
|
||||
clickhouse_type: "UInt32",
|
||||
meaning: "Output token count",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "input",
|
||||
clickhouse_column: "Input",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized input payload",
|
||||
},
|
||||
NormalizedFieldDefinition {
|
||||
name: "output",
|
||||
clickhouse_column: "Output",
|
||||
clickhouse_type: "String",
|
||||
meaning: "Normalized output payload",
|
||||
},
|
||||
];
|
||||
|
||||
trait SpanNormalizer {
|
||||
fn matches(&self, scope_name: &str, attributes: &BTreeMap<String, String>) -> bool;
|
||||
fn consumed_attributes(&self, attributes: &BTreeMap<String, String>) -> [&'static str; 2];
|
||||
|
|
@ -106,7 +41,7 @@ trait SpanNormalizer {
|
|||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, DecodeError>;
|
||||
) -> Result<NormalizedSpan, Error>;
|
||||
fn display_name(&self, _attributes: &BTreeMap<String, String>) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
|
@ -136,27 +71,27 @@ fn first<'a>(attributes: &'a BTreeMap<String, String>, left: &str, right: &str)
|
|||
}
|
||||
}
|
||||
|
||||
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, DecodeError> {
|
||||
fn tokens(attributes: &BTreeMap<String, String>, key: &str) -> Result<u32, Error> {
|
||||
let value = attr(attributes, key).trim();
|
||||
if value.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
match value.parse::<i128>() {
|
||||
Ok(number) if (0..=u32::MAX as i128).contains(&number) => Ok(number as u32),
|
||||
Ok(_) => Err(DecodeError::TokenCountOutOfRange),
|
||||
Ok(_) => Err(Error::TokenCountOutOfRange),
|
||||
Err(_)
|
||||
if value
|
||||
.trim_start_matches(['+', '-'])
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_digit()) =>
|
||||
{
|
||||
Err(DecodeError::TokenCountOutOfRange)
|
||||
Err(Error::TokenCountOutOfRange)
|
||||
}
|
||||
Err(_) => Ok(0),
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), DecodeError> {
|
||||
fn usage_tokens(attributes: &BTreeMap<String, String>) -> Result<(u32, u32), Error> {
|
||||
Ok((
|
||||
tokens(attributes, "gen_ai.usage.input_tokens")?,
|
||||
tokens(attributes, "gen_ai.usage.output_tokens")?,
|
||||
|
|
@ -223,7 +158,7 @@ pub fn normalize(
|
|||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
events: &[DecodedEvent],
|
||||
) -> Result<Normalization, DecodeError> {
|
||||
) -> Result<Normalization, Error> {
|
||||
let normalizers: [&dyn SpanNormalizer; 4] = [
|
||||
&ClaudeCodeNormalizer,
|
||||
&LangSmithNormalizer,
|
||||
|
|
@ -257,11 +192,11 @@ pub fn normalize(
|
|||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{NORMALIZED_FIELD_DEFINITIONS, ObservationType, normalize};
|
||||
use super::{ObservationType, normalize};
|
||||
|
||||
#[rstest]
|
||||
#[case::langsmith("langsmith", [("langsmith.span.kind", "llm"), ("openinference.span.kind", "TOOL")], ObservationType::Llm)]
|
||||
|
|
@ -286,25 +221,6 @@ mod tests {
|
|||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn field_definitions_match_serialized_normalized_span() {
|
||||
let fields = normalize("", "root", "", &BTreeMap::new(), &[])
|
||||
.expect("valid tokens")
|
||||
.span;
|
||||
let serialized = serde_json::to_value(fields).expect("serializable fields");
|
||||
let keys: BTreeSet<_> = serialized
|
||||
.as_object()
|
||||
.expect("field object")
|
||||
.keys()
|
||||
.map(String::as_str)
|
||||
.collect();
|
||||
let mapped: BTreeSet<_> = NORMALIZED_FIELD_DEFINITIONS
|
||||
.iter()
|
||||
.map(|field| field.name)
|
||||
.collect();
|
||||
assert_eq!(keys, mapped);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn token_counts_accept_surrounding_whitespace() {
|
||||
let attributes =
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{NormalizedSpan, ObservationType, SpanNormalizer, attr, tokens, usage_tokens};
|
||||
use crate::{DecodeError, otlp::DecodedEvent};
|
||||
use crate::{Error, otlp::DecodedEvent};
|
||||
|
||||
pub(super) struct OpenInferenceNormalizer;
|
||||
|
||||
|
|
@ -20,7 +20,7 @@ impl SpanNormalizer for OpenInferenceNormalizer {
|
|||
parent_span_id: &str,
|
||||
attributes: &BTreeMap<String, String>,
|
||||
_events: &[DecodedEvent],
|
||||
) -> Result<NormalizedSpan, DecodeError> {
|
||||
) -> Result<NormalizedSpan, Error> {
|
||||
let (usage_input, usage_output) = usage_tokens(attributes)?;
|
||||
let observation_type = match attr(attributes, "openinference.span.kind")
|
||||
.to_ascii_uppercase()
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ use serde::{
|
|||
};
|
||||
|
||||
use super::limits::{Budget, MAX_ATTRIBUTES};
|
||||
use crate::DecodeError;
|
||||
use crate::Error;
|
||||
|
||||
struct AttributeWriter<'a> {
|
||||
body: Vec<u8>,
|
||||
|
|
@ -32,9 +32,9 @@ impl Write for AttributeWriter<'_> {
|
|||
pub(super) fn attributes(
|
||||
values: Vec<KeyValue>,
|
||||
budget: &mut Budget,
|
||||
) -> Result<BTreeMap<String, String>, DecodeError> {
|
||||
) -> Result<BTreeMap<String, String>, Error> {
|
||||
if values.len() > MAX_ATTRIBUTES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
values
|
||||
.into_iter()
|
||||
|
|
@ -59,8 +59,8 @@ pub(super) fn attributes(
|
|||
budget,
|
||||
};
|
||||
serde_json::to_writer(&mut writer, &AttributeJson(value.as_ref()))
|
||||
.map_err(|_| DecodeError::TooLarge)?;
|
||||
String::from_utf8(writer.body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
.map_err(|_| Error::TooLarge)?;
|
||||
String::from_utf8(writer.body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
};
|
||||
Ok((entry.key, text))
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ use std::fmt;
|
|||
use prost::encoding::{DecodeContext, WireType, decode_key, decode_varint, skip_field};
|
||||
use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor};
|
||||
|
||||
use crate::{DecodeError, Shared};
|
||||
use crate::{Error, Shared};
|
||||
|
||||
pub(super) const MAX_DEPTH: usize = 32;
|
||||
pub(super) const MAX_NODES: usize = 65_536;
|
||||
|
|
@ -12,7 +12,7 @@ pub(super) const MAX_ATTRIBUTES: usize = 256;
|
|||
pub(super) const MAX_EVENTS: usize = 256;
|
||||
pub(super) const MAX_DECODED_SPAN_BYTES: usize = 16 * 1024 * 1024;
|
||||
|
||||
pub(super) fn json_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
||||
pub(super) fn json_preflight(payload: &[u8]) -> Result<(), Error> {
|
||||
let mut nodes = 0;
|
||||
let mut exceeded = false;
|
||||
let mut decoder = serde_json::Deserializer::from_slice(payload);
|
||||
|
|
@ -24,9 +24,9 @@ pub(super) fn json_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
|||
.deserialize(&mut decoder)
|
||||
.and_then(|()| decoder.end());
|
||||
if exceeded {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
result.map_err(|_| DecodeError::InvalidPayload)
|
||||
result.map_err(|_| Error::InvalidPayload)
|
||||
}
|
||||
|
||||
struct JsonBudget<'a> {
|
||||
|
|
@ -146,7 +146,7 @@ impl MessageKind {
|
|||
}
|
||||
}
|
||||
|
||||
pub(super) fn protobuf_preflight(payload: &[u8]) -> Result<(), DecodeError> {
|
||||
pub(super) fn protobuf_preflight(payload: &[u8]) -> Result<(), Error> {
|
||||
scan_message(payload, MessageKind::Export, 0, &mut 0)
|
||||
}
|
||||
|
||||
|
|
@ -155,27 +155,27 @@ fn scan_message(
|
|||
kind: MessageKind,
|
||||
depth: usize,
|
||||
nodes: &mut usize,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
if depth > MAX_DEPTH {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
while !payload.is_empty() {
|
||||
*nodes += 1;
|
||||
if *nodes > MAX_NODES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
let (tag, wire) = decode_key(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let (tag, wire) = decode_key(&mut payload).map_err(|_| Error::InvalidPayload)?;
|
||||
if let (WireType::LengthDelimited, Some(child)) = (wire, kind.child(tag)) {
|
||||
let length = decode_varint(&mut payload).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let length = usize::try_from(length).map_err(|_| DecodeError::InvalidPayload)?;
|
||||
let length = decode_varint(&mut payload).map_err(|_| Error::InvalidPayload)?;
|
||||
let length = usize::try_from(length).map_err(|_| Error::InvalidPayload)?;
|
||||
let (message, rest) = payload
|
||||
.split_at_checked(length)
|
||||
.ok_or(DecodeError::InvalidPayload)?;
|
||||
.ok_or(Error::InvalidPayload)?;
|
||||
scan_message(message, child, depth + 1, nodes)?;
|
||||
payload = rest;
|
||||
} else {
|
||||
skip_field(wire, tag, &mut payload, DecodeContext::default())
|
||||
.map_err(|_| DecodeError::InvalidPayload)?;
|
||||
.map_err(|_| Error::InvalidPayload)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
|
|
@ -194,7 +194,7 @@ impl Budget {
|
|||
&mut self,
|
||||
value: &Shared<T>,
|
||||
allocated_bytes: impl FnOnce(&T) -> usize,
|
||||
) -> Result<Shared<T>, DecodeError> {
|
||||
) -> Result<Shared<T>, Error> {
|
||||
let cloned = value.clone();
|
||||
if !value.shares_storage_with(&cloned) {
|
||||
self.consume(allocated_bytes(value))?;
|
||||
|
|
@ -202,11 +202,8 @@ impl Budget {
|
|||
Ok(cloned)
|
||||
}
|
||||
|
||||
pub(super) fn consume(&mut self, bytes: usize) -> Result<(), DecodeError> {
|
||||
self.remaining = self
|
||||
.remaining
|
||||
.checked_sub(bytes)
|
||||
.ok_or(DecodeError::TooLarge)?;
|
||||
pub(super) fn consume(&mut self, bytes: usize) -> Result<(), Error> {
|
||||
self.remaining = self.remaining.checked_sub(bytes).ok_or(Error::TooLarge)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ mod wire;
|
|||
use serde::Serialize;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::{DecodeError, NormalizedSpan, Shared};
|
||||
use crate::{Error, NormalizedSpan, Shared};
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct DecodedEvent {
|
||||
|
|
@ -35,10 +35,7 @@ pub struct DecodedSpan {
|
|||
pub consumed_attributes: [&'static str; 2],
|
||||
}
|
||||
|
||||
pub fn decode_otlp(
|
||||
body: &[u8],
|
||||
content_type: Option<&str>,
|
||||
) -> Result<Vec<DecodedSpan>, DecodeError> {
|
||||
pub fn decode_otlp(body: &[u8], content_type: Option<&str>) -> Result<Vec<DecodedSpan>, Error> {
|
||||
let request = wire::decode(body, content_type)?;
|
||||
span::flatten(request)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,11 +11,11 @@ use super::{
|
|||
limits::{Budget, MAX_ATTRIBUTES, MAX_DECODED_SPAN_BYTES, MAX_EVENTS, MAX_SPANS},
|
||||
};
|
||||
use crate::{
|
||||
DecodeError, Shared,
|
||||
Error, Shared,
|
||||
normalize::{CLAUDE_CODE_AGENT, CLAUDE_CODE_SCOPE, normalize},
|
||||
};
|
||||
|
||||
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, DecodeError> {
|
||||
pub(super) fn flatten(request: ExportTraceServiceRequest) -> Result<Vec<DecodedSpan>, Error> {
|
||||
let mut budget = Budget::new(MAX_DECODED_SPAN_BYTES);
|
||||
let mut spans = Vec::new();
|
||||
for resource in request.resource_spans {
|
||||
|
|
@ -28,7 +28,7 @@ fn append_resource(
|
|||
resource: ResourceSpans,
|
||||
budget: &mut Budget,
|
||||
spans: &mut Vec<DecodedSpan>,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
let attributes = Shared::new(attributes(
|
||||
resource
|
||||
.resource
|
||||
|
|
@ -47,17 +47,17 @@ fn append_scope(
|
|||
resource: &Shared<BTreeMap<String, String>>,
|
||||
budget: &mut Budget,
|
||||
spans: &mut Vec<DecodedSpan>,
|
||||
) -> Result<(), DecodeError> {
|
||||
) -> Result<(), Error> {
|
||||
let scope = scope_spans.scope.unwrap_or_default();
|
||||
if scope.attributes.len() > MAX_ATTRIBUTES {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
budget.consume(scope.name.len() + scope.version.len())?;
|
||||
let scope_name: Shared<String> = scope.name.into();
|
||||
let scope_version: Shared<String> = scope.version.into();
|
||||
for span in scope_spans.spans {
|
||||
if spans.len() >= MAX_SPANS {
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
validate_span(&span)?;
|
||||
budget.consume(
|
||||
|
|
@ -85,7 +85,7 @@ fn valid_id(value: &[u8], length: usize) -> bool {
|
|||
value.len() == length && value.iter().any(|byte| *byte != 0)
|
||||
}
|
||||
|
||||
fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
||||
fn validate_span(span: &Span) -> Result<(), Error> {
|
||||
if !valid_id(&span.trace_id, 16)
|
||||
|| !valid_id(&span.span_id, 8)
|
||||
|| (!span.parent_span_id.is_empty() && !valid_id(&span.parent_span_id, 8))
|
||||
|
|
@ -97,7 +97,7 @@ fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
|||
.iter()
|
||||
.any(|link| !valid_id(&link.trace_id, 16) || !valid_id(&link.span_id, 8))
|
||||
{
|
||||
return Err(DecodeError::InvalidPayload);
|
||||
return Err(Error::InvalidPayload);
|
||||
}
|
||||
if span.events.len() > MAX_EVENTS
|
||||
|| span.links.len() > MAX_EVENTS
|
||||
|
|
@ -111,7 +111,7 @@ fn validate_span(span: &Span) -> Result<(), DecodeError> {
|
|||
.iter()
|
||||
.any(|event| event.attributes.len() > MAX_ATTRIBUTES)
|
||||
{
|
||||
return Err(DecodeError::TooLarge);
|
||||
return Err(Error::TooLarge);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -126,7 +126,7 @@ fn decoded_span(
|
|||
scope_name: &Shared<String>,
|
||||
scope_version: &Shared<String>,
|
||||
budget: &mut Budget,
|
||||
) -> Result<DecodedSpan, DecodeError> {
|
||||
) -> Result<DecodedSpan, Error> {
|
||||
let status = span.status.unwrap_or_default();
|
||||
let parent_span_id = hex_bytes(&span.parent_span_id);
|
||||
let span_attributes = attributes(span.attributes, budget)?;
|
||||
|
|
@ -140,7 +140,7 @@ fn decoded_span(
|
|||
attributes: attributes(event.attributes, budget)?,
|
||||
})
|
||||
})
|
||||
.collect::<Result<Vec<_>, DecodeError>>()?;
|
||||
.collect::<Result<Vec<_>, Error>>()?;
|
||||
let normalization = normalize(
|
||||
scope_name.as_ref(),
|
||||
&span.name,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use opentelemetry_proto::tonic::collector::trace::v1::ExportTraceServiceRequest;
|
|||
use prost::Message;
|
||||
|
||||
use super::limits::{json_preflight, protobuf_preflight};
|
||||
use crate::DecodeError;
|
||||
use crate::Error;
|
||||
|
||||
#[derive(strum::EnumString)]
|
||||
#[strum(ascii_case_insensitive)]
|
||||
|
|
@ -19,7 +19,7 @@ enum OtlpMediaType {
|
|||
pub(super) fn decode(
|
||||
body: &[u8],
|
||||
content_type: Option<&str>,
|
||||
) -> Result<ExportTraceServiceRequest, DecodeError> {
|
||||
) -> Result<ExportTraceServiceRequest, Error> {
|
||||
let media_type = content_type
|
||||
.unwrap_or("application/x-protobuf")
|
||||
.split(';')
|
||||
|
|
@ -27,16 +27,16 @@ pub(super) fn decode(
|
|||
.unwrap_or_default()
|
||||
.trim()
|
||||
.parse::<OtlpMediaType>()
|
||||
.map_err(|_| DecodeError::InvalidPayload)?;
|
||||
.map_err(|_| Error::InvalidPayload)?;
|
||||
|
||||
let request = match media_type {
|
||||
OtlpMediaType::Json => {
|
||||
json_preflight(body)?;
|
||||
serde_json::from_slice(body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
serde_json::from_slice(body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
OtlpMediaType::Protobuf => {
|
||||
protobuf_preflight(body)?;
|
||||
ExportTraceServiceRequest::decode(body).map_err(|_| DecodeError::InvalidPayload)?
|
||||
ExportTraceServiceRequest::decode(body).map_err(|_| Error::InvalidPayload)?
|
||||
}
|
||||
};
|
||||
Ok(request)
|
||||
|
|
|
|||
|
|
@ -1,346 +1,23 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
pub mod named;
|
||||
|
||||
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(Clone, Copy, Debug, Eq, PartialEq, strum::EnumString, strum::Display, strum::AsRefStr)]
|
||||
#[strum(serialize_all = "snake_case")]
|
||||
pub enum ReadQuery {
|
||||
ListTraces,
|
||||
TraceSpans,
|
||||
TraceIdentity,
|
||||
SpanDetail,
|
||||
SpanError,
|
||||
SpendByResponseIds,
|
||||
Availability,
|
||||
Agents,
|
||||
Sample,
|
||||
Content,
|
||||
Evidence,
|
||||
}
|
||||
|
||||
#[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)
|
||||
);
|
||||
impl ReadQuery {
|
||||
pub fn parse(value: &str) -> Result<Self, crate::InvalidQuery> {
|
||||
value.parse().map_err(|_| crate::InvalidQuery)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
152
litellm-rust/crates/traces/src/query/named.rs
Normal file
152
litellm-rust/crates/traces/src/query/named.rs
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ReadAccessParams {
|
||||
pub all_teams: u8,
|
||||
pub user_id: String,
|
||||
pub team_ids: Vec<String>,
|
||||
pub api_key_hash: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub start_ms: i64,
|
||||
pub end_ms: i64,
|
||||
pub cursor_ms: i64,
|
||||
pub cursor_trace_id: String,
|
||||
pub limit: u32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct ListTracesRow {
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub team_id: String,
|
||||
pub api_key_hash: String,
|
||||
pub user_id: String,
|
||||
pub name: String,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub status: String,
|
||||
pub start_ms: i64,
|
||||
pub duration_ms: i64,
|
||||
pub span_count: u64,
|
||||
pub agent_count: u64,
|
||||
pub agent_invocations: u64,
|
||||
#[serde(default)]
|
||||
pub agent_names: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub frameworks: Vec<String>,
|
||||
pub llm_calls: u64,
|
||||
pub tool_calls: u64,
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub models: Vec<String>,
|
||||
pub error_count: u64,
|
||||
pub request_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceSpansParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceSpansRow {
|
||||
pub span_id: String,
|
||||
pub parent_span_id: String,
|
||||
pub name: String,
|
||||
#[serde(rename = "type")]
|
||||
pub kind: String,
|
||||
pub agent: String,
|
||||
#[serde(default)]
|
||||
pub framework: String,
|
||||
pub status: String,
|
||||
pub status_message: String,
|
||||
pub error_truncated: u8,
|
||||
pub start_ns: i64,
|
||||
pub duration_ns: u64,
|
||||
pub service: String,
|
||||
pub input_preview: String,
|
||||
pub model: String,
|
||||
pub input_tokens: u32,
|
||||
pub output_tokens: u32,
|
||||
pub litellm_request_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key_hash: String,
|
||||
pub user_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanDetailParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub span_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanDetailRow {
|
||||
pub span_id: String,
|
||||
pub input: String,
|
||||
pub output: String,
|
||||
pub attributes: BTreeMap<String, String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
pub trace_ref: String,
|
||||
pub span_id: String,
|
||||
pub error_offset: u64,
|
||||
pub error_version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpanErrorRow {
|
||||
pub span_id: String,
|
||||
pub message: String,
|
||||
pub total_chars: u64,
|
||||
pub version: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub response_ids: Vec<String>,
|
||||
pub start_ms: i64,
|
||||
pub end_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct SpendByResponseIdsRow {
|
||||
pub request_id: String,
|
||||
pub response_id: String,
|
||||
pub team_id: String,
|
||||
pub api_key: String,
|
||||
pub user: String,
|
||||
pub spend: f64,
|
||||
pub start_ms: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceIdentityParams {
|
||||
#[serde(flatten)]
|
||||
pub access: ReadAccessParams,
|
||||
pub trace_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Serialize)]
|
||||
pub struct TraceIdentityRow {
|
||||
pub trace_ref: String,
|
||||
}
|
||||
|
|
@ -1,15 +1,6 @@
|
|||
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"];
|
||||
use crate::InvalidScope;
|
||||
|
||||
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
|
||||
|
|
@ -18,6 +9,11 @@ pub enum QueryScope {
|
|||
Team {
|
||||
team_id: String,
|
||||
},
|
||||
Logs {
|
||||
user_id: String,
|
||||
team_ids: Vec<String>,
|
||||
api_key_hash: String,
|
||||
},
|
||||
Key {
|
||||
team_id: String,
|
||||
api_key_hash: String,
|
||||
|
|
@ -25,176 +21,21 @@ pub enum QueryScope {
|
|||
}
|
||||
|
||||
impl QueryScope {
|
||||
fn validate(&self) -> Result<(), QueryAccessError> {
|
||||
pub fn validate(&self) -> Result<(), InvalidScope> {
|
||||
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,
|
||||
Self::Logs {
|
||||
user_id,
|
||||
team_ids,
|
||||
api_key_hash,
|
||||
} => format!(
|
||||
"{team} = {} AND {key} = {}",
|
||||
literal(team_id),
|
||||
literal(api_key_hash)
|
||||
),
|
||||
} if (!user_id.is_empty() || !team_ids.is_empty() || !api_key_hash.is_empty())
|
||||
&& team_ids.iter().all(|team| !team.is_empty()) =>
|
||||
{
|
||||
Ok(())
|
||||
}
|
||||
_ => Err(InvalidScope),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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('\'', "\\'"))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,70 +0,0 @@
|
|||
use litellm_http::Client;
|
||||
use litellm_migrate::Migration;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::Connection;
|
||||
use crate::Error;
|
||||
|
||||
const SCHEMA_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
|
||||
const MIGRATIONS: &[Migration] = litellm_migrate::migrate!("migrations");
|
||||
|
||||
pub fn schema_statements(database: &str, retention_days: u32) -> Result<Vec<String>, Error> {
|
||||
if database.is_empty()
|
||||
|| !database
|
||||
.bytes()
|
||||
.all(|c| c.is_ascii_alphanumeric() || c == b'_')
|
||||
|| retention_days == 0
|
||||
{
|
||||
return Err(Error::InvalidSchema);
|
||||
}
|
||||
let database = format!("`{database}`");
|
||||
Ok(
|
||||
std::iter::once(format!("CREATE DATABASE IF NOT EXISTS {database}"))
|
||||
.chain(MIGRATIONS.iter().map(|migration| {
|
||||
migration
|
||||
.sql
|
||||
.replace("{database}", &database)
|
||||
.replace("{retention_days}", &retention_days.to_string())
|
||||
}))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
pub async fn ensure_schema(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
database: &str,
|
||||
retention_days: u32,
|
||||
) -> Result<(), Error> {
|
||||
ensure_schema_with_timeout(
|
||||
client,
|
||||
connection,
|
||||
database,
|
||||
retention_days,
|
||||
SCHEMA_REQUEST_TIMEOUT,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn ensure_schema_with_timeout(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
database: &str,
|
||||
retention_days: u32,
|
||||
request_timeout: Duration,
|
||||
) -> Result<(), Error> {
|
||||
for statement in schema_statements(database, retention_days)? {
|
||||
let response = client
|
||||
.post(connection.url().clone())
|
||||
.timeout(request_timeout)
|
||||
.body(statement)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|_| Error::Transport)?;
|
||||
if !response.status().is_success() {
|
||||
return Err(Error::SchemaFailed(response.status().as_u16()));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
|
@ -1,76 +0,0 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use litellm_http::Client;
|
||||
|
||||
use crate::{Connection, Error, Parameter, execute_read};
|
||||
|
||||
pub enum ReadQuery {
|
||||
ListTraces,
|
||||
TraceSpans,
|
||||
SpanDetail,
|
||||
SpanError,
|
||||
SpendByResponseIds,
|
||||
}
|
||||
|
||||
impl ReadQuery {
|
||||
pub fn parse(value: &str) -> Result<Self, Error> {
|
||||
match value {
|
||||
"list_traces" => Ok(Self::ListTraces),
|
||||
"trace_spans" => Ok(Self::TraceSpans),
|
||||
"span_detail" => Ok(Self::SpanDetail),
|
||||
"span_error" => Ok(Self::SpanError),
|
||||
"spend_by_response_ids" => Ok(Self::SpendByResponseIds),
|
||||
_ => Err(Error::InvalidQuery),
|
||||
}
|
||||
}
|
||||
|
||||
fn sql(&self) -> &'static str {
|
||||
match self {
|
||||
Self::ListTraces => include_str!("../query/list_traces.sql"),
|
||||
Self::TraceSpans => include_str!("../query/trace_spans.sql"),
|
||||
Self::SpanDetail => include_str!("../query/span_detail.sql"),
|
||||
Self::SpanError => include_str!("../query/span_error.sql"),
|
||||
Self::SpendByResponseIds => include_str!("../query/spend_by_response_ids.sql"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub enum LensQuery {
|
||||
Availability,
|
||||
Agents,
|
||||
Sample,
|
||||
Content,
|
||||
Evidence,
|
||||
}
|
||||
|
||||
impl LensQuery {
|
||||
pub fn parse(name: &str) -> Result<Self, Error> {
|
||||
match name {
|
||||
"availability" => Ok(Self::Availability),
|
||||
"agents" => Ok(Self::Agents),
|
||||
"sample" => Ok(Self::Sample),
|
||||
"content" => Ok(Self::Content),
|
||||
"evidence" => Ok(Self::Evidence),
|
||||
_ => Err(Error::InvalidQuery),
|
||||
}
|
||||
}
|
||||
pub fn sql(self) -> &'static str {
|
||||
match self {
|
||||
Self::Availability => include_str!("../query/lens_availability.sql"),
|
||||
Self::Agents => include_str!("../query/lens_agents.sql"),
|
||||
Self::Sample => include_str!("../query/lens_sample.sql"),
|
||||
Self::Content => include_str!("../query/lens_content.sql"),
|
||||
Self::Evidence => include_str!("../query/lens_evidence.sql"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn execute_named_read(
|
||||
client: &Client,
|
||||
connection: &Connection,
|
||||
query: ReadQuery,
|
||||
parameters: &BTreeMap<String, Parameter>,
|
||||
) -> Result<String, Error> {
|
||||
execute_read(client, connection, query.sql(), parameters).await
|
||||
}
|
||||
1614
litellm-rust/crates/traces/tests/fixtures/deeplite_auth_error.json
vendored
Normal file
1614
litellm-rust/crates/traces/tests/fixtures/deeplite_auth_error.json
vendored
Normal file
File diff suppressed because one or more lines are too long
9812
litellm-rust/crates/traces/tests/fixtures/deeplite_swarm.json
vendored
Normal file
9812
litellm-rust/crates/traces/tests/fixtures/deeplite_swarm.json
vendored
Normal file
File diff suppressed because one or more lines are too long
39
litellm-rust/crates/traces/tests/fixtures/query_alternate.json
vendored
Normal file
39
litellm-rust/crates/traces/tests/fixtures/query_alternate.json
vendored
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
{
|
||||
"resourceSpans": [
|
||||
{
|
||||
"resource": {
|
||||
"attributes": [
|
||||
{
|
||||
"key": "service.name",
|
||||
"value": {
|
||||
"stringValue": "fixture"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"scopeSpans": [
|
||||
{
|
||||
"scope": {
|
||||
"name": "fixture"
|
||||
},
|
||||
"spans": [
|
||||
{
|
||||
"traceId": "01010101010101010101010101010101",
|
||||
"spanId": "0404040404040404",
|
||||
"parentSpanId": "",
|
||||
"name": "alternate",
|
||||
"kind": 1,
|
||||
"startTimeUnixNano": "1735689601000000000",
|
||||
"endTimeUnixNano": "1735689602000000000",
|
||||
"attributes": [],
|
||||
"status": {
|
||||
"code": 1,
|
||||
"message": ""
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
85
litellm-rust/crates/traces/tests/fixtures/query_children.json
vendored
Normal file
85
litellm-rust/crates/traces/tests/fixtures/query_children.json
vendored
Normal file
|
|
@ -0,0 +1,85 @@
|
|||
{
|
||||
"resourceSpans": [
|
||||
{
|
||||
"resource": {
|
||||
"attributes": [
|
||||
{
|
||||
"key": "service.name",
|
||||
"value": {
|
||||
"stringValue": "fixture"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"scopeSpans": [
|
||||
{
|
||||
"scope": {
|
||||
"name": "fixture"
|
||||
},
|
||||
"spans": [
|
||||
{
|
||||
"traceId": "01010101010101010101010101010101",
|
||||
"spanId": "0202020202020202",
|
||||
"parentSpanId": "0101010101010101",
|
||||
"name": "completion",
|
||||
"kind": 1,
|
||||
"startTimeUnixNano": "1735689600100000000",
|
||||
"endTimeUnixNano": "1735689600600000000",
|
||||
"attributes": [
|
||||
{
|
||||
"key": "gen_ai.operation.name",
|
||||
"value": {
|
||||
"stringValue": "chat"
|
||||
}
|
||||
},
|
||||
{
|
||||
"key": "gen_ai.response.id",
|
||||
"value": {
|
||||
"stringValue": "response-shared"
|
||||
}
|
||||
},
|
||||
{
|
||||
"key": "gen_ai.usage.input_tokens",
|
||||
"value": {
|
||||
"stringValue": "12"
|
||||
}
|
||||
},
|
||||
{
|
||||
"key": "gen_ai.usage.output_tokens",
|
||||
"value": {
|
||||
"stringValue": "6"
|
||||
}
|
||||
}
|
||||
],
|
||||
"status": {
|
||||
"code": 1,
|
||||
"message": ""
|
||||
}
|
||||
},
|
||||
{
|
||||
"traceId": "01010101010101010101010101010101",
|
||||
"spanId": "0303030303030303",
|
||||
"parentSpanId": "0101010101010101",
|
||||
"name": "lookup",
|
||||
"kind": 1,
|
||||
"startTimeUnixNano": "1735689600700000000",
|
||||
"endTimeUnixNano": "1735689600800000000",
|
||||
"attributes": [
|
||||
{
|
||||
"key": "gen_ai.operation.name",
|
||||
"value": {
|
||||
"stringValue": "execute_tool"
|
||||
}
|
||||
}
|
||||
],
|
||||
"status": {
|
||||
"code": 2,
|
||||
"message": "lookup timed out"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
46
litellm-rust/crates/traces/tests/fixtures/query_other_team.json
vendored
Normal file
46
litellm-rust/crates/traces/tests/fixtures/query_other_team.json
vendored
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
{
|
||||
"resourceSpans": [
|
||||
{
|
||||
"resource": {
|
||||
"attributes": [
|
||||
{
|
||||
"key": "service.name",
|
||||
"value": {
|
||||
"stringValue": "fixture"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"scopeSpans": [
|
||||
{
|
||||
"scope": {
|
||||
"name": "fixture"
|
||||
},
|
||||
"spans": [
|
||||
{
|
||||
"traceId": "01010101010101010101010101010101",
|
||||
"spanId": "0505050505050505",
|
||||
"parentSpanId": "",
|
||||
"name": "other-team",
|
||||
"kind": 1,
|
||||
"startTimeUnixNano": "1735689602000000000",
|
||||
"endTimeUnixNano": "1735689603000000000",
|
||||
"attributes": [
|
||||
{
|
||||
"key": "gen_ai.response.id",
|
||||
"value": {
|
||||
"stringValue": "response-shared"
|
||||
}
|
||||
}
|
||||
],
|
||||
"status": {
|
||||
"code": 1,
|
||||
"message": ""
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
46
litellm-rust/crates/traces/tests/fixtures/query_root.json
vendored
Normal file
46
litellm-rust/crates/traces/tests/fixtures/query_root.json
vendored
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
{
|
||||
"resourceSpans": [
|
||||
{
|
||||
"resource": {
|
||||
"attributes": [
|
||||
{
|
||||
"key": "service.name",
|
||||
"value": {
|
||||
"stringValue": "fixture"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
"scopeSpans": [
|
||||
{
|
||||
"scope": {
|
||||
"name": "fixture"
|
||||
},
|
||||
"spans": [
|
||||
{
|
||||
"traceId": "01010101010101010101010101010101",
|
||||
"spanId": "0101010101010101",
|
||||
"parentSpanId": "",
|
||||
"name": "review",
|
||||
"kind": 1,
|
||||
"startTimeUnixNano": "1735689600000000000",
|
||||
"endTimeUnixNano": "1735689602000000000",
|
||||
"attributes": [
|
||||
{
|
||||
"key": "gen_ai.input.messages",
|
||||
"value": {
|
||||
"stringValue": "Review the change"
|
||||
}
|
||||
}
|
||||
],
|
||||
"status": {
|
||||
"code": 1,
|
||||
"message": ""
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
|
|
@ -1,11 +1,35 @@
|
|||
use litellm_traces::decode_otlp;
|
||||
use litellm_traces::{ObservationType, Shared};
|
||||
use opentelemetry_proto::tonic::trace::v1::Span;
|
||||
use rstest::rstest;
|
||||
|
||||
const FIXTURE: &[u8] = include_bytes!(
|
||||
"../../../../tests/test_litellm/tracing/fixtures/langsmith_deep_agent_export.json"
|
||||
);
|
||||
|
||||
#[rstest]
|
||||
#[case::root(include_bytes!("fixtures/query_root.json"), ObservationType::Agent, 0, 0)]
|
||||
#[case::children(include_bytes!("fixtures/query_children.json"), ObservationType::Llm, 12, 6)]
|
||||
#[case::alternate(include_bytes!("fixtures/query_alternate.json"), ObservationType::Agent, 0, 0)]
|
||||
#[case::other_team(include_bytes!("fixtures/query_other_team.json"), ObservationType::Agent, 0, 0)]
|
||||
fn query_fixtures_decode_and_normalize(
|
||||
#[case] body: &[u8],
|
||||
#[case] observation_type: ObservationType,
|
||||
#[case] input_tokens: u32,
|
||||
#[case] output_tokens: u32,
|
||||
) {
|
||||
let spans = decode_otlp(body, Some("application/json")).unwrap();
|
||||
let first = &spans[0];
|
||||
assert_eq!(first.normalized.observation_type, observation_type);
|
||||
assert_eq!(first.normalized.input_tokens, input_tokens);
|
||||
assert_eq!(first.normalized.output_tokens, output_tokens);
|
||||
assert!(
|
||||
spans
|
||||
.iter()
|
||||
.all(|span| span.resource_attributes["service.name"] == "fixture")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::json(FIXTURE, Some("application/json"))]
|
||||
fn decodes_neutral_spans(#[case] body: &[u8], #[case] content_type: Option<&str>) {
|
||||
|
|
@ -133,7 +157,7 @@ fn rejects_ids_and_timestamps_that_cannot_be_stored(
|
|||
};
|
||||
assert!(matches!(
|
||||
decode_otlp(&request_with(span).encode_to_vec(), None),
|
||||
Err(litellm_traces::DecodeError::InvalidPayload)
|
||||
Err(litellm_traces::Error::InvalidPayload)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -216,7 +240,7 @@ fn nested_values_are_serialized_once(span: opentelemetry_proto::tonic::trace::v1
|
|||
fn rejects_json_structure_before_building_a_tree(#[case] body: Vec<u8>) {
|
||||
assert!(matches!(
|
||||
decode_otlp(&body, Some("application/json")),
|
||||
Err(litellm_traces::DecodeError::TooLarge)
|
||||
Err(litellm_traces::Error::TooLarge)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -252,7 +276,7 @@ fn protobuf_preflight_rejects_expansion_before_prost_allocates(
|
|||
let body = request.encode_to_vec();
|
||||
assert!(matches!(
|
||||
decode_otlp(&body, None),
|
||||
Err(litellm_traces::DecodeError::TooLarge)
|
||||
Err(litellm_traces::Error::TooLarge)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -310,7 +334,7 @@ fn unique_attribute_expansion_still_respects_decoded_budget(
|
|||
assert!(body.len() < 16 * 1024 * 1024);
|
||||
assert!(matches!(
|
||||
decode_otlp(&body, None),
|
||||
Err(litellm_traces::DecodeError::TooLarge)
|
||||
Err(litellm_traces::Error::TooLarge)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -338,7 +362,7 @@ fn escaped_attribute_expansion_is_bounded_below_four_mib(
|
|||
assert!(body.len() < 4 * 1024 * 1024);
|
||||
assert!(matches!(
|
||||
decode_otlp(&body, None),
|
||||
Err(litellm_traces::DecodeError::TooLarge)
|
||||
Err(litellm_traces::Error::TooLarge)
|
||||
));
|
||||
}
|
||||
|
||||
|
|
@ -387,6 +411,279 @@ fn normalizes_langsmith_fixture() {
|
|||
assert!(tool.normalized.output.starts_with("Based on my research"));
|
||||
}
|
||||
|
||||
fn decode_normalization(
|
||||
span: Span,
|
||||
scope: &str,
|
||||
attributes: &[(&str, &str)],
|
||||
) -> Result<litellm_traces::DecodedSpan, litellm_traces::Error> {
|
||||
use opentelemetry_proto::tonic::{
|
||||
collector::trace::v1::ExportTraceServiceRequest,
|
||||
common::v1::{AnyValue, InstrumentationScope, KeyValue, any_value::Value},
|
||||
trace::v1::{ResourceSpans, ScopeSpans},
|
||||
};
|
||||
use prost::Message;
|
||||
|
||||
let request = ExportTraceServiceRequest {
|
||||
resource_spans: vec![ResourceSpans {
|
||||
scope_spans: vec![ScopeSpans {
|
||||
scope: Some(InstrumentationScope {
|
||||
name: scope.to_owned(),
|
||||
..Default::default()
|
||||
}),
|
||||
spans: vec![Span {
|
||||
attributes: attributes
|
||||
.iter()
|
||||
.map(|(key, value)| KeyValue {
|
||||
key: (*key).to_owned(),
|
||||
value: Some(AnyValue {
|
||||
value: Some(Value::StringValue((*value).to_owned())),
|
||||
}),
|
||||
..Default::default()
|
||||
})
|
||||
.collect(),
|
||||
..span
|
||||
}],
|
||||
..Default::default()
|
||||
}],
|
||||
..Default::default()
|
||||
}],
|
||||
};
|
||||
decode_otlp(&request.encode_to_vec(), None)
|
||||
.map(|spans| spans.into_iter().next().expect("one synthetic span"))
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::agent("invoke_agent", false, ObservationType::Agent)]
|
||||
#[case::chat("chat", false, ObservationType::Llm)]
|
||||
#[case::completion("text_completion", false, ObservationType::Llm)]
|
||||
#[case::content("generate_content", false, ObservationType::Llm)]
|
||||
#[case::tool("execute_tool", false, ObservationType::Tool)]
|
||||
#[case::unknown_root("unknown", true, ObservationType::Agent)]
|
||||
#[case::unknown_child("unknown", false, ObservationType::Chain)]
|
||||
#[case::missing_root("", true, ObservationType::Agent)]
|
||||
#[case::missing_child("", false, ObservationType::Chain)]
|
||||
fn genai_operations_and_parentage_classify_spans(
|
||||
span: Span,
|
||||
#[case] operation: &str,
|
||||
#[case] root: bool,
|
||||
#[case] expected: ObservationType,
|
||||
) {
|
||||
let decoded = decode_normalization(
|
||||
Span {
|
||||
parent_span_id: if root { vec![] } else { vec![3; 8] },
|
||||
..span
|
||||
},
|
||||
"",
|
||||
&[("gen_ai.operation.name", operation)],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(decoded.normalized.observation_type, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::primary("request-model", "messages-in", "messages-out", ["request-model", "messages-in", "messages-out"])]
|
||||
#[case::fallback("", "", "", ["response-model", "tool-in", "tool-out"])]
|
||||
#[case::independent_fallback("request-model", "", "messages-out", ["request-model", "tool-in", "messages-out"])]
|
||||
fn genai_fields_and_consumed_attributes_follow_the_same_fallback(
|
||||
span: Span,
|
||||
#[case] model: &str,
|
||||
#[case] input: &str,
|
||||
#[case] output: &str,
|
||||
#[case] expected: [&str; 3],
|
||||
) {
|
||||
let decoded = decode_normalization(
|
||||
span,
|
||||
"",
|
||||
&[
|
||||
("gen_ai.request.model", model),
|
||||
("gen_ai.response.model", "response-model"),
|
||||
("gen_ai.input.messages", input),
|
||||
("gen_ai.output.messages", output),
|
||||
("gen_ai.tool.call.arguments", "tool-in"),
|
||||
("gen_ai.tool.call.result", "tool-out"),
|
||||
("gen_ai.agent.name", "test-agent"),
|
||||
("gen_ai.response.id", "response-1"),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let fields = &decoded.normalized;
|
||||
assert_eq!(
|
||||
[
|
||||
fields.model.as_str(),
|
||||
fields.input.as_str(),
|
||||
fields.output.as_str()
|
||||
],
|
||||
expected
|
||||
);
|
||||
assert_eq!(fields.agent_name, "test-agent");
|
||||
assert_eq!(fields.litellm_request_id, "response-1");
|
||||
assert_eq!(
|
||||
fields.input,
|
||||
decoded.attributes[decoded.consumed_attributes[0]]
|
||||
);
|
||||
assert_eq!(
|
||||
fields.output,
|
||||
decoded.attributes[decoded.consumed_attributes[1]]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::specific(&[("llm.token_count.prompt", "5"), ("llm.token_count.completion", "9")], 5, 9)]
|
||||
#[case::fallback(&[], 17, 23)]
|
||||
#[case::mixed(&[("llm.token_count.prompt", "5")], 5, 23)]
|
||||
#[case::empty_specific(&[("llm.token_count.prompt", "")], 0, 23)]
|
||||
fn openinference_fields_override_genai_and_usage_falls_back_per_field(
|
||||
span: Span,
|
||||
#[case] token_attributes: &[(&str, &str)],
|
||||
#[case] input_tokens: u32,
|
||||
#[case] output_tokens: u32,
|
||||
) {
|
||||
let attributes = [
|
||||
("openinference.span.kind", "lLm"),
|
||||
("gen_ai.operation.name", "execute_tool"),
|
||||
("llm.model_name", "inference-model"),
|
||||
("gen_ai.request.model", "other-model"),
|
||||
("agent.name", "inference-agent"),
|
||||
("input.value", "inference-input"),
|
||||
("output.value", "inference-output"),
|
||||
("gen_ai.input.messages", "other-input"),
|
||||
("gen_ai.output.messages", "other-output"),
|
||||
("gen_ai.usage.input_tokens", "17"),
|
||||
("gen_ai.usage.output_tokens", "23"),
|
||||
];
|
||||
let combined = attributes
|
||||
.iter()
|
||||
.chain(token_attributes)
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
let decoded = decode_normalization(span, "", &combined).unwrap();
|
||||
let fields = &decoded.normalized;
|
||||
assert_eq!(fields.observation_type, ObservationType::Llm);
|
||||
assert_eq!(fields.model, "inference-model");
|
||||
assert_eq!(fields.agent_name, "inference-agent");
|
||||
assert_eq!(fields.input, "inference-input");
|
||||
assert_eq!(fields.output, "inference-output");
|
||||
assert_eq!(
|
||||
(fields.input_tokens, fields.output_tokens),
|
||||
(input_tokens, output_tokens)
|
||||
);
|
||||
assert_eq!(decoded.consumed_attributes, ["input.value", "output.value"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::scope("langsmith", &[], ObservationType::Agent)]
|
||||
#[case::attribute("other", &[("langsmith.span.kind", "llm")], ObservationType::Llm)]
|
||||
fn langsmith_dispatch_overrides_other_conventions(
|
||||
span: Span,
|
||||
#[case] scope: &str,
|
||||
#[case] convention_attributes: &[(&str, &str)],
|
||||
#[case] observation_type: ObservationType,
|
||||
) {
|
||||
let attributes = [
|
||||
("openinference.span.kind", "TOOL"),
|
||||
("gen_ai.operation.name", "execute_tool"),
|
||||
("langsmith.metadata.lc_agent_name", "test-agent"),
|
||||
(
|
||||
"gen_ai.prompt",
|
||||
r#"{"messages":[{"type":"human","content":"hello"}]}"#,
|
||||
),
|
||||
("gen_ai.completion", "{}"),
|
||||
("input.value", "other-input"),
|
||||
];
|
||||
let combined = attributes
|
||||
.iter()
|
||||
.chain(convention_attributes)
|
||||
.copied()
|
||||
.collect::<Vec<_>>();
|
||||
let decoded = decode_normalization(span, scope, &combined).unwrap();
|
||||
assert_eq!(decoded.normalized.observation_type, observation_type);
|
||||
assert_eq!(decoded.normalized.agent_name, "test-agent");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<serde_json::Value>(&decoded.normalized.input).unwrap(),
|
||||
serde_json::json!([{"role": "user", "content": "hello"}]),
|
||||
);
|
||||
assert_eq!(
|
||||
decoded.consumed_attributes,
|
||||
["gen_ai.prompt", "gen_ai.completion"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::flat(r#"{"messages":[{"type":"human","content":"hello"}]}"#)]
|
||||
#[case::nested(r#"{"messages":[[{"kwargs":{"type":"human","content":"hello"}}],[{"type":"human","content":"ignored batch"}]]}"#)]
|
||||
fn langsmith_llm_messages_preserve_visible_content_and_tool_calls(
|
||||
span: Span,
|
||||
#[case] prompt: &str,
|
||||
) {
|
||||
let completion = r#"{
|
||||
"generations": [[{"message": {"kwargs": {
|
||||
"type": "ai",
|
||||
"content": [
|
||||
{"type": "text", "text": "first"},
|
||||
{"type": "thinking", "thinking": "hidden"},
|
||||
{"type": "tool_use", "id": "call-1"},
|
||||
{"type": "text", "text": "second"}
|
||||
],
|
||||
"tool_calls": [{"name": "search", "args": {"query": "hello"}, "id": "call-1"}],
|
||||
"response_metadata": {"id": "response-1"}
|
||||
}}}]]
|
||||
}"#;
|
||||
let decoded = decode_normalization(
|
||||
span,
|
||||
"langsmith",
|
||||
&[
|
||||
("langsmith.span.kind", "llm"),
|
||||
("gen_ai.prompt", prompt),
|
||||
("gen_ai.completion", completion),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
let input: serde_json::Value = serde_json::from_str(&decoded.normalized.input).unwrap();
|
||||
let output: serde_json::Value = serde_json::from_str(&decoded.normalized.output).unwrap();
|
||||
assert_eq!(
|
||||
input,
|
||||
serde_json::json!([{"role": "user", "content": "hello"}])
|
||||
);
|
||||
assert_eq!(
|
||||
output,
|
||||
serde_json::json!({
|
||||
"role": "assistant",
|
||||
"content": "first\n\nsecond",
|
||||
"tool_calls": [{"name": "search", "args": {"query": "hello"}, "id": "call-1"}],
|
||||
})
|
||||
);
|
||||
assert_eq!(decoded.normalized.litellm_request_id, "response-1");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::string(r#""result""#, "result")]
|
||||
#[case::wrapped(r#"{"output":{"content":"result"}}"#, "result")]
|
||||
#[case::command(
|
||||
r#"{"output":{"update":{"messages":[{"content":"ignored"},{"content":"result"}]}}}"#,
|
||||
"result"
|
||||
)]
|
||||
#[case::object(r#"{"output":{"count":2}}"#, r#"{"count": 2}"#)]
|
||||
#[case::null(r#"{"output":null}"#, "null")]
|
||||
fn langsmith_tool_output_unwraps_supported_shapes(
|
||||
span: Span,
|
||||
#[case] completion: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let decoded = decode_normalization(
|
||||
span,
|
||||
"langsmith",
|
||||
&[
|
||||
("langsmith.span.kind", "tool"),
|
||||
("gen_ai.prompt", "raw-tool-input"),
|
||||
("gen_ai.completion", completion),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(decoded.normalized.observation_type, ObservationType::Tool);
|
||||
assert_eq!(decoded.normalized.input, "raw-tool-input");
|
||||
assert_eq!(decoded.normalized.output, expected);
|
||||
}
|
||||
|
||||
const CLAUDE_AGENT_SDK_FIXTURE: &[u8] =
|
||||
include_bytes!("../../../../tests/test_litellm/tracing/fixtures/claude_agent_sdk_export.json");
|
||||
const CLAUDE_AGENT_SDK_DETAILED_FIXTURE: &[u8] = include_bytes!(
|
||||
|
|
|
|||
34
litellm-rust/crates/traces/tests/query.rs
Normal file
34
litellm-rust/crates/traces/tests/query.rs
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
use litellm_traces::{InvalidQuery, ReadQuery};
|
||||
use rstest::rstest;
|
||||
|
||||
#[rstest]
|
||||
#[case::list_traces("list_traces", ReadQuery::ListTraces)]
|
||||
#[case::trace_spans("trace_spans", ReadQuery::TraceSpans)]
|
||||
#[case::span_detail("span_detail", ReadQuery::SpanDetail)]
|
||||
#[case::span_error("span_error", ReadQuery::SpanError)]
|
||||
#[case::identity("trace_identity", ReadQuery::TraceIdentity)]
|
||||
#[case::spend("spend_by_response_ids", ReadQuery::SpendByResponseIds)]
|
||||
#[case::availability("availability", ReadQuery::Availability)]
|
||||
#[case::agents("agents", ReadQuery::Agents)]
|
||||
#[case::sample("sample", ReadQuery::Sample)]
|
||||
#[case::content("content", ReadQuery::Content)]
|
||||
#[case::evidence("evidence", ReadQuery::Evidence)]
|
||||
fn names_select_the_public_query(#[case] name: &str, #[case] query: ReadQuery) {
|
||||
assert_eq!(ReadQuery::parse(name).unwrap(), query);
|
||||
assert_eq!(query.as_ref(), name);
|
||||
assert_eq!(query.to_string(), name);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unknown("unknown")]
|
||||
#[case::case_sensitive("List_Traces")]
|
||||
#[case::whitespace(" list_traces")]
|
||||
#[case::empty("")]
|
||||
fn invalid_names_preserve_the_public_error(#[case] name: &str) {
|
||||
let error = ReadQuery::parse(name).unwrap_err();
|
||||
assert!(matches!(error, InvalidQuery));
|
||||
assert_eq!(error.to_string(), "unknown ClickHouse read query");
|
||||
}
|
||||
|
||||
#[path = "query/named.rs"]
|
||||
mod named;
|
||||
63
litellm-rust/crates/traces/tests/query/named.rs
Normal file
63
litellm-rust/crates/traces/tests/query/named.rs
Normal file
|
|
@ -0,0 +1,63 @@
|
|||
use litellm_traces::query::named::*;
|
||||
use rstest::rstest;
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
fn round_trip<T: DeserializeOwned + Serialize>(wire: Value) {
|
||||
let contract: T = serde_json::from_value(wire.clone()).unwrap();
|
||||
assert_eq!(serde_json::to_value(contract).unwrap(), wire);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::admin(vec![], "")]
|
||||
#[case::multiple_teams(vec!["team-a", "team-b"], "")]
|
||||
#[case::key(vec!["team-a"], "key")]
|
||||
#[case::teamless_key(vec![], "key")]
|
||||
fn named_requests_preserve_all_access_cases(#[case] teams: Vec<&str>, #[case] key: &str) {
|
||||
let access = json!({"all_teams": u8::from(teams.is_empty() && key.is_empty()), "user_id": "", "team_ids": teams, "api_key_hash": key});
|
||||
round_trip::<ReadAccessParams>(access.clone());
|
||||
let request = |specific: Value| {
|
||||
Value::Object(
|
||||
access
|
||||
.as_object()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.chain(specific.as_object().unwrap())
|
||||
.map(|(key, value)| (key.clone(), value.clone()))
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
round_trip::<ListTracesParams>(request(
|
||||
json!({"start_ms": -1, "end_ms": 10, "cursor_ms": 0, "cursor_trace_id": "", "limit": 100}),
|
||||
));
|
||||
round_trip::<TraceIdentityParams>(request(json!({"trace_id": "trace"})));
|
||||
round_trip::<TraceSpansParams>(request(json!({"trace_id": "trace", "trace_ref": "ref"})));
|
||||
round_trip::<SpanDetailParams>(request(
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span"}),
|
||||
));
|
||||
round_trip::<SpanErrorParams>(request(
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "span_id": "span", "error_offset": u64::MAX, "error_version": "version"}),
|
||||
));
|
||||
round_trip::<SpendByResponseIdsParams>(request(
|
||||
json!({"response_ids": ["response"], "start_ms": -1, "end_ms": 10}),
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn result_contracts_preserve_public_field_names() {
|
||||
round_trip::<ListTracesRow>(
|
||||
json!({"trace_id": "trace", "trace_ref": "ref", "team_id": "team", "api_key_hash": "key", "user_id": "user", "name": "agent", "service": "service", "input_preview": "input", "status": "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"]}),
|
||||
);
|
||||
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"}),
|
||||
);
|
||||
round_trip::<SpanDetailRow>(
|
||||
json!({"span_id": "span", "input": "input", "output": "output", "attributes": {"count": "42"}}),
|
||||
);
|
||||
round_trip::<SpanErrorRow>(
|
||||
json!({"span_id": "span", "message": "error", "total_chars": u64::MAX, "version": "version"}),
|
||||
);
|
||||
round_trip::<SpendByResponseIdsRow>(
|
||||
json!({"request_id": "request", "response_id": "response", "team_id": "team", "api_key": "key", "user": "user", "spend": 0.125, "start_ms": -1}),
|
||||
);
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue