feat: allow for specifying model on CreateMessage reqs

Author: skeptrunedev

server/src/data/models.rs

@@ -4619,6 +4619,7 @@ impl ApiKeyRequestParams {
             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,
+            model: payload.model,
         }
     }
 
@@ -9506,6 +9507,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             typo_options: Option<TypoOptions>,
             pub only_include_docs_used: Option<bool>,
             pub number_of_messages_to_include: Option<u64>,
+            pub model: Option<String>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9548,6 +9550,7 @@ impl<'de> Deserialize<'de> for CreateMessageReqPayload {
             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,
+            model: helper.model,
         })
     }
 }
@@ -9584,6 +9587,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             pub use_agentic_search: Option<bool>,
             pub only_include_docs_used: Option<bool>,
             pub number_of_messages_to_include: Option<u64>,
+            pub model: Option<String>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9623,6 +9627,7 @@ impl<'de> Deserialize<'de> for RegenerateMessageReqPayload {
             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,
+            model: helper.model,
         })
     }
 }
@@ -9663,6 +9668,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             pub use_agentic_search: Option<bool>,
             pub only_include_docs_used: Option<bool>,
             pub number_of_messages_to_include: Option<u64>,
+            pub model: Option<String>,
         }
 
         let mut helper = Helper::deserialize(deserializer)?;
@@ -9706,6 +9712,7 @@ impl<'de> Deserialize<'de> for EditMessageReqPayload {
             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,
+            model: helper.model,
         })
     }
 }

server/src/handlers/message_handler.rs

@@ -138,6 +138,8 @@ pub struct CreateMessageReqPayload {
     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>,
+    /// Model name to use for the completion. If not specified, this defaults to the dataset's model.
+    pub model: Option<String>,
 }
 
 /// Create message
@@ -451,6 +453,8 @@ pub struct RegenerateMessageReqPayload {
     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>,
+    /// Model name to use for the completion. If not specified, this defaults to the dataset's model.
+    pub model: Option<String>,
 }
 
 #[derive(Serialize, Debug, ToSchema)]
@@ -509,6 +513,8 @@ pub struct EditMessageReqPayload {
     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>,
+    /// Model name to use for the completion. If not specified, this defaults to the dataset's model.
+    pub model: Option<String>,
 }
 
 impl From<EditMessageReqPayload> for CreateMessageReqPayload {
@@ -540,6 +546,7 @@ impl From<EditMessageReqPayload> for CreateMessageReqPayload {
             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,
+            model: data.model,
         }
     }
 }
@@ -573,6 +580,7 @@ impl From<RegenerateMessageReqPayload> for CreateMessageReqPayload {
             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,
+            model: data.model,
         }
     }
 }

server/src/operators/message_operator.rs

@@ -854,7 +854,10 @@ pub async fn stream_response(
     };
 
     let rag_prompt = dataset_config.RAG_PROMPT.clone();
-    let chosen_model = dataset_config.LLM_DEFAULT_MODEL.clone();
+    let chosen_model = create_message_req_payload
+        .model
+        .clone()
+        .unwrap_or(dataset_config.LLM_DEFAULT_MODEL.clone());
 
     let (search_event, score_chunks) = get_rag_chunks_query(
         create_message_req_payload.clone(),
@@ -1894,8 +1897,10 @@ pub async fn stream_response_with_agentic_search(
         messages_len
     };
 
-    let chosen_model = dataset_config.LLM_DEFAULT_MODEL.clone();
-
+    let chosen_model = create_message_req_payload
+        .model
+        .clone()
+        .unwrap_or(dataset_config.LLM_DEFAULT_MODEL.clone());
     let tools = vec![ChatCompletionTool {
         r#type: ChatCompletionToolType::Function,
         function: ChatCompletionFunction {