DEV Community

BAOFUFAN
BAOFUFAN

Posted on

I Spent a Whole Weekend Fixing LangChain Memory Bugs with Pytest — Here’s What I Learned

Last Friday, just before clocking out, our product manager drifted over with a spooky silence: “Our AI customer service bot… why does it keep forgetting what the user said in the previous turn? We fixed this last month, and now it’s broken again.” I opened the logs and — sure enough — the conversation memory had been silently wiped in a specific branch. That was the third time this month.

In retrospect, the root cause wasn’t some hidden flaw in LangChain’s memory mechanism. It was much simpler: we had no automated regression suite for conversational memory. Every time we tweaked the prompt, swapped the model, or adjusted a chain parameter, we’d manually run a handful of exchanges. That’s nowhere near enough to cover edge cases in multi-turn dialogues. So I dug in and built a reusable memory regression test suite from scratch with pytest. This article walks through that journey — and all the “the official docs won’t tell you this” pitfalls I hit along the way.

Breaking Down the Problem: Why Manual Memory Tests Always Miss Something

Memory in an LLM-powered application is essentially a state store that gets written to and read from continuously during a conversation. It’s not like a traditional API where input-output pairs are deterministic. The same conversation history can produce different responses with a different model, a different temperature, or even the same model at a different time. This leads directly to two pain points:

  1. Non-deterministic memory content – you can’t simply assert “after round N the memory must contain this exact sentence,” because the LLM may rephrase things naturally.
  2. Memory boundaries rarely get triggered – token-limit truncation, summarization, multi-user isolation… all of these logic branches are almost impossible to reach through short manual chats.

A common workaround is mocking the LLM, but that defeats the purpose of regression testing. What we actually want to catch is whether real-world memory behavior degrades when the LLM environment changes. So the goal was clear: run multi-turn conversations against a real LLM (or a stable local model) and then perform semantic-level assertions on the memory state — not character-level comparisons.

Solution Design: Why Not Just Use LangChain’s Built-in Tools?

The LangChain open-source ecosystem offers little in the way of off-the-shelf memory test suites. I had a few paths in front of me:

  • Option A: LangSmith’s evaluation features. It’s paid, heavily dependent on cloud services, and not ideal for fast local CI/CD feedback loops.
  • Option B: Hand-roll a bunch of unittest scripts. They work, but they scale poorly. Parameterizing one dialogue scenario with multiple memory types means writing a ton of repetitive boilerplate.
  • Option C: pytest + fixtures + custom assertions. Pytest’s fixture system is a natural fit for managing shared chain and memory objects. Combined with a few semantic matching helpers, you can write highly readable cases like “the memory should contain the user’s address information.”

I went with Option C. The architecture ended up with three layers:

  • Fixture layer – initializes the LLM (I used a local Ollama model, qwen2:7b, for deterministic, reproducible behavior) and various memory types (ConversationBufferMemory, ConversationSummaryMemory, etc.).
  • Step layer – wraps a single “user input → model output → memory update” turn and returns a snapshot of the current state.
  • Assert layer – uses the LLM again to perform simple yes/no judgments, or extracts key entities for semantic assertions.

Core Implementation: From Fixtures to Semantic Assertions — Step by Step

1. Fixtures and a basic conversation helper — solving “how to reuse models and memory”

The conftest.py below extracts model and memory initialization so each test case can inject different parameters on demand.

# conftest.py
import pytest
from langchain_community.chat_models import ChatOllama
from langchain.memory import ConversationBufferMemory, ConversationSummaryMemory
from langchain.chains import ConversationChain

@pytest.fixture
def ollama_model():
    # Fixed parameters to ensure reproducible local regression results
    return ChatOllama(
        model="qwen2:7b",
        temperature=0,       # Deterministic output for comparison
        top_p=0.1,
    )

@pytest.fixture(params=["buffer", "summary"])
def memory(request, ollama_model):
    """Parametrized fixture: covers multiple Memory implementations in one go"""
    if request.param == "buffer":
        return ConversationBufferMemory(return_messages=True)
    elif request.param == "summary":
        # Use the same ollama_model for summarization
        return ConversationSummaryMemory(
            llm=ollama_model,
            return_messages=True,
        )
Enter fullscreen mode Exit fullscreen mode

Next comes a generic conversation function that takes a fully built chain and a user input, runs one interaction turn, and returns the updated memory contents.

# conftest.py (continued)
def run_turn(chain: ConversationChain, user_input: str) -> dict:
    """Perform one conversation turn and return the list of messages in memory"""
    _ = chain.run(user_input)                # Ignore the model's reply
    # Read all history messages directly from chain.memory
    messages = chain.memory.load_memory_variables({})["history"]
    return {
        "message_count": len(messages),
        "last_human_msg": user_input,
        "history_text": " ".join([m.content for m in messages]),
    }
Enter fullscreen mode Exit fullscreen mode

2. The first real test: verifying that fundamental memory persists across turns

The code below tackles the most essential regression assertion — “what the user said earlier, does it actually still live somewhere in memory in later turns?”

# test_memory_retention.py
import pytest
from langchain.chains import ConversationChain

@pytes
Enter fullscreen mode Exit fullscreen mode

Top comments (0)