Add support for json based content types in Message (#480)

crates/brightstaff/src/router/llm_router.rs

@@ -1,14 +1,13 @@
 use std::sync::Arc;
 
 use common::{
-    api::open_ai::{ChatCompletionsResponse, Message},
+    api::open_ai::{ChatCompletionsResponse, ContentType, Message},
     configuration::LlmProvider,
     consts::ARCH_PROVIDER_HINT_HEADER,
-    utils::shorten_string,
 };
 use hyper::header;
 use thiserror::Error;
-use tracing::{info, warn};
+use tracing::{debug, info, warn};
 
 use super::router_model::RouterModel;
 
@@ -59,9 +58,9 @@ impl RouterService {
             .collect::<Vec<String>>()
             .join("\n");
 
-        info!(
+        debug!(
             "llm_providers from config with usage: {}...",
-            shorten_string(&llm_providers_with_usage_yaml.replace("\n", "\\n"))
+            llm_providers_with_usage_yaml.replace("\n", "\\n")
         );
 
         let router_model = Arc::new(super::router_model_v1::RouterModelV1::new(
@@ -83,7 +82,6 @@ impl RouterService {
         messages: &[Message],
         trace_parent: Option<String>,
     ) -> Result<Option<String>> {
-
         if !self.llm_usage_defined {
             return Ok(None);
         }
@@ -91,8 +89,14 @@ impl RouterService {
         let router_request = self.router_model.generate_request(messages);
 
         info!(
-            "router_request: {}",
-            shorten_string(&serde_json::to_string(&router_request).unwrap()),
+            "sending request to arch-router model: {}, endpoint: {}",
+            self.router_model.get_model_name(),
+            self.router_url
+        );
+
+        debug!(
+            "arch request body: {}",
+            &serde_json::to_string(&router_request).unwrap(),
         );
 
         let mut llm_route_request_headers = header::HeaderMap::new();
@@ -113,6 +117,7 @@ impl RouterService {
             );
         }
 
+        let start_time = std::time::Instant::now();
         let res = self
             .client
             .post(&self.router_url)
@@ -122,6 +127,7 @@ impl RouterService {
             .await?;
 
         let body = res.text().await?;
+        let router_response_time = start_time.elapsed();
 
         let chat_completion_response: ChatCompletionsResponse = match serde_json::from_str(&body) {
             Ok(response) => response,
@@ -138,14 +144,18 @@ impl RouterService {
             }
         };
 
-        let selected_llm = self.router_model.parse_response(
-            chat_completion_response.choices[0]
-                .message
-                .content
-                .as_ref()
-                .unwrap(),
-        )?;
-
-        Ok(selected_llm)
+        if let Some(ContentType::Text(content)) =
+            &chat_completion_response.choices[0].message.content
+        {
+            info!(
+                "router response: {}, response time: {}ms",
+                content.replace("\n", "\\n"),
+                router_response_time.as_millis()
+            );
+            let selected_llm = self.router_model.parse_response(content)?;
+            Ok(selected_llm)
+        } else {
+            Ok(None)
+        }
     }
 }

crates/brightstaff/src/router/router_model.rs

@@ -12,4 +12,5 @@ pub type Result<T> = std::result::Result<T, RoutingModelError>;
 pub trait RouterModel: Send + Sync {
     fn generate_request(&self, messages: &[Message]) -> ChatCompletionsRequest;
     fn parse_response(&self, content: &str) -> Result<Option<String>>;
+    fn get_model_name(&self) -> String;
 }

crates/brightstaff/src/router/router_model_v1.rs

@@ -1,9 +1,8 @@
 use common::{
-    api::open_ai::{ChatCompletionsRequest, Message},
+    api::open_ai::{ChatCompletionsRequest, ContentType, Message},
     consts::{SYSTEM_ROLE, USER_ROLE},
 };
 use serde::{Deserialize, Serialize};
-use tracing::info;
 
 use super::router_model::{RouterModel, RoutingModelError};
 
@@ -68,7 +67,7 @@ impl RouterModel for RouterModelV1 {
         ChatCompletionsRequest {
             model: self.routing_model.clone(),
             messages: vec![Message {
-                content: Some(message),
+                content: Some(ContentType::Text(message)),
                 role: USER_ROLE.to_string(),
                 model: None,
                 tool_calls: None,
@@ -86,10 +85,6 @@ impl RouterModel for RouterModelV1 {
             return Ok(None);
         }
         let router_resp_fixed = fix_json_response(content);
-        info!(
-            "router response (fixed): {}",
-            router_resp_fixed.replace("\n", "\\n")
-        );
         let router_response: LlmRouterResponse = serde_json::from_str(router_resp_fixed.as_str())?;
 
         let selected_llm = router_response.route.unwrap_or_default().to_string();
@@ -100,6 +95,10 @@ impl RouterModel for RouterModelV1 {
 
         Ok(Some(selected_llm))
     }
+
+    fn get_model_name(&self) -> String {
+        self.routing_model.clone()
+    }
 }
 
 fn fix_json_response(body: &str) -> String {
@@ -172,22 +171,28 @@ user: "seattle"
         let messages = vec![
             Message {
                 role: "system".to_string(),
-                content: Some("You are a helpful assistant.".to_string()),
+                content: Some(ContentType::Text(
+                    "You are a helpful assistant.".to_string(),
+                )),
                 ..Default::default()
             },
             Message {
                 role: "user".to_string(),
-                content: Some("Hello, I want to book a flight.".to_string()),
+                content: Some(ContentType::Text(
+                    "Hello, I want to book a flight.".to_string(),
+                )),
                 ..Default::default()
             },
             Message {
                 role: "assistant".to_string(),
-                content: Some("Sure, where would you like to go?".to_string()),
+                content: Some(ContentType::Text(
+                    "Sure, where would you like to go?".to_string(),
+                )),
                 ..Default::default()
             },
             Message {
                 role: "user".to_string(),
-                content: Some("seattle".to_string()),
+                content: Some(ContentType::Text("seattle".to_string())),
                 ..Default::default()
             },
         ];
@@ -198,7 +203,7 @@ user: "seattle"
 
         println!("Prompt: {}", prompt);
 
-        assert_eq!(expected_prompt, prompt);
+        assert_eq!(expected_prompt, prompt.to_string());
     }
 }