Drop Batch associated type from DetectionBatcher and generics, drop detector_id from DetectionStream, integrate single detection stream optimization

Signed-off-by: declark1 <44146800+declark1@users.noreply.github.com>

Author: declark1

src/orchestrator/common/tasks.rs

@@ -190,9 +190,9 @@ pub async fn text_contents_detections(
     ctx: Arc<Context>,
     headers: HeaderMap,
     detectors: HashMap<String, DetectorParams>,
-    input_id: InputId,
+    input_id: u32,
     inputs: Vec<(usize, String)>,
-) -> Result<(InputId, Detections), Error> {
+) -> Result<(u32, Detections), Error> {
     let chunkers = get_chunker_ids(&ctx, &detectors)?;
     let chunk_map = chunks(ctx.clone(), chunkers, inputs).await?;
     let inputs = detectors
@@ -249,7 +249,7 @@ pub async fn text_contents_detection_streams(
     ctx: Arc<Context>,
     headers: HeaderMap,
     detectors: HashMap<String, DetectorParams>,
-    input_id: InputId,
+    input_id: u32,
     input_rx: mpsc::Receiver<Result<(usize, String), Error>>, // (message_index, text)
 ) -> Result<Vec<DetectionStream>, Error> {
     // Create chunk streams
@@ -294,14 +294,8 @@ pub async fn text_contents_detection_streams(
                                         .filter(|detection| detection.score >= threshold)
                                         .collect::<Detections>();
                                     // Send to detection channel
-                                    let _ = detection_tx
-                                        .send(Ok((
-                                            input_id,
-                                            detector_id.clone(),
-                                            chunk,
-                                            detections,
-                                        )))
-                                        .await;
+                                    let _ =
+                                        detection_tx.send(Ok((input_id, chunk, detections))).await;
                                 }
                                 Err(error) => {
                                     // Send error to detection channel
@@ -985,9 +979,7 @@ mod test {
 
         let mut fake_detector_stream = detection_streams.swap_remove(0);
         let mut results = Vec::with_capacity(1);
-        while let Some(Ok((_input_id, _detector_id, _chunk, detections))) =
-            fake_detector_stream.next().await
-        {
+        while let Some(Ok((_input_id, _chunk, detections))) = fake_detector_stream.next().await {
             results.push(detections);
         }
         assert_eq!(results.len(), 1);

src/orchestrator/handlers/streaming_classification_with_gen.rs

@@ -40,8 +40,7 @@ use crate::{
         Context, Error, Orchestrator,
         common::{self, validate_detectors},
         types::{
-            Chunk, DetectionBatchStream, DetectionStream, Detections, GenerationStream,
-            MaxProcessedIndexBatcher,
+            Chunk, DetectionBatchStream, Detections, GenerationStream, MaxProcessedIndexBatcher,
         },
     },
 };
@@ -254,12 +253,6 @@ async fn handle_output_detection(
         let generations = generations.clone();
         async move {
             match detection_streams {
-                Ok(mut detection_streams) if detection_streams.len() == 1 => {
-                    // Process single detection stream, batching not applicable
-                    let detection_stream = detection_streams.swap_remove(0);
-                    process_detection_stream(trace_id, generations, detection_stream, response_tx)
-                        .await;
-                }
                 Ok(detection_streams) => {
                     // Create detection batch stream
                     let detection_batch_stream = DetectionBatchStream::new(
@@ -335,47 +328,17 @@ async fn forward_generation_stream(
     info!(%trace_id, "task completed: generation stream closed");
 }
 
-/// Consumes a detection stream, builds responses, and sends them to a response channel.
-#[instrument(skip_all)]
-async fn process_detection_stream(
-    trace_id: TraceId,
-    generations: Arc<RwLock<Vec<ClassifiedGeneratedTextStreamResult>>>,
-    mut detection_stream: DetectionStream,
-    response_tx: mpsc::Sender<Result<ClassifiedGeneratedTextStreamResult, Error>>,
-) {
-    while let Some(result) = detection_stream.next().await {
-        match result {
-            Ok((_, _detector_id, chunk, detections)) => {
-                // Create response for this batch with output detections
-                let response = output_detection_response(&generations, chunk, detections).unwrap();
-                // Send message to response channel
-                if response_tx.send(Ok(response)).await.is_err() {
-                    info!(%trace_id, "task completed: client disconnected");
-                    return;
-                }
-            }
-            Err(error) => {
-                error!(%trace_id, %error, "task failed: error received from detection stream");
-                // Send error to response channel and terminate
-                let _ = response_tx.send(Err(error)).await;
-                return;
-            }
-        }
-    }
-    info!(%trace_id, "task completed: detection stream closed");
-}
-
 /// Consumes a detection batch stream, builds responses, and sends them to a response channel.
 #[instrument(skip_all)]
 async fn process_detection_batch_stream(
     trace_id: TraceId,
     generations: Arc<RwLock<Vec<ClassifiedGeneratedTextStreamResult>>>,
-    mut detection_batch_stream: DetectionBatchStream<MaxProcessedIndexBatcher>,
+    mut detection_batch_stream: DetectionBatchStream,
     response_tx: mpsc::Sender<Result<ClassifiedGeneratedTextStreamResult, Error>>,
 ) {
     while let Some(result) = detection_batch_stream.next().await {
         match result {
-            Ok((chunk, detections)) => {
+            Ok((_, chunk, detections)) => {
                 // Create response for this batch with output detections
                 let response = output_detection_response(&generations, chunk, detections).unwrap();
                 // Send message to response channel

src/orchestrator/handlers/streaming_content_detection.rs

@@ -30,7 +30,7 @@ use crate::{
     orchestrator::{
         Context, Error, Orchestrator,
         common::{self, validate_detectors},
-        types::{BoxStream, DetectionBatchStream, DetectionStream, MaxProcessedIndexBatcher},
+        types::{BoxStream, DetectionBatchStream, MaxProcessedIndexBatcher},
     },
 };
 
@@ -133,11 +133,6 @@ async fn handle_detection(
     tokio::spawn(
         async move {
             match detection_streams {
-                Ok(mut detection_streams) if detection_streams.len() == 1 => {
-                    // Process single detection stream, batching not applicable
-                    let detection_stream = detection_streams.swap_remove(0);
-                    process_detection_stream(trace_id, detection_stream, response_tx).await;
-                }
                 Ok(detection_streams) => {
                     // Create detection batch stream
                     let detection_batch_stream = DetectionBatchStream::new(
@@ -177,48 +172,16 @@ async fn handle_detection(
     );
 }
 
-/// Consumes a detection stream, builds responses, and sends them to a response channel.
-#[instrument(skip_all)]
-async fn process_detection_stream(
-    trace_id: TraceId,
-    mut detection_stream: DetectionStream,
-    response_tx: mpsc::Sender<Result<StreamingContentDetectionResponse, Error>>,
-) {
-    while let Some(result) = detection_stream.next().await {
-        match result {
-            Ok((_, _detector_id, chunk, detections)) => {
-                let response = StreamingContentDetectionResponse {
-                    start_index: chunk.start as u32,
-                    processed_index: chunk.end as u32,
-                    detections: detections.into(),
-                };
-                // Send message to response channel
-                if response_tx.send(Ok(response)).await.is_err() {
-                    info!(%trace_id, "task completed: client disconnected");
-                    return;
-                }
-            }
-            Err(error) => {
-                error!(%trace_id, %error, "task failed: error received from detection stream");
-                // Send error to response channel and terminate
-                let _ = response_tx.send(Err(error)).await;
-                return;
-            }
-        }
-    }
-    info!(%trace_id, "task completed: detection stream closed");
-}
-
 /// Consumes a detection batch stream, builds responses, and sends them to a response channel.
 #[instrument(skip_all)]
 async fn process_detection_batch_stream(
     trace_id: TraceId,
-    mut detection_batch_stream: DetectionBatchStream<MaxProcessedIndexBatcher>,
+    mut detection_batch_stream: DetectionBatchStream,
     response_tx: mpsc::Sender<Result<StreamingContentDetectionResponse, Error>>,
 ) {
     while let Some(result) = detection_batch_stream.next().await {
         match result {
-            Ok((chunk, detections)) => {
+            Ok((_, chunk, detections)) => {
                 let response = StreamingContentDetectionResponse {
                     start_index: chunk.start as u32,
                     processed_index: chunk.end as u32,

src/orchestrator/types.rs

@@ -37,12 +37,11 @@ use crate::{
 
 pub type ChunkerId = String;
 pub type DetectorId = String;
-pub type InputId = u32;
 
 pub type BoxStream<T> = Pin<Box<dyn Stream<Item = T> + Send>>;
 pub type ChunkStream = BoxStream<Result<Chunk, Error>>;
 pub type InputStream = BoxStream<Result<(usize, String), Error>>;
-pub type DetectionStream = BoxStream<Result<(InputId, DetectorId, Chunk, Detections), Error>>;
+pub type DetectionStream = BoxStream<Result<(u32, Chunk, Detections), Error>>;
 pub type GenerationStream = BoxStream<(usize, Result<ClassifiedGeneratedTextStreamResult, Error>)>;
 pub type ChatCompletionStream = BoxStream<(usize, Result<Option<ChatCompletionChunk>, Error>)>;
 pub type CompletionStream = BoxStream<(usize, Result<Option<Completion>, Error>)>;