Entity memory in LangChain: tracking facts across turns

Entity memory solves a different problem than buffer or summary memory: it maintains a structured knowledge graph of entities and their attributes across turns, updating incrementally as new information arrives. This guide walks through setting up langchain entity memory tracking facts in a production context, covering store selection, extraction tuning, conflict resolution, and persistence.

What entity memory actually does

ConversationEntityMemory maintains two data structures: an entity store mapping entity names to summarized descriptions, and a chat history buffer for recent context. On each turn, an extraction chain runs against the new input plus recent history, identifies entities, and updates their summaries. The updated entity summaries are then injected into the prompt as a “current entities” section.

This differs from ConversationSummaryMemory in a critical way: summary memory compresses the entire conversation into a single narrative. Entity memory preserves discrete, queryable facts — “user prefers dark roast coffee,” “project deadline is March 15” — that can be retrieved and updated independently.

The tradeoff: extraction requires an additional LLM call per turn, and the quality depends entirely on your extraction prompt and the model’s ability to follow it.

Setting up the entity store

LangChain ships with three store implementations. Choose based on your persistence and scaling needs:

from langchain.memory import ConversationEntityMemory
from langchain.memory.entity import InMemoryEntityStore, RedisEntityStore
from langchain_community.llms import OpenAI

# Development only — lost on restart
store = InMemoryEntityStore()

# Production: survives restarts, works across workers
# Requires redis-py and a Redis instance
store = RedisEntityStore(redis_url="redis://localhost:6379/0")

# Custom: implement BaseEntityStore for Postgres, DynamoDB, etc.

The store interface is minimal — get(entity), set(entity, summary), delete(entity), clear() — so wrapping an existing database is straightforward. If you’re already using n4n.ai for model routing, the same Redis cluster can back your entity store with no additional infrastructure.

Initialize memory with the store and an LLM for extraction:

llm = OpenAI(temperature=0, model="gpt-3.5-turbo-instruct")

memory = ConversationEntityMemory(
    llm=llm,
    entity_store=store,
    k=3,  # recent turns to include in extraction context
    return_messages=True,
)

k controls how many recent messages feed the extraction chain. Higher values catch cross-turn references but increase token cost and latency. Start with 3–5.

Wiring it into a chain

The memory object plugs into any chain that accepts a memory parameter. For LCEL (LangChain Expression Language), use RunnableWithMessageHistory:

from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_openai import ChatOpenAI

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful assistant. Current entities:\n{entities}"),
    MessagesPlaceholder(variable_name="history"),
    ("human", "{input}"),
])

chain = prompt | ChatOpenAI(model="gpt-4o-mini", temperature=0)

chain_with_memory = RunnableWithMessageHistory(
    chain,
    lambda session_id: memory.chat_memory,
    input_messages_key="input",
    history_messages_key="history",
)

# Invoke
config = {"configurable": {"session_id": "user-123"}}
response = chain_with_memory.invoke({"input": "I'm planning a trip to Kyoto next month"}, config=config)

Note: ConversationEntityMemory exposes chat_memory (a BaseChatMessageHistory) for the history placeholder, while the entity summaries are injected via the {entities} variable automatically populated by the memory’s load_memory_variables.

Controlling extraction with prompts

The default extraction prompt is generic. For domain-specific fact tracking, override it:

from langchain.memory.entity import ENTITY_MEMORY_CONVERSATION_TEMPLATE
from langchain_core.prompts import PromptTemplate

custom_prompt = PromptTemplate.from_template("""
You are an entity extraction system for a travel planning assistant.
Extract entities and their attributes from the conversation below.
Focus on: destinations, dates, budgets, preferences, constraints, people.
Ignore: greetings, filler, meta-commentary.

Current entities:
{entities}

Conversation:
{history}
Human: {input}

Output JSON only:
{{
  "entities": {{
    "entity_name": "updated summary including all known attributes",
    ...
  }}
}}
""")

