Per-chunk entity attribution in koi_entity_chunk_links

[DarrenZal]

Summary

The new link_entities_to_chunks write-path (eaf803f) fans every doc-level entity across every chunk of the document. For a typical Notion page (18 entities × 19 chunks = 342 link rows), every chunk appears to contain every entity. Historical koi_entity_chunk_links data (614K rows from before Feb 10) had genuine per-chunk attribution — we should return to that fidelity.

Proposed work

1. Evaluate whether per-chunk precision is worth the cost vs. current doc-level fanout
2. Implement either:
   • A. Per-chunk LLM extraction (highest fidelity, N× LLM cost, N ≈ 20 for Notion)
   • B. Offset-match: after doc-level extraction, substring-search each entity in each chunk's text and only link where it actually appears (~free, misses synonyms/pronouns)
3. Populate char_offset (currently NULL) for chunk-level context in UI
4. Adjust ranking in koi-query-api.ts to use per-chunk entity density where available

Option B is likely the right MVP. It compounds well with improved entity extraction (#19).

Dependencies

Blocked by #19 (entity extraction quality) — not strictly, but the combined improvement is much more valuable than either alone.

References

• Current write-path: eaf803f — link_entities_to_chunks in src/core/koi_event_bridge_v2.py
• Backfill script: scripts/backfill_entity_chunk_links.py
• Search consumer: koi-query-api.ts line ~1163 (entity_memories CTE)

Task key: other-2026-04-14-per-chunk-entity-attribution
Vault task: Tasks/2026-04-14-per-chunk-entity-attribution.md

[DarrenZal]

Additional context for investigators

Option B (offset-match) — concrete algorithm sketch

After trigger_kg_extraction writes the doc-level entities to koi_kg_extractions.entities:

```python
async def link_entities_to_chunks_precise(conn, document_rid: str) -> int:
extraction = await conn.fetchrow(
"SELECT entities FROM koi_kg_extractions "
"WHERE memory_rid=$1 AND extraction_type='passA' "
"ORDER BY created_at DESC LIMIT 1",
document_rid,
)
if not extraction or not extraction["entities"]:
return 0
entities = extraction["entities"]
if isinstance(entities, str):
entities = json.loads(entities)

chunks = await conn.fetch(
    "SELECT id, rid, content->>'text' AS text FROM koi_memories "
    "WHERE rid LIKE \$1 AND superseded_at IS NULL",
    f"{document_rid}#chunk%",
)
if not chunks:
    return 0

link_rows = []
for chunk in chunks:
    text_lower = (chunk["text"] or "").lower()
    for ent in entities:
        name = (ent.get("name") or "").strip()
        if not name:
            continue
        name_lower = name.lower()

        # Find first occurrence; skip chunks where the entity does not appear
        offset = text_lower.find(name_lower)
        if offset < 0:
            # TODO: alias/synonym fallback (e.g. resolve via entity_registry)
            continue

        link_rows.append((
            name, name_lower, ent.get("type", "Unknown"),
            ent.get("rid"), str(chunk["id"]),
            int(chunk["rid"].split("#chunk")[1]) if "#chunk" in chunk["rid"] else None,
            chunk["rid"],
            float(ent.get("confidence", 0.8)),
            offset,          # char_offset — new!
        ))

if not link_rows:
    return 0
await conn.executemany(
    "INSERT INTO koi_entity_chunk_links "
    "(entity_name, entity_name_lower, entity_type, entity_uri, "
    " chunk_rid, chunk_index, document_rid, confidence, char_offset) "
    "VALUES (\$1, \$2, \$3, \$4, \$5, \$6, \$7, \$8, \$9)",
    link_rows,
)
return len(link_rows)

```

Why offset-match misses some cases (and why that's OK)

• Pronoun/anaphora: "Becca shared the framework. She mentioned X." The entity "Becca" only appears in chunk 0; chunk 1 references her as "She". Doc-level extraction correctly got "Becca", but offset-match won't link chunk 1.
• Synonym/alias: Extractor extracted "Claims Engine"; chunk text says "claims system". Offset-match misses. Fix: join with entity_registry.aliases array and match any alias.
• Normalization: "SeaTrees" extracted, chunk says "Sea Trees" (with space). Fix: apply the same lightweight normalization used by entity_registry.canonical_name comparison.

All three failure modes reduce recall relative to the current coarse approach. The quality-of-ranking gain (chunks that actually contain the term rank higher) should dominate the recall loss for typical queries, but it's worth A/B-ing.

A/B methodology

1. Pick a query set (e.g. 20 queries used in koi-query-api evals)
2. For each query, capture top-20 chunks under coarse vs precise
3. Measure:
   • Precision @ 10: manual judgment of whether each chunk really discusses the query
   • Coverage (MRR): is there any chunk of a relevant document in top-k
   • Chunk-level precision: of chunks containing the query's main entity, what fraction actually reference it in text

Interaction with #19

Entity-extraction improvements in #19 compound here:

• More entities extracted → more links created
• Better canonical naming → offset match works more often (less variation in how the term appears)
• Aliased extraction → offset fallback via entity_registry.aliases needs less work

Recommended order: land #19's alias normalization first, then ship offset-match here with alias-aware fallback.

Table schema note

koi_entity_chunk_links.char_offset already exists in the schema but is currently NULL for all rows. Populating it is essentially free once we're doing offset-matching.