refactor(llm): update LLM-related ops to use async execution

- Modify LLM-related ops to use async_execute instead of execute
- Replace chat method with achat for asynchronous LLM calls
- Update method signatures and return types to support async operations
- This change affects multiple files across the project

reme_ai/retrieve/personal/retrieve_memory_op.py

@@ -1,15 +1,14 @@
 from typing import List
 
-from flowllm import C
+from flowllm import C, BaseLLMOp
 from flowllm.schema.vector_node import VectorNode
 from loguru import logger
 
 from reme_ai.schema.memory import BaseMemory, vector_node_to_memory
-from reme_ai.vector_store import RecallVectorStoreOp
 
 
 @C.register_op()
-class RetrieveMemoryOp(RecallVectorStoreOp):
+class RetrieveMemoryOp(BaseLLMOp):
     """
     Retrieves memories based on specified criteria such as status, type, and timestamp.
     Processes these memories concurrently, sorts them by similarity, and logs the activity,

reme_ai/retrieve/task/build_query_op.py

@@ -9,7 +9,7 @@ from reme_ai.schema import Message, Role
 class BuildQueryOp(BaseLLMOp):
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         if "query" in self.context:
             query = self.context.query
 
@@ -17,7 +17,7 @@ class BuildQueryOp(BaseLLMOp):
             if self.op_params.get("enable_llm_build", True):
                 execution_process = merge_messages_content(self.context.messages)
                 prompt = self.prompt_format(prompt_name="query_build", execution_process=execution_process)
-                message = self.llm.chat(messages=[Message(role=Role.USER, content=prompt)])
+                message = await self.llm.achat(messages=[Message(role=Role.USER, content=prompt)])
                 query = message.content
 
             else:

reme_ai/retrieve/task/merge_memory_op.py

@@ -9,7 +9,7 @@ from reme_ai.schema.memory import BaseMemory
 @C.register_op()
 class MergeMemoryOp(BaseOp):
 
-    def execute(self):
+    async def async_execute(self):
         memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
 
         if not memory_list:

reme_ai/retrieve/task/rerank_memory_op.py

@@ -17,7 +17,7 @@ class RerankMemoryOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Execute rerank operation"""
         memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
         retrieval_query: str = self.context.query
@@ -36,7 +36,7 @@ class RerankMemoryOp(BaseLLMOp):
 
         # Step 1: LLM reranking (optional)
         if enable_llm_rerank:
-            memory_list = self._llm_rerank(retrieval_query, memory_list)
+            memory_list = await self._llm_rerank(retrieval_query, memory_list)
             logger.info(f"After LLM reranking: {len(memory_list)} memories")
 
         # Step 2: Score-based filtering (optional)
@@ -51,7 +51,7 @@ class RerankMemoryOp(BaseLLMOp):
         # Store results in context
         self.context.response.metadata["memory_list"] = reranked_memories
 
-    def _llm_rerank(self, query: str, candidates: List[BaseMemory]) -> List[BaseMemory]:
+    async def _llm_rerank(self, query: str, candidates: List[BaseMemory]) -> List[BaseMemory]:
         """LLM-based reranking of candidate experiences"""
         if not candidates:
             return candidates
@@ -65,7 +65,7 @@ class RerankMemoryOp(BaseLLMOp):
             candidates=candidates_text,
             num_candidates=len(candidates))
 
-        response = self.llm.chat([Message(role=Role.USER, content=prompt)])
+        response = await self.llm.achat([Message(role=Role.USER, content=prompt)])
 
         # Parse reranking results
         reranked_indices = self._parse_rerank_response(response.content)

reme_ai/retrieve/task/rewrite_memory_op.py

@@ -17,7 +17,7 @@ class RewriteMemoryOp(BaseLLMOp):
     """
     file_path: str = __file__
 
-    def execute(self):
+    async def async_execute(self):
         """Execute rewrite operation"""
         memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
         query: str = self.context.query
@@ -32,13 +32,13 @@ class RewriteMemoryOp(BaseLLMOp):
         logger.info(f"Generating context from {len(memory_list)} memories")
 
         # Generate initial context message
-        rewritten_memory = self._generate_context_message(query, messages, memory_list)
+        rewritten_memory = await self._generate_context_message(query, messages, memory_list)
 
         # Store results in context
         self.context.response.answer = rewritten_memory
         self.context.response.metadata["memory_list"] = [memory.model_dump() for memory in memory_list]
 
-    def _generate_context_message(self, query: str, messages: List[Message], memories: List[BaseMemory]) -> str:
+    async def _generate_context_message(self, query: str, messages: List[Message], memories: List[BaseMemory]) -> str:
         """Generate context message from retrieved memories"""
         if not memories:
             return ""
@@ -49,7 +49,7 @@ class RewriteMemoryOp(BaseLLMOp):
             formatted_memories = self._format_memories_for_context(memories)
 
             if self.op_params.get("enable_llm_rewrite", True):
-                context_content = self._rewrite_context(query, formatted_memories, messages)
+                context_content = await self._rewrite_context(query, formatted_memories, messages)
             else:
                 context_content = formatted_memories
 
@@ -59,7 +59,7 @@ class RewriteMemoryOp(BaseLLMOp):
             logger.error(f"Error generating context message: {e}")
             return self._format_memories_for_context(memories)
 
-    def _rewrite_context(self, query: str, context_content: str, messages: List[Message]) -> str:
+    async def _rewrite_context(self, query: str, context_content: str, messages: List[Message]) -> str:
         """LLM-based context rewriting to make experiences more relevant and actionable"""
         if not context_content:
             return context_content
@@ -74,7 +74,7 @@ class RewriteMemoryOp(BaseLLMOp):
                 current_context=current_context,
                 original_context=context_content)
 
-            response = self.llm.chat([Message(role=Role.USER, content=prompt)])
+            response = await self.llm.achat([Message(role=Role.USER, content=prompt)])
 
             # Extract rewritten context
             rewritten_context = self._parse_json_response(response.content, "rewritten_context")

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.
         

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