danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #376 from ggozad/feat/vfs-search

Browse source 
Analysis skill: structural VFS, multimodal citations, drop research
This commit is contained in:
Yiorgis Gozadinos 2026-05-20 14:17:31 +03:00 committed by GitHub
commit ae17e53f97
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
128 changed files with 75798 additions and 17471 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
CHANGELOG.md
View file

@ -1,14 +1,58 @@
# Changelog
## [Unreleased]
### Added
- `heading_level` and `tree_depth` on `DocumentItem`, populated by `extract_items` and persisted on `document_items`. 0.48.0 migration backfills existing rows from each doc's docling structure blob.
- `toc.json` in the analysis sandbox VFS at `/documents/{id}/toc.json`. Nested tree on HTML/markdown sources, flat sibling list on PDFs. Each node carries `{self_ref, level, title, page_numbers, item_range, chunk_ids, children}`. `chunk_ids` aggregates the citable chunks across the section's `item_range`, so the analysis skill can `cite()` a section from one VFS read instead of falling back to a corpus-wide `search()` that risks cross-document hits.
- `chunk_ids` on every `items.jsonl` row (the citable chunks that contain that item).
- `picture_refs` on sandbox `search()` result dicts and on `Citation` (subset of `doc_item_refs` starting with `#/pictures/`).
- `picture_captions: dict[str, str]` on `SearchResult`, populated alongside `image_data` and rendered as a labelled line in `format_for_agent` for picture-bearing chunks.
- Chat TUI renders picture citations inline via `textual_image.widget.Image` inside the existing `CitationWidget`.
- CLI citation panel renders `picture_refs` inline via `textual_image.renderable.Image` next to the text preview. `format_citations_rich` is async and takes an optional `HaikuRAG` client; without one, figures fall back to `[Figure: <ref>]` markers.
- BTree scalar indexes on `document_items.{document_id, position, self_ref}`. The 0.48.0 migration creates them on existing DBs. Per-doc lookups go from full-table scans (~100–300 ms) to point queries (~3–21 ms) on small/medium corpora.
- Per-doc lazy cache for `items.jsonl` and `toc.json` in the analysis sandbox. First read fetches; subsequent reads of either file in the same `execute_code` session hit a serialized cache. One DB fetch per doc per session.
- `cite` tool now accepts chunk_ids that resolve via the chunks table, not only chunk_ids from a prior `search()` result. Lets the model cite directly from `items.jsonl` / `toc.json` rows. The hallucination guard (`ModelRetry` on chunk_ids that don't exist in the DB) is preserved.
- `AppConfig.evaluations` (`EvaluationsConfig`) with an optional `judge: ModelConfig`. Lets the eval CLI pin the LLM-as-judge per-yaml — including a custom `base_url` for any OpenAI-compatible endpoint (vLLM, LM Studio) without env-var routing.
### Removed
- **Multi-agent research workflow.** Removes `agents/research/` (graph, state, deps, models, prompts), `client.research`, the CLI `research` command, the MCP `research_question` tool, `ResearchConfig`, `AppConfig.research`, `PromptsConfig.synthesis`, and the corresponding wiring in `chat/__init__.py` and `client/downloads.py`. The pydantic-graph workflow was a three-node loop whose differentiators vs the rag skill (planner step, structured `ResearchReport`, iteration bound) were either redundant with `qa.max_searches` or sat on pre-cite-tool legacy. Multi-step questions go through `client.ask` (rag skill).
- `llm()` from the analysis sandbox. Sandbox externals are now `search` and `list_documents` only.
- `list_documents` top-level tool from the analysis skill (still available as `await list_documents()` inside `execute_code`).
- `documents=` kwarg on `client.analyze` (and the `--document` flag on `haiku-rag analyze` / MCP `analyze` tool). The pre-loaded `documents` Python variable inside the sandbox is no longer populated. Use `filter=` (SQL WHERE clause) to scope analysis to specific documents.
- `AnalysisResult.program`. The per-execution programs are still tracked on `AnalysisState.executions` (the analysis skill's `execute_code` tool populates it); consumers that need the executed code should pull it from the skill state instead of the function return value.
- `--cite` flag on `haiku-rag ask`. Citations always render after the answer now.
- `system_prompt` kwarg on `client.ask`. No production caller used it; `config.prompts.domain_preamble` already covers the preamble use case.
- Standalone QA agent (`haiku.rag.agents.qa.*`) and analysis agent module (`haiku.rag.agents.analysis.agent`, `haiku.rag.agents.analysis.prompts`). Also drops `RawAnalysisResult`, `CodeExecution`, `AnalysisDeps`, and the dead `documents=` preload path in `Sandbox`.
- `prompts.qa` config field.
- `evaluations optimize` subcommand and GEPA prompt-optimization. Drops `gepa` dep.
- `--target qa` from `evaluations run`. Default is now `rag-skill`.
- `--judge-model` flag from `evaluations run`. Set the judge in `config.evaluations.judge` instead.
- `position` field on `toc.json` nodes (redundant with `item_range[0]`).
- `position` and `tree_depth` from `items.jsonl` row serialization. Both fields are still persisted on `DocumentItem`; they are no longer surfaced to the sandbox.
### Changed
- `haiku.rag.agents.analysis` moved to `haiku.rag.sandbox`. Public surface: `from haiku.rag.sandbox import Sandbox, SandboxResult, AnalysisContext, AnalysisResult`.
- `Citation` and `resolve_citations` moved to `haiku.rag.store.models.citation` (was `haiku.rag.agents.research.models`), peer to the other output domain models.
- `search.limit` default lowered from `10` to `5`.
- Search result formatter surfaces picture captions on a labelled line when a chunk's expanded refs include pictures.
- Picture bytes attached to `search()` results are bounded to the pre-expansion chunk's `doc_item_refs`. Section expansion that sweeps in adjacent picture-bearing items no longer pulls their bytes into the response. Observed ~16× reduction on tool-response payload sizes; eliminates a class of cross-figure contamination in the agent's view.
- rag-analysis SKILL.md steers structural lookups ("which section X", "list sections of Y", "summarize section Z") to read `/documents/{id}/toc.json` first and cite the matching node's `chunk_ids` directly, instead of calling `search()`. Empirically validated on the ORB multimodal and Wix corpora: locator and orientation questions now cite the correct document instead of falling back to cross-document search hits.
- CLI citation panel compacted: 300-char text preview (no full chunk dump), `[N] Title (URI) — pp. — §Section` header, dimmed `doc: <id> chunk: <id>` footer. Green "Citations" label matches the green "Answer:" label.
- `AnalysisConfig.model` defaults to `None` (was `ollama:gpt-oss/no-thinking/temp=0`). Resolves via `config.analysis.model or config.qa.model`.
- `client.ask` and `client.analyze` route through the rag and rag-analysis skills internally.
- Bump `docling>=2.93.0` and `docling-core>=2.75.0`.
- Bump `pydantic-ai-slim>=1.96.0`. Migrate off deprecated APIs: AG-UI imports use `pydantic_ai.ui.ag_ui`, docs/CLI examples use the explicit `openai-chat:` model prefix, and `Agent(retries=)` is split into `tool_retries=` + `output_retries=`.
- Bump `pydantic-monty>=0.0.17`. Migrate off deprecated `pydantic_monty.run_repl_async(repl, ...)` to `repl.feed_run_async(...)`.
- Cap `transformers<5.0.0` in the `mxbai` extra: `mxbai-rerank>=0.1.6` calls `tokenizer.prepare_for_model` which transformers 5 removed.
- Refresh the rest of the lockfile to latest within current constraints (pydantic, pydantic-ai, rich, ruff, ty, pytest, torch, textual, textual-image, watchfiles, pre-commit, datasets, and transitives).
### Fixed
- Chat TUI's state-edit screen syntax-highlights JSON instead of falling back to plain text. Adds `tree-sitter` + `tree-sitter-json` to the `[tui]` extra.
## [0.47.0] - 2026-05-14
### Added

20
README.md
View file

@ -11,14 +11,13 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
- **Question answering** — QA agents with citations (page numbers, section headings)
- **Question answering** — RAG skill with citations (page numbers, section headings)
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA/Research: any model supported by Pydantic AI
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA: any model supported by Pydantic AI
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
- **CLI & Python API** — Full functionality from command line or code
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
@ -63,9 +62,6 @@ haiku-rag search "attention mechanism"
# Ask questions with citations
haiku-rag ask "What datasets were used for evaluation?" --cite
# Research mode — iterative planning and search
haiku-rag research "What are the limitations of the approach?"
# Analyze — complex analytical tasks via code execution
haiku-rag analyze "How many documents mention transformers?"
@ -83,7 +79,7 @@ See [Configuration](https://ggozad.github.io/haiku.rag/configuration/) for custo
```python
from haiku.rag.client import HaikuRAG
async with HaikuRAG("research.lancedb", create=True) as rag:
async with HaikuRAG("knowledge.lancedb", create=True) as rag:
# Index documents
await rag.create_document_from_source("paper.pdf")
await rag.create_document_from_source("https://arxiv.org/pdf/1706.03762")
@ -100,7 +96,7 @@ async with HaikuRAG("research.lancedb", create=True) as rag:
print(f" [{cite.chunk_id}] p.{cite.page_numbers}: {cite.content[:80]}")
```
For research agents and chat, see the [Agents docs](https://ggozad.github.io/haiku.rag/agents/).
For details on the skills the client wraps, see the [Skills docs](https://ggozad.github.io/haiku.rag/skills/).
## MCP Server
@ -123,7 +119,7 @@ Add to your Claude Desktop configuration:
}
```
Provides tools for document management, search, QA, and research directly in your AI assistant.
Provides tools for document management, search, QA, and analysis directly in your AI assistant.
## Examples
@ -141,8 +137,8 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - YAML configuration
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA and research agents
- [Analysis Agent](https://ggozad.github.io/haiku.rag/agents/analysis/) - Complex analytical tasks via code execution
- [Skills](https://ggozad.github.io/haiku.rag/skills/) - The RAG and analysis skills the client wraps
- [Analysis](https://ggozad.github.io/haiku.rag/agents/analysis/) - Complex analytical tasks via code execution
- [Applications](https://ggozad.github.io/haiku.rag/apps/) - Chat TUI, web app, and inspector
- [Server](https://ggozad.github.io/haiku.rag/server/) - File monitoring and MCP
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration

42
docs/agents/analysis.md
View file

@ -1,6 +1,6 @@
# Analysis Agent
# Analysis
The analysis agent enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
The analysis skill enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
- **Aggregation**: "How many documents mention security vulnerabilities?"
- **Computation**: "What's the average revenue across all quarterly reports?"
@ -9,10 +9,10 @@ The analysis agent enables complex analytical tasks by writing and executing Pyt
## How It Works
1. The agent receives a question
1. The skill receives a question
2. It writes Python code to explore the knowledge base
3. Code executes in a sandboxed Python interpreter with access to search, LLM, and a virtual filesystem of documents
4. The agent iterates: run code, examine results, refine approach
3. Code executes in a sandboxed Python interpreter with access to search and a virtual filesystem of documents
4. The skill iterates: run code, examine results, refine approach
5. Final answer is synthesized from the gathered data
## CLI Usage
@ -21,11 +21,8 @@ The analysis agent enables complex analytical tasks by writing and executing Pyt
# Basic usage
haiku-rag analyze "How many documents are in the database?"
# With document filter (restricts what the agent can access)
# With document filter (restricts what the skill can access)
haiku-rag analyze "Summarize the key points" --filter "uri LIKE '%report%'"
# Pre-load specific documents
haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document "Q2 Report"
```
## Python Usage
@ -34,24 +31,20 @@ haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document
from haiku.rag.client import HaikuRAG
async with HaikuRAG(path_to_db) as client:
# Basic question
result = await client.analyze("How many documents mention 'security'?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
print(result.answer)
for citation in result.citations:
print(citation.uri, citation.title)
# With filter (agent can only see filtered documents)
# With filter (skill can only see filtered documents)
result = await client.analyze(
"What is the total revenue?",
filter="title LIKE '%Financial%'"
)
# Pre-load specific documents
result = await client.analyze(
"Compare the conclusions",
documents=["Report A", "Report B"]
)
```
The executed Python program(s) for each turn are not on `AnalysisResult` itself; they live on `AnalysisState.executions` while the skill runs.
## Sandbox Capabilities
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with:
@ -60,9 +53,8 @@ The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https:
| Function | Description |
|----------|-------------|
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` for cross-referencing with `items.jsonl` |
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` and `picture_refs` for cross-referencing with `items.jsonl` |
| `list_documents()` | List all documents in the knowledge base |
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
### Document Filesystem
@ -81,13 +73,11 @@ All documents are mounted as a virtual filesystem at `/documents/`. The agent us
Search results include `doc_item_refs` (e.g. `["#/texts/5", "#/tables/0"]`) that match `self_ref` values in `items.jsonl`, enabling navigation from search hits to document structure.
When documents are pre-loaded via the `documents` parameter, they are also injected as a `documents` variable accessible in the sandbox code.
### Python Features
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, file I/O via `pathlib.Path`, and the `json`, `re`, `math` modules.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, use `import re` or string methods.
### Security
@ -101,10 +91,10 @@ Code executes in an isolated interpreter with:
## Context Filter
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
The `filter` parameter restricts what documents the skill can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
```python
# Agent can only see documents with "confidential" in the URI
# Skill can only see documents with "confidential" in the URI
result = await client.analyze(
"Summarize all findings",
filter="uri LIKE '%confidential%'"

163
docs/agents/index.md
View file

@ -1,163 +0,0 @@
# Agents
Three agentic flows are provided by haiku.rag:
- **Simple QA Agent** — a focused question answering agent
- **Research Graph** — a multi-step research workflow with question decomposition
- **Analysis Agent** — complex analytical tasks via sandboxed Python code execution (see [Analysis Agent](analysis.md))
For multi-turn conversational RAG, haiku.rag provides [skills](../skills/index.md) built on [haiku.skills](https://github.com/ggozad/haiku.skills). The skills bundle search, Q&A, analysis, and research tools with session state management.
See [QA and Research Configuration](../configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings.
## Simple QA Agent
The simple QA agent answers a single question using the knowledge base. It retrieves relevant chunks, optionally expands context around them, and asks the model to answer strictly based on that context.
Key points:
- Uses a single `search_documents` tool to fetch relevant chunks
- Can be run with or without inline citations in the prompt
- Returns a plain string answer
**CLI usage:**
```bash
haiku-rag ask "What is climate change?"
# With citations
haiku-rag ask "What is climate change?" --cite
```
**Python usage:**
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import ModelConfig
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
async with HaikuRAG(path_to_db) as client:
agent = QuestionAnswerAgent(
client=client,
model_config=ModelConfig(provider="openai", name="gpt-4o-mini"),
)
answer, citations = await agent.answer("What is climate change?")
print(answer)
```
## Research Graph
The research workflow is implemented as a typed pydantic-graph. It uses an iterative feedback loop where the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize.
```mermaid
---
title: Research graph
---
stateDiagram-v2
state plan_next_decision <<choice>>
[*] --> plan_next
plan_next --> plan_next_decision
plan_next_decision --> search_one: Has next question
plan_next_decision --> synthesize: Complete or max iterations
search_one --> plan_next: Answer added to context
synthesize --> [*]
note right of plan_next
Uses prior_answers from previous iterations.
Uses a different prompt when prior answers exist.
end note
```
The graph receives a `ResearchContext` containing:
- `original_question` — the user's question
- `qa_responses` — prior answers from previous iterations (injected as `<prior_answers>` XML)
When prior answers are provided, the planner uses a context-aware prompt that evaluates whether existing evidence is sufficient. If it is, the planner marks `is_complete=True` and the graph skips directly to synthesis without any searches.
**Key nodes:**
- **plan_next**: Evaluates gathered evidence and either proposes the next question to investigate or marks research as complete. Uses a context-aware prompt when prior answers exist, allowing it to skip research entirely.
- **search_one**: Answers a single question using the knowledge base (up to 3 search calls per question). Each answer is added to `ResearchContext.qa_responses` for the next planning iteration.
- **synthesize**: Generates the final output from all gathered evidence.
**Iterative flow:**
- Each iteration: planner evaluates context → proposes one question → search answers it → loop back
- Planner can decompose complex questions (e.g., "benefits and drawbacks" → start with "benefits")
- Prior answers let the planner skip redundant searches
- Loop terminates when planner marks `is_complete=True` or `max_iterations` is reached
### CLI Usage
```bash
# Basic usage
haiku-rag research "How does haiku.rag organize and query documents?"
# With document filter
haiku-rag research "What are the key findings?" --filter "uri LIKE '%report%'"
```
### Python Usage
**Basic example:**
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=Config)
context = ResearchContext(original_question="What are the main features?")
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
report = await graph.run(state=state, deps=deps)
print(report.title)
print(report.executive_summary)
```
**With custom config:**
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ModelConfig, ResearchConfig
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
custom_config = AppConfig(
research=ResearchConfig(
model=ModelConfig(provider="openai", name="gpt-4o-mini"),
max_iterations=5,
max_concurrency=3,
)
)
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=custom_config)
context = ResearchContext(original_question="What are the main features?")
state = ResearchState.from_config(context=context, config=custom_config)
deps = ResearchDeps(client=client)
report = await graph.run(state=state, deps=deps)
```
### Filtering Documents
Restrict searches to specific documents via the `search_filter` parameter:
```python
# Set filter before running the graph
state = ResearchState.from_config(context=context, config=Config)
state.search_filter = "id IN ('doc-123', 'doc-456')"
report = await graph.run(state=state, deps=deps)
```
The filter applies to all search operations in the graph. See [Filtering Search Results](../python.md#filtering-search-results) for available filter columns and syntax.

2
docs/apps.md
View file

@ -185,7 +185,7 @@ Search uses hybrid (vector + full-text) search across all chunks.
### Context Expansion
Press `c` while viewing a chunk to see the expanded context that would be provided to the QA agent:
Press `c` while viewing a chunk to see the expanded context that would be provided to the rag skill:
- Section-aware expansion: expands to fill the current document section
- Noise filtering: footnotes, page headers/footers excluded from structured documents

8
docs/benchmarks.md
View file

@ -61,9 +61,9 @@ evaluations run repliqa --config /path/to/haiku.rag.yaml --db /path/to/custom.la
- `--skip-qa` - Skip QA benchmark
- `--limit N` - Limit number of test cases
- `--name NAME` - Override the evaluation name
- `--judge-model PROVIDER:NAME` - Override the LLM judge model. Defaults to `ollama:qwen3.6` so the judge stays stable when the QA / skill model changes.
- `--target {qa,rag-skill,analysis-skill}` - Choose what to benchmark (default: `qa`). `rag-skill` and `analysis-skill` run the corresponding [skill](skills/index.md) end-to-end against the same datasets and judge as the QA agent.
- `--skill-model PROVIDER:NAME` - Override the skill model independently from the judge (default: `config.qa.model`). Only valid with skill targets.
- `--judge-model PROVIDER:NAME` - Override the LLM judge model. Defaults to `ollama:qwen3.6` so the judge stays stable when the answering model changes.
- `--target {rag-skill,analysis-skill}` - Choose which [skill](skills/index.md) to benchmark end-to-end against the same datasets and judge (default: `rag-skill`).
- `--skill-model PROVIDER:NAME` - Override the skill model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-skill`).
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
@ -88,7 +88,7 @@ If no config file is specified, the script searches standard locations: `./haiku
### QA Accuracy
For question-answering evaluation, `pydantic-evals` coordinates an LLM judge to determine whether answers are correct. The default judge is `ollama:qwen3.6` — pinned so changes to the QA or skill model don't change the judge underneath. Override per run with `--judge-model provider:name`. Accuracy is the fraction of correctly answered questions.
For question-answering evaluation, `pydantic-evals` coordinates an LLM judge to determine whether answers are correct. The default judge is `ollama:qwen3.6` — pinned so changes to the skill model don't change the judge underneath. Override per run with `--judge-model provider:name`. Accuracy is the fraction of correctly answered questions.
We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.

39
docs/cli.md
View file

@ -171,21 +171,15 @@ Ask questions about your documents:
haiku-rag ask "Who is the author of haiku.rag?"
```
Ask questions with citations showing source documents:
```bash
haiku-rag ask "Who is the author of haiku.rag?" --cite
```
Filter to specific documents:
```bash
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
```
The QA agent searches your documents for relevant information and provides a comprehensive answer. When available, citations use the document title; otherwise they fall back to the URI.
`ask` runs the [rag skill](skills/index.md) and always renders citations under the answer. When available, citations use the document title; otherwise they fall back to the URI.
Flags:
- `--cite`: Include citations showing which documents were used
- `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
## Chat
@ -237,26 +231,6 @@ The inspector provides:
See [Applications](apps.md#inspector) for details.
## Research
Run the multi-step research graph:
```bash
haiku-rag research "How does haiku.rag organize and query documents?"
```
Filter to specific documents:
```bash
haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'"
```
Flags:
- `--filter` / `-f`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results))
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
## Analyze
Answer complex analytical questions via code execution:
@ -271,18 +245,11 @@ Filter to specific documents:
haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial%'"
```
Pre-load specific documents for comparison:
```bash
haiku-rag analyze "Compare the conclusions" --document "Report A" --document "Report B"
```
Flags:
- `--filter` / `-f`: SQL WHERE clause to restrict document access
- `--document` / `-d`: Pre-load a document by title or ID (can repeat)
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
See [Analysis](agents/analysis.md) for details on capabilities and configuration.
## Create Skill
@ -523,7 +490,7 @@ This command downloads:
- Docling OCR/conversion models
- HuggingFace tokenizer (for chunking)
- Ollama models referenced in your configuration (embeddings, QA, research, rerank)
- Ollama models referenced in your configuration (embeddings, QA, rerank)
Progress is displayed in real-time with download status and progress bars for Ollama model pulls.

15
docs/configuration/index.md
View file

@ -90,15 +90,6 @@ qa:
temperature: 0.3
max_searches: 3
research:
model:
provider: "" # Empty to use qa settings
name: ""
enable_thinking: false
temperature: 0.3
max_iterations: 3
max_concurrency: 1
search:
limit: 10 # Default number of results to return
max_context_chars: 10000 # Maximum characters in expanded context
@ -106,9 +97,7 @@ search:
vector_refine_factor: 30
prompts:
domain_preamble: "" # Prepended to all agent prompts
qa: null # Custom QA agent prompt (null = use default)
synthesis: null # Custom research synthesis prompt (null = use default)
domain_preamble: "" # Prepended to skill instructions
processing:
converter: docling-local # docling-local or docling-serve
@ -190,7 +179,7 @@ This is useful for:
For detailed configuration of specific topics, see:
- **[Providers](providers.md)** - Model settings and provider-specific configuration (embeddings, reranking)
- **[Search and Question Answering](qa-research.md)** - Search settings, question answering, and research workflows
- **[Search and Question Answering](qa.md)** - Search settings and question answering
- **[Document Processing](processing.md)** - Document conversion, chunking, and file monitoring
- **[Storage](storage.md)** - Database, remote storage, and vector indexing
- **[Prompts](prompts.md)** - Customize agent prompts for your domain

2
docs/configuration/processing.md
View file

@ -289,7 +289,7 @@ processing:
chunk_size: 256 # Maximum tokens per chunk
```
Context expansion settings (for enriching search results with surrounding content) are configured in the `search` section. See [Search Settings](qa-research.md#search-settings).
Context expansion settings (for enriching search results with surrounding content) are configured in the `search` section. See [Search Settings](qa.md#search-settings).
## File Monitoring

74
docs/configuration/prompts.md
View file

@ -1,36 +1,30 @@
# Prompt Customization
Customize the prompts used by haiku.rag's AI agents to better match your domain and use case.
Customize the prompts used by haiku.rag's skills to better match your domain and use case.
## Configuration
```yaml
prompts:
# Domain context prepended to all agent prompts
# Domain context prepended to skill instructions
domain_preamble: |
This knowledge base contains technical documentation for the Helios solar panel
system, including installation manuals, maintenance procedures, and safety guidelines.
Questions about "the system" or unqualified specs refer to the Helios panel.
# Full replacement for QA agent prompt (optional)
qa: null
# Full replacement for research synthesis prompt (optional)
synthesis: null
# VLM prompt for image description during conversion (optional)
picture_description: null # Uses default prompt
```
## Domain Preamble
The `domain_preamble` field provides **domain context** that is prepended to all agent prompts — the main agent, skill subagents, and internal agents (QA, research planning, search, evaluation, and synthesis). Use this to:
The `domain_preamble` field provides **domain context** prepended to the rag and rag-analysis skill instructions. Use this to:
- Describe what the knowledge base contains
- Clarify domain-specific terminology
- Provide context that helps agents interpret ambiguous queries
- Provide context that helps the model interpret ambiguous queries
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Behavioral guidance (tone, response style, formatting rules) belongs in the agent's system prompt or custom `prompts.qa`.
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Behavioral guidance (tone, response style, formatting rules) lives in the skill's SKILL.md — fork the skill via `haiku-rag create-skill` to customize behavior.
**Example:**
@ -42,62 +36,6 @@ prompts:
"Deployment" refers to Acme's managed deployment service, not general CI/CD.
```
## Custom QA Prompt
Replace the default QA agent prompt entirely by setting `prompts.qa`. The prompt should instruct the agent how to:
1. Use the `search_documents` tool to find relevant content
2. Interpret search results with scores and metadata
3. Cite sources using chunk IDs
4. Handle insufficient information
**Example:**
```yaml
prompts:
qa: |
You are a concise technical assistant. Answer questions using only the knowledge base.
Process:
1. Search for relevant documents using the search_documents tool
2. Review results ordered by relevance (rank 1 = most relevant)
3. Provide a brief, direct answer based on retrieved content
Guidelines:
- Use only information from search results
- Include chunk IDs in cited_chunks for sources you use
- If information is insufficient, say so clearly
- Be concise - avoid unnecessary elaboration
```
## Custom Synthesis Prompt
Replace the research report synthesis prompt by setting `prompts.synthesis`. This controls how the multi-agent research workflow generates its final report.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `sources_summary`.
**Example:**
```yaml
prompts:
synthesis: |
Generate a research report based on the gathered evidence.
Output format:
- title: 5-12 word title
- executive_summary: 3-5 sentence overview
- main_findings: 4-8 bullet points of key findings
- conclusions: 2-4 bullet points
- recommendations: 2-5 actionable recommendations
- limitations: 1-3 limitations or gaps
- sources_summary: Brief description of sources used
Guidelines:
- Base all content strictly on collected evidence
- Be specific and objective
- Avoid meta-commentary like "This report covers..."
```
## Picture Description Prompt
Customize the prompt used when generating VLM descriptions for embedded images during document conversion. This prompt is sent to the configured Vision Language Model for each image.
@ -130,8 +68,6 @@ from haiku.rag.config.models import PromptsConfig
config = AppConfig(
prompts=PromptsConfig(
domain_preamble="This knowledge base contains Acme Corp product documentation and API references.",
qa=None, # Use default QA prompt
synthesis=None, # Use default synthesis prompt
picture_description="Describe this image for search indexing.",
)
)

10
docs/configuration/providers.md
View file

@ -7,7 +7,7 @@ haiku.rag supports multiple AI providers for embeddings, question answering, and
## Model Settings
Configure model behavior for `qa` and `research` workflows. These settings apply to any provider that supports them.
Configure model behavior for the `qa` and `analysis` skills. These settings apply to any provider that supports them.
### Basic Settings
@ -22,7 +22,7 @@ qa:
**Available options:**
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA, research, and title generation; 0.0 for analysis and picture description.
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA and title generation; 0.0 for analysis and picture description.
- Lower (0.0-0.3): Deterministic, focused responses
- Medium (0.4-0.7): Balanced
- Higher (0.8-1.0+): Creative, varied responses
@ -39,10 +39,6 @@ The `enable_thinking` setting controls whether models use explicit reasoning ste
qa:
model:
enable_thinking: true # Better grounded answers
research:
model:
enable_thinking: true # Deeper reasoning
```
**Values:**
@ -64,7 +60,7 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
**When to use:**
- Enable for QA, research, complex reasoning, and mathematical problems
- Enable for QA, complex reasoning, and mathematical problems
- Disable for speed-critical applications, title generation, and simple tasks
### Raw Provider Pass-through

37
docs/configuration/qa-research.md → docs/configuration/qa.md
View file

@ -10,7 +10,7 @@ search:
max_context_chars: 10000 # Maximum characters in expanded context
```
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers) — this naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
@ -20,7 +20,7 @@ Context expansion is automatic and section-aware. For structured documents (with
## Question Answering Configuration
Configure the QA workflow:
Configure the rag skill (used by `client.ask`, `haiku-rag ask`, and the MCP `ask_question` tool):
```yaml
qa:
@ -29,38 +29,17 @@ qa:
name: gpt-oss
enable_thinking: true
temperature: 0.3 # Default: 0.3
vision: false # Set true for vision-capable QA models
vision: false # Set true for vision-capable models
max_searches: 3 # Maximum search tool calls per question
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
- **model.vision**: Set to `true` for vision-capable QA models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The agent's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`; otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures--embedder--qa-model-how-the-pieces-compose) for the full matrix.
- **max_searches**: Maximum number of search tool calls the QA agent can make per question (default: 3)
## Research Configuration
Configure the multi-agent research workflow:
```yaml
research:
model:
provider: "" # Empty to use qa settings
name: "" # Empty to use qa model
enable_thinking: false
temperature: 0.3 # Default: 0.3
max_iterations: 3
max_concurrency: 1
```
- **model**: LLM configuration. Leave provider/model empty to inherit from `qa` (see [Providers](providers.md#model-settings))
- **max_iterations**: Maximum planning/search iterations (default: 3)
- **max_concurrency**: Concurrent search operations (default: 1)
The research workflow uses an iterative feedback loop: the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize. This continues until the planner marks research as complete or `max_iterations` is reached.
- **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The skill's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`; otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures--embedder--qa-model-how-the-pieces-compose) for the full matrix.
- **max_searches**: Maximum number of search tool calls the rag skill can make per question (default: 3)
## Analysis Configuration
Configure the analysis agent:
Configure the analysis skill:
```yaml
analysis:
@ -72,8 +51,8 @@ analysis:
max_output_chars: 50000 # Truncate output after this many chars
```
- **model**: LLM configuration (see [Providers](providers.md#model-settings))
- **model**: LLM configuration (see [Providers](providers.md#model-settings)). When unset, falls back to `qa.model`.
- **code_timeout**: Maximum seconds for each code execution (default: 60)
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
See [Analysis Agent](../agents/analysis.md) for usage details.
See [Analysis](../agents/analysis.md) for usage details.

2
docs/configuration/storage.md
View file

@ -123,7 +123,7 @@ search:
vector_refine_factor: 30 # Re-ranking factor for accuracy
```
For search behavior settings (`limit`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
For search behavior settings (`limit`, `max_context_chars`), see [Search and Question Answering](qa.md#search-settings).
- **vector_index_metric**: Distance metric for vector similarity:
- `cosine`: Cosine similarity (default, best for most embeddings)

11
docs/index.md
View file

@ -8,14 +8,13 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
- **Question answering** — QA agents with citations (page numbers, section headings)
- **Question answering** — RAG skill with citations (page numbers, section headings)
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text via pydantic-ai `BinaryContent` when `qa.model.vision = true`
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA/Research: any model supported by Pydantic AI
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA: any model supported by Pydantic AI
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
- **CLI & Python API** — Full functionality from command line or code
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)
@ -67,8 +66,8 @@ haiku-rag chat # Interactive conversation mode
- [CLI](cli.md) - Command line interface usage
- [Python](python.md) - Python API reference
- [Custom Pipelines](custom-pipelines.md) - Build custom processing workflows
- [Agents](agents/index.md) - QA, chat, and research agents
- [Analysis Agent](agents/analysis.md) - Complex analytical tasks via code execution
- [Skills](skills/index.md) - The RAG and analysis skills the client wraps
- [Analysis](agents/analysis.md) - Complex analytical tasks via code execution
- [Applications](apps.md) - Chat TUI, web app, and inspector
- [Server](server.md) - File monitoring and server mode
- [MCP](mcp.md) - Model Context Protocol integration

4
docs/mcp.md
View file

@ -52,10 +52,6 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients like
- `cite` (optional): Include source citations (default: false)
- `deep` (optional): Use multi-agent deep QA for complex questions (default: false)
- **`research_question`** - Run multi-agent research on complex topics
- `question` (required): The research question
- Returns a structured research report with findings, conclusions, and sources
- **`analyze`** - Answer complex analytical questions via code execution
- `question` (required): The question to answer
- `filter` (optional): SQL WHERE clause to restrict document access

32
docs/python.md
View file

@ -421,19 +421,6 @@ for cite in citations:
print(f" [{cite.chunk_id}] {cite.document_title or cite.document_uri}")
```
Customize the QA agent's behavior with a custom system prompt:
```python
custom_prompt = """You are a technical support expert for WIX.
Answer questions based on the knowledge base documents provided.
Be concise and helpful."""
answer, citations = await client.ask(
"How do I create a blog?",
system_prompt=custom_prompt
)
```
Filter to specific documents:
```python
@ -443,11 +430,11 @@ answer, citations = await client.ask(
)
```
The QA agent searches your documents for relevant information and uses the configured LLM to generate an answer. The method returns a tuple of `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references.
`client.ask` runs the [rag skill](skills/index.md) under the hood and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, and document references.
The QA provider and model are configured in `haiku.rag.yaml` or can be passed directly to the client (see [Configuration](configuration/index.md)).
See also: [Agents](agents/index.md) for details on the QA agent and the multi‑agent research workflow.
See also: [Skills](skills/index.md) for details on the skills the client wraps.
## Analysis
@ -456,25 +443,20 @@ Answer complex analytical questions via code execution:
```python
# Aggregation across documents
result = await client.analyze("Which quarter had the highest revenue?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
print(result.answer)
for citation in result.citations:
print(citation.uri, citation.title)
# Computation within a document set
result = await client.analyze(
"What is the average deal size mentioned in these contracts?",
filter="uri LIKE '%contracts%'"
)
# Multi-document comparison
result = await client.analyze(
"What changed between these two versions of the policy?",
documents=["Policy v1.0", "Policy v2.0"]
)
```
The analysis agent writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
`client.analyze` runs the [analysis skill](skills/index.md), which writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
See [Analysis](agents/analysis.md) for details on capabilities and configuration.
## Building Custom Agents

2
docs/tools.md
View file

@ -6,7 +6,7 @@ For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset`
## Low-Level Toolsets
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are used internally by the QA agent and can be composed into custom agents.
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are the same primitives the rag and rag-analysis skills compose, and can be reused to build custom agents.
### RAGDeps Protocol

43
docs/tuning.md
View file

@ -20,11 +20,11 @@ Larger embedding models produce better representations at the cost of slower ind
### Reranking
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision — on the Wix benchmark, adding `mxbai-rerank-base-v2` raised MAP from 0.34 to 0.39 on HTML content. See [Search Settings](configuration/qa-research.md#search-settings) for how reranking integrates with search.
When configured, a cross-encoder reranker re-scores 10x the requested candidates and returns the top results. This adds latency but improves precision — on the Wix benchmark, adding `mxbai-rerank-base-v2` raised MAP from 0.34 to 0.39 on HTML content. See [Search Settings](configuration/qa.md#search-settings) for how reranking integrates with search.
### Search Settings
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa-research.md#search-settings).
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa.md#search-settings).
Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
@ -32,9 +32,7 @@ Context expansion is automatic and section-aware — search results are expanded
Model and temperature selection affect answer quality directly — see [Providers](configuration/providers.md#model-settings) for options.
`domain_preamble` prepends domain context to all agent prompts — including the main agent, skill subagents, and internal agents (QA, research). Use it to describe what the knowledge base contains and clarify domain-specific terminology. For full prompt replacement, set `prompts.qa` directly. See [Prompt Customization](configuration/prompts.md).
For automated prompt optimization, see [Prompt Optimization (GEPA)](#prompt-optimization-gepa) below.
`domain_preamble` prepends domain context to the rag and rag-analysis skill instructions. Use it to describe what the knowledge base contains and clarify domain-specific terminology. See [Prompt Customization](configuration/prompts.md).
## What Requires a Rebuild
@ -66,38 +64,3 @@ evaluations run <dataset> --limit 50
```
See [Benchmarks](benchmarks.md) for dataset details, methodology, and baseline results.
## Prompt Optimization (GEPA)
The `evaluations optimize` command uses GEPA (Generalized Evolutionary Prompt Algorithm) to evolve the QA system prompt. It evaluates candidates on minibatches scored by an LLM judge, reflects on failures, proposes mutations, and accepts improvements.
```bash
# Basic optimization
evaluations optimize wix
# Constrained run
evaluations optimize repliqa --limit 40 --num-candidates 30
# Save result
evaluations optimize wix --output optimized_prompt.txt
```
| Option | Default | Description |
|--------|---------|-------------|
| `--limit` | all cases | QA cases to use (split 50/50 train/val) |
| `--num-candidates` | `50` | Number of candidate prompts to evaluate |
| `--output` | — | Save optimized prompt to file |
| `--config` | auto | haiku.rag YAML config path |
| `--db` | auto | Database path override |
| `--judge-model` | `config.qa.model` | LLM judge as `provider:name` |
| `--reflect-model` | `config.qa.model` | Reflection LLM as `provider:name` |
Apply the result in your config:
```yaml
prompts:
qa: |
Your optimized prompt text here...
```
Or programmatically: `get_qa_agent(client, config, system_prompt=optimized_prompt)`.

2
docs/tutorial.md
View file

@ -206,6 +206,6 @@ The following people are presenting talks at PyCon Finland 2025:
- **[Chat](apps.md#chat-tui)** - Interactive conversations with `haiku-rag chat`
- **[CLI Reference](cli.md)** - All available commands and options
- **[Python API](python.md)** - Use haiku.rag in your Python applications
- **[Agents](agents/index.md)** - Deep QA and multi-agent research workflows
- **[Skills](skills/index.md)** - The RAG and analysis skills the client wraps
- **[Configuration](configuration/index.md)** - Complete YAML configuration reference
- **[Server Mode](server.md)** - File monitoring and MCP server

28
evaluations/README.md
View file

@ -6,7 +6,7 @@ This package is not published to PyPI and is only used for development and testi
## Overview
Contains evaluation scripts for benchmarking RAG retrieval and QA performance, plus GEPA-based prompt optimization. Available datasets:
Contains evaluation scripts for benchmarking RAG retrieval and QA performance. Available datasets:
- RepliQA (`repliqa`)
- WiX (`wix`)
@ -41,10 +41,11 @@ evaluations run repliqa --skip-qa
evaluations run repliqa --limit 100
```
### Benchmarking the skills
### Choosing the target
By default `evaluations run` benchmarks the QA agent. Pass `--target` to
benchmark the RAG or analysis skill instead, against the same datasets and judge:
`evaluations run` benchmarks `--target rag-skill` by default. Use
`--target analysis-skill` to benchmark the analysis skill against the same
datasets and judge:
```bash
evaluations run wix --target rag-skill
@ -52,9 +53,10 @@ evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
```
`--skill-model "provider:name"` overrides the skill model independently from
the judge (defaults to `qa.model`). For skill targets, a citation retrieval
metric (`cited_mrr` / `cited_map`) is computed alongside QA accuracy from the
URIs the skill registered via the `cite` tool.
the judge (defaults to `qa.model`, or `analysis.model` when set for the
analysis-skill target). A citation retrieval metric (`cited_mrr` / `cited_map`)
is computed alongside QA accuracy from the URIs the skill registered via the
`cite` tool.
### Pre-built Databases
@ -73,18 +75,6 @@ evaluations upload repliqa
evaluations upload all
```
### Prompt Optimization
Optimize QA system prompts using GEPA (Generalized Evolutionary Prompt Algorithm):
```bash
evaluations optimize wix
evaluations optimize repliqa --limit 40 --num-candidates 30
evaluations optimize wix --output optimized_prompt.txt
```
See [Tuning docs](https://ggozad.github.io/haiku.rag/tuning/#prompt-optimization-gepa) for details on applying results.
## Database Storage
By default, evaluation databases are stored in the haiku.rag data directory:

137
evaluations/evaluations/benchmark.py
View file

@ -28,7 +28,6 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
from haiku.rag.config.models import ModelConfig
from haiku.rag.logging import configure_cli_logging
from haiku.rag.agents.qa import get_qa_agent
from haiku.rag.utils import get_model, parse_model_option
_CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
@ -36,8 +35,8 @@ _CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
MAPEvaluator: CitationMAPEvaluator,
}
Target = Literal["qa", "rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("qa", "rag-skill", "analysis-skill")
Target = Literal["rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
# their QA model does not inadvertently change the judge — keeps cross-run
@ -59,7 +58,7 @@ def build_experiment_metadata(
test_cases: int,
config: AppConfig,
judge_config: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_config: ModelConfig | None = None,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
@ -350,7 +349,7 @@ async def run_qa_benchmark(
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
@ -363,13 +362,16 @@ async def run_qa_benchmark(
]
judge_config = judge_model or DEFAULT_JUDGE_MODEL
skill_config = (skill_model or config.qa.model) if target != "qa" else None
if target == "analysis-skill":
# Mirror the skill-code resolver: explicit analysis.model wins,
# else fall back to qa.model.
skill_config = skill_model or config.analysis.model or config.qa.model
else:
skill_config = skill_model or config.qa.model
db = spec.db_path(db_path)
citation_evaluator: Evaluator | None = None
if target != "qa":
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
evaluators: list[Evaluator] = [
LLMJudge(
@ -409,32 +411,21 @@ async def run_qa_benchmark(
metadata=experiment_metadata,
)
if target == "qa":
async with HaikuRAG(db, config=config) as rag:
qa = get_qa_agent(rag, config)
skill_factory = _skill_factory_for_target(target)
resolved_skill_model = get_model(skill_config, config)
async def answer_question(question: str) -> str:
answer, _ = await qa.answer(question)
return answer
async def answer_question(question: str) -> str:
result = await run_skill_question(
skill_factory=skill_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
)
set_eval_attribute("cited_uris", result.cited_uris)
return result.answer
report = await _evaluate(answer_question)
else:
skill_factory = _skill_factory_for_target(target)
assert skill_config is not None
resolved_skill_model = get_model(skill_config, config)
async def answer_question(question: str) -> str:
result = await run_skill_question(
skill_factory=skill_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
)
set_eval_attribute("cited_uris", result.cited_uris)
return result.answer
report = await _evaluate(answer_question)
report = await _evaluate(answer_question)
passing_cases = sum(
1
@ -498,7 +489,7 @@ async def evaluate_dataset(
vacuum_interval: int = 100,
multimodal_only: bool = False,
judge_model: ModelConfig | None = None,
target: Target = "qa",
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
) -> None:
if not skip_db:
@ -600,22 +591,17 @@ def run(
"--multimodal-only",
help="Only evaluate queries requiring image understanding.",
),
judge_model: str | None = typer.Option(
None,
"--judge-model",
help="Judge model as 'provider:name'. Defaults to ollama:qwen3.6.",
),
target: str = typer.Option(
"qa",
"rag-skill",
"--target",
help="What to benchmark: qa | rag-skill | analysis-skill.",
help="What to benchmark: rag-skill | analysis-skill.",
),
skill_model: str | None = typer.Option(
None,
"--skill-model",
help=(
"Skill model as 'provider:name'. Used when --target is rag-skill or "
"analysis-skill. Defaults to qa.model from the config."
"Skill model as 'provider:name'. Defaults to qa.model (or "
"analysis.model when --target is analysis-skill) from the config."
),
),
) -> None:
@ -626,12 +612,8 @@ def run(
f"Unknown target {target!r}. Choose from: {', '.join(TARGETS)}"
)
target_value = cast(Target, target)
judge_model_config = parse_model_option(judge_model) if judge_model else None
judge_model_config = app_config.evaluations.judge
skill_model_config = parse_model_option(skill_model) if skill_model else None
if target_value == "qa" and skill_model_config is not None:
raise typer.BadParameter(
"--skill-model is only valid when --target is rag-skill or analysis-skill."
)
asyncio.run(
evaluate_dataset(
@ -652,63 +634,6 @@ def run(
)
@app.command()
def optimize(
dataset: str = typer.Argument(..., help="Dataset key to optimize prompt for."),
config: Path | None = typer.Option(
None, "--config", help="Path to haiku.rag YAML config file."
),
db: Path | None = typer.Option(None, "--db", help="Override the database path."),
limit: int | None = typer.Option(
None, "--limit", help="Limit QA cases (split 50/50 into train/val)."
),
num_candidates: int = typer.Option(
50, "--num-candidates", help="Number of candidate prompts to evaluate."
),
output: Path | None = typer.Option(
None, "--output", help="Save optimized prompt to file."
),
judge_model: str | None = typer.Option(
None,
"--judge-model",
help="Judge model as 'provider:name'. Defaults to ollama:qwen3.6.",
),
reflect_model: str | None = typer.Option(
None,
"--reflect-model",
help="Reflect model as 'provider:name' (e.g. 'anthropic:claude-sonnet-4-20250514').",
),
) -> None:
"""Optimize QA system prompt using GEPA evolutionary optimization."""
from evaluations.optimization import run_optimization
spec = _resolve_dataset(dataset)
app_config = _load_config(config)
corpus = spec.qa_loader()
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
cases: list[Case[str, str, dict[str, str]]] = [
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
for index, doc in enumerate(corpus, start=1)
]
judge_model_config = parse_model_option(judge_model) if judge_model else None
reflect_model_config = parse_model_option(reflect_model) if reflect_model else None
run_optimization(
spec=spec,
config=app_config,
cases=cases,
num_candidates=num_candidates,
db_path=db,
output=output,
judge_model=judge_model_config,
reflect_model=reflect_model_config,
)
@app.command()
def download(
dataset: str = typer.Argument(..., help="Dataset key or 'all' to download all."),

292
evaluations/evaluations/optimization.py
View file

@ -1,292 +0,0 @@
import asyncio
import logging
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from pydantic_ai.models import Model
from pydantic_evals import Case
from pydantic_evals.evaluators.llm_as_a_judge import judge_input_output_expected
from gepa.core.adapter import EvaluationBatch
from evaluations.config import DatasetSpec
from haiku.rag.agents.qa import QuestionAnswerAgent, get_qa_agent
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.utils import get_model
logger = logging.getLogger(__name__)
OPTIMIZATION_SCORING_RUBRIC = """You are evaluating the quality of an answer to a question,
comparing it against a reference answer.
Score on a scale of 0.0 to 1.0:
- 1.0: The answer is factually correct, complete, and concise. It covers all key points
from the reference answer without contradictions or significant omissions.
- 0.7-0.9: The answer is mostly correct and addresses the core question, but may miss
some secondary details or include minor inaccuracies.
- 0.4-0.6: The answer is partially correct — it addresses some aspects of the question
but misses key information or contains notable inaccuracies.
- 0.1-0.3: The answer is mostly incorrect or fails to address the core question,
though it may contain some tangentially relevant information.
- 0.0: The answer is completely wrong, irrelevant, or empty.
GUIDELINES:
- Focus on factual correctness relative to the reference answer
- Ignore differences in phrasing, style, or formatting
- A concise correct answer scores higher than a verbose partially correct one
- "I cannot find enough information" when the reference has an answer scores 0.0
"""
@dataclass
class EvalTrajectory:
"""Per-case evaluation result for GEPA reflection."""
question: str
expected_answer: str
actual_answer: str | None
score: float
judge_reason: str | None = None
QACase = Case[str, str, dict[str, str]]
@dataclass
class QAPromptAdapter:
"""GEPA adapter that evaluates QA prompt candidates against a dataset.
Implements the GEPAAdapter protocol:
- evaluate(): Run QA agent with candidate prompt, score with LLMJudge
- make_reflective_dataset(): Build failure records for the GEPA proposer
"""
config: AppConfig
db_path: Path
judge_model: Model
def evaluate(
self,
batch: list[QACase],
candidate: dict[str, str],
capture_traces: bool = False,
) -> EvaluationBatch[EvalTrajectory, str | None]:
instructions = candidate["instructions"]
return asyncio.run(
self._evaluate_with_setup(batch, instructions, capture_traces)
)
async def _evaluate_with_setup(
self,
batch: list[QACase],
instructions: str,
capture_traces: bool,
) -> EvaluationBatch[EvalTrajectory, str | None]:
async with HaikuRAG(self.db_path, config=self.config) as rag:
qa = get_qa_agent(rag, self.config, system_prompt=instructions)
return await self._evaluate_async(batch, qa, capture_traces)
async def _evaluate_async(
self,
batch: list[QACase],
qa: QuestionAnswerAgent,
capture_traces: bool,
) -> EvaluationBatch[EvalTrajectory, str | None]:
outputs: list[str | None] = []
scores: list[float] = []
trajectories: list[EvalTrajectory] | None = [] if capture_traces else None
for case in batch:
question = case.inputs
expected = case.expected_output or ""
try:
answer, _ = await qa.answer(question)
except Exception:
logger.warning(
"QA agent failed for question: %s", question, exc_info=True
)
answer = None
if answer is not None:
score, reason = await self._judge(question, answer, expected)
else:
score, reason = 0.0, "QA agent failed to produce an answer"
outputs.append(answer)
scores.append(score)
if capture_traces and trajectories is not None:
trajectories.append(
EvalTrajectory(
question=question,
expected_answer=expected,
actual_answer=answer,
score=score,
judge_reason=reason,
)
)
return EvaluationBatch(
outputs=outputs,
scores=scores,
trajectories=trajectories,
)
async def _judge(
self, question: str, answer: str, expected: str
) -> tuple[float, str | None]:
"""Score an answer using pydantic-evals LLMJudge with float scoring."""
result = await judge_input_output_expected(
inputs=question,
output=answer,
expected_output=expected,
rubric=OPTIMIZATION_SCORING_RUBRIC,
model=self.judge_model,
)
return result.score, result.reason
def make_reflective_dataset(
self,
candidate: dict[str, str],
eval_batch: EvaluationBatch[EvalTrajectory, str | None],
components_to_update: list[str],
) -> Mapping[str, Sequence[Mapping[str, Any]]]:
if eval_batch.trajectories is None:
return {}
records: list[dict[str, Any]] = []
for traj in eval_batch.trajectories:
records.append(
{
"Inputs": {"question": traj.question},
"Generated Outputs": {
"answer": traj.actual_answer or "(no answer)"
},
"Feedback": (
f"Expected answer: {traj.expected_answer}\n"
f"Score: {traj.score:.2f}\n"
f"Judge reasoning: {traj.judge_reason or 'N/A'}"
),
}
)
return {"instructions": records}
propose_new_texts = None
class ReflectionLM:
"""LanguageModel implementation for GEPA's ReflectiveMutationProposer.
Wraps a pydantic-ai Agent to satisfy GEPA's LanguageModel protocol.
"""
def __init__(self, model_config: ModelConfig, config: AppConfig) -> None:
from pydantic_ai import Agent
model = get_model(model_config, config)
self._agent: Agent[None, str] = Agent(model=model, output_type=str)
def __call__(self, prompt: str | list[dict[str, Any]]) -> str:
if isinstance(prompt, list):
text = "\n".join(
f"{msg.get('role', 'user')}: {msg.get('content', '')}" for msg in prompt
)
else:
text = prompt
result = self._agent.run_sync(text)
return result.output
# Cases per GEPA reflection minibatch (used for budget calculation)
REFLECTION_MINIBATCH_SIZE = 3
def run_optimization(
spec: DatasetSpec,
config: AppConfig,
cases: list[QACase],
num_candidates: int,
db_path: Path | None = None,
output: Path | None = None,
judge_model: ModelConfig | None = None,
reflect_model: ModelConfig | None = None,
) -> dict[str, Any]:
"""Run GEPA optimization and return results summary."""
from rich.console import Console
console = Console()
judge_config = judge_model or config.qa.model
judge = get_model(judge_config, config)
db = spec.db_path(db_path)
adapter = QAPromptAdapter(
config=config,
db_path=db,
judge_model=judge,
)
reflect_config = reflect_model or config.qa.model
reflection_lm = ReflectionLM(reflect_config, config)
seed_prompt = config.prompts.qa or QA_SYSTEM_PROMPT
seed_candidate = {"instructions": seed_prompt}
mid = len(cases) // 2
trainset = cases[:mid]
valset = cases[mid:]
# Budget: initial valset eval + per-candidate worst case
# (each candidate: 2 minibatch evals + full valset if accepted)
max_metric_calls = len(valset) + num_candidates * (
2 * REFLECTION_MINIBATCH_SIZE + len(valset)
)
console.print(f"Optimizing prompt for dataset: {spec.key}", style="bold magenta")
console.print(
f"Train: {len(trainset)}, Val: {len(valset)}, "
f"Candidates: {num_candidates}, Budget: {max_metric_calls} eval calls"
)
console.print(f"Seed prompt length: {len(seed_prompt)} chars")
from gepa import optimize as gepa_optimize
result = gepa_optimize(
seed_candidate=seed_candidate,
trainset=trainset,
valset=valset,
adapter=adapter,
reflection_lm=reflection_lm,
max_metric_calls=max_metric_calls,
display_progress_bar=True,
)
best_score = result.val_aggregate_scores[result.best_idx]
best_prompt = result.best_candidate
if isinstance(best_prompt, dict):
best_prompt = best_prompt["instructions"]
total_calls = result.total_metric_calls or "unknown"
console.print("\n=== Optimization Results ===", style="bold cyan")
console.print(f"Total metric calls: {total_calls}")
console.print(f"Candidates explored: {result.num_candidates}")
console.print(f"Best score: {best_score:.4f}")
console.print(f"\nOptimized prompt:\n{best_prompt}")
if output:
output.write_text(best_prompt)
console.print(f"\nSaved to: {output}", style="green")
return {
"best_score": best_score,
"best_prompt": best_prompt,
"total_calls": total_calls,
"num_candidates": result.num_candidates,
}

2
evaluations/evaluations/skill_runner.py
View file

@ -5,7 +5,7 @@ from typing import Protocol, cast
from pydantic_ai.models import Model
from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.citation import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.skills import run_skill

1
evaluations/pyproject.toml
View file

@ -14,7 +14,6 @@ dependencies = [
"huggingface_hub>=0.20.0",
"typer>=0.21.0,<0.22.0",
"python-dotenv>=1.2.2",
"gepa>=0.1.0",
]
[project.scripts]

20
evaluations/tests/test_benchmark.py
View file

@ -131,8 +131,7 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.HaikuRAG"),
patch("evaluations.benchmark.get_qa_agent"),
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -142,7 +141,7 @@ class TestRunQaBenchmarkJudgeModel:
judge_model=custom_judge,
)
mock_get_model.assert_called_once_with(custom_judge, AppConfig())
mock_get_model.assert_any_call(custom_judge, AppConfig())
@pytest.mark.asyncio
async def test_defaults_to_pinned_judge_model(self, tmp_path: Path) -> None:
@ -150,8 +149,7 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.HaikuRAG"),
patch("evaluations.benchmark.get_qa_agent"),
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -160,7 +158,7 @@ class TestRunQaBenchmarkJudgeModel:
db_path=tmp_path / "test.lancedb",
)
mock_get_model.assert_called_once_with(DEFAULT_JUDGE_MODEL, AppConfig())
mock_get_model.assert_any_call(DEFAULT_JUDGE_MODEL, AppConfig())
class TestEvaluateDatasetJudgeModel:
@ -197,11 +195,11 @@ class TestEvaluateDatasetJudgeModel:
class TestExperimentMetadataTargets:
def test_default_target_is_qa(self) -> None:
def test_default_target_is_rag_skill(self) -> None:
result = build_experiment_metadata(
dataset_key="test", test_cases=1, config=AppConfig()
)
assert result["target"] == "qa"
assert result["target"] == "rag-skill"
assert "skill_provider" not in result
assert "skill_model" not in result
@ -255,7 +253,7 @@ class TestEvaluateDatasetTarget:
assert mock_qa.call_args[1]["skill_model"] is skill
@pytest.mark.asyncio
async def test_default_target_is_qa(self) -> None:
async def test_default_target_is_rag_skill(self) -> None:
with patch(
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
) as mock_qa:
@ -269,7 +267,7 @@ class TestEvaluateDatasetTarget:
name=None,
db_path=None,
)
assert mock_qa.call_args[1]["target"] == "qa"
assert mock_qa.call_args[1]["target"] == "rag-skill"
assert mock_qa.call_args[1]["skill_model"] is None
@ -323,7 +321,7 @@ class TestRunQaBenchmarkSkillTarget:
assert _skill_factory_for_target("rag-skill") is rag_factory
assert _skill_factory_for_target("analysis-skill") is analysis_factory
with pytest.raises(ValueError, match="not a skill target"):
_skill_factory_for_target("qa") # type: ignore[arg-type]
_skill_factory_for_target("unknown") # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
class TestCitationEvaluatorWiring:

462
evaluations/tests/test_optimization.py
View file

@ -1,462 +0,0 @@
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic_ai.models.test import TestModel
from pydantic_evals import Case
from gepa.core.adapter import EvaluationBatch
from evaluations.config import DatasetSpec
from evaluations.optimization import (
EvalTrajectory,
QAPromptAdapter,
ReflectionLM,
run_optimization,
)
from haiku.rag.config.models import AppConfig, ModelConfig
@pytest.fixture
def sample_cases() -> list[Case[str, str, dict[str, str]]]:
return [
Case(
name="q1",
inputs="What is X?",
expected_output="X is a thing.",
metadata={"case_index": "1"},
),
Case(
name="q2",
inputs="How does Y work?",
expected_output="Y works by Z.",
metadata={"case_index": "2"},
),
]
@pytest.fixture
def adapter(tmp_path: Path) -> QAPromptAdapter:
return QAPromptAdapter(
config=AppConfig(),
db_path=tmp_path / "test.lancedb",
judge_model=MagicMock(),
)
class TestMakeReflectiveDataset:
def test_builds_records_from_trajectories(self, adapter: QAPromptAdapter) -> None:
trajectories = [
EvalTrajectory(
question="What is X?",
expected_answer="X is a thing.",
actual_answer="X is wrong.",
score=0.2,
judge_reason="Factually incorrect",
),
EvalTrajectory(
question="How does Y?",
expected_answer="Y works by Z.",
actual_answer="Y works by Z.",
score=1.0,
judge_reason=None,
),
]
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=["X is wrong.", "Y works by Z."],
scores=[0.2, 1.0],
trajectories=trajectories,
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert "instructions" in result
records = result["instructions"]
assert len(records) == 2
assert records[0]["Inputs"]["question"] == "What is X?"
assert records[0]["Generated Outputs"]["answer"] == "X is wrong."
assert "Expected answer: X is a thing." in records[0]["Feedback"]
assert "Score: 0.20" in records[0]["Feedback"]
assert "Factually incorrect" in records[0]["Feedback"]
assert records[1]["Inputs"]["question"] == "How does Y?"
assert records[1]["Generated Outputs"]["answer"] == "Y works by Z."
assert "Score: 1.00" in records[1]["Feedback"]
assert "N/A" in records[1]["Feedback"]
def test_returns_empty_when_no_trajectories(self, adapter: QAPromptAdapter) -> None:
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=[], scores=[], trajectories=None
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert result == {}
def test_none_answer_becomes_no_answer(self, adapter: QAPromptAdapter) -> None:
trajectories = [
EvalTrajectory(
question="What?",
expected_answer="Answer.",
actual_answer=None,
score=0.0,
judge_reason="Failed",
),
]
eval_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=[None],
scores=[0.0],
trajectories=trajectories,
)
result = adapter.make_reflective_dataset(
{"instructions": "test"}, eval_batch, ["instructions"]
)
assert result["instructions"][0]["Generated Outputs"]["answer"] == "(no answer)"
class TestEvaluateAsync:
@pytest.mark.asyncio
async def test_returns_scores_and_outputs(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
stub_qa = AsyncMock()
stub_qa.answer = AsyncMock(return_value=("X is a thing.", []))
adapter._judge = AsyncMock(return_value=(0.85, "Good answer")) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
result = await adapter._evaluate_async(
sample_cases, stub_qa, capture_traces=False
)
assert len(result.outputs) == 2
assert len(result.scores) == 2
assert all(o == "X is a thing." for o in result.outputs)
assert all(s == 0.85 for s in result.scores)
assert result.trajectories is None
@pytest.mark.asyncio
async def test_populates_trajectories_when_captured(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
stub_qa = AsyncMock()
stub_qa.answer = AsyncMock(return_value=("An answer.", []))
adapter._judge = AsyncMock(return_value=(0.9, "Almost perfect")) # type: ignore[method-assign] # ty: ignore[invalid-assignment]
result = await adapter._evaluate_async(
sample_cases, stub_qa, capture_traces=True
)
assert result.trajectories is not None
assert len(result.trajectories) == 2
traj = result.trajectories[0]
assert traj.question == "What is X?"
assert traj.expected_answer == "X is a thing."
assert traj.actual_answer == "An answer."
assert traj.score == 0.9
assert traj.judge_reason == "Almost perfect"
@pytest.mark.asyncio
async def test_handles_qa_failure(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
stub_qa = AsyncMock()
stub_qa.answer = AsyncMock(side_effect=RuntimeError("LLM down"))
result = await adapter._evaluate_async(
sample_cases, stub_qa, capture_traces=True
)
assert all(o is None for o in result.outputs)
assert all(s == 0.0 for s in result.scores)
assert result.trajectories is not None
assert all(t.actual_answer is None for t in result.trajectories)
assert all(
t.judge_reason == "QA agent failed to produce an answer"
for t in result.trajectories
)
class TestReflectionLM:
def test_handles_string_prompt(self) -> None:
test_model = TestModel(custom_output_text="Reflected response")
with patch("evaluations.optimization.get_model", return_value=test_model):
lm = ReflectionLM(model_config=AppConfig().qa.model, config=AppConfig())
result = lm("test prompt")
assert result == "Reflected response"
def test_formats_chat_messages_into_string(self) -> None:
test_model = TestModel(custom_output_text="Chat response")
prompts_received: list[str] = []
with patch("evaluations.optimization.get_model", return_value=test_model):
lm = ReflectionLM(model_config=AppConfig().qa.model, config=AppConfig())
original_run_sync = lm._agent.run_sync
def capturing_run_sync(prompt: str, **kwargs: Any) -> Any:
prompts_received.append(prompt)
return original_run_sync(prompt, **kwargs)
lm._agent.run_sync = capturing_run_sync # type: ignore[method-assign] # ty: ignore[invalid-assignment]
messages: list[dict[str, Any]] = [
{"role": "system", "content": "You are helpful."},
{"role": "user", "content": "Hello"},
]
result = lm(messages)
assert result == "Chat response"
assert len(prompts_received) == 1
assert "system: You are helpful." in prompts_received[0]
assert "user: Hello" in prompts_received[0]
class TestEvaluateSync:
def test_delegates_to_evaluate_async(
self,
adapter: QAPromptAdapter,
sample_cases: list[Case[str, str, dict[str, str]]],
) -> None:
expected_batch: EvaluationBatch[EvalTrajectory, str | None] = EvaluationBatch(
outputs=["answer1", "answer2"],
scores=[0.9, 0.8],
trajectories=None,
)
with patch.object(
adapter,
"_evaluate_with_setup",
new_callable=AsyncMock,
return_value=expected_batch,
) as mock_eval:
result = adapter.evaluate(
sample_cases, {"instructions": "my prompt"}, capture_traces=True
)
mock_eval.assert_called_once_with(sample_cases, "my prompt", True)
assert result is expected_batch
class TestProposalAttribute:
def test_propose_new_texts_is_none(self, adapter: QAPromptAdapter) -> None:
assert adapter.propose_new_texts is None
def _make_cases(n: int) -> list[Case[str, str, dict[str, str]]]:
return [
Case(
name=f"q{i}",
inputs=f"Question {i}?",
expected_output=f"Answer {i}.",
metadata={"case_index": str(i)},
)
for i in range(1, n + 1)
]
@pytest.fixture
def gepa_mock_result() -> MagicMock:
mock_result = MagicMock()
mock_result.best_idx = 0
mock_result.val_aggregate_scores = [0.95]
mock_result.best_candidate = {"instructions": "optimized prompt"}
mock_result.total_metric_calls = 10
mock_result.num_candidates = 3
return mock_result
class TestRunOptimization:
def _make_spec(self, db_path: Path) -> DatasetSpec:
return DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
def test_returns_results(self, tmp_path: Path, gepa_mock_result: MagicMock) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=gepa_mock_result),
):
result = run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
)
assert result["best_score"] == 0.95
assert result["best_prompt"] == "optimized prompt"
assert result["total_calls"] == 10
assert result["num_candidates"] == 3
def test_saves_output_file(
self, tmp_path: Path, gepa_mock_result: MagicMock
) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
output_path = tmp_path / "prompt.txt"
gepa_mock_result.val_aggregate_scores = [0.85]
gepa_mock_result.best_candidate = {"instructions": "saved prompt"}
gepa_mock_result.total_metric_calls = 5
gepa_mock_result.num_candidates = 2
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=5,
db_path=tmp_path / "test.lancedb",
output=output_path,
)
assert output_path.read_text() == "saved prompt"
def test_uses_default_prompt_when_spec_has_none(self, tmp_path: Path) -> None:
spec = DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
cases = _make_cases(4)
mock_result = MagicMock()
mock_result.best_idx = 0
mock_result.val_aggregate_scores = [0.5]
mock_result.best_candidate = "fallback prompt"
mock_result.total_metric_calls = 1
mock_result.num_candidates = 1
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=mock_result) as mock_gepa,
):
result = run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=1,
db_path=tmp_path / "test.lancedb",
)
# When best_candidate is a string (not dict), it should be used directly
assert result["best_prompt"] == "fallback prompt"
# Verify seed_candidate used QA_SYSTEM_PROMPT (not None)
call_kwargs = mock_gepa.call_args[1]
seed = call_kwargs["seed_candidate"]
assert seed["instructions"] is not None
assert len(seed["instructions"]) > 0
def test_splits_cases_into_train_and_val(self, tmp_path: Path) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(10)
mock_result = MagicMock()
mock_result.best_idx = 0
mock_result.val_aggregate_scores = [0.7]
mock_result.best_candidate = {"instructions": "prompt"}
mock_result.total_metric_calls = 50
mock_result.num_candidates = 1
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=mock_result) as mock_gepa,
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=5,
db_path=tmp_path / "test.lancedb",
)
call_kwargs = mock_gepa.call_args[1]
assert len(call_kwargs["trainset"]) == 5
assert len(call_kwargs["valset"]) == 5
# Budget = valset_size + num_candidates * (2*minibatch + valset_size)
assert call_kwargs["max_metric_calls"] == 5 + 5 * (2 * 3 + 5)
def test_uses_custom_reflect_model(
self, tmp_path: Path, gepa_mock_result: MagicMock
) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
reflect_model = ModelConfig(
provider="anthropic", name="claude-sonnet-4-20250514"
)
with (
patch("evaluations.optimization.get_model"),
patch("evaluations.optimization.ReflectionLM") as mock_rlm,
patch("gepa.optimize", return_value=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
reflect_model=reflect_model,
)
mock_rlm.assert_called_once_with(reflect_model, AppConfig())
def test_uses_custom_judge_model(
self, tmp_path: Path, gepa_mock_result: MagicMock
) -> None:
spec = self._make_spec(tmp_path / "test.lancedb")
cases = _make_cases(4)
judge_model = ModelConfig(provider="openai", name="gpt-4o")
with (
patch("evaluations.optimization.get_model") as mock_get_model,
patch("evaluations.optimization.ReflectionLM"),
patch("gepa.optimize", return_value=gepa_mock_result),
):
run_optimization(
spec=spec,
config=AppConfig(),
cases=cases,
num_candidates=10,
db_path=tmp_path / "test.lancedb",
judge_model=judge_model,
)
mock_get_model.assert_called_once_with(judge_model, AppConfig())

4
evaluations/tests/test_skill_runner.py
View file

@ -6,7 +6,7 @@ import pytest
from pydantic_ai.models.test import TestModel
from evaluations.skill_runner import SkillRunResult, run_skill_question
from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.citation import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.embeddings import EmbedderWrapper
@ -296,7 +296,7 @@ class TestRunSkillQuestionEndToEnd:
db_path=rag_db,
config=app_config,
question="What is machine learning?",
skill_model=TestModel(),
skill_model=TestModel(call_tools=["search"]),
)
assert isinstance(result, SkillRunResult)

21
haiku_rag_slim/haiku/rag/agents/analysis/__init__.py
View file

@ -1,21 +0,0 @@
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,
RawAnalysisResult,
)
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
__all__ = [
"ANALYSIS_SYSTEM_PROMPT",
"AnalysisContext",
"AnalysisDeps",
"RawAnalysisResult",
"AnalysisResult",
"CodeExecution",
"Sandbox",
"SandboxResult",
"create_analysis_agent",
]

60
haiku_rag_slim/haiku/rag/agents/analysis/agent.py
View file

@ -1,60 +0,0 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
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, RawAnalysisResult]:
"""Create an analysis agent with code execution capability.
The analysis agent can write and execute Python code in a sandboxed
environment to solve problems that require computation, aggregation,
or complex traversal across documents.
Args:
config: Application configuration.
Returns:
A pydantic-ai Agent configured for analysis execution.
"""
model = get_model(config.analysis.model, config)
agent: Agent[AnalysisDeps, RawAnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model,
deps_type=AnalysisDeps,
output_type=RawAnalysisResult,
instructions=ANALYSIS_SYSTEM_PROMPT,
tool_retries=3,
output_retries=3,
)
@agent.tool
async def execute_code(ctx: RunContext[AnalysisDeps], code: str) -> CodeExecution:
"""Execute Python code in a sandboxed interpreter.
The code has access to search() and llm() functions, and a
virtual filesystem at /documents/ with document content and structure.
Use print() to output results.
Args:
code: Python code to execute.
Returns:
Structured result with success status, stdout, and stderr.
"""
result = await ctx.deps.sandbox.execute(code)
execution = CodeExecution(
code=code,
stdout=result.stdout,
stderr=result.stderr,
success=result.success,
)
return execution
return agent

23
haiku_rag_slim/haiku/rag/agents/analysis/dependencies.py
View file

@ -1,23 +0,0 @@
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from haiku.rag.store.models import Document
if TYPE_CHECKING:
from haiku.rag.agents.analysis.sandbox import Sandbox
@dataclass
class AnalysisContext:
"""Mutable context accumulating data during analysis execution."""
documents: list[Document] | None = None
filter: str | None = None
@dataclass
class AnalysisDeps:
"""Dependencies for analysis agent."""
sandbox: "Sandbox"
context: AnalysisContext = field(default_factory=AnalysisContext)

27
haiku_rag_slim/haiku/rag/agents/analysis/models.py
View file

@ -1,27 +0,0 @@
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."""
code: str = Field(description="The Python code that was executed")
stdout: str = Field(description="Standard output captured during execution")
stderr: str = Field(description="Standard error captured during execution")
success: bool = Field(description="Whether execution completed without error")
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")
class AnalysisResult(BaseModel):
"""Result from analysis execution with resolved citations."""
answer: str
program: str
citations: list[Citation] = Field(default_factory=list)

127
haiku_rag_slim/haiku/rag/agents/analysis/prompts.py
View file

@ -1,127 +0,0 @@
ANALYSIS_SYSTEM_PROMPT = """You are an analysis agent that solves complex research questions by writing and executing Python code.
You MUST use the `execute_code` tool to run Python code. The functions and filesystem described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
## Available Functions
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
- results = await search("query") ✓ CORRECT
- import search ✗ WRONG - will fail
- results = search("query") ✗ WRONG - must use await
### await search(query, limit=10) -> list[dict]
Search the knowledge base using hybrid search (vector + full-text).
Results are automatically expanded with surrounding context (adjacent paragraphs, complete tables, section content).
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
### await list_documents() -> list[dict]
List all documents in the knowledge base.
Returns list of dicts with keys: id, title, uri, created_at
### await llm(prompt) -> str
Call an LLM directly with the given prompt. Returns the response as a string.
Use this for classification, summarization, extraction, or any task where you
already have the content and just need LLM reasoning.
## Document Filesystem
All documents in the knowledge base are available as files under `/documents/`. Use `from pathlib import Path` and standard file I/O to access them.
### Directory structure
```
/documents/
{document_id}/
metadata.json # {"id", "title", "uri", "created_at"}
content.txt # Full document text
items.jsonl # Structured document items (one JSON object per line)
```
### metadata.json
Small file with document metadata. Use to discover and identify documents.
```python
from pathlib import Path
import json
for doc_dir in Path('/documents').iterdir():
meta = json.loads((doc_dir / 'metadata.json').read_text())
print(meta['title'], meta['uri'])
```
### content.txt
Full text content of the document. Use for regex, keyword search, or full-text analysis.
```python
content = Path(f'/documents/{doc_id}/content.txt').read_text()
```
### items.jsonl
Structured document items as JSONL. Each line is a JSON object with:
- `position`: sequential position in the document
- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote", etc.
- `text`: rendered content (tables are markdown with `|` columns)
- `page_numbers`: list of page numbers where the item appears
Use items.jsonl to find tables, section headers, or specific structural elements:
```python
import json
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
for line in items_text.strip().split(chr(10)):
item = json.loads(line)
if item['label'] == 'table':
print(f"Table on page {item['page_numbers']}: {item['text'][:100]}")
```
## Cross-referencing search results with items
Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl. Use this to navigate from a search hit to the surrounding document structure:
```python
results = await search("revenue", limit=5)
r = results[0]
doc_id = r['document_id']
refs = set(r['doc_item_refs'])
import json
items_text = Path(f'/documents/{doc_id}/items.jsonl').read_text()
for line in items_text.strip().split(chr(10)):
item = json.loads(line)
if item['self_ref'] in refs:
print(f"Matched: {item['label']} on page {item['page_numbers']}")
```
## Pre-loaded Documents Variable
If documents were pre-loaded for this session, a `documents` variable is available:
```python
# documents is a list of dicts with keys: id, title, uri, content
for doc in documents:
print(doc['title'], len(doc['content']))
```
Check if it exists with: `try: documents ... except NameError: ...`
## Available Python Features
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `filter()`, `getattr()`, `sorted()`/`.sort(key=...)`, try/except, and the `json`, `re`, `math` modules. File I/O via `pathlib.Path` is supported for the `/documents/` filesystem.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements.
## Strategy Guide
1. **Search First**: Start with `search()` to find relevant content. Results include expanded context and `doc_item_refs` for cross-referencing.
2. **Discover Documents**: Use `list_documents()` to see what's in the knowledge base.
3. **Use items.jsonl for Structure**: Find tables, section headers, or specific elements by label and page number. Tables are pre-rendered as markdown.
4. **Use content.txt for Full Text**: When you need the complete document text (e.g., for regex across the whole document).
5. **Iterate**: Run code, examine results, refine your approach. Don't try to solve everything in one execution.
6. **Use llm() for Reasoning**: When you have content and need classification, summarization, or extraction, use `llm()` rather than writing complex parsing logic.
## Output Format
Your final response MUST be valid JSON matching this exact schema:
```json
{"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.
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."""

324
haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py
View file

@ -1,324 +0,0 @@
import asyncio
import atexit
import concurrent.futures
import json
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, 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
@dataclass
class SandboxResult:
"""Result of executing code in the sandbox."""
stdout: str
stderr: str
success: bool
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
atexit.register(_executor.shutdown, wait=False)
def _run_async(coro: Any) -> Any:
"""Run an async coroutine from a sync context (CallbackFile read)."""
return _executor.submit(asyncio.run, coro).result()
class Sandbox:
"""Execute code in a sandboxed Python interpreter.
Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
External functions (search, llm) are called by Monty code using ``await``
and resolved asynchronously on the host.
Documents are exposed via a virtual filesystem at ``/documents/{id}/``.
The interpreter uses a REPL session — variables persist across
``execute()`` calls within the same Sandbox instance.
sandbox = Sandbox(db_path, config, context)
result = await sandbox.execute("x = await search('query')")
result = await sandbox.execute("print(x[0]['content'])") # x persists
"""
_db_path: Path
_config: AppConfig
_context: AnalysisContext
_search_results: "list[SearchResult]"
_items_cache: dict[str, str] | None
_repl: MontyRepl | None
_vfs: OSAccess | None
def __init__(
self,
db_path: Path,
config: AppConfig,
context: AnalysisContext,
):
self._db_path = db_path
self._config = config
self._context = context
self._search_results = []
self._items_cache = None
self._repl = None
self._vfs = None
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
db_path = self._db_path
config = self._config
context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(query, limit=limit, filter=context.filter)
expanded = await rag.expand_context(results)
self._search_results.extend(expanded)
return [
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
"doc_item_refs": r.doc_item_refs,
"labels": r.labels,
}
for r in expanded
]
async def list_documents() -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
docs = await rag.list_documents(filter=context.filter)
return [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"created_at": str(d.created_at),
}
for d in docs
]
async def llm(prompt: str) -> str:
from pydantic_ai import Agent
from haiku.rag.utils import get_model
model = get_model(config.analysis.model, config)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
return {
"search": search,
"list_documents": list_documents,
"llm": llm,
}
async def _build_vfs(self) -> OSAccess:
"""Build the virtual filesystem with document data.
Mounts per-document directories with:
- metadata.json: MemoryFile (eager, small)
- content.txt: CallbackFile (lazy, can be large)
- items.jsonl: CallbackFile (lazy, can be large)
"""
from haiku.rag.client import HaikuRAG
db_path = self._db_path
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)
doc_ids = [doc.id for doc in docs if doc.id]
def _load_items_cache() -> dict[str, str]:
"""Bulk-fetch all document items in one query, serialize to JSONL."""
async def _fetch() -> dict[str, str]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
grouped = await rag.document_item_repository.get_all_items_grouped(
doc_ids
)
result: dict[str, str] = {}
for did, items in grouped.items():
lines = []
for item in items:
lines.append(
json.dumps(
{
"position": item.position,
"self_ref": item.self_ref,
"label": item.label,
"text": item.text,
"page_numbers": item.page_numbers,
},
ensure_ascii=False,
)
)
result[did] = "\n".join(lines)
return result
return _run_async(_fetch())
sandbox = self
def _make_items_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_items(_path: "PurePosixPath") -> str:
if sandbox._items_cache is None:
sandbox._items_cache = _load_items_cache()
return sandbox._items_cache.get(did, "")
return read_items
for doc in docs:
if not doc.id:
continue
doc_id: str = doc.id
doc_dir = f"/documents/{doc_id}"
metadata = json.dumps(
{
"id": doc_id,
"title": doc.title,
"uri": doc.uri,
"created_at": str(doc.created_at),
},
ensure_ascii=False,
)
files.append(MemoryFile(f"{doc_dir}/metadata.json", metadata))
def _make_content_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(
db_path, config=config, read_only=True
) as rag:
content = await rag.document_repository.get_content(did)
return content or ""
return _run_async(_fetch())
return read_content
files.append(
CallbackFile(
f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id),
write=_deny_write,
)
)
files.append(
CallbackFile(
f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id),
write=_deny_write,
)
)
return OSAccess(files)
async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]:
"""Initialize the REPL session and VFS on first use."""
if self._repl is None:
self._vfs = await self._build_vfs()
self._repl = MontyRepl(
limits={
"max_duration_secs": self._config.analysis.code_timeout,
},
)
if self._context.documents:
await self._repl.feed_run_async(
"pass",
inputs={
"documents": [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"content": d.content,
}
for d in self._context.documents
]
},
external_functions=self._build_external_functions(),
os=self._vfs,
)
assert self._repl is not None and self._vfs is not None
return self._repl, self._vfs
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty REPL.
Variables persist across calls within the same Sandbox instance.
"""
repl, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions()
stdout_lines: list[str] = []
def print_callback(_stream: Literal["stdout"], text: str) -> None:
stdout_lines.append(text)
max_chars = self._config.analysis.max_output_chars
try:
output = await repl.feed_run_async(
code,
external_functions=external_fns,
print_callback=print_callback,
os=vfs,
)
except (
pydantic_monty.MontySyntaxError,
pydantic_monty.MontyRuntimeError,
) as e:
stdout = "".join(stdout_lines)
if len(stdout) > max_chars:
stdout = stdout[:max_chars] + "\n... (output truncated)"
return SandboxResult(stdout=stdout, stderr=str(e), success=False)
stdout = "".join(stdout_lines)
if output is not None:
stdout_with_output = f"{stdout}{output}" if stdout else str(output)
else:
stdout_with_output = stdout
if len(stdout_with_output) > max_chars:
stdout_with_output = (
stdout_with_output[:max_chars] + "\n... (output truncated)"
)
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)

35
haiku_rag_slim/haiku/rag/agents/qa/__init__.py
View file

@ -1,35 +0,0 @@
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, Config
from haiku.rag.utils import build_prompt
def get_qa_agent(
client: HaikuRAG,
config: AppConfig = Config,
system_prompt: str | None = None,
) -> QuestionAnswerAgent:
"""Factory function to get a QA agent based on the configuration.
Args:
client: HaikuRAG client instance.
config: Configuration to use. Defaults to global Config.
system_prompt: Optional custom system prompt (overrides config).
Returns:
A configured QuestionAnswerAgent instance.
"""
# Determine the base prompt: explicit > config > default
if system_prompt is None:
system_prompt = config.prompts.qa or QA_SYSTEM_PROMPT
# Prepend system_context if configured
system_prompt = build_prompt(system_prompt, config)
return QuestionAnswerAgent(
client=client,
model_config=config.qa.model,
config=config,
system_prompt=system_prompt,
)

81
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -1,81 +0,0 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.agents.research.models import (
Citation,
RawSearchAnswer,
resolve_citations,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.utils import get_model
@dataclass
class _QARunDeps:
client: HaikuRAG
class QuestionAnswerAgent:
def __init__(
self,
client: HaikuRAG,
model_config: ModelConfig,
config: AppConfig | None = None,
system_prompt: str | None = None,
):
self._client = client
self._config = config or Config
self._model_config = model_config
self._system_prompt = system_prompt or QA_SYSTEM_PROMPT
async def answer(
self, question: str, filter: str | None = None
) -> tuple[str, list[Citation]]:
"""Answer a question using the RAG system.
Args:
question: The question to answer
filter: SQL WHERE clause to filter documents
Returns:
Tuple of (answer text, list of resolved citations)
"""
accumulated_results: list[SearchResult] = []
max_searches = self._config.qa.max_searches
search_toolset = create_search_toolset(
self._config,
base_filter=filter,
tool_name="search",
on_results=accumulated_results.extend,
max_searches=max_searches,
)
# Agent created per-call: toolset varies with filter, and Agent
# construction is pure Python (no IO).
model = get_model(self._model_config, self._config)
try:
system_prompt = self._system_prompt.format(max_searches=max_searches)
except KeyError:
system_prompt = self._system_prompt
agent: Agent[_QARunDeps, RawSearchAnswer] = Agent( # ty: ignore[invalid-assignment]
model=model,
deps_type=_QARunDeps,
output_type=RawSearchAnswer,
instructions=system_prompt,
toolsets=[search_toolset],
tool_retries=3,
output_retries=3,
)
deps = _QARunDeps(client=self._client)
result = await agent.run(question, deps=deps)
output = result.output
citations = resolve_citations(output.cited_chunks, accumulated_results)
return output.answer, citations

41
haiku_rag_slim/haiku/rag/agents/qa/prompts.py
View file

@ -1,41 +0,0 @@
QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base.
Process:
1. Call search with relevant keywords from the question
2. Review the results ordered by relevance
3. If needed, perform follow-up searches with different keywords (max {max_searches} total)
4. Provide a concise answer based strictly on the retrieved content
The search tool returns results like:
[chunk_abc123] [rank 1 of 5]
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[chunk_def456] [rank 2 of 5]
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and rank position (rank 1 = most relevant)
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
When a result is of `Type: picture` (a figure or diagram), the search tool may also attach the picture itself as image content alongside the text — use it for visual reasoning when answering. Reference it by its chunk_id like any other source.
IMPORTANT: You MUST include in cited_chunks the COMPLETE IDs of every chunk you reference. Copy the full ID string without brackets — e.g. "5ae52166-5329-42e9-b6a5-756fc0cb7200" not "[5ae52166]" or "5ae52166". Never truncate IDs. Never leave cited_chunks empty if you found relevant content.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge
- Use the Source and Type metadata to understand context
- If multiple results are relevant, synthesize them coherently
- Be concise and direct - avoid elaboration unless asked
- Results are ordered by relevance, with rank 1 being most relevant
- If the search tool tells you the search limit is reached, stop searching immediately and answer with what you have
- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer an answer from tangentially related content.
"""

7
haiku_rag_slim/haiku/rag/agents/research/__init__.py
View file

@ -1,7 +0,0 @@
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.agents.research.models import (
Citation,
IterativePlanResult,
ResearchReport,
SearchAnswer,
)

34
haiku_rag_slim/haiku/rag/agents/research/dependencies.py
View file

@ -1,34 +0,0 @@
from typing import TYPE_CHECKING, Any
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models import SearchResult
if TYPE_CHECKING:
from haiku.rag.agents.research.models import SearchAnswer
class ResearchContext(BaseModel):
"""Context shared across research agents."""
original_question: str = Field(description="The original research question")
qa_responses: list[Any] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
def add_qa_response(self, qa: "SearchAnswer") -> None:
"""Add a structured QA response."""
self.qa_responses.append(qa)
class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context."""
model_config = {"arbitrary_types_allowed": True}
client: HaikuRAG = Field(description="RAG client for document operations")
context: ResearchContext = Field(description="Shared research context")
search_results: list[SearchResult] = Field(
default_factory=list, description="Search results for citation resolution"
)

280
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -1,280 +0,0 @@
import asyncio
from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.agents.research.models import (
IterativePlanResult,
RawSearchAnswer,
ResearchReport,
SearchAnswer,
)
from haiku.rag.agents.research.prompts import (
ITERATIVE_PLAN_PROMPT,
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
SEARCH_PROMPT,
SYNTHESIS_PROMPT,
)
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import build_prompt, get_model
def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for prompts."""
context_data: dict[str, object] = {}
context_data["question"] = context.original_question
if context.qa_responses:
context_data["prior_answers"] = [
{
"question": qa.query,
"answer": qa.answer,
"confidence": qa.confidence,
"source": qa.primary_source,
}
for qa in context.qa_responses
]
return format_as_xml(context_data, root_tag="context")
async def _iterative_plan_logic(
state: ResearchState,
deps: ResearchDeps,
config: AppConfig,
) -> IterativePlanResult:
"""Evaluate context and decide next question or mark complete."""
has_prior_answers = bool(state.context.qa_responses)
# If max iterations reached, skip LLM and mark complete
if state.iterations >= state.max_iterations:
return IterativePlanResult(
is_complete=True,
next_question=None,
reasoning=f"Max iterations ({state.max_iterations}) reached.",
)
model_config = config.research.model
if has_prior_answers:
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, config)
else:
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
model = get_model(model_config, config)
plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model=model,
output_type=IterativePlanResult,
instructions=effective_prompt,
tool_retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
# Build prompt based on current state
if has_prior_answers:
context_xml = format_context_for_prompt(state.context)
prompt = (
f"Review the gathered evidence and decide whether to continue or synthesize.\n\n"
f"{context_xml}"
)
else:
context_xml = format_context_for_prompt(state.context)
prompt = f"Plan the research investigation.\n\n{context_xml}"
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
result = await plan_agent.run(prompt, deps=agent_deps)
# Enforce: if no prior answers, must have a next_question to investigate
if not has_prior_answers:
if result.output.is_complete or not result.output.next_question:
return IterativePlanResult(
is_complete=False,
next_question=result.output.next_question
or state.context.original_question,
reasoning=result.output.reasoning,
)
return result.output
async def _search_one_step_logic(
state: ResearchState,
deps: ResearchDeps,
config: AppConfig,
search_prompt: str,
sub_q: str,
) -> SearchAnswer:
"""Answer a single question using the knowledge base."""
model_config = config.research.model
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
model = get_model(model_config, config)
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model=model,
output_type=RawSearchAnswer,
instructions=search_prompt,
tool_retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
"""Search the knowledge base for relevant documents."""
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
ctx2.deps.search_results = results
total = len(results)
parts = [
r.format_for_agent(rank=i + 1, total=total)
for i, r in enumerate(results)
]
if not parts:
return f"No relevant information found for: {query}"
return "\n\n".join(parts)
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
# Increment iterations after each search completes
state.iterations += 1
if raw_answer:
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
state.context.add_qa_response(answer)
return answer
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
def build_research_graph(
config: AppConfig = Config,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the iterative research graph.
Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters)
Returns:
Configured research graph with iterative planning
"""
model_config = config.research.model
search_prompt = build_prompt(SEARCH_PROMPT, config)
synthesis_prompt = build_prompt(
config.prompts.synthesis or SYNTHESIS_PROMPT, config
)
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ResearchReport,
)
@g.step
async def plan_next(
ctx: StepContext[ResearchState, ResearchDeps, None | SearchAnswer],
) -> IterativePlanResult:
"""Evaluate context and decide next question or complete."""
return await _iterative_plan_logic(ctx.state, ctx.deps, config)
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single question using the knowledge base."""
try:
return await _search_one_step_logic(
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
)
except Exception as e:
return SearchAnswer(
query=ctx.inputs,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
model = get_model(model_config, config)
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model=model,
output_type=ResearchReport,
instructions=synthesis_prompt,
tool_retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
# Build graph edges: iterative loop
#
# START -> plan_next -> [decision]
# |
# [is_complete or max_iterations] -> synthesize -> END
# |
# [has next_question] -> search_one -> plan_next (loop)
def extract_question(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> str:
"""Extract next_question from IterativePlanResult."""
return ctx.inputs.next_question or ""
g.add(
g.edge_from(g.start_node).to(plan_next),
g.edge_from(plan_next).to(
g.decision()
.branch(
g.match(
IterativePlanResult,
matches=lambda r: not r.is_complete and r.next_question is not None,
)
.label("Continue research")
.transform(extract_question)
.to(search_one)
)
.branch(
g.match(IterativePlanResult).label("Done researching").to(synthesize)
)
),
g.edge_from(search_one).to(plan_next),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

132
haiku_rag_slim/haiku/rag/agents/research/models.py
View file

@ -1,132 +0,0 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class IterativePlanResult(BaseModel):
"""Output from iterative planning step."""
is_complete: bool = Field(
description="Whether research is complete and can be synthesized"
)
next_question: str | None = Field(
default=None, description="Next question to investigate, if not complete"
)
reasoning: str = Field(description="Brief explanation of the decision")
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding.
Used by research graph and chat applications. The optional index field
supports UI display ordering in chat contexts.
"""
index: int | None = None
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
class RawSearchAnswer(BaseModel):
"""Answer to a search query with chunk references."""
query: str = Field(..., description="The question that was answered")
answer: str = Field(..., description="The answer to the question")
cited_chunks: list[str] = Field(
default_factory=list,
description="Complete chunk IDs from search results (e.g. '5ae52166-5329-42e9-b6a5-756fc0cb7200'). Copy the full UUID without brackets. Must not be empty when providing an answer.",
)
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
class SearchAnswer(RawSearchAnswer):
"""Answer to a search query with resolved citations."""
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)
@property
def primary_source(self) -> str | None:
"""Get primary source title from citations."""
if not self.citations:
return None
first = self.citations[0]
return first.document_title or first.document_uri
@classmethod
def from_raw(
cls,
raw: RawSearchAnswer,
search_results: "list[SearchResult]",
) -> "SearchAnswer":
"""Create SearchAnswer from RawSearchAnswer with resolved citations."""
citations = resolve_citations(raw.cited_chunks, search_results)
return cls(
query=raw.query,
answer=raw.answer,
cited_chunks=raw.cited_chunks,
confidence=raw.confidence,
citations=citations,
)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for raw_id in cited_chunk_ids:
chunk_id = raw_id.strip("[]")
r = by_id.get(chunk_id)
if not r:
continue
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
)
)
return citations
class ResearchReport(BaseModel):
"""Final research report structure."""
title: str = Field(description="Concise title for the research")
executive_summary: str = Field(description="Brief overview of key findings")
main_findings: list[str] = Field(
description="Primary research findings with supporting evidence"
)
conclusions: list[str] = Field(description="Evidence-based conclusions")
limitations: list[str] = Field(
description="Limitations of the current research", default=[]
)
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
sources_summary: str = Field(
description="Summary of sources used and their reliability"
)

114
haiku_rag_slim/haiku/rag/agents/research/prompts.py
View file

@ -1,114 +0,0 @@
ITERATIVE_PLAN_PROMPT = """You are the research orchestrator planning the investigation.
Your task:
1. Analyze the original question
2. Propose the first question to investigate
For simple questions, investigate them directly. For composite or complex questions,
you may decompose into a focused sub-question. For example:
- "What are the benefits and drawbacks of X?" → Start with "What are the benefits of X?"
- Ambiguous references should be resolved using background context if available
Output requirements:
- Set is_complete=False (you are just starting the investigation)
- Set next_question to the question to investigate
- Provide brief reasoning explaining your choice
The question must be standalone and self-contained:
- Include concrete entities, scope, and any qualifiers
- Avoid ambiguous pronouns (it/they/this/that)"""
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
You have access to context that may include:
- <prior_answers>: Previous Q&A pairs with confidence scores
Your task:
1. Review the provided evidence carefully
2. Assess whether it sufficiently answers the original question
3. Decide whether to continue research or synthesize
Decision criteria:
- Set is_complete=True if the evidence adequately answers the question
- Set is_complete=False with a next_question if important gaps remain
If not complete, propose exactly ONE high-value follow-up question in next_question:
- Focus on the most critical gap not covered by prior_answers
- The question must be standalone and self-contained
- Avoid repeating questions that have already been answered
- Include concrete entities, scope, and any qualifiers
Provide brief reasoning explaining your decision."""
SEARCH_PROMPT = """You are a search and question-answering specialist.
Process:
1. Call search_and_answer with relevant keywords from the question.
2. Review the results ordered by relevance.
3. If needed, perform follow-up searches with different keywords (max 3 total).
4. Provide a concise answer based strictly on the retrieved content.
The search tool returns results like:
[9bde5847-44c9-400a-8997-0e6b65babf92] [rank 1 of 5]
Source: "Document Title" > Section > Subsection
Type: paragraph
Content:
The actual text content here...
[d5a63c82-cb40-439f-9b2e-de7d177829b7] [rank 2 of 5]
Source: "Another Document"
Type: table
Content:
| Column 1 | Column 2 |
...
Each result includes:
- chunk_id in brackets and rank position (rank 1 = most relevant)
- Source: document title and section hierarchy (when available)
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
Output format:
- query: Echo the question you are answering
- answer: Your concise answer based on the retrieved content
- cited_chunks: List of plain strings containing only the chunk UUIDs (not objects)
- confidence: A score from 0.0 to 1.0 indicating answer confidence
IMPORTANT: Use the EXACT, COMPLETE chunk ID (full UUID). Do NOT truncate IDs.
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge.
- Use the Source and Type metadata to understand context.
- If multiple results are relevant, synthesize them coherently.
- If information is insufficient, say so clearly.
- Be concise and direct; avoid meta commentary about the process.
- Results are ordered by relevance, with rank 1 being most relevant."""
SYNTHESIS_PROMPT = """You are a synthesis specialist producing the final
research report that directly answers the original question.
Goals:
1. Directly answer the research question using gathered evidence.
2. Present findings clearly and concisely.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
Report guidelines (map to output fields):
- title: concise (5-12 words), informative.
- executive_summary: 3-5 sentences that DIRECTLY ANSWER the original question.
Write the actual answer, not a description of what the report contains.
BAD: "This report examines the topic and presents findings..."
GOOD: "The system requires configuration X and supports features Y and Z..."
- main_findings: list of plain strings, 4-8 one-sentence bullets reflecting evidence.
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- sources_summary: single string listing sources with document paths and page numbers.
All list fields must contain plain strings only, not objects.
Style:
- Base all content solely on the collected evidence.
- Be professional, objective, and specific.
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
Instead, state the actual information directly."""

59
haiku_rag_slim/haiku/rag/agents/research/state.py
View file

@ -1,59 +0,0 @@
import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.client import HaikuRAG
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@dataclass
class ResearchDeps:
"""Dependencies for research graph execution."""
client: HaikuRAG
semaphore: asyncio.Semaphore | None = None
class ResearchState(BaseModel):
"""Research graph state model."""
model_config = {"arbitrary_types_allowed": True}
context: ResearchContext = Field(
description="Shared research context with questions and QA responses"
)
iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
max_concurrency: int = Field(
default=1, description="Maximum concurrent search operations", ge=1
)
search_filter: str | None = Field(
default=None, description="SQL WHERE clause to filter search results"
)
@classmethod
def from_config(
cls,
context: ResearchContext,
config: "AppConfig",
max_iterations: int | None = None,
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
Args:
context: The ResearchContext containing the question
config: The AppConfig object
max_iterations: Override max iterations (None uses config default)
"""
return cls(
context=context,
max_iterations=max_iterations
if max_iterations is not None
else config.research.max_iterations,
max_concurrency=config.research.max_concurrency,
)

100
haiku_rag_slim/haiku/rag/app.py
View file

@ -16,7 +16,6 @@ from rich.progress import (
TextColumn,
TransferSpeedColumn,
)
from rich.syntax import Syntax
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, Config
@ -443,14 +442,12 @@ class HaikuRAGApp: # pragma: no cover
async def ask(
self,
question: str,
cite: bool = False,
filter: str | None = None,
):
"""Ask a question using the RAG system.
Args:
question: The question to ask
cite: Include citations in the answer
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
@ -465,21 +462,20 @@ class HaikuRAGApp: # pragma: no cover
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
if cite and citations:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
for renderable in await format_citations_rich(
citations, client=self.client
):
self.console.print(renderable)
async def analyze(
self,
question: str,
document: str | None = None,
filter: str | None = None,
):
"""Answer a question using the analysis agent with code execution.
"""Answer a question using the rag-analysis skill.
Args:
question: The question to answer
document: Optional document ID or title to pre-load
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
@ -488,95 +484,21 @@ class HaikuRAGApp: # pragma: no cover
read_only=self.read_only,
before=self.before,
) as self.client:
documents = [document] if document else None
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print(
"[dim]Running analysis agent with code execution...[/dim]"
"[dim]Running analysis skill with code execution...[/dim]"
)
self.console.print()
result = await self.client.analyze(
question, documents=documents, filter=filter
)
result = await self.client.analyze(question, filter=filter)
self.console.print("[bold yellow]Program:[/bold yellow]")
self.console.print(Syntax(result.program, "python"))
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(result.answer))
async def research(
self,
question: str,
filter: str | None = None,
):
"""Run research via the pydantic-graph pipeline.
Args:
question: The research question
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as client:
self.console.print("[bold cyan]Starting research[/bold cyan]")
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
report = await client.research(question=question, filter=filter)
if report is None:
self.console.print("[red]Research did not produce a report.[/red]")
return
# Display the report
self.console.print("[bold green]Research Report[/bold green]")
self.console.rule()
# Title and Executive Summary
self.console.print(f"[bold]{report.title}[/bold]")
self.console.print()
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
self.console.print(report.executive_summary)
self.console.print()
# Main Findings
if report.main_findings:
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"• {finding}")
self.console.print()
# Conclusions
if report.conclusions:
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
for conclusion in report.conclusions:
self.console.print(f"• {conclusion}")
self.console.print()
# Recommendations
if report.recommendations:
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
for rec in report.recommendations:
self.console.print(f"• {rec}")
self.console.print()
# Limitations
if report.limitations:
self.console.print("[bold yellow]Limitations:[/bold yellow]")
for limitation in report.limitations:
self.console.print(f"• {limitation}")
self.console.print()
# Sources Summary
if report.sources_summary:
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
for renderable in await format_citations_rich(
result.citations, client=self.client
):
self.console.print(renderable)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(

1
haiku_rag_slim/haiku/rag/chat/__init__.py
View file

@ -36,7 +36,6 @@ def run_chat(
if model:
model_config = parse_model_option(model)
config.qa.model = model_config
config.research.model = model_config
config.analysis.model = model_config
enabled = skills or ["rag"]

22
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -339,8 +339,26 @@ class ChatApp(App):
for cid in cited_ids:
if cid in citation_index:
citations.append(citation_index[cid])
if citations:
await chat_history.add_citations(citations)
if not citations:
return
picture_bytes: dict[str, list[bytes]] = {}
if self.client is not None:
for citation in citations:
refs = list(citation.picture_refs or [])
if not refs:
continue
blobs: list[bytes] = []
for ref in refs:
data = await self.client.document_item_repository.get_picture_bytes(
citation.document_id, ref
)
if data:
blobs.append(data)
if blobs:
picture_bytes[citation.chunk_id] = blobs
await chat_history.add_citations(citations, picture_bytes=picture_bytes)
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if analysis_state:

47
haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py
View file

@ -1,12 +1,15 @@
from io import BytesIO
from typing import TYPE_CHECKING, Any
from PIL import Image as PILImage
from textual.containers import Horizontal, VerticalScroll
from textual.css.query import NoMatches
from textual.message import Message
from textual.widgets import Collapsible, LoadingIndicator, Markdown, Static
from textual.widgets.markdown import MarkdownStream
from textual_image.widget import Image as TextualImage
from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.citation import Citation
if TYPE_CHECKING:
from textual.app import ComposeResult
@ -119,7 +122,12 @@ class CitationWidget(Collapsible):
super().__init__()
self.widget = widget
def __init__(self, citation: Citation, **kwargs) -> None:
def __init__(
self,
citation: Citation,
picture_bytes: list[bytes] | None = None,
**kwargs,
) -> None:
title = f"[{citation.index}] {citation.document_title or citation.document_uri}"
if citation.page_numbers:
pages = ", ".join(map(str, citation.page_numbers[:3]))
@ -131,7 +139,13 @@ class CitationWidget(Collapsible):
if len(content) > 500:
content = content[:500] + "..."
children: list[Markdown | Static] = [Markdown(content)]
children: list[Any] = [Markdown(content)]
for blob in picture_bytes or []:
try:
pil = PILImage.open(BytesIO(blob))
except Exception:
continue
children.append(TextualImage(pil, classes="citation-image"))
if citation.headings:
headings = " > ".join(citation.headings[:3])
children.append(Static(f"Section: {headings}", classes="citation-metadata"))
@ -333,6 +347,14 @@ class ChatHistory(VerticalScroll):
text-style: italic;
}
CitationWidget .citation-image {
width: auto;
height: auto;
max-width: 100%;
max-height: 30;
margin: 1 0;
}
/* Program */
ProgramWidget {
margin: 0 0 0 2;
@ -414,13 +436,26 @@ class ChatHistory(VerticalScroll):
widget.mark_completed()
widget.add_class("complete")
async def add_citations(self, citations: list[Citation]) -> None:
"""Add citations inline after a response."""
async def add_citations(
self,
citations: list[Citation],
picture_bytes: dict[str, list[bytes]] | None = None,
) -> None:
"""Add citations inline after a response.
``picture_bytes`` maps citation ``chunk_id`` → list of raw PNG bytes,
one per entry in the citation's ``picture_refs``. Pre-fetched by the
caller (typically the chat app's post-response hook) so widget
construction stays synchronous.
"""
if not citations:
return
await self.mount(SourcesHeader(len(citations)))
picture_bytes = picture_bytes or {}
for citation in citations:
widget = CitationWidget(citation)
widget = CitationWidget(
citation, picture_bytes=picture_bytes.get(citation.chunk_id)
)
await self.mount(widget)
self.scroll_end(animate=False)

34
haiku_rag_slim/haiku/rag/cli.py
View file

@ -369,11 +369,6 @@ def ask( # pragma: no cover
"--db",
help="Path to the LanceDB database file",
),
cite: bool = typer.Option(
False,
"--cite",
help="Include citations in the response",
),
filter: str | None = typer.Option(
None,
"--filter",
@ -385,13 +380,12 @@ def ask( # pragma: no cover
asyncio.run(
app.ask(
question=question,
cite=cite,
filter=filter,
)
)
@_cli.command("analyze", help="Answer questions using code execution (analysis agent)")
@_cli.command("analyze", help="Answer questions using the rag-analysis skill")
def analyze( # pragma: no cover
question: str = typer.Argument(
help="The question to answer",
@ -401,12 +395,6 @@ def analyze( # pragma: no cover
"--db",
help="Path to the LanceDB database file",
),
document: str | None = typer.Option(
None,
"--document",
"-d",
help="Document ID or title to pre-load for analysis",
),
filter: str | None = typer.Option(
None,
"--filter",
@ -418,31 +406,11 @@ def analyze( # pragma: no cover
asyncio.run(
app.analyze(
question=question,
document=document,
filter=filter,
)
)
@_cli.command("research", help="Run multi-agent research and output a concise report")
def research( # pragma: no cover
question: str = typer.Argument(..., help="The research question to investigate"),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
),
):
app = create_app(db)
asyncio.run(app.research(question=question, filter=filter))
@_cli.command("settings", help="Display current configuration settings")
def settings(): # pragma: no cover
config = get_config()

26
haiku_rag_slim/haiku/rag/client/__init__.py
View file

@ -30,11 +30,8 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from PIL import Image as PILImage
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import (
Citation,
ResearchReport,
)
from haiku.rag.sandbox import AnalysisResult
from haiku.rag.store.models.citation import Citation
logger = logging.getLogger(__name__)
@ -368,35 +365,20 @@ class HaikuRAG:
async def ask(
self,
question: str,
system_prompt: str | None = None,
filter: str | None = None,
) -> "tuple[str, list[Citation]]":
from haiku.rag.client.agents import ask
return await ask(self, question, system_prompt, filter)
async def research(
self,
question: str,
*,
filter: str | None = None,
max_iterations: int | None = None,
) -> "ResearchReport":
from haiku.rag.client.agents import research
return await research(
self, question, filter=filter, max_iterations=max_iterations
)
return await ask(self, question, filter)
async def analyze(
self,
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "AnalysisResult":
from haiku.rag.client.agents import analyze
return await analyze(self, question, documents, filter)
return await analyze(self, question, filter)
async def visualize_chunk(self, chunk: Chunk) -> list:
from haiku.rag.client.search import visualize_chunk

139
haiku_rag_slim/haiku/rag/client/agents.py
View file

@ -1,140 +1,75 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import Citation, ResearchReport
from haiku.rag.client import HaikuRAG
from haiku.rag.sandbox import AnalysisResult
from haiku.rag.store.models.citation import Citation
async def ask(
client: "HaikuRAG",
question: str,
system_prompt: str | None = None,
filter: str | None = None,
) -> "tuple[str, list[Citation]]":
"""Ask a question using the configured QA agent.
"""Ask a question against the knowledge base via the rag skill.
Args:
client: The HaikuRAG client.
question: The question to ask.
system_prompt: Optional custom system prompt for the QA agent.
filter: SQL WHERE clause to filter documents.
Returns:
Tuple of (answer text, list of resolved citations).
"""
from haiku.rag.agents.qa import get_qa_agent
from haiku.rag.skills.rag import RAGState, create_skill
from haiku.rag.utils import get_model
from haiku.skills import run_skill
qa_agent = get_qa_agent(client, config=client._config, system_prompt=system_prompt)
return await qa_agent.answer(question, filter=filter)
async def research(
client: "HaikuRAG",
question: str,
*,
filter: str | None = None,
max_iterations: int | None = None,
) -> "ResearchReport":
"""Run multi-agent research to investigate a question.
Args:
client: The HaikuRAG client.
question: The research question to investigate.
filter: SQL WHERE clause to filter documents.
max_iterations: Override max iterations (None uses config default).
Returns:
ResearchReport with structured findings.
"""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
graph = build_research_graph(config=client._config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context, config=client._config, max_iterations=max_iterations
)
state.search_filter = filter
deps = ResearchDeps(client=client)
return await graph.run(state=state, deps=deps)
skill = create_skill(db_path=client.store.db_path, config=client._config)
state = RAGState(document_filter=filter)
model = get_model(client._config.qa.model, client._config)
answer, _, _ = await run_skill(model, skill, question, state=state)
citations = [
state.citation_index[cid]
for cid in state.citations
if cid in state.citation_index
]
return answer, citations
async def analyze(
client: "HaikuRAG",
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "AnalysisResult":
"""Answer a question using the analysis agent with code execution.
"""Answer a question against the knowledge base via the rag-analysis skill.
The analysis agent can write and execute Python code in a sandboxed
environment to solve problems that require computation, aggregation, or
complex traversal across documents.
The analysis skill exposes ``search``, ``execute_code``, and ``cite`` tools.
The driving model decides when to reach for code (structural traversal,
computation, aggregation) versus a direct ``search → cite → answer``.
Args:
client: The HaikuRAG client.
question: The question to answer.
documents: Optional list of document IDs or titles to pre-load.
filter: SQL WHERE clause to filter documents during searches.
Returns:
AnalysisResult with the answer and the final consolidated program.
AnalysisResult with the answer and resolved citations.
"""
from haiku.rag.agents.analysis import (
AnalysisContext,
AnalysisDeps,
Sandbox,
create_analysis_agent,
)
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import Citation
context = AnalysisContext(filter=filter)
if documents:
loaded_docs = []
for doc_ref in documents:
doc = await client.resolve_document(doc_ref)
if doc:
loaded_docs.append(doc)
context.documents = loaded_docs if loaded_docs else None
sandbox = Sandbox(
db_path=client.store.db_path,
config=client._config,
context=context,
)
deps = AnalysisDeps(
sandbox=sandbox,
context=context,
)
agent = create_analysis_agent(client._config)
result = await agent.run(question, deps=deps)
output = result.output
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,
citations=citations,
from haiku.rag.sandbox import AnalysisResult
from haiku.rag.skills.analysis import AnalysisState, create_skill
from haiku.rag.utils import get_model
from haiku.skills import run_skill
skill = create_skill(db_path=client.store.db_path, config=client._config)
state = AnalysisState(document_filter=filter)
model = get_model(
client._config.analysis.model or client._config.qa.model, client._config
)
answer, _, _ = await run_skill(model, skill, question, state=state)
citations = [
state.citation_index[cid]
for cid in state.citations
if cid in state.citation_index
]
return AnalysisResult(answer=answer, citations=citations)

2
haiku_rag_slim/haiku/rag/client/downloads.py
View file

@ -101,8 +101,6 @@ async def download_models(
required_models.add(config.embeddings.model.name)
if config.qa.model.provider == "ollama":
required_models.add(config.qa.model.name)
if config.research.model.provider == "ollama":
required_models.add(config.research.model.name)
if config.reranking.model and config.reranking.model.provider == "ollama":
required_models.add(config.reranking.model.name)
pic_desc = config.processing.conversion_options.picture_description

27
haiku_rag_slim/haiku/rag/client/search.py
View file

@ -3,6 +3,7 @@ from typing import TYPE_CHECKING
from haiku.rag.reranking import get_reranker
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
if TYPE_CHECKING:
from PIL import Image as PILImage
@ -93,7 +94,9 @@ def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
seen: dict[tuple[str | None, str], int] = {}
keep: list[bool] = [True] * len(results)
for i, r in enumerate(results):
if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith("#/pictures/"):
if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith(
PICTURE_REF_PREFIX
):
key = (r.document_id, r.doc_item_refs[0])
prior = seen.get(key)
if prior is None:
@ -111,13 +114,13 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
Groups results by document_id and batches one picture-bytes lookup per
document so a result set spanning N documents costs N reads, not one per
picture. Only refs starting with ``#/pictures/`` are queried.
picture. Only refs starting with ``PICTURE_REF_PREFIX`` are queried.
"""
by_doc: dict[str, list[SearchResult]] = {}
for r in results:
if not r.document_id:
continue
if not any(ref.startswith("#/pictures/") for ref in r.doc_item_refs):
if not any(ref.startswith(PICTURE_REF_PREFIX) for ref in r.doc_item_refs):
continue
by_doc.setdefault(r.document_id, []).append(r)
@ -126,7 +129,7 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
seen: set[str] = set()
for r in doc_results:
for ref in r.doc_item_refs:
if ref.startswith("#/pictures/") and ref not in seen:
if ref.startswith(PICTURE_REF_PREFIX) and ref not in seen:
wanted.append(ref)
seen.add(ref)
if not wanted:
@ -136,14 +139,23 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
)
if not bytes_by_ref:
continue
captions_by_ref = await client.document_item_repository.get_text_for_refs(
doc_id, list(bytes_by_ref.keys())
)
for r in doc_results:
attached: dict[str, str] = {}
captions: dict[str, str] = {}
for ref in r.doc_item_refs:
blob = bytes_by_ref.get(ref)
if blob:
attached[ref] = base64.b64encode(blob).decode("ascii")
caption = captions_by_ref.get(ref)
if caption:
captions[ref] = caption
if attached:
r.image_data = attached
if captions:
r.picture_captions = captions
async def expand_context(
@ -192,9 +204,10 @@ async def expand_context(
expanded_results.extend(expanded)
expanded_results.sort(key=lambda r: r.score, reverse=True)
# expand_with_items rebuilds SearchResult objects, so attach picture bytes
# to the fresh set — picture self_refs may have grown via section expansion.
await _populate_image_data(client, expanded_results)
# image_data and picture_captions are preserved through expansion by
# expand_with_items — we deliberately do not re-attach bytes for refs
# introduced by section expansion, so the multimodal payload stays
# bounded by what was originally retrieved.
return expanded_results

2
haiku_rag_slim/haiku/rag/config/__init__.py
View file

@ -17,7 +17,6 @@ from haiku.rag.config.models import (
ProvidersConfig,
QAConfig,
RerankingConfig,
ResearchConfig,
S3MonitorEntry,
StorageConfig,
)
@ -37,7 +36,6 @@ __all__ = [
"ProvidersConfig",
"QAConfig",
"RerankingConfig",
"ResearchConfig",
"S3MonitorEntry",
"StorageConfig",
"find_config_file",

50
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -102,28 +102,15 @@ class QAConfig(BaseModel):
max_searches: int = 5
class ResearchConfig(BaseModel):
model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
name="gpt-oss",
enable_thinking=False,
temperature=0.3,
)
)
max_iterations: int = 3
max_concurrency: int = 1
class AnalysisConfig(BaseModel):
model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
name="gpt-oss",
enable_thinking=False,
temperature=0.0,
)
)
"""Driving model + sandbox limits for the analysis skill.
``model`` defaults to ``None``, meaning "no override — use ``qa.model``."
Consumers resolve via ``config.analysis.model or config.qa.model``. Set
explicitly when the analysis workload wants a different model from QA
(e.g. a stronger model for computational tasks)."""
model: ModelConfig | None = None
code_timeout: float = 60.0
max_output_chars: int = 50_000
@ -215,7 +202,7 @@ class ProcessingConfig(BaseModel):
class SearchConfig(BaseModel):
limit: int = 10
limit: int = 5
max_context_chars: int = 10000
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
vector_refine_factor: int = 30
@ -241,8 +228,6 @@ class ProvidersConfig(BaseModel):
class PromptsConfig(BaseModel):
domain_preamble: str = ""
qa: str | None = None
synthesis: str | None = None
picture_description: str = (
"Describe this image for a blind user. "
"State the image type (screenshot, chart, photo, etc.), "
@ -251,6 +236,19 @@ class PromptsConfig(BaseModel):
)
class EvaluationsConfig(BaseModel):
"""Settings consumed only by the `evaluations` package."""
judge: ModelConfig | None = Field(
default=None,
description=(
"Judge model for `evaluations run`'s LLM-as-judge step. "
"ModelConfig's base_url lets the judge point at any "
"OpenAI-compatible endpoint."
),
)
class AppConfig(BaseModel):
environment: str = "production"
storage: StorageConfig = Field(default_factory=StorageConfig)
@ -259,9 +257,11 @@ class AppConfig(BaseModel):
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
qa: QAConfig = Field(default_factory=QAConfig)
research: ResearchConfig = Field(default_factory=ResearchConfig)
analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
prompts: PromptsConfig = Field(default_factory=PromptsConfig)
evaluations: "EvaluationsConfig" = Field(
default_factory=lambda: EvaluationsConfig()
)

16
haiku_rag_slim/haiku/rag/context.py
View file

@ -251,6 +251,20 @@ async def expand_with_items(
if r.headings:
all_headings.extend(h for h in r.headings if h not in all_headings)
# Carry image_data and picture_captions through expansion so that
# only pictures from the originally retrieved chunks get attached.
# Pictures swept in by section expansion are referenced in `refs`
# for cross-referencing but their bytes are not re-fetched —
# otherwise a single search can balloon the response with adjacent
# figures the model did not actually retrieve.
merged_image_data: dict[str, str] = {}
merged_captions: dict[str, str] = {}
for r in original_results:
if r.image_data:
merged_image_data.update(r.image_data)
if r.picture_captions:
merged_captions.update(r.picture_captions)
first = original_results[0]
# Expansion should never return less content than the original chunk.
@ -272,6 +286,8 @@ async def expand_with_items(
page_numbers=sorted(pages) or first.page_numbers,
headings=all_headings or None,
labels=sorted(labels) or first.labels,
image_data=merged_image_data or None,
picture_captions=merged_captions,
)
)

35
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -3,7 +3,6 @@ from typing import Any
from fastmcp import FastMCP
from haiku.rag.agents.research.models import ResearchReport
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, Config
from haiku.rag.store.models import Document, SearchResult
@ -202,42 +201,19 @@ def create_mcp_server(
except Exception as e:
return f"Error answering question: {e!s}"
@mcp.tool()
async def research_question(
question: str,
) -> ResearchReport | None:
"""Run multi-agent research to investigate a complex question.
The research process uses multiple agents to plan, search, evaluate, and synthesize
information iteratively until confidence threshold is met or max iterations reached.
Args:
question: The research question to investigate.
Returns:
A research report with findings, or None if an error occurred.
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
return await rag.research(question=question)
except Exception:
return None
@mcp.tool()
async def analyze(
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex questions using code execution (analysis agent).
"""Answer complex questions using the rag-analysis skill.
Use this for questions requiring computation, aggregation, or
complex traversal across documents. The agent can write Python
code to search, analyze, and compute answers.
structural traversal across documents. The skill can write and
execute Python code in a sandboxed interpreter.
Args:
question: The question to answer.
document: Optional document ID or title to pre-load for analysis.
filter: Optional SQL WHERE clause to filter documents.
Returns:
@ -245,10 +221,9 @@ def create_mcp_server(
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
documents = [document] if document else None
result = await rag.analyze(question, documents=documents, filter=filter)
result = await rag.analyze(question, filter=filter)
return result.answer
except Exception as e:
return f"Error running analysis agent: {e!s}"
return f"Error running analysis skill: {e!s}"
return mcp

10
haiku_rag_slim/haiku/rag/sandbox/__init__.py Normal file
View file

@ -0,0 +1,10 @@
from haiku.rag.sandbox.dependencies import AnalysisContext
from haiku.rag.sandbox.models import AnalysisResult
from haiku.rag.sandbox.sandbox import Sandbox, SandboxResult
__all__ = [
"AnalysisContext",
"AnalysisResult",
"Sandbox",
"SandboxResult",
]

8
haiku_rag_slim/haiku/rag/sandbox/dependencies.py Normal file
View file

@ -0,0 +1,8 @@
from dataclasses import dataclass
@dataclass
class AnalysisContext:
"""Mutable context accumulating data during analysis execution."""
filter: str | None = None

14
haiku_rag_slim/haiku/rag/sandbox/models.py Normal file
View file

@ -0,0 +1,14 @@
from pydantic import BaseModel, Field
from haiku.rag.store.models.citation import Citation
class AnalysisResult(BaseModel):
"""Result from analysis execution with resolved citations.
Executed code is tracked on ``AnalysisState.executions`` (populated by the
analysis skill's ``execute_code`` tool). Consumers that need the program
should pull it from the skill state."""
answer: str
citations: list[Citation] = Field(default_factory=list)

448
haiku_rag_slim/haiku/rag/sandbox/sandbox.py Normal file
View file

@ -0,0 +1,448 @@
import asyncio
import atexit
import concurrent.futures
import json
import os
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from pydantic_monty import CallbackFile, MemoryFile, MontyRepl, OSAccess
from haiku.rag.config.models import AppConfig
from haiku.rag.sandbox.dependencies import AnalysisContext
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem
if TYPE_CHECKING:
from pathlib import PurePosixPath
@dataclass
class SandboxResult:
"""Result of executing code in the sandbox."""
stdout: str
stderr: str
success: bool
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)
atexit.register(_executor.shutdown, wait=False)
def _run_async(coro: Any) -> Any:
"""Run an async coroutine from a sync context (CallbackFile read)."""
return _executor.submit(asyncio.run, coro).result()
def _build_toc(
items: list["DocumentItem"],
chunk_index: dict[str, list[str]],
) -> list[dict[str, Any]]:
"""Build a nested section tree from items in position order.
Each ``section_header`` with ``heading_level > 0`` becomes a node. Nesting
follows the explicit levels: a header pops the stack until the top is at
a strictly shallower level, then becomes a child of that top (or a root).
``item_range = [position, end_exclusive]`` where ``end_exclusive`` is the
position of the next header whose level is the same or shallower (i.e.
the next sibling or ancestor that ends this section), or the total item
count if no such header exists.
``chunk_ids`` aggregates the chunks covered by all items in the section's
``item_range`` (deduped, order preserved). Pass directly to ``cite()`` to
ground a section-scoped answer without a corpus-wide ``search()`` call.
Items without a section_header label (or with ``heading_level == 0``) are
skipped. When all section_headers carry the same level the output is a
flat sibling list (see docling-project/docling#2121 for an upstream case
where every PDF section_header is emitted at level=1).
"""
# Defensive: every consumer is supposed to pass items in position order,
# but the end_exclusive lookahead below silently miscomputes section
# boundaries if it's not — better to sort once than trust the caller.
items = sorted(items, key=lambda i: i.position)
headers: list[DocumentItem] = [
i for i in items if i.label == "section_header" and i.heading_level > 0
]
if not headers:
return []
total = max((i.position for i in items), default=-1) + 1
items_by_position: dict[int, DocumentItem] = {i.position: i for i in items}
ends: list[int] = []
for idx, h in enumerate(headers):
end = total
for j in range(idx + 1, len(headers)):
if headers[j].heading_level <= h.heading_level:
end = headers[j].position
break
ends.append(end)
roots: list[dict[str, Any]] = []
stack: list[tuple[int, dict[str, Any]]] = []
for h, end in zip(headers, ends, strict=True):
seen: set[str] = set()
chunk_ids: list[str] = []
for pos in range(h.position, end):
item = items_by_position.get(pos)
if item is None:
continue
for cid in chunk_index.get(item.self_ref, []):
if cid not in seen:
seen.add(cid)
chunk_ids.append(cid)
node: dict[str, Any] = {
"self_ref": h.self_ref,
"level": h.heading_level,
"title": h.text,
"page_numbers": list(h.page_numbers),
"item_range": [h.position, end],
"chunk_ids": chunk_ids,
"children": [],
}
while stack and stack[-1][0] >= h.heading_level:
stack.pop()
(stack[-1][1]["children"] if stack else roots).append(node)
stack.append((h.heading_level, node))
return roots
class Sandbox:
"""Execute code in a sandboxed Python interpreter.
Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
External functions (search, list_documents) are called by Monty code
using ``await`` and resolved asynchronously on the host. Documents are
exposed via a virtual filesystem at ``/documents/{id}/``.
The interpreter uses a REPL session — variables persist across
``execute()`` calls within the same Sandbox instance.
sandbox = Sandbox(db_path, config, context)
result = await sandbox.execute("x = await search('query')")
result = await sandbox.execute("print(x[0]['content'])") # x persists
"""
_db_path: Path
_config: AppConfig
_context: AnalysisContext
_search_results: "list[SearchResult]"
_doc_items: dict[str, list["DocumentItem"]]
_doc_chunk_index: dict[str, dict[str, list[str]]]
_items_jsonl_cache: dict[str, str]
_toc_json_cache: dict[str, str]
_repl: MontyRepl | None
_vfs: OSAccess | None
def __init__(
self,
db_path: Path,
config: AppConfig,
context: AnalysisContext,
):
self._db_path = db_path
self._config = config
self._context = context
self._search_results = []
self._doc_items = {}
self._doc_chunk_index = {}
self._items_jsonl_cache = {}
self._toc_json_cache = {}
self._repl = None
self._vfs = None
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
db_path = self._db_path
config = self._config
context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
# Picture bytes are deliberately not attached to in-code search
# results: the Monty interpreter has no PIL/base64/hashlib, so the
# agent's Python can't do anything with them. The driving model
# gets figures through the top-level `search` tool when the
# question is visual; in-code search is for structural work.
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(query, limit=limit, filter=context.filter)
expanded = await rag.expand_context(results)
self._search_results.extend(expanded)
out: list[dict[str, Any]] = []
for r in expanded:
picture_refs = [
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
]
out.append(
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
"doc_item_refs": r.doc_item_refs,
"labels": r.labels,
"picture_refs": picture_refs,
}
)
return out
async def list_documents() -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
docs = await rag.list_documents(filter=context.filter)
return [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"created_at": str(d.created_at),
}
for d in docs
]
return {
"search": search,
"list_documents": list_documents,
}
async def _build_vfs(self) -> OSAccess:
"""Build the virtual filesystem with document data.
Mounts per-document directories with:
- metadata.json: MemoryFile (eager, small)
- content.txt: CallbackFile (lazy, can be large)
- items.jsonl: CallbackFile (lazy, bulk-cached)
- toc.json: CallbackFile (lazy, bulk-cached)
"""
from haiku.rag.client import HaikuRAG
db_path = self._db_path
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)
doc_titles = {doc.id: doc.title for doc in docs if doc.id}
sandbox = self
def _get_items(did: str) -> list[DocumentItem]:
"""Fetch items for one doc, cached on the sandbox."""
cached = sandbox._doc_items.get(did)
if cached is not None:
return cached
async def _fetch() -> list[DocumentItem]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
return await rag.document_item_repository.get_all_items(did)
items = _run_async(_fetch())
sandbox._doc_items[did] = items
return items
def _get_chunk_index(did: str) -> dict[str, list[str]]:
"""Fetch the self_ref → chunk_ids index for one doc, cached."""
cached = sandbox._doc_chunk_index.get(did)
if cached is not None:
return cached
async def _fetch() -> dict[str, list[str]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
index = (
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
[did]
)
)
return index.get(did, {})
chunk_index = _run_async(_fetch())
sandbox._doc_chunk_index[did] = chunk_index
return chunk_index
def _make_items_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_items(_path: "PurePosixPath") -> str:
cached = sandbox._items_jsonl_cache.get(did)
if cached is not None:
return cached
items = _get_items(did)
chunk_index = _get_chunk_index(did)
jsonl = "\n".join(
json.dumps(
{
"self_ref": item.self_ref,
"label": item.label,
"text": item.text,
"page_numbers": item.page_numbers,
"heading_level": item.heading_level,
"chunk_ids": chunk_index.get(item.self_ref, []),
},
ensure_ascii=False,
)
for item in items
)
sandbox._items_jsonl_cache[did] = jsonl
return jsonl
return read_items
def _make_toc_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_toc(_path: "PurePosixPath") -> str:
cached = sandbox._toc_json_cache.get(did)
if cached is not None:
return cached
items = _get_items(did)
chunk_index = _get_chunk_index(did)
toc = json.dumps(
{
"doc_id": did,
"title": doc_titles.get(did),
"tree": _build_toc(items, chunk_index),
},
ensure_ascii=False,
)
sandbox._toc_json_cache[did] = toc
return toc
return read_toc
for doc in docs:
if not doc.id:
continue
doc_id: str = doc.id
doc_dir = f"/documents/{doc_id}"
metadata = json.dumps(
{
"id": doc_id,
"title": doc.title,
"uri": doc.uri,
"created_at": str(doc.created_at),
},
ensure_ascii=False,
)
files.append(MemoryFile(f"{doc_dir}/metadata.json", metadata))
def _make_content_reader(
did: str,
) -> Callable[["PurePosixPath"], str]:
def read_content(_path: "PurePosixPath") -> str:
async def _fetch() -> str:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(
db_path, config=config, read_only=True
) as rag:
content = await rag.document_repository.get_content(did)
return content or ""
return _run_async(_fetch())
return read_content
files.append(
CallbackFile(
f"{doc_dir}/content.txt",
read=_make_content_reader(doc_id),
write=_deny_write,
)
)
files.append(
CallbackFile(
f"{doc_dir}/items.jsonl",
read=_make_items_reader(doc_id),
write=_deny_write,
)
)
# HAIKU_RAG_DISABLE_TOC is an evaluation-time toggle for measuring
# whether toc.json's outline view earns its place in the VFS.
# Production callers should leave it unset.
if not os.environ.get("HAIKU_RAG_DISABLE_TOC"):
files.append(
CallbackFile(
f"{doc_dir}/toc.json",
read=_make_toc_reader(doc_id),
write=_deny_write,
)
)
return OSAccess(files)
async def _ensure_initialized(self) -> tuple[MontyRepl, OSAccess]:
"""Initialize the REPL session and VFS on first use."""
if self._repl is None:
self._vfs = await self._build_vfs()
self._repl = MontyRepl(
limits={
"max_duration_secs": self._config.analysis.code_timeout,
},
)
assert self._repl is not None and self._vfs is not None
return self._repl, self._vfs
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty REPL.
Variables persist across calls within the same Sandbox instance.
"""
repl, vfs = await self._ensure_initialized()
external_fns = self._build_external_functions()
stdout_lines: list[str] = []
def print_callback(_stream: Literal["stdout"], text: str) -> None:
stdout_lines.append(text)
max_chars = self._config.analysis.max_output_chars
try:
output = await repl.feed_run_async(
code,
external_functions=external_fns,
print_callback=print_callback,
os=vfs,
)
except (
pydantic_monty.MontySyntaxError,
pydantic_monty.MontyRuntimeError,
) as e:
stdout = "".join(stdout_lines)
if len(stdout) > max_chars:
stdout = stdout[:max_chars] + "\n... (output truncated)"
return SandboxResult(stdout=stdout, stderr=str(e), success=False)
stdout = "".join(stdout_lines)
if output is not None:
stdout_with_output = f"{stdout}{output}" if stdout else str(output)
else:
stdout_with_output = stdout
if len(stdout_with_output) > max_chars:
stdout_with_output = (
stdout_with_output[:max_chars] + "\n... (output truncated)"
)
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)

2
haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2
View file

@ -26,7 +26,7 @@ Retrieve a document by ID, title, or URI. Partial matches work.
{% if "execute_code" in tool_names %}
### execute_code
Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, `await llm()`, and a virtual filesystem at `/documents/` with document content and structure.
Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, and a virtual filesystem at `/documents/` with document content and structure.
{% endif %}
{% if "cite" in tool_names %}

2
haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
View file

@ -6,7 +6,7 @@ from haiku.rag.config.models import AppConfig
from haiku.skills.models import Skill
from haiku.skills.parser import parse_skill_md
{% if "cite" in tool_names %}
from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.citation import Citation
{% endif %}
{% if "search" in tool_names %}
from haiku.rag.store.models.chunk import SearchResult

5
haiku_rag_slim/haiku/rag/skills/_deps.py
View file

@ -8,8 +8,8 @@ from haiku.rag.config.models import AppConfig
from haiku.skills.state import SkillRunDeps
if TYPE_CHECKING:
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
from haiku.rag.sandbox import Sandbox
@dataclass
@ -60,9 +60,8 @@ def make_rag_lifespan(db_path: Path, config: AppConfig):
def make_analysis_lifespan(db_path: Path, config: AppConfig):
@asynccontextmanager
async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]:
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
from haiku.rag.sandbox import AnalysisContext, Sandbox
doc_filter = getattr(deps.state, "document_filter", None)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:

81
haiku_rag_slim/haiku/rag/skills/_tools.py
View file

@ -2,14 +2,14 @@ from pathlib import Path
from typing import Any
from pydantic import BaseModel
from pydantic_ai import RunContext
from pydantic_ai import ModelRetry, RunContext
from pydantic_ai.messages import ToolReturn
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
from haiku.rag.tools.search import build_binary_parts_from_results
@ -155,12 +155,18 @@ def create_skill_tools(
config: AppConfig,
state_type: type[BaseModel],
tool_names: list[str],
model: ModelConfig,
) -> dict[str, Any]:
"""Create tool closures for a skill.
Returns a dict mapping tool name to async callable.
Each tool extracts state from RunContext, calls the shared implementation,
and updates state.
and updates state. ``model`` is the driving model for the skill (e.g.
``config.qa.model`` for the RAG skill, or
``config.analysis.model or config.qa.model`` for the analysis skill,
which defaults to ``None`` and inherits QA's model when unconfigured);
its ``vision`` flag gates picture-bytes attachment on the ``search``
tool.
"""
tools: dict[str, Any] = {}
@ -173,10 +179,9 @@ def create_skill_tools(
"""Search the knowledge base using hybrid search (vector + full-text).
Returns ranked results with content and metadata. When picture
content is in the result set and the configured QA model is
vision-capable (``qa.model.vision = true``), picture bytes are
attached as ``BinaryContent`` parts so the model sees figures
alongside text.
content is in the result set and the driving skill model is
vision-capable, picture bytes are attached as ``BinaryContent``
parts so the model sees figures alongside text.
Args:
query: The search query.
@ -199,7 +204,7 @@ def create_skill_tools(
if state:
state.searches[query] = results
if not config.qa.model.vision:
if not model.vision:
return formatted
binary_parts = build_binary_parts_from_results(results)
@ -242,9 +247,10 @@ def create_skill_tools(
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter.
The code has access to search(), list_documents(), llm() functions
and a virtual filesystem at /documents/ with document content and
structure (metadata.json, content.txt, items.jsonl per document).
The code has access to search() and list_documents() functions
and a virtual filesystem at /documents/ with document content
and structure (metadata.json, content.txt, items.jsonl, toc.json
per document).
Use print() to output results. Variables persist between calls
within the same skill invocation.
@ -289,26 +295,65 @@ def create_skill_tools(
async def cite(ctx: RunContext[RAGRunDeps], chunk_ids: list[str]) -> str:
"""Register chunk IDs as citations for your answer.
Call this after searching, with the chunk_id values from search
results that support your answer.
Accepts chunk_ids from search results AND from direct file reads
(items.jsonl, toc.json). Verbatim copies only — chunk_ids that
don't exist in the database trigger a retry.
Args:
chunk_ids: List of chunk_id values from search results.
chunk_ids: List of chunk_id values from search results or VFS reads.
"""
from haiku.rag.agents.research.models import resolve_citations
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import resolve_citations
state = _get_state(ctx, state_type)
if not state:
return "No state available."
all_results = []
if not chunk_ids:
return "Registered 0 citations (empty chunk_ids)."
all_results: list[SearchResult] = []
for results_list in state.searches.values():
all_results.extend(results_list)
citations = resolve_citations(chunk_ids, all_results)
resolved_ids = {c.chunk_id for c in citations}
missing = [
cid.strip("[]")
for cid in chunk_ids
if cid.strip("[]") not in resolved_ids
]
if missing:
rag = _require_rag(ctx)
synthetic: list[SearchResult] = []
doc_cache: dict[str, Any] = {}
for cid in missing:
chunk = await rag.get_chunk_by_id(cid)
if chunk is None or not chunk.document_id:
continue
did = chunk.document_id
if did in doc_cache:
doc = doc_cache[did]
else:
doc = await rag.get_document_by_id(did)
doc_cache[did] = doc
chunk.document_uri = doc.uri if doc else None
chunk.document_title = doc.title if doc else None
synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
if synthetic:
citations.extend(resolve_citations(missing, synthetic))
if citations:
_register_citations(state, citations)
return f"Registered {len(citations)} citation(s)."
return f"Registered {len(citations)} citation(s)."
raise ModelRetry(
f"None of the supplied chunk_ids {list(chunk_ids)} could be "
"resolved. Copy chunk_ids verbatim from `search` results or "
"from the `chunk_ids` field on items.jsonl / toc.json rows — "
"never reconstruct, abbreviate, or paraphrase them."
)
tools["cite"] = cite

5
haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -4,10 +4,10 @@ from pathlib import Path
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 CodeExecutionEntry
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@ -77,7 +77,8 @@ def create_skill(
db_path,
config,
AnalysisState,
["search", "list_documents", "execute_code", "cite"],
["search", "execute_code", "cite"],
model=config.analysis.model or config.qa.model,
)
extras = create_skill_extras(db_path, config)

56
haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
View file

@ -10,7 +10,7 @@ description: >
# Analysis
You solve complex analytical questions by writing and executing Python code against the knowledge base.
You answer questions over a document knowledge base. Most questions are answered directly with `search → cite → answer`. Reach for `execute_code` when a question requires computation, aggregation, or structural traversal that a single search cannot deliver.
## Tools
@ -18,21 +18,25 @@ You solve complex analytical questions by writing and executing Python code agai
Execute Python code in a sandboxed interpreter. Variables persist between calls — you can build state incrementally. Use `print()` to output results.
Inside the code, these functions are available (use `await`):
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels, picture_refs (subset of doc_item_refs labeled `picture`)
- `await list_documents()` → list of dicts with keys: id, title, uri, created_at
- `await llm(prompt)` → string response from an LLM (for classification, summarization, extraction)
Available modules: `json`, `re`, `math`, `pathlib`
Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
### search
Search the knowledge base directly (outside code execution). Use for initial exploration before writing code. Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
### list_documents
List available documents. Use to discover what's in the knowledge base.
Search the knowledge base directly (outside code execution). Each result has a `Type:` (paragraph, table, code, list_item, picture). When the Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text — use it directly to answer questions about figures, diagrams, charts, screenshots.
### cite
Register chunk IDs as citations. Call after your analysis with chunk_id values from search results that support your answer.
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer.
Chunk IDs come from two places:
- The `chunk_id` field on `search` / `await search(...)` results
- The `chunk_ids` field on `items.jsonl` rows (when you ground via direct file reads)
Do NOT cite `self_ref` (`#/texts/N` style refs), `position`, or any other identifier-shaped field. They are not chunk IDs and the tool will reject them. Copy chunk IDs verbatim — they are opaque UUIDs.
Every answer that uses search or file-read evidence must be backed by `cite`.
## Document Filesystem (inside execute_code)
@ -43,8 +47,11 @@ All documents are mounted as a virtual filesystem at `/documents/`:
metadata.json # {"id", "title", "uri", "created_at"}
content.txt # Full document text
items.jsonl # Structured items (one JSON object per line)
toc.json # Section tree derived from heading_level
```
`{document_id}` is an internal identifier, not the user-facing `uri` (filename, URL, etc.). When you only know a document by its URI or title, use `await list_documents()` to enumerate ids and match against `uri` / `title` — that's a single call to the host. Iterating `/documents/` and reading every `metadata.json` works too but is much slower on portal-scale corpora.
### Reading files
Always use `Path.read_text()` — do NOT use `open()` or `with` statements (they are not supported).
@ -74,28 +81,37 @@ Document metadata: `id`, `title`, `uri`, `created_at`.
Full text content. Use for regex or keyword search across a whole document.
### items.jsonl
Structured document items. Each line is a JSON object with:
- `position`: sequential position in the document
- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote"
Structured document items. One JSON object per line. The row's **line index** is the item's position — `item_range` values in `toc.json` are line-slice bounds into this file.
Each row carries:
- `self_ref`: item reference (e.g. `"#/texts/5"`, `"#/tables/0"`) — used to cross-reference with `doc_item_refs` from search results
- `label`: item type — one of `"section_header"`, `"text"`, `"table"`, `"list_item"`, `"caption"`, `"formula"`, `"picture"`, `"code"`, `"footnote"`
- `text`: rendered content (tables are markdown with `|` columns)
- `page_numbers`: list of page numbers where the item appears
- `chunk_ids`: chunks that contain this item — pass to `cite()` to ground an answer that read this item directly
- `heading_level`: H-level for `section_header` rows; `0` on non-header rows
### toc.json
Section tree derived from `heading_level`: `{"doc_id", "title", "tree": [...]}` where each node has `{self_ref, level, title, page_numbers, item_range: [start, end_exclusive], chunk_ids, children}`. `item_range` is a line slice into `items.jsonl` — `items[start:end]`. `chunk_ids` aggregates the citable chunks across all items in the section — pass directly to `cite()` to ground a section-scoped answer without a corpus-wide `search()` call. `tree: []` for docs with no headers.
### Cross-referencing search results with items
Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl.
Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in `items.jsonl`. To find which section a hit lives in: locate the item by `self_ref`, take its line index, and walk `toc.json` to find the deepest node whose `item_range` contains that index.
## Strategy
1. Use `search` tool first to understand what's in the knowledge base
2. Use `execute_code` to write analysis code
3. Iterate: run code, examine output, refine approach
4. Call `cite` with chunk IDs from search results you referenced
1. Search first.
2. If the top results contain the answer, call `cite` with the supporting chunk_ids and write a concise answer.
3. Reach for `execute_code` when search results are insufficient or when the task requires computation, aggregation, traversal across documents, or section-scoped reading. From inside code you can search again with different terms, or read `items.jsonl` / `toc.json` / `content.txt` directly from the document filesystem.
4. For questions about a *known document's* structure ("which section contains X", "list the sections of doc Y", "summarise section Z"), read `/documents/{id}/toc.json` first. Each node carries `item_range` (a slice into `items.jsonl`) and `chunk_ids` (citable). Prefer this over `search()` for in-document navigation — `search()` ranks across the whole corpus and can return chunks from unrelated documents.
5. Call `cite` with the chunk_ids that ground your answer before writing the final response.
You MUST call `cite` with at least one chunk ID before producing your final answer, **unless** you are refusing for lack of information. Answers without citations are considered ungrounded. In a refusal case do **not** call `cite` — there is nothing to cite.
## Important
- Variables persist between `execute_code` calls — you can search in one call and process results in the next
- Use `print()` to output results — the output is your only feedback
- Always execute code to answer questions — don't just describe what code would do
- Use `await` for all async functions inside execute_code (search, list_documents, llm)
- When you write code, execute it — don't describe what code would do. But not every question needs code; simple lookups are best answered by `search → cite`.
- Use `await` for all async functions inside execute_code (search, list_documents)
- Use `Path.read_text()` to read files — do NOT use `open()`, `with` statements, or `collections` module
- Do NOT include chunk IDs or UUIDs in your answer text — use the `cite` tool separately
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.

6
haiku_rag_slim/haiku/rag/skills/rag.py
View file

@ -4,9 +4,9 @@ from pathlib import Path
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@ -88,7 +88,9 @@ def create_skill(
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
tools = create_skill_tools(db_path, config, RAGState, _RAG_TOOLS)
tools = create_skill_tools(
db_path, config, RAGState, _RAG_TOOLS, model=config.qa.model
)
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()

2
haiku_rag_slim/haiku/rag/skills/rag/SKILL.md
View file

@ -30,6 +30,8 @@ Retrieve a document by ID, title, or URI. Partial matches work. Use when the use
### cite
Register the chunk IDs that ground your answer. Call this BEFORE writing your final answer, with the `chunk_id` values from search results that support each claim. Every answer that uses search results must be backed by `cite`.
Use chunk_ids exactly as they appear in the search response — copy the full UUID verbatim. Do not abbreviate, paraphrase, or reconstruct chunk_ids from memory; the tool matches them as opaque strings.
## How to answer questions
1. Call `search` with relevant keywords from the question

2
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -141,6 +141,8 @@ class DocumentItemRecord(LanceModel):
text: str = Field(default="")
page_numbers: str = Field(default="[]")
picture_data: bytes | None = None
heading_level: int = Field(default=0)
tree_depth: int = Field(default=0)
def get_document_items_arrow_schema() -> pa.Schema:

10
haiku_rag_slim/haiku/rag/store/models/chunk.py
View file

@ -138,6 +138,7 @@ class SearchResult(BaseModel):
headings: list[str] | None = None
labels: list[str] = []
image_data: dict[str, str] | None = None
picture_captions: dict[str, str] = {}
@classmethod
def from_chunk(
@ -198,6 +199,15 @@ class SearchResult(BaseModel):
if primary_label:
parts.append(f"Type: {primary_label}")
# Surface picture captions when present. Order matches the binary
# attachments emitted by build_binary_parts_from_results, so the model
# can correlate caption ↔ attached image by position (BinaryContent
# identifiers don't survive serialization to the OpenAI vision API).
if self.picture_captions:
for self_ref, caption in self.picture_captions.items():
if caption:
parts.append(f"Figure caption ({self_ref}): {caption}")
# The actual content
parts.append(f"Content:\n{self.content}")

62
haiku_rag_slim/haiku/rag/store/models/citation.py Normal file
View file

@ -0,0 +1,62 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
if TYPE_CHECKING:
from haiku.rag.store.models import SearchResult
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding.
Used by the rag and analysis skills and rendered by the CLI / chat
application. The optional index field supports UI display ordering.
``picture_refs`` lists the ``self_ref`` values of picture items in the
cited chunk. Empty for text-only citations. UIs can fetch the picture
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
and render them alongside the text content.
"""
index: int | None = None
document_id: str
chunk_id: str
document_uri: str
document_title: str | None = None
page_numbers: list[int] = Field(default_factory=list)
headings: list[str] | None = None
content: str
picture_refs: list[str] = Field(default_factory=list)
def resolve_citations(
cited_chunk_ids: list[str],
search_results: "list[SearchResult]",
) -> list[Citation]:
"""Resolve chunk IDs to full Citation objects with metadata."""
by_id = {r.chunk_id: r for r in search_results if r.chunk_id}
citations = []
for raw_id in cited_chunk_ids:
chunk_id = raw_id.strip("[]")
r = by_id.get(chunk_id)
if not r:
continue
picture_refs = [
ref for ref in r.doc_item_refs if ref.startswith(PICTURE_REF_PREFIX)
]
citations.append(
Citation(
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers,
headings=r.headings,
content=r.content,
picture_refs=picture_refs,
)
)
return citations

13
haiku_rag_slim/haiku/rag/store/models/document_item.py
View file

@ -7,6 +7,9 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, NodeItem, PictureItem
PICTURE_REF_PREFIX = "#/pictures/"
class DocumentItem(BaseModel):
document_id: str
position: int
@ -15,6 +18,8 @@ class DocumentItem(BaseModel):
text: str = ""
page_numbers: list[int] = []
picture_data: bytes | None = None
heading_level: int = 0
tree_depth: int = 0
def _picture_description_text(item: "PictureItem") -> str | None:
@ -110,12 +115,12 @@ def extract_items(
fall-back lookup against ``existing_picture_data`` (keyed by ``self_ref``)
preserves the bytes that were captured at original ingest time.
"""
from docling_core.types.doc.document import PictureItem
from docling_core.types.doc.document import PictureItem, SectionHeaderItem
existing = existing_picture_data or {}
items: list[DocumentItem] = []
for position, (item, _level) in enumerate(docling_doc.iterate_items()):
for position, (item, level) in enumerate(docling_doc.iterate_items()):
label = getattr(item, "label", None)
label_str = str(label.value) if hasattr(label, "value") else str(label or "")
@ -134,6 +139,8 @@ def extract_items(
if picture_data is None:
picture_data = existing.get(item.self_ref)
heading_level = item.level if isinstance(item, SectionHeaderItem) else 0
items.append(
DocumentItem(
document_id=document_id,
@ -143,6 +150,8 @@ def extract_items(
text=text,
page_numbers=sorted(page_numbers),
picture_data=picture_data,
heading_level=heading_level,
tree_depth=level,
)
)

32
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -351,6 +351,38 @@ class ChunkRepository:
chunks.sort(key=lambda c: c.order)
return chunks
async def get_chunk_ids_by_self_ref_grouped(
self, document_ids: list[str]
) -> dict[str, dict[str, list[str]]]:
"""For each document, build a self_ref → [chunk_id, ...] index.
One query across all requested documents. The map lets items.jsonl
rows expose which chunks contain them, so callers can bridge from an
item to a `cite`-acceptable chunk_id without a separate search.
"""
from haiku.rag.utils import escape_sql_string
if not document_ids:
return {}
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
rows = await (
self.store.chunks_table.query()
.select(["id", "document_id", "metadata"])
.where(f"document_id IN ({safe_ids})")
.to_list()
)
index: dict[str, dict[str, list[str]]] = {}
for row in rows:
did = row["document_id"]
md = json.loads(row.get("metadata") or "{}")
refs = md.get("doc_item_refs") or []
doc_index = index.setdefault(did, {})
for ref in refs:
doc_index.setdefault(ref, []).append(row["id"])
return index
async def count_by_document_id(self, document_id: str) -> int:
"""Count the number of chunks for a specific document."""
df = await (

36
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

@ -14,6 +14,8 @@ _METADATA_COLUMNS = [
"label",
"text",
"page_numbers",
"heading_level",
"tree_depth",
]
@ -31,6 +33,8 @@ class DocumentItemRepository:
label=row.get("label", ""),
text=row.get("text", ""),
page_numbers=json.loads(row.get("page_numbers", "[]")),
heading_level=row.get("heading_level", 0) or 0,
tree_depth=row.get("tree_depth", 0) or 0,
)
async def create_items(self, document_id: str, items: list[DocumentItem]) -> None:
@ -48,6 +52,8 @@ class DocumentItemRepository:
text=item.text,
page_numbers=json.dumps(item.page_numbers),
picture_data=item.picture_data,
heading_level=item.heading_level,
tree_depth=item.tree_depth,
)
for item in items
]
@ -200,3 +206,33 @@ class DocumentItemRepository:
if data:
result[row["self_ref"]] = data
return result
async def get_text_for_refs(
self, document_id: str, refs: list[str]
) -> dict[str, str]:
"""Fetch the ``text`` field for multiple self_refs within a single document.
Returns ``{self_ref: text}`` for refs whose text is non-empty. Used
alongside ``get_pictures_for_chunk`` to label figures in agent-facing
search results: picture items carry their VLM-generated caption in
the ``text`` field, and the OpenAI vision message format has no
identifier on binary parts, so the caption text is the only signal a
model can use to correlate a description with the picture it sees.
"""
if not refs:
return {}
safe_id = escape_sql_string(document_id)
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
rows = await (
self.store.document_items_table.query()
.select(["self_ref", "text"])
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
.to_list()
)
result: dict[str, str] = {}
for row in rows:
text = row.get("text") or ""
if text:
result[row["self_ref"]] = text
return result

4
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -87,6 +87,9 @@ from haiku.rag.store.upgrades.v0_40_0 import (
from haiku.rag.store.upgrades.v0_45_0 import (
upgrade_extract_picture_bytes as upgrade_0_45_0_extract_picture_bytes,
)
from haiku.rag.store.upgrades.v0_48_0 import (
upgrade_backfill_heading_hierarchy as upgrade_0_48_0_heading_hierarchy,
)
upgrades.append(upgrade_0_20_0_docling)
upgrades.append(upgrade_0_23_1_contextualize)
@ -94,3 +97,4 @@ upgrades.append(upgrade_0_25_0_compress)
upgrades.append(upgrade_0_38_0_split_pages)
upgrades.append(upgrade_0_40_0_document_items)
upgrades.append(upgrade_0_45_0_extract_picture_bytes)
upgrades.append(upgrade_0_48_0_heading_hierarchy)

171
haiku_rag_slim/haiku/rag/store/upgrades/v0_48_0.py Normal file
View file

@ -0,0 +1,171 @@
import logging
import pyarrow as pa
from lancedb.index import BTree
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.upgrades import Upgrade
from haiku.rag.utils import escape_sql_string
logger = logging.getLogger(__name__)
PROGRESS_INTERVAL = 10
async def _ensure_columns(store: Store) -> None:
"""Add heading_level + tree_depth (int64) columns if missing. Idempotent."""
arrow_schema = await store.document_items_table.schema()
existing = {f.name for f in arrow_schema}
new_fields = []
if "heading_level" not in existing:
new_fields.append(pa.field("heading_level", pa.int64())) # pragma: no cover
if "tree_depth" not in existing:
new_fields.append(pa.field("tree_depth", pa.int64())) # pragma: no cover
if new_fields:
# Pre-0.48.0 DBs only — fresh DBs declare both columns in _init_tables.
await store.document_items_table.add_columns(
pa.schema(new_fields)
) # pragma: no cover
async def _ensure_indexes(store: Store) -> None:
"""Ensure BTree scalar indexes exist on the hot document_items lookup columns.
Fresh DBs created via ``_init_tables`` get these on first creation, but
DBs that predate that code (or were downloaded as pre-built artifacts) were
table-scanning every per-doc query — visible as ~100–300 ms even for small
docs. Built before the heading_level backfill so the per-doc WHERE clauses
in the backfill loop benefit from it.
"""
for column in ("document_id", "position", "self_ref"):
await store.document_items_table.create_index(
column, config=BTree(), replace=True
)
async def _apply_backfill_heading_hierarchy(store: Store) -> None:
"""Add heading_level + tree_depth columns and backfill from docling structure.
For each document with a docling_document blob, decompress it, re-run
extract_items, and update the existing items rows with the two new ints.
Documents without docling (plain-text adds) are skipped; their rows keep
the column-add default (NULL, materialised as 0 by the repository).
Idempotent: re-running on migrated rows yields identical values.
"""
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.store.compression import decompress_json
from haiku.rag.store.models.document_item import extract_items
await _ensure_columns(store)
await _ensure_indexes(store)
ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist()
ids = [row["id"] for row in ids]
total = len(ids)
logger.info("Backfilling heading_level + tree_depth across %d documents", total)
backfilled = 0
skipped = 0
for idx, doc_id in enumerate(ids, 1):
safe_id = escape_sql_string(doc_id)
rows = await (
store.documents_table.query()
.select(["id", "docling_document"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
blob = rows[0].get("docling_document")
if not blob or not isinstance(blob, bytes):
skipped += 1
continue
try:
docling_doc = DoclingDocument.model_validate_json(decompress_json(blob))
fresh_items = extract_items(doc_id, docling_doc)
except Exception: # pragma: no cover
logger.warning(
"Failed to re-extract items for %s; skipping", doc_id, exc_info=True
)
skipped += 1
continue
if not fresh_items: # pragma: no cover
skipped += 1
continue
existing_rows = await (
store.document_items_table.query()
.where(f"document_id = '{safe_id}'")
.to_list()
)
existing_by_ref = {r["self_ref"]: r for r in existing_rows}
# If the stored rows and a fresh extract disagree on count, the
# docling parser version has drifted. Skip and let `rebuild` reconcile.
if len(existing_rows) != len(fresh_items): # pragma: no cover
logger.warning(
"Item count drift for %s (stored=%d, fresh=%d); skipping",
doc_id,
len(existing_rows),
len(fresh_items),
)
skipped += 1
continue
records: list[DocumentItemRecord] = []
for item in fresh_items:
row = existing_by_ref.get(item.self_ref)
if row is None: # pragma: no cover
continue
records.append(
DocumentItemRecord(
document_id=doc_id,
position=row["position"],
self_ref=item.self_ref,
label=row.get("label", ""),
text=row.get("text", ""),
page_numbers=row.get("page_numbers", "[]"),
picture_data=row.get("picture_data"),
heading_level=item.heading_level,
tree_depth=item.tree_depth,
)
)
if records:
await (
store.document_items_table.merge_insert(["document_id", "self_ref"])
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(records)
)
backfilled += 1
if idx % PROGRESS_INTERVAL == 0 or idx == total:
logger.info(
"Progress: %d/%d (%d backfilled, %d skipped)",
idx,
total,
backfilled,
skipped,
)
logger.info(
"Backfill complete: %d backfilled, %d skipped of %d",
backfilled,
skipped,
total,
)
upgrade_backfill_heading_hierarchy = Upgrade(
version="0.48.0",
apply=_apply_backfill_heading_hierarchy,
description=(
"Backfill heading_level + tree_depth on document_items, "
"ensure BTree indexes on document_id / position / self_ref"
),
)

145
haiku_rag_slim/haiku/rag/utils.py
View file

@ -11,8 +11,9 @@ from packaging.version import Version, parse
if TYPE_CHECKING:
from rich.console import RenderableType
from haiku.rag.agents.research.models import Citation
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.store.models.citation import Citation
def parse_model_option(value: str) -> "ModelConfig":
@ -338,10 +339,34 @@ def format_bytes(num_bytes: int) -> str:
return f"{size:.1f} PB"
CITATION_PREVIEW_CHARS = 300
def _citation_pages(c: "Citation") -> str | None:
if not c.page_numbers:
return None
if len(c.page_numbers) == 1:
return f"p. {c.page_numbers[0]}"
return f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}"
def _citation_section(c: "Citation") -> str | None:
if c.headings:
return c.headings[-1]
return None
def _citation_label(c: "Citation") -> str:
if c.document_title and c.document_uri:
return f"{c.document_title} ({c.document_uri})"
return c.document_title or c.document_uri
def format_citations(citations: "list[Citation]") -> str:
"""Format citations as plain text with preserved formatting.
Used by things like the MCP server where Rich renderables are not available.
Pictures referenced by the chunk are surfaced as ``[Figure: <ref>]`` markers.
"""
if not citations:
return ""
@ -353,34 +378,42 @@ def format_citations(citations: "list[Citation]") -> str:
title = c.document_title or c.document_uri
header = f"[{idx}] {title}"
# Location info
location_parts = []
if c.page_numbers:
if len(c.page_numbers) == 1:
location_parts.append(f"p. {c.page_numbers[0]}")
else:
location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}")
if c.headings:
location_parts.append(f"Section: {c.headings[-1]}")
pages = _citation_pages(c)
if pages:
location_parts.append(pages)
section = _citation_section(c)
if section:
location_parts.append(f"Section: {section}")
source = c.document_uri
if location_parts:
source += f" - {', '.join(location_parts)}"
lines.append(f"{header} {source}")
for ref in c.picture_refs:
lines.append(f"[Figure: {ref}]")
lines.append(c.content)
lines.append("")
return "\n".join(lines)
def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]":
"""Format citations as Rich renderables.
async def format_citations_rich(
citations: "list[Citation]",
client: "HaikuRAG | None" = None,
) -> "list[RenderableType]":
"""Format citations as Rich renderables for terminal display.
Returns a list of Rich Panel objects for direct console printing,
with content rendered as markdown for syntax highlighting.
Each citation becomes a Panel with a compact header (``[N] Title (URI) — locator``),
a body holding any referenced figures followed by a truncated text preview, and
a dimmed footer that exposes the document and chunk IDs.
When ``client`` is supplied, picture bytes for ``picture_refs`` are fetched and
rendered inline via ``textual_image``. Without a client, picture refs appear as
``[Figure: <ref>]`` text markers.
"""
from rich.markdown import Markdown
from rich.console import Group
from rich.panel import Panel
from rich.text import Text
@ -388,35 +421,49 @@ def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]"
return []
renderables: list[RenderableType] = []
renderables.append(Text("Citations", style="bold"))
renderables.append(Text(""))
renderables.append(Text("Citations", style="bold green"))
renderables.append(Text(""))
for c in citations:
# Build header with IDs
header = Text()
header.append("doc: ", style="dim")
header.append(c.document_id, style="cyan")
header.append(" chunk: ", style="dim")
header.append(c.chunk_id, style="cyan")
for i, c in enumerate(citations):
if i > 0:
renderables.append(Text(""))
idx = c.index if c.index is not None else (i + 1)
# Location info for subtitle
location_parts = []
if c.page_numbers:
if len(c.page_numbers) == 1:
location_parts.append(f"p. {c.page_numbers[0]}")
else:
location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}")
if c.headings:
location_parts.append(f"Section: {c.headings[-1]}")
header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"]
pages = _citation_pages(c)
if pages:
header_parts.append(pages)
section = _citation_section(c)
if section:
header_parts.append(f"§{section}")
header = Text(" — ".join(header_parts), style="bold")
body: list[RenderableType] = []
for ref in c.picture_refs:
image_renderable = await _render_picture(client, c.document_id, ref)
body.append(
image_renderable
if image_renderable
else Text(f"[Figure: {ref}]", style="italic dim")
)
preview = c.content
if len(preview) > CITATION_PREVIEW_CHARS:
preview = preview[:CITATION_PREVIEW_CHARS].rstrip() + "…"
body.append(Text(preview))
footer = Text()
footer.append("doc: ", style="dim")
footer.append(c.document_id, style="dim cyan")
footer.append(" chunk: ", style="dim")
footer.append(c.chunk_id, style="dim cyan")
subtitle = c.document_uri
if c.document_title:
subtitle = f"{c.document_title} ({c.document_uri})"
if location_parts:
subtitle += f" - {', '.join(location_parts)}"
panel = Panel(
Markdown(c.content),
Group(*body),
title=header,
subtitle=subtitle,
title_align="left",
subtitle=footer,
subtitle_align="left",
border_style="dim",
)
@ -425,6 +472,28 @@ def format_citations_rich(citations: "list[Citation]") -> "list[RenderableType]"
return renderables
async def _render_picture(
client: "HaikuRAG | None", document_id: str, ref: str
) -> "RenderableType | None":
"""Fetch a picture and return a Rich renderable, or None on failure/no client."""
if client is None:
return None
from io import BytesIO
from PIL import Image as PILImage
from textual_image.renderable import Image as RichImage
data = await client.document_item_repository.get_picture_bytes(document_id, ref)
if not data:
return None
try:
pil = PILImage.open(BytesIO(data))
pil.load()
except Exception:
return None
return RichImage(pil)
def get_default_data_dir() -> Path:
"""Get the user data directory for the current system platform.

9
haiku_rag_slim/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.47.0"
version = "0.48.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -54,7 +54,12 @@ zeroentropy = ["zeroentropy>=0.1.0a11"]
jina = ["transformers>=4.40.0", "torch>=2.0.0"]
cross-encoder = ["sentence-transformers>=3.0.0"]
# TUI (chat and inspect commands)
tui = ["textual>=8.2.4", "textual-image>=0.8.5"]
tui = [
"textual>=8.2.4",
"textual-image>=0.8.5",
"tree-sitter>=0.25.2",
"tree-sitter-json>=0.24.8",
]
# Model providers (delegated to pydantic-ai-slim)
anthropic = ["pydantic-ai-slim[anthropic]"]
groq = ["pydantic-ai-slim[groq]"]

6
mkdocs.yml
View file

@ -62,7 +62,7 @@ nav:
- Configuration:
- configuration/index.md
- Providers: configuration/providers.md
- Search and Question Answering: configuration/qa-research.md
- Search and Question Answering: configuration/qa.md
- Document Processing: configuration/processing.md
- Storage: configuration/storage.md
- Prompts: configuration/prompts.md
@ -70,9 +70,7 @@ nav:
- Python: python.md
- Custom Pipelines: custom-pipelines.md
- Tuning: tuning.md
- Agents:
- agents/index.md
- Analysis Agent: agents/analysis.md
- Analysis: agents/analysis.md
- Skills:
- skills/index.md
- RAG: skills/rag.md

8
pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.47.0"
version = "0.48.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -30,7 +30,7 @@ classifiers = [
]
dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.47.0",
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui,cross-encoder]==0.48.0",
]
[project.scripts]
@ -38,8 +38,8 @@ haiku-rag = "haiku.rag.cli:cli"
[project.optional-dependencies]
tui = ["textual>=8.2.4"]
s3 = ["haiku.rag-slim[s3]==0.47.0"]
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.47.0"]
s3 = ["haiku.rag-slim[s3]==0.48.0"]
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.48.0"]
[build-system]
requires = ["hatchling"]

0
tests/agents/analysis/__init__.py
View file

311
tests/agents/analysis/test_agent.py
View file

@ -1,311 +0,0 @@
from pathlib import Path
import pytest
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 CodeExecution, RawAnalysisResult
from haiku.rag.config import AppConfig, Config
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_analysis")
class TestCreateAnalysisAgent:
def test_creates_agent(self):
agent = create_analysis_agent(Config)
assert isinstance(agent, Agent)
assert agent.deps_type is AnalysisDeps
assert agent.output_type is RawAnalysisResult
def test_agent_has_execute_code_tool(self):
agent = create_analysis_agent(Config)
tool_names = list(agent._function_toolset.tools.keys())
assert "execute_code" in tool_names
class TestCodeExecutionModel:
def test_code_execution_has_correct_fields(self):
"""Test that CodeExecution has all expected fields."""
execution = CodeExecution(
code="print('hello')",
stdout="hello\n",
stderr="",
success=True,
)
assert execution.code == "print('hello')"
assert execution.stdout == "hello\n"
assert execution.stderr == ""
assert execution.success is True
class TestClientAnalysisIntegration:
"""Integration tests for client.analyze() method."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
"""Test analysis agent can count documents.
Agent program:
docs = list_documents(limit=1000)
print(len(docs))
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document("First document about cats.", title="Doc 1")
await client.create_document("Second document about dogs.", title="Doc 2")
await client.create_document("Third document about birds.", title="Doc 3")
result = await client.analyze("How many documents are in the database?")
assert "3" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
"""Test analysis agent can perform aggregation across documents.
Agent program:
import re
revs = {}
for d in ['Q1 Report', 'Q2 Report', 'Q3 Report']:
content = get_document(d)
if content:
vals = re.findall(r'\\$([\\d,]+)', content)
if vals:
rev = sum(int(v.replace(',', '')) for v in vals)
else:
rev = None
else:
rev = None
revs[d] = rev
print(revs)
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"Sales report Q1: Revenue was $100,000.", title="Q1 Report"
)
await client.create_document(
"Sales report Q2: Revenue was $150,000.", title="Q2 Report"
)
await client.create_document(
"Sales report Q3: Revenue was $200,000.", title="Q3 Report"
)
result = await client.analyze(
"What is the total revenue across all quarterly reports?"
)
assert "450" in result.answer or "450,000" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
"""Test analysis agent respects filter parameter.
Agent program:
docs = list_documents(limit=1000)
print(len(docs))
print(docs[:5])
The filter is applied via context, so list_documents() only sees "Cats".
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document("Cat document.", title="Cats")
await client.create_document("Dog document.", title="Dogs")
await client.create_document("Bird document.", title="Birds")
result = await client.analyze(
"How many documents are available?",
filter="title = 'Cats'",
)
assert "1" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_search_and_identify_source(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can search and identify source documents."""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"The quick brown fox jumps over the lazy dog.",
title="Animal Facts",
)
result = await client.analyze(
"Search for content about animals and tell me "
"which document it came from."
)
assert "Animal Facts" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_semantic_analysis_with_llm(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can use llm() for semantic analysis combined with computation.
Agent program:
docs = list_documents(limit=100)
print(len(docs))
print([d['title'] for d in docs[:20]])
sentiments = {}
for title in ['Q1 Update', 'Q2 Update', 'Q3 Update']:
content = get_document(title)
if content:
result = llm(f"Classify sentiment as positive/negative/mixed: {content}")
sentiments[title] = result
print(sentiments)
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"The new product launch exceeded expectations. Sales grew 40% "
"and customer feedback has been overwhelmingly positive. "
"Team morale is at an all-time high.",
title="Q1 Update",
)
await client.create_document(
"We faced significant challenges this quarter. Supply chain issues "
"caused delays, and we missed our revenue target by 15%. "
"Several key employees left the company.",
title="Q2 Update",
)
await client.create_document(
"Mixed results this quarter. While product quality improved, "
"marketing campaigns underperformed. Revenue was flat compared "
"to last year but customer retention increased.",
title="Q3 Update",
)
result = await client.analyze(
"Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?"
)
# Should identify: Q1=positive, Q2=negative, Q3=mixed
assert "positive" in result.answer.lower()
assert "negative" in result.answer.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test analysis agent can use search() to find content and extract information.
Agent program:
results = search("document element types", limit=20)
print(len(results))
for r in results[:5]:
print(r['document_title'], r['chunk_id'], r['score'])
print(r['content'][:200])
results = search("DocBank element types", limit=10)
...
"""
from haiku.rag.client import HaikuRAG
pdf_path = Path("tests/data/doclaynet.pdf")
config = AppConfig()
config.processing.conversion_options.do_ocr = False
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path)
result = await client.analyze(
"Search for content about document element types or labels. "
"What are all the different document element types mentioned? "
"List them all."
)
# The doclaynet.pdf defines exactly 11 class labels for document elements
# Normalize Unicode hyphens (U+2011 non-breaking hyphen) to regular hyphens
answer_lower = result.answer.lower().replace("\u2011", "-")
expected_labels = [
"caption",
"footnote",
"formula",
"list-item",
"page-footer",
"page-header",
"picture",
"section-header",
"table",
"text",
"title",
]
# Check that the agent found at least 6 of the 11 labels
# (LLM summaries may not always include all labels)
found_labels = [
label
for label in expected_labels
if label in answer_lower or label.replace("-", " ") in answer_lower
]
assert len(found_labels) >= 6, (
f"Expected at least 6 labels, found {len(found_labels)}: {found_labels}"
)
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_analyze_with_preloaded_documents(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can use pre-loaded documents variable.
Agent program:
if 'documents' in dir():
for doc in documents:
print(doc['title'], len(doc['content']))
else:
print('No preloaded documents')
"""
from haiku.rag.client import HaikuRAG
config = AppConfig()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document(
"The company was founded in 1985 by Jane Smith.",
title="Company History",
)
await client.create_document(
"Our mission is to make technology accessible to everyone.",
title="Mission Statement",
)
result = await client.analyze(
"Using the pre-loaded documents variable, "
"tell me when was the company founded and what is their mission?",
documents=["Company History", "Mission Statement"],
)
assert "1985" in result.answer
assert (
"accessible" in result.answer.lower()
or "technology" in result.answer.lower()
)

32
tests/agents/analysis/test_models.py
View file

@ -1,32 +0,0 @@
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
class TestCodeExecution:
def test_create_successful_execution(self):
execution = CodeExecution(
code="print('hello')",
stdout="hello\n",
stderr="",
success=True,
)
assert execution.code == "print('hello')"
assert execution.stdout == "hello\n"
assert execution.stderr == ""
assert execution.success is True
def test_create_failed_execution(self):
execution = CodeExecution(
code="1/0",
stdout="",
stderr="ZeroDivisionError: division by zero",
success=False,
)
assert execution.success is False
assert "ZeroDivisionError" in execution.stderr
class TestAnalysisResult:
def test_create_result(self):
result = AnalysisResult(answer="The answer is 42", program="print(42)")
assert result.answer == "The answer is 42"
assert result.program == "print(42)"

69
tests/agents/qa/test_qa.py
View file

@ -1,69 +0,0 @@
from pathlib import Path
import pytest
from datasets import Dataset
from evaluations.evaluators import LLMJudge
from haiku.rag.agents.qa.agent import QuestionAnswerAgent
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import ModelConfig
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_qa")
@pytest.mark.asyncio
async def test_get_qa_agent_factory(temp_db_path):
"""Test get_qa_agent factory function creates a properly configured agent."""
from haiku.rag.agents.qa import get_qa_agent
async with HaikuRAG(temp_db_path, create=True) as client:
agent = get_qa_agent(client, Config)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
# Verify internal client is set correctly
assert agent._client is client
@pytest.mark.asyncio
async def test_get_qa_agent_with_custom_prompt(temp_db_path):
"""Test get_qa_agent factory with custom system prompt."""
from haiku.rag.agents.qa import get_qa_agent
async with HaikuRAG(temp_db_path, create=True) as client:
custom_prompt = "You are a custom QA assistant."
agent = get_qa_agent(client, Config, system_prompt=custom_prompt)
assert agent is not None
assert isinstance(agent, QuestionAnswerAgent)
assert agent._system_prompt == custom_prompt
@pytest.mark.vcr()
async def test_qa_ollama(allow_model_requests, qa_corpus: Dataset, temp_db_path):
"""Test Ollama QA with LLM judge (VCR recorded)."""
async with HaikuRAG(temp_db_path, create=True) as client:
qa = QuestionAnswerAgent(
client,
ModelConfig(provider="ollama", name="gpt-oss", enable_thinking=True),
)
llm_judge = LLMJudge()
doc = qa_corpus[1]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
question = doc["question"]
expected_answer = doc["answer"]
answer, _ = await qa.answer(question)
is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer)
assert is_equivalent, (
f"Generated answer not equivalent to expected answer.\nQuestion: {question}\nGenerated: {answer}\nExpected: {expected_answer}"
)

159
tests/agents/research/test_models.py
View file

@ -1,159 +0,0 @@
from haiku.rag.agents.research.models import Citation, SearchAnswer, resolve_citations
from haiku.rag.store.models import SearchResult
class TestCitation:
"""Tests for unified Citation class."""
def test_citation_without_index(self):
"""Test Citation can be created without index (research graph use case)."""
citation = Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title="Test Document",
page_numbers=[1, 2],
headings=["Introduction"],
content="Test content",
)
assert citation.document_id == "doc-1"
assert citation.chunk_id == "chunk-1"
assert citation.document_uri == "test.md"
assert citation.document_title == "Test Document"
assert citation.page_numbers == [1, 2]
assert citation.headings == ["Introduction"]
assert citation.content == "Test content"
assert citation.index is None
def test_citation_with_index(self):
"""Test Citation can be created with index (chat use case)."""
citation = Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
content="Test content",
)
assert citation.index == 1
assert citation.document_id == "doc-1"
def test_citation_index_defaults_to_none(self):
"""Test Citation index defaults to None."""
citation = Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
content="Test content",
)
assert citation.index is None
def test_citation_serialization_includes_index_when_set(self):
"""Test Citation serialization includes index when set."""
citation = Citation(
index=2,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
content="Test content",
)
data = citation.model_dump()
assert data["index"] == 2
def test_citation_deserialization_from_dict_with_index(self):
"""Test Citation can be deserialized from dict with index (AG-UI state sync)."""
data = {
"index": 1,
"document_id": "doc-1",
"chunk_id": "chunk-1",
"document_uri": "test.md",
"document_title": "Test Doc",
"page_numbers": [1, 2],
"headings": ["Intro"],
"content": "Test content",
}
citation = Citation.model_validate(data)
assert citation.index == 1
assert citation.document_id == "doc-1"
class TestSearchAnswerPrimarySource:
"""Tests for SearchAnswer.primary_source property."""
def test_primary_source_returns_title_when_available(self):
"""Test primary_source returns first citation's title."""
answer = SearchAnswer(
query="test query",
answer="test answer",
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title="Test Document",
content="content",
),
],
)
assert answer.primary_source == "Test Document"
def test_primary_source_returns_uri_when_no_title(self):
"""Test primary_source returns URI when title is None."""
answer = SearchAnswer(
query="test query",
answer="test answer",
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title=None,
content="content",
),
],
)
assert answer.primary_source == "test.md"
def test_primary_source_returns_none_when_no_citations(self):
"""Test primary_source returns None when no citations."""
answer = SearchAnswer(
query="test query",
answer="test answer",
citations=[],
)
assert answer.primary_source is None
class TestResolveCitations:
"""Tests for resolve_citations function."""
def _make_result(self, chunk_id: str) -> SearchResult:
return SearchResult(
content="test content",
score=1.0,
chunk_id=chunk_id,
document_id="doc-1",
document_uri="test.md",
document_title="Test Doc",
)
def test_resolves_exact_ids(self):
results = [self._make_result("abc123")]
citations = resolve_citations(["abc123"], results)
assert len(citations) == 1
assert citations[0].chunk_id == "abc123"
def test_strips_brackets_from_ids(self):
results = [self._make_result("abc123")]
citations = resolve_citations(["[abc123]"], results)
assert len(citations) == 1
assert citations[0].chunk_id == "abc123"
def test_skips_unmatched_ids(self):
results = [self._make_result("abc123")]
citations = resolve_citations(["nonexistent"], results)
assert len(citations) == 0
def test_empty_cited_chunks(self):
results = [self._make_result("abc123")]
citations = resolve_citations([], results)
assert len(citations) == 0

31
tests/agents/research/test_plan_prompt_selection.py
View file

@ -1,31 +0,0 @@
from haiku.rag.agents.research.prompts import (
ITERATIVE_PLAN_PROMPT,
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
)
def test_iterative_plan_prompt_proposes_first_question():
"""ITERATIVE_PLAN_PROMPT should instruct to propose the first question."""
assert "first question" in ITERATIVE_PLAN_PROMPT.lower()
assert "is_complete=False" in ITERATIVE_PLAN_PROMPT
def test_iterative_plan_prompt_with_context_evaluates_evidence():
"""ITERATIVE_PLAN_PROMPT_WITH_CONTEXT should evaluate prior answers."""
assert "prior_answers" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
assert (
"evaluat" in ITERATIVE_PLAN_PROMPT_WITH_CONTEXT.lower()
) # matches evaluate/evaluating
def test_prompt_selection_uses_context_prompt_with_prior_answers():
"""When prior_answers exist, should use ITERATIVE_PLAN_PROMPT_WITH_CONTEXT."""
has_prior_answers = True
effective_plan_prompt = (
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT
if has_prior_answers
else ITERATIVE_PLAN_PROMPT
)
assert effective_plan_prompt == ITERATIVE_PLAN_PROMPT_WITH_CONTEXT

109
tests/agents/research/test_research_graph.py
View file

@ -1,109 +0,0 @@
from pathlib import Path
import pytest
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import ResearchReport
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(
Path(__file__).parent.parent.parent / "cassettes" / "test_research_graph"
)
@pytest.mark.vcr()
async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
"""Test research graph with real LLM calls recorded via VCR."""
graph = build_research_graph()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = qa_corpus[0]
await client.create_document(
content=doc["document_extracted"], uri=doc["document_id"]
)
state = ResearchState(
context=ResearchContext(original_question=doc["question"]),
max_iterations=1,
max_concurrency=1,
)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
assert result is not None
assert isinstance(result, ResearchReport)
assert result.title
assert result.executive_summary
def test_iterative_plan_result_model():
"""Test IterativePlanResult model validation."""
from haiku.rag.agents.research.models import IterativePlanResult
# Test complete state
complete = IterativePlanResult(
is_complete=True,
next_question=None,
reasoning="All aspects covered.",
)
assert complete.is_complete is True
assert complete.next_question is None
# Test continue state
continue_result = IterativePlanResult(
is_complete=False,
next_question="What are the specific requirements?",
reasoning="Need more details.",
)
assert continue_result.is_complete is False
assert continue_result.next_question == "What are the specific requirements?"
def test_build_research_graph_returns_graph():
"""Test build_research_graph returns a valid Graph instance."""
from pydantic_graph.beta import Graph
graph = build_research_graph()
assert graph is not None
assert isinstance(graph, Graph)
def test_format_context_for_prompt_basic():
"""Test format_context_for_prompt with basic context."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(original_question="What is X?")
result = format_context_for_prompt(context)
assert "<context>" in result
assert "What is X?" in result
def test_format_context_for_prompt_with_prior_answers():
"""Test format_context_for_prompt includes prior_answers."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
from haiku.rag.agents.research.models import SearchAnswer
context = ResearchContext(original_question="Main question?")
context.add_qa_response(
SearchAnswer(
query="Sub question?",
answer="The answer is here.",
confidence=0.9,
)
)
result = format_context_for_prompt(context)
assert "<prior_answers>" in result
assert "Sub question?" in result
assert "The answer is here." in result

125
tests/agents/research/test_search_filter.py
View file

@ -1,125 +0,0 @@
from pathlib import Path
import pytest
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_search_filter")
@pytest.fixture
async def client_with_docs(temp_db_path):
"""Create a client with two distinct documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add two documents with distinct content
doc1 = await client.create_document(
"Document about cats: Cats are small furry mammals that purr.",
title="Cat Facts",
)
doc2 = await client.create_document(
"Document about dogs: Dogs are loyal companions that bark.",
title="Dog Facts",
)
yield client, doc1.id, doc2.id
@pytest.mark.vcr()
async def test_search_filter_restricts_results(client_with_docs):
"""Test that search_filter restricts search to specified documents."""
client, doc1_id, doc2_id = client_with_docs
# Search without filter - should find both
results_all = await client.search("animals mammals companions")
assert len(results_all) >= 1
# Search with filter for doc1 only
filter_doc1 = f"id = '{doc1_id}'"
results_filtered = await client.search(
"animals mammals companions", filter=filter_doc1
)
# All results should be from doc1
for result in results_filtered:
assert result.document_id == doc1_id
@pytest.mark.vcr()
@pytest.mark.asyncio
async def test_research_graph_uses_search_filter(
allow_model_requests, client_with_docs
):
"""Test that research graph passes search_filter to search operations."""
client, doc1_id, doc2_id = client_with_docs
# Track search calls to verify filter is passed
search_calls = []
original_search = client.search
async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
search_calls.append({"query": query, "filter": filter})
return await original_search(query, limit, search_type, filter)
client.search = tracking_search
graph = build_research_graph()
# Create state with search_filter
filter_clause = f"id = '{doc1_id}'"
state = ResearchState(
context=ResearchContext(original_question="Tell me about animals"),
max_iterations=1,
search_filter=filter_clause,
)
deps = ResearchDeps(client=client)
await graph.run(state=state, deps=deps)
# Verify search was called with the filter
assert len(search_calls) > 0, "Expected search to be called"
for call in search_calls:
assert call["filter"] == filter_clause, (
f"Expected filter '{filter_clause}', got '{call['filter']}'"
)
@pytest.mark.vcr()
@pytest.mark.asyncio
async def test_search_filter_none_searches_all(allow_model_requests, client_with_docs):
"""Test that search_filter=None searches all documents."""
client, doc1_id, doc2_id = client_with_docs
# Track search calls
search_calls = []
original_search = client.search
async def tracking_search(query, limit=None, search_type="hybrid", filter=None):
search_calls.append({"query": query, "filter": filter})
return await original_search(query, limit, search_type, filter)
client.search = tracking_search
graph = build_research_graph()
# Create state without search_filter (None)
state = ResearchState(
context=ResearchContext(original_question="Tell me about animals"),
max_iterations=1,
search_filter=None,
)
deps = ResearchDeps(client=client)
await graph.run(state=state, deps=deps)
# Verify search was called with None filter
assert len(search_calls) > 0, "Expected search to be called"
for call in search_calls:
assert call["filter"] is None, f"Expected filter None, got '{call['filter']}'"

1827
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_aggregation.yaml
View file

File diff suppressed because one or more lines are too long

578
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_count_documents.yaml
View file

File diff suppressed because one or more lines are too long

2677
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_search_and_extract.yaml
View file

File diff suppressed because one or more lines are too long

550
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_search_and_identify_source.yaml
View file

File diff suppressed because one or more lines are too long

1055
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_semantic_analysis_with_llm.yaml
View file

File diff suppressed because one or more lines are too long

582
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

806
tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_with_preloaded_documents.yaml
View file

File diff suppressed because one or more lines are too long

851
tests/cassettes/test_client/test_client_ask.yaml
View file

File diff suppressed because one or more lines are too long

2003
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_aggregation.yaml Normal file
View file

File diff suppressed because one or more lines are too long

3395
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_count_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

3126
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_search_and_extract.yaml Normal file
View file

File diff suppressed because one or more lines are too long

1844
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_search_and_identify_source.yaml Normal file
View file

File diff suppressed because one or more lines are too long

60512
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_with_filter.yaml Normal file
View file

File diff suppressed because one or more lines are too long

Some files were not shown because too many files have changed in this diff Show more