Agent 记忆系统（三）：向量数据库、知识图谱与分层记忆架构

# 构建 Agent 情景记忆系统
# 使用 ChromaDB + OpenAI Embedding 实现

import chromadb
from openai import OpenAI
from datetime import datetime
from typing import List, Optional
import json

class EpisodicMemory:
    """Agent 情景记忆系统
    
    核心功能：
    - 存储交互事件（带时间戳和元数据）
    - 语义检索相关记忆
    - 自动过期清理
    """
    
    def __init__(
        self,
        collection_name: str = "agent_memories",
        persist_directory: str = "./memory_db",
        embedding_model: str = "text-embedding-3-small",
        similarity_threshold: float = 0.7,
        max_memories: int = 10000,
    ):
        self.client = chromadb.PersistentClient(path=persist_directory)
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        self.openai = OpenAI()
        self.embedding_model = embedding_model
        self.threshold = similarity_threshold
        self.max_memories = max_memories
    
    def _embed(self, text: str) -> List[float]:
        """生成文本嵌入向量"""
        response = self.openai.embeddings.create(
            model=self.embedding_model,
            input=text,
        )
        return response.data[0].embedding
    
    def store_memory(
        self,
        user_message: str,
        agent_response: str,
        user_id: str = "default",
        tags: Optional[List[str]] = None,
    ) -> str:
        """存储一次交互事件到情景记忆
        
        参数：
            user_message: 用户消息
            agent_response: Agent 回复
            user_id: 用户标识（支持多用户记忆隔离）
            tags: 自定义标签（如 "project:ai-master", "topic:python"）
        
        返回：
            memory_id: 记忆的唯一标识
        """
        memory_id = f"mem_{self.collection.count() + 1}"
        timestamp = datetime.now().isoformat()
        
        # 将交互内容组合为可嵌入的文本
        document = f"用户: {user_message}\nAgent: {agent_response}"
        embedding = self._embed(document)
        
        metadata = {
            "user_id": user_id,
            "timestamp": timestamp,
            "user_message": user_message,
            "agent_response": agent_response,
            "tags": json.dumps(tags or []),
        }
        
        self.collection.add(
            ids=[memory_id],
            embeddings=[embedding],
            documents=[document],
            metadatas=[metadata],
        )
        
        # 自动清理过期记忆
        if self.collection.count() > self.max_memories:
            self._prune_old_memories()
        
        return memory_id
    
    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        user_id: Optional[str] = None,
        tags: Optional[List[str]] = None,
    ) -> List[dict]:
        """检索相关的情景记忆
        
        参数：
            query: 查询文本
            top_k: 返回结果数量
            user_id: 过滤特定用户的记忆
            tags: 过滤特定标签的记忆
        """
        query_embedding = self._embed(query)
        
        where_filter = {}
        if user_id:
            where_filter["user_id"] = user_id
        
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where=where_filter if where_filter else None,
            include=["documents", "metadatas", "distances"],
        )
        
        memories = []
        for i in range(len(results["ids"][0])):
            distance = results["distances"][0][i]
            similarity = 1 - distance  # 余弦距离转相似度
            
            if similarity < self.threshold:
                continue
            
            metadata = results["metadatas"][0][i]
            memories.append({
                "similarity": round(similarity, 3),
                "timestamp": metadata.get("timestamp"),
                "user_message": metadata.get("user_message"),
                "agent_response": metadata.get("agent_response"),
                "tags": json.loads(metadata.get("tags", "[]")),
            })
        
        return memories
    
    def _prune_old_memories(self):
        """清理最早的记忆，保持总数不超过限制"""
        count = self.collection.count()
        if count <= self.max_memories:
            return
        
        # 删除最早的一半记忆
        all_ids = self.collection.get()["ids"]
        ids_to_delete = all_ids[: count - self.max_memories]
        if ids_to_delete:
            self.collection.delete(ids=ids_to_delete)
    
    def get_memory_stats(self) -> dict:
        """获取记忆统计信息"""
        return {
            "total_memories": self.collection.count(),
            "max_memories": self.max_memories,
            "similarity_threshold": self.threshold,
        }


# 使用示例
memory = EpisodicMemory(
    persist_directory="./my_agent_memory",
    similarity_threshold=0.7,
    max_memories=10000,
)

