Avoid citation meta inside LLM context

haiku_rag_slim/haiku/rag/graph/common/models.py

@@ -29,17 +29,17 @@ class ResearchPlan(BaseModel):
 class Citation(BaseModel):
     """Resolved citation with full metadata for display/visual grounding."""
 
-    document_id: str = ""
-    chunk_id: str = ""
-    document_uri: str = ""
+    document_id: str
+    chunk_id: str
+    document_uri: str
     document_title: str | None = None
     page_numbers: list[int] = Field(default_factory=list)
     headings: list[str] | None = None
-    content: str = ""
+    content: str
 
 
-class SearchAnswer(BaseModel):
-    """Structured answer from a search operation."""
+class RawSearchAnswer(BaseModel):
+    """Answer to a search query with chunk references."""
 
     query: str = Field(..., description="The question that was answered")
     answer: str = Field(..., description="The answer to the question")
@@ -53,11 +53,32 @@ class SearchAnswer(BaseModel):
         ge=0.0,
         le=1.0,
     )
+
+
+class SearchAnswer(RawSearchAnswer):
+    """Answer to a search query with resolved citations."""
+
     citations: list[Citation] = Field(
         default_factory=list,
-        description="Resolved citations with full metadata (populated after search)",
+        description="Resolved citations with full metadata",
     )
 
+    @classmethod
+    def from_raw(
+        cls,
+        raw: RawSearchAnswer,
+        search_results: "list[SearchResult]",
+    ) -> "SearchAnswer":
+        """Create SearchAnswer from RawSearchAnswer with resolved citations."""
+        citations = resolve_citations(raw.cited_chunks, search_results)
+        return cls(
+            query=raw.query,
+            answer=raw.answer,
+            cited_chunks=raw.cited_chunks,
+            confidence=raw.confidence,
+            citations=citations,
+        )
+
 
 def resolve_citations(
     cited_chunk_ids: list[str],

haiku_rag_slim/haiku/rag/graph/common/nodes.py

@@ -13,7 +13,7 @@ from haiku.rag.config import Config
 from haiku.rag.config.models import AppConfig, ModelConfig
 from haiku.rag.graph.agui.emitter import AGUIEmitter
 from haiku.rag.graph.common import get_model
-from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
+from haiku.rag.graph.common.models import RawSearchAnswer, ResearchPlan, SearchAnswer
 from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
 from haiku.rag.store.models import SearchResult
 
@@ -24,9 +24,7 @@ class GraphContext(Protocol):
     original_question: str
     sub_questions: list[str]
 
-    def add_qa_response(
-        self, qa: SearchAnswer, search_results: list[SearchResult]
-    ) -> None:
+    def add_qa_response(self, qa: SearchAnswer) -> None:
         """Add a QA response to context."""
         ...
 
@@ -221,7 +219,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
 
     agent = Agent(
         model=get_model(model_config, config),
-        output_type=ToolOutput(SearchAnswer, max_retries=3),
+        output_type=ToolOutput(RawSearchAnswer, max_retries=3),
         instructions=SEARCH_AGENT_PROMPT,
         retries=3,
         deps_type=deps_type,
@@ -257,9 +255,11 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
 
     try:
         result = await agent.run(sub_q, deps=agent_deps)
-        answer = result.output
-        if answer:
-            state.context.add_qa_response(answer, agent_deps.search_results)
+        raw_answer = result.output
+        if raw_answer:
+            # Convert RawSearchAnswer to SearchAnswer with resolved citations
+            answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
+            state.context.add_qa_response(answer)
             # State updated with new answer - emit state update and narrate
             if deps.agui_emitter:
                 deps.agui_emitter.update_state(state)
@@ -279,7 +279,9 @@ async def _do_search[AgentDepsT: GraphAgentDeps](
                         "confidence": answer.confidence,
                     },
                 )
-        return answer
+            return answer
+        # Return empty SearchAnswer if no result
+        return SearchAnswer(query=sub_q, answer="", confidence=0.0)
     except Exception as e:
         if handle_exceptions:
             # Narrate the error

haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py

@@ -1,7 +1,7 @@
 from pydantic import BaseModel, Field
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.common.models import SearchAnswer, resolve_citations
+from haiku.rag.graph.common.models import SearchAnswer
 from haiku.rag.store.models import SearchResult
 
 
@@ -14,11 +14,8 @@ class DeepQAContext(BaseModel):
         default_factory=list, description="QA pairs collected during answering"
     )
 
-    def add_qa_response(
-        self, qa: SearchAnswer, search_results: list[SearchResult]
-    ) -> None:
-        """Add a QA response with resolved citations."""
-        qa.citations = resolve_citations(qa.cited_chunks, search_results)
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        """Add a QA response (citations already resolved)."""
         self.qa_responses.append(qa)
 
 

haiku_rag_slim/haiku/rag/graph/research/dependencies.py

@@ -3,7 +3,7 @@ from collections.abc import Iterable
 from pydantic import BaseModel, Field, PrivateAttr
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.graph.common.models import SearchAnswer, resolve_citations
+from haiku.rag.graph.common.models import SearchAnswer
 from haiku.rag.graph.research.models import (
     GapRecord,
     InsightAnalysis,
@@ -38,11 +38,8 @@ class ResearchContext(BaseModel):
         self._insights_by_id = {ins.id: ins for ins in self.insights}
         self._gaps_by_id = {gap.id: gap for gap in self.gaps}
 
-    def add_qa_response(
-        self, qa: SearchAnswer, search_results: list[SearchResult]
-    ) -> None:
-        """Add a structured QA response with resolved citations."""
-        qa.citations = resolve_citations(qa.cited_chunks, search_results)
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        """Add a structured QA response (citations already resolved)."""
         self.qa_responses.append(qa)
 
     def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:

haiku_rag_slim/haiku/rag/qa/agent.py

@@ -6,11 +6,7 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.config.models import ModelConfig
 from haiku.rag.graph.common import get_model
-from haiku.rag.graph.common.models import (
-    Citation,
-    SearchAnswer,
-    resolve_citations,
-)
+from haiku.rag.graph.common.models import Citation, RawSearchAnswer, resolve_citations
 from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
 from haiku.rag.store.models import SearchResult
 
@@ -34,7 +30,7 @@ class QuestionAnswerAgent:
         self._agent = Agent(
             model=model_obj,
             deps_type=Dependencies,
-            output_type=ToolOutput(SearchAnswer, max_retries=3),
+            output_type=ToolOutput(RawSearchAnswer, max_retries=3),
             instructions=system_prompt or QA_SYSTEM_PROMPT,
             retries=3,
         )

tests/graph/test_deep_qa.py

@@ -59,7 +59,7 @@ async def test_deep_qa_context_operations():
         cited_chunks=["chunk_1", "chunk_2"],
         confidence=0.9,
     )
-    context.add_qa_response(qa, search_results=[])
+    context.add_qa_response(qa)
     assert len(context.qa_responses) == 1
     assert context.qa_responses[0].query == "Sub Q1"