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

reme_ai/service/personal_memory_service.py

@@ -0,0 +1,128 @@
+import asyncio
+from typing import Optional, Dict, Any, List
+
+from flowllm import C
+from flowllm.flow import BaseToolFlow
+from flowllm.schema.flow_response import FlowResponse
+from loguru import logger
+from pydantic import Field, BaseModel
+
+from reme_ai.config.config_parser import ConfigParser
+from reme_ai.schema.memory import PersonalMemory
+from reme_ai.service.base_memory_service import BaseMemoryService
+
+
+class PersonalMemoryService(BaseMemoryService):
+
+    async def start(self) -> None:
+        C.set_service_config(parser=ConfigParser, config_name="config=default").init_by_service_config()
+
+    async def stop(self) -> None:
+        C.stop_by_service_config()
+
+    async def health(self) -> bool:
+        return True
+
+    async def add_memory(self, user_id: str, messages: list, session_id: Optional[str] = None) -> None:
+        summary_flow: BaseToolFlow = C.flow_dict["summary_personal_memory"]
+
+        new_messages: List[dict] = []
+        for message in messages:
+            if isinstance(message, dict):
+                new_messages.append(message)
+            elif isinstance(message, BaseModel):
+                new_messages.append(message.model_dump())
+            else:
+                raise ValueError(f"Invalid message type={type(message)}")
+
+        kwargs = {
+            "workspace_id": user_id,
+            "trajectories": [
+                {"messages": messages, "score": 1.0}
+            ]
+        }
+
+        result: FlowResponse = await summary_flow(**kwargs)
+        memory_list: List[PersonalMemory] = result.metadata.get("memory_list", [])
+        for memory in memory_list:
+            memory_id = memory.memory_id
+            self.add_session_memory_id(session_id, memory_id)
+            logger.info(f"[personal_memory_service] user_id={user_id} session_id={session_id} add memory: {memory}")
+
+    async def search_memory(self, user_id: str, messages: list, filters: Optional[Dict[str, Any]] = Field(
+        description="Associated filters for the messages, "
+                    "such as top_k, score etc.",
+        default=None,
+    )) -> list:
+
+        retrieve_flow: BaseToolFlow = C.flow_dict["retrieve_personal_memory"]
+
+        new_messages: List[dict] = []
+        for message in messages:
+            if isinstance(message, dict):
+                new_messages.append(message)
+            elif isinstance(message, BaseModel):
+                new_messages.append(message.model_dump())
+            else:
+                raise ValueError(f"Invalid message type={type(message)}")
+
+        # Extract query from the last message
+        query = messages[-1]["content"] if messages else ""
+
+        kwargs = {
+            "workspace_id": user_id,
+            "query": query,
+            "top_k": filters.get("top_k", 1) if filters else 1
+        }
+
+        result: FlowResponse = await retrieve_flow(**kwargs)
+        logger.info(f"[personal_memory_service] user_id={user_id} search result: {result.model_dump_json()}")
+
+        return [result.answer]
+
+    async def list_memory(self, user_id: str, filters: Optional[Dict[str, Any]] = Field(
+        description="Associated filters for the messages, "
+                    "such as top_k, score etc.",
+        default=None,
+    )) -> list:
+        vector_store_flow: BaseToolFlow = C.flow_dict["vector_store"]
+        result = await vector_store_flow(workspace_id=user_id, action="list")
+        logger.info(f"[personal_memory_service] list_memory result: {result}")
+
+        result = result.metadata["action_result"]
+        for i, line in enumerate(result):
+            logger.info(f"[personal_memory_service] list memory.{i}={line}")
+        return result
+
+    async def delete_memory(self, user_id: str, session_id: Optional[str] = None) -> None:
+        delete_ids = self.session_id_dict.get(session_id, [])
+        if not delete_ids:
+            return
+
+        vector_store_flow: BaseToolFlow = C.flow_dict["vector_store"]
+        result = await vector_store_flow(workspace_id=user_id, action="delete_ids", memory_ids=delete_ids)
+        result = result.metadata["action_result"]
+        logger.info(f"[personal_memory_service] delete memory result={result}")
+
+
+async def main():
+    async with PersonalMemoryService() as service:
+        logger.info("========== start personal memory service ==========")
+
+        await service.add_memory(user_id="u_12345",
+                                 messages=[{"content": "I really enjoy playing tennis on weekends"}],
+                                 session_id="s_123456")
+
+        await service.search_memory(user_id="u_12345",
+                                    messages=[{"content": "What do I like to do for fun?"}],
+                                    filters={"top_k": 1})
+
+        await service.list_memory(user_id="u_12345")
+        await service.delete_memory(user_id="u_12345", session_id="s_123456")
+        await service.list_memory(user_id="u_12345")
+
+        logger.info("========== end personal memory service ==========")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

