LangChain Memory Cross-Talk: 3 Tests That Saved My Sanity

#programming #python

At 2:17 AM, my boss dropped a screenshot into the team chat: our AI assistant had sent last week’s proposal for Client A to Client B, word for word. Pulling the trace in LangSmith, a vector search had recalled a memory fragment that didn’t belong — similarity 0.83, just over the threshold. It was supposed to be from a totally different session. After fixing the bug in the middle of the night, I stared at vectorstore.similarity_search(score_threshold=0.8) and, for the first time, lost trust in my memory layer.

The very next day I made a rule: the memory store gets a full suite of automated tests — not optional, mandatory. Here are the three biggest pitfalls I hit and a reusable integration testing blueprint you can adopt right now.

Problem Breakdown: Why Testing an LLM Memory Layer Is So Hard

In traditional backend work, testing a storage layer is just CRUD assertions. But LLM memory brings three counter‑intuitive challenges:

Non‑determinism: The order of documents returned by a vector search can jitter because of floating‑point variation. similarity_search scores can differ by 0.01 between your local machine and CI, just due to CPU architecture.
Hidden side effects: LangChain’s ConversationBufferMemory or VectorStoreRetrieverMemory injects hidden keys into the context. You look at the docs, see memory.chat_memory.messages, and think you’ve cleared it — but the underlying vector store still holds the embeddings. The next retrieval brings those old memories right back.
Write amplification under concurrency: When multiple user threads write memories simultaneously, a mis‑configured connection pool or missing upsert can lead to a performance cliff or silently dropped messages.

Mocking the vector store or only testing the buffer layer catches none of this. Real confidence comes from integration tests that spin up a real vector database and run multi‑turn, concurrent interactions.

Solution Design: Real Dependencies, Deterministic Data

Tooling: pytest + pytest‑asyncio + Chroma (or pgvector) started via Docker. I skipped FAISS because it produces binary files that behave differently across environments; I skipped pure mocks because they hide connection‑pool leaks and serialization bugs — the real‑world problems.

Architecture:

Use docker‑compose or testcontainers to launch a Chroma container once per pytest session. All tests share the same vector store but use isolated collections.
Never use a real LLM for embeddings in tests. Instead, use FakeEmbeddings (fixed vectors) or pre‑built embedding arrays so that retrieval results are exactly predictable.
Design test cases around the three core risks: correctness (consistent recall after CRUD), concurrency safety (no lost or overwritten writes), and performance baseline (P99 latency, connection pool health).

Core Implementation: Three Tests for Peaceful Sleep

1. Memory Consistency Test — Fix the Cross‑Talk Bug

This test verifies that after writing a batch of messages to a given session, other sessions cannot retrieve them, while the same session recalls everything as expected. We use FakeEmbeddings to map text to fixed‑size vectors, making similarity scores fully deterministic.

# test_memory_consistency.py
import asyncio
import pytest
from langchain.vectorstores import Chroma
from langchain.embeddings import FakeEmbeddings

# 固定向量长度，确保测试可重复
FAKE_EMBED_SIZE = 384

@pytest.fixture(scope="session")
def vectorstore():
    """ 启动 Chroma 客户端，每个测试使用独立 collection 避免污染 """
    embeddings = FakeEmbeddings(size=FAKE_EMBED_SIZE)
    vs = Chroma(
        collection_name="test_memory",
        embedding_function=embeddings,
        persist_directory="/tmp/test_chroma"  # 可挂载到容器内路径
    )
    yield vs
    # 测试结束后彻底清空，防止残留
    vs.delete_collection()

def test_cross_session_isolation(vectorstore):
    """ 不同 session 的记忆必须严格隔离 """
    session_a, session_b = "user_123", "user_456"
    # 写入 session_a 的消息
    vectorstore.add_texts(
        texts=["我叫张三，预算50万"],
        metadatas=[{"session_id": session_a}]
    )
    # 检索 session_b 的记忆——不应有结果
    results_b = vectorstore.similarity_search(
        "预算", k=3,
        filter={"session_id": session_b}
    )
    assert len(results_b) == 0, "session B 不应该召回 A 的记忆"

    # session_a 自己应该能召回
    results_a = vectorstore.similarity_search(
        "预算", k=3,
        filter={"session_id": session_a}
    )
    assert any("50万" in doc.page_content for doc in results_a)

2. Concurrent Write Safety Test — Stop Lost Updates

Using asyncio.gather, we simulate multiple coroutines writing to the same session simultaneously and then assert that the total number of documents for that session equals the number of concurrent writers. This exposes a common pitfall: LangChain’s add_texts calls collection.add under the hood; without unique IDs, concurrent requests can each insert the full set of embeddings, leading to duplicates or OOM.

# test_concurrent_writes.py
import asyncio
import pytest
from langchain.vectorstores import Chroma
from lang

DEV Community