memory_box.py

"""
记忆盒子模块 - Memory Box Module

提供显式记忆（用户主动添加）和隐式记忆（智能体自动总结）管理功能。
"""

import sqlite3
import json
import uuid
from datetime import datetime
from typing import List, Dict, Optional, Any
import os
import logging

logger = logging.getLogger(__name__)


class MemoryBoxError(Exception):
    """记忆盒子相关错误"""
    pass


class MemoryBox:
    """记忆盒子类 - 管理用户画像和偏好"""
    
    def __init__(self, db_path: str = 'memory_box.db'):
        """
        初始化记忆盒子
        
        Args:
            db_path: 数据库文件路径
        """
        self.db_path = db_path
        self._init_database()
        
    def _init_database(self):
        """初始化数据库表"""
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 显式记忆表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS explicit_memories (
                        id TEXT PRIMARY KEY,
                        session_id TEXT NOT NULL,
                        content TEXT NOT NULL,
                        category TEXT DEFAULT 'general',
                        created_at TEXT NOT NULL,
                        updated_at TEXT NOT NULL,
                        enabled INTEGER DEFAULT 1
                    )
                ''')
                
                # 隐式记忆表（用户画像和偏好）
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS implicit_memories (
                        id TEXT PRIMARY KEY,
                        session_id TEXT NOT NULL,
                        memory_type TEXT NOT NULL,
                        content TEXT NOT NULL,
                        summary TEXT,
                        confidence REAL DEFAULT 0.5,
                        interaction_count INTEGER DEFAULT 0,
                        created_at TEXT NOT NULL,
                        updated_at TEXT NOT NULL,
                        enabled INTEGER DEFAULT 1
                    )
                ''')
                
                # 记忆更新历史表
                cursor.execute('''
                    CREATE TABLE IF NOT EXISTS memory_update_history (
                        id TEXT PRIMARY KEY,
                        session_id TEXT NOT NULL,
                        update_type TEXT NOT NULL,
                        conversation_count INTEGER,
                        summary TEXT,
                        created_at TEXT NOT NULL
                    )
                ''')
                
                conn.commit()
                logger.info(f"记忆盒子数据库初始化成功: {self.db_path}")
                
        except Exception as e:
            logger.error(f"数据库初始化失败: {str(e)}")
            raise MemoryBoxError(f"数据库初始化失败: {str(e)}")
    
    def add_explicit_memory(self, session_id: str, content: str, category: str = 'general') -> str:
        """
        添加显式记忆
        
        Args:
            session_id: 会话ID
            content: 记忆内容
            category: 记忆分类（general, preference, profile, skill等）
            
        Returns:
            记忆ID
        """
        try:
            memory_id = str(uuid.uuid4())
            now = datetime.now().isoformat()
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO explicit_memories 
                    (id, session_id, content, category, created_at, updated_at, enabled)
                    VALUES (?, ?, ?, ?, ?, ?, 1)
                ''', (memory_id, session_id, content, category, now, now))
                conn.commit()
                
            logger.info(f"添加显式记忆成功: {memory_id} (session: {session_id})")
            return memory_id
            
        except Exception as e:
            logger.error(f"添加显式记忆失败: {str(e)}")
            raise MemoryBoxError(f"添加显式记忆失败: {str(e)}")
    
    def update_explicit_memory(self, memory_id: str, content: str = None, 
                              category: str = None, enabled: bool = None) -> bool:
        """
        更新显式记忆
        
        Args:
            memory_id: 记忆ID
            content: 新内容
            category: 新分类
            enabled: 是否启用
            
        Returns:
            是否成功
        """
        try:
            updates = []
            params = []
            
            if content is not None:
                updates.append("content = ?")
                params.append(content)
            if category is not None:
                updates.append("category = ?")
                params.append(category)
            if enabled is not None:
                updates.append("enabled = ?")
                params.append(1 if enabled else 0)
                
            if not updates:
                return True
                
            updates.append("updated_at = ?")
            params.append(datetime.now().isoformat())
            params.append(memory_id)
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                query = f"UPDATE explicit_memories SET {', '.join(updates)} WHERE id = ?"
                cursor.execute(query, params)
                conn.commit()
                
            logger.info(f"更新显式记忆成功: {memory_id}")
            return True
            
        except Exception as e:
            logger.error(f"更新显式记忆失败: {str(e)}")
            return False
    
    def delete_explicit_memory(self, memory_id: str) -> bool:
        """
        删除显式记忆
        
        Args:
            memory_id: 记忆ID
            
        Returns:
            是否成功
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute("DELETE FROM explicit_memories WHERE id = ?", (memory_id,))
                conn.commit()
                
            logger.info(f"删除显式记忆成功: {memory_id}")
            return True
            
        except Exception as e:
            logger.error(f"删除显式记忆失败: {str(e)}")
            return False
    
    def get_explicit_memories(self, session_id: str, enabled_only: bool = True) -> List[Dict[str, Any]]:
        """
        获取显式记忆列表
        
        Args:
            session_id: 会话ID
            enabled_only: 是否仅返回启用的记忆
            
        Returns:
            记忆列表
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                query = "SELECT * FROM explicit_memories WHERE session_id = ?"
                params = [session_id]
                
                if enabled_only:
                    query += " AND enabled = 1"
                    
                query += " ORDER BY created_at DESC"
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                columns = [desc[0] for desc in cursor.description]
                memories = []
                for row in rows:
                    memory = dict(zip(columns, row))
                    memory['enabled'] = bool(memory.get('enabled', 1))
                    memories.append(memory)
                    
                return memories
                
        except Exception as e:
            logger.error(f"获取显式记忆失败: {str(e)}")
            return []
    
    def add_or_update_implicit_memory(self, session_id: str, memory_type: str, 
                                     content: str, summary: str = None,
                                     confidence: float = 0.5) -> str:
        """
        添加或更新隐式记忆
        
        Args:
            session_id: 会话ID
            memory_type: 记忆类型（user_profile, preference, behavior_pattern等）
            content: 记忆内容（详细）
            summary: 记忆摘要（精简）
            confidence: 置信度（0-1）
            
        Returns:
            记忆ID
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                # 检查是否已存在相同类型的记忆
                cursor.execute('''
                    SELECT id, interaction_count FROM implicit_memories 
                    WHERE session_id = ? AND memory_type = ? AND enabled = 1
                ''', (session_id, memory_type))
                
                existing = cursor.fetchone()
                now = datetime.now().isoformat()
                
                if existing:
                    # 更新现有记忆
                    memory_id = existing[0]
                    interaction_count = existing[1] + 1
                    
                    cursor.execute('''
                        UPDATE implicit_memories 
                        SET content = ?, summary = ?, confidence = ?, 
                            interaction_count = ?, updated_at = ?
                        WHERE id = ?
                    ''', (content, summary, confidence, interaction_count, now, memory_id))
                else:
                    # 创建新记忆
                    memory_id = str(uuid.uuid4())
                    cursor.execute('''
                        INSERT INTO implicit_memories 
                        (id, session_id, memory_type, content, summary, confidence, 
                         interaction_count, created_at, updated_at, enabled)
                        VALUES (?, ?, ?, ?, ?, ?, 1, ?, ?, 1)
                    ''', (memory_id, session_id, memory_type, content, summary, 
                          confidence, now, now))
                
                conn.commit()
                
            logger.info(f"{'更新' if existing else '添加'}隐式记忆成功: {memory_id} (type: {memory_type})")
            return memory_id
            
        except Exception as e:
            logger.error(f"添加/更新隐式记忆失败: {str(e)}")
            raise MemoryBoxError(f"添加/更新隐式记忆失败: {str(e)}")
    
    def get_implicit_memories(self, session_id: str, enabled_only: bool = True) -> List[Dict[str, Any]]:
        """
        获取隐式记忆列表
        
        Args:
            session_id: 会话ID
            enabled_only: 是否仅返回启用的记忆
            
        Returns:
            记忆列表
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                query = "SELECT * FROM implicit_memories WHERE session_id = ?"
                params = [session_id]
                
                if enabled_only:
                    query += " AND enabled = 1"
                    
                query += " ORDER BY confidence DESC, updated_at DESC"
                
                cursor.execute(query, params)
                rows = cursor.fetchall()
                
                columns = [desc[0] for desc in cursor.description]
                memories = []
                for row in rows:
                    memory = dict(zip(columns, row))
                    memory['enabled'] = bool(memory.get('enabled', 1))
                    memories.append(memory)
                    
                return memories
                
        except Exception as e:
            logger.error(f"获取隐式记忆失败: {str(e)}")
            return []
    
    def get_memory_context(self, session_id: str, max_chars: int = 500) -> str:
        """
        获取用于提示词的记忆上下文（精简版）
        
        Args:
            session_id: 会话ID
            max_chars: 最大字符数
            
        Returns:
            格式化的记忆上下文
        """
        try:
            explicit_memories = self.get_explicit_memories(session_id, enabled_only=True)
            implicit_memories = self.get_implicit_memories(session_id, enabled_only=True)
            
            if not explicit_memories and not implicit_memories:
                return ""
            
            context_parts = []
            
            # 添加隐式记忆（用户画像和偏好）
            if implicit_memories:
                implicit_texts = []
                for mem in implicit_memories[:3]:  # 取前3个最重要的
                    text = mem.get('summary') or mem.get('content', '')
                    if text:
                        implicit_texts.append(f"- {text}")
                
                if implicit_texts:
                    context_parts.append("【用户画像】\n" + "\n".join(implicit_texts))
            
            # 添加显式记忆
            if explicit_memories:
                explicit_texts = []
                for mem in explicit_memories[:5]:  # 取前5个
                    content = mem.get('content', '')
                    if content:
                        explicit_texts.append(f"- {content}")
                
                if explicit_texts:
                    context_parts.append("【用户备注】\n" + "\n".join(explicit_texts))
            
            full_context = "\n\n".join(context_parts)
            
            # 如果超过最大字符数，进行裁剪
            if len(full_context) > max_chars:
                full_context = full_context[:max_chars] + "..."
            
            return full_context
            
        except Exception as e:
            logger.error(f"获取记忆上下文失败: {str(e)}")
            return ""
    
    def record_update_history(self, session_id: str, update_type: str, 
                             conversation_count: int, summary: str = None) -> str:
        """
        记录记忆更新历史
        
        Args:
            session_id: 会话ID
            update_type: 更新类型（auto_summary, manual_edit等）
            conversation_count: 对话轮次
            summary: 更新摘要
            
        Returns:
            历史记录ID
        """
        try:
            history_id = str(uuid.uuid4())
            now = datetime.now().isoformat()
            
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                cursor.execute('''
                    INSERT INTO memory_update_history 
                    (id, session_id, update_type, conversation_count, summary, created_at)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (history_id, session_id, update_type, conversation_count, summary, now))
                conn.commit()
                
            return history_id
            
        except Exception as e:
            logger.error(f"记录更新历史失败: {str(e)}")
            return ""
    
    def get_all_memories(self, session_id: str) -> Dict[str, List[Dict[str, Any]]]:
        """
        获取所有记忆（显式+隐式）
        
        Args:
            session_id: 会话ID
            
        Returns:
            包含显式和隐式记忆的字典
        """
        return {
            'explicit': self.get_explicit_memories(session_id),
            'implicit': self.get_implicit_memories(session_id)
        }
    
    def clear_session_memories(self, session_id: str, memory_type: str = 'all') -> bool:
        """
        清空会话的记忆
        
        Args:
            session_id: 会话ID
            memory_type: 记忆类型（'explicit', 'implicit', 'all'）
            
        Returns:
            是否成功
        """
        try:
            with sqlite3.connect(self.db_path) as conn:
                cursor = conn.cursor()
                
                if memory_type in ['explicit', 'all']:
                    cursor.execute("DELETE FROM explicit_memories WHERE session_id = ?", (session_id,))
                
                if memory_type in ['implicit', 'all']:
                    cursor.execute("DELETE FROM implicit_memories WHERE session_id = ?", (session_id,))
                
                conn.commit()
                
            logger.info(f"清空会话记忆成功: {session_id} (type: {memory_type})")
            return True
            
        except Exception as e:
            logger.error(f"清空会话记忆失败: {str(e)}")
            return False