reme_ai/service/task_memory_service.py

@@ -0,0 +1,126 @@
+import asyncio
+from typing import Optional, Dict, Any, List
+
+from flowllm import C
+from flowllm.flow import BaseToolFlow
+from flowllm.schema.flow_response import FlowResponse
+from loguru import logger
+from pydantic import Field, BaseModel
+
+from reme_ai.config.config_parser import ConfigParser
+from reme_ai.schema.memory import TaskMemory
+from reme_ai.service.base_memory_service import BaseMemoryService
+
+
+class TaskMemoryService(BaseMemoryService):
+
+    async def start(self) -> None:
+        C.set_service_config(parser=ConfigParser, config_name="config=default").init_by_service_config()
+
+    async def stop(self) -> None:
+        C.stop_by_service_config()
+
+    async def health(self) -> bool:
+        return True
+
+    async def add_memory(self, user_id: str, messages: list, session_id: Optional[str] = None) -> None:
+        summary_flow: BaseToolFlow = C.flow_dict["summary_task_memory"]
+
+        new_messages: List[dict] = []
+        for message in messages:
+            if isinstance(message, dict):
+                new_messages.append(message)
+            elif isinstance(message, BaseModel):
+                new_messages.append(message.model_dump())
+            else:
+                raise ValueError(f"Invalid message type={type(message)}")
+
+        kwargs = {
+            "workspace_id": user_id,
+            "trajectories": [
+                {"messages": messages, "score": 1.0}
+            ]
+        }
+
+        result: FlowResponse = await summary_flow(**kwargs)
+        memory_list: List[TaskMemory] = result.metadata.get("memory_list", [])
+        for memory in memory_list:
+            memory_id = memory.memory_id
+            self.add_session_memory_id(session_id, memory_id)
+            logger.info(f"[task_memory_service] user_id={user_id} session_id={session_id} add memory: {memory}")
+
+    async def search_memory(self, user_id: str, messages: list, filters: Optional[Dict[str, Any]] = Field(
+        description="Associated filters for the messages, "
+                    "such as top_k, score etc.",
+        default=None,
+    )) -> list:
+
+        retrieve_flow: BaseToolFlow = C.flow_dict["retrieve_task_memory"]
+
+        new_messages: List[dict] = []
+        for message in messages:
+            if isinstance(message, dict):
+                new_messages.append(message)
+            elif isinstance(message, BaseModel):
+                new_messages.append(message.model_dump())
+            else:
+                raise ValueError(f"Invalid message type={type(message)}")
+
+        kwargs = {
+            "workspace_id": user_id,
+            "messages": messages,
+            "top_k": filters.get("top_k", 1)
+        }
+
+        result: FlowResponse = await retrieve_flow(**kwargs)
+        logger.info(f"[task_memory_service] user_id={user_id} add result: {result.model_dump_json()}")
+
+        return [result.answer]
+
+    async def list_memory(self, user_id: str, filters: Optional[Dict[str, Any]] = Field(
+        description="Associated filters for the messages, "
+                    "such as top_k, score etc.",
+        default=None,
+    )) -> list:
+        vector_store_flow: BaseToolFlow = C.flow_dict["vector_store"]
+        result = await vector_store_flow(workspace_id=user_id, action="list")
+        print("list_memory result:", result)
+
+
+        result = result.metadata["action_result"]
+        for i, line in enumerate(result):
+            logger.info(f"[task_memory_service] list memory.{i}={line}")
+        return result
+
+    async def delete_memory(self, user_id: str, session_id: Optional[str] = None) -> None:
+        delete_ids = self.session_id_dict.get(session_id, [])
+        if not delete_ids:
+            return
+
+        vector_store_flow: BaseToolFlow = C.flow_dict["vector_store"]
+        result = await vector_store_flow(workspace_id=user_id, action="delete_ids", memory_ids=delete_ids)
+        result = result.metadata["action_result"]
+        logger.info(f"[task_memory_service] delete memory result={result}")
+
+
+async def main():
+    async with TaskMemoryService() as service:
+        logger.info("========== start task memory service ==========")
+
+        await service.add_memory(user_id="u_123456",
+                                 messages=[{"content": "please use web search tool to search financial news:"}],
+                                 session_id="s_123456")
+
+        await service.search_memory(user_id="u_123456",
+                                    messages=[{"content": "please use web search tool to search financial news"}],
+                                    filters={"top_k": 1})
+
+        await service.list_memory(user_id="u_123456")
+        await service.delete_memory(user_id="u_123456", session_id="s_123456")
+        await service.list_memory(user_id="u_123456")
+
+        logger.info("========== end task memory service ==========")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

