AI Agent 记忆系统设计与向量数据库实战

# HNSW 性能基准测试（2026 年 2 月数据） # 数据集：100 万条 1536 维 OpenAI embeddings import hnswlib import numpy as np import time # 初始化索引 p = hnswlib.Index(space='cosine', dim=1536) p.init_index(max_elements=1000000, ef_construction=200, M=16) # 批量添加数据 data = np.random.rand(1000000, 1536).astype('float32') p.add_items(data, range(1000000)) # 性能测试 p.set_ef(50) # ef 参数影响精度/速度平衡 query = np.random.rand(1, 1536).astype('float32') start = time.time() labels, distances = p.knn_query(query, k=10) print(f"P99 延迟：{(time.time()-start)*1000:.2f}ms") # 结果：~0.85ms（单线程 CPU）

# requirements.txt langgraph>=0.2.0 chromadb>=0.5.0 openai>=1.12.0 numpy>=1.24.0 tiktoken>=0.5.0 pydantic>=2.0.0

from dataclasses import dataclass, field from datetime import datetime, timedelta from typing import List, Dict, Optional, Any from enum import Enum import chromadb import numpy as np import tiktoken class MemoryType(Enum): WORKING = "working" # 工作记忆 SEMANTIC = "semantic" # 语义记忆 EPISODIC = "episodic" # 情景记忆 PROCEDURAL = "procedural" # 程序记忆 @dataclass class Memory: """记忆单元基类""" id: str content: str created_at: datetime = field(default_factory=datetime.utcnow) importance: float = 0.5 # 重要性分数 0-1 access_count: int = 0 # 访问次数 metadata: Dict[str, Any] = field(default_factory=dict) @dataclass class EpisodicMemory(Memory): """情景记忆：带时间戳的事件记录""" event_type: str = "conversation" participants: List[str] = field(default_factory=list) outcome: Optional[str] = None class MemoryController: """分层记忆系统控制器""" def __init__( self, embedding_model: str = "text-embedding-3-small", chroma_path: str = "./chroma_db", token_budget: int = 4000 ): self.token_budget = token_budget self.encoding = tiktoken.encoding_for_model("gpt-4") # 初始化 ChromaDB self.chroma_client = chromadb.PersistentClient(path=chroma_path) self.semantic_collection = self.chroma_client.get_or_create_collection( name="semantic_memory", metadata={"hnsw:space": "cosine"} ) self.episodic_collection = self.chroma_client.get_or_create_collection( name="episodic_memory", metadata={"hnsw:space": "cosine"} ) # 内存中的记忆存储 self.working_memory: List[Memory] = [] self.procedural_memory: Dict[str, List[str]] = {} # Embedding 缓存 self._embedding_cache: Dict[str, List[float]] = {} def _get_embedding(self, text: str) -> List[float]: """获取文本的 embedding 向量（带缓存）""" if text in self._embedding_cache: return self._embedding_cache[text] # 实际项目中调用 OpenAI API 或本地模型 # 此处简化为随机向量 embedding = np.random.randn(1536).tolist() self._embedding_cache[text] = embedding return embedding def _count_tokens(self, text: str) -> int: """计算 token 数量""" return len(self.encoding.encode(text)) def add_semantic_memory( self, content: str, category: str, metadata: Optional[Dict] = None ) -> str: """添加语义记忆（事实知识）""" memory_id = f"sem_{datetime.utcnow().timestamp()}" embedding = self._get_embedding(content) full_metadata = { "category": category, "created_at": datetime.utcnow().isoformat(), **(metadata or {}) } self.semantic_collection.upsert( ids=[memory_id], embeddings=[embedding], metadatas=[full_metadata], documents=[content] ) return memory_id def add_episodic_memory( self, content: str, event_type: str = "conversation", participants: Optional[List[str]] = None, outcome: Optional[str] = None ) -> str: """添加情景记忆（事件记录）""" memory_id = f"epi_{datetime.utcnow().timestamp()}" embedding = self._get_embedding(content) full_metadata = { "event_type": event_type, "created_at": datetime.utcnow().isoformat(), "participants": ",".join(participants or []), "outcome": outcome or "" } self.episodic_collection.upsert( ids=[memory_id], embeddings=[embedding], metadatas=[full_metadata], documents=[content] ) return memory_id def retrieve_semantic( self, query: str, top_k: int = 5, category_filter: Optional[str] = None ) -> List[Dict]: """检索语义记忆""" query_embedding = self._get_embedding(query) where_clause = None if category_filter: where_clause = {"category": category_filter} results = self.semantic_collection.query( query_embeddings=[query_embedding], n_results=top_k, where=where_clause, include=["documents", "metadatas", "distances"] ) return [ { "id": results["ids"][0][i], "content": results["documents"][0][i], "metadata": results["metadatas"][0][i], "distance": results["distances"][0][i] } for i in range(len(results["ids"][0])) ] def retrieve_episodic( self, query: str, top_k: int = 5, time_range: Optional[tuple] = None, apply_time_decay: bool = True ) -> List[Dict]: """检索情景记忆（支持时间范围过滤和时间衰减）""" query_embedding = self._get_embedding(query) results = self.episodic_collection.query( query_embeddings=[query_embedding], n_results=top_k * 2, # 先多取一些用于后续过滤 include=["documents", "metadatas", "distances"] ) memories = [] now = datetime.utcnow() for i in range(len(results["ids"][0])): metadata = results["metadatas"][0][i] created_at = datetime.fromisoformat(metadata["created_at"]) # 时间范围过滤 if time_range: start, end = time_range if not (start <= created_at <= end): continue memory = { "id": results["ids"][0][i], "content": results["documents"][0][i], "metadata": metadata, "distance": results["distances"][0][i], "created_at": created_at } # 应用时间衰减（24 小时半衰期） if apply_time_decay: age_hours = (now - created_at).total_seconds() / 3600 decay_factor = np.exp(-age_hours / 24) memory["score"] = (1 - memory["distance"]) * decay_factor else: memory["score"] = 1 - memory["distance"] memories.append(memory) # 按分数排序，取 top_k memories.sort(key=lambda x: x["score"], reverse=True) return memories[:top_k] def build_context( self, query: str, semantic_top_k: int = 3, episodic_top_k: int = 3, max_tokens: Optional[int] = None ) -> str: """构建 LLM 上下文（自动在 token 预算内优化）""" max_tokens = max_tokens or self.token_budget # 检索相关记忆 semantic_memories = self.retrieve_semantic(query, top_k=semantic_top_k) episodic_memories = self.retrieve_episodic(query, top_k=episodic_top_k) # 构建上下文 context_parts = [] current_tokens = 0 # 添加语义记忆 for memory in semantic_memories: content = f"[知识库] {memory['content']}" tokens = self._count_tokens(content) if current_tokens + tokens <= max_tokens * 0.5: context_parts.append(content) current_tokens += tokens # 添加情景记忆 for memory in episodic_memories: content = f"[历史记录] {memory['content']}" tokens = self._count_tokens(content) if current_tokens + tokens <= max_tokens: context_parts.append(content) current_tokens += tokens return "\n\n".join(context_parts) def consolidate_old_memories( self, days_threshold: int = 7, summary_model: str = "gpt-4o-mini" ): """压缩旧情景记忆为语义摘要（定期执行）""" cutoff = datetime.utcnow() - timedelta(days=days_threshold) # 获取旧记忆 old_memories = self.retrieve_episodic( query="", top_k=100, time_range=(datetime.min, cutoff), apply_time_decay=False ) if len(old_memories) < 5: return # 数据太少，无需压缩 # 调用 LLM 生成摘要 # 实际项目中调用 OpenAI API summary = f"压缩 {len(old_memories)} 条旧记忆为语义知识" # 添加为语义记忆 self.add_semantic_memory( content=summary, category="consolidated_summary", metadata={"source_count": len(old_memories)} ) # 删除旧记忆（简化版，实际应考虑保留重要记忆） for memory in old_memories: self.episodic_collection.delete(ids=[memory["id"]])

