Recover chunk ids mistyped from search results

Resolve a cited id that misses exactly to the nearest id the run
retrieved, above a 0.75 similarity cutoff.

CHANGELOG.md

@@ -8,6 +8,7 @@
 ### Fixed
 
 - A capability that reaches its request limit keeps its cite tool for two further model requests while its other tools are removed, so an exhausted run can still register citations.
+- A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff, recovering ids mistyped from search results.
 - Dotfiles are parsed as their actual format instead of a single unstructured text block. Docling ignores the extension of a name starting with a dot, so converters strip leading dots from the name they hand it.
 
 ### Documentation

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 from dataclasses import dataclass, field, replace
+from difflib import get_close_matches
 from pathlib import Path
 from typing import Any, cast
 
@@ -30,6 +31,27 @@ from haiku.rag.tools.search import build_binary_parts_from_results
 CITATION_GRACE_REQUESTS = 2
 """Requests the cite tool outlives this capability's other tools by."""
 
+CHUNK_ID_MATCH_CUTOFF = 0.75
+"""Similarity a cited chunk id needs to be treated as a corrupted known id.
+
+Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or
+duplicating a character or a whole group stays above 0.75, so the gap is wide.
+"""
+
+
+def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str:
+    """Recover a chunk id the model damaged while transcribing it.
+
+    Models copying opaque UUIDs drop and duplicate characters and whole
+    hyphen-separated groups. Candidates are limited to ids the run actually
+    retrieved, so a wrong match needs both a near miss and a same-run neighbour.
+    Ids that match nothing are returned unchanged for the caller to report.
+    """
+    if not known_ids or chunk_id in known_ids:
+        return chunk_id
+    match = get_close_matches(chunk_id, known_ids, n=1, cutoff=CHUNK_ID_MATCH_CUTOFF)
+    return match[0] if match else chunk_id
+
 
 def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
     if db_path is not None:
@@ -293,11 +315,11 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         state = cast(Any, self.state)
         for results in state.searches.values():
             all_results.extend(results)
-        citations = resolve_citations(chunk_ids, all_results)
+        known_ids = [result.chunk_id for result in all_results if result.chunk_id]
+        requested = [_nearest_known_id(cid.strip("[]"), known_ids) for cid in chunk_ids]
+        citations = resolve_citations(requested, all_results)
         resolved = {citation.chunk_id for citation in citations}
-        missing = [
-            cid.strip("[]") for cid in chunk_ids if cid.strip("[]") not in resolved
-        ]
+        missing = [cid for cid in requested if cid not in resolved]
 
         if missing:
             async with self.rag_lock:

tests/capabilities/test_capabilities.py

@@ -353,6 +353,44 @@ async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path):
     assert capability.state.citations == ["chunk-1"]
 
 
+@pytest.mark.asyncio
+async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
+    """Models mistype opaque UUIDs; near misses resolve to the retrieved id."""
+    true_id = "b8e25ea1-0bb3-48b1-8fea-2ac1f148bf7c"
+    unrelated = "9c2cd07e-5a3f-45a6-968d-cbd6f06ab57b"
+    capability = create_rag(db_path=temp_db_path, config=AppConfig())
+    capability.state = RAGState(
+        searches={
+            "q": [
+                SearchResult(
+                    content="evidence",
+                    score=1.0,
+                    chunk_id=true_id,
+                    document_id="doc-1",
+                    document_uri="test://document",
+                )
+            ]
+        }
+    )
+    client = AsyncMock()
+    client.get_chunk_by_id.return_value = None
+    capability.rag = client
+
+    dropped_char = "b8e25ea1-0bb3-48b1-8fea-2ac1f148bf7"
+    dropped_group = "0bb3-48b1-8fea-2ac1f148bf7c"
+
+    assert await capability._cite([dropped_char]) == "Registered 1 citation(s)."
+    assert await capability._cite([dropped_group]) == "Registered 1 citation(s)."
+    assert capability.state.citations == [true_id]
+
+    # An unrelated UUID is never attributed to a retrieved neighbour.
+    with pytest.raises(ModelRetry, match=unrelated):
+        await capability._cite([unrelated])
+
+    assert capability.state.citations == [true_id]
+    client.get_chunk_by_id.assert_awaited_once_with(unrelated)
+
+
 @pytest.mark.asyncio
 async def test_analysis_records_new_sandbox_search_results(temp_db_path):
     capability = create_analysis(db_path=temp_db_path, config=AppConfig())