add citation support to analysis agent

haiku_rag_slim/haiku/rag/agents/analysis/__init__.py

@@ -1,6 +1,10 @@
 from haiku.rag.agents.analysis.agent import create_analysis_agent
 from haiku.rag.agents.analysis.dependencies import AnalysisContext, AnalysisDeps
-from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
+from haiku.rag.agents.analysis.models import (
+    AnalysisResult,
+    CodeExecution,
+    RawAnalysisResult,
+)
 from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
 from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
 
@@ -8,6 +12,7 @@ __all__ = [
     "ANALYSIS_SYSTEM_PROMPT",
     "AnalysisContext",
     "AnalysisDeps",
+    "RawAnalysisResult",
     "AnalysisResult",
     "CodeExecution",
     "Sandbox",

haiku_rag_slim/haiku/rag/agents/analysis/agent.py

@@ -1,13 +1,13 @@
 from pydantic_ai import Agent, RunContext
 
 from haiku.rag.agents.analysis.dependencies import AnalysisDeps
-from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
+from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
 from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
 from haiku.rag.config.models import AppConfig
 from haiku.rag.utils import get_model
 
 
-def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, AnalysisResult]:
+def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, RawAnalysisResult]:
     """Create an analysis agent with code execution capability.
 
     The analysis agent can write and execute Python code in a sandboxed
@@ -22,10 +22,10 @@ def create_analysis_agent(config: AppConfig) -> Agent[AnalysisDeps, AnalysisResu
     """
     model = get_model(config.analysis.model, config)
 
-    agent: Agent[AnalysisDeps, AnalysisResult] = Agent(  # type: ignore[assignment]  # ty: ignore[invalid-assignment]
+    agent: Agent[AnalysisDeps, RawAnalysisResult] = Agent(  # type: ignore[assignment]  # ty: ignore[invalid-assignment]
         model,
         deps_type=AnalysisDeps,
-        output_type=AnalysisResult,
+        output_type=RawAnalysisResult,
         instructions=ANALYSIS_SYSTEM_PROMPT,
         retries=3,
     )

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

@@ -1,5 +1,7 @@
 from pydantic import BaseModel, Field
 
+from haiku.rag.agents.research.models import Citation
+
 
 class CodeExecution(BaseModel):
     """Result of executing a code block in the analysis sandbox."""
@@ -10,8 +12,20 @@ class CodeExecution(BaseModel):
     success: bool = Field(description="Whether execution completed without error")
 
 
-class AnalysisResult(BaseModel):
+class RawAnalysisResult(BaseModel):
-    """Result from analysis agent execution."""
+    """Raw result from the analysis agent (LLM output)."""
 
     answer: str = Field(description="The answer to the user's question")
     program: str = Field(description="The final consolidated program")
+    cited_chunks: list[str] = Field(
+        default_factory=list,
+        description="Chunk IDs from search results that informed the answer. Copy full UUIDs from search result chunk_id fields.",
+    )
+
+
+class AnalysisResult(BaseModel):
+    """Result from analysis execution with resolved citations."""
+
+    answer: str
+    program: str
+    citations: list[Citation] = Field(default_factory=list)

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

@@ -116,12 +116,13 @@ Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available)
 
 Your final response MUST be valid JSON matching this exact schema:
 ```json
-{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
+{"answer": "Your answer here", "program": "Your final program here", "cited_chunks": ["chunk-id-1", "chunk-id-2"]}
 ```
 
 - `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
 - `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
+- `cited_chunks`: List of chunk_id values from search results that informed your answer. Copy the full UUID strings from the `chunk_id` field of search results you used.
 
-Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
+Do NOT return arbitrary JSON structures. Always use the exact format above.
 
 You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py

@@ -10,6 +10,7 @@ from pydantic_monty import CallbackFile, MemoryFile, OSAccess
 
 from haiku.rag.agents.analysis.dependencies import AnalysisContext
 from haiku.rag.config.models import AppConfig
+from haiku.rag.store.models.chunk import SearchResult
 
 if TYPE_CHECKING:
     from pathlib import PurePosixPath
@@ -47,6 +48,7 @@ class Sandbox:
     _client: "HaikuRAG"
     _config: AppConfig
     _context: AnalysisContext
+    _search_results: "list[SearchResult]"
 
     def __init__(
         self,
@@ -57,6 +59,7 @@ class Sandbox:
         self._client = client
         self._config = config
         self._context = context
+        self._search_results = []
 
     def _build_external_functions(self) -> dict[str, Any]:
         """Build async external functions for the Monty interpreter."""
@@ -67,6 +70,7 @@ class Sandbox:
         async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
             results = await client.search(query, limit=limit, filter=context.filter)
             expanded = await client.expand_context(results)
+            self._search_results.extend(expanded)
             return [
                 {
                     "chunk_id": r.chunk_id,

haiku_rag_slim/haiku/rag/chat/app.py

@@ -30,6 +30,7 @@ from textual.worker import Worker
 from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import get_config
+from haiku.rag.skills.analysis import AnalysisState
 from haiku.rag.skills.rag import RAGState, get_agent_preamble
 from haiku.skills.agent import (
     SkillToolset,
@@ -51,6 +52,7 @@ if TYPE_CHECKING:
 
 
 RAG_STATE_NAMESPACE = "rag"
+ANALYSIS_STATE_NAMESPACE = "analysis"
 
 
 class ChatApp(App):
@@ -319,15 +321,15 @@ class ChatApp(App):
             chat_input.focus()
 
     async def _show_citations(self, chat_history: "ChatHistory") -> None:
-        """Show citations from the RAG state after an agent response."""
+        """Show citations from skill states after an agent response."""
         if not self._toolset:
             return
-        rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
+        citations = []
-        if rag_state is None:
+        for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
-            return
+            state = self._toolset.get_namespace(namespace)
-        citations = getattr(rag_state, "citations", [])
+            if state:
+                citations.extend(getattr(state, "citations", []))
         if citations:
-            # Show only new citations (since last response)
             await chat_history.add_citations(citations)
 
     async def action_clear_chat(self) -> None:
@@ -420,9 +422,11 @@ class ChatApp(App):
         self._document_filter = event.selected
 
         if self._toolset:
+            doc_filter = build_multi_document_filter(self._document_filter)
             rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
             if isinstance(rag_state, RAGState):
-                rag_state.document_filter = build_multi_document_filter(
+                rag_state.document_filter = doc_filter
-                    self._document_filter
+            analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
-                )
+            if isinstance(analysis_state, AnalysisState):
-                self._state = self._toolset.build_state_snapshot()
+                analysis_state.document_filter = doc_filter
+            self._state = self._toolset.build_state_snapshot()

haiku_rag_slim/haiku/rag/client.py

@@ -1239,10 +1239,19 @@ class HaikuRAG:
             context=context,
         )
 
+        from haiku.rag.agents.analysis.models import AnalysisResult
+        from haiku.rag.agents.research.models import resolve_citations
+
         agent = create_analysis_agent(self._config)
         result = await agent.run(question, deps=deps)
 
-        return result.output
+        output = result.output
+        citations = resolve_citations(output.cited_chunks, sandbox._search_results)
+        return AnalysisResult(
+            answer=output.answer,
+            program=output.program,
+            citations=citations,
+        )
 
     async def visualize_chunk(self, chunk: Chunk) -> list:
         """Render page images with bounding box highlights for a chunk.

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -210,7 +210,7 @@ async def skill_analyze(
     question: str,
     document: str | None = None,
     document_filter: str | None = None,
-) -> tuple[str, str, str | None]:
+) -> tuple[str, str, str | None, "list[Citation]"]:
     from haiku.rag.client import HaikuRAG
 
     async with HaikuRAG(db_path, config=config, read_only=True) as rag:
@@ -222,7 +222,7 @@ async def skill_analyze(
         if result.program:
             output += f"\n\nProgram:\n{result.program}"
 
-    return output, result.answer, result.program
+    return output, result.answer, result.program, result.citations
 
 
 def update_documents_state(
@@ -246,6 +246,15 @@ def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> An
     return None
 
 
+def _append_citations(state: Any, citations: "list[Citation]") -> None:
+    """Index and append citations to a skill state's citations list."""
+    next_index = len(state.citations) + 1
+    for citation in citations:
+        citation.index = next_index
+        next_index += 1
+    state.citations.extend(citations)
+
+
 def create_skill_extras(
     db_path: Path,
     config: AppConfig,
@@ -406,11 +415,7 @@ def create_skill_tools(
             )
 
             if state:
-                next_index = len(state.citations) + 1
+                _append_citations(state, citations)
-                for citation in citations:
-                    citation.index = next_index
-                    next_index += 1
-                state.citations.extend(citations)
                 state.qa_history.append(
                     QAHistoryEntry(
                         question=question, answer=answer, citations=citations
@@ -476,9 +481,11 @@ def create_skill_tools(
                 question: The question to answer.
                 document: Optional document ID or title to pre-load for analysis.
             """
+            from haiku.rag.utils import format_citations
+
             state = _get_state(ctx, state_type)
             state_filter = state.document_filter if state else None
-            output, answer, program = await skill_analyze(
+            output, answer, program, citations = await skill_analyze(
                 db_path,
                 config,
                 question,
@@ -493,6 +500,11 @@ def create_skill_tools(
                         program=program,
                     )
                 )
+                if citations:
+                    _append_citations(state, citations)
+
+            if citations:
+                output += "\n\n" + format_citations(citations)
 
             return output
 

haiku_rag_slim/haiku/rag/skills/analysis.py

@@ -2,8 +2,9 @@ import os
 from functools import cache
 from pathlib import Path
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
+from haiku.rag.agents.research.models import Citation
 from haiku.rag.config.models import AppConfig
 from haiku.rag.skills._tools import AnalysisEntry
 from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
@@ -13,6 +14,7 @@ from haiku.skills.parser import parse_skill_md
 class AnalysisState(BaseModel):
     document_filter: str | None = None
     analyses: list[AnalysisEntry] = []
+    citations: list[Citation] = Field(default_factory=list)
 
 
 STATE_TYPE = AnalysisState

tests/agents/analysis/test_agent.py

@@ -5,7 +5,7 @@ from pydantic_ai import Agent
 
 from haiku.rag.agents.analysis.agent import create_analysis_agent
 from haiku.rag.agents.analysis.dependencies import AnalysisDeps
-from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
+from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
 from haiku.rag.config import AppConfig, Config
 
 
@@ -19,7 +19,7 @@ class TestCreateAnalysisAgent:
         agent = create_analysis_agent(Config)
         assert isinstance(agent, Agent)
         assert agent.deps_type is AnalysisDeps
-        assert agent.output_type is AnalysisResult
+        assert agent.output_type is RawAnalysisResult
 
     def test_agent_has_execute_code_tool(self):
         agent = create_analysis_agent(Config)