Agent 记忆系统（一）：向量数据库与知识图谱

from dataclasses import dataclass, field
from typing import List, Dict, Any

@dataclass
class AgentMemory:
    """Agent 三层记忆架构"""
    # 短期记忆：当前对话上下文（受限于上下文窗口）
    short_term: List[Dict[str, str]] = field(default_factory=list)
    # 工作记忆：结构化任务状态
    working: Dict[str, Any] = field(default_factory=dict)
    # 长期记忆：外部存储引用（向量数据库 ID）
    long_term_refs: List[str] = field(default_factory=list)

    def add_to_working(self, key: str, value: Any) -> None:
        """更新工作记忆"""
        self.working[key] = value

    def consolidate(self) -> None:
        """记忆整合：将工作记忆中的重要信息转入长期记忆"""
        for key, value in self.working.items():
            if self._is_important(key, value):
                self._store_long_term(key, value)

    def _is_important(self, key: str, value: Any) -> bool:
        return key.startswith("fact_") or key.startswith("preference_")

    def _store_long_term(self, key: str, value: Any) -> None:
        # 实际实现中会写入向量数据库
        ref_id = f"mem_{hash(key)}"
        self.long_term_refs.append(ref_id)

# 记忆生命周期管理示例
class MemoryManager:
    """管理记忆的创建、检索、遗忘和整合"""

    def __init__(self, max_short_term: int = 20):
        self.max_short = max_short_term
        self.short_buffer: List[str] = []
        self.decay_rate = 0.95  # 记忆衰减率

    def store(self, info: str) -> None:
        """存入短期记忆，超限时触发遗忘"""
        self.short_buffer.append(info)
        if len(self.short_buffer) > self.max_short:
            self._forget_oldest()

    def retrieve(self, query: str, top_k: int = 3) -> List[str]:
        """从短期记忆中检索（简单关键词匹配）"""
        query_words = set(query.lower().split())
        scored = []
        for item in self.short_buffer:
            overlap = len(query_words & set(item.lower().split()))
            scored.append((overlap, item))
        scored.sort(reverse=True)
        return [item for _, item in scored[:top_k]]

    def _forget_oldest(self) -> str:
        """遗忘最早的条目（先进先出策略）"""
        return self.short_buffer.pop(0)

import chromadb
from chromadb.utils import embedding_functions

# 初始化 Chroma 客户端（持久化到磁盘）
client = chromadb.PersistentClient(path="./memory_db")

# 使用 Sentence-Transformers 嵌入模型
embedder = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# 创建集合（Collection）
collection = client.create_collection(
    name="agent_memories",
    embedding_function=embedder,
    metadata={"hnsw:space": "cosine"}  # 余弦相似度
)

# 添加文档
collection.add(
    documents=[
        "用户喜欢 Python 编程",
        "项目截止日期是 2026 年 5 月",
        "部署环境是 AWS us-east-1"
    ],
    metadatas=[
        {"type": "preference", "category": "dev"},
        {"type": "fact", "category": "project"},
        {"type": "fact", "category": "infra"}
    ],
    ids=["mem_001", "mem_002", "mem_003"]
)

# 相似度检索
results = collection.query(
    query_texts=["用户的编程偏好是什么？"],
    n_results=2
)
print(results["documents"])  # 返回最匹配的文档

import faiss
import numpy as np

# FAISS：高性能向量索引
# 创建 768 维向量的索引（使用内积相似度）
dimension = 768
index = faiss.IndexFlatIP(dimension)  # 内积（等价于余弦，假设已归一化）

# 生成模拟数据
np.random.seed(42)
vectors = np.random.randn(1000, dimension).astype(np.float32)
# L2 归一化
faiss.normalize_L2(vectors)

# 添加向量到索引
index.add(vectors)
print(f"索引中的向量数: {index.ntotal}")

# 查询：搜索最相似的 5 个向量
query = np.random.randn(1, dimension).astype(np.float32)
faiss.normalize_L2(query)
distances, indices = index.search(query, k=5)

