From bf3fa875aa101e7620bd38ba2a74160fdea1a1b6 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 6 Sep 2025 21:42:55 +0800 Subject: [PATCH] refactor(vector_store): update memory operations to use async search - Update load_today_memory_op.py to use async version of _retrieve_today_memories - Update memory_deduplication_op.py to use async version of _deduplicate_task_memories and _get_existing_task_memory_embeddings - Update retrieve_memory_op.py to use async_search instead of search --- reme_ai/retrieve/personal/retrieve_memory_op.py | 2 +- reme_ai/summary/personal/load_today_memory_op.py | 6 +++--- reme_ai/summary/task/memory_deduplication_op.py | 10 +++++----- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/reme_ai/retrieve/personal/retrieve_memory_op.py b/reme_ai/retrieve/personal/retrieve_memory_op.py index 6f96815b..25283759 100644 --- a/reme_ai/retrieve/personal/retrieve_memory_op.py +++ b/reme_ai/retrieve/personal/retrieve_memory_op.py @@ -23,7 +23,7 @@ class RetrieveMemoryOp(BaseLLMOp): assert query, "query should be not empty!" workspace_id: str = self.context.workspace_id - nodes: List[VectorNode] = self.vector_store.search(query=query, workspace_id=workspace_id, top_k=top_k) + nodes: List[VectorNode] = await self.vector_store.async_search(query=query, workspace_id=workspace_id, top_k=top_k) memory_list: List[BaseMemory] = [] memory_content_list: List[str] = [] for node in nodes: diff --git a/reme_ai/summary/personal/load_today_memory_op.py b/reme_ai/summary/personal/load_today_memory_op.py index 46113196..dbc83f83 100644 --- a/reme_ai/summary/personal/load_today_memory_op.py +++ b/reme_ai/summary/personal/load_today_memory_op.py @@ -36,7 +36,7 @@ class LoadTodayMemoryOp(BaseLLMOp): logger.info(f"Loading today's memories for user: {user_name} (top_k: {top_k})") # Get today's memories from vector store - today_memories = self._retrieve_today_memories(workspace_id, user_name, top_k) + today_memories = await self._retrieve_today_memories(workspace_id, user_name, top_k) if not today_memories: logger.info("No memories found for today") @@ -47,7 +47,7 @@ class LoadTodayMemoryOp(BaseLLMOp): self.context.today_memories = today_memories logger.info(f"Final today's memory list size: {len(today_memories)}") - def _retrieve_today_memories(self, workspace_id: str, user_name: str, top_k: int) -> List[BaseMemory]: + async def _retrieve_today_memories(self, workspace_id: str, user_name: str, top_k: int) -> List[BaseMemory]: """ Retrieve memories from today using vector store with date filtering. @@ -74,7 +74,7 @@ class LoadTodayMemoryOp(BaseLLMOp): } # Search vector store with date filter - nodes: List[VectorNode] = self.vector_store.search( + nodes: List[VectorNode] = await self.vector_store.async_search( query="", # Empty query to get all results for today workspace_id=workspace_id, top_k=top_k, diff --git a/reme_ai/summary/task/memory_deduplication_op.py b/reme_ai/summary/task/memory_deduplication_op.py index 2a475cb9..be884eda 100644 --- a/reme_ai/summary/task/memory_deduplication_op.py +++ b/reme_ai/summary/task/memory_deduplication_op.py @@ -22,7 +22,7 @@ class MemoryDeduplicationOp(BaseOp): logger.info(f"Starting deduplication for {len(task_memories)} task memories") # Perform deduplication - deduplicated_task_memories = self._deduplicate_task_memories(task_memories) + deduplicated_task_memories = await self._deduplicate_task_memories(task_memories) logger.info( f"Deduplication complete: {len(deduplicated_task_memories)} deduplicated task memories out of {len(task_memories)}") @@ -30,7 +30,7 @@ class MemoryDeduplicationOp(BaseOp): # Update context self.context.memory_list = deduplicated_task_memories - def _deduplicate_task_memories(self, task_memories: List[BaseMemory]) -> List[BaseMemory]: + async def _deduplicate_task_memories(self, task_memories: List[BaseMemory]) -> List[BaseMemory]: """Remove duplicate task memories""" if not task_memories: return task_memories @@ -41,7 +41,7 @@ class MemoryDeduplicationOp(BaseOp): unique_task_memories = [] # Get existing task memory embeddings - existing_embeddings = self._get_existing_task_memory_embeddings(workspace_id) + existing_embeddings = await self._get_existing_task_memory_embeddings(workspace_id) for task_memory in task_memories: # Generate embedding for current task memory @@ -67,14 +67,14 @@ class MemoryDeduplicationOp(BaseOp): return unique_task_memories - def _get_existing_task_memory_embeddings(self, workspace_id: str) -> List[List[float]]: + async def _get_existing_task_memory_embeddings(self, workspace_id: str) -> List[List[float]]: """Get embeddings of existing task memories""" try: if not hasattr(self.context, 'vector_store') or not self.context.vector_store or not workspace_id: return [] # Query existing task memory nodes - existing_nodes = self.context.vector_store.search( + existing_nodes = await self.context.vector_store.async_search( query="...", # Empty query to get all workspace_id=workspace_id, top_k=self.op_params.get("max_existing_task_memories", 1000)