From 81bb960d8c2f6384a93a92c6fd93470f30296040 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Tue, 8 Sep 2026 14:57:57 -0700 Subject: [PATCH] refactor(rust): encapsulate trace parity in python bridge --- litellm-rust/Cargo.lock | 4 +- litellm-rust/crates/ai-gateway/Cargo.toml | 2 - litellm-rust/crates/ai-gateway/README.md | 3 +- litellm-rust/crates/ai-gateway/src/lib.rs | 2 - litellm-rust/crates/core/Cargo.toml | 3 - litellm-rust/crates/core/src/lib.rs | 2 - .../crates/core/src/observability/mod.rs | 59 ----------------- litellm-rust/crates/python-bridge/Cargo.toml | 9 ++- litellm-rust/crates/python-bridge/src/lib.rs | 5 +- .../src/routes/audio_transcription.rs | 4 +- .../python-bridge/src/routes/definition.rs | 4 +- .../src/routes/gateway_messages.rs | 30 --------- .../crates/python-bridge/src/routes/mod.rs | 5 +- .../src/trace_parity/collector.rs} | 63 +++++++++++++------ .../src/trace_parity/gateway.rs} | 35 ++++++++--- .../mod.rs} | 7 ++- .../strategies/trace_parity/AGENTS.md | 2 +- 17 files changed, 98 insertions(+), 141 deletions(-) delete mode 100644 litellm-rust/crates/core/src/observability/mod.rs delete mode 100644 litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs rename litellm-rust/crates/{core/src/observability/function_trace.rs => python-bridge/src/trace_parity/collector.rs} (77%) rename litellm-rust/crates/{ai-gateway/src/trace_parity.rs => python-bridge/src/trace_parity/gateway.rs} (68%) rename litellm-rust/crates/python-bridge/src/{function_trace.rs => trace_parity/mod.rs} (88%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 10dc3cad13d..badf1dcaf96 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1466,7 +1466,6 @@ dependencies = [ "tokio", "tokio-tungstenite 0.24.0", "tracing", - "tracing-subscriber", ] [[package]] @@ -1484,6 +1483,7 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "axum", "criterion", "litellm-ai-gateway", "litellm-core", @@ -1493,7 +1493,9 @@ dependencies = [ "serde", "serde_json", "tokio", + "tower", "tracing", + "tracing-subscriber", ] [[package]] diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index 9400416f22a..c22bf0edb88 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -24,7 +24,6 @@ futures-util.workspace = true serde_json.workspace = true axum = { workspace = true, features = ["ws"], optional = true } serde.workspace = true -tower = { version = "0.5.3", features = ["util"], optional = true } [features] default = [] @@ -32,7 +31,6 @@ server = ["dep:axum", "dep:litellm-gateway-auth"] # Build the gateway's config from the proxy YAML via an embedded Python # interpreter (links libpython; requires `litellm` importable at runtime). python-config = ["litellm-config/python"] -trace-parity = ["server", "dep:tower", "litellm-core/observability"] [dev-dependencies] futures-channel = "0.3" diff --git a/litellm-rust/crates/ai-gateway/README.md b/litellm-rust/crates/ai-gateway/README.md index 48659bc571e..9d1e4bff5cc 100644 --- a/litellm-rust/crates/ai-gateway/README.md +++ b/litellm-rust/crates/ai-gateway/README.md @@ -18,7 +18,8 @@ deployment, and adapts frames while core dials OpenAI and splices the session. Dependency direction is acyclic: config depends on core, the gateway depends on config and core, and the Python bridge depends on core and Python interop. Its -optional `trace-parity` diagnostics also depend on the gateway. +private `trace_parity` module owns trace collection and gateway fixtures, enabled +by the bridge’s `trace-parity` feature and the gateway’s `server` feature - **Client endpoint:** `wss:///v1/realtime?model=` (WebSocket) - **Auth:** `Authorization: Bearer $LITELLM_MASTER_KEY` (fails closed if unset) diff --git a/litellm-rust/crates/ai-gateway/src/lib.rs b/litellm-rust/crates/ai-gateway/src/lib.rs index 21e31d748ce..e97e29f17d3 100644 --- a/litellm-rust/crates/ai-gateway/src/lib.rs +++ b/litellm-rust/crates/ai-gateway/src/lib.rs @@ -14,7 +14,5 @@ pub mod io; pub mod routes; #[cfg(feature = "server")] pub mod state; -#[cfg(feature = "trace-parity")] -pub mod trace_parity; mod constants; diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 8b1b3e7dfdf..0697b11a410 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -19,7 +19,6 @@ thiserror.workspace = true tracing.workspace = true tokio = { workspace = true, features = ["rt", "sync", "time"] } tokio-tungstenite.workspace = true -tracing-subscriber = { workspace = true, optional = true } litellm-auth-aws = { workspace = true, optional = true } [features] @@ -27,10 +26,8 @@ default = [] bedrock-auth = [ "dep:litellm-auth-aws", ] -observability = ["dep:tracing-subscriber"] [dev-dependencies] futures-channel = "0.3" rstest.workspace = true tokio = { workspace = true, features = ["io-util", "macros", "net", "rt-multi-thread"] } -tracing-subscriber.workspace = true diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index c826adb6347..12dde12fc7e 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -7,8 +7,6 @@ pub mod http_utils; pub mod integrations; pub mod lifecycle; pub mod messages; -#[cfg(any(feature = "observability", test))] -pub mod observability; pub mod ocr; pub mod providers; pub mod realtime; diff --git a/litellm-rust/crates/core/src/observability/mod.rs b/litellm-rust/crates/core/src/observability/mod.rs deleted file mode 100644 index 3f9da8e2bb4..00000000000 --- a/litellm-rust/crates/core/src/observability/mod.rs +++ /dev/null @@ -1,59 +0,0 @@ -use tracing::span::Id; -use tracing::{Level, Metadata, Subscriber}; -use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn}; -use tracing_subscriber::layer::Context; -use tracing_subscriber::registry::LookupSpan; - -use crate::constants::FUNCTION_TRACE_TARGET; - -pub mod function_trace; - -pub use function_trace::{FunctionTrace, FunctionTraceEvent}; - -pub fn function_trace_filter() -> FilterFn) -> bool> { - filter_fn(|metadata| { - metadata.is_span() - && metadata.target() == FUNCTION_TRACE_TARGET - && *metadata.level() == Level::TRACE - }) - .with_max_level_hint(LevelFilter::TRACE) -} - -pub fn span_depth(context: &Context<'_, S>, id: &Id) -> usize -where - S: Subscriber + for<'lookup> LookupSpan<'lookup>, -{ - context - .span(id) - .map(|span| span.scope().skip(1).count()) - .unwrap_or_default() -} - -#[cfg(test)] -mod tests { - use tracing::instrument::WithSubscriber; - - use super::*; - - #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] - async fn instrumented_with_literal_target() {} - - #[tokio::test] - async fn literal_instrument_target_matches_filter_constant() { - assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace"); - - let trace = FunctionTrace::default(); - instrumented_with_literal_target() - .with_subscriber(trace.dispatcher()) - .await; - - let events = trace.events(); - assert_eq!(events.len(), 1); - assert_eq!(events[0].id, 0); - assert_eq!(events[0].parent_id, None); - assert_eq!(events[0].function, "instrumented_with_literal_target"); - assert_eq!(events[0].module_path, Some(module_path!())); - assert_eq!(events[0].file, Some(file!())); - assert!(events[0].line.is_some()); - } -} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 0b8861961c9..900be52909f 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,11 +17,16 @@ panic-test = [] trace-parity = [ "dep:tracing", "dep:litellm-ai-gateway", - "litellm-core/observability", - "litellm-ai-gateway/trace-parity", + "dep:tracing-subscriber", + "dep:axum", + "dep:tower", + "litellm-ai-gateway/server", ] [dependencies] +tracing-subscriber = { workspace = true, optional = true } +axum = { workspace = true, optional = true } +tower = { version = "0.5.3", features = ["util"], optional = true } tracing = { workspace = true, optional = true } litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false, optional = true } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 1433ef89542..d2f92028700 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -3,11 +3,11 @@ mod diagnostics; mod driver; mod errors; -#[cfg(feature = "trace-parity")] -mod function_trace; mod marshal; mod retained; mod routes; +#[cfg(feature = "trace-parity")] +mod trace_parity; use pyo3::prelude::*; @@ -88,6 +88,7 @@ mod tests { "amessages", "chat_completions", "achat_completions", + "chat_completions_decline", "gateway_messages", ] ); diff --git a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs index 5d9b36fdaf5..520d8c40362 100644 --- a/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs +++ b/litellm-rust/crates/python-bridge/src/routes/audio_transcription.rs @@ -150,7 +150,7 @@ mod trace { })?; litellm_python_interop::run_sync( py, - crate::function_trace::capture(future), + crate::trace_parity::capture(future), core_error_to_pyerr, ) } @@ -181,7 +181,7 @@ mod trace { })?; litellm_python_interop::run_async( py, - crate::function_trace::capture(future), + crate::trace_parity::capture(future), core_error_to_pyerr, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/definition.rs b/litellm-rust/crates/python-bridge/src/routes/definition.rs index 4fd409b41e3..d76c51d2dba 100644 --- a/litellm-rust/crates/python-bridge/src/routes/definition.rs +++ b/litellm-rust/crates/python-bridge/src/routes/definition.rs @@ -79,7 +79,7 @@ mod tests { let future = prepare_echo(EchoInputs { value })?; litellm_python_interop::run_sync( py, - crate::function_trace::capture(future), + crate::trace_parity::capture(future), map_error, ) } @@ -89,7 +89,7 @@ mod tests { let future = prepare_echo(EchoInputs { value })?; litellm_python_interop::run_async( py, - crate::function_trace::capture(future), + crate::trace_parity::capture(future), map_error, ) } diff --git a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs b/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs deleted file mode 100644 index 8fc44bfeacd..00000000000 --- a/litellm-rust/crates/python-bridge/src/routes/gateway_messages.rs +++ /dev/null @@ -1,30 +0,0 @@ -use pyo3::prelude::*; -use serde_json::Value; - -use crate::errors::core_error_to_pyerr; -use litellm_python_interop::run_async; - -#[pyfunction] -fn gateway_messages<'py>( - py: Python<'py>, - model_alias: String, - provider_model: String, - api_base: String, - #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, -) -> PyResult> { - let future = litellm_ai_gateway::trace_parity::messages_request( - model_alias, - provider_model, - api_base, - body, - ); - run_async( - py, - crate::function_trace::capture(future), - core_error_to_pyerr, - ) -} - -pub(super) fn register_trace(module: &Bound<'_, PyModule>) -> PyResult<()> { - super::definition::add_function(module, wrap_pyfunction!(gateway_messages, module)?) -} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 9b6e0c8d159..607625229a9 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -2,9 +2,6 @@ use pyo3::prelude::*; mod definition; -#[cfg(feature = "trace-parity")] -mod gateway_messages; - mod audio_transcription; mod chat_completions; mod messages; @@ -24,7 +21,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { audio_transcription::register_trace(&trace)?; messages::register_trace(&trace)?; chat_completions::register_trace(&trace)?; - gateway_messages::register_trace(&trace)?; + crate::trace_parity::register_gateway(&trace)?; module.add_submodule(&trace)?; } Ok(()) diff --git a/litellm-rust/crates/core/src/observability/function_trace.rs b/litellm-rust/crates/python-bridge/src/trace_parity/collector.rs similarity index 77% rename from litellm-rust/crates/core/src/observability/function_trace.rs rename to litellm-rust/crates/python-bridge/src/trace_parity/collector.rs index 2031e35901c..fa5c2ec35f3 100644 --- a/litellm-rust/crates/core/src/observability/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/trace_parity/collector.rs @@ -3,32 +3,42 @@ use std::sync::{Arc, Mutex}; use serde::Serialize; use tracing::span::{Attributes, Id}; -use tracing::{Dispatch, Subscriber}; +use tracing::{Dispatch, Level, Metadata, Subscriber}; +use tracing_subscriber::filter::{FilterFn, LevelFilter, filter_fn}; use tracing_subscriber::layer::Context; use tracing_subscriber::prelude::*; use tracing_subscriber::registry::LookupSpan; use tracing_subscriber::{Layer, Registry}; -use super::function_trace_filter; +use litellm_core::constants::FUNCTION_TRACE_TARGET; + +fn function_trace_filter() -> FilterFn) -> bool> { + filter_fn(|metadata| { + metadata.is_span() + && metadata.target() == FUNCTION_TRACE_TARGET + && *metadata.level() == Level::TRACE + }) + .with_max_level_hint(LevelFilter::TRACE) +} #[derive(Clone, Debug, PartialEq, Serialize)] -pub struct FunctionTraceEvent { - pub id: usize, - pub parent_id: Option, - pub function: &'static str, - pub module_path: Option<&'static str>, - pub file: Option<&'static str>, - pub line: Option, +pub(super) struct FunctionTraceEvent { + id: usize, + parent_id: Option, + function: &'static str, + module_path: Option<&'static str>, + file: Option<&'static str>, + line: Option, } #[derive(Clone, Default)] -pub struct FunctionTrace { +pub(super) struct FunctionTrace { events: Arc>>, span_events: Arc>>, } impl FunctionTrace { - pub fn dispatcher(&self) -> Dispatch { + pub(super) fn dispatcher(&self) -> Dispatch { Dispatch::new( Registry::default().with( FunctionTraceLayer { @@ -39,7 +49,7 @@ impl FunctionTrace { ) } - pub fn events(&self) -> Vec { + pub(super) fn events(&self) -> Vec { self.events .lock() .unwrap_or_else(|error| error.into_inner()) @@ -90,9 +100,8 @@ where #[cfg(test)] mod tests { - use crate::constants::FUNCTION_TRACE_TARGET; - use super::*; + use tracing::instrument::WithSubscriber; fn event( id: usize, @@ -129,8 +138,6 @@ mod tests { #[tokio::test] async fn concurrent_futures_keep_separate_traces_across_yields() { - use tracing::instrument::WithSubscriber; - let first = FunctionTrace::default(); let second = FunctionTrace::default(); let outside = FunctionTrace::default(); @@ -161,8 +168,6 @@ mod tests { #[tokio::test] async fn concurrent_siblings_keep_the_same_parent() { - use tracing::instrument::WithSubscriber; - let trace = FunctionTrace::default(); concurrent_parent() .with_subscriber(trace.dispatcher()) @@ -212,4 +217,26 @@ mod tests { vec![event(0, None, "outer"), event(1, Some(0), "inner")] ); } + + #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] + async fn instrumented_with_literal_target() {} + + #[tokio::test] + async fn literal_instrument_target_matches_filter_constant() { + assert_eq!(FUNCTION_TRACE_TARGET, "litellm::function_trace"); + + let trace = FunctionTrace::default(); + instrumented_with_literal_target() + .with_subscriber(trace.dispatcher()) + .await; + + let events = trace.events(); + assert_eq!(events.len(), 1); + assert_eq!(events[0].id, 0); + assert_eq!(events[0].parent_id, None); + assert_eq!(events[0].function, "instrumented_with_literal_target"); + assert_eq!(events[0].module_path, Some(module_path!())); + assert_eq!(events[0].file, Some(file!())); + assert!(events[0].line.is_some()); + } } diff --git a/litellm-rust/crates/ai-gateway/src/trace_parity.rs b/litellm-rust/crates/python-bridge/src/trace_parity/gateway.rs similarity index 68% rename from litellm-rust/crates/ai-gateway/src/trace_parity.rs rename to litellm-rust/crates/python-bridge/src/trace_parity/gateway.rs index 614852c541d..6408ffb0b74 100644 --- a/litellm-rust/crates/ai-gateway/src/trace_parity.rs +++ b/litellm-rust/crates/python-bridge/src/trace_parity/gateway.rs @@ -1,27 +1,28 @@ -//! Harness-only in-process adapters. Never mounted as production routes. - use std::sync::Arc; use axum::body::{Body, to_bytes}; use axum::http::header::{AUTHORIZATION, CONTENT_TYPE}; use axum::http::{Request, StatusCode}; +use litellm_ai_gateway::io::realtime_pool::RealtimePool; +use litellm_ai_gateway::routes; +use litellm_ai_gateway::state::AppState; use litellm_core::Error; use litellm_core::router::{Deployment, LiteLLMParams, Router as ModelRouter}; +use litellm_python_interop::run_async; +use pyo3::prelude::*; use serde::Serialize; use serde_json::Value; use tower::ServiceExt; -use crate::io::realtime_pool::RealtimePool; -use crate::routes; -use crate::state::AppState; +use crate::errors::core_error_to_pyerr; #[derive(Debug, Serialize)] -pub struct GatewayResponse { - pub status: u16, - pub body: Value, +struct GatewayResponse { + status: u16, + body: Value, } -pub async fn messages_request( +async fn messages_request( model_alias: String, provider_model: String, api_base: String, @@ -63,3 +64,19 @@ pub async fn messages_request( body, }) } + +#[pyfunction] +fn gateway_messages<'py>( + py: Python<'py>, + model_alias: String, + provider_model: String, + api_base: String, + #[pyo3(from_py_with = litellm_python_interop::from_py)] body: Value, +) -> PyResult> { + let future = messages_request(model_alias, provider_model, api_base, body); + run_async(py, super::capture(future), core_error_to_pyerr) +} + +pub(crate) fn register_gateway(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(gateway_messages, module)?) +} diff --git a/litellm-rust/crates/python-bridge/src/function_trace.rs b/litellm-rust/crates/python-bridge/src/trace_parity/mod.rs similarity index 88% rename from litellm-rust/crates/python-bridge/src/function_trace.rs rename to litellm-rust/crates/python-bridge/src/trace_parity/mod.rs index bc3c962f7a3..be26b9d7a75 100644 --- a/litellm-rust/crates/python-bridge/src/function_trace.rs +++ b/litellm-rust/crates/python-bridge/src/trace_parity/mod.rs @@ -1,10 +1,15 @@ +mod collector; +mod gateway; + use std::fmt::Display; use std::future::Future; -use litellm_core::observability::{FunctionTrace, FunctionTraceEvent}; use serde::Serialize; use tracing::instrument::WithSubscriber; +use self::collector::{FunctionTrace, FunctionTraceEvent}; +pub(crate) use gateway::register_gateway; + #[derive(Serialize)] pub(crate) struct TracedResponse { #[serde(skip_serializing_if = "Option::is_none")] diff --git a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md index bb7cb8c91d8..b862acf8fef 100644 --- a/tests/rust-python-harness/strategies/trace_parity/AGENTS.md +++ b/tests/rust-python-harness/strategies/trace_parity/AGENTS.md @@ -1 +1 @@ -Maps Python profiler frames onto feature-gated Rust span names via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response. +Maps Python profiler frames onto Rust span names captured by the bridge’s `trace-parity` feature via an explicit per-case mapping (Rust span name is the identity) and compares steps, order, and nesting of both live traces against a replayed provider response.