# Implementing Fine-Grained Context Control for AI Chats in Open Notebook

> Learn to implement fine-grained context control for AI chats in Open Notebook. Discover how the ContextBuilder utility assembles weighted context based on frontend-defined inclusion rules.

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

---

**Open Notebook exposes a decoupled `ContextBuilder` utility that assembles weighted, token-aware context from specific sources, notes, and insights based on fine-grained inclusion rules defined by the frontend.**

Open Notebook is a privacy-first research assistant that requires precise control over what data reaches the language model. The implementation of **fine-grained context control for AI chats** centers on a pure Python builder class in [`open_notebook/utils/context_builder.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/context_builder.py) that transforms user selections into optimized LLM payloads while respecting token budgets and content priorities.

## The ContextBuilder Core Architecture

The subsystem revolves around the **`ContextBuilder`** class, which operates independently of FastAPI-specific code. This decoupling allows the same logic to run in unit tests, background jobs, or external services.

The builder accepts a **`ContextConfig`** dataclass that defines inclusion maps for sources and notes, boolean flags for insights, and an optional token ceiling. When initialized, the builder stores these parameters and prepares an internal list of `ContextItem` objects.

Key methods driving the assembly process include:

- **`__init__`** (lines 65-78) – Validates parameters and creates a default `ContextConfig` when none is supplied.
- **`build()`** (lines 105-138) – Orchestrates the pipeline: adds content, processes custom hooks, removes duplicates, sorts by priority, truncates to fit token budgets, and returns a JSON-compatible dictionary.
- **`_add_source_context()`** (lines 42-82) – Fetches `Source` records and determines inclusion levels ("insights", "full content", or exclusion).
- **`_add_notebook_context()`** (lines 110-147) – Aggregates notebook-level sources and notes by iterating over `ContextConfig` mappings.
- **`_add_note_context()`** (lines 154-176) – Pulls specific `Note` content based on inclusion rules.

## ContextConfig and ContextItem Data Models

Two dataclasses define the schema for fine-grained control.

**`ContextItem`** (lines 21-30) wraps each piece of context with:

- `id` and `type` (source, note, or insight)
- `content` (raw text or structured data)
- `priority` weight for sorting
- Token count calculated on-the-fly via `token_utils.token_count`

**`ContextConfig`** (lines 38-48) holds the user-defined rules:

- `sources`: Dict mapping source IDs to inclusion levels ("insights", "full content", "short")
- `notes`: Dict mapping note IDs to inclusion levels
- `include_insights`: Global boolean toggle
- `include_notes`: Global boolean toggle
- `max_tokens`: Optional ceiling to prevent context overflow
- `priority_weights`: Customizable weights (defaults: source=100, insight=75, note=50)

## The Context Assembly Pipeline

The `build()` method executes a deterministic pipeline to ensure consistent context construction.

First, the builder invokes `_add_notebook_context()` or source-specific helpers to populate the internal item list. Each addition calculates token counts immediately using the utility function.

Next, **`remove_duplicates()`** (lines 351-362) eliminates redundant entries that might appear across overlapping sources and notes. Then **`prioritize()`** (lines 315-318) sorts the list by the assigned priority weights in descending order.

Finally, **`truncate_to_fit()`** (lines 321-350) enforces the `max_tokens` budget by iterating from the lowest priority upward and dropping items until the total token count falls within limits. This guarantees that the most important context survives when budgets are tight.

## API Endpoints for Frontend Integration

Two FastAPI routers expose the builder to the React frontend, enabling real-time **fine-grained context control** through HTTP requests.

**`/chat/context`** in [`api/routers/chat.py`](https://github.com/lfnovo/open-notebook/blob/main/api/routers/chat.py) (lines 425-531) serves as the primary entry point. The frontend sends a payload containing `notebook_id` and a `context_config` dictionary that mirrors the `ContextConfig` schema. The endpoint delegates to `build_notebook_context()` or `build_mixed_context()` convenience wrappers and returns the flattened context as JSON.

**`/notebooks/{notebook_id}/context`** in [`api/routers/context.py`](https://github.com/lfnovo/open-notebook/blob/main/api/routers/context.py) (lines 12-115) provides a low-level debugging endpoint that returns raw context for a specific notebook without initiating a chat turn.

Both endpoints utilize the same builder logic, ensuring the UI can let users toggle individual sources, select inclusion levels per item, and set token limits while the backend reliably enforces these constraints.

## Token Budget and Priority Management

The system implements intelligent cost control through explicit priority weights and truncation logic.

When a `max_tokens` value is provided, the builder calculates running totals after deduplication. The **`truncate_to_fit()`** method removes lowest-priority items first, ensuring high-value sources (default weight 100) survive over notes (weight 50) when budgets are constrained.

This approach prevents model context window overflows while maintaining the semantic relevance of the retained content. The frontend can expose sliders or numeric inputs for token budgets, confident that the backend will respect the ceiling.

## Practical Implementation Examples

### Frontend TypeScript Integration

```typescript
// UI gathers user selections
const payload = {
  notebook_id: selectedNotebook.id,
  context_config: {
    sources: {
      "source:123": "insights",          // include only insights
      "source:456": "full content",      // include full text
    },
    notes: {
      "note:789": "full content",        // include note text
    },
    include_insights: true,
    include_notes: true,
    max_tokens: 3000,
  },
};

await fetch(`${API_URL}/chat/context`, {
  method: "POST",
  headers: { "Content-Type": "application/json" },
  body: JSON.stringify(payload),
});

```

This request triggers the flow through [`api/routers/chat.py`](https://github.com/lfnovo/open-notebook/blob/main/api/routers/chat.py) → `ContextBuilder` → assembled JSON respecting the supplied maps.

### Direct Python Usage for Background Jobs

```python
from open_notebook.utils.context_builder import build_notebook_context, ContextConfig

cfg = ContextConfig(
    sources={"source:abc": "short"},
    notes={"note:def": "full content"},
    include_insights=False,
    max_tokens=2000,
)

context = await build_notebook_context(
    notebook_id="notebook:xyz", 
    context_config=cfg
)

# Returns dict with sources, notes, insights filtered and truncated

```

This pattern mirrors the API internals, allowing the same builder to run in Celery tasks or scheduled jobs outside the HTTP request cycle.

### Extending with Custom Inclusion Levels

```python
class MetadataContextBuilder(ContextBuilder):
    async def _add_source_context(self, source_id: str, inclusion_level: str):
        if inclusion_level == "metadata":
            source = await Source.get(source_id)
            meta = {"title": source.title, "url": source.url}
            self.add_item(
                ContextItem(
                    id=source.id,
                    type="source",
                    content=meta,
                    priority=self.context_config.priority_weights["source"],
                )
            )
            return
        await super()._add_source_context(source_id, inclusion_level)

```

This subclass demonstrates the extensible hook points designed for future growth, such as adding "metadata only" inclusion levels without modifying core builder logic.

## Summary

- **Fine-grained context control** in Open Notebook relies on the `ContextBuilder` class in [`open_notebook/utils/context_builder.py`](https://github.com/lfnovo/open-notebook/blob/main/open_notebook/utils/context_builder.py) to assemble LLM payloads.
- **`ContextConfig`** and **`ContextItem`** dataclasses define inclusion rules and weighted content pieces with automatic token counting.
- The **`build()`** method orchestrates deduplication, prioritization, and token-budget truncation to prevent context overflow.
- **`/chat/context`** and **`/notebooks/{id}/context`** endpoints expose this functionality to frontend applications.
- Default priority weights (sources: 100, insights: 75, notes: 50) ensure important context survives when **`truncate_to_fit()`** enforces `max_tokens` limits.
- The decoupled architecture allows reuse in background jobs, unit tests, and custom workflows through convenience wrappers like `build_notebook_context()`.

## Frequently Asked Questions

### How does Open Notebook prevent token limit errors when building context?

The `ContextBuilder` calculates token counts for each `ContextItem` using `token_utils.token_count`. When a `max_tokens` value is provided in `ContextConfig`, the **`truncate_to_fit()`** method (lines 321-350) removes lowest-priority items first until the total falls within budget. This ensures high-priority sources survive while preventing context window overflows.

### Can I include only specific insights from a source without the full text?

Yes. Set the source's inclusion level to `"insights"` in the `ContextConfig.sources` dictionary. The **`_add_source_context()`** method checks this level and calls `Source.get_insights()` rather than retrieving the full content, adding only the generated insights to the context list.

### What happens if the same content appears in both a source and a note?

The **`remove_duplicates()`** method (lines 351-362) eliminates redundant entries by comparing item IDs during the build process. This ensures that overlapping content between sources and notes only consumes tokens once, optimizing the available context budget.

### Is the ContextBuilder usable outside of HTTP API calls?

Yes. The class contains no FastAPI-specific dependencies, making it suitable for background tasks like podcast generation or scheduled reports. Use the convenience wrappers `build_notebook_context()`, `build_source_context()`, or `build_mixed_context()` (lines 422-492) to invoke the builder directly from Python scripts or Celery workers.