qwen_agent/agent/memori_manager.py

"""
Memori 连接和实例管理器
负责管理 Memori 客户端实例的创建、缓存和生命周期
"""

import asyncio
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

from psycopg_pool import AsyncConnectionPool
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session

from .memori_config import MemoriConfig

logger = logging.getLogger("app")


class MemoriManager:
    """
    Memori 连接和实例管理器

    主要功能：
    1. 管理 Memori 实例的创建和缓存
    2. 支持多租户隔离（entity_id + process_id）
    3. 处理数据库连接和会话管理
    4. 提供记忆召回和存储接口
    """

    def __init__(
        self,
        db_pool: Optional[AsyncConnectionPool] = None,
        db_url: Optional[str] = None,
        api_key: Optional[str] = None,
    ):
        """初始化 MemoriManager

        Args:
            db_pool: PostgreSQL 异步连接池（与 Checkpointer 共享）
            db_url: 数据库连接 URL（如果不使用连接池）
            api_key: Memori API 密钥（用于高级增强功能）
        """
        self._db_pool = db_pool
        self._db_url = db_url
        self._api_key = api_key

        # 缓存 Memori 实例: key = f"{entity_id}:{process_id}"
        self._instances: Dict[str, Any] = {}
        self._sync_engines: Dict[str, Any] = {}
        self._initialized = False

    async def initialize(self) -> None:
        """初始化 MemoriManager

        创建数据库表结构（如果不存在）
        """
        if self._initialized:
            return

        logger.info("Initializing MemoriManager...")

        try:
            # 创建第一个 Memori 实例来初始化表结构
            if self._db_pool or self._db_url:
                db_url = self._db_url or getattr(self._db_pool, "_url", None)
                if db_url:
                    await self._build_schema(db_url)

            self._initialized = True
            logger.info("MemoriManager initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize MemoriManager: {e}")
            # 不抛出异常，允许系统在没有 Memori 的情况下运行

    async def _build_schema(self, db_url: str) -> None:
        """构建 Memori 数据库表结构

        Args:
            db_url: 数据库连接 URL
        """
        try:
            from memori import Memori

            # 创建同步引擎用于初始化
            engine = create_engine(db_url)
            SessionLocal = sessionmaker(bind=engine)

            # 创建 Memori 实例并构建表结构
            mem = Memori(conn=SessionLocal)
            mem.config.storage.build()

            logger.info("Memori schema built successfully")
        except ImportError:
            logger.warning("memori package not available, skipping schema build")
        except Exception as e:
            logger.error(f"Failed to build Memori schema: {e}")

    def _get_sync_session(self, db_url: str) -> Session:
        """获取同步数据库会话（Memori 需要）

        Args:
            db_url: 数据库连接 URL

        Returns:
            SQLAlchemy Session
        """
        if db_url not in self._sync_engines:
            from sqlalchemy import create_engine
            from sqlalchemy.orm import sessionmaker

            engine = create_engine(db_url, pool_pre_ping=True)
            self._sync_engines[db_url] = sessionmaker(bind=engine)

        return self._sync_engines[db_url]()

    async def get_memori(
        self,
        entity_id: str,
        process_id: str,
        session_id: str,
        config: Optional[MemoriConfig] = None,
    ) -> Any:
        """获取或创建 Memori 实例

        Args:
            entity_id: 实体 ID（通常是 user_identifier）
            process_id: 进程 ID（通常是 bot_id）
            session_id: 会话 ID
            config: Memori 配置

        Returns:
            Memori 实例
        """
        cache_key = f"{entity_id}:{process_id}"

        # 检查缓存
        if cache_key in self._instances:
            memori_instance = self._instances[cache_key]
            # 更新会话
            memori_instance.config.session_id = session_id
            return memori_instance

        # 创建新实例
        memori_instance = await self._create_memori_instance(
            entity_id=entity_id,
            process_id=process_id,
            session_id=session_id,
            config=config,
        )

        # 缓存实例
        self._instances[cache_key] = memori_instance
        return memori_instance

    async def _create_memori_instance(
        self,
        entity_id: str,
        process_id: str,
        session_id: str,
        config: Optional[MemoriConfig] = None,
    ) -> Any:
        """创建新的 Memori 实例

        Args:
            entity_id: 实体 ID
            process_id: 进程 ID
            session_id: 会话 ID
            config: Memori 配置

        Returns:
            Memori 实例
        """
        try:
            from memori import Memori
        except ImportError:
            logger.error("memori package not installed")
            raise RuntimeError("memori package is required but not installed")

        # 获取数据库连接 URL
        db_url = self._db_url
        if self._db_pool and hasattr(self._db_pool, "_url"):
            db_url = str(self._db_pool._url)

        if not db_url:
            raise ValueError("Either db_pool or db_url must be provided")

        # 创建同步会话（Memori 目前需要同步连接）
        session_factory = self._get_sync_session(db_url)

        # 创建 Memori 实例
        mem = Memori(conn=session_factory)

        # 设置 API 密钥（如果提供）
        if self._api_key or (config and config.api_key):
            api_key = config.api_key if config else self._api_key
            mem.config.api_key = api_key

        # 设置 attribution
        mem.attribution(entity_id=entity_id, process_id=process_id)

        # 设置会话
        mem.config.session_id = session_id

        # 配置召回参数
        if config:
            mem.config.recall_facts_limit = config.semantic_search_top_k
            mem.config.recall_relevance_threshold = config.semantic_search_threshold
            mem.config.recall_embeddings_limit = config.semantic_search_embeddings_limit

        logger.info(
            f"Created Memori instance: entity={entity_id}, process={process_id}, session={session_id}"
        )

        return mem

    async def recall_memories(
        self,
        query: str,
        entity_id: str,
        process_id: str,
        session_id: str,
        config: Optional[MemoriConfig] = None,
    ) -> List[Dict[str, Any]]:
        """召回相关记忆

        Args:
            query: 查询文本
            entity_id: 实体 ID
            process_id: 进程 ID
            session_id: 会话 ID
            config: Memori 配置

        Returns:
            记忆列表，每个记忆包含 content, similarity 等字段
        """
        try:
            mem = await self.get_memori(entity_id, process_id, session_id, config)

            # 调用 recall 进行语义搜索
            results = mem.recall(query=query, limit=config.semantic_search_top_k if config else 5)

            # 转换为统一格式
            memories = []
            for result in results:
                memory = {
                    "content": result.get("content", ""),
                    "similarity": result.get("similarity", 0.0),
                    "fact_type": result.get("fact_type", "unknown"),
                    "created_at": result.get("created_at"),
                }
                # 过滤低相关度记忆
                threshold = config.semantic_search_threshold if config else 0.7
                if memory["similarity"] >= threshold:
                    memories.append(memory)

            logger.info(f"Recalled {len(memories)} memories for query: {query[:50]}...")
            return memories

        except Exception as e:
            logger.error(f"Failed to recall memories: {e}")
            return []

    async def wait_for_augmentation(
        self,
        entity_id: str,
        process_id: str,
        session_id: str,
        timeout: Optional[float] = None,
    ) -> None:
        """等待后台增强任务完成

        Args:
            entity_id: 实体 ID
            process_id: 进程 ID
            session_id: 会话 ID
            timeout: 超时时间（秒）
        """
        try:
            mem = await self.get_memori(entity_id, process_id, session_id)

            if timeout:
                # 在线程池中运行同步的 wait()
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, lambda: mem.augmentation.wait(timeout=timeout))
            else:
                # 无限等待
                loop = asyncio.get_event_loop()
                await loop.run_in_executor(None, mem.augmentation.wait)

        except Exception as e:
            logger.error(f"Failed to wait for augmentation: {e}")

    def clear_cache(self, entity_id: Optional[str] = None, process_id: Optional[str] = None) -> None:
        """清除缓存的 Memori 实例

        Args:
            entity_id: 实体 ID（如果为 None，清除所有）
            process_id: 进程 ID（如果为 None，清除所有）
        """
        if entity_id is None and process_id is None:
            self._instances.clear()
            logger.info("Cleared all Memori instances from cache")
        else:
            keys_to_remove = []
            for key in self._instances:
                e_id, p_id = key.split(":")
                if entity_id and e_id != entity_id:
                    continue
                if process_id and p_id != process_id:
                    continue
                keys_to_remove.append(key)

            for key in keys_to_remove:
                del self._instances[key]

            logger.info(f"Cleared {len(keys_to_remove)} Memori instances from cache")

    async def close(self) -> None:
        """关闭管理器并清理资源"""
        logger.info("Closing MemoriManager...")

        # 清理缓存的实例
        self._instances.clear()

        # 关闭同步引擎
        for engine in self._sync_engines.values():
            try:
                engine.dispose()
            except Exception as e:
                logger.error(f"Error closing engine: {e}")

        self._sync_engines.clear()
        self._initialized = False

        logger.info("MemoriManager closed")


# 全局单例
_global_manager: Optional[MemoriManager] = None


def get_memori_manager() -> MemoriManager:
    """获取全局 MemoriManager 单例

    Returns:
        MemoriManager 实例
    """
    global _global_manager
    if _global_manager is None:
        _global_manager = MemoriManager()
    return _global_manager


async def init_global_memori(
    db_pool: Optional[AsyncConnectionPool] = None,
    db_url: Optional[str] = None,
    api_key: Optional[str] = None,
) -> MemoriManager:
    """初始化全局 MemoriManager

    Args:
        db_pool: PostgreSQL 连接池
        db_url: 数据库连接 URL
        api_key: Memori API 密钥

    Returns:
        MemoriManager 实例
    """
    manager = get_memori_manager()
    manager._db_pool = db_pool
    manager._db_url = db_url
    manager._api_key = api_key
    await manager.initialize()
    return manager


async def close_global_memori() -> None:
    """关闭全局 MemoriManager"""
    global _global_manager
    if _global_manager is not None:
        await _global_manager.close()