Stream content endpoint (foundation-model-stack#272)

* wip: /stream-content endpoint

Signed-off-by: Mateus Devino <[email protected]>

* wip: initial streaming and request parsing

Signed-off-by: Mateus Devino <[email protected]>

* wip: refactor stream error handling

Signed-off-by: Mateus Devino <[email protected]>

* Improve error handling

Signed-off-by: Mateus Devino <[email protected]>

* tweaks

Signed-off-by: declark1 <[email protected]>

* Move handle_streaming_content_detection handler to separate file

Signed-off-by: declark1 <[email protected]>

* Rename response 'detections' field

Signed-off-by: Mateus Devino <[email protected]>

* Fix error handling when 'content' is missing on first request

Signed-off-by: Mateus Devino <[email protected]>

* Fix generation streaming endpoint link

Signed-off-by: Mateus Devino <[email protected]>

* Fix error handling - second stream message onwards

Signed-off-by: Mateus Devino <[email protected]>

* Return an actual ndjson

Signed-off-by: Mateus Devino <[email protected]>

* stream client binary (for testing)

Signed-off-by: Mateus Devino <[email protected]>

* nd_json for error

Signed-off-by: Mateus Devino <[email protected]>

* Refactor detector extraction from first frame

Signed-off-by: Mateus Devino <[email protected]>

* wip: adapt existing code to new endpoint

Signed-off-by: Mateus Devino <[email protected]>

* Adapt aggregator logic

Signed-off-by: Mateus Devino <[email protected]>

* Forward results to output stream

Signed-off-by: Mateus Devino <[email protected]>

* Fix error handling for second frame onward

Signed-off-by: Mateus Devino <[email protected]>

* Remove unused code

Signed-off-by: Mateus Devino <[email protected]>

* Rename generation variables

Signed-off-by: Mateus Devino <[email protected]>

* Re-add request validation

Signed-off-by: Mateus Devino <[email protected]>

* Remove unused code

Signed-off-by: Mateus Devino <[email protected]>

* Re-add missing aggregator test

Signed-off-by: Mateus Devino <[email protected]>

* Remove unneeded derives

Signed-off-by: Mateus Devino <[email protected]>

* Add copyright note to streaming_content_detection/aggregator.rs

Signed-off-by: Mateus Devino <[email protected]>

* Nest imports on streaming_content_detection.rs

Signed-off-by: Mateus Devino <[email protected]>

* Rename request to frame

Signed-off-by: Mateus Devino <[email protected]>

* Remove unused variable

Signed-off-by: Mateus Devino <[email protected]>

* Replace errors Debug with Display

Signed-off-by: Mateus Devino <[email protected]>

* Reuse streaming/aggregator.rs types

Signed-off-by: Mateus Devino <[email protected]>

* Simplify aggregator logic

Signed-off-by: Mateus Devino <[email protected]>

* Refactor nd_json conversion logic

Signed-off-by: Mateus Devino <[email protected]>

* Rename streaming_output_detection_task()

Signed-off-by: Mateus Devino <[email protected]>

* Drop validation for extra detectors on second frame onwards

Signed-off-by: Mateus Devino <[email protected]>

* Fix import formatting

Signed-off-by: Mateus Devino <[email protected]>

* Add comments for a possible refactor

Signed-off-by: Mateus Devino <[email protected]>

* Replace Tracker Default implementation with a derive macro

Signed-off-by: Mateus Devino <[email protected]>

---------

Signed-off-by: Mateus Devino <[email protected]>
Signed-off-by: declark1 <[email protected]>
Co-authored-by: declark1 <[email protected]>

Author: mdevino

src/models.rs

@@ -1105,6 +1105,34 @@ pub struct EvidenceObj {
     pub evidence: Option<Vec<Evidence>>,
 }
 
+/// Stream content detection stream request
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+#[cfg_attr(test, derive(Default))]
+pub struct StreamingContentDetectionRequest {
+    pub detectors: Option<HashMap<String, DetectorParams>>,
+    pub content: String,
+}
+
+impl StreamingContentDetectionRequest {
+    /// validates stream messages
+    pub fn validate(&self) -> Result<(), ValidationError> {
+        if self.content.is_empty() {
+            return Err(ValidationError::Invalid(
+                "`content` cannot be empty".to_string(),
+            ));
+        }
+        Ok(())
+    }
+}
+
+/// Stream content detection response
+#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct StreamingContentDetectionResponse {
+    pub detections: Vec<ContentAnalysisResponse>,
+    pub processed_index: u32,
+    pub start_index: u32,
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/orchestrator.rs

@@ -17,11 +17,13 @@
 
 pub mod errors;
 pub use errors::Error;
+use futures::Stream;
 pub mod chat_completions_detection;
 pub mod streaming;
+pub mod streaming_content_detection;
 pub mod unary;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, pin::Pin, sync::Arc};
 
 use axum::http::header::HeaderMap;
 use opentelemetry::trace::TraceId;
