feat(retriever): add token-aware context truncation to cap documents passed to LLM

[GovindhKishore]

Summary

Adds token-aware context truncation to HybridRetriever to cap the number of documents and tokens passed to the LLM. Fixes the unbounded document passing problem described in #138.

Motivation

create_stuff_documents_chain in rag_chain.py stuffs ALL retrieved documents into the LLM context with zero token awareness. With two databases installed (Reactome + UniProt), this results in 60-100 documents per query reaching the LLM, causing:

• Verbose, unfocused answers from too much noisy context
• High token cost depending on the number of tokens in a document and the current OpenAI API cost.
• "Lost in the middle" degradation - LLMs ignore middle documents in long contexts.

Changes

New - src/util/context_truncator.py

• truncate_to_token_limit() - truncates a ranked document list to fit within a token budget and document count limit
• Uses tiktoken for exact GPT-4o token counting - already installed via langchain-openai, zero new dependencies
• Processes documents in order (best first) and stops when either limit is hit - least relevant documents always removed first
• Guarantees at least one document is always returned even if it exceeds the token budget, preventing empty context edge case

Modified - src/retrievers/csv_chroma.py

• retrieve_documents() - returns truncate_to_token_limit(subdirectory_docs) instead of raw subdirectory_docs
• aretrieve_documents() - same change applied to async retrieval path
• Added import for truncate_to_token_limit from util.context_truncator

Modified - config_default.yml

• Added retriever.context_truncation block with max_docs: 15 and max_tokens: 12000
• Config wiring to HybridRetriever can be added later - current defaults are hardcoded

Why These Default Values

max_docs: 15
  15 high quality ranked docs provides sufficient context for any question
  beyond this quality degrades due to lost-in-the-middle effect

max_tokens: 12000
  leaves room for system prompt + chat history + answer
  within GPT-4o's 128k context window

Expected Impact

Limits context passed to the LLM to a maximum of 15 documents and 12,000 tokens per query, down from an unbounded 60-100 documents. This reduces token usage significantly and avoids the "lost in the middle" quality degradation that occurs with excessively long contexts.

Note: Exact token counts per document depend on installed database versions and chunk sizes. A follow-up evaluation with real embeddings will quantify the precise reduction.

Note: GPT-4o supports a 128k context window but the 12,000 token limit is intentional - it reserves space for system prompt, chat history, and model output, while avoiding the "lost in the middle" quality degradation that occurs with excessively long contexts.

Interaction With Other PRs

Truncation runs at the end of retrieve_documents() and aretrieve_documents() - after WRR ranking and after FlashRank reranking (PR #116). This means truncation always removes the least relevant documents from last, preserving quality ordering established by both ranking stages.

BM25 + Vector retrieval
      ->
weighted_reciprocal_rank
      ->
FlashRank reranking (PR #116)
      ->
truncate_to_token_limit()   
      ->
create_stuff_documents_chain

Related

• Closes fix(retriever): HybridRetriever passes unlimited documents to LLM with no token budget #138
• Complements PR Add FlashRank reranker to HybridRetriever to improve retrieval quality #116 (FlashRank reranking)
• Complements PR Improve response precision by refining system prompts for Reactome and UniProt retrievers #121 (prompt precision)

[adamjohnwright]

Thank you — the token-aware framing is right, and better than the document count that landed.

main now caps fused documents per collection (#167), which took the prompt from ~222 documents / ~32k tokens to 40 / ~6.3k. But your point stands: measured over six real questions, those 40 documents range 5,796–9,531 tokens, a 1.6x spread. A document count is a loose proxy for the thing that actually costs money and context.

Not adopting it yet, for one reason: 9,531 tokens is 7% of the 128k window, so a second limit would tighten a bound that is not currently binding. Recorded with the measurements in #172 so it is picked up when it matters.

The more valuable half of your PR is putting the limits in config.yml rather than a constant — that is the right shape, and it matches a need we have for different budgets per surface (chat vs. search-results integration). It is blocked on wiring AgentGraph to the YAML config, the same plumbing #151 proposes, so those want doing together.

Noting also that tiktoken is already in the lockfile via langchain-openai, so this costs no new dependency when it is taken up.

Closing as recorded rather than rejected — the analysis is in #172 with your name on it.