qwen_agent/agent/mem0_config.py

"""
Mem0 configuration data class.
Used to manage configuration parameters for the Mem0 long-term memory system.
"""

from datetime import datetime
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Optional

# Avoid circular imports.
if TYPE_CHECKING:
    from langchain_core.language_models import BaseChatModel


# Cache the loaded prompt template.
_fact_extraction_prompt_template: Optional[str] = None


def _load_fact_extraction_prompt() -> str:
    """Load the prompt template from prompt/FACT_RETRIEVAL_PROMPT.md.

    Returns:
        str: Prompt template content.
    """
    global _fact_extraction_prompt_template

    if _fact_extraction_prompt_template is not None:
        return _fact_extraction_prompt_template

    # Get the project root directory, assuming prompt/ is under the project root.
    project_root = Path(__file__).parent.parent
    prompt_file = project_root / "prompt" / "FACT_RETRIEVAL_PROMPT.md"

    try:
        _fact_extraction_prompt_template = prompt_file.read_text(encoding="utf-8")
    except FileNotFoundError:
        # If the file does not exist, return the default prompt.
        _fact_extraction_prompt_template = (
            "Extract relevant facts about the user from the conversation.\n"
            "Today's date is {current_time}.\n"
            "Return response in JSON format: {\"facts\": [\"fact1\", \"fact2\"]}"
        )

    return _fact_extraction_prompt_template


@dataclass
class Mem0Config:
    """Mem0 long-term memory configuration class."""

    # Feature flags
    enabled: bool = False

    # Semantic search configuration
    semantic_search_top_k: int = 20

    memory_prompt_template: str = (
        "\n\n=== 相关记忆 ===\n"
        "以下是从历史对话中检索到的相关信息，可以帮助你更好地回答用户问题：\n"
        "{memories}\n"
        "==================\n"
    )

    # Multi-tenant configuration
    user_id: Optional[str] = None  # User identifier
    agent_id: Optional[str] = None  # Bot identifier
    session_id: Optional[str] = None  # Session identifier

    @property
    def bot_id(self) -> str:
        """Compatibility bot_id property required by execute_hooks."""
        return self.agent_id or ""

    # LLM instance, used for Mem0 memory extraction and enrichment
    llm_instance: Optional["BaseChatModel"] = None  # LangChain LLM instance

    def get_attribution_tuple(self) -> tuple[str, str]:
        """Get the tuple (user_id, agent_id) required for attribution.

        Returns:
            (user_id, agent_id) tuple
        """
        if not self.user_id or not self.agent_id:
            raise ValueError("user_id and agent_id are required for attribution")
        return (self.user_id, self.agent_id)

    def is_enabled(self) -> bool:
        """Check whether Mem0 is enabled.

        Returns:
            bool: Whether it is enabled.
        """
        return self.enabled

    def get_memory_prompt(self, memories: list[str]) -> str:
        """Generate the injected prompt from the memory list.

        Args:
            memories: List of memory contents.

        Returns:
            str: Formatted memory prompt.
        """
        if not memories:
            return ""

        memory_text = "\n".join(f"- {m}" for m in memories)
        return self.memory_prompt_template.format(memories=memory_text)

    def get_custom_fact_extraction_prompt(self) -> str:
        """Get the custom memory extraction prompt.

        Reads the template from prompt/FACT_RETRIEVAL_PROMPT.md and replaces
        {current_time} with the current date.

        Returns:
            str: Custom memory extraction prompt.
        """
        template = _load_fact_extraction_prompt()
        current_date = datetime.now().strftime("%Y-%m-%d")
        return template.format(current_time=current_date)

    async def get_custom_fact_extraction_prompt_async(self) -> str:
        """Asynchronously get the custom memory extraction prompt.

        Supports injection from the PreMemoryPrompt hook. It reads the default
        template from prompt/FACT_RETRIEVAL_PROMPT.md, then executes
        PreMemoryPrompt hooks. If a hook returns content, that content replaces
        the entire template. Finally, it replaces {current_time} with the current date.

        Returns:
            str: Custom memory extraction prompt.
        """
        from agent.plugin_hook_loader import execute_hooks

        template = await execute_hooks('PreMemoryPrompt', self) or _load_fact_extraction_prompt()
        return template.format(current_time=datetime.now().strftime("%Y-%m-%d"))

    def with_session(self, session_id: str) -> "Mem0Config":
        """Create a copy of the configuration with a new session_id.

        Args:
            session_id: New session ID.

        Returns:
            A new Mem0Config instance.
        """
        new_config = Mem0Config(**self.__dict__)
        new_config.session_id = session_id
        return new_config