memory = ConversationEntityMemory(
    llm=llm,
    entity_store=store,
    extractor_prompt=custom_prompt,
    k=4,
)

Key design decisions in the prompt:

• Output format: Force JSON. The default parser expects a specific structure; custom prompts must match or you’ll need a custom parser.
• Scope restriction: Explicitly list what to extract. Without this, the model extracts every noun phrase, bloating the store.
• Update semantics: “Updated summary including all known attributes” tells the model to merge, not replace. The default prompt handles this poorly — it often drops old attributes.

Handling updates and conflicts

Entity memory merges by rewriting the entire entity summary on each extraction. This creates two problems:

1. Attribute drift: The model may subtly change wording or drop details across updates.
2. Contradictions: If the user says “I hate sushi” then later “I love sushi,” the last extraction wins silently.

Mitigation strategies:

# 1. Add explicit conflict detection to your extraction prompt
conflict_prompt = PromptTemplate.from_template("""
...
If the new information CONTRADICTS an existing entity attribute, 
include "conflict": true and "conflict_detail": "description" in that entity's output.
""")

# 2. Post-process: compare old vs new summaries before committing
class ValidatedEntityStore(InMemoryEntityStore):
    def set(self, entity: str, summary: str) -> None:
        old = self.get(entity)
        if old and self._semantically_different(old, summary):
            # Log, alert, or queue for human review
            logger.warning(f"Entity '{entity}' changed significantly: {old} -> {summary}")
        super().set(entity, summary)
    
    def _semantically_different(self, a: str, b: str) -> bool:
        # Cheap heuristic: length change > 50% or key terms dropped
        # For production, use an embedding similarity check
        return abs(len(a) - len(b)) / max(len(a), len(b)) > 0.5

For high-stakes domains (medical, legal, financial), don’t rely on automatic merging. Stage extracted facts in a pending state and require confirmation.

Persistence and production concerns

Redis configuration

import redis
from langchain.memory.entity import RedisEntityStore

redis_client = redis.Redis(
    host="localhost",
    port=6379,
    db=0,
    decode_responses=True,
    max_connections=20,
    socket_timeout=2,
    socket_connect_timeout=2,
)

store = RedisEntityStore(redis_client=redis_client, key_prefix="entity:")

Use a dedicated Redis database (not shared with caching) and set TTLs if entities should expire:

# Extend RedisEntityStore to add TTL
class TTLRedisEntityStore(RedisEntityStore):
    def __init__(self, *args, ttl_seconds: int = 86400 * 30, **kwargs):
        super().__init__(*args, **kwargs)
        self.ttl = ttl_seconds
    
    def set(self, entity: str, summary: str) -> None:
        key = f"{self.key_prefix}{entity}"
        self.redis.set(key, summary, ex=self.ttl)

Token budget management

Entity summaries grow unbounded. A single entity accumulating 50 turns of detail can blow your context window. Cap summaries in the extractor prompt:

Keep each entity summary under 80 words. Prioritize recent and decision-relevant facts.

Or enforce it in code:

class CappedEntityStore(InMemoryEntityStore):
    MAX_WORDS = 100
    
    def set(self, entity: str, summary: str) -> None:
        words = summary.split()
        if len(words) > self.MAX_WORDS:
            summary = " ".join(words[-self.MAX_WORDS:]) + " [truncated]"
        super().set(entity, summary)

Concurrency

ConversationEntityMemory is not thread-safe. In a multi-worker deployment (Gunicorn, uWSGI, Kubernetes), each worker gets its own memory instance but shares the Redis store. The race condition: two workers read the same entity, both update, last write wins.

Fix options:

• Optimistic locking: Add a version field to entity values, retry on conflict.
• Single-writer: Route all chat for a session to the same worker (sticky sessions).
• Accept it: For most chat applications, occasional lost updates are tolerable.

Common pitfalls

Pitfall 1: Extracting the world

Default extraction pulls every proper noun. After 20 turns you have 200 entities — “Tuesday,” “the weather,” “my dog” — most useless.

Fix: Restrict extraction via prompt (see above) and post-filter:

ALLOWED_TYPES = {"destination", "date", "budget", "preference", "person", "constraint"}

def filter_entities(entities: dict) -> dict:
    # Requires your extraction prompt to include a "type" field per entity
    return {k: v for k, v in entities.items() if v.get("type") in ALLOWED_TYPES}

Pitfall 2: Entity name fragmentation

“Kyoto,” “Kyoto, Japan,” and “Kyoto city” become three separate entities.

Fix: Normalize in the extraction prompt:

Canonicalize entity names: use "Kyoto" not "Kyoto, Japan". Merge variants.

Or normalize at write time:

CANONICAL = {
    "kyoto, japan": "Kyoto",
    "kyoto city": "Kyoto",
    "tokyo": "Tokyo",
}

class NormalizedStore(InMemoryEntityStore):
    def _norm(self, entity: str) -> str:
        return CANONICAL.get(entity.lower(), entity)
    
    def get(self, entity: str) -> str | None:
        return super().get(self._norm(entity))
    
    def set(self, entity: str, summary: str) -> None:
        super().set(self._norm(entity), summary)

Pitfall 3: Stale entities polluting context

A user mentioned “Paris trip” three months ago. The entity persists and gets injected into every new conversation, confusing the model.

Fix: TTL (see above) or session-scoped stores. For multi-session users, use a composite key: user:{id}:session:{session_id}.

Pitfall 4: Extraction latency

Each turn adds 500–1500ms for the extraction call. In streaming UX, this blocks the first token.

Fix: Run extraction asynchronously, decoupled from the response:

import asyncio
from concurrent.futures import ThreadPoolExecutor

executor = ThreadPoolExecutor(max_workers=4)

async def extract_async(memory, input_text, history):
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(executor, memory.save_context, {"input": input_text}, {"output": ""})

# In your handler:
async def handle_message(input_text, session_id):
    # Start extraction in background
    asyncio.create_task(extract_async(memory, input_text, recent_history))
    # Stream response immediately using existing memory state
    async for chunk in chain_with_memory.astream({"input": input_text}, config=config):
        yield chunk

The user sees the response immediately; entity updates apply on the next turn. Acceptable for most cases.

When to use something else

Entity memory isn’t universal. Consider alternatives:

Entity memory shines when: conversations span 10–100+ turns, facts are user-specific and evolving, and you need to query “what do we know about X?” at any point.

Testing your extraction

Write evals for the extraction prompt. Not optional.

import json
from langchain_core.outputs import Generation

TEST_CASES = [
    {
        "input": "I want to go to Tokyo in April, budget around $3000",
        "history": [],
        "expected_entities": {
            "Tokyo": {"type": "destination", "attributes": {"month": "April"}},
            "budget": {"type": "budget", "attributes": {"amount": 3000, "currency": "USD"}},
        },
    },
    {
        "input": "Actually make it May, and I prefer boutique hotels",
        "history": [("human", "I want to go to Tokyo in April"), ("ai", "Great!")],
        "expected_entities": {
            "Tokyo": {"type": "destination", "attributes": {"month": "May"}},
            "budget": {"type": "budget", "attributes": {"amount": 3000, "currency": "USD"}},
            "accommodation": {"type": "preference", "attributes": {"style": "boutique hotels"}},
        },
    },
]

def eval_extraction(memory, test_cases):
    for tc in test_cases:
        # Manually run the extraction chain
        entities = memory.entity_store.store  # current state
        # ... invoke extractor ...
        # Compare extracted vs expected
        pass

Test for: correct entity identification, attribute merging (not replacement), conflict handling, and canonicalization.

Entity memory gives you a queryable fact layer over conversation history. The moving parts — store, extraction prompt, merge logic, persistence — each need deliberate choices. Start with InMemoryEntityStore and a restrictive extraction prompt, add Redis and TTL when you deploy, and instrument the entity update path so you can audit what the model actually learns.