1 files changed, 44 insertions, 15 deletions

diff --git a/src/config/input.rs b/src/config/input.rs
index ae94799..56ae5ed 100644
--- a/src/config/input.rs
+++ b/src/config/input.rs
@@ -1,11 +1,11 @@
 use super::{role::Role, session::Session, GlobalConfig};
 
 use crate::client::{
-    init_client, list_models, ChatCompletionsData, Client, ImageUrl, Message, MessageContent,
+    init_client, list_chat_models, ChatCompletionsData, Client, ImageUrl, Message, MessageContent,
     MessageContentPart, MessageRole, Model,
 };
 use crate::function::{ToolCallResult, ToolResults};
-use crate::utils::{base64_encode, sha256};
+use crate::utils::{base64_encode, sha256, AbortSignal};
 
 use anyhow::{bail, Context, Result};
 use fancy_regex::Regex;
@@ -29,9 +29,11 @@ lazy_static! {
 pub struct Input {
     config: GlobalConfig,
     text: String,
+    patch_text: Option<String>,
     medias: Vec<String>,
     data_urls: HashMap<String, String>,
     tool_call: Option<ToolResults>,
+    rag: Option<String>,
     context: InputContext,
 }
 
@@ -40,9 +42,11 @@ impl Input {
         Self {
             config: config.clone(),
             text: text.to_string(),
+            patch_text: None,
             medias: Default::default(),
             data_urls: Default::default(),
             tool_call: None,
+            rag: None,
             context: context.unwrap_or_else(|| InputContext::from_config(config)),
         }
     }
@@ -92,9 +96,11 @@ impl Input {
         Ok(Self {
             config: config.clone(),
             text: texts.join("\n"),
+            patch_text: None,
             medias,
             data_urls,
             tool_call: Default::default(),
+            rag: None,
             context: context.unwrap_or_else(|| InputContext::from_config(config)),
         })
     }
@@ -108,13 +114,41 @@ impl Input {
     }
 
     pub fn text(&self) -> String {
-        self.text.clone()
+        match self.patch_text.clone() {
+            Some(text) => text,
+            None => self.text.clone(),
+        }
     }
 
     pub fn set_text(&mut self, text: String) {
         self.text = text;
     }
 
+    pub async fn maybe_embeddings(&mut self, abort_signal: AbortSignal) -> Result<()> {
+        if self.text.is_empty() {
+            return Ok(());
+        }
+        if !self.text.is_empty() {
+            let rag = self.config.read().rag.clone();
+            if let Some(rag) = rag {
+                let top_k = self.config.read().rag_top_k;
+                let embeddings = rag.search(&self.text, top_k, abort_signal).await?;
+                let text = self.config.read().rag_template(&embeddings, &self.text);
+                self.patch_text = Some(text);
+                self.rag = Some(rag.name().to_string());
+            }
+        }
+        Ok(())
+    }
+
+    pub fn rag(&self) -> Option<&str> {
+        self.rag.as_deref()
+    }
+
+    pub fn clear_patch_text(&mut self) {
+        self.patch_text.take();
+    }
+
     pub fn merge_tool_call(
         mut self,
         output: String,
@@ -134,7 +168,7 @@ impl Input {
         let model = self.config.read().model.clone();
         if let Some(model_id) = self.role().and_then(|v| v.model_id.clone()) {
             if model.id() != model_id {
-                if let Some(model) = list_models(&self.config.read())
+                if let Some(model) = list_chat_models(&self.config.read())
                     .into_iter()
                     .find(|v| v.id() == model_id)
                 {
@@ -158,7 +192,7 @@ impl Input {
             bail!("The current model does not support vision.");
         }
         let messages = self.build_messages()?;
-        self.config.read().model.max_input_tokens_limit(&messages)?;
+        self.config.read().model.guard_max_input_tokens(&messages)?;
         let (temperature, top_p) = if let Some(session) = self.session(&self.config.read().session)
         {
             (session.temperature(), session.top_p())
@@ -262,12 +296,12 @@ impl Input {
 
     pub fn render(&self) -> String {
         if self.medias.is_empty() {
-            return self.text.clone();
+            return self.text();
         }
         let text = if self.text.is_empty() {
-            self.text.to_string()
+            String::new()
         } else {
-            format!(" -- {}", self.text)
+            format!(" -- {}", self.text())
         };
         let files: Vec<String> = self
             .medias
@@ -280,7 +314,7 @@ impl Input {
 
     pub fn message_content(&self) -> MessageContent {
         if self.medias.is_empty() {
-            MessageContent::Text(self.text.clone())
+            MessageContent::Text(self.text())
         } else {
             let mut list: Vec<MessageContentPart> = self
                 .medias
@@ -291,12 +325,7 @@ impl Input {
                 })
                 .collect();
             if !self.text.is_empty() {
-                list.insert(
-                    0,
-                    MessageContentPart::Text {
-                        text: self.text.clone(),
-                    },
-                );
+                list.insert(0, MessageContentPart::Text { text: self.text() });
             }
             MessageContent::Array(list)
         }