scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat(rust): merge provider debug logging
Some checks failed
LiteLLM Rust / rustfmt, clippy, test (push) Has been cancelled
Details

Browse source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Ishaan Jaff 2026-07-30 21:30:10 -07:00
commit d8df153f94
35 changed files with 1783 additions and 134 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
litellm-rust/Cargo.lock generated
View file

@ -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"

2
litellm-rust/Cargo.toml
View file

@ -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"

5
litellm-rust/crates/ai-gateway/src/main.rs
View file

@ -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());

13
litellm-rust/crates/ai-gateway/src/ocr/common_utils.rs
View file

@ -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| {

25
litellm-rust/crates/ai-gateway/src/ocr/handler.rs
View file

@ -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<Value> {
if request.config.response_handling() != OcrResponseHandling::AzureDocumentIntelligencePoll {
let response = execute_json::<Value>(
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),
});
}

1
litellm-rust/crates/ai-gateway/src/ocr/hooks.rs
View file

@ -111,6 +111,7 @@ impl OcrLifecycleHooks {
body,
upstream_headers,
timeout: request.timeout,
logger: request.logger,
})
}

10
litellm-rust/crates/ai-gateway/src/ocr/prepare.rs
View file

@ -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),

8
litellm-rust/crates/ai-gateway/src/ocr/tests.rs
View file

@ -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");

4
litellm-rust/crates/ai-gateway/src/ocr/types.rs
View file

@ -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<Arc<dyn CustomGuardrail>>,
pub request_metadata: RequestMetadata,
pub litellm_call_id: Option<&'a str>,
pub logging_sink: Option<Arc<dyn LogSink>>,
}
pub(crate) struct PreparedOcrRequest {
@ -34,6 +36,7 @@ pub(crate) struct PreparedOcrRequest {
pub(crate) extra_headers: Option<Map<String, Value>>,
pub(crate) optional_params: Map<String, Value>,
pub(crate) timeout: Option<Duration>,
pub(crate) logger: Arc<CallLogger>,
}
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<Duration>,
pub(crate) logger: Arc<CallLogger>,
}

18
litellm-rust/crates/ai-gateway/src/realtime/streaming.rs
View file

@ -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();
}
}

29
litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs
View file

@ -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<Value>,
) -> Result<Response, MessagesRouteError> {
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<Response, MessagesRouteError> {
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,
}
}

6
litellm-rust/crates/ai-gateway/src/routes/messages/service.rs
View file

@ -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<Router>,
body: Value,
extra_headers: Option<Map<String, Value>>,
logging_sink: Option<Arc<dyn LogSink>>,
) -> CoreResult<MessagesResponse> {
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);

23
litellm-rust/crates/ai-gateway/src/routes/realtime/service.rs
View file

@ -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).

1
litellm-rust/crates/ai-gateway/src/routes/responses/mod.rs
View file

@ -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,
}
}

2
litellm-rust/crates/ai-gateway/src/state.rs
View file

@ -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<RealtimePool>,
pub logging_sink: Option<Arc<dyn LogSink>>,
}

4
litellm-rust/crates/core/CLAUDE.md
View file

@ -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.

4
litellm-rust/crates/core/Cargo.toml
View file

@ -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 }

2
litellm-rust/crates/core/src/lib.rs
View file

@ -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};

222
litellm-rust/crates/core/src/logging/console.rs Normal file
View file

@ -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<Box<dyn Write + Send>>,
}
impl ConsoleDebugHook {
pub fn from_env() -> Self {
Self::with_writer(Box::new(std::io::stderr()))
}
pub fn with_writer(writer: Box<dyn Write + Send>) -> Self {
Self::with_writer_and_mode(writer, matches!(*render_mode(), RenderMode::Pretty))
}
pub fn with_writer_and_mode(writer: Box<dyn Write + Send>, 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<Arc<dyn LogSink>> {
enabled.then(|| Arc::new(ConsoleDebugHook::from_env()) as Arc<dyn LogSink>)
}
fn render_mode() -> &'static RenderMode {
static MODE: OnceLock<RenderMode> = 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<Mutex<Vec<u8>>>);
impl Write for Buffer {
fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
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::<serde_json::Value>(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\""));
}
}

260
litellm-rust/crates/core/src/logging/events.rs Normal file
View file

@ -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<u16>,
pub kind: &'static str,
pub message: String,
pub body: Option<ResponseBody>,
}
#[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<String, String>,
pub body: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_truncated: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_original_bytes: Option<usize>,
}
#[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<String, String>,
pub body: Value,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_truncated: Option<bool>,
#[serde(skip_serializing_if = "Option::is_none")]
pub body_original_bytes: Option<usize>,
}
#[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<String>,
}
#[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<u16>,
pub kind: &'static str,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub body: Option<Value>,
}
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<String>,
) -> 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<String>,
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"));
}
}