from langgraph.graph import StateGraph, START, END from langgraph.checkpoint.memory import MemorySaver from typing import TypedDict, Annotated import operator class AgentState(TypedDict): """定义 Agent 状态""" messages: Annotated[List[str], operator.add] user_preferences: Dict[str, Any] current_task: str # 初始化持久化存储 memory = MemorySaver() # 构建图 builder = StateGraph(AgentState) def process_node(state: AgentState): """处理节点""" return {"messages": ["处理完成"]} builder.add_node("processor", process_node) builder.add_edge(START, "processor") builder.add_edge("processor", END) # 编译时传入 checkpoint graph = builder.compile(checkpointer=memory) # 执行时指定 thread_id（会话标识） config = {"configurable": {"thread_id": "user_123"}} # 第一次执行 result1 = graph.invoke({ "messages": [], "user_preferences": {"language": "Python"}, "current_task": "build_agent" }, config) # 第二次执行（同一 thread_id 会继承之前的状态） result2 = graph.invoke({ "messages": ["继续执行"], "user_preferences": {}, "current_task": "test_agent" }, config)

from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, CheckpointMetadata, SerializerProtocol ) from langgraph.serde.jsonplus import JsonPlusSerializer import sqlite3 import json class SqliteSaver(BaseCheckpointSaver): """SQLite 持久化 Checkpoint 实现""" serializer: SerializerProtocol = field(default_factory=JsonPlusSerializer) def __init__(self, db_path: str): self.conn = sqlite3.connect(db_path) self._setup_tables() def _setup_tables(self): """创建数据表""" self.conn.execute(""" CREATE TABLE IF NOT EXISTS checkpoints ( thread_id TEXT PRIMARY KEY, checkpoint_id TEXT, checkpoint BLOB, metadata BLOB, updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) self.conn.commit() def get_tuple(self, config: Dict) -> Optional[tuple]: """获取 checkpoint""" thread_id = config["configurable"]["thread_id"] cursor = self.conn.execute( "SELECT checkpoint_id, checkpoint, metadata FROM checkpoints WHERE thread_id = ?", (thread_id,) ) row = cursor.fetchone() if not row: return None return ( row[0], self.serializer.loads(row[1]), self.serializer.loads(row[2]) ) def put( self, config: Dict, checkpoint: Checkpoint, metadata: CheckpointMetadata ) -> Dict: """保存 checkpoint""" thread_id = config["configurable"]["thread_id"] checkpoint_id = checkpoint["id"] self.conn.execute( """ INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_id, checkpoint, metadata) VALUES (?, ?, ?, ?) """, ( thread_id, checkpoint_id, self.serializer.dumps(checkpoint), self.serializer.dumps(metadata) ) ) self.conn.commit() return {"configurable": {"thread_id": thread_id}}

from neo4j import GraphDatabase class GraphMemory: """基于 Neo4j 的图记忆系统""" def __init__(self, uri: str, user: str, password: str): self.driver = GraphDatabase.driver(uri, auth=(user, password)) def close(self): self.driver.close() def add_entity(self, entity_id: str, entity_type: str, properties: Dict): """添加实体节点""" with self.driver.session() as session: props_str = ", ".join(f"{k}: $props.{k}" for k in properties.keys()) query = f""" MERGE (e:{entity_type} {{id: $id}}) SET e += $props RETURN e """ session.run(query, id=entity_id, props=properties) def add_relation( self, from_entity: str, relation: str, to_entity: str, properties: Optional[Dict] = None ): """添加关系边""" with self.driver.session() as session: props_set = f"SET r += $props" if properties else "" query = f""" MATCH (a), (b) WHERE a.id = $from AND b.id = $to MERGE (a)-[r:{relation}]->(b) {props_set} RETURN r """ session.run(query, from=from_entity, to=to_entity, props=properties or {}) def query_path( self, start_entity: str, end_entity: str, max_hops: int = 3 ) -> List[Dict]: """查询两个实体之间的路径（多跳推理）""" with self.driver.session() as session: query = """ MATCH path = shortestPath( (start {id: $start})-[*..max_hops]-(end {id: $end}) ) RETURN path """ result = session.run(query, start=start_entity, end=end_entity, max_hops=max_hops) return list(result)