mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-24 00:51:43 +00:00
61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
"""Memory retrieval operation for personal memories.
|
|
|
|
This module provides functionality to retrieve memories from a vector store
|
|
based on query similarity and score thresholds.
|
|
"""
|
|
|
|
from typing import List
|
|
|
|
from flowllm.core.context import C
|
|
from flowllm.core.op import BaseAsyncOp
|
|
from flowllm.core.schema import VectorNode
|
|
from loguru import logger
|
|
|
|
from reme_ai.schema.memory import BaseMemory, vector_node_to_memory
|
|
|
|
|
|
@C.register_op()
|
|
class RetrieveMemoryOp(BaseAsyncOp):
|
|
"""
|
|
Retrieves memories based on specified criteria such as status, type, and timestamp.
|
|
Processes these memories concurrently, sorts them by similarity, and logs the activity,
|
|
facilitating efficient memory retrieval operations within a given scope.
|
|
"""
|
|
|
|
async def async_execute(self):
|
|
"""
|
|
Executes the memory retrieval operation.
|
|
|
|
This method:
|
|
1. Retrieves memories from vector store based on query similarity
|
|
2. Removes duplicate memories based on content
|
|
3. Filters memories by score threshold if specified
|
|
4. Stores the retrieved memories in context metadata
|
|
"""
|
|
recall_key: str = self.op_params.get("recall_key", "query")
|
|
top_k: int = self.context.get("top_k", 3)
|
|
|
|
query: str = self.context[recall_key]
|
|
assert query, "query should be not empty!"
|
|
|
|
workspace_id: str = self.context.workspace_id
|
|
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:
|
|
memory: BaseMemory = vector_node_to_memory(node)
|
|
if memory.content not in memory_content_list:
|
|
memory_list.append(memory)
|
|
memory_content_list.append(memory.content)
|
|
logger.info(f"retrieve memory.size={len(memory_list)}")
|
|
|
|
threshold_score: float | None = self.op_params.get("threshold_score", None)
|
|
if threshold_score is not None:
|
|
memory_list = [mem for mem in memory_list if mem.score >= threshold_score or mem.score is None]
|
|
logger.info(f"after filter by threshold_score size={len(memory_list)}")
|
|
|
|
self.context.response.metadata["memory_list"] = memory_list
|