Skip to content

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
mainfrom
feature/res-945-extraction-batch-token-cap
Open

fix(extraction): cap known-entity coreference context to stay under model input-token ceiling (RES-945)#336
jhkchan wants to merge 1 commit into
mainfrom
feature/res-945-extraction-batch-token-cap

Conversation

@jhkchan

@jhkchan jhkchan commented Jul 17, 2026

Copy link
Copy Markdown
Member

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 to batch_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) in batch_processor.py:

  • Keeps the most-connected entities first (alias count = how often an entity was referenced/merged = its coreference value), drops the long tail.
  • Deterministic name tie-break → stable prompt (cache-friendly, testable).
  • Applied at both registry fetch sites: initial load and the post-persist reload.
  • New config extraction_known_entities_max (default 500; 0 disables → 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/missing aliases safe, 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

…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

jhkchan commented Jul 17, 2026

Copy link
Copy Markdown
Member Author

Independent adversarial review — APPROVE for the Gemini/1M fix; one P1 for the Qwen path.

Verified: both get_all_canonical() fetch sites are capped (initial load + post-persist reload), settings.extraction_known_entities_max is in scope at both, no other call site bypasses it, batch_pipeline consumes the capped snapshot (doesn't re-fetch). Helper is deterministic, off-by-one correct (len==max keeps all), <=0 disables (identity), robust to missing/None/non-list aliases. Message content (12k) and thread context (200 chars) are already bounded, so the entity registry was indeed the sole unbounded contributor — the fix structurally closes the 1M overflow.

  • [P1] The "500 fits Qwen 64k" claim is overstated. The cap keeps the most-aliased entities first, and the retained head is exactly where aliases concentrate. Measured json.dumps of 500 entities: ~16k tokens @3 aliases, but ~58k @20 aliases and ~106k @40 — which, added to the 12k message budget + schema, exceeds Qwen's 64k window. This does not threaten the Gemini fix (safe under 1M in every case), but the stated rationale that F2 unblocks F1 "for exactly this reason" isn't satisfied in the high-alias regime. Before enabling no-cloud I'll convert this to a serialized-token-budget cap (keep most-prominent until a token budget is hit) rather than a raw count.
  • [P2] test_config_default_fits_both_model_windows only asserts == 500; it doesn't serialize/measure — rename or make it real.
  • [P2] Prominence = alias count ≠ cross-batch recurrence; a frequently-mentioned single-spelling entity (0 aliases) can be dropped before a twice-spelled one. Follow-up: rank by reference count if available.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant