remove cited_chunks from analysis agent, hoist _deny_write out of loop

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

@@ -17,10 +17,6 @@ class RawAnalysisResult(BaseModel):
 
     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):

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

@@ -116,12 +116,11 @@ 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 answer here", "program": "Your final program here", "cited_chunks": ["chunk-id-1", "chunk-id-2"]}
+{"answer": "Your answer here", "program": "Your final program here"}
 ```
 
 - `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 above.
 

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

@@ -141,6 +141,9 @@ class Sandbox:
         config = self._config
         files: list[MemoryFile | CallbackFile] = []
 
+        def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
+            raise PermissionError(f"Document files are read-only: {_path}")
+
         async with HaikuRAG(db_path, config=config, read_only=True) as rag:
             docs = await rag.list_documents(filter=self._context.filter)
 
@@ -211,9 +214,6 @@ class Sandbox:
 
                 return read_items
 
-            def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
-                raise PermissionError(f"Document files are read-only: {_path}")
-
             files.append(
                 CallbackFile(
                     f"{doc_dir}/content.txt",

haiku_rag_slim/haiku/rag/client.py

@@ -1240,13 +1240,29 @@ class HaikuRAG:
         )
 
         from haiku.rag.agents.analysis.models import AnalysisResult
-        from haiku.rag.agents.research.models import resolve_citations
+        from haiku.rag.agents.research.models import Citation
 
         agent = create_analysis_agent(self._config)
         result = await agent.run(question, deps=deps)
 
         output = result.output
-        citations = resolve_citations(output.cited_chunks, sandbox._search_results)
+        seen: set[str] = set()
+        citations: list[Citation] = []
+        for sr in sandbox._search_results:
+            if sr.chunk_id and sr.chunk_id not in seen:
+                seen.add(sr.chunk_id)
+                citations.append(
+                    Citation(
+                        index=len(seen),
+                        document_id=sr.document_id or "",
+                        chunk_id=sr.chunk_id,
+                        document_uri=sr.document_uri or "",
+                        document_title=sr.document_title,
+                        page_numbers=sr.page_numbers,
+                        headings=sr.headings,
+                        content=sr.content,
+                    )
+                )
         return AnalysisResult(
             answer=output.answer,
             program=output.program,