feat(tracing): add scoped SQL queries and schema-aware help

litellm-rust/Cargo.toml

@@ -65,6 +65,7 @@ litellm-token-counter-tiktoken = { path = "crates/token-counter-tiktoken" }
 litellm-host-python = { path = "crates/host-python" }
 litellm-python-compat = { path = "crates/python-compat" }
 
+askama = { version = "0.16.1", default-features = false, features = ["derive", "std"] }
 tracing = "0.1"
 axum = { version = "0.8.9", default-features = false, features = ["http1", "tokio", "multipart"] }
 axum-login = "0.18.0"

litellm-rust/crates/python-bridge/src/routes/traces.rs

@@ -3,7 +3,9 @@ use std::collections::BTreeMap;
 use litellm_host_python::{FromPythonCache, ToPythonCache};
 use litellm_http::ClientVariant;
 use litellm_storage_clickhouse::Storage;
-use litellm_traces::{Error, InsertTable, Parameter, ReadQuery, Shared};
+use litellm_traces::{
+    Error, InsertTable, Parameter, QueryAccessError, QueryReaders, QueryScope, ReadQuery, Shared,
+};
 use prost::Message;
 use pyo3::{
     exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
@@ -46,9 +48,25 @@ fn map_error(error: Error) -> PyErr {
     }
 }
 
+fn map_sql_error(error: Error) -> PyErr {
+    match error {
+        Error::QueryFailed(400 | 404) => PyValueError::new_err(error.to_string()),
+        error => map_error(error),
+    }
+}
+
+fn map_query_access_error(error: QueryAccessError) -> PyErr {
+    match error {
+        QueryAccessError::Storage(error) => map_sql_error(error),
+        QueryAccessError::InvalidScope => PyValueError::new_err(error.to_string()),
+        error => PyRuntimeError::new_err(error.to_string()),
+    }
+}
+
 #[pyclass]
 pub struct NativeTraceStorage {
     storage: Storage,
+    query_readers: QueryReaders,
 }
 
 #[pymethods]
@@ -57,8 +75,13 @@ impl NativeTraceStorage {
     #[pyo3(signature = (database, url, reader_url = None))]
     fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
         litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
+        let storage = Storage::new(database, url, reader_url).map_err(map_error)?;
         Ok(Self {
-            storage: Storage::new(database, url, reader_url).map_err(map_error)?,
+            query_readers: QueryReaders::new(
+                storage.writer().clone(),
+                storage.database().to_owned(),
+            ),
+            storage,
         })
     }
 
@@ -107,6 +130,52 @@ impl NativeTraceStorage {
         )
     }
 
+    fn query_sql<'py>(
+        &self,
+        py: Python<'py>,
+        sql: String,
+        #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
+        secret: String,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        if sql.trim().is_empty() {
+            return Err(map_error(Error::EmptySql));
+        }
+        let readers = self.query_readers.clone();
+        let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
+        crate::execution::run_async(
+            py,
+            async move {
+                let _permit = readers.acquire()?;
+                let connection = readers.connection(&client, &scope, &secret).await?;
+                litellm_traces::query_sql(&client, &connection, &sql)
+                    .await
+                    .map_err(QueryAccessError::Storage)
+            },
+            map_query_access_error,
+        )
+    }
+
+    fn query_help<'py>(
+        &self,
+        py: Python<'py>,
+        #[pyo3(from_py_with = litellm_host_python::from_py_argument)] scope: QueryScope,
+        secret: String,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let readers = self.query_readers.clone();
+        let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
+        crate::execution::run_async(
+            py,
+            async move {
+                let _permit = readers.acquire()?;
+                let connection = readers.connection(&client, &scope, &secret).await?;
+                litellm_traces::query_help(&client, &connection)
+                    .await
+                    .map_err(QueryAccessError::Storage)
+            },
+            map_query_access_error,
+        )
+    }
+
     fn lens_query<'py>(
         &self,
         py: Python<'py>,

litellm-rust/crates/traces/Cargo.toml

@@ -6,9 +6,13 @@ license.workspace = true
 repository.workspace = true
 
 [dependencies]
+askama.workspace = true
 base64.workspace = true
 flate2.workspace = true
+futures-util.workspace = true
+hmac = "0.12.1"
 indexmap = { version = "2", features = ["serde"] }
+moka.workspace = true
 opentelemetry-proto = { workspace = true, features = ["gen-tonic-messages", "trace", "with-serde"] }
 prost.workspace = true
 time = { workspace = true, features = ["formatting"] }
@@ -20,6 +24,8 @@ serde = { workspace = true, features = ["rc"] }
 serde_json.workspace = true
 strum.workspace = true
 thiserror.workspace = true
+tokio.workspace = true
+url.workspace = true
 
 [dev-dependencies]
 criterion.workspace = true

litellm-rust/crates/traces/src/error.rs

@@ -7,3 +7,23 @@ pub enum DecodeError {
     #[error("OTLP token count is outside the storage range")]
     TokenCountOutOfRange,
 }
+
+#[derive(Debug, thiserror::Error)]
+pub enum QueryAccessError {
+    #[error("trace SQL queries require a configured proxy master key")]
+    MissingSecret,
+    #[error("invalid trace query scope")]
+    InvalidScope,
+    #[error("trace SQL query concurrency limit exceeded")]
+    Busy,
+    #[error(
+        "ClickHouse reader provisioning failed with HTTP status {0}; the configured connection must be allowed to manage users, row policies, and SELECT grants on the trace tables"
+    )]
+    ProvisionFailed(u16),
+    #[error("ClickHouse reader provisioning transport failed")]
+    ProvisionTransport,
+    #[error(transparent)]
+    Storage(#[from] litellm_storage_clickhouse::Error),
+    #[error(transparent)]
+    Cached(#[from] std::sync::Arc<QueryAccessError>),
+}

litellm-rust/crates/traces/src/lib.rs

@@ -2,17 +2,22 @@ mod error;
 mod insert;
 mod normalize;
 mod otlp;
+mod query;
+mod query_access;
 mod schema;
 mod shared;
 mod sql;
 
-pub use error::DecodeError;
+pub use error::{DecodeError, QueryAccessError};
 pub use insert::{InsertRow, InsertTable, encode_rows, insert_rows, insert_shared_rows};
 pub use litellm_storage_clickhouse::{Connection, Error, Parameter, execute_read};
 pub use normalize::{
     NORMALIZED_FIELD_DEFINITIONS, NormalizedFieldDefinition, NormalizedSpan, ObservationType,
 };
 pub use otlp::{DecodedSpan, decode_otlp};
+pub use query_access::{QueryReaders, QueryScope};
 pub use schema::{ensure_schema, schema_statements};
 pub use shared::{Shared, SharedIdentity};
 pub use sql::{LensQuery, ReadQuery, execute_named_read};
+
+pub use query::{query_help, query_sql};