198
litellm-rust/crates/core/src/logging/http.rs Normal file
View file

@ -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<CallLogger>,
pub model: String,
pub stream: bool,
pub url: String,
pub headers: Vec<(String, String)>,
pub body: Value,
pub timeout: Option<Duration>,
}
pub async fn execute_json<T: DeserializeOwned>(
client: &reqwest::Client,
request: JsonRequest,
) -> CoreResult<T> {
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::<Vec<_>>();
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::<Value>(&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<reqwest::Response> {
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::<Vec<_>>();
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)
}

159
litellm-rust/crates/core/src/logging/mod.rs Normal file
View file

@ -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<Arc<dyn LogSink>>,
started: Instant,
bytes_received: AtomicUsize,
frames_received: AtomicUsize,
events_decoded: AtomicUsize,
}
impl CallLogger {
pub fn new(context: &CallLifecycleContext, sink: Option<Arc<dyn LogSink>>) -> 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<u16>,
kind: &'static str,
message: String,
body: Option<ResponseBody>,
) {
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<String>) {
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<Mutex<Vec<LogEvent>>>);
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"));
}
}

204
litellm-rust/crates/core/src/logging/redaction.rs Normal file
View file

@ -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<bool>,
pub body_original_bytes: Option<usize>,
}
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<String, String> {
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::<Vec<_>>();
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::<Map<_, _>>(),
),
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::<String>();
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
);
}
}

115
litellm-rust/crates/core/src/logging/stream.rs Normal file
View file

@ -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<S, E>(
stream: S,
logger: Arc<CallLogger>,
) -> impl Stream<Item = Result<Bytes, E>>
where
S: Stream<Item = Result<Bytes, E>> + 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<Mutex<Vec<LogEvent>>>);
impl LogSink for Sink {
fn emit(&self, event: &LogEvent) {
self.0.lock().expect("lock").push(event.clone());
}
}
fn logger(sink: Sink) -> Arc<CallLogger> {
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::<Vec<_>>()
.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::<Vec<_>>()
.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);
}
}

2
litellm-rust/crates/core/src/messages/client.rs
View file

@ -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<reqwest::Client> = OnceLock::new();
CLIENT.get_or_init(|| {
reqwest::Client::builder()

9
litellm-rust/crates/core/src/messages/common_utils.rs
View file

@ -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> {

101
litellm-rust/crates/core/src/messages/handler.rs
View file

@ -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<AnthropicMessagesResponse> {
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::<AnthropicMessagesResponse>(
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<reqwest::Response> {
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<super::types::MessagesStreamResponse> {
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,
})
}

6
litellm-rust/crates/core/src/messages/mod.rs
View file

@ -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<AnthropicMessagesResponse> {
println!("ENTERED RUST MESSAGES");
execute_messages_provider_call(prepare_messages_call(request)?).await
}
pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult<reqwest::Response> {
pub async fn messages_stream(request: MessagesRequest<'_>) -> CoreResult<MessagesStreamResponse> {
execute_messages_provider_stream(prepare_messages_call(request)?).await
}

19
litellm-rust/crates/core/src/messages/prepare.rs
View file

@ -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)),
})
}

305
litellm-rust/crates/core/src/messages/tests.rs
View file

@ -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<Mutex<Vec<LogEvent>>>);
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::<Vec<_>>(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::<Vec<_>>()
.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(_)))
);
}

11
litellm-rust/crates/core/src/messages/types.rs
View file

@ -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<Map<String, Value>>,
pub timeout: Option<Duration>,
pub litellm_call_id: Option<&'a str>,
pub logging_sink: Option<Arc<dyn LogSink>>,
}
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<Duration>,
pub(super) logger: Arc<CallLogger>,
}
pub struct MessagesStreamResponse {
pub provider: String,
pub response: reqwest::Response,
pub logger: Arc<CallLogger>,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]

9
litellm-rust/crates/core/src/utils.rs Normal file
View file

@ -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)")
}

13
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -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<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
debug: bool,
) -> PyResult<Py<PyAny>> {
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<String>,
extra_headers: Option<Py<PyAny>>,
timeout_seconds: Option<f64>,
debug: bool,
) -> PyResult<Bound<'_, PyAny>> {
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)?;

26
litellm/rust_bridge/messages.py
View file

@ -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,

80
tests/test_litellm/anthropic_interface/test_rust_bridge_messages.py
View file

@ -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",