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

replace haiku.skills with native Pydantic AI capabilities

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-07-17 17:12:25 +03:00
parent ec78641cc4
commit 9deb1f2bd4
No known key found for this signature in database
88 changed files with 18723 additions and 87462 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
CHANGELOG.md
View file

@ -28,6 +28,21 @@
### Added
- `hotpotqa` evaluation dataset.
- Native deferred Pydantic AI `RAGCapability` and `AnalysisCapability` implementations under `haiku.rag.capabilities`, with namespaced host state and lazy per-run database and sandbox resources.
- Prior-turn RAG and analysis tool results are compacted before model requests while current-turn evidence remains intact.
### Changed
- Require `pydantic-ai-slim>=2.11,<3`; the `vertexai` optional extra now installs Pydantic AI's `google` extra.
- The chat TUI consumes native Pydantic AI stream events. The web example uses the standard `AGUIAdapter` and emits one final state snapshot instead of forwarding sub-agent activity and per-tool state events.
- Chat capability selection is now `haiku-rag chat --capability/-c {rag,analysis}`. Migrate from `--skill/-s`.
- Evaluation targets are now `rag-capability` and `analysis-capability`, and the model override is `--capability-model`. Migrate from `rag-skill`, `analysis-skill`, and `--skill-model`.
### Removed
- The `haiku.skills` dependency, `haiku.rag.skills` modules, Python entry-point discovery, and sub-agent execution layer. Migrate `create_skill(...)` plus `SkillToolset` usage to `haiku.rag.capabilities.*.create_capability(...)` passed through `Agent(capabilities=[...])`.
- The `haiku-rag create-skill` package generator. Compose native capabilities directly and package application-specific instructions and data in the consuming project.
- Legacy sub-agent `ActivitySnapshotEvent` plumbing and per-tool `StateDeltaEvent` generation.
### Changed

4
README.md
View file

