add fallback message support

Author: RobGeada

config/config.yaml

@@ -1,8 +1,7 @@
 orchestrator:
-  host: localhost
-  port: 8085
+  host: guardrails-orchestrator-http-model-namespace.apps.rosa.trustyai-rob2.n1ai.p3.openshiftapps.com
 detectors:
-  - name: regex
+  - name: regex_language
     detector_params:
       regex:
         - email
@@ -11,6 +10,7 @@ detectors:
 routes:
   - name: pii
     detectors:
-      - regex
+      - regex_language
+    fallback_message: "I'm sorry Dave, I'm afraid I can't do that."
   - name: passthrough
     detectors:

src/api.rs

@@ -0,0 +1,72 @@
+use std::collections::HashMap;
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Serialize)]
+pub(crate) struct OrchestratorDetector {
+    pub(crate) input: HashMap<String, serde_json::Value>,
+    pub(crate) output: HashMap<String, serde_json::Value>,
+    // implement when output is completed, also need to see about splitting detectors in config to input/output
+    // output: HashMap<String, serde_json::Value>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub(crate) struct GenerationMessage {
+    pub(crate) content: String,
+    pub(crate) refusal: Option<String>,
+    pub(crate) role: String,
+    pub(crate) tool_calls: Option<serde_json::Value>,
+    pub(crate) audio: Option<serde_json::Value>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub(crate) struct GenerationChoice {
+    pub(crate) finish_reason: String,
+    pub(crate) index: u32,
+    pub(crate) message: GenerationMessage,
+    pub(crate) logprobs: Option<serde_json::Value>,
+}
+
+
+#[derive(Serialize, Deserialize, Debug)]
+struct DetectionResult {
+    start: serde_json::Value,
+    end: u32,
+    text: String,
+    detection_type: String,
+    detection: String,
+    detector_id: String,
+    score: f64,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct InputDetection {
+    message_index: u16,
+    results: Option<Vec<DetectionResult>>
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct OutputDetection {
+    choice_index: u32,
+    results: Option<Vec<DetectionResult>>
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub(crate) struct Detections {
+    input: Option<Vec<InputDetection>>,
+    output: Option<Vec<OutputDetection>>,
+}
+
+
+#[derive(Serialize, Deserialize, Debug)]
+pub(crate) struct OrchestratorResponse {
+    id: String,
+    pub(crate) choices: Vec<GenerationChoice>,
+    created: u64,
+    model: String,
+    service_tier: Option<String>,
+    system_fingerprint: Option<String>,
+    object: Option<String>,
+    usage: serde_json::Value,
+    pub(crate) detections: Option<Detections>,
+    pub(crate) warnings: Option<Vec<HashMap<String, String>>>
+}

src/config.rs

@@ -12,19 +12,22 @@ pub struct GatewayConfig {
 #[derive(Debug, Deserialize, Clone)]
 pub struct OrchestratorConfig {
     pub host: String,
-    pub port: u16,
+    pub port: Option<u16>,
 }
 
 #[derive(Debug, Deserialize, Clone)]
 pub struct DetectorConfig {
     pub name: String,
+    pub input: bool,
+    pub output: bool,
     pub detector_params: Option<serde_json::Value>,
 }
 
 #[derive(Debug, Deserialize, Clone)]
 pub struct RouteConfig {
     pub name: String,
     pub detectors: Vec<String>,
+    pub fallback_message: Option<String>,
 }
 
 pub fn read_config(path: &str) -> GatewayConfig {
@@ -62,7 +65,7 @@ mod tests {
         let gc = GatewayConfig {
             orchestrator: OrchestratorConfig {
                 host: "localhost".to_string(),
-                port: 1234,
+                port: Some(1234),
             },
             detectors: vec![DetectorConfig {
                 name: "regex".to_string(),

src/main.rs

@@ -3,38 +3,42 @@ use std::{
     env,
     net::{IpAddr, SocketAddr},
 };
-
+use std::ptr::null;
 use axum::{http::StatusCode, response::IntoResponse, routing::post, Json, Router};
 use config::{validate_registered_detectors, DetectorConfig, GatewayConfig};
-use serde::Serialize;
 use serde_json::json;
 use serde_json::{Map, Value};
 use tower_http::trace::{self, TraceLayer};
 use tracing::Level;
+use crate::api::{GenerationChoice, GenerationMessage, OrchestratorResponse};
+use crate::config::RouteConfig;
 
 mod config;
+mod api;
 
-#[derive(Debug, Serialize)]
-struct OrchestratorDetector {
-    input: HashMap<String, serde_json::Value>,
-    // implement when output is completed, also need to see about splitting detectors in config to input/output
-    // output: HashMap<String, serde_json::Value>,
-}
 
 fn get_orchestrator_detectors(
     detectors: Vec<String>,
     detector_config: Vec<DetectorConfig>,
-) -> OrchestratorDetector {
+) -> api::OrchestratorDetector {
     let mut input_detectors = HashMap::new();
+    let mut output_detectors = HashMap::new();
 
     for detector in detector_config {
         if detectors.contains(&detector.name) && detector.detector_params.is_some() {
-            input_detectors.insert(detector.name, detector.detector_params.unwrap());
+            let detector_params = detector.detector_params.unwrap();
+            if detector.input {
+                input_detectors.insert(detector.name.clone(), detector_params.clone());
+            }
+            if detector.output {
+                output_detectors.insert(detector.name, detector_params);
+            }
         }
     }
 
-    OrchestratorDetector {
+    api::OrchestratorDetector {
         input: input_detectors,
+        output: output_detectors,
     }
 }
 
@@ -59,10 +63,11 @@ async fn main() {
         let gateway_config = gateway_config.clone();
         let detectors = route.detectors.clone();
         let path = format!("/{}/v1/chat/completions", route.name);
+        let fallback_message = route.fallback_message.clone();
         app = app.route(
             &path,
             post(|Json(payload): Json<serde_json::Value>| {
-                handle_generation(Json(payload), detectors, gateway_config)
+                handle_generation(Json(payload), detectors, gateway_config, fallback_message)
             }),
         );
         tracing::info!("exposed endpoints: {}", path);
@@ -86,27 +91,66 @@ async fn main() {
     axum::serve(listener, app).await.unwrap();
 }
 
+
+async fn handle_orchestrator_payload_parsing(orchestrator_response: &mut OrchestratorResponse, route_fallback_message: Option<String>) {
+    if route_fallback_message.is_some() && orchestrator_response.detections.is_some() {
+        let fallback_generation = GenerationMessage {
+            content: route_fallback_message.clone().unwrap(),
+            refusal: None,
+            role: String::from("assistant"),
+            tool_calls: None,
+            audio: None,
+        };
+
+        orchestrator_response.choices = vec![GenerationChoice {
+            message: fallback_generation,
+            finish_reason: String::from("stop"),
+            index: 0,
+            logprobs: None
+        }];
+    }
+}
+
 async fn handle_generation(
     Json(mut payload): Json<serde_json::Value>,
     detectors: Vec<String>,
     gateway_config: GatewayConfig,
+    route_fallback_message: Option<String>,
 ) -> Result<impl IntoResponse, (StatusCode, String)> {
     let orchestrator_detectors = get_orchestrator_detectors(detectors, gateway_config.detectors);
 
     let mut payload = payload.as_object_mut();
 
-    let url = format!(
-        "http://{}:{}/api/v2/chat/completions-detection",
-        gateway_config.orchestrator.host, gateway_config.orchestrator.port
-    );
+    let url;
+    if gateway_config.orchestrator.port.is_some() {
+        url = format!(
+            "http://{}:{}/api/v2/chat/completions-detection",
+            gateway_config.orchestrator.host, gateway_config.orchestrator.port.unwrap()
+        );
+    } else {
+        url = format!(
+            "https://{}/api/v2/chat/completions-detection",
+            gateway_config.orchestrator.host
+        );
+    }
+
 
     payload.as_mut().unwrap().insert(
         "detectors".to_string(),
         serde_json::to_value(&orchestrator_detectors).unwrap(),
     );
     let response_payload = orchestrator_post_request(payload, &url).await;
-
-    Ok(Json(json!(response_payload.unwrap())).into_response())
+    if response_payload.is_ok() {
+        let response_unwrapped = response_payload.unwrap();
+        println!("{}", response_unwrapped);
+
+        let mut response : OrchestratorResponse = serde_json::from_value(response_unwrapped).unwrap();
+        handle_orchestrator_payload_parsing(&mut response, route_fallback_message).await;
+        Ok(Json(json!(response)).into_response())
+    } else {
+        //println!("{:#?}", response_payload.err().unwrap().to_string());
+        Err((StatusCode::INTERNAL_SERVER_ERROR, response_payload.err().unwrap().to_string()))
+    }
 }
 
 async fn orchestrator_post_request(
@@ -117,6 +161,5 @@ async fn orchestrator_post_request(
     let response = client.post(url).json(&payload).send();
 
     let json = response.await?.json().await?;
-    println!("{:#?}", json);
     Ok(json)
 }