🐍 Day09 Practice Template

"""
Day 9 练习：实现带记忆和 RAG 的 Agent
验收标准：
1. 实现了短期记忆（最近 N 轮对话历史）
2. 对话历史超出限制时有裁剪策略
3. 实现了长期记忆 / RAG（向量检索）
4. 能对新输入做语义检索，找到相关历史/文档
5. 多轮对话：Agent 能"记住"之前说过的话

安装依赖：
    pip install chromadb sentence-transformers
    # 或者
    pip install faiss-cpu numpy
"""

import asyncio
import json
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass, field
from abc import ABC, abstractmethod


# ==================== 任务 1：实现短期记忆（对话历史） ====================

@dataclass
class ConversationMessage:
    """对话消息"""
    role: str  # "user" | "assistant" | "system" | "tool"
    content: str
    timestamp: float = field(default_factory=lambda: __import__("time").time())


class ShortTermMemory:
    """
    短期记忆管理器
    
    需求：
    1. 存储最近 N 条消息（默认 20 条）
    2. 超出限制时自动裁剪最早的消息
    3. 可以在 messages 和格式化文本之间转换
    4. 保留 system message 不被裁剪
    """
    
    def __init__(self, max_messages: int = 20):
        self.max_messages = max_messages
        self.messages: List[ConversationMessage] = []
    
    def add(self, role: str, content: str):
        """添加消息"""
        # TODO 1.1：实现添加消息
        # 提示：
        # 1. 创建 ConversationMessage
        # 2. 追加到列表
        # 3. 如果超出限制，裁剪（但保留 system message）
        pass
    
    def get_recent(self, n: Optional[int] = None) -> List[ConversationMessage]:
        """获取最近 N 条消息"""
        # TODO 1.2：实现获取消息
        pass
    
    def to_api_format(self) -> List[Dict]:
        """转换为 LLM API 格式"""
        # TODO 1.3：实现格式转换
        # [{"role": "user", "content": "..."}, ...]
        pass
    
    def clear(self):
        """清除历史（保留 system message）"""
        # TODO 1.4：实现清除
        pass


# ==================== 任务 2：实现长期记忆 / 向量存储 ====================

class VectorStore(ABC):
    """向量存储抽象类"""
    
    @abstractmethod
    def add(self, text: str, metadata: Optional[Dict] = None):
        """添加文本到向量存储"""
        pass
    
    @abstractmethod
    def search(self, query: str, top_k: int = 5) -> List[Tuple[str, Dict, float]]:
        """搜索最相关的文本，返回 (文本, 元数据, 相似度)"""
        pass


class SimpleVectorStore(VectorStore):
    """
    简单的向量存储（不依赖外部库，用于学习）
    
    使用方式：
    1. 存储（文本, 向量, 元数据）
    2. 检索时计算相似度（cosine similarity）
    
    正式项目请用 ChromaDB 或 FAISS
    """
    
    def __init__(self):
        self.documents: List[Tuple[str, List[float], Dict]] = []
    
    def add(self, text: str, metadata: Optional[Dict] = None, embedding: Optional[List[float]] = None):
        """添加文本"""
        # TODO 2.1：实现添加
        # 提示：先存文本，embedding 可以后续计算
        pass
    
    def search(self, query: str, top_k: int = 3) -> List[Tuple[str, Dict, float]]:
        """搜索最相关的文本（基于关键词匹配的简单实现）"""
        # TODO 2.2：实现简单搜索
        # 提示：
        # 简单版：基于关键词重叠度打分
        # 正式版：用 embedding + cosine similarity
        pass


def simple_keyword_search(query: str, documents: List[str], top_k: int = 3) -> List[Tuple[str, float]]:
    """简单的关键词搜索（不依赖 embedding 模型）"""
    query_words = set(query.lower().split())
    
    scores = []
    for doc in documents:
        doc_words = set(doc.lower().split())
        # Jaccard 相似度
        intersection = len(query_words & doc_words)
        union = len(query_words | doc_words)
        score = intersection / union if union > 0 else 0
        scores.append((doc, score))
    
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]


# ==================== 任务 3：实现 RAG 检索模块 ====================

class RAGRetriever:
    """
    RAG 检索器
    
    流程：
    1. 知识库管理：添加文档
    2. 检索：根据用户问题，找到最相关的知识片段
    3. 返回检索结果供 LLM 使用
    """
    
    def __init__(self, vector_store: VectorStore):
        self.vector_store = vector_store
        self.knowledge_base: List[str] = []
    
    def add_knowledge(self, text: str):
        """添加知识到知识库"""
        # TODO 3.1：实现知识添加
        pass
    
    def retrieve(self, query: str, top_k: int = 3) -> str:
        """
        检索相关知识
        
        返回：格式化的知识文本（可直接拼接到 Prompt）
        """
        # TODO 3.2：实现检索
        # 提示：
        # 1. 搜索知识库
        # 2. 格式化结果
        # 3. 返回拼接后的文本
        pass


