用 FastAPI + pgvector 构建 RAG 记忆系统

在开发 Psycho 这款心理陪伴 AI 时，我需要让 AI 记住用户的历史偏好、情绪状态和重要事件。本文记录了我如何用 pgvector 实现长期记忆检索。

为什么需要向量记忆

大语言模型的上下文窗口是有限的。用户每次对话不可能把几个月的历史全塞进去。但如果完全忘掉历史，AI 就会显得冷漠、重复。

向量记忆的思路是：

1. 每次对话结束后，提取关键事件和情绪状态存到数据库
2. 下次对话前，根据当前话题做语义检索，拉出最相关的几条记忆
3. 把这些记忆注入到 system prompt，让 AI 表现得像真的"记得你"

环境准备

# PostgreSQL + pgvector 扩展
docker run -d \
  -e POSTGRES_PASSWORD=yourpassword \
  -p 5432:5432 \
  pgvector/pgvector:pg16

# Python 依赖
pip install fastapi sqlalchemy asyncpg pgvector openai

数据库建表

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE memories (
  id          BIGSERIAL PRIMARY KEY,
  user_id     TEXT NOT NULL,
  content     TEXT NOT NULL,
  embedding   vector(1536),
  created_at  TIMESTAMPTZ DEFAULT now()
);

CREATE INDEX ON memories USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);

写入记忆

from openai import AsyncOpenAI
from sqlalchemy.ext.asyncio import AsyncSession

client = AsyncOpenAI()

async def save_memory(db: AsyncSession, user_id: str, content: str):
    resp = await client.embeddings.create(
        model="text-embedding-3-small",
        input=content,
    )
    embedding = resp.data[0].embedding

    await db.execute(
        "INSERT INTO memories (user_id, content, embedding) VALUES (:uid, :content, :emb)",
        {"uid": user_id, "content": content, "emb": embedding},
    )
    await db.commit()

检索记忆

async def retrieve_memories(db: AsyncSession, user_id: str, query: str, top_k: int = 5):
    resp = await client.embeddings.create(
        model="text-embedding-3-small",
        input=query,
    )
    query_emb = resp.data[0].embedding

    rows = await db.execute(
        """
        SELECT content, 1 - (embedding <=> :emb) AS similarity
        FROM memories
        WHERE user_id = :uid
        ORDER BY embedding <=> :emb
        LIMIT :k
        """,
        {"emb": query_emb, "uid": user_id, "k": top_k},
    )
    return rows.fetchall()

效果

在 Psycho 中，这套记忆系统让 AI 能记住「用户上周提到的工作压力」「用户喜欢晚上聊天」等细节，对话质量提升非常明显。

---

如果你在做类似的项目，欢迎联系交流。