# Context Builder Implementation in Open Notebook's source_chat.py: Architecture and Usage

> Explore the context builder in open_notebook's source_chat.py. Learn how it asynchronously builds context, manages token limits, and prepares data for LangChain models. Discover its architecture and usage.

- Repository: [Luis Novo/open-notebook](https://github.com/lfnovo/open-notebook)
- Tags: architecture
- Published: 2026-06-23

---

**The context builder implementation in [`open_notebook/graphs/source_chat.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/graphs/source_chat.py) orchestrates asynchronous context assembly through the `call_model_with_source_context` node, which leverages the `ContextBuilder` class to gather source text, enforce token budgets, and render Jinja2 templates before provisioning LangChain models.**

Open Notebook's source-specific chat workflow demonstrates a modular approach to context management through the `source_chat` graph. This implementation leverages a reusable `ContextBuilder` class defined in [`open_notebook/utils/context_builder.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/context_builder.py) to asynchronously gather relevant data while respecting token limits, clearly separating data assembly from prompt rendering and model provisioning. Understanding this architecture reveals how the system maintains clean separation of concerns between graph orchestration, context building, and LLM invocation.

## Source Chat Graph Architecture

The `source_chat` graph in [`open_notebook/graphs/source_chat.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/graphs/source_chat.py) implements a LangGraph state machine with a single processing node: `call_model_with_source_context`. This node coordinates four distinct phases of the context building pipeline.

### The call_model_with_source_context Node

The `call_model_with_source_context` function executes a four-step workflow that transforms raw state into LLM-ready prompts:

**Context Gathering** – It instantiates a `ContextBuilder` with the target `source_id` and requests a fully assembled context payload. The builder retrieves the source's main text, optionally includes extracted insights, and respects configuration flags such as `include_notes` and `max_tokens`.

**Prompt Rendering** – The gathered context passes through `_format_source_context` and injects into the `source_chat/system` Jinja2 template. This produces a system message containing structured source data, insight summaries, and context indicators.

**Model Provisioning** – `provision_langchain_model` from [`open_notebook/ai/provision.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/ai/provision.py) initializes the appropriate LangChain model instance, accepting optional model overrides from the graph state and configuring provider-specific parameters.

**Response Cleaning** – The raw LLM output undergoes post-processing via `clean_thinking_content` (from [`open_notebook/utils/text_utils.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/text_utils.py)) to strip `<thinking>` tags and extract clean text content before returning the enriched state.

## Inside the ContextBuilder Implementation

The `ContextBuilder` class serves as the backbone of the context assembly process. Unlike graph-specific logic, this utility resides in [`open_notebook/utils/context_builder.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/context_builder.py) and supports multiple workflows including notebook-wide search and transformation graphs.

### Parameter Extraction and Source Loading

The builder initializes with explicit parameters: `source_id`, `notebook_id`, `include_insights`, `include_notes`, `max_tokens`, and an optional `ContextConfig` object. When `source_id` is present, the `_add_source_context` method fetches the `Source` record and calls `source.get_context()` to retrieve content (short or long form), then optionally `source.get_insights()` to attach extracted intelligence.

Each content piece becomes a `ContextItem` instance with assigned priority derived from `ContextConfig.priority_weights`. The builder supports notebook-level context through `_add_notebook_context`, which iterates across notebook sources and notes while delegating to specialized helpers.

### Deduplication and Prioritization

After collecting potential context items, the builder performs two critical filtering operations:

- **Deduplication** – Removes duplicate IDs to prevent context bloat from overlapping sources or notes
- **Priority Sorting** – Orders items by priority weight (highest first) to ensure critical information survives token budget cuts

### Token Budget Enforcement

When `max_tokens` is specified, the `truncate_to_fit` method iteratively removes lowest-priority `ContextItem` instances until the cumulative token count falls within budget. This greedy prioritization algorithm guarantees that the most relevant context remains while respecting model context windows.

The final `_format_response` method groups items by type (sources, insights, notes) and attaches metadata including item counts, token totals, and configuration flags, returning a dictionary that `source_chat` receives as `context_data`.

## Practical Implementation Examples

### Direct ContextBuilder Usage

Access the context builder directly for custom scripts or alternative graph implementations:

```python
from open_notebook.utils.context_builder import build_source_context

async def demo_source_context(source_id: str):
    # Build context including source text and insights, capped at 30,000 tokens

    ctx = await build_source_context(
        source_id=source_id,
        include_insights=True,
        max_tokens=30_000,
    )
    print("Formatted context snippet:")
    print(ctx["sources"][0][:500])          # First 500 characters of source text

    print("Insights count:", len(ctx["insights"]))

```

### Simulating source_chat Internals

Replicate the internal workflow of `source_chat` for debugging or customization:

```python
from open_notebook.graphs.source_chat import call_model_with_source_context
from open_notebook.utils.context_builder import ContextBuilder
from open_notebook.ai.provision import provision_langchain_model
from langchain_core.messages import SystemMessage

async def run_source_chat(source_id: str, user_messages):
    # 1. Build context asynchronously

    builder = ContextBuilder(
        source_id=source_id,
        include_insights=True,
        include_notes=False,
        max_tokens=50_000,
    )
    context_data = await builder.build()

    # 2. Render system prompt using Prompter

    prompt_data = {
        "source": context_data["sources"][0] if context_data["sources"] else None,
        "insights": context_data["insights"],
        "context": context_data,  # Formatted via _format_source_context internally

        "context_indicators": {
            "sources": [source_id],
            "insights": [i["id"] for i in context_data["insights"]],
            "notes": [],
        },
    }
    
    # 3. Provision model and invoke

    # Note: Actual implementation uses Prompter class for template rendering

    payload = [SystemMessage(content=str(prompt_data))] + user_messages
    model = await provision_langchain_model(
        str(payload),
        model_id=None,
        task="chat",
        max_tokens=8192,
    )
    return model.invoke(payload)

```

### Extending with Custom Parameters

The `ContextBuilder` accepts arbitrary `custom_*` parameters for future extensibility:

```python
builder = ContextBuilder(
    source_id="source:abc123",
    include_insights=True,
    custom_highlight="important",   # Ignored by default implementation,

                                   # but subclasses can override _process_custom_params

)
await builder.build()

```

## Summary

- **Graph Orchestration**: The `source_chat` graph uses `call_model_with_source_context` in [`open_notebook/graphs/source_chat.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/graphs/source_chat.py) to coordinate context building, prompt rendering, and model invocation in a single LangGraph node.

- **Async Context Assembly**: `ContextBuilder` in [`open_notebook/utils/context_builder.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/context_builder.py) asynchronously loads source content, deduplicates entries, applies priority weights, and enforces token budgets through `truncate_to_fit`.

- **Modular Design**: The architecture cleanly separates **context gathering** (ContextBuilder), **prompt preparation** (Jinja2 templates), and **model provisioning** (`provision_langchain_model`), enabling independent modification of each component.

- **Token Management**: The builder implements intelligent truncation that preserves high-priority `ContextItem` instances while fitting content within specified `max_tokens` limits.

## Frequently Asked Questions

### How does ContextBuilder handle async database calls without blocking the graph?

The `ContextBuilder` performs all data loading operations asynchronously. When `source_chat` invokes the builder, it either runs in a fresh event loop or uses a thread-pooled loop if already inside an async context. This guarantees that slow database queries for source content or insights do not block the LangGraph state machine execution.

### What determines the priority of ContextItem instances during truncation?

Priority derives from the `ContextConfig.priority_weights` mapping, which assigns numeric weights to different content types (sources, insights, notes). During `truncate_to_fit`, the builder sorts items by descending priority and removes lowest-priority items first until the token budget is satisfied. Higher priority items survive budget cuts, ensuring critical context reaches the LLM.

### Can I use ContextBuilder for notebook-wide context instead of single sources?

Yes. Instantiate `ContextBuilder` with `notebook_id` instead of (or alongside) `source_id` to trigger `_add_notebook_context`, which iterates across all notebook sources and notes. This enables the same token-budgeted, prioritized context assembly for multi-source chat workflows or notebook transformation pipelines.

### Where is the system prompt template defined for source_chat?

The system prompt template resides at `open_notebook/prompt_templates/source_chat/system.jinja`. The `Prompter` class renders this Jinja2 template using the formatted context data, injecting source text, insight summaries, and context indicators into the LLM's system message before provisioning occurs.