🚧 using summarization api based on beta usage ratio

server/config/default.toml

@@ -5,6 +5,7 @@ port = 3030
 max_query_length = 300
 max_sources = 10
 max_search_context = 5
+beta_usage_ratio = 0.0
 
 [summarizer]
 model = ""

server/src/llms/summarizer.rs

@@ -49,7 +49,7 @@ pub struct SummarizerStreamOutput {
     pub generated_text: Option<String>,
 }
 
-fn prepare_context_string(
+fn prepare_llm_context_string(
     settings: &SummarizerSettings,
     summarizer_input: SummarizerInput,
 ) -> SummarizerAPIInput {
@@ -69,13 +69,13 @@ fn prepare_context_string(
 }
 
 #[tracing::instrument(level = "debug", ret, err)]
-pub async fn generate_text_stream(
+pub async fn generate_text_with_llm(
     settings: SummarizerSettings,
     summarizer_input: SummarizerInput,
     update_processor: api_models::UpdateResultProcessor,
     tx: Sender<api_models::SearchByIdResponse>,
 ) -> crate::Result<()> {
-    let summarizer_input = prepare_context_string(&settings, summarizer_input);
+    let summarizer_input = prepare_llm_context_string(&settings, summarizer_input);
     let client = Client::new();
 
     let response = client

server/src/rag/models.rs

@@ -28,4 +28,5 @@ pub struct SearchSettings {
     pub max_query_length: u16,
     pub max_sources: u8,
     pub max_search_context: u8,
+    pub beta_usage_ratio: f32,
 }

server/src/rag/post_process.rs

@@ -3,6 +3,7 @@ use crate::proto::Embeddings;
 use crate::rag::utils;
 use crate::search::api_models;
 use crate::settings::Settings;
+use rand::Rng;
 use regex::Regex;
 use std::cmp::Ordering;
 use tokio::sync::mpsc::Sender;
@@ -44,17 +45,32 @@ pub async fn summarize_search_results(
     stream_regex: Regex,
     tx: Sender<api_models::SearchByIdResponse>,
 ) -> crate::Result<()> {
-    summarizer::generate_text_with_openai(
-        settings.openai,
-        summarizer::SummarizerInput {
-            query: search_query_request.query,
-            retrieved_result: search_response,
-        },
-        update_processor,
-        stream_regex,
-        tx,
-    )
-    .await?;
+    let random_number = rand::thread_rng().gen_range(0.0..1.0);
+
+    if random_number < settings.search.beta_usage_ratio {
+        summarizer::generate_text_with_llm(
+            settings.summarizer,
+            summarizer::SummarizerInput {
+                query: search_query_request.query,
+                retrieved_result: search_response,
+            },
+            update_processor,
+            tx,
+        )
+        .await?;
+    } else {
+        summarizer::generate_text_with_openai(
+            settings.openai,
+            summarizer::SummarizerInput {
+                query: search_query_request.query,
+                retrieved_result: search_response,
+            },
+            update_processor,
+            stream_regex,
+            tx,
+        )
+        .await?;
+    }
 
     Ok(())
 }