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:
jinli.yl 2025-09-06 21:42:55 +08:00
parent a9188edeb2
commit bf3fa875aa
3 changed files with 9 additions and 9 deletions

View file

@ -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:

View file

@ -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,

View file

@ -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)