@ -96,7 +96,7 @@ async with HaikuRAG("knowledge.lancedb", create=True) as rag:
print(f" [{cite.chunk_id}] p.{cite.page_numbers}: {cite.content[:80]}")
```
For details on the skills the client wraps, see the [Skills docs](https://ggozad.github.io/haiku.rag/skills/).
For direct agent composition, see the [capabilities documentation](https://ggozad.github.io/haiku.rag/capabilities/).
## MCP Server
@ -137,7 +137,7 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - YAML reference
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Skills](https://ggozad.github.io/haiku.rag/skills/) - The RAG and analysis skills the client wraps
- [Capabilities](https://ggozad.github.io/haiku.rag/capabilities/) - Native Pydantic AI RAG and analysis capabilities
- [Tuning](https://ggozad.github.io/haiku.rag/tuning/) - Retrieval and answer-quality tuning
- [Ingester](https://ggozad.github.io/haiku.rag/ingester/) - Production ingester for continuous indexing from FS, HTTP, S3, and WebDAV
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration

63
app/backend/main.py
View file

@ -2,7 +2,9 @@ import asyncio
import logging
import os
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from ag_ui.core import EventType, StateSnapshotEvent
from dotenv import find_dotenv, load_dotenv
@ -16,18 +18,12 @@ from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability
from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rag import create_skill, get_agent_preamble
from haiku.rag.telemetry import configure as configure_telemetry
from haiku.rag.utils import get_model
from haiku.skills import (
SkillDeps,
SkillToolset,
run_agui_stream,
)
from haiku.skills.prompts import build_system_prompt
load_dotenv(find_dotenv(usecwd=True))
@ -75,17 +71,18 @@ async def get_client() -> HaikuRAG:
return _client
# Create skill, toolset, and agent
skill = create_skill(db_path=db_path, config=Config)
toolset = SkillToolset(skills=[skill])
@dataclass
class AppDeps:
state: dict[str, Any] = field(default_factory=dict)
capability = create_capability(db_path=db_path, config=Config)
agent = Agent(
get_model(Config.qa.model, Config),
instructions=build_system_prompt(
toolset.skill_catalog, preamble=get_agent_preamble(Config)
),
toolsets=[toolset],
deps_type=SkillDeps,
instructions=AGENT_PREAMBLE,
capabilities=[capability],
deps_type=AppDeps,
)
@ -98,33 +95,21 @@ async def stream_chat(request: Request) -> Response:
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
incoming_state = run_input.state if isinstance(run_input.state, dict) else {}
incoming_state.setdefault("rag", RAGState().model_dump(mode="json"))
deps = AppDeps(state=incoming_state)
async def event_stream():
async with run_agui_stream(
adapter, toolset=toolset, deps=SkillDeps(state=incoming_state)
) as stream:
# Emit a STATE_SNAPSHOT after RUN_STARTED so the client holds every
# namespace object before any STATE_DELTA patches into it. Without it,
# the first `add /rag/<field>/...` fails against a missing parent.
if incoming_state:
toolset.restore_state_snapshot(incoming_state)
snapshot = StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=toolset.build_state_snapshot(),
)
async def with_final_state():
async for event in adapter.run_stream(deps=deps):
if getattr(event, "type", None) == EventType.RUN_FINISHED:
yield StateSnapshotEvent(
type=EventType.STATE_SNAPSHOT,
snapshot=deps.state,
)
yield event
async def with_state_snapshot():
emitted = False
async for event in stream:
yield event
if not emitted and getattr(event, "type", None) == (
EventType.RUN_STARTED
):
yield snapshot
emitted = True
async for chunk in adapter.encode_stream(with_state_snapshot()):
yield chunk
async for chunk in adapter.encode_stream(with_final_state()):
yield chunk
return StreamingResponse(
event_stream(),

2
app/backend/pyproject.toml
View file

@ -6,7 +6,7 @@ requires-python = ">=3.12"
dependencies = [
"starlette>=0.50.0",
"uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=1.81.0",
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.11.0,<3.0.0",
"python-dotenv>=1.2.1",
"haiku.rag-slim>=0.68.0",
"logfire[pydantic-ai]>=3.17.0",

143
app/frontend/components/Chat.tsx
View file

@ -32,7 +32,7 @@ import DbInfo from "./DbInfo";
import DocumentFilter from "./DocumentFilter";
import SessionManager from "./SessionManager";
// Must match state_namespace from haiku.rag.skills.rag
// Must match RAGCapability.state_namespace.
const AGUI_STATE_KEY = "rag";
// AG-UI state is namespaced under AGUI_STATE_KEY
@ -115,24 +115,6 @@ function MessageIcon() {
);
}
function FileIcon() {
return (
<svg
width="14"
height="14"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
strokeWidth="2"
strokeLinecap="round"
strokeLinejoin="round"
>
<path d="M15 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7Z" />
<path d="M14 2v4a2 2 0 0 0 2 2h4" />
</svg>
);
}
function ToolCallIndicator({
toolName,
status,
@ -146,13 +128,9 @@ function ToolCallIndicator({
const getToolIcon = () => {
switch (toolName) {
case "search":
case "rag_search":
return <SearchIcon />;
case "get_document":
return <FileIcon />;
case "execute_skill":
case "execute_code":
case "cite":
case "rag_cite":
return <MessageIcon />;
default:
return <SearchIcon />;
@ -161,18 +139,10 @@ function ToolCallIndicator({
const getToolLabel = () => {
switch (toolName) {
case "search":
case "rag_search":
return "Search";
case "get_document":
return "Document";
case "execute_skill":
return "Skill";
case "execute_code":
return "Code";
case "cite":
case "rag_cite":
return "Cite";
case "list_documents":
return "Documents";
default:
return toolName;
}
@ -180,34 +150,12 @@ function ToolCallIndicator({
const getDescription = () => {
switch (toolName) {
case "execute_skill": {
const skill = args.skill_name as string | undefined;
const request = args.request as string | undefined;
return (
<span className="tool-query">
{skill ? `${skill}: ` : ""}
{request ?? "Processing..."}
</span>
);
}
case "search": {
case "rag_search": {
const query = args.query as string;
return <span className="tool-query">{query}</span>;
}
case "get_document":
return <span className="tool-query">{args.query as string}</span>;
case "execute_code": {
const code = args.code as string | undefined;
return (
<span className="tool-query">
{code ? code.slice(0, 80) : "Running code..."}
</span>
);
}
case "cite":
case "rag_cite":
return <span className="tool-query">Registering citations</span>;
case "list_documents":
return <span className="tool-query">Listing documents</span>;
default:
return <span>Processing...</span>;
}
@ -234,39 +182,6 @@ function ToolCallIndicator({
);
}
// Render an activity message from a skill sub-agent tool call/result
function ActivityIndicator({
message,
isComplete,
}: {
// biome-ignore lint/suspicious/noExplicitAny: AG-UI activity message shape
message: any;
isComplete: boolean;
}) {
const content = message.content ?? {};
const toolName = content.tool_name ?? "tool";
let args: Record<string, unknown> = {};
if (content.args) {
try {
args =
typeof content.args === "string"
? JSON.parse(content.args)
: content.args;
} catch {
// ignore parse errors
}
}
return (
<ToolCallIndicator
toolName={toolName}
status={isComplete ? "complete" : "loading"}
args={args}
/>
);
}
// Context for sharing chat state with the message view
const ChatStateContext = createContext<RAGState | null>(null);
@ -298,21 +213,6 @@ function MessageViewWithCitations({
const ragState = useContext(ChatStateContext);
const latestCitations = ragState ? getLatestCitations(ragState) : [];
// Collect completed tool_call_ids from skill_tool_result activity messages
const completedToolCallIds = useMemo(() => {
const ids = new Set<string>();
for (const msg of messages) {
if (
msg.role === "activity" &&
msg.activityType === "skill_tool_result" &&
msg.content?.tool_call_id
) {
ids.add(msg.content.tool_call_id);
}
}
return ids;
}, [messages]);
const cursor = isRunning ? (
<div key="cursor" className="streaming-cursor">
<span className="dot" />
@ -322,8 +222,7 @@ function MessageViewWithCitations({
) : null;
// CopilotChatMessageView renders one element per user/assistant message.
// We interleave activity indicators (skill sub-agent tool calls) and
// optionally inject CitationBlocks after assistant responses that
// Inject CitationBlocks after assistant responses that
// followed tool calls.
return (
<CopilotChatMessageView messages={messages} isRunning={isRunning}>
@ -345,24 +244,6 @@ function MessageViewWithCitations({
seenToolCalls = true;
}
// Activity messages are not rendered by CopilotKit —
// render them ourselves without consuming messageElements
if (msg.role === "activity") {
if (msg.activityType === "skill_tool_call") {
const toolCallId = msg.content?.tool_call_id;
result.push(
<ActivityIndicator
key={`activity-${msg.id}`}
message={msg}
isComplete={
toolCallId ? completedToolCallIds.has(toolCallId) : false
}
/>,
);
}
continue;
}
if (msg.role !== "user" && msg.role !== "assistant") continue;
if (elemIdx < messageElements.length) {
@ -438,8 +319,7 @@ function ChatContentInner({
useEffect(() => {
if (agent.messages.length > 0) return;
const session = getSession(sessionId);
// Seed the namespaced AG-UI state so the backend's first STATE_DELTA
// (e.g. add /rag/searches/...) has a namespace object to patch into.
// Seed state for the capability; the backend replaces it after each run.
agent.setState({
[AGUI_STATE_KEY]: normalizeRAGState(session?.ragState),
});
@ -466,10 +346,7 @@ function ChatContentInner({
}
}, [JSON.stringify(agent.messages), ragState, sessionId]);
// Deduplicate messages by id, keeping the last occurrence.
// haiku.skills 0.10.0+ sends activity snapshots with the same id
// and replace=true — CopilotKit doesn't deduplicate these, so we
// must do it to avoid React duplicate key warnings.
// Deduplicate messages by id to avoid React duplicate key warnings.
// biome-ignore lint/correctness/useExhaustiveDependencies: stable identity via agent ref
const messages = useMemo(() => {
const seen = new Map<string, number>();

2
app/frontend/lib/sessionStorage.ts
View file

@ -11,7 +11,7 @@ export interface Citation {
doc_item_refs?: string[];
}
// Matches RAGState from the backend skill
// Matches RAGState from the backend capability.
export interface RAGState {
citation_index: Record<string, Citation>;
citations: string[];

34
docs/benchmarks.md
View file

@ -1,6 +1,6 @@
# Benchmarks
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, and Wix are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the rag and rag-analysis skills.
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, and Wix are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities.
## Running Evaluations
@ -36,7 +36,7 @@ Active datasets:
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-skill` | ~2 GB |
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB |
| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB |
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
@ -64,8 +64,8 @@ evaluations run wix --config /path/to/haiku.rag.yaml --db /path/to/custom.lanced
- `--skip-qa` - Skip QA benchmark
- `--limit N` - Limit number of test cases
- `--name NAME` - Override the evaluation name
- `--target {rag-skill,analysis-skill}` - Choose which [skill](skills/index.md) to benchmark end-to-end (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`).
- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
@ -93,15 +93,15 @@ evaluations:
### QA Accuracy
`pydantic-evals` coordinates an LLM judge to determine whether the skill's answer is correct. The default judge is `ollama:qwen3.6`, pinned so changes to the skill model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions.
`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.6`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). 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.
### Citation Retrieval
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
This is computed alongside QA accuracy from the same skill run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the skill grounded its answer on it.
This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it.
## Current results
@ -130,12 +130,12 @@ Two approaches are benchmarked separately:
##### QA accuracy + citation retrieval
| Embedding Model | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
| Embedding Model | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------|
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 |
*Measured on haiku.rag v0.52.0, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases.*
@ -152,7 +152,7 @@ Two approaches are benchmarked separately:
##### QA accuracy + citation retrieval
| Embedding Model | VLM | Skill model | Cases | QA accuracy | Mean `cited_map` |
| Embedding Model | VLM | Capability model | Cases | QA accuracy | Mean `cited_map` |
|------------------------------------------|----------------------|------------------------------|------:|-------------|------------------|
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.80 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 2836 | 0.96 | 0.81 |
@ -165,9 +165,9 @@ Two approaches are benchmarked separately:
##### QA accuracy + citation retrieval
| Embedding Model | Reranker | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
| Embedding Model | Reranker | Target | Capability model | Cases | QA accuracy | Mean `cited_map` |
|----------------------|------------------------|------------------|------------------------------|------:|-------------|------------------|
| `qwen3-embedding:4b` | `mxbai-rerank-base-v2` | `analysis-skill` | `vllm:Qwen3.6-35B-A3B-NVFP4` | 7939 | 0.77 | 0.78 |
| `qwen3-embedding:4b` | `mxbai-rerank-base-v2` | `analysis-capability` | `vllm:Qwen3.6-35B-A3B-NVFP4` | 7939 | 0.77 | 0.78 |
*Measured on haiku.rag v0.55.0, deterministic Number-Match scoring (ε=0.01), 2560-dim `qwen3-embedding:4b` (vLLM) with `mxbai-rerank-base-v2`. 341 / 8281 cases excluded as nulls (analysis spirals from the request limit and in-generation loops). Accuracy and `cited_map` are over the 7939 scored cases. Mean 16.0s/case.*
@ -197,9 +197,9 @@ The reranker's contribution is larger here than on the single-doc datasets: hybr
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
`evaluations run wix --target rag-capability` runs the RAG capability end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the capability registered via the `cite` tool against the gold `expected_uris`.
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
| Capability model | Reranker | QA accuracy | Mean `cited_map` |
|------------------------------|------------------------|-------------|------------------|
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |

47
docs/capabilities/analysis.md Normal file
View file

@ -0,0 +1,47 @@
# Analysis Capability
`AnalysisCapability` adds search, citations, and sandboxed Python computation over the document corpus. Use it for counts, aggregation, comparison, structural traversal, and section-scoped reading.
It is deferred by default, keeping its substantial instructions and tool schemas out of context until the model chooses to load it.
## Tools
| Tool | Purpose |
|---|---|
| `analysis_search(query, limit?)` | Search the corpus for evidence. |
| `analysis_execute_code(code)` | Run Python against the virtual document filesystem. |
| `analysis_cite(chunk_ids)` | Register retrieved or filesystem-derived chunk IDs. |
The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`.
## Compose with RAG
```python
from pydantic_ai import Agent
from haiku.rag.capabilities.analysis import create_capability as analysis
from haiku.rag.capabilities.rag import create_capability as rag
agent = Agent(
"openai:gpt-5",
capabilities=[
rag(db_path="my.lancedb"),
analysis(db_path="my.lancedb"),
],
)
```
For the high-level convenience API:
```python
from haiku.rag.client import HaikuRAG
async with HaikuRAG("my.lancedb") as client:
result = await client.analyze("Which quarter had the highest revenue?")
print(result.answer)
```
## State
When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, and citations. Per-run searches and executions reset automatically; the filter and citation index persist.
The capability lazily opens both LanceDB and the sandbox only after it is loaded and a tool requires them. Resources close at the end of the agent run.

49
docs/capabilities/index.md Normal file
View file

@ -0,0 +1,49 @@
# Capabilities
haiku.rag provides two native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
| Capability | Use it for |
|---|---|
| [`RAGCapability`](rag.md) | Grounded document search and citations. |
| [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. |
Both capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability.
## Compose an agent
```python
from pydantic_ai import Agent
from haiku.rag.capabilities.rag import create_capability
rag = create_capability(db_path="my.lancedb")
agent = Agent("openai:gpt-5", capabilities=[rag])
result = await agent.run("What does the knowledge base say about X?")
print(result.output)
```
Attach both capabilities when an agent should choose between retrieval and computation:
```python
from haiku.rag.capabilities.analysis import create_capability as analysis
from haiku.rag.capabilities.rag import create_capability as rag
agent = Agent(
"openai:gpt-5",
capabilities=[rag(db_path="my.lancedb"), analysis(db_path="my.lancedb")],
)
```
## State
Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol.
Applications serving AG-UI should adapt the agent with Pydantic AI's `AGUIAdapter`. Native model and tool events require no haiku.rag-specific bridge.
## Database path
Both factories resolve their database in this order:
1. The `db_path` argument.
2. `HAIKU_RAG_DB`.
3. `config.storage.data_dir / "haiku.rag.lancedb"`.

53
docs/capabilities/rag.md Normal file
View file

@ -0,0 +1,53 @@
# RAG Capability
`RAGCapability` adds grounded document search and citations to a Pydantic AI agent. It is deferred by default, so its instructions and tools do not consume model context until loaded.
## Tools
| Tool | Purpose |
|---|---|
| `rag_search(query, limit?)` | Hybrid vector and full-text search with context expansion. |
| `rag_cite(chunk_ids)` | Register exact result chunk IDs as answer citations. |
The distinct `rag_` prefix lets this capability coexist with analysis and other search providers.
## Create and compose
```python
from pydantic_ai import Agent
from haiku.rag.capabilities.rag import create_capability
rag = create_capability(db_path="my.lancedb")
agent = Agent("openai:gpt-5", capabilities=[rag])
result = await agent.run("What safety equipment does the manual require?")
print(result.output)
```
`create_capability` accepts `db_path`, `config`, and `defer_loading`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary.
## State
When agent dependencies expose a `state` dictionary, the capability maintains a `RAGState` under `"rag"`:
```python
class RAGState(BaseModel):
citation_index: dict[str, Citation]
citations: list[str]
document_filter: str | None
searches: dict[str, list[SearchResult]]
```
`document_filter` persists between runs. Current citations and searches reset for each run, while the citation index remains available to the host application.
State is ordinary application state; the capability does not depend on AG-UI. An AG-UI application can expose it using Pydantic AI's standard adapter.
## Context management
Large RAG tool results from earlier user turns are replaced with a short marker before model requests. Tool-call pairing and current-turn evidence are retained. This prevents long conversations from repeatedly sending old retrieved content.
## Domain context and vision
`prompts.domain_preamble` is prepended to the packaged capability instructions. When the selected QA model has `vision: true`, picture results are attached to search returns as `BinaryContent`.
See [Search and question answering](../configuration/qa.md) and [picture processing](../configuration/processing.md#picture-handling).

14
docs/chat.md
View file

@ -13,11 +13,11 @@ haiku-rag chat --db /path/to/database.lancedb
haiku-rag chat --model openai:gpt-4o
```
![Chat TUI session against the rag-analysis skill](img/chat-qa.png)
![Chat TUI session with the analysis capability](img/chat-qa.png)
## How it works
The chat is a Pydantic AI agent with the `rag` [skill](skills/rag.md) attached by default. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and a live indicator of which tool is running.
The chat is a Pydantic AI agent with the deferred [RAG capability](capabilities/rag.md) attached by default. Each turn the agent decides which tool to call next, runs hybrid search against your documents, expands context around the hits, may issue further searches, and answers with citations. You see streaming text and native tool events directly from Pydantic AI.
The session is in-memory for the lifetime of the TUI. Conversation history is kept across turns so follow-up questions reuse prior context. Citations are tracked per turn and inspectable via the command palette. Clearing the chat resets the session and the agent's memory.
@ -51,26 +51,26 @@ haiku-rag visualize <chunk_id>
| Database info | Document and chunk counts, storage stats |
| View state | Current session state, citations, and intermediate tool results |
## Skills
## Capabilities
The default skill is `rag`. Enable `analysis` when the question needs computation, aggregation, comparison across documents, or section-scoped reading that a single search can't deliver:
The default capability is `rag`. Enable `analysis` when the question needs computation, aggregation, comparison across documents, or section-scoped reading that a single search can't deliver:
```bash
# both skills (the agent routes between them)
# both capabilities (the agent routes between them)
haiku-rag chat -s rag -s analysis
# analysis only
haiku-rag chat -s analysis
```
The `analysis` skill mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like:
The `analysis` capability mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like:
- "How many of these documents mention X?"
- "Summarize Section 5 of paper Y."
- "Compare the experimental sections across these three reports."
- "Which section discusses the proof of Theorem 4.10?"
For everyday Q&A, the rag skill alone is faster and cheaper. Attaching both lets the agent pick. See [Analysis skill](skills/analysis.md) for the full sandbox capabilities and worked code patterns.
For everyday Q&A, RAG alone is faster and cheaper. Attaching both lets the agent pick. See [Analysis capability](capabilities/analysis.md).
## Document filter

64
docs/cli.md
View file

@ -161,7 +161,7 @@ Filter to specific documents:
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
```
`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.
`ask` runs the [RAG capability](capabilities/rag.md) and always renders citations under the answer. When available, citations use the document title, otherwise they fall back to the URI.
Flags:
@ -185,7 +185,7 @@ Flags:
- `--filter` / `-f`: SQL WHERE clause to restrict document access
See [Analysis skill](skills/analysis.md) for details on capabilities and configuration.
See [Analysis capability](capabilities/analysis.md) for details and configuration.
## Chat
@ -204,7 +204,7 @@ haiku-rag chat -s rag -s analysis
Flags:
- `--skill` / `-s`: Skills to enable. `rag` (default), `analysis`. Can be repeated for multiple skills.
- `--capability` / `-c`: Capabilities to enable. `rag` (default), `analysis`. Can be repeated.
The chat interface provides:
@ -482,64 +482,6 @@ haiku-rag init-config [output_path]
If no path is specified, creates `haiku.rag.yaml` in the current directory.
## Create Skill
Generate a standalone skill package with an embedded database:
```bash
haiku-rag create-skill --name myskill --db /path/to/database.lancedb
```
The generated package is a pip-installable Python package that registers as a `haiku.skills` entry point.
### Options
| Flag | Description | Default |
|------|-------------|---------|
| `--name` | Skill name (lowercase alphanumeric and hyphens, required) | — |
| `--db` | Path to LanceDB database to embed (required) | — |
| `--description` | Skill description | Standard RAG description |
| `--tools` | Comma-separated tool names, or `all` | `all` |
| `--preamble` | Custom preamble for skill instructions | Standard RAG preamble |
| `--config-file` | Path to `haiku.rag.yaml` to embed | None |
| `--output` / `-o` | Output directory | Current directory |
### Available Tools
`cite`, `execute_code`, `get_document`, `list_documents`, `search`
### Example
```bash
# Generate a skill with specific tools and custom preamble
haiku-rag create-skill \
--name medic \
--db /path/to/medic.lancedb \
--tools search,cite \
--config-file /path/to/haiku.rag.yaml \
--description "Military medic knowledge base" \
--preamble "You are a military medic expert."
# Install the generated package
uv pip install -e ./medic-skill
# Use with haiku-skills
haiku-skills chat --use-entrypoints --skill medic
```
### Generated Package Structure
```
{name}-skill/
├── pyproject.toml
└── {name}_skill/
├── __init__.py # create_skill() entry point
├── SKILL.md # Skill metadata and instructions
└── assets/
├── {name}.lancedb/ # Embedded database
└── haiku.rag.yaml # Optional config
```
## Tags
A tag names the current database state. It is a logical snapshot composed of one LanceDB tag on each of the five tables, created from a single version snapshot.

4
docs/configuration/prompts.md
View file

@ -1,6 +1,6 @@
# Prompt Customization
Customize the prompts used by haiku.rag's skills to better match your domain and use case.
Customize the prompts used by haiku.rag's capabilities to match your domain.
## Configuration
@ -24,7 +24,7 @@ The `domain_preamble` field provides **domain context** prepended to the rag and
- Clarify domain-specific terminology
- 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) lives in the skill's SKILL.md. Fork the skill via `haiku-rag create-skill` to customize behavior.
**Important:** `domain_preamble` is for domain context, not behavioral instructions. Descriptions of subject matter, terminology, and content scope belong here. Applications can add behavioral guidance through normal Pydantic AI agent instructions.
**Example:**

2
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 the `qa` and `analysis` skills. These settings apply to any provider that supports them.
Configure model behavior for the `qa` and `analysis` capabilities. These settings apply to any provider that supports them.
### Basic Settings

2
docs/configuration/qa.md
View file

@ -60,4 +60,4 @@ analysis:
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
- **max_executions**: Maximum `execute_code` calls per question before the skill is told to answer from what it has (default: 15)
See [Analysis skill](../skills/analysis.md) for usage details.
See [Analysis capability](../capabilities/analysis.md) for usage details.

8
docs/overview.md
View file

@ -1,6 +1,6 @@
# Overview
haiku.rag is an agentic RAG that runs locally and scales to production. Index PDFs, web pages, or whole directories. Ask questions and get cited answers. Build agents, skills, and MCP integrations on top.
haiku.rag is an agentic RAG that runs locally and scales to production. Index PDFs, web pages, or whole directories. Ask questions and get cited answers. Build agents, capabilities, and MCP integrations on top.
haiku.rag is open-source first. The defaults run open models through [Ollama](https://ollama.com/) so the full pipeline works without external API keys. Any provider Pydantic AI supports works in its place.
@ -19,7 +19,7 @@ haiku-rag add-src ~/Documents/some-paper.pdf
haiku-rag chat
```
The chat TUI is one way to interact with the database. `haiku-rag ask` and `haiku-rag search` cover one-shot CLI usage. Python integrations, skills, and the MCP server work against the same database.
The chat TUI is one way to interact with the database. `haiku-rag ask` and `haiku-rag search` cover one-shot CLI usage. Python integrations, capabilities, and the MCP server work against the same database.
## What it does
@ -29,14 +29,14 @@ The chat TUI is one way to interact with the database. `haiku-rag ask` and `haik
**Answer.** RAG skill with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis skill with a sandboxed Python interpreter for aggregation and computation across documents.
**Integrate.** Use it from Python, the CLI, the [MCP server](mcp.md), or as composable [skills](skills/index.md) built on haiku.skills. Skills bundle tools, prompts, and state for use inside any Pydantic AI agent.
**Integrate.** Use it from Python, the CLI, the [MCP server](mcp.md), or through composable native Pydantic AI [capabilities](capabilities/index.md).
**Operate.** Embedded LanceDB by default. Also runs on S3, GCS, Azure, or LanceDB Cloud. Time-travel queries via LanceDB versioning. The [`haiku-ingester`](ingester.md) service runs continuously for production deployments.
## Where to go next
- [Quickstart](tutorial.md): install, index, chat.
- [Skills](skills/index.md): the rag and rag-analysis skills you compose into Pydantic AI agents.
- [Capabilities](capabilities/index.md): native RAG and analysis capabilities for Pydantic AI agents.
- [Python API](python.md): use haiku.rag from code.
- [MCP server](mcp.md): expose haiku.rag to Claude Desktop or other AI assistants.
- [Tuning](tuning.md): improve retrieval quality.

10
docs/python.md
View file

@ -328,11 +328,11 @@ answer, citations = await client.ask(
)
```
`client.ask` runs the [rag skill](skills/index.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, and the document's metadata (`document_meta`), so UIs can render metadata keys such as a public source URL alongside the citation.
`client.ask` runs the [RAG capability](capabilities/rag.md) and returns `(answer_text, list[Citation])`. Citations include page numbers, section headings, document references, and the document's metadata (`document_meta`), so UIs can render metadata keys such as a public source URL alongside the citation.
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: [Skills](skills/index.md) for details on the skills the client wraps.
See also: [Capabilities](capabilities/index.md) for direct agent composition.
## Analysis
@ -352,13 +352,13 @@ result = await client.analyze(
)
```
`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.
`client.analyze` runs the [analysis capability](capabilities/analysis.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 skill](skills/analysis.md) for details on capabilities and configuration.
See [Analysis capability](capabilities/analysis.md) for details and configuration.
## Building custom agents
`client.ask` and `client.analyze` are the convenience wrappers. To build your own Pydantic AI agent against the same database, attach the rag and rag-analysis skills directly with `SkillToolset`. See [Skills](skills/index.md) for the full story and worked examples.
`client.ask` and `client.analyze` are convenience wrappers. To build your own Pydantic AI agent, attach the native RAG and analysis capabilities directly. See [Capabilities](capabilities/index.md).
For the low-level toolset factories under `haiku.rag.tools` (one rung below the skill abstraction), see [Toolsets](tools.md).

199
docs/skills/analysis.md
View file

@ -1,199 +0,0 @@
# Analysis Skill
Plain RAG (search → cite → answer) works for questions whose answer sits in a chunk or two: "Who wrote this?", "What does X say about Y?". It struggles when the answer requires touching the whole corpus, reading a specific section in full, or doing arithmetic on the data.
The analysis skill (`rag-analysis`) gives the agent a second tool (`execute_code`) that runs Python in a sandboxed interpreter against a structured view of your documents. The agent can search, read, count, slice, and compare without leaving the tool call. Citations work the same way as the rag skill.
`client.analyze`, `haiku-rag analyze`, the MCP `analyze` tool, and the chat TUI (when `-s analysis` is enabled) all run through this skill.
## When to use it
Reach for the analysis skill when the question needs more than a search:
- **Aggregation across the corpus.** "How many documents mention security vulnerabilities?"
- **Section-scoped reading.** "Summarize Section 5 of paper Y."
- **Structural comparison.** "Do both papers have an Experimental Results section?"
- **Computation on retrieved data.** "What's the average revenue across these quarterly reports?"
- **Multi-step chains.** Search, filter the results in Python, search again, aggregate, all in one tool call.
For everyday Q&A, the [RAG skill](rag.md) is faster and cheaper. Attach both and the agent routes.
## How it works
Two things make the agent's programs short and the resulting analyses tractable:
1. **Search and document listing are awaitable inside the code.** `await search(query)` returns the same hits the rag skill sees: chunk IDs, text, source metadata, picture refs. The agent can immediately filter, sort, count, or follow up with another search without exiting the tool call.
2. **Every document is mounted as a virtual filesystem at `/documents/{id}/`.** The agent reads four files per document: identifiers and metadata, full text, a list of structured items (paragraphs, tables, figures, headings), and a section tree built from the document's headings. The structure exposes what search alone hides. The agent can navigate from a search hit to the section it lives in, slice a single section instead of pulling the whole document, or scan a document's text directly when keyword precision matters.
A search hit is always a starting point. The agent reads structure around it, drills into the right section, and cites the chunks it actually used. Chunk IDs from search results and chunk IDs surfaced through the VFS are both accepted by `cite`.
### Sandbox guarantees
The interpreter is [pydantic-monty](https://github.com/pydantic/monty), isolated from the host:
- **Virtual filesystem only.** `/documents/` is the entire FS.
- **No network.** HTTP, sockets, and the `requests` family are unavailable.
- **Limited imports.** Only `json`, `re`, `math`, `pathlib`.
- **Execution timeout** (default 60s, configurable via `analysis.code_timeout`).
- **Output truncation** (default 50000 chars, configurable via `analysis.max_output_chars`).
- **Execution budget** (default 15 calls, configurable via `analysis.max_executions`). Past the budget, `execute_code` returns a notice telling the skill to answer from what it has instead of running more code.
Variables persist between `execute_code` calls within one invocation, so the agent can build state step by step. A fresh sandbox is built per `client.analyze` call.
## Tools
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search with context expansion. Same as the RAG skill's `search`. |
| `execute_code(code)` | Run Python in a sandboxed interpreter with VFS access. |
| `cite(chunk_ids)` | Register chunk IDs as citations. Call before producing the final answer. |
`list_documents` isn't exposed as a top-level tool but is available inside `execute_code` as `await list_documents()`.
## State
`AnalysisState` lives under the `"analysis"` namespace:
```python
class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = []
citation_index: dict[str, Citation] = {}
citations: list[str] = []
searches: dict[str, list[SearchResult]] = {}
```
- **document_filter** — SQL WHERE clause applied to `search` and the VFS. The LLM can't bypass it: both views are scoped.
- **executions** — Each `execute_code` call appends an entry with code, stdout, stderr, success. Cleared at the start of each invocation.
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations.
- **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared per-invocation.
- **searches** — Search results from both the `search` tool and sandbox-internal searches. Cleared per-invocation.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path="my.lancedb")
```
| Parameter | Default | Description |
|-----------|---------|-------------|
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
| `config` | `None` | `AppConfig` instance. Falls back to `get_config()`. |
## Use it
### From `client.analyze`
```python
from haiku.rag.client import HaikuRAG
async with HaikuRAG("my.lancedb") as client:
result = await client.analyze("How many documents mention 'security'?")
print(result.answer)
for citation in result.citations:
print(citation.uri, citation.title)
```
`client.analyze` runs the skill end-to-end and returns an `AnalysisResult` with `answer` and `citations`. The executed Python programs live on `AnalysisState.executions` during the run, not on the returned result.
### Combine with the RAG skill
```python
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
from haiku.skills.agent import SkillToolset
rag = create_rag_skill(db_path="my.lancedb")
analysis = create_analysis_skill(db_path="my.lancedb")
toolset = SkillToolset(skills=[rag, analysis])
```
The agent routes Q&A to the rag skill and computational questions to rag-analysis.
## What the agent actually writes
You don't write these programs yourself. The agent does, inside `execute_code`. Seeing the shape helps when you tune prompts, debug a run via `AnalysisState.executions`, or design a custom skill.
**Aggregate across the corpus.** *"How many documents mention security vulnerabilities?"*
```python
hits = await search("security vulnerability", limit=50)
doc_ids = {h['document_id'] for h in hits}
print(f"{len(doc_ids)} documents mention security vulnerabilities")
# Cite the top hit per document
seen = set()
for hit in hits:
if hit['document_id'] not in seen:
seen.add(hit['document_id'])
await cite(hit['chunk_id'])
```
**Read one section in depth.** *"Summarize Section 5."*
```python
from pathlib import Path
import json
doc_id = "..." # from a prior search or list_documents
toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text())
section = next(n for n in toc['tree'] if n['title'].startswith('5'))
start, end = section['item_range']
lines = Path(f'/documents/{doc_id}/items.jsonl').read_text().splitlines()[start:end]
for line in lines:
print(json.loads(line)['text'])
await cite(section['chunk_ids'])
```
The section node already aggregates the chunks underneath it, so the agent cites the whole section without a separate search.
**Compare structure across documents.** *"Do both papers have an Experimental Results section?"*
```python
from pathlib import Path
import json
for doc_id in ["doc-a-id", "doc-b-id"]:
toc = json.loads(Path(f'/documents/{doc_id}/toc.json').read_text())
print(f"\n=== {toc['title']} ===")
for node in toc['tree']:
if 'experiment' in node['title'].lower():
print(f" {node['title']} (pages {node['page_numbers']})")
await cite(node['chunk_ids'])
```
## Context filter
The `filter` parameter is enforced at the deps layer. The LLM can't bypass it: both the VFS and search results are scoped to the filter.
```python
result = await client.analyze(
"Summarize all findings",
filter="uri LIKE '%confidential%'"
)
```
Useful for scoping to a corpus subset, enforcing access control, or restricting context.
## Configuration
```yaml
analysis:
model:
provider: anthropic
name: claude-sonnet-4-20250514
code_timeout: 60.0 # Max seconds per code execution
max_output_chars: 50000 # Truncate output after this many chars
max_executions: 15 # Max execute_code calls per question
```
When `analysis.model` is unset, the skill falls back to `qa.model`.
See [Search and question answering](../configuration/qa.md#analysis-configuration) for the full set.

124
docs/skills/custom.md
View file

@ -1,124 +0,0 @@
# Custom Skills
The two skills haiku.rag ships work against any LanceDB database. When you want a *domain-specific* skill that bundles its own data, prompt, and tool surface (for example, a "recipes" skill that knows about cooking and ships with a recipes database), generate one with `haiku-rag create-skill`.
The generated package is a regular pip-installable Python package that registers as a `haiku.skills` entry point. Any haiku.skills-aware host (haiku.skills CLI, your own agent, the AG-UI adapter) discovers it automatically.
## When to use a custom skill
- The model should consult a specific knowledge base for a specific kind of question, alongside other skills.
- You want a different instruction prompt than the generic `rag` skill (different tone, refusal style, domain rules).
- You want to ship a knowledge base plus its prompt as one distributable unit.
- You're running multiple skills against different databases in the same agent.
If you just want to point a haiku.rag database at your own model and prompt, configure `haiku.rag.yaml` and use the built-in `rag` skill. No custom package needed.
## Generate
```bash
haiku-rag create-skill \
--name recipes \
--db /path/to/recipes.lancedb \
--tools search,cite \
--description "Recipe and cooking knowledge base" \
--preamble "You are a culinary expert helping with recipes and cooking techniques."
```
Then install and use:
```bash
uv pip install -e ./recipes-skill
haiku-skills list --use-entrypoints
# recipes — Recipe and cooking knowledge base
haiku-skills chat --use-entrypoints --skill recipes
```
### Flags
| Flag | Description | Default |
|------|-------------|---------|
| `--name` | Skill name (lowercase alphanumeric and hyphens). Required. | — |
| `--db` | Path to the LanceDB database to embed. Required. | — |
| `--description` | One-line skill description. The agent reads this to decide when to invoke. | Standard RAG description |
| `--tools` | Comma-separated tool subset, or `all`. | `all` |
| `--preamble` | Custom preamble for the skill's instructions. | Standard RAG preamble |
| `--config-file` | Path to a `haiku.rag.yaml` to embed alongside the database. | None |
| `--output` / `-o` | Output directory. | Current directory |
### Available tools
`cite`, `execute_code`, `get_document`, `list_documents`, `search`.
Drop `execute_code` from `--tools` if the skill shouldn't run sandboxed Python. That gives you a search-and-cite-only skill with no analysis capabilities.
## Anatomy of a generated skill
```
{name}-skill/
├── pyproject.toml
└── {name}_skill/
├── __init__.py # create_skill() entry point
├── SKILL.md # Skill metadata and instructions
└── assets/
├── {name}.lancedb/ # The embedded database
└── haiku.rag.yaml # Optional config (only if --config-file passed)
```
- **`SKILL.md`** carries the instruction prompt the agent will follow. The frontmatter includes the skill name and description. Everything below is the prompt body. Edit this to change behavior.
- **`__init__.py`** exposes `create_skill()` (the entry point) and `visualize_chunk()` for rendering visual grounding.
- **`assets/{name}.lancedb/`** is the database, shipped inside the package.
- **`assets/haiku.rag.yaml`** (optional) pins provider settings the skill needs.
The package can be installed locally with `uv pip install -e .` or published to PyPI.
## Generating visual grounding from a custom skill
Each generated skill exposes a `visualize_chunk()` function that returns the chunk's bounding boxes rendered onto its source page:
```python
from recipes_skill import visualize_chunk
images = await visualize_chunk(chunk_id)
# images is a list of PIL.Image objects, one per page the chunk covers
images[0].save("citation.png")
```
Pass chunk IDs from skill citations or search results. Same prerequisites as elsewhere in haiku.rag: documents need stored page images, and the chunk must come from a PDF or other docling-converted source.
## Multi-skill agents
Each generated skill is self-contained with its own database and instructions. Compose multiple skills in one agent and the model routes between them via their descriptions:
```python
from recipes_skill import create_skill as create_recipes_skill
from medic_skill import create_skill as create_medic_skill
from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent
recipes = create_recipes_skill()
medic = create_medic_skill()
toolset = SkillToolset(skills=[recipes, medic])
agent = Agent(
"openai-chat:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
await agent.run("What's the optimal temperature for braising short ribs?")
# Routes to recipes
await agent.run("What's the field treatment for tension pneumothorax?")
# Routes to medic
```
Each skill maintains state under its own namespace (`recipes`, `medic`, …), so citations and searches don't collide.
## Writing a skill from scratch
`create-skill` is the convenience path. If you need full control over the tools, state model, or instruction loading, write the skill against [haiku.skills](https://github.com/ggozad/haiku.skills) directly. The generated package in `{name}_skill/__init__.py` is a good reference. It composes haiku.rag's `_tools` factory with a `haiku.skills.Skill` and registers under the `haiku.skills` entry point group in `pyproject.toml`.
See the haiku.skills repository for the full Skill contract.

105
docs/skills/index.md
View file

@ -1,105 +0,0 @@
# Skills
Skills put haiku.rag in front of a model. A skill bundles tools, an instruction prompt, and managed state into a unit that drops into any Pydantic AI agent via `SkillToolset`. haiku.rag ships two skills and supports custom skills.
Built on [haiku.skills](https://github.com/ggozad/haiku.skills).
## Available skills
| Skill | What it does | Reach for it when |
|-------|--------------|-------------------|
| [`rag`](rag.md) | Search, retrieve, and cite content from a knowledge base. | The model needs to find and quote evidence from documents. |
| [`rag-analysis`](analysis.md) | Same as `rag`, plus a sandboxed Python interpreter mounting every document as a virtual filesystem. | The question requires computation, aggregation, structural traversal, or section-scoped reading. |
To ship your own skill (bundled with its own database), see [Custom skills](custom.md).
## Your first agent
```python
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent
rag = create_skill(db_path="my.lancedb")
toolset = SkillToolset(skills=[rag])
agent = Agent(
"openai-chat:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
result = await agent.run("What does the knowledge base say about X?")
print(result.output)
```
The skill searches, cites, and answers. You supply the model and the question.
To run analysis against the same database, swap in the `rag-analysis` skill or attach both:
```python
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
rag = create_rag_skill(db_path="my.lancedb")
analysis = create_analysis_skill(db_path="my.lancedb")
toolset = SkillToolset(skills=[rag, analysis])
```
The agent reads each skill's description and routes questions itself. See the individual skill pages for the tool surface, state model, and worked examples.
## State
Each skill manages its own state under a dedicated namespace. State is synced via the AG-UI protocol when using `AGUIAdapter`.
```python
rag_state = toolset.get_namespace("rag")
analysis_state = toolset.get_namespace("analysis")
```
Both state models track citations, the current document filter, and per-turn searches. Analysis state also carries the sandbox execution log. See [RAG skill: state](rag.md#state) and [Analysis skill: state](analysis.md#state).
## Database path resolution
Both skills resolve the database path in the same order:
1. `db_path` argument passed to `create_skill()`
2. `HAIKU_RAG_DB` environment variable
3. Config default (`config.storage.data_dir / "haiku.rag.lancedb"`)
## AG-UI streaming for web apps
For browser apps, use pydantic-ai's `AGUIAdapter` to stream tool calls, text, and state deltas:
```python
from pydantic_ai.ui.ag_ui import AGUIAdapter
adapter = AGUIAdapter(agent=agent, run_input=run_input)
event_stream = adapter.run_stream()
sse_event_stream = adapter.encode_stream(event_stream)
```
See the [Web application](../apps.md) reference implementation.
## Exposing via MCP
To use a skill from Claude Desktop or another MCP-aware client, run the MCP server:
```bash
haiku-rag mcp --stdio
```
The server exposes the skill tools (search, ask, analyze) over MCP. See [MCP](../mcp.md).
## Discovery
Skills are registered as Python entry points under `haiku.skills`. They are discovered automatically:
```bash
haiku-skills list --use-entrypoints
# rag — Search, retrieve and analyze documents using RAG.
# rag-analysis — Analyze documents using code execution in a sandboxed interpreter.
```
This is what makes custom skills installable as plain pip packages. See [Custom skills](custom.md).

179
docs/skills/rag.md
View file

@ -1,179 +0,0 @@
# RAG Skill
The `rag` skill answers questions over a knowledge base with hybrid search, structure-aware context expansion, and explicit citations. `client.ask`, `haiku-rag ask`, the MCP `ask_question` tool, and the chat TUI all run through this skill.
## When to use it
- The model needs to find and quote evidence from a document corpus.
- You want citations under every answer.
- You're building a Q&A agent, a documentation chatbot, or any RAG-style integration.
If the question requires *computation* over the corpus (counts, aggregates, comparisons, section-scoped reading), reach for the [Analysis skill](analysis.md) instead, or attach both.
## Tools
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with section-aware context expansion. Returns `chunk_id`, content, `doc_item_refs`, `picture_refs`, `picture_captions`, source metadata. |
| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer. The agent calls this before writing the final response. |
For corpus enumeration or full-document reads, reach for the [Analysis skill](analysis.md), which exposes `await list_documents()` and a `/documents/{id}/content.txt` virtual filesystem inside `execute_code`. Both are also available as opt-in tools when building a [custom skill](custom.md).
## State
The skill manages a `RAGState` under the `"rag"` namespace:
```python
class RAGState(BaseModel):
citation_index: dict[str, Citation] = {}
citations: list[str] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
```
- **citation_index** — All citations indexed by chunk ID. Accumulates across invocations so historical chunk IDs stay resolvable in UI scrollback.
- **citations** — Chunk IDs registered via `cite` during the current invocation. Deduplicated, cleared at the start of each invocation.
- **document_filter** — SQL WHERE clause applied to `search`. Persists across invocations.
- **searches** — Search results keyed by query string. Cleared at the start of each invocation.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path="my.lancedb")
```
| Parameter | Default | Description |
|-----------|---------|-------------|
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
| `config` | `None` | `AppConfig` instance. Falls back to `get_config()`. |
## Examples
### Minimal agent
```python
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent
rag = create_skill(db_path="my.lancedb")
toolset = SkillToolset(skills=[rag])
agent = Agent(
"openai-chat:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
result = await agent.run("What does the manual say about safety procedures?")
print(result.output)
# Inspect what the model cited
state = toolset.get_namespace("rag")
for chunk_id in state.citations:
citation = state.citation_index[chunk_id]
print(f"- {citation.document_title}: {citation.content[:100]}…")
```
### Domain customization
Set a domain preamble in `haiku.rag.yaml` and the skill picks it up:
```yaml
prompts:
domain_preamble: |
The knowledge base contains the operations manual for the Helios solar array.
"The array" or unqualified specs refer to Helios. Terminology like "string"
refers to a series-connected panel chain, not text.
```
To scope a session to a subset of documents, set the filter on the namespace state:
```python
state = toolset.get_namespace("rag")
state.document_filter = "uri LIKE '%helios/v4/%'"
result = await agent.run("What's the maintenance interval for the inverters?")
```
The filter applies to every `search` call for the rest of the session, and the model can't bypass it from inside.
### Combining with the analysis skill
Attach both skills and the agent routes between them:
```python
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
rag = create_rag_skill(db_path="my.lancedb")
analysis = create_analysis_skill(db_path="my.lancedb")
toolset = SkillToolset(skills=[rag, analysis])
agent = Agent(
"openai-chat:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
# Q&A → uses rag
await agent.run("What safety equipment is required on-site?")
# Computational question → uses rag-analysis
await agent.run("How many checklists mention torque specifications?")
```
### Streaming to a web frontend
Wrap the agent with `AGUIAdapter` to stream tool calls, text deltas, and state changes to a CopilotKit-style frontend:
```python
from pydantic_ai.ui.ag_ui import AGUIAdapter
adapter = AGUIAdapter(agent=agent, run_input=run_input)
sse_stream = adapter.encode_stream(adapter.run_stream())
```
See the [Web application](../apps.md) reference implementation for the full Starlette + Next.js setup.
### Exposing via MCP
To call the skill from Claude Desktop (or any MCP client), run the MCP server:
```bash
haiku-rag mcp --stdio
```
The exposed `ask_question` tool runs this skill. See [MCP](../mcp.md) for the configuration block.
## Configuration
The skill picks up its model and search behavior from the standard config sections:
```yaml
qa:
model:
provider: ollama
name: gpt-oss
enable_thinking: true
temperature: 0.3
vision: false # set true for vision-capable QA models
max_searches: 3
search:
limit: 5
max_context_chars: 5000
```
See [Search and question answering](../configuration/qa.md) for every knob.
## Vision support
When `qa.model.vision: true` is set, the skill's `search` tool attaches picture bytes to its tool returns as `BinaryContent`. The model can then read figures, diagrams, and screenshots directly alongside the surrounding text. Requires `processing.pictures != none` so the bytes exist on disk. See the [pictures × embedder × QA model matrix](../configuration/processing.md#picture-handling) for the combinations that make sense.
## Customizing the skill prompt
The skill's instruction prompt lives in `SKILL.md` inside the package. For behavior changes (different phrasing, refusal style, additional rules), the supported path is to fork the skill with `haiku-rag create-skill` and edit the generated `SKILL.md`. The `domain_preamble` field above is for *what the corpus is about*, not for *how the agent should behave*. See [Custom skills](custom.md).

6
docs/tools.md
View file

@ -1,12 +1,12 @@
# Toolsets
haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggozad/haiku.skills) skills. For most integrations, see [Skills](skills/index.md).
For agent integrations, use the native Pydantic AI [capabilities](capabilities/index.md). This page documents the lower-level toolsets used by other haiku.rag surfaces.
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories that the skills themselves compose.
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories used across haiku.rag.
## Low-Level Toolsets
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.
For advanced use cases, individual toolset factories are available in `haiku.rag.tools` and can be reused to build custom agents.
### RAGDeps Protocol

2
docs/tutorial.md
View file

@ -81,6 +81,6 @@ haiku-rag ask "Who wrote haiku.rag?"
- [Chat](chat.md): sessions, citations, and the full TUI.
- [CLI reference](cli.md): every command.
- [Python API](python.md): use haiku.rag in your own code.
- [Skills](skills/index.md): the rag and rag-analysis skills the client wraps.
- [Capabilities](capabilities/index.md): native RAG and analysis components used by the client.
- [Tuning](tuning.md): better retrieval.
- [Configuration](configuration/index.md): every setting.

16
evaluations/README.md
View file

@ -42,24 +42,24 @@ evaluations run wix --limit 100
### Choosing the target
`evaluations run` benchmarks `--target rag-skill` by default. Use
`--target analysis-skill` to benchmark the analysis skill against the same
`evaluations run` benchmarks `--target rag-capability` by default. Use
`--target analysis-capability` to benchmark the analysis capability against the same
datasets and judge:
```bash
evaluations run wix --target rag-skill
evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
evaluations run wix --target rag-capability
evaluations run wix --target analysis-capability --capability-model ollama:gpt-oss
```
`--skill-model "provider:name"` overrides the skill model independently from
`--capability-model "provider:name"` overrides the capability model independently from
the judge (defaults to `qa.model`, or `analysis.model` when set for the
analysis-skill target). A citation retrieval metric (`cited_map`) is computed
alongside QA accuracy from the URIs the skill registered via the `cite` tool.
analysis-capability target). A citation retrieval metric (`cited_map`) is computed
alongside QA accuracy from the URIs the capability registered via the `cite` tool.
### Debugging runs in Logfire
With `LOGFIRE_TOKEN` set, runs ship spans under `service_name = 'evals'`. The
`debug-evals` skill in `.claude/skills/` turns these into ready-made Logfire
`debug-evals` capability in `.claude/capabilities/` turns these into ready-made Logfire
queries (recent runs, per-case pass rate and `cited_map`, failing and slowest
cases) for use from Claude Code.

2
evaluations/configs/t2_finqa.yaml
View file

@ -1,6 +1,6 @@
# Reference config for the `t2_finqa` pre-built evaluation database.
# T²-RAGBench (FinQA) financial QA, scored by exact numeric match.
# Run: evaluations run t2_finqa --skip-db --target analysis-skill --config configs/t2_finqa.yaml
# Run: evaluations run t2_finqa --skip-db --target analysis-capability --config configs/t2_finqa.yaml
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
environment: development

87
evaluations/evaluations/benchmark.py
View file

@ -20,7 +20,7 @@ from evaluations.evaluators import (
CitationMAPEvaluator,
MAPEvaluator,
)
from evaluations.skill_runner import SkillFactory, run_skill_question
from evaluations.capability_runner import CapabilityFactory, run_capability_question
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
@ -28,8 +28,8 @@ from haiku.rag.logging import configure_cli_logging
from haiku.rag.telemetry import configure as configure_telemetry
from haiku.rag.utils import get_model, parse_model_option
Target = Literal["rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
Target = Literal["rag-capability", "analysis-capability"]
TARGETS: tuple[Target, ...] = ("rag-capability", "analysis-capability")
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
# their QA model does not inadvertently change the judge — keeps cross-run
@ -52,8 +52,8 @@ def build_experiment_metadata(
test_cases: int,
config: AppConfig,
judge_config: ModelConfig | None = None,
target: Target = "rag-skill",
skill_config: ModelConfig | None = None,
target: Target = "rag-capability",
capability_config: ModelConfig | None = None,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
metadata: dict[str, Any] = {
@ -87,14 +87,14 @@ def build_experiment_metadata(
"judge_enable_thinking": judge_config.enable_thinking,
}
)
if skill_config is not None:
if capability_config is not None:
metadata.update(
{
"skill_provider": skill_config.provider,
"skill_model": skill_config.name,
"skill_temperature": skill_config.temperature,
"skill_max_tokens": skill_config.max_tokens,
"skill_enable_thinking": skill_config.enable_thinking,
"capability_provider": capability_config.provider,
"capability_model": capability_config.name,
"capability_temperature": capability_config.temperature,
"capability_max_tokens": capability_config.max_tokens,
"capability_enable_thinking": capability_config.enable_thinking,
}
)
return metadata
@ -221,6 +221,7 @@ async def run_retrieval_benchmark(
metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper()
dataset = EvalDataset(
name=f"{spec.key}-retrieval",
cases=cases,
evaluators=[evaluator],
)
@ -286,16 +287,16 @@ async def run_retrieval_benchmark(
}
def _skill_factory_for_target(target: Target) -> SkillFactory:
if target == "rag-skill":
from haiku.rag.skills.rag import create_skill
def _capability_factory_for_target(target: Target) -> CapabilityFactory:
if target == "rag-capability":
from haiku.rag.capabilities.rag import create_capability
return create_skill
if target == "analysis-skill":
from haiku.rag.skills.analysis import create_skill
return create_capability
if target == "analysis-capability":
from haiku.rag.capabilities.analysis import create_capability
return create_skill
raise ValueError(f"target {target!r} is not a skill target")
return create_capability
raise ValueError(f"target {target!r} is not a capability target")
def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None:
@ -352,8 +353,8 @@ async def run_qa_benchmark(
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
target: Target = "rag-capability",
capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
@ -367,12 +368,12 @@ async def run_qa_benchmark(
]
judge_config = judge_model or DEFAULT_JUDGE_MODEL
if target == "analysis-skill":
# Mirror the skill-code resolver: explicit analysis.model wins,
if target == "analysis-capability":
# Mirror the capability-code resolver: explicit analysis.model wins,
# else fall back to qa.model.
skill_config = skill_model or config.analysis.model or config.qa.model
capability_config = capability_model or config.analysis.model or config.qa.model
else:
skill_config = skill_model or config.qa.model
capability_config = capability_model or config.qa.model
db = spec.db_path(db_path)
_attach_relevant_uris(cases, spec, limit)
@ -409,7 +410,7 @@ async def run_qa_benchmark(
config=config,
judge_config=judge_config,
target=target,
skill_config=skill_config,
capability_config=capability_config,
)
async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
@ -421,16 +422,16 @@ async def run_qa_benchmark(
metadata=experiment_metadata,
)
skill_factory = _skill_factory_for_target(target)
resolved_skill_model = get_model(skill_config, config)
capability_factory = _capability_factory_for_target(target)
resolved_capability_model = get_model(capability_config, config)
async def answer_question(question: str) -> str:
result = await run_skill_question(
skill_factory=skill_factory,
result = await run_capability_question(
capability_factory=capability_factory,
db_path=db,
config=config,
question=question,
skill_model=resolved_skill_model,
capability_model=resolved_capability_model,
)
set_eval_attribute("cited_uris", result.cited_uris)
return result.answer
@ -510,8 +511,8 @@ async def evaluate_dataset(
vacuum_interval: int = 100,
multimodal_only: bool = False,
judge_model: ModelConfig | None = None,
target: Target = "rag-skill",
skill_model: ModelConfig | None = None,
target: Target = "rag-capability",
capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None,
) -> None:
if not skip_db:
@ -543,7 +544,7 @@ async def evaluate_dataset(
db_path=db_path,
judge_model=judge_model,
target=target,
skill_model=skill_model,
capability_model=capability_model,
case_ids=case_ids,
)
@ -622,16 +623,16 @@ def run(
help="Only evaluate queries requiring image understanding.",
),
target: str = typer.Option(
"rag-skill",
"rag-capability",
"--target",
help="What to benchmark: rag-skill | analysis-skill.",
help="What to benchmark: rag-capability | analysis-capability.",
),
skill_model: str | None = typer.Option(
capability_model: str | None = typer.Option(
None,
"--skill-model",
"--capability-model",
help=(
"Skill model as 'provider:name'. Defaults to qa.model (or "
"analysis.model when --target is analysis-skill) from the config."
"Capability model as 'provider:name'. Defaults to qa.model (or "
"analysis.model when --target is analysis-capability) from the config."
),
),
filter_ids: Path | None = typer.Option(
@ -651,7 +652,9 @@ def run(
)
target_value = cast(Target, target)
judge_model_config = app_config.evaluations.judge
skill_model_config = parse_model_option(skill_model) if skill_model else None
capability_model_config = (
parse_model_option(capability_model) if capability_model else None
)
asyncio.run(
evaluate_dataset(
@ -667,7 +670,7 @@ def run(
multimodal_only=multimodal_only,
judge_model=judge_model_config,
target=target_value,
skill_model=skill_model_config,
capability_model=capability_model_config,
case_ids=_load_case_ids(filter_ids),
)
)

79
evaluations/evaluations/skill_runner.py → evaluations/evaluations/capability_runner.py
View file

@ -1,17 +1,18 @@
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Protocol, cast
from typing import Any, Protocol, cast
from pydantic_ai import Agent
from pydantic_ai.models import Model
from pydantic_ai.usage import UsageLimits
from haiku.rag.store.models.citation import Citation
from haiku.rag.capabilities import RAGCapabilityBase
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.skills import run_skill
from haiku.skills.models import Skill
from haiku.rag.store.models.citation import Citation
SkillFactory = Callable[..., Skill]
CapabilityFactory = Callable[..., RAGCapabilityBase[Any]]
class _RagLikeState(Protocol):
@ -22,7 +23,7 @@ class _RagLikeState(Protocol):
@dataclass
class SkillRunResult:
class CapabilityRunResult:
answer: str
cited_uris: list[str] = field(default_factory=list)
cited_chunk_ids: list[str] = field(default_factory=list)
@ -31,38 +32,60 @@ class SkillRunResult:
n_executions: int = 0
async def run_skill_question(
skill_factory: SkillFactory,
@dataclass
class _EvalDeps:
state: dict[str, Any] = field(default_factory=dict)
async def run_capability_question(
capability_factory: CapabilityFactory,
db_path: Path,
config: AppConfig,
question: str,
skill_model: str | Model,
capability_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
) -> SkillRunResult:
"""Run a single question through a skill and return answer + retrieval data.
) -> CapabilityRunResult:
"""Run a single question through a capability and return answer + retrieval data.
Builds the skill via ``skill_factory(db_path=..., config=...)`` and
invokes it with a fresh state instance derived from
``skill.state_type``. After the run, citations and searched documents
Builds a native capability via ``capability_factory(db_path=..., config=...)``.
After the run, citations and searched documents
are extracted from the state for downstream eval scoring.
The skill must produce a state with RAG-skill-shaped fields (citation
The capability must produce a state with RAG-capability-shaped fields (citation
index, searches, optional document filter) — i.e. ``RAGState`` or
``AnalysisState`` from ``haiku.rag.skills``.
``AnalysisState`` from ``haiku.rag.capabilities``.
"""
skill = skill_factory(db_path=db_path, config=config)
if request_limit is not None:
skill.request_limit = request_limit
if skill.state_type is None:
raise ValueError(f"Skill {skill.metadata.name!r} has no state_type")
state = skill.state_type()
capability = capability_factory(
db_path=db_path,
config=config,
defer_loading=False,
)
state = capability.state_type()
typed = cast(_RagLikeState, state)
if document_filter is not None:
typed.document_filter = document_filter
answer, _, _ = await run_skill(skill_model, skill, question, state=state)
deps = _EvalDeps(state={capability.state_namespace: state.model_dump(mode="json")})
agent = Agent(
capability_model,
deps_type=_EvalDeps,
capabilities=[capability],
)
effective_request_limit = (
request_limit if request_limit is not None else capability.default_request_limit
)
agent_result = await agent.run(
question,
deps=deps,
usage_limits=(
UsageLimits(request_limit=effective_request_limit)
if effective_request_limit is not None
else None
),
)
state = capability.state_type.model_validate(deps.state[capability.state_namespace])
typed = cast(_RagLikeState, state)
cited_chunk_ids: list[str] = list(typed.citations)
seen_cited: set[str] = set()
@ -78,8 +101,8 @@ async def run_skill_question(
seen_searched: set[str] = set()
searched_uris: list[str] = []
for results in typed.searches.values():
for result in results:
uri = result.document_uri
for search_result in results:
uri = search_result.document_uri
if uri and uri not in seen_searched:
seen_searched.add(uri)
searched_uris.append(uri)
@ -87,8 +110,8 @@ async def run_skill_question(
executions = getattr(state, "executions", None)
n_executions = len(executions) if executions is not None else 0
return SkillRunResult(
answer=answer,
return CapabilityRunResult(
answer=agent_result.output,
cited_uris=cited_uris,
cited_chunk_ids=cited_chunk_ids,
searched_uris=searched_uris,

2
evaluations/evaluations/evaluators/citation.py
View file

@ -15,7 +15,7 @@ def _relevant_uris(ctx: EvaluatorContext) -> set[str]:
@dataclass
class CitationMAPEvaluator(Evaluator):
"""Average precision over the URIs the skill cited via the `cite` tool.
"""Average precision over the URIs the capability cited via the `cite` tool.
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from

2
evaluations/evaluations/submission.py
View file

@ -19,7 +19,7 @@ def _format_number(value: float) -> str:
def extract_prediction(output: str | None) -> str:
"""Pull the primary numeric answer from a skill output, for submission.
"""Pull the primary numeric answer from a capability output, for submission.
Restricts to a declared ``ANSWER:`` line when present (via ``_answer_segment``)
so reasoning numbers don't leak. Strips ``$`` and thousands separators,

100
evaluations/tests/test_benchmark.py
View file

@ -131,7 +131,9 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
patch(
"evaluations.benchmark.run_capability_question", new_callable=AsyncMock
),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -149,7 +151,9 @@ class TestRunQaBenchmarkJudgeModel:
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.benchmark.run_skill_question", new_callable=AsyncMock),
patch(
"evaluations.benchmark.run_capability_question", new_callable=AsyncMock
),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -195,27 +199,29 @@ class TestEvaluateDatasetJudgeModel:
class TestExperimentMetadataTargets:
def test_default_target_is_rag_skill(self) -> None:
def test_default_target_is_rag_capability(self) -> None:
result = build_experiment_metadata(
dataset_key="test", test_cases=1, config=AppConfig()
)
assert result["target"] == "rag-skill"
assert "skill_provider" not in result
assert "skill_model" not in result
assert result["target"] == "rag-capability"
assert "capability_provider" not in result
assert "capability_model" not in result
def test_skill_target_includes_skill_config(self) -> None:
skill = ModelConfig(provider="ollama", name="gpt-oss-large", temperature=0.2)
def test_capability_target_includes_capability_config(self) -> None:
capability = ModelConfig(
provider="ollama", name="gpt-oss-large", temperature=0.2
)
result = build_experiment_metadata(
dataset_key="test",
test_cases=1,
config=AppConfig(),
target="rag-skill",
skill_config=skill,
target="rag-capability",
capability_config=capability,
)
assert result["target"] == "rag-skill"
assert result["skill_provider"] == "ollama"
assert result["skill_model"] == "gpt-oss-large"
assert result["skill_temperature"] == 0.2
assert result["target"] == "rag-capability"
assert result["capability_provider"] == "ollama"
assert result["capability_model"] == "gpt-oss-large"
assert result["capability_temperature"] == 0.2
class TestEvaluateDatasetTarget:
@ -230,8 +236,8 @@ class TestEvaluateDatasetTarget:
)
@pytest.mark.asyncio
async def test_threads_target_and_skill_model(self) -> None:
skill = ModelConfig(provider="ollama", name="gpt-oss")
async def test_threads_target_and_capability_model(self) -> None:
capability = ModelConfig(provider="ollama", name="gpt-oss")
with patch(
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
) as mock_qa:
@ -244,16 +250,16 @@ class TestEvaluateDatasetTarget:
limit=None,
name=None,
db_path=None,
target="rag-skill",
skill_model=skill,
target="rag-capability",
capability_model=capability,
)
mock_qa.assert_called_once()
assert mock_qa.call_args[1]["target"] == "rag-skill"
assert mock_qa.call_args[1]["skill_model"] is skill
assert mock_qa.call_args[1]["target"] == "rag-capability"
assert mock_qa.call_args[1]["capability_model"] is capability
@pytest.mark.asyncio
async def test_default_target_is_rag_skill(self) -> None:
async def test_default_target_is_rag_capability(self) -> None:
with patch(
"evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
) as mock_qa:
@ -267,11 +273,11 @@ class TestEvaluateDatasetTarget:
name=None,
db_path=None,
)
assert mock_qa.call_args[1]["target"] == "rag-skill"
assert mock_qa.call_args[1]["skill_model"] is None
assert mock_qa.call_args[1]["target"] == "rag-capability"
assert mock_qa.call_args[1]["capability_model"] is None
class TestRunQaBenchmarkSkillTarget:
class TestRunQaBenchmarkCapabilityTarget:
def _spec(self, tmp_path: Path) -> DatasetSpec:
return DatasetSpec(
key="test",
@ -283,17 +289,19 @@ class TestRunQaBenchmarkSkillTarget:
)
@pytest.mark.asyncio
async def test_rag_skill_target_uses_run_skill_question(
async def test_rag_capability_target_uses_run_capability_question(
self, tmp_path: Path
) -> None:
from evaluations.skill_runner import SkillRunResult
from evaluations.capability_runner import CapabilityRunResult
skill_run = AsyncMock(return_value=SkillRunResult(answer="from skill"))
capability_run = AsyncMock(
return_value=CapabilityRunResult(answer="from capability")
)
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch(
"evaluations.benchmark.run_skill_question", new=skill_run
) as mock_run_skill,
"evaluations.benchmark.run_capability_question", new=capability_run
) as mock_run_capability,
patch("evaluations.benchmark.HaikuRAG") as mock_haiku,
):
mock_get_model.return_value = "fake-model"
@ -301,27 +309,31 @@ class TestRunQaBenchmarkSkillTarget:
self._spec(tmp_path),
AppConfig(),
db_path=tmp_path / "test.lancedb",
target="rag-skill",
target="rag-capability",
)
# When target is rag-skill, HaikuRAG context manager is NOT entered
# (the skill manages its own client via lifespan).
# When target is rag-capability, HaikuRAG context manager is NOT entered
# (the capability manages its own client via lifespan).
mock_haiku.assert_not_called()
# skill model defaults to qa.model when not provided
skill_call = mock_get_model.call_args_list[-1]
assert skill_call[0][0] == AppConfig().qa.model
assert mock_run_skill is skill_run
# capability model defaults to qa.model when not provided
capability_call = mock_get_model.call_args_list[-1]
assert capability_call[0][0] == AppConfig().qa.model
assert mock_run_capability is capability_run
@pytest.mark.asyncio
async def test_analysis_skill_target_resolves_factory(self, tmp_path: Path) -> None:
from evaluations.benchmark import _skill_factory_for_target
from haiku.rag.skills.analysis import create_skill as analysis_factory
from haiku.rag.skills.rag import create_skill as rag_factory
async def test_analysis_capability_target_resolves_factory(
self, tmp_path: Path
) -> None:
from evaluations.benchmark import _capability_factory_for_target
from haiku.rag.capabilities.analysis import (
create_capability as analysis_factory,
)
from haiku.rag.capabilities.rag import create_capability as rag_factory
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("unknown") # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
assert _capability_factory_for_target("rag-capability") is rag_factory
assert _capability_factory_for_target("analysis-capability") is analysis_factory
with pytest.raises(ValueError, match="not a capability target"):
_capability_factory_for_target("unknown") # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
class TestCitationEvaluatorWiring:

58
evaluations/tests/test_capability_runner.py Normal file
View file

@ -0,0 +1,58 @@
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch
import pytest
from pydantic_ai.models.test import TestModel
from evaluations.capability_runner import run_capability_question
from haiku.rag.capabilities.analysis import create_capability as create_analysis
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig
async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
result = await run_capability_question(
create_rag,
tmp_path / "rag.lancedb",
AppConfig(),
"hello",
TestModel(call_tools=[]),
document_filter="uri = 'manual.pdf'",
)
assert result.answer == "success (no tool calls)"
assert result.cited_uris == []
assert result.n_searches == 0
async def test_runs_analysis_capability_without_legacy_capability_layer(tmp_path):
result = await run_capability_question(
create_analysis,
tmp_path / "rag.lancedb",
AppConfig(),
"hello",
TestModel(call_tools=[]),
request_limit=5,
)
assert result.answer == "success (no tool calls)"
assert result.n_executions == 0
@pytest.mark.parametrize(("override", "expected"), [(None, 30), (5, 5)])
async def test_analysis_capability_applies_request_limit(tmp_path, override, expected):
with patch(
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
) as run:
run.return_value = SimpleNamespace(output="done")
await run_capability_question(
create_analysis,
tmp_path / "rag.lancedb",
AppConfig(),
"hello",
TestModel(call_tools=[]),
request_limit=override,
)
assert run.call_args.kwargs["usage_limits"].request_limit == expected

313
evaluations/tests/test_skill_runner.py
View file

@ -1,313 +0,0 @@
import random
from pathlib import Path
from typing import Any
import pytest
from pydantic_ai.models.test import TestModel
from evaluations.skill_runner import SkillRunResult, run_skill_question
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
from haiku.rag.skills.analysis import (
AnalysisState,
create_skill as create_analysis_skill,
)
from haiku.rag.skills.rag import RAGState, create_skill as create_rag_skill
from haiku.rag.store.models.chunk import SearchResult
VECTOR_DIM = 2560
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch: pytest.MonkeyPatch):
"""Deterministic embeddings so search is reproducible."""
async def fake_embed_query(self, text):
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(VECTOR_DIM)]
async def fake_embed_documents(self, texts):
result = []
for t in texts:
random.seed(hash(t) % (2**32))
result.append([random.random() for _ in range(VECTOR_DIM)])
return result
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
@pytest.fixture
def app_config():
return AppConfig(environment="skill-runner-test")
@pytest.fixture
async def rag_db(tmp_path: Path):
"""A small two-document database."""
db_path = tmp_path / "test.lancedb"
async with HaikuRAG(db_path, create=True) as rag:
await rag.create_document(
"Artificial intelligence is transforming healthcare and finance.",
title="AI Overview",
uri="test://ai",
)
await rag.create_document(
"Machine learning includes supervised, unsupervised, and reinforcement.",
title="ML Basics",
uri="test://ml",
)
return db_path
class TestRunSkillQuestionMocked:
"""Verify the runner reads state correctly without going through a real skill loop."""
async def test_extracts_cited_and_searched_uris(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
state.citation_index["c1"] = Citation(
chunk_id="c1",
document_id="d1",
document_uri="test://doc-a",
document_title="A",
content="alpha",
)
state.citation_index["c2"] = Citation(
chunk_id="c2",
document_id="d2",
document_uri="test://doc-b",
document_title="B",
content="beta",
)
state.citations = ["c1", "c2"]
state.searches["q1"] = [
SearchResult(content="x", score=0.9, document_uri="test://doc-a"),
SearchResult(content="y", score=0.8, document_uri="test://doc-c"),
]
state.searches["q2"] = [
SearchResult(content="z", score=0.7, document_uri="test://doc-a"),
]
return "answer", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="anything?",
skill_model=TestModel(),
)
assert isinstance(result, SkillRunResult)
assert result.answer == "answer"
assert result.cited_chunk_ids == ["c1", "c2"]
assert result.cited_uris == ["test://doc-a", "test://doc-b"]
assert result.searched_uris == ["test://doc-a", "test://doc-c"]
assert result.n_searches == 2
async def test_skips_chunks_missing_from_index(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
state.citation_index["c1"] = Citation(
chunk_id="c1",
document_id="d1",
document_uri="test://doc-a",
content="a",
)
state.citations = ["c1", "missing"]
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert result.cited_chunk_ids == ["c1", "missing"]
assert result.cited_uris == ["test://doc-a"]
async def test_document_filter_is_set_on_state(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["filter"] = state.document_filter
captured["state_type"] = type(state)
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
document_filter="uri = 'test://ai'",
)
assert captured["filter"] == "uri = 'test://ai'"
assert captured["state_type"] is RAGState
async def test_request_limit_override(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["request_limit"] = skill.request_limit
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
request_limit=42,
)
assert captured["request_limit"] == 42
async def test_request_limit_unset_leaves_skill_default(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["request_limit"] = skill.request_limit
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert captured["request_limit"] is None
async def test_analysis_skill_uses_analysis_state(
self, monkeypatch: pytest.MonkeyPatch, app_config: AppConfig, rag_db: Path
) -> None:
captured: dict = {}
async def fake_run_skill(
model: Any,
skill: Any,
request: str,
state: Any = None,
event_sink: Any = None,
) -> tuple[str, list[Any], list[Any]]:
captured["state_type"] = type(state)
return "ok", [], []
monkeypatch.setattr("evaluations.skill_runner.run_skill", fake_run_skill)
await run_skill_question(
skill_factory=create_analysis_skill,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
assert captured["state_type"] is AnalysisState
async def test_raises_when_skill_has_no_state_type(
self, app_config: AppConfig, rag_db: Path
) -> None:
from haiku.skills.models import Skill, SkillMetadata, SkillSource
def factory(*, db_path, config) -> Skill:
return Skill(
metadata=SkillMetadata(name="bare", description="No state."),
source=SkillSource.ENTRYPOINT,
instructions="Do nothing.",
)
with pytest.raises(ValueError, match="no state_type"):
await run_skill_question(
skill_factory=factory,
db_path=rag_db,
config=app_config,
question="?",
skill_model=TestModel(),
)
class TestRunSkillQuestionEndToEnd:
"""Real skill loop against a real LanceDB. Verifies the wiring beyond mocks."""
async def test_rag_skill_runs_against_real_db(
self,
allow_model_requests: None,
app_config: AppConfig,
rag_db: Path,
) -> None:
result = await run_skill_question(
skill_factory=create_rag_skill,
db_path=rag_db,
config=app_config,
question="What is machine learning?",
skill_model=TestModel(call_tools=["search"]),
)
assert isinstance(result, SkillRunResult)
assert result.answer
assert result.n_searches >= 1
assert all(uri.startswith("test://") for uri in result.searched_uris)
@pytest.fixture
def allow_model_requests():
import pydantic_ai.models
with pydantic_ai.models.override_allow_model_requests(True):
yield

4
examples/README.md
View file

@ -16,7 +16,7 @@ See `docker/README.md` for setup instructions.
**Script:** `custom_agent.py`
Uses the RAG skill with `SkillToolset` to build a conversational agent.
Uses the deferred RAG capability to build a conversational agent.
```bash
uv run python examples/custom_agent.py /path/to/db.lancedb
@ -26,7 +26,7 @@ uv run python examples/custom_agent.py /path/to/db.lancedb
**Script:** `custom_agent_agui.py`
A Starlette app that serves an AG-UI streaming endpoint using the RAG skill with `SkillToolset`.
A Starlette app that adapts a native RAG-capable agent to AG-UI.
```bash
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000

15
examples/custom_agent.py
View file

@ -1,7 +1,6 @@
"""Custom agent using the haiku.rag RAG skill.
"""Custom agent using the native haiku.rag RAG capability.
Demonstrates how to use the RAG skill with haiku.skills SkillToolset
to build a conversational agent.
Demonstrates composing a deferred Pydantic AI capability into an agent.
Requirements:
- An Ollama instance running locally (default embedder)
@ -18,19 +17,15 @@ from pathlib import Path
from pydantic_ai import Agent
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt
from haiku.rag.capabilities.rag import create_capability
async def main(db_path: str) -> None:
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
capability = create_capability(db_path=Path(db_path))
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
capabilities=[capability],
)
print("Custom agent ready. Ctrl+C to exit.\n")

17
examples/custom_agent_agui.py
View file

@ -1,7 +1,7 @@
"""Custom agent with AG-UI streaming.
A Starlette app that serves an AG-UI streaming endpoint using the
haiku.rag RAG skill with haiku.skills SkillToolset.
haiku.rag's native Pydantic AI RAG capability.
Requirements:
- An Ollama instance running locally (default embedder)
@ -23,9 +23,7 @@ from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset, run_agui_stream
from haiku.skills.prompts import build_system_prompt
from haiku.rag.capabilities.rag import create_capability
db_path = os.environ.get("DB_PATH")
if not db_path:
@ -34,13 +32,11 @@ if not db_path:
)
sys.exit(1)
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
capability = create_capability(db_path=Path(db_path))
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
capabilities=[capability],
)
@ -52,9 +48,8 @@ async def stream_chat(request: Request) -> Response:
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
async def event_stream():
async with run_agui_stream(adapter, toolset=toolset) as stream:
async for chunk in adapter.encode_stream(stream):
yield chunk
async for chunk in adapter.encode_stream(adapter.run_stream()):
yield chunk
return StreamingResponse(
event_stream(),

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

@ -620,7 +620,7 @@ class HaikuRAGApp: # pragma: no cover
question: str,
filter: str | None = None,
):
"""Answer a question using the rag-analysis skill.
"""Answer a question using the analysis capability.
Args:
question: The question to answer
@ -634,7 +634,7 @@ class HaikuRAGApp: # pragma: no cover
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print(
"[dim]Running analysis skill with code execution...[/dim]"
"[dim]Running analysis capability with code execution...[/dim]"
)
self.console.print()

13
haiku_rag_slim/haiku/rag/capabilities/__init__.py Normal file
View file

@ -0,0 +1,13 @@
"""Native Pydantic AI capabilities provided by haiku.rag."""
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
from haiku.rag.capabilities.rag import RAGCapability, RAGState
__all__ = [
"AnalysisCapability",
"AnalysisState",
"RAGCapability",
"RAGCapabilityBase",
"RAGState",
]

248
haiku_rag_slim/haiku/rag/capabilities/_base.py Normal file
View file

@ -0,0 +1,248 @@
import asyncio
import os
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any, cast
from pydantic import BaseModel
from pydantic_ai import ModelRetry, RunContext
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import (
ModelMessage,
ModelRequest,
ToolReturn,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.run import AgentRunResult
from pydantic_ai.toolsets import AgentToolset
from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation, resolve_citations
from haiku.rag.tools.search import build_binary_parts_from_results
def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
if db_path is not None:
return db_path
if env_db := os.environ.get("HAIKU_RAG_DB"):
return Path(env_db).expanduser()
return config.storage.data_dir / "haiku.rag.lancedb"
def _clear_invocation_state(state: BaseModel) -> None:
for field_name in ("citations", "searches", "executions"):
value = getattr(state, field_name, None)
if hasattr(value, "clear"):
value.clear()
def _compact_old_tool_returns(
messages: list[ModelMessage], tool_names: frozenset[str]
) -> list[ModelMessage]:
"""Remove bulky prior-turn evidence while retaining the current turn.
Tool call and return parts remain paired; only the old return payload is
replaced. This keeps provider histories valid and preserves all evidence
gathered since the most recent user prompt.
"""
latest_user_message = -1
for index, message in enumerate(messages):
if isinstance(message, ModelRequest) and any(
isinstance(part, UserPromptPart) for part in message.parts
):
latest_user_message = index
if latest_user_message < 0:
return messages
compacted = list(messages)
for index, message in enumerate(messages[:latest_user_message]):
if not isinstance(message, ModelRequest):
continue
parts = [
replace(
part,
content="[Prior-turn RAG tool output removed; citations remain in state.]",
)
if isinstance(part, ToolReturnPart) and part.tool_name in tool_names
else part
for part in message.parts
]
if parts != message.parts:
compacted[index] = replace(message, parts=parts)
return compacted
@dataclass
class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
db_path: Path
config: AppConfig
state_type: type[StateT]
state_namespace: str
instruction_text: str
model: ModelConfig
tool_names: frozenset[str]
default_request_limit: int | None = None
state: StateT | None = field(default=None, repr=False)
outer_state: dict[str, Any] | None = field(default=None, repr=False)
rag: HaikuRAG | None = field(default=None, repr=False)
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
search_count: int = field(default=0, repr=False)
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
outer = getattr(ctx.deps, "state", None)
outer_state = outer if isinstance(outer, dict) else None
raw_state = outer_state.get(self.state_namespace) if outer_state else None
state = self.state_type.model_validate(raw_state or {})
_clear_invocation_state(state)
run_capability = replace(
self,
state=state,
outer_state=outer_state,
rag=None,
rag_lock=asyncio.Lock(),
resource_lock=asyncio.Lock(),
search_count=0,
)
run_capability._sync_state()
return run_capability
def get_instructions(self) -> str:
if self.config.prompts.domain_preamble:
return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}"
return self.instruction_text
async def before_model_request(
self, ctx: RunContext[Any], request_context: ModelRequestContext
) -> ModelRequestContext:
request_context.messages = _compact_old_tool_returns(
request_context.messages, self.tool_names
)
return request_context
async def after_run(
self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
) -> AgentRunResult[Any]:
await self._close()
return result
async def on_run_error(
self, ctx: RunContext[Any], *, error: BaseException
) -> AgentRunResult[Any]:
await self._close()
raise error
async def _ensure_rag(self) -> HaikuRAG:
if self.rag is None:
async with self.resource_lock:
if self.rag is None:
rag = HaikuRAG(self.db_path, config=self.config, read_only=True)
await rag.__aenter__()
self.rag = rag
return self.rag
async def _close(self) -> None:
if self.rag is not None:
await self.rag.__aexit__(None, None, None)
self.rag = None
def _sync_state(self) -> None:
if self.outer_state is not None and self.state is not None:
self.outer_state[self.state_namespace] = self.state.model_dump(mode="json")
async def _with_state(self, operation: Any) -> Any:
"""Execute an operation and copy its state back to the host dependencies."""
result = await operation
self._sync_state()
return result
async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
assert self.state is not None
self.search_count += 1
if self.search_count > self.config.qa.max_searches:
return (
"Search limit reached. Answer the question using "
"the results you already have."
)
async with self.rag_lock:
formatted, results = await search_corpus(
await self._ensure_rag(),
query,
limit=limit,
document_filter=getattr(self.state, "document_filter", None),
)
state = cast(Any, self.state)
state.searches[query] = results
if self.model.vision and (parts := build_binary_parts_from_results(results)):
return ToolReturn(return_value=formatted, content=parts)
return formatted
async def _cite(self, chunk_ids: list[str]) -> str:
assert self.state is not None
if not chunk_ids:
return "Registered 0 citations (empty chunk_ids)."
all_results: list[SearchResult] = []
state = cast(Any, self.state)
for results in state.searches.values():
all_results.extend(results)
citations = resolve_citations(chunk_ids, all_results)
resolved = {citation.chunk_id for citation in citations}
missing = [
cid.strip("[]") for cid in chunk_ids if cid.strip("[]") not in resolved
]
if missing:
async with self.rag_lock:
rag = await self._ensure_rag()
synthetic: list[SearchResult] = []
documents: dict[str, Any] = {}
for chunk_id in missing:
chunk = await rag.get_chunk_by_id(chunk_id)
if chunk is None or not chunk.document_id:
continue
document = documents.get(chunk.document_id)
if chunk.document_id not in documents:
document = await rag.get_document_by_id(chunk.document_id)
documents[chunk.document_id] = document
chunk.document_uri = document.uri if document else None
chunk.document_title = document.title if document else None
chunk.document_meta = document.metadata if document else {}
synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
citations.extend(resolve_citations(missing, synthetic))
if not citations:
raise ModelRetry(
f"None of the supplied chunk_ids {list(chunk_ids)} could be resolved. "
"Copy chunk_ids verbatim from search results."
)
self._register_citations(citations)
return f"Registered {len(citations)} citation(s)."
def _register_citations(self, citations: list[Citation]) -> None:
assert self.state is not None
state = cast(Any, self.state)
next_index = len(state.citation_index) + 1
for citation in citations:
if citation.chunk_id not in state.citation_index:
citation.index = next_index
next_index += 1
state.citation_index[citation.chunk_id] = citation
if citation.chunk_id not in state.citations:
state.citations.append(citation.chunk_id)
def get_toolset(self) -> AgentToolset[Any] | None:
raise NotImplementedError
__all__ = [
"CodeExecutionEntry",
"RAGCapabilityBase",
"resolve_db_path",
]

33
haiku_rag_slim/haiku/rag/capabilities/_tools.py Normal file
View file

@ -0,0 +1,33 @@
from pydantic import BaseModel
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.chunk import SearchResult
class CodeExecutionEntry(BaseModel):
code: str
stdout: str
stderr: str = ""
success: bool = True
async def search_corpus(
rag: HaikuRAG,
query: str,
limit: int | None = None,
document_filter: str | None = None,
) -> tuple[str, list[SearchResult]]:
"""Search and context-expand results for a capability tool."""
results = await rag.search(query, limit=limit, filter=document_filter)
results = await rag.expand_context(results)
formatted = "\n\n---\n\n".join(
result.format_for_agent(rank=index + 1, total=len(results))
for index, result in enumerate(results)
)
return formatted, list(results)
__all__ = [
"CodeExecutionEntry",
"search_corpus",
]

157
haiku_rag_slim/haiku/rag/capabilities/analysis.py Normal file
View file

@ -0,0 +1,157 @@
from dataclasses import dataclass, field
from functools import cache
from pathlib import Path
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
from haiku.rag.capabilities._base import (
CodeExecutionEntry,
RAGCapabilityBase,
resolve_db_path,
)
from haiku.rag.config.models import AppConfig
from haiku.rag.sandbox import AnalysisContext, Sandbox
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
STATE_NAMESPACE = "analysis"
_CAPABILITY_ID = "haiku-rag-analysis"
_TOOL_NAMES = frozenset({"analysis_search", "analysis_execute_code", "analysis_cite"})
_instructions_path = Path(__file__).parent / "instructions" / "analysis.md"
class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
@cache
def instructions() -> str:
return _instructions_path.read_text().strip()
@dataclass
class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
"""Deferred capability for sandboxed computation over a RAG corpus."""
sandbox: Sandbox | None = field(default=None, repr=False)
execute_count: int = field(default=0, repr=False)
async def for_run(self, ctx: RunContext[Any]) -> "AnalysisCapability":
capability = await super().for_run(ctx)
assert isinstance(capability, AnalysisCapability)
capability.sandbox = None
capability.execute_count = 0
return capability
async def _ensure_sandbox(self) -> Sandbox:
if self.sandbox is None:
rag = await self._ensure_rag()
assert self.state is not None
self.sandbox = Sandbox(
db_path=self.db_path,
config=self.config,
context=AnalysisContext(filter=self.state.document_filter),
rag=rag,
lock=self.rag_lock,
)
return self.sandbox
async def _close(self) -> None:
if self.sandbox is not None:
self.sandbox.close()
self.sandbox = None
await super()._close()
async def _execute_code(self, code: str) -> str:
assert self.state is not None
self.execute_count += 1
if self.execute_count > self.config.analysis.max_executions:
return (
"Code-execution limit reached. Give your final answer now from what "
"you already have; do not call analysis_execute_code again."
)
sandbox = await self._ensure_sandbox()
result = await sandbox.execute(code)
if sandbox._search_results:
existing = self.state.searches.get("_sandbox", [])
seen = {item.chunk_id for item in existing}
for item in sandbox._search_results:
if item.chunk_id not in seen:
existing.append(item)
seen.add(item.chunk_id)
self.state.searches["_sandbox"] = existing
self.state.executions.append(
CodeExecutionEntry(
code=code,
stdout=result.stdout,
stderr=result.stderr,
success=result.success,
)
)
if result.success:
return result.stdout or "No output."
return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
def get_toolset(self) -> FunctionToolset[Any]:
async def analysis_search(
ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn:
"""Search the knowledge base for evidence to analyze."""
return await self._with_state(self._search(query, limit))
async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any:
"""Execute Python against the sandboxed document filesystem."""
return await self._with_state(self._execute_code(code))
async def analysis_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
"""Register exact retrieved chunk IDs as citations for the answer."""
return await self._with_state(self._cite(chunk_ids))
return FunctionToolset(
[analysis_search, analysis_execute_code, analysis_cite], id=_CAPABILITY_ID
)
def create_capability(
db_path: Path | None = None,
config: AppConfig | None = None,
*,
defer_loading: bool = True,
) -> AnalysisCapability:
"""Create a native Pydantic AI analysis capability."""
if config is None:
from haiku.rag.config import get_config
config = get_config()
return AnalysisCapability(
db_path=resolve_db_path(db_path, config),
config=config,
state_type=AnalysisState,
state_namespace=STATE_NAMESPACE,
instruction_text=instructions(),
model=config.analysis.model or config.qa.model,
tool_names=_TOOL_NAMES,
default_request_limit=30,
id=_CAPABILITY_ID,
description=(
"Analyze the haiku.rag corpus with search and sandboxed Python code."
),
defer_loading=defer_loading,
)
__all__ = [
"AnalysisCapability",
"AnalysisState",
"STATE_NAMESPACE",
"create_capability",
"instructions",
]

46
haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md → haiku_rag_slim/haiku/rag/capabilities/instructions/analysis.md
View file

@ -1,25 +1,15 @@
---
name: rag-analysis
description: >
Computational analysis of the knowledge base via code execution in a sandboxed Python interpreter.
Use for questions requiring counting, aggregation, statistics, data traversal,
comparison across documents, or any task best answered by writing Python code.
Examples: "how many pages?", "compare table 3 across documents",
"calculate average word count", "extract all email addresses".
---
# Analysis
You answer questions over a document knowledge base. Two common workflows:
- **`search → cite → answer`** when the answer is grounded on specific document content. Call `cite` with the supporting chunk_ids before writing the answer.
- **`execute_code → answer`** when the answer is a count, aggregation, listing, or structural computation over the corpus (e.g. "how many documents?", "average page count"). No `cite` is needed when no specific chunks support the answer.
- **`analysis_search → analysis_cite → answer`** when the answer is grounded on specific document content. Call `analysis_cite` with the supporting chunk_ids before writing the answer.
- **`analysis_execute_code → answer`** when the answer is a count, aggregation, listing, or structural computation over the corpus (e.g. "how many documents?", "average page count"). No `analysis_cite` is needed when no specific chunks support the answer.
You can mix the two. The rule: cite when grounded on retrieved evidence; don't fabricate citations for corpus-level computation.
## Tools
### execute_code
### analysis_execute_code
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`):
@ -29,13 +19,13 @@ Inside the code, these functions are available (use `await`):
Available modules: `json`, `re`, `math`, `pathlib`
Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
### search
### analysis_search
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 the chunk IDs that ground your answer. **You must call `cite` before writing any final answer that uses retrieved evidence — search results, items.jsonl rows, toc.json nodes, or content.txt content.** Skipping `cite` leaves the answer ungrounded and is treated as a failure.
### analysis_cite
Register the chunk IDs that ground your answer. **You must call `analysis_cite` before writing any final answer that uses retrieved evidence — search results, items.jsonl rows, toc.json nodes, or content.txt content.** Skipping `analysis_cite` leaves the answer ungrounded and is treated as a failure.
`cite` is **not** required when your answer is a corpus-level computation that doesn't draw on specific chunks — counts, aggregations, listings, averages across documents. Don't fabricate citations for these.
`analysis_cite` is **not** required when your answer is a corpus-level computation that doesn't draw on specific chunks — counts, aggregations, listings, averages across documents. Don't fabricate citations for these.
Chunk IDs come from two places:
- The `chunk_id` field on `search` / `await search(...)` results
@ -93,11 +83,11 @@ Each row carries:
- `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
- `chunk_ids`: chunks that contain this item — pass to `analysis_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.
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 `analysis_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`. 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.
@ -105,19 +95,19 @@ Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) tha
## Strategy
1. Search first.
2. Identify the chunk_ids from the search results that support your answer and call `cite` with them. Then 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.
2. Identify the chunk_ids from the search results that support your answer and call `analysis_cite` with them. Then write a concise answer.
3. Reach for `analysis_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. Before writing your final response, call `cite` with the chunk_ids that ground your answer.
5. Before writing your final response, call `analysis_cite` with the chunk_ids that ground your answer.
You MUST call `cite` with at least one chunk ID before producing your final answer **when your answer is grounded on retrieved evidence**. Skip `cite` in two cases: (a) you are refusing for lack of information, or (b) your answer is a corpus-level computation (count, aggregation, listing) that doesn't draw on specific chunks. In those cases do **not** fabricate citations.
You MUST call `analysis_cite` with at least one chunk ID before producing your final answer **when your answer is grounded on retrieved evidence**. Skip `analysis_cite` in two cases: (a) you are refusing for lack of information, or (b) your answer is a corpus-level computation (count, aggregation, listing) that doesn't draw on specific chunks. In those cases do **not** fabricate citations.
## Important
- Variables persist between `execute_code` calls — you can search in one call and process results in the next
- Variables persist between `analysis_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
- 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)
- 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 `analysis_search → analysis_cite`.
- Use `await` for all async functions inside `analysis_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 — your answer should read naturally. Use the `cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `cite` tool call registers a citation.
- **Before you write your final answer, invoke the `cite` tool with the supporting chunk_ids.** This is the last tool call before answering whenever your answer draws on retrieved evidence.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `analysis_cite` tool separately to register citations. `cite{...}` markdown-style inline references do nothing; only an actual `analysis_cite` tool call registers a citation.
- **Before you write your final answer, invoke the `analysis_cite` tool with the supporting chunk_ids.** This is the last tool call before answering whenever your answer draws on retrieved evidence.

21
haiku_rag_slim/haiku/rag/skills/rag/SKILL.md → haiku_rag_slim/haiku/rag/capabilities/instructions/rag.md
View file

@ -1,8 +1,3 @@
---
name: rag
description: Search, retrieve and analyze documents using RAG (Retrieval Augmented Generation).
---
# RAG
You are a RAG assistant with access to a document knowledge base.
@ -10,7 +5,7 @@ Use your tools to search and answer questions. Never make up information — alw
## Tools
### search
### rag_search
Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content.
Each result includes:
@ -21,20 +16,20 @@ Each result includes:
When a result's Type is `picture`, the corresponding figure may also be attached to the tool response as an image alongside the text. Use the image directly to answer questions about figures, diagrams, charts, screenshots.
### 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`.
### rag_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 `rag_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
1. Call `rag_search` with relevant keywords from the question
2. Review the results — they are ordered by relevance (rank 1 = best match)
3. If needed, search again with different keywords (you have a limited number of searches)
4. Identify the chunk IDs that support your answer and call `cite` with them
4. Identify the chunk IDs that support your answer and call `rag_cite` with them
5. Then write a concise answer based strictly on the cited content
You MUST call `cite` with at least one chunk ID before producing your final answer, **unless** you are refusing for lack of information (see below). Answers without citations are considered ungrounded.
You MUST call `rag_cite` with at least one chunk ID before producing your final answer, **unless** you are refusing for lack of information (see below). Answers without citations are considered ungrounded.
## Guidelines
@ -43,8 +38,8 @@ You MUST call `cite` with at least one chunk ID before producing your final answ
- If multiple results are relevant, synthesize them coherently
- Be concise and direct — avoid elaboration unless asked
- If the search tool tells you the search limit is reached, stop searching 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 from tangentially related content. In this refusal case do **not** call `cite` — there is nothing to cite.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
- 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 from tangentially related content. In this refusal case do **not** call `rag_cite` — there is nothing to cite.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `rag_cite` tool separately to register citations.
## When search returns irrelevant results

97
haiku_rag_slim/haiku/rag/capabilities/rag.py Normal file
View file

@ -0,0 +1,97 @@
from dataclasses import dataclass
from functools import cache
from pathlib import Path
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from pydantic_ai.messages import ToolReturn
from pydantic_ai.toolsets import FunctionToolset
from haiku.rag.capabilities._base import (
RAGCapabilityBase,
resolve_db_path,
)
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.citation import Citation
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. NEVER make up information - always use capabilities to get facts from the knowledge base
3. When a capability returns citations, always include them in your response
"""
STATE_NAMESPACE = "rag"
_CAPABILITY_ID = "haiku-rag"
_TOOL_NAMES = frozenset({"rag_search", "rag_cite"})
_instructions_path = Path(__file__).parent / "instructions" / "rag.md"
class RAGState(BaseModel):
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
@cache
def instructions() -> str:
return _instructions_path.read_text().strip()
@dataclass
class RAGCapability(RAGCapabilityBase[RAGState]):
"""Deferred, native Pydantic AI capability for grounded RAG queries."""
def get_toolset(self) -> FunctionToolset[Any]:
async def rag_search(
ctx: RunContext[Any], query: str, limit: int | None = None
) -> str | ToolReturn:
"""Search the knowledge base using hybrid vector and full-text search."""
return await self._with_state(self._search(query, limit))
async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
"""Register exact search-result chunk IDs as citations for the answer."""
return await self._with_state(self._cite(chunk_ids))
return FunctionToolset([rag_search, rag_cite], id=_CAPABILITY_ID)
def create_capability(
db_path: Path | None = None,
config: AppConfig | None = None,
*,
defer_loading: bool = True,
) -> RAGCapability:
"""Create a native Pydantic AI RAG capability."""
if config is None:
from haiku.rag.config import get_config
config = get_config()
return RAGCapability(
db_path=resolve_db_path(db_path, config),
config=config,
state_type=RAGState,
state_namespace=STATE_NAMESPACE,
instruction_text=instructions(),
model=config.qa.model,
tool_names=_TOOL_NAMES,
id=_CAPABILITY_ID,
description=(
"Search the haiku.rag knowledge base and cite evidence for grounded answers."
),
defer_loading=defer_loading,
)
__all__ = [
"AGENT_PREAMBLE",
"RAGCapability",
"RAGState",
"STATE_NAMESPACE",
"create_capability",
"instructions",
]

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

@ -5,7 +5,7 @@ def run_chat(
db_path: Path | None = None,
read_only: bool = False,
model: str | None = None,
skills: list[str] | None = None,
capabilities: list[str] | None = None,
) -> None:
"""Run the chat TUI.
@ -13,7 +13,7 @@ def run_chat(
db_path: Path to the LanceDB database. If None, uses default from config.
read_only: Whether to open the database in read-only mode.
model: Model to use for the chat.
skills: Skills to enable ("rag", "analysis"). Defaults to ["rag"].
capabilities: Capabilities to enable ("rag", "analysis"). Defaults to ["rag"].
"""
try:
from haiku.rag.chat.app import ChatApp
@ -24,7 +24,6 @@ def run_chat(
from haiku.rag.config import get_config
from haiku.rag.utils import get_model, parse_model_option
from haiku.skills.models import Skill
config = get_config()
if db_path is None:
@ -35,22 +34,22 @@ def run_chat(
config.qa.model = model_config
config.analysis.model = model_config
enabled = skills or ["rag"]
skill_list: list[Skill] = []
enabled = capabilities or ["rag"]
capability_list = []
if "rag" in enabled:
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.capabilities.rag import create_capability
skill_list.append(create_rag_skill(db_path=db_path, config=config))
capability_list.append(create_capability(db_path=db_path, config=config))
if "analysis" in enabled:
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
from haiku.rag.capabilities.analysis import create_capability
skill_list.append(create_analysis_skill(db_path=db_path, config=config))
capability_list.append(create_capability(db_path=db_path, config=config))
app = ChatApp(
db_path,
skills=skill_list,
capabilities=capability_list,
read_only=read_only,
model=model or get_model(config.qa.model, config),
)

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

@ -1,43 +1,34 @@
import asyncio
import json
import uuid
from collections.abc import Iterable
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
import textual_image.widget # noqa: F401 - import early for renderer detection
from ag_ui.core import (
ActivitySnapshotEvent,
AssistantMessage,
EventType,
RunAgentInput,
StateDeltaEvent,
TextMessageContentEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallStartEvent,
UserMessage,
)
from jsonpatch import JsonPatch
from pydantic_ai import Agent
from pydantic_ai.ui.ag_ui import AGUIAdapter
from pydantic_ai.messages import (
FunctionToolCallEvent,
FunctionToolResultEvent,
PartDeltaEvent,
PartEndEvent,
PartStartEvent,
TextPart,
TextPartDelta,
)
from pydantic_ai.run import AgentRunResultEvent
from textual.app import App, SystemCommand
from textual.binding import Binding
from textual.widgets import Footer, Header, Input
from textual.worker import Worker
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.analysis import AnalysisState
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.skills.analysis import AnalysisState
from haiku.rag.skills.rag import RAGState, get_agent_preamble
from haiku.rag.telemetry import configure as configure_telemetry
from haiku.skills.agent import (
SkillToolset,
run_agui_stream,
)
from haiku.skills.models import Skill
from haiku.skills.prompts import build_system_prompt
configure_telemetry(service_name="haiku-rag")
@ -49,6 +40,11 @@ RAG_STATE_NAMESPACE = "rag"
ANALYSIS_STATE_NAMESPACE = "analysis"
@dataclass
class ChatDeps:
state: dict[str, Any] = field(default_factory=dict)
class ChatApp(App):
"""Textual TUI for conversational RAG."""
@ -82,27 +78,24 @@ class ChatApp(App):
def __init__(
self,
db_path: Path,
skills: list[Skill],
capabilities: list[RAGCapabilityBase[Any]],
read_only: bool = False,
model: str | None = None,
) -> None:
super().__init__()
self.db_path = db_path
self._skills = skills
self._capabilities = capabilities
self.read_only = read_only
self._model = model
self.client: HaikuRAG | None = None
self.config = get_config()
self._toolset: SkillToolset | None = None
self._agent: Agent[None, str] | None = None
self._agent: Agent[ChatDeps, str] | None = None
self._messages: list[Any] = []
self._state: dict[str, Any] = {}
self._is_processing = False
self._current_worker: Worker[None] | None = None
self._document_filter: list[str] = []
# Stable per-launch id so multi-turn chats land in one Logfire
# conversation. AGUIAdapter reads run_input.thread_id and exports it
# as the `gen_ai.conversation.id` OTel attribute on every agent run.
# Stable per-launch id for multi-turn model and telemetry correlation.
self._conversation_id = str(uuid.uuid4())
def compose(self) -> "ComposeResult":
@ -135,11 +128,6 @@ class ChatApp(App):
"Show database information",
self.action_show_info,
)
yield SystemCommand(
"View state",
"Show the current session state",
self.action_view_state,
)
async def on_mount(self) -> None:
"""Initialize the app when mounted."""
@ -150,16 +138,17 @@ class ChatApp(App):
)
await self.client.__aenter__()
self._toolset = SkillToolset(skills=self._skills)
self._agent = Agent(
self._model,
instructions=build_system_prompt(
self._toolset.skill_catalog,
preamble=get_agent_preamble(self.config),
),
toolsets=[self._toolset],
deps_type=ChatDeps,
instructions=AGENT_PREAMBLE,
capabilities=self._capabilities,
)
self._state = self._toolset.build_state_snapshot()
self._state = {}
for capability in self._capabilities:
self._state[capability.state_namespace] = (
capability.state_type().model_dump(mode="json")
)
self.query_one(Input).focus()
@ -179,14 +168,6 @@ class ChatApp(App):
chat_history = self.query_one(ChatHistory)
await chat_history.add_message("user", user_message)
self._messages.append(
UserMessage(
id=str(uuid.uuid4()),
role="user",
content=user_message,
)
)
self._is_processing = True
self.query_one(Input).disabled = True
self._current_worker = self.run_worker(
@ -195,113 +176,58 @@ class ChatApp(App):
async def _run_agent(self, user_message: str) -> None:
"""Run the agent in a background worker."""
if not self._agent or not self._toolset:
if not self._agent:
return
chat_history = self.query_one(ChatHistory)
await chat_history.show_thinking()
run_input = RunAgentInput(
thread_id=self._conversation_id,
run_id=str(uuid.uuid4()),
messages=self._messages,
state=self._state,
tools=[],
context=[],
forwarded_props={},
)
adapter = AGUIAdapter(agent=self._agent, run_input=run_input)
message = None
accumulated_text = ""
tool_args_deltas: dict[str, str] = {}
deps = ChatDeps(state=self._state)
try:
async with run_agui_stream(adapter, toolset=self._toolset) as stream:
async with self._agent.run_stream_events(
user_message,
message_history=self._messages,
conversation_id=self._conversation_id,
deps=deps,
) as stream:
async for event in stream:
if event.type == EventType.TEXT_MESSAGE_START:
if isinstance(event, PartStartEvent) and isinstance(
event.part, TextPart
):
chat_history.hide_thinking()
message = await chat_history.add_message("assistant")
accumulated_text = ""
elif event.type == EventType.TEXT_MESSAGE_CONTENT:
assert isinstance(event, TextMessageContentEvent)
accumulated_text += event.delta
if event.part.content:
await message.append_delta(event.part.content)
elif isinstance(event, PartDeltaEvent) and isinstance(
event.delta, TextPartDelta
):
if message:
await message.append_delta(event.delta)
await message.append_delta(event.delta.content_delta)
chat_history.scroll_end(animate=False)
elif event.type == EventType.TEXT_MESSAGE_END:
elif isinstance(event, PartEndEvent) and isinstance(
event.part, TextPart
):
if message:
await message.finish_stream()
self._messages.append(
AssistantMessage(
id=str(uuid.uuid4()),
role="assistant",
content=accumulated_text,
)
)
await self._show_citations_and_programs(chat_history)
elif event.type == EventType.TOOL_CALL_START:
assert isinstance(event, ToolCallStartEvent)
elif isinstance(event, FunctionToolCallEvent):
part = event.part
chat_history.hide_thinking()
await chat_history.add_tool_call(
event.tool_call_id, event.tool_call_name
part.tool_call_id, part.tool_name
)
chat_history.update_tool_args(
part.tool_call_id, part.args_as_dict()
)
tool_args_deltas[event.tool_call_id] = ""
await chat_history.show_thinking("Executing tasks...")
elif event.type == EventType.TOOL_CALL_ARGS:
assert isinstance(event, ToolCallArgsEvent)
tool_args_deltas[event.tool_call_id] = (
tool_args_deltas.get(event.tool_call_id, "") + event.delta
)
try:
args = json.loads(tool_args_deltas[event.tool_call_id])
chat_history.update_tool_args(event.tool_call_id, args)
except json.JSONDecodeError:
pass
elif event.type == EventType.TOOL_CALL_END:
assert isinstance(event, ToolCallEndEvent)
chat_history.mark_tool_complete(event.tool_call_id)
elif event.type == EventType.ACTIVITY_SNAPSHOT:
assert isinstance(event, ActivitySnapshotEvent)
content = event.content
if event.activity_type == "skill_tool_call":
tool_call_id = content["tool_call_id"]
skill_name = content.get("skill", "")
tool_name = content["tool_name"]
display_name = (
f"{skill_name} → {tool_name}"
if skill_name
else tool_name
)
args_str = content.get("args", "{}")
chat_history.hide_thinking()
await chat_history.add_tool_call(tool_call_id, display_name)
try:
args = json.loads(args_str)
chat_history.update_tool_args(tool_call_id, args)
except json.JSONDecodeError:
pass
await chat_history.show_thinking("Working...")
elif event.activity_type == "skill_tool_result":
tool_call_id = content["tool_call_id"]
chat_history.mark_tool_complete(tool_call_id)
elif event.type == EventType.STATE_DELTA:
assert isinstance(event, StateDeltaEvent)
patch = JsonPatch(event.delta)
self._state = patch.apply(self._state)
self._toolset.restore_state_snapshot(self._state)
elif event.type == EventType.STATE_SNAPSHOT:
self._state = getattr(event, "snapshot", self._state)
self._toolset.restore_state_snapshot(self._state)
elif event.type == EventType.RUN_FINISHED:
elif isinstance(event, FunctionToolResultEvent):
chat_history.mark_tool_complete(event.part.tool_call_id)
elif isinstance(event, AgentRunResultEvent):
self._messages = event.result.all_messages()
self._state = deps.state
chat_history.hide_thinking()
elif event.type == EventType.RUN_ERROR:
chat_history.hide_thinking()
error_msg = getattr(event, "message", "Unknown error")
await chat_history.add_message(
"assistant", f"Error: {error_msg}"
)
await self._show_citations_and_programs(chat_history)
except asyncio.CancelledError:
chat_history.hide_thinking()
@ -321,14 +247,14 @@ class ChatApp(App):
chat_input.focus()
async def _show_citations_and_programs(self, chat_history: "ChatHistory") -> None:
"""Show citations and programs from skill states after an agent response."""
if not self._toolset:
return
"""Show citations and programs from capability states after a response."""
citations = []
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
state = self._toolset.get_namespace(namespace)
if not state:
state_data = self._state.get(namespace)
if not state_data:
continue
state_type = RAGState if namespace == RAG_STATE_NAMESPACE else AnalysisState
state = state_type.model_validate(state_data)
cited_ids = getattr(state, "citations", [])
citation_index = getattr(state, "citation_index", {})
for cid in cited_ids:
@ -355,8 +281,8 @@ class ChatApp(App):
await chat_history.add_citations(citations, picture_bytes=picture_bytes)
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if analysis_state:
if analysis_data := self._state.get(ANALYSIS_STATE_NAMESPACE):
analysis_state = AnalysisState.model_validate(analysis_data)
executions = getattr(analysis_state, "executions", [])
successful = [e for e in executions if e.success]
if successful:
@ -367,9 +293,10 @@ class ChatApp(App):
chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages()
self._messages.clear()
# Reset state
if self._toolset:
self._state = self._toolset.build_state_snapshot()
self._state = {
capability.state_namespace: capability.state_type().model_dump(mode="json")
for capability in self._capabilities
}
# Cleared chat starts a fresh Logfire conversation.
self._conversation_id = str(uuid.uuid4())
@ -429,24 +356,6 @@ class ChatApp(App):
await self.push_screen(InfoModal(self.client, self.db_path))
def action_view_state(self) -> None:
"""Show the current session state."""
from haiku.skills.chat.app import StateScreen
self.push_screen(StateScreen(self._state, on_save=self._apply_state_edit))
def _apply_state_edit(self, new_state: dict[str, Any]) -> None:
if self._toolset is None:
return
if not isinstance(new_state, dict):
raise ValueError("state must be a JSON object")
for namespace, data in new_state.items():
current = self._toolset.get_namespace(namespace)
if current is not None:
type(current).model_validate(data)
self._toolset.restore_state_snapshot(new_state)
self._state = self._toolset.build_state_snapshot()
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection."""
chat_history = self.query_one(ChatHistory)
@ -476,12 +385,12 @@ class ChatApp(App):
self._document_filter = event.selected
if self._toolset:
doc_filter = build_multi_document_filter(self._document_filter)
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if isinstance(rag_state, RAGState):
rag_state.document_filter = doc_filter
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if isinstance(analysis_state, AnalysisState):
analysis_state.document_filter = doc_filter
self._state = self._toolset.build_state_snapshot()
doc_filter = build_multi_document_filter(self._document_filter)
for namespace, state_type in (
(RAG_STATE_NAMESPACE, RAGState),
(ANALYSIS_STATE_NAMESPACE, AnalysisState),
):
if namespace in self._state:
state = state_type.model_validate(self._state[namespace])
state.document_filter = doc_filter
self._state[namespace] = state.model_dump(mode="json")

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

@ -82,21 +82,12 @@ class ToolCallWidget(Static):
yield Static(desc, classes="tool-desc")
def _build_description(self) -> str:
if self.tool_name == "execute_skill":
skill = self.args.get("skill_name", "")
request = self.args.get("request", "...")
prefix = f"{skill}: " if skill else ""
return f'{prefix}"{request}"'
elif self.tool_name == "search":
if self.tool_name in {"rag_search", "analysis_search"}:
query = self.args.get("query", "...")
return f'"{query}"'
elif self.tool_name == "ask":
question = self.args.get("question", "...")
return f'"{question}"'
elif self.tool_name == "get_document":
query = self.args.get("query", "...")
return f'"{query}"'
elif self.args:
if self.tool_name == "analysis_execute_code":
return str(self.args.get("code", "..."))[:120]
if self.args:
return str(self.args)
return ""

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

@ -362,7 +362,7 @@ def ask( # pragma: no cover
)
@_cli.command("analyze", help="Answer questions using the rag-analysis skill")
@_cli.command("analyze", help="Answer questions using the analysis capability")
def analyze( # pragma: no cover
question: str = typer.Argument(
help="The question to answer",
@ -746,24 +746,24 @@ def chat( # pragma: no cover
"--model",
help="Model to use for the chat (e.g. openai-chat:gpt-4o)",
),
skill: list[str] | None = typer.Option(
capability: list[str] | None = typer.Option(
None,
"--skill",
"-s",
help="Skills to enable: rag, analysis (can repeat, default: rag)",
"--capability",
"-c",
help="Capabilities to enable: rag, analysis (can repeat, default: rag)",
),
):
"""Launch the chat TUI for conversational RAG."""
from haiku.rag.chat import run_chat
db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
skills = skill if skill else ["rag"]
capabilities = capability if capability else ["rag"]
run_chat(
db_path,
read_only=True,
model=model,
skills=skills,
capabilities=capabilities,
)
@ -803,78 +803,5 @@ def mcp(
)
@_cli.command(
"create-skill",
help="Generate a standalone skill package with an embedded or remote database",
)
def create_skill_cmd( # pragma: no cover
name: str = typer.Option(
...,
"--name",
help="Skill name (lowercase alphanumeric and hyphens)",
),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database to embed (omit for remote storage)",
),
description: str | None = typer.Option(
None,
"--description",
help="Skill description (default: standard RAG description)",
),
tools: str = typer.Option(
"all",
"--tools",
help="Comma-separated tool names, or 'all'",
),
preamble: str | None = typer.Option(
None,
"--preamble",
help="Custom preamble for the skill instructions",
),
config_file: Path | None = typer.Option(
None,
"--config-file",
help="Path to haiku.rag.yaml to embed in the skill",
),
output: Path = typer.Option(
Path("."),
"--output",
"-o",
help="Output directory for the generated package",
),
):
"""Generate a standalone haiku.skills package with an embedded database."""
from haiku.rag.skill_generator import (
AVAILABLE_TOOLS,
DEFAULT_DESCRIPTION,
generate_skill,
)
if description is None:
description = DEFAULT_DESCRIPTION
if tools.strip().lower() == "all":
tool_names = sorted(AVAILABLE_TOOLS)
else:
tool_names = [t.strip() for t in tools.split(",") if t.strip()]
try:
result = generate_skill(
db_path=db,
output_dir=output,
name=name,
description=description,
tool_names=tool_names,
config_path=config_file,
preamble=preamble,
)
typer.echo(f"Skill generated: {result}")
except ValueError as e:
typer.echo(f"Error: {e}", err=True)
raise typer.Exit(1)
if __name__ == "__main__": # pragma: no cover
cli()

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

@ -1,4 +1,8 @@
from typing import TYPE_CHECKING
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from pydantic_ai import Agent
from pydantic_ai.usage import UsageLimits
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
@ -6,12 +10,17 @@ if TYPE_CHECKING:
from haiku.rag.store.models.citation import Citation
@dataclass
class _AgentDeps:
state: dict[str, Any] = field(default_factory=dict)
async def ask(
client: "HaikuRAG",
question: str,
filter: str | None = None,
) -> "tuple[str, list[Citation]]":
"""Ask a question against the knowledge base via the rag skill.
"""Ask a question against the knowledge base via the RAG capability.
Args:
client: The HaikuRAG client.
@ -21,20 +30,36 @@ async def ask(
Returns:
Tuple of (answer text, list of resolved citations).
"""
from haiku.rag.skills.rag import RAGState, create_skill
from haiku.rag.capabilities.rag import (
AGENT_PREAMBLE,
RAGState,
create_capability,
)
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 = RAGState(document_filter=filter)
capability = create_capability(
db_path=client.store.db_path,
config=client._config,
defer_loading=False,
)
deps = _AgentDeps(
state={"rag": RAGState(document_filter=filter).model_dump(mode="json")}
)
model = get_model(client._config.qa.model, client._config)
answer, _, _ = await run_skill(model, skill, question, state=state)
agent = Agent(
model,
deps_type=_AgentDeps,
instructions=AGENT_PREAMBLE,
capabilities=[capability],
)
result = await agent.run(question, deps=deps)
state = RAGState.model_validate(deps.state["rag"])
citations = [
state.citation_index[cid]
for cid in state.citations
if cid in state.citation_index
]
return answer, citations
return result.output, citations
async def analyze(
@ -42,9 +67,9 @@ async def analyze(
question: str,
filter: str | None = None,
) -> "AnalysisResult":
"""Answer a question against the knowledge base via the rag-analysis skill.
"""Answer a question using the analysis capability.
The analysis skill exposes ``search``, ``execute_code``, and ``cite`` tools.
The capability exposes search, code execution, and citation tools.
The driving model decides when to reach for code (structural traversal,
computation, aggregation) versus a direct ``search → cite → answer``.
@ -56,20 +81,37 @@ async def analyze(
Returns:
AnalysisResult with the answer and resolved citations.
"""
from haiku.rag.capabilities.analysis import AnalysisState, create_capability
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)
capability = create_capability(
db_path=client.store.db_path,
config=client._config,
defer_loading=False,
)
deps = _AgentDeps(
state={
"analysis": AnalysisState(document_filter=filter).model_dump(mode="json")
}
)
model = get_model(
client._config.analysis.model or client._config.qa.model, client._config
)
answer, _, _ = await run_skill(model, skill, question, state=state)
agent = Agent(
model,
deps_type=_AgentDeps,
capabilities=[capability],
)
result = await agent.run(
question,
deps=deps,
usage_limits=UsageLimits(request_limit=capability.default_request_limit),
)
state = AnalysisState.model_validate(deps.state["analysis"])
citations = [
state.citation_index[cid]
for cid in state.citations
if cid in state.citation_index
]
return AnalysisResult(answer=answer, citations=citations)
return AnalysisResult(answer=result.output, citations=citations)

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

@ -99,7 +99,7 @@ class QAConfig(BaseModel):
class AnalysisConfig(BaseModel):
"""Driving model + sandbox limits for the analysis skill.
"""Driving model and sandbox limits for the analysis capability.
``model`` defaults to ``None``, meaning "no override — use ``qa.model``."
Consumers resolve via ``config.analysis.model or config.qa.model``. Set

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

@ -210,10 +210,10 @@ def create_mcp_server(
question: str,
filter: str | None = None,
) -> str:
"""Answer complex questions using the rag-analysis skill.
"""Answer complex questions using the analysis capability.
Use this for questions requiring computation, aggregation, or
structural traversal across documents. The skill can write and
structural traversal across documents. The capability can write and
execute Python code in a sandboxed interpreter.
Args:
@ -228,6 +228,6 @@ def create_mcp_server(
result = await rag.analyze(question, filter=filter)
return result.answer
except Exception as e:
return f"Error running analysis skill: {e!s}"
return f"Error running analysis capability: {e!s}"
return mcp

4
haiku_rag_slim/haiku/rag/sandbox/models.py
View file

@ -7,8 +7,8 @@ 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."""
analysis capability's code-execution tool). Consumers that need the program
should pull it from capability state."""
answer: str
citations: list[Citation] = Field(default_factory=list)

159
haiku_rag_slim/haiku/rag/skill_generator/__init__.py
View file

@ -1,159 +0,0 @@
import pathlib
import shutil
from importlib.metadata import version
from jinja2 import Environment, PackageLoader
AVAILABLE_TOOLS: set[str] = {
"list_documents",
"get_document",
"search",
"execute_code",
"cite",
}
DEFAULT_PREAMBLE = (
"You are a RAG (Retrieval Augmented Generation) assistant "
"with access to a document knowledge base.\n"
"Use your tools to search and answer questions. "
"Never make up information — always use tools to get facts "
"from the knowledge base."
)
DEFAULT_DESCRIPTION = (
"Search, retrieve and analyze documents using RAG (Retrieval Augmented Generation)."
)
def _get_env() -> Environment:
return Environment(
loader=PackageLoader("haiku.rag.skill_generator", "templates"),
autoescape=False,
keep_trailing_newline=True,
lstrip_blocks=True,
trim_blocks=True,
)
def validate_metadata(name: str, description: str) -> None:
from haiku.skills import SkillMetadata
SkillMetadata(name=name, description=description)
def validate_tools(tools: list[str]) -> None:
if not tools:
raise ValueError("tools must contain at least one tool")
unknown = set(tools) - AVAILABLE_TOOLS
if unknown:
raise ValueError(
f"Unknown tools: {', '.join(sorted(unknown))}."
f" Available: {', '.join(sorted(AVAILABLE_TOOLS))}"
)
def validate_db_path(db_path: pathlib.Path) -> None:
if not db_path.exists():
raise ValueError(f"db_path does not exist: {db_path}")
if not db_path.is_dir():
raise ValueError(f"db_path is not a directory: {db_path}")
def validate_output_dir(output_dir: pathlib.Path, name: str) -> None:
if not output_dir.exists():
raise ValueError(f"output_dir does not exist: {output_dir}")
target = output_dir / f"{name}-skill"
if target.exists():
raise ValueError(f"Target directory already exists: {target}")
def render_templates(
output_dir: pathlib.Path,
name: str,
description: str,
tool_names: list[str],
preamble: str | None = None,
remote: bool = False,
) -> pathlib.Path:
if preamble is None:
preamble = DEFAULT_PREAMBLE
pkg_name = name.replace("-", "_")
rag_version = version("haiku.rag-slim")
env = _get_env()
context = {
"name": name,
"pkg_name": pkg_name,
"description": description,
"tool_names": tool_names,
"preamble": preamble,
"rag_version": rag_version,
"remote": remote,
}
result_dir = output_dir / f"{name}-skill"
pkg_dir = result_dir / f"{pkg_name}_skill"
assets_dir = pkg_dir / "assets"
assets_dir.mkdir(parents=True)
# Render pyproject.toml
template = env.get_template("pyproject.toml.j2")
(result_dir / "pyproject.toml").write_text(template.render(context))
# Render README.md
template = env.get_template("README.md.j2")
(result_dir / "README.md").write_text(template.render(context))
# Render __init__.py
template = env.get_template("__init__.py.j2")
(pkg_dir / "__init__.py").write_text(template.render(context))
# Render SKILL.md
template = env.get_template("SKILL.md.j2")
(pkg_dir / "SKILL.md").write_text(template.render(context))
return result_dir
def generate_skill(
db_path: pathlib.Path | None,
output_dir: pathlib.Path,
name: str,
description: str,
tool_names: list[str],
config_path: pathlib.Path | None = None,
preamble: str | None = None,
) -> pathlib.Path:
validate_metadata(name, description)
validate_tools(tool_names)
if db_path is None:
if config_path is None:
raise ValueError(
"config_path is required when db_path is not provided "
"(remote storage needs connection config)"
)
else:
validate_db_path(db_path)
validate_output_dir(output_dir, name)
result = render_templates(
output_dir=output_dir,
name=name,
description=description,
tool_names=tool_names,
preamble=preamble,
remote=db_path is None,
)
pkg_name = name.replace("-", "_")
assets_dir = result / f"{pkg_name}_skill" / "assets"
if db_path is not None:
shutil.copytree(db_path, assets_dir / f"{name}.lancedb")
if config_path is not None:
shutil.copy2(config_path, assets_dir / "haiku.rag.yaml")
return result

11
haiku_rag_slim/haiku/rag/skill_generator/templates/README.md.j2
View file

@ -1,11 +0,0 @@
# {{ name }}-skill
{{ description }}
This skill package was generated by [`haiku-rag create-skill`](https://ggozad.github.io/haiku.rag/).
## Installation
```bash
pip install {{ name }}-skill/
```

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

@ -1,64 +0,0 @@
---
name: {{ name }}
description: {{ description }}
---
# {{ name }}
{{ preamble }}
## Tools
{% if "search" in tool_names %}
### search
Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content. Use for answering questions, finding passages, exploring topics. 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.
{% endif %}
{% if "list_documents" in tool_names %}
### list_documents
List all documents in the knowledge base.
{% endif %}
{% if "get_document" in tool_names %}
### get_document
Retrieve a document by ID, title, or URI. Partial matches work.
{% endif %}
{% 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()`, and a virtual filesystem at `/documents/` with document content and structure.
{% endif %}
{% if "cite" in tool_names %}
### cite
Register chunk IDs as citations. Call after formulating your answer with chunk_id values from search results that support it. Do NOT include chunk IDs in your answer text.
{% endif %}
{% if "search" in tool_names %}
## How to answer questions
1. Call `search` with relevant keywords from the question
2. Review results — they are ordered by relevance (rank 1 = best match)
3. If needed, search again with different keywords (up to 3-4 searches total)
4. Synthesize a concise answer based strictly on the retrieved content
{% if "cite" in tool_names %}
5. Call `cite` with the chunk IDs you referenced
{% endif %}
## Guidelines
- Base answers strictly on retrieved content — do not use external knowledge
- Be concise and direct — avoid elaboration unless asked
- If results don't match the question, report that the knowledge base lacks the information
{% if "cite" in tool_names %}
- Do NOT include chunk IDs or UUIDs in your answer text — use the `cite` tool separately
{% endif %}
{% endif %}
{% if "get_document" in tool_names %}
## When the user mentions a specific document
If the user says "search in [doc]", "find in [doc]", or "answer from [doc]":
- Use **get_document** or **list_documents** first to identify the document
- Then search for the topic
{% endif %}

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

@ -1,80 +0,0 @@
from pathlib import Path
from pydantic import BaseModel, Field
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.store.models.citation import Citation
{% endif %}
{% if "search" in tool_names %}
from haiku.rag.store.models.chunk import SearchResult
{% endif %}
{% if "execute_code" in tool_names %}
from haiku.rag.skills._tools import CodeExecutionEntry
{% endif %}
_TOOL_NAMES = {{ tool_names | tojson }}
_ASSETS_DIR = Path(__file__).resolve().parent / "assets"
{% if remote %}
_DB_PATH = None
{% else %}
_DB_PATH = _ASSETS_DIR / "{{ name }}.lancedb"
{% endif %}
_CONFIG_PATH = _ASSETS_DIR / "haiku.rag.yaml"
class SkillState(BaseModel):
{% if "cite" in tool_names %}
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
{% endif %}
document_filter: str | None = None
{% if "search" in tool_names %}
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
{% endif %}
{% if "execute_code" in tool_names %}
executions: list[CodeExecutionEntry] = Field(default_factory=list)
{% endif %}
def _get_config():
if _CONFIG_PATH.exists():
from haiku.rag.config import AppConfig, load_yaml_config
return AppConfig.model_validate(load_yaml_config(_CONFIG_PATH))
from haiku.rag.config import get_config
return get_config()
def create_skill(
db_path: Path | None = None,
config: AppConfig | None = None,
) -> Skill:
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
metadata, instructions = parse_skill_md(Path(__file__).parent / "SKILL.md")
if config is None:
config = _get_config()
if db_path is None:
db_path = _DB_PATH
tools = create_skill_tools(db_path, config, SkillState, _TOOL_NAMES)
extras = create_skill_extras(db_path, config)
if config.prompts.domain_preamble and instructions:
instructions = f"{config.prompts.domain_preamble}\n\n{instructions}"
return Skill(
metadata=metadata,
instructions=instructions,
tools=list(tools.values()),
extras=extras,
state_type=SkillState,
state_namespace="{{ name }}",
)

20
haiku_rag_slim/haiku/rag/skill_generator/templates/pyproject.toml.j2
View file

@ -1,20 +0,0 @@
[build-system]
requires = ["setuptools>=61.0"]
build-backend = "setuptools.build_meta"
[project]
name = "{{ name }}-skill"
version = "0.1.0"
description = "{{ description }}"
readme = "README.md"
requires-python = ">=3.12"
dependencies = [
"haiku.rag-slim >= {{ rag_version }}",
"haiku-skills >= 0.12.0",
]
[project.entry-points."haiku.skills"]
{{ name }} = "{{ pkg_name }}_skill:create_skill"
[tool.setuptools.package-data]
{{ pkg_name }}_skill = ["SKILL.md", "assets/**/*"]

0
haiku_rag_slim/haiku/rag/skills/__init__.py
View file

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

@ -1,92 +0,0 @@
import asyncio
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
from haiku.rag.config.models import AppConfig
from haiku.skills.state import SkillRunDeps
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
from haiku.rag.sandbox import Sandbox
@dataclass
class RAGRunDeps(SkillRunDeps):
rag: "HaikuRAG | None" = None
# pydantic-ai runs a turn's tool calls concurrently; LanceDB's per-connection
# state cannot take two in-flight operations at once, so every use of ``rag``
# (skill tools and the analysis sandbox) serializes through this lock. Always
# present so serialization is never accidentally skipped.
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
search_count: int = 0
@dataclass
class AnalysisRunDeps(RAGRunDeps):
sandbox: "Sandbox | None" = None
execute_count: int = 0
def _reset_invocation_state(state: Any) -> None:
"""Clear state fields scoped to a single invocation.
Keeps ``citation_index`` (accumulates resolved citations across the session
for lookup) and ``document_filter`` (session-level). Clears ``citations``,
``searches``, and (for analysis) ``executions``.
"""
if state is None:
return
citations = getattr(state, "citations", None)
if citations is not None:
citations.clear()
searches = getattr(state, "searches", None)
if searches is not None:
searches.clear()
executions = getattr(state, "executions", None)
if executions is not None:
executions.clear()
def make_rag_lifespan(db_path: Path, config: AppConfig):
@asynccontextmanager
async def lifespan(deps: RAGRunDeps) -> AsyncIterator[None]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
deps.rag = rag
deps.search_count = 0
_reset_invocation_state(deps.state)
yield
return lifespan
def make_analysis_lifespan(db_path: Path, config: AppConfig):
@asynccontextmanager
async def lifespan(deps: AnalysisRunDeps) -> AsyncIterator[None]:
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:
deps.rag = rag
deps.search_count = 0
deps.execute_count = 0
sandbox = Sandbox(
db_path=db_path,
config=config,
context=AnalysisContext(filter=doc_filter),
rag=rag,
lock=deps.rag_lock,
)
deps.sandbox = sandbox
_reset_invocation_state(deps.state)
try:
yield
finally:
sandbox.close()
return lifespan

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

@ -1,406 +0,0 @@
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from pydantic_ai import ModelRetry, RunContext
from pydantic_ai.messages import ToolReturn
from haiku.rag.client import HaikuRAG
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
class CodeExecutionEntry(BaseModel):
code: str
stdout: str
stderr: str = ""
success: bool = True
async def skill_search(
rag: HaikuRAG,
query: str,
limit: int | None = None,
document_filter: str | None = None,
) -> tuple[str, list[SearchResult]]:
results = await rag.search(query, limit=limit, filter=document_filter)
results = await rag.expand_context(results)
formatted = "\n\n---\n\n".join(
r.format_for_agent(rank=i + 1, total=len(results))
for i, r in enumerate(results)
)
return formatted, list(results)
async def skill_list_documents(
rag: HaikuRAG,
filter: str | None = None,
) -> list[dict[str, Any]]:
documents = await rag.list_documents(filter=filter)
return [
{
"id": doc.id,
"title": doc.title,
"uri": doc.uri,
"metadata": doc.metadata,
"created_at": str(doc.created_at),
"updated_at": str(doc.updated_at),
}
for doc in documents
]
async def skill_get_document(
rag: HaikuRAG,
query: str,
) -> dict[str, Any] | None:
document = await rag.resolve_document(query)
if document is None:
return None
return {
"id": document.id,
"content": document.content,
"title": document.title,
"uri": document.uri,
"metadata": document.metadata,
"created_at": str(document.created_at),
"updated_at": str(document.updated_at),
}
def _get_state(ctx: RunContext[RAGRunDeps], state_type: type[BaseModel]) -> Any:
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, state_type):
return ctx.deps.state
return None
def _require_rag(ctx: RunContext[RAGRunDeps]) -> HaikuRAG:
assert ctx.deps is not None and ctx.deps.rag is not None, (
"RAGRunDeps.rag is not set — skill lifespan must run before tools."
)
return ctx.deps.rag
@asynccontextmanager
async def _serialized(ctx: RunContext[RAGRunDeps]) -> AsyncIterator[None]:
"""Serialize access to the shared connection through the run's lock.
pydantic-ai runs a turn's tool calls concurrently and LanceDB's
per-connection state cannot take two in-flight operations at once. The lock
is always present (``RAGRunDeps`` creates one by default); the no-op branch
is a guard for a missing deps/lock.
"""
lock = ctx.deps.rag_lock if ctx.deps is not None else None
if lock is None:
yield
else:
async with lock:
yield
def _register_citations(state: Any, citations: "list[Citation]") -> None:
"""Add citations to the index and record cited chunk IDs for this invocation."""
next_index = len(state.citation_index) + 1
for citation in citations:
cid = citation.chunk_id
if cid not in state.citation_index:
citation.index = next_index
next_index += 1
state.citation_index[cid] = citation
if cid not in state.citations:
state.citations.append(cid)
def create_skill_extras(
db_path: Path,
config: AppConfig,
) -> dict[str, Any]:
"""Create non-tool utility functions bound to a specific database.
Returns a dict of values that can be attached to a Skill's extras:
Keys:
- 'db_path': path to the LanceDB used to configure the skill
- 'config': config passed to (or derived for) the skill
- 'list_documents': returns info for documents in the database
- 'visualize_chunk': returns visualizations for chunks in the database
"""
async def visualize_chunk(chunk_id: str | list[str]) -> list:
from haiku.rag.client import HaikuRAG
chunk_ids = [chunk_id] if isinstance(chunk_id, str) else chunk_id
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
chunks = []
for cid in chunk_ids:
chunk = await rag.get_chunk_by_id(cid)
if chunk is not None:
chunks.append(chunk)
if not chunks:
return []
return await rag.visualize_chunk(chunks)
async def list_documents(
limit: int | None = None,
offset: int | None = None,
filter: str | None = None,
) -> list[dict[str, Any]]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = await rag.list_documents(limit, offset, filter=filter)
return [
{
"id": doc.id,
"title": doc.title,
"uri": doc.uri,
"metadata": doc.metadata,
"created_at": str(doc.created_at),
"updated_at": str(doc.updated_at),
}
for doc in documents
]
return {
"db_path": db_path,
"config": config,
"visualize_chunk": visualize_chunk,
"list_documents": list_documents,
}
def create_skill_tools(
db_path: Path,
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. ``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] = {}
if "search" in tool_names:
max_searches = config.qa.max_searches
async def search(
ctx: RunContext[RAGRunDeps], query: str, limit: int | None = None
) -> str | ToolReturn:
"""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 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.
limit: Maximum number of results.
"""
ctx.deps.search_count += 1
if ctx.deps.search_count > max_searches:
return (
"Search limit reached. Answer the question using "
"the results you already have."
)
state = _get_state(ctx, state_type)
async with _serialized(ctx):
formatted, results = await skill_search(
_require_rag(ctx),
query,
limit=limit,
document_filter=state.document_filter if state else None,
)
if state:
state.searches[query] = results
if not model.vision:
return formatted
binary_parts = build_binary_parts_from_results(results)
if binary_parts:
return ToolReturn(return_value=formatted, content=binary_parts)
return formatted
tools["search"] = search
if "list_documents" in tool_names:
async def list_documents(
ctx: RunContext[RAGRunDeps],
) -> list[dict[str, Any]]:
"""List all documents in the knowledge base."""
state = _get_state(ctx, state_type)
async with _serialized(ctx):
return await skill_list_documents(
_require_rag(ctx),
filter=state.document_filter if state else None,
)
tools["list_documents"] = list_documents
if "get_document" in tool_names:
async def get_document(
ctx: RunContext[RAGRunDeps], query: str
) -> dict[str, Any] | None:
"""Retrieve a document by ID, title, or URI.
Args:
query: Document ID, title, or URI to look up.
"""
async with _serialized(ctx):
return await skill_get_document(_require_rag(ctx), query)
tools["get_document"] = get_document
if "execute_code" in tool_names:
max_executions = config.analysis.max_executions
async def execute_code(ctx: RunContext[AnalysisRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter.
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.
Args:
code: Python code to execute.
"""
ctx.deps.execute_count += 1
if ctx.deps.execute_count > max_executions:
return (
"Code-execution limit reached. Give your final answer now "
"from what you already have; do not call execute_code again."
)
assert ctx.deps is not None and ctx.deps.sandbox is not None, (
"AnalysisRunDeps.sandbox is not set — skill lifespan must run before execute_code."
)
sandbox = ctx.deps.sandbox
result = await sandbox.execute(code)
state = _get_state(ctx, state_type)
if state and sandbox._search_results:
existing = state.searches.get("_sandbox", [])
seen = {r.chunk_id for r in existing}
for sr in sandbox._search_results:
if sr.chunk_id not in seen:
existing.append(sr)
seen.add(sr.chunk_id)
state.searches["_sandbox"] = existing
if state:
state.executions.append(
CodeExecutionEntry(
code=code,
stdout=result.stdout,
stderr=result.stderr,
success=result.success,
)
)
if result.success:
return result.stdout if result.stdout else "No output."
return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
tools["execute_code"] = execute_code
if "cite" in tool_names:
async def cite(ctx: RunContext[RAGRunDeps], chunk_ids: list[str]) -> str:
"""Register chunk IDs as citations for 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 or VFS reads.
"""
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."
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:
async with _serialized(ctx):
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
chunk.document_meta = doc.metadata if doc else {}
synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
if synthetic:
citations.extend(resolve_citations(missing, synthetic))
if citations:
_register_citations(state, citations)
resolved_ids = {c.chunk_id for c in citations}
unresolved = [cid for cid in missing if cid not in resolved_ids]
if unresolved:
return (
f"Registered {len(citations)} citation(s); "
f"ignored {len(unresolved)} unresolvable id(s): "
f"{unresolved}. Copy chunk_ids verbatim from `search` "
"results or items.jsonl / toc.json rows and cite again."
)
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
return tools

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

@ -1,101 +0,0 @@
import os
from functools import cache
from pathlib import Path
from pydantic import BaseModel, Field
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
class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
STATE_TYPE = AnalysisState
STATE_NAMESPACE = "analysis"
_skill_path = Path(__file__).parent / "rag-analysis"
@cache
def skill_metadata() -> SkillMetadata:
metadata, _ = parse_skill_md(_skill_path / "SKILL.md")
return metadata
@cache
def instructions() -> str | None:
_, instr = parse_skill_md(_skill_path / "SKILL.md")
return instr
def state_metadata() -> StateMetadata:
return StateMetadata(
namespace=STATE_NAMESPACE,
type=STATE_TYPE,
schema=STATE_TYPE.model_json_schema(),
)
def create_skill(
db_path: Path | None = None,
config: AppConfig | None = None,
) -> Skill:
"""Create an analysis skill for computational document analysis.
Args:
db_path: Path to the LanceDB database. Resolved from:
1. This argument
2. HAIKU_RAG_DB environment variable
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
if config is None:
config = get_config()
if db_path is None:
env_db = os.environ.get("HAIKU_RAG_DB")
if env_db:
db_path = Path(env_db).expanduser()
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
tools = create_skill_tools(
db_path,
config,
AnalysisState,
["search", "execute_code", "cite"],
model=config.analysis.model or config.qa.model,
)
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()
if config.prompts.domain_preamble and skill_instructions:
skill_instructions = f"{config.prompts.domain_preamble}\n\n{skill_instructions}"
return Skill(
metadata=skill_metadata(),
source=SkillSource.ENTRYPOINT,
path=_skill_path,
instructions=skill_instructions,
tools=list(tools.values()),
extras=extras,
state_type=STATE_TYPE,
state_namespace=STATE_NAMESPACE,
deps_type=AnalysisRunDeps,
lifespan=make_analysis_lifespan(db_path, config),
request_limit=30,
)

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

@ -1,111 +0,0 @@
import os
from functools import cache
from pathlib import Path
from pydantic import BaseModel, Field
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
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. NEVER make up information - always use skills to get facts from the knowledge base
3. When a skill returns citations, always include them in your response
"""
_RAG_TOOLS = ["search", "cite"]
def get_agent_preamble(config: AppConfig) -> str:
"""Build the main agent preamble, prepending domain_preamble if configured."""
if config.prompts.domain_preamble:
return f"{config.prompts.domain_preamble}\n\n{AGENT_PREAMBLE}"
return AGENT_PREAMBLE
class RAGState(BaseModel):
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[str] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
STATE_TYPE = RAGState
STATE_NAMESPACE = "rag"
_skill_path = Path(__file__).parent / "rag"
@cache
def skill_metadata() -> SkillMetadata:
metadata, _ = parse_skill_md(_skill_path / "SKILL.md")
return metadata
@cache
def instructions() -> str | None:
_, instr = parse_skill_md(_skill_path / "SKILL.md")
return instr
def state_metadata() -> StateMetadata:
return StateMetadata(
namespace=STATE_NAMESPACE,
type=STATE_TYPE,
schema=STATE_TYPE.model_json_schema(),
)
def create_skill(
db_path: Path | None = None,
config: AppConfig | None = None,
) -> Skill:
"""Create a RAG skill for searching and analyzing documents.
Args:
db_path: Path to the LanceDB database. Resolved from:
1. This argument
2. HAIKU_RAG_DB environment variable
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
if config is None:
config = get_config()
if db_path is None:
env_db = os.environ.get("HAIKU_RAG_DB")
if env_db:
db_path = Path(env_db).expanduser()
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
tools = create_skill_tools(
db_path, config, RAGState, _RAG_TOOLS, model=config.qa.model
)
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()
if config.prompts.domain_preamble and skill_instructions:
skill_instructions = f"{config.prompts.domain_preamble}\n\n{skill_instructions}"
return Skill(
metadata=skill_metadata(),
source=SkillSource.ENTRYPOINT,
path=_skill_path,
instructions=skill_instructions,
tools=list(tools.values()),
extras=extras,
state_type=STATE_TYPE,
state_namespace=STATE_NAMESPACE,
deps_type=RAGRunDeps,
lifespan=make_rag_lifespan(db_path, config),
)

2
haiku_rag_slim/haiku/rag/store/models/citation.py
View file

@ -11,7 +11,7 @@ if TYPE_CHECKING:
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
Used by the RAG and analysis capabilities 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

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

@ -134,9 +134,8 @@ def get_model(
# Apply thinking control only for reasoning models (o-series, gpt-5)
profile = cast(OpenAIModelProfile, openai_model_profile(model))
if (
model_config.enable_thinking is not None
and profile.openai_supports_encrypted_reasoning_content
if model_config.enable_thinking is not None and profile.get(
"openai_supports_reasoning", False
):
if model_config.enable_thinking is False:
openai_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")

10
haiku_rag_slim/pyproject.toml
View file

@ -23,15 +23,13 @@ classifiers = [
dependencies = [
"docling-core>=2.82.0,<3.0.0",
"haiku.skills>=0.18.0",
"httpx>=0.28.1",
"jinja2>=3.1.0",
"jsonpatch>=1.33",
"fastmcp>=3.3.0",
"lancedb==0.34.0",
"pathspec>=1.0.4",
"pydantic>=2.12.5",
"pydantic-ai-slim[openai,logfire,ag-ui]>=1.100.0",
"pydantic-ai-slim[openai,logfire,ag-ui]>=2.11.0,<3.0.0",
"pydantic-monty>=0.0.17",
"pypdfium2>=5.0",
"python-dotenv>=1.2.2",
@ -76,11 +74,7 @@ groq = ["pydantic-ai-slim[groq]"]
google = ["pydantic-ai-slim[google]"]
mistral = ["pydantic-ai-slim[mistral]"]
bedrock = ["pydantic-ai-slim[bedrock]"]
vertexai = ["pydantic-ai-slim[vertexai]"]
[project.entry-points."haiku.skills"]
rag = "haiku.rag.skills.rag:create_skill"
rag-analysis = "haiku.rag.skills.analysis:create_skill"
vertexai = ["pydantic-ai-slim[google]"]
[project.scripts]
haiku-rag = "haiku.rag.cli:cli"

28
tests/skills/conftest.py → tests/capabilities/conftest.py
View file

@ -1,13 +1,11 @@
import random
from unittest.mock import MagicMock, patch
from unittest.mock import patch
import pytest
from pydantic_ai import RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.skills._deps import AnalysisRunDeps, RAGRunDeps
VECTOR_DIM = 2560
@ -25,26 +23,6 @@ async def _fake_embed_documents(self, texts: list[str]) -> list[list[float]]:
return [_seeded_vector(t) for t in texts]
def _make_ctx(state=None, rag=None, sandbox=None):
"""Create a mock RunContext with RAGRunDeps (or AnalysisRunDeps when state is AnalysisState)."""
from haiku.rag.skills.analysis import AnalysisState
ctx = MagicMock(spec=RunContext)
if isinstance(state, AnalysisState) or sandbox is not None:
ctx.deps = AnalysisRunDeps(state=state, rag=rag, sandbox=sandbox)
else:
ctx.deps = RAGRunDeps(state=state, rag=rag)
return ctx
def _get_tool(skill, name):
"""Get a tool function from a skill by name."""
for tool in skill.tools:
if callable(tool) and tool.__name__ == name:
return tool
raise ValueError(f"Tool {name!r} not found in skill")
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch):
"""Monkeypatch the embedder to return deterministic vectors."""
@ -54,7 +32,7 @@ def mock_embedder(monkeypatch):
@pytest.fixture
def test_app_config():
return AppConfig(environment="skills-test")
return AppConfig(environment="capabilities-test")
@pytest.fixture(scope="session")
@ -66,7 +44,7 @@ async def rag_db(tmp_path_factory):
vectors use the same seeded fakes as ``mock_embedder`` so search stays
consistent with query-time embeddings.
"""
db_path = tmp_path_factory.mktemp("skills_rag_db") / "rag.lancedb"
db_path = tmp_path_factory.mktemp("capabilities_rag_db") / "rag.lancedb"
with (
patch.object(EmbedderWrapper, "embed_query", _fake_embed_query),
patch.object(EmbedderWrapper, "embed_documents", _fake_embed_documents),

232
tests/capabilities/test_capabilities.py Normal file
View file

@ -0,0 +1,232 @@
from dataclasses import dataclass, field
from typing import Any
import pytest
from pydantic_ai import Agent, RunContext
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
TextPart,
ToolCallPart,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from haiku.rag.capabilities._base import _compact_old_tool_returns
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
from haiku.rag.capabilities.analysis import create_capability as create_analysis
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig, PromptsConfig
@dataclass
class Deps:
state: dict[str, Any] = field(default_factory=dict)
def make_context(deps: Deps) -> RunContext[Deps]:
return RunContext(
deps=deps,
model=TestModel(),
usage=RunUsage(),
run_id="test-run",
)
def test_rag_capability_api(temp_db_path):
capability = create_rag(db_path=temp_db_path, config=AppConfig())
assert isinstance(capability, RAGCapability)
assert capability.id == "haiku-rag"
assert capability.defer_loading is True
assert set(capability.get_toolset().tools) == {"rag_search", "rag_cite"}
assert capability.state_type is RAGState
assert capability.state_namespace == "rag"
def test_analysis_capability_api(temp_db_path):
capability = create_analysis(db_path=temp_db_path, config=AppConfig())
assert isinstance(capability, AnalysisCapability)
assert capability.id == "haiku-rag-analysis"
assert capability.defer_loading is True
assert set(capability.get_toolset().tools) == {
"analysis_search",
"analysis_execute_code",
"analysis_cite",
}
assert capability.state_type is AnalysisState
def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
config = AppConfig(
prompts=PromptsConfig(domain_preamble="The corpus contains solar manuals.")
)
capability = create_rag(db_path=temp_db_path, config=config)
assert capability.get_instructions().startswith(
"The corpus contains solar manuals.\n\n# RAG"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("factory", "agent_instructions", "heading"),
[
(create_rag, AGENT_PREAMBLE, "# RAG"),
(create_analysis, None, "# Analysis"),
],
)
async def test_capability_instructions_are_injected_once(
temp_db_path, factory, agent_instructions, heading
):
domain = "The corpus contains solar manuals."
seen_instructions = []
def model_function(_messages, info):
seen_instructions.append(info.instructions or "")
return ModelResponse(parts=[TextPart("done")])
config = AppConfig(prompts=PromptsConfig(domain_preamble=domain))
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
instructions=agent_instructions,
capabilities=[
factory(
db_path=temp_db_path,
config=config,
defer_loading=False,
)
],
)
await agent.run("Answer", deps=Deps())
assert seen_instructions[0].count(domain) == 1
assert seen_instructions[0].count(heading) == 1
@pytest.mark.asyncio
async def test_capability_isolated_per_run_and_round_trips_state(temp_db_path):
capability = create_rag(db_path=temp_db_path, config=AppConfig())
deps = Deps(
state={
"rag": RAGState(
document_filter="uri = 'manual.pdf'",
citations=["old"],
searches={"old": []},
).model_dump(mode="json")
}
)
run_capability = await capability.for_run(make_context(deps))
assert run_capability is not capability
assert run_capability.state is not None
assert run_capability.state.document_filter == "uri = 'manual.pdf'"
assert run_capability.state.citations == []
assert run_capability.state.searches == {}
assert deps.state["rag"]["document_filter"] == "uri = 'manual.pdf'"
@pytest.mark.asyncio
async def test_native_agent_composition_initializes_host_state(temp_db_path):
capability = create_rag(
db_path=temp_db_path,
config=AppConfig(),
defer_loading=False,
)
deps = Deps()
agent = Agent(
TestModel(call_tools=[]),
deps_type=Deps,
capabilities=[capability],
)
result = await agent.run("Hello", deps=deps)
assert result.output == "success (no tool calls)"
assert deps.state["rag"] == RAGState().model_dump(mode="json")
@pytest.mark.asyncio
async def test_deferred_capability_loads_native_tools(temp_db_path):
seen_instructions = []
loaded_payloads = []
def model_function(messages, info):
seen_instructions.append(info.instructions or "")
loaded_payloads.extend(
str(part.content)
for message in messages
for part in message.parts
if isinstance(part, ToolReturnPart) and part.tool_name == "load_capability"
)
loaded = any(
isinstance(part, ToolReturnPart) and part.tool_name == "load_capability"
for message in messages
for part in message.parts
)
if not loaded:
return ModelResponse(
parts=[ToolCallPart("load_capability", {"id": "haiku-rag"})]
)
return ModelResponse(parts=[TextPart("loaded")])
agent = Agent(
FunctionModel(model_function),
deps_type=Deps,
capabilities=[create_rag(db_path=temp_db_path, config=AppConfig())],
)
result = await agent.run("Use RAG", deps=Deps())
assert result.output == "loaded"
assert "# RAG" not in seen_instructions[0]
assert "# RAG" in loaded_payloads[0]
assert "rag_search" in loaded_payloads[0]
def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept():
messages = [
ModelRequest(parts=[UserPromptPart("old question")]),
ModelResponse(parts=[ToolCallPart("rag_search", {}, "old-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "large old evidence", "old-call")]
),
ModelRequest(parts=[UserPromptPart("current question")]),
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
),
]
compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
old_return = compacted[2].parts[0]
current_return = compacted[5].parts[0]
assert isinstance(old_return, ToolReturnPart)
assert "removed" in str(old_return.content)
assert isinstance(current_return, ToolReturnPart)
assert current_return.content == "current evidence"
def test_tool_results_are_unchanged_when_history_has_no_user_prompt():
messages = [
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
),
]
compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
assert compacted is messages
current_return = compacted[1].parts[0]
assert isinstance(current_return, ToolReturnPart)
assert current_return.content == "current evidence"

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

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_client/test_client_create_document_from_directory_with_filters.yaml
View file

File diff suppressed because one or more lines are too long

5460
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_aggregation.yaml
View file

File diff suppressed because one or more lines are too long

3130
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_count_documents.yaml
View file

File diff suppressed because one or more lines are too long

28243
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_search_and_extract.yaml
View file

File diff suppressed because one or more lines are too long

1905
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_search_and_identify_source.yaml
View file

File diff suppressed because one or more lines are too long

60229
tests/cassettes/test_client_analyze/TestClientAnalysisIntegration.test_analyze_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

92
tests/cassettes/test_context_enhancement/test_max_items_limit_caps_expansion.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_search_tools/TestSearchToolExecution.test_search_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

39
tests/chat/test_chat_app.py
View file

@ -1,11 +1,11 @@
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock, patch
import pytest
from typer.testing import CliRunner
from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.cli import _cli as cli
from haiku.rag.skills.rag import RAGState
runner = CliRunner()
@ -45,42 +45,23 @@ def _make_app(db_path: Path, mock_client: AsyncMock | None = None):
if mock_client is None:
mock_client = _make_mock_client()
skill = MagicMock()
skill.state_type = None
skill.state_namespace = None
skill.tools = []
skill.toolsets = []
skill.resources = []
skill.metadata = MagicMock()
skill.metadata.name = "rag"
skill.metadata.description = "RAG skill"
return ChatApp(
db_path=db_path,
skills=[skill],
capabilities=[create_capability(db_path=db_path)],
read_only=True,
), mock_client
def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None):
"""Create a ChatApp with a skill that has RAGState."""
"""Create a ChatApp with a RAG capability and state."""
from haiku.rag.chat.app import ChatApp
from haiku.skills.models import Skill, SkillMetadata, SkillSource
if mock_client is None:
mock_client = _make_mock_client()
skill = Skill(
metadata=SkillMetadata(name="rag", description="RAG skill"),
source=SkillSource.ENTRYPOINT,
tools=[],
state_type=RAGState,
state_namespace="rag",
)
return ChatApp(
db_path=db_path,
skills=[skill],
capabilities=[create_capability(db_path=db_path)],
read_only=True,
), mock_client
@ -289,8 +270,7 @@ async def test_show_citations_renders_from_flat_state(temp_db_path: Path):
with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
async with app.run_test() as pilot:
rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE)
assert isinstance(rag_state, RAGState)
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
citation = Citation(
index=1,
@ -303,6 +283,7 @@ async def test_show_citations_renders_from_flat_state(temp_db_path: Path):
)
rag_state.citation_index["chunk1"] = citation
rag_state.citations.append("chunk1")
app._state[RAG_STATE_NAMESPACE] = rag_state.model_dump(mode="json")
chat_history = app.query_one(ChatHistory)
await app._show_citations_and_programs(chat_history)
@ -331,8 +312,7 @@ async def test_document_filter_updates_rag_state(temp_db_path: Path):
)
# RAGState.document_filter should be set
rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE)
assert rag_state is not None
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
expected_filter = build_multi_document_filter(selected)
assert rag_state.document_filter == expected_filter
@ -354,12 +334,13 @@ async def test_document_filter_cleared_when_empty(temp_db_path: Path):
app.on_document_filter_modal_filter_changed(
DocumentFilterModal.FilterChanged(["AI Overview"])
)
rag_state = app._toolset.get_namespace(RAG_STATE_NAMESPACE)
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
assert rag_state.document_filter is not None
# Then clear it
app.on_document_filter_modal_filter_changed(
DocumentFilterModal.FilterChanged([])
)
rag_state = RAGState.model_validate(app._state[RAG_STATE_NAMESPACE])
assert rag_state.document_filter is None
assert app._state["rag"]["document_filter"] is None

2
tests/sandbox/test_sandbox.py
View file

@ -419,7 +419,7 @@ class TestSandboxVFS:
class TestSandboxHeldConnection:
"""VFS reads must work while another connection to the same DB stays open.
Mirrors the analysis-skill lifespan, which keeps a read-only connection open
Mirrors the analysis-capability lifespan, which keeps a read-only connection open
for the whole turn while sandboxed code reads the document VFS.
"""

0
tests/skills/__init__.py
View file

406
tests/skills/test_analysis.py
View file

@ -1,406 +0,0 @@
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.analysis import (
STATE_NAMESPACE,
STATE_TYPE,
AnalysisState,
instructions,
skill_metadata,
state_metadata,
)
from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
class TestAnalysisModuleAPI:
def test_state_type_is_analysis_state(self):
assert STATE_TYPE is AnalysisState
def test_state_namespace(self):
assert STATE_NAMESPACE == "analysis"
def test_state_metadata_returns_state_metadata(self):
result = state_metadata()
assert isinstance(result, StateMetadata)
assert result.namespace == "analysis"
assert result.type is AnalysisState
assert result.schema == AnalysisState.model_json_schema()
def test_skill_metadata_returns_skill_metadata(self):
result = skill_metadata()
assert isinstance(result, SkillMetadata)
assert result.name == "rag-analysis"
def test_instructions_returns_string(self):
result = instructions()
assert isinstance(result, str)
assert len(result) > 0
def test_constants_match_create_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.state_type is STATE_TYPE
assert skill.state_namespace == STATE_NAMESPACE
assert skill.metadata == skill_metadata()
assert skill.instructions == instructions()
class TestAnalysisSkillCreation:
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag-analysis"
assert skill.metadata.description
assert skill.instructions
def test_create_skill_sets_request_limit_backstop(
self, test_app_config, temp_db_path
):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.request_limit == 30
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {"search", "execute_code", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import AnalysisState, create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is AnalysisState
assert skill._state_namespace == "analysis"
def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.extras["config"] is test_app_config
assert skill.extras["db_path"] is temp_db_path
assert "visualize_chunk" in skill.extras
assert "list_documents" in skill.extras
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.analysis import create_skill
skill = create_skill()
assert skill.metadata.name == "rag-analysis"
class TestDomainPreambleInAnalysisSkillInstructions:
def test_create_skill_without_domain_preamble(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import create_skill, instructions
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.instructions == instructions()
def test_create_skill_with_domain_preamble(self, temp_db_path):
from haiku.rag.config.models import PromptsConfig
from haiku.rag.skills.analysis import create_skill, instructions
config = AppConfig(
prompts=PromptsConfig(
domain_preamble="This knowledge base contains Helios solar panel documentation."
)
)
skill = create_skill(config=config, db_path=temp_db_path)
assert skill.instructions is not None
assert skill.instructions.startswith(
"This knowledge base contains Helios solar panel documentation."
)
base_instructions = instructions()
assert base_instructions is not None
assert base_instructions in skill.instructions
class TestExecuteCodeTool:
async def test_execute_code_returns_output(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(ctx, code="print('hello')")
assert "hello" in result
async def test_execute_code_updates_state(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="print('hello')")
assert len(state.executions) == 1
assert state.executions[0].code == "print('hello')"
assert state.executions[0].success is True
assert "hello" in state.executions[0].stdout
async def test_execute_code_reports_errors(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(ctx, code="x = 1/0")
assert "Error" in result
assert "ZeroDivisionError" in result
assert state.executions[0].success is False
async def test_execute_code_rate_limited(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
config = AppConfig()
config.analysis.max_executions = 2
skill = create_skill(db_path=rag_db, config=config)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="print('first')")
await execute_code(ctx, code="print('second')")
result = await execute_code(ctx, code="print('third')")
assert "limit reached" in result.lower()
assert ctx.deps.execute_count == 3
assert len(state.executions) == 2
async def test_execute_code_applies_document_filter(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state, sandbox=sandbox_factory(filter=state.document_filter))
result = await execute_code(
ctx, code="docs = await list_documents()\nprint(len(docs))"
)
assert "1" in result
async def test_execute_code_accumulates_search_results(
self, rag_db, sandbox_factory
):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(
ctx, code="results = await search('intelligence')\nprint(len(results))"
)
assert "_sandbox" in state.searches
assert len(state.searches["_sandbox"]) > 0
async def test_cite_picture_chunk_records_picture_refs(self, rag_db):
"""Citing a picture chunk populates ``Citation.picture_refs`` from the
docling ``#/pictures/...`` self_ref prefix.
``labels`` is a deduplicated set after expand_context and is not
aligned with ``doc_item_refs``, so the prefix is the only reliable
signal.
"""
from haiku.rag.skills.analysis import AnalysisState, create_skill
from haiku.rag.store.models.chunk import SearchResult
expanded_picture_hit = SearchResult(
chunk_id="picture-chunk-abc",
content="picture + surrounding prose",
document_id="doc-1",
document_uri="test://doc-1",
document_title="Doc 1",
score=1.0,
page_numbers=[1, 2],
headings=None,
doc_item_refs=[
"#/texts/3",
"#/pictures/0",
"#/texts/4",
"#/pictures/1",
],
labels=["caption", "picture", "text"],
)
text_result = SearchResult(
chunk_id="text-chunk-xyz",
content="prose",
document_id="doc-1",
document_uri="test://doc-1",
document_title="Doc 1",
score=0.7,
page_numbers=[1],
headings=None,
doc_item_refs=["#/texts/4"],
labels=["text"],
)
skill = create_skill(db_path=rag_db)
cite = _get_tool(skill, "cite")
state = AnalysisState()
state.searches["q1"] = [expanded_picture_hit, text_result]
ctx = _make_ctx(state)
await cite(ctx, chunk_ids=["picture-chunk-abc", "text-chunk-xyz"])
assert state.citations == ["picture-chunk-abc", "text-chunk-xyz"]
picture_citation = state.citation_index["picture-chunk-abc"]
assert picture_citation.picture_refs == ["#/pictures/0", "#/pictures/1"]
text_citation = state.citation_index["text-chunk-xyz"]
assert text_citation.picture_refs == []
async def test_execute_code_vfs_write_denied(self, rag_db, sandbox_factory):
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
result = await execute_code(
ctx,
code=(
"from pathlib import Path\n"
"import json\n"
"dirs = list(Path('/documents').iterdir())\n"
"p = dirs[0] / 'content.txt'\n"
"p.write_text('hacked')"
),
)
assert "Error" in result
assert "read-only" in result
async def test_execute_code_variables_persist_within_invocation(
self, rag_db, sandbox_factory
):
"""Same sandbox across two calls → vars persist (one skill invocation)."""
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state, sandbox=sandbox_factory())
await execute_code(ctx, code="x = 42")
result = await execute_code(ctx, code="print(x * 2)")
assert "84" in result
async def test_execute_code_isolated_across_invocations(
self, rag_db, sandbox_factory
):
"""Different Sandbox instances → no cross-invocation leak."""
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
execute_code = _get_tool(skill, "execute_code")
ctx1 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
await execute_code(ctx1, code="secret = 'do not leak'")
ctx2 = _make_ctx(AnalysisState(), sandbox=sandbox_factory())
result = await execute_code(ctx2, code="print(secret)")
assert not result.startswith("do not leak")
assert "Error" in result or "NameError" in result
class TestAnalysisLifespan:
async def test_opens_client_and_sandbox_per_invocation(self, rag_db):
from haiku.rag.sandbox import Sandbox
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
deps = AnalysisRunDeps()
async with lifespan(deps):
assert deps.rag is not None
assert deps.rag.is_read_only
assert deps.search_count == 0
assert deps.execute_count == 0
assert isinstance(deps.sandbox, Sandbox)
docs = await deps.rag.list_documents()
assert len(docs) == 2
async def test_lifespan_reads_document_filter_from_state(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
from haiku.rag.skills.analysis import AnalysisState
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
state = AnalysisState(document_filter="title = 'AI Overview'")
deps = AnalysisRunDeps(state=state)
async with lifespan(deps):
assert deps.sandbox is not None
assert deps.sandbox._context.filter == "title = 'AI Overview'"
async def test_lifespan_resets_counts_per_invocation(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
deps = AnalysisRunDeps(search_count=7, execute_count=42)
async with lifespan(deps):
assert deps.search_count == 0
assert deps.execute_count == 0
async def test_skill_has_lifespan_and_deps_type(
self, test_app_config, temp_db_path
):
from haiku.rag.skills._deps import AnalysisRunDeps
from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.deps_type is AnalysisRunDeps
assert skill.lifespan is not None
async def test_run_skill_end_to_end_opens_client_and_sandbox(
self, allow_model_requests, rag_db
):
"""Full sub-agent path: lifespan builds client + sandbox, tools see them."""
from pydantic_ai.models.test import TestModel
from haiku.rag.skills.analysis import create_skill
from haiku.skills.agent import run_skill
skill = create_skill(db_path=rag_db)
result, *_ = await run_skill(TestModel(), skill, "Print the document count.")
assert result
async def test_lifespan_clears_executions_citations_searches(self, rag_db):
from haiku.rag.skills._deps import AnalysisRunDeps, make_analysis_lifespan
from haiku.rag.skills._tools import CodeExecutionEntry
from haiku.rag.skills.analysis import AnalysisState
from haiku.rag.store.models.citation import Citation
config = AppConfig()
lifespan = make_analysis_lifespan(rag_db, config)
state = AnalysisState(
document_filter="title = 'AI Overview'",
executions=[CodeExecutionEntry(code="prior", stdout="", success=True)],
citation_index={
"c1": Citation(
index=1,
chunk_id="c1",
document_id="d1",
document_title="t",
document_uri="u",
content="x",
page_numbers=[],
headings=[],
)
},
citations=["c1"],
searches={"prior": []},
)
deps = AnalysisRunDeps(state=state)
async with lifespan(deps):
assert state.executions == []
assert state.citations == []
assert state.searches == {}
assert "c1" in state.citation_index
assert state.document_filter == "title = 'AI Overview'"

513
tests/skills/test_rag.py
View file

@ -1,513 +0,0 @@
import asyncio
import pytest
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rag import (
STATE_NAMESPACE,
STATE_TYPE,
RAGState,
instructions,
skill_metadata,
state_metadata,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
class TestRAGModuleAPI:
def test_state_type_is_rag_state(self):
assert STATE_TYPE is RAGState
def test_state_namespace(self):
assert STATE_NAMESPACE == "rag"
def test_state_metadata_returns_state_metadata(self):
result = state_metadata()
assert isinstance(result, StateMetadata)
assert result.namespace == "rag"
assert result.type is RAGState
assert result.schema == RAGState.model_json_schema()
def test_skill_metadata_returns_skill_metadata(self):
result = skill_metadata()
assert isinstance(result, SkillMetadata)
assert result.name == "rag"
def test_instructions_returns_string(self):
result = instructions()
assert isinstance(result, str)
assert len(result) > 0
def test_constants_match_create_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.state_type is STATE_TYPE
assert skill.state_namespace == STATE_NAMESPACE
assert skill.metadata == skill_metadata()
assert skill.instructions == instructions()
class TestGetAgentPreamble:
def test_without_domain_preamble(self):
from haiku.rag.skills.rag import AGENT_PREAMBLE, get_agent_preamble
config = AppConfig()
assert get_agent_preamble(config) == AGENT_PREAMBLE
def test_with_domain_preamble(self):
from haiku.rag.config.models import PromptsConfig
from haiku.rag.skills.rag import AGENT_PREAMBLE, get_agent_preamble
config = AppConfig(
prompts=PromptsConfig(
domain_preamble="This knowledge base contains Helios solar panel documentation."
)
)
result = get_agent_preamble(config)
assert result.startswith(
"This knowledge base contains Helios solar panel documentation."
)
assert AGENT_PREAMBLE in result
class TestDomainPreambleInSkillInstructions:
def test_create_skill_without_domain_preamble(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill, instructions
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.instructions == instructions()
def test_create_skill_with_domain_preamble(self, temp_db_path):
from haiku.rag.config.models import PromptsConfig
from haiku.rag.skills.rag import create_skill, instructions
config = AppConfig(
prompts=PromptsConfig(
domain_preamble="This knowledge base contains Helios solar panel documentation."
)
)
skill = create_skill(config=config, db_path=temp_db_path)
assert skill.instructions is not None
assert skill.instructions.startswith(
"This knowledge base contains Helios solar panel documentation."
)
base_instructions = instructions()
assert base_instructions is not None
assert base_instructions in skill.instructions
class TestRAGSkillCreation:
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag"
assert skill.metadata.description
assert skill.instructions
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {"search", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is RAGState
assert skill._state_namespace == "rag"
def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.extras["config"] is test_app_config
assert skill.extras["db_path"] is temp_db_path
assert "visualize_chunk" in skill.extras
assert "list_documents" in skill.extras
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.rag import create_skill
skill = create_skill()
assert skill.metadata.name == "rag"
class TestSkillExtras:
async def test_list_documents_returns_all(self, test_app_config, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=rag_db)
list_docs = skill.extras["list_documents"]
results = await list_docs()
assert len(results) == 2
assert all(k in results[0] for k in ("id", "title", "uri", "metadata"))
async def test_list_documents_with_filter(self, test_app_config, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=rag_db)
list_docs = skill.extras["list_documents"]
results = await list_docs(filter="title = 'AI Overview'")
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
class TestSearchTool:
async def test_search_returns_formatted_string(self, rag_db, rag_client):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx(rag=rag_client)
result = await search(ctx, query="artificial intelligence")
assert isinstance(result, str)
assert len(result) > 0
async def test_concurrent_searches_serialize_on_lock(
self, rag_db, rag_client, monkeypatch
):
"""Tool calls in a turn run concurrently; the shared lock keeps only one
connection operation in flight (pydantic-ai borrow guard)."""
from haiku.rag.skills import _tools
from haiku.rag.skills.rag import create_skill
original = _tools.skill_search
inflight = {"now": 0, "max": 0}
async def tracking_search(*args, **kwargs):
inflight["now"] += 1
inflight["max"] = max(inflight["max"], inflight["now"])
try:
await asyncio.sleep(0.05) # widen the window an overlap would use
return await original(*args, **kwargs)
finally:
inflight["now"] -= 1
monkeypatch.setattr(_tools, "skill_search", tracking_search)
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx(rag=rag_client) # one ctx -> one shared lock, as in a turn
results = await asyncio.gather(
search(ctx, query="artificial intelligence"),
search(ctx, query="machine learning"),
)
assert all(isinstance(r, str) and r for r in results)
assert inflight["max"] == 1
async def test_search_updates_state(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="artificial intelligence")
assert "artificial intelligence" in state.searches
results = state.searches["artificial intelligence"]
assert len(results) > 0
assert isinstance(results[0], SearchResult)
async def test_search_applies_document_filter_from_state(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state, rag=rag_client)
result = await search(ctx, query="artificial intelligence")
assert "AI Overview" in result
assert "ML Basics" not in result
async def test_search_without_state(self, rag_db, rag_client):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx(state=None, rag=rag_client)
result = await search(ctx, query="artificial intelligence")
assert isinstance(result, str)
async def test_search_rate_limited(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill
config = AppConfig()
config.qa.max_searches = 2
skill = create_skill(db_path=rag_db, config=config)
search = _get_tool(skill, "search")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="first")
await search(ctx, query="second")
result = await search(ctx, query="third")
assert "Search limit reached" in result
assert ctx.deps.search_count == 3
assert len(state.searches) == 2
class TestCiteTool:
async def test_cite_registers_citations(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="artificial intelligence")
chunk_ids = [
sr.chunk_id
for results in state.searches.values()
for sr in results
if sr.chunk_id
][:2]
result = await cite(ctx, chunk_ids=chunk_ids)
assert "Registered" in result
assert len(state.citations) == 2
assert all(cid in state.citations for cid in chunk_ids)
assert all(cid in state.citation_index for cid in chunk_ids)
async def test_cite_deduplicates_in_index(self, rag_db, rag_client):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="artificial intelligence")
chunk_ids = [
sr.chunk_id
for results in state.searches.values()
for sr in results
if sr.chunk_id
][:1]
await cite(ctx, chunk_ids=chunk_ids)
await cite(ctx, chunk_ids=chunk_ids)
assert len(state.citation_index) == 1
assert len(state.citations) == 1
async def test_cite_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
cite = _get_tool(skill, "cite")
ctx = _make_ctx(state=None)
result = await cite(ctx, chunk_ids=["nonexistent"])
assert "No state" in result
async def test_cite_raises_modelretry_when_chunk_ids_unresolved(
self, rag_db, rag_client
):
"""When supplied chunk_ids don't match any search result, cite raises
ModelRetry so pydantic-ai prompts the model to retry with valid ids
instead of silently registering zero citations."""
from pydantic_ai import ModelRetry
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="artificial intelligence")
assert state.searches, "fixture should have produced some search results"
with pytest.raises(ModelRetry) as exc_info:
await cite(ctx, chunk_ids=["372c9ddf-not-a-real-id"])
message = str(exc_info.value)
assert "verbatim" in message
assert "372c9ddf-not-a-real-id" in message
async def test_cite_raises_modelretry_when_chunk_id_does_not_exist(
self, rag_db, rag_client
):
"""With no prior search() and a chunk_id that doesn't exist in the DB,
cite raises ModelRetry — the DB-fallback path tried and found nothing."""
from pydantic_ai import ModelRetry
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
assert not state.searches
with pytest.raises(ModelRetry) as exc_info:
await cite(ctx, chunk_ids=["nonexistent-chunk-id"])
message = str(exc_info.value)
assert "verbatim" in message
assert "nonexistent-chunk-id" in message
async def test_cite_reports_unresolved_ids_on_partial_success(
self, rag_db, rag_client
):
"""A cite call mixing valid and bogus ids registers the valid ones
and names the rejected ids so the model can re-cite the rest."""
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
await search(ctx, query="artificial intelligence")
valid_id = next(
sr.chunk_id
for results in state.searches.values()
for sr in results
if sr.chunk_id
)
result = await cite(ctx, chunk_ids=[valid_id, "6.43", "6.51.2"])
assert "Registered 1 citation(s)" in result
assert "6.43" in result
assert "6.51.2" in result
assert "verbatim" in result
assert len(state.citations) == 1
assert valid_id in state.citations
async def test_cite_returns_message_when_chunk_ids_empty(self, rag_db):
"""An empty chunk_ids list is a no-op, not a retry trigger."""
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state)
result = await cite(ctx, chunk_ids=[])
assert "0" in result
async def test_cite_accepts_chunk_id_from_db_without_prior_search(
self, rag_db, rag_client
):
"""chunk_ids sourced from items.jsonl / toc.json are valid citations.
The skill calls cite directly with chunk_ids it read from the VFS;
no search() has been recorded in state.searches. cite must look the
chunk up in the DB and build a Citation with full document context.
"""
from haiku.rag.skills.rag import RAGState, create_skill
doc = await rag_client.get_document_by_uri("test://ai-overview")
assert doc, "fixture should have the ai-overview document"
doc_id = doc.id
chunks = await rag_client.chunk_repository.get_by_document_id(doc_id)
assert chunks, "fixture document should have chunks"
chunk_id = chunks[0].id
skill = create_skill(db_path=rag_db)
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state, rag=rag_client)
assert not state.searches, "this test exercises the no-prior-search path"
result = await cite(ctx, chunk_ids=[chunk_id])
assert "Registered 1 citation(s)." == result
assert chunk_id in state.citations
registered = state.citation_index[chunk_id]
assert registered.document_id == doc_id
assert registered.document_uri # uri must be populated from doc lookup
assert registered.document_meta.get("topic") == "ai"
class TestLifespan:
async def test_opens_one_client_per_invocation(self, rag_db):
"""Lifespan opens one HaikuRAG client, available on ctx.deps.rag throughout."""
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
config = AppConfig()
lifespan = make_rag_lifespan(rag_db, config)
deps = RAGRunDeps()
async with lifespan(deps):
assert deps.rag is not None
assert deps.rag.is_read_only
assert deps.search_count == 0
docs = await deps.rag.list_documents()
assert len(docs) == 2
# after exit the client has been closed; field still references it
assert deps.rag is not None
async def test_search_count_resets_per_invocation(self, rag_db):
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
config = AppConfig()
lifespan = make_rag_lifespan(rag_db, config)
deps = RAGRunDeps(search_count=42)
async with lifespan(deps):
assert deps.search_count == 0
deps2 = RAGRunDeps(search_count=5)
async with lifespan(deps2):
assert deps2.search_count == 0
def test_skill_has_lifespan_and_deps_type(self, test_app_config, temp_db_path):
from haiku.rag.skills._deps import RAGRunDeps
from haiku.rag.skills.rag import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.deps_type is RAGRunDeps
assert skill.lifespan is not None
async def test_run_skill_end_to_end_opens_and_closes_client(
self, allow_model_requests, rag_db
):
"""Full sub-agent path: lifespan opens the client, tools see it, lifespan closes it."""
from pydantic_ai.models.test import TestModel
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import run_skill
skill = create_skill(db_path=rag_db)
result, *_ = await run_skill(TestModel(), skill, "List the documents.")
assert result
async def test_lifespan_clears_citations_and_searches_but_keeps_index(self, rag_db):
from haiku.rag.skills._deps import RAGRunDeps, make_rag_lifespan
from haiku.rag.skills.rag import RAGState
from haiku.rag.store.models.citation import Citation
config = AppConfig()
lifespan = make_rag_lifespan(rag_db, config)
state = RAGState(
document_filter="title = 'AI Overview'",
citation_index={
"c1": Citation(
index=1,
chunk_id="c1",
document_id="d1",
document_title="t",
document_uri="u",
content="x",
page_numbers=[],
headings=[],
)
},
citations=["c1"],
searches={"prior": []},
)
deps = RAGRunDeps(state=state)
async with lifespan(deps):
assert state.citations == []
assert state.searches == {}
assert "c1" in state.citation_index # preserved for cross-turn lookup
assert state.document_filter == "title = 'AI Overview'"

3
tests/test_client.py
View file

@ -1133,14 +1133,13 @@ async def test_delete_missing_id_returns_false_without_vacuum(temp_db_path):
@pytest.mark.vcr()
async def test_client_ask(allow_model_requests, temp_db_path):
"""Test asking questions returns answer and citations (VCR recorded)."""
"""Test asking questions through the native RAG capability."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Create a test document for the agent to search
await client.create_document(
content="Python is a high-level programming language.", uri="test.txt"
)
# Use real QA agent with VCR-recorded responses
answer, citations = await client.ask("What is Python?")
# Should return a valid response

4
tests/test_client_analyze.py
View file

@ -12,7 +12,7 @@ def vcr_cassette_dir():
class TestClientAnalysisIntegration:
"""Integration tests for client.analyze() through the rag-analysis skill."""
"""Integration tests for client.analyze() through AnalysisCapability."""
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -65,7 +65,7 @@ class TestClientAnalysisIntegration:
filter="title = 'Cats'",
)
assert "1" in result.answer
assert "1" in result.answer or "one" in result.answer.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()

77
tests/test_picture_in_context.py
View file

@ -12,6 +12,7 @@ from pydantic_ai.messages import BinaryContent, ToolReturn
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from haiku.rag.capabilities.rag import RAGState, create_capability
from haiku.rag.client import HaikuRAG
from haiku.rag.client.search import _populate_image_data
from haiku.rag.config import AppConfig, Config
@ -982,16 +983,7 @@ async def test_search_tool_returns_plain_string_when_no_pictures():
@pytest.mark.asyncio
async def test_skill_search_tool_uses_skill_model_vision_flag(tmp_path):
"""The skill-level ``search`` tool gates picture attachment on the model
passed to ``create_skill_tools`` (per-skill driving model), not on
``config.qa.model.vision``. Verifies the per-skill plumbing by setting
qa.model.vision=False and analysis.model.vision=True simultaneously."""
from haiku.rag.client import HaikuRAG
from haiku.rag.skills._deps import RAGRunDeps
from haiku.rag.skills._tools import create_skill_tools
from haiku.rag.skills.rag import RAGState
async def test_rag_capability_attaches_images_for_vision_model(temp_db_path):
picture_result = SearchResult(
content="A figure",
score=1.0,
@ -1001,58 +993,21 @@ async def test_skill_search_tool_uses_skill_model_vision_flag(tmp_path):
labels=["picture"],
image_data={"#/pictures/0": PICTURE_B64},
)
from haiku.rag.config.models import ModelConfig
fake_client = AsyncMock()
fake_client.search = AsyncMock(return_value=[picture_result])
fake_client.expand_context = AsyncMock(return_value=[picture_result])
config = AppConfig()
config.qa.model.vision = False
# Explicitly set analysis.model (it defaults to None = inherit from qa).
config.analysis.model = ModelConfig(
provider="openai", name="vision-model", vision=True
config.qa.model.vision = True
capability = create_capability(
db_path=temp_db_path,
config=config,
defer_loading=False,
)
capability.state = RAGState()
capability.rag = fake_client
async with HaikuRAG(tmp_path / "db.lancedb", create=True) as rag:
rag.search = AsyncMock(return_value=[picture_result]) # type: ignore[method-assign]
rag.expand_context = AsyncMock(return_value=[picture_result]) # type: ignore[method-assign]
result = await capability._search("anything", None)
# rag skill: should NOT attach binaries (qa.model.vision is False)
rag_tools = create_skill_tools(
tmp_path / "db.lancedb",
config,
RAGState,
["search"],
model=config.qa.model,
)
deps = RAGRunDeps(state=RAGState(), rag=rag, emit=lambda _e: None)
ctx = RunContext(
deps=deps,
model=TestModel(),
usage=RunUsage(),
run_id="run-1",
)
result_rag = await rag_tools["search"](ctx, "anything")
assert isinstance(result_rag, str), (
"rag skill with qa.model.vision=False must return plain text"
)
# analysis skill: SHOULD attach binaries (analysis.model.vision is True)
analysis_tools = create_skill_tools(
tmp_path / "db.lancedb",
config,
RAGState, # state shape doesn't matter for this assertion
["search"],
model=config.analysis.model or config.qa.model,
)
deps2 = RAGRunDeps(state=RAGState(), rag=rag, emit=lambda _e: None)
ctx2 = RunContext(
deps=deps2,
model=TestModel(),
usage=RunUsage(),
run_id="run-2",
)
result_analysis = await analysis_tools["search"](ctx2, "anything")
assert isinstance(result_analysis, ToolReturn), (
"analysis skill with analysis.model.vision=True must wrap binaries"
)
assert result_analysis.content is not None
assert any(isinstance(p, BinaryContent) for p in result_analysis.content)
assert isinstance(result, ToolReturn)
assert result.content is not None
assert any(isinstance(part, BinaryContent) for part in result.content)

42
tests/test_skill_extras.py
View file

@ -1,42 +0,0 @@
"""Tests for non-tool utilities from ``haiku.rag.skills._tools.create_skill_extras``."""
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
from haiku.rag.skills._tools import create_skill_extras
from haiku.rag.store.models.chunk import Chunk
async def _seed_chunk(db_path) -> str:
docling_doc = DoclingDocument(name="extras-test")
chunk = Chunk(
content="Some content.",
metadata={"doc_item_refs": ["#/texts/0"], "page_numbers": [1]},
order=0,
embedding=[0.1] * 2560,
)
async with HaikuRAG(db_path, create=True) as client:
doc = await client.import_document(docling_doc, [chunk], uri="test://extras")
stored = await client.chunk_repository.get_by_document_id(doc.id)
return stored[0].id
async def test_extras_visualize_chunk_accepts_str_and_list(temp_db_path):
chunk_id = await _seed_chunk(temp_db_path)
extras = create_skill_extras(temp_db_path, AppConfig())
visualize_chunk = extras["visualize_chunk"]
# A document imported without page images yields no visualizations, but the
# str and list inputs must both resolve the chunk and reach visualize_chunk.
assert await visualize_chunk(chunk_id) == []
assert await visualize_chunk([chunk_id]) == []
async def test_extras_visualize_chunk_unknown_id_returns_empty(temp_db_path):
await _seed_chunk(temp_db_path)
extras = create_skill_extras(temp_db_path, AppConfig())
visualize_chunk = extras["visualize_chunk"]
assert await visualize_chunk("does-not-exist") == []
assert await visualize_chunk(["does-not-exist"]) == []

497
tests/test_skill_generator.py
View file

@ -1,497 +0,0 @@
import shutil
import subprocess
import zipfile
from pathlib import Path
import pytest
import yaml
from haiku.rag.skill_generator import (
AVAILABLE_TOOLS,
generate_skill,
render_templates,
validate_db_path,
validate_metadata,
validate_output_dir,
validate_tools,
)
class TestAvailableTools:
def test_all_tools_present(self):
assert AVAILABLE_TOOLS == {
"list_documents",
"get_document",
"search",
"execute_code",
"cite",
}
class TestValidateMetadata:
def test_valid(self):
validate_metadata("my-recipes", "A skill.")
def test_rejects_invalid_name(self):
with pytest.raises(ValueError):
validate_metadata("Bad_Name!", "A skill.")
class TestValidateTools:
def test_valid_single_tool(self):
validate_tools(["search"])
def test_valid_multiple_tools(self):
validate_tools(["list_documents", "get_document", "search", "cite"])
def test_valid_all_tools(self):
validate_tools(list(AVAILABLE_TOOLS))
def test_rejects_empty(self):
with pytest.raises(ValueError, match="at least one"):
validate_tools([])
def test_rejects_unknown_tool(self):
with pytest.raises(ValueError, match="Unknown"):
validate_tools(["search", "bogus"])
class TestValidateDbPath:
def test_valid_path(self, tmp_path):
db_path = tmp_path / "test.lancedb"
db_path.mkdir()
validate_db_path(db_path)
def test_rejects_nonexistent(self, tmp_path):
db_path = tmp_path / "nonexistent.lancedb"
with pytest.raises(ValueError, match="does not exist"):
validate_db_path(db_path)
def test_rejects_file(self, tmp_path):
db_path = tmp_path / "test.lancedb"
db_path.touch()
with pytest.raises(ValueError, match="not a directory"):
validate_db_path(db_path)
class TestValidateOutputDir:
def test_valid_output_dir(self, tmp_path):
validate_output_dir(tmp_path, "recipes")
def test_rejects_nonexistent(self, tmp_path):
output_dir = tmp_path / "nonexistent"
with pytest.raises(ValueError, match="does not exist"):
validate_output_dir(output_dir, "recipes")
def test_rejects_existing_target(self, tmp_path):
target = tmp_path / "recipes-skill"
target.mkdir()
with pytest.raises(ValueError, match="already exists"):
validate_output_dir(tmp_path, "recipes")
class TestRenderTemplates:
def test_output_structure(self, tmp_path):
result = render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["list_documents", "get_document", "search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assert result.is_dir()
assert (result / "README.md").is_file()
pkg = result / "recipes_skill"
assert (pkg / "__init__.py").is_file()
assert (pkg / "SKILL.md").is_file()
assert (pkg / "assets").is_dir()
def test_dashed_name_uses_underscores_for_python(self, tmp_path):
result = render_templates(
output_dir=tmp_path,
name="my-recipes",
description="A recipe skill.",
tool_names=["search", "cite"],
)
assert result == tmp_path / "my-recipes-skill"
pkg = result / "my_recipes_skill"
assert (pkg / "__init__.py").is_file()
init = (pkg / "__init__.py").read_text()
assert '"my-recipes.lancedb"' in init
assert 'state_namespace="my-recipes"' in init
toml = (result / "pyproject.toml").read_text()
assert 'name = "my-recipes-skill"' in toml
assert 'my-recipes = "my_recipes_skill:create_skill"' in toml
def test_tool_names_list_matches_selection(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "cite"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert '["search", "cite"]' in content
def test_create_skill_tools_called(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=list(AVAILABLE_TOOLS),
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "create_skill_tools(db_path, config, SkillState, _TOOL_NAMES)" in content
def test_tool_names_in_init(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
assert '"search"' in content
assert '"cite"' in content
def test_pyproject_toml(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
)
toml = tmp_path / "recipes-skill" / "pyproject.toml"
content = toml.read_text()
assert 'name = "recipes-skill"' in content
assert 'description = "A recipe skill."' in content
assert 'readme = "README.md"' in content
assert 'recipes = "recipes_skill:create_skill"' in content
assert "haiku.rag-slim >= " in content
assert "[tool.setuptools.package-data]" in content
assert 'recipes_skill = ["SKILL.md", "assets/**/*"]' in content
def test_skill_md_conditionals(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search"],
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "### search" in content
assert "### cite" not in content
assert "### list_documents" not in content
assert "### execute_code" not in content
def test_skill_md_includes_all_selected(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "execute_code", "cite"],
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "search" in content
assert "execute_code" in content
assert "cite" in content
def test_custom_preamble(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search"],
preamble="You are a docs expert.",
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "You are a docs expert." in content
def test_state_namespace_is_skill_name(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
assert 'state_namespace="recipes"' in content
def test_execute_code_state_fields(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "execute_code"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "executions" in content
def test_imports_from_shared_tools(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "execute_code", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
assert (
"from haiku.rag.skills._tools import create_skill_extras, create_skill_tools"
in content
)
def test_readme(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
)
readme = tmp_path / "recipes-skill" / "README.md"
content = readme.read_text()
assert "recipes" in content
assert "haiku-rag" in content
def test_extras_in_create_skill(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "create_skill_extras" in content
assert "extras=extras" in content
def test_domain_preamble_applied_to_instructions(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "config.prompts.domain_preamble" in content
assert (
'instructions = f"{config.prompts.domain_preamble}\\n\\n{instructions}"'
in content
)
def test_create_skill_accepts_optional_params(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "db_path: Path | None = None" in content
assert "config: AppConfig | None = None" in content
def test_generated_python_is_valid(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=list(AVAILABLE_TOOLS),
)
pkg = tmp_path / "recipes-skill" / "recipes_skill"
for py_file in pkg.glob("*.py"):
source = py_file.read_text()
compile(source, str(py_file), "exec")
def _make_fake_lancedb(path):
path.mkdir()
(path / "data.lance").touch()
return path
class TestGenerateSkill:
def test_end_to_end(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
result = generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assets = result / "recipes_skill" / "assets"
assert (assets / "recipes.lancedb").is_dir()
assert (assets / "recipes.lancedb" / "data.lance").is_file()
assert not (assets / "haiku.rag.yaml").exists()
def test_with_config(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
config_file = tmp_path / "haiku.rag.yaml"
config_file.write_text("storage:\n data_dir: /tmp\n")
result = generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
config_path=config_file,
)
assets = result / "recipes_skill" / "assets"
assert (assets / "haiku.rag.yaml").is_file()
assert (assets / "haiku.rag.yaml").read_text() == (
"storage:\n data_dir: /tmp\n"
)
def test_rejects_invalid_name(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
with pytest.raises(ValueError):
generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="Bad_Name!",
description="A skill.",
tool_names=["search"],
)
def test_rejects_invalid_tools(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
with pytest.raises(ValueError, match="Unknown"):
generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A skill.",
tool_names=["bogus"],
)
def test_rejects_nonexistent_db(self, tmp_path):
with pytest.raises(ValueError, match="does not exist"):
generate_skill(
db_path=tmp_path / "nope.lancedb",
output_dir=tmp_path,
name="recipes",
description="A skill.",
tool_names=["search"],
)
def test_rejects_existing_target(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
(tmp_path / "recipes-skill").mkdir()
with pytest.raises(ValueError, match="already exists"):
generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A skill.",
tool_names=["search"],
)
def test_with_preamble(self, tmp_path):
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
result = generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
preamble="You are a recipe expert.",
)
skill_md = result / "recipes_skill" / "SKILL.md"
content = skill_md.read_text()
assert "You are a recipe expert." in content
def test_wheel_includes_package_data(self, tmp_path):
if not shutil.which("uv"):
pytest.skip("uv not available")
db_path = _make_fake_lancedb(tmp_path / "test.lancedb")
config_file = tmp_path / "haiku.rag.yaml"
config_file.write_text("storage:\n data_dir: /tmp\n")
result = generate_skill(
db_path=db_path,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
config_path=config_file,
)
dist_dir = result / "dist"
subprocess.check_call(
[
"uv",
"run",
"--with",
"build",
"--no-project",
"python",
"-m",
"build",
"--wheel",
"--outdir",
str(dist_dir),
str(result),
]
)
wheels = list(dist_dir.glob("*.whl"))
assert len(wheels) == 1
with zipfile.ZipFile(wheels[0]) as zf:
names = zf.namelist()
assert any(n.endswith("SKILL.md") for n in names)
assert any("assets/" in n and n.endswith("data.lance") for n in names)
assert any(n.endswith("haiku.rag.yaml") for n in names)
def _make_remote_config(tmp_path: Path) -> Path:
config_file = tmp_path / "haiku.rag.yaml"
config_file.write_text(
yaml.dump(
{
"lancedb": {
"uri": "s3://my-bucket/haiku-rag",
"storage_options": {
"endpoint": "http://minio:9000",
"region": "us-east-1",
},
}
}
)
)
return config_file
class TestGenerateSkillRemote:
def test_remote_skips_copytree(self, tmp_path):
config_file = _make_remote_config(tmp_path)
result = generate_skill(
db_path=None,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "cite"],
config_path=config_file,
)
assets = result / "recipes_skill" / "assets"
# No bundled database
assert not (assets / "recipes.lancedb").exists()
# Config must be copied
assert (assets / "haiku.rag.yaml").is_file()
def test_remote_requires_config_path(self, tmp_path):
with pytest.raises(ValueError, match="config_path.*required"):
generate_skill(
db_path=None,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
)

319
tests/test_skill_tools.py
View file

@ -1,319 +0,0 @@
"""Tests for skill tool closures from ``haiku.rag.skills._tools.create_skill_tools``.
These cover the vision toggle on the skill ``search`` tool: when the configured
QA model is vision-capable, picture bytes from search results must reach the
sub-agent as ``BinaryContent`` parts (so a vision model can read figures).
When the QA model is not vision-capable, the same search must return plain
text only.
"""
import base64
from io import BytesIO
from pathlib import Path
from unittest.mock import AsyncMock
import pytest
from PIL import Image as PILImageModule
from pydantic_ai import RunContext
from pydantic_ai.messages import BinaryContent, ToolReturn
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from haiku.rag.config import AppConfig
from haiku.rag.skills._deps import RAGRunDeps
from haiku.rag.skills._tools import create_skill_tools
from haiku.rag.skills.rag import RAGState
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document import Document
def _make_png(color: str = "red") -> bytes:
buf = BytesIO()
PILImageModule.new("RGB", (4, 4), color).save(buf, "PNG")
return buf.getvalue()
PICTURE_BYTES = _make_png("red")
PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")
def _picture_result() -> SearchResult:
return SearchResult(
content="A diagram of the layout",
score=1.0,
chunk_id="chunk-1",
document_id="doc-1",
doc_item_refs=["#/pictures/0"],
labels=["picture"],
image_data={"#/pictures/0": PICTURE_B64},
)
def _text_result() -> SearchResult:
return SearchResult(
content="Some surrounding paragraph text",
score=0.9,
chunk_id="chunk-2",
document_id="doc-1",
doc_item_refs=["#/texts/3"],
labels=["paragraph"],
image_data=None,
)
def _make_ctx(rag, state: RAGState) -> RunContext[RAGRunDeps]:
deps = RAGRunDeps(state=state, rag=rag)
return RunContext(
deps=deps,
model=TestModel(),
usage=RunUsage(),
run_id="run-1",
)
def _build_search_tool(config: AppConfig):
tools = create_skill_tools(
db_path=Path("/tmp/unused.lancedb"),
config=config,
state_type=RAGState,
tool_names=["search"],
model=config.qa.model,
)
return tools["search"]
def _fake_rag(results: list[SearchResult]) -> AsyncMock:
rag = AsyncMock()
rag.search = AsyncMock(return_value=results)
rag.expand_context = AsyncMock(return_value=results)
return rag
@pytest.mark.asyncio
async def test_skill_search_attaches_binary_content_when_vision_capable():
"""vision=True + picture in results → ToolReturn carries BinaryContent."""
config = AppConfig()
config.qa.model.vision = True
search = _build_search_tool(config)
rag = _fake_rag([_picture_result()])
ctx = _make_ctx(rag, RAGState())
result = await search(ctx, "diagram")
assert isinstance(result, ToolReturn)
assert isinstance(result.return_value, str)
assert "rank 1" in result.return_value
assert result.content is not None
assert len(result.content) == 1
part = result.content[0]
assert isinstance(part, BinaryContent)
assert part.data == PICTURE_BYTES
assert part.media_type == "image/png"
assert part.identifier == "#/pictures/0"
@pytest.mark.asyncio
async def test_skill_search_returns_plain_string_when_not_vision_capable():
"""vision=False (the default) + picture in results → plain text only.
The picture bytes must not reach a text-only model — providers behave
inconsistently with image content (Ollama silently accepts and the
model hallucinates)."""
config = AppConfig()
assert config.qa.model.vision is False
search = _build_search_tool(config)
rag = _fake_rag([_picture_result()])
ctx = _make_ctx(rag, RAGState())
result = await search(ctx, "diagram")
assert isinstance(result, str)
assert "rank 1" in result
@pytest.mark.asyncio
async def test_skill_search_returns_plain_string_when_no_pictures():
"""vision=True + no pictures in results → no ToolReturn wrapper, just
text. The wrapper is only needed when there's actually image content
to carry."""
config = AppConfig()
config.qa.model.vision = True
search = _build_search_tool(config)
rag = _fake_rag([_text_result()])
ctx = _make_ctx(rag, RAGState())
result = await search(ctx, "paragraph")
assert isinstance(result, str)
assert "rank 1" in result
@pytest.mark.asyncio
async def test_skill_search_records_results_into_state():
"""Whether or not the QA model is vision-capable, the SearchResult
list must land in state.searches[query] so cite/visualize_chunk can
look chunks up later."""
config = AppConfig()
search = _build_search_tool(config)
rag = _fake_rag([_picture_result(), _text_result()])
state = RAGState()
ctx = _make_ctx(rag, state)
await search(ctx, "anything")
assert "anything" in state.searches
assert len(state.searches["anything"]) == 2
@pytest.mark.asyncio
async def test_skill_search_dedups_picture_bytes_by_self_ref():
"""When two search results reference the same picture self_ref (e.g. a
text chunk and a synthetic picture chunk), the BinaryContent list
must include that picture exactly once. Otherwise the model receives
duplicate image content and pays double the image-token cost."""
config = AppConfig()
config.qa.model.vision = True
other = SearchResult(
content="Surrounding text mentioning the figure",
score=0.8,
chunk_id="chunk-3",
document_id="doc-1",
doc_item_refs=["#/texts/2", "#/pictures/0"],
labels=["paragraph", "picture"],
image_data={"#/pictures/0": PICTURE_B64},
)
search = _build_search_tool(config)
rag = _fake_rag([_picture_result(), other])
ctx = _make_ctx(rag, RAGState())
result = await search(ctx, "figure")
assert isinstance(result, ToolReturn)
assert result.content is not None
assert len(result.content) == 1
assert result.content[0].identifier == "#/pictures/0" # type: ignore[attr-defined]
@pytest.mark.asyncio
async def test_skill_search_keeps_same_self_ref_from_different_documents():
"""``#/pictures/0`` in document A and ``#/pictures/0`` in document B
are different figures. Dedup must key on ``(document_id, self_ref)``;
keying on ``self_ref`` alone would drop document B's bytes."""
other_bytes = _make_png("blue")
other_b64 = base64.b64encode(other_bytes).decode("ascii")
doc_a = SearchResult(
content="Figure in doc A",
score=1.0,
chunk_id="chunk-a",
document_id="doc-A",
doc_item_refs=["#/pictures/0"],
labels=["picture"],
image_data={"#/pictures/0": PICTURE_B64},
)
doc_b = SearchResult(
content="Figure in doc B",
score=0.9,
chunk_id="chunk-b",
document_id="doc-B",
doc_item_refs=["#/pictures/0"],
labels=["picture"],
image_data={"#/pictures/0": other_b64},
)
config = AppConfig()
config.qa.model.vision = True
search = _build_search_tool(config)
rag = _fake_rag([doc_a, doc_b])
ctx = _make_ctx(rag, RAGState())
result = await search(ctx, "figure")
assert isinstance(result, ToolReturn)
assert result.content is not None
assert len(result.content) == 2
payloads = {part.data for part in result.content} # type: ignore[attr-defined]
assert PICTURE_BYTES in payloads
assert other_bytes in payloads
def _build_tool(config: AppConfig, name: str):
tools = create_skill_tools(
db_path=Path("/tmp/unused.lancedb"),
config=config,
state_type=RAGState,
tool_names=[name],
model=config.qa.model,
)
return tools[name]
@pytest.mark.asyncio
async def test_list_documents_tool_returns_shaped_dicts():
config = AppConfig()
list_documents = _build_tool(config, "list_documents")
rag = AsyncMock()
rag.list_documents = AsyncMock(
return_value=[
Document(id="d1", content="x", title="AI", uri="test://ai"),
Document(id="d2", content="y", title="ML", uri="test://ml"),
]
)
ctx = _make_ctx(rag, RAGState())
results = await list_documents(ctx)
assert [r["title"] for r in results] == ["AI", "ML"]
assert all(
set(r.keys()) == {"id", "title", "uri", "metadata", "created_at", "updated_at"}
for r in results
)
@pytest.mark.asyncio
async def test_list_documents_tool_forwards_document_filter_from_state():
config = AppConfig()
list_documents = _build_tool(config, "list_documents")
rag = AsyncMock()
rag.list_documents = AsyncMock(return_value=[])
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(rag, state)
await list_documents(ctx)
rag.list_documents.assert_awaited_once_with(filter="title = 'AI Overview'")
@pytest.mark.asyncio
async def test_get_document_tool_returns_shaped_dict():
config = AppConfig()
get_document = _build_tool(config, "get_document")
rag = AsyncMock()
rag.resolve_document = AsyncMock(
return_value=Document(id="d1", content="full text", title="AI", uri="test://ai")
)
ctx = _make_ctx(rag, RAGState())
result = await get_document(ctx, "AI")
assert result is not None
assert result["content"] == "full text"
assert result["title"] == "AI"
@pytest.mark.asyncio
async def test_get_document_tool_returns_none_when_missing():
config = AppConfig()
get_document = _build_tool(config, "get_document")
rag = AsyncMock()
rag.resolve_document = AsyncMock(return_value=None)
ctx = _make_ctx(rag, RAGState())
result = await get_document(ctx, "nonexistent")
assert result is None

141
uv.lock
View file

@ -225,7 +225,7 @@ wheels = [
[[package]]
name = "anthropic"
version = "0.102.0"
version = "0.118.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
@ -237,9 +237,9 @@ dependencies = [
{ name = "sniffio" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b9/47/cb2a71f70431fb09af4db83e3ea89eb2dd8e0e348d27af53ed32e6c599dd/anthropic-0.102.0.tar.gz", hash = "sha256:96f747cad11886c4ae12d4080131b94eebd68b202bd2190fe27959031bb1fa9c", size = 763697, upload-time = "2026-05-13T18:12:41.624Z" }
sdist = { url = "https://files.pythonhosted.org/packages/07/ff/bbad57650babf07ae9761f1d974fcee73430387b5d3aa0ddb395234906e8/anthropic-0.118.0.tar.gz", hash = "sha256:acb3c43b7e7592cf635fa21930e67a99d32641d68ccd53f625b52f71d7870d55", size = 994705, upload-time = "2026-07-22T16:43:57.019Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/87/75/0f6c603594876413bc858a00e7cc0d80a0cc14edf5c7b959a3ea6ec45e44/anthropic-0.102.0-py3-none-any.whl", hash = "sha256:ab96540bbd4b0f36564252d955a86f8abbe4f00944a24bc9931acc9b139bab6f", size = 763070, upload-time = "2026-05-13T18:12:43.474Z" },
{ url = "https://files.pythonhosted.org/packages/91/9f/d46e77054d7efb61c85ef5e8c7e95927cc89988b6491b46d22af60d63bda/anthropic-0.118.0-py3-none-any.whl", hash = "sha256:524d835869b8e374510b3bcc552e28dbdafc035a61d63fc0088c4035286ef8fe", size = 1010922, upload-time = "2026-07-22T16:43:55.461Z" },
]
[[package]]
@ -1411,15 +1411,15 @@ wheels = [
[[package]]
name = "genai-prices"
version = "0.0.56"
version = "0.0.72"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "httpx" },
{ name = "httpx2" },
{ name = "pydantic" },
]
sdist = { url = "https://files.pythonhosted.org/packages/44/6b/94b3018a672c7775edfb485f0fed8f6068fba75e49b067e8a1ac5eb96764/genai_prices-0.0.56.tar.gz", hash = "sha256:ac24b16a84d0ab97539bfa48dfa4649689de8e3ce71c12ebacef29efb1998045", size = 65872, upload-time = "2026-03-20T20:33:00.732Z" }
sdist = { url = "https://files.pythonhosted.org/packages/6f/c8/2549fa8ceaaf0bd61114cf392250a107be657233723bc4e65cb91dba934f/genai_prices-0.0.72.tar.gz", hash = "sha256:a7e481d0ea85922fcf48df6864f2491fc81a4f03dcf0ecbd9aa8c2c6d9fcdbe8", size = 82753, upload-time = "2026-07-22T21:11:03.997Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a3/f6/8ef7e4c286deb2709d11ca96a5237caae3ef4876ab3c48095856cfd2df30/genai_prices-0.0.56-py3-none-any.whl", hash = "sha256:dbe86be8f3f556bed1b72209ed36851fec8b01793b3b220f42921a4e7da945f6", size = 68966, upload-time = "2026-03-20T20:33:02.555Z" },
{ url = "https://files.pythonhosted.org/packages/22/2e/1e2c31666f7e4e2d0327a80f1c920fd8764ccd27672c4bc389d79219ece6/genai_prices-0.0.72-py3-none-any.whl", hash = "sha256:21281d8df34d9bfbb736e0763a60ceda02226938cd47379ae437a44fcf8ba23d", size = 85238, upload-time = "2026-07-22T21:11:02.911Z" },
]
[[package]]
@ -1672,10 +1672,8 @@ source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "docling-core" },
{ name = "fastmcp" },
{ name = "haiku-skills" },
{ name = "httpx" },
{ name = "jinja2" },
{ name = "jsonpatch" },
{ name = "lancedb" },
{ name = "pathspec" },
{ name = "pydantic" },
@ -1738,7 +1736,7 @@ tui = [
{ name = "tree-sitter-json" },
]
vertexai = [
{ name = "pydantic-ai-slim", extra = ["vertexai"] },
{ name = "pydantic-ai-slim", extra = ["google"] },
]
voyageai = [
{ name = "pydantic-ai-slim", extra = ["voyageai"] },
@ -1757,10 +1755,8 @@ requires-dist = [
{ name = "fastapi", marker = "extra == 'ingester'", specifier = ">=0.125" },
{ name = "fastmcp", specifier = ">=3.3.0" },
{ name = "haiku-rag-slim", extras = ["s3"], marker = "extra == 'ingester'", editable = "haiku_rag_slim" },
{ name = "haiku-skills", specifier = ">=0.18.0" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "jinja2", specifier = ">=3.1.0" },
{ name = "jsonpatch", specifier = ">=1.33" },
{ name = "lancedb", specifier = "==0.34.0" },
{ name = "obstore", marker = "extra == 's3'", specifier = ">=0.9,<0.10" },
{ name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.6.0.66,<5.0.0.0" },
@ -1769,10 +1765,10 @@ requires-dist = [
{ name = "pydantic-ai-slim", extras = ["anthropic"], marker = "extra == 'anthropic'" },
{ name = "pydantic-ai-slim", extras = ["bedrock"], marker = "extra == 'bedrock'" },
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'" },
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'vertexai'" },
{ name = "pydantic-ai-slim", extras = ["groq"], marker = "extra == 'groq'" },
{ name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" },
{ name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=1.100.0" },
{ name = "pydantic-ai-slim", extras = ["vertexai"], marker = "extra == 'vertexai'" },
{ name = "pydantic-ai-slim", extras = ["openai", "logfire", "ag-ui"], specifier = ">=2.11.0,<3.0.0" },
{ name = "pydantic-ai-slim", extras = ["voyageai"], marker = "extra == 'voyageai'" },
{ name = "pydantic-monty", specifier = ">=0.0.17" },
{ name = "pypdfium2", specifier = ">=5.0" },
@ -1795,23 +1791,6 @@ requires-dist = [
]
provides-extras = ["docling", "s3", "voyageai", "cohere", "zeroentropy", "jina", "cross-encoder", "ingester", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
[[package]]
name = "haiku-skills"
version = "0.18.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "ag-ui-protocol" },
{ name = "jsonpatch" },
{ name = "pydantic" },
{ name = "pydantic-ai-slim", extra = ["ag-ui", "mcp", "openai"] },
{ name = "pyyaml" },
{ name = "skills-ref" },
]
sdist = { url = "https://files.pythonhosted.org/packages/64/a4/41cddc90280bea6d4150ed3361fe1cc804a0d110ca11319a78bf9880cfe9/haiku_skills-0.18.0.tar.gz", hash = "sha256:ac7033bafa79059800515a69d5e60a6fb6f0659174694bfceec473308add85b1", size = 189494, upload-time = "2026-06-29T08:00:29.782Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/bd/b9/33ed400bd24ba0083ccc352053aae536b7a6bdbcd000171b7683bb756fd4/haiku_skills-0.18.0-py3-none-any.whl", hash = "sha256:cfeb0cd2bd8d297ccd607cf9d94f635c1a9d411676bb3ac8516b149d372154a8", size = 34044, upload-time = "2026-06-29T08:00:28.7Z" },
]
[[package]]
name = "hf-xet"
version = "1.5.2"
@ -1849,6 +1828,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl", hash = "sha256:2d400746a40668fc9dec9810239072b40b4484b640a8c38fd654a024c7a1bf55", size = 78784, upload-time = "2025-04-24T22:06:20.566Z" },
]
[[package]]
name = "httpcore2"
version = "2.8.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "h11" },
{ name = "truststore" },
]
sdist = { url = "https://files.pythonhosted.org/packages/ce/50/a33260c2a959a7f12d3cf6b96bb35c7028c5d1904fa7b2f5850fa175af01/httpcore2-2.8.0.tar.gz", hash = "sha256:44be3730e7cd4fd206478488c00f57c0960b5e00ee9f7099ef18bb2b0c5cdb79", size = 66779, upload-time = "2026-07-23T11:54:40.443Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/12/b7/76d6baa8c087cf5170f36dbeb92471e77c5c48f8f7ca8e5ba4ad613ea492/httpcore2-2.8.0-py3-none-any.whl", hash = "sha256:60ee2f9963ca942955ffd44553fc1111f00286283f9cc047d74e19888e061ace", size = 82638, upload-time = "2026-07-23T11:54:37.976Z" },
]
[[package]]
name = "httptools"
version = "0.7.1"
@ -1902,6 +1894,22 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/d2/fd/6668e5aec43ab844de6fc74927e155a3b37bf40d7c3790e49fc0406b6578/httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc", size = 8960, upload-time = "2025-10-10T21:48:21.158Z" },
]
[[package]]
name = "httpx2"
version = "2.8.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
{ name = "httpcore2" },
{ name = "idna" },
{ name = "truststore" },
{ name = "typing-extensions", marker = "python_full_version < '3.13'" },
]
sdist = { url = "https://files.pythonhosted.org/packages/31/ca/928387a80e8e6ebbfa7ed020f1412e99038092657d63d801fe935f56677c/httpx2-2.8.0.tar.gz", hash = "sha256:7137d9278553773c672856112c36062ae691146212ab748a6a8362479b19bdee", size = 94485, upload-time = "2026-07-23T11:54:41.363Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e6/a7/aea061b450d51cf448c1ccc8e1fc6f04c0d8fdc97c17524c7f835018bb92/httpx2-2.8.0-py3-none-any.whl", hash = "sha256:9617a6af8354667c94dc06a1b7b88a32db618ffb1f0649f8be31c0639dc58227", size = 90232, upload-time = "2026-07-23T11:54:39.222Z" },
]
[[package]]
name = "huggingface-hub"
version = "1.24.0"
@ -2626,7 +2634,7 @@ wheels = [
[[package]]
name = "mistralai"
version = "2.1.3"
version = "2.7.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "eval-type-backport" },
@ -2638,9 +2646,9 @@ dependencies = [
{ name = "python-dateutil" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/7d/98/5fe39d514c19477f06a07e088ce4a44c2e60ac9deebefb9e2c8ed8ef87d2/mistralai-2.1.3.tar.gz", hash = "sha256:0c5de4855b043cd0582406d5c1ddfd91e176f484a158e6ee0b4a0054231be266", size = 331929, upload-time = "2026-03-23T15:00:29.579Z" }
sdist = { url = "https://files.pythonhosted.org/packages/7d/e0/bc0188444f122083e1cd554a33440589f4c7a55bab57563333f004cf4934/mistralai-2.7.1.tar.gz", hash = "sha256:1b67a224a9387b33d8ca381d24cd8d075dbe6650af305d763d71c4cad06b2ddf", size = 530696, upload-time = "2026-07-21T13:12:51.041Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/15/7c/f91a26bf469c1cff57325379afa112baeb113ac577d28e69dd408cee5745/mistralai-2.1.3-py3-none-any.whl", hash = "sha256:26daac3bdc69fc2dd58f2c421710eb34131be7883b44a9ea81904a6306e6a90a", size = 754931, upload-time = "2026-03-23T15:00:30.934Z" },
{ url = "https://files.pythonhosted.org/packages/45/1a/c502cad6ebb3363911ad30ec0318c216dcac4f5be485651b7672bdcd3ff7/mistralai-2.7.1-py3-none-any.whl", hash = "sha256:25a78835e52eb44ca8c10046dbb1635545c53c40a06397075df246c82f12b186", size = 1266550, upload-time = "2026-07-21T13:12:49.5Z" },
]
[[package]]
@ -3092,7 +3100,7 @@ wheels = [
[[package]]
name = "openai"
version = "2.29.0"
version = "2.47.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
@ -3104,9 +3112,9 @@ dependencies = [
{ name = "tqdm" },
{ name = "typing-extensions" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b4/15/203d537e58986b5673e7f232453a2a2f110f22757b15921cbdeea392e520/openai-2.29.0.tar.gz", hash = "sha256:32d09eb2f661b38d3edd7d7e1a2943d1633f572596febe64c0cd370c86d52bec", size = 671128, upload-time = "2026-03-17T17:53:49.599Z" }
sdist = { url = "https://files.pythonhosted.org/packages/bf/61/9aeef14de759306e85175126d3d6d56ee4f5072a9512c6c171d58d02a62d/openai-2.47.0.tar.gz", hash = "sha256:4e205548acd4304f235b86202269912e55bc88270b15d2a051fa2b53b90343a6", size = 1089906, upload-time = "2026-07-22T17:47:29.723Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/d0/b1/35b6f9c8cf9318e3dbb7146cc82dab4cf61182a8d5406fc9b50864362895/openai-2.29.0-py3-none-any.whl", hash = "sha256:b7c5de513c3286d17c5e29b92c4c98ceaf0d775244ac8159aeb1bddf840eb42a", size = 1141533, upload-time = "2026-03-17T17:53:47.348Z" },
{ url = "https://files.pythonhosted.org/packages/41/69/26b032059273ad798d18fbcdbe369e871181841fd8bcb5caee32b7510039/openai-2.47.0-py3-none-any.whl", hash = "sha256:b3a1a7ad974092427ccb46d89f8852bdb67866680bcabeecc3ff5a3fdd71b15b", size = 1639987, upload-time = "2026-07-22T17:47:27.873Z" },
]
[[package]]
@ -3818,7 +3826,7 @@ email = [
[[package]]
name = "pydantic-ai-slim"
version = "1.102.0"
version = "2.16.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "genai-prices" },
@ -3829,9 +3837,9 @@ dependencies = [
{ name = "pydantic-graph" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/e2/3e/14980440e8f0532535e1fbe936fec5f8d8e7bc6cafa81f6f3c51b1884fe5/pydantic_ai_slim-1.102.0.tar.gz", hash = "sha256:0b8f2b70fa2b40efcbd09d341a346934fc4e46622ae281f858c6bfd3d0d3152b", size = 739988, upload-time = "2026-05-23T01:14:32.808Z" }
sdist = { url = "https://files.pythonhosted.org/packages/76/6a/3048579c646f4cea7966889009a1d7eef1365f9126393945d39cefad9e83/pydantic_ai_slim-2.16.0.tar.gz", hash = "sha256:36d17cb12edd72ffc62f9e06cc49ac5f23cb77cf6b665c4b1fd11c00dbe6a852", size = 889343, upload-time = "2026-07-23T02:47:20.696Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/b4/2e/089df86adaf904dd97a1b139d29fe728af0e41430d747f5b6315df3b0c1e/pydantic_ai_slim-1.102.0-py3-none-any.whl", hash = "sha256:f9fa9c3fb58a76f85522f78d1037d201b424de46d532263ed780b3730060449f", size = 919311, upload-time = "2026-05-23T01:14:23.464Z" },
{ url = "https://files.pythonhosted.org/packages/2e/95/2c79b9f8e875562bae8141078af27428120e8415d493bb07ea36ca5905f9/pydantic_ai_slim-2.16.0-py3-none-any.whl", hash = "sha256:7cad27fb8f45ce4af4e8da83d7f206a3e1038dbc7535625c0ce5518fe378a55d", size = 1077057, upload-time = "2026-07-23T02:47:12.633Z" },
]
[package.optional-dependencies]
@ -3857,9 +3865,6 @@ groq = [
logfire = [
{ name = "logfire", extra = ["httpx"] },
]
mcp = [
{ name = "fastmcp-slim", extra = ["client"] },
]
mistral = [
{ name = "mistralai" },
]
@ -3867,10 +3872,6 @@ openai = [
{ name = "openai" },
{ name = "tiktoken" },
]
vertexai = [
{ name = "google-auth" },
{ name = "requests" },
]
voyageai = [
{ name = "voyageai", marker = "python_full_version < '3.14'" },
]
@ -3948,7 +3949,7 @@ wheels = [
[[package]]
name = "pydantic-evals"
version = "1.102.0"
version = "2.16.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "anyio" },
@ -3958,14 +3959,14 @@ dependencies = [
{ name = "pyyaml" },
{ name = "rich" },
]
sdist = { url = "https://files.pythonhosted.org/packages/2a/2a/2f0a18e170dc1db4b32120bea9e1162ef196c1f453db823878f5eaf7b8bb/pydantic_evals-1.102.0.tar.gz", hash = "sha256:711a6335d24a11c324e5a5c7758b12dfd77209f885ab2501d7eedb9dd5b75b18", size = 78557, upload-time = "2026-05-23T01:14:34.447Z" }
sdist = { url = "https://files.pythonhosted.org/packages/52/f5/7c2bf8ce45da52ed70c030e91ce43747317f61bd0615a40e91f1ec515c70/pydantic_evals-2.16.0.tar.gz", hash = "sha256:717e9615c7688650cdc716046f1d8edfbe0bac1748286c0a2744e3689fb518c7", size = 85147, upload-time = "2026-07-23T02:47:22.089Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/e2/fd/2281c166b2c5cedab003b12bf8a630656cb5a9bbd552e4981ee190570d15/pydantic_evals-1.102.0-py3-none-any.whl", hash = "sha256:579edd6f7056d0fe52e03c7004377a0b9c42264c60a370258235fb0750fe20a2", size = 93529, upload-time = "2026-05-23T01:14:25.559Z" },
{ url = "https://files.pythonhosted.org/packages/f0/f6/d7f20f976f81d073856138bd046bf1ce5dc96561391526753c1c9076d33c/pydantic_evals-2.16.0-py3-none-any.whl", hash = "sha256:6705b427ea7c77d7f6b7d152ced6f86728931eceb9c57a8aa07ab8225fb1a934", size = 100439, upload-time = "2026-07-23T02:47:14.523Z" },
]
[[package]]
name = "pydantic-graph"
version = "1.102.0"
version = "2.16.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "httpx" },
@ -3973,9 +3974,9 @@ dependencies = [
{ name = "pydantic" },
{ name = "typing-inspection" },
]
sdist = { url = "https://files.pythonhosted.org/packages/51/37/4265a1a63eddf35a5aa621c9b2355525bdeae3eb59c3954b165fbfe31404/pydantic_graph-1.102.0.tar.gz", hash = "sha256:e285bd7115e4e92676eaf0a5e7e6faa64cda8c4819f67923a118c50666b909ab", size = 62584, upload-time = "2026-05-23T01:14:36.056Z" }
sdist = { url = "https://files.pythonhosted.org/packages/c3/22/6b6426e14607275f6b15f2a0c4530514c48ff01f86c53bbaac4041d09df9/pydantic_graph-2.16.0.tar.gz", hash = "sha256:f71e5c8e78a4ce56bc044861178e506ebd99881717ae0993924c457f5de6230e", size = 43979, upload-time = "2026-07-23T02:47:23.328Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/a4/49/5597c52d50114440047dd4ce4f6505e32ee336f43267639907d1a17648ee/pydantic_graph-1.102.0-py3-none-any.whl", hash = "sha256:b1a28314adc4abca4db02cf095d064782ec5712e0847ce7a6b79a3c84bf1fc01", size = 80100, upload-time = "2026-05-23T01:14:27.583Z" },
{ url = "https://files.pythonhosted.org/packages/80/c0/362a7fb50562d7b51e3d02d454826e0a7b5652f27e5be8e835ddfaf84da6/pydantic_graph-2.16.0-py3-none-any.whl", hash = "sha256:99c25852c436d4d510d1ecdcb53ce4a0b11aa4d1d9d81472727587f2b5a3bc13", size = 51661, upload-time = "2026-07-23T02:47:15.921Z" },
]
[[package]]
@ -4897,19 +4898,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" },
]
[[package]]
name = "skills-ref"
version = "0.1.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "click" },
{ name = "strictyaml" },
]
sdist = { url = "https://files.pythonhosted.org/packages/23/42/943d3ba8b097af7068b7178563a5062ad8a977982f4a7b4f67facfc575e9/skills_ref-0.1.1.tar.gz", hash = "sha256:6b400ca6e0049be62dca0167ff943ba2745fd67efb37fbba4d0ee341fccd2695", size = 93519, upload-time = "2026-01-10T13:23:41.423Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/af/25/36a43c3a61fb6cc3984e6ad5e556929b8ae71c95eba615dae4cf2f427964/skills_ref-0.1.1-py3-none-any.whl", hash = "sha256:d35db5bb8de71ae301daf5ca9cb71f8a555e8c6f83a6d40e46a5bc09f8f461b5", size = 12918, upload-time = "2026-01-10T13:23:40.106Z" },
]
[[package]]
name = "sniffio"
version = "1.3.1"
@ -5000,18 +4988,6 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/ec/bb/2799cc2ede3ed41131f8975621e7213dfc7ef4acbbaadfa440f32500c370/starlette-1.3.1-py3-none-any.whl", hash = "sha256:c7372aae11c3c3f26a42df7bd626cec2f47d03483d261d369516a615a53714c6", size = 73632, upload-time = "2026-06-12T09:23:10.017Z" },
]
[[package]]
name = "strictyaml"
version = "1.7.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "python-dateutil" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b3/08/efd28d49162ce89c2ad61a88bd80e11fb77bc9f6c145402589112d38f8af/strictyaml-1.7.3.tar.gz", hash = "sha256:22f854a5fcab42b5ddba8030a0e4be51ca89af0267961c8d6cfa86395586c407", size = 115206, upload-time = "2023-03-10T12:50:27.062Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/96/7c/a81ef5ef10978dd073a854e0fa93b5d8021d0594b639cc8f6453c3c78a1d/strictyaml-1.7.3-py3-none-any.whl", hash = "sha256:fb5c8a4edb43bebb765959e420f9b3978d7f1af88c80606c03fb420888f5d1c7", size = 123917, upload-time = "2023-03-10T12:50:17.242Z" },
]
[[package]]
name = "sympy"
version = "1.14.0"
@ -5419,6 +5395,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/f0/ac/229b7d4589d2e5937310e72c6d46e89599d16a4a12b479ffa1499fee8eb8/triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68", size = 197824404, upload-time = "2026-06-17T19:53:42.772Z" },
]
[[package]]
name = "truststore"
version = "0.10.4"
source = { registry = "https://pypi.org/simple" }
sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" },
]
[[package]]
name = "ty"
version = "0.0.28"

9
zensical.toml
View file

@ -19,11 +19,10 @@ nav = [
{ CLI = "cli.md" },
{ Chat = "chat.md" },
] },
{ Skills = [
"skills/index.md",
{ "RAG skill" = "skills/rag.md" },
{ "Analysis skill" = "skills/analysis.md" },
{ "Custom skills" = "skills/custom.md" },
{ Capabilities = [
"capabilities/index.md",
{ "RAG capability" = "capabilities/rag.md" },
{ "Analysis capability" = "capabilities/analysis.md" },
] },
{ Configure = [
"configuration/index.md",