evoagentx/agents/long_term_memory_agent.py

import json
import asyncio
from uuid import uuid4
from pydantic import Field
from datetime import datetime
from typing import Optional, List, Tuple, Dict, Union

from evoagentx.agents import Agent
from evoagentx.core.parser import Parser
from evoagentx.models import BaseLLM
from evoagentx.core.logging import logger
from evoagentx.models import OpenAILLMConfig
from evoagentx.storages.base import StorageHandler
from evoagentx.core.message import Message, MessageType
from evoagentx.memory.memory_manager import MemoryManager
from evoagentx.memory.long_term_memory import LongTermMemory
from evoagentx.actions.action import Action, ActionInput, ActionOutput
from evoagentx.rag.rag_config import RAGConfig


class MemoryActionInput(ActionInput):
    user_prompt: str = Field(description="The user's input prompt")
    conversation_id: Optional[str] = Field(default=None, description="ID for tracking conversation")
    top_k: Optional[int] = Field(default=5, description="Number of memory results to retrieve")
    metadata_filters: Optional[Dict] = Field(default=None, description="Filters for memory retrieval")


class MemoryActionOutput(ActionOutput):
    response: str = Field(description="The agent's response based on memory and prompt")


class MemoryAction(Action):
    def __init__(
        self,
        name: str = "MemoryAction",
        description: str = "Action that processes user input with long-term memory context",
        prompt: str = "Based on the following context and user prompt, provide a relevant response:\n\nContext: {context}\n\nUser Prompt: {user_prompt}\n\n",
        inputs_format: ActionInput = None,
        outputs_format: ActionOutput = None,
        **kwargs
    ):
        inputs_format = inputs_format or MemoryActionInput
        outputs_format = outputs_format or MemoryActionOutput
        super().__init__(
            name=name,
            description=description,
            prompt=prompt,
            inputs_format=inputs_format,
            outputs_format=outputs_format,
            **kwargs
        )

    def execute(self, llm: BaseLLM | None = None, 
                inputs: Dict | None = None, 
                sys_msg: str | None = None, 
                return_prompt: bool = False, 
                memory_manager: Optional[MemoryManager] = None,
                **kwargs
    ) -> Parser | Tuple[Parser | str] | None:
        return asyncio.run(self.async_execute(llm, inputs, sys_msg, return_prompt, memory_manager, **kwargs))

    async def async_execute(
        self,
        llm: Optional["BaseLLM"] = None,
        inputs: Optional[Dict] = None,
        sys_msg: Optional[str] = None,
        return_prompt: bool = False,
        memory_manager: Optional[MemoryManager] = None,
        **kwargs
    ) -> Union[MemoryActionOutput, tuple]:
        if not memory_manager:
            logger.error("MemoryManager is required for MemoryAction execution")
            raise ValueError("MemoryManager is required for MemoryAction")

        action_input = self.inputs_format(**inputs)
        user_prompt = action_input.user_prompt
        conversation_id = action_input.conversation_id
        if not conversation_id:
            conversation_id = str(uuid4())
            logger.warning("No conversation_id provided; generated a new UUID4 for this session")
        top_k = action_input.top_k
        metadata_filters = action_input.metadata_filters

        message = await memory_manager.create_conversation_message(
            user_prompt=user_prompt,
            conversation_id=conversation_id,
            top_k=top_k,
            metadata_filters=metadata_filters
        )

        action_input_attrs = self.inputs_format.get_attrs()
        action_input_data = {attr: getattr(action_input, attr, "undefined") for attr in action_input_attrs}
        action_input_data["context"] = message.content
        prompt = self.prompt.format(**action_input_data)
        logger.info(f"The New Created Message by LongTermMemory:\n\n{prompt}")

        output = await llm.async_generate(
            prompt=prompt,
            system_message=sys_msg,
            parser=self.outputs_format,
            parse_mode='str'
        )
        
        response_message = Message(
            content=output.content,
            msg_type=MessageType.RESPONSE,
            timestamp=datetime.now().isoformat(),
            conversation_id=conversation_id,
            memory_ids=message.memory_ids
        )
        memory_ids = await memory_manager.handle_memory(
            action="add",
            data=response_message,
        )

        # Prepare the final output
        final_output = self.outputs_format(
            response=output.content,
            memory_ids=memory_ids
        )

        if return_prompt:
            return final_output, prompt
        return final_output