print(f"Top-5 相似度: {distances[0]}")
print(f"Top-5 索引号: {indices[0]}")

# 使用 Sentence-Transformers 评估嵌入模型
from sentence_transformers import SentenceTransformer
import numpy as np

# 加载模型（中文优化版 BGE）
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

# 测试语义相似度
sentences = [
    "如何学习 Python 编程？",
    "Python 入门教程推荐",
    "今天天气真不错",
    "Python 编程语言的学习路径"
]

# 生成嵌入向量
embeddings = model.encode(sentences, normalize_embeddings=True)

# 计算余弦相似度矩阵
similarity = embeddings @ embeddings.T

# 打印相似度矩阵
for i, s1 in enumerate(sentences):
    for j, s2 in enumerate(sentences):
        if i < j:
            print(f"{s1} ↔ {s2}: {similarity[i][j]:.4f}")

# 预期输出：句1和句4（语义最相关）相似度应最高

# 批量嵌入与向量化流水线
from sentence_transformers import SentenceTransformer
import json
import os

class EmbeddingPipeline:
    """批量嵌入处理流水线"""

    def __init__(self, model_name: str, batch_size: int = 64):
        self.model = SentenceTransformer(model_name)
        self.batch_size = batch_size
        self.dimension = self.model.get_sentence_embedding_dimension()

    def embed_batch(self, texts: list[str]) -> np.ndarray:
        """批量生成嵌入（自动分块）"""
        all_embeddings = []
        for i in range(0, len(texts), self.batch_size):
            batch = texts[i:i + self.batch_size]
            embeddings = self.model.encode(
                batch,
                show_progress_bar=False,
                normalize_embeddings=True
            )
            all_embeddings.extend(embeddings.tolist())
        return np.array(all_embeddings, dtype=np.float32)

    def get_stats(self) -> dict:
        return {
            "model": self.model.tokenizer.name_or_path,
            "dimension": self.dimension,
            "batch_size": self.batch_size
        }

# 使用示例
pipeline = EmbeddingPipeline("BAAI/bge-large-zh-v1.5", batch_size=32)
print(pipeline.get_stats())
# {'model': 'BAAI/bge-large-zh-v1.5', 'dimension': 1024, 'batch_size': 32}

# 完整 RAG 管道实现
import chromadb
from sentence_transformers import SentenceTransformer
from openai import OpenAI

class RAGPipeline:
    """检索增强生成管道"""

    def __init__(self, db_path: str, llm_model: str = "gpt-4o"):
        self.client = chromadb.PersistentClient(path=db_path)
        self.embedder = SentenceTransformer("BAAI/bge-large-zh-v1.5")
        self.llm = OpenAI()
        self.llm_model = llm_model
        self.collection = self.client.get_or_create_collection(
            name="knowledge",
            metadata={"hnsw:space": "cosine"}
        )

    def ingest(self, documents: list[str], ids: list[str]) -> None:
        """文档入库"""
        embeddings = self.embedder.encode(
            documents, normalize_embeddings=True
        ).tolist()
        self.collection.add(
            documents=documents,
            embeddings=embeddings,
            ids=ids
        )

    def query(self, question: str, top_k: int = 3) -> str:
        """RAG 查询"""
        # 1. 检索
        query_vec = self.embedder.encode(
            [question], normalize_embeddings=True
        ).tolist()
        results = self.collection.query(
            query_embeddings=query_vec, n_results=top_k
        )
        context = "\n".join(results["documents"][0])

        # 2. 增强生成
        prompt = f"""请基于以下参考资料回答问题。如果资料中没有相关信息，请如实告知。

参考资料：
{context}

问题：{question}

回答："""
        response = self.llm.chat.completions.create(
            model=self.llm_model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        )
        return response.choices[0].message.content

