diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/__init__.py b/haiku_rag_slim/haiku/rag/agents/analysis/__init__.py index 0d0cca81..f7507d7d 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/__init__.py +++ b/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", diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/agent.py b/haiku_rag_slim/haiku/rag/agents/analysis/agent.py index 76ef9140..47e49d5c 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/agent.py +++ b/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, ) diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/models.py b/haiku_rag_slim/haiku/rag/agents/analysis/models.py index 0c8e86b4..413c25b0 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/models.py +++ b/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): - """Result from analysis agent execution.""" +class RawAnalysisResult(BaseModel): + """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) diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py b/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py index de4a253b..a26b70cf 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/prompts.py +++ b/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.""" diff --git a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py b/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py index 77a4dc87..c39ec1b3 100644 --- a/haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 834b66a0..6e0d36ca 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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) - if rag_state is None: - return - citations = getattr(rag_state, "citations", []) + citations = [] + for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE): + state = self._toolset.get_namespace(namespace) + 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( - self._document_filter - ) - self._state = self._toolset.build_state_snapshot() + rag_state.document_filter = doc_filter + analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE) + if isinstance(analysis_state, AnalysisState): + analysis_state.document_filter = doc_filter + self._state = self._toolset.build_state_snapshot() diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 5f76736c..c8473dd3 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index e675c563..36e9eb32 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/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 - for citation in citations: - citation.index = next_index - next_index += 1 - state.citations.extend(citations) + _append_citations(state, 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 diff --git a/haiku_rag_slim/haiku/rag/skills/analysis.py b/haiku_rag_slim/haiku/rag/skills/analysis.py index 9648d65f..022ce778 100644 --- a/haiku_rag_slim/haiku/rag/skills/analysis.py +++ b/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 diff --git a/tests/agents/analysis/test_agent.py b/tests/agents/analysis/test_agent.py index 98e0eb51..3f2ec6ff 100644 --- a/tests/agents/analysis/test_agent.py +++ b/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)