How to chunk documents for long-running research agents

Long-running research agents fail silently when context is fed in arbitrary fragments. To build reliable systems, you need a disciplined method to chunk documents research agents can actually use: split on structure, respect token limits, and track provenance so the agent can revisit sources.

Step 1: Analyze document structure before splitting

Naive character splitting destroys section boundaries and mixes unrelated topics. Parse the source format and keep its logical hierarchy. For Markdown, split on headings; for HTML, use the DOM; for PDF, extract text per page but preserve any embedded headings.

A minimal Markdown splitter:

import re

def split_markdown_by_heading(text: str, level: int = 2) -> list[dict]:
    """Return sections with heading and body."""
    pattern = re.compile(rf'^(#{{{level}}})\s+(.*)$', re.MULTILINE)
    parts = pattern.split(text)
    # parts: ['', '#', 'Heading', 'body', '#', 'Heading2', 'body2', ...]
    sections = []
    for i in range(1, len(parts), 3):
        heading = parts[i+1].strip()
        body = parts[i+2].strip()
        sections.append({"heading": heading, "body": body})
    return sections

doc = open("report.md").read()
sections = split_markdown_by_heading(doc, level=2)
print(f"Got {len(sections)} sections")

If you receive HTML, use bs4.BeautifulSoup with soup.find_all(['h2','p']). The point is to never flatten a document into one string before chunking.

Step 2: Choose a token budget per chunk

Pick a chunk size that fits your retrieval strategy and model limits. For most research agents using 8k–32k context windows, 384–768 tokens per chunk with 64–128 token overlap works well. Overlap prevents sentences from being cut at boundaries.

Count tokens with the same tokenizer your embedding or completion model uses. OpenAI models use cl100k_base:

import tiktoken

enc = tiktoken.get_encoding("cl100k_base")

def token_len(text: str) -> int:
    return len(enc.encode(text))

MAX_TOKENS = 512
OVERLAP = 64

If you later switch to a different model family, change the encoding accordingly. Mismatched token counts cause silent truncation.

Step 3: Implement structure-aware chunking

Split each section by tokens only if it exceeds the budget. Keep headings prepended to every chunk so the agent knows context.

def chunk_section(heading: str, body: str, max_tokens: int, overlap: int) -> list[str]:
    chunks = []
    tokens = enc.encode(body)
    if len(tokens) <= max_tokens:
        chunks.append(f"{heading}\n{body}")
        return chunks
    
    start = 0
    while start < len(tokens):
        end = min(start + max_tokens, len(tokens))
        chunk_tokens = tokens[start:end]
        chunk_text = enc.decode(chunk_tokens)
        chunks.append(f"{heading}\n{chunk_text}")
        if end == len(tokens):
            break
        start += max_tokens - overlap
    return chunks

all_chunks = []
for sec in sections:
    all_chunks.extend(chunk_section(sec["heading"], sec["body"], MAX_TOKENS, OVERLAP))
print(f"Total chunks: {len(all_chunks)}")

This two-stage approach lets you chunk documents research agents will retrieve without losing the section label that anchors meaning.

Step 4: Attach metadata for retrieval and deduplication

A chunk without provenance is useless for a research agent that must cite sources. Store at least: chunk id, source document id, section heading, token count, and the raw text.

import hashlib, json

def make_chunk_records(src_id: str, chunks: list[str]) -> list[dict]:
    records = []
    for i, text in enumerate(chunks):
        cid = hashlib.sha1(f"{src_id}:{i}:{text[:32]}".encode()).hexdigest()[:12]
        records.append({
            "chunk_id": cid,
            "doc_id": src_id,
            "token_count": token_len(text),
            "text": text
        })
    return records

records = make_chunk_records("report-2024", all_chunks)

Add embedding vectors later if you do semantic search. Keep the metadata flat to simplify filtering.

Step 5: Store chunks in a lightweight index

For a single-machine agent, JSONL is enough to start. You can upgrade to SQLite or a vector DB without changing the chunking logic.

with open("chunks.jsonl", "w") as f:
    for r in records:
        f.write(json.dumps(r) + "\n")

To retrieve, load and filter:

def retrieve(query_tokens: set, top_k: int = 5) -> list[dict]:
    hits = []
    with open("chunks.jsonl") as f:
        for line in f:
            r = json.loads(line)
            score = sum(1 for t in query_tokens if t in r["text"])
            hits.append((score, r))
    hits.sort(key=lambda x: x[0], reverse=True)
    return [r for _, r in hits[:top_k]]

In production, replace the loop with FAISS or pgvector. The chunk schema stays identical.

Step 6: Feed chunks to the research agent with context rotation

A long-running agent cannot hold every chunk in context. Implement a sliding window: keep the most recent K chunks plus the current task state. Summarize evicted chunks into a running brief.

from openai import OpenAI

client = OpenAI()  # or any OpenAI-compatible endpoint

def agent_step(question: str, retrieved: list[dict], history: str) -> str:
    context = "\n---\n".join([r["text"] for r in retrieved])
    prompt = f"History summary:\n{history}\n\nContext:\n{context}\n\nQuestion: {question}"
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=256
    )
    return resp.choices[0].message.content

# Loop over research steps
history = ""
for step_q in ["What were the Q2 risks?", "How did they compare to Q1?"]:
    retrieved = retrieve(set(step_q.lower().split()), top_k=3)
    answer = agent_step(step_q, retrieved, history)
    history += f"\nQ: {step_q}\nA: {answer}"

This pattern lets you chunk documents research agents consume incrementally, rather than dumping a whole PDF into the prompt.

Step 7: Scale LLM calls with fallback and metering

A research agent running for hours will hit rate limits. If you self-host retries, you waste engineering time. Route completion calls through an OpenAI-compatible gateway that automates failover. n4n.ai exposes one endpoint covering 240+ models and automatically falls back when a provider is degraded, so the agent keeps working instead of throwing 429s.

# Point the client at the gateway; same API shape
client = OpenAI(
    base_url="https://api.n4n.ai/v1",
    api_key=os.environ["N4N_KEY"]
)
# provider cache-control hints are forwarded automatically

Per-token metering on such a gateway lets you attribute cost to each research task without building your own accounting middleware.

How to verify success

Run these checks before shipping:

1. Token audit – Assert every chunk’s token_count ≤ MAX_TOKENS. Print max observed.
2. Provenance – Pick a random chunk_id, confirm doc_id and heading exist in source.
3. Agent completion – Execute the agent on a 200-page document with 20 sequential questions. It should finish without context_length_exceeded errors.
4. Retrieval recall – Build a small gold set: for 10 questions, mark the source section. Verify the correct chunk appears in retrieve(top_k=5) at least 8 times.
5. Fallback test – Temporarily block the primary model; confirm the gateway serves a fallback and the agent output remains structurally valid.

If all five pass, your pipeline to chunk documents research agents rely on is sound. Tune MAX_TOKENS and OVERLAP based on observed retrieval recall, not on intuition.