mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
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
This commit is contained in:
parent
a9188edeb2
commit
bf3fa875aa
3 changed files with 9 additions and 9 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue