from dataclasses import dataclass, field from functools import cache from pathlib import Path from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import RunContext, ToolFailed from pydantic_ai.messages import ToolReturn from pydantic_ai.toolsets import FunctionToolset if TYPE_CHECKING: from haiku.rag.client import HaikuRAG from haiku.rag.capabilities._base import ( CodeExecutionEntry, RAGCapabilityBase, resolve_db_path, ) from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation STATE_NAMESPACE = "analysis" _CAPABILITY_ID = "haiku-rag-analysis" _TOOL_NAMES = frozenset({"analysis_search", "analysis_execute_code", "analysis_cite"}) _instructions_path = Path(__file__).parent / "instructions" / "analysis.md" class AnalysisState(BaseModel): document_filter: str | None = None executions: list[CodeExecutionEntry] = Field(default_factory=list) citation_index: dict[str, Citation] = Field(default_factory=dict) citations: list[str] = Field(default_factory=list) evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) searches: dict[str, list[SearchResult]] = Field(default_factory=dict) @cache def instructions() -> str: return _instructions_path.read_text().strip() def _recovery_hint(stderr: str) -> str: """Name the workaround for sandbox limits models trip over repeatedly. The instructions already say file objects are not iterable, and models write ``for line in open(...)`` regardless. Carrying the fix in the error gives them something to act on for the retry. """ if "TextIOWrapper" in stderr and "not iterable" in stderr: return ( "\n\nHint: file objects cannot be iterated here. Read lines with " '.readlines() or .read().split("\\n").' ) return "" @dataclass class AnalysisCapability(RAGCapabilityBase[AnalysisState]): """Deferred capability for sandboxed computation over a RAG corpus.""" sandbox: Sandbox | None = field(default=None, repr=False) execute_count: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "AnalysisCapability": capability = await super().for_run(ctx) assert isinstance(capability, AnalysisCapability) capability.sandbox = None capability.execute_count = 0 return capability async def _ensure_sandbox(self) -> Sandbox: if self.sandbox is None: rag = await self._ensure_rag() assert self.state is not None self.sandbox = Sandbox( db_path=self.db_path, config=self.config, context=AnalysisContext(filter=self.state.document_filter), rag=rag, lock=self.rag_lock, ) return self.sandbox async def _close(self) -> None: if self.sandbox is not None: await self.sandbox.close() self.sandbox = None await super()._close() def evidence_tool_names(self) -> set[str]: # Searching from inside the sandbox does not count against # `qa.max_searches`, so code execution outlives a spent search budget # as a way to reach new evidence. return super().evidence_tool_names() | {"analysis_execute_code"} def _spent_tool_names(self) -> set[str]: spent = super()._spent_tool_names() if self.execute_count >= self.config.analysis.max_executions: spent.add("analysis_execute_code") return spent async def _execute_code(self, code: str) -> str: assert self.state is not None self.execute_count += 1 if self.execute_count > self.config.analysis.max_executions: raise ToolFailed( "Code-execution limit reached. Give your final answer now from what " "you already have; do not call analysis_execute_code again." ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) if result.success or result.stdout: self._note_evidence() if sandbox._search_results: existing = self.state.searches.get("_sandbox", []) seen = {item.chunk_id for item in existing} for item in sandbox._search_results: if item.chunk_id not in seen: existing.append(item) seen.add(item.chunk_id) self.state.searches["_sandbox"] = existing self.state.executions.append( CodeExecutionEntry( code=code, stdout=result.stdout, stderr=result.stderr, success=result.success, ) ) if not result.success: raise ToolFailed( f"{result.stderr}{_recovery_hint(result.stderr)}" f"\n\nOutput: {result.stdout}" ) return result.stdout or "No output." def get_toolset(self) -> FunctionToolset[Any]: async def analysis_search( ctx: RunContext[Any], query: str, limit: int | None = None ) -> str | ToolReturn: """Search the knowledge base for evidence to analyze.""" return await self._with_state(self._search(query, limit)) async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any: """Execute Python against the sandboxed document filesystem.""" return await self._with_state(self._execute_code(code)) async def analysis_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any: """Register exact retrieved chunk IDs as citations for the answer.""" return await self._with_state(self._cite(chunk_ids)) return FunctionToolset( [analysis_search, analysis_execute_code, analysis_cite], id=_CAPABILITY_ID, max_retries=3, sequential=True, ) def create_capability( db_path: Path | None = None, config: AppConfig | None = None, *, defer_loading: bool = True, rag: "HaikuRAG | None" = None, request_limit: int | None = 30, vision: bool | None = None, ) -> AnalysisCapability: """Create a native Pydantic AI analysis capability. ``vision`` gates whether picture chunks are attached to search results as images, and should reflect the model the hosting agent actually runs. Defaults to ``config.analysis.model.vision`` (falling back to ``config.qa.model.vision``). """ if config is None: from haiku.rag.config import get_config config = get_config() analysis_model = config.analysis.model or config.qa.model return AnalysisCapability( db_path=resolve_db_path(db_path, config), config=config, borrowed_rag=rag, state_type=AnalysisState, state_namespace=STATE_NAMESPACE, instruction_text=instructions(), vision=analysis_model.vision if vision is None else vision, tool_names=_TOOL_NAMES, request_limit=request_limit, id=_CAPABILITY_ID, description=( "Analyze the haiku.rag corpus with search and sandboxed Python code. " "Use for counting, aggregation, statistics, data traversal, comparison " "across documents, and other tasks best solved by writing Python code." ), defer_loading=defer_loading, ) __all__ = [ "AnalysisCapability", "AnalysisState", "STATE_NAMESPACE", "create_capability", "instructions", ]