# 文档切分策略对比
class DocumentChunker:
    """多种文档切分策略"""

    @staticmethod
    def fixed_size(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
        """固定大小切分（带重叠）"""
        chunks = []
        words = text.split()
        for i in range(0, len(words), chunk_size - overlap):
            chunk = " ".join(words[i:i + chunk_size])
            if len(chunk.split()) > 10:  # 过滤过短的块
                chunks.append(chunk)
        return chunks

    @staticmethod
    def semantic_split(text: str, model_name: str = "BAAI/bge-large-zh-v1.5") -> list[str]:
        """语义切分：按段落/句子边界，确保语义完整"""
        import re
        # 按句子边界切分
        sentences = re.split(r'[。！？；\n]', text)
        sentences = [s.strip() for s in sentences if s.strip()]

        chunks = []
        current = ""
        model = SentenceTransformer(model_name)

        for sent in sentences:
            if len(current) + len(sent) > 500:
                chunks.append(current)
                current = sent
            else:
                current += "。" + sent if current else sent
        if current:
            chunks.append(current)
        return chunks

    @staticmethod
    def hierarchical(text: str, max_chunk: int = 2000) -> list[str]:
        """层次化切分：先按章节，再按段落"""
        sections = re.split(r'#+\s+', text)
        chunks = []
        for section in sections:
            if len(section) <= max_chunk:
                chunks.append(section.strip())
            else:
                sub = DocumentChunker.fixed_size(section, chunk_size=500)
                chunks.extend(sub)
        return chunks

# 使用 NetworkX 构建知识图谱
import networkx as nx
import json

class KnowledgeGraph:
    """Agent 知识图谱"""

    def __init__(self):
        self.graph = nx.DiGraph()

    def add_entity(self, name: str, entity_type: str, **properties) -> None:
        """添加实体节点"""
        self.graph.add_node(name, type=entity_type, **properties)

    def add_relation(self, source: str, target: str, relation: str, **properties) -> None:
        """添加关系边"""
        self.graph.add_edge(source, target, relation=relation, **properties)

    def find_related(self, entity: str, max_depth: int = 2) -> list:
        """查找相关实体（BFS 遍历）"""
        if entity not in self.graph:
            return []
        visited = set()
        results = []
        queue = [(entity, 0)]
        while queue:
            node, depth = queue.pop(0)
            if depth > max_depth or node in visited:
                continue
            visited.add(node)
            for neighbor in self.graph.neighbors(node):
                edge_data = self.graph[node][neighbor]
                results.append({
                    "source": node,
                    "relation": edge_data["relation"],
                    "target": neighbor,
                    "depth": depth + 1
                })
                queue.append((neighbor, depth + 1))
        return results

    def to_dict(self) -> dict:
        return nx.node_link_data(self.graph)

# GraphRAG：结合知识图谱与向量检索
class GraphRAG:
    """图增强检索生成"""

    def __init__(self, vector_db, knowledge_graph: KnowledgeGraph):
        self.vector_db = vector_db
        self.kg = knowledge_graph
        self.llm = OpenAI()

    def hybrid_query(self, question: str) -> dict:
        """混合查询：向量检索 + 图谱推理"""
        # 1. 向量检索（语义匹配）
        vec_results = self.vector_db.query(
            question, top_k=3
        )

        # 2. 图谱推理（关系查询）
        entities = self._extract_entities(question)
        graph_results = []
        for entity in entities:
            related = self.kg.find_related(entity, max_depth=2)
            graph_results.extend(related)

        # 3. 合并结果构建上下文
        context_parts = [
            f"[语义检索] {doc}"
            for doc in vec_results.get("documents", [[]])[0]
        ]
        for rel in graph_results:
            context_parts.append(
                f"[图谱推理] {rel['source']} "
                f"-{rel['relation']}→ "
                f"{rel['target']}"
            )

        return {
            "context": "\n".join(context_parts),
            "sources": {
                "vector": vec_results,
                "graph": graph_results
            }
        }

    def _extract_entities(self, text: str) -> list:
        """简单实体提取（实际应用中用 NER 模型）"""
        return [n for n in self.kg.graph.nodes() if n in text]

# 记忆摘要与压缩
class MemoryCompressor:
    """记忆压缩器"""

    def __init__(self, llm, max_items: int = 50):
        self.llm = llm
        self.max_items = max_items
        self.memory_buffer: list[dict] = []

    def add(self, entry: str, importance: float = 0.5) -> None:
        """添加记忆条目"""
        self.memory_buffer.append({
            "content": entry,
            "importance": importance,
            "timestamp": time.time()
        })
        # 超出阈值时触发压缩
        if len(self.memory_buffer) > self.max_items:
            self.compress()

    def compress(self, target_ratio: float = 0.3) -> list[dict]:
        """将记忆压缩为摘要"""
        # 按重要性排序
        sorted_mem = sorted(
            self.memory_buffer,
            key=lambda x: x["importance"],
            reverse=True
        )

        # 保留最重要的条目
        keep_count = max(int(self.max_items * target_ratio), 5)
        kept = sorted_mem[:keep_count]

        # 其余条目生成摘要
        to_summarize = sorted_mem[keep_count:]
        texts = [m["content"] for m in to_summarize]
        summary = self._llm_summarize(texts)

        compressed = kept + [{
            "content": summary,
            "importance": 0.3,
            "timestamp": time.time(),
            "is_summary": True
        }]

        self.memory_buffer = compressed
        return compressed

    def _llm_summarize(self, texts: list[str]) -> str:
        prompt = f"请将以下 {len(texts)} 条信息浓缩为 3-5 条核心要点：\n"
        for i, t in enumerate(texts):
            prompt += f"{i+1}. {t}\n"
        response = self.llm.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1
        )
        return response.choices[0].message.content

