At 1 a.m., I was jolted awake by an alert: the production AI assistant's prompt tokens had blown up again. I opened the logs and, lo and behold, three long-term memories for the same user had been stuffed into the context at once — "User likes American-style coffee", "User drinks American-style coffee without sugar every day", "User prefers American-style coffee, no sugar".
This wasn't a bug — it was our memory deduplication strategy running wild.
Breaking Down the Problem
We're building a long-term memory system for an AI personal assistant. After each conversation, the model extracts important information into memories and stores them in PostgreSQL. Before the next conversation, it retrieves relevant memories and assembles them into the prompt. The problem was in two places:
- Deduplication only looked at text similarity. We used Jaccard similarity + keyword hashing. Two memories would slip through as long as they were phrased differently. A user saying "I run three times a week" and "I'm used to jogging three times a week" have low text overlap, but they're semantically identical.
-
The decay strategy had no automated verification. The memory weight decay formula was documented, but in the actual code
decay_factorwas sometimes0.9, sometimes adatetimesubtraction error, and it had never been covered by tests.
The result: duplicate memories piled up, the context window got stuffed, and search results were full of semantically identical gibberish. The usual approach was to manually run a few SQL queries before each deploy to check whether nearest-neighbor queries returned the right results. But honestly, who has time to manually test every day? By the time we noticed, the duplication rate was already 18%.
Design Choices
The technology choice was a no-brainer: PostgreSQL + pgvector. The reason is practical — we already use Postgres for structured data in production. Adding Qdrant or Milvus would mean another piece of infrastructure and more operational overhead. pgvector extends the existing Postgres instance, so transactions, backups, and permissions all carry over.
For testing, we chose pytest + testcontainers. Why not mock pgvector? Because the biggest risk in a memory system is getting the vector functions wrong in SQL. Mocking would defeat the purpose. We use testcontainers to spin up a real pgvector/pgvector:pg16 container in CI, run the tests, and tear it down. Environment consistency is guaranteed.
The architecture is simple: an upsert_memory function handles deduplicated insertion, and apply_decay handles time-based weight decay. Tests cover these two core functions and run automatically on every PR.
Core Implementation
The first code block solves the test environment problem: use testcontainers to launch a pgvector container, create the table, and register the vector type. Without a clean, reproducible Postgres environment, all subsequent tests would be castles in the air.
# conftest.py
import pytest
from testcontainers.postgres import PostgresContainer
import psycopg
from pgvector.psycopg import register_vector
@pytest.fixture(scope="session")
def postgres_url():
# 使用 pgvector 官方镜像,pg16 版本
with PostgresContainer("pgvector/pgvector:pg16") as postgres:
postgres.with_env("POSTGRES_PASSWORD", "testpass")
yield postgres.get_connection_url()
@pytest.fixture(scope="session")
def conn(postgres_url):
# psycopg 3 连接
with psycopg.connect(postgres_url) as conn:
# 关键:注册 vector 类型,否则无法返回 embedding
register_vector(conn)
with conn.cursor() as cur:
cur.execute("""
CREATE TABLE memories (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
user_id TEXT NOT NULL,
content TEXT NOT NULL,
embedding vector(1536) NOT NULL,
importance FLOAT DEFAULT 1.0,
decay_rate FLOAT DEFAULT 0.05,
last_access_at TIMESTAMPTZ DEFAULT now(),
created_at TIMESTAMPTZ DEFAULT now()
)
""")
# 用 IVFFlat 索引,适合大数据量近似搜索
cur.execute("""
CREATE INDEX ON memories
USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 100)
""")
conn.commit()
yield conn
Note that the register_vector step isn't emphasized in the official docs, but without it, SELECT embedding will throw unsupported type.
The second code block solves the deduplicated insertion problem: when a new memory comes in, first query the nearest vector for the same user. If the distance is below the threshold, update the old memory; otherwise, insert a new row.
# memory_dedup.py
import uuid
import numpy as np
import psycopg
from pgvector.psycopg import register_vector
def upsert_memory(conn, user_id: str, content: str, embedding: list[float],
threshold: float = 0.3) -> str:
"""
去重插入记忆。pgvector 的 cosine distance = 1 - cosine similarity,
所以距离越小越相似。threshold 0.3 等价于相似度 0.7。
返回记忆 ID:重复返回旧 ID,否则返
Top comments (0)