feature: set n messages to include w/ context window

Author: densumesh

server/src/data/models.rs

@@ -3324,7 +3324,7 @@ impl DatasetConfiguration {
                         query_parameter_description: Some("The search query to find relevant information".to_string()),
                         price_filter_description: Some("The price filter to use for the search".to_string()),
                         max_price_option_description: Some("The maximum price to filter by".to_string()),
-                        min_price_option_description: Some("The minimum price to filter by".to_string()),
+                        min_price_option_description: Some("The minimum price to filter by".to_string()),   
                     }),
                 }))
                 .as_object()
@@ -4618,6 +4618,7 @@ impl ApiKeyRequestParams {
             metadata: payload.metadata,
             use_agentic_search: payload.use_agentic_search,
             only_include_docs_used: payload.only_include_docs_used,
+            number_of_messages_to_include: payload.number_of_messages_to_include,
         }
     }
 
@@ -9504,6 +9505,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             remove_stop_words: Option<bool>,
             typo_options: Option<TypoOptions>,
             pub only_include_docs_used: Option<bool>,
+            pub number_of_messages_to_include: Option<u64>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9545,6 +9547,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             typo_options: helper.typo_options,
             use_agentic_search: helper.use_agentic_search,
             only_include_docs_used: helper.only_include_docs_used,
+            number_of_messages_to_include: helper.number_of_messages_to_include,
         })
     }
 }
@@ -9580,6 +9583,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             pub rag_context: Option<String>,
             pub use_agentic_search: Option<bool>,
             pub only_include_docs_used: Option<bool>,
+            pub number_of_messages_to_include: Option<u64>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9618,6 +9622,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             rag_context: helper.rag_context,
             use_agentic_search: helper.use_agentic_search,
             only_include_docs_used: helper.only_include_docs_used,
+            number_of_messages_to_include: helper.number_of_messages_to_include,
         })
     }
 }
@@ -9657,6 +9662,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             pub rag_context: Option<String>,
             pub use_agentic_search: Option<bool>,
             pub only_include_docs_used: Option<bool>,
+            pub number_of_messages_to_include: Option<u64>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9699,6 +9705,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             rag_context: helper.rag_context,
             use_agentic_search: helper.use_agentic_search,
             only_include_docs_used: helper.only_include_docs_used,
+            number_of_messages_to_include: helper.number_of_messages_to_include,
         })
     }
 }

server/src/handlers/message_handler.rs

@@ -136,6 +136,8 @@ pub struct CreateMessageReqPayload {
     pub rag_context: Option<String>,
     /// If true, the search will be conducted using llm tool calling. If not specified, this defaults to false.
     pub use_agentic_search: Option<bool>,
+    /// Number of messages to include in the context window. If not specified, this defaults to 10.
+    pub number_of_messages_to_include: Option<u64>,
 }
 
 /// Create message
@@ -243,6 +245,11 @@ pub async fn create_message(
     // remove chunks from the previous messages
     previous_messages = previous_messages
         .into_iter()
+        .take(
+            create_message_data
+                .number_of_messages_to_include
+                .unwrap_or(10) as usize,
+        )
         .map(|message| {
             let mut message = message;
             if message.role == "assistant" {
@@ -442,6 +449,8 @@ pub struct RegenerateMessageReqPayload {
     pub rag_context: Option<String>,
     /// If true, the search will be conducted using llm tool calling. If not specified, this defaults to false.
     pub use_agentic_search: Option<bool>,
+    /// Number of messages to include in the context window. If not specified, this defaults to 10.
+    pub number_of_messages_to_include: Option<u64>,
 }
 
 #[derive(Serialize, Debug, ToSchema)]
@@ -498,6 +507,8 @@ pub struct EditMessageReqPayload {
     pub rag_context: Option<String>,
     /// If true, the search will be conducted using llm tool calling. If not specified, this defaults to false.
     pub use_agentic_search: Option<bool>,
+    /// Number of messages to include in the context window. If not specified, this defaults to 10.
+    pub number_of_messages_to_include: Option<u64>,
 }
 
 impl From<EditMessageReqPayload> for CreateMessageReqPayload {
@@ -528,6 +539,7 @@ impl From<EditMessageReqPayload> for CreateMessageReqPayload {
             rag_context: data.rag_context,
             use_agentic_search: data.use_agentic_search,
             only_include_docs_used: data.only_include_docs_used,
+            number_of_messages_to_include: data.number_of_messages_to_include,
         }
     }
 }
@@ -560,6 +572,7 @@ impl From<RegenerateMessageReqPayload> for CreateMessageReqPayload {
             rag_context: data.rag_context,
             use_agentic_search: data.use_agentic_search,
             only_include_docs_used: data.only_include_docs_used,
+            number_of_messages_to_include: data.number_of_messages_to_include,
         }
     }
 }
@@ -732,6 +745,11 @@ pub async fn regenerate_message_patch(
     // remove citations from the previous messages
     previous_messages = previous_messages
         .into_iter()
+        .take(
+            create_message_data
+                .number_of_messages_to_include
+                .unwrap_or(10) as usize,
+        )
         .map(|message| {
             let mut message = message;
             if message.role == "assistant" {