# 分层记忆管理：热记忆 vs 冷记忆
class TieredMemory:
    """分层记忆：热（最近/重要）vs 冷（历史/归档）"""

    def __init__(self, hot_limit: int = 100):
        self.hot_memory: list[dict] = []  # 热记忆：快速访问
        self.cold_memory: list[dict] = []  # 冷记忆：归档存储
        self.hot_limit = hot_limit

    def store(self, content: str, importance: float = 0.5) -> str:
        """存储记忆（自动分层）"""
        entry = {
            "id": str(uuid.uuid4()),
            "content": content,
            "importance": importance,
            "access_count": 0,
            "last_access": time.time(),
            "created_at": time.time()
        }
        self.hot_memory.append(entry)
        self._evict_if_needed()
        return entry["id"]

    def retrieve(self, query: str, top_k: int = 5) -> list[dict]:
        """从热记忆中检索"""
        scored = []
        for entry in self.hot_memory:
            # 混合评分：语义 + 时效性 + 重要性 + 访问频率
            semantic_score = self._semantic_match(query, entry["content"])
            recency = time.time() - entry["created_at"]
            recency_score = 1.0 / (1.0 + recency / 86400)  # 天为单位
            final_score = (
                0.5 * semantic_score
                + 0.2 * recency_score
                + 0.2 * entry["importance"]
                + 0.1 * min(entry["access_count"] / 10, 1.0)
            )
            scored.append((final_score, entry))
        scored.sort(reverse=True)
        results = [entry for _, entry in scored[:top_k]]
        for entry in results:
            entry["access_count"] += 1
            entry["last_access"] = time.time()
        return results

    def _evict_if_needed(self) -> None:
        """将最不活跃的条目移入冷记忆"""
        while len(self.hot_memory) > self.hot_limit:
            self.hot_memory.sort(key=lambda x: x["last_access"])
            self.cold_memory.append(self.hot_memory.pop(0))

# 完整记忆增强 Agent 系统
class MemoryAgent:
    """集成向量检索 + 知识图谱 + 记忆管理的 Agent"""

    def __init__(self, config: dict):
        # 存储层
        self.vector_db = ChromaVectorDB(config["db_path"])
        self.knowledge_graph = KnowledgeGraph()
        self.memory_store = TieredMemory(hot_limit=100)

        # 嵌入与 LLM
        self.embedder = SentenceTransformer("BAAI/bge-large-zh-v1.5")
        self.llm = OpenAI(model=config["llm_model"])

        # 记忆管理器
        self.compressor = MemoryCompressor(self.llm, max_items=80)

        # Agent 状态
        self.session_memory: list[dict] = []
        self.user_profile: dict = {}

    def chat(self, user_message: str) -> dict:
        """完整的记忆增强对话流程"""
        # 1. 存储到短期记忆
        self.session_memory.append({"role": "user", "content": user_message})

        # 2. 向量检索（RAG）
        vec_results = self.vector_db.query(user_message, top_k=3)

        # 3. 图谱推理（GraphRAG）
        graph_results = self._graph_query(user_message)

        # 4. 构建增强上下文
        context = self._build_context(
            user_message, vec_results, graph_results
        )

        # 5. LLM 生成回复
        response = self._generate(context)

        # 6. 更新记忆
        self._update_memory(user_message, response)

        return {"response": response, "sources": vec_results}

    def _build_context(self, query, vec_results, graph_results) -> str:
        """构建增强提示词"""
        parts = []
        parts.append("你是用户的 AI 助理，拥有长期记忆和外部知识。")
        if vec_results:
            parts.append("\n[参考资料]")
            for doc in vec_results.get("documents", [[]])[0]:
                parts.append(f"- {doc}")
        if graph_results:
            parts.append("\n[知识图谱信息]")
            for rel in graph_results:
                parts.append(
                    f"- {rel['source']} {rel['relation']} {rel['target']}"
                )
        if self.session_memory:
            parts.append("\n[最近对话]")
            for msg in self.session_memory[-5:]:
                parts.append(f"{msg['role']}: {msg['content']}")
        parts.append(f"\n用户问题: {query}")
        return "\n".join(parts)

    def _update_memory(self, query: str, response: str) -> None:
        """对话后更新长期记忆"""
        self.session_memory.append({"role": "assistant", "content": response})
        self.compressor.add(f"用户: {query}\nAI: {response}", importance=0.6)
        self.memory_store.store(
            f"{query} -> {response}",
            importance=self._estimate_importance(query)
        )

# Agent 配置文件与启动
import yaml
import asyncio

class AgentConfig:
    """Agent 系统配置"""

    DEFAULT_CONFIG = {
        "db_path": "./agent_memory",
        "llm_model": "gpt-4o",
        "embedder_model": "BAAI/bge-large-zh-v1.5",
        "vector_db": "chroma",
        "max_context_length": 4000,
        "top_k_retrieval": 3,
        "memory_compress_threshold": 80,
        "enable_graph_rag": True,
        "enable_memory_compression": True,
        "user_profile_path": "./user_profile.json",
        "log_level": "INFO"
    }

    @classmethod
    def load(cls, config_path: str = "agent_config.yaml") -> dict:
        """从 YAML 文件加载配置"""
        try:
            with open(config_path, "r") as f:
                user_config = yaml.safe_load(f)
            return {cls.DEFAULT_CONFIG, user_config}
        except FileNotFoundError:
            return cls.DEFAULT_CONFIG

    @classmethod
    def save(cls, config: dict, path: str) -> None:
        with open(path, "w") as f:
            yaml.dump(config, f, default_flow_style=False)

# 启动示例
def run_agent():
    config = AgentConfig.load()
    agent = MemoryAgent(config)

    print("🍪 记忆增强 Agent 已启动")
    print("输入 'quit' 退出，输入 'memory' 查看记忆状态")

    while True:
        user_input = input("\n你: ").strip()
        if user_input.lower() == "quit":
            break
        if user_input.lower() == "memory":
            print(f"会话记忆: {len(agent.session_memory)} 条")
            print(f"热记忆: {len(agent.memory_store.hot_memory)} 条")
            print(f"冷记忆: {len(agent.memory_store.cold_memory)} 条")
            continue

        result = agent.chat(user_input)
        print(f"AI: {result['response']}")