refactor: use comment filed in annotated to pass metric-related information (#1385)

tedzhouhk · web-flow · commit 227a0e71f06d · 2025-06-11T17:50:29.000Z
diff --git a/lib/engines/mistralrs/src/lib.rs b/lib/engines/mistralrs/src/lib.rs
@@ -418,9 +418,6 @@ impl
                             id: None,
                             data: Some(delta),
                             event: None,
-                            chunk_tokens: None,
-                            input_tokens: None,
-                            output_tokens: None,
                             comment: None,
                         };
                         yield ann;
@@ -585,9 +582,6 @@ impl AsyncEngine<SingleIn<NvCreateCompletionRequest>, ManyOut<Annotated<Completi
                             id: None,
                             data: Some(inner),
                             event: None,
-                            chunk_tokens: None,
-                            input_tokens: None,
-                            output_tokens: None,
                             comment: None,
                         };
                         yield ann;
diff --git a/lib/llm/src/engines.rs b/lib/llm/src/engines.rs
@@ -202,15 +202,15 @@ impl
                 let response = NvCreateChatCompletionStreamResponse {
                     inner,
                 };
-                yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, chunk_tokens: None, input_tokens: None, output_tokens: None, comment: None };
+                yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None };
                 id += 1;
             }
 
             let inner = deltas.create_choice(0, None, Some(async_openai::types::FinishReason::Stop), None);
             let response = NvCreateChatCompletionStreamResponse {
                 inner,
             };
-            yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, chunk_tokens: None, input_tokens: None, output_tokens: None, comment: None };
+            yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None };
         };
 
         Ok(ResponseStream::new(Box::pin(output), ctx))
@@ -234,11 +234,11 @@ impl AsyncEngine<SingleIn<NvCreateCompletionRequest>, ManyOut<Annotated<Completi
             for c in chars_string.chars() {
                 tokio::time::sleep(*TOKEN_ECHO_DELAY).await;
                 let response = deltas.create_choice(0, Some(c.to_string()), None);
-                yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, chunk_tokens: None, input_tokens: None, output_tokens: None, comment: None };
+                yield Annotated{ id: Some(id.to_string()), data: Some(response), event: None, comment: None };
                 id += 1;
             }
             let response = deltas.create_choice(0, None, Some("stop".to_string()));
-            yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, chunk_tokens: None, input_tokens: None, output_tokens: None, comment: None };
+            yield Annotated { id: Some(id.to_string()), data: Some(response), event: None, comment: None };
 
         };
 
diff --git a/lib/llm/src/http/service/openai.rs b/lib/llm/src/http/service/openai.rs
@@ -27,6 +27,7 @@ use super::{
     service_v2, RouteDoc,
 };
 
