DEV Community

BAOFUFAN
BAOFUFAN

Posted on

How I Spent 5 Hours Fixing Flaky LLM Memory Recall — and Locked It Down with pytest

At 4:30 a.m., the CI pipeline went red again. I opened the details: the assertion recalled_memories == expected_top3 failed — “the third memory is a completely different one.” Yet last week, this test case was green. Even more maddening: running the unit test locally, the failure rate was about 20%, with no pattern at all.

Our long‑term memory module had been live for two weeks, and recall stability was basically a roll of the dice. Every deployment felt like opening a mystery box. This post digs deep into that rabbit hole and shows how I used pytest automation to return determinism to LLM memory storage.

Breaking down the problem: why are LLM memory recall tests so fragile?

Here’s the scenario: a user has a multi‑turn conversation with an LLM, and the system needs to retrieve key historical information from a “long‑term memory store” — things like “the user is allergic to cats” or “book the Hilton for the next business trip to Tokyo.” The technical approach is standard: after each conversation turn, an embedding model generates a vector and stores it in a vector database. When a new conversation starts, the current query is vectorised with the same model, the top‑K most similar memories are retrieved, and they’re injected into the prompt.

The test objective was also straightforward: given a known set of memories and a fixed query, two calls to the recall interface must return exactly the same list of memory IDs. In pytest, it looked like this:

def test_recall_consistency(memory_store, sample_query):
    result1 = memory_store.search(sample_query, top_k=3)
    result2 = memory_store.search(sample_query, top_k=3)
    assert result1 == result2
Enter fullscreen mode Exit fullscreen mode

Yet it kept failing intermittently. I broke down the root causes into three layers:

  1. Non‑determinism in the embedding model itself: even with the same model (e.g., all-MiniLM-L6-v2), the same text can produce tiny floating‑point differences (on the order of 1e-6) across machines or even across different calls in the same process. That’s enough for 0.8532 and 0.8531 to swap places in a top‑K ranking.
  2. Unstable ordering from the vector database: even when similarity scores are identical, many databases (Milvus, Chroma, Qdrant included) return documents with the same score in an “implementation‑defined” or even random order, unless you explicitly add ORDER BY id.
  3. Tests were too dependent on a real model: loading sentence-transformers directly in tests is slow and introduces environmental differences — the floating‑point unit of a different CPU in CI can affect results.

The usual advice is “just don’t test it, rely on manual regression” or “increase the threshold and accept order wobble.” But in a memory scenario, a changed order means changed context inside the prompt, which can lead to completely different LLM answers. We must test, and we need 100% reliability.

Design strategy: keep non‑determinism out of the tests

Our chosen strategy was not to “tolerate” uncertainty, but to eliminate it entirely inside tests. The idea is simple:

  • Replace the real embedding model with a deterministic fake: do not call any neural network during tests; use predefined fixed vectors.
  • Enforce a stable sort on the vector database results: after retrieval, apply a secondary sort by (score, id).
  • Provide a configurable MemoryStore through pytest fixtures: production code uses the real model; test code swaps in a DeterministicEmbedder via dependency injection.

Why not other approaches?

  • Option A: using pytest.approx for floating‑point approximate comparison. That doesn’t work — we assert on “sets” and “order,” not exact scores. Approximate comparison won’t fix order drift.
  • Option B: pinning the model version and fixing random seeds in CI. Even with fixed seeds, PyTorch’s deterministic mode still shows differences across GPUs/CPUs. And it’s slow — tests would take 30 seconds each, which is unacceptable.
  • Option C: mocking the database return directly. Too thin — it wouldn’t test retrieval logic itself (top‑k truncation, score filtering). Better to build a lightweight in‑memory vector store plus deterministic vectors, striking the right balance between realism and speed.

Overall architecture: a pytest fixture creates an InMemoryVectorStore (pure Python, no external dependencies) combined with a DeterministicEmbedder (maps text to pseudo‑vectors of fixed length via hashing). The business‑logic MemoryService remains unaware of the swap. The test case only cares about “identical inputs must produce identical outputs”.

Core implementation: locking down tests with pytest, step by step

1. Build a deterministic embedder so the same text always yields the same vector

This code solves the problem of “the same text producing different vectors across calls.” Using Python’s built‑in hashlib, we map text to a fixed‑dimension floating‑point vector. As long as the text doesn’t change, the vector never changes.

