diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index afe60136c67..d2a548a5bad 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -567,6 +567,17 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" +[[package]] +name = "colored_json" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e35980a1b846f8e3e359fd18099172a0857140ba9230affc4f71348081e039b6" +dependencies = [ + "serde", + "serde_json", + "yansi", +] + [[package]] name = "const-oid" version = "0.10.2" @@ -1265,6 +1276,9 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "bytes", + "colored_json", + "futures-util", "rand 0.8.7", "reqwest", "serde", @@ -1272,6 +1286,7 @@ dependencies = [ "sha2 0.10.9", "thiserror 2.0.19", "tokio", + "url", ] [[package]] @@ -2598,6 +2613,12 @@ version = "0.13.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" +[[package]] +name = "yansi" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfe53a6657fd280eaa890a3bc59152892ffa3e30101319d168b781ed6529b049" + [[package]] name = "yoke" version = "0.8.3" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 889d8da9895..6ba82490bad 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -32,3 +32,5 @@ base64 = "0.22" bytes = "1" aws-smithy-eventstream = "0.60.3" aws-smithy-types = "1.6.1" +colored_json = "5.0" +url = "2.5" diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index abfceb0e866..9bb7a9e3188 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -20,6 +20,7 @@ use litellm_ai_gateway::integrations::custom_logger::CustomLogger; use litellm_ai_gateway::integrations::litellm_python_proxy_api::LiteLLMPythonProxyAPILogger; #[cfg(feature = "python-config")] use litellm_ai_gateway::python; +use litellm_core::logging::console::hook; /// Bind to localhost by default so the gateway is not a public, unauthenticated /// provider proxy out of the box. Override with `HOST` (e.g. `0.0.0.0`). @@ -67,11 +68,15 @@ async fn main() { ); } + let debug_logging = std::env::var("LITELLM_LOG") + .ok() + .is_some_and(|value| value.eq_ignore_ascii_case("DEBUG")); let state = AppState { router, master_key, loggers: Arc::new(loggers), realtime_pool, + logging_sink: hook(debug_logging), }; let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string()); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs index 9bc2818b6e7..c146ef26228 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs @@ -20,19 +20,10 @@ use litellm_core::providers::vertex_ai::ocr::transformation::{ use crate::client::http_client; -const ERROR_BODY_MAX_CHARS: usize = 256; const AZURE_DOCUMENT_INTELLIGENCE_POLL_TIMEOUT_SECS: u64 = 120; const DEFAULT_MAX_IMAGE_URL_DOWNLOAD_SIZE_MB: f64 = 50.0; const MAX_SAFE_FETCH_REDIRECTS: usize = 10; -pub(super) fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") -} - pub(super) fn ocr_provider_config( provider: &str, model: &str, @@ -269,7 +260,7 @@ pub(super) async fn convert_document_url_to_data_uri(document: Value) -> CoreRes let body = response.text().await.unwrap_or_default(); return Err(CoreError::Http { status: status.as_u16(), - body: truncate_error_body(&body), + body: litellm_core::utils::truncate_error_body(&body), }); } let content_type = response @@ -383,7 +374,7 @@ pub(super) async fn poll_document_intelligence( if !status.is_success() { return Err(CoreError::Http { status: status.as_u16(), - body: truncate_error_body(&text), + body: litellm_core::utils::truncate_error_body(&text), }); } let response_json: Value = serde_json::from_str(&text).map_err(|err| { diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 1de34eb400e..52b6d27a944 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -1,13 +1,34 @@ use litellm_core::CoreResult; use litellm_core::error::CoreError; +use litellm_core::logging::http::{JsonRequest, execute_json}; use litellm_core::ocr::transformation::OcrResponseHandling; use serde_json::Value; -use super::common_utils::{poll_document_intelligence, truncate_error_body}; +use super::common_utils::poll_document_intelligence; use super::types::ProviderOcrRequest; use crate::client::http_client; pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> CoreResult { + if request.config.response_handling() != OcrResponseHandling::AzureDocumentIntelligencePoll { + let response = execute_json::( + http_client(), + JsonRequest { + logger: request.logger, + model: request.model.clone(), + stream: false, + url: request.url, + headers: request.upstream_headers, + body: request.body, + timeout: request.timeout, + }, + ) + .await?; + return Ok(request + .config + .transform_ocr_response(&request.model, response)? + .into_json()); + } + let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { request_builder = request_builder.header(key, value); @@ -57,7 +78,7 @@ pub(crate) async fn execute_ocr_provider_call(request: ProviderOcrRequest) -> Co if !status.is_success() { return Err(CoreError::Http { status: status.as_u16(), - body: truncate_error_body(&text), + body: litellm_core::utils::truncate_error_body(&text), }); } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index ffe2e0122c0..f9be8d131b4 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -111,6 +111,7 @@ impl OcrLifecycleHooks { body, upstream_headers, timeout: request.timeout, + logger: request.logger, }) } diff --git a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs index 6231393c889..1211c59a3d1 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/prepare.rs @@ -1,3 +1,4 @@ +use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; @@ -7,6 +8,8 @@ use super::hooks::OcrLifecycleHooks; use super::types::{OcrRequest, PreparedOcrRequest}; use crate::integrations::custom_guardrail::CustomGuardrailRunner; use crate::integrations::custom_logger::CustomLoggerRunner; +use litellm_core::call_lifecycle::CallLifecycleContext; +use litellm_core::logging::CallLogger; pub(crate) struct PreparedOcrCall { pub(crate) request: PreparedOcrRequest, @@ -26,6 +29,12 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { let model = provider_info.model.to_string(); let custom_llm_provider = provider_info.custom_llm_provider.to_string(); + let context = CallLifecycleContext::new( + "ocr", + request.model, + custom_llm_provider.clone(), + call_id.clone(), + ); PreparedOcrCall { request: PreparedOcrRequest { model, @@ -37,6 +46,7 @@ pub(crate) fn prepare_ocr_call(request: OcrRequest<'_>) -> PreparedOcrCall { extra_headers: request.extra_headers, optional_params: request.optional_params, timeout: request.timeout, + logger: Arc::new(CallLogger::new(&context, request.logging_sink)), }, hooks: OcrLifecycleHooks::new( CustomLoggerRunner::new(request.callbacks), diff --git a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs index bb2a6b06501..68f1a195310 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/tests.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/tests.rs @@ -7,7 +7,7 @@ use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; -use super::common_utils::{has_header, ocr_provider_config, string_headers, truncate_error_body}; +use super::common_utils::{has_header, ocr_provider_config, string_headers}; use super::{OcrRequest, ocr}; use crate::integrations::custom_guardrail::{ CustomGuardrail, GuardrailContext, GuardrailDecision, GuardrailError, GuardrailEventHook, @@ -17,6 +17,7 @@ use crate::integrations::custom_logger::{ CallbackTiming, CallbackValue, CustomLogger, LogFuture, ModelCallDetails, }; use crate::integrations::types::RequestMetadata; +use litellm_core::utils::truncate_error_body; async fn read_http_headers(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -326,6 +327,7 @@ async fn ocr_lifecycle_runs_pre_during_and_success_hooks() { ..Default::default() }, litellm_call_id: Some("ocr-call-1"), + logging_sink: None, }) .await .expect("ocr request succeeds"); @@ -391,6 +393,7 @@ async fn ocr_lifecycle_runs_failure_hook_on_provider_error() { guardrails: Vec::new(), request_metadata: RequestMetadata::default(), litellm_call_id: Some("ocr-call-2"), + logging_sink: None, }) .await .expect_err("provider error propagates"); @@ -435,6 +438,7 @@ async fn ocr_lifecycle_pre_call_block_skips_provider_socket() { guardrails: vec![guardrail.clone()], request_metadata: RequestMetadata::default(), litellm_call_id: Some("ocr-call-3"), + logging_sink: None, }) .await .expect_err("guardrail blocks request"); @@ -505,6 +509,7 @@ async fn ocr_does_not_duplicate_authorization_header_when_header_is_supplied() { guardrails: Vec::new(), request_metadata: RequestMetadata::default(), litellm_call_id: None, + logging_sink: None, }) .await .expect("ocr request succeeds"); @@ -574,6 +579,7 @@ async fn document_intelligence_poll_uses_resolved_subscription_key() { guardrails: Vec::new(), request_metadata: RequestMetadata::default(), litellm_call_id: None, + logging_sink: None, }) .await .expect("document intelligence request succeeds"); diff --git a/litellm-rust/crates/ai-gateway/src/ocr/types.rs b/litellm-rust/crates/ai-gateway/src/ocr/types.rs index bde734a4dd1..e70804c1cc3 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/types.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/types.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use litellm_core::call_lifecycle::{CallLifecycleContext, CallLifecycleRequest}; +use litellm_core::logging::{CallLogger, LogSink}; use litellm_core::ocr::transformation::OcrProviderConfig; use serde_json::{Map, Value}; @@ -22,6 +23,7 @@ pub struct OcrRequest<'a> { pub guardrails: Vec>, pub request_metadata: RequestMetadata, pub litellm_call_id: Option<&'a str>, + pub logging_sink: Option>, } pub(crate) struct PreparedOcrRequest { @@ -34,6 +36,7 @@ pub(crate) struct PreparedOcrRequest { pub(crate) extra_headers: Option>, pub(crate) optional_params: Map, pub(crate) timeout: Option, + pub(crate) logger: Arc, } impl CallLifecycleRequest for PreparedOcrRequest { @@ -54,4 +57,5 @@ pub(crate) struct ProviderOcrRequest { pub(crate) body: Value, pub(crate) upstream_headers: Vec<(String, String)>, pub(crate) timeout: Option, + pub(crate) logger: Arc, } diff --git a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs index c32e727de54..edd9338b4f2 100644 --- a/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs +++ b/litellm-rust/crates/ai-gateway/src/realtime/streaming.rs @@ -106,16 +106,16 @@ impl RealTimeStreaming { /// `litellm_call_id`, replacing the gateway-generated fallback. fn on_session(&mut self, event: &RealtimeEvent) { let session = event.data.get("session").and_then(Value::as_object); - if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) { - if !id.is_empty() { - self.id = id.to_string(); - self.litellm_call_id = id.to_string(); - } + if let Some(id) = session.and_then(|s| s.get("id")).and_then(Value::as_str) + && !id.is_empty() + { + self.id = id.to_string(); + self.litellm_call_id = id.to_string(); } - if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) { - if !model.is_empty() { - self.model = model.to_string(); - } + if let Some(model) = session.and_then(|s| s.get("model")).and_then(Value::as_str) + && !model.is_empty() + { + self.model = model.to_string(); } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index ea167761319..c197bd190e6 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -14,6 +14,7 @@ use bytes::Bytes; use futures_util::StreamExt; use futures_util::stream::{self, BoxStream}; use litellm_core::CoreError; +use litellm_core::logging::stream::count_forwarded_stream; use serde_json::{Map, Value}; use crate::auth::RequireMasterKey; @@ -34,21 +35,23 @@ async fn handle( Json(body): Json, ) -> Result { let extra_headers = forwarded_headers(&headers)?; - match service::run(&state.router, body, extra_headers) + match service::run(&state.router, body, extra_headers, state.logging_sink) .await .map_err(MessagesRouteError::from)? { service::MessagesResponse::Json(body) => Ok(Json(body).into_response()), - service::MessagesResponse::Stream { provider, response } => { - stream_response(provider, response) - } + service::MessagesResponse::Stream(response) => stream_response(response), } } fn stream_response( - provider: String, - upstream: reqwest::Response, + upstream: litellm_core::messages::types::MessagesStreamResponse, ) -> Result { + let litellm_core::messages::types::MessagesStreamResponse { + provider, + response: upstream, + logger, + } = upstream; let is_bedrock = provider == BEDROCK_MESSAGES_PROVIDER; let content_type = if is_bedrock { HeaderValue::from_static("text/event-stream") @@ -79,13 +82,12 @@ fn stream_response( .map(|result| result.map_err(|error| std::io::Error::other(error.to_string()))) .boxed() }; - response - .body(Body::from_stream(body_stream)) - .map_err(|error| { - MessagesRouteError(CoreError::InvalidResponse(format!( - "failed to build streaming response: {error}" - ))) - }) + let content = count_forwarded_stream(body_stream, logger); + response.body(Body::from_stream(content)).map_err(|error| { + MessagesRouteError(CoreError::InvalidResponse(format!( + "failed to build streaming response: {error}" + ))) + }) } #[allow(dead_code)] @@ -240,6 +242,7 @@ mod tests { master_key: master_key.map(Arc::from), loggers: Arc::new(Vec::new()), realtime_pool: RealtimePool::disabled(), + logging_sink: None, } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs index 5f4c5fe8de4..351007bdfd8 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/service.rs @@ -1,3 +1,4 @@ +use litellm_core::logging::LogSink; use std::sync::Arc; use litellm_core::constants::ANTHROPIC_MESSAGES_PROVIDER; @@ -9,13 +10,14 @@ use serde_json::{Map, Value}; pub(crate) enum MessagesResponse { Json(Value), - Stream(reqwest::Response), + Stream(litellm_core::messages::types::MessagesStreamResponse), } pub async fn run( router: &Arc, body: Value, extra_headers: Option>, + logging_sink: Option>, ) -> CoreResult { let model = body .get("model") @@ -51,6 +53,8 @@ pub async fn run( custom_llm_provider, extra_headers, timeout: None, + litellm_call_id: None, + logging_sink, }; if request.body.get("stream").and_then(Value::as_bool) == Some(true) { return messages_stream(request).await.map(MessagesResponse::Stream); diff --git a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs index 4ae8cfe7379..864d830c580 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs @@ -51,18 +51,17 @@ where provider_model, params.api_key.as_deref(), params.api_base.as_deref(), - ) { - if let Some(handoff) = pool.take(&key) { - return crate::io::realtime::realtime_warm( - provider_model, - handoff, - idle_timeout, - observe, - client_in, - client_out, - ) - .await; - } + ) && let Some(handoff) = pool.take(&key) + { + return crate::io::realtime::realtime_warm( + provider_model, + handoff, + idle_timeout, + observe, + client_in, + client_out, + ) + .await; } // Cold path: fresh dial (the original behavior). diff --git a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs index a94853e106d..770bfb8e65d 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs @@ -318,6 +318,7 @@ mod tests { master_key: Some(Arc::from("master-key")), loggers: Arc::new(Vec::new()), realtime_pool: RealtimePool::disabled(), + logging_sink: None, } } diff --git a/litellm-rust/crates/ai-gateway/src/state.rs b/litellm-rust/crates/ai-gateway/src/state.rs index 3b61d8309ea..0ba1300b3bb 100644 --- a/litellm-rust/crates/ai-gateway/src/state.rs +++ b/litellm-rust/crates/ai-gateway/src/state.rs @@ -4,6 +4,7 @@ use crate::io::realtime_pool::RealtimePool; use litellm_core::router::Router; use crate::integrations::custom_logger::CustomLogger; +use litellm_core::logging::LogSink; /// Shared application state handed to every route handler. #[derive(Clone)] @@ -18,4 +19,5 @@ pub struct AppState { /// (`RealtimePool::disabled()`) when `REALTIME_POOL_SIZE=0`, in which case /// every realtime connect fresh-dials exactly as before. pub realtime_pool: Arc, + pub logging_sink: Option>, } diff --git a/litellm-rust/crates/core/CLAUDE.md b/litellm-rust/crates/core/CLAUDE.md index 5d36305ded5..4dc1349c7ed 100644 --- a/litellm-rust/crates/core/CLAUDE.md +++ b/litellm-rust/crates/core/CLAUDE.md @@ -26,7 +26,9 @@ Not allowed: - Filesystem, database, or cache access. - Config file reading or rollout state; the host resolves those and passes them in. Env reads are limited to credential fallback in a route's `prepare.rs`. -- Logging callbacks, tracing spans, spend writes, or customer callbacks. +- Logging callbacks, tracing spans, spend writes, or customer callbacks, except + the `logging` module's debug sink, which owns the renderer env vars and stderr + output. - Provider-specific branching that belongs in `providers`. - Panics for user/provider-controlled input. diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index ab8050734f2..51700143509 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -8,8 +8,12 @@ repository.workspace = true [dependencies] rand.workspace = true reqwest.workspace = true +bytes.workspace = true +futures-util.workspace = true +url.workspace = true serde.workspace = true serde_json.workspace = true +colored_json.workspace = true thiserror.workspace = true sha2.workspace = true aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 51ea19750ea..e80ea45fe1e 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -3,6 +3,7 @@ pub mod caching; pub mod call_lifecycle; pub mod constants; pub mod error; +pub mod logging; pub mod messages; pub mod ocr; pub mod providers; @@ -10,5 +11,6 @@ pub mod realtime; pub mod responses; pub mod router; pub mod routing_utils; +pub mod utils; pub use error::{CoreError, CoreResult}; diff --git a/litellm-rust/crates/core/src/logging/console.rs b/litellm-rust/crates/core/src/logging/console.rs new file mode 100644 index 00000000000..1046a1ae31a --- /dev/null +++ b/litellm-rust/crates/core/src/logging/console.rs @@ -0,0 +1,222 @@ +use std::io::{IsTerminal, Write}; +use std::sync::OnceLock; +use std::sync::{Arc, Mutex}; + +use colored_json::{ColorMode, ColoredFormatter, Output, PrettyFormatter}; + +use super::{LogEvent, LogSink}; + +#[derive(Clone, Copy)] +enum RenderMode { + Compact, + Pretty, +} + +pub struct ConsoleDebugHook { + mode: RenderMode, + color_mode: ColorMode, + output: Mutex>, +} + +impl ConsoleDebugHook { + pub fn from_env() -> Self { + Self::with_writer(Box::new(std::io::stderr())) + } + + pub fn with_writer(writer: Box) -> Self { + Self::with_writer_and_mode(writer, matches!(*render_mode(), RenderMode::Pretty)) + } + + pub fn with_writer_and_mode(writer: Box, pretty: bool) -> Self { + Self { + mode: if pretty { + RenderMode::Pretty + } else { + RenderMode::Compact + }, + color_mode: ColorMode::Auto(Output::StdErr).eval(), + output: Mutex::new(writer), + } + } +} + +pub fn hook(enabled: bool) -> Option> { + enabled.then(|| Arc::new(ConsoleDebugHook::from_env()) as Arc) +} + +fn render_mode() -> &'static RenderMode { + static MODE: OnceLock = OnceLock::new(); + MODE.get_or_init(|| { + if std::env::var("JSON_LOGS") + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(false) + || !std::io::stderr().is_terminal() + { + RenderMode::Compact + } else { + RenderMode::Pretty + } + }) +} + +fn header(event: &LogEvent) -> String { + match event { + LogEvent::Request(value) => { + format!("provider.request {} {}", value.call_id, value.provider) + } + LogEvent::Response(value) => format!( + "provider.response {} {} status={} duration_ms={}", + value.call_id, value.provider, value.status, value.duration_ms + ), + LogEvent::StreamStarted(value) => format!( + "provider.stream.started {} {} status={}", + value.call_id, value.provider, value.status + ), + LogEvent::StreamCompleted(value) => format!( + "provider.stream.completed {} {} duration_ms={}", + value.call_id, value.provider, value.duration_ms + ), + LogEvent::Error(value) => format!( + "provider.error {} {}{} duration_ms={}", + value.call_id, + value.provider, + value + .status + .map_or(String::new(), |status| format!(" status={status}")), + value.duration_ms + ), + } +} + +fn decorate(value: &str, color_mode: ColorMode, code: &str) -> String { + if color_mode == ColorMode::On { + format!("\x1b[{code}m{value}\x1b[0m") + } else { + value.to_string() + } +} + +impl LogSink for ConsoleDebugHook { + fn emit(&self, event: &LogEvent) { + let Ok(mut output) = self.output.lock() else { + return; + }; + let Ok(json) = serde_json::to_string(event) else { + return; + }; + match self.mode { + RenderMode::Compact => { + let _ = writeln!(output, "{json}"); + } + RenderMode::Pretty => { + let pretty = serde_json::to_string_pretty(event).unwrap_or(json); + let _ = writeln!( + output, + "{}", + decorate(&header(event), self.color_mode, "36") + ); + let rendered = if self.color_mode == ColorMode::Off { + pretty + } else { + ColoredFormatter::new(PrettyFormatter::new()) + .to_colored_json(event, self.color_mode) + .unwrap_or(pretty) + }; + let _ = writeln!(output, "{rendered}"); + let _ = writeln!( + output, + "{}", + decorate( + "────────────────────────────────────────", + self.color_mode, + "2" + ) + ); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use serde_json::json; + + use crate::logging::ProviderRequestEvent; + + use super::*; + struct Buffer(Arc>>); + + impl Write for Buffer { + fn write(&mut self, bytes: &[u8]) -> std::io::Result { + self.0.lock().expect("buffer lock").extend_from_slice(bytes); + Ok(bytes.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } + } + + #[test] + fn compact_output_is_canonical_json() { + let buffer = Arc::new(Mutex::new(Vec::new())); + let hook = ConsoleDebugHook::with_writer_and_mode(Box::new(Buffer(buffer.clone())), false); + let event = LogEvent::Request(ProviderRequestEvent { + source: "litellm-rust", + call_id: "call_01".to_string(), + provider: "anthropic".to_string(), + model: "claude".to_string(), + stream: false, + method: "POST", + url: "https://example.test".to_string(), + headers: Default::default(), + body: json!({"prompt": "visible"}), + body_truncated: None, + body_original_bytes: None, + }); + let expected = serde_json::to_value(&event).expect("event serializes"); + hook.emit(&event); + let output = + String::from_utf8(buffer.lock().expect("buffer lock").clone()).expect("output is utf8"); + assert_eq!( + serde_json::from_str::(output.trim()).expect("output is JSON"), + expected + ); + } + + #[test] + fn pretty_output_has_header_separator_and_indented_payload() { + let buffer = Arc::new(Mutex::new(Vec::new())); + let hook = ConsoleDebugHook::with_writer_and_mode(Box::new(Buffer(buffer.clone())), true); + let event = LogEvent::Request(ProviderRequestEvent { + source: "litellm-rust", + call_id: "call_01".to_string(), + provider: "anthropic".to_string(), + model: "claude".to_string(), + stream: false, + method: "POST", + url: "https://example.test".to_string(), + headers: Default::default(), + body: json!({"prompt": "visible"}), + body_truncated: None, + body_original_bytes: None, + }); + hook.emit(&event); + let output = + String::from_utf8(buffer.lock().expect("buffer lock").clone()).expect("output is utf8"); + assert!(!output.contains('\x1b')); + let payload = &output + [output.find('{').expect("payload starts")..=output.rfind('}').expect("payload ends")]; + let expected = serde_json::to_string_pretty(&event).expect("event pretty serializes"); + assert_eq!(payload, expected); + assert!( + payload.find("\"event\"").expect("event key") + < payload.find("\"body\"").expect("body key") + ); + assert!(output.contains("provider.request call_01 anthropic")); + assert!(output.contains("────────────────")); + assert!(output.contains("\n \"event\"")); + } +} diff --git a/litellm-rust/crates/core/src/logging/events.rs b/litellm-rust/crates/core/src/logging/events.rs new file mode 100644 index 00000000000..88cc7d8d645 --- /dev/null +++ b/litellm-rust/crates/core/src/logging/events.rs @@ -0,0 +1,260 @@ +use serde::Serialize; +use serde_json::Value; +use std::collections::BTreeMap; + +pub use super::redaction::BodySnapshot; +use super::redaction::{redact_headers, redact_url, snapshot_json}; + +#[derive(Clone, Debug, Serialize)] +#[serde(tag = "event")] +pub enum LogEvent { + #[serde(rename = "provider.request")] + Request(ProviderRequestEvent), + #[serde(rename = "provider.response")] + Response(ProviderResponseEvent), + #[serde(rename = "provider.stream.started")] + StreamStarted(ProviderStreamStartedEvent), + #[serde(rename = "provider.stream.completed")] + StreamCompleted(ProviderStreamCompletedEvent), + #[serde(rename = "provider.error")] + Error(ProviderErrorEvent), +} + +pub struct RequestEventInput { + pub model: String, + pub stream: bool, + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Value, +} + +pub struct ResponseEventInput { + pub call_id: String, + pub provider: String, + pub status: u16, + pub duration_ms: u128, + pub headers: Vec<(String, String)>, + pub body: ResponseBody, +} + +pub struct ErrorEventInput { + pub call_id: String, + pub provider: String, + pub duration_ms: u128, + pub status: Option, + pub kind: &'static str, + pub message: String, + pub body: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ProviderRequestEvent { + pub source: &'static str, + pub call_id: String, + pub provider: String, + pub model: String, + pub stream: bool, + pub method: &'static str, + pub url: String, + pub headers: BTreeMap, + pub body: Value, + #[serde(skip_serializing_if = "Option::is_none")] + pub body_truncated: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub body_original_bytes: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ProviderResponseEvent { + pub source: &'static str, + pub call_id: String, + pub provider: String, + pub status: u16, + pub duration_ms: u128, + pub headers: BTreeMap, + pub body: Value, + #[serde(skip_serializing_if = "Option::is_none")] + pub body_truncated: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub body_original_bytes: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ProviderStreamStartedEvent { + pub source: &'static str, + pub call_id: String, + pub provider: String, + pub status: u16, + pub content_type: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ProviderStreamCompletedEvent { + pub source: &'static str, + pub call_id: String, + pub provider: String, + pub duration_ms: u128, + pub bytes_received: usize, + pub frames_received: usize, + pub events_decoded: usize, +} + +#[derive(Clone, Debug, Serialize)] +pub struct ProviderErrorEvent { + pub source: &'static str, + pub call_id: String, + pub provider: String, + pub duration_ms: u128, + pub status: Option, + pub kind: &'static str, + pub message: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub body: Option, +} + +pub(crate) fn request_event( + call_id: String, + provider: String, + input: RequestEventInput, +) -> LogEvent { + let snapshot = snapshot_json(input.body); + LogEvent::Request(ProviderRequestEvent { + source: "litellm-rust", + call_id, + provider, + model: input.model, + stream: input.stream, + method: "POST", + url: redact_url(&input.url), + headers: redact_headers(&input.headers), + body: snapshot.body, + body_truncated: snapshot.body_truncated, + body_original_bytes: snapshot.body_original_bytes, + }) +} + +pub(crate) fn response_event(input: ResponseEventInput) -> LogEvent { + let snapshot = input.body.snapshot(); + LogEvent::Response(ProviderResponseEvent { + source: "litellm-rust", + call_id: input.call_id, + provider: input.provider, + status: input.status, + duration_ms: input.duration_ms, + headers: redact_headers(&input.headers), + body: snapshot.body, + body_truncated: snapshot.body_truncated, + body_original_bytes: snapshot.body_original_bytes, + }) +} + +pub(crate) fn error_event(input: ErrorEventInput) -> LogEvent { + LogEvent::Error(ProviderErrorEvent { + source: "litellm-rust", + call_id: input.call_id, + provider: input.provider, + duration_ms: input.duration_ms, + status: input.status, + kind: input.kind, + message: input.message, + body: input.body.map(|body| body.snapshot().body), + }) +} + +pub(crate) fn stream_started( + call_id: String, + provider: String, + status: u16, + content_type: Option, +) -> LogEvent { + LogEvent::StreamStarted(ProviderStreamStartedEvent { + source: "litellm-rust", + call_id, + provider, + status, + content_type, + }) +} + +pub(crate) fn stream_completed( + call_id: String, + provider: String, + duration_ms: u128, + bytes_received: usize, + frames_received: usize, + events_decoded: usize, +) -> LogEvent { + LogEvent::StreamCompleted(ProviderStreamCompletedEvent { + source: "litellm-rust", + call_id, + provider, + duration_ms, + bytes_received, + frames_received, + events_decoded, + }) +} + +#[derive(Clone, Debug)] +pub enum ResponseBody { + Json(Value), + Binary { + media_type: Option, + bytes: usize, + }, +} + +impl ResponseBody { + fn snapshot(self) -> BodySnapshot { + match self { + Self::Json(value) => snapshot_json(value), + Self::Binary { media_type, bytes } => BodySnapshot { + body: serde_json::json!({"media_type": media_type, "bytes": bytes}), + body_truncated: None, + body_original_bytes: None, + }, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn redacts_credentials_recursively() { + let event = request_event( + "call_01".to_string(), + "anthropic".to_string(), + RequestEventInput { + model: "claude".to_string(), + stream: false, + url: "https://example.test?signature=secret&x=ok".to_string(), + headers: vec![("Authorization".to_string(), "Bearer secret".to_string())], + body: serde_json::json!({"nested": {"token": "secret"}, "prompt": "visible"}), + }, + ); + let json = serde_json::to_string(&event).expect("serializes"); + assert!(!json.contains("secret")); + assert!(json.contains("visible")); + assert!(json.contains("[REDACTED]")); + } + + #[test] + fn redact_url_preserves_queryless_encoded_paths() { + assert_eq!( + redact_url("https://example.test/v1%3A0/invoke"), + "https://example.test/v1%3A0/invoke" + ); + } + + #[test] + fn redact_url_preserves_non_secret_query_params() { + let redacted = redact_url( + "https://example.test/invoke?X-Amz-Signature=sig&X-Amz-Credential=cred&foo=bar", + ); + assert!(redacted.contains("X-Amz-Signature=%5BREDACTED%5D")); + assert!(redacted.contains("X-Amz-Credential=%5BREDACTED%5D")); + assert!(redacted.contains("foo=bar")); + } +} diff --git a/litellm-rust/crates/core/src/logging/http.rs b/litellm-rust/crates/core/src/logging/http.rs new file mode 100644 index 00000000000..028be43a1a9 --- /dev/null +++ b/litellm-rust/crates/core/src/logging/http.rs @@ -0,0 +1,198 @@ +use std::sync::Arc; +use std::time::Duration; + +use serde::de::DeserializeOwned; +use serde_json::Value; + +use crate::CoreResult; +use crate::error::CoreError; + +use super::{CallLogger, ResponseBody}; + +pub struct JsonRequest { + pub logger: Arc, + pub model: String, + pub stream: bool, + pub url: String, + pub headers: Vec<(String, String)>, + pub body: Value, + pub timeout: Option, +} + +pub async fn execute_json( + client: &reqwest::Client, + request: JsonRequest, +) -> CoreResult { + let body_bytes = serde_json::to_vec(&request.body) + .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + request + .logger + .request_about_to_be_sent(super::RequestEventInput { + model: request.model, + stream: request.stream, + url: request.url.clone(), + headers: request.headers.clone(), + body: request.body, + }); + let builder = request.headers.iter().fold( + client.post(&request.url).body(body_bytes), + |builder, (name, value)| builder.header(name, value), + ); + let builder = match request.timeout { + Some(timeout) => builder.timeout(timeout), + None => builder, + }; + let response = builder.send().await.map_err(|error| { + request + .logger + .failure(None, "network_error", error.to_string(), None); + CoreError::Network(error.to_string()) + })?; + let status = response.status(); + let headers = response + .headers() + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect::>(); + let media_type = headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) + .map(|(_, value)| value.clone()); + let text = response.text().await.map_err(|error| { + request.logger.failure( + Some(status.as_u16()), + "network_error", + error.to_string(), + None, + ); + CoreError::Network(error.to_string()) + })?; + if !status.is_success() { + let body = serde_json::from_str(&text) + .map(ResponseBody::Json) + .unwrap_or(ResponseBody::Binary { + media_type: media_type.clone(), + bytes: text.len(), + }); + request.logger.failure( + Some(status.as_u16()), + "http_error", + format!("provider returned HTTP {}", status.as_u16()), + Some(body), + ); + return Err(CoreError::Http { + status: status.as_u16(), + body: crate::utils::truncate_error_body(&text), + }); + } + let value = serde_json::from_str::(&text).map_err(|error| { + request.logger.failure( + Some(status.as_u16()), + "invalid_json", + error.to_string(), + Some(ResponseBody::Binary { + media_type: media_type.clone(), + bytes: text.len(), + }), + ); + CoreError::InvalidResponse(format!("invalid provider response JSON: {error}")) + })?; + let typed = T::deserialize(&value); + request + .logger + .response_received(status.as_u16(), headers, ResponseBody::Json(value)); + let typed = typed.map_err(|error| { + request.logger.failure( + Some(status.as_u16()), + "invalid_json", + error.to_string(), + None, + ); + CoreError::InvalidResponse(format!("invalid provider response: {error}")) + })?; + Ok(typed) +} + +pub async fn execute_stream( + client: &reqwest::Client, + request: JsonRequest, +) -> CoreResult { + let body_bytes = serde_json::to_vec(&request.body) + .map_err(|error| CoreError::InvalidRequest(error.to_string()))?; + request + .logger + .request_about_to_be_sent(super::RequestEventInput { + model: request.model, + stream: true, + url: request.url.clone(), + headers: request.headers.clone(), + body: request.body, + }); + let builder = request.headers.iter().fold( + client.post(&request.url).body(body_bytes), + |builder, (name, value)| builder.header(name, value), + ); + let builder = match request.timeout { + Some(timeout) => builder.timeout(timeout), + None => builder, + }; + let response = builder.send().await.map_err(|error| { + request + .logger + .failure(None, "network_error", error.to_string(), None); + CoreError::Network(error.to_string()) + })?; + let status = response.status(); + if !status.is_success() { + let headers = response + .headers() + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|value| (name.to_string(), value.to_string())) + }) + .collect::>(); + let text = response.text().await.map_err(|error| { + request.logger.failure( + Some(status.as_u16()), + "network_error", + error.to_string(), + None, + ); + CoreError::Network(error.to_string()) + })?; + let body = serde_json::from_str(&text) + .map(ResponseBody::Json) + .unwrap_or(ResponseBody::Binary { + media_type: headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("content-type")) + .map(|(_, value)| value.clone()), + bytes: text.len(), + }); + request.logger.failure( + Some(status.as_u16()), + "http_error", + format!("provider returned HTTP {}", status.as_u16()), + Some(body), + ); + return Err(CoreError::Http { + status: status.as_u16(), + body: crate::utils::truncate_error_body(&text), + }); + } + let content_type = response + .headers() + .get("content-type") + .and_then(|value| value.to_str().ok()) + .map(str::to_string); + request.logger.stream_started(status.as_u16(), content_type); + Ok(response) +} diff --git a/litellm-rust/crates/core/src/logging/mod.rs b/litellm-rust/crates/core/src/logging/mod.rs new file mode 100644 index 00000000000..ba6913c412a --- /dev/null +++ b/litellm-rust/crates/core/src/logging/mod.rs @@ -0,0 +1,159 @@ +//! Provider debug events are enabled by `litellm._turn_on_debug()` in Python +//! or `LITELLM_LOG=DEBUG` in the standalone gateway. `JSON_LOGS` selects compact +//! output and `NO_COLOR` disables terminal colors. Prompt and response content +//! remains visible and may contain sensitive application data. + +mod redaction; + +pub mod console; +pub mod events; +pub mod http; +pub mod stream; + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Instant; + +use crate::call_lifecycle::CallLifecycleContext; + +pub trait LogSink: Send + Sync { + fn emit(&self, event: &LogEvent); +} + +pub struct CallLogger { + context: CallLifecycleContext, + sink: Option>, + started: Instant, + bytes_received: AtomicUsize, + frames_received: AtomicUsize, + events_decoded: AtomicUsize, +} + +impl CallLogger { + pub fn new(context: &CallLifecycleContext, sink: Option>) -> Self { + Self { + context: context.clone(), + sink, + started: Instant::now(), + bytes_received: AtomicUsize::new(0), + frames_received: AtomicUsize::new(0), + events_decoded: AtomicUsize::new(0), + } + } + + pub fn request_about_to_be_sent(&self, input: events::RequestEventInput) { + self.emit(events::request_event( + self.context.litellm_call_id.clone(), + self.context.custom_llm_provider.clone(), + input, + )); + } + + pub fn response_received( + &self, + status: u16, + headers: Vec<(String, String)>, + body: ResponseBody, + ) { + self.emit(events::response_event(events::ResponseEventInput { + call_id: self.context.litellm_call_id.clone(), + provider: self.context.custom_llm_provider.clone(), + status, + duration_ms: self.started.elapsed().as_millis(), + headers, + body, + })); + } + + pub fn failure( + &self, + status: Option, + kind: &'static str, + message: String, + body: Option, + ) { + self.emit(events::error_event(events::ErrorEventInput { + call_id: self.context.litellm_call_id.clone(), + provider: self.context.custom_llm_provider.clone(), + duration_ms: self.started.elapsed().as_millis(), + status, + kind, + message, + body, + })); + } + + pub fn stream_started(&self, status: u16, content_type: Option) { + self.emit(events::stream_started( + self.context.litellm_call_id.clone(), + self.context.custom_llm_provider.clone(), + status, + content_type, + )); + } + + pub fn stream_finished(&self) { + self.emit(events::stream_completed( + self.context.litellm_call_id.clone(), + self.context.custom_llm_provider.clone(), + self.started.elapsed().as_millis(), + self.bytes_received.load(Ordering::Relaxed), + self.frames_received.load(Ordering::Relaxed), + self.events_decoded.load(Ordering::Relaxed), + )); + } + + pub fn stream_chunk_observed(&self, bytes: usize, events: usize) { + self.bytes_received.fetch_add(bytes, Ordering::Relaxed); + self.frames_received.fetch_add(1, Ordering::Relaxed); + self.events_decoded.fetch_add(events, Ordering::Relaxed); + } + + fn emit(&self, event: LogEvent) { + if let Some(sink) = &self.sink { + sink.emit(&event); + } + } +} + +pub use events::{ + BodySnapshot, ErrorEventInput, LogEvent, ProviderErrorEvent, ProviderRequestEvent, + ProviderResponseEvent, ProviderStreamCompletedEvent, ProviderStreamStartedEvent, + RequestEventInput, ResponseBody, ResponseEventInput, +}; + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use super::*; + + #[derive(Clone, Default)] + struct RecordingSink(Arc>>); + + impl LogSink for RecordingSink { + fn emit(&self, event: &LogEvent) { + self.0.lock().expect("recording lock").push(event.clone()); + } + } + + #[test] + fn logger_uses_lifecycle_context_and_redacts_events() { + let sink = RecordingSink::default(); + let context = CallLifecycleContext::new("messages", "claude", "anthropic", "req_123"); + let logger = CallLogger::new(&context, Some(Arc::new(sink.clone()))); + logger.request_about_to_be_sent(events::RequestEventInput { + model: "claude".to_string(), + stream: false, + url: "https://example.test/v1/messages".to_string(), + headers: vec![("authorization".to_string(), "Bearer secret".to_string())], + body: serde_json::json!({"token": "secret", "prompt": "visible"}), + }); + + let events = sink.0.lock().expect("recording lock"); + let serialized = serde_json::to_string(&events[0]).expect("event serializes"); + assert!(serialized.contains("\"call_id\":\"req_123\"")); + assert!(!serialized.contains("secret")); + assert!(serialized.contains("visible")); + } +} diff --git a/litellm-rust/crates/core/src/logging/redaction.rs b/litellm-rust/crates/core/src/logging/redaction.rs new file mode 100644 index 00000000000..8019bb6f56c --- /dev/null +++ b/litellm-rust/crates/core/src/logging/redaction.rs @@ -0,0 +1,204 @@ +use std::collections::BTreeMap; + +use serde_json::{Map, Value}; + +pub const PROVIDER_DEBUG_BODY_MAX_BYTES: usize = 64 * 1024; + +pub struct BodySnapshot { + pub body: Value, + pub body_truncated: Option, + pub body_original_bytes: Option, +} + +pub fn snapshot_json(value: Value) -> BodySnapshot { + let redacted = redact_value(value); + let serialized = serde_json::to_vec(&redacted).unwrap_or_default(); + if serialized.len() <= PROVIDER_DEBUG_BODY_MAX_BYTES { + return BodySnapshot { + body: redacted, + body_truncated: None, + body_original_bytes: None, + }; + } + BodySnapshot { + body: Value::String( + String::from_utf8_lossy(&serialized[..PROVIDER_DEBUG_BODY_MAX_BYTES]).into_owned(), + ), + body_truncated: Some(true), + body_original_bytes: Some(serialized.len()), + } +} + +pub fn redact_headers(headers: &[(String, String)]) -> BTreeMap { + headers + .iter() + .map(|(name, value)| { + let value = if is_credential_name(name) { + "[REDACTED]".to_string() + } else { + value.clone() + }; + (name.clone(), value) + }) + .collect() +} + +pub fn redact_url(url: &str) -> String { + let Ok(mut parsed) = url::Url::parse(url) else { + return url.to_string(); + }; + if !parsed.username().is_empty() { + let _ = parsed.set_username("[REDACTED]"); + } + if parsed.password().is_some() { + let _ = parsed.set_password(Some("[REDACTED]")); + } + let Some(_) = parsed.query() else { + return parsed.to_string(); + }; + let pairs = parsed + .query_pairs() + .map(|(key, value)| { + let value = if is_credential_name(&key) { + "[REDACTED]" + } else { + value.as_ref() + }; + (key.into_owned(), value.to_string()) + }) + .collect::>(); + parsed.query_pairs_mut().clear().extend_pairs(pairs); + parsed.to_string() +} + +fn redact_value(value: Value) -> Value { + match value { + Value::Object(map) => Value::Object( + map.into_iter() + .map(|(key, value)| { + if is_secret_key(&key) { + (key, Value::String("[REDACTED]".to_string())) + } else { + (key, redact_value(value)) + } + }) + .collect::>(), + ), + Value::Array(values) => Value::Array(values.into_iter().map(redact_value).collect()), + other => other, + } +} + +fn is_secret_key(key: &str) -> bool { + is_credential_name(key) +} + +fn is_credential_name(name: &str) -> bool { + let normalized = name + .chars() + .filter(|character| *character != '-' && *character != '_') + .flat_map(char::to_lowercase) + .collect::(); + matches!( + normalized.as_str(), + "authorization" + | "proxyauthorization" + | "xapikey" + | "apikey" + | "xamzsecuritytoken" + | "cookie" + | "setcookie" + | "xamzsignature" + | "xamzcredential" + | "key" + | "accesstoken" + | "signature" + | "secret" + | "password" + | "token" + | "clientsecret" + | "awssecretaccesskey" + | "awsaccesskeyid" + | "awssessiontoken" + ) || [ + "apikey", + "secret", + "token", + "password", + "credential", + "signature", + "authorization", + ] + .iter() + .any(|marker| normalized.contains(marker)) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::{PROVIDER_DEBUG_BODY_MAX_BYTES, redact_headers, redact_url, snapshot_json}; + + #[test] + fn redacts_explicit_sensitive_headers() { + let headers = redact_headers(&[ + ("authorization".to_string(), "Bearer secret".to_string()), + ("cookie".to_string(), "session-secret".to_string()), + ("content-type".to_string(), "application/json".to_string()), + ]); + assert_eq!(headers["authorization"], "[REDACTED]"); + assert_eq!(headers["cookie"], "[REDACTED]"); + assert_eq!(headers["content-type"], "application/json"); + } + + #[test] + fn redacts_credential_marker_headers_and_preserves_ordinary_headers() { + let headers = redact_headers(&[ + ("x-goog-api-key".to_string(), "google-secret".to_string()), + ("X_Custom_Token".to_string(), "custom-secret".to_string()), + ("content-type".to_string(), "application/json".to_string()), + ]); + assert_eq!(headers["x-goog-api-key"], "[REDACTED]"); + assert_eq!(headers["X_Custom_Token"], "[REDACTED]"); + assert_eq!(headers["content-type"], "application/json"); + } + + #[test] + fn redacts_credential_marker_query_parameters_and_userinfo() { + let url = redact_url( + "https://user:password@example.test/invoke?token=secret&client_secret=hidden&keep=value", + ); + assert!(url.contains("%5BREDACTED%5D:%5BREDACTED%5D@example.test")); + assert!(url.contains("token=%5BREDACTED%5D")); + assert!(url.contains("client_secret=%5BREDACTED%5D")); + assert!(url.contains("keep=value")); + } + + #[test] + fn redact_url_preserves_queryless_urls_without_trailing_question_mark() { + assert_eq!( + redact_url("https://example.test/v1%3A0/invoke"), + "https://example.test/v1%3A0/invoke" + ); + } + + #[test] + fn redacts_nested_body_keys() { + let snapshot = snapshot_json(json!({ + "outer": {"token": "secret", "visible": "keep"}, + "items": [{"password": "hidden"}] + })); + assert_eq!(snapshot.body["outer"]["token"], "[REDACTED]"); + assert_eq!(snapshot.body["outer"]["visible"], "keep"); + assert_eq!(snapshot.body["items"][0]["password"], "[REDACTED]"); + } + + #[test] + fn records_body_truncation_metadata() { + let snapshot = snapshot_json(json!({"content": "x".repeat(PROVIDER_DEBUG_BODY_MAX_BYTES)})); + assert_eq!(snapshot.body_truncated, Some(true)); + assert!( + snapshot.body_original_bytes.expect("original bytes") > PROVIDER_DEBUG_BODY_MAX_BYTES + ); + } +} diff --git a/litellm-rust/crates/core/src/logging/stream.rs b/litellm-rust/crates/core/src/logging/stream.rs new file mode 100644 index 00000000000..0ba09c86845 --- /dev/null +++ b/litellm-rust/crates/core/src/logging/stream.rs @@ -0,0 +1,115 @@ +use std::fmt::Display; +use std::sync::Arc; + +use bytes::Bytes; +use futures_util::Stream; +use futures_util::StreamExt; + +use super::CallLogger; + +pub fn count_forwarded_stream( + stream: S, + logger: Arc, +) -> impl Stream> +where + S: Stream> + Send + 'static, + E: Display + Send + 'static, +{ + futures_util::stream::unfold( + (Box::pin(stream), logger, false, false), + |(mut stream, logger, trailing_newline, failed)| async move { + match stream.next().await { + None => { + if !failed { + logger.stream_finished(); + } + None + } + Some(Ok(bytes)) => { + let events = bytes + .iter() + .enumerate() + .filter(|(index, byte)| { + **byte == b'\n' + && ((*index > 0 && bytes[*index - 1] == b'\n') + || (*index == 0 && trailing_newline)) + }) + .count(); + let trailing_newline = bytes.last().copied() == Some(b'\n'); + logger.stream_chunk_observed(bytes.len(), events); + Some((Ok(bytes), (stream, logger, trailing_newline, failed))) + } + Some(Err(error)) => { + logger.failure(None, "stream_error", error.to_string(), None); + Some((Err(error), (stream, logger, trailing_newline, true))) + } + } + }, + ) +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use bytes::Bytes; + use futures_util::StreamExt; + + use crate::call_lifecycle::CallLifecycleContext; + + use super::super::{CallLogger, LogEvent, LogSink}; + use super::count_forwarded_stream; + + #[derive(Clone, Default)] + struct Sink(Arc>>); + + impl LogSink for Sink { + fn emit(&self, event: &LogEvent) { + self.0.lock().expect("lock").push(event.clone()); + } + } + + fn logger(sink: Sink) -> Arc { + Arc::new(CallLogger::new( + &CallLifecycleContext::new("messages", "model", "anthropic", "call"), + Some(Arc::new(sink)), + )) + } + + #[tokio::test] + async fn counts_only_delimiters_and_handles_split_boundaries() { + let sink = Sink::default(); + let logger = logger(sink.clone()); + let stream = futures_util::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\n")), + Ok(Bytes::from_static(b"\ndata: b\nx\n")), + ]); + let chunks = count_forwarded_stream(stream, logger) + .collect::>() + .await; + assert_eq!(chunks.len(), 2); + let events = sink.0.lock().expect("lock"); + let LogEvent::StreamCompleted(completed) = events.last().expect("completion") else { + panic!("expected completion"); + }; + assert_eq!(completed.events_decoded, 1); + } + + #[tokio::test] + async fn newline_after_non_delimiter_does_not_count_from_carry() { + let sink = Sink::default(); + let logger = logger(sink.clone()); + let stream = futures_util::stream::iter([ + Ok::<_, std::io::Error>(Bytes::from_static(b"data: a\n")), + Ok(Bytes::from_static(b"x\n")), + ]); + let _ = count_forwarded_stream(stream, logger) + .collect::>() + .await; + let events = sink.0.lock().expect("lock"); + let LogEvent::StreamCompleted(completed) = events.last().expect("completion") else { + panic!("expected completion"); + }; + assert_eq!(completed.events_decoded, 0); + } +} diff --git a/litellm-rust/crates/core/src/messages/client.rs b/litellm-rust/crates/core/src/messages/client.rs index 6281270b964..2d2fac44aa7 100644 --- a/litellm-rust/crates/core/src/messages/client.rs +++ b/litellm-rust/crates/core/src/messages/client.rs @@ -3,7 +3,7 @@ use std::time::Duration; use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS}; -pub(super) fn http_client() -> &'static reqwest::Client { +pub(crate) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() diff --git a/litellm-rust/crates/core/src/messages/common_utils.rs b/litellm-rust/crates/core/src/messages/common_utils.rs index 9dcfcaa71e3..daa53bcb527 100644 --- a/litellm-rust/crates/core/src/messages/common_utils.rs +++ b/litellm-rust/crates/core/src/messages/common_utils.rs @@ -1,20 +1,11 @@ use serde_json::{Map, Value}; -use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; use crate::error::{CoreError, CoreResult, json_type_name}; use crate::providers::anthropic::messages::transformation::ANTHROPIC_MESSAGES_CONFIG; use crate::providers::azure_ai::messages::transformation::AZURE_ANTHROPIC_MESSAGES_CONFIG; use super::transformation::AnthropicMessagesProviderConfig; -pub(super) fn truncate_error_body(body: &str) -> String { - if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS { - return body.to_string(); - } - let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect(); - format!("{truncated}... (truncated)") -} - pub(super) fn messages_provider_config( provider: &str, ) -> Option<&'static dyn AnthropicMessagesProviderConfig> { diff --git a/litellm-rust/crates/core/src/messages/handler.rs b/litellm-rust/crates/core/src/messages/handler.rs index a7fc15f8876..be98da80df2 100644 --- a/litellm-rust/crates/core/src/messages/handler.rs +++ b/litellm-rust/crates/core/src/messages/handler.rs @@ -1,78 +1,49 @@ -use crate::constants::ANTHROPIC_MESSAGES_PROVIDER; -use crate::error::{CoreError, CoreResult}; +use crate::error::CoreResult; +use crate::logging::http::{JsonRequest, execute_json, execute_stream}; use super::client::http_client; -use super::common_utils::truncate_error_body; use super::types::{AnthropicMessagesResponse, ProviderMessagesRequest}; pub(super) async fn execute_messages_provider_call( request: ProviderMessagesRequest, ) -> CoreResult { - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - println!("IN EXECUTE MESSAGES PROVIDER CALL"); - - let response = request_builder - .send() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - let status = response.status(); - let text = response - .text() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - - if !status.is_success() { - return Err(CoreError::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - - let response = serde_json::from_str(&text).map_err(|err| { - CoreError::InvalidResponse(format!("invalid messages response JSON: {err}")) - })?; + let response = execute_json::( + http_client(), + JsonRequest { + logger: request.logger, + model: request.model.clone(), + stream: false, + url: request.url, + headers: request.upstream_headers, + body: request.body, + timeout: request.timeout, + }, + ) + .await?; request.config.transform_response(&request.model, response) } pub(super) async fn execute_messages_provider_stream( request: ProviderMessagesRequest, -) -> CoreResult { - if request.provider != ANTHROPIC_MESSAGES_PROVIDER { - return Err(CoreError::InvalidRequest( - "streaming messages is not supported for this provider".to_string(), - )); - } - - let mut request_builder = http_client().post(&request.url).json(&request.body); - for (key, value) in &request.upstream_headers { - request_builder = request_builder.header(key, value); - } - if let Some(duration) = request.timeout { - request_builder = request_builder.timeout(duration); - } - - let response = request_builder - .send() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - let status = response.status(); - if !status.is_success() { - let text = response - .text() - .await - .map_err(|err| CoreError::Network(err.to_string()))?; - return Err(CoreError::Http { - status: status.as_u16(), - body: truncate_error_body(&text), - }); - } - Ok(response) +) -> CoreResult { + let provider = request.provider; + let logger = request.logger.clone(); + let response = execute_stream( + http_client(), + JsonRequest { + logger: logger.clone(), + model: request.model, + stream: true, + url: request.url, + headers: request.upstream_headers, + body: request.body, + timeout: request.timeout, + }, + ) + .await?; + Ok(super::types::MessagesStreamResponse { + provider, + response, + logger, + }) } diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index 873a0b41e18..8d14ffca0e8 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -8,7 +8,7 @@ //! can splice the event stream to its own caller. mod client; -mod common_utils; +pub(crate) mod common_utils; mod handler; mod prepare; pub mod transformation; @@ -18,14 +18,14 @@ use crate::error::CoreResult; use handler::{execute_messages_provider_call, execute_messages_provider_stream}; use prepare::prepare_messages_call; -use types::{AnthropicMessagesResponse, MessagesRequest}; +use types::{AnthropicMessagesResponse, MessagesRequest, MessagesStreamResponse}; pub async fn messages(request: MessagesRequest<'_>) -> CoreResult { println!("ENTERED RUST MESSAGES"); execute_messages_provider_call(prepare_messages_call(request)?).await } -pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult { +pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult { execute_messages_provider_stream(prepare_messages_call(request)?).await } diff --git a/litellm-rust/crates/core/src/messages/prepare.rs b/litellm-rust/crates/core/src/messages/prepare.rs index 94b5b1eaed7..2791a26db69 100644 --- a/litellm-rust/crates/core/src/messages/prepare.rs +++ b/litellm-rust/crates/core/src/messages/prepare.rs @@ -1,4 +1,8 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::call_lifecycle::CallLifecycleContext; use crate::error::{CoreError, CoreResult}; +use crate::logging::CallLogger; use crate::routing_utils::provider::{CustomLlmProvider, get_custom_llm_provider}; use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; @@ -50,6 +54,9 @@ pub(super) fn prepare_messages_call( headers.push((name.to_string(), value.to_string())); } } + if !has_header(&headers, "content-type") { + headers.push(("content-type".to_string(), "application/json".to_string())); + } let url = config.complete_url(request.api_base, &model, &env_lookup)?; let typed_request = serde_json::from_value(request.body).map_err(|err| { @@ -62,6 +69,17 @@ pub(super) fn prepare_messages_call( )) })?; + let call_id = request + .litellm_call_id + .map(str::to_owned) + .unwrap_or_else(|| { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_nanos()); + format!("messages-{nanos}") + }); + let context = + CallLifecycleContext::new("messages", model.clone(), provider.to_string(), call_id); Ok(ProviderMessagesRequest { provider: provider.to_string(), model, @@ -70,5 +88,6 @@ pub(super) fn prepare_messages_call( body, upstream_headers: headers, timeout: request.timeout, + logger: std::sync::Arc::new(CallLogger::new(&context, request.logging_sink)), }) } diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index 9fc1763683b..0fbffc73807 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -1,16 +1,18 @@ +use std::sync::{Arc, Mutex}; use std::time::Duration; +use crate::logging::{LogEvent, LogSink}; +use futures_util::StreamExt; use serde_json::{Map, Value, json}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use crate::error::CoreError; +use crate::utils::truncate_error_body; -use super::common_utils::{ - has_bearer_auth, has_header, messages_provider_config, string_headers, truncate_error_body, -}; -use super::messages; +use super::common_utils::{has_bearer_auth, has_header, messages_provider_config, string_headers}; use super::types::MessagesRequest; +use super::{messages, messages_stream}; async fn read_http_request(socket: &mut TcpStream) -> String { let mut request = Vec::new(); @@ -150,6 +152,8 @@ async fn messages_round_trip_builds_azure_request_and_passes_response_through() custom_llm_provider: Some("azure_ai"), extra_headers: None, timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect("messages request succeeds"); @@ -206,6 +210,8 @@ async fn messages_round_trip_builds_native_anthropic_request() { custom_llm_provider: Some("anthropic"), extra_headers: None, timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect("messages request succeeds"); @@ -259,6 +265,8 @@ async fn messages_does_not_duplicate_auth_when_x_api_key_supplied() { custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect("messages request succeeds"); @@ -313,6 +321,8 @@ async fn messages_forwards_entra_id_bearer_without_requiring_api_key() { custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect("entra id request succeeds without api key"); @@ -337,6 +347,8 @@ async fn messages_requires_auth_when_no_key_and_no_header() { custom_llm_provider: Some("azure_ai"), extra_headers: None, timeout: Some(Duration::from_millis(50)), + litellm_call_id: None, + logging_sink: None, }) .await .expect_err("missing auth errors"); @@ -375,6 +387,8 @@ async fn messages_ignores_malformed_authorization_and_uses_api_key() { custom_llm_provider: Some("azure_ai"), extra_headers: Some(headers), timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect("falls back to api key"); @@ -416,6 +430,8 @@ async fn messages_maps_provider_error_status_to_http_error() { custom_llm_provider: Some("azure_ai"), extra_headers: None, timeout: Some(Duration::from_secs(5)), + litellm_call_id: None, + logging_sink: None, }) .await .expect_err("provider error propagates"); @@ -433,9 +449,290 @@ async fn messages_rejects_unsupported_provider() { custom_llm_provider: Some("openai"), extra_headers: None, timeout: Some(Duration::from_millis(50)), + litellm_call_id: None, + logging_sink: None, }) .await .expect_err("unsupported provider errors"); assert!(matches!(err, CoreError::InvalidProvider(provider) if provider == "openai")); } + +#[derive(Clone, Default)] +struct RecordingSink(Arc>>); + +impl LogSink for RecordingSink { + fn emit(&self, event: &LogEvent) { + self.0.lock().expect("recording lock").push(event.clone()); + } +} + +#[tokio::test] +async fn messages_debug_logs_transformed_request_and_redacted_response() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let request = read_http_request(&mut socket).await; + let response = r#"{"id":"msg_1","type":"message","role":"assistant","content":[{"type":"text","text":"visible"}],"model":"claude","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1},"token":"response-secret"}"#; + socket + .write_all(write_response(response).as_bytes()) + .await + .expect("writes response"); + request + }); + let sink = RecordingSink::default(); + let response = messages(MessagesRequest { + model: "claude", + body: json!({ + "model": "claude", + "max_tokens": 8, + "messages": [{"role": "user", "content": "prompt visible", "token": "request-secret"}] + }), + api_key: Some("request-secret"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: Some( + json!({ + "Authorization": "Bearer bearer-secret", + "X-Amz-Security-Token": "session-secret" + }) + .as_object() + .expect("headers") + .clone(), + ), + timeout: Some(Duration::from_secs(5)), + litellm_call_id: Some("debug-success"), + logging_sink: Some(Arc::new(sink.clone())), + }) + .await + .expect("messages succeeds"); + assert_eq!(response.content[0]["text"], "visible"); + let request = server.await.expect("server task"); + let events = sink.0.lock().expect("recording lock"); + let serialized = serde_json::to_string(&*events).expect("events serialize"); + assert!(serialized.contains("prompt visible")); + assert!(!serialized.contains("request-secret")); + assert!(!serialized.contains("bearer-secret")); + assert!(!serialized.contains("session-secret")); + assert!(!serialized.contains("response-secret")); + let LogEvent::Request(request_event) = &events[0] else { + panic!("request event first"); + }; + assert!(request_event.url.ends_with("/v1/messages")); + assert_eq!(request_event.headers["content-type"], "application/json"); + let sent_body = request.split_once("\r\n\r\n").expect("request body").1; + assert!(sent_body.contains("prompt visible")); + let response_event = events + .iter() + .find_map(|event| match event { + LogEvent::Response(value) => Some(value), + _ => None, + }) + .expect("response event"); + assert_eq!(response_event.status, 200); + assert!(response_event.duration_ms < 10_000); +} + +#[tokio::test] +async fn messages_debug_wrong_shape_emits_response_then_failure() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + socket + .write_all(write_response(r#"{"id":"wrong-shape"}"#).as_bytes()) + .await + .expect("writes response"); + }); + let sink = RecordingSink::default(); + let err = messages(MessagesRequest { + model: "claude", + body: json!({"model": "claude", "max_tokens": 8, "messages": []}), + api_key: Some("request-secret"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + litellm_call_id: Some("wrong-shape"), + logging_sink: Some(Arc::new(sink.clone())), + }) + .await + .expect_err("wrong response shape errors"); + assert!(matches!(err, CoreError::InvalidResponse(_))); + + server.await.expect("server task"); + let events = sink.0.lock().expect("recording lock"); + assert!(matches!(events[0], LogEvent::Request(_))); + assert!(matches!(events[1], LogEvent::Response(_))); + assert!(matches!(events[2], LogEvent::Error(_))); + assert_eq!(events.len(), 3); +} + +#[tokio::test] +async fn messages_debug_http_failure_emits_one_error_and_none_emits_nothing() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + let body = r#"{"token":"provider-secret","error":"no"}"#; + let response = format!( + "HTTP/1.1 401 Unauthorized\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}", + body.len(), + body + ); + socket + .write_all(response.as_bytes()) + .await + .expect("writes response"); + }); + let sink = RecordingSink::default(); + let err = messages(MessagesRequest { + model: "claude", + body: json!({"model": "claude", "max_tokens": 8, "messages": []}), + api_key: Some("secret"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + litellm_call_id: Some("debug-error"), + logging_sink: Some(Arc::new(sink.clone())), + }) + .await + .expect_err("request fails"); + assert!(matches!(err, CoreError::Http { status: 401, .. })); + server.await.expect("server task"); + { + let events = sink.0.lock().expect("recording lock"); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, LogEvent::Error(_))) + .count(), + 1 + ); + assert!( + !events + .iter() + .any(|event| matches!(event, LogEvent::Response(_))) + ); + } + + let none_sink = messages(MessagesRequest { + model: "claude", + body: json!({"model": "claude", "max_tokens": 8, "messages": []}), + api_key: Some("secret"), + api_base: Some("http://127.0.0.1:1"), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_millis(20)), + litellm_call_id: Some("debug-disabled"), + logging_sink: None, + }) + .await + .expect_err("disabled test endpoint fails"); + assert!(matches!(none_sink, CoreError::Network(_))); +} + +#[tokio::test] +async fn messages_debug_stream_counts_clean_and_interrupted_streams() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + let body = b"data: one\n\ndata: two\n\n"; + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n", + body.len() + ); + socket + .write_all(response.as_bytes()) + .await + .expect("headers"); + socket.write_all(body).await.expect("body"); + }); + let sink = RecordingSink::default(); + let upstream = messages_stream(MessagesRequest { + model: "claude", + body: json!({"model": "claude", "max_tokens": 8, "stream": true, "messages": []}), + api_key: Some("secret"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + litellm_call_id: Some("debug-stream"), + logging_sink: Some(Arc::new(sink.clone())), + }) + .await + .expect("stream starts"); + let chunks = crate::logging::stream::count_forwarded_stream( + upstream.response.bytes_stream(), + upstream.logger, + ); + let forwarded = futures_util::StreamExt::collect::>(chunks).await; + assert!(!forwarded.is_empty()); + assert!(forwarded.iter().all(Result::is_ok)); + server.await.expect("server task"); + let events = sink.0.lock().expect("recording lock"); + let completed = events + .iter() + .find_map(|event| match event { + LogEvent::StreamCompleted(value) => Some(value), + _ => None, + }) + .expect("completion"); + assert!(completed.bytes_received > 0); + assert!(completed.frames_received > 0); + assert!(completed.events_decoded > 0); +} + +#[tokio::test] +async fn messages_debug_interrupted_stream_emits_error_without_completion() { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let _ = read_http_request(&mut socket).await; + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: 100\r\nconnection: close\r\n\r\nshort", + ) + .await + .expect("writes partial response"); + }); + let sink = RecordingSink::default(); + let upstream = messages_stream(MessagesRequest { + model: "claude", + body: json!({"model": "claude", "max_tokens": 8, "stream": true, "messages": []}), + api_key: Some("secret"), + api_base: Some(&format!("http://{addr}")), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(5)), + litellm_call_id: Some("debug-interrupted"), + logging_sink: Some(Arc::new(sink.clone())), + }) + .await + .expect("stream starts"); + let results = crate::logging::stream::count_forwarded_stream( + upstream.response.bytes_stream(), + upstream.logger, + ) + .collect::>() + .await; + assert!(results.iter().any(Result::is_err)); + server.await.expect("server task"); + let events = sink.0.lock().expect("recording lock"); + assert!(events.iter().any(|event| matches!( + event, + LogEvent::Error(value) if value.kind == "stream_error" + ))); + assert!( + !events + .iter() + .any(|event| matches!(event, LogEvent::StreamCompleted(_))) + ); +} diff --git a/litellm-rust/crates/core/src/messages/types.rs b/litellm-rust/crates/core/src/messages/types.rs index bd620f4b4c6..69eda1825f6 100644 --- a/litellm-rust/crates/core/src/messages/types.rs +++ b/litellm-rust/crates/core/src/messages/types.rs @@ -1,5 +1,7 @@ +use std::sync::Arc; use std::time::Duration; +use crate::logging::{CallLogger, LogSink}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; @@ -13,6 +15,8 @@ pub struct MessagesRequest<'a> { pub custom_llm_provider: Option<&'a str>, pub extra_headers: Option>, pub timeout: Option, + pub litellm_call_id: Option<&'a str>, + pub logging_sink: Option>, } pub(super) struct ProviderMessagesRequest { @@ -23,6 +27,13 @@ pub(super) struct ProviderMessagesRequest { pub(super) body: Value, pub(super) upstream_headers: Vec<(String, String)>, pub(super) timeout: Option, + pub(super) logger: Arc, +} + +pub struct MessagesStreamResponse { + pub provider: String, + pub response: reqwest::Response, + pub logger: Arc, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] diff --git a/litellm-rust/crates/core/src/utils.rs b/litellm-rust/crates/core/src/utils.rs new file mode 100644 index 00000000000..c5beac28823 --- /dev/null +++ b/litellm-rust/crates/core/src/utils.rs @@ -0,0 +1,9 @@ +use crate::constants::MESSAGES_ERROR_BODY_MAX_CHARS; + +pub fn truncate_error_body(body: &str) -> String { + if body.chars().count() <= MESSAGES_ERROR_BODY_MAX_CHARS { + return body.to_string(); + } + let truncated: String = body.chars().take(MESSAGES_ERROR_BODY_MAX_CHARS).collect(); + format!("{truncated}... (truncated)") +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index f0cc26a0cca..e1a7d1c402f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -7,6 +7,7 @@ use litellm_ai_gateway::io::audio_transcription::{ use litellm_ai_gateway::io::ocr::{OcrRequest, ocr as run_ocr}; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use litellm_core::error::CoreError; +use litellm_core::logging::console::hook; use litellm_core::messages::messages as run_messages; use litellm_core::messages::types::{AnthropicMessagesResponse, MessagesRequest}; use pyo3::exceptions::{PyRuntimeError, PyValueError}; @@ -204,6 +205,7 @@ fn ocr( guardrails: Vec::new(), request_metadata: Default::default(), litellm_call_id: None, + logging_sink: None, })) }); @@ -249,6 +251,7 @@ fn aocr( guardrails: Vec::new(), request_metadata: Default::default(), litellm_call_id: None, + logging_sink: None, }) .await .map_err(core_error_to_pyerr)?; @@ -364,7 +367,7 @@ fn marshal_messages_inputs( } #[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, debug=false))] #[allow(clippy::too_many_arguments)] fn messages( py: Python<'_>, @@ -375,6 +378,7 @@ fn messages( custom_llm_provider: Option, extra_headers: Option>, timeout_seconds: Option, + debug: bool, ) -> PyResult> { let (body, extra_headers, timeout) = marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; @@ -388,6 +392,8 @@ fn messages( custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, timeout, + litellm_call_id: None, + logging_sink: hook(debug), })) }); @@ -398,7 +404,7 @@ fn messages( } #[pyfunction] -#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None))] +#[pyo3(signature = (model, body, api_key=None, api_base=None, custom_llm_provider=None, extra_headers=None, timeout_seconds=None, debug=false))] #[allow(clippy::too_many_arguments)] fn amessages( py: Python<'_>, @@ -409,6 +415,7 @@ fn amessages( custom_llm_provider: Option, extra_headers: Option>, timeout_seconds: Option, + debug: bool, ) -> PyResult> { let (body, extra_headers, timeout) = marshal_messages_inputs(py, body, extra_headers, timeout_seconds)?; @@ -422,6 +429,8 @@ fn amessages( custom_llm_provider: custom_llm_provider.as_deref(), extra_headers, timeout, + litellm_call_id: None, + logging_sink: hook(debug), }) .await .map_err(core_error_to_pyerr)?; diff --git a/litellm/rust_bridge/messages.py b/litellm/rust_bridge/messages.py index 5abb21879d3..893dbb7d9bc 100644 --- a/litellm/rust_bridge/messages.py +++ b/litellm/rust_bridge/messages.py @@ -6,6 +6,8 @@ from dataclasses import dataclass from typing import Awaitable, Final, Protocol, Union, cast import httpx +import inspect +from litellm._logging import _is_debugging_on from litellm.rust_bridge.timeouts import timeout_to_seconds @@ -20,6 +22,7 @@ class RustMessages(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout_seconds: float | None, + debug: bool, ) -> dict[str, object]: raise NotImplementedError @@ -34,6 +37,7 @@ class RustAmessages(Protocol): custom_llm_provider: str | None, extra_headers: dict[str, object] | None, timeout_seconds: float | None, + debug: bool, ) -> Awaitable[dict[str, object]]: raise NotImplementedError @@ -100,6 +104,17 @@ def messages( rust_messages = load_rust_messages() if rust_messages is None: return None + if "debug" in inspect.signature(rust_messages).parameters: + return rust_messages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + debug=_is_debugging_on(), + ) return rust_messages( model=model, body=body, @@ -124,6 +139,17 @@ async def amessages( rust_amessages = load_rust_amessages() if rust_amessages is None: return None + if "debug" in inspect.signature(rust_amessages).parameters: + return await rust_amessages( + model=model, + body=body, + api_key=api_key, + api_base=api_base, + custom_llm_provider=custom_llm_provider, + extra_headers=extra_headers, + timeout_seconds=timeout_to_seconds(timeout), + debug=_is_debugging_on(), + ) return await rust_amessages( model=model, body=body, diff --git a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py index fbd7e36e298..31a30369595 100644 --- a/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py +++ b/tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py @@ -7,6 +7,7 @@ import httpx import pytest import litellm +from litellm import _logging from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, @@ -89,6 +90,44 @@ class RecordingAsyncMessages: return dict(FAKE_MESSAGES_RESPONSE) +class DebugRecordingMessages: + def __init__(self) -> None: + self.debug: bool | None = None + + def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + debug: bool, + ) -> dict[str, object]: + self.debug = debug + return dict(FAKE_MESSAGES_RESPONSE) + + +class DebugRecordingAsyncMessages: + def __init__(self) -> None: + self.debug: bool | None = None + + async def __call__( + self, + model: str, + body: dict[str, object], + api_key: str | None, + api_base: str | None, + custom_llm_provider: str | None, + extra_headers: dict[str, object] | None, + timeout_seconds: float | None, + debug: bool, + ) -> dict[str, object]: + self.debug = debug + return dict(FAKE_MESSAGES_RESPONSE) + + class ExplodingAsyncMessages: def __init__(self) -> None: self.calls = 0 @@ -195,6 +234,26 @@ def test_messages_wrapper_forwards_args_and_converts_timeout(): } +def test_messages_wrapper_forwards_debug_to_new_bridge(): + bridge = DebugRecordingMessages() + litellm.use_litellm_rust(True, messages=bridge) + litellm._turn_on_debug() + try: + response = rust_messages.messages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers=None, + timeout=12.5, + ) + finally: + _logging._disable_debugging() + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.debug is True + + @pytest.mark.asyncio async def test_amessages_wrapper_forwards_args(): bridge = RecordingAsyncMessages() @@ -215,6 +274,27 @@ async def test_amessages_wrapper_forwards_args(): assert bridge.calls[0]["timeout_seconds"] == 12.5 +@pytest.mark.asyncio +async def test_amessages_wrapper_forwards_debug_to_new_bridge(): + bridge = DebugRecordingAsyncMessages() + litellm.use_litellm_rust(True, amessages=bridge) + litellm._turn_on_debug() + try: + response = await rust_messages.amessages( + model="claude-sonnet-4-5", + body=REQUEST_BODY, + api_key="sk-azure", + api_base="https://resource.services.ai.azure.com/anthropic", + custom_llm_provider="azure_ai", + extra_headers=None, + timeout=12.5, + ) + finally: + _logging._disable_debugging() + assert response == FAKE_MESSAGES_RESPONSE + assert bridge.debug is True + + def _gate(**overrides): kwargs = { "custom_llm_provider": "azure_ai",