AI Agent 记忆系统设计：短期记忆 + 长期记忆 + 向量检索实战

凌晨两点，某电商平台的"双十一"大促刚刚结束。客服系统的 AI Agent 处理了超过 50 万次用户咨询，其中既有"这款手机支持分期吗"的简单查询，也有需要跨店铺比价、查询历史订单、结合用户偏好推荐配件的复杂多轮对话。更关键的是，这些对话数据在次日自动沉淀为知识库的一部分——当用户再次咨询同类问题时，Agent 能在 200ms 内调取出相关历史记录给出个性化回复。

这就是本文要解决的工程问题：如何设计一套高效、可扩展的 AI Agent 记忆系统，同时兼顾响应速度与成本控制？

为什么 AI Agent 需要三层记忆架构

在真实业务场景中，AI Agent 面临的核心矛盾是：上下文窗口有限（成本高），但用户对话历史无限（数据多）。三层记忆架构是业界主流解决方案：

短期记忆（Session Memory）：存储当前对话上下文，使用 KV Cache 或固定长度窗口，适合 20 轮以内的单次会话
长期记忆（Persistent Memory）：基于 PostgreSQL/MySQL 的结构化存储，存储用户画像、偏好设置、关键交互结论
向量记忆（Vector Memory）：基于 Embedding + 向量数据库的语义检索，存储非结构化知识、文档、历史对话摘要

三者协同工作的逻辑如下：用户发起请求时，Agent 首先从短期记忆读取当前对话上下文，再从向量记忆检索相关历史知识，最后根据长期记忆中的用户画像调整回复策略。这种架构在 HolySheep API 的低成本加持下，单次请求成本可控制在 $0.001 以下。

系统架构设计与核心实现

1. 环境准备与依赖安装

# requirements.txt
openai==1.12.0
psycopg2-binary==2.9.9
redis==5.0.1
chromadb==0.4.22
sentence-transformers==2.4.0
pydantic==2.6.0
python-dotenv==1.0.1

安装命令
pip install -r requirements.txt

2. 核心记忆模块实现

import os
import json
import redis
import chromadb
from openai import OpenAI
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
import psycopg2
from psycopg2.extras import RealDictCursor

HolySheep API 配置（汇率优势：¥1=$1，注册送免费额度）
client = OpenAI(
    api_key="YOUR_HOLYSHEEP_API_KEY",
    base_url="https://api.holysheep.ai/v1"  # 国内直连 <50ms
)

class MemorySystem:
    """三层记忆系统：短期 + 长期 + 向量"""
    
    def __init__(self, redis_url="redis://localhost:6379/0"):
        # 短期记忆：Redis KV 存储（TTL 30分钟）
        self.redis_client = redis.from_url(redis_url, decode_responses=True)
        self.session_ttl = 1800  # 30分钟
        
        # 长期记忆：PostgreSQL（用户画像 + 结构化数据）
        self.db_config = {
            "host": os.getenv("PG_HOST", "localhost"),
            "database": os.getenv("PG_DATABASE", "agent_memory"),
            "user": os.getenv("PG_USER", "postgres"),
            "password": os.getenv("PG_PASSWORD", "")
        }
        
        # 向量记忆：ChromaDB（语义检索）
        self.vector_db = chromadb.Client()
        self.collection = self.vector_db.get_or_create_collection(
            name="agent_memories",
            metadata={"hnsw:space": "cosine"}
        )
        
        # Embedding 模型
        self._embedding_model = "text-embedding-3-small"
    
    def _get_embedding(self, text: str) -> List[float]:
        """调用 HolySheep API 生成向量嵌入"""
        response = client.embeddings.create(
            model=self._embedding_model,
            input=text
        )
        return response.data[0].embedding
    
    # ============ 短期记忆操作 ============
    
    def save_short_term(self, session_id: str, role: str, content: str) -> None:
        """保存短期对话记忆到 Redis"""
        key = f"session:{session_id}"
        message = json.dumps({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })
        self.redis_client.rpush(key, message)
        self.redis_client.expire(key, self.session_ttl)
    
    def get_short_term(self, session_id: str, max_messages: int = 10) -> List[Dict]:
        """获取短期记忆（最近 N 条）"""
        key = f"session:{session_id}"
        messages = self.redis_client.lrange(key, -max_messages, -1)
        return [json.loads(m) for m in messages]
    
    def clear_short_term(self, session_id: str) -> None:
        """清理会话记忆"""
        self.redis_client.delete(f"session:{session_id}")
    
    # ============ 长期记忆操作 ============
    
    def _get_db_connection(self):
        return psycopg2.connect(**self.db_config, cursor_factory=RealDictCursor)
    
    def save_long_term(self, user_id: str, memory_type: str, data: Dict[str, Any]) -> int:
        """保存长期记忆到 PostgreSQL"""
        with self._get_db_connection() as conn:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO user_memories (user_id, memory_type, data, created_at)
                    VALUES (%s, %s, %s, NOW())
                    RETURNING id
                """, (user_id, memory_type, json.dumps(data)))
                result = cur.fetchone()
                conn.commit()
                return result['id']
    
    def get_long_term(self, user_id: str, memory_type: Optional[str] = None) -> List[Dict]:
        """获取用户长期记忆"""
        with self._get_db_connection() as conn:
            with conn.cursor() as cur:
                if memory_type:
                    cur.execute("""
                        SELECT * FROM user_memories 
                        WHERE user_id = %s AND memory_type = %s
                        ORDER BY created_at DESC LIMIT 50
                    """, (user_id, memory_type))
                else:
                    cur.execute("""
                        SELECT * FROM user_memories 
                        WHERE user_id = %s
                        ORDER BY created_at DESC LIMIT 50
                    """, (user_id,))
                return cur.fetchall()
    
    # ============ 向量记忆操作 ============
    
    def save_vector_memory(self, user_id: str, content: str, metadata: Dict = None) -> str:
        """保存向量记忆"""
        embedding = self._get_embedding(content)
        doc_id = f"{user_id}_{datetime.now().strftime('%Y%m%d%H%M%S')}"
        
        self.collection.add(
            ids=[doc_id],
            embeddings=[embedding],
            documents=[content],
            metadatas=[{
                "user_id": user_id,
                "created_at": datetime.now().isoformat(),
                **(metadata or {})
            }]
        )
        return doc_id
    
    def search_vector_memory(self, user_id: str, query: str, top_k: int = 5) -> List[Dict]:
        """语义检索向量记忆"""
        query_embedding = self._get_embedding(query)
        
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where={"user_id": user_id}  # 过滤当前用户
        )
        
        return [
            {
                "id": results['ids'][0][i],
                "content": results['documents'][0][i],
                "distance": results['distances'][0][i],
                "metadata": results['metadatas'][0][i]
            }
            for i in range(len(results['ids'][0]))
        ]
    
    # ============ 记忆融合（核心方法）============
    
    def build_context(self, session_id: str, user_id: str, current_query: str) -> str:
        """构建完整上下文：三层记忆融合"""
        context_parts = []
        
        # 1. 短期记忆：当前对话上下文
        short_term = self.get_short_term(session_id, max_messages=8)
        if short_term:
            context_parts.append("【当前会话历史】")
            for msg in short_term:
                context_parts.append(f"{msg['role']}: {msg['content']}")
        
        # 2. 向量记忆：检索相关历史知识
        relevant_memories = self.search_vector_memory(user_id, current_query, top_k=3)
        if relevant_memories:
            context_parts.append("\n【相关历史经验】")
            for mem in relevant_memories:
                if mem['distance'] < 0.7:  # 相似度阈值
                    context_parts.append(f"- {mem['content']}")
        
        # 3. 长期记忆：用户画像与偏好
        user_profile = self.get_long_term(user_id, memory_type="profile")
        if user_profile:
            profile = user_profile[0]['data']
            context_parts.append(f"\n【用户画像】偏好: {profile.get('preferences', '暂无')}")
        
        return "\n".join(context_parts) if context_parts else ""


数据库初始化 SQL
INIT_SQL = """
CREATE TABLE IF NOT EXISTS user_memories (
    id SERIAL PRIMARY KEY,
    user_id VARCHAR(128) NOT NULL,
    memory_type VARCHAR(64) NOT NULL,
    data JSONB NOT NULL,
    created_at TIMESTAMP DEFAULT NOW()
);

CREATE INDEX IF NOT EXISTS idx_user_memories_user_id ON user_memories(user_id);
CREATE INDEX IF NOT EXISTS idx_user_memories_type ON user_memories(memory_type);
"""

3. Agent 集成与调用示例

from memory_system import MemorySystem

memory = MemorySystem()

def agent_response(session_id: str, user_id: str, user_message: str) -> str:
    """完整 Agent 响应流程"""
    
    # 步骤1：构建上下文（三层记忆融合）
    context = memory.build_context(session_id, user_id, user_message)
    
    # 步骤2：构造提示词
    system_prompt = """你是一个专业的电商客服助手。请根据上下文信息，准确回答用户问题。
    如果涉及用户历史偏好，请优先考虑用户习惯。如果问题超出范围，请礼貌引导。"""
    
    # 步骤3：保存当前对话到短期记忆
    memory.save_short_term(session_id, "user", user_message)
    
    # 步骤4：调用 HolySheep API（GPT-4.1 $8/MTok 或 Claude Sonnet 4.5 $15/MTok）
    response = client.chat.completions.create(
        model="gpt-4.1",  # 或 "claude-sonnet-4.5"
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"{context}\n\n当前问题: {user_message}"} if context else user_message
        ],
        temperature=0.7,
        max_tokens=500
    )
    
    assistant_message = response.choices[0].message.content
    
    # 步骤5：保存助手回复到短期记忆
    memory.save_short_term(session_id, "assistant", assistant_message)
    
    # 步骤6：关键信息沉淀到长期记忆（如用户确认了收货地址）
    if "确认" in user_message and "地址" in user_message:
        memory.save_long_term(
            user_id=user_id,
            memory_type="preference",
            data={"shipping_address": user_message.split("地址:")[1].strip() if "地址:" in user_message else None}
        )
    
    return assistant_message


示例调用
if __name__ == "__main__":
    # 模拟电商场景
    session_id = "sess_20241111_001"
    user_id = "user_12345"
    
    # 第一轮对话
    resp1 = agent_response(session_id, user_id, "我想买一台笔记本电脑，预算 8000 元")
    print(f"Agent: {resp1}")
    
    # 第二轮对话（Agent 已记住前文）
    resp2 = agent_response(session_id, user_id, "帮我看看有哪些品牌可以选择？")
    print(f"Agent: {resp2}")
    
    # 保存关键知识到向量记忆（用于未来检索）
    memory.save_vector_memory(
        user_id=user_id,
        content="用户偏好：8000元预算笔记本电脑，关注性能和性价比",
        metadata={"category": "product_preference", "budget": 8000}
    )

性能优化与成本控制

在实际生产环境中，记忆系统的性能瓶颈通常在向量检索和 API 调用两个环节。以下是针对 HolySheep API 的优化策略：

1. Embedding 批量处理

def batch_embed_texts(texts: List[str], batch_size: int = 100) -> List[List[float]]:
    """批量生成 Embedding，降低 API 调用次数"""
    all_embeddings = []
    
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=batch
        )
        all_embeddings.extend([item.embedding for item in response.data])
        
    return all_embeddings

离线批量处理历史对话（节省 70% 成本）
def migrate_historical_conversations(conversations: List[Dict]):
    """将历史对话批量向量化入库"""
    texts = [conv['content'] for conv in conversations]
    embeddings = batch_embed_texts(texts)
    
    memory.collection.add(
        ids=[f"legacy_{conv['id']}" for conv in conversations],
        embeddings=embeddings,
        documents=texts,
        metadatas=[{"user_id": conv['user_id'], "source": "legacy"} for conv in conversations]
    )

2. 缓存策略

import hashlib
from functools import lru_cache

class CachedMemorySystem(MemorySystem):
    """带缓存的记忆系统，减少重复 API 调用"""
    
    def __init__(self, *args, cache_ttl=3600, **kwargs):
        super().__init__(*args, **kwargs)
        self.cache_ttl = cache_ttl  # 缓存 1 小时
    
    def _cache_key(self, text: str) -> str:
        return f"emb:{hashlib.md5(text.encode()).hexdigest()}"
    
    def _get_embedding_cached(self, text: str) -> List[float]:
        cache_key = self._cache_key(text)
        cached = self.redis_client.get(cache_key)
        
        if cached:
            return json.loads(cached)
        
        # 调用 HolySheep API
        embedding = super()._get_embedding(text)
        self.redis_client.setex(cache_key, self.cache_ttl, json.dumps(embedding))
        
        return embedding
    
    # 重写方法使用缓存
    def search_vector_memory(self, user_id: str, query: str, top_k: int = 5):
        # 复用缓存的 embedding
        query_embedding = self._get_embedding_cached(query)
        # ... 其余逻辑相同

常见报错排查

报错1：Redis 连接超时 "ConnectionError: Error 111 connecting to localhost:6379"

原因：Redis 服务未启动或端口被占用

解决：

# 检查 Redis 状态
sudo systemctl status redis-server

如果未启动，手动启动
sudo systemctl start redis-server

或使用 Docker 快速启动
docker run -d -p 6379:6379 redis:alpine

报错2：向量检索返回空结果 "ChromaDB returned empty result"

原因：向量数据库中尚无该用户的数据，或 embedding 模型返回失败

解决：

# 检查 collection 是否存在
print(memory.collection.count())

检查 embedding 是否正常生成
test_emb = memory._get_embedding("测试文本")
print(f"Embedding 维度: {len(test_emb)}")  # 应该是 1536 或 1024

确认用户 ID 过滤条件正确
results = memory.collection.get(where={"user_id": user_id})
print(f"该用户已有记录数: {len(results['ids'])}")

报错3：PostgreSQL 权限错误 "permission denied for table user_memories"

原因：数据库用户缺少表的读写

AI Agent 记忆系统设计：短期记忆 + 长期记忆 + 向量检索实战

为什么 AI Agent 需要三层记忆架构

系统架构设计与核心实现

1. 环境准备与依赖安装

安装命令

2. 核心记忆模块实现

HolySheep API 配置（汇率优势：¥1=$1，注册送免费额度）

数据库初始化 SQL

3. Agent 集成与调用示例

示例调用

性能优化与成本控制

1. Embedding 批量处理

离线批量处理历史对话（节省 70% 成本）

2. 缓存策略

常见报错排查

报错1：Redis 连接超时 "ConnectionError: Error 111 connecting to localhost:6379"

如果未启动，手动启动

或使用 Docker 快速启动

报错2：向量检索返回空结果 "ChromaDB returned empty result"

检查 embedding 是否正常生成

确认用户 ID 过滤条件正确

报错3：PostgreSQL 权限错误 "permission denied for table user_memories"

相关资源

相关文章

为什么 AI Agent 需要三层记忆架构

系统架构设计与核心实现

1. 环境准备与依赖安装

安装命令

2. 核心记忆模块实现

HolySheep API 配置（汇率优势：¥1=$1，注册送免费额度）

数据库初始化 SQL

3. Agent 集成与调用示例

示例调用

性能优化与成本控制

1. Embedding 批量处理

离线批量处理历史对话（节省 70% 成本）

2. 缓存策略

常见报错排查

报错1：Redis 连接超时 "ConnectionError: Error 111 connecting to localhost:6379"

如果未启动，手动启动

或使用 Docker 快速启动

报错2：向量检索返回空结果 "ChromaDB returned empty result"

检查 embedding 是否正常生成

确认用户 ID 过滤条件正确

报错3：PostgreSQL 权限错误 "permission denied for table user_memories"

相关资源

相关文章

🔥 推荐使用 HolySheep AI