From 413f36a9185063f02fa4a9247df408aaad92f2e2 Mon Sep 17 00:00:00 2001
From: Adil Hafeez <adil.hafeez@gmail.com>
Date: Thu, 22 May 2025 23:44:53 -0700
Subject: [PATCH] refactor a bit

---
 crates/common/src/api/hallucination.rs    | 12 ++++++------
 crates/prompt_gateway/src/http_context.rs |  5 ++---
 2 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/crates/common/src/api/hallucination.rs b/crates/common/src/api/hallucination.rs
index 44e252fa..a7caba67 100644
--- a/crates/common/src/api/hallucination.rs
+++ b/crates/common/src/api/hallucination.rs
@@ -36,22 +36,22 @@ pub fn extract_messages_for_hallucination(messages: &[Message]) -> Vec<String> {
         for message in messages.iter().rev() {
             if let Some(model) = message.model.as_ref() {
                 if !model.starts_with(ARCH_MODEL_PREFIX) {
-                    if let Some(ContentType::Text(content)) = &message.content {
-                        if !content.starts_with(HALLUCINATION_TEMPLATE) {
+                    if let Some(content) = &message.content {
+                        if !content.to_string().starts_with(HALLUCINATION_TEMPLATE) {
                             break;
                         }
                     }
                 }
             }
             if message.role == USER_ROLE {
-                if let Some(ContentType::Text(content)) = &message.content {
-                    user_messages.push(content.clone());
+                if let Some(content) = &message.content {
+                    user_messages.push(content.to_string());
                 }
             }
         }
     } else if let Some(message) = messages.last() {
-        if let Some(ContentType::Text(content)) = &message.content {
-            user_messages.push(content.clone());
+        if let Some(content) = &message.content {
+            user_messages.push(content.to_string());
         }
     }
     user_messages.reverse(); // Reverse to maintain the original order
diff --git a/crates/prompt_gateway/src/http_context.rs b/crates/prompt_gateway/src/http_context.rs
index d1c6670f..bb673208 100644
--- a/crates/prompt_gateway/src/http_context.rs
+++ b/crates/prompt_gateway/src/http_context.rs
@@ -2,7 +2,6 @@ use crate::stream_context::{ResponseHandlerType, StreamCallContext, StreamContex
 use common::{
     api::open_ai::{
         self, ArchState, ChatCompletionStreamResponse, ChatCompletionTool, ChatCompletionsRequest,
-        ContentType,
     },
     consts::{
         ARCH_FC_MODEL_NAME, ARCH_INTERNAL_CLUSTER_NAME, ARCH_ROUTING_HEADER,
@@ -238,12 +237,12 @@ impl HttpContext for StreamContext {
             Duration::from_secs(5),
         );
 
-        if let Some(ContentType::Text(content)) =
+        if let Some(content) =
             self.user_prompt.as_ref().unwrap().content.as_ref()
         {
             let call_context = StreamCallContext {
                 response_handler_type: ResponseHandlerType::ArchFC,
-                user_message: Some(content.clone()),
+                user_message: Some(content.to_string()),
                 prompt_target_name: None,
                 request_body: self.chat_completions_request.as_ref().unwrap().clone(),
                 similarity_scores: None,