diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/models.py b/haiku_rag_slim/haiku/rag/agents/analysis/models.py index 413c25b0..6b88e027 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/models.py +++ b/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): diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py index a26b70cf..53147c9b 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index ad5288fa..90976a49 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index d83ce3fc..3f0639eb 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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,