At 2 a.m. an alert yanked me out of bed. I pulled up the logs and my blood pressure went through the roof: a user asked “Where is my order?” and the AI agent replied “What’s your favorite color?” I had trained this agent myself, run over 20 rounds of dialogue before deployment, and everything looked flawless. Right then I knew—it had to be the memory system again. Six hours later I finally pinpointed the root cause, and it made me want to throw my keyboard out the window: LangChain’s Memory silently drops messages under certain edge cases, and manual testing never covered them. In this post I’ll share the complete memory-consistency testing setup I built with Pytest—along with the debugging journey—so you can save those six hours.
Problem Breakdown: Why “20 rounds worked” still exploded
Our scenario is an order-inquiry agent. Users provide their phone number and order ID over multiple turns, the agent calls a lookup tool, and the conversation might branch into phone-number changes, coupon inquiries, and the like. We used LangChain’s ConversationBufferMemory to store history and injected it into the prompt via {chat_history}.
During the first week in production we caught occasional reports that “the agent seems to forget who I am,” but reproducing it was tough. That night we finally pulled the full-chain logs and spotted a pattern:
- When a tool returned really long content (e.g., a hundreds-of-characters coupon list), by the time the next user message arrived, only the last two messages survived in memory. The phone number and order ID from earlier rounds simply evaporated. Deprived of context, the agent would blurt out a random fallback reply.
The root cause? We had set k=3 on the Memory to control token usage, meaning “keep only the last 3 turns of interaction.” But in ConversationBufferWindowMemory with return_messages=True, when the total token length exceeded the model limit, LangChain would perform a secondary trim—at a higher priority—that completely ignored the k value, chopping history arbitrarily. This behavior wasn’t documented anywhere.
Why did our regular tests miss it? We used to spin up a real LLM, run 20-scripted conversations, and spot-check outputs manually. Slow, expensive, and because of LLM randomness we couldn’t write assertions. Those tests could only catch “completely broken” bugs, never the silent memory corruption.
Solution Design: Make Pytest the Black-Box Judge of Memory
The core idea is dead simple: mock all LLM calls and test only the Memory read/write behavior inside the Agent flow. Tests execute in milliseconds, cost zero, and let you parametrize the heck out of edge cases.
Technical choices:
-
Pytest — parametrization and fixtures are far more elegant than unittest.
@pytest.mark.parametrizelets a single function run dozens of cases. -
unittest.mock — patch
ChatOpenAIand similar LLM calls to make the agent’s “thinking” fully deterministic. - Why not LangSmith or other platforms? — They’re great, but require internet access, introduce latency, and aren’t suitable for fast local TDD red-green loops. Budget was tight, too.
- Why not integration tests? — Integration tests ensure components interact correctly, but memory-logic edge cases need to be violently enumerated with unit tests. They complement each other.
Architecture-wise, we designed a set of fixtures: mock_llm → inject a custom reply sequence → create AgentExecutor or ConversationChain → run multiple turns → assert the contents of memory.chat_memory.messages. This puts the memory component in the interrogation room all by itself.
Core Implementation: Three Steps to a Memory-Consistency Test
Step 1: Build an obedient Mock LLM
This snippet solves the “uncontrollable LLM” problem. The mock returns preset replies one by one, simulating an AI that always follows the script.
from unittest.mock import MagicMock, patch
from langchain.schema import AIMessage, HumanMessage
from langchain.chat_models import ChatOpenAI
import pytest
def create_mock_llm(responses):
"""创建一个顺序返回指定回复的mock LLM"""
mock = MagicMock(spec=ChatOpenAI)
# 让mock的invoke方法依次返回AIMessage
mock.invoke.side_effect = [AIMessage(content=resp) for resp in responses]
mock.return_value = mock # 支持链式调用
return mock
Step 2: Wire up a Chain with Memory, run multiple turns
This code verifies that under normal flow, memory accumulates correctly. A pytest fixture initializes the chain; the test function feeds three user messages back to back, then asserts the correct number of messages in memory.
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory
@pytest.fixture
def chain_with_memory():
# 注意:return_messages=True 让memory存储消息对象,便于断言
memory = ConversationBufferMemory(return_messages=True)
mock_llm = create_mock_llm(["你好!", "你叫什么?", "我叫小智"])
chain = ConversationChain(llm=mock_llm, memory=memory)
return chain
def test_memory_accumulates_correctly(chain_with_memory):
chain = chain_with_memory
chain.predict(input="你好")
chain.predict(input="我叫小明")
chain.predict(input="帮我查订单")
messages = chain.memory.chat_memory.messages
# 期望3轮对话 = 6条消息(Human/AI交替)
assert len(messages) == 6
# 第一条用户消息应始终保留
assert messages[0].content == "你好"
# 最后一条AI消息应为我预设的"我叫小智"
assert messages[-1].content == "我叫小智"
Step 3: Reproduce the bug that cost me 6 hours
This test reveals the boundary vulnerability in k truncation. We deliberately set k=2 (keep only the last 2 rounds) and run three rounds of dialogue. If early information still exists, the configuration failed to take effect; if the third user message is present but the second is missing, it indicates the truncation logic has a bug—exactly the silent corruption we saw in production. By writing this as a fast unit test, we caught the exact scenario that took me all night to find in logs, and we can now run it in seconds on every commit.
Top comments (0)