class ImplicitMemoryAnalyzer:
    """隐式记忆分析器 - 自动分析和总结用户画像"""
    
    def __init__(self, siliconflow_token: str, model: str = "Qwen/Qwen2.5-7B-Instruct"):
        """
        初始化分析器
        
        Args:
            siliconflow_token: API Token
            model: 使用的模型
        """
        self.api_token = siliconflow_token
        self.model = model
        self.api_url = "https://api.siliconflow.cn/v1/chat/completions"
    
    def analyze_conversation_history(self, conversation_history: List[Dict[str, str]]) -> Dict[str, Any]:
        """
        分析对话历史，提取用户画像和偏好
        
        Args:
            conversation_history: 对话历史列表
            
        Returns:
            分析结果
        """
        try:
            import requests
            
            # 构建分析提示词
            conversation_text = "\n".join([
                f"{'用户' if msg['role'] == 'user' else '助手'}: {msg['content']}"
                for msg in conversation_history[-10:]  # 只分析最近10轮对话
            ])
            
            prompt = f"""请分析以下对话历史，总结用户的画像和偏好。请以简洁的方式输出，格式如下：

用户画像：[一句话描述用户的基本特征]
兴趣偏好：[用户感兴趣的话题或领域]
交互风格：[用户的沟通风格和习惯]
专业背景：[如果能识别出用户的专业领域]

对话历史：
{conversation_text}

请直接输出分析结果，不要有其他解释。"""

            headers = {
                "Authorization": f"Bearer {self.api_token}",
                "Content-Type": "application/json"
            }
            
            data = {
                "model": self.model,
                "messages": [{"role": "user", "content": prompt}],
                "temperature": 0.3,
                "max_tokens": 300
            }
            
            response = requests.post(self.api_url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                analysis = result['choices'][0]['message']['content']
                
                return {
                    'success': True,
                    'analysis': analysis,
                    'summary': analysis[:200]  # 取前200字符作为摘要
                }
            else:
                logger.error(f"分析失败: {response.status_code}")
                return {'success': False, 'error': f"API错误: {response.status_code}"}
                
        except Exception as e:
            logger.error(f"分析对话历史失败: {str(e)}")
            return {'success': False, 'error': str(e)}