ReMe/reme_ai/retrieve/task/rewrite_memory_op.py

"""Memory rewriting operation module.

This module provides functionality to rewrite and format retrieved memories
into context messages that can be used by LLMs for task completion.
"""

import json
import re
from typing import List

from flowllm.core.context import C
from flowllm.core.enumeration import Role
from flowllm.core.op import BaseAsyncOp
from flowllm.core.schema import Message
from loguru import logger

from reme_ai.schema.memory import BaseMemory


@C.register_op()
class RewriteMemoryOp(BaseAsyncOp):
    """Generate and rewrite context messages from reranked experiences.

    This operation takes reranked memories and formats them into context messages
    that can be used by LLMs. It optionally uses LLM-based rewriting to make
    the context more relevant and actionable for the current task.
    """

    file_path: str = __file__

    async def async_execute(self):
        """Execute the memory rewrite operation.

        Retrieves memories from context metadata, formats them, and optionally
        rewrites them using LLM to make them more relevant for the current query.
        Stores the rewritten context in the response answer field.
        """
        memory_list: List[BaseMemory] = self.context.response.metadata["memory_list"]
        query: str = self.context.query
        messages: List[Message] = [Message(**x) if isinstance(x, dict) else x for x in self.context.get("messages", [])]

        if not memory_list:
            logger.info("No reranked memories to rewrite")
            self.context.response.answer = ""
            return

        logger.info(f"Generating context from {len(memory_list)} memories")

        # Generate initial context message
        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]

    async def _generate_context_message(self, query: str, messages: List[Message], memories: List[BaseMemory]) -> str:
        """Generate context message from retrieved memories.

        Args:
            query: The current query string.
            messages: List of conversation messages for context.
            memories: List of retrieved memories to format.

        Returns:
            Formatted context string, optionally rewritten by LLM.
        """
        if not memories:
            return ""

        try:
            logger.info("memories")
            # Format retrieved memories
            formatted_memories = self._format_memories_for_context(memories)

            if self.op_params.get("enable_llm_rewrite", True):
                context_content = await self._rewrite_context(query, formatted_memories, messages)
            else:
                context_content = formatted_memories

            return context_content

        except Exception as e:
            logger.error(f"Error generating context message: {e}")
            return self._format_memories_for_context(memories)

    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.

        Args:
            query: The current query string.
            context_content: The formatted context content to rewrite.
            messages: List of conversation messages for additional context.

        Returns:
            Rewritten context string optimized for the current task.
        """
        if not context_content:
            return context_content

        try:
            # Extract current context
            current_context = self._extract_context(messages)

            prompt = self.prompt_format(
                prompt_name="memory_rewrite_prompt",
                current_query=query,
                current_context=current_context,
                original_context=context_content,
            )

            response = await self.llm.achat([Message(role=Role.USER, content=prompt)])

            # Extract rewritten context
            rewritten_context = self._parse_json_response(response.content, "rewritten_context")

            if rewritten_context and rewritten_context.strip():
                logger.info("Context successfully rewritten for current task")
                return rewritten_context.strip()

            return context_content

        except Exception as e:
            logger.error(f"Error in context rewriting: {e}")
            return context_content

    @staticmethod
    def _format_memories_for_context(memories: List[BaseMemory]) -> str:
        """Format memories for context generation.

        Args:
            memories: List of memories to format.

        Returns:
            Formatted string containing all memories with their conditions and content.
        """
        formatted_memories = []

        for i, memory in enumerate(memories, 1):
            condition = memory.when_to_use
            memory_content = memory.content
            memory_text = f"Memory {i} :\n When to use: {condition}\n Content: {memory_content}\n"

            formatted_memories.append(memory_text)

        return "\n".join(formatted_memories)

    @staticmethod
    def _extract_context(messages: List[Message]) -> str:
        """Extract relevant context from messages.

        Args:
            messages: List of conversation messages.

        Returns:
            Formatted string containing recent conversation context.
        """
        if not messages:
            return ""

        context_parts = []

        # Add recent messages if available
        recent_messages = messages[-3:]  # Last 3 messages
        message_summaries = []
        for message in recent_messages:
            content = message.content[:300] + "..." if len(message.content) > 300 else message.content
            message_summaries.append(f"- {message.role.value}: {content}")

        if message_summaries:
            context_parts.append("Recent conversation:\n" + "\n".join(message_summaries))

        return "\n\n".join(context_parts)

    @staticmethod
    def _parse_json_response(response: str, key: str) -> str:
        """Parse JSON response to extract specific key.

        Args:
            response: The response string that may contain JSON.
            key: The key to extract from the JSON object.

        Returns:
            The value associated with the key, or the response string if parsing fails.
        """
        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 key in parsed:
                    return parsed[key]

            # Fallback: try to parse the entire response as JSON
            parsed = json.loads(response)
            if isinstance(parsed, dict) and key in parsed:
                return parsed[key]

        except json.JSONDecodeError:
            logger.warning(f"Failed to parse JSON response for key '{key}', using raw response")
            # If JSON parsing fails, return the response as-is for fallback
            return response.strip()

        return ""