5. RAG and Memory - A Comprehensive Dive

Imports and setup

Defines a create_kb function.

Splits a manually created Document using RecursiveCharacterTextSplitter with:

The document includes metadata such as source="blank_chain.md".

Uses split_documents() because the input is a Document object.

Creates a vector store from the chunks using vectorstore_from_documents(…​), the embeddings model, and a persistence directory like ./temp.

Returns the vector store.

Calls create_kb() to get the vector store.

Builds a retriever with:

Initializes a chat model with temperature=0.2.

Creates a prompt that instructs the model to answer only from the given context and say “I don’t know” if unsure.

Defines a helper like format_docs(docs) to join retrieved chunks into a single string for the prompt.

Creates a chain with two inputs:

Pipes the inputs through:

Produces a final string answer.

Invokes the chain with sample questions like:

Prints the responses.

The core setup stays the same: knowledge base/vector store, retriever, and LLM are reused.

The main change is that the system now returns sources or citations along with the answer.

Q&A bots

enterprise search

trustworthy AI systems

verify answers

trace them back to original documents

trust the results more easily

The pipeline uses:

The prompt instructs the model to:

The chain works by:

Questions covered by the knowledge base produce normal answers.

Questions outside the knowledge base trigger the fallback response.

It reduces hallucinations by preventing the model from guessing.

It makes the system more honest, reliable, and useful in real-world situations where users may ask unsupported questions.

Set up the vector store, retriever, and LLM.

Create an LLM chain extractor to act as the compressor.

Wrap the base retriever with a Contextual Compression Retriever using:

Run the query and compare:

In simple cases, compression may not seem dramatic because the documents are already focused.

In more complex documents, the reduction is much clearer:

The output keeps only the information needed to answer the question, such as framework names like LangChain and LangGraph.

It adds extra LLM calls during retrieval, which increases latency and cost.

For “What is Postgres?”, the ensemble gives the best combined result.

For “What database stores vectors?”, both retrievers identify relevant vector database content.

For “asset transactions”, BM25 succeeds where semantic search drifts off-topic.

For “How do I store AI model outputs for later retrieval?” and “fast similarity lookup embeddings”, both retrievers contribute useful signals.

BM25 is strong for exact keyword matching

Semantic search is strong for intent and meaning

Ensemble retrieval combines both to improve accuracy and robustness

Split documents into:

Search is done over the small child chunks

The system returns the full parent chunk to the LLM

Use an in-memory vector store for embeddings

Use an in-memory document store for parent documents

Name the collection something like parent-child-demo

Build the retriever with:

Small chunks

Large chunks

accurate search from small chunks

complete context from large chunks

first, find the most relevant small child chunk

then, return its corresponding larger parent document

Chat model setup using init_chat_model or ChatOpenAI

Prompt template with:

Message history storage with:

RunnableWithMessageHistory to automatically load and save messages for each session

StrOutputParser to format model output

remembering the user’s name

remembering what the user is learning

Use one shared LLM

Build a prompt that accepts:

Create a chain from the prompt and LLM

Store conversation histories in a dictionary

Add a helper function that gets or creates a session’s history

Wrap the chain so:

User A says: “My favorite language is Python.”

User B says: “I love JavaScript.”

“What is my favorite language?”

User A → Python

User B → JavaScript

remember past messages

keep conversations separate by user

answer based on each user’s own history

Recent context stays exact

Older context is preserved in compressed form

Token usage remains bounded, preventing context overflow

old info → summarized

new info → kept verbatim

Use RunnableWithMessageHistory and SQLChatMessageHistory to store conversation history in a local SQLite database.

Each chat session is identified by a session ID, so messages are saved and retrieved per user/session.

The chatbot can remember preferences across restarts, such as:

To make this work, you:

To confirm memory is truly saved, you can:

You can also inspect the SQLite database directly to see stored human and AI messages.

After about 10 messages, the conversation can be summarized automatically.

The summary can be stored as memory, while older raw messages may be pruned if desired.

remembers preferences,

persists across restarts,

stores memory locally,

and can later be extended with automatic summarization.

Introduces structured output models:

Builds an AIResearchAssistant class that bundles the three core RAG components:

The constructor sets defaults like:

Adds document ingestion methods:

Includes inspection utilities:

Confirms persistence and indexing through tests and cleanup steps

Add a ChatOpenAI model to the assistant.

Build a retriever that uses similarity search and returns the top 4 relevant chunks.

Test the retriever to confirm it returns relevant document fragments.

Add a function to format retrieved documents into plain-text context for the LLM.

Create an ask method that:

The current retriever only matches the query using similar words.

This works, but it can miss relevant chunks if the answer uses different terminology.

Multi-query retrieval fixes this by using an LLM to generate several semantically similar queries from the original question.

Basic retriever: simple similarity search, returns about four documents.

Advanced retriever: multi-query retrieval, which expands the original question into multiple related searches.

Update _build_retriever in AIResearchAssistant.

Add a use_advanced flag.

If false, use the basic retriever.

If true, use the multi-query retriever.

Update ask to pass this flag through and handle retrieval the same way afterward.

With debugging enabled, the advanced retriever shows generated alternate queries.

It returns more relevant and unique chunks than the basic retriever.

This gives the model richer context and usually improves answer quality.

The current ask function returns a plain string, which is hard to use programmatically.

A ResearchResponse data model already exists to solve this by structuring outputs with fields like:

A new function, ask_structured, is introduced right before ask.

Inside ask_structured, the LLM is wrapped with with_structured_output(ResearchResponse) so the model returns data in the schema format.

This makes it easy to access individual parts of the response directly in code, such as response.answer or response.sources.

Plain text is fine for display, but structured output is better for downstream processing.

It makes it easier to extract answers, confidence scores, sources, and follow-up questions.

The project now includes a full RAG pipeline with:

Document ingestion was intentionally left out for simplicity, but should be implemented separately in a real project: