fix(extraction): cap known-entity coreference context to stay under model input-token ceiling (RES-945)

src/beever_atlas/infra/config.py

@@ -396,6 +396,26 @@ def public_bot_base(self) -> str:
         le=60,
         description="Hard cap on messages per batch. Derived from observed bench data — successful batches had ≤65 msgs; failure cluster started at 89. Prevents output truncation at the source.",
     )
+    # RES-945 — the known-entity registry injected into every extraction batch
+    # (``get_all_canonical()``) grows unbounded with the corpus. On the RLP
+    # 36k-doc run it ballooned the serialised prompt past Gemini's 1,048,576-token
+    # input ceiling, so LiteLLM rejected every batch and extraction produced 0
+    # facts. Cap how many known entities are carried as cross-batch coreference
+    # context; the most-connected entities (most aliases) are kept first so
+    # coreference quality degrades gracefully. Set to 0 to disable the cap
+    # (legacy unbounded behaviour).
+    #
+    # Default 500 (~15-20k serialised tokens) is sized to fit BOTH windows: it
+    # leaves ample headroom under Gemini's 1,048,576-token input 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 head
+    # of the entity-frequency distribution (the few hundred most-referenced
+    # orgs/people/projects) carries almost all real cross-batch coreference.
+    extraction_known_entities_max: int = Field(
+        default=500,
+        ge=0,
+        description="Max canonical entities injected into an extraction batch prompt as coreference context. Prevents the entity registry from pushing the prompt past the model input-token ceiling on large corpora. Sized to fit both the Gemini 1M and self-hosted Qwen 64k windows. 0 disables the cap.",
+    )
     llm_outage_breaker_threshold: int = Field(
         default=3,
         ge=1,

src/beever_atlas/services/batch_processor.py

@@ -158,6 +158,42 @@ def _thread_aware_batches(messages: list[Any], batch_size: int) -> list[list[Any
     return batches
 
 
+def _cap_known_entities(
+    entities: list[dict[str, Any]], max_entities: int
+) -> list[dict[str, Any]]:
+    """Bound the known-entity registry injected into an extraction batch prompt.
+
+    RES-945: ``entity_registry.get_all_canonical()`` grows without limit as a
+    channel ingests more documents. The full list is JSON-serialised into every
+    extraction batch prompt as cross-batch coreference context, so on a large
+    corpus (RLP: 36k docs) the prompt overran Gemini's 1,048,576-token input
+    ceiling and LiteLLM rejected every batch → 0 facts extracted.
+
+    Keep the most-connected entities first — an entity with more aliases has
+    been referenced/merged more often, so it is the most valuable coreference
+    anchor — and drop the long tail. Ties break on name for deterministic
+    output (important for cache stability and tests). ``max_entities <= 0``
+    disables the cap (legacy unbounded behaviour).
+    """
+    if max_entities <= 0 or len(entities) <= max_entities:
+        return entities
+
+    def _prominence(entity: dict[str, Any]) -> tuple[int, str]:
+        aliases = entity.get("aliases") or []
+        alias_count = len(aliases) if isinstance(aliases, (list, tuple, set)) else 0
+        # Negative alias_count → most-aliased first; name ascending for ties.
+        return (-alias_count, str(entity.get("name") or ""))
+
+    capped = sorted(entities, key=_prominence)[:max_entities]
+    logger.info(
+        "BatchProcessor: capped known-entity coreference context %d→%d "
+        "(extraction_known_entities_max) to stay under the model input-token ceiling",
+        len(entities),
+        len(capped),
+    )
+    return capped
+
+
 def _summarize_exception(exc: Exception) -> str:
     """Create a compact, actionable error message for logs and sync status."""
     if isinstance(exc, ExceptionGroup):
@@ -516,7 +552,10 @@ async def process_messages(
             else None
         )
 
-        known_entities: list[dict[str, Any]] = await stores.entity_registry.get_all_canonical()
+        known_entities: list[dict[str, Any]] = _cap_known_entities(
+            await stores.entity_registry.get_all_canonical(),
+            settings.extraction_known_entities_max,
+        )
         cumulative_timings: dict[str, float] = {}
 
         # ── Bounded-concurrency batch execution ───────────────────────────────
@@ -2044,7 +2083,10 @@ async def _tagged(idx: int, b: list[Any]) -> tuple[int, Any]:
                             cumulative_timings.get(stage_key, 0.0) + duration
                         )
                     if entities_persisted:
-                        known_entities = await stores.entity_registry.get_all_canonical()
+                        known_entities = _cap_known_entities(
+                            await stores.entity_registry.get_all_canonical(),
+                            settings.extraction_known_entities_max,
+                        )
             processed_so_far += len(batch)
             await stores.mongodb.update_sync_progress(
                 job_id=sync_job_id,

tests/services/test_batch_processor_known_entities_cap.py

@@ -0,0 +1,60 @@
+"""RES-945 — the known-entity registry injected into every extraction batch must
+be bounded so a large corpus cannot push the serialised prompt past the model
+input-token ceiling (which made LiteLLM reject every batch → 0 facts on the RLP
+36k-doc run).
+
+Covers ``_cap_known_entities`` (keep most-connected entities first, deterministic,
+0 disables) and the config default ``extraction_known_entities_max``.
+"""
+
+from __future__ import annotations
+
+from beever_atlas.infra.config import Settings
+from beever_atlas.services.batch_processor import _cap_known_entities
+
+
+def _ent(name: str, alias_count: int) -> dict:
+    return {"name": name, "type": "PERSON", "aliases": [f"{name}-{i}" for i in range(alias_count)]}
+
+
+def test_under_limit_returns_input_unchanged() -> None:
+    ents = [_ent("a", 1), _ent("b", 2)]
+    assert _cap_known_entities(ents, 10) is ents
+
+
+def test_caps_to_max_keeping_most_aliased_first() -> None:
+    ents = [_ent("low", 0), _ent("high", 9), _ent("mid", 4)]
+    capped = _cap_known_entities(ents, 2)
+    assert [e["name"] for e in capped] == ["high", "mid"]  # dropped the 0-alias tail
+
+
+def test_tie_break_is_deterministic_by_name() -> None:
+    # All same alias count → must fall back to name-ascending, stable across runs.
+    ents = [_ent("charlie", 3), _ent("alpha", 3), _ent("bravo", 3)]
+    capped = _cap_known_entities(ents, 2)
+    assert [e["name"] for e in capped] == ["alpha", "bravo"]
+
+
+def test_zero_disables_cap() -> None:
+    ents = [_ent(str(i), i) for i in range(50)]
+    assert _cap_known_entities(ents, 0) is ents
+
+
+def test_negative_disables_cap() -> None:
+    ents = [_ent("x", 1)]
+    assert _cap_known_entities(ents, -1) is ents
+
+
+def test_missing_or_malformed_aliases_does_not_raise() -> None:
+    ents = [{"name": "n", "type": "ORG"}, {"name": "m", "aliases": None}, _ent("k", 5)]
+    capped = _cap_known_entities(ents, 2)
+    # The 5-alias entity is the most prominent and must survive; no KeyError/TypeError.
+    assert "k" in [e["name"] for e in capped]
+    assert len(capped) == 2
+
+
+def test_config_default_fits_both_model_windows() -> None:
+    # 500 canonical entities serialise to ~15-20k tokens — safe under Gemini's
+    # 1,048,576 window AND the self-hosted Qwen 64k window (RES-944 / F1).
+    s = Settings()
+    assert s.extraction_known_entities_max == 500