reme-ai 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

reme_ai/summary/task/memory_deduplication_op.py

@@ -10,7 +10,7 @@ from reme_ai.schema.memory import BaseMemory
 class MemoryDeduplicationOp(BaseOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Remove duplicate task memories"""
         # Get task memories to deduplicate
         task_memories: List[BaseMemory] = self.context.memory_list
@@ -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)

reme_ai/summary/task/memory_validation_op.py

@@ -3,6 +3,8 @@ import re
 from typing import List, Dict, Any
 
 from flowllm import C, BaseLLMOp
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message as FlowMessage
 from loguru import logger
 
 from reme_ai.schema import Message
@@ -13,7 +15,7 @@ from reme_ai.schema.memory import BaseMemory
 class MemoryValidationOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Validate quality of extracted task memories"""
 
         task_memories: List[BaseMemory] = []
@@ -31,7 +33,7 @@ class MemoryValidationOp(BaseLLMOp):
         validated_task_memories = []
 
         for task_memory in task_memories:
-            validation_result = self._validate_single_task_memory(task_memory)
+            validation_result = await self._validate_single_task_memory(task_memory)
             if validation_result and validation_result.get("is_valid", False):
                 task_memory.score = validation_result.get("score", 0.0)
                 validated_task_memories.append(task_memory)
@@ -45,13 +47,13 @@ class MemoryValidationOp(BaseLLMOp):
         self.context.response.answer = json.dumps([x.model_dump() for x in validated_task_memories])
         self.context.response.metadata["memory_list"] = validated_task_memories
 
-    def _validate_single_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
+    async def _validate_single_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
         """Validate single task memory"""
-        validation_info = self._llm_validate_task_memory(task_memory)
+        validation_info = await self._llm_validate_task_memory(task_memory)
         logger.info(f"Validating: {validation_info}")
         return validation_info
 
-    def _llm_validate_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
+    async def _llm_validate_task_memory(self, task_memory: BaseMemory) -> Dict[str, Any]:
         """Validate task memory using LLM"""
         try:
             prompt = self.prompt_format(
@@ -96,7 +98,7 @@ class MemoryValidationOp(BaseLLMOp):
                         "reason": f"Parse error: {str(e_inner)}"
                     }
 
-            return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_validation)
+            return await self.llm.achat(messages=[FlowMessage(role=Role.USER, content=prompt)], callback_fn=parse_validation)
 
         except Exception as e:
             logger.error(f"LLM validation failed: {e}")

reme_ai/summary/task/simple_comparative_summary_op.py

@@ -2,6 +2,8 @@ import json
 from typing import List, Dict
 
 from flowllm import C, BaseLLMOp
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message as FlowMessage
 from loguru import logger
 
 from reme_ai.schema import Message, Trajectory
@@ -13,7 +15,7 @@ from reme_ai.utils.op_utils import merge_messages_content
 class SimpleComparativeSummaryOp(BaseLLMOp):
     file_path: str = __file__
 
-    def compare_summary_trajectory(self, trajectory_a: Trajectory, trajectory_b: Trajectory) -> List[BaseMemory]:
+    async def compare_summary_trajectory(self, trajectory_a: Trajectory, trajectory_b: Trajectory) -> List[BaseMemory]:
         summary_prompt = self.prompt_format(prompt_name="summary_prompt",
                                             execution_process_a=merge_messages_content(trajectory_a.messages),
                                             execution_process_b=merge_messages_content(trajectory_b.messages),
@@ -42,9 +44,9 @@ class SimpleComparativeSummaryOp(BaseLLMOp):
                 logger.exception(f"parse content failed!\n{content}")
                 raise e
 
-        return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content)
+        return await self.llm.achat(messages=[FlowMessage(role=Role.USER, content=summary_prompt)], callback_fn=parse_content)
 
-    def execute(self):
+    async def async_execute(self):
         trajectories: list = self.context.get("trajectories", [])
         trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
 
@@ -61,7 +63,7 @@ class SimpleComparativeSummaryOp(BaseLLMOp):
                 continue
 
             if task_trajectories[0].score > task_trajectories[-1].score:
-                task_memories = self.compare_summary_trajectory(trajectory_a=task_trajectories[0],
+                task_memories = await self.compare_summary_trajectory(trajectory_a=task_trajectories[0],
                                                                 trajectory_b=task_trajectories[-1])
                 memory_list.extend(task_memories)
 

reme_ai/summary/task/simple_summary_op.py

@@ -2,6 +2,8 @@ import json
 from typing import List
 
 from flowllm import C, BaseLLMOp
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message as FlowMessage
 from loguru import logger
 
 from reme_ai.schema import Message, Trajectory
@@ -13,7 +15,7 @@ from reme_ai.utils.op_utils import merge_messages_content
 class SimpleSummaryOp(BaseLLMOp):
     file_path: str = __file__
 
-    def summary_trajectory(self, trajectory: Trajectory) -> List[BaseMemory]:
+    async def summary_trajectory(self, trajectory: Trajectory) -> List[BaseMemory]:
         execution_process = merge_messages_content(trajectory.messages)
         success_score_threshold: float = self.op_params.get("success_score_threshold", 0.9)
         logger.info(f"success_score_threshold={success_score_threshold}")
@@ -49,15 +51,15 @@ class SimpleSummaryOp(BaseLLMOp):
                 logger.exception(f"parse content failed!\n{content}")
                 raise e
 
-        return self.llm.chat(messages=[Message(content=summary_prompt)], callback_fn=parse_content)
+        return await self.llm.achat(messages=[FlowMessage(role=Role.USER, content=summary_prompt)], callback_fn=parse_content)
 
-    def execute(self):
+    async def async_execute(self):
         trajectories: list = self.context.trajectories
         trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
 
         memory_list: List[BaseMemory] = []
         for trajectory in trajectories:
-            memories = self.summary_trajectory(trajectory)
+            memories = await self.summary_trajectory(trajectory)
             if memories:
                 memory_list.extend(memories)
 

reme_ai/summary/task/success_extraction_op.py

@@ -1,6 +1,8 @@
 from typing import List
 
 from flowllm import C, BaseLLMOp
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message as FlowMessage
 from loguru import logger
 
 from reme_ai.schema import Message, Trajectory
@@ -12,7 +14,7 @@ from reme_ai.utils.op_utils import merge_messages_content, parse_json_experience
 class SuccessExtractionOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Extract task memories from successful trajectories"""
         success_trajectories: List[Trajectory] = self.context.success_trajectories
 
@@ -29,11 +31,11 @@ class SuccessExtractionOp(BaseLLMOp):
             if "segments" in trajectory.metadata:
                 # Process segmented step sequences
                 for segment in trajectory.metadata["segments"]:
-                    task_memories = self._extract_success_task_memory_from_steps(segment, trajectory)
+                    task_memories = await self._extract_success_task_memory_from_steps(segment, trajectory)
                     success_task_memories.extend(task_memories)
             else:
                 # Process entire trajectory
-                task_memories = self._extract_success_task_memory_from_steps(trajectory.messages, trajectory)
+                task_memories = await self._extract_success_task_memory_from_steps(trajectory.messages, trajectory)
                 success_task_memories.extend(task_memories)
 
         logger.info(f"Extracted {len(success_task_memories)} success task memories")
@@ -41,7 +43,7 @@ class SuccessExtractionOp(BaseLLMOp):
         # Add task memories to context
         self.context.success_task_memories = success_task_memories
 
-    def _extract_success_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
+    async def _extract_success_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
         """Extract task memory from successful step sequences"""
         step_content = merge_messages_content(steps)
         context = get_trajectory_context(trajectory, steps)
@@ -70,4 +72,4 @@ class SuccessExtractionOp(BaseLLMOp):
 
             return task_memories
 
-        return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_task_memories)
+        return await self.llm.achat(messages=[FlowMessage(role=Role.USER, content=prompt)], callback_fn=parse_task_memories)

reme_ai/summary/task/trajectory_preprocess_op.py

@@ -1,4 +1,3 @@
-import json
 from typing import List, Dict
 
 from flowllm import C, BaseOp
@@ -11,18 +10,10 @@ from reme_ai.schema import Trajectory
 class TrajectoryPreprocessOp(BaseOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Preprocess trajectories: validate and classify"""
         trajectories: list = self.context.get("trajectories", [])
-        # trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
-        new_trajectories: List[Trajectory] = []
-        for x in trajectories:
-            if isinstance(x, dict):
-                x["messages"] = self._modify_tool_calls(x["messages"])
-                new_trajectories.append(Trajectory(**x))
-            else:
-                new_trajectories.append(x)
-        trajectories = new_trajectories
+        trajectories: List[Trajectory] = [Trajectory(**x) if isinstance(x, dict) else x for x in trajectories]
 
         # Classify trajectories
         classified = self._classify_trajectories(trajectories)
@@ -53,24 +44,4 @@ class TrajectoryPreprocessOp(BaseOp):
             'success': success_trajectories,
             'failure': failure_trajectories,
             'all': trajectories
-        }
-
-    def _modify_tool_calls(self, messages: List[Dict]) -> List[Dict]:
-        new_messages = []
-
-        for msg in messages:
-            if 'tool_calls' in msg:
-                processed_tool_calls = []
-                for tool_call in msg['tool_calls']:
-                    tool_type = tool_call.get("type", "function")
-                    nested_data = tool_call.get(tool_type, {})
-                    tool_call.update({
-                        "arguments": json.loads(nested_data.get("arguments", "")),
-                        "name": nested_data.get("name", "")
-                    })
-                    tool_call.pop(tool_type)
-                    processed_tool_calls.append(tool_call)
-                msg['tool_calls'] = processed_tool_calls
-            new_messages.append(msg)
-
-        return new_messages
+        }