# ==================== 任务 4：实现带记忆的 Agent ====================

class MemoryAgent:
    """
    带记忆的 AI Agent
    
    组件：
    - ShortTermMemory：短期对话历史
    - RAGRetriever：长期记忆/知识检索
    - 工具系统：和之前一样
    
    每次对话时：
    1. 检索相关记忆/知识
    2. 将检索结果注入 System Prompt
    3. 处理用户输入
    4. 更新短期记忆
    """
    
    def __init__(self):
        self.short_memory = ShortTermMemory(max_messages=20)
        self.vector_store = SimpleVectorStore()
        self.rag = RAGRetriever(self.vector_store)
        
        # 添加一些初始知识
        self.rag.add_knowledge("Python 是一种高级编程语言，广泛用于 AI 开发。")
        self.rag.add_knowledge("AI Agent 是一种能自主决策的智能体，核心是 ReAct 循环。")
        self.rag.add_knowledge("LangGraph 是 LangChain 的图工作流框架，用于构建 AI Agent。")
    
    async def chat(self, user_input: str) -> str:
        """处理用户输入，返回回答"""
        # TODO 4.1：实现带记忆的对话
        # 流程：
        # 1. 检索相关知识（RAG）
        # 2. 构建 System Prompt（包含相关知识）
        # 3. 构建 Messages（短期记忆 + 当前输入）
        # 4. 调用 LLM（先用模拟）
        # 5. 将回答加入短期记忆
        # 6. 返回回答
        pass
    
    def remember_fact(self, fact: str):
        """记住一个事实（长期记忆）"""
        self.rag.add_knowledge(fact)


# ==================== 任务 5：实现上下文窗口管理 ====================

def compress_conversation(
    messages: List[ConversationMessage],
    max_tokens: int = 4000,
) -> List[ConversationMessage]:
    """
    对话压缩策略
    
    当对话历史超出 token 限制时：
    1. 方案 A：滑动窗口（只保留最近 N 轮）
    2. 方案 B：摘要压缩（调用 LLM 将历史总结为一段话）
    3. 方案 C：智能裁剪（保留关键信息，丢弃冗余）

    这里实现最简单的方案 A
    """
    # TODO 5.1：实现对话压缩
    # 提示：
    # 1. 估算每条消息的 token 数（粗略：1 个中文字 ≈ 2 tokens）
    # 2. 从最新消息开始往前保留
    # 3. 确保 system message 不被丢弃
    pass


# ==================== 主函数：测试 ====================

async def main():
    print("=" * 60)
    print("Day 9 练习：记忆系统 + RAG")
    print("=" * 60)
    
    # 测试 1：短期记忆
    print("\n📝 测试 1：短期记忆（添加 + 裁剪）")
    memory = ShortTermMemory(max_messages=5)
    memory.add("system", "你是一个 AI 助手")
    memory.add("user", "我叫张三")
    memory.add("assistant", "你好张三！")
    memory.add("user", "Python 是什么？")
    memory.add("assistant", "Python 是一种编程语言。")
    memory.add("user", "它难学吗？")  # 触发裁剪
    
    消息 = memory.to_api_format()
    print(f"  存储了 {len(memory.messages)} 条消息")
    for m in 消息:
        print(f"    [{m['role']}] {m['content'][:40]}...")
    
    # 测试 2：关键词搜索
    print("\n📝 测试 2：关键词搜索")
    docs = ["Python 编程", "AI Agent 开发", "天气查询", "LangGraph 框架"]
    results = simple_keyword_search("Python Agent 开发", docs, top_k=3)
    for doc, score in results:
        print(f"  📄 {doc}: 相似度 {score:.2f}")
    
    # 测试 3：RAG 检索
    print("\n📝 测试 3：RAG 检索")
    retriever = RAGRetriever(SimpleVectorStore())
    retriever.add_knowledge("Python 是 AI 开发的主力语言")
    retriever.add_knowledge("AI Agent 的核心是 ReAct 循环")
    retriever.add_knowledge("今天北京天气晴朗")
    
    result = retriever.retrieve("AI Agent 的核心是什么？", top_k=2)
    print(f"  检索结果：\n{result}")
    
    # 测试 4：多轮对话记忆
    print("\n📝 测试 4：多轮对话记忆")
    agent = MemoryAgent()
    
    response = await agent.chat("我叫张三，我喜欢 Python")
    print(f"  用户：我叫张三，我喜欢 Python")
    print(f"  助手：{response}")
    
    response = await agent.chat("我叫什么名字？我喜欢什么？")
    print(f"\n  用户：我叫什么名字？我喜欢什么？")
    print(f"  助手：{response}")
    print(f"\n  注：如果 Agent 答对了，说明记忆系统正常工作 ✅")
    
    print("\n" + "=" * 60)
    print("✅ 测试完成！现在运行 python day09_practice_validator.py 验收")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())