feat(server): rag_context field for adding better context

Author: cdxker

server/src/data/models.rs

@@ -4514,6 +4514,7 @@ impl ApiKeyRequestParams {
     ) -> CreateMessageReqPayload {
         CreateMessageReqPayload {
             new_message_content: payload.new_message_content,
+            rag_context: payload.rag_context,
             topic_id: payload.topic_id,
             user_id: payload.user_id,
             sort_options: payload.sort_options,
@@ -9417,6 +9418,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             pub only_include_docs_used: Option<bool>,
             pub currency: Option<String>,
             metadata: Option<serde_json::Value>,
+            pub rag_context: Option<String>,
             #[serde(flatten)]
             other: std::collections::HashMap<String, serde_json::Value>,
             use_quote_negated_terms: Option<bool>,
@@ -9457,6 +9459,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             context_options,
             no_result_message: helper.no_result_message,
             metadata: helper.metadata,
+            rag_context: helper.rag_context,
             only_include_docs_used: helper.only_include_docs_used,
             use_quote_negated_terms: helper.use_quote_negated_terms,
             remove_stop_words: helper.remove_stop_words,
@@ -9494,6 +9497,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             pub use_quote_negated_terms: Option<bool>,
             pub remove_stop_words: Option<bool>,
             pub typo_options: Option<TypoOptions>,
+            pub rag_context: Option<String>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9530,6 +9534,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             use_quote_negated_terms: helper.use_quote_negated_terms,
             remove_stop_words: helper.remove_stop_words,
             typo_options: helper.typo_options,
+            rag_context: helper.rag_context,
         })
     }
 }
@@ -9567,6 +9572,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             pub use_quote_negated_terms: Option<bool>,
             pub remove_stop_words: Option<bool>,
             pub typo_options: Option<TypoOptions>,
+            pub rag_context: Option<String>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9607,6 +9613,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             use_quote_negated_terms: helper.use_quote_negated_terms,
             remove_stop_words: helper.remove_stop_words,
             typo_options: helper.typo_options,
+            rag_context: helper.rag_context,
         })
     }
 }

server/src/handlers/message_handler.rs

@@ -129,6 +129,8 @@ pub struct CreateMessageReqPayload {
     pub typo_options: Option<TypoOptions>,
     /// Metadata is any metadata you want to associate w/ the event that is created from this request
     pub metadata: Option<serde_json::Value>,
+    /// Overrides what the way chunks are placed into the context window
+    pub rag_context: Option<String>,
 }
 
 /// Create message
@@ -414,6 +416,8 @@ pub struct RegenerateMessageReqPayload {
     pub typo_options: Option<TypoOptions>,
     /// Metadata is any metadata you want to associate w/ the event that is created from this request
     pub metadata: Option<serde_json::Value>,
+    /// Overrides what the way chunks are placed into the context window
+    pub rag_context: Option<String>,
 }
 
 #[derive(Serialize, Debug, ToSchema)]
@@ -466,6 +470,8 @@ pub struct EditMessageReqPayload {
     pub typo_options: Option<TypoOptions>,
     /// Metadata is any metadata you want to associate w/ the event that is created from this request
     pub metadata: Option<serde_json::Value>,
+    /// Overrides what the way chunks are placed into the context window
+    pub rag_context: Option<String>,
 }
 
 impl From<EditMessageReqPayload> for CreateMessageReqPayload {
@@ -494,6 +500,7 @@ impl From<EditMessageReqPayload> for CreateMessageReqPayload {
             remove_stop_words: data.remove_stop_words,
             typo_options: data.typo_options,
             metadata: data.metadata,
+            rag_context: data.rag_context,
         }
     }
 }
@@ -524,6 +531,7 @@ impl From<RegenerateMessageReqPayload> for CreateMessageReqPayload {
             remove_stop_words: data.remove_stop_words,
             typo_options: data.typo_options,
             metadata: data.metadata,
+            rag_context: data.rag_context,
         }
     }
 }

server/src/operators/message_operator.rs

@@ -837,20 +837,25 @@ pub async fn stream_response(
             .streaming(response_stream));
     }
 
-    let rag_content = score_chunks
-        .iter()
-        .enumerate()
-        .map(|(idx, score_chunk)| {
-            json!({
-                "doc": idx + 1,
-                "text": convert_html_to_text(&(ChunkMetadata::from(score_chunk.chunk.clone()).chunk_html.clone().unwrap_or_default())),
-                "num_value": ChunkMetadata::from(score_chunk.chunk.clone()).num_value.map(|x| format!("{} {}", create_message_req_payload.currency.clone().unwrap_or("".to_string()), x)).unwrap_or("".to_string()),
-                "link": ChunkMetadata::from(score_chunk.chunk.clone()).link.clone().unwrap_or_default()
+    let rag_content = match create_message_req_payload.rag_context {
+        Some(rag_context) => rag_context,
+        None => {
+            score_chunks
+            .iter()
+            .enumerate()
+            .map(|(idx, score_chunk)| {
+                json!({
+                    "doc": idx + 1,
+                    "text": convert_html_to_text(&(ChunkMetadata::from(score_chunk.chunk.clone()).chunk_html.clone().unwrap_or_default())),
+                    "num_value": ChunkMetadata::from(score_chunk.chunk.clone()).num_value.map(|x| format!("{} {}", create_message_req_payload.currency.clone().unwrap_or("".to_string()), x)).unwrap_or("".to_string()),
+                    "link": ChunkMetadata::from(score_chunk.chunk.clone()).link.clone().unwrap_or_default()
+                })
+                .to_string()
             })
-            .to_string()
-        })
-        .collect::<Vec<String>>()
-        .join("\n\n");
+            .collect::<Vec<String>>()
+            .join("\n\n")
+        }
+    };
 
     let user_message = match &openai_messages
         .last()