Merge pull request #305 from ggozad/fix/always-cite-in-ask

Enforce the QA agent to always return citations when answer is based on search results

CHANGELOG.md

@@ -12,6 +12,10 @@
 - **QA search cap**: Replace dead `max_iterations`/`max_concurrency` config with `max_searches` (default: 3). The QA agent now enforces a per-run search limit, reducing average response time from ~30s to ~15s while maintaining accuracy. The limit resets per agent run so toolsets can be safely reused.
 - **Default search limit**: Increased from 5 to 10 results per search query for better coverage.
 
+### Fixed
+
+- **QA citations**: Strengthened prompt to clarify chunk ID format (complete IDs without brackets). `resolve_citations` now strips `[]` from IDs, handling models that copy brackets from search result formatting.
+
 ## [0.33.1] - 2026-03-06
 
 ### Changed

haiku_rag_slim/haiku/rag/agents/qa/prompts.py

@@ -26,7 +26,7 @@ Each result includes:
 - Type: content type like paragraph, table, code, list_item (when available)
 - Content: the actual text
 
-In your response, include the chunk IDs you used in cited_chunks.
+IMPORTANT: You MUST include in cited_chunks the COMPLETE IDs of every chunk you reference. Copy the full ID string without brackets — e.g. "5ae52166-5329-42e9-b6a5-756fc0cb7200" not "[5ae52166]" or "5ae52166". Never truncate IDs. Never leave cited_chunks empty if you found relevant content.
 
 Guidelines:
 - Base answers strictly on retrieved content - do not use external knowledge

haiku_rag_slim/haiku/rag/agents/research/models.py

@@ -93,7 +93,8 @@ def resolve_citations(
     by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
 
     citations = []
-    for chunk_id in cited_chunk_ids:
+    for raw_id in cited_chunk_ids:
+        chunk_id = raw_id.strip("[]")
         r = by_id.get(chunk_id)
         if not r:
             continue

tests/agents/research/test_models.py

@@ -1,4 +1,5 @@
-from haiku.rag.agents.research.models import Citation, SearchAnswer
+from haiku.rag.agents.research.models import Citation, SearchAnswer, resolve_citations
+from haiku.rag.store.models import SearchResult
 
 
 class TestCitation:
@@ -120,3 +121,39 @@ class TestSearchAnswerPrimarySource:
             citations=[],
         )
         assert answer.primary_source is None
+
+
+class TestResolveCitations:
+    """Tests for resolve_citations function."""
+
+    def _make_result(self, chunk_id: str) -> SearchResult:
+        return SearchResult(
+            content="test content",
+            score=1.0,
+            chunk_id=chunk_id,
+            document_id="doc-1",
+            document_uri="test.md",
+            document_title="Test Doc",
+        )
+
+    def test_resolves_exact_ids(self):
+        results = [self._make_result("abc123")]
+        citations = resolve_citations(["abc123"], results)
+        assert len(citations) == 1
+        assert citations[0].chunk_id == "abc123"
+
+    def test_strips_brackets_from_ids(self):
+        results = [self._make_result("abc123")]
+        citations = resolve_citations(["[abc123]"], results)
+        assert len(citations) == 1
+        assert citations[0].chunk_id == "abc123"
+
+    def test_skips_unmatched_ids(self):
+        results = [self._make_result("abc123")]
+        citations = resolve_citations(["nonexistent"], results)
+        assert len(citations) == 0
+
+    def test_empty_cited_chunks(self):
+        results = [self._make_result("abc123")]
+        citations = resolve_citations([], results)
+        assert len(citations) == 0