Merge pull request #508 from ggozad/fix/cite-partial-feedback

cite reports unresolvable chunk ids on partial success

CHANGELOG.md

@@ -5,6 +5,10 @@
 
 - `hotpotqa` evaluation dataset.
 
+### Changed
+
+- `cite` skill tool names unresolvable chunk ids in its response when at least one id resolves.
+
 ### Fixed
 
 - Document deletion no longer raises `ConfigMismatchError` on embedding config drift.

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -382,6 +382,15 @@ def create_skill_tools(
 
             if citations:
                 _register_citations(state, citations)
+                resolved_ids = {c.chunk_id for c in citations}
+                unresolved = [cid for cid in missing if cid not in resolved_ids]
+                if unresolved:
+                    return (
+                        f"Registered {len(citations)} citation(s); "
+                        f"ignored {len(unresolved)} unresolvable id(s): "
+                        f"{unresolved}. Copy chunk_ids verbatim from `search` "
+                        "results or items.jsonl / toc.json rows and cite again."
+                    )
                 return f"Registered {len(citations)} citation(s)."
 
             raise ModelRetry(

tests/skills/test_rag.py

@@ -355,6 +355,35 @@ class TestCiteTool:
         assert "verbatim" in message
         assert "nonexistent-chunk-id" in message
 
+    async def test_cite_reports_unresolved_ids_on_partial_success(
+        self, rag_db, rag_client
+    ):
+        """A cite call mixing valid and bogus ids registers the valid ones
+        and names the rejected ids so the model can re-cite the rest."""
+        from haiku.rag.skills.rag import RAGState, create_skill
+
+        skill = create_skill(db_path=rag_db)
+        search = _get_tool(skill, "search")
+        cite = _get_tool(skill, "cite")
+        state = RAGState()
+        ctx = _make_ctx(state, rag=rag_client)
+
+        await search(ctx, query="artificial intelligence")
+        valid_id = next(
+            sr.chunk_id
+            for results in state.searches.values()
+            for sr in results
+            if sr.chunk_id
+        )
+
+        result = await cite(ctx, chunk_ids=[valid_id, "6.43", "6.51.2"])
+        assert "Registered 1 citation(s)" in result
+        assert "6.43" in result
+        assert "6.51.2" in result
+        assert "verbatim" in result
+        assert len(state.citations) == 1
+        assert valid_id in state.citations
+
     async def test_cite_returns_message_when_chunk_ids_empty(self, rag_db):
         """An empty chunk_ids list is a no-op, not a retry trigger."""
         from haiku.rag.skills.rag import RAGState, create_skill