# 存储记忆
mem_id = memory.store_memory(
    user_message="帮我写一个 Python 的快速排序算法",
    agent_response="快速排序是一种分治算法，以下是 Python 实现：...",
    user_id="user_001",
    tags=["python", "algorithm", "sorting"],
)
print(f"记忆已存储: {mem_id}")

# 检索相关记忆
results = memory.retrieve(
    query="排序算法的实现方法",
    top_k=3,
    user_id="user_001",
)
for r in results:
    print(f"相似度: {r['similarity']}")
    print(f"用户: {r['user_message'][:50]}...")

# 构建 Agent 语义记忆知识图谱
# 使用 Neo4j + Python 实现结构化知识存储

from neo4j import GraphDatabase
from typing import List, Tuple

class KnowledgeGraphMemory:
    """Agent 结构化语义记忆系统
    
    核心功能：
    - 存储实体和关系（三元组）
    - 多跳推理查询
    - 知识更新和去重
    """
    
    def __init__(self, uri: str = "bolt://localhost:7687", 
                 username: str = "neo4j", password: str = "password"):
        self.driver = GraphDatabase.driver(uri, auth=(username, password))
    
    def close(self):
        self.driver.close()
    
    def add_fact(self, subject: str, relation: str, obj: str, 
                 source: str = "agent_derived", confidence: float = 1.0):
        """添加一个事实三元组到知识图谱
        
        参数：
            subject: 主体实体
            relation: 关系类型
            obj: 客体实体
            source: 知识来源
            confidence: 置信度（0-1）
        """
        with self.driver.session() as session:
            session.run("""
                MERGE (s:Entity {name: $subject})
                MERGE (o:Entity {name: $obj})
                MERGE (s)-[r:REL {type: $relation}]->(o)
                SET r.source = $source,
                    r.confidence = $confidence,
                    r.updated_at = timestamp()
            """,
                subject=subject, obj=obj, 
                source=source, confidence=confidence,
            )
    
    def query_hops(self, start_entity: str, max_hops: int = 2) -> List[dict]:
        """从指定实体出发进行多跳查询
        
        这相当于 Agent 的「联想记忆」——
        从一个已知实体出发，探索与之相关的其他知识。
        """
        with self.driver.session() as session:
            result = session.run("""
                MATCH path = (start:Entity {name: $start})-[*1..$hops]-(related)
                RETURN path
                LIMIT 50
            """, start=start_entity, hops=max_hops)
            
            facts = []
            for record in result:
                path = record["path"]
                nodes = [n["name"] for n in path.nodes]
                rels = [r.type for r in path.relationships]
                facts.append({
                    "path": " -> ".join(nodes),
                    "relations": rels,
                    "depth": len(rels),
                })
            return facts
    
    def find_connections(self, entity1: str, entity2: str) -> List[dict]:
        """查找两个实体之间的关联路径
        
        这是知识图谱最强大的能力之一——
        发现两个看似不相关的实体之间的隐含联系。
        """
        with self.driver.session() as session:
            result = session.run("""
                MATCH path = shortestPath(
                    (a:Entity {name: $e1})-[*1..4]-(b:Entity {name: $e2})
                )
                RETURN path
            """, e1=entity1, e2=entity2)
            
            connections = []
            for record in result:
                path = record["path"]
                nodes = [n["name"] for n in path.nodes]
                rels = [r.type for r in path.relationships]
                connections.append({
                    "path": " -> ".join(nodes),
                    "relations": rels,
                })
            return connections


# 使用示例
kg = KnowledgeGraphMemory()

# 添加知识
kg.add_fact("Python", "用于", "机器学习", confidence=0.95)
kg.add_fact("Python", "用于", "Web开发", confidence=0.9)
kg.add_fact("TensorFlow", "使用", "Python", confidence=0.98)
kg.add_fact("PyTorch", "使用", "Python", confidence=0.98)
kg.add_fact("机器学习", "属于", "人工智能", confidence=1.0)

# 多跳查询：Python 相关知识
facts = kg.query_hops("Python", max_hops=2)
for f in facts:
    print(f"路径: {f['path']} (深度: {f['depth']})")

# 查找关联：TensorFlow 和 人工智能 有什么关系？
connections = kg.find_connections("TensorFlow", "人工智能")
for c in connections:
    print(f"连接: {c['path']}")

kg.close()