reme_ai/summary/personal/contra_repeat_op.py

@@ -24,7 +24,7 @@ class ContraRepeatOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """
         Executes the primary routine of the ContraRepeatOp which involves:
         1. Gets memory list from context
@@ -82,7 +82,7 @@ class ContraRepeatOp(BaseLLMOp):
         logger.info(f"contra_repeat_prompt={full_prompt}")
 
         # Call LLM
-        response = self.llm.chat([Message(role=Role.USER, content=full_prompt)])
+        response = await self.llm.achat([Message(role=Role.USER, content=full_prompt)])
 
         # Return if empty
         if not response or not response.content:

reme_ai/summary/personal/get_observation_op.py

@@ -16,7 +16,7 @@ class GetObservationOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Extract personal observations from chat messages"""
         # Get messages from context - guaranteed to exist by flow input
         messages: List[Message] = self.context.messages
@@ -34,7 +34,7 @@ class GetObservationOp(BaseLLMOp):
         logger.info(f"Extracting observations from {len(filtered_messages)} filtered messages")
 
         # Extract observations using LLM
-        observation_memories = self._extract_observations_from_messages(filtered_messages)
+        observation_memories = await self._extract_observations_from_messages(filtered_messages)
 
         # Store results in context using standardized key
         self.context.observation_memories = observation_memories
@@ -58,7 +58,7 @@ class GetObservationOp(BaseLLMOp):
         logger.info(f"Filtered messages from {len(messages)} to {len(filtered_messages)}")
         return filtered_messages
 
-    def _extract_observations_from_messages(self, filtered_messages: List[Message]) -> List[BaseMemory]:
+    async def _extract_observations_from_messages(self, filtered_messages: List[Message]) -> List[BaseMemory]:
         """Extract observations from filtered messages using LLM"""
         user_name = self.context.get("user_name", "user")
 
@@ -113,7 +113,7 @@ class GetObservationOp(BaseLLMOp):
             return observation_memories
 
         # Use LLM chat with callback function
-        return self.llm.chat(messages=[Message(content=full_prompt)], callback_fn=parse_observations)
+        return await self.llm.achat(messages=[Message(content=full_prompt)], callback_fn=parse_observations)
 
     @staticmethod
     def parse_observation_response(response_text: str) -> List[dict]:

reme_ai/summary/personal/get_observation_with_time_op.py

@@ -16,7 +16,7 @@ class GetObservationWithTimeOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Extract personal observations with time information from chat messages"""
         # Get messages from context - guaranteed to exist by flow input
         messages: List[Message] = self.context.messages
@@ -34,7 +34,7 @@ class GetObservationWithTimeOp(BaseLLMOp):
         logger.info(f"Extracting observations with time from {len(filtered_messages)} filtered messages")
 
         # Extract observations using LLM
-        observation_memories_with_time = self._extract_observations_with_time_from_messages(filtered_messages)
+        observation_memories_with_time = await self._extract_observations_with_time_from_messages(filtered_messages)
 
         # Store results in context using standardized key
         self.context.observation_memories_with_time = observation_memories_with_time
@@ -58,7 +58,7 @@ class GetObservationWithTimeOp(BaseLLMOp):
         logger.info(f"Filtered messages from {len(messages)} to {len(filtered_messages)}")
         return filtered_messages
 
-    def _extract_observations_with_time_from_messages(self, filtered_messages: List[Message]) -> List[BaseMemory]:
+    async def _extract_observations_with_time_from_messages(self, filtered_messages: List[Message]) -> List[BaseMemory]:
         """Extract observations with time information from filtered messages using LLM"""
         user_name = self.context.get("user_name", "user")
 
@@ -123,7 +123,7 @@ class GetObservationWithTimeOp(BaseLLMOp):
             return observation_memories
 
         # Use LLM chat with callback function
-        return self.llm.chat(messages=[Message(content=full_prompt)], callback_fn=parse_observations)
+        return await self.llm.achat(messages=[Message(content=full_prompt)], callback_fn=parse_observations)
 
     def _get_colon_word(self) -> str:
         """Get language-specific colon word"""

reme_ai/summary/personal/get_reflection_subject_op.py

@@ -39,7 +39,7 @@ class GetReflectionSubjectOp(BaseLLMOp):
             }
         )
 
-    def execute(self):
+    async def async_execute(self):
         """
         Generate reflection subjects (topics) from personal memories for insight extraction.
 
