reme-ai 0.1.0__py3-none-any.whl

reme_ai/retrieve/personal/fuse_rerank_op.py

@@ -0,0 +1,180 @@
+from typing import Dict, List
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.constants.common_constants import EXTRACT_TIME_DICT
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class FuseRerankOp(BaseLLMOp):
+    """
+    Reranks the memory nodes by scores, types, and temporal relevance. Formats the top-K reranked nodes to print.
+    """
+    file_path: str = __file__
+
+    @staticmethod
+    def match_memory_time(extract_time_dict: Dict[str, str], memory: BaseMemory):
+        """
+        Determines whether the memory is relevant based on time matching.
+        """
+        if extract_time_dict:
+            match_event_flag = True
+            for k, v in extract_time_dict.items():
+                event_value = memory.metadata.get(f"event_{k}", "")
+                if event_value in ["-1", v]:
+                    continue
+                else:
+                    match_event_flag = False
+                    break
+
+            match_msg_flag = True
+            for k, v in extract_time_dict.items():
+                msg_value = memory.metadata.get(f"msg_{k}", "")
+                if msg_value == v:
+                    continue
+                else:
+                    match_msg_flag = False
+                    break
+        else:
+            match_event_flag = False
+            match_msg_flag = False
+
+        memory.metadata["match_event_flag"] = str(int(match_event_flag))
+        memory.metadata["match_msg_flag"] = str(int(match_msg_flag))
+        return match_event_flag, match_msg_flag
+
+    def execute(self):
+        """
+        Executes the reranking process on memories considering their scores, types, and temporal relevance.
+
+        This method performs the following steps:
+        1. Retrieves extraction time data and a list of memories from the context.
+        2. Reranks memories based on a combination of their original score, type,
+           and temporal alignment with extracted events/messages.
+        3. Selects the top-K reranked memories according to the predefined threshold.
+        4. Formats the final list of memories for output.
+        5. Sets both response.answer and response.metadata["memory_list"]
+        """
+        # Get operation parameters
+        fuse_score_threshold = self.op_params.get("fuse_score_threshold", 0.1)
+        fuse_ratio_dict = self.op_params.get("fuse_ratio_dict", {
+            "conversation": 0.5,
+            "observation": 1,
+            "obs_customized": 1.2,
+            "insight": 2.0
+        })
+        fuse_time_ratio = self.op_params.get("fuse_time_ratio", 2.0)
+        output_memory_max_count = self.op_params.get("output_memory_max_count", 5)
+
+        # Parse input parameters from the context
+        extract_time_dict: Dict[str, str] = self.context.get(EXTRACT_TIME_DICT, {})
+        memory_list: List[BaseMemory] = self.context.response.metadata.get("memory_list", [])
+
+        # Check if memories are available; warn and return if not
+        if not memory_list:
+            logger.warning("No memories available for fuse reranking")
+            self.context.response.answer = ""
+            self.context.response.metadata["memory_list"] = []
+            return
+
+        logger.info(f"Fuse reranking {len(memory_list)} memories with time dict: {bool(extract_time_dict)}")
+
+        # Perform reranking based on score, type, and time relevance
+        reranked_memories = self._apply_fuse_reranking(
+            memory_list, extract_time_dict, fuse_score_threshold,
+            fuse_ratio_dict, fuse_time_ratio
+        )
+
+        # Sort and select top-k memories
+        reranked_memories = sorted(reranked_memories,
+                                   key=lambda x: x.score or 0.0,
+                                   reverse=True)[:output_memory_max_count]
+
+        logger.info(f"Final reranked memories: {len(reranked_memories)}")
+
+        # Format memories for output
+        formatted_memories = self._format_memories_for_output(reranked_memories)
+
+        # Store results in context - both answer and metadata as required
+        self.context.response.metadata["memory_list"] = reranked_memories
+        self.context.response.answer = "\n".join(formatted_memories)
+
+    def _apply_fuse_reranking(self,
+                              memory_list: List[BaseMemory],
+                              extract_time_dict: Dict[str, str],
+                              fuse_score_threshold: float,
+                              fuse_ratio_dict: Dict[str, float],
+                              fuse_time_ratio: float) -> List[BaseMemory]:
+        """Apply fuse reranking logic to memories"""
+        reranked_memories = []
+
+        for memory in memory_list:
+            # Skip memories below the fuse score threshold
+            memory_score = memory.score or 0.0
+            if memory_score < fuse_score_threshold:
+                continue
+
+            # Calculate type-based adjustment factor
+            memory_type = memory.metadata.get("memory_type", "default")
+            if memory_type not in fuse_ratio_dict:
+                logger.debug(f"Memory type '{memory_type}' not in fuse_ratio_dict, using default 0.1")
+            type_ratio: float = fuse_ratio_dict.get(memory_type, 0.1)
+
+            # Determine time relevance adjustment factor
+            match_event_flag, match_msg_flag = self.match_memory_time(extract_time_dict, memory)
+            time_ratio: float = fuse_time_ratio if match_event_flag or match_msg_flag else 1.0
+
+            # Apply reranking score adjustments
+            original_score = memory_score
+            memory.score = memory_score * type_ratio * time_ratio
+
+            logger.debug(f"Memory reranked: {original_score:.3f} -> {memory.score:.3f} "
+                         f"(type={type_ratio}, time={time_ratio})")
+
+            reranked_memories.append(memory)
+
+        return reranked_memories
+
+    def _format_memories_for_output(self, memories: List[BaseMemory]) -> List[str]:
+        """Format memories for final output"""
+        formatted_memories = []
+
+        for memory in memories:
+            # Log reranking details
+            logger.info(f"Final memory: Score={memory.score:.3f}, "
+                        f"Event={memory.metadata.get('match_event_flag', '0')}, "
+                        f"Msg={memory.metadata.get('match_msg_flag', '0')}, "
+                        f"Content={memory.content[:50]}...")
+
+            # Format memory with timestamp if available
+            formatted_content = self._format_memory_with_timestamp(memory, self.language)
+            formatted_memories.append(formatted_content)
+
+        return formatted_memories
+
+    @staticmethod
+    def _format_memory_with_timestamp(memory, language: str = "en") -> str:
+        """
+        Format memory content with timestamp if available.
+
+        Args:
+            memory: Memory object
+            language: Language for formatting
+
+        Returns:
+            Formatted memory content string
+        """
+        try:
+            if hasattr(memory, 'timestamp') and memory.timestamp:
+                from reme_ai.utils.datetime_handler import DatetimeHandler
+                dt_handler = DatetimeHandler(memory.timestamp)
+                datetime_str = dt_handler.datetime_format("%Y-%m-%d %H:%M:%S")
+                weekday = dt_handler.get_dt_info_dict(language)["weekday"]
+                return f"[{datetime_str} {weekday}] {memory.content}"
+            else:
+                return memory.content
+        except Exception as e:
+            logger.warning(f"Failed to format memory with timestamp: {e}")
+            return memory.content