reme_ai/summary/task/trajectory_segmentation_op.py

@@ -3,6 +3,8 @@ import re
 from typing import List
 
 from flowllm import C, BaseLLMOp
+from flowllm.enumeration.role import Role
+from flowllm.schema.message import Message as FlowMessage
 from loguru import logger
 
 from reme_ai.schema import Message, Trajectory
@@ -12,7 +14,7 @@ from reme_ai.schema import Message, Trajectory
 class TrajectorySegmentationOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Segment trajectories into meaningful steps"""
         # Get trajectories from context
         all_trajectories: List[Trajectory] = self.context.get("all_trajectories", [])
@@ -30,7 +32,7 @@ class TrajectorySegmentationOp(BaseLLMOp):
         # Add segmentation info to trajectories
         segmented_count = 0
         for trajectory in target_trajectories:
-            segments = self._llm_segment_trajectory(trajectory)
+            segments = await self._llm_segment_trajectory(trajectory)
             trajectory.metadata["segments"] = segments
             segmented_count += 1
 
@@ -51,7 +53,7 @@ class TrajectorySegmentationOp(BaseLLMOp):
         else:
             return all_trajectories
 
-    def _llm_segment_trajectory(self, trajectory: Trajectory) -> List[List[Message]]:
+    async def _llm_segment_trajectory(self, trajectory: Trajectory) -> List[List[Message]]:
         """Use LLM for trajectory segmentation"""
         trajectory_content = self._format_trajectory_content(trajectory)
 
@@ -80,7 +82,7 @@ class TrajectorySegmentationOp(BaseLLMOp):
 
             return segments if segments else [trajectory.messages]
 
-        return self.llm.chat(messages=[Message(content=prompt)], callback_fn=parse_segmentation,
+        return await self.llm.achat(messages=[FlowMessage(role=Role.USER, content=prompt)], callback_fn=parse_segmentation,
                              default_value=[trajectory.messages])
 
     @staticmethod

reme_ai/utils/datetime_handler.py

@@ -2,8 +2,8 @@ import datetime
 import re
 from typing import List
 
+from reme_ai.constants.language_constants import LanguageEnum
 from reme_ai.constants.language_constants import WEEKDAYS, DATATIME_WORD_LIST, MONTH_DICT
-from reme_ai.enumeration.language_constants import LanguageEnum
 
 
 class DatetimeHandler(object):

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

reme_ai/vector_store/recall_vector_store_op.py

@@ -10,15 +10,15 @@ 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)
+        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] = 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:

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:

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
 

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

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,16 @@ 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 == "list":
+            def node_to_memory(node: VectorNode) -> dict:
+                return vector_node_to_memory(node).model_dump()
+
+            result = await self.vector_store.async_iter_workspace_nodes(workspace_id=workspace_id)
+            result = [node_to_memory(node) for node in result]
 
         elif action == "load":
             path: str = self.context.path
@@ -41,15 +48,11 @@ 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}")
 
-        # Store results in context
-        if isinstance(result, dict):
-            self.context.response.metadata["action_result"] = result
-        else:
-            self.context.response.metadata["action_result"] = {"result": str(result)}
+        self.context.response.metadata["action_result"] = result