AI Agent Memory Storage Pitfalls: One pgvector Index Parameter That Cost Me 4 Hours Late at Night

At 1:30 AM, I was woken by an alert message—the production AI agent's memory retrieval suddenly returned duplicate content, and users started complaining. When I opened Grafana, I saw that similarity query latency had spiked from 20ms to 800ms. My first thought was that the database connection pool was exhausted. But when I logged into PostgreSQL and checked pg_stat_activity, I found a bunch of queries running, but the execution plan showed Seq Scan on agent_memory. Where was the index? I had clearly created an IVFFlat index.

Breaking Down the Problem

Our team builds long-term memory for an AI agent, storing embeddings with pgvector. Each memory is a 384-dimensional vector, and queries use cosine similarity to retrieve the top-k results. During local development, the data volume was small, so a full table scan didn't feel like a problem. But once production data reached 1 million rows, the index not being used directly dragged down queries.

The root cause is that IVFFlat indexes have two hidden conditions: first, they won't be used at all when the data volume is insufficient; second, the ivfflat.probes parameter must be set during queries, otherwise the default value may be low enough that the optimizer gives up on the index. Manual testing with just a few SQL queries can't cover this kind of data distribution and parameter combination. So we decided to automate regression tests using pytest + pgvector to simulate the production scenario.

Solution Design

We went straight to setting up an automated regression test suite with pytest + pgvector. The test environment uses Docker to spin up a real PostgreSQL 16 + pgvector image—no mocking, because you can't mock index behavior. We use numpy to generate controllable random vectors, insert enough data (at least several thousand rows) to simulate production distribution, and then assert the top-k results, similarity ordering, and whether the query plan uses the index.

Why not other approaches? SQLite doesn't support pgvector; using SQLAlchemy's in-memory fake database would mask SQL execution plan issues; pure SQL scripts can't be automated for regression and CI integration. We chose pytest because its fixture mechanism handles database lifecycle and test data isolation well.

Core Implementation

First Code: conftest.py Solves Test Database Initialization

This code connects to a real pgvector instance, creates the extension and table schema, inserts random vector data, and builds the IVFFlat index. The key point: the index must be created after data insertion, otherwise the list data will be inaccurate.

# conftest.py
import os
import pytest
import psycopg2
from pgvector.psycopg2 import register_vector

@pytest.fixture(scope="session")
def db_conn():
    # 连接真实 pgvector 实例（Docker 启动）
    conn = psycopg2.connect(
        host=os.getenv("PGHOST", "localhost"),
        port=int(os.getenv("PGPORT", 5432)),
        user=os.getenv("PGUSER", "postgres"),
        password=os.getenv("PGPASSWORD", "postgres"),
        dbname=os.getenv("PGDATABASE", "postgres")
    )
    register_vector(conn)
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS agent_memory (
                id SERIAL PRIMARY KEY,
                content TEXT NOT NULL,
                embedding vector(3)  -- 演示用 3 维，实际可改 384
            )
        """)
        # 先清空，避免上次测试残留
        cur.execute("TRUNCATE agent_memory RESTART IDENTITY")
        # 插入足够多的随机向量，让 IVFFlat 索引有数据可训练
        cur.execute("""
            INSERT INTO agent_memory (content, embedding)
            SELECT 'memory_' || i,
                   ARRAY[
                       (random() * 2 - 1)::real,
                       (random() * 2 - 1)::real,
                       (random() * 2 - 1)::real
                   ]::vector
            FROM generate_series(1, 5000) AS i
        """)
        # 关键：IVFFlat 必须在数据插入后再建，否则列表数据不准
        cur.execute("""
            CREATE INDEX IF NOT EXISTS agent_memory_embedding_idx
            ON agent_memory USING ivfflat (embedding vector_cosine_ops)
            WITH (lists = 100)
        """)
    conn.commit()
    yield conn
    conn.close()

Second Code: First Test Case in test_memory_store.py

This test addresses the assertion of similarity search result correctness. It inserts two known vectors—one similar and one dissimilar—then queries and asserts whether the ordering matches expectations.

# test_memory_store.py
import numpy as np
import psycopg2
from pgvector.psycopg2 import register_vector

def test_similarity_search_returns_correct_order(db_conn):
    register_vector(db_conn)
    target = np.array([0.1, 0.2, 0.3], dtype=np.float32)
    similar_vec = target + np.array([0.01, 0.0