diff --git a/reme_ai/vector_store/delete_memory_op.py b/reme_ai/vector_store/delete_memory_op.py index e1c227b4..df7e6484 100644 --- a/reme_ai/vector_store/delete_memory_op.py +++ b/reme_ai/vector_store/delete_memory_op.py @@ -8,7 +8,7 @@ from flowllm.schema.vector_node import VectorNode class DeleteMemoryOp(BaseLLMOp): file_path: str = __file__ - def execute(self): + async def async_execute(self): workspace_id: str = self.context.workspace_id freq_threshold: int = self.context.freq_threshold utility_threshold: float = self.context.utility_threshold diff --git a/reme_ai/vector_store/recall_vector_store_op.py b/reme_ai/vector_store/recall_vector_store_op.py index bb375bde..048027e5 100644 --- a/reme_ai/vector_store/recall_vector_store_op.py +++ b/reme_ai/vector_store/recall_vector_store_op.py @@ -10,7 +10,7 @@ from reme_ai.schema.memory import BaseMemory, vector_node_to_memory @C.register_op() class RecallVectorStoreOp(BaseLLMOp): - def execute(self): + async def async_execute(self): recall_key: str = self.op_params.get("recall_key", "query") top_k: int = self.op_params.get("top_k", 3) @@ -18,7 +18,7 @@ class RecallVectorStoreOp(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/vector_store/update_memory_freq_op.py b/reme_ai/vector_store/update_memory_freq_op.py index d4b181a8..bc2dff15 100644 --- a/reme_ai/vector_store/update_memory_freq_op.py +++ b/reme_ai/vector_store/update_memory_freq_op.py @@ -10,7 +10,7 @@ from reme_ai.schema.memory import BaseMemory, dict_to_memory class UpdateMemoryFreqOp(BaseOp): file_path: str = __file__ - def execute(self): + async def async_execute(self): memory_dicts: List[dict] = self.context.memory_dicts if not memory_dicts: diff --git a/reme_ai/vector_store/update_memory_utility_op.py b/reme_ai/vector_store/update_memory_utility_op.py index 48d5b21b..999cacbf 100644 --- a/reme_ai/vector_store/update_memory_utility_op.py +++ b/reme_ai/vector_store/update_memory_utility_op.py @@ -10,7 +10,7 @@ from reme_ai.schema.memory import BaseMemory class UpdateMemoryUtilityOp(BaseOp): file_path: str = __file__ - def execute(self): + async def async_execute(self): memory_dicts: List[dict] = self.context.memory_dicts update_utility = self.context.update_utility diff --git a/reme_ai/vector_store/update_vector_store_op.py b/reme_ai/vector_store/update_vector_store_op.py index e41af0f6..1fd2487c 100644 --- a/reme_ai/vector_store/update_vector_store_op.py +++ b/reme_ai/vector_store/update_vector_store_op.py @@ -11,18 +11,18 @@ from reme_ai.schema.memory import BaseMemory @C.register_op() class UpdateVectorStoreOp(BaseLLMOp): - def execute(self): + async def async_execute(self): workspace_id: str = self.context.workspace_id deleted_memory_ids: List[str] = self.context.response.metadata.get("deleted_memory_ids", []) if deleted_memory_ids: - self.vector_store.delete(node_ids=deleted_memory_ids, workspace_id=workspace_id) + await self.vector_store.async_delete(node_ids=deleted_memory_ids, workspace_id=workspace_id) logger.info(f"delete memory_ids={json.dumps(deleted_memory_ids, indent=2)}") insert_memory_list: List[BaseMemory] = self.context.response.metadata.get("memory_list", []) if insert_memory_list: insert_nodes: List[VectorNode] = [x.to_vector_node() for x in insert_memory_list] - self.vector_store.insert(nodes=insert_nodes, workspace_id=workspace_id) + await self.vector_store.async_insert(nodes=insert_nodes, workspace_id=workspace_id) logger.info(f"insert insert_node.size={len(insert_nodes)}") # Store results in context diff --git a/reme_ai/vector_store/vector_store_action_op.py b/reme_ai/vector_store/vector_store_action_op.py index 9bd7a17c..ba4f722f 100644 --- a/reme_ai/vector_store/vector_store_action_op.py +++ b/reme_ai/vector_store/vector_store_action_op.py @@ -7,22 +7,22 @@ from reme_ai.schema.memory import vector_node_to_memory, dict_to_memory, BaseMem @C.register_op() class VectorStoreActionOp(BaseLLMOp): - def execute(self): + async def async_execute(self): workspace_id: str = self.context.workspace_id action: str = self.context.action - + result = "" if action == "copy": src_workspace_id: str = self.context.src_workspace_id - result = self.vector_store.copy_workspace(src_workspace_id=src_workspace_id, - dest_workspace_id=workspace_id) + result = await self.vector_store.async_copy_workspace(src_workspace_id=src_workspace_id, + dest_workspace_id=workspace_id) elif action == "delete": - if self.vector_store.exist_workspace(workspace_id): - result = self.vector_store.delete_workspace(workspace_id=workspace_id) + if await self.vector_store.async_exist_workspace(workspace_id): + result = await self.vector_store.async_delete_workspace(workspace_id=workspace_id) elif action == "delete_ids": memory_ids: list = self.context.memory_ids - result = self.vector_store.delete(workspace_id=workspace_id, node_ids=memory_ids) + result = await self.vector_store.async_delete(workspace_id=workspace_id, node_ids=memory_ids) elif action == "dump": path: str = self.context.path @@ -30,9 +30,9 @@ class VectorStoreActionOp(BaseLLMOp): def node_to_memory(node: VectorNode) -> dict: return vector_node_to_memory(node).model_dump() - result = self.vector_store.dump_workspace(workspace_id=workspace_id, - path=path, - callback_fn=node_to_memory) + result = await self.vector_store.async_dump_workspace(workspace_id=workspace_id, + path=path, + callback_fn=node_to_memory) elif action == "load": path: str = self.context.path @@ -41,9 +41,9 @@ class VectorStoreActionOp(BaseLLMOp): memory: BaseMemory = dict_to_memory(memory_dict=memory_dict) return memory.to_vector_node() - result = self.vector_store.load_workspace(workspace_id=workspace_id, - path=path, - callback_fn=memory_dict_to_node) + result = await self.vector_store.async_load_workspace(workspace_id=workspace_id, + path=path, + callback_fn=memory_dict_to_node) else: raise ValueError(f"invalid action={action}")