@@ -44,7 +46,8 @@ use crate::{
     models::{
         ChatDetectionHttpRequest, ContextDocsHttpRequest, DetectionOnGeneratedHttpRequest,
         DetectorParams, GenerationWithDetectionHttpRequest, GuardrailsConfig,
-        GuardrailsHttpRequest, GuardrailsTextGenerationParameters, TextContentDetectionHttpRequest,
+        GuardrailsHttpRequest, GuardrailsTextGenerationParameters,
+        StreamingContentDetectionRequest, TextContentDetectionHttpRequest,
     },
 };
 
@@ -490,6 +493,31 @@ impl ChatCompletionsDetectionTask {
     }
 }
 
+pub struct StreamingContentDetectionTask {
+    pub trace_id: TraceId,
+    pub headers: HeaderMap,
+    pub detectors: HashMap<String, DetectorParams>,
+    pub input_stream:
+        Pin<Box<dyn Stream<Item = Result<StreamingContentDetectionRequest, Error>> + Send>>,
+}
+
+impl StreamingContentDetectionTask {
+    pub fn new(
+        trace_id: TraceId,
+        headers: HeaderMap,
+        input_stream: Pin<
+            Box<dyn Stream<Item = Result<StreamingContentDetectionRequest, Error>> + Send>,
+        >,
+    ) -> Self {
+        Self {
+            trace_id,
+            headers,
+            detectors: HashMap::default(),
+            input_stream,
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

src/orchestrator/errors.rs

@@ -14,8 +14,7 @@
  limitations under the License.
 
 */
-
-use crate::clients;
+use crate::{clients, models::ValidationError};
 
 /// Orchestrator errors.
 #[derive(Debug, Clone, PartialEq, thiserror::Error)]
@@ -34,10 +33,14 @@ pub enum Error {
     ChatGenerateRequestFailed { id: String, error: clients::Error },
     #[error("tokenize request failed for `{id}`: {error}")]
     TokenizeRequestFailed { id: String, error: clients::Error },
+    #[error("validation error: {0}")]
+    Validation(String),
     #[error("{0}")]
     Other(String),
     #[error("cancelled")]
     Cancelled,
+    #[error("json deserialization error: {0}")]
+    JsonError(String),
 }
 
 impl From<tokio::task::JoinError> for Error {
@@ -49,3 +52,15 @@ impl From<tokio::task::JoinError> for Error {
         }
     }
 }
+
+impl From<serde_json::Error> for Error {
+    fn from(value: serde_json::Error) -> Self {
+        Self::JsonError(value.to_string())
+    }
+}
+
+impl From<ValidationError> for Error {
+    fn from(value: ValidationError) -> Self {
+        Self::Validation(value.to_string())
+    }
+}

src/orchestrator/streaming.rs

@@ -15,7 +15,7 @@
 
 */
 
-mod aggregator;
+pub mod aggregator;
 
 use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration};
 

src/orchestrator/streaming/aggregator.rs

@@ -370,7 +370,7 @@ impl GenerationActorHandle {
 }
 
 #[derive(Debug, Clone)]
-struct TrackerEntry {
+pub struct TrackerEntry {
     pub chunk: Chunk,
     pub detections: Vec<Detections>,
 }
@@ -384,8 +384,8 @@ impl TrackerEntry {
     }
 }
 
-#[derive(Debug, Clone)]
-struct Tracker {
+#[derive(Debug, Clone, Default)]
+pub struct Tracker {
     state: BTreeMap<Span, TrackerEntry>,
 }