mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
feat(rust): add structured route lifecycle tracing (#43466)
Co-authored-by: Yujong Lee <yujong@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
8ab124309f
commit
876539e1b3
33 changed files with 1338 additions and 140 deletions
22
litellm-rust/.agents/skills/rust-tracing/SKILL.md
Normal file
22
litellm-rust/.agents/skills/rust-tracing/SKILL.md
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
---
|
||||
name: rust-tracing
|
||||
description: Add or change Rust diagnostic tracing in litellm-rust, including route spans, subscriber layers, and Python logger delivery
|
||||
---
|
||||
|
||||
# Rust tracing
|
||||
|
||||
Use upstream `tracing` throughout Rust, including `#[tracing::instrument]`, events, and span propagation. Centralize collection and delivery infrastructure in `crates/tracing`. Direct upstream imports still reach our configured subscriber; re-exporting macros does not control delivery. Do not introduce Rust `log` or `pyo3-log` for this path
|
||||
|
||||
`litellm-tracing` owns shared subscriber layers, span field collection, and diagnostic processing. Keep adapters composable as `tracing_subscriber::Layer`s, with `Logger` providing host setup. Runtime-specific delivery belongs in the host bridge. The Python bridge delivers directly to the existing Python SDK logger, preserving its handlers, filtering, redaction, and request correlation. Keep Python dependencies out of `crates/tracing`
|
||||
|
||||
Hosts configure subscribers. Keep Python execution scoped to its captured dispatch rather than installing a process-wide subscriber. Propagate both span context and dispatch across spawned work and returned streams
|
||||
|
||||
In core, instrument execution shared by native calls and hosted machines. Use consistent route, model, provider, streaming, and outcome fields. Put status recording at shared provider boundaries instead of scattering basic logging through handlers. Keep upstream HTTP status separate from route success
|
||||
|
||||
Use `skip_all` and explicitly selected fields. Basic tracing excludes bodies, credentials, headers, and raw error strings. Avoid automatic `ret` or `err` capture of sensitive values. Keep payload diagnostics separate and subject to existing redaction
|
||||
|
||||
A returned stream retains its route span until exhaustion, error, or drop, with exactly one terminal outcome. Builder construction does not start a trace. Never hold a span entry guard across an await. Diagnostic tracing remains separate from lifecycle callbacks and `CustomLogger` dispatch
|
||||
|
||||
Use `litellm_tracing::sink_layer` to compose a sink with other subscriber layers. It inherits span fields into events and emits span-close summaries with elapsed time. Test observable records, concurrent isolation, dynamic filtering, sensitive-field exclusion, and stream cancellation when changing this behavior
|
||||
|
||||
Consult the [tracing API](https://docs.rs/tracing/latest/tracing/) and [subscriber layers](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/layer/index.html) for implementation details
|
||||
|
|
@ -1,5 +1,7 @@
|
|||
# Rust workspace rules
|
||||
|
||||
For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md)
|
||||
|
||||
## Test placement
|
||||
|
||||
- Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;`
|
||||
|
|
|
|||
3
litellm-rust/Cargo.lock
generated
3
litellm-rust/Cargo.lock
generated
|
|
@ -3395,6 +3395,7 @@ dependencies = [
|
|||
"time",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"url",
|
||||
"veil",
|
||||
"wiremock",
|
||||
|
|
@ -3596,6 +3597,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"veil",
|
||||
"webpki-roots",
|
||||
]
|
||||
|
|
@ -3698,6 +3700,7 @@ dependencies = [
|
|||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
"tracing",
|
||||
"url",
|
||||
"veil",
|
||||
"wiremock",
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ litellm-auth-aws.workspace = true
|
|||
litellm-http.workspace = true
|
||||
litellm-llms.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
tracing.workspace = true
|
||||
moka.workspace = true
|
||||
mime_guess = "2.0.5"
|
||||
rand.workspace = true
|
||||
|
|
|
|||
|
|
@ -31,9 +31,24 @@ impl AudioTranscriptionRoute {
|
|||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "audio_transcription",
|
||||
model = request.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub async fn execute(&self, request: AudioTranscriptionRequest<'_>) -> Result<Value, Error> {
|
||||
let request =
|
||||
prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?;
|
||||
execute_audio_transcription_provider_call(&self.http, &self.auth, request).await
|
||||
crate::diagnostic::unary(async {
|
||||
let request =
|
||||
prepare_audio_transcription_provider_call(request, self.secrets.as_ref()).await?;
|
||||
crate::diagnostic::provider(&request.model, &request.custom_llm_provider);
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<Value, Error>> = Box::pin(
|
||||
execute_audio_transcription_provider_call(&self.http, &self.auth, request),
|
||||
);
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ fn provider_config(provider: LlmProviders) -> Option<&'static dyn BaseAudioTrans
|
|||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub async fn prepare_audio_transcription_provider_call(
|
||||
request: AudioTranscriptionRequest<'_>,
|
||||
secrets: &dyn SecretSource,
|
||||
|
|
|
|||
|
|
@ -68,17 +68,33 @@ impl ChatCompletionsRoute {
|
|||
litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "chat_completions",
|
||||
model = %request.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
request: ChatCompletionsRequest<'_>,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let resolved = resolve_request(request)?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
.await?;
|
||||
let prepared = prepare_provider_request(resolved, snapshot)?;
|
||||
handler::execute(&self.http, &self.auth, prepared, hooks).await
|
||||
crate::diagnostic::unary(async {
|
||||
let resolved = resolve_request(request)?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
.await?;
|
||||
let prepared = prepare_provider_request(resolved, snapshot)?;
|
||||
crate::diagnostic::provider(&prepared.model, &prepared.custom_llm_provider);
|
||||
let execute: futures_util::future::BoxFuture<
|
||||
'_,
|
||||
Result<ChatCompletionsResponse, Error>,
|
||||
> = Box::pin(handler::execute(&self.http, &self.auth, prepared, hooks));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -96,6 +96,7 @@ fn validate_environment(
|
|||
})
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) fn prepare_provider_request(
|
||||
request: ResolvedChatCompletionsRequest<'_>,
|
||||
secrets: Secrets,
|
||||
|
|
|
|||
324
litellm-rust/crates/core/src/diagnostic.rs
Normal file
324
litellm-rust/crates/core/src/diagnostic.rs
Normal file
|
|
@ -0,0 +1,324 @@
|
|||
use std::{
|
||||
future::Future,
|
||||
pin::Pin,
|
||||
task::{Context, Poll},
|
||||
};
|
||||
|
||||
use futures_util::{Stream, stream::BoxStream};
|
||||
use litellm_host::call::CallOutput;
|
||||
use litellm_tracing::Logger;
|
||||
use tracing::Span;
|
||||
|
||||
struct Completion {
|
||||
span: Span,
|
||||
outcome: &'static str,
|
||||
}
|
||||
|
||||
impl Completion {
|
||||
fn new(name: &str) -> Self {
|
||||
let current = Span::current();
|
||||
Self {
|
||||
span: if current
|
||||
.metadata()
|
||||
.is_some_and(|metadata| metadata.name() == name)
|
||||
{
|
||||
current
|
||||
} else {
|
||||
Span::none()
|
||||
},
|
||||
outcome: "cancelled",
|
||||
}
|
||||
}
|
||||
|
||||
fn finish(mut self, outcome: &'static str) {
|
||||
self.outcome = outcome;
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Completion {
|
||||
fn drop(&mut self) {
|
||||
self.span.record("outcome", self.outcome);
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn provider(model: &str, provider: &str) {
|
||||
let span = Span::current();
|
||||
span.record("resolved_model", model);
|
||||
span.record("provider", provider);
|
||||
}
|
||||
|
||||
pub(crate) async fn unary<R, E>(execute: impl Future<Output = Result<R, E>>) -> Result<R, E> {
|
||||
operation("litellm.route", execute).await
|
||||
}
|
||||
|
||||
pub(crate) async fn operation<R, E>(
|
||||
name: &str,
|
||||
execute: impl Future<Output = Result<R, E>>,
|
||||
) -> Result<R, E> {
|
||||
let completion = Completion::new(name);
|
||||
let result = execute.await;
|
||||
completion.finish(if result.is_ok() { "success" } else { "failure" });
|
||||
result
|
||||
}
|
||||
|
||||
pub(crate) async fn call<R, H, C, E>(
|
||||
execute: impl Future<Output = Result<CallOutput<R, H, C, E>, E>>,
|
||||
) -> Result<CallOutput<R, H, C, E>, E>
|
||||
where
|
||||
C: Send + 'static,
|
||||
E: Send + 'static,
|
||||
{
|
||||
let completion = Completion::new("litellm.route");
|
||||
match execute.await {
|
||||
Err(error) => {
|
||||
completion.finish("failure");
|
||||
Err(error)
|
||||
}
|
||||
Ok(CallOutput::Complete(response)) => {
|
||||
completion.span.record("stream", false);
|
||||
completion.finish("success");
|
||||
Ok(CallOutput::Complete(response))
|
||||
}
|
||||
Ok(CallOutput::Stream { head, chunks }) => {
|
||||
completion.span.record("stream", true);
|
||||
Ok(CallOutput::Stream {
|
||||
head,
|
||||
chunks: Box::pin(TracedStream {
|
||||
state: Some(StreamState {
|
||||
chunks,
|
||||
completion,
|
||||
logger: Logger::current(),
|
||||
}),
|
||||
}),
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct StreamState<C, E> {
|
||||
chunks: BoxStream<'static, Result<C, E>>,
|
||||
completion: Completion,
|
||||
logger: Logger,
|
||||
}
|
||||
|
||||
impl<C, E> StreamState<C, E> {
|
||||
fn close(self, outcome: &'static str) {
|
||||
self.logger
|
||||
.scope(|| self.completion.span.in_scope(|| drop(self.chunks)));
|
||||
self.completion.finish(outcome);
|
||||
}
|
||||
}
|
||||
|
||||
struct TracedStream<C, E> {
|
||||
state: Option<StreamState<C, E>>,
|
||||
}
|
||||
|
||||
impl<C, E> Stream for TracedStream<C, E> {
|
||||
type Item = Result<C, E>;
|
||||
|
||||
fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
let Some(state) = self.state.as_mut() else {
|
||||
return Poll::Ready(None);
|
||||
};
|
||||
let next = state.logger.scope(|| {
|
||||
state
|
||||
.completion
|
||||
.span
|
||||
.in_scope(|| state.chunks.as_mut().poll_next(context))
|
||||
});
|
||||
let outcome = match &next {
|
||||
Poll::Ready(None) => "success",
|
||||
Poll::Ready(Some(Err(_))) => "failure",
|
||||
_ => return next,
|
||||
};
|
||||
if let Some(state) = self.state.take() {
|
||||
state.close(outcome);
|
||||
}
|
||||
next
|
||||
}
|
||||
}
|
||||
|
||||
impl<C, E> Drop for TracedStream<C, E> {
|
||||
fn drop(&mut self) {
|
||||
if let Some(state) = self.state.take() {
|
||||
state.close("cancelled");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{sync::mpsc, task::Context};
|
||||
|
||||
use futures_util::{StreamExt, task::noop_waker_ref};
|
||||
use litellm_tracing::{Metadata, Record, Sink};
|
||||
use rstest::{fixture, rstest};
|
||||
use serde_json::{Value, json};
|
||||
|
||||
use super::*;
|
||||
|
||||
struct Capture(mpsc::Sender<Value>);
|
||||
|
||||
impl Sink for Capture {
|
||||
fn enabled(&self, metadata: &Metadata<'_>) -> bool {
|
||||
*metadata.level() <= tracing::Level::INFO
|
||||
}
|
||||
fn emit(&self, record: &Record) {
|
||||
self.0.send(Value::Object(record.fields.clone())).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
#[fixture]
|
||||
fn logger() -> (Logger, mpsc::Receiver<Value>) {
|
||||
let (sender, receiver) = mpsc::channel();
|
||||
(Logger::new(Capture(sender)), receiver)
|
||||
}
|
||||
|
||||
struct Chunks(std::vec::IntoIter<Result<u8, &'static str>>);
|
||||
|
||||
impl Stream for Chunks {
|
||||
type Item = Result<u8, &'static str>;
|
||||
fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||
tracing::info!(event = "poll");
|
||||
Poll::Ready(self.0.next())
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Chunks {
|
||||
fn drop(&mut self) {
|
||||
tracing::info!(event = "drop");
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "litellm.route",
|
||||
skip_all,
|
||||
fields(route = "fixture", stream, outcome)
|
||||
)]
|
||||
async fn streamed() -> Result<CallOutput<(), (), u8, &'static str>, &'static str> {
|
||||
call(async {
|
||||
Ok(CallOutput::Stream {
|
||||
head: (),
|
||||
chunks: Box::pin(Chunks(vec![Ok(1), Err("broken"), Ok(2)].into_iter())),
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn stream_errors_finish_once_and_poll_and_drop_use_the_captured_context(
|
||||
logger: (Logger, mpsc::Receiver<Value>),
|
||||
) {
|
||||
let (logger, records) = logger;
|
||||
let CallOutput::Stream { mut chunks, .. } = logger.instrument(streamed()).await.unwrap()
|
||||
else {
|
||||
panic!()
|
||||
};
|
||||
assert!(records.try_recv().is_err());
|
||||
tokio::spawn(async move {
|
||||
assert_eq!(chunks.next().await, Some(Ok(1)));
|
||||
assert_eq!(chunks.next().await, Some(Err("broken")));
|
||||
assert_eq!(chunks.next().await, None);
|
||||
let emitted = records.try_iter().collect::<Vec<_>>();
|
||||
assert_eq!(emitted.len(), 4);
|
||||
assert!(emitted.iter().all(|record| record["route"] == "fixture"));
|
||||
assert_eq!(emitted[0]["event"], "poll");
|
||||
assert_eq!(emitted[1]["event"], "poll");
|
||||
assert_eq!(emitted[2]["event"], "drop");
|
||||
assert_eq!(emitted[3]["outcome"], "failure");
|
||||
drop(chunks);
|
||||
assert!(records.try_recv().is_err());
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(route = "waiting", outcome))]
|
||||
async fn waiting(streaming: bool) {
|
||||
if streaming {
|
||||
let _: Result<CallOutput<(), (), u8, ()>, ()> = call(std::future::pending()).await;
|
||||
} else {
|
||||
let _: Result<(), ()> = unary(std::future::pending()).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unary(false)]
|
||||
#[case::streaming(true)]
|
||||
fn cancellation_before_headers_closes_the_span(
|
||||
logger: (Logger, mpsc::Receiver<Value>),
|
||||
#[case] streaming: bool,
|
||||
) {
|
||||
let (logger, records) = logger;
|
||||
let mut future = Box::pin(logger.instrument(waiting(streaming)));
|
||||
assert!(
|
||||
future
|
||||
.as_mut()
|
||||
.poll(&mut Context::from_waker(noop_waker_ref()))
|
||||
.is_pending()
|
||||
);
|
||||
assert!(records.try_recv().is_err());
|
||||
drop(future);
|
||||
let summary = records.try_recv().unwrap();
|
||||
assert_eq!(summary["outcome"], "cancelled");
|
||||
assert_eq!(summary["route"], "waiting");
|
||||
assert!(records.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn dropped_stream_teardown_uses_its_original_logger(
|
||||
logger: (Logger, mpsc::Receiver<Value>),
|
||||
) {
|
||||
let (logger, records) = logger;
|
||||
let output = logger.instrument(streamed()).await.unwrap();
|
||||
Logger::default().scope(|| drop(output));
|
||||
let emitted = records.try_iter().collect::<Vec<_>>();
|
||||
assert_eq!(emitted.len(), 2);
|
||||
assert_eq!(
|
||||
emitted[0],
|
||||
json!({"route":"fixture", "stream":true, "event":"drop"})
|
||||
);
|
||||
assert_eq!(emitted[1]["outcome"], "cancelled");
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "disabled", level = "debug", skip_all, fields(outcome))]
|
||||
async fn disabled_child() {
|
||||
let _: Result<(), ()> = operation("disabled", async { Err(()) }).await;
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_filtered_operation_does_not_overwrite_its_parent_outcome(
|
||||
logger: (Logger, mpsc::Receiver<Value>),
|
||||
) {
|
||||
let (logger, records) = logger;
|
||||
logger
|
||||
.instrument(async {
|
||||
let parent = tracing::info_span!("parent", outcome = "original");
|
||||
tracing::Instrument::instrument(disabled_child(), parent).await;
|
||||
})
|
||||
.await;
|
||||
let summary = records.try_recv().unwrap();
|
||||
assert_eq!(summary["span_name"], "parent");
|
||||
assert_eq!(summary["outcome"], "original");
|
||||
assert!(records.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(stream = true, outcome))]
|
||||
async fn completed() -> Result<CallOutput<(), (), u8, ()>, ()> {
|
||||
call(async { Ok(CallOutput::Complete(())) }).await
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn streaming_mode_reflects_the_returned_output(logger: (Logger, mpsc::Receiver<Value>)) {
|
||||
let (logger, records) = logger;
|
||||
logger.instrument(completed()).await.unwrap();
|
||||
let summary = records.try_recv().unwrap();
|
||||
assert_eq!(summary["stream"], false);
|
||||
assert_eq!(summary["outcome"], "success");
|
||||
assert!(records.try_recv().is_err());
|
||||
}
|
||||
}
|
||||
|
|
@ -1,3 +1,5 @@
|
|||
mod diagnostic;
|
||||
|
||||
pub mod audio_transcription;
|
||||
pub mod chat_completions;
|
||||
pub mod constants;
|
||||
|
|
|
|||
|
|
@ -199,20 +199,20 @@ fn decoded_chunks(
|
|||
}
|
||||
|
||||
fn log_request_body(provider: &str, stream: bool, body: &serde_json::Value) {
|
||||
litellm_tracing::debug!(provider, stream, body = %body, "provider request");
|
||||
tracing::debug!(provider, stream, body = %body, "provider request");
|
||||
}
|
||||
|
||||
fn log_response_body(body: &str) {
|
||||
litellm_tracing::debug!(body, "provider response body");
|
||||
tracing::debug!(body, "provider response body");
|
||||
}
|
||||
|
||||
fn log_error_body(status: u16, body: &str) {
|
||||
litellm_tracing::debug!(status, body, "provider error body");
|
||||
tracing::debug!(status, body, "provider error body");
|
||||
}
|
||||
|
||||
fn log_chunk(provider: &str, stage: &str, data: &bytes::Bytes) {
|
||||
let chunk = ByteChunk::new(data);
|
||||
litellm_tracing::debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk");
|
||||
tracing::debug!(provider, stage, encoding = chunk.encoding(), chunk = %chunk, "stream chunk");
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
|
|||
|
|
@ -39,12 +39,26 @@ impl MessagesRoute {
|
|||
litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "messages",
|
||||
model = %call.body.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = call.body.params.stream == Some(true),
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
call: MessagesCall,
|
||||
hooks: &impl litellm_host::hooks::RouteHooks<Error>,
|
||||
) -> Result<MessagesResponse, Error> {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
handler::execute(&self.http, &self.auth, request, hooks).await
|
||||
crate::diagnostic::call(async {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
crate::diagnostic::provider(&request.body.model, request.provider.as_str());
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<MessagesResponse, Error>> =
|
||||
Box::pin(handler::execute(&self.http, &self.auth, request, hooks));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ pub(super) struct ProviderMessagesRequest {
|
|||
pub(super) api_key: Option<SecretValue>,
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) async fn prepare(
|
||||
call: MessagesCall,
|
||||
secrets: &dyn SecretSource,
|
||||
|
|
|
|||
|
|
@ -28,31 +28,47 @@ impl OcrRoute {
|
|||
litellm_host::lifecycle::observe_unary(hooks.observer(), self.run(request, hooks)).await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "ocr",
|
||||
model = %request.model,
|
||||
resolved_model = %request.model,
|
||||
provider = <&str>::from(request.config.provider()),
|
||||
stream = false,
|
||||
outcome
|
||||
))]
|
||||
pub(super) async fn run(
|
||||
&self,
|
||||
request: LiteLLMOcrRequest,
|
||||
hooks: &impl RouteHooks<Error>,
|
||||
) -> Result<LiteLLMOcrResponse, Error> {
|
||||
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
|
||||
let prepared = prepare_request_document(request).await?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
|
||||
Box::pin(perform_ocr_request(
|
||||
&self.client,
|
||||
prepared,
|
||||
hooks,
|
||||
caller_document,
|
||||
));
|
||||
execute.await
|
||||
crate::diagnostic::unary(async {
|
||||
let caller_document = matches!(&request.document, OcrDocumentInput::Document(_));
|
||||
let prepared = prepare_request_document(request).await?;
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<LiteLLMOcrResponse, Error>> =
|
||||
Box::pin(perform_ocr_request(
|
||||
&self.client,
|
||||
prepared,
|
||||
hooks,
|
||||
caller_document,
|
||||
));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
async fn prepare_request_document(
|
||||
request: LiteLLMOcrRequest<OcrDocumentInput>,
|
||||
) -> Result<ResolvedOcrRequest, Error> {
|
||||
if let OcrDocumentInput::Document(_) = &request.document {
|
||||
return request.map_document(super::document::prepare_document);
|
||||
}
|
||||
tokio::task::spawn_blocking(move || request.map_document(super::document::prepare_document))
|
||||
.await
|
||||
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
|
||||
let logger = litellm_tracing::Logger::current();
|
||||
let span = tracing::Span::current();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
logger.scope(|| span.in_scope(|| request.map_document(super::document::prepare_document)))
|
||||
})
|
||||
.await
|
||||
.map_err(|error| Error::DocumentTask(Arc::new(error)))?
|
||||
}
|
||||
|
|
|
|||
|
|
@ -4,11 +4,19 @@ use litellm_http::outbound::OutboundRequest;
|
|||
use litellm_llms::base_llm::auth::Authenticated;
|
||||
use serde_json::Value;
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "litellm.provider.send",
|
||||
level = "debug",
|
||||
skip_all,
|
||||
fields(status)
|
||||
)]
|
||||
pub(crate) async fn send(
|
||||
request: OutboundRequest,
|
||||
client: &litellm_http::Client,
|
||||
) -> Result<reqwest::Response, reqwest::Error> {
|
||||
request.send(client).await
|
||||
request.send(client).await.inspect(|response| {
|
||||
tracing::Span::current().record("status", response.status().as_u16());
|
||||
})
|
||||
}
|
||||
|
||||
/// Header credentials are already in `headers`; SigV4 is applied here, over the
|
||||
|
|
|
|||
|
|
@ -41,12 +41,29 @@ impl ResponsesRoute {
|
|||
litellm_host::lifecycle::observe_call(hooks.observer(), self.run(call, hooks)).await
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(
|
||||
route = "responses",
|
||||
model = %call.model,
|
||||
provider,
|
||||
resolved_model,
|
||||
stream = call.optional_params.get("stream").and_then(serde_json::Value::as_bool).unwrap_or(false),
|
||||
outcome
|
||||
))]
|
||||
async fn run(
|
||||
&self,
|
||||
call: ResponsesCall,
|
||||
hooks: &impl RouteHooks<Error>,
|
||||
) -> Result<ResponsesOutput, Error> {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
handler::execute(&self.http, &self.auth, request, hooks).await
|
||||
crate::diagnostic::call(async {
|
||||
let request = prepare::prepare(call, self.secrets.as_ref()).await?;
|
||||
crate::diagnostic::provider(
|
||||
&request.context.model,
|
||||
&request.context.custom_llm_provider,
|
||||
);
|
||||
let execute: futures_util::future::BoxFuture<'_, Result<ResponsesOutput, Error>> =
|
||||
Box::pin(handler::execute(&self.http, &self.auth, request, hooks));
|
||||
execute.await
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ use super::{
|
|||
types::{ProviderResponsesRequest, ResponsesCall},
|
||||
};
|
||||
|
||||
#[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)]
|
||||
pub(super) async fn prepare(
|
||||
call: ResponsesCall,
|
||||
secrets: &dyn SecretSource,
|
||||
|
|
|
|||
|
|
@ -29,83 +29,121 @@ pub struct ResponsesWebSocketConnection {
|
|||
}
|
||||
|
||||
impl ResponsesWebSocketConnection {
|
||||
#[tracing::instrument(
|
||||
name = "litellm.websocket.connect_url",
|
||||
level = "debug",
|
||||
skip_all,
|
||||
fields(outcome)
|
||||
)]
|
||||
pub async fn connect_url(
|
||||
url: &str,
|
||||
headers: &HashMap<String, String>,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<Self, Error> {
|
||||
let mut request = url.into_client_request().map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})?;
|
||||
for (name, value) in headers {
|
||||
let header_name = name
|
||||
.parse::<HeaderName>()
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
request.headers_mut().insert(header_name, header_value);
|
||||
}
|
||||
let connect = connect_upstream(request);
|
||||
let result = match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(
|
||||
"Responses WebSocket connection timed out".into(),
|
||||
))
|
||||
})?,
|
||||
None => connect.await,
|
||||
};
|
||||
let (socket, _) = result.map_err(|error| match *error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => {
|
||||
Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
})
|
||||
}
|
||||
other => Error::Transport(litellm_http::transport::Error::Network(other.to_string())),
|
||||
})?;
|
||||
Ok(Self {
|
||||
socket: Arc::new(Mutex::new(Some(socket))),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn send_text(&self, text: String) -> Result<(), Error> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Err(Error::Transport(litellm_http::transport::Error::Network(
|
||||
"Responses WebSocket is closed".into(),
|
||||
)));
|
||||
};
|
||||
socket.send(Message::Text(text)).await.map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Ok(None);
|
||||
};
|
||||
match socket.next().await {
|
||||
Some(Ok(Message::Text(text))) => Ok(Some(text)),
|
||||
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
|
||||
.map(Some)
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string().into())),
|
||||
Some(Ok(Message::Close(_))) | None => Ok(None),
|
||||
Some(Ok(_)) => Ok(None),
|
||||
Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network(
|
||||
error.to_string(),
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn close(&self) -> Result<(), Error> {
|
||||
let mut socket = self.socket.lock().await;
|
||||
if let Some(socket) = socket.as_mut() {
|
||||
socket.close(None).await.map_err(|error| {
|
||||
crate::diagnostic::operation("litellm.websocket.connect_url", async {
|
||||
let mut request = url.into_client_request().map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})?;
|
||||
}
|
||||
*socket = None;
|
||||
Ok(())
|
||||
for (name, value) in headers {
|
||||
let header_name = name
|
||||
.parse::<HeaderName>()
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?;
|
||||
request.headers_mut().insert(header_name, header_value);
|
||||
}
|
||||
let connect = connect_upstream(request);
|
||||
let result = match timeout {
|
||||
Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(
|
||||
"Responses WebSocket connection timed out".into(),
|
||||
))
|
||||
})?,
|
||||
None => connect.await,
|
||||
};
|
||||
let (socket, _) = result.map_err(|error| match *error {
|
||||
tokio_tungstenite::tungstenite::Error::Http(response) => {
|
||||
Error::Transport(litellm_http::transport::Error::Http {
|
||||
status: response.status().as_u16(),
|
||||
body: String::new(),
|
||||
})
|
||||
}
|
||||
other => {
|
||||
Error::Transport(litellm_http::transport::Error::Network(other.to_string()))
|
||||
}
|
||||
})?;
|
||||
Ok(Self {
|
||||
socket: Arc::new(Mutex::new(Some(socket))),
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "litellm.websocket.send_text",
|
||||
level = "debug",
|
||||
skip_all,
|
||||
fields(outcome)
|
||||
)]
|
||||
pub async fn send_text(&self, text: String) -> Result<(), Error> {
|
||||
crate::diagnostic::operation("litellm.websocket.send_text", async {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Err(Error::Transport(litellm_http::transport::Error::Network(
|
||||
"Responses WebSocket is closed".into(),
|
||||
)));
|
||||
};
|
||||
socket.send(Message::Text(text)).await.map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "litellm.websocket.recv_text",
|
||||
level = "debug",
|
||||
skip_all,
|
||||
fields(outcome)
|
||||
)]
|
||||
pub async fn recv_text(&self) -> Result<Option<String>, Error> {
|
||||
crate::diagnostic::operation("litellm.websocket.recv_text", async {
|
||||
let mut socket = self.socket.lock().await;
|
||||
let Some(socket) = socket.as_mut() else {
|
||||
return Ok(None);
|
||||
};
|
||||
match socket.next().await {
|
||||
Some(Ok(Message::Text(text))) => Ok(Some(text)),
|
||||
Some(Ok(Message::Binary(bytes))) => String::from_utf8(bytes.to_vec())
|
||||
.map(Some)
|
||||
.map_err(|error| Error::InvalidResponse(error.to_string().into())),
|
||||
Some(Ok(Message::Close(_))) | None => Ok(None),
|
||||
Some(Ok(_)) => Ok(None),
|
||||
Some(Err(error)) => Err(Error::Transport(litellm_http::transport::Error::Network(
|
||||
error.to_string(),
|
||||
))),
|
||||
}
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
#[tracing::instrument(
|
||||
name = "litellm.websocket.close",
|
||||
level = "debug",
|
||||
skip_all,
|
||||
fields(outcome)
|
||||
)]
|
||||
pub async fn close(&self) -> Result<(), Error> {
|
||||
crate::diagnostic::operation("litellm.websocket.close", async {
|
||||
let mut socket = self.socket.lock().await;
|
||||
if let Some(socket) = socket.as_mut() {
|
||||
socket.close(None).await.map_err(|error| {
|
||||
Error::Transport(litellm_http::transport::Error::Network(error.to_string()))
|
||||
})?;
|
||||
}
|
||||
*socket = None;
|
||||
Ok(())
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -250,3 +250,31 @@ async fn an_unreadable_success_body_is_an_invalid_response(
|
|||
|
||||
assert!(matches!(error, Error::InvalidResponse(_)), "{error:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn transcription_records_route_and_resolved_provider(
|
||||
request: AudioTranscriptionRequest<'static>,
|
||||
traces: TraceCapture,
|
||||
) {
|
||||
let upstream = upstream([transcript_response("hello")]).await;
|
||||
let base = upstream.uri();
|
||||
let model = request.model;
|
||||
traces
|
||||
.logger()
|
||||
.instrument(transcribe(AudioTranscriptionRequest {
|
||||
api_base: Some(&base),
|
||||
..request
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["route"], "audio_transcription");
|
||||
assert_eq!(summaries[0]["model"], model);
|
||||
assert_eq!(summaries[0]["resolved_model"], model);
|
||||
assert_eq!(summaries[0]["provider"], "bedrock");
|
||||
assert_eq!(summaries[0]["outcome"], "success");
|
||||
assert_eq!(summaries[0]["stream"], false);
|
||||
assert!(!format!("{:?}", traces.records()).contains("secret-key"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -434,3 +434,33 @@ async fn a_post_call_hook_failure_never_looks_safe_to_retry(
|
|||
assert_eq!(*source, Error::InvalidRequest("callback rejected".into()));
|
||||
assert_eq!(received(&upstream).await.len(), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn completed_chat_records_route_and_resolved_provider(
|
||||
request: ChatCompletionsRequest<'static>,
|
||||
traces: TraceCapture,
|
||||
) {
|
||||
let upstream = upstream([anthropic_response(ANTHROPIC_MESSAGE)]).await;
|
||||
let base = upstream.uri();
|
||||
let model = request.model;
|
||||
traces
|
||||
.logger()
|
||||
.instrument(complete(ChatCompletionsRequest {
|
||||
api_base: Some(&base),
|
||||
..request
|
||||
}))
|
||||
.await
|
||||
.unwrap();
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["route"], "chat_completions");
|
||||
assert_eq!(summaries[0]["model"], model);
|
||||
assert_eq!(summaries[0]["provider"], "anthropic");
|
||||
assert_eq!(
|
||||
summaries[0]["resolved_model"],
|
||||
only_request(&upstream).await.json()["model"]
|
||||
);
|
||||
assert_eq!(summaries[0]["outcome"], "success");
|
||||
assert_eq!(summaries[0]["stream"], false);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -273,3 +273,34 @@ fn a_body_that_does_not_parse_is_an_invalid_request(#[case] raw: Value) {
|
|||
"{error:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn message_route_summary_excludes_payload_diagnostics(
|
||||
call: MessagesCall,
|
||||
traces: TraceCapture,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
let model = call.body.model.clone();
|
||||
traces
|
||||
.logger()
|
||||
.instrument(run_message(MessagesCall {
|
||||
api_key: Some("private-key-sentinel".into()),
|
||||
api_base: Some(upstream.uri()),
|
||||
..call
|
||||
}))
|
||||
.await;
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["route"], "messages");
|
||||
assert_eq!(summaries[0]["model"], model);
|
||||
assert_eq!(
|
||||
summaries[0]["resolved_model"],
|
||||
only_request(&upstream).await.json()["model"]
|
||||
);
|
||||
assert_eq!(summaries[0]["provider"], "anthropic");
|
||||
assert_eq!(summaries[0]["outcome"], "success");
|
||||
assert_eq!(summaries[0]["stream"], false);
|
||||
assert!(summaries[0].get("body").is_none());
|
||||
assert!(!format!("{:?}", traces.records()).contains("private-key-sentinel"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::sync::Mutex;
|
||||
use std::sync::{Mutex, mpsc};
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, TryStreamExt};
|
||||
|
|
@ -7,6 +7,7 @@ use litellm_core::messages::{
|
|||
route::{Messages, MessagesStreamHead},
|
||||
};
|
||||
use litellm_host::protocol::Demand;
|
||||
use litellm_tracing::{Logger, Metadata, Record, Sink};
|
||||
use rstest::rstest;
|
||||
use tokio::{
|
||||
io::{AsyncReadExt, AsyncWriteExt},
|
||||
|
|
@ -28,6 +29,20 @@ enum Seen {
|
|||
Deliver(Bytes),
|
||||
}
|
||||
|
||||
struct TraceSink(mpsc::Sender<(String, Value)>);
|
||||
|
||||
impl Sink for TraceSink {
|
||||
fn enabled(&self, metadata: &Metadata<'_>) -> bool {
|
||||
metadata.target().starts_with("litellm_core::messages")
|
||||
}
|
||||
|
||||
fn emit(&self, record: &Record) {
|
||||
self.0
|
||||
.send((record.message.clone(), Value::Object(record.fields.clone())))
|
||||
.unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
/// Projects like `LocalMessagesHost`, records every stream op in the order the route
|
||||
/// performs it, and detaches after `detach_after` ops.
|
||||
struct RecordingStreamHost {
|
||||
|
|
@ -161,6 +176,37 @@ async fn upstream_headers_are_on_the_stream_head_before_the_first_chunk(call: Me
|
|||
assert_eq!(delivered, SSE_BODY.as_bytes());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn debug_trace_keeps_provider_input_and_every_stream_chunk(call: MessagesCall) {
|
||||
let upstream = upstream([sse_response()]).await;
|
||||
let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX);
|
||||
let (sender, receiver) = mpsc::channel();
|
||||
|
||||
Logger::new(TraceSink(sender))
|
||||
.instrument(stream_through(&host))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let records: Vec<(String, Value)> = receiver.try_iter().collect();
|
||||
let request = records
|
||||
.iter()
|
||||
.find(|(message, _)| message == "provider request")
|
||||
.unwrap();
|
||||
let body: Value = serde_json::from_str(request.1["body"].as_str().unwrap()).unwrap();
|
||||
assert_eq!(body["messages"][0]["content"], "hi");
|
||||
assert_eq!(request.1["stream"], true);
|
||||
let chunks: String = records
|
||||
.iter()
|
||||
.filter(|(message, fields)| {
|
||||
message == "stream chunk" && fields["stage"] == "provider_response"
|
||||
})
|
||||
.map(|(_, fields)| fields["chunk"].as_str().unwrap())
|
||||
.collect();
|
||||
assert_eq!(chunks, SSE_BODY);
|
||||
assert!(!format!("{records:?}").contains("sk-ant"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::at_open(1)]
|
||||
#[case::after_the_first_chunk(2)]
|
||||
|
|
|
|||
|
|
@ -370,3 +370,34 @@ async fn direct_execution_uses_hooks_without_a_machine() {
|
|||
]
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::native(false)]
|
||||
#[case::hosted(true)]
|
||||
#[tokio::test]
|
||||
async fn ocr_records_one_route_summary_across_both_execution_paths(
|
||||
traces: TraceCapture,
|
||||
#[case] hosted: bool,
|
||||
) {
|
||||
let upstream = upstream([pages_response()]).await;
|
||||
let request = ocr_request("mistral/model", &upstream.uri(), json!({}));
|
||||
let model = request.model.clone();
|
||||
traces
|
||||
.logger()
|
||||
.instrument(async {
|
||||
if hosted {
|
||||
perform_with(LocalOcrHost::new(request)).await
|
||||
} else {
|
||||
perform(request).await
|
||||
}
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["route"], "ocr");
|
||||
assert_eq!(summaries[0]["model"], model);
|
||||
assert_eq!(summaries[0]["provider"], "mistral");
|
||||
assert_eq!(summaries[0]["outcome"], "success");
|
||||
assert_eq!(summaries[0]["stream"], false);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -231,3 +231,212 @@ async fn unsupported_providers_fail_before_secrets_or_transport(
|
|||
assert!(secrets.requested().is_empty());
|
||||
assert!(received(&upstream).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::success(200, json!({"id":"response-1", "model":"test-model", "output":[]}), "success")]
|
||||
#[case::invalid_response(200, json!("private-response-sentinel"), "failure")]
|
||||
#[case::upstream_error(429, json!({"error":"private-response-sentinel"}), "failure")]
|
||||
#[tokio::test]
|
||||
async fn route_tracing_covers_native_and_hosted_outcomes(
|
||||
call: ResponsesCall,
|
||||
traces: TraceCapture,
|
||||
#[case] status: u16,
|
||||
#[case] body: serde_json::Value,
|
||||
#[case] outcome: &str,
|
||||
#[values(false, true)] hosted: bool,
|
||||
) {
|
||||
let upstream = upstream([status_response(status, body)]).await;
|
||||
let host = RecordingCall::<Responses>::new(ResponsesCall {
|
||||
api_base: Some(upstream.uri()),
|
||||
input: json!("private-prompt-sentinel"),
|
||||
api_key: Some("private-key-sentinel".into()),
|
||||
..call
|
||||
});
|
||||
let route = responses_route(no_secrets());
|
||||
let result = traces
|
||||
.logger()
|
||||
.instrument(async {
|
||||
if hosted {
|
||||
litellm_host::in_process::run_hosted(
|
||||
route.clone().machine(host.request().unwrap()),
|
||||
host.runtime(),
|
||||
)
|
||||
.await
|
||||
.map(|_| ())
|
||||
} else {
|
||||
route
|
||||
.execute(host.request().unwrap(), &())
|
||||
.await
|
||||
.map(|_| ())
|
||||
}
|
||||
})
|
||||
.await;
|
||||
assert_eq!(result.is_ok(), outcome == "success");
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
let [summary] = summaries.as_slice() else {
|
||||
panic!("expected one route summary: {summaries:?}")
|
||||
};
|
||||
assert_eq!(summary["route"], "responses");
|
||||
assert_eq!(summary["model"], "openai/test-model");
|
||||
assert_eq!(summary["resolved_model"], "test-model");
|
||||
assert_eq!(summary["provider"], "openai");
|
||||
assert_eq!(summary["stream"], false);
|
||||
assert_eq!(summary["outcome"], outcome);
|
||||
assert!(summary["duration_ms"].as_f64().unwrap() >= 0.0);
|
||||
let sends = traces.summaries("litellm.provider.send");
|
||||
assert_eq!(sends.len(), 1);
|
||||
assert_eq!(sends[0]["status"], status);
|
||||
assert!(!format!("{:?}", traces.records()).contains("private-"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::exhausted(true, "success")]
|
||||
#[case::dropped(false, "cancelled")]
|
||||
#[tokio::test]
|
||||
async fn stream_trace_survives_handoff_and_closes_before_the_stream_object_is_dropped(
|
||||
call: ResponsesCall,
|
||||
traces: TraceCapture,
|
||||
#[case] exhaust: bool,
|
||||
#[case] outcome: &str,
|
||||
) {
|
||||
let upstream = upstream([ResponseTemplate::new(200).set_body_string("stream-bytes")]).await;
|
||||
let output = traces
|
||||
.logger()
|
||||
.instrument(async {
|
||||
responses_route(no_secrets())
|
||||
.execute(
|
||||
ResponsesCall {
|
||||
api_base: Some(upstream.uri()),
|
||||
optional_params: json!({"stream":true}).as_object().unwrap().clone(),
|
||||
..call
|
||||
},
|
||||
&(),
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(traces.summaries("litellm.route").is_empty());
|
||||
let ResponsesOutput::Stream { mut chunks, .. } = output else {
|
||||
panic!("expected stream")
|
||||
};
|
||||
let captured = traces.clone();
|
||||
tokio::spawn(async move {
|
||||
if exhaust {
|
||||
while chunks.try_next().await.unwrap().is_some() {}
|
||||
assert_eq!(captured.summaries("litellm.route").len(), 1);
|
||||
assert!(chunks.try_next().await.unwrap().is_none());
|
||||
}
|
||||
drop(chunks);
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["stream"], true);
|
||||
assert_eq!(summaries[0]["outcome"], outcome);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn preparation_failure_is_traced_but_unpolled_builders_are_not(
|
||||
call: ResponsesCall,
|
||||
traces: TraceCapture,
|
||||
) {
|
||||
traces.logger().scope(|| {
|
||||
drop(responses_route(no_secrets()).execute(
|
||||
ResponsesCall {
|
||||
model: "unknown/model".into(),
|
||||
..call
|
||||
},
|
||||
&(),
|
||||
))
|
||||
});
|
||||
assert!(traces.records().is_empty());
|
||||
let result = traces
|
||||
.logger()
|
||||
.instrument(async {
|
||||
responses_route(no_secrets())
|
||||
.execute(
|
||||
ResponsesCall {
|
||||
model: "unknown/model".into(),
|
||||
..self::call()
|
||||
},
|
||||
&(),
|
||||
)
|
||||
.await
|
||||
})
|
||||
.await;
|
||||
assert!(result.is_err());
|
||||
let summaries = traces.summaries("litellm.route");
|
||||
assert_eq!(summaries.len(), 1);
|
||||
assert_eq!(summaries[0]["outcome"], "failure");
|
||||
assert!(traces.summaries("litellm.provider.send").is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credentials(
|
||||
traces: TraceCapture,
|
||||
) {
|
||||
use futures_util::{SinkExt, StreamExt};
|
||||
use litellm_core::responses::websocket::ResponsesWebSocketConnection;
|
||||
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
let (socket, _) = listener.accept().await.unwrap();
|
||||
let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
|
||||
let message = socket.next().await.unwrap().unwrap();
|
||||
socket.send(message).await.unwrap();
|
||||
let _ = socket.next().await;
|
||||
});
|
||||
traces
|
||||
.logger()
|
||||
.instrument(async {
|
||||
let connection = ResponsesWebSocketConnection::connect_url(
|
||||
&format!("ws://{address}/responses"),
|
||||
&std::collections::HashMap::from([(
|
||||
"authorization".into(),
|
||||
"private-key-sentinel".into(),
|
||||
)]),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
connection
|
||||
.send_text("private-frame-sentinel".into())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
connection.recv_text().await.unwrap().as_deref(),
|
||||
Some("private-frame-sentinel")
|
||||
);
|
||||
connection.close().await.unwrap();
|
||||
assert!(
|
||||
connection
|
||||
.send_text("private-frame-sentinel".into())
|
||||
.await
|
||||
.is_err()
|
||||
);
|
||||
})
|
||||
.await;
|
||||
server.await.unwrap();
|
||||
let sends = traces.summaries("litellm.websocket.send_text");
|
||||
assert_eq!(sends.len(), 2);
|
||||
assert_eq!(sends[0]["outcome"], "success");
|
||||
assert_eq!(sends[1]["outcome"], "failure");
|
||||
assert_eq!(
|
||||
traces.summaries("litellm.websocket.connect_url")[0]["outcome"],
|
||||
"success"
|
||||
);
|
||||
assert_eq!(
|
||||
traces.summaries("litellm.websocket.recv_text")[0]["outcome"],
|
||||
"success"
|
||||
);
|
||||
assert_eq!(
|
||||
traces.summaries("litellm.websocket.close")[0]["outcome"],
|
||||
"success"
|
||||
);
|
||||
assert!(!format!("{:?}", traces.records()).contains("private-"));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -350,3 +350,41 @@ where
|
|||
self.events.0.lock().unwrap().push(event.clone());
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct TraceCapture(Arc<Mutex<Vec<Value>>>);
|
||||
|
||||
impl TraceCapture {
|
||||
pub fn logger(&self) -> litellm_tracing::Logger {
|
||||
litellm_tracing::Logger::new(self.clone())
|
||||
}
|
||||
|
||||
pub fn records(&self) -> Vec<Value> {
|
||||
self.0.lock().unwrap().clone()
|
||||
}
|
||||
|
||||
pub fn summaries(&self, name: &str) -> Vec<Value> {
|
||||
self.records()
|
||||
.into_iter()
|
||||
.filter(|record| record["span_name"] == name)
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl litellm_tracing::Sink for TraceCapture {
|
||||
fn enabled(&self, metadata: &litellm_tracing::Metadata<'_>) -> bool {
|
||||
metadata.target().starts_with("litellm_core")
|
||||
}
|
||||
|
||||
fn emit(&self, record: &litellm_tracing::Record) {
|
||||
self.0
|
||||
.lock()
|
||||
.unwrap()
|
||||
.push(Value::Object(record.fields.clone()));
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::fixture]
|
||||
pub fn traces() -> TraceCapture {
|
||||
TraceCapture::default()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ serde.workspace = true
|
|||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tokio.workspace = true
|
||||
tracing.workspace = true
|
||||
veil.workspace = true
|
||||
webpki-roots.workspace = true
|
||||
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ fn tls_config() -> Result<Arc<ClientConfig>, Box<tokio_tungstenite::tungstenite:
|
|||
Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built)))
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)]
|
||||
pub async fn connect_upstream<R>(
|
||||
request: R,
|
||||
) -> Result<(UpstreamWebSocket, Response), Box<tokio_tungstenite::tungstenite::Error>>
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ tokio = { workspace = true, features = ["rt", "sync"] }
|
|||
url.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
tracing.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-secrets-aws.workspace = true
|
||||
serde.workspace = true
|
||||
|
|
|
|||
|
|
@ -66,6 +66,19 @@ fn warning(py: Python<'_>) {
|
|||
});
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "litellm.route", skip_all, fields(route = "fixture", outcome))]
|
||||
async fn traced_operation(_secret: &str) -> PyResult<()> {
|
||||
tokio::task::yield_now().await;
|
||||
tracing::info!("inside route");
|
||||
tracing::Span::current().record("outcome", "success");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn span_warning(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
|
||||
super::run_async_value(py, traced_operation("private-key-sentinel"))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn levels(py: Python<'_>) {
|
||||
super::capture(py).scope(|| {
|
||||
|
|
@ -109,7 +122,7 @@ fn http_warning(py: Python<'_>) -> PyResult<()> {
|
|||
crate::http::call_config(py, &PyDict::new(py), false).map(|_| ())
|
||||
}
|
||||
|
||||
#[test]
|
||||
#[rstest::rstest]
|
||||
fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication() {
|
||||
if std::env::var_os("LITELLM_LOGGER_TEST_PROCESS").is_none() {
|
||||
let output = Command::new(std::env::current_exe().unwrap())
|
||||
|
|
@ -156,6 +169,9 @@ fn native_events_reach_python_with_levels_context_reentry_and_http_deduplication
|
|||
locals
|
||||
.set_item("warning", wrap_pyfunction!(warning, py).unwrap())
|
||||
.unwrap();
|
||||
locals
|
||||
.set_item("span_warning", wrap_pyfunction!(span_warning, py).unwrap())
|
||||
.unwrap();
|
||||
locals
|
||||
.set_item(
|
||||
"asynchronous_warning",
|
||||
|
|
@ -263,6 +279,25 @@ try:
|
|||
('warn', logging.WARNING), ('error', logging.ERROR),
|
||||
]
|
||||
|
||||
before_spans = len(capture.records)
|
||||
async def traced_request():
|
||||
session = session_id_var.set('span-session')
|
||||
trace = trace_id_var.set('span-trace')
|
||||
try:
|
||||
await span_warning()
|
||||
finally:
|
||||
trace_id_var.reset(trace)
|
||||
session_id_var.reset(session)
|
||||
asyncio.run(traced_request())
|
||||
span_records = capture.records[before_spans:]
|
||||
assert [r.getMessage() for r in span_records] == ['inside route', 'span closed']
|
||||
assert all(r.session_id == 'span-session' and r.trace_id == 'span-trace' for r in span_records)
|
||||
assert all(r.rust_fields['route'] == 'fixture' for r in span_records)
|
||||
assert span_records[1].rust_fields['outcome'] == 'success'
|
||||
assert span_records[1].rust_fields['span_name'] == 'litellm.route'
|
||||
assert span_records[1].rust_fields['duration_ms'] >= 0
|
||||
assert 'private-key-sentinel' not in repr([r.rust_fields for r in span_records])
|
||||
|
||||
before = len(capture.records)
|
||||
litellm.ssl_ecdh_curve = 'logger-test-unsupported-curve'
|
||||
http_warning()
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# Native diagnostic tracing
|
||||
|
||||
`litellm-tracing` connects standard `tracing` events to a host-provided `Sink`. It has no Python dependency and does not install a global subscriber
|
||||
`litellm-tracing` connects standard `tracing` spans and events to a host-provided `Sink`. It has no Python dependency. Hosts choose scoped dispatch or explicitly install a global subscriber
|
||||
|
||||
Use the exported `debug!`, `info!`, `warn!`, and `error!` macros in native code. A host creates a `Logger` with its sink, uses `scope` for synchronous operations, and wraps futures with `instrument`. Instrument spawned futures explicitly because thread-local subscribers do not automatically follow spawned work
|
||||
Use upstream `tracing` macros and `#[tracing::instrument(skip_all, fields(...))]` in native code. A host creates a `Logger` with its sink, uses `scope` for synchronous operations, and wraps futures with `instrument`. Propagate both spans and dispatch into spawned work and returned streams. `Logger::current()` captures the current dispatch. Existing event macro re-exports remain available
|
||||
|
||||
Bindings implement `litellm_tracing::Sink` to connect events to their host runtime:
|
||||
|
||||
|
|
@ -17,6 +17,8 @@ Pass the implementation to `litellm_tracing::Logger::new(sink)`, then call `logg
|
|||
|
||||
Records retain event metadata, the message, and typed event fields. Sink filtering runs for each event so runtime level changes take effect. Logging from inside a sink is suppressed to prevent recursion
|
||||
|
||||
`sink_layer(sink)` exposes the same adapter as a composable `tracing_subscriber::Layer`, with an independent dynamic filter. It records span field updates and inherits fields from outer to inner spans, with event fields taking precedence. Closing a span emits `span closed` at the span's level with `span_name` and monotonic `duration_ms`. Delivery rechecks the sink filter. Core records route outcomes and retains route spans until a stream ends, fails, or is dropped
|
||||
|
||||
The Python bridge scopes native execution to a sink that uses LiteLLM's existing Python logger. It preserves request correlation, redacts before delivering to handlers, maps Rust trace events to Python debug, and reports handler failures through `sys.unraisablehook`. It accepts LiteLLM targets only, keeping dependency wire diagnostics out of the application logger
|
||||
|
||||
Python consumers continue using `litellm._logging` and its existing loggers, filters, formatters, and context setters. Catalog dispatch selects the processing backend for both Python and native diagnostics. The pure `Processor` takes explicit settings and never emits events
|
||||
|
|
|
|||
143
litellm-rust/crates/tracing/src/layer.rs
Normal file
143
litellm-rust/crates/tracing/src/layer.rs
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
use std::{sync::Arc, time::Instant};
|
||||
|
||||
use serde_json::Map;
|
||||
use tracing::{
|
||||
Dispatch, Event, Metadata, Subscriber,
|
||||
span::{Attributes, Id, Record as SpanRecord},
|
||||
};
|
||||
use tracing_subscriber::{
|
||||
Layer, Registry, filter::dynamic_filter_fn, layer::Context, prelude::*, registry::LookupSpan,
|
||||
};
|
||||
|
||||
use crate::{Emitting, Record, Sink};
|
||||
|
||||
pub fn sink_layer<S, R>(sink: S) -> impl Layer<R>
|
||||
where
|
||||
S: Sink,
|
||||
R: Subscriber + for<'a> LookupSpan<'a>,
|
||||
{
|
||||
let sink = Arc::new(sink);
|
||||
let filter_sink = sink.clone();
|
||||
Output(sink).with_filter(dynamic_filter_fn(move |metadata, _| {
|
||||
enabled(filter_sink.as_ref(), metadata)
|
||||
}))
|
||||
}
|
||||
|
||||
pub(crate) fn dispatch(sink: impl Sink) -> Dispatch {
|
||||
let sink = Arc::new(sink);
|
||||
let filter_sink = sink.clone();
|
||||
Dispatch::new(
|
||||
Registry::default()
|
||||
.with(Output(sink))
|
||||
.with(dynamic_filter_fn(move |metadata, _| {
|
||||
enabled(filter_sink.as_ref(), metadata)
|
||||
})),
|
||||
)
|
||||
}
|
||||
|
||||
fn enabled(sink: &impl Sink, metadata: &Metadata<'_>) -> bool {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return false;
|
||||
};
|
||||
sink.enabled(metadata)
|
||||
}
|
||||
|
||||
struct Output<S>(Arc<S>);
|
||||
|
||||
struct SpanData {
|
||||
record: Record,
|
||||
started: Instant,
|
||||
}
|
||||
|
||||
impl<S, R> Layer<R> for Output<S>
|
||||
where
|
||||
S: Sink,
|
||||
R: Subscriber + for<'a> LookupSpan<'a>,
|
||||
{
|
||||
fn on_new_span(&self, attributes: &Attributes<'_>, id: &Id, context: Context<'_, R>) {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return;
|
||||
};
|
||||
let Some(span) = context.span(id) else {
|
||||
return;
|
||||
};
|
||||
let mut extensions = span.extensions_mut();
|
||||
if extensions.get_mut::<SpanData>().is_some() {
|
||||
return;
|
||||
}
|
||||
let mut record = Record {
|
||||
metadata: attributes.metadata(),
|
||||
message: String::new(),
|
||||
fields: Map::new(),
|
||||
};
|
||||
attributes.record(&mut record);
|
||||
extensions.insert(SpanData {
|
||||
record,
|
||||
started: Instant::now(),
|
||||
});
|
||||
}
|
||||
|
||||
fn on_record(&self, id: &Id, values: &SpanRecord<'_>, context: Context<'_, R>) {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return;
|
||||
};
|
||||
let Some(span) = context.span(id) else {
|
||||
return;
|
||||
};
|
||||
if let Some(data) = span.extensions_mut().get_mut::<SpanData>() {
|
||||
values.record(&mut data.record);
|
||||
}
|
||||
}
|
||||
|
||||
fn on_event(&self, event: &Event<'_>, context: Context<'_, R>) {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return;
|
||||
};
|
||||
let mut record = Record {
|
||||
metadata: event.metadata(),
|
||||
message: String::new(),
|
||||
fields: Map::new(),
|
||||
};
|
||||
if let Some(scope) = context.event_scope(event) {
|
||||
for span in scope.from_root() {
|
||||
if let Some(data) = span.extensions().get::<SpanData>() {
|
||||
record.fields.extend(data.record.fields.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
event.record(&mut record);
|
||||
self.0.emit(&record);
|
||||
}
|
||||
|
||||
fn on_close(&self, id: Id, context: Context<'_, R>) {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return;
|
||||
};
|
||||
let Some(span) = context.span(&id) else {
|
||||
return;
|
||||
};
|
||||
if !self.0.enabled(span.metadata()) {
|
||||
return;
|
||||
}
|
||||
let mut record = Record {
|
||||
metadata: span.metadata(),
|
||||
message: "span closed".into(),
|
||||
fields: Map::new(),
|
||||
};
|
||||
for ancestor in span.scope().from_root() {
|
||||
if let Some(data) = ancestor.extensions().get::<SpanData>() {
|
||||
record.fields.extend(data.record.fields.clone());
|
||||
}
|
||||
}
|
||||
let extensions = span.extensions();
|
||||
let Some(data) = extensions.get::<SpanData>() else {
|
||||
return;
|
||||
};
|
||||
record.fields.insert("span_name".into(), span.name().into());
|
||||
record.fields.insert(
|
||||
"duration_ms".into(),
|
||||
(data.started.elapsed().as_secs_f64() * 1000.0).into(),
|
||||
);
|
||||
self.0.emit(&record);
|
||||
}
|
||||
}
|
||||
|
|
@ -8,15 +8,15 @@ use std::{
|
|||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use serde_json::{Map, Value};
|
||||
use tracing::{
|
||||
Dispatch, Event, Subscriber,
|
||||
Dispatch,
|
||||
field::{Field, Visit},
|
||||
subscriber::Interest,
|
||||
};
|
||||
use tracing_subscriber::{Layer, Registry, layer::Context, prelude::*};
|
||||
|
||||
mod layer;
|
||||
mod processing;
|
||||
mod redaction;
|
||||
|
||||
pub use layer::sink_layer;
|
||||
pub use processing::{DiagnosticInput, DiagnosticOutput, Policy, Processor};
|
||||
pub use redaction::{REDACTED, SecretRedactor};
|
||||
pub use tracing::{Level, Metadata, debug, error, info, trace, warn};
|
||||
|
|
@ -66,7 +66,13 @@ pub struct Logger {
|
|||
impl Logger {
|
||||
pub fn new(sink: impl Sink) -> Self {
|
||||
Self {
|
||||
dispatch: Dispatch::new(Registry::default().with(Output(sink))),
|
||||
dispatch: layer::dispatch(sink),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn current() -> Self {
|
||||
Self {
|
||||
dispatch: tracing::dispatcher::get_default(Clone::clone),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -108,37 +114,9 @@ impl Drop for Emitting {
|
|||
}
|
||||
}
|
||||
|
||||
struct Output<S>(S);
|
||||
|
||||
impl<S: Sink, R: Subscriber> Layer<R> for Output<S> {
|
||||
fn register_callsite(&self, _: &'static Metadata<'static>) -> Interest {
|
||||
Interest::sometimes()
|
||||
}
|
||||
|
||||
fn enabled(&self, metadata: &Metadata<'_>, _: Context<'_, R>) -> bool {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return false;
|
||||
};
|
||||
self.0.enabled(metadata)
|
||||
}
|
||||
|
||||
fn on_event(&self, event: &Event<'_>, _: Context<'_, R>) {
|
||||
let Some(_guard) = Emitting::enter() else {
|
||||
return;
|
||||
};
|
||||
let mut record = Record {
|
||||
metadata: event.metadata(),
|
||||
message: String::new(),
|
||||
fields: Map::new(),
|
||||
};
|
||||
event.record(&mut record);
|
||||
self.0.emit(&record);
|
||||
}
|
||||
}
|
||||
|
||||
impl Record {
|
||||
fn field(&mut self, field: &Field, value: Value) {
|
||||
if field.name() == "message" {
|
||||
if field.name() == "message" && self.metadata.is_event() {
|
||||
self.message = match value {
|
||||
Value::String(message) => message,
|
||||
value => value.to_string(),
|
||||
|
|
|
|||
|
|
@ -137,3 +137,145 @@ fn byte_chunk_logging_preserves_exact_bytes(#[case] bytes: &[u8], #[case] encodi
|
|||
};
|
||||
assert_eq!(recovered, bytes);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn spans_inherit_fields_accept_updates_and_close_once() {
|
||||
let (sender, receiver) = mpsc::channel();
|
||||
let logger = Logger::new(Output {
|
||||
enabled: Arc::new(AtomicBool::new(true)),
|
||||
sender,
|
||||
});
|
||||
logger.scope(|| {
|
||||
let route = tracing::info_span!(
|
||||
"route",
|
||||
route = "messages",
|
||||
model = "requested",
|
||||
message = "span message",
|
||||
outcome = tracing::field::Empty
|
||||
);
|
||||
route.in_scope(|| {
|
||||
let child = tracing::debug_span!(
|
||||
"provider",
|
||||
model = "resolved",
|
||||
status = tracing::field::Empty
|
||||
);
|
||||
child.record("status", 201_u64);
|
||||
child.in_scope(|| info!(model = "event", "response"));
|
||||
route.record("outcome", "success");
|
||||
});
|
||||
});
|
||||
let records = receiver.try_iter().collect::<Vec<_>>();
|
||||
assert_eq!(records.len(), 3);
|
||||
assert_eq!(records[0].0, "response");
|
||||
assert_eq!(
|
||||
records[0].1,
|
||||
json!({"route": "messages", "model": "event", "status": 201, "message": "span message"})
|
||||
);
|
||||
assert_eq!(records[1].1["span_name"], "provider");
|
||||
assert_eq!(records[1].1["model"], "resolved");
|
||||
assert_eq!(records[1].1["status"], 201);
|
||||
assert_eq!(records[2].1["span_name"], "route");
|
||||
assert_eq!(records[2].1["model"], "requested");
|
||||
assert_eq!(records[2].1["outcome"], "success");
|
||||
assert!(
|
||||
records[2].1["duration_ms"].as_f64().unwrap()
|
||||
>= records[1].1["duration_ms"].as_f64().unwrap()
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn sink_layers_filter_independently_and_recheck_span_delivery() {
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
let (first_sender, first) = mpsc::channel();
|
||||
let (second_sender, second) = mpsc::channel();
|
||||
let enabled = Arc::new(AtomicBool::new(false));
|
||||
let subscriber = tracing_subscriber::Registry::default()
|
||||
.with(litellm_tracing::sink_layer(Output {
|
||||
enabled: enabled.clone(),
|
||||
sender: first_sender,
|
||||
}))
|
||||
.with(litellm_tracing::sink_layer(Output {
|
||||
enabled: Arc::new(AtomicBool::new(true)),
|
||||
sender: second_sender,
|
||||
}));
|
||||
tracing::subscriber::with_default(subscriber, || {
|
||||
info!("second only");
|
||||
enabled.store(true, Ordering::Relaxed);
|
||||
let span = tracing::info_span!("shared", route = "responses");
|
||||
span.in_scope(|| info!("both"));
|
||||
enabled.store(false, Ordering::Relaxed);
|
||||
drop(span);
|
||||
});
|
||||
assert_eq!(
|
||||
first.try_iter().map(|record| record.0).collect::<Vec<_>>(),
|
||||
["both"]
|
||||
);
|
||||
let records = second.try_iter().collect::<Vec<_>>();
|
||||
assert_eq!(
|
||||
records
|
||||
.iter()
|
||||
.map(|record| record.0.as_str())
|
||||
.collect::<Vec<_>>(),
|
||||
["second only", "both", "span closed"]
|
||||
);
|
||||
assert_eq!(records[1].1["route"], "responses");
|
||||
assert_eq!(records[2].1["route"], "responses");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn disabled_spans_do_not_evaluate_fields() {
|
||||
let (sender, receiver) = mpsc::channel();
|
||||
let logger = Logger::new(Output {
|
||||
enabled: Arc::new(AtomicBool::new(false)),
|
||||
sender,
|
||||
});
|
||||
let evaluated = AtomicBool::new(false);
|
||||
logger.scope(|| {
|
||||
let _span = tracing::info_span!(
|
||||
"disabled",
|
||||
field = {
|
||||
evaluated.store(true, Ordering::Relaxed);
|
||||
"value"
|
||||
}
|
||||
);
|
||||
assert!(_span.is_disabled());
|
||||
});
|
||||
assert!(!evaluated.load(Ordering::Relaxed));
|
||||
assert!(receiver.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "worker", skip_all, fields(id))]
|
||||
async fn traced_worker(id: u64) {
|
||||
tracing::Span::current().record("id", id);
|
||||
tokio::task::yield_now().await;
|
||||
info!("inside worker");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn instrumented_futures_keep_span_fields_in_their_captured_sink() {
|
||||
let tasks = (0..2)
|
||||
.map(|id| {
|
||||
let (sender, receiver) = mpsc::channel();
|
||||
let logger = Logger::new(Output {
|
||||
enabled: Arc::new(AtomicBool::new(true)),
|
||||
sender,
|
||||
});
|
||||
(
|
||||
id,
|
||||
tokio::spawn(logger.instrument(traced_worker(id))),
|
||||
receiver,
|
||||
)
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
for (id, task, receiver) in tasks {
|
||||
task.await.unwrap();
|
||||
let records = receiver.try_iter().collect::<Vec<_>>();
|
||||
assert_eq!(records.len(), 2);
|
||||
assert_eq!(records[0].0, "inside worker");
|
||||
assert_eq!(records[0].1["id"], id);
|
||||
assert_eq!(records[1].1["id"], id);
|
||||
assert_eq!(records[1].1["span_name"], "worker");
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue