refactor(proxy): inject tracing receiver and access context

litellm-rust/Cargo.toml

@@ -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-storage-clickhouse = { path = "crates/storage-clickhouse" }
 litellm-core = { path = "crates/core" }
 litellm-gateway-mcp = { path = "crates/gateway-mcp" }
 litellm-gateway = { path = "crates/gateway" }

litellm-rust/crates/python-bridge/Cargo.toml

@@ -22,6 +22,7 @@ tiktoken = ["litellm-token-counter/tiktoken"]
 fancy-regex.workspace = true
 litellm-tracing.workspace = true
 litellm-traces.workspace = true
+litellm-storage-clickhouse.workspace = true
 litellm-host.workspace = true
 bytes.workspace = true
 futures-util.workspace = true

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

@@ -1,7 +1,8 @@
 use std::collections::BTreeMap;
 
 use litellm_http::ClientVariant;
-use litellm_traces::{Connection, Error, InsertTable, Parameter, ReadQuery};
+use litellm_storage_clickhouse::Storage;
+use litellm_traces::{Error, InsertTable, Parameter, ReadQuery};
 use pyo3::{
     exceptions::{PyOverflowError, PyRuntimeError, PyValueError},
     prelude::*,
@@ -27,9 +28,7 @@ fn map_error(error: Error) -> PyErr {
 
 #[pyclass]
 pub struct NativeTraceStorage {
-    database: String,
-    writer: Connection,
-    reader: Option<Connection>,
+    storage: Storage,
 }
 
 #[pymethods]
@@ -39,12 +38,7 @@ impl NativeTraceStorage {
     fn new(database: String, url: &str, reader_url: Option<&str>) -> PyResult<Self> {
         litellm_traces::schema_statements(&database, 1, 1).map_err(map_error)?;
         Ok(Self {
-            writer: Connection::writer(url).map_err(map_error)?,
-            reader: reader_url
-                .map(|value| Connection::reader(value, &database))
-                .transpose()
-                .map_err(map_error)?,
-            database,
+            storage: Storage::new(database, url, reader_url).map_err(map_error)?,
         })
     }
 
@@ -55,8 +49,8 @@ impl NativeTraceStorage {
         spend_log_retention_days: u32,
     ) -> PyResult<Bound<'py, PyAny>> {
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
-        let connection = self.writer.clone();
-        let database = self.database.clone();
+        let connection = self.storage.writer().clone();
+        let database = self.storage.database().to_owned();
         crate::execution::run_async(
             py,
             async move {
@@ -83,8 +77,8 @@ impl NativeTraceStorage {
     ) -> PyResult<Bound<'py, PyAny>> {
         let table = InsertTable::parse(table).map_err(map_error)?;
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
-        let connection = self.writer.clone();
-        let database = self.database.clone();
+        let connection = self.storage.writer().clone();
+        let database = self.storage.database().to_owned();
         crate::execution::run_async(
             py,
             async move {
@@ -104,7 +98,7 @@ impl NativeTraceStorage {
         >,
     ) -> PyResult<Bound<'py, PyAny>> {
         let query = litellm_traces::LensQuery::parse(name).map_err(map_error)?;
-        let connection = self.reader.clone().ok_or_else(|| {
+        let connection = self.storage.reader().cloned().ok_or_else(|| {
             PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
         })?;
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;
@@ -127,7 +121,7 @@ impl NativeTraceStorage {
         >,
     ) -> PyResult<Bound<'py, PyAny>> {
         let query = ReadQuery::parse(query).map_err(map_error)?;
-        let connection = self.reader.clone().ok_or_else(|| {
+        let connection = self.storage.reader().cloned().ok_or_else(|| {
             PyRuntimeError::new_err("Trace reads require a separate ClickHouse reader URL")
         })?;
         let client = crate::http::host_client(py, ClientVariant::NoRedirect)?;

litellm-rust/crates/storage-clickhouse/Cargo.toml

@@ -0,0 +1,20 @@
+[package]
+name = "litellm-storage-clickhouse"
+version = "0.1.0"
+description = "Shared ClickHouse connection and HTTP storage for LiteLLM features"
+edition.workspace = true
+license.workspace = true
+repository.workspace = true
+
+[dependencies]
+flate2.workspace = true
+litellm-http.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+thiserror.workspace = true
+url.workspace = true
+
+[dev-dependencies]
+litellm-http = { workspace = true, features = ["test-support"] }
+rstest.workspace = true
+tokio.workspace = true

litellm-rust/crates/storage-clickhouse/README.md

@@ -0,0 +1,5 @@
+# ClickHouse storage
+
+`litellm-storage-clickhouse` exports `Storage`, a shared writer connection and optional reader connection for one ClickHouse 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

litellm-rust/crates/storage-clickhouse/src/error.rs

@@ -0,0 +1,29 @@
+#[derive(Debug, thiserror::Error)]
+pub enum Error {
+    #[error("invalid ClickHouse insert row")]
+    InvalidRow,
+    #[error("invalid ClickHouse insert table")]
+    InvalidTable,
+    #[error("invalid ClickHouse HTTP URL")]
+    InvalidUrl,
+    #[error("database must be a nonempty SQL identifier and retention must be positive")]
+    InvalidSchema,
+    #[error("SQL query must not be empty")]
+    EmptySql,
+    #[error("unknown ClickHouse read query")]
+    InvalidQuery,
+    #[error("ClickHouse query failed with HTTP status {0}")]
+    QueryFailed(u16),
+    #[error("ClickHouse insert failed with HTTP status {0}")]
+    InsertFailed(u16),
+    #[error("ClickHouse insert exceeds the encoded size limit")]
+    InsertTooLarge,
+    #[error("ClickHouse schema setup failed with HTTP status {0}")]
+    SchemaFailed(u16),
+    #[error("ClickHouse query exceeded the response size limit")]
+    ResponseTooLarge,
+    #[error("ClickHouse returned an invalid or failed JSON query response")]
+    InvalidResponse,
+    #[error("ClickHouse query transport failed")]
+    Transport,
+}

litellm-rust/crates/storage-clickhouse/src/insert.rs

@@ -0,0 +1,70 @@
+use std::{io::Write, time::Duration};
+
+use flate2::{Compression, write::GzEncoder};
+use litellm_http::Client;
+
+use crate::{Connection, Error, valid_identifier};
+
+const INSERT_TIMEOUT: Duration = Duration::from_secs(30);
+
+pub async fn insert_encoded_rows(
+    client: &Client,
+    connection: &Connection,
+    database: &str,
+    table: &str,
+    token: &str,
+    encoded: &str,
+) -> Result<(), Error> {
+    if !valid_identifier(database) {
+        return Err(Error::InvalidSchema);
+    }
+    if !valid_identifier(table) {
+        return Err(Error::InvalidTable);
+    }
+    let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
+    encoder
+        .write_all(encoded.as_bytes())
+        .map_err(|_| Error::InvalidRow)?;
+    let body = encoder.finish().map_err(|_| Error::InvalidRow)?;
+    let mut url = connection.url().clone();
+    let existing_pairs: Vec<(String, String)> = url
+        .query_pairs()
+        .filter(|(key, _)| {
+            !matches!(
+                key.as_ref(),
+                "query"
+                    | "async_insert"
+                    | "async_insert_deduplicate"
+                    | "wait_for_async_insert"
+                    | "input_format_skip_unknown_fields"
+                    | "date_time_input_format"
+            )
+        })
+        .map(|(key, value)| (key.into_owned(), value.into_owned()))
+        .collect();
+    url.query_pairs_mut()
+        .clear()
+        .extend_pairs(existing_pairs)
+        .append_pair(
+            "query",
+            &format!("INSERT INTO `{database}`.{} FORMAT JSONEachRow", table),
+        )
+        .append_pair("insert_deduplication_token", token)
+        .append_pair("async_insert", "1")
+        .append_pair("async_insert_deduplicate", "1")
+        .append_pair("wait_for_async_insert", "1")
+        .append_pair("input_format_skip_unknown_fields", "0")
+        .append_pair("date_time_input_format", "best_effort");
+    let response = client
+        .post(url)
+        .timeout(INSERT_TIMEOUT)
+        .header("Content-Encoding", "gzip")
+        .body(body)
+        .send()
+        .await
+        .map_err(|_| Error::Transport)?;
+    if !response.status().is_success() {
+        return Err(Error::InsertFailed(response.status().as_u16()));
+    }
+    Ok(())
+}

litellm-rust/crates/storage-clickhouse/src/lib.rs

@@ -0,0 +1,127 @@
+mod error;
+mod insert;
+mod read;
+
+pub use error::Error;
+pub use insert::insert_encoded_rows;
+pub use read::{Parameter, execute_read};
+use url::Url;
+
+#[derive(Clone)]
+pub struct Connection {
+    url: Url,
+}
+
+impl Connection {
+    pub fn parse(value: &str) -> Result<Self, Error> {
+        let url = Url::parse(value).map_err(|_| Error::InvalidUrl)?;
+        if !matches!(url.scheme(), "http" | "https") || url.host().is_none() {
+            return Err(Error::InvalidUrl);
+        }
+        Ok(Self { url })
+    }
+
+    pub fn configured(
+        url: &str,
+        database: &str,
+        user: &str,
+        password: &str,
+    ) -> Result<Self, Error> {
+        let mut connection = Self::parse(url)?;
+        connection
+            .url
+            .set_username(user)
+            .map_err(|_| Error::InvalidUrl)?;
+        connection
+            .url
+            .set_password(Some(password))
+            .map_err(|_| Error::InvalidUrl)?;
+        let pairs: Vec<_> = connection
+            .url
+            .query_pairs()
+            .filter(|(key, _)| !matches!(key.as_ref(), "database" | "user" | "password"))
+            .map(|(key, value)| (key.into_owned(), value.into_owned()))
+            .collect();
+        connection
+            .url
+            .query_pairs_mut()
+            .clear()
+            .extend_pairs(pairs)
+            .append_pair("database", database);
+        Ok(connection)
+    }
+
+    pub fn writer(url: &str) -> Result<Self, Error> {
+        let mut connection = Self::parse(url)?;
+        let pairs: Vec<_> = connection
+            .url
+            .query_pairs()
+            .filter(|(key, _)| !matches!(key.as_ref(), "database" | "readonly" | "query"))
+            .map(|(key, value)| (key.into_owned(), value.into_owned()))
+            .collect();
+        connection.url.query_pairs_mut().clear().extend_pairs(pairs);
+        Ok(connection)
+    }
+
+    pub fn reader(url: &str, database: &str) -> Result<Self, Error> {
+        let mut connection = Self::parse(url)?;
+        let pairs: Vec<_> = connection
+            .url
+            .query_pairs()
+            .filter(|(key, _)| key != "database")
+            .map(|(key, value)| (key.into_owned(), value.into_owned()))
+            .collect();
+        connection
+            .url
+            .query_pairs_mut()
+            .clear()
+            .extend_pairs(pairs)
+            .append_pair("database", database);
+        Ok(connection)
+    }
+
+    pub fn url(&self) -> &Url {
+        &self.url
+    }
+}
+
+#[derive(Clone)]
+pub struct Storage {
+    database: String,
+    writer: Connection,
+    reader: Option<Connection>,
+}
+
+impl Storage {
+    pub fn new(database: String, url: &str, reader_url: Option<&str>) -> Result<Self, Error> {
+        if !valid_identifier(&database) {
+            return Err(Error::InvalidSchema);
+        }
+        Ok(Self {
+            writer: Connection::writer(url)?,
+            reader: reader_url
+                .map(|value| Connection::reader(value, &database))
+                .transpose()?,
+            database,
+        })
+    }
+
+    pub fn database(&self) -> &str {
+        &self.database
+    }
+
+    pub fn writer(&self) -> &Connection {
+        &self.writer
+    }
+
+    pub fn reader(&self) -> Option<&Connection> {
+        self.reader.as_ref()
+    }
+}
+
+pub(crate) fn valid_identifier(value: &str) -> bool {
+    !value.is_empty()
+        && value
+            .bytes()
+            .all(|c| c.is_ascii_alphanumeric() || c == b'_')
+}