import hashlib
import numpy as np

class DeterministicEmbedder:
    """Convert arbitrary text into a fixed, pseudo‑random vector — fully reproducible."""
    def __init__(self, dim: int = 384):
        self.dim = dim

    def embed(self, text: str) -> np.ndarray:
        # Use SHA256 over text + a fixed seed to guarantee identical results every time
        hash_bytes = hashlib.sha256(f"seed_v1:{text}".encode()).digest()
        # Fill vector to the required dimension by cycling through hash bytes
        vec = np.zeros(self.dim, dtype=np.float32)
        for i in range(self.dim):
            # Normalise byte to [-1, 1]
            vec[i] = (hash_bytes[i % len(hash_bytes)] / 127.5) - 1.0
        # Normalise for cosine similarity (dot product)
        norm = np.linalg.norm(vec)
        return vec / norm

    def embed_batch(self, texts: list[str]) -> list[np.ndarray]:
        return [self.embed(t) for t in texts]
Enter fullscreen mode Exit fullscreen mode

Now “user is allergic to cats” always produces exactly the same unique vector. Even if CI changes machines a thousand times, this vector won’t budge.

2. Build an in‑memory vector store with forced stable ordering

Next we write a lightweight InMemoryVectorStore that replaces a real database. This code addresses “random database order”: before returning results, we sort by (-score, id), guaranteeing stability when scores are identical by using the memory ID’s lexicographic order.

from dataclasses import dataclass, field
from typing import List, Tuple

@dataclass
class MemoryRecord:
    id: str
    vector: np.ndarray

class InMemoryVectorStore:
    def __init__(self, embedder):
        self.embedder = embedder
        self.records: List[MemoryRecord] = []

    def add(self, memory_id: str, text: str):
        vec = self.embedder.embed(text)
        self.records.append(MemoryRecord(id=memory_id, vector=vec))

    def search(self, query: str, top_k: int = 3) -> List[str]:
        query_vec = self.embedder.embed(query)
        scores = []
        for rec in self.records:
            # Cosine similarity via dot product (vectors are normalised)
            sim = np.dot(query_vec, rec.vector)
            scores.append((sim, rec.id))
        # Sort by score descending, then by id ascending to ensure stability
        scores.sort(key=lambda x: (-x[0], x[1]))
        return [doc_id for _, doc_id in scores[:top_k]]
Enter fullscreen mode Exit fullscreen mode

3. Wiring it up: a pytest fixture that makes the test deterministic

Finally, we create a fixture that assembles the deterministic embedder and in‑memory store. The service under test receives this store and can be tested for perfect recall consistency.

import pytest

@pytest.fixture
def deterministic_memory_store():
    embedder = DeterministicEmbedder(dim=384)
    store = InMemoryVectorStore(embedder)
    # Pre‑seed with known memories
    store.add("mem-1", "User is allergic to cats")
    store.add("mem-2", "Loves sushi, especially salmon")
    store.add("mem-3", "Next trip: Tokyo, prefers Hilton")
    return store

def test_recall_consistency_now_deterministic(deterministic_memory_store):
    query = "What do I need to know about the Tokyo trip?"
    result1 = deterministic_memory_store.search(query, top_k=3)
    result2 = deterministic_memory_store.search(query, top_k=3)
    assert result1 == result2
    # Optionally, assert the exact expected IDs
    assert result1 == ["mem-3", "mem-1", "mem-2"]  # example stable order
Enter fullscreen mode Exit fullscreen mode

This test will now never flake. The same text always maps to the same vector, and the store always sorts identically. CI is green, and every deployment no longer feels like a gamble.

Key takeaways

  • Test the logic, not the model. Replace non‑deterministic components (neural embeddings) with deterministic substitutes so that you can verify correctness of retrieval, ordering, and truncation.
  • Enforce stable ordering in your test doubles. Real vector DBs don’t guarantee order on equal scores; your tests should impose one ((score, id)) so that assertions are meaningful.
  • Keep tests fast and environment‑free. By using pure Python in‑memory structures, the test suite remains quick and never breaks because of a different GPU driver or CPU model.
  • pytest fixtures are perfect for swapping implementations. Inject a fast, deterministic fake at test time while keeping production code untouched.

If your LLM memory pipeline has been haunted by CI ghosts, give this approach a try. You’ll reclaim your nights and finally trust that green pipeline.

Top comments (0)