@@ -85,7 +85,7 @@ class GetReflectionSubjectOp(BaseLLMOp):
             return
 
         # Generate reflection subjects using LLM
-        insight_memories = self._generate_reflection_subjects(
+        insight_memories = await self._generate_reflection_subjects(
             memory_contents, existing_subjects, user_name, reflect_num_questions
         )
 
@@ -93,7 +93,7 @@ class GetReflectionSubjectOp(BaseLLMOp):
         self.context.response.metadata["insight_memories"] = insight_memories
         logger.info(f"Generated {len(insight_memories)} new reflection subject memories")
 
-    def _generate_reflection_subjects(self, memory_contents: List[str], existing_subjects: List[str],
+    async def _generate_reflection_subjects(self, memory_contents: List[str], existing_subjects: List[str],
                                       user_name: str, num_questions: int) -> List[BaseMemory]:
         """
         Generate new reflection subjects using LLM analysis of memory contents.
@@ -148,7 +148,7 @@ class GetReflectionSubjectOp(BaseLLMOp):
             return insight_memories
 
         # Generate subjects using LLM
-        return self.llm.chat(messages=[Message(content=full_prompt)], callback_fn=parse_reflection_response)
+        return await self.llm.achat(messages=[Message(content=full_prompt)], callback_fn=parse_reflection_response)
 
     def get_language_value(self, value_dict: dict):
         """Get language-specific value from dictionary"""

reme_ai/summary/personal/info_filter_op.py

@@ -16,7 +16,7 @@ class InfoFilterOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Filter messages based on information content scores"""
         # Get messages from context - guaranteed to exist by flow input
         trajectories: list = self.context.trajectories
@@ -45,7 +45,7 @@ class InfoFilterOp(BaseLLMOp):
         logger.info(f"Filtering {len(info_messages)} messages for information content")
 
         # Filter messages using LLM
-        filtered_memories = self._filter_messages_with_llm(info_messages, user_name, preserved_scores)
+        filtered_memories = await self._filter_messages_with_llm(info_messages, user_name, preserved_scores)
 
         # Store results in context using standardized key
         self.context.messages = filtered_memories
@@ -81,7 +81,7 @@ class InfoFilterOp(BaseLLMOp):
         logger.info(f"Filtered messages from {len(messages)} to {len(info_messages)}")
         return info_messages
 
-    def _filter_messages_with_llm(self, info_messages: List[Message], user_name: str, preserved_scores: str) -> List[
+    async def _filter_messages_with_llm(self, info_messages: List[Message], user_name: str, preserved_scores: str) -> List[
         PersonalMemory]:
         """Filter messages using LLM to score information content"""
 
@@ -146,7 +146,7 @@ class InfoFilterOp(BaseLLMOp):
             return filtered_memories
 
         # Use LLM chat with callback function
-        return self.llm.chat(messages=[Message(content=full_prompt)], callback_fn=parse_and_filter)
+        return await self.llm.achat(messages=[Message(content=full_prompt)], callback_fn=parse_and_filter)
 
     def _get_colon_word(self) -> str:
         """Get language-specific colon word"""

reme_ai/summary/personal/load_today_memory_op.py

@@ -16,7 +16,7 @@ class LoadTodayMemoryOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """
         Load today's memories from vector store and perform deduplication.
         
@@ -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,12 +74,11 @@ class LoadTodayMemoryOp(BaseLLMOp):
             }
 
             # Search vector store with date filter
-            nodes: List[VectorNode] = self.vector_store.search(
-                query="",  # Empty query to get all results for today
+            nodes: List[VectorNode] = await self.vector_store.async_search(
+                query=" ",
                 workspace_id=workspace_id,
                 top_k=top_k,
-                filter_dict=filter_dict
-            )
+                filter_dict=filter_dict)
 
             logger.info(f"Vector store returned {len(nodes)} nodes for today")
 

reme_ai/summary/personal/long_contra_repeat_op.py

@@ -19,7 +19,7 @@ class LongContraRepeatOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """
         Analyze memories for contradictions and redundancies, resolving conflicts.
         
@@ -62,13 +62,13 @@ class LongContraRepeatOp(BaseLLMOp):
         logger.info(f"Processing {len(sorted_memories)} memories for contradictions and redundancies")
 
         # Analyze and resolve contradictions
-        filtered_memories = self._analyze_and_resolve_conflicts(sorted_memories)
+        filtered_memories = await self._analyze_and_resolve_conflicts(sorted_memories)
 
         # Store results in context
         self.context.response.metadata["memory_list"] = filtered_memories
         logger.info(f"Conflict resolution: {len(sorted_memories)} -> {len(filtered_memories)} memories")
 
-    def _analyze_and_resolve_conflicts(self, memories: List[BaseMemory]) -> List[BaseMemory]:
+    async def _analyze_and_resolve_conflicts(self, memories: List[BaseMemory]) -> List[BaseMemory]:
         """
         Analyze memories for contradictions and redundancies using LLM.
         
@@ -104,7 +104,7 @@ class LongContraRepeatOp(BaseLLMOp):
         logger.info(f"Contradiction analysis prompt length: {len(full_prompt)} chars")
 
         # Get LLM analysis
-        response = self.llm.chat([Message(role=Role.USER, content=full_prompt)])
+        response = await self.llm.achat([Message(role=Role.USER, content=full_prompt)])
 
         if not response or not response.content:
             logger.warning("Empty response from LLM, keeping all memories")

reme_ai/summary/personal/update_insight_op.py

@@ -17,7 +17,7 @@ class UpdateInsightOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """
         Update insight values based on new observation memories.
         
@@ -65,7 +65,7 @@ class UpdateInsightOp(BaseLLMOp):
         # Update each selected insight
         updated_insights = []
         for insight_memory, relevance_score, relevant_observations in top_insights:
-            updated_insight = self._update_insight_with_observations(
+            updated_insight = await self._update_insight_with_observations(
                 insight_memory, relevant_observations, user_name
             )
             if updated_insight:
@@ -135,7 +135,7 @@ class UpdateInsightOp(BaseLLMOp):
 
         return intersection / union if union > 0 else 0.0
 
-    def _update_insight_with_observations(self, insight_memory: PersonalMemory,
+    async def _update_insight_with_observations(self, insight_memory: PersonalMemory,
                                           relevant_observations: List[PersonalMemory],
                                           user_name: str) -> PersonalMemory:
         """
@@ -208,7 +208,7 @@ class UpdateInsightOp(BaseLLMOp):
 
         # Use LLM chat with callback function
         try:
-            return self.llm.chat(messages=[Message(content=full_prompt)], callback_fn=parse_update_response)
+            return await self.llm.achat(messages=[Message(content=full_prompt)], callback_fn=parse_update_response)
         except Exception as e:
             logger.error(f"Error updating insight: {e}")
             return insight_memory

reme_ai/summary/task/__init__.py

@@ -2,7 +2,6 @@ from .comparative_extraction_op import ComparativeExtractionOp
 from .failure_extraction_op import FailureExtractionOp
 from .memory_deduplication_op import MemoryDeduplicationOp
 from .memory_validation_op import MemoryValidationOp
-from .pdf_preprocess_op_wrapper import PDFPreprocessOp
 from .simple_comparative_summary_op import SimpleComparativeSummaryOp
 from .simple_summary_op import SimpleSummaryOp
 from .success_extraction_op import SuccessExtractionOp

reme_ai/summary/task/comparative_extraction_op.py

@@ -1,6 +1,8 @@
 from typing import List, Tuple, Optional
 
 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 ComparativeExtractionOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Extract comparative task memories by comparing different scoring trajectories"""
         all_trajectories: List[Trajectory] = self.context.get("all_trajectories", [])
         success_trajectories: List[Trajectory] = self.context.get("success_trajectories", [])
@@ -26,7 +28,7 @@ class ComparativeExtractionOp(BaseLLMOp):
             if highest_traj and lowest_traj and highest_traj.score > lowest_traj.score:
                 logger.info(
                     f"Extracting soft comparative task memories: highest ({highest_traj.score:.2f}) vs lowest ({lowest_traj.score:.2f})")
-                soft_task_memories = self._extract_soft_comparative_task_memory(highest_traj, lowest_traj)
+                soft_task_memories = await self._extract_soft_comparative_task_memory(highest_traj, lowest_traj)
                 comparative_task_memories.extend(soft_task_memories)
 
         # Hard comparison: success vs failure (if similarity search is enabled)
@@ -37,7 +39,7 @@ class ComparativeExtractionOp(BaseLLMOp):
             logger.info(f"Found {len(similar_pairs)} similar pairs for hard comparison")
 
             for success_steps, failure_steps, similarity_score in similar_pairs:
-                hard_task_memories = self._extract_hard_comparative_task_memory(success_steps, failure_steps,
+                hard_task_memories = await self._extract_hard_comparative_task_memory(success_steps, failure_steps,
                                                                                 similarity_score)
                 comparative_task_memories.extend(hard_task_memories)
 
@@ -73,7 +75,7 @@ class ComparativeExtractionOp(BaseLLMOp):
         """Get trajectory score"""
         return trajectory.score
 
-    def _extract_soft_comparative_task_memory(self, higher_traj: Trajectory, lower_traj: Trajectory) -> List[
+    async def _extract_soft_comparative_task_memory(self, higher_traj: Trajectory, lower_traj: Trajectory) -> List[
         BaseMemory]:
         """Extract soft comparative task memory (high score vs low score)"""
         higher_steps = self._get_trajectory_steps(higher_traj)
@@ -105,9 +107,9 @@ class ComparativeExtractionOp(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)
 
-    def _extract_hard_comparative_task_memory(self, success_steps: List[Message],
+    async def _extract_hard_comparative_task_memory(self, success_steps: List[Message],
                                               failure_steps: List[Message], similarity_score: float) -> List[
         BaseMemory]:
         """Extract hard comparative task memory (success vs failure)"""
@@ -134,7 +136,7 @@ class ComparativeExtractionOp(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)
 
     @staticmethod
     def _get_trajectory_steps(trajectory: Trajectory) -> List[Message]:

reme_ai/summary/task/failure_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 FailureExtractionOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Extract task memories from failed trajectories"""
         failure_trajectories: List[Trajectory] = self.context.get("failure_trajectories", [])
 
@@ -29,11 +31,11 @@ class FailureExtractionOp(BaseLLMOp):
             if hasattr(trajectory, 'segments') and trajectory.segments:
                 # Process segmented step sequences
                 for segment in trajectory.segments:
-                    task_memories = self._extract_failure_task_memory_from_steps(segment, trajectory)
+                    task_memories = await self._extract_failure_task_memory_from_steps(segment, trajectory)
                     failure_task_memories.extend(task_memories)
             else:
                 # Process entire trajectory
-                task_memories = self._extract_failure_task_memory_from_steps(trajectory.messages, trajectory)
+                task_memories = await self._extract_failure_task_memory_from_steps(trajectory.messages, trajectory)
                 failure_task_memories.extend(task_memories)
 
         logger.info(f"Extracted {len(failure_task_memories)} failure task memories")
@@ -41,7 +43,7 @@ class FailureExtractionOp(BaseLLMOp):
         # Add task memories to context
         self.context.failure_task_memories = failure_task_memories
 
-    def _extract_failure_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
+    async def _extract_failure_task_memory_from_steps(self, steps: List[Message], trajectory: Trajectory) -> List[BaseMemory]:
         """Extract task memory from failed step sequences"""
         step_content = merge_messages_content(steps)
         context = get_trajectory_context(trajectory, steps)
@@ -70,4 +72,4 @@ class FailureExtractionOp(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)