From afa8ef7a635fadc570ce945114563c410265a4dd Mon Sep 17 00:00:00 2001
From: cotran <cotran2@utexas.edu>
Date: Wed, 23 Oct 2024 17:15:37 -0700
Subject: [PATCH] fix

---
 crates/common/src/consts.rs                 |  1 +
 crates/prompt_gateway/src/hallucination.rs  | 10 ++++++----
 crates/prompt_gateway/src/stream_context.rs |  7 +++----
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/crates/common/src/consts.rs b/crates/common/src/consts.rs
index 8f59e981..533b8a69 100644
--- a/crates/common/src/consts.rs
+++ b/crates/common/src/consts.rs
@@ -25,3 +25,4 @@ pub const ARCH_INTERNAL_CLUSTER_NAME: &str = "arch_internal";
 pub const ARCH_UPSTREAM_HOST_HEADER: &str = "x-arch-upstream";
 pub const ARCH_LLM_UPSTREAM_LISTENER: &str = "arch_llm_listener";
 pub const ARCH_MODEL_PREFIX: &str = "Arch";
+pub const HALLUCINATION_TEMPLATE: &str = "It seems I’m missing some information. Could you provide the following details";
diff --git a/crates/prompt_gateway/src/hallucination.rs b/crates/prompt_gateway/src/hallucination.rs
index 62b119ac..c4425957 100644
--- a/crates/prompt_gateway/src/hallucination.rs
+++ b/crates/prompt_gateway/src/hallucination.rs
@@ -1,6 +1,6 @@
 use common::{
     common_types::open_ai::Message,
-    consts::{ARCH_MODEL_PREFIX, ASSISTANT_ROLE, USER_ROLE},
+    consts::{ARCH_MODEL_PREFIX, USER_ROLE, HALLUCINATION_TEMPLATE},
 };
 
 pub fn extract_messages_for_hallucination(messages: &Vec<Message>) -> Vec<String> {
@@ -18,9 +18,11 @@ pub fn extract_messages_for_hallucination(messages: &Vec<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 message.role == ASSISTANT_ROLE {
-                        break;
-                    }
+                  if let Some(content) = &message.content {
+                      if !content.starts_with(HALLUCINATION_TEMPLATE) {
+                          break;
+                      }
+                  }
                 }
             }
             if message.role == USER_ROLE {
diff --git a/crates/prompt_gateway/src/stream_context.rs b/crates/prompt_gateway/src/stream_context.rs
index 1a2665c9..f9b217d3 100644
--- a/crates/prompt_gateway/src/stream_context.rs
+++ b/crates/prompt_gateway/src/stream_context.rs
@@ -12,7 +12,7 @@ use common::common_types::{
 };
 use common::configuration::{Overrides, PromptGuards, PromptTarget};
 use common::consts::{
-    ARCH_FC_INTERNAL_HOST, ARCH_FC_MODEL_NAME, ARCH_FC_REQUEST_TIMEOUT_MS, ARCH_INTERNAL_CLUSTER_NAME, ARCH_MESSAGES_KEY, ARCH_MODEL_PREFIX, ARCH_STATE_HEADER, ARCH_UPSTREAM_HOST_HEADER, ASSISTANT_ROLE, DEFAULT_EMBEDDING_MODEL, DEFAULT_HALLUCINATED_THRESHOLD, DEFAULT_INTENT_MODEL, DEFAULT_PROMPT_TARGET_THRESHOLD, EMBEDDINGS_INTERNAL_HOST, HALLUCINATION_INTERNAL_HOST, REQUEST_ID_HEADER, SYSTEM_ROLE, TOOL_ROLE, USER_ROLE, ZEROSHOT_INTERNAL_HOST
+    ARCH_FC_INTERNAL_HOST, ARCH_FC_MODEL_NAME, ARCH_FC_REQUEST_TIMEOUT_MS, ARCH_INTERNAL_CLUSTER_NAME, ARCH_MESSAGES_KEY, ARCH_MODEL_PREFIX, ARCH_STATE_HEADER, ARCH_UPSTREAM_HOST_HEADER, ASSISTANT_ROLE, DEFAULT_EMBEDDING_MODEL, HALLUCINATION_TEMPLATE, DEFAULT_HALLUCINATED_THRESHOLD, DEFAULT_INTENT_MODEL, DEFAULT_PROMPT_TARGET_THRESHOLD, EMBEDDINGS_INTERNAL_HOST, HALLUCINATION_INTERNAL_HOST, REQUEST_ID_HEADER, SYSTEM_ROLE, TOOL_ROLE, USER_ROLE, ZEROSHOT_INTERNAL_HOST
 };
 use common::embeddings::{
     CreateEmbeddingRequest, CreateEmbeddingRequestInput, CreateEmbeddingResponse,
@@ -323,12 +323,11 @@ impl StreamContext {
 
         if !keys_with_low_score.is_empty() {
             let response =
-                "It seems I’m missing some information. Could you provide the following details: "
-                    .to_string()
+                    HALLUCINATION_TEMPLATE.to_string()
                     + &keys_with_low_score.join(", ")
                     + " ?";
             let message = Message {
-                role: SYSTEM_ROLE.to_string(),
+                role: ASSISTANT_ROLE.to_string(),
                 content: Some(response),
                 model: Some(ARCH_FC_MODEL_NAME.to_string()),
                 tool_calls: None,