How the ContextBuilder Constructs RAG Prompts for Chat Workflows in Open Notebook
The ContextBuilder assembles Retrieval-Augmented Generation (RAG) prompts by collecting source documents and insights, deduplicating and prioritizing them against token limits, and formatting them into structured sections that are injected into a Jinja2 system prompt template.
The lfnovo/open-notebook repository implements a sophisticated RAG pipeline for chat workflows using the ContextBuilder utility. This component serves as the engine that retrieves, filters, and formats contextual data before injection into Large Language Model (LLM) prompts. Understanding how the context_builder constructs RAG prompts reveals how the system balances information density with token constraints while maintaining semantic relevance.
The Three-Stage RAG Construction Pipeline
The chat workflows (specifically the source-chat graph) delegate context assembly to the generic ContextBuilder, which operates through three distinct stages: raw item collection, intelligent post-processing, and LLM-ready formatting.
Stage 1: Collecting Raw Items
The process begins with ContextBuilder.build(), which gathers sources, notes, and insights according to workflow parameters such as source_id, include_insights, and max_tokens.
In open_notebook/utils/context_builder.py, the builder loads source objects and their associated insights through _add_source_context (lines 42-80). When a notebook ID is present, it additionally pulls notebook-wide sources and notes via _add_notebook_context (lines 118-150).
Stage 2: Post-Processing and Optimization
After collection, the raw items undergo three critical optimization steps to ensure the context fits within model constraints while preserving the most relevant information:
- Deduplication (
remove_duplicates, lines 52-64): Eliminates duplicate items by ID to prevent token waste on redundant content. - Prioritization (
prioritize, lines 15-18): Reorders items using thepriority_weightsconfiguration, ensuring high-value insights appear first. - Truncation (
truncate_to_fit, lines 30-46): Cuts the list to respect themax_tokenslimit, preventing context window overflow.
The _format_response method (lines 90-100) returns a structured dictionary containing four top-level keys: sources, notes, insights, and metadata (which includes token counts for budget tracking).
Stage 3: Formatting for LLM Consumption
The chat graph converts the structured context into a human-readable string via _format_source_context in open_notebook/graphs/source_chat.py (lines 90-140). This function creates distinct markdown sections such as "## SOURCE CONTENT" and "## SOURCE INSIGHTS", inserting titles, IDs, and truncated full-text snippets.
The formatted string is then merged into a Jinja2 system-prompt template (source_chat/system) using Prompter.render (lines 27-32 of source_chat.py). Finally, the constructed system prompt is concatenated with chat history and passed to provision_langchain_model (lines 33-43) for inference.
Integration with Chat Workflows
Because ContextBuilder operates asynchronously, the chat workflow executes it within a temporary event loop. The builder accepts configuration through the ContextConfig dataclass, allowing workflows to specify boolean flags for including insights or notes, along with explicit token budgets.
from open_notebook.utils.context_builder import ContextBuilder, ContextConfig
async def get_notebook_rag(notebook_id: str):
cfg = ContextConfig(max_tokens=30_000, include_insights=True, include_notes=True)
builder = ContextBuilder(notebook_id=notebook_id, context_config=cfg)
return await builder.build()
Within the synchronous graph execution, the workflow spawns a fresh event loop to run the builder:
def build_context():
# Run the async builder in a fresh event loop
new_loop = asyncio.new_event_loop()
asyncio.set_event_loop(new_loop)
cb = ContextBuilder(
source_id=state["source_id"],
include_insights=True,
include_notes=False,
max_tokens=50_000,
)
return new_loop.run_until_complete(cb.build())
Formatting and Prompt Injection
The raw context dictionary undergoes transformation into the final RAG prompt through manual string construction that respects the token budget:
def _format_source_context(ctx: dict) -> str:
parts = []
if ctx.get("sources"):
parts.append("## SOURCE CONTENT")
for src in ctx["sources"]:
parts.append(f"**Source ID:** {src.get('id')}")
parts.append(f"**Title:** {src.get('title')}")
txt = src.get("full_text", "")
if len(txt) > 5_000:
txt = txt[:5_000] + "... [truncated]"
parts.append(f"**Content:**\n{txt}\n")
if ctx.get("insights"):
parts.append("## SOURCE INSIGHTS")
for ins in ctx["insights"]:
parts.append(f"**Insight ID:** {ins.get('id')}")
parts.append(f"**Type:** {ins.get('insight_type')}")
parts.append(f"**Content:** {ins.get('content')}\n")
return "\n".join(parts)
The formatted context is then injected into the system prompt template:
prompt_data = {
"source": source.model_dump() if source else None,
"insights": [i.model_dump() for i in insights],
"context": formatted_context,
"context_indicators": context_indicators,
}
system_prompt = Prompter(prompt_template="source_chat/system").render(data=prompt_data)
payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
Key Implementation Files
open_notebook/utils/context_builder.py: Core builder implementing collection, deduplication, prioritization, and token truncation.open_notebook/graphs/source_chat.py: Chat workflow that invokes the builder, formats context, and manages LLM interaction.open_notebook/ai/provision.py: Provides theprovision_langchain_modelhelper for model instantiation.open_notebook/utils/__init__.py: Re-exportsContextBuilderfor convenient imports throughout the application.
Summary
- The
ContextBuilderoperates asynchronously to gather sources, notes, and insights from both individual sources and notebook-wide contexts. - Post-processing includes ID-based deduplication, weighted prioritization, and strict token budgeting through
truncate_to_fit. - The builder returns a structured dictionary with
sources,notes,insights, andmetadatakeys. - Chat workflows format this structure into markdown sections before injecting it into Jinja2 templates via
Prompter.render. - The final RAG prompt combines the formatted context with chat history and is sent to the provisioned LangChain model.
Frequently Asked Questions
How does the ContextBuilder handle token limits when constructing RAG prompts?
The builder implements a truncate_to_fit method (lines 30-46 in context_builder.py) that automatically truncates the collected items to respect the max_tokens parameter specified in ContextConfig. This ensures the final prompt never exceeds the model's context window, protecting against token overflow errors while preserving the highest-priority items first.
Can I customize which items are prioritized in the RAG context?
Yes. The prioritization logic utilizes a priority_weights configuration (defined in lines 15-18 of context_builder.py) that assigns weights to different content types. The prioritize method uses these weights to reorder items before truncation, ensuring the most relevant content (such as high-confidence insights) appears at the beginning of the context.
Why does the ContextBuilder use a temporary event loop in the chat workflow?
The ContextBuilder is implemented as an asynchronous class to support non-blocking database queries and file I/O. Because LangGraph workflows (in source_chat.py) are synchronous, the workflow creates a temporary event loop using asyncio.new_event_loop() to execute the async builder.build() method, then returns the result to the synchronous graph execution.
What is the difference between _add_source_context and _add_notebook_context?
_add_source_context (lines 42-80) loads a specific Source object and its associated insights by ID, optimizing for single-document chat workflows. _add_notebook_context (lines 118-150) expands the scope to include all sources and notes associated with a given notebook ID, enabling broader conversational context that spans multiple documents.
Have a question about this repo?
These articles cover the highlights, but your codebase questions are specific. Give your agent direct access to the source. Share this with your agent to get started:
curl -s "https://instagit.com/install.md" Maintain an open-source project? Get it listed too →