reme_ai/retrieve/personal/print_memory_op.py

@@ -0,0 +1,131 @@
+from typing import List
+
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class PrintMemoryOp(BaseOp):
+    """
+    Formats the memories to print.
+    """
+    file_path: str = __file__
+
+    def execute(self):
+        """
+        Executes the primary function, it involves:
+        1. Fetches the memories.
+        2. Formats them for printing.
+        3. Set the formatted string back into the context
+        """
+        # Get memory list from context
+        memory_list: List[BaseMemory] = self.context.response.metadata.get("memory_list", [])
+
+        if not memory_list:
+            logger.info("No memories to print")
+            self.context.response.answer = "No memories found."
+            return
+
+        logger.info(f"Formatting {len(memory_list)} memories for printing")
+
+        # Format memories for printing
+        formatted_memories = self._format_memories_for_print(memory_list)
+
+        # Store result in context
+        self.context.response.answer = formatted_memories
+        logger.info(f"Formatted memories: {formatted_memories}")
+
+    @staticmethod
+    def _format_memories_for_print(memories: List[BaseMemory]) -> str:
+        """Format memories for printing"""
+        if not memories:
+            return "No memories available."
+
+        formatted_memories = []
+
+        for i, memory in enumerate(memories, 1):
+            memory_text = f"Memory {i}:\n"
+            memory_text += f"  When to use: {memory.when_to_use}\n"
+            memory_text += f"  Content: {memory.content}\n"
+
+            # Add additional metadata if available
+            if hasattr(memory, 'metadata') and memory.metadata:
+                metadata_items = []
+                for key, value in memory.metadata.items():
+                    if key not in ['when_to_use', 'content']:
+                        metadata_items.append(f"{key}: {value}")
+                if metadata_items:
+                    memory_text += f"  Metadata: {', '.join(metadata_items)}\n"
+
+            formatted_memories.append(memory_text)
+
+        return "\n".join(formatted_memories)
+
+    @staticmethod
+    def format_memories_for_output(memories: List) -> str:
+        """
+        Format memory list for output string.
+        
+        Args:
+            memories: List of memory objects
+            
+        Returns:
+            Formatted string
+        """
+        if not memories:
+            return ""
+
+        formatted_parts = []
+        for i, memory in enumerate(memories, 1):
+            when_to_use = getattr(memory, 'when_to_use', '') or memory.get('when_to_use', '')
+            content = getattr(memory, 'content', '') or memory.get('content', '')
+
+            part = f"Memory {i}:\n"
+            if when_to_use:
+                part += f"When to use: {when_to_use}\n"
+            if content:
+                part += f"Content: {content}\n"
+
+            formatted_parts.append(part)
+
+        return "\n".join(formatted_parts)
+
+    @staticmethod
+    def format_memories_for_simple_output(memories: List) -> str:
+        """
+        Format memory list for simple flow output.
+        
+        Args:
+            memories: List of memory objects
+            
+        Returns:
+            Formatted string suitable for response.answer
+        """
+        if not memories:
+            return "No relevant memories found."
+
+        content_parts = ["Previous Memory"]
+
+        for memory in memories:
+            # Safely get field values
+            when_to_use = getattr(memory, 'when_to_use', '') or memory.get('when_to_use', '')
+            content = getattr(memory, 'content', '') or memory.get('content', '')
+
+            # Skip memories with empty content
+            if not content:
+                continue
+
+            # Format individual memory
+            memory_text = f"- when_to_use: {when_to_use}\n  content: {content}"
+            content_parts.append(memory_text)
+
+        # If no valid memories, return empty message
+        if len(content_parts) == 1:  # Only title
+            return "No relevant memories with valid content found."
+
+        content_parts.append("\nPlease consider the helpful parts from these in answering the question, "
+                             "to make the response more comprehensive and substantial.")
+
+        return "\n".join(content_parts)

reme_ai/retrieve/personal/read_message_op.py

@@ -0,0 +1,52 @@
+from typing import List
+
+from flowllm import C, BaseOp
+from flowllm.schema.message import Message
+from loguru import logger
+
+
+@C.register_op()
+class ReadMessageOp(BaseOp):
+    """
+    Fetches unmemorized chat messages.
+    """
+    file_path: str = __file__
+
+    def execute(self):
+        """
+        Executes the primary function to fetch unmemorized chat messages.
+        """
+        # Get chat messages from context
+        chat_messages = self.context.chat_messages
+        target_name = self.context.target_name
+        contextual_msg_max_count = self.op_params.get('contextual_msg_max_count', 10)
+
+        chat_messages_not_memorized: List[List[Message]] = []
+        for messages in chat_messages:
+            if not messages:
+                continue
+
+            if hasattr(messages[0], 'memorized') and messages[0].memorized:
+                continue
+
+            contain_flag = False
+
+            for msg in messages:
+                if hasattr(msg, 'role_name') and msg.role_name == target_name:
+                    contain_flag = True
+                    break
+
+            if contain_flag:
+                chat_messages_not_memorized.append(messages)
+
+        chat_message_scatter = []
+        for messages in chat_messages_not_memorized[-contextual_msg_max_count:]:
+            chat_message_scatter.extend(messages)
+
+        # Sort by time_created if available
+        if chat_message_scatter and hasattr(chat_message_scatter[0], 'time_created'):
+            chat_message_scatter.sort(key=lambda _: _.time_created)
+
+        # Store result in context
+        self.context.chat_messages = chat_message_scatter
+        logger.info(f"Retrieved {len(chat_message_scatter)} unmemorized chat messages")

reme_ai/retrieve/personal/retrieve_memory_op.py

@@ -0,0 +1,13 @@
+from flowllm import C
+
+from reme_ai.vector_store import RecallVectorStoreOp
+
+
+@C.register_op()
+class RetrieveMemoryOp(RecallVectorStoreOp):
+    """
+    Retrieves memories based on specified criteria such as status, type, and timestamp.
+    Processes these memories concurrently, sorts them by similarity, and logs the activity,
+    facilitating efficient memory retrieval operations within a given scope.
+    """
+    file_path: str = __file__

reme_ai/retrieve/personal/semantic_rank_op.py

