diff --git a/CHANGELOG.md b/CHANGELOG.md index 10ed91bb..4a760209 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/agents/qa/prompts.py b/haiku_rag_slim/haiku/rag/agents/qa/prompts.py index 6c6c126f..f235f9ff 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/prompts.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/agents/research/models.py b/haiku_rag_slim/haiku/rag/agents/research/models.py index a986769f..6d8fe8dd 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/models.py +++ b/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 diff --git a/tests/agents/research/test_models.py b/tests/agents/research/test_models.py index 148f06ca..53440e20 100644 --- a/tests/agents/research/test_models.py +++ b/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