Memory Augmented RAG - Conversational Context Across Turns

Overview

Memory Augmented RAG solves the statelessness problem in multi-turn RAG systems. Standard RAG treats every query in isolation, so follow-up questions lose all prior context. Memory Augmented RAG maintains short-term conversational history and optionally a long-term persistent memory store, and rewrites each new query to be self-contained before retrieval.

Pattern #18 in the 37 RAG Patterns Collection

flowchart TD
  Q["User query — turn N"] --> PR["Pronoun resolution: replace 'it', 'they', 'that' using last 3 turns"]
  PR --> EQ[Enriched self-contained query]
  EQ --> VS[Vector search documents top-4]
  EQ --> HS[Search conversation history window last 5 turns]
  VS --> P["Build prompt: HISTORY + DOCS + QUESTION"]
  HS --> P
  P --> LLM[Generate]
  LLM --> ANS[Answer]
  ANS --> MEM[(Update sliding window — drop oldest turn if full)]

The Problem

Each RAG query is stateless by default:

Turn 1:
  User: "Explain the Transformer architecture."
  RAG:  Retrieves 5 relevant chunks → generates good answer

Turn 2:
  User: "What are the limitations of that approach?"
  RAG:  query = "What are the limitations of that approach?"
        → "that approach" resolves to nothing
        → retrieves generic "AI limitations" docs
        → generates irrelevant answer ❌

Turn 3:
  User: "How did the authors address those in later work?"
  RAG:  query = "How did the authors address those in later work?"
        → "the authors" and "those" are unresolvable pronouns ❌

Without memory, co-references and ellipsis in follow-up questions produce completely wrong retrievals.

Solution

Maintain conversation history and use it to rewrite follow-up queries into standalone questions before retrieval:

Turn 2 — with memory:
  History: ["User: Explain Transformer architecture", "AI: ..."]
  Raw query: "What are the limitations of that approach?"
  Rewritten: "What are the limitations of the Transformer architecture?"
  → Retrieves correct documents ✅

Memory Types

Short-Term Memory (Sliding Window)

Keeps the last N turns in working memory. Fast and lightweight; loses context after the window slides past.

python

from collections import deque

class ShortTermMemory:
    def __init__(self, window_size: int = 6):
        self.history = deque(maxlen=window_size)

    def add(self, role: str, content: str):
        self.history.append({"role": role, "content": content})

    def get_context(self) -> str:
        return "\n".join(f"{m['role']}: {m['content']}" for m in self.history)

Long-Term Memory (Summarised Store)

Periodically compresses older turns into a persistent summary. Survives across sessions.

python

def compress_history(history: list[dict], llm) -> str:
    """Summarise a block of conversation turns into a compact memory entry."""
    turns = "\n".join(f"{m['role']}: {m['content']}" for m in history)
    prompt = f"""Summarise the following conversation into 2-3 concise sentences,
preserving key entities, decisions, and facts discussed.

Conversation:
{turns}

Summary:"""
    return llm.invoke(prompt)

Implementation

python

class MemoryAugmentedRAG:
    def __init__(self, vector_store, llm, window_size: int = 6):
        self.vector_store = vector_store
        self.llm = llm
        self.short_term = ShortTermMemory(window_size)
        self.long_term_summary = ""
        self.turn_count = 0
        self.compress_every = 10  # compress after every 10 turns

    def rewrite_query(self, raw_query: str) -> str:
        """Use conversation history to rewrite follow-up into a standalone question."""
        history_ctx = self.short_term.get_context()
        if not history_ctx:
            return raw_query  # first turn — no rewriting needed

        prompt = f"""Given the following conversation history and a follow-up question,
rewrite the follow-up question to be fully self-contained. Do not answer it.

Conversation history:
{history_ctx}

Follow-up question: {raw_query}

Standalone question:"""
        return self.llm.invoke(prompt)

    def extract_entities(self, text: str) -> list[str]:
        """Track key entities for retrieval prioritisation."""
        prompt = f"List the main named entities (people, products, concepts) in: {text}"
        return self.llm.invoke(prompt, schema={"type": "array", "items": {"type": "string"}})

    def query(self, user_input: str) -> str:
        # Step 1: Rewrite to standalone query
        standalone_query = self.rewrite_query(user_input)

        # Step 2: Extract entities for boosted retrieval
        entities = self.extract_entities(standalone_query)

        # Step 3: Retrieve — boost results that mention tracked entities
        candidates = self.vector_store.search(standalone_query, top_k=10)
        entity_boosted = sorted(
            candidates,
            key=lambda d: sum(e.lower() in d["content"].lower() for e in entities),
            reverse=True,
        )
        top_docs = entity_boosted[:5]

        # Step 4: Build context — include long-term summary if available
        context_parts = []
        if self.long_term_summary:
            context_parts.append(f"Background from earlier: {self.long_term_summary}")
        context_parts.extend(d["content"] for d in top_docs)
        context = "\n\n".join(context_parts)

        # Step 5: Generate
        response = self.llm.generate(
            f"Context: {context}\n\nConversation history:\n"
            f"{self.short_term.get_context()}\n\nUser: {user_input}"
        )

        # Step 6: Update memory
        self.short_term.add("user", user_input)
        self.short_term.add("assistant", response)
        self.turn_count += 1

        # Periodically compress to long-term memory
        if self.turn_count % self.compress_every == 0:
            self.long_term_summary = compress_history(
                list(self.short_term.history), self.llm
            )

        return response

Complete Notebook

Run the full implementation:

Memory Augmented RAG AWS Notebook

Includes:

▸Short-term sliding window and long-term summary memory
▸Query rewriting with conversation history
▸Entity tracking and retrieval boosting
▸Multi-session persistence with DynamoDB
▸Benchmark on follow-up question accuracy

Performance

Metric	Without Memory	With Memory
Follow-up accuracy	43%	91%
Co-reference resolution	28%	87%
Avg tokens/query	800	1,050
Latency	~300ms	~420ms

Memory raises follow-up accuracy from 43% to 91% at the cost of ~120ms extra latency.

When to Use

Best For

▸Chatbots and virtual assistants — users expect follow-up questions to work naturally
▸Customer support systems — agents need full conversation context to resolve issues
▸Multi-turn research assistants — users iteratively refine their understanding of a topic
▸Personalised tutoring — retaining what a learner has already covered across sessions

Combine With

▸Corrective RAG (#13) — validate that memory-conditioned retrievals are actually relevant before generation
▸Streaming RAG (#32) — stream tokens in real time while memory and retrieval run asynchronously for better UX

Next Steps

▸Try it: Run the notebook
▸Persist across sessions: Store the long-term summary and entity list in DynamoDB keyed by user ID
▸Add quality checks: Combine with Corrective RAG (#13) to verify retrieved chunks are relevant after query rewriting
▸Improve UX: Pair with Streaming RAG (#32) so responses begin appearing while retrieval runs

Part of: 37 RAG Patterns Collection See also: Complete Pattern Index

Memory Augmented RAG - Conversational Context Across Turns

Overview

The Problem

Solution

Memory Types

Short-Term Memory (Sliding Window)

Long-Term Memory (Summarised Store)

Implementation

Complete Notebook

Performance

When to Use

Best For

Combine With

Related Papers

Next Steps

Related Concepts

Go Deeper With Live Instruction