@@ -0,0 +1,170 @@
+import json
+import re
+from typing import List
+
+from flowllm import C, BaseLLMOp
+from loguru import logger
+
+from reme_ai.schema import Message, Role
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class SemanticRankOp(BaseLLMOp):
+    """
+    The SemanticRankOp class processes queries by retrieving memory nodes,
+    removing duplicates, ranking them based on semantic relevance using a model,
+    assigning scores, sorting the nodes, and storing the ranked nodes back,
+    while logging relevant information.
+    """
+    file_path: str = __file__
+
+    def execute(self):
+        """
+        Executes the primary workflow of the SemanticRankOp which includes:
+        - Retrieves query and memory list from context.
+        - Removes duplicate memories.
+        - Ranks memories semantically using LLM.
+        - Assigns scores to memories.
+        - Sorts memories by score.
+        - Saves the ranked memories back to context.
+
+        If no memories are retrieved or if the ranking fails,
+        appropriate warnings are logged.
+        """
+        # Get memory list from context - previous op guarantees this exists
+        memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
+        query: str = self.context.query
+
+        # Get parameters from op_params
+        enable_ranker: bool = self.op_params.get("enable_ranker", True)
+        output_memory_max_count: int = self.op_params.get("output_memory_max_count", 10)
+
+        if not memory_list:
+            logger.warning("Memory list is empty!")
+            return
+
+        logger.info(f"Semantic ranking {len(memory_list)} memories for query: {query[:100]}...")
+
+        if not enable_ranker or len(memory_list) <= output_memory_max_count:
+            # Use original scores if ranker is disabled or memory count is small
+            logger.info("Skipping semantic ranking - using original scores")
+        else:
+            # Remove duplicates based on content
+            memory_dict = {memory.content.strip(): memory for memory in memory_list if memory.content.strip()}
+            memory_list = list(memory_dict.values())
+            logger.info(f"After deduplication: {len(memory_list)} memories")
+
+            # Perform semantic ranking using LLM
+            ranked_memories = self._semantic_rank_memories(query, memory_list)
+            if ranked_memories:
+                memory_list = ranked_memories
+
+        # Sort by score
+        memory_list = sorted(memory_list, key=lambda m: m.score, reverse=True)
+
+        # Log top ranked memories
+        logger.info(f"Semantic ranking completed for query: {query[:50]}...")
+        for i, memory in enumerate(memory_list[:5]):  # Log top 5
+            logger.info(f"Top {i + 1}: Score={memory.score:.3f}, Content={memory.content[:80]}...")
+
+        # Save ranked memories back to context
+        self.context.response.metadata["memory_list"] = memory_list
+
+    def _semantic_rank_memories(self, query: str, memories: List[BaseMemory]) -> List[BaseMemory]:
+        """
+        Use LLM to semantically rank memories based on relevance to the query
+        """
+        if not memories:
+            return memories
+
+        # Format memories for ranking
+        formatted_memories = SemanticRankOp.format_memories_for_llm_ranking(memories)
+
+        # Create prompt for semantic ranking
+        prompt = f"""Given the query: "{query}"
+
+Please rank the following memories by their semantic relevance to the query. 
+Rate each memory on a scale of 0.0 to 1.0 where 1.0 is most relevant.
+
+Memories:
+{formatted_memories}
+
+Please respond in JSON format:
+{{"rankings": [{{"index": 0, "score": 0.8}}, {{"index": 1, "score": 0.6}}, ...]}}"""
+
+        response = self.llm.chat([Message(role=Role.USER, content=prompt)])
+
+        if not response or not response.content:
+            logger.warning("LLM ranking failed, using original order")
+            return memories
+
+        # Parse and apply ranking results
+        rankings = SemanticRankOp.parse_llm_ranking_response(response.content)
+
+        if rankings:
+            applied_count = SemanticRankOp.apply_semantic_scores_to_memories(memories, rankings)
+            logger.info(f"Successfully applied semantic rankings to {applied_count} memories")
+        else:
+            logger.warning("Failed to parse ranking results")
+
+        return memories
+
+    @staticmethod
+    def parse_llm_ranking_response(response: str) -> List[dict]:
+        """Parse LLM ranking response to extract rankings."""
+        try:
+            # Try to extract JSON blocks
+            json_pattern = r'```json\s*([\s\S]*?)\s*```'
+            json_blocks = re.findall(json_pattern, response)
+
+            if json_blocks:
+                parsed = json.loads(json_blocks[0])
+                if isinstance(parsed, dict) and "rankings" in parsed:
+                    return parsed["rankings"]
+
+            # Fallback: try to parse the entire response as JSON
+            parsed = json.loads(response)
+            if isinstance(parsed, dict) and "rankings" in parsed:
+                return parsed["rankings"]
+
+        except json.JSONDecodeError:
+            logger.warning("Failed to parse ranking response as JSON")
+
+        return []
+
+    @staticmethod
+    def apply_semantic_scores_to_memories(memories: List, rankings: List[dict]) -> int:
+        """Apply semantic ranking scores to memory objects."""
+        applied_count = 0
+
+        for ranking in rankings:
+            idx = ranking.get("index", -1)
+            score = ranking.get("score", 0.0)
+
+            if 0 <= idx < len(memories):
+                # Set score on memory object
+                if hasattr(memories[idx], 'score'):
+                    memories[idx].score = score
+                    applied_count += 1
+                else:
+                    # Add score as metadata if score attribute doesn't exist
+                    if not hasattr(memories[idx], 'metadata'):
+                        memories[idx].metadata = {}
+                    memories[idx].metadata['semantic_score'] = score
+                    applied_count += 1
+
+        return applied_count
+
+    @staticmethod
+    def format_memories_for_llm_ranking(memories: List) -> str:
+        """Format memories for LLM ranking input."""
+        formatted_memories = []
+
+        for i, memory in enumerate(memories):
+            memory_text = f"Memory {i}:\n"
+            memory_text += f"When to use: {memory.when_to_use}\n"
+            memory_text += f"Content: {memory.content}\n"
+            formatted_memories.append(memory_text)
+
+        return "\n---\n".join(formatted_memories)

reme_ai/retrieve/personal/set_query_op.py

@@ -0,0 +1,37 @@
+import datetime
+from typing import Tuple
+
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.constants.common_constants import QUERY_WITH_TS
+
+
+@C.register_op()
+class SetQueryOp(BaseOp):
+    """
+    The `SetQueryOp` class is responsible for setting a query and its associated timestamp
+    into the context, utilizing either provided parameters or details from the context.
+    """
+
+    def execute(self):
+        """
+        Executes the operation's primary function, which involves determining the query and its
+        timestamp, then storing these values within the context.
+
+        Input requirement: self.context.query must exist (flow input requirement)
+        """
+        # Flow guarantees query exists - use it directly
+        query: str = self.context.query
+        timestamp: int = int(datetime.datetime.now().timestamp())
+
+        # Set timestamp if provided in op_params
+        _timestamp = self.op_params.get("timestamp")
+        if _timestamp and isinstance(_timestamp, int):
+            timestamp = _timestamp
+
+        # Store the query and its timestamp in the context
+        query_with_ts: Tuple[str, int] = (query, timestamp)
+        self.context[QUERY_WITH_TS] = query_with_ts
+
+        logger.info(f"Set query with timestamp: query='{query}', timestamp={timestamp}")

reme_ai/retrieve/task/__init__.py

@@ -0,0 +1,4 @@
+from .build_query_op import BuildQueryOp
+from .merge_memory_op import MergeMemoryOp
+from .rerank_memory_op import RerankMemoryOp
+from .rewrite_memory_op import RewriteMemoryOp

reme_ai/retrieve/task/build_query_op.py

@@ -0,0 +1,38 @@
+from flowllm import C, BaseLLMOp
+from flowllm.utils.llm_utils import merge_messages_content
+from loguru import logger
+
+from reme_ai.schema import Message, Role
+
+
+@C.register_op()
+class BuildQueryOp(BaseLLMOp):
+    file_path: str = __file__
+
+    def execute(self):
+        if "query" in self.context:
+            query = self.context.query
+
+        elif "messages" in self.context:
+            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)])
+                query = message.content
+
+            else:
+                context_parts = []
+                message_summaries = []
+                for message in self.context.messages[-3:]:  # Last 3 messages
+                    content = message.content[:200] + "..." if len(message.content) > 200 else message.content
+                    message_summaries.append(f"- {message.role.value}: {content}")
+                if message_summaries:
+                    context_parts.append("Recent messages:\n" + "\n".join(message_summaries))
+
+                query = "\n\n".join(context_parts)
+
+        else:
+            raise RuntimeError("query or messages is required!")
+
+        logger.info(f"build.query={query}")
+        self.context.query = query

reme_ai/retrieve/task/merge_memory_op.py

@@ -0,0 +1,27 @@
+from typing import List
+
+from flowllm import C, BaseOp
+from loguru import logger
+
+from reme_ai.schema.memory import BaseMemory
+
+
+@C.register_op()
+class MergeMemoryOp(BaseOp):
+
+    def execute(self):
+        memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
+
+        if not memory_list:
+            return
+
+        content_collector = ["Previous Memory"]
+        for memory in memory_list:
+            if not memory.content:
+                continue
+
+            content_collector.append(f"- {memory.when_to_use} {memory.content}\n")
+        content_collector.append("Please consider the helpful parts from these in answering the question, "
+                                 "to make the response more comprehensive and substantial.")
+        self.context.response.answer = "\n".join(content_collector)
+        logger.info(f"response.answer={self.context.response.answer}")