feat(memory): Add VoyageAI embedding support for memory

Author: User

huly-coder.yaml

@@ -61,6 +61,19 @@ permission_mode: manual_approval
 #---------------------------------------
 workspace: ./target/workspace
 
+#---------------------------------------
+# Memory Embedding Configuration
+#---------------------------------------
+# Supported embedding providers:
+# - fastembed (default)
+# - voyage_ai
+#     type: voyage_ai
+#     api_key: xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+#     model: voyage-3.5-lite
+#     dimensions: 256
+memory_embedding:
+  type: fastembed
+
 #---------------------------------------
 # Web Interaction Configuration
 #---------------------------------------

src/agent/mod.rs

@@ -19,7 +19,7 @@ use crate::tools::execute_command::tools::TerminateCommandTool;
 use crate::tools::execute_command::ProcessRegistry;
 use crate::tools::list_files::ListFilesTool;
 use crate::tools::memory;
-use crate::tools::memory::Entity;
+use crate::tools::memory::indexer::MemoryIndexer;
 use crate::tools::memory::MemoryManager;
 use crate::tools::read_file::ReadFileTool;
 use crate::tools::replace_in_file::ReplaceInFileTool;
@@ -40,7 +40,6 @@ use rig::agent::AgentBuilder;
 use rig::completion::CompletionError;
 use rig::completion::CompletionModel;
 use rig::completion::CompletionResponse;
-use rig::embeddings::EmbeddingsBuilder;
 use rig::message::AssistantContent;
 use rig::message::ImageMediaType;
 use rig::message::Message;
@@ -50,8 +49,6 @@ use rig::message::UserContent;
 use rig::tool::Tool;
 use rig::tool::ToolError;
 use rig::tool::ToolSetError;
-use rig::vector_store::in_memory_store::InMemoryVectorIndex;
-use rig::vector_store::in_memory_store::InMemoryVectorStore;
 use rig::OneOrMany;
 use serde::Deserialize;
 use serde::Serialize;