class MemoryAgent(Agent):
    memory_manager: Optional[MemoryManager] = Field(default=None, description="Manager for long-term memory operations")
    inputs: List[Dict] = Field(default_factory=list, description="Input specifications for the memory action")
    outputs: List[Dict] = Field(default_factory=list, description="Output specifications for the memory action")

    def __init__(
        self,
        name: str = "MemoryAgent",
        description: str = "An agent that uses long-term memory to provide context-aware responses",
        inputs: Optional[List[Dict]] = None,
        outputs: Optional[List[Dict]] = None,
        llm_config: Optional[OpenAILLMConfig] = None,
        storage_handler: Optional[StorageHandler] = None,
        rag_config: Optional[RAGConfig] = None,
        conversation_id: Optional[str] = None,
        system_prompt: Optional[str] = None,
        prompt: str = "Based on the following context and user prompt, provide a relevant response:\n\nContext: {context}\n\nUser Prompt: {user_prompt}",
        **kwargs
    ):
        # Define inputs and outputs inspired by CustomizeAgent
        inputs = inputs or []
        outputs = outputs or []

        # Initialize base Agent with provided parameters
        super().__init__(
            name=name,
            description=description,
            llm_config=llm_config,
            system_prompt=system_prompt,
            storage_handler=storage_handler,
            inputs=inputs,
            outputs=outputs,
            **kwargs
        )

        self.long_term_memory = LongTermMemory(
            storage_handler=storage_handler,
            rag_config=rag_config,
            default_corpus_id=conversation_id
        )
        self.memory_manager = MemoryManager(
            memory=self.long_term_memory,
            llm=llm_config.get_llm() if llm_config else None,
            use_llm_management=True
        )

        # Initialize inputs and outputs
        self.inputs = inputs
        self.outputs = outputs

        # Initialize actions list and add MemoryAction
        self.actions = []
        self._action_map = {}
        memory_action = MemoryAction(
            name="MemoryAction",
            description="Action that processes user input with long-term memory context",
            prompt=prompt,
            inputs_format=MemoryActionInput,
            outputs_format=MemoryActionOutput
        )
        self.add_action(memory_action)

    def _create_output_message(
        self,
        action_output,
        action_name: str,
        action_input_data: Optional[Dict],
        prompt: str,
        return_msg_type: MessageType = MessageType.RESPONSE,
        **kwargs
    ) -> Message:
        msg = super()._create_output_message(
            action_output=action_output,
            action_name=action_name,
            action_input_data=action_input_data,
            prompt=prompt,
            return_msg_type=return_msg_type,
            **kwargs
        )

        if action_input_data and "user_prompt" in action_input_data:
            user_msg = Message(
                content=action_input_data["user_prompt"],
                msg_type=MessageType.REQUEST,
                conversation_id=msg.conversation_id
            )
            asyncio.create_task(self.memory_manager.handle_memory(action="add", data=user_msg))

        response_msg = Message(
            content=action_output.response if hasattr(action_output, "response") else str(action_output),
            msg_type=MessageType.RESPONSE,
            conversation_id=msg.conversation_id
        )
        asyncio.create_task(self.memory_manager.handle_memory(action="add", data=response_msg))

        return msg

    async def async_execute(
        self,
        action_name: str,
        msgs: Optional[List[Message]] = None,
        action_input_data: Optional[Dict] = None,
        return_msg_type: Optional[MessageType] = MessageType.RESPONSE,
        return_action_input_data: Optional[bool] = False,
        **kwargs
    ) -> Union[Message, Tuple[Message, Dict]]:
        """
        Execute an action asynchronously with memory management.

        Args:
            action_name: Name of the action to execute
            msgs: Optional list of messages providing context
            action_input_data: Optional input data for the action
            return_msg_type: Message type for the return message
            return_action_input_data: Whether to return the action input data
            **kwargs: Additional parameters

        Returns:
            Message or tuple: The execution result, optionally with input data
        """
        action, action_input_data = self._prepare_execution(
            action_name=action_name,
            msgs=msgs,
            action_input_data=action_input_data,
            **kwargs
        )

        # Execute action with memory_manager
        execution_results = await action.async_execute(
            llm=self.llm,
            inputs=action_input_data,
            sys_msg=self.system_prompt,
            return_prompt=True,
            memory_manager=self.memory_manager,
            **kwargs
        )
        action_output, prompt = execution_results

        message = self._create_output_message(
            action_output=action_output,
            prompt=prompt,
            action_name=action_name,
            return_msg_type=return_msg_type,
            action_input_data=action_input_data,
            **kwargs
        )
        if return_action_input_data:
            return message, action_input_data
        return message

    def execute(
        self,
        action_name: str,
        msgs: Optional[List[Message]] = None,
        action_input_data: Optional[Dict] = None,
        return_msg_type: Optional[MessageType] = MessageType.RESPONSE,
        return_action_input_data: Optional[bool] = False,
        **kwargs
    ) -> Union[Message, Tuple[Message, Dict]]:
        """
        Execute an action synchronously with memory management.

        Args:
            action_name: Name of the action to execute
            msgs: Optional list of messages providing context
            action_input_data: Optional input data for the action
            return_msg_type: Message type for the return message
            return_action_input_data: Whether to return the action input data
            **kwargs: Additional parameters

        Returns:
            Message or tuple: The execution result, optionally with input data
        """
        action, action_input_data = self._prepare_execution(
            action_name=action_name,
            msgs=msgs,
            action_input_data=action_input_data,
            **kwargs
        )

        # Execute action with memory_manager
        execution_results = action.execute(
            llm=self.llm,
            inputs=action_input_data,
            sys_msg=self.system_prompt,
            return_prompt=True,
            memory_manager=self.memory_manager,
            **kwargs
        )
        action_output, prompt = execution_results

        message = self._create_output_message(
            action_output=action_output,
            prompt=prompt,
            action_name=action_name,
            return_msg_type=return_msg_type,
            action_input_data=action_input_data,
            **kwargs
        )
        if return_action_input_data:
            return message, action_input_data
        return message

    def chat(
        self,
        user_prompt: str,
        *,
        conversation_id: Optional[str] = None,
        top_k: Optional[int] = None,
        metadata_filters: Optional[dict] = None,
        return_message: bool = True,
        **kwargs
    ):
        action_input_data = {
            "user_prompt": user_prompt,
            "conversation_id": conversation_id or self._default_conversation_id(),
            "top_k": top_k if top_k is not None else 3,
            "metadata_filters": metadata_filters or {},
        }
        msg = self.execute(
            action_name="MemoryAction",
            action_input_data=action_input_data,
            return_msg_type=MessageType.RESPONSE,
            **kwargs
        )
        return msg if return_message else (getattr(msg, "content", None) or str(msg))


    async def async_chat(
        self,
        user_prompt: str,
        *,
        conversation_id: Optional[str] = None,
        top_k: Optional[int] = None,
        metadata_filters: Optional[dict] = None,
        return_message: bool = True,
        **kwargs
    ):
        action_input_data = {
            "user_prompt": user_prompt,
            "conversation_id": conversation_id or self._default_conversation_id(),
            "top_k": top_k if top_k is not None else 3,
            "metadata_filters": metadata_filters or {},
        }
        msg = await self.async_execute(
            action_name="MemoryAction",
            action_input_data=action_input_data,
            return_msg_type=MessageType.RESPONSE,
            **kwargs
        )
        return msg if return_message else (getattr(msg, "content", None) or str(msg))


    def _default_conversation_id(self) -> str:
        """
        Session scope: By default, a new uuid4() is returned (new session).
        User/global scope: Reuse LongTermMemory.default_corpus_id (stable namespace).
        Note: The final ID is still uniformly managed by MemoryAgent._prepare_execution() (which will override based on the scope).
        """
        scope = getattr(self, "conversation_scope", "session")
        if scope == "session":
            return str(uuid4())
        return getattr(getattr(self, "long_term_memory", None), "default_corpus_id", None) or "global_corpus"
    
    async def interactive_chat(
        self,
        conversation_id: Optional[str] = None,
        top_k: int = 3,
        metadata_filters: Optional[dict] = None
    ):
        """
        In interactive chat, each round of input will:
        1. Retrieve from memory
        2. Generate a response based on historical context
        3. Write the input/output to long-term memory and refresh the index 
        """
        conversation_id = conversation_id or self._default_conversation_id()
        metadata_filters = metadata_filters or {}

        print("💬 MemoryAgent has been started (type 'exit' to quit)\n")

        while True:
            user_prompt = input("You: ").strip()
            if user_prompt.lower() in ["exit", "quit"]:
                print("🔚 Conversation ended")
                break

            # Retrieve historical context
            retrieved_memories = await self.memory_manager.handle_memory(
                action="search",
                user_prompt=user_prompt,
                top_k=top_k,
                metadata_filters=metadata_filters
            )

            context_texts = []
            for msg, _ in retrieved_memories:
                if hasattr(msg, "content") and msg.content:
                    context_texts.append(msg.content)
            context_str = "\n".join(context_texts)

            # if context_str:
            #     print(f"📖 Retrieved context from memory:\n{context_str}\n")

            # Concatenate the historical context into the user input and invoke async_chat
            full_prompt = f"Context:\n{context_str}\n\nUser: {user_prompt}" if context_str else user_prompt
            msg = await self.async_chat(
                user_prompt=full_prompt,
                conversation_id=conversation_id,
                top_k=top_k,
                metadata_filters=metadata_filters
            )

            print(f"Agent: {msg.content}\n")

            # Refresh the index to ensure it can be retrieved in the next round
            if hasattr(self.memory_manager, "handle_memory_flush"):
                await self.memory_manager.handle_memory_flush()
            else:
                await asyncio.sleep(0.1)



    def save_module(self, path: str, ignore: List[str] = ["llm", "llm_config", "memory_manager"], **kwargs) -> str:
        """
        Save the agent's configuration to a JSON file, excluding memory_manager by default.

        Args:
            path: File path to save the configuration
            ignore: List of keys to exclude from the saved configuration
            **kwargs: Additional parameters for saving

        Returns:
            str: The path where the configuration was saved
        """
        return super().save_module(path=path, ignore=ignore, **kwargs)

    @classmethod
    def from_file(cls, path: str, llm_config: OpenAILLMConfig, storage_handler: Optional[StorageHandler] = None, rag_config: Optional[RAGConfig] = None, **kwargs) -> "MemoryAgent":
        """
        Load a MemoryAgent from a JSON configuration file.

        Args:
            path: Path to the JSON configuration file
            llm_config: LLM configuration
            storage_handler: Optional storage handler
            rag_config: Optional RAG configuration
            **kwargs: Additional parameters

        Returns:
            MemoryAgent: The loaded agent instance
        """
        with open(path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return cls(
            name=config.get("name", "MemoryAgent"),
            description=config.get("description", "An agent that uses long-term memory"),
            llm_config=llm_config,
            storage_handler=storage_handler,
            rag_config=rag_config,
            system_prompt=config.get("system_prompt"),
            prompt=config.get("prompt"),
            use_long_term_memory=config.get("use_long_term_memory", True),
            **kwargs
        )