fix(extraction): cap known-entity coreference context to stay under model input-token ceiling (RES-945) - #336
Open
jhkchan wants to merge 1 commit into
Open
Conversation
…odel input-token ceiling (RES-945) Every extraction batch injects the full canonical entity registry (`entity_registry.get_all_canonical()`) as cross-batch coreference context, JSON-serialised into the prompt. That list grows without bound as a channel ingests more documents. On the RLP 36k-doc run it pushed the serialised prompt past Gemini's 1,048,576-token INPUT ceiling, so LiteLLM rejected every batch and extraction produced **0 facts** — even though the message batcher itself already bounds message content to `batch_max_prompt_tokens`. Add `_cap_known_entities`: keep the most-connected entities first (more aliases = referenced/merged more often = the most valuable coreference anchor), drop the long tail, deterministic name tie-break for cache stability. Applied at both registry fetch sites in the batch processor (initial load + post-persist reload). New config `extraction_known_entities_max` (default 2000; 0 disables → legacy unbounded behaviour). Adds tests/services/test_batch_processor_known_entities_cap.py: under-limit passthrough, most-aliased-kept, deterministic tie-break, 0/negative disables, malformed-aliases safe, config default. 7 passed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PMMM6KQXmzAEA42UxpyMUm
jhkchan
force-pushed
the
feature/res-945-extraction-batch-token-cap
branch
from
July 17, 2026 03:09
dfa3599 to
8a73c32
Compare
Member
Author
|
Independent adversarial review — APPROVE for the Gemini/1M fix; one P1 for the Qwen path. Verified: both
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem (RES-945)
Extraction produced 0 facts on the RLP 36k-doc corpus (and even on a scoped 300-doc channel). Root cause: every extraction batch injects the full canonical entity registry (
entity_registry.get_all_canonical()) as cross-batch coreference context, JSON-serialised into the prompt. That list is unbounded — it grows with every document a channel ingests. At scale it pushed the serialised prompt past Gemini's 1,048,576-token input ceiling, so LiteLLM rejected every batch → no facts.Note the message batcher (
adaptive_batcher.token_aware_batches) already bounds message content tobatch_max_prompt_tokens; this was the separate entity-registry payload injected on top of it, which had no cap.Fix
_cap_known_entities(entities, max)inbatch_processor.py:extraction_known_entities_max(default 500;0disables → legacy unbounded behaviour).Why 500 (fits both model windows)
500 canonical entities serialise to ~15-20k tokens. That leaves ample headroom under Gemini's 1M ceiling and under the self-hosted Qwen 64k window that the no-cloud path (RES-944 / F1) targets — so this cap does not reintroduce the overflow on Qwen (the ticket flags that F2 blocks F1 for exactly this reason). The head of the entity-frequency distribution (the few hundred most-referenced orgs/people/projects) carries almost all real cross-batch coreference.
Tests
tests/services/test_batch_processor_known_entities_cap.py: under-limit passthrough, most-aliased-kept, deterministic tie-break, 0/negative disables cap, malformed/missingaliasessafe, config default fits both windows. 7 passed locally.Scope note
This unblocks extraction at corpus scale by bounding the injected context; it does not change which facts the model extracts (coreference degrades gracefully rather than failing hard at 0 facts). A fully token-budget-aware cap derived from the resolved model window is a natural follow-up once F1 lands the per-endpoint model resolver.
Part of epic RES-943 (RLP full-corpus scale + no-cloud gaps).
🤖 Generated with Claude Code
https://claude.ai/code/session_01PMMM6KQXmzAEA42UxpyMUm