feat(server): allow users to pass in custom model for rag on specified chunks

Author: skeptrunedev

clients/ts-sdk/openapi.json

@@ -14589,6 +14589,11 @@
             "description": "Metadata is any metadata you want to associate w/ the event that is created from this request",
             "nullable": true
           },
+          "model": {
+            "type": "string",
+            "description": "Model to use for the completion. If not specified, the default model configured for the dataset will be used.",
+            "nullable": true
+          },
           "presence_penalty": {
             "type": "number",
             "format": "float",

clients/ts-sdk/src/types.gen.ts

@@ -2264,6 +2264,10 @@ export type GenerateOffChunksReqPayload = {
      * Metadata is any metadata you want to associate w/ the event that is created from this request
      */
     metadata?: unknown;
+    /**
+     * Model to use for the completion. If not specified, the default model configured for the dataset will be used.
+     */
+    model?: (string) | null;
     /**
      * Presence penalty is a number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics. Default is 0.7.
      */

server/src/handlers/chunk_handler.rs

@@ -2523,6 +2523,8 @@ pub struct GenerateOffChunksReqPayload {
     pub image_config: Option<ImageConfig>,
     /// Context options to use for the completion. If not specified, all options will default to false.
     pub context_options: Option<ContextOptions>,
+    /// Model to use for the completion. If not specified, the default model configured for the dataset will be used.
+    pub model: Option<String>,
 }
 
 /// RAG on Specified Chunks
@@ -2584,7 +2586,10 @@ pub async fn generate_off_chunks(
 
     let base_url = dataset_config.LLM_BASE_URL;
     let rag_prompt = dataset_config.RAG_PROMPT.clone();
-    let chosen_model = dataset_config.LLM_DEFAULT_MODEL;
+    let chosen_model = data
+        .model
+        .clone()
+        .unwrap_or(dataset_config.LLM_DEFAULT_MODEL);
 
     let base_url = if base_url.is_empty() {
         "https://openrouter.ai/api/v1".into()

server/src/handlers/message_handler.rs

@@ -289,6 +289,8 @@ pub async fn create_message(
             redis_pool,
             dataset_config,
             create_message_data,
+            #[cfg(feature = "hallucination-detection")]
+            hallucination_detector,
         )
         .await
     } else {
@@ -704,6 +706,8 @@ pub async fn regenerate_message_patch(
                 redis_pool,
                 dataset_config,
                 create_message_data,
+                #[cfg(feature = "hallucination-detection")]
+                hallucination_detector,
             )
             .await;
         } else {
@@ -792,6 +796,8 @@ pub async fn regenerate_message_patch(
             redis_pool,
             dataset_config,
             create_message_data,
+            #[cfg(feature = "hallucination-detection")]
+            hallucination_detector,
         )
         .await
     } else {

server/src/operators/message_operator.rs

@@ -1698,7 +1698,7 @@ pub async fn stream_response_with_agentic_search(
     redis_pool: web::Data<RedisPool>,
     dataset_config: DatasetConfiguration,
     create_message_req_payload: CreateMessageReqPayload,
-    #[cfg(feature = "hallucination-detection")] hallucination_detector: web::Data<
+    #[cfg(feature = "hallucination-detection")] _hallucination_detector: web::Data<
         HallucinationDetector,
     >,
 ) -> Result<HttpResponse, actix_web::Error> {

yarn.lock

@@ -7047,6 +7047,10 @@ postcss@^8.3.11, postcss@^8.4.32, postcss@^8.4.43, postcss@^8.4.45, postcss@^8.4
   version "8.5.3"
   resolved "https://registry.npmjs.org/postcss/-/postcss-8.5.3.tgz"
   integrity sha512-dle9A3yYxlBSrt8Fu+IpjGT8SY8hN0mlaA6GY8t0P5PjIOZemULz/E2Bnm/2dcUOena75OTNkHI76uZBNUUq3A==
+  dependencies:
+    nanoid "^3.3.8"
+    picocolors "^1.1.1"
+    source-map-js "^1.2.1"
 
 preact@^10.25.1:
   version "10.26.6"
@@ -7826,7 +7830,7 @@ solid-use@^0.9.0:
   resolved "https://registry.npmjs.org/solid-use/-/solid-use-0.9.1.tgz"
   integrity sha512-UwvXDVPlrrbj/9ewG9ys5uL2IO4jSiwys2KPzK4zsnAcmEl7iDafZWW1Mo4BSEWOmQCGK6IvpmGHo1aou8iOFw==
 
-source-map-js@^1.0.1, source-map-js@^1.2.1:
+source-map-js@^1.0.1:
   version "1.2.1"
   resolved "https://registry.npmjs.org/source-map-js/-/source-map-js-1.2.1.tgz"
   integrity sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA==