At 2 AM, I was jolted out of sleep by an angry user call: “Your AI has the memory of a goldfish — it forgets what I just said. Multi-turn conversations are completely broken!” After digging in, we found the culprit: during the last release, someone accidentally changed one parameter in the semantic search logic for memory storage. But our test suite never caught it, because we were still relying on manual Postman clicks. Each full regression cycle took over two hours — there was simply no way to cover everything.
This wasn’t one person’s oversight. It was a collapse of the entire testing strategy: trying to manually verify a memory system driven by vector similarity is essentially guesswork. In this article, I’ll share how we built an automated test harness on top of pgvector, turning the reliability of multi-turn memory persistence and semantic search accuracy from “lottery-level quality” into something that breaks right inside CI.
Why manual testing can’t catch memory bugs
The scenario is typical — smart customer service, role-playing chatbots, knowledge-base Q&A. All these AI products lean on a long-term memory module: conversation history gets embedded, stored in pgvector, and then retrieved by semantic search to assemble the prompt for the LLM. Two core quality attributes are at play:
- Persistence: After every turn, newly formed memories must be written to the database. They can’t disappear due to restarts, concurrency, or exceptions.
- Semantic search accuracy: For a given query, the top‑k returned memories should closely match what a human would deem “relevant”.
Manual testing fails on both fronts:
- Persistence: You send a request via Postman, then manually check the database. That’s doable once, but with multi-turn, cross-conversation, and concurrent scenarios, the number of combinations explodes — and humans make mistakes.
-
Search accuracy: You eyeball the top‑5 memory snippets and judge relevancy. But vector search is approximate (ANN). Results depend on the index type, the
probesparameter, and data distribution. The same query can return a different order at different times. Spot-checking one or two examples won’t reveal real recall fluctuations.
Even worse, most unit tests use SQLite or mock pgvector, completely sidestepping the actual index behavior. That’s testing nothing. We needed a fully automated regression suite running on a production-like pgvector instance, capable of surfacing any memory defect introduced by a code change — all within 10 minutes.
Design: pytest transactional rollback vs. Docker ephemeral database
The core stack was clear:
- Database: PostgreSQL 15+ with the pgvector extension, using the exact same version and configuration as development — no behavioral gaps.
-
Test framework: pytest, with SQLAlchemy and the
pgvectorPython client. - Isolation strategy: Each test case runs inside a transaction. When the test finishes, the transaction rolls back, leaving the database clean for the next case.
Why not spin up a fresh Docker‑compose database for every run? Because pgvector’s IVFFlat index needs a representative amount of data to train its lists. In an ephemeral DB, you create the index and immediately run queries — the lists aren’t fully trained, so recall results vary wildly. Tests become flaky, passing sometimes and failing others. A flickering test is far worse than no test at all.
Instead of scattering raw SQL across application code, we encapsulated all pgvector operations in a MemoryStore class. The tests would then validate this class’s public methods — exercising both the SQL logic and, indirectly, the index behavior.
Architecturally, the test suite is split into two layers:
- Contract tests: verify CRUD operations for memories, ensuring persistence.
- Semantic accuracy regression tests: using a hand‑labeled dataset of “conversation–memory” pairs, automatically evaluate top‑k recall. If a code change causes accuracy to drop below a threshold, CI fails immediately.
Core implementation: three tests that speak for themselves
1. Make MemoryStore testable first
This code shields the underlying SQL while still letting tests hit the real pgvector. We define the memory table with SQLAlchemy and use pgvector.sqlalchemy.Vector.
# memory_store.py
import numpy as np
from sqlalchemy import create_engine, Column, String, DateTime, func
from sqlalchemy.orm import declarative_base, sessionmaker
from pgvector.sqlalchemy import Vector
Base = declarative_base()
class Memory(Base):
__tablename__ = "ai_memories"
id = Column(String, primary_key=True)
user_id = Column(String, nullable=False)
content = Column(String, nullable=False)
embedding = Column(Vector(1536)) # OpenAI embedding dimension
created_at = Column(DateTime, server_default=func.now())
class MemoryStore:
def __init__(self, db_url: str):
self.engine = create_engine(db_url)
Base.metadata.create_all(self.engine)
self.Session = sessionmaker(bind=self.engine)
def add_memory(self, user_id: str, mem_id: str, content: str, embedding: list[float]):
with self.Session() as session:
mem = Memory(id=mem_id, user_id=user_id, content=content, embedding=embedding)
session.add(mem)
session.commit()
def search_memories(self, user_id: str, query_embedding: list[float], top_k: int = 5, probes: int = 3):
# Critical: setting probes influences recall, runs as SET LOCAL ivfflat.probes
with self.Session() as session:
session.execute("SET LOCAL ivfflat.probes = :probes", {"probes": probes})
results = (
# actual query construction here ...
)
(The full implementation continues with the search query and the rating evaluation, but the important part is that we are now talking directly to pgvector’s index.)
Top comments (0)