diff --git a/CHANGELOG.md b/CHANGELOG.md index 164d728d..c0876c50 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 4471d162..66e09efd 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/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: diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 5f8146bb..45dc36cf 100644 --- a/tests/capabilities/test_capabilities.py +++ b/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())