✨ Update to /text/contents API for detector (#62)

* 🔧 Use text contents endpoint

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 🚧 WIP detector API update

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 💡🏷️ Add comments for new contents API types

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 🚧 Updated content analysis handling

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 🚧 Remove clone on chunks

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 💡 Add comments

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 🚧 Update by reference

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* ♻️ Use from

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* 💡 Add TODO comment

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

* Update handle_detection_task and index_codepoints, rename DetectorClient method

---------

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>
Co-authored-by: declark1 <daniel.clark@ibm.com>

Author: evaline-ju

config/config.yaml

@@ -13,7 +13,7 @@ chunkers:
 detectors:
     hap-en:
         service:
-            hostname: https://localhost/api/v1/detector
+            hostname: https://localhost/api/v1/text/contents # full url / endpoint currently expected
             port: 8080
             tls: caikit
         chunker_id: en_regex

src/clients/detector.rs

@@ -28,11 +28,11 @@ impl DetectorClient {
             .clone())
     }
 
-    pub async fn classify(
+    pub async fn text_contents(
         &self,
         model_id: &str,
-        request: DetectorRequest,
-    ) -> Result<DetectorResponse, Error> {
+        request: ContentAnalysisRequest,
+    ) -> Result<Vec<Vec<ContentAnalysisResponse>>, Error> {
         let client = self.client(model_id)?;
         let url = client.base_url().as_str();
         let response = client
@@ -47,39 +47,71 @@ impl DetectorClient {
     }
 }
 
-#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
-pub struct DetectorRequest {
-    pub text: String,
-    pub parameters: HashMap<String, serde_json::Value>,
+/// Request for text content analysis
+/// Results of this request will contain analysis / detection of each of the provided documents
+/// in the order they are present in the `contents` object.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ContentAnalysisRequest {
+    /// Field allowing users to provide list of documents for analysis
+    pub contents: Vec<String>,
 }
 
-impl DetectorRequest {
-    pub fn new(text: String, parameters: HashMap<String, serde_json::Value>) -> Self {
-        Self { text, parameters }
+impl ContentAnalysisRequest {
+    pub fn new(contents: Vec<String>) -> ContentAnalysisRequest {
+        ContentAnalysisRequest { contents }
     }
 }
 
-#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
-pub struct Detection {
+/// Evidence type
+#[derive(Clone, Copy, Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum EvidenceType {
+    Url,
+    Title,
+}
+
+/// Source of the evidence e.g. url
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct Evidence {
+    /// Evidence source
+    pub source: String,
+}
+
+/// Evidence in response
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct EvidenceObj {
+    /// Type field signifying the type of evidence provided
+    #[serde(rename = "type")]
+    pub r#type: EvidenceType,
+    /// Evidence currently only containing source
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub evidence: Option<Evidence>,
+}
+
+/// Response of text content analysis endpoint
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ContentAnalysisResponse {
+    /// Start index of detection
     pub start: usize,
+    /// End index of detection
     pub end: usize,
-    pub text: String,
+    /// Relevant detection class
     pub detection: String,
+    /// Detection type or aggregate detection label
     pub detection_type: String,
+    /// Score of detection
     pub score: f64,
+    /// Optional, any applicable evidences for detection
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub evidences: Option<Vec<EvidenceObj>>,
 }
 
-#[derive(Clone, Default, Debug, PartialEq, Serialize, Deserialize)]
-pub struct DetectorResponse {
-    pub detections: Vec<Detection>,
-}
-
-impl From<Detection> for crate::models::TokenClassificationResult {
-    fn from(value: Detection) -> Self {
+impl From<ContentAnalysisResponse> for crate::models::TokenClassificationResult {
+    fn from(value: ContentAnalysisResponse) -> Self {
         Self {
             start: value.start as u32,
             end: value.end as u32,
-            word: value.text,
+            word: "".to_string(), // TODO: fill in when provided in the detector API in the next iteration
             entity: value.detection,
             entity_group: value.detection_type,
             score: value.score,

src/orchestrator.rs

@@ -12,7 +12,7 @@ use uuid::Uuid;
 
 use crate::{
     clients::{
-        self, detector::DetectorRequest, ChunkerClient, DetectorClient, GenerationClient,
+        self, detector::ContentAnalysisRequest, ChunkerClient, DetectorClient, GenerationClient,
         NlpClient, TgisClient,
     },
     config::{GenerationProvider, OrchestratorConfig},
@@ -198,7 +198,7 @@ async fn chunk_and_detect(
             Ok::<String, Error>(chunker_id)
         })
         .collect::<Result<Vec<_>, Error>>()?;
-    // Spawn chunking tasks, returning a map of chunker_id->chunks.
+    // Spawn chunking tasks, returning a map of chunker_id->chunks
     let chunks = chunk(ctx.clone(), chunker_ids, text_with_offsets).await?;
     // Spawn detection tasks
     let detections = detect(ctx.clone(), detectors, chunks).await?;
@@ -319,69 +319,57 @@ async fn handle_chunk_task(
     Ok((chunker_id, chunks))
 }
 
-/// Sends a buffered, concurrent stream of requests to a detector service.
+/// Sends a request to a detector service.
 async fn handle_detection_task(
     ctx: Arc<Context>,
     detector_id: String,
     default_threshold: f32,
     detector_params: DetectorParams,
     chunks: Vec<Chunk>,
 ) -> Result<Vec<TokenClassificationResult>, Error> {
-    let detections = stream::iter(chunks)
-        .map(|chunk| {
-            let ctx = ctx.clone();
-            let detector_id = detector_id.clone();
-            let detector_params = detector_params.clone();
-            async move {
-                // NOTE: The detector request is expected to change and not actually
-                // take parameters. Any parameters will be ignored for now
-                // ref. https://github.com/foundation-model-stack/fms-guardrails-orchestrator/issues/37
-                let request = DetectorRequest::new(chunk.text.clone(), detector_params.clone());
-                debug!(
-                    %detector_id,
-                    ?request,
-                    "sending detector request"
-                );
-                let response = ctx
-                    .detector_client
-                    .classify(&detector_id, request)
-                    .await
-                    .map_err(|error| Error::DetectorRequestFailed {
-                        detector_id: detector_id.clone(),
-                        error,
-                    })?;
-                debug!(
-                    %detector_id,
-                    ?response,
-                    "received detector response"
-                );
-                // Filter results based on threshold (if applicable) here
-                let results = response
-                    .detections
-                    .into_iter()
-                    .filter_map(|detection| {
-                        let mut result: TokenClassificationResult = detection.into();
-                        result.start += chunk.offset as u32;
-                        result.end += chunk.offset as u32;
-                        let threshold = detector_params
-                            .get("threshold")
-                            .and_then(|v| v.as_f64())
-                            .unwrap_or(default_threshold as f64);
-                        (result.score >= threshold).then_some(result)
-                    })
-                    .collect::<Vec<_>>();
-                Ok::<Vec<TokenClassificationResult>, Error>(results)
-            }
-        })
-        .buffered(5)
-        .collect::<Vec<_>>()
+    let detector_id = detector_id.clone();
+    let threshold = detector_params
+        .get("threshold")
+        .and_then(|v| v.as_f64())
+        .unwrap_or(default_threshold as f64);
+    let contents = chunks.iter().map(|chunk| chunk.text.clone()).collect();
+    let request = ContentAnalysisRequest::new(contents);
+    debug!(
+        %detector_id,
+        ?request,
+        "sending detector request"
+    );
+    let response = ctx
+        .detector_client
+        .text_contents(&detector_id, request)
         .await
+        .map_err(|error| Error::DetectorRequestFailed {
+            detector_id: detector_id.clone(),
+            error,
+        })?;
+    debug!(
+        %detector_id,
+        ?response,
+        "received detector response"
+    );
+    let results = chunks
         .into_iter()
-        .collect::<Result<Vec<_>, Error>>()?
-        .into_iter()
-        .flatten()
+        .zip(response)
+        .flat_map(|(chunk, response)| {
+            response
+                .into_iter()
+                .filter_map(|resp| {
+                    let mut result: TokenClassificationResult = resp.into();
+                    result.word =
+                        index_codepoints(&chunk.text, result.start as usize, result.end as usize);
+                    result.start += chunk.offset as u32;
+                    result.end += chunk.offset as u32;
+                    (result.score >= threshold).then_some(result)
+                })
+                .collect::<Vec<_>>()
+        })
         .collect::<Vec<_>>();
-    Ok(detections)
+    Ok::<Vec<TokenClassificationResult>, Error>(results)
 }
 
 /// Sends tokenize request to a generation service.
@@ -591,6 +579,12 @@ async fn generate(
     }
 }
 
+/// Get codepoints of text between start and end indices
+fn index_codepoints(text: &str, start: usize, end: usize) -> String {
+    let chars = text.chars().collect::<Vec<_>>();
+    chars[start..end].iter().collect()
+}
+
 /// Applies masks to input text, returning (offset, masked_text) pairs.
 fn apply_masks(text: &str, masks: &[(usize, usize)]) -> Vec<(usize, String)> {
     let chars = text.chars().collect::<Vec<_>>();
@@ -708,4 +702,12 @@ mod tests {
         ];
         assert_eq!(text_with_offsets, expected_text_with_offsets)
     }
+
+    #[test]
+    fn test_index_codepoints() {
+        let s = "Hello world";
+        assert_eq!(index_codepoints(s, 0, 5), "Hello");
+        let s = "哈囉世界";
+        assert_eq!(index_codepoints(s, 3, 4), "界");
+    }
 }