@@ -138,41 +135,27 @@ struct AgentContext {
     state: Arc<RwLock<AgentState>>,
     sender: mpsc::UnboundedSender<AgentOutputEvent>,
     process_registry: Arc<RwLock<ProcessRegistry>>,
-    memory_index: Arc<RwLock<InMemoryVectorIndex<rig_fastembed::EmbeddingModel, Entity>>>,
+    memory_index: Arc<RwLock<MemoryIndexer>>,
     system_prompt_token_count: u32,
     current_input_tokens: u32,
     current_completion_tokens: u32,
 }
 
 impl Agent {
     pub fn new(
-        data_dir: &str,
+        _data_dir: &str,
         config: Config,
+        memory: Arc<RwLock<MemoryManager>>,
         sender: mpsc::UnboundedSender<AgentOutputEvent>,
     ) -> Self {
         Self {
             config,
             sender,
-            memory: Arc::new(RwLock::new(MemoryManager::new(data_dir, false))),
+            memory,
             process_registry: Arc::new(RwLock::new(ProcessRegistry::default())),
         }
     }
 
-    pub async fn init_memory_index(
-        &mut self,
-    ) -> InMemoryVectorIndex<rig_fastembed::EmbeddingModel, Entity> {
-        let documents = self.memory.read().await.entities().clone();
-        let client = rig_fastembed::Client::new();
-        let model = client.embedding_model(&rig_fastembed::FastembedModel::AllMiniLML6V2);
-        let embeddings = EmbeddingsBuilder::new(model.clone())
-            .documents(documents)
-            .unwrap()
-            .build()
-            .await
-            .unwrap();
-        InMemoryVectorStore::from_documents(embeddings.into_iter()).index(model)
-    }
-
     fn add_static_tools<M>(
         agent_builder: AgentBuilder<M>,
         context: BuildAgentContext<'_>,
@@ -418,7 +401,7 @@ impl Agent {
         data_dir: &str,
         receiver: mpsc::UnboundedReceiver<AgentControlEvent>,
         messages: Vec<Message>,
-        memory_index: InMemoryVectorIndex<rig_fastembed::EmbeddingModel, Entity>,
+        memory_index: Arc<RwLock<MemoryIndexer>>,
     ) {
         tracing::info!(
             "Run agent: {:?} : {}",
@@ -468,7 +451,6 @@ impl Agent {
             .unwrap();
 
         let messages = Arc::new(RwLock::new(messages));
-        let memory_index = Arc::new(RwLock::new(memory_index));
         let sender = self.sender.clone();
         let state = Arc::new(RwLock::new(state));
         let config_state = Arc::new(RwLock::new(AgentConfigState::new(data_dir)));

src/agent/utils.rs

@@ -6,14 +6,12 @@ use std::sync::Arc;
 
 use itertools::Itertools;
 use rig::message::{Message, UserContent};
-use rig::vector_store::in_memory_store::InMemoryVectorIndex;
-use rig::vector_store::VectorStoreIndex;
 use tokio::sync::RwLock;
 
 use crate::config::Config;
 use crate::templates::{ENV_DETAILS, SYSTEM_PROMPT};
 use crate::tools::execute_command::ProcessRegistry;
-use crate::tools::memory::{self, Entity};
+use crate::tools::memory::indexer::MemoryIndexer;
 use crate::HISTORY_PATH;
 
 pub const MAX_FILES: usize = 10000;
@@ -51,9 +49,7 @@ pub async fn prepare_system_prompt(config: &Config) -> String {
 
 pub async fn add_env_message<'a>(
     msg: &'a mut Message,
-    memory_index: Arc<
-        tokio::sync::RwLock<InMemoryVectorIndex<rig_fastembed::EmbeddingModel, memory::Entity>>,
-    >,
+    memory_index: Arc<RwLock<MemoryIndexer>>,
     data_dir: &'a Path,
     workspace: &'a Path,
     process_registry: Arc<RwLock<ProcessRegistry>>,
@@ -98,8 +94,7 @@ pub async fn add_env_message<'a>(
             _ => "",
         };
         if !txt.is_empty() {
-            let res: Vec<(f64, String, Entity)> = memory_index.top_n(txt, 10).await.unwrap();
-            let result: Vec<_> = res.into_iter().map(|(_, _, entity)| entity).collect();
+            let result = memory_index.search(txt, 10).await.unwrap();
             memory_entries = serde_yaml::to_string(&result).unwrap();
         }
 

src/config.rs

@@ -92,6 +92,17 @@ pub struct Appearance {
     pub user_name: String,
 }
 
+#[derive(Debug, Deserialize, Clone)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum EmbeddingProvider {
+    VoyageAi {
+        api_key: String,
+        model: String,
+        dimensions: usize,
+    },
+    Fastembed,
+}
+
 #[derive(Debug, Deserialize, Clone)]
 pub struct Config {
     pub provider: ProviderKind,
@@ -106,6 +117,7 @@ pub struct Config {
     pub mcp: Option<McpConfig>,
     pub web_search: Option<WebSearchProvider>,
     pub web_fetch: Option<WebFetchProvider>,
+    pub memory_embedding: EmbeddingProvider,
 }
 
 impl Config {

src/main.rs

@@ -5,6 +5,7 @@ use std::io::stdout;
 use std::panic::set_hook;
 use std::panic::take_hook;
 use std::path::Path;
+use std::sync::Arc;
 
 use crossterm::execute;
 use crossterm::terminal::disable_raw_mode;
@@ -15,13 +16,16 @@ use providers::model_info::model_info;
 use ratatui::prelude::CrosstermBackend;
 use ratatui::DefaultTerminal;
 use ratatui::Terminal;
+use tokio::sync::RwLock;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
 use tracing_subscriber::Layer;
 
 use self::config::Config;
 use crate::agent::AgentControlEvent;
 use crate::agent::AgentOutputEvent;
+use crate::tools::memory::indexer::MemoryIndexer;
+use crate::tools::memory::MemoryManager;
 use clap::Parser;
 
 mod agent;
@@ -145,8 +149,14 @@ async fn main() -> color_eyre::Result<()> {
     let model_info = model_info(&args.data, &config).await?;
     tracing::info!("Model info: {:?}", model_info);
 
-    let mut agent = agent::Agent::new(&args.data, config.clone(), output_sender);
-    let memory_index = agent.init_memory_index().await;
+    let memory_index = Arc::new(RwLock::new(MemoryIndexer::new(data_dir, &config)));
+    let memory = Arc::new(RwLock::new(MemoryManager::new(
+        &args.data,
+        memory_index.clone(),
+        false,
+    )));
+    let mut agent = agent::Agent::new(&args.data, config.clone(), memory.clone(), output_sender);
+    memory_index.write().await.init(memory.clone()).await?;
 
     let messages = history.clone();
     let data_dir = args.data.clone();