Add chat content detector (input detection only) (foundation-model-stack#276)

* ✨ Add new datatype for internal conversion of chat content objects
* ♻️ Refactor some implementations and add request conversion
* ✨ Add detectors config in openai request
* 🧑‍💻 Add detection response and warning in openai chat completion response
* 🐛 Push detector processing module missed earlier
* ♻️ Make certain data models available public
* ♻️ Refactor chunking mechanism
* 🚧 Iterate over chunking for chat completion content detection
* 🐛 Fixed compiler error for iterating over chat messages
* 🚧 Add content detection function with datatype transformations
* 🚧 Add input detection handling to main function call
* 🐛 Fix validation error propagation for filtering messages
* 🐛 Fix error handling for chunker
* ⚰️ Remove print lines and deadcode
* 🐛 Remove orchestrator detection result tagged response
* 🐛 Fix detections and results typo
* 🔧 Make detector optional in request and add some error handling todos
* 🎨 Fix formatting and clippy warnings
* ♻️ Refactor part of detection logic for chat completion to allow proper message indexing
* :white_check_marks: Add unit test for message filtering
* 🚚 Rename filter_chat_messages function to preprocess
* 🔧 Fix dummy user id generation and add system to allowed message role
* 🔧 Fix according to review suggestions
* Apply suggestions from code review
Co-authored-by: Evaline Ju <[email protected]>
Co-authored-by: Dan Clark <[email protected]>
Signed-off-by: Gaurav Kumbhat <[email protected]>

* 🎨🚚 Rename and modify optional parameter as per review suggestions
* ⚡ Remove need to clone and perform inplace sorting
* ⚡✨ Address review suggestions and fixes
* Replace warning type from string to predefined type
* Add input detection error to warning
* Replace detection result sorting to previous method and remove results cloning
* Remove extra cloning for filter_chat_message function
* Detector message traversal optimizations
* Add empty message validation

* Apply suggestions from code review
Co-authored-by: Dan Clark <[email protected]>
Signed-off-by: Gaurav Kumbhat <[email protected]>

* 🎨 replace match statement with let-else to rustify
* Apply suggestions from code review
Co-authored-by: Dan Clark <[email protected]>
Signed-off-by: Gaurav Kumbhat <[email protected]>

* :white_check_marks: Fix unit test and rename filter message function
* ➕ Add uuid in cargo file
* ♻️ Refactor request object handling to avoid cloning

---------

Signed-off-by: gkumbhat <[email protected]>
Signed-off-by: Gaurav Kumbhat <[email protected]>
Co-authored-by: Evaline Ju <[email protected]>
Co-authored-by: Dan Clark <[email protected]>

Author: 3 people

Cargo.lock

@@ -55,6 +55,7 @@ tracing = "0.1.41"
 tracing-opentelemetry = "0.28.0"
 tracing-subscriber = { version = "0.3.19", features = ["json", "env-filter"] }
 url = "2.5.4"
+uuid = { version = "1.10.0", features = ["v4"] }
 
 [build-dependencies]
 tonic-build = "0.12.3"

Cargo.toml

@@ -27,7 +27,11 @@ use tokio::sync::mpsc;
 use tracing::{info, instrument};
 
 use super::{create_http_client, http::HttpClientExt, Client, Error, HttpClient};
-use crate::{config::ServiceConfig, health::HealthCheckResult};
+use crate::{
+    config::ServiceConfig,
+    health::HealthCheckResult,
+    models::{DetectorParams, GuardrailDetection, InputWarningReason},
+};
 
 const DEFAULT_PORT: u16 = 8080;
 
@@ -157,13 +161,13 @@ impl HttpClientExt for OpenAiClient {
 
 #[derive(Debug)]
 pub enum ChatCompletionsResponse {
-    Unary(ChatCompletion),
+    Unary(Box<ChatCompletion>),
     Streaming(mpsc::Receiver<Result<Option<ChatCompletionChunk>, Error>>),
 }
 
 impl From<ChatCompletion> for ChatCompletionsResponse {
     fn from(value: ChatCompletion) -> Self {
-        Self::Unary(value)
+        Self::Unary(Box::new(value))
     }
 }
 
@@ -275,6 +279,23 @@ pub struct ChatCompletionsRequest {
     pub skip_special_tokens: Option<bool>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub spaces_between_special_tokens: Option<bool>,
+
+    // Detectors
+    // Note: We are making it optional, since this structure also gets used to
+    // form request for chat completions. And downstream server, might choose to
+    // reject extra parameters.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub detectors: Option<DetectorConfig>,
+}
+
+/// Structure to contain parameters for detectors.
+#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct DetectorConfig {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub input: Option<HashMap<String, DetectorParams>>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub output: Option<HashMap<String, DetectorParams>>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -386,7 +407,7 @@ pub struct Message {
     pub tool_call_id: Option<String>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 #[serde(untagged)]
 pub enum Content {
     /// The text contents of the message.
@@ -430,7 +451,7 @@ impl From<Vec<String>> for Content {
     }
 }
 
-#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
 pub enum ContentType {
     #[serde(rename = "text")]
     #[default]
@@ -439,7 +460,7 @@ pub enum ContentType {
     ImageUrl,
 }
 
-#[derive(Debug, Default, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize, PartialEq)]
 pub struct ContentPart {
     /// The type of the content part.
     #[serde(rename = "type")]
@@ -455,7 +476,7 @@ pub struct ContentPart {
     pub refusal: Option<String>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
 pub struct ImageUrl {
     /// Either a URL of the image or the base64 encoded image data.
     pub url: String,
@@ -485,7 +506,7 @@ pub struct Function {
 }
 
 /// Represents a chat completion response returned by model, based on the provided input.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct ChatCompletion {
     /// A unique identifier for the chat completion.
     pub id: String,
@@ -506,6 +527,12 @@ pub struct ChatCompletion {
     /// This field is only included if the `service_tier` parameter is specified in the request.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub service_tier: Option<String>,
+    /// Result of running different guardrail detectors
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub detections: Option<ChatDetections>,
+    /// Optional warnings
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub warnings: Vec<OrchestratorWarning>,
 }
 
 /// A chat completion choice.
@@ -621,7 +648,7 @@ pub struct ChatCompletionDelta {
 }
 
 /// Usage statistics for a completion.
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Serialize, Deserialize)]
 pub struct Usage {
     /// Number of tokens in the prompt.
     pub prompt_tokens: u32,
@@ -665,3 +692,44 @@ pub struct OpenAiError {
     pub param: Option<String>,
     pub code: u16,
 }
+
+/// Guardrails detection results.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ChatDetections {
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub input: Vec<InputDetectionResult>,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub output: Vec<OutputDetectionResult>,
+}
+
+/// Guardrails detection result for application on input.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct InputDetectionResult {
+    pub message_index: usize,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub results: Vec<GuardrailDetection>,
+}
+
+/// Guardrails detection result for application output.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OutputDetectionResult {
+    choice_index: usize,
+    #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    results: Vec<GuardrailDetection>,
+}
+
+/// Warnings generated by guardrails.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct OrchestratorWarning {
+    r#type: InputWarningReason,
+    message: String,
+}
+
+impl OrchestratorWarning {
+    pub fn new(warning_type: InputWarningReason, message: &str) -> Self {
+        Self {
+            r#type: warning_type,
+            message: message.to_string(),
+        }
+    }
+}

src/clients/openai.rs

@@ -908,6 +908,13 @@ pub struct DetectionResult {
     pub evidence: Option<Vec<EvidenceObj>>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(untagged)]
+pub enum GuardrailDetection {
+    ContentAnalysisResponse(ContentAnalysisResponse),
+    ClassificationResult(DetectionResult),
+}
+
 /// The request format expected in the /api/v2/text/context endpoint.
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ContextDocsHttpRequest {

src/models.rs

@@ -19,6 +19,7 @@ pub mod errors;
 pub use errors::Error;
 use futures::Stream;
 pub mod chat_completions_detection;
+pub mod detector_processing;
 pub mod streaming;
 pub mod streaming_content_detection;
 pub mod unary;