+use crate::preprocessor::LLMMetricAnnotation;
 use crate::protocols::openai::embeddings::{NvCreateEmbeddingRequest, NvCreateEmbeddingResponse};
 use crate::protocols::openai::{
     chat_completions::NvCreateChatCompletionResponse, completions::CompletionResponse,
@@ -500,6 +501,12 @@ fn process_event_converter<T: Serialize>(
 ) -> Result<Event, axum::Error> {
     let annotated = annotated.0;
 
+    // update metrics
+    if let Ok(Some(metrics)) = LLMMetricAnnotation::from_annotation(&annotated) {
+        response_collector.observe_current_osl(metrics.output_tokens);
+        response_collector.observe_response(metrics.input_tokens, metrics.chunk_tokens);
+    }
+
     let mut event = Event::default();
 
     if let Some(data) = annotated.data {
@@ -516,16 +523,6 @@ fn process_event_converter<T: Serialize>(
         event = event.event(msg);
     }
 
-    if let Some(osl) = annotated.output_tokens {
-        response_collector.observe_current_osl(osl);
-    }
-
-    if let Some(isl) = annotated.input_tokens {
-        if let Some(chunk_tokens) = annotated.chunk_tokens {
-            response_collector.observe_response(isl, chunk_tokens);
-        }
-    }
-
     if let Some(comments) = annotated.comment {
         for comment in comments {
             event = event.comment(comment);
diff --git a/lib/llm/src/preprocessor.rs b/lib/llm/src/preprocessor.rs
@@ -59,6 +59,41 @@ pub use crate::protocols::common::llm_backend::{BackendOutput, PreprocessedReque
 
 pub const ANNOTATION_FORMATTED_PROMPT: &str = "formatted_prompt";
 pub const ANNOTATION_TOKEN_IDS: &str = "token_ids";
+pub const ANNOTATION_LLM_METRICS: &str = "llm_metrics";
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct LLMMetricAnnotation {
+    pub input_tokens: usize,
+    pub output_tokens: usize,
+    pub chunk_tokens: usize,
+}
+
+impl LLMMetricAnnotation {
+    /// Convert this metrics struct to an Annotated event
+    pub fn to_annotation<T>(&self) -> Result<Annotated<T>, serde_json::Error> {
+        Annotated::from_annotation(ANNOTATION_LLM_METRICS, self)
+    }
+
+    /// Extract LLM metrics from an Annotated event, if present
+    pub fn from_annotation<T>(
+        annotation: &Annotated<T>,
+    ) -> Result<Option<LLMMetricAnnotation>, Box<dyn std::error::Error>> {
+        if annotation.event.is_none() {
+            return Ok(None);
+        }
+        if annotation.event.as_ref().unwrap() != ANNOTATION_LLM_METRICS {
+            return Ok(None);
+        }
+        let comments = annotation
+            .comment
+            .as_ref()
+            .ok_or("missing comments block")?;
+        if comments.len() != 1 {
+            return Err("malformed comments block - expected exactly 1 comment".into());
+        }
+        let metrics: LLMMetricAnnotation = serde_json::from_str(&comments[0])?;
+        Ok(Some(metrics))
+    }
+}
 
 pub struct OpenAIPreprocessor {
     mdcsum: String,
@@ -251,9 +286,20 @@ impl OpenAIPreprocessor {
                             .map_err(|e| e.to_string())
                     });
 
-                    response.chunk_tokens = Some(chunk_tokens);
-                    response.input_tokens = Some(isl);
-                    response.output_tokens = Some(current_osl);
+                    // Create LLM metrics annotation
+                    let llm_metrics = LLMMetricAnnotation {
+                        input_tokens: isl,
+                        output_tokens: current_osl,
+                        chunk_tokens,
+                    };
+
+                    if let Ok(metrics_annotated) = llm_metrics.to_annotation::<()>() {
+                        // Only set event if not already set to avoid overriding existing events (like errors)
+                        if response.event.is_none() {
+                            response.event = metrics_annotated.event;
+                        }
+                        response.comment = metrics_annotated.comment;
+                    }
 
                     tracing::trace!(
                         request_id = inner.context.id(),
diff --git a/lib/llm/src/protocols/codec.rs b/lib/llm/src/protocols/codec.rs
@@ -118,9 +118,6 @@ where
             data,
             id: value.id,
             event: value.event,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: value.comments,
         })
     }
diff --git a/lib/llm/src/protocols/openai/chat_completions/aggregator.rs b/lib/llm/src/protocols/openai/chat_completions/aggregator.rs
@@ -284,9 +284,6 @@ mod tests {
             data: Some(data),
             id: Some("test_id".to_string()),
             event: None,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: None,
         }
     }
@@ -430,9 +427,6 @@ mod tests {
             data: Some(data),
             id: Some("test_id".to_string()),
             event: None,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: None,
         };
         let stream = Box::pin(stream::iter(vec![annotated_delta]));
diff --git a/lib/llm/src/protocols/openai/completions/aggregator.rs b/lib/llm/src/protocols/openai/completions/aggregator.rs
@@ -205,9 +205,6 @@ mod tests {
             }),
             id: Some("test_id".to_string()),
             event: None,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: None,
         }
     }
@@ -317,9 +314,6 @@ mod tests {
             }),
             id: Some("test_id".to_string()),
             event: None,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: None,
         };
 
diff --git a/lib/runtime/src/protocols/annotated.rs b/lib/runtime/src/protocols/annotated.rs
@@ -37,12 +37,6 @@ pub struct Annotated<R> {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub event: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub chunk_tokens: Option<usize>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub input_tokens: Option<usize>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub output_tokens: Option<usize>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     pub comment: Option<Vec<String>>,
 }
 
@@ -53,9 +47,6 @@ impl<R> Annotated<R> {
             data: None,
             id: None,
             event: Some("error".to_string()),
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: Some(vec![error]),
         }
     }
@@ -66,9 +57,6 @@ impl<R> Annotated<R> {
             data: Some(data),
             id: None,
             event: None,
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: None,
         }
     }
@@ -84,9 +72,6 @@ impl<R> Annotated<R> {
             data: None,
             id: None,
             event: Some(name.into()),
-            chunk_tokens: None,
-            input_tokens: None,
-            output_tokens: None,
             comment: Some(vec![serde_json::to_string(value)?]),
         })
     }
@@ -122,9 +107,6 @@ impl<R> Annotated<R> {
             data,
             id: self.id,
             event: self.event,
-            chunk_tokens: self.chunk_tokens,
-            input_tokens: self.input_tokens,
-            output_tokens: self.output_tokens,
             comment: self.comment,
         }
     }
@@ -140,9 +122,6 @@ impl<R> Annotated<R> {
                 data,
                 id: self.id,
                 event: self.event,
-                chunk_tokens: self.chunk_tokens,
-                input_tokens: self.input_tokens,
-                output_tokens: self.output_tokens,
                 comment: self.comment,
             },
             Err(e) => Annotated::from_error(e),
