Add service to stream custom otel traces to otel-collector (#262)

crates/prompt_gateway/src/http_context.rs

@@ -1,4 +1,7 @@
-use std::{collections::HashMap, time::Duration};
+use std::{
+    collections::HashMap,
+    time::{Duration, SystemTime, UNIX_EPOCH},
+};
 
 use common::{
     common_types::{
@@ -13,7 +16,9 @@ use common::{
         HEALTHZ_PATH, REQUEST_ID_HEADER, TOOL_ROLE, TRACE_PARENT_HEADER, USER_ROLE,
     },
     errors::ServerError,
-    http::{CallArgs, Client}, pii::obfuscate_auth_header,
+    http::{CallArgs, Client},
+    pii::obfuscate_auth_header,
+    tracing::{Event, Span},
 };
 use http::StatusCode;
 use log::{debug, trace, warn};
@@ -239,7 +244,7 @@ impl HttpContext for StreamContext {
 
     fn on_http_response_body(&mut self, body_size: usize, end_of_stream: bool) -> Action {
         trace!(
-            "recv [S={}] bytes={} end_stream={}",
+            "on_http_response_body: recv [S={}] bytes={} end_stream={}",
             self.context_id,
             body_size,
             end_of_stream
@@ -250,6 +255,55 @@ impl HttpContext for StreamContext {
             return Action::Continue;
         }
 
+        if self.time_to_first_token.is_none() {
+            self.time_to_first_token = Some(
+                SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_nanos(),
+            );
+        }
+
+        if end_of_stream && body_size == 0 {
+            if let Some(traceparent) = self.traceparent.as_ref() {
+                let since_the_epoch_ns = SystemTime::now()
+                    .duration_since(UNIX_EPOCH)
+                    .unwrap()
+                    .as_nanos();
+
+                let traceparent_tokens = traceparent.split("-").collect::<Vec<&str>>();
+                if traceparent_tokens.len() != 4 {
+                    warn!("traceparent header is invalid: {}", traceparent);
+                    return Action::Continue;
+                }
+                let parent_trace_id = traceparent_tokens[1];
+                let parent_span_id = traceparent_tokens[2];
+                let mut trace_data = common::tracing::TraceData::new();
+                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);
+                // send trace_data to http tracing endpoint
+            }
+            return Action::Continue;
+        }
+
         let body = if self.streaming_response {
             let streaming_chunk = match self.get_http_response_body(0, body_size) {
                 Some(chunk) => chunk,