chore: support multiple source search in chat (#7906)

* chore: support multiple source search in chat

* chore: clippy

Author: appflowy

frontend/rust-lib/flowy-ai/src/ai_manager.rs

@@ -217,7 +217,6 @@ impl AIManager {
       let summary = select_chat_summary(&mut conn, chat_id).unwrap_or_default();
 
       let model = self.get_active_model(&chat_id.to_string()).await;
-      trace!("[AI Plugin] notify open chat: {}", chat_id);
       self
         .local_ai
         .open_chat(&workspace_id, chat_id, &model.name, rag_ids, summary)
@@ -240,7 +239,7 @@ impl AIManager {
       .await
       {
         Ok(settings) => {
-          local_ai.set_rag_ids(&chat_id, &settings.rag_ids);
+          local_ai.set_rag_ids(&chat_id, &settings.rag_ids).await;
           let rag_ids = settings
             .rag_ids
             .into_iter()
@@ -712,7 +711,7 @@ impl AIManager {
 
     let user_service = self.user_service.clone();
     let external_service = self.external_service.clone();
-    self.local_ai.set_rag_ids(chat_id, &rag_ids);
+    self.local_ai.set_rag_ids(chat_id, &rag_ids).await;
 
     let rag_ids = rag_ids
       .into_iter()

frontend/rust-lib/flowy-ai/src/embeddings/store.rs

@@ -1,6 +1,7 @@
 use crate::embeddings::document_indexer::split_text_into_chunks;
 use crate::embeddings::embedder::{Embedder, OllamaEmbedder};
 use crate::embeddings::indexer::{EmbeddingModel, IndexerProvider};
+use crate::local_ai::chat::retriever::MultipleSourceRetrieverStore;
 use async_trait::async_trait;
 use flowy_ai_pub::cloud::CollabType;
 use flowy_ai_pub::entities::{RAG_IDS, SOURCE_ID};
@@ -9,10 +10,8 @@ use flowy_sqlite_vec::db::VectorSqliteDB;
 use flowy_sqlite_vec::entities::{EmbeddedContent, SqliteEmbeddedDocument};
 use futures::stream::{self, StreamExt};
 use langchain_rust::llm::client::OllamaClient;
-use langchain_rust::{
-  schemas::Document,
-  vectorstore::{VecStoreOptions, VectorStore},
-};
+use langchain_rust::schemas::Document;
+use langchain_rust::vectorstore::{VecStoreOptions, VectorStore};
 use ollama_rs::generation::embeddings::request::{EmbeddingsInput, GenerateEmbeddingsRequest};
 use serde_json::Value;
 use std::collections::HashMap;
@@ -86,6 +85,80 @@ impl SqliteVectorStore {
   }
 }
 
+#[async_trait]
+impl MultipleSourceRetrieverStore for SqliteVectorStore {
+  fn retriever_name(&self) -> &'static str {
+    "Sqlite Multiple Source Retriever"
+  }
+
+  async fn read_documents(
+    &self,
+    workspace_id: &Uuid,
+    query: &str,
+    limit: usize,
+    rag_ids: &[String],
+    score_threshold: f32,
+    _full_search: bool,
+  ) -> FlowyResult<Vec<Document>> {
+    let vector_db = match self.vector_db.upgrade() {
+      Some(db) => db,
+      None => return Err(FlowyError::internal().with_context("Vector database not initialized")),
+    };
+
+    // Create embedder and generate embedding for query
+    let embedder = self.create_embedder()?;
+    let request = GenerateEmbeddingsRequest::new(
+      embedder.model().name().to_string(),
+      EmbeddingsInput::Single(query.to_string()),
+    );
+
+    let embedding = embedder.embed(request).await?.embeddings;
+    if embedding.is_empty() {
+      return Ok(Vec::new());
+    }
+
+    debug_assert!(embedding.len() == 1);
+    let query_embedding = embedding.first().unwrap();
+
+    // Perform similarity search in the database
+    let results = vector_db
+      .search_with_score(
+        &workspace_id.to_string(),
+        rag_ids,
+        query_embedding,
+        limit as i32,
+        score_threshold,
+      )
+      .await?;
+
+    trace!(
+      "[VectorStore] Found {} results for query:{}, rag_ids: {:?}, score_threshold: {}",
+      results.len(),
+      query,
+      rag_ids,
+      score_threshold
+    );
+
+    // Convert results to Documents
+    let documents = results
+      .into_iter()
+      .map(|result| {
+        let mut metadata = HashMap::new();
+
+        if let Some(map) = result.metadata.as_ref().and_then(|v| v.as_object()) {
+          for (key, value) in map {
+            metadata.insert(key.clone(), value.clone());
+          }
+        }
+
+        Document::new(result.content).with_metadata(metadata)
+      })
+      .collect();
+
+    Ok(documents)
+  }
+}
+
 #[async_trait]
 impl VectorStore for SqliteVectorStore {
   type Options = VecStoreOptions<Value>;
@@ -215,74 +288,23 @@ impl VectorStore for SqliteVectorStore {
 
     // Return empty result if workspace_id is missing
     let workspace_id = match workspace_id {
-      Some(id) => id.to_string(),
+      Some(id) => id,
       None => {
         warn!("[VectorStore] Missing workspace_id in filters. Returning empty result.");
         return Ok(Vec::new());
       },
     };
 
-    // Get the vector database
-    let vector_db = match self.vector_db.upgrade() {
-      Some(db) => db,
-      None => return Err("Vector database not initialized".into()),
-    };
-
-    // Create embedder and generate embedding for query
-    let embedder = self.create_embedder()?;
-    let request = GenerateEmbeddingsRequest::new(
-      embedder.model().name().to_string(),
-      EmbeddingsInput::Single(query.to_string()),
-    );
-
-    let embedding = match embedder.embed(request).await {
-      Ok(result) => result.embeddings,
-      Err(e) => return Err(Box::new(e)),
-    };
-
-    if embedding.is_empty() {
-      return Ok(Vec::new());
-    }
-
-    let score_threshold = opt.score_threshold.unwrap_or(0.4);
-    debug_assert!(embedding.len() == 1);
-    let query_embedding = embedding.first().unwrap();
-
-    // Perform similarity search in the database
-    let results = vector_db
-      .search_with_score(
+    self
+      .read_documents(
         &workspace_id,
+        query,
+        limit,
         &rag_ids,
-        query_embedding,
-        limit as i32,
-        score_threshold,
+        opt.score_threshold.unwrap_or(0.4),
+        true,
       )
-      .await?;
-
-    trace!(
-      "[VectorStore] Found {} results for query:{}, rag_ids: {:?}, score_threshold: {}",
-      results.len(),
-      query,
-      rag_ids,
-      score_threshold
-    );
-
-    // Convert results to Documents
-    let documents = results
-      .into_iter()
-      .map(|result| {
-        let mut metadata = HashMap::new();
-
-        if let Some(map) = result.metadata.as_ref().and_then(|v| v.as_object()) {
-          for (key, value) in map {
-            metadata.insert(key.clone(), value.clone());
-          }
-        }
-
-        Document::new(result.content).with_metadata(metadata)
-      })
-      .collect();
-
-    Ok(documents)
+      .await
+      .map_err(|err| Box::new(err) as Box<dyn Error>)
   }
 }

frontend/rust-lib/flowy-ai/src/local_ai/chat/chains/context_question_chain.rs

@@ -1,6 +1,7 @@
 use crate::local_ai::chat::llm::LLMOllama;
 use crate::SqliteVectorStore;
 use flowy_error::{FlowyError, FlowyResult};
+use flowy_sqlite_vec::entities::EmbeddedContent;
 use langchain_rust::language_models::llm::LLM;
 use langchain_rust::prompt::TemplateFormat;
 use langchain_rust::prompt::{PromptFromatter, PromptTemplate};
@@ -10,6 +11,7 @@ use ollama_rs::generation::parameters::{FormatType, JsonStructure};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_json::json;
+use std::fmt::Debug;
 use tracing::trace;
 use uuid::Uuid;
 
@@ -60,13 +62,13 @@ pub struct ContextQuestion {
   pub object_id: String,
 }
 
-pub struct RelatedQuestionChain {
+pub struct ContextRelatedQuestionChain {
   workspace_id: Uuid,
   llm: LLMOllama,
   store: SqliteVectorStore,
 }
 
-impl RelatedQuestionChain {
+impl ContextRelatedQuestionChain {
   pub fn new(workspace_id: Uuid, ollama: LLMOllama, store: SqliteVectorStore) -> Self {
     let format = FormatType::StructuredJson(JsonStructure::new::<ContextQuestionsResponse>());
     Self {
@@ -76,25 +78,16 @@ impl RelatedQuestionChain {
     }
   }
 
-  pub async fn generate_questions(&self, rag_ids: &[String]) -> FlowyResult<Vec<ContextQuestion>> {
-    trace!(
-      "[embedding] Generating context related questions for RAG IDs: {:?}",
-      rag_ids
-    );
-
-    let context = self
-      .store
-      .select_all_embedded_content(&self.workspace_id.to_string(), rag_ids, 3)
-      .await?;
-
-    trace!(
-      "[embedding] Generating related questions base on: {:?}",
-      context,
-    );
-
-    let context_str = json!(context).to_string();
+  pub async fn generate_questions_from_context<T>(
+    &self,
+    rag_ids: &[T],
+    context: &str,
+  ) -> FlowyResult<Vec<ContextQuestion>>
+  where
+    T: AsRef<str>,
+  {
     let input_variables = prompt_args! {
-        "context" => context_str,
+        "context" => context,
     };
 
     let template = PromptTemplate::new(
@@ -116,8 +109,42 @@ impl RelatedQuestionChain {
     // filter out questions that are not in the rag_ids
     parsed_result
       .questions
-      .retain(|v| rag_ids.contains(&v.object_id));
+      .retain(|v| rag_ids.iter().any(|id| id.as_ref() == v.object_id));
 
     Ok(parsed_result.questions)
   }
+
+  pub async fn generate_questions<T>(
+    &self,
+    rag_ids: &[T],
+  ) -> FlowyResult<(String, Vec<ContextQuestion>)>
+  where
+    T: AsRef<str> + Debug,
+  {
+    trace!(
+      "[embedding] Generating context related questions for RAG IDs: {:?}",
+      rag_ids
+    );
+
+    let rag_ids_str: Vec<String> = rag_ids.iter().map(|id| id.as_ref().to_string()).collect();
+    let context = self
+      .store
+      .select_all_embedded_content(&self.workspace_id.to_string(), &rag_ids_str, 3)
+      .await?;
+
+    trace!(
+      "[embedding] Generating related questions base on: {:?}",
+      context,
+    );
+
+    let context_str = embedded_documents_to_context_str(context);
+    self
+      .generate_questions_from_context(rag_ids, &context_str)
+      .await
+      .map(|questions| (context_str, questions))
+  }
+}
+
+pub fn embedded_documents_to_context_str(documents: Vec<EmbeddedContent>) -> String {
+  json!(documents).to_string()
 }