refactor: fully complete state re-architecting

src/rag/mod.rs

@@ -20,6 +20,30 @@ use std::{
 };
 use tokio::time::sleep;
 
+const RAG_TEMPLATE: &str = r#"Answer the query based on the context while respecting the rules. (user query, some textual context and rules, all inside xml tags)
+
+<context>
+__CONTEXT__
+</context>
+
+<sources>
+__SOURCES__
+</sources>
+
+<rules>
+- If you don't know, just say so.
+- If you are not sure, ask for clarification.
+- Answer in the same language as the user query.
+- If the context appears unreadable or of poor quality, tell the user then answer as best as you can.
+- If the answer is not in the context but you think you know the answer, explain that to the user then answer with your own knowledge.
+- Answer directly and without using xml tags.
+- When using information from the context, cite the relevant source from the <sources> section.
+</rules>
+
+<user_query>
+__INPUT__
+</user_query>"#;
+
 pub struct Rag {
     app_config: Arc<AppConfig>,
     name: String,
@@ -318,6 +342,29 @@ impl Rag {
         Ok((embeddings, sources, ids))
     }
 
+    pub async fn search_with_template(
+        &self,
+        app: &AppConfig,
+        text: &str,
+        abort_signal: AbortSignal,
+    ) -> Result<String> {
+        let (reranker_model, top_k) = self.get_config();
+        let (embeddings, sources, ids) = self
+            .search(text, top_k, reranker_model.as_deref(), abort_signal)
+            .await?;
+        let rag_template = app.rag_template.as_deref().unwrap_or(RAG_TEMPLATE);
+        let text = if embeddings.is_empty() {
+            text.to_string()
+        } else {
+            rag_template
+                .replace("__CONTEXT__", &embeddings)
+                .replace("__SOURCES__", &sources)
+                .replace("__INPUT__", text)
+        };
+        self.set_last_sources(&ids);
+        Ok(text)
+    }
+
     fn resolve_source(&self, id: &DocumentId) -> String {
         let (file_index, _) = id.split();
         self.data