diff --git a/crates/common/src/tracing.rs b/crates/common/src/tracing.rs index 4d398333..31961d8b 100644 --- a/crates/common/src/tracing.rs +++ b/crates/common/src/tracing.rs @@ -41,6 +41,63 @@ pub struct Span { pub end_time_unix_nano: String, pub kind: u32, pub attributes: Vec, + pub events: Option>, +} + +impl Span { + pub fn new( + name: String, + parent_trace_id: String, + parent_span_id: Option, + start_time_unix_nano: u128, + end_time_unix_nano: u128, + ) -> Self { + Span { + trace_id: parent_trace_id, + span_id: get_random_span_id(), + parent_span_id, + name, + start_time_unix_nano: format!("{}", start_time_unix_nano), + end_time_unix_nano: format!("{}", end_time_unix_nano), + kind: 0, + attributes: Vec::new(), + events: None, + } + } + + pub fn add_attribute(&mut self, key: String, value: String) { + self.attributes.push(Attribute { + key, + value: AttributeValue { + string_value: Some(value), + }, + }); + } + + pub fn add_event(&mut self, event: Event) { + if self.events.is_none() { + self.events = Some(Vec::new()); + } + self.events.as_mut().unwrap().push(event); + } +} + +#[derive(Serialize, Deserialize, Debug)] +pub struct Event { + #[serde(rename = "timeUnixNano")] + pub time_unix_nano: String, + pub name: String, + pub attributes: Vec, +} + +impl Event { + pub fn new(name: String, time_unix_nano: u128) -> Self { + Event { + time_unix_nano: format!("{}", time_unix_nano), + name, + attributes: Vec::new(), + } + } } #[derive(Serialize, Deserialize, Debug)] diff --git a/crates/prompt_gateway/src/http_context.rs b/crates/prompt_gateway/src/http_context.rs index ae51fdc9..2ef10025 100644 --- a/crates/prompt_gateway/src/http_context.rs +++ b/crates/prompt_gateway/src/http_context.rs @@ -18,7 +18,7 @@ use common::{ errors::ServerError, http::{CallArgs, Client}, pii::obfuscate_auth_header, - tracing::{get_random_span_id, Span}, + tracing::{Event, Span}, }; use http::StatusCode; use log::{debug, trace, warn}; @@ -243,9 +243,11 @@ impl HttpContext for StreamContext { } fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> Action { - debug!( + trace!( "on_http_response_body: recv [S={}] bytes={} end_stream={}", - self.context_id, body_size, end_of_stream + self.context_id, + body_size, + end_of_stream ); if !self.is_chat_completions_request { @@ -277,27 +279,23 @@ impl HttpContext for StreamContext { let parent_trace_id = traceparent_tokens[1]; let parent_span_id = traceparent_tokens[2]; let mut trace_data = common::tracing::TraceData::new(); - trace_data.add_span(Span { - trace_id: parent_trace_id.to_string(), - parent_span_id: Some(parent_span_id.to_string()), - span_id: format!("{}", get_random_span_id()), - name: "total_time".to_string(), - start_time_unix_nano: format!("{}", self.start_upstream_llm_request_time), - end_time_unix_nano: format!("{}", since_the_epoch_ns), - kind: 1, - attributes: vec![], - }); - - trace_data.add_span(Span { - trace_id: parent_trace_id.to_string(), - parent_span_id: Some(parent_span_id.to_string()), - span_id: format!("{}", get_random_span_id()), - name: "time_to_first_token".to_string(), - start_time_unix_nano: format!("{}", self.start_upstream_llm_request_time), - end_time_unix_nano: format!("{}", self.time_to_first_token.unwrap()), - kind: 1, - attributes: vec![], - }); + let mut llm_span = Span::new( + "upstream_llm_time".to_string(), + parent_trace_id.to_string(), + Some(parent_span_id.to_string()), + self.start_upstream_llm_request_time, + since_the_epoch_ns, + ); + if let Some(prompt) = self.user_prompt.as_ref() { + if let Some(content) = prompt.content.as_ref() { + llm_span.add_attribute("user_prompt".to_string(), content.to_string()); + } + } + llm_span.add_event(Event::new( + "time_to_first_token".to_string(), + self.time_to_first_token.unwrap(), + )); + trace_data.add_span(llm_span); let trace_data_str = serde_json::to_string(&trace_data).unwrap(); debug!("upstream_llm trace details: {}", trace_data_str);