From 6bb3ab90fcc640787f354bf318c25e2e493e4a82 Mon Sep 17 00:00:00 2001 From: "jinli.yl" Date: Sat, 6 Sep 2025 20:27:25 +0800 Subject: [PATCH] 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 | 5 ++--- reme_ai/retrieve/task/build_query_op.py | 4 ++-- reme_ai/retrieve/task/merge_memory_op.py | 2 +- reme_ai/retrieve/task/rerank_memory_op.py | 8 ++++---- reme_ai/retrieve/task/rewrite_memory_op.py | 12 ++++++------ reme_ai/summary/personal/contra_repeat_op.py | 4 ++-- reme_ai/summary/personal/get_observation_op.py | 8 ++++---- .../summary/personal/get_observation_with_time_op.py | 8 ++++---- .../summary/personal/get_reflection_subject_op.py | 8 ++++---- reme_ai/summary/personal/info_filter_op.py | 8 ++++---- reme_ai/summary/personal/load_today_memory_op.py | 2 +- reme_ai/summary/personal/long_contra_repeat_op.py | 8 ++++---- reme_ai/summary/personal/update_insight_op.py | 8 ++++---- 13 files changed, 42 insertions(+), 43 deletions(-) diff --git a/reme_ai/retrieve/personal/retrieve_memory_op.py b/reme_ai/retrieve/personal/retrieve_memory_op.py index d22f133a..6f96815b 100644 --- a/reme_ai/retrieve/personal/retrieve_memory_op.py +++ b/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, diff --git a/reme_ai/retrieve/task/build_query_op.py b/reme_ai/retrieve/task/build_query_op.py index 0cf80a88..80e11b92 100644 --- a/reme_ai/retrieve/task/build_query_op.py +++ b/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: diff --git a/reme_ai/retrieve/task/merge_memory_op.py b/reme_ai/retrieve/task/merge_memory_op.py index 75b776fb..ab1e8e88 100644 --- a/reme_ai/retrieve/task/merge_memory_op.py +++ b/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: diff --git a/reme_ai/retrieve/task/rerank_memory_op.py b/reme_ai/retrieve/task/rerank_memory_op.py index ded5b62a..fec119a8 100644 --- a/reme_ai/retrieve/task/rerank_memory_op.py +++ b/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) diff --git a/reme_ai/retrieve/task/rewrite_memory_op.py b/reme_ai/retrieve/task/rewrite_memory_op.py index aa36f304..eae0b0a1 100644 --- a/reme_ai/retrieve/task/rewrite_memory_op.py +++ b/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") diff --git a/reme_ai/summary/personal/contra_repeat_op.py b/reme_ai/summary/personal/contra_repeat_op.py index f4643a0f..9330e45c 100644 --- a/reme_ai/summary/personal/contra_repeat_op.py +++ b/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: diff --git a/reme_ai/summary/personal/get_observation_op.py b/reme_ai/summary/personal/get_observation_op.py index d126e41c..6c7abd1a 100644 --- a/reme_ai/summary/personal/get_observation_op.py +++ b/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]: diff --git a/reme_ai/summary/personal/get_observation_with_time_op.py b/reme_ai/summary/personal/get_observation_with_time_op.py index 0b474db2..cd3056b3 100644 --- a/reme_ai/summary/personal/get_observation_with_time_op.py +++ b/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""" diff --git a/reme_ai/summary/personal/get_reflection_subject_op.py b/reme_ai/summary/personal/get_reflection_subject_op.py index 232c8c1d..328963de 100644 --- a/reme_ai/summary/personal/get_reflection_subject_op.py +++ b/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""" diff --git a/reme_ai/summary/personal/info_filter_op.py b/reme_ai/summary/personal/info_filter_op.py index 97d054d4..6bc7d186 100644 --- a/reme_ai/summary/personal/info_filter_op.py +++ b/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""" diff --git a/reme_ai/summary/personal/load_today_memory_op.py b/reme_ai/summary/personal/load_today_memory_op.py index a4bfef15..46113196 100644 --- a/reme_ai/summary/personal/load_today_memory_op.py +++ b/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. diff --git a/reme_ai/summary/personal/long_contra_repeat_op.py b/reme_ai/summary/personal/long_contra_repeat_op.py index 6c1dfc35..bf77837d 100644 --- a/reme_ai/summary/personal/long_contra_repeat_op.py +++ b/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") diff --git a/reme_ai/summary/personal/update_insight_op.py b/reme_ai/summary/personal/update_insight_op.py index 57ccfafa..b20e2882 100644 --- a/reme_ai/summary/personal/update_insight_op.py +++ b/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