Add detector_id in response for orchestrator apis

src/clients.rs

@@ -31,8 +31,7 @@ use ginepro::LoadBalancedChannel;
 use hyper_timeout::TimeoutConnector;
 use hyper_util::rt::TokioExecutor;
 use tonic::{metadata::MetadataMap, Request};
-use tower::timeout::TimeoutLayer;
-use tower::ServiceBuilder;
+use tower::{timeout::TimeoutLayer, ServiceBuilder};
 use tracing::{debug, instrument, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 use url::Url;

src/clients/detector/text_contents.rs

@@ -132,6 +132,8 @@ pub struct ContentAnalysisResponse {
     pub detection: String,
     /// Detection type or aggregate detection label
     pub detection_type: String,
+    /// ID of Detector
+    pub detector_id: Option<String>,
     /// Score of detection
     pub score: f64,
     /// Optional, any applicable evidence for detection
@@ -147,6 +149,7 @@ impl From<ContentAnalysisResponse> for crate::models::TokenClassificationResult
             word: value.text,
             entity: value.detection,
             entity_group: value.detection_type,
+            detector_id: value.detector_id,
             score: value.score,
             token_count: None,
         }

src/clients/http.rs

@@ -26,8 +26,7 @@ use hyper_rustls::HttpsConnector;
 use hyper_timeout::TimeoutConnector;
 use hyper_util::client::legacy::connect::HttpConnector;
 use serde::{de::DeserializeOwned, Serialize};
-use tower::timeout::Timeout;
-use tower::Service;
+use tower::{timeout::Timeout, Service};
 use tower_http::{
     classify::{
         NeverClassifyEos, ServerErrorsAsFailures, ServerErrorsFailureClass, SharedClassifier,

src/clients/otel_grpc.rs

@@ -15,21 +15,23 @@
 
 */
 
-use crate::utils::trace::{current_trace_id, with_traceparent_header};
-use http::{Request, Response, StatusCode};
-use pin_project_lite::pin_project;
 use std::{
     error::Error,
     future::Future,
     pin::Pin,
     task::{Context, Poll},
 };
+
+use http::{Request, Response, StatusCode};
+use pin_project_lite::pin_project;
 use tokio::time::Instant;
 use tonic::client::GrpcService;
 use tower::Layer;
 use tracing::{error, info, info_span, Span};
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
+use crate::utils::trace::{current_trace_id, with_traceparent_header};
+
 // Adapted from https://github.com/davidB/tracing-opentelemetry-instrumentation-sdk/tree/main/tonic-tracing-opentelemetry
 /// Layer for grpc (tonic client):
 ///

src/models.rs

@@ -459,6 +459,9 @@ pub struct TokenClassificationResult {
     /// Aggregate label, if applicable
     pub entity_group: String,
 
+    /// id of detector (model) responsible for result(s)
+    pub detector_id: Option<String>,
+
     /// Confidence-like score of this classification prediction in [0, 1]
     pub score: f64,
 
@@ -894,6 +897,9 @@ pub struct DetectionResult {
     // The detection class
     pub detection: String,
 
+    // The id of the detector
+    pub detector_id: Option<String>,
+
     // The confidence level in the detection class
     pub score: f64,
 

src/orchestrator/streaming.rs

@@ -427,7 +427,8 @@ async fn detection_task(
                                             .into_iter()
                                             .flat_map(|r| {
                                                 r.into_iter().filter_map(|resp| {
-                                                    let result: TokenClassificationResult = resp.into();
+                                                    let mut result: TokenClassificationResult = resp.into();
+                                                    result.detector_id = Some(detector_id.clone()); // attach detector_id to the result
                                                     (result.score >= threshold).then_some(result)
                                                 })
                                             })

src/orchestrator/streaming/aggregator.rs

@@ -458,13 +458,15 @@ mod tests {
         text: &str,
         detection: &str,
         detection_type: &str,
+        detector_id: &str,
     ) -> TokenClassificationResult {
         TokenClassificationResult {
             start: span.0 as u32,
             end: span.1 as u32,
             word: text.to_string(),
             entity: detection.to_string(),
             entity_group: detection_type.to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.99,
             token_count: None,
         }
@@ -498,11 +500,11 @@ mod tests {
             let partial_span = (chunk_token.start + 2, chunk_token.end - 2);
 
             let (detector_tx1, detector_rx1) = mpsc::channel(1);
-            let detection = get_detection_obj(whole_span, text, "has_HAP", "HAP");
+            let detection = get_detection_obj(whole_span, text, "has_HAP", "HAP", "en-hap");
             let _ = detector_tx1.send((chunk.clone(), vec![detection])).await;
 
             let (detector_tx2, detector_rx2) = mpsc::channel(1);
-            let detection = get_detection_obj(partial_span, text, "email_ID", "PII");
+            let detection = get_detection_obj(partial_span, text, "email_ID", "PII", "en-pii");
             let _ = detector_tx2.send((chunk.clone(), vec![detection])).await;
 
             // Push HAP after PII to make sure detection ordering is not coincidental

src/orchestrator/streaming_content_detection.rs

@@ -27,27 +27,29 @@
 ///////////////////////////////////////////////////////////////////////////////////
 mod aggregator;
 
-use aggregator::Aggregator;
 use std::{collections::HashMap, pin::Pin, sync::Arc, time::Duration};
 
+use aggregator::Aggregator;
 use futures::{future::try_join_all, stream::Peekable, Stream, StreamExt, TryStreamExt};
 use hyper::HeaderMap;
 use tokio::sync::{broadcast, mpsc};
 use tokio_stream::wrappers::{BroadcastStream, ReceiverStream};
 use tracing::{debug, error, info, instrument, warn};
 
 use super::{streaming::Detections, Context, Error, Orchestrator, StreamingContentDetectionTask};
-use crate::clients::{
-    chunker::{tokenize_whole_doc_stream, ChunkerClient, DEFAULT_CHUNKER_ID},
-    detector::ContentAnalysisRequest,
-    TextContentsDetectorClient,
-};
-use crate::models::{
-    DetectorParams, StreamingContentDetectionRequest, StreamingContentDetectionResponse,
-    TokenClassificationResult,
+use crate::{
+    clients::{
+        chunker::{tokenize_whole_doc_stream, ChunkerClient, DEFAULT_CHUNKER_ID},
+        detector::ContentAnalysisRequest,
+        TextContentsDetectorClient,
+    },
+    models::{
+        DetectorParams, StreamingContentDetectionRequest, StreamingContentDetectionResponse,
+        TokenClassificationResult,
+    },
+    orchestrator::{get_chunker_ids, streaming::Chunk},
+    pb::caikit::runtime::chunkers,
 };
-use crate::orchestrator::{get_chunker_ids, streaming::Chunk};
-use crate::pb::caikit::runtime::chunkers;
 
 type ContentInputStream =
     Pin<Box<dyn Stream<Item = Result<StreamingContentDetectionRequest, Error>> + Send>>;
@@ -419,7 +421,8 @@ async fn detection_task(
                                             .into_iter()
                                             .flat_map(|r| {
                                                 r.into_iter().filter_map(|resp| {
-                                                    let result: TokenClassificationResult = resp.into();
+                                                    let mut result: TokenClassificationResult = resp.into();
+                                                    result.detector_id = Some(detector_id.clone());
                                                     (result.score >= threshold).then_some(result)
                                                 })
                                             })

src/orchestrator/streaming_content_detection/aggregator.rs

@@ -158,6 +158,7 @@ impl AggregationActor {
                             text: r.word,
                             detection: r.entity,
                             detection_type: r.entity_group,
+                            detector_id: r.detector_id,
                             score: r.score,
                             evidence: None,
                         })
@@ -206,13 +207,15 @@ mod tests {
         text: &str,
         detection: &str,
         detection_type: &str,
+        detector_id: &str,
     ) -> TokenClassificationResult {
         TokenClassificationResult {
             start: span.0 as u32,
             end: span.1 as u32,
             word: text.to_string(),
             entity: detection.to_string(),
             entity_group: detection_type.to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.99,
             token_count: None,
         }
@@ -246,11 +249,11 @@ mod tests {
             let partial_span = (chunk_token.start + 2, chunk_token.end - 2);
 
             let (detector_tx1, detector_rx1) = mpsc::channel(1);
-            let detection = get_detection_obj(whole_span, text, "has_HAP", "HAP");
+            let detection = get_detection_obj(whole_span, text, "has_HAP", "HAP", "en-hap");
             let _ = detector_tx1.send((chunk.clone(), vec![detection])).await;
 
             let (detector_tx2, detector_rx2) = mpsc::channel(1);
-            let detection = get_detection_obj(partial_span, text, "email_ID", "PII");
+            let detection = get_detection_obj(partial_span, text, "email_ID", "PII", "en-pii");
             let _ = detector_tx2.send((chunk.clone(), vec![detection])).await;
 
             // Push HAP after PII to make sure detection ordering is not coincidental

src/orchestrator/unary.rs

@@ -696,6 +696,7 @@ pub async fn detect(
                 .into_iter()
                 .filter_map(|resp| {
                     let mut result: TokenClassificationResult = resp.into();
+                    result.detector_id = Some(detector_id.clone()); // attach detector_id to the result, as ContentAnalsyisResponse does not contain one at this point
                     result.start += chunk.offset as u32;
                     result.end += chunk.offset as u32;
                     (result.score >= threshold).then_some(result)
@@ -756,6 +757,7 @@ pub async fn detect_content(
                 .filter_map(|mut resp| {
                     resp.start += chunk.offset;
                     resp.end += chunk.offset;
+                    resp.detector_id = Some(detector_id.clone()); // add detector_id
                     (resp.score >= threshold).then_some(resp)
                 })
                 .collect::<Vec<_>>()
@@ -803,6 +805,11 @@ pub async fn detect_for_generation(
             results
                 .into_iter()
                 .filter(|detection| detection.score > threshold)
+                .map(|mut detection| {
+                    //add detector_id
+                    detection.detector_id = Some(detector_id.clone());
+                    detection
+                })
                 .collect()
         })
         .map_err(|error| Error::DetectorRequestFailed {
@@ -844,6 +851,11 @@ pub async fn detect_for_chat(
             results
                 .into_iter()
                 .filter(|detection| detection.score > threshold)
+                .map(|mut detection| {
+                    //add detector_id
+                    detection.detector_id = Some(detector_id.clone());
+                    detection
+                })
                 .collect()
         })
         .map_err(|error| Error::DetectorRequestFailed {
@@ -899,6 +911,11 @@ pub async fn detect_for_context(
             results
                 .into_iter()
                 .filter(|detection| detection.score > threshold)
+                .map(|mut detection| {
+                    //add detector_id
+                    detection.detector_id = Some(detector_id.clone());
+                    detection
+                })
                 .collect()
         })
         .map_err(|error| Error::DetectorRequestFailed {
@@ -1131,6 +1148,7 @@ mod tests {
             word: second_sentence.clone(),
             entity: "has_HAP".to_string(),
             entity_group: "hap".to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.9,
             token_count: None,
         }];
@@ -1151,6 +1169,7 @@ mod tests {
                 text: first_sentence.clone(),
                 detection: "has_HAP".to_string(),
                 detection_type: "hap".to_string(),
+                detector_id: Some(detector_id.to_string()),
                 score: 0.1,
                 evidence: Some(vec![]),
             }],
@@ -1160,6 +1179,7 @@ mod tests {
                 text: second_sentence.clone(),
                 detection: "has_HAP".to_string(),
                 detection_type: "hap".to_string(),
+                detector_id: Some(detector_id.to_string()),
                 score: 0.9,
                 evidence: Some(vec![]),
             }],
@@ -1300,6 +1320,7 @@ mod tests {
         let expected_response: Vec<DetectionResult> = vec![DetectionResult {
             detection_type: "relevance".to_string(),
             detection: "is_relevant".to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.9,
             evidence: Some(
                 [EvidenceObj {
@@ -1325,6 +1346,7 @@ mod tests {
         .then_return(Ok(vec![DetectionResult {
             detection_type: "relevance".to_string(),
             detection: "is_relevant".to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.9,
             evidence: Some(
                 [EvidenceObj {
@@ -1393,6 +1415,7 @@ mod tests {
         .then_return(Ok(vec![DetectionResult {
             detection_type: "relevance".to_string(),
             detection: "is_relevant".to_string(),
+            detector_id: Some(detector_id.to_string()),
             score: 0.1,
             evidence: None,
         }]));

src/utils/tls.rs

@@ -1,3 +1,5 @@
+use std::{fs::File, io, path::PathBuf, sync::Arc};
+
 use http_serde::http::StatusCode;
 use hyper_rustls::ConfigBuilderExt;
 use rustls::{
@@ -6,7 +8,6 @@ use rustls::{
     ClientConfig, DigitallySignedStruct, SignatureScheme,
 };
 use serde::Deserialize;
-use std::{fs::File, io, path::PathBuf, sync::Arc};
 
 use crate::{clients, config::TlsConfig};
 

src/utils/trace.rs

@@ -40,8 +40,10 @@ use tracing::{error, info, info_span, Span};
 use tracing_opentelemetry::{MetricsLayer, OpenTelemetrySpanExt};
 use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Layer};
 
-use crate::args::{LogFormat, OtlpProtocol, TracingConfig};
-use crate::clients::http::TracedResponse;
+use crate::{
+    args::{LogFormat, OtlpProtocol, TracingConfig},
+    clients::http::TracedResponse,
+};
 
 #[derive(Debug, thiserror::Error)]
 pub enum TracingError {