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

Merge pull request #342 from ggozad/feat/analysis

Browse source 
Refactor analysis sandbox with document VFS and expanded search
This commit is contained in:
Yiorgis Gozadinos 2026-04-20 16:34:41 +03:00 committed by GitHub
commit 2cd4a85f5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
96 changed files with 2402 additions and 2939 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

39
CHANGELOG.md
View file

@ -1,6 +1,45 @@
# Changelog
## [Unreleased]
### Added
- **Document virtual filesystem in analysis sandbox**: Documents mounted at `/documents/{id}/` with `metadata.json` (eager), `content.txt` (lazy), and `items.jsonl` (lazy). Standard Python `pathlib.Path` for browsing and reading document content and structure.
- **`execute_code` skill tool**: Direct code execution in the sandbox, surfaced as individual AG-UI events in the chat TUI. Items VFS uses a lazy bulk cache (~1s for 1000 documents vs 60s+ per-document queries).
- **`cite` skill tool**: Explicit citation registration with per-turn tracking via `citation_index` and `citations` fields in state
- **`--skill` flag for chat TUI**: `haiku-rag chat -s rag -s analysis` to enable specific skills
- **`--model` overrides all agents**: Chat, QA, research, and analysis agents all use the specified model
- **Collapsible program display in chat TUI**: Analysis code execution results shown as expandable code blocks
### Changed
- **BREAKING: Flatten skill architecture**: Skill sub-agents now call `search`, `execute_code`, `cite`, `list_documents`, `get_document` directly — every tool call surfaces as an AG-UI event. Removes the 3rd agent layer where `ask`/`analyze`/`research` spawned inner agents whose tool calls were invisible.
- **BREAKING: Rename RLM agent to analysis agent** throughout:
- `agents/rlm/` → `agents/analysis/`, all classes renamed (`RLMResult` → `AnalysisResult`, etc.)
- `client.rlm()` → `client.analyze()`
- CLI: `haiku-rag rlm` → `haiku-rag analyze`
- MCP: `rlm_question` → `analyze`
- Config: `rlm:` → `analysis:` in YAML, `RLMConfig` → `AnalysisConfig`
- Skill entrypoint: `rag-rlm` → `rag-analysis`
- **Analysis sandbox `search()` returns expanded results** with `doc_item_refs` and `labels` for cross-referencing with `items.jsonl`
- **`list_documents` skill tool** takes no parameters — returns all documents
- **Per-turn citation tracking**: `citation_index: dict[str, Citation]` (deduplicated) + `citations: list[list[str]]` (per-turn chunk IDs) replaces flat citation list
- **Search rate limiting**: Skill search tool enforces `config.qa.max_searches`
- **Context expansion respects section boundaries**: Sections within the char budget are returned whole regardless of item count. Too-large sections expand bounded by section edges. Adjacent sections no longer merge — only overlapping ranges do.
- **Visualization shows full expanded section**: `visualize_chunk` expands context before resolving bounding boxes, so all pages the section spans get highlighted.
### Removed
- **`ask` skill tool**: Replaced by direct `search` + `cite` — the skill sub-agent searches and answers directly
- **`analyze` skill tool**: Replaced by direct `execute_code` + `search` + `cite`
- **`research` skill tool**: Removed from skill layer (still available via CLI `haiku-rag research` and MCP)
- **`get_document()`, `get_docling_document()`**: Removed from analysis sandbox — replaced by VFS
- **`get_chunk()`**: Removed from analysis sandbox — search results include expanded context
- **`create_analysis_toolset()`**: Removed unused `tools/analysis.py` module
- **`qa_history`, `reports` from skill state**: Conversational context handled by the outer chat agent
- **`combine_filters`, `build_document_filter`**: Removed from public API
- **`max_context_items`**: Removed from `SearchConfig` — `max_context_chars` is the sole expansion constraint
- **`QAHistoryEntry`, `tools/qa.py`**: Removed unused QA history model and relevance threshold
## [0.40.1] - 2026-04-17
### Fixed

8
README.md
View file

@ -11,7 +11,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Question answering** — QA agents with citations (page numbers, section headings)
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
- **RLM agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM. QA/Research: any model supported by Pydantic AI
@ -62,8 +62,8 @@ haiku-rag ask "What datasets were used for evaluation?" --cite
# Research mode — iterative planning and search
haiku-rag research "What are the limitations of the approach?"
# RLM mode — complex analytical tasks via code execution
haiku-rag rlm "How many documents mention transformers?"
# Analyze — complex analytical tasks via code execution
haiku-rag analyze "How many documents mention transformers?"
# Interactive chat — multi-turn conversations with memory
haiku-rag chat
@ -138,7 +138,7 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA and research agents
- [RLM Agent](https://ggozad.github.io/haiku.rag/rlm/) - Complex analytical tasks via code execution
- [Analysis Agent](https://ggozad.github.io/haiku.rag/agents/analysis/) - Complex analytical tasks via code execution
- [Applications](https://ggozad.github.io/haiku.rag/apps/) - Chat TUI, web app, and inspector
- [Server](https://ggozad.github.io/haiku.rag/server/) - File monitoring and MCP
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration

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

@ -21,8 +21,8 @@ import { FilterIcon } from "../lib/icons";
import type { RAGState } from "../lib/sessionStorage";
import {
createSession,
deriveCitationsHistory,
getActiveSessionId,
getLatestCitations,
getSession,
normalizeRAGState,
updateSessionMessages,
@ -148,11 +148,11 @@ function ToolCallIndicator({
switch (toolName) {
case "search":
return <SearchIcon />;
case "ask":
return <MessageIcon />;
case "get_document":
return <FileIcon />;
case "execute_skill":
case "execute_code":
case "cite":
return <MessageIcon />;
default:
return <SearchIcon />;
@ -163,16 +163,16 @@ function ToolCallIndicator({
switch (toolName) {
case "search":
return "Search";
case "ask":
return "Ask";
case "get_document":
return "Document";
case "execute_skill":
return "Skill";
case "analyze":
return "Analyze";
case "research":
return "Research";
case "execute_code":
return "Code";
case "cite":
return "Cite";
case "list_documents":
return "Documents";
default:
return toolName;
}
@ -194,16 +194,20 @@ function ToolCallIndicator({
const query = args.query as string;
return <span className="tool-query">{query}</span>;
}
case "ask": {
const question = args.question as string;
return <span className="tool-query">{question}</span>;
}
case "get_document":
return <span className="tool-query">{args.query as string}</span>;
case "analyze":
return <span className="tool-query">{args.question as string}</span>;
case "research":
return <span className="tool-query">{args.question 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":
return <span className="tool-query">Registering citations</span>;
case "list_documents":
return <span className="tool-query">Listing documents</span>;
default:
return <span>Processing...</span>;
}
@ -292,7 +296,7 @@ function MessageViewWithCitations({
isRunning?: boolean;
}) {
const ragState = useContext(ChatStateContext);
const citationsHistory = ragState ? deriveCitationsHistory(ragState) : [];
const latestCitations = ragState ? getLatestCitations(ragState) : [];
// Collect completed tool_call_ids from skill_tool_result activity messages
const completedToolCallIds = useMemo(() => {
@ -326,7 +330,6 @@ function MessageViewWithCitations({
{({ messageElements }) => {
const result: React.ReactNode[] = [];
let elemIdx = 0;
let citIdx = 0;
let seenToolCalls = false;
for (const msg of messages) {
@ -368,19 +371,15 @@ function MessageViewWithCitations({
}
// After an assistant text response that followed tool calls,
// inject the next citations entry (one per turn)
// show citations from the latest turn
if (msg.role === "assistant" && msg.content && seenToolCalls) {
if (citIdx < citationsHistory.length) {
const citations = citationsHistory[citIdx];
if (citations?.length) {
result.push(
<CitationBlock
key={`citations-${citIdx}`}
citations={citations}
/>,
);
}
citIdx++;
if (latestCitations.length > 0) {
result.push(
<CitationBlock
key={`citations-${i}`}
citations={latestCitations}
/>,
);
}
seenToolCalls = false;
}

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

@ -9,33 +9,12 @@ export interface Citation {
content: string;
}
export interface QAHistoryEntry {
question: string;
answer: string;
citations: Citation[];
}
export interface DocumentInfo {
id: string;
title: string;
uri: string;
created: string;
}
export interface ResearchEntry {
question: string;
title: string;
executive_summary: string;
}
// Matches RAGState from the backend skill
export interface RAGState {
citations: Citation[];
qa_history: QAHistoryEntry[];
citation_index: Record<string, Citation>;
citations: string[][];
document_filter: string | null;
searches: Record<string, unknown[]>;
documents: DocumentInfo[];
reports: ResearchEntry[];
}
export interface StoredMessage {
@ -59,20 +38,20 @@ const ACTIVE_SESSION_KEY = "haiku.rag.activeSession";
export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
return {
citation_index: state?.citation_index ?? {},
citations: state?.citations ?? [],
qa_history: state?.qa_history ?? [],
document_filter: state?.document_filter ?? null,
searches: state?.searches ?? {},
documents: state?.documents ?? [],
reports: state?.reports ?? [],
};
}
// Derive per-turn citation arrays from qa_history
export function deriveCitationsHistory(state: RAGState): Citation[][] {
return state.qa_history
.filter((entry) => entry.citations?.length > 0)
.map((entry) => entry.citations);
export function getLatestCitations(state: RAGState): Citation[] {
const turns = state.citations;
if (turns.length === 0) return [];
const latestIds = turns[turns.length - 1];
return latestIds
.map((id) => state.citation_index[id])
.filter((c): c is Citation => c !== undefined);
}
export function getAllSessions(): StoredSession[] {

127
docs/agents/analysis.md Normal file
View file

@ -0,0 +1,127 @@
# Analysis Agent
The analysis agent enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
- **Aggregation**: "How many documents mention security vulnerabilities?"
- **Computation**: "What's the average revenue across all quarterly reports?"
- **Multi-document analysis**: "Compare the key findings between Report A and Report B"
- **Structured data extraction**: "Extract all dollar amounts and compute totals"
## How It Works
1. The agent receives a question
2. It writes Python code to explore the knowledge base
3. Code executes in a sandboxed Python interpreter with access to search, LLM, and a virtual filesystem of documents
4. The agent iterates: run code, examine results, refine approach
5. Final answer is synthesized from the gathered data
## CLI Usage
```bash
# Basic usage
haiku-rag analyze "How many documents are in the database?"
# With document filter (restricts what the agent can access)
haiku-rag analyze "Summarize the key points" --filter "uri LIKE '%report%'"
# Pre-load specific documents
haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document "Q2 Report"
```
## Python Usage
```python
from haiku.rag.client import HaikuRAG
async with HaikuRAG(path_to_db) as client:
# Basic question
result = await client.analyze("How many documents mention 'security'?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
# With filter (agent can only see filtered documents)
result = await client.analyze(
"What is the total revenue?",
filter="title LIKE '%Financial%'"
)
# Pre-load specific documents
result = await client.analyze(
"Compare the conclusions",
documents=["Report A", "Report B"]
)
```
## Sandbox Capabilities
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with:
### Functions
| Function | Description |
|----------|-------------|
| `search(query, limit)` | Hybrid search (vector + full-text) with automatic context expansion. Returns `doc_item_refs` for cross-referencing with `items.jsonl` |
| `list_documents()` | List all documents in the knowledge base |
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
### Document Filesystem
All documents are mounted as a virtual filesystem at `/documents/`. The agent uses standard Python `pathlib.Path` to browse and read files:
```
/documents/{document_id}/
metadata.json # {id, title, uri, created_at}
content.txt # Full document text
items.jsonl # Structured items: position, self_ref, label, text, page_numbers
```
- **`metadata.json`** — Loaded eagerly (small). Use `Path('/documents').iterdir()` to discover documents.
- **`content.txt`** — Lazy-loaded on first read. Full document text for regex or keyword search.
- **`items.jsonl`** — Lazy-loaded on first read. One JSON object per line with structured document elements. Tables are pre-rendered as markdown. Labels include `section_header`, `text`, `table`, `list_item`, `caption`, `formula`, `picture`, `code`, `footnote`, etc.
Search results include `doc_item_refs` (e.g. `["#/texts/5", "#/tables/0"]`) that match `self_ref` values in `items.jsonl`, enabling navigation from search hits to document structure.
When documents are pre-loaded via the `documents` parameter, they are also injected as a `documents` variable accessible in the sandbox code.
### Python Features
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, file I/O via `pathlib.Path`, and the `json`, `re`, `math` modules.
Not supported: most imports (only `json`, `re`, `math`, `pathlib` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
### Security
Code executes in an isolated interpreter with:
- **Virtual filesystem only**: The `/documents/` filesystem is sandboxed — no access to the real filesystem
- **No network access**: Code cannot make HTTP requests or open sockets
- **No imports**: Only `json`, `re`, `math`, and `pathlib` modules are available
- **Execution timeout**: Configurable limit (default 60s)
- **Output truncation**: Large outputs are truncated to prevent memory issues
## Context Filter
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM — both the VFS and search results are scoped to the filter:
```python
# Agent can only see documents with "confidential" in the URI
result = await client.analyze(
"Summarize all findings",
filter="uri LIKE '%confidential%'"
)
```
This is useful for scoping to specific document sets, enforcing access control, or limiting context for focused analysis.
## Configuration
Analysis settings can be configured in `haiku.rag.yaml`:
```yaml
analysis:
model:
provider: anthropic
name: claude-sonnet-4-20250514
code_timeout: 60.0 # Max seconds for code execution
max_output_chars: 50000 # Truncate output after this many chars
```

2
docs/agents/index.md
View file

@ -4,7 +4,7 @@ Three agentic flows are provided by haiku.rag:
- **Simple QA Agent** — a focused question answering agent
- **Research Graph** — a multi-step research workflow with question decomposition
- **RLM Agent** — complex analytical tasks via sandboxed Python code execution (see [RLM Agent](rlm.md))
- **Analysis Agent** — complex analytical tasks via sandboxed Python code execution (see [Analysis Agent](analysis.md))
For multi-turn conversational RAG, haiku.rag provides [skills](../skills/index.md) built on [haiku.skills](https://github.com/ggozad/haiku.skills). The skills bundle search, Q&A, analysis, and research tools with session state management.

111
docs/agents/rlm.md
View file

@ -1,111 +0,0 @@
# RLM Agent (Recursive Language Model)
The RLM agent enables complex analytical tasks by writing and executing Python code in a sandboxed environment. It solves problems that traditional RAG struggles with:
- **Aggregation**: "How many documents mention security vulnerabilities?"
- **Computation**: "What's the average revenue across all quarterly reports?"
- **Multi-document analysis**: "Compare the key findings between Report A and Report B"
- **Structured data extraction**: "Extract all dollar amounts and compute totals"
## How It Works
1. The agent receives a question
2. It writes Python code to explore the knowledge base
3. Code executes in a sandboxed Python interpreter with access to knowledge base functions
4. The agent iterates: run code, examine results, refine approach
5. Final answer is synthesized from the gathered data
## CLI Usage
```bash
# Basic usage
haiku-rag rlm "How many documents are in the database?"
# With document filter (restricts what the agent can access)
haiku-rag rlm "Summarize the key points" --filter "uri LIKE '%report%'"
# Pre-load specific documents
haiku-rag rlm "Compare these two reports" --document "Q1 Report" --document "Q2 Report"
```
## Python Usage
```python
from haiku.rag.client import HaikuRAG
async with HaikuRAG(path_to_db) as client:
# Basic question
result = await client.rlm("How many documents mention 'security'?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
# With filter (agent can only see filtered documents)
result = await client.rlm(
"What is the total revenue?",
filter="title LIKE '%Financial%'"
)
# Pre-load specific documents
result = await client.rlm(
"Compare the conclusions",
documents=["Report A", "Report B"]
)
```
## Sandbox Capabilities
The agent's code runs in a sandboxed Python interpreter ([pydantic-monty](https://github.com/pydantic/monty)) with access to these knowledge base functions:
| Function | Description |
|----------|-------------|
| `search(query, limit)` | Hybrid search (vector + full-text) returning matching chunks with scores |
| `list_documents(limit, offset)` | List documents in the knowledge base |
| `get_document(id_or_title)` | Get full text content of a document |
| `get_chunk(chunk_id)` | Get a chunk with metadata (headings, page numbers, labels) for citations |
| `get_docling_document(document_id)` | Get the DoclingDocument structure as a dict (texts, tables, pictures) |
| `llm(prompt)` | Call an LLM for classification, summarization, or extraction |
When documents are pre-loaded via the `documents` parameter, they are injected as a `documents` variable accessible in the sandbox code.
### Python Features
The interpreter supports a subset of Python: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `filter()`, `getattr()`, try/except, and the `json`, `re`, `math` modules.
Not supported: most imports (only `json`, `re`, `math` are available), class definitions, generators/yield, match statements, decorators, `with` statements. For pattern matching, the agent can use `import re`, string methods, or the `llm()` function.
### Security
Code executes in an isolated interpreter with:
- **No filesystem access**: Code cannot read or write files
- **No network access**: Code cannot make HTTP requests or open sockets
- **No imports**: Only `json`, `re`, and `math` modules are available
- **Execution timeout**: Configurable limit (default 60s)
- **Output truncation**: Large outputs are truncated to prevent memory issues
## Context Filter
The `filter` parameter restricts what documents the agent can access. Unlike tool parameters, the filter is applied automatically and cannot be bypassed by the LLM:
```python
# Agent can only see documents with "confidential" in the URI
result = await client.rlm(
"Summarize all findings",
filter="uri LIKE '%confidential%'"
)
```
This is useful for scoping to specific document sets, enforcing access control, or limiting context for focused analysis.
## Configuration
RLM settings can be configured in `haiku.rag.yaml`:
```yaml
rlm:
model:
provider: anthropic
name: claude-sonnet-4-20250514
code_timeout: 60.0 # Max seconds for code execution
max_output_chars: 50000 # Truncate output after this many chars
```

9
docs/apps.md
View file

@ -14,6 +14,12 @@ Conversational RAG from the terminal with streaming responses and session memory
```bash
haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
# Enable analysis skill (code execution)
haiku-rag chat -s rag -s analysis
# Analysis only
haiku-rag chat -s analysis
```
### Interface
@ -46,7 +52,6 @@ Press `Ctrl+P` to open the command palette:
### Session Management
- Conversation history is maintained in memory for the session
- Previous Q/A pairs are automatically used as context for follow-up questions via the `ask` tool
- Citations are tracked per response and can be inspected
- Document filter restricts all searches to selected documents
- Clearing chat resets session state
@ -61,7 +66,7 @@ Browser-based conversational RAG with a CopilotKit frontend.
- Expandable citations with source documents, pages, and headings
- Visual grounding to view chunk source locations in documents
- Document filter to restrict searches to selected documents
- Session state view for inspecting accumulated Q&A history, citations, and documents
- Session state view for inspecting citations and search results
### Quick Start

233
docs/architecture.md
View file

@ -1,233 +0,0 @@
# Architecture
High-level overview of haiku.rag components and data flow.
## System Overview
```mermaid
flowchart TB
subgraph Sources["Document Sources"]
Files[Files]
URLs[URLs]
Text[Text]
end
subgraph Processing["Processing Pipeline"]
Converter[Converter]
Chunker[Chunker]
Embedder[Embedder]
end
subgraph Storage["Storage Layer"]
LanceDB[(LanceDB)]
end
subgraph Agents["Agent Layer"]
QA[QA Agent]
Skill[RAG Skill]
Research[Research Graph]
RLM[RLM Agent]
end
subgraph Apps["Applications"]
CLI[CLI]
ChatTUI[Chat TUI]
WebApp[Web App]
Inspector[Inspector]
MCP[MCP Server]
end
Sources --> Converter
Converter --> Chunker
Chunker --> Embedder
Embedder --> LanceDB
LanceDB --> Agents
Agents --> Apps
```
## Core Components
### Storage Layer
LanceDB provides vector storage with full-text search capabilities:
- **DocumentRecord** - Document metadata and full content
- **ChunkRecord** - Text chunks with embeddings and structural metadata
- **SettingsRecord** - Database configuration and version info
Repositories handle CRUD operations:
- `DocumentRepository` - Create, read, update, delete documents
- `ChunkRepository` - Chunk management and hybrid search
- `SettingsRepository` - Configuration persistence
### Processing Pipeline
```mermaid
flowchart LR
Source[Source] --> Converter
Converter --> DoclingDoc[DoclingDocument]
DoclingDoc --> Chunker
Chunker --> Chunks[Chunks]
Chunks --> Embedder
Embedder --> Vectors[Vectors]
Vectors --> DB[(LanceDB)]
```
**Converters** transform sources into DoclingDocuments:
- `docling-local` - Local Docling processing
- `docling-serve` - Remote processing via docling-serve
**Chunkers** split documents into semantic chunks:
- Preserves document structure (tables, lists, code blocks)
- Maintains provenance (page numbers, headings)
- Configurable chunk size
**Embedders** generate vector representations:
| Provider | Models |
|----------|--------|
| Ollama | nomic-embed-text, mxbai-embed-large |
| OpenAI | text-embedding-3-small, text-embedding-3-large |
| VoyageAI | voyage-3, voyage-code-3 |
| vLLM | Any compatible model |
| LM Studio | Any compatible model |
### Agent Layer
Three agent types and a RAG skill for different use cases:
```mermaid
flowchart TB
subgraph QA["QA Agent"]
Q1[Question] --> S1[Search]
S1 --> A1[Answer]
end
subgraph Skill["RAG Skill"]
Q2[Question] --> Tools[Tool Selection]
Tools --> S2[Search / Ask / Analyze]
S2 --> A2[Answer]
A2 --> State[RAG State]
State -.-> Q2
end
subgraph Research["Research Graph"]
Q3[Question] --> Plan[Plan Next]
Plan --> SearchOne[Search One]
SearchOne --> Eval[Evaluate]
Eval -->|Continue| Plan
Eval -->|Done| Synthesize[Synthesize]
end
subgraph RLM["RLM Agent"]
Q4[Question] --> Code[Write Code]
Code --> Execute[Execute]
Execute --> Examine[Examine Results]
Examine -->|Iterate| Code
Examine -->|Done| A4[Answer]
end
```
**QA Agent** - Single-turn question answering:
- Searches for relevant chunks
- Expands context around results
- Generates answer with optional citations
**RAG Skill** - Multi-turn conversational RAG via [haiku.skills](https://github.com/ggozad/haiku.skills):
- Bundles search, list_documents, get_document, ask, analyze, and research tools
- Managed `RAGState` for session state (citations, QA history, document filters)
- Integrates with any pydantic-ai agent via `SkillToolset`
- Powers both the Chat TUI and web application
**Research Graph** - Iterative research workflow:
- Proposes one question at a time, evaluates the answer, then decides whether to continue
- Prior answers let the planner skip redundant searches
- Synthesizes structured report
**RLM Agent** - Complex analytical tasks via code execution:
- Writes Python code to explore the knowledge base
- Executes in sandboxed environment
- Handles aggregation, computation, multi-document analysis
- Iterates until answer is found
### Applications
| Application | Interface | Use Case |
|-------------|-----------|----------|
| CLI | Command line | Scripts, one-off queries, batch processing |
| Chat TUI | Terminal | Interactive conversations |
| Web App | Browser | Team collaboration, visual interface |
| Inspector | Terminal | Database exploration, debugging |
| MCP Server | Protocol | AI assistant integration |
## Data Flow
### Document Ingestion
```mermaid
sequenceDiagram
participant User
participant CLI
participant Converter
participant Chunker
participant Embedder
participant DB as LanceDB
User->>CLI: add-src document.pdf
CLI->>Converter: Convert to DoclingDocument
Converter-->>CLI: DoclingDocument
CLI->>Chunker: Split into chunks
Chunker-->>CLI: Chunks with metadata
CLI->>Embedder: Generate embeddings
Embedder-->>CLI: Vectors
CLI->>DB: Store document + chunks
DB-->>User: Document ID
```
### Search and QA
```mermaid
sequenceDiagram
participant User
participant Agent
participant Embedder
participant DB as LanceDB
participant LLM
User->>Agent: Ask question
Agent->>Embedder: Embed query
Embedder-->>Agent: Query vector
Agent->>DB: Hybrid search
DB-->>Agent: Relevant chunks
Agent->>Agent: Expand context
Agent->>LLM: Generate answer
LLM-->>Agent: Answer + citations
Agent-->>User: Response
```
## Configuration
Configuration flows through the system:
```
CLI args → Environment variables → haiku.rag.yaml → Defaults
```
Key configuration areas:
- **Storage** - Database path, vacuum settings
- **Embeddings** - Provider, model, dimensions
- **Processing** - Chunk size, converter, chunker
- **Search** - Limits, context expansion
- **QA/Research** - Model, iterations, concurrency
- **Providers** - Ollama, vLLM, docling-serve URLs
See [Configuration](configuration/index.md) for details.

21
docs/cli.md
View file

@ -165,11 +165,18 @@ Launch an interactive chat session for multi-turn conversations:
```bash
haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
# Enable analysis skill (code execution)
haiku-rag chat -s rag -s analysis
```
!!! note
Requires the `tui` extra: `pip install haiku.rag-slim[tui]` (included in full `haiku.rag` package)
Flags:
- `--skill` / `-s`: Skills to enable — `rag` (default), `analysis`. Can be repeated for multiple skills.
The chat interface provides:
- Streaming responses with real-time tool execution
@ -220,24 +227,24 @@ Flags:
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
## RLM (Recursive Language Model)
## Analyze
Answer complex analytical questions via code execution:
```bash
haiku-rag rlm "How many documents mention security?"
haiku-rag analyze "How many documents mention security?"
```
Filter to specific documents:
```bash
haiku-rag rlm "What is the total revenue?" --filter "title LIKE '%Financial%'"
haiku-rag analyze "What is the total revenue?" --filter "title LIKE '%Financial%'"
```
Pre-load specific documents for comparison:
```bash
haiku-rag rlm "Compare the conclusions" --document "Report A" --document "Report B"
haiku-rag analyze "Compare the conclusions" --document "Report A" --document "Report B"
```
Flags:
@ -245,7 +252,7 @@ Flags:
- `--filter` / `-f`: SQL WHERE clause to restrict document access
- `--document` / `-d`: Pre-load a document by title or ID (can repeat)
See [RLM Agent](agents/rlm.md) for details on capabilities and configuration.
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
## Create Skill
@ -271,7 +278,7 @@ The generated package is a pip-installable Python package that registers as a `h
### Available Tools
`analyze`, `ask`, `get_document`, `list_documents`, `research`, `search`
`cite`, `execute_code`, `get_document`, `list_documents`, `search`
### Example
@ -280,7 +287,7 @@ The generated package is a pip-installable Python package that registers as a `h
haiku-rag create-skill \
--name medic \
--db /path/to/medic.lancedb \
--tools search,ask \
--tools search,cite \
--config-file /path/to/haiku.rag.yaml \
--description "Military medic knowledge base" \
--preamble "You are a military medic expert."

1
docs/configuration/index.md
View file

@ -99,7 +99,6 @@ research:
search:
limit: 10 # Default number of results to return
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
vector_index_metric: cosine # cosine, l2, or dot
vector_refine_factor: 30

2
docs/configuration/providers.md
View file

@ -22,7 +22,7 @@ qa:
**Available options:**
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA, research, and title generation; 0.0 for RLM and picture description.
- **temperature**: Sampling temperature (0.0-1.0+). Defaults vary by task: 0.3 for QA, research, and title generation; 0.0 for analysis and picture description.
- Lower (0.0-0.3): Deterministic, focused responses
- Medium (0.4-0.7): Balanced
- Higher (0.8-1.0+): Creative, varied responses

10
docs/configuration/qa-research.md
View file

@ -7,12 +7,10 @@ Configure search behavior and context expansion:
```yaml
search:
limit: 10 # Default number of results to return
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
```
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
- **max_context_items**: Limits how many document items (paragraphs, list items, etc.) can be included in expanded context. Default: 10.
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000.
Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers) — this naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded.
@ -58,12 +56,12 @@ research:
The research workflow uses an iterative feedback loop: the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize. This continues until the planner marks research as complete or `max_iterations` is reached.
## RLM Configuration
## Analysis Configuration
Configure the RLM (Recursive Language Model) agent:
Configure the analysis agent:
```yaml
rlm:
analysis:
model:
provider: anthropic
name: claude-sonnet-4-20250514
@ -76,4 +74,4 @@ rlm:
- **code_timeout**: Maximum seconds for each code execution (default: 60)
- **max_output_chars**: Truncate code output after this many characters (default: 50000)
See [RLM Agent](../agents/rlm.md) for usage details.
See [Analysis Agent](../agents/analysis.md) for usage details.

2
docs/configuration/storage.md
View file

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

5
docs/index.md
View file

@ -8,7 +8,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Question answering** — QA agents with citations (page numbers, section headings)
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize
- **RLM agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Analysis agent** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM. QA/Research: any model supported by Pydantic AI
@ -59,13 +59,12 @@ haiku-rag chat # Interactive conversation mode
- [Getting started](tutorial.md) - Tutorial
- [Installation](installation.md) - Install haiku.rag with different providers
- [Architecture](architecture.md) - System overview and data flow
- [Configuration](configuration/index.md) - Environment variables and settings
- [CLI](cli.md) - Command line interface usage
- [Python](python.md) - Python API reference
- [Custom Pipelines](custom-pipelines.md) - Build custom processing workflows
- [Agents](agents/index.md) - QA, chat, and research agents
- [RLM Agent](agents/rlm.md) - Complex analytical tasks via code execution
- [Analysis Agent](agents/analysis.md) - Complex analytical tasks via code execution
- [Applications](apps.md) - Chat TUI, web app, and inspector
- [Server](server.md) - File monitoring and server mode
- [MCP](mcp.md) - Model Context Protocol integration

2
docs/mcp.md
View file

@ -50,7 +50,7 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients like
- `question` (required): The research question
- Returns a structured research report with findings, conclusions, and sources
- **`rlm_question`** - Answer complex analytical questions via code execution
- **`analyze`** - Answer complex analytical questions via code execution
- `question` (required): The question to answer
- `filter` (optional): SQL WHERE clause to restrict document access
- `document` (optional): Document title/ID to pre-load (can repeat)

13
docs/python.md
View file

@ -376,7 +376,6 @@ Context expansion is automatic and section-aware. For structured documents (with
Configuration:
- **search.max_context_items**: Maximum items in expanded context. Default: 10.
- **search.max_context_chars**: Maximum characters in expanded context. Default: 10000.
**Smart Merging**: When expanded results overlap within the same document, they are automatically merged into a single result with continuous content and the highest relevance score.
@ -420,32 +419,32 @@ The QA provider and model are configured in `haiku.rag.yaml` or can be passed di
See also: [Agents](agents/index.md) for details on the QA agent and the multi‑agent research workflow.
## RLM (Recursive Language Model)
## Analysis
Answer complex analytical questions via code execution:
```python
# Aggregation across documents
result = await client.rlm("Which quarter had the highest revenue?")
result = await client.analyze("Which quarter had the highest revenue?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
# Computation within a document set
result = await client.rlm(
result = await client.analyze(
"What is the average deal size mentioned in these contracts?",
filter="uri LIKE '%contracts%'"
)
# Multi-document comparison
result = await client.rlm(
result = await client.analyze(
"What changed between these two versions of the policy?",
documents=["Policy v1.0", "Policy v2.0"]
)
```
The RLM agent writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
The analysis agent writes and executes Python code in a sandboxed environment to solve problems that traditional RAG struggles with: aggregation, computation, and multi-document analysis.
See [RLM Agent](agents/rlm.md) for details on capabilities and configuration.
See [Analysis Agent](agents/analysis.md) for details on capabilities and configuration.
## Building Custom Agents

67
docs/skills/analysis.md Normal file
View file

@ -0,0 +1,67 @@
# Analysis Skill
The analysis skill provides computational analysis via code execution. It writes and runs Python code in a sandboxed interpreter to answer questions that require computation, aggregation, or data traversal.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=db_path, config=config)
```
| 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. If None, uses `get_config()`. |
## Tools
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with context expansion |
| `list_documents()` | List all documents in the knowledge base |
| `execute_code(code)` | Execute Python code in a sandboxed interpreter with VFS access |
| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer |
## State
The skill manages an `AnalysisState` under the `"analysis"` namespace:
```python
class AnalysisState(BaseModel):
document_filter: str | None = None
executions: list[CodeExecutionEntry] = []
citation_index: dict[str, Citation] = {}
citations: list[list[str]] = []
searches: dict[str, list[SearchResult]] = {}
```
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls.
- **executions** — Each `execute_code` call appends a `CodeExecutionEntry` with code, stdout, stderr, and success status.
- **citation_index** / **citations** — Same per-turn citation tracking as the RAG skill.
- **searches** — Search results from both the `search` tool and sandbox-internal searches.
## Usage with RAG Skill
Combine both skills to give the agent full RAG + analysis capabilities:
```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
from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent
rag = create_rag_skill(db_path=db_path)
analysis = create_analysis_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, analysis])
agent = Agent(
"openai:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
```
See the [Analysis Agent](../agents/analysis.md) documentation for details on how the underlying sandbox works.

8
docs/skills/index.md
View file

@ -7,7 +7,7 @@ haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggoz
| Skill | Description |
|-------|-------------|
| [`rag`](rag.md) | Search, retrieve, and answer questions from the knowledge base |
| [`rag-rlm`](rlm.md) | Computational analysis via code execution |
| [`rag-analysis`](analysis.md) | Computational analysis via code execution |
## Discovery
@ -16,7 +16,7 @@ Skills are registered as Python entrypoints under `haiku.skills`. They are disco
```bash
haiku-skills list --use-entrypoints
# rag — Search, retrieve and analyze documents using RAG.
# rag-rlm — Analyze documents using code execution in a sandboxed interpreter.
# rag-analysis — Analyze documents using code execution in a sandboxed interpreter.
```
## Usage
@ -47,7 +47,7 @@ Use `create-skill` to generate a standalone skill package with an embedded datab
haiku-rag create-skill \
--name recipes \
--db /path/to/recipes.lancedb \
--tools search,ask \
--tools search,cite \
--description "Recipe knowledge base" \
--preamble "You are a recipe expert."
```
@ -88,7 +88,7 @@ Each skill manages its own state under a dedicated namespace. State is automatic
```python
rag_state = toolset.get_namespace("rag")
rlm_state = toolset.get_namespace("rlm")
analysis_state = toolset.get_namespace("analysis")
```
See the individual skill pages for state model details.

21
docs/skills/rag.md
View file

@ -1,6 +1,6 @@
# RAG Skill
The RAG skill is the primary way to use haiku.rag tools. It bundles search, Q&A, document browsing, and research into a single skill with managed state.
The RAG skill is the primary way to use haiku.rag tools. It bundles search, document browsing, and citation management into a single skill with managed state.
## `create_skill(db_path?, config?)`
@ -20,10 +20,9 @@ skill = create_skill(db_path=db_path, config=config)
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with context expansion |
| `list_documents(limit?, offset?, filter?)` | Paginated document listing |
| `list_documents()` | List all documents in the knowledge base |
| `get_document(query)` | Retrieve a document by ID, title, or URI |
| `ask(question)` | Q&A with citations via the QA agent |
| `research(question)` | Deep multi-agent research producing comprehensive reports |
| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer |
## State
@ -31,17 +30,13 @@ The skill manages a `RAGState` under the `"rag"` namespace:
```python
class RAGState(BaseModel):
citations: list[Citation] = []
qa_history: list[QAHistoryEntry] = []
citation_index: dict[str, Citation] = {}
citations: list[list[str]] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
documents: list[DocumentInfo] = []
reports: list[ResearchEntry] = []
```
- **citations** — Accumulated citations from `ask` calls, with sequential indexing across calls.
- **qa_history** — Questions and answers from `ask` calls. Prior Q&A is used as context for follow-up questions when embeddings are similar.
- **document_filter** — SQL WHERE clause applied to `search`, `list_documents`, `ask`, and `research` calls. Set this to scope queries to specific documents.
- **citation_index** — All citations indexed by chunk ID (deduplicated across turns).
- **citations** — Per-turn lists of chunk IDs registered via the `cite` tool.
- **document_filter** — SQL WHERE clause applied to `search` and `list_documents` calls. Set this to scope queries to specific documents.
- **searches** — Search results keyed by query string.
- **documents** — Documents seen via `list_documents` or `get_document` (deduplicated by ID).
- **reports** — Research reports from `research` calls.

70
docs/skills/rlm.md
View file

@ -1,70 +0,0 @@
# RLM Skill
The RLM (Recursive Language Model) skill provides computational analysis via code execution. It writes and runs Python code in a sandboxed interpreter to answer questions that require computation, aggregation, or data traversal.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.rlm import create_skill
skill = create_skill(db_path=db_path, config=config)
```
| 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. If None, uses `get_config()`. |
## Tools
| Tool | Purpose |
|------|---------|
| `analyze(question, document?, filter?)` | Answer analytical questions using code execution |
**Parameters:**
- `question` — The analytical question to answer.
- `document` — Optional document ID or title to pre-load for analysis.
- `filter` — Optional SQL WHERE clause to filter documents.
## State
The skill manages an `RLMState` under the `"rlm"` namespace:
```python
class RLMState(BaseModel):
document_filter: str | None = None
analyses: list[AnalysisEntry] = []
class AnalysisEntry(BaseModel):
question: str
answer: str
program: str | None = None
```
- **document_filter** — SQL WHERE clause applied to `analyze` calls (combined with any explicit `filter` parameter). Set this to scope analysis to specific documents.
- **analyses** — Each `analyze` call appends an `AnalysisEntry` with the question, answer, and executed program.
## Usage with RAG Skill
Combine both skills to give the agent full RAG + analysis capabilities:
```python
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.skills.rlm import create_skill as create_rlm_skill
from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent
rag = create_rag_skill(db_path=db_path)
rlm = create_rlm_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, rlm])
agent = Agent(
"openai:gpt-4o",
instructions=build_system_prompt(toolset.skill_catalog),
toolsets=[toolset],
)
```
See the [RLM Agent](../agents/rlm.md) documentation for details on how the underlying agent works.

20
docs/tools.md
View file

@ -59,26 +59,8 @@ docs = create_document_toolset(config)
- `get_document(query)` — Retrieve a document by title or URI.
- `summarize_document(query)` — Generate an LLM summary of a document's content.
### Analysis Toolset
`create_analysis_toolset()` provides computational analysis via the RLM agent.
```python
from haiku.rag.tools import create_analysis_toolset
analysis = create_analysis_toolset(config)
```
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | `AppConfig` |
| `base_filter` | `None` | SQL WHERE clause applied to searches |
| `tool_name` | `"analyze"` | Name of the tool exposed to the agent |
## Filter Helpers
`haiku.rag.tools.filters` provides utilities for building SQL filters:
- **`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593").
- **`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic.
- **`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`.
- **`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic. Matches against both `uri` and `title`, case-insensitive.

2
docs/tuning.md
View file

@ -26,7 +26,7 @@ When configured, a cross-encoder reranker re-scores 10x the requested candidates
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa-research.md#search-settings).
Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_items` and `max_context_chars` cap expansion to prevent context bloat.
Context expansion is automatic and section-aware — search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
## Tuning Generation

1
evaluations/evaluations/benchmark.py
View file

@ -49,7 +49,6 @@ def build_experiment_metadata(
"embedder_dim": config.embeddings.model.vector_dim,
"chunk_size": config.processing.chunk_size,
"search_limit": config.search.limit,
"max_context_items": config.search.max_context_items,
"max_context_chars": config.search.max_context_chars,
"rerank_provider": config.reranking.model.provider
if config.reranking.model

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

@ -0,0 +1,21 @@
from haiku.rag.agents.analysis.agent import create_analysis_agent
from haiku.rag.agents.analysis.dependencies import AnalysisContext, AnalysisDeps
from haiku.rag.agents.analysis.models import (
AnalysisResult,
CodeExecution,
RawAnalysisResult,
)
from haiku.rag.agents.analysis.prompts import ANALYSIS_SYSTEM_PROMPT
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
__all__ = [
"ANALYSIS_SYSTEM_PROMPT",
"AnalysisContext",
"AnalysisDeps",
"RawAnalysisResult",
"AnalysisResult",
"CodeExecution",
"Sandbox",
"SandboxResult",
"create_analysis_agent",
]

59
haiku_rag_slim/haiku/rag/agents/analysis/agent.py Normal file
View file

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

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

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

16
haiku_rag_slim/haiku/rag/agents/rlm/models.py → haiku_rag_slim/haiku/rag/agents/analysis/models.py
View file

@ -1,8 +1,10 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
class CodeExecution(BaseModel):
"""Result of executing a code block in the RLM sandbox."""
"""Result of executing a code block in the analysis sandbox."""
code: str = Field(description="The Python code that was executed")
stdout: str = Field(description="Standard output captured during execution")
@ -10,8 +12,16 @@ class CodeExecution(BaseModel):
success: bool = Field(description="Whether execution completed without error")
class RLMResult(BaseModel):
"""Result from RLM agent execution."""
class RawAnalysisResult(BaseModel):
"""Raw result from the analysis agent (LLM output)."""
answer: str = Field(description="The answer to the user's question")
program: str = Field(description="The final consolidated program")
class AnalysisResult(BaseModel):
"""Result from analysis execution with resolved citations."""
answer: str
program: str
citations: list[Citation] = Field(default_factory=list)

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

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

316
haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py Normal file
View file

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

16
haiku_rag_slim/haiku/rag/agents/rlm/__init__.py
View file

@ -1,16 +0,0 @@
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.rlm.prompts import RLM_SYSTEM_PROMPT
from haiku.rag.agents.rlm.sandbox import Sandbox, SandboxResult
__all__ = [
"CodeExecution",
"RLMContext",
"RLMDeps",
"RLMResult",
"RLM_SYSTEM_PROMPT",
"Sandbox",
"SandboxResult",
"create_rlm_agent",
]

59
haiku_rag_slim/haiku/rag/agents/rlm/agent.py
View file

@ -1,59 +0,0 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.rlm.dependencies import RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.rlm.prompts import RLM_SYSTEM_PROMPT
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import get_model
def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
"""Create an RLM agent with code execution capability.
The RLM (Recursive Language Model) agent can write and execute Python code
in a sandboxed environment to solve problems that require computation,
aggregation, or complex traversal across documents.
Args:
config: Application configuration.
Returns:
A pydantic-ai Agent configured for RLM execution.
"""
model = get_model(config.rlm.model, config)
agent: Agent[RLMDeps, RLMResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model,
deps_type=RLMDeps,
output_type=RLMResult,
instructions=RLM_SYSTEM_PROMPT,
retries=3,
)
@agent.tool
async def execute_code(ctx: RunContext[RLMDeps], code: str) -> CodeExecution:
"""Execute Python code in a sandboxed interpreter.
The code has access to haiku.rag functions (search, list_documents,
get_document, get_chunk, llm).
Use print() to output results.
Args:
code: Python code to execute.
Returns:
Structured result with success status, stdout, and stderr.
"""
result = await ctx.deps.sandbox.execute(code)
execution = CodeExecution(
code=code,
stdout=result.stdout,
stderr=result.stderr,
success=result.success,
)
return execution
return agent

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

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

123
haiku_rag_slim/haiku/rag/agents/rlm/prompts.py
View file

@ -1,123 +0,0 @@
RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do.
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
- results = await search("query") ✓ CORRECT
- import search ✗ WRONG - will fail
- results = search("query") ✗ WRONG - must use await
## Available Functions
### await search(query, limit=10) -> list[dict]
Search the knowledge base using hybrid search (vector + full-text).
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
### await list_documents(limit=10, offset=0) -> list[dict]
List available documents in the knowledge base.
Returns list of dicts with keys: id, title, uri, created_at
### await get_document(id_or_title) -> str | None
Get the full text content of a document by ID, title, or URI.
Returns the document content as a string, or None if not found.
### await get_chunk(chunk_id) -> dict | None
Get a specific chunk by its ID (from search results).
Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
Use this to retrieve full chunk details and metadata for citation.
### await get_docling_document(document_id) -> dict | None
Get the full document structure as a dict (DoclingDocument format).
Use `list_documents()` or search results to get document IDs first.
- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
- `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
- `pictures`: list of figures/images with metadata
- `pages`: page dimensions and metadata
### await llm(prompt) -> str
Call an LLM directly with the given prompt. Returns the response as a string.
Use this for classification, summarization, extraction, or any task where you
already have the content and just need LLM reasoning.
## Pre-loaded Documents Variable
If documents were pre-loaded for this session, a `documents` variable is available:
```python
# documents is a list of dicts with keys: id, title, uri, content
for doc in documents:
print(doc['title'], len(doc['content']))
```
Check if it exists with: `try: documents ... except NameError: ...`
## Available Python Features
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `filter()`, `getattr()`, `sorted()`/`.sort(key=...)`, try/except, and the `json`, `re`, `math` modules.
Not supported: most imports (only `json`, `re`, `math` are available), class definitions, generators/yield, match statements, decorators, `with` statements.
For pattern matching or text extraction, use `import re`, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
## Strategy Guide
1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
## Example Patterns
### Counting documents matching a condition
```python
docs = await list_documents(limit=100)
count = 0
for doc in docs:
content = await get_document(doc['id'])
if content and 'keyword' in content.lower():
count += 1
print(f"Found in: {doc['title']}")
print(f"Total: {count}")
```
### Extracting data with regex
```python
import re
numbers = []
results = await search("financial data", limit=20)
for r in results:
amounts = re.findall(r'\\$([\\d,]+)', r['content'])
for a in amounts:
numbers.append(int(a.replace(',', '')))
if numbers:
print(f"Average: {sum(numbers) / len(numbers)}")
```
### Extracting tables from a document
```python
docs = await list_documents(limit=10)
for d in docs:
doc = await get_docling_document(d['id'])
if doc:
tables = doc.get('tables', [])
if tables:
print(f"{d['title']}: {len(tables)} table(s)")
for i, table in enumerate(tables):
grid = table.get('data', {}).get('grid', [])
for row in grid:
cells = [cell.get('text', '') for cell in row]
print(f" Table {i}: {cells}")
```
## Output Format
Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

205
haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py
View file

@ -1,205 +0,0 @@
import json
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.config.models import AppConfig
from haiku.rag.store.compression import decompress_json
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
@dataclass
class SandboxResult:
"""Result of executing code in the sandbox."""
stdout: str
stderr: str
success: bool
class Sandbox:
"""Execute code in a sandboxed Python interpreter.
Uses pydantic-monty, a minimal secure Python interpreter written in Rust.
External functions (search, list_documents, etc.) are called by Monty code
using ``await`` and resolved asynchronously on the host.
sandbox = Sandbox(client, config, context)
result = await sandbox.execute("print('hello')")
"""
_client: "HaikuRAG"
_config: AppConfig
_context: RLMContext
def __init__(
self,
client: "HaikuRAG",
config: AppConfig,
context: RLMContext,
):
self._client = client
self._config = config
self._context = context
def _build_external_functions(self) -> dict[str, Any]:
"""Build async external functions for the Monty interpreter."""
client = self._client
config = self._config
context = self._context
async def search(query: str, limit: int = 10) -> list[dict[str, Any]]:
results = await client.search(query, limit=limit, filter=context.filter)
return [
{
"chunk_id": r.chunk_id,
"content": r.content,
"document_id": r.document_id,
"document_title": r.document_title,
"document_uri": r.document_uri,
"score": r.score,
"page_numbers": r.page_numbers,
"headings": r.headings,
}
for r in results
]
async def list_documents(
limit: int = 10, offset: int = 0
) -> list[dict[str, Any]]:
docs = await client.list_documents(
limit=limit, offset=offset, filter=context.filter
)
return [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"created_at": str(d.created_at),
}
for d in docs
]
async def get_document(id_or_title: str) -> str | None:
doc = await client.resolve_document(id_or_title)
return doc.content if doc else None
async def get_chunk(chunk_id: str) -> dict[str, Any] | None:
chunk = await client.get_chunk_by_id(chunk_id)
if not chunk:
return None
meta = chunk.get_chunk_metadata()
doc_title = chunk.document_title
if not doc_title and chunk.document_id:
doc = await client.get_document_by_id(chunk.document_id)
if doc:
doc_title = doc.title
return {
"chunk_id": chunk.id,
"content": chunk.content,
"document_id": chunk.document_id,
"document_title": doc_title,
"headings": meta.headings,
"page_numbers": meta.page_numbers,
"labels": meta.labels,
}
async def get_docling_document(
document_id: str,
) -> dict[str, Any] | None:
doc = await client.get_document_by_id(document_id)
if not doc or not doc.docling_document:
return None
json_str = decompress_json(doc.docling_document)
return json.loads(json_str)
async def llm(prompt: str) -> str:
from pydantic_ai import Agent
from haiku.rag.utils import get_model
model = get_model(config.rlm.model, config)
agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt)
return result.output
return {
"search": search,
"list_documents": list_documents,
"get_document": get_document,
"get_chunk": get_chunk,
"get_docling_document": get_docling_document,
"llm": llm,
}
async def execute(self, code: str) -> SandboxResult:
"""Execute Python code in the Monty interpreter."""
external_fns = self._build_external_functions()
input_names: list[str] = []
inputs: dict[str, Any] | None = None
if self._context.documents:
input_names.append("documents")
inputs = {
"documents": [
{
"id": d.id,
"title": d.title,
"uri": d.uri,
"content": d.content,
}
for d in self._context.documents
]
}
try:
monty = pydantic_monty.Monty(
code,
inputs=input_names,
)
except (
pydantic_monty.MontySyntaxError,
pydantic_monty.MontyRuntimeError,
) as e:
return SandboxResult(stdout="", stderr=str(e), success=False)
stdout_lines: list[str] = []
def print_callback(_stream: Literal["stdout"], text: str) -> None:
stdout_lines.append(text)
max_chars = self._config.rlm.max_output_chars
limits: pydantic_monty.ResourceLimits = {
"max_duration_secs": self._config.rlm.code_timeout,
}
try:
output = await pydantic_monty.run_monty_async(
monty,
inputs=inputs,
external_functions=external_fns,
limits=limits,
print_callback=print_callback,
)
except pydantic_monty.MontyRuntimeError as e:
stdout = "".join(stdout_lines)
if len(stdout) > max_chars:
stdout = stdout[:max_chars] + "\n... (output truncated)"
return SandboxResult(stdout=stdout, stderr=str(e), success=False)
stdout = "".join(stdout_lines)
if output is not None:
stdout_with_output = f"{stdout}{output}" if stdout else str(output)
else:
stdout_with_output = stdout
if len(stdout_with_output) > max_chars:
stdout_with_output = (
stdout_with_output[:max_chars] + "\n... (output truncated)"
)
return SandboxResult(stdout=stdout_with_output, stderr="", success=True)

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

@ -446,13 +446,13 @@ class HaikuRAGApp: # pragma: no cover
for renderable in format_citations_rich(citations):
self.console.print(renderable)
async def rlm(
async def analyze(
self,
question: str,
document: str | None = None,
filter: str | None = None,
):
"""Answer a question using the RLM agent with code execution.
"""Answer a question using the analysis agent with code execution.
Args:
question: The question to answer
@ -469,10 +469,14 @@ class HaikuRAGApp: # pragma: no cover
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[dim]Running RLM agent with code execution...[/dim]")
self.console.print(
"[dim]Running analysis agent with code execution...[/dim]"
)
self.console.print()
result = await self.client.rlm(question, documents=documents, filter=filter)
result = await self.client.analyze(
question, documents=documents, filter=filter
)
self.console.print("[bold yellow]Program:[/bold yellow]")
self.console.print(Syntax(result.program, "python"))

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

@ -7,6 +7,7 @@ def run_chat(
read_only: bool = False,
before: datetime | None = None,
model: str | None = None,
skills: list[str] | None = None,
) -> None:
"""Run the chat TUI.
@ -15,6 +16,7 @@ def run_chat(
read_only: Whether to open the database in read-only mode.
before: Query database as it existed before this datetime.
model: Model to use for the chat.
skills: Skills to enable ("rag", "analysis"). Defaults to ["rag"].
"""
try:
from haiku.rag.chat.app import ChatApp
@ -24,18 +26,35 @@ def run_chat(
) from e
from haiku.rag.config import get_config
from haiku.rag.skills.rag import create_skill
from haiku.rag.utils import get_model
from haiku.rag.utils import get_model, parse_model_option
from haiku.skills.models import Skill
config = get_config()
if db_path is None:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
skill = create_skill(db_path=db_path, config=config)
if model:
model_config = parse_model_option(model)
config.qa.model = model_config
config.research.model = model_config
config.analysis.model = model_config
enabled = skills or ["rag"]
skill_list: list[Skill] = []
if "rag" in enabled:
from haiku.rag.skills.rag import create_skill as create_rag_skill
skill_list.append(create_rag_skill(db_path=db_path, config=config))
if "analysis" in enabled:
from haiku.rag.skills.analysis import create_skill as create_analysis_skill
skill_list.append(create_analysis_skill(db_path=db_path, config=config))
app = ChatApp(
db_path,
skill=skill,
skills=skill_list,
read_only=read_only,
before=before,
model=model or get_model(config.qa.model, config),

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

@ -30,6 +30,7 @@ from textual.worker import Worker
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.skills.analysis import AnalysisState
from haiku.rag.skills.rag import RAGState, get_agent_preamble
from haiku.skills.agent import (
SkillToolset,
@ -51,6 +52,7 @@ if TYPE_CHECKING:
RAG_STATE_NAMESPACE = "rag"
ANALYSIS_STATE_NAMESPACE = "analysis"
class ChatApp(App):
@ -86,14 +88,14 @@ class ChatApp(App):
def __init__(
self,
db_path: Path,
skill: Skill,
skills: list[Skill],
read_only: bool = False,
before: datetime | None = None,
model: str | None = None,
) -> None:
super().__init__()
self.db_path = db_path
self._skill = skill
self._skills = skills
self.read_only = read_only
self.before = before
self._model = model
@ -153,7 +155,7 @@ class ChatApp(App):
)
await self.client.__aenter__()
self._toolset = SkillToolset(skills=[self._skill])
self._toolset = SkillToolset(skills=self._skills)
self._agent = Agent(
self._model,
instructions=build_system_prompt(
@ -241,8 +243,7 @@ class ChatApp(App):
content=accumulated_text,
)
)
# Show citations from RAG state
await self._show_citations(chat_history)
await self._show_citations_and_programs(chat_history)
elif event.type == EventType.TOOL_CALL_START:
assert isinstance(event, ToolCallStartEvent)
chat_history.hide_thinking()
@ -318,18 +319,32 @@ class ChatApp(App):
chat_input.disabled = False
chat_input.focus()
async def _show_citations(self, chat_history: "ChatHistory") -> None:
"""Show citations from the RAG state after an agent response."""
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
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if rag_state is None:
return
citations = getattr(rag_state, "citations", [])
citations = []
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
state = self._toolset.get_namespace(namespace)
if not state:
continue
citation_turns = getattr(state, "citations", [])
citation_index = getattr(state, "citation_index", {})
if citation_turns:
latest_ids = citation_turns[-1]
for cid in latest_ids:
if cid in citation_index:
citations.append(citation_index[cid])
if citations:
# Show only new citations (since last response)
await chat_history.add_citations(citations)
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if analysis_state:
executions = getattr(analysis_state, "executions", [])
successful = [e for e in executions if e.success]
if successful:
await chat_history.add_program(successful[-1].code)
async def action_clear_chat(self) -> None:
"""Clear the chat history and reset session."""
chat_history = self.query_one(ChatHistory)
@ -420,9 +435,11 @@ class ChatApp(App):
self._document_filter = event.selected
if self._toolset:
doc_filter = build_multi_document_filter(self._document_filter)
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if isinstance(rag_state, RAGState):
rag_state.document_filter = build_multi_document_filter(
self._document_filter
)
self._state = self._toolset.build_state_snapshot()
rag_state.document_filter = doc_filter
analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
if isinstance(analysis_state, AnalysisState):
analysis_state.document_filter = doc_filter
self._state = self._toolset.build_state_snapshot()

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

@ -133,6 +133,25 @@ class CitationWidget(Collapsible):
event.stop()
class ProgramWidget(Collapsible):
"""Inline expandable program code block."""
def __init__(self, program: str, **kwargs) -> None:
content = f"```python\n{program}\n```"
super().__init__(
Markdown(content),
title="Program",
collapsed=True,
**kwargs,
)
def on_key(self, event: "Key") -> None:
"""Handle Enter to toggle expand/collapse."""
if event.key == "enter":
self.collapsed = not self.collapsed
event.stop()
class ThinkingWidget(Static):
"""Thinking indicator shown while agent is processing."""
@ -290,6 +309,22 @@ class ChatHistory(VerticalScroll):
text-style: italic;
}
/* Program */
ProgramWidget {
margin: 0 0 0 2;
background: $surface;
}
ProgramWidget > CollapsibleTitle {
padding: 0 1;
color: $text-muted;
}
ProgramWidget Contents {
padding: 1 2;
background: $panel;
}
/* Thinking indicator */
ThinkingWidget {
margin: 1 0 0 4;
@ -365,6 +400,13 @@ class ChatHistory(VerticalScroll):
await self.mount(widget)
self.scroll_end(animate=False)
async def add_program(self, program: str) -> None:
"""Add a collapsible program block after a response."""
if not program:
return
await self.mount(ProgramWidget(program))
self.scroll_end(animate=False)
async def show_thinking(self, text: str = "Thinking...") -> None:
"""Show the thinking indicator."""
try:

95
haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py
View file

@ -1,95 +0,0 @@
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.screen import ModalScreen
from textual.widgets import Button, Markdown, Static
class ContextModal(ModalScreen):
"""Modal screen for viewing session Q&A history."""
BINDINGS = [
Binding("escape", "cancel", "Close", show=False),
Binding("ctrl+o", "cancel", "Close", show=False),
]
CSS = """
ContextModal {
align: center middle;
background: rgba(0, 0, 0, 0.5);
}
#context-container {
width: 70;
height: auto;
max-height: 32;
background: $surface;
border: tall $primary;
padding: 1 2;
}
#context-header {
height: auto;
margin-bottom: 1;
}
#context-description {
height: auto;
margin-bottom: 1;
color: $text-muted;
}
#context-content {
height: 1fr;
max-height: 16;
scrollbar-gutter: stable;
}
#button-row {
height: auto;
margin-top: 1;
align: right middle;
}
#button-row Button {
margin-left: 1;
min-width: 10;
}
"""
def __init__(self, qa_history: list | None = None) -> None:
super().__init__()
self._qa_history = qa_history or []
def compose(self) -> ComposeResult:
with Vertical(id="context-container"):
yield Static("[bold]Session Context[/bold]", id="context-header")
yield Static(
"Questions and answers from this session.",
id="context-description",
)
with VerticalScroll(id="context-content"):
yield Markdown(self._get_content())
with Horizontal(id="button-row"):
yield Button("Close", id="cancel-btn", variant="primary")
def _get_content(self) -> str:
if not self._qa_history:
return "*No questions asked yet.*"
parts = []
for entry in self._qa_history:
q = getattr(entry, "question", str(entry))
a = getattr(entry, "answer", "")
parts.append(f"**Q:** {q}\n\n**A:** {a}")
return "\n\n---\n\n".join(parts)
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
if event.button.id == "cancel-btn":
self.action_cancel()
def action_cancel(self) -> None:
"""Cancel and close."""
self.app.pop_screen()

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

@ -368,8 +368,8 @@ def ask( # pragma: no cover
)
@_cli.command("rlm", help="Answer questions using code execution (RLM agent)")
def rlm( # pragma: no cover
@_cli.command("analyze", help="Answer questions using code execution (analysis agent)")
def analyze( # pragma: no cover
question: str = typer.Argument(
help="The question to answer",
),
@ -393,7 +393,7 @@ def rlm( # pragma: no cover
):
app = create_app(db)
asyncio.run(
app.rlm(
app.analyze(
question=question,
document=document,
filter=filter,
@ -645,17 +645,25 @@ def chat( # pragma: no cover
"--model",
help="Model to use for the chat (e.g. openai:gpt-4o)",
),
skill: list[str] | None = typer.Option(
None,
"--skill",
"-s",
help="Skills 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"]
run_chat(
db_path,
read_only=_read_only,
before=_before,
model=model,
skills=skills,
)

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

@ -30,11 +30,11 @@ from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import (
Citation,
ResearchReport,
)
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__)
@ -1105,7 +1105,6 @@ class HaikuRAG:
"""
from haiku.rag.context import expand_with_items
max_items = self._config.search.max_context_items
max_chars = self._config.search.max_context_chars
# Group by document_id for efficient processing
@ -1132,7 +1131,6 @@ class HaikuRAG:
self.document_item_repository,
doc_id,
doc_results,
max_items,
max_chars,
)
expanded_results.extend(expanded)
@ -1192,17 +1190,17 @@ class HaikuRAG:
return await graph.run(state=state, deps=deps)
async def rlm(
async def analyze(
self,
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "RLMResult":
"""Answer a question using the RLM agent with code execution.
) -> "AnalysisResult":
"""Answer a question using the analysis agent with code execution.
The RLM (Recursive Language Model) agent can write and execute Python
code in a sandboxed environment to solve problems that require
computation, aggregation, or complex traversal across documents.
The analysis agent can write and execute Python code in a sandboxed
environment to solve problems that require computation, aggregation,
or complex traversal across documents.
Args:
question: The question to answer.
@ -1210,16 +1208,16 @@ class HaikuRAG:
filter: SQL WHERE clause to filter documents during searches.
Returns:
RLMResult with the answer and the final consolidated program.
AnalysisResult with the answer and the final consolidated program.
"""
from haiku.rag.agents.rlm import (
RLMContext,
RLMDeps,
from haiku.rag.agents.analysis import (
AnalysisContext,
AnalysisDeps,
Sandbox,
create_rlm_agent,
create_analysis_agent,
)
context = RLMContext(filter=filter)
context = AnalysisContext(filter=filter)
if documents:
loaded_docs = []
@ -1230,26 +1228,51 @@ class HaikuRAG:
context.documents = loaded_docs if loaded_docs else None
sandbox = Sandbox(
client=self,
db_path=self.store.db_path,
config=self._config,
context=context,
)
deps = RLMDeps(
deps = AnalysisDeps(
sandbox=sandbox,
context=context,
)
agent = create_rlm_agent(self._config)
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import Citation
agent = create_analysis_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output
output = result.output
seen: set[str] = set()
citations: list[Citation] = []
for sr in sandbox._search_results:
if sr.chunk_id and sr.chunk_id not in seen:
seen.add(sr.chunk_id)
citations.append(
Citation(
index=len(seen),
document_id=sr.document_id or "",
chunk_id=sr.chunk_id,
document_uri=sr.document_uri or "",
document_title=sr.document_title,
page_numbers=sr.page_numbers,
headings=sr.headings,
content=sr.content,
)
)
return AnalysisResult(
answer=output.answer,
program=output.program,
citations=citations,
)
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.
Gets the DoclingDocument from the chunk's document, resolves bounding boxes
from chunk metadata, and renders all pages that contain bounding boxes with
yellow/orange highlight overlays.
Expands the chunk's context to find the full section, then resolves
bounding boxes from all items in the expanded range. This ensures
visualization covers all pages the expanded content spans.
Args:
chunk: The chunk to visualize.
@ -1262,6 +1285,8 @@ class HaikuRAG:
from PIL import ImageDraw
from haiku.rag.store.models.chunk import ChunkMetadata
# Get the document structure (from cache if available)
if not chunk.document_id:
return []
@ -1274,9 +1299,23 @@ class HaikuRAG:
if not docling_doc:
return []
# Resolve bounding boxes from chunk metadata
# Expand context to get all doc_item_refs in the section
chunk_meta = chunk.get_chunk_metadata()
bounding_boxes = chunk_meta.resolve_bounding_boxes(docling_doc)
if chunk_meta.doc_item_refs:
search_result = SearchResult(
content=chunk.content,
score=1.0,
chunk_id=chunk.id,
document_id=chunk.document_id,
doc_item_refs=chunk_meta.doc_item_refs,
page_numbers=chunk_meta.page_numbers,
)
expanded = await self.expand_context([search_result])
refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
meta = ChunkMetadata(doc_item_refs=refs)
else:
meta = chunk_meta
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
if not bounding_boxes:
return []

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

@ -96,7 +96,7 @@ class ResearchConfig(BaseModel):
max_concurrency: int = 1
class RLMConfig(BaseModel):
class AnalysisConfig(BaseModel):
model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
@ -174,7 +174,6 @@ class ProcessingConfig(BaseModel):
class SearchConfig(BaseModel):
limit: int = 10
max_context_items: int = 10
max_context_chars: int = 10000
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
vector_refine_factor: int = 30
@ -219,7 +218,7 @@ class AppConfig(BaseModel):
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
qa: QAConfig = Field(default_factory=QAConfig)
research: ResearchConfig = Field(default_factory=ResearchConfig)
rlm: RLMConfig = Field(default_factory=RLMConfig)
analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)

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

@ -6,12 +6,16 @@ the document_items table. The algorithm adapts to document structure:
For STRUCTURED documents (containing section_header or title labels):
1. Resolve matched doc_item_refs to positions in the items table
2. Find section boundaries around each match (section_header/title labels)
3. If the section fits within the budget, include it entirely
4. If the section exceeds the budget, OR the section is too small (under
20% of max_context_chars), expand item-by-item from the match center
outward, skipping noise labels. This lets small sections (e.g., a
title+authors area) grow into the next section's content.
5. Merge overlapping ranges from multiple results in the same document
3. If the section fits within the char budget, include it entirely
4. If the section exceeds the char budget, expand item-by-item from the
match center outward, bounded by section edges
5. If the section is too small (under 20% of max_context_chars), expand
item-by-item crossing into adjacent sections until the budget is filled.
This lets small sections (e.g., title+authors) grow into neighboring
content.
6. Merge overlapping ranges from multiple results in the same document.
Adjacent but non-overlapping ranges stay separate to preserve section
independence.
For UNSTRUCTURED documents (no section headers):
Expand outward item-by-item from the match center until the character
@ -20,7 +24,6 @@ For UNSTRUCTURED documents (no section headers):
In both cases:
- max_context_chars caps total characters per expanded result
- max_context_items caps total items per expanded result
- Noise labels (footnote, page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents
- Results without doc_item_refs pass through unexpanded
@ -42,7 +45,7 @@ _MIN_SECTION_BUDGET_RATIO = 0.2
def _merge_ranges(
ranges: list[tuple[int, int, SearchResult]],
) -> list[tuple[int, int, list[SearchResult]]]:
"""Merge overlapping or adjacent ranges."""
"""Merge overlapping ranges. Adjacent but non-overlapping ranges stay separate."""
if not ranges:
return []
@ -55,7 +58,7 @@ def _merge_ranges(
)
for min_idx, max_idx, result in sorted_ranges[1:]:
if cur_max >= min_idx - 1: # Overlapping or adjacent
if cur_max >= min_idx: # Truly overlapping
cur_max = max(cur_max, max_idx)
cur_results.append(result)
else:
@ -69,27 +72,32 @@ def _merge_ranges(
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
max_items: int,
max_chars: int,
skip_noise: bool = False,
lo_bound: int = 0,
hi_bound: int | None = None,
) -> tuple[int, int]:
"""Expand item-by-item outward from center until budget is filled.
"""Expand item-by-item outward from center until char budget is filled.
When skip_noise is True, noise labels are excluded from char counting
(used in structured documents so footnotes don't consume budget).
lo_bound and hi_bound constrain expansion (e.g., to section edges).
"""
if hi_bound is None:
hi_bound = len(items) - 1
lo = hi = center_idx
center_is_noise = skip_noise and items[center_idx].label in _NOISE_LABELS
char_count = 0 if center_is_noise else len(items[center_idx].text)
while char_count < max_chars and hi - lo + 1 < max_items:
while char_count < max_chars:
grew = False
if lo > 0:
if lo > lo_bound:
lo -= 1
if not (skip_noise and items[lo].label in _NOISE_LABELS):
char_count += len(items[lo].text)
grew = True
if hi < len(items) - 1 and char_count < max_chars:
if hi < hi_bound and char_count < max_chars:
hi += 1
if not (skip_noise and items[hi].label in _NOISE_LABELS):
char_count += len(items[hi].text)
@ -104,7 +112,6 @@ def _find_expansion_range(
items: list[DocumentItem],
matched_positions: set[int],
has_sections: bool,
max_items: int,
max_chars: int,
) -> tuple[int, int]:
"""Find the expansion range for matched positions within a window of items."""
@ -113,7 +120,7 @@ def _find_expansion_range(
center_idx = matched_indices[len(matched_indices) // 2]
if not has_sections:
return _expand_outward(items, center_idx, max_items, max_chars)
return _expand_outward(items, center_idx, max_chars)
# Build section spans: [(start_idx, end_idx), ...]
headers = [
@ -140,23 +147,34 @@ def _find_expansion_range(
if items[i].label not in _NOISE_LABELS
)
# Section fits nicely in the budget — return it as-is
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
if min_useful <= sec_chars <= max_chars and sec_end - sec_start + 1 <= max_items:
if sec_chars <= max_chars and sec_chars >= min_useful:
# Section fits in char budget — return it regardless of item count
return (items[sec_start].position, items[sec_end].position)
# Section is too large or too small — expand item-by-item from center.
# For too-large sections this stays within budget.
# For too-small sections (e.g., title+authors) this naturally grows
# into adjacent sections until the budget is filled.
return _expand_outward(items, center_idx, max_items, max_chars, skip_noise=True)
if sec_chars > max_chars:
# Section too large — expand outward bounded by section edges
return _expand_outward(
items,
center_idx,
max_chars,
skip_noise=True,
lo_bound=sec_start,
hi_bound=sec_end,
)
# Section too small (e.g., title+authors) — expand across boundaries
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
_WINDOW_MARGIN = 100
async def expand_with_items(
document_item_repository: DocumentItemRepository,
document_id: str,
results: list[SearchResult],
max_items: int,
max_chars: int,
) -> list[SearchResult]:
"""Expand results using the document_items table."""
@ -172,7 +190,7 @@ async def expand_with_items(
# wide enough to find section boundaries (the nearest section_header/title
# above and below the match).
all_positions = sorted(ref_positions.values())
window_margin = max_items * 10
window_margin = _WINDOW_MARGIN
window_start = max(0, min(all_positions) - window_margin)
window_end = max(all_positions) + window_margin
window_items = await document_item_repository.get_items_in_range(
@ -194,9 +212,7 @@ async def expand_with_items(
passthrough.append(result)
continue
lo, hi = _find_expansion_range(
window_items, matched, has_sections, max_items, max_chars
)
lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
ranges.append((lo, hi, result))
merged = _merge_ranges(ranges)

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

@ -183,12 +183,12 @@ def create_mcp_server(
return None
@mcp.tool()
async def rlm_question(
async def analyze(
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex questions using code execution (RLM agent).
"""Answer complex questions using code execution (analysis agent).
Use this for questions requiring computation, aggregation, or
complex traversal across documents. The agent can write Python
@ -205,9 +205,9 @@ def create_mcp_server(
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
documents = [document] if document else None
result = await rag.rlm(question, documents=documents, filter=filter)
result = await rag.analyze(question, documents=documents, filter=filter)
return result.answer
except Exception as e:
return f"Error running RLM agent: {e!s}"
return f"Error running analysis agent: {e!s}"
return mcp

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

@ -8,9 +8,8 @@ AVAILABLE_TOOLS: set[str] = {
"list_documents",
"get_document",
"search",
"ask",
"research",
"analyze",
"execute_code",
"cite",
}
DEFAULT_PREAMBLE = (

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

@ -7,48 +7,58 @@ description: {{ description }}
{{ preamble }}
## How to decide which tool to use
{% if "ask" in tool_names %}
## Tools
{% if "search" in tool_names %}
**Default rule:** If the user is asking a question, use **ask**. Only use **search** when the user explicitly wants to browse or find passages.
### 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.
{% endif %}
{% if "list_documents" in tool_names %}
- **list_documents** — Use when the user wants to browse or see what documents are available (e.g., "what documents do you have?", "show me the documents", "list available docs").
### list_documents
List all documents in the knowledge base.
{% endif %}
{% if "get_document" in tool_names %}
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
### get_document
Retrieve a document by ID, title, or URI. Partial matches work.
{% endif %}
{% if "search" in tool_names %}
- **search** — Use when the user wants to browse, explore, or find specific passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns all matching results as sources.
{% if "execute_code" in tool_names %}
### execute_code
Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, `await llm()`, and a virtual filesystem at `/documents/` with document content and structure.
{% endif %}
{% if "ask" in tool_names %}
- **ask** — Use for factual questions that need a synthesized answer (e.g., "what is DocLayNet?", "explain the methodology"). Searches, synthesizes, and returns only the chunks actually used as citations. Always include the citations in your response.
{% endif %}
{% if "research" in tool_names %}
- **research** — Deep multi-agent research that produces comprehensive reports. **Only use when the user explicitly requests deep research** (e.g., "do a deep research on X", "research this topic thoroughly"). Never call this tool on your own — it is slow and expensive.
{% endif %}
{% if "analyze" in tool_names %}
- **analyze** — Use for complex analytical questions that require computation, aggregation, or data traversal across documents (e.g., "how many pages?", "compare table 3 across documents", "calculate average word count"). Executes Python code in a sandboxed interpreter.
{% 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 %}
## When search returns irrelevant results
## How to answer questions
If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
{% if "ask" in tool_names %}
- Use **ask** if the question is factual
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 %}
- Report that the knowledge base doesn't contain relevant information
{% 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]":
- Extract the **topic** as the `query`/`question` parameter
- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
Examples:
- "search for embeddings in the ML paper" -> first identify "ML paper", then search for "embeddings"
- "what does the DocLayNet paper say about annotations?" -> ask with question="what are the annotation methods?"
- Use **get_document** or **list_documents** first to identify the document
- Then search for the topic
{% endif %}

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

@ -5,23 +5,14 @@ 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 "ask" in tool_names or "research" in tool_names %}
{% if "cite" in tool_names %}
from haiku.rag.agents.research.models import Citation
{% endif %}
{% if "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
from haiku.rag.tools.document import DocumentInfo
{% endif %}
{% if "ask" in tool_names %}
from haiku.rag.tools.qa import QAHistoryEntry
{% endif %}
{% if "search" in tool_names %}
from haiku.rag.store.models.chunk import SearchResult
{% endif %}
{% if "research" in tool_names %}
from haiku.rag.skills._tools import ResearchEntry
{% endif %}
{% if "analyze" in tool_names %}
from haiku.rag.skills._tools import AnalysisEntry
{% if "execute_code" in tool_names %}
from haiku.rag.skills._tools import CodeExecutionEntry
{% endif %}
_TOOL_NAMES = {{ tool_names | tojson }}
@ -36,24 +27,16 @@ _CONFIG_PATH = _ASSETS_DIR / "haiku.rag.yaml"
class SkillState(BaseModel):
{% if "ask" in tool_names or "research" in tool_names %}
citations: list[Citation] = Field(default_factory=list)
{% endif %}
{% if "ask" in tool_names %}
qa_history: list[QAHistoryEntry] = Field(default_factory=list)
{% 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 "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
documents: list[DocumentInfo] = Field(default_factory=list)
{% endif %}
{% if "research" in tool_names %}
reports: list[ResearchEntry] = Field(default_factory=list)
{% endif %}
{% if "analyze" in tool_names %}
analyses: list[AnalysisEntry] = Field(default_factory=list)
{% if "execute_code" in tool_names %}
executions: list[CodeExecutionEntry] = Field(default_factory=list)
{% endif %}

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

@ -7,59 +7,14 @@ from pydantic_ai import RunContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.filters import combine_filters
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.state import SkillRunDeps
class ResearchEntry(BaseModel):
question: str
title: str
executive_summary: str
class AnalysisEntry(BaseModel):
question: str
answer: str
program: str | None = None
async def find_relevant_prior_qa(
qa_history: list[QAHistoryEntry],
query: str,
config: AppConfig,
) -> list[QAHistoryEntry]:
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.utils import cosine_similarity
if not qa_history:
return []
embedder = get_embedder(config)
query_embedding = await embedder.embed_query(query)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
qa_history[idx].question_embedding = new_embeddings[i]
matches = []
for qa in qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(query_embedding, qa.question_embedding)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matches.append(qa)
return matches
class CodeExecutionEntry(BaseModel):
code: str
stdout: str
stderr: str = ""
success: bool = True
async def skill_search(
@ -89,14 +44,12 @@ async def skill_search(
async def skill_list_documents(
db_path: Path,
config: AppConfig,
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)
documents = await rag.list_documents(filter=filter)
return [
{
"id": doc.id,
@ -132,119 +85,26 @@ async def skill_get_document(
}
async def skill_ask(
db_path: Path,
config: AppConfig,
question: str,
qa_history: list[QAHistoryEntry] | None = None,
document_filter: str | None = None,
) -> tuple[str, list[Citation]]:
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import format_citations
ask_question = question
if qa_history:
matches = await find_relevant_prior_qa(qa_history, question, config)
if matches:
prior_parts = []
for qa in matches:
part = f"Q: {qa.question}\nA: {qa.answer}"
if qa.citations:
part += "\n" + format_citations(qa.citations)
prior_parts.append(part)
ask_question = (
"Context from prior questions in this session:\n\n"
+ "\n\n---\n\n".join(prior_parts)
+ "\n\n---\n\nCurrent question: "
+ question
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
answer, citations = await rag.ask(
ask_question,
filter=document_filter,
)
return answer, citations
async def skill_research(
db_path: Path,
config: AppConfig,
question: str,
document_filter: str | None = None,
) -> tuple[str, str, str]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
report = await rag.research(question, filter=document_filter)
parts = [
f"# {report.title}",
f"\n## Executive Summary\n{report.executive_summary}",
]
if report.main_findings:
parts.append("\n## Main Findings")
for finding in report.main_findings:
parts.append(f"- {finding}")
if report.conclusions:
parts.append("\n## Conclusions")
for conclusion in report.conclusions:
parts.append(f"- {conclusion}")
if report.limitations:
parts.append("\n## Limitations")
for limitation in report.limitations:
parts.append(f"- {limitation}")
if report.recommendations:
parts.append("\n## Recommendations")
for rec in report.recommendations:
parts.append(f"- {rec}")
parts.append(f"\n## Sources\n{report.sources_summary}")
formatted = "\n".join(parts)
return formatted, report.title, report.executive_summary
async def skill_analyze(
db_path: Path,
config: AppConfig,
question: str,
document: str | None = None,
filter: str | None = None,
) -> tuple[str, str, str | None]:
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = [document] if document else None
result = await rag.rlm(question, documents=documents, filter=filter)
output = result.answer
if result.program:
output += f"\n\nProgram:\n{result.program}"
return output, result.answer, result.program
def update_documents_state(
documents_state: list[DocumentInfo],
doc_dicts: list[dict[str, Any]],
) -> None:
for doc_dict in doc_dicts:
doc_info = DocumentInfo(
id=str(doc_dict["id"]),
title=doc_dict["title"] or "Untitled",
uri=doc_dict.get("uri") or "",
created=doc_dict.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in documents_state):
documents_state.append(doc_info)
def _get_state(ctx: RunContext[SkillRunDeps], 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 _register_citations(state: Any, citations: "list[Citation]") -> None:
"""Add citations to the index and record the turn's chunk IDs."""
chunk_ids = []
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
chunk_ids.append(cid)
state.citations.append(chunk_ids)
def create_skill_extras(
db_path: Path,
config: AppConfig,
@ -313,6 +173,8 @@ def create_skill_tools(
tools: dict[str, Any] = {}
if "search" in tool_names:
max_searches = config.qa.max_searches
search_counts: dict[str, int] = {}
async def search(
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
@ -325,6 +187,14 @@ def create_skill_tools(
query: The search query.
limit: Maximum number of results.
"""
rid = ctx.run_id or ""
search_counts[rid] = search_counts.get(rid, 0) + 1
if search_counts[rid] > max_searches:
return (
"Search limit reached. Answer the question using "
"the results you already have."
)
state = _get_state(ctx, state_type)
formatted, results = await skill_search(
db_path,
@ -343,25 +213,14 @@ def create_skill_tools(
async def list_documents(
ctx: RunContext[SkillRunDeps],
limit: int | None = None,
offset: int | None = None,
) -> list[dict[str, Any]]:
"""List documents in the knowledge base with optional pagination.
Args:
limit: Maximum number of documents to return.
offset: Number of documents to skip.
"""
"""List all documents in the knowledge base."""
state = _get_state(ctx, state_type)
result = await skill_list_documents(
db_path,
config,
limit,
offset,
filter=state.document_filter if state else None,
)
if state:
update_documents_state(state.documents, result)
return result
tools["list_documents"] = list_documents
@ -376,124 +235,86 @@ def create_skill_tools(
Args:
query: Document ID, title, or URI to look up.
"""
result = await skill_get_document(db_path, config, query)
if result is not None:
state = _get_state(ctx, state_type)
if state:
update_documents_state(state.documents, [result])
return result
return await skill_get_document(db_path, config, query)
tools["get_document"] = get_document
if "ask" in tool_names:
if "execute_code" in tool_names:
async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Ask a question and get an answer with citations from the knowledge base.
async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
"""Execute Python code in a sandboxed interpreter.
The code has access to search(), list_documents(), llm() functions
and a virtual filesystem at /documents/ with document content and
structure (metadata.json, content.txt, items.jsonl per document).
Use print() to output results. Each call runs in a fresh
interpreter — variables do not persist between calls.
Args:
question: The question to ask.
code: Python code to execute.
"""
from haiku.rag.utils import format_citations
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
state = _get_state(ctx, state_type)
answer, citations = await skill_ask(
db_path,
config,
question,
qa_history=state.qa_history if state else None,
document_filter=state.document_filter if state else None,
)
doc_filter = state.document_filter if state else None
context = AnalysisContext(filter=doc_filter)
sandbox = Sandbox(db_path=db_path, config=config, context=context)
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:
next_index = len(state.citations) + 1
for citation in citations:
citation.index = next_index
next_index += 1
state.citations.extend(citations)
state.qa_history.append(
QAHistoryEntry(
question=question, answer=answer, citations=citations
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[SkillRunDeps], chunk_ids: list[str]) -> str:
"""Register chunk IDs as citations for your answer.
Call this after searching, with the chunk_id values from search
results that support your answer.
Args:
chunk_ids: List of chunk_id values from search results.
"""
from haiku.rag.agents.research.models import resolve_citations
state = _get_state(ctx, state_type)
if not state:
return "No state available."
all_results = []
for results_list in state.searches.values():
all_results.extend(results_list)
citations = resolve_citations(chunk_ids, all_results)
if citations:
answer += "\n\n" + format_citations(citations)
_register_citations(state, citations)
return f"Registered {len(citations)} citation(s)."
return answer
tools["ask"] = ask
if "research" in tool_names:
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Conduct deep multi-agent research on a question.
Iteratively searches, analyzes, and synthesizes information from the
knowledge base to produce a comprehensive research report.
Only use when the user explicitly requests deep research.
Args:
question: The research question to investigate.
"""
state = _get_state(ctx, state_type)
formatted, title, executive_summary = await skill_research(
db_path,
config,
question,
document_filter=state.document_filter if state else None,
)
if state:
state.reports.append(
ResearchEntry(
question=question,
title=title,
executive_summary=executive_summary,
)
)
state.qa_history.append(
QAHistoryEntry(question=question, answer=executive_summary)
)
return formatted
tools["research"] = research
if "analyze" in tool_names:
async def analyze(
ctx: RunContext[SkillRunDeps],
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex analytical questions using code execution.
Use this for questions requiring computation, aggregation, or
data traversal across documents.
Args:
question: The question to answer.
document: Optional document ID or title to pre-load for analysis.
filter: Optional SQL WHERE clause to filter documents.
"""
state = _get_state(ctx, state_type)
state_filter = state.document_filter if state else None
effective_filter = combine_filters(state_filter, filter)
output, answer, program = await skill_analyze(
db_path, config, question, document=document, filter=effective_filter
)
if state:
state.analyses.append(
AnalysisEntry(
question=question,
answer=answer,
program=program,
)
)
return output
tools["analyze"] = analyze
tools["cite"] = cite
return tools

28
haiku_rag_slim/haiku/rag/skills/rlm.py → haiku_rag_slim/haiku/rag/skills/analysis.py
View file

@ -2,23 +2,28 @@ import os
from functools import cache
from pathlib import Path
from pydantic import BaseModel
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.skills._tools import AnalysisEntry
from haiku.rag.skills._tools import CodeExecutionEntry
from haiku.rag.store.models.chunk import SearchResult
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
class RLMState(BaseModel):
class AnalysisState(BaseModel):
document_filter: str | None = None
analyses: list[AnalysisEntry] = []
executions: list[CodeExecutionEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
STATE_TYPE = RLMState
STATE_NAMESPACE = "rlm"
STATE_TYPE = AnalysisState
STATE_NAMESPACE = "analysis"
_skill_path = Path(__file__).parent / "rag-rlm"
_skill_path = Path(__file__).parent / "rag-analysis"
@cache
@ -45,7 +50,7 @@ def create_skill(
db_path: Path | None = None,
config: AppConfig | None = None,
) -> Skill:
"""Create an RLM analysis skill for computational document analysis.
"""Create an analysis skill for computational document analysis.
Args:
db_path: Path to the LanceDB database. Resolved from:
@ -67,7 +72,12 @@ def create_skill(
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
tools = create_skill_tools(db_path, config, RLMState, ["analyze"])
tools = create_skill_tools(
db_path,
config,
AnalysisState,
["search", "list_documents", "execute_code", "cite"],
)
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()

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

@ -0,0 +1,101 @@
---
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 solve complex analytical questions by writing and executing Python code against the knowledge base.
## Tools
### execute_code
Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — write self-contained code. Use `print()` to output results.
Inside the code, these functions are available (use `await`):
- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
- `await list_documents()` → list of dicts with keys: id, title, uri, created_at
- `await llm(prompt)` → string response from an LLM (for classification, summarization, extraction)
Available modules: `json`, `re`, `math`, `pathlib`
Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
### search
Search the knowledge base directly (outside code execution). Use for initial exploration before writing code.
### list_documents
List available documents. Use to discover what's in the knowledge base.
### cite
Register chunk IDs as citations. Call after your analysis with chunk_id values from search results that support your answer.
## Document Filesystem (inside execute_code)
All documents are mounted as a virtual filesystem at `/documents/`:
```
/documents/{document_id}/
metadata.json # {"id", "title", "uri", "created_at"}
content.txt # Full document text
items.jsonl # Structured items (one JSON object per line)
```
### Reading files
Always use `Path.read_text()` — do NOT use `open()` or `with` statements (they are not supported).
```python
from pathlib import Path
import json
# Discover documents
for doc_dir in Path('/documents').iterdir():
meta = json.loads((doc_dir / 'metadata.json').read_text())
print(meta['title'])
# Read full text
content = Path(f'/documents/{doc_id}/content.txt').read_text()
# Read and parse items
for line in Path(f'/documents/{doc_id}/items.jsonl').read_text().strip().split(chr(10)):
item = json.loads(line)
if item['label'] == 'table':
print(item['text'][:200])
```
### metadata.json
Document metadata: `id`, `title`, `uri`, `created_at`.
### content.txt
Full text content. Use for regex or keyword search across a whole document.
### items.jsonl
Structured document items. Each line is a JSON object with:
- `position`: sequential position in the document
- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote"
- `text`: rendered content (tables are markdown with `|` columns)
- `page_numbers`: list of page numbers where the item appears
### 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.
## Strategy
1. Use `search` tool first to understand what's in the knowledge base
2. Use `execute_code` to write analysis code
3. Iterate: run code, examine output, refine approach
4. Call `cite` with chunk IDs from search results you referenced
## Important
- Each `execute_code` call runs in a fresh interpreter — write self-contained code blocks
- Use `print()` to output results — the output is your only feedback
- Always execute code to answer questions — don't just describe what code would do
- Use `await` for all async functions inside execute_code (search, list_documents, llm)
- Use `Path.read_text()` to read files — do NOT use `open()`, `with` statements, or `collections` module
- Do NOT include chunk IDs or UUIDs in your answer text — use the `cite` tool separately

13
haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md
View file

@ -1,13 +0,0 @@
---
name: rag-rlm
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".
---
# RLM Analysis
Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in a sandboxed Python interpreter.

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

@ -6,10 +6,7 @@ from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.skills._tools import ResearchEntry
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@ -21,7 +18,7 @@ CRITICAL RULES:
3. When a skill returns citations, always include them in your response
"""
_RAG_TOOLS = ["search", "list_documents", "get_document", "ask", "research"]
_RAG_TOOLS = ["search", "list_documents", "get_document", "cite"]
def get_agent_preamble(config: AppConfig) -> str:
@ -32,12 +29,10 @@ def get_agent_preamble(config: AppConfig) -> str:
class RAGState(BaseModel):
citations: list[Citation] = Field(default_factory=list)
qa_history: list[QAHistoryEntry] = Field(default_factory=list)
citation_index: dict[str, Citation] = Field(default_factory=dict)
citations: list[list[str]] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
documents: list[DocumentInfo] = Field(default_factory=list)
reports: list[ResearchEntry] = Field(default_factory=list)
STATE_TYPE = RAGState

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

@ -5,31 +5,55 @@ description: Search, retrieve and analyze documents using RAG (Retrieval Augment
# RAG
You are a RAG (Retrieval Augmented Generation) assistant with access to a document knowledge base.
You are a RAG assistant with access to a document knowledge base.
Use your tools to search and answer questions. Never make up information — always use tools to get facts from the knowledge base.
## How to decide which tool to use
## Tools
**Default rule:** If the user is asking a question, use **ask**. Only use **search** when the user explicitly wants to browse or find passages.
### search
Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content.
- **list_documents** — Use when the user wants to browse or see what documents are available (e.g., "what documents do you have?", "show me the documents", "list available docs").
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
- **search** — Use when the user wants to browse, explore, or find specific passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns all matching results as sources.
- **ask** — Use for factual questions that need a synthesized answer (e.g., "what is DocLayNet?", "explain the methodology"). Searches, synthesizes, and returns only the chunks actually used as citations. Always include the citations in your response.
- **research** — Deep multi-agent research that produces comprehensive reports. **Only use when the user explicitly requests deep research** (e.g., "do a deep research on X", "research this topic thoroughly"). Never call this tool on your own — it is slow and expensive.
Each result includes:
- `chunk_id` in brackets and rank position (rank 1 = most relevant)
- Source: document title and section hierarchy
- Type: content type (paragraph, table, code, list_item)
- Content: the actual text
## When search returns irrelevant results
### list_documents
List available documents in the knowledge base. Use when the user wants to browse what's available.
If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
- Use **ask** if the question is factual
- Report that the knowledge base doesn't contain relevant information
### get_document
Retrieve a document by ID, title, or URI. Partial matches work. Use when the user wants the full content of a specific document.
### cite
Register chunk IDs as citations for your answer. Call this AFTER formulating your answer, with the `chunk_id` values from search results that support it.
## How to answer questions
1. Call `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. Synthesize a concise answer based strictly on the retrieved content
5. Call `cite` with the chunk IDs of search results that informed your answer
## Guidelines
- Base answers strictly on retrieved content — do not use external knowledge
- Use the Source and Type metadata to understand context
- If multiple results are relevant, synthesize them coherently
- Be concise and direct — avoid elaboration unless asked
- 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.
- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
## When the user mentions a specific document
If the user says "search in [doc]", "find in [doc]", or "answer from [doc]":
- Extract the **topic** as the `query`/`question` parameter
- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
- Use `get_document` or `list_documents` first to identify the document
- Then search for the topic
Examples:
- "search for embeddings in the ML paper" -> first identify "ML paper", then search for "embeddings"
- "what does the DocLayNet paper say about annotations?" -> ask with question="what are the annotation methods?"
## When search returns irrelevant results
If your first search returns results that clearly don't match the question:
- Try one more search with different keywords
- If still irrelevant, report that the knowledge base doesn't contain relevant information

14
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -100,6 +100,20 @@ class DocumentRepository:
return self._record_to_document(results[0])
async def get_content(self, entity_id: str) -> str | None:
"""Get only the text content of a document (skips docling blobs)."""
safe_id = escape_sql_string(entity_id)
results = list(
self.store.documents_table.search()
.select(["content"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not results:
return None
return results[0]["content"]
_DOCLING_COLUMNS = ["id", "docling_document", "docling_version"]
async def get_docling_data(self, entity_id: str) -> Document | None:

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

@ -40,6 +40,44 @@ class DocumentItemRepository:
]
self.store.document_items_table.add(records)
async def get_all_items(self, document_id: str) -> list[DocumentItem]:
"""Get all items for a document, sorted by position."""
safe_id = escape_sql_string(document_id)
rows = (
self.store.document_items_table.search()
.where(f"document_id = '{safe_id}'")
.to_list()
)
items = [self._record_to_item(row) for row in rows]
items.sort(key=lambda x: x.position)
return items
async def get_all_items_grouped(
self, document_ids: list[str] | None = None
) -> dict[str, list[DocumentItem]]:
"""Get all items grouped by document_id in a single query.
Args:
document_ids: If provided, only fetch items for these documents.
If None, fetches all items.
Returns:
Dict mapping document_id to sorted list of DocumentItem.
"""
query = self.store.document_items_table.search()
if document_ids is not None:
safe_ids = ", ".join(f"'{escape_sql_string(did)}'" for did in document_ids)
query = query.where(f"document_id IN ({safe_ids})")
rows = query.to_list()
grouped: dict[str, list[DocumentItem]] = {}
for row in rows:
item = self._record_to_item(row)
grouped.setdefault(item.document_id, []).append(item)
for items in grouped.values():
items.sort(key=lambda x: x.position)
return grouped
async def get_items_in_range(
self, document_id: str, start: int, end: int
) -> list[DocumentItem]:

14
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

@ -1,23 +1,11 @@
from haiku.rag.tools.analysis import AnalysisResult, create_analysis_toolset
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
)
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
from haiku.rag.tools.filters import build_multi_document_filter
from haiku.rag.tools.search import create_search_toolset
__all__ = [
"AnalysisResult",
"PRIOR_ANSWER_RELEVANCE_THRESHOLD",
"QAHistoryEntry",
"RAGDeps",
"build_document_filter",
"build_multi_document_filter",
"combine_filters",
"create_analysis_toolset",
"create_document_toolset",
"create_search_toolset",
]

87
haiku_rag_slim/haiku/rag/tools/analysis.py
View file

@ -1,87 +0,0 @@
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.sandbox import Sandbox
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import (
build_document_filter,
combine_filters,
)
class AnalysisResult(BaseModel):
"""Result from the analysis toolset (RLM execution)."""
answer: str = Field(description="The answer produced by analysis")
code_executed: bool = Field(
default=True,
description="Whether code was executed to produce this answer",
)
def create_analysis_toolset(
config: AppConfig,
base_filter: str | None = None,
tool_name: str = "analyze",
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with code analysis capabilities via RLM agent.
Args:
config: Application configuration.
base_filter: Optional base SQL WHERE clause applied to searches.
tool_name: Name for the analyze tool. Defaults to "analyze".
Returns:
FunctionToolset with an analyze tool.
"""
async def analyze( # pragma: no cover
ctx: RunContext[RAGDeps],
task: str,
document_name: str | None = None,
) -> AnalysisResult:
"""Execute a computational task via code execution.
Uses the RLM (Recursive Language Model) agent to write and execute
Python code to answer the task.
Args:
task: A specific, actionable instruction describing what to compute.
document_name: Optional document name/title to focus on.
Returns:
AnalysisResult with answer and execution metadata.
"""
client = ctx.deps.client
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(base_filter, doc_filter)
rlm_context = RLMContext(filter=effective_filter)
sandbox = Sandbox(
client=client,
config=config,
context=rlm_context,
)
deps = RLMDeps(
sandbox=sandbox,
context=rlm_context,
)
rlm_agent = create_rlm_agent(config)
result = await rlm_agent.run(task, deps=deps)
program = result.output.program
return AnalysisResult(
answer=result.output.answer,
code_executed=bool(program),
)
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(analyze, name=tool_name)
return toolset

17
haiku_rag_slim/haiku/rag/tools/filters.py
View file

@ -1,4 +1,4 @@
def build_document_filter(document_name: str) -> str:
def _build_document_filter(document_name: str) -> str:
"""Build SQL filter for document name matching.
Matches against both uri and title fields, case-insensitive.
@ -19,20 +19,7 @@ def build_multi_document_filter(document_names: list[str]) -> str | None:
"""
if not document_names:
return None
filters = [build_document_filter(name) for name in document_names]
filters = [_build_document_filter(name) for name in document_names]
if len(filters) == 1:
return filters[0]
return " OR ".join(f"({f})" for f in filters)
def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
"""Combine two SQL filters with AND logic.
Returns None if both filters are None.
"""
filters = [f for f in [filter1, filter2] if f]
if not filters:
return None
if len(filters) == 1:
return filters[0]
return f"({filters[0]}) AND ({filters[1]})"

32
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -1,32 +0,0 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation, SearchAnswer
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
class QAHistoryEntry(BaseModel):
"""A Q&A pair with optional cached embedding for similarity matching."""
question: str
answer: str
confidence: float = 0.9
citations: list[Citation] = Field(default_factory=list)
question_embedding: list[float] | None = Field(default=None, exclude=True)
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
def to_search_answer(self) -> SearchAnswer:
"""Convert to SearchAnswer for research graph context."""
return SearchAnswer(
query=self.question,
answer=self.answer,
confidence=self.confidence,
cited_chunks=[c.chunk_id for c in self.citations],
citations=self.citations,
)

2
haiku_rag_slim/pyproject.toml
View file

@ -62,7 +62,7 @@ vertexai = ["pydantic-ai-slim[vertexai]"]
[project.entry-points."haiku.skills"]
rag = "haiku.rag.skills.rag:create_skill"
rag-rlm = "haiku.rag.skills.rlm:create_skill"
rag-analysis = "haiku.rag.skills.analysis:create_skill"
[project.scripts]
haiku-rag = "haiku.rag.cli:cli"

5
mkdocs.yml
View file

@ -59,7 +59,6 @@ nav:
- index.md
- Getting started: tutorial.md
- Installation: installation.md
- Architecture: architecture.md
- Configuration:
- configuration/index.md
- Providers: configuration/providers.md
@ -73,11 +72,11 @@ nav:
- Tuning: tuning.md
- Agents:
- agents/index.md
- RLM Agent: agents/rlm.md
- Analysis Agent: agents/analysis.md
- Skills:
- skills/index.md
- RAG: skills/rag.md
- RLM: skills/rlm.md
- Analysis: skills/analysis.md
- Toolsets: tools.md
- Applications: apps.md
- Server: server.md

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

13
tests/agents/rlm/conftest.py → tests/agents/analysis/conftest.py
View file

@ -1,7 +1,7 @@
import pytest
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import Sandbox
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
@ -14,8 +14,9 @@ async def empty_client(temp_db_path):
@pytest.fixture
async def sandbox(empty_client):
async def sandbox(temp_db_path):
"""Create a Monty sandbox for testing."""
config = AppConfig()
context = RLMContext()
return Sandbox(client=empty_client, config=config, context=context)
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
context = AnalysisContext()
return Sandbox(db_path=temp_db_path, config=config, context=context)

73
tests/agents/rlm/test_agent.py → tests/agents/analysis/test_agent.py
View file

@ -3,26 +3,26 @@ from pathlib import Path
import pytest
from pydantic_ai import Agent
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.analysis.agent import create_analysis_agent
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import CodeExecution, RawAnalysisResult
from haiku.rag.config import AppConfig, Config
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_rlm")
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_analysis")
class TestCreateRLMAgent:
class TestCreateAnalysisAgent:
def test_creates_agent(self):
agent = create_rlm_agent(Config)
agent = create_analysis_agent(Config)
assert isinstance(agent, Agent)
assert agent.deps_type is RLMDeps
assert agent.output_type is RLMResult
assert agent.deps_type is AnalysisDeps
assert agent.output_type is RawAnalysisResult
def test_agent_has_execute_code_tool(self):
agent = create_rlm_agent(Config)
agent = create_analysis_agent(Config)
tool_names = list(agent._function_toolset.tools.keys())
assert "execute_code" in tool_names
@ -42,13 +42,13 @@ class TestCodeExecutionModel:
assert execution.success is True
class TestClientRLMIntegration:
"""Integration tests for client.rlm() method."""
class TestClientAnalysisIntegration:
"""Integration tests for client.analyze() method."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_count_documents(self, allow_model_requests, temp_db_path):
"""Test RLM agent can count documents.
async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
"""Test analysis agent can count documents.
Agent program:
docs = list_documents(limit=1000)
@ -63,14 +63,14 @@ class TestClientRLMIntegration:
await client.create_document("Second document about dogs.", title="Doc 2")
await client.create_document("Third document about birds.", title="Doc 3")
result = await client.rlm("How many documents are in the database?")
result = await client.analyze("How many documents are in the database?")
assert "3" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_aggregation(self, allow_model_requests, temp_db_path):
"""Test RLM agent can perform aggregation across documents.
async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
"""Test analysis agent can perform aggregation across documents.
Agent program:
import re
@ -103,7 +103,7 @@ class TestClientRLMIntegration:
"Sales report Q3: Revenue was $200,000.", title="Q3 Report"
)
result = await client.rlm(
result = await client.analyze(
"What is the total revenue across all quarterly reports?"
)
@ -111,8 +111,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_with_filter(self, allow_model_requests, temp_db_path):
"""Test RLM agent respects filter parameter.
async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
"""Test analysis agent respects filter parameter.
Agent program:
docs = list_documents(limit=1000)
@ -130,7 +130,7 @@ class TestClientRLMIntegration:
await client.create_document("Dog document.", title="Dogs")
await client.create_document("Bird document.", title="Birds")
result = await client.rlm(
result = await client.analyze(
"How many documents are available?",
filter="title = 'Cats'",
)
@ -139,15 +139,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_search_and_get_chunk(self, allow_model_requests, temp_db_path):
"""Test RLM agent can search and use get_chunk for citations.
Agent program:
results = search("content", limit=5)
for r in results:
chunk = get_chunk(r['chunk_id'])
print(chunk['document_title'], chunk['chunk_id'])
"""
async def test_analyze_search_and_identify_source(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can search and identify source documents."""
from haiku.rag.client import HaikuRAG
config = AppConfig()
@ -158,7 +153,7 @@ class TestClientRLMIntegration:
title="Animal Facts",
)
result = await client.rlm(
result = await client.analyze(
"Search for content about animals and tell me "
"which document it came from."
)
@ -167,10 +162,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_semantic_analysis_with_llm(
async def test_analyze_semantic_analysis_with_llm(
self, allow_model_requests, temp_db_path
):
"""Test RLM agent can use llm() for semantic analysis combined with computation.
"""Test analysis agent can use llm() for semantic analysis combined with computation.
Agent program:
docs = list_documents(limit=100)
@ -209,7 +204,7 @@ class TestClientRLMIntegration:
title="Q3 Update",
)
result = await client.rlm(
result = await client.analyze(
"Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?"
)
@ -220,8 +215,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test RLM agent can use search() to find content and extract information.
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test analysis agent can use search() to find content and extract information.
Agent program:
results = search("document element types", limit=20)
@ -242,7 +237,7 @@ class TestClientRLMIntegration:
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path)
result = await client.rlm(
result = await client.analyze(
"Search for content about document element types or labels. "
"What are all the different document element types mentioned? "
"List them all."
@ -277,10 +272,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_with_preloaded_documents(
async def test_analyze_with_preloaded_documents(
self, allow_model_requests, temp_db_path
):
"""Test RLM agent can use pre-loaded documents variable.
"""Test analysis agent can use pre-loaded documents variable.
Agent program:
if 'documents' in dir():
@ -303,7 +298,7 @@ class TestClientRLMIntegration:
title="Mission Statement",
)
result = await client.rlm(
result = await client.analyze(
"Using the pre-loaded documents variable, "
"tell me when was the company founded and what is their mission?",
documents=["Company History", "Mission Statement"],

6
tests/agents/rlm/test_models.py → tests/agents/analysis/test_models.py
View file

@ -1,4 +1,4 @@
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
class TestCodeExecution:
@ -25,8 +25,8 @@ class TestCodeExecution:
assert "ZeroDivisionError" in execution.stderr
class TestRLMResult:
class TestAnalysisResult:
def test_create_result(self):
result = RLMResult(answer="The answer is 42", program="print(42)")
result = AnalysisResult(answer="The answer is 42", program="print(42)")
assert result.answer == "The answer is 42"
assert result.program == "print(42)"

374
tests/agents/rlm/test_sandbox.py → tests/agents/analysis/test_sandbox.py
View file

@ -2,8 +2,8 @@ from pathlib import Path
import pytest
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import Sandbox, SandboxResult
from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.agents.analysis.sandbox import Sandbox, SandboxResult
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import Document
@ -74,12 +74,12 @@ class TestSandboxErrors:
assert result.stderr != ""
class TestSandboxHaikuRAG:
"""Test haiku.rag functions in sandbox."""
class TestSandboxListDocuments:
"""Test list_documents function in sandbox."""
@pytest.mark.asyncio
async def test_list_documents_empty(self, sandbox):
"""Test list_documents returns empty list for empty database."""
"""list_documents returns empty list for empty database."""
result = await sandbox.execute(
"docs = await list_documents()\nprint(type(docs).__name__, len(docs))"
)
@ -89,7 +89,7 @@ class TestSandboxHaikuRAG:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_with_data(self, temp_db_path):
"""Test list_documents returns documents when populated."""
"""list_documents returns documents when populated."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
@ -98,8 +98,8 @@ class TestSandboxHaikuRAG:
title="Test Document",
)
context = RLMContext()
sb = Sandbox(client=client, config=config, context=context)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"docs = await list_documents()\n"
"print(len(docs))\n"
@ -109,6 +109,10 @@ class TestSandboxHaikuRAG:
assert "1" in result.stdout
assert "Test Document" in result.stdout
class TestSandboxSearch:
"""Test search function in sandbox."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_search_with_data(self, temp_db_path):
@ -121,8 +125,8 @@ class TestSandboxHaikuRAG:
title="Animals",
)
context = RLMContext()
sb = Sandbox(client=client, config=config, context=context)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('fox', limit=5)\n"
"print(len(results))\n"
@ -134,68 +138,52 @@ class TestSandboxHaikuRAG:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_get_document(self, temp_db_path):
"""Test get_document function."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
)
context = RLMContext()
sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute(
f"content = await get_document('{doc.id}')\n"
"print('foxes' in content.lower() if content else 'None')"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
async def test_get_document_not_found(self, sandbox):
"""Test get_document returns None for missing document."""
result = await sandbox.execute(
"content = await get_document('nonexistent-id')\nprint(content is None)"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_get_chunk(self, temp_db_path):
"""Test get_chunk function returns chunk with metadata."""
async def test_search_returns_doc_item_refs_and_labels(self, temp_db_path):
"""Search results include doc_item_refs and labels."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
content="The quick brown fox jumps over the lazy dog.",
uri="test://animals",
title="Animals",
)
context = RLMContext()
sb = Sandbox(client=client, config=config, context=context)
# First search to get a chunk_id
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('foxes', limit=1)\n"
"chunk_id = results[0]['chunk_id']\n"
"chunk = await get_chunk(chunk_id)\n"
"print(chunk['document_title'])\n"
"print('content' in chunk)"
"results = await search('fox', limit=1)\n"
"r = results[0]\n"
"print('doc_item_refs' in r)\n"
"print('labels' in r)\n"
"print(type(r['doc_item_refs']).__name__)\n"
"print(type(r['labels']).__name__)"
)
assert result.success
assert "Fox Document" in result.stdout
assert "True" in result.stdout
assert "True\nTrue" in result.stdout
assert "list\nlist" in result.stdout
@pytest.mark.asyncio
async def test_get_chunk_not_found(self, sandbox):
"""Test get_chunk returns None for missing chunk."""
result = await sandbox.execute(
"chunk = await get_chunk('nonexistent-id')\nprint(chunk is None)"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.vcr()
async def test_search_returns_expanded_content(self, temp_db_path):
"""search() returns context-expanded results."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="The quick brown fox jumps over the lazy dog.",
uri="test://animals",
title="Animals",
)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"results = await search('fox', limit=1)\n"
"print(type(results[0]['content']).__name__)\n"
"print('fox' in results[0]['content'].lower())"
)
assert result.success
assert "str" in result.stdout
assert "True" in result.stdout
class TestSandboxExternalFunctionEdgeCases:
@ -254,38 +242,151 @@ class TestSandboxOutputTruncation:
"""Test output truncation behavior."""
@pytest.mark.asyncio
async def test_truncate_stdout_on_runtime_error(self, empty_client):
async def test_truncate_stdout_on_runtime_error(self, temp_db_path):
"""Test stdout is truncated when a runtime error occurs after large output."""
config = AppConfig()
config.rlm.max_output_chars = 20
context = RLMContext()
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('a' * 100)\nx = 1/0")
assert not result.success
assert "ZeroDivisionError" in result.stderr
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute("print('a' * 100)\nx = 1/0")
assert not result.success
assert "ZeroDivisionError" in result.stderr
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
@pytest.mark.asyncio
async def test_truncate_successful_output(self, empty_client):
async def test_truncate_successful_output(self, temp_db_path):
"""Test output is truncated on successful execution with large output."""
config = AppConfig()
config.rlm.max_output_chars = 20
context = RLMContext()
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('b' * 100)")
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
config.analysis.max_output_chars = 20
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute("print('b' * 100)")
assert result.success
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
class TestSandboxVFS:
"""Test virtual filesystem for document access."""
@pytest.mark.asyncio
async def test_empty_database_has_no_documents(self, sandbox):
"""Empty database has no document directories."""
result = await sandbox.execute(
"from pathlib import Path\nprint(Path('/documents').exists())"
)
assert result.success
assert result.stdout.endswith("... (output truncated)")
assert len(result.stdout) < 100
class TestSandboxContextFilter:
"""Test context filter is applied."""
# /documents dir may or may not exist when empty, both are valid
# The key is it doesn't error
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_filter_applied_to_list_documents(self, temp_db_path):
"""Test that context filter is passed to list_documents."""
async def test_iterdir_discovers_documents(self, temp_db_path):
"""Path('/documents').iterdir() lists document directories."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
content="Test content",
uri="test://doc1",
title="Test Document",
)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"dirs = list(Path('/documents').iterdir())\n"
"print(len(dirs))\n"
"print(dirs[0].is_dir())"
)
assert result.success
assert "1" in result.stdout
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_metadata_json(self, temp_db_path):
"""metadata.json contains document title and uri."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Test content",
uri="test://doc1",
title="Test Document",
)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
f"meta = json.loads(Path('/documents/{doc.id}/metadata.json').read_text())\n"
"print(meta['title'])\n"
"print(meta['uri'])"
)
assert result.success
assert "Test Document" in result.stdout
assert "test://doc1" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_content_txt(self, temp_db_path):
"""content.txt returns full document text (lazy loaded)."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Content about foxes and dogs.",
uri="test://doc",
title="Fox Document",
)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
f"content = Path('/documents/{doc.id}/content.txt').read_text()\n"
"print('foxes' in content.lower())"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_items_jsonl(self, temp_db_path):
"""items.jsonl returns document items as JSONL (lazy loaded)."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="The quick brown fox jumps over the lazy dog.",
uri="test://animals",
title="Animals",
)
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"from pathlib import Path\n"
"import json\n"
f"text = Path('/documents/{doc.id}/items.jsonl').read_text()\n"
"lines = text.strip().split('\\n')\n"
"print(len(lines) > 0)\n"
"item = json.loads(lines[0])\n"
"print('position' in item)\n"
"print('self_ref' in item)\n"
"print('label' in item)\n"
"print('text' in item)\n"
"print('page_numbers' in item)"
)
assert result.success
assert result.stdout.count("True") == 6
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_context_filter_limits_vfs(self, temp_db_path):
"""Context filter restricts which documents appear in VFS."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
await client.create_document(
@ -299,13 +400,15 @@ class TestSandboxContextFilter:
title="Private Doc",
)
context = RLMContext(filter="uri LIKE 'public://%'")
sb = Sandbox(client=client, config=config, context=context)
context = AnalysisContext(filter="uri LIKE 'public://%'")
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"docs = await list_documents()\n"
"print(len(docs))\n"
"if docs:\n"
" print(docs[0]['title'])"
"from pathlib import Path\n"
"import json\n"
"dirs = list(Path('/documents').iterdir())\n"
"print(len(dirs))\n"
"meta = json.loads((dirs[0] / 'metadata.json').read_text())\n"
"print(meta['title'])"
)
assert result.success
assert "1" in result.stdout
@ -324,61 +427,25 @@ class TestSandboxPreloadedDocuments:
assert "NameError" in result.stderr
@pytest.mark.asyncio
async def test_documents_variable_available_with_preload(self, empty_client):
async def test_documents_variable_available_with_preload(self, temp_db_path):
"""documents variable is available when context.documents is set."""
config = AppConfig()
docs = [
Document(id="1", content="Content A", title="Doc A", uri="a://1"),
Document(id="2", content="Content B", title="Doc B", uri="b://2"),
]
context = RLMContext(documents=docs)
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute(
"print(len(documents))\n"
"print(documents[0]['title'])\n"
"print(documents[1]['title'])"
)
assert result.success
assert "2" in result.stdout
assert "Doc A" in result.stdout
assert "Doc B" in result.stdout
class TestSandboxDoclingDocument:
"""Test get_docling_document() external function."""
@pytest.mark.asyncio
async def test_returns_none_for_missing_document(self, sandbox):
"""get_docling_document returns None for a non-existent document."""
result = await sandbox.execute(
"doc = await get_docling_document('nonexistent-id')\nprint(doc is None)"
)
assert result.success
assert "True" in result.stdout
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_returns_dict_for_document_with_docling_data(self, temp_db_path):
"""get_docling_document returns a dict for a document with docling data."""
config = AppConfig()
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(
content="Docling processed content",
uri="test://docling",
title="Docling Doc",
)
context = RLMContext()
sb = Sandbox(client=client, config=config, context=context)
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
docs = [
Document(id="1", content="Content A", title="Doc A", uri="a://1"),
Document(id="2", content="Content B", title="Doc B", uri="b://2"),
]
context = AnalysisContext(documents=docs)
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
f"doc = await get_docling_document('{doc.id}')\n"
"print(type(doc).__name__)\n"
"print(doc['name'])\n"
"print('texts' in doc)"
"print(len(documents))\n"
"print(documents[0]['title'])\n"
"print(documents[1]['title'])"
)
assert result.success
assert "dict" in result.stdout
assert "True" in result.stdout
assert "2" in result.stdout
assert "Doc A" in result.stdout
assert "Doc B" in result.stdout
class TestSandboxLLM:
@ -386,14 +453,15 @@ class TestSandboxLLM:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_llm_function(self, allow_model_requests, empty_client):
async def test_llm_function(self, allow_model_requests, temp_db_path):
"""Test llm() calls the model and returns a string."""
config = AppConfig()
context = RLMContext()
sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute(
"answer = await llm('What is 2 + 2? Reply with just the number.')\n"
"print(answer)"
)
assert result.success
assert "4" in result.stdout
async with HaikuRAG(temp_db_path, create=True):
config = AppConfig()
context = AnalysisContext()
sb = Sandbox(db_path=temp_db_path, config=config, context=context)
result = await sb.execute(
"answer = await llm('What is 2 + 2? Reply with just the number.')\n"
"print(answer)"
)
assert result.success
assert "4" in result.stdout

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_aggregation.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_aggregation.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_count_documents.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_count_documents.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_search_and_extract.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_search_and_extract.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_search_and_get_chunk.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_search_and_identify_source.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_semantic_analysis_with_llm.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_semantic_analysis_with_llm.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_with_filter.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_with_filter.yaml
View file

0
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_with_preloaded_documents.yaml → tests/cassettes/test_analysis/TestClientAnalysisIntegration.test_analyze_with_preloaded_documents.yaml
View file

42
tests/cassettes/test_sandbox/TestSandboxDoclingDocument.test_returns_dict_for_document_with_docling_data.yaml
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_get_chunk.yaml
View file

File diff suppressed because one or more lines are too long

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_list_documents_with_data.yaml → tests/cassettes/test_sandbox/TestSandboxListDocuments.test_list_documents_with_data.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_search_with_data.yaml → tests/cassettes/test_sandbox/TestSandboxSearch.test_search_returns_doc_item_refs_and_labels.yaml
View file

82
tests/cassettes/test_sandbox/TestSandboxSearch.test_search_returns_expanded_content.yaml Normal file
View file

File diff suppressed because one or more lines are too long

82
tests/cassettes/test_sandbox/TestSandboxSearch.test_search_with_data.yaml Normal file
View file

File diff suppressed because one or more lines are too long

0
tests/cassettes/test_sandbox/TestSandboxHaikuRAG.test_get_document.yaml → tests/cassettes/test_sandbox/TestSandboxVFS.test_content_txt.yaml
View file

0
tests/cassettes/test_sandbox/TestSandboxContextFilter.test_filter_applied_to_list_documents.yaml → tests/cassettes/test_sandbox/TestSandboxVFS.test_context_filter_limits_vfs.yaml
View file

82
tests/cassettes/test_sandbox/TestSandboxVFS.test_items_jsonl.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestSandboxVFS.test_iterdir_discovers_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

42
tests/cassettes/test_sandbox/TestSandboxVFS.test_metadata_json.yaml Normal file
View file

File diff suppressed because one or more lines are too long

4
tests/chat/test_chat_app.py
View file

@ -57,7 +57,7 @@ def _make_app(db_path: Path, mock_client: AsyncMock | None = None):
return ChatApp(
db_path=db_path,
skill=skill,
skills=[skill],
read_only=True,
), mock_client
@ -80,7 +80,7 @@ def _make_app_with_state(db_path: Path, mock_client: AsyncMock | None = None):
return ChatApp(
db_path=db_path,
skill=skill,
skills=[skill],
read_only=True,
), mock_client

194
tests/skills/test_analysis.py Normal file
View file

@ -0,0 +1,194 @@
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_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", "list_documents", "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):
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)
result = await execute_code(ctx, code="print('hello')")
assert "hello" in result
async def test_execute_code_updates_state(self, rag_db):
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)
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):
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)
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_applies_document_filter(self, rag_db):
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)
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):
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)
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_execute_code_vfs_write_denied(self, rag_db):
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)
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

462
tests/skills/test_rag.py
View file

@ -1,7 +1,3 @@
from unittest.mock import AsyncMock
from haiku.rag.agents.research.models import Citation, ResearchReport
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rag import (
STATE_NAMESPACE,
@ -12,8 +8,6 @@ from haiku.rag.skills.rag import (
state_metadata,
)
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
@ -116,13 +110,7 @@ class TestRAGSkillCreation:
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",
"list_documents",
"get_document",
"ask",
"research",
}
assert tool_names == {"search", "list_documents", "get_document", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill
@ -139,8 +127,6 @@ class TestRAGSkillCreation:
assert skill.extras["db_path"] is temp_db_path
assert "visualize_chunk" in skill.extras
assert "list_documents" in skill.extras
assert callable(skill.extras["visualize_chunk"])
assert callable(skill.extras["list_documents"])
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
@ -169,44 +155,6 @@ class TestSkillExtras:
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
async def test_visualize_chunk_unknown_returns_empty(
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)
visualize = skill.extras["visualize_chunk"]
result = await visualize("nonexistent-chunk-id")
assert result == []
async def test_visualize_chunk_returns_images(
self,
test_app_config,
rag_db,
monkeypatch,
):
from haiku.rag.client import HaikuRAG
from haiku.rag.skills.rag import create_skill
monkeypatch.setattr(
HaikuRAG, "visualize_chunk", AsyncMock(return_value=["img1"])
)
skill = create_skill(config=test_app_config, db_path=rag_db)
visualize = skill.extras["visualize_chunk"]
# Get a real chunk_id from the db
async with HaikuRAG(rag_db, read_only=True) as rag:
docs = await rag.list_documents()
doc = await rag.get_document_by_id(docs[0].id)
chunks = await rag.chunk_repository.get_by_document_id(doc.id)
chunk_id = str(chunks[0].id)
result = await visualize(chunk_id)
assert result == ["img1"]
class TestSearchTool:
async def test_search_returns_formatted_string(self, rag_db):
@ -252,6 +200,23 @@ class TestSearchTool:
result = await search(ctx, query="artificial intelligence")
assert isinstance(result, str)
async def test_search_rate_limited(self, rag_db):
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)
ctx.run_id = "test-run"
await search(ctx, query="first")
await search(ctx, query="second")
result = await search(ctx, query="third")
assert "Search limit reached" in result
assert len(state.searches) == 2
class TestListDocumentsTool:
async def test_list_documents_returns_results(self, rag_db):
@ -264,18 +229,6 @@ class TestListDocumentsTool:
assert isinstance(results, list)
assert len(results) == 2
async def test_list_documents_updates_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
state = RAGState()
ctx = _make_ctx(state)
await list_docs(ctx)
assert len(state.documents) == 2
assert isinstance(state.documents[0], DocumentInfo)
assert state.documents[0].id is not None
async def test_list_documents_applies_document_filter_from_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
@ -287,17 +240,6 @@ class TestListDocumentsTool:
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
async def test_list_documents_no_duplicates_in_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
state = RAGState()
ctx = _make_ctx(state)
await list_docs(ctx)
await list_docs(ctx)
assert len(state.documents) == 2
class TestGetDocumentTool:
async def test_get_document_by_title(self, rag_db):
@ -310,18 +252,6 @@ class TestGetDocumentTool:
assert result is not None
assert result["title"] == "AI Overview"
async def test_get_document_updates_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
state = RAGState()
ctx = _make_ctx(state)
await get_doc(ctx, query="AI Overview")
assert len(state.documents) == 1
assert isinstance(state.documents[0], DocumentInfo)
assert state.documents[0].title == "AI Overview"
async def test_get_document_not_found(self, rag_db):
from haiku.rag.skills.rag import create_skill
@ -332,343 +262,57 @@ class TestGetDocumentTool:
assert result is None
class TestAskTool:
async def test_ask_returns_answer_with_citations(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
document_title="AI Overview",
content="AI is transforming industries.",
)
]
monkeypatch.setattr(
HaikuRAG,
"ask",
AsyncMock(return_value=("AI transforms industries worldwide.", citations)),
)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
ctx = _make_ctx()
result = await ask(ctx, question="What is AI?")
assert isinstance(result, str)
assert "AI transforms industries" in result
async def test_ask_updates_state(self, rag_db, monkeypatch):
class TestCiteTool:
async def test_cite_registers_citations(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
content="AI content",
)
]
monkeypatch.setattr(
HaikuRAG,
"ask",
AsyncMock(return_value=("AI transforms industries.", citations)),
)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state)
await ask(ctx, question="What is AI?")
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) == 1
assert len(state.qa_history) == 1
assert isinstance(state.qa_history[0], QAHistoryEntry)
assert state.qa_history[0].question == "What is AI?"
assert len(state.citations[0]) == 2
assert all(cid in state.citation_index for cid in chunk_ids)
async def test_ask_assigns_citation_indices(self, rag_db, monkeypatch):
async def test_cite_deduplicates_in_index(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
first_citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://doc1",
content="First.",
),
Citation(
document_id="d2",
chunk_id="c2",
document_uri="test://doc2",
content="Second.",
),
]
second_citations = [
Citation(
document_id="d3",
chunk_id="c3",
document_uri="test://doc3",
content="Third.",
),
]
call_count = 0
async def mock_ask(self, question, **kwargs):
nonlocal call_count
call_count += 1
if call_count == 1:
return ("Answer 1", first_citations)
return ("Answer 2", second_citations)
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
search = _get_tool(skill, "search")
cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state)
await ask(ctx, question="First question")
assert state.citations[0].index == 1
assert state.citations[1].index == 2
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 ask(ctx, question="Second question")
assert state.citations[2].index == 3
await cite(ctx, chunk_ids=chunk_ids)
await cite(ctx, chunk_ids=chunk_ids)
assert len(state.citation_index) == 1
assert len(state.citations) == 2
async def test_ask_applies_document_filter_from_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
captured_kwargs = {}
async def mock_ask(self, question, **kwargs):
captured_kwargs.update(kwargs)
return ("Answer.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
await ask(ctx, question="What is AI?")
assert captured_kwargs.get("filter") == "title = 'AI Overview'"
async def test_ask_includes_prior_qa_context(self, rag_db, monkeypatch):
import random
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
captured_questions = []
async def mock_ask(self, question, **kwargs):
captured_questions.append(question)
return ("Answer about AI.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
# Pre-compute the embedding the fake embedder will produce for "Tell me about AI"
query_text = "Tell me about AI"
random.seed(hash(query_text) % (2**32))
query_embedding = [random.random() for _ in range(VECTOR_DIM)]
prior_citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
document_title="AI Overview",
content="AI content from source.",
)
]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is artificial intelligence?",
answer="AI is the simulation of human intelligence by machines.",
question_embedding=query_embedding,
citations=prior_citations,
),
]
)
ctx = _make_ctx(state)
await ask(ctx, question=query_text)
# rag.ask() should receive augmented question with prior context
assert len(captured_questions) == 1
augmented = captured_questions[0]
assert "Context from prior questions" in augmented
assert "What is artificial intelligence?" in augmented
assert "AI is the simulation" in augmented
assert "AI Overview" in augmented
assert query_text in augmented
# State should store the original question, not the augmented one
assert state.qa_history[-1].question == query_text
async def test_ask_embeds_prior_qa_on_demand(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
captured_questions = []
async def mock_ask(self, question, **kwargs):
captured_questions.append(question)
return ("Answer about AI.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
# Use the same question text for the prior QA entry and query so
# their fake embeddings are identical (cosine similarity = 1.0).
prior_question = "Tell me about AI"
query_text = prior_question
# Leave question_embedding=None to exercise the lazy embedding path
state = RAGState(
qa_history=[
QAHistoryEntry(
question=prior_question,
answer="AI is the simulation of human intelligence by machines.",
question_embedding=None,
),
]
)
ctx = _make_ctx(state)
await ask(ctx, question=query_text)
# The lazy embedding should have populated question_embedding
assert state.qa_history[0].question_embedding is not None
assert len(state.qa_history[0].question_embedding) == VECTOR_DIM
# The augmented question should include prior context
assert len(captured_questions) == 1
assert "Context from prior questions" in captured_questions[0]
assert prior_question in captured_questions[0]
async def test_ask_no_prior_qa_context_when_irrelevant(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
captured_questions = []
async def mock_ask(self, question, **kwargs):
captured_questions.append(question)
return ("Answer.", [])
monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
# Use orthogonal embedding — won't match the fake embedder's output
orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is the weather?",
answer="It is sunny today.",
question_embedding=orthogonal,
),
]
)
ctx = _make_ctx(state)
await ask(ctx, question="Explain quantum computing")
# rag.ask() should receive the original question unchanged
assert len(captured_questions) == 1
assert captured_questions[0] == "Explain quantum computing"
class TestResearchTool:
async def test_research_returns_report(self, rag_db, monkeypatch):
async def test_cite_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
report = ResearchReport(
title="AI Research",
executive_summary="AI is transforming industries.",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Multiple sources consulted.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
ctx = _make_ctx()
result = await research(ctx, question="What is AI?")
assert isinstance(result, str)
assert "AI Research" in result
async def test_research_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
report = ResearchReport(
title="AI Research",
executive_summary="AI is transforming industries.",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Multiple sources consulted.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
state = RAGState()
ctx = _make_ctx(state)
await research(ctx, question="What is AI?")
assert len(state.reports) == 1
assert state.reports[0].question == "What is AI?"
assert len(state.qa_history) == 1
assert state.qa_history[0].question == "What is AI?"
assert state.qa_history[0].answer == "AI is transforming industries."
async def test_research_applies_document_filter_from_state(
self, rag_db, monkeypatch
):
from haiku.rag.skills.rag import RAGState, create_skill
captured_kwargs = {}
report = ResearchReport(
title="AI Research",
executive_summary="Summary.",
main_findings=["Finding"],
conclusions=["Conclusion"],
sources_summary="Sources.",
)
async def mock_research(self, question, **kwargs):
captured_kwargs.update(kwargs)
return report
monkeypatch.setattr(HaikuRAG, "research", mock_research)
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
state = RAGState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
await research(ctx, question="What is AI?")
assert captured_kwargs.get("filter") == "title = 'AI Overview'"
async def test_research_without_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
report = ResearchReport(
title="AI Research",
executive_summary="Summary.",
main_findings=["Finding"],
conclusions=["Conclusion"],
sources_summary="Sources.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
cite = _get_tool(skill, "cite")
ctx = _make_ctx(state=None)
result = await research(ctx, question="What is AI?")
assert isinstance(result, str)
result = await cite(ctx, chunk_ids=["nonexistent"])
assert "No state" in result

226
tests/skills/test_rlm.py
View file

@ -1,226 +0,0 @@
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.models import RLMResult
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rlm import (
STATE_NAMESPACE,
STATE_TYPE,
RLMState,
instructions,
skill_metadata,
state_metadata,
)
from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
class TestRLMModuleAPI:
def test_state_type_is_rlm_state(self):
assert STATE_TYPE is RLMState
def test_state_namespace(self):
assert STATE_NAMESPACE == "rlm"
def test_state_metadata_returns_state_metadata(self):
result = state_metadata()
assert isinstance(result, StateMetadata)
assert result.namespace == "rlm"
assert result.type is RLMState
assert result.schema == RLMState.model_json_schema()
def test_skill_metadata_returns_skill_metadata(self):
result = skill_metadata()
assert isinstance(result, SkillMetadata)
assert result.name == "rag-rlm"
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.rlm 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 TestRLMSkillCreation:
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag-rlm"
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.rlm 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 == {"analyze"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import RLMState, create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is RLMState
assert skill._state_namespace == "rlm"
def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm 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
assert callable(skill.extras["visualize_chunk"])
assert callable(skill.extras["list_documents"])
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.rlm import create_skill
skill = create_skill()
assert skill.metadata.name == "rag-rlm"
class TestDomainPreambleInRLMSkillInstructions:
def test_create_skill_without_domain_preamble(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm 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.rlm 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 TestAnalyzeTool:
async def test_analyze_returns_result(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import create_skill
monkeypatch.setattr(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
ctx = _make_ctx()
result = await analyze(ctx, question="How many documents?")
assert isinstance(result, str)
assert "42" in result
assert "print(42)" in result
async def test_analyze_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import RLMState, create_skill
monkeypatch.setattr(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
state = RLMState()
ctx = _make_ctx(state)
await analyze(ctx, question="How many documents?")
assert len(state.analyses) == 1
assert state.analyses[0].question == "How many documents?"
assert state.analyses[0].answer == "42"
assert state.analyses[0].program == "print(42)"
async def test_analyze_applies_document_filter_from_state(
self, rag_db, monkeypatch
):
from haiku.rag.skills.rlm import RLMState, create_skill
captured_kwargs = {}
async def mock_rlm(self, question, **kwargs):
captured_kwargs.update(kwargs)
return RLMResult(answer="42", program="print(42)")
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
state = RLMState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
await analyze(ctx, question="How many documents?")
assert captured_kwargs.get("filter") == "title = 'AI Overview'"
async def test_analyze_combines_state_filter_with_explicit_filter(
self, rag_db, monkeypatch
):
from haiku.rag.skills.rlm import RLMState, create_skill
captured_kwargs = {}
async def mock_rlm(self, question, **kwargs):
captured_kwargs.update(kwargs)
return RLMResult(answer="Result", program="code()")
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
state = RLMState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
await analyze(
ctx,
question="Count pages",
filter="uri LIKE '%test%'",
)
result_filter = captured_kwargs["filter"]
assert isinstance(result_filter, str)
assert "title = 'AI Overview'" in result_filter
assert "uri LIKE '%test%'" in result_filter
async def test_analyze_with_document_and_filter(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import create_skill
captured_kwargs = {}
async def mock_rlm(self, question, **kwargs):
captured_kwargs.update(kwargs)
return RLMResult(answer="Result", program="code()")
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
ctx = _make_ctx()
await analyze(
ctx,
question="Count pages",
document="AI Overview",
filter="title = 'AI Overview'",
)
assert captured_kwargs.get("documents") == ["AI Overview"]
assert captured_kwargs.get("filter") == "title = 'AI Overview'"

93
tests/test_context.py
View file

@ -49,11 +49,12 @@ class TestMergeRanges:
assert len(merged) == 1
assert merged[0] == (0, 15, [r1, r2])
def test_adjacent(self):
def test_adjacent_stay_separate(self):
r1, r2 = _result(), _result()
merged = _merge_ranges([(0, 5, r1), (6, 10, r2)])
assert len(merged) == 1
assert merged[0] == (0, 10, [r1, r2])
assert len(merged) == 2
assert merged[0] == (0, 5, [r1])
assert merged[1] == (6, 10, [r2])
def test_sorts_by_position(self):
r1, r2 = _result(), _result()
@ -65,7 +66,7 @@ class TestMergeRanges:
class TestExpandOutward:
def test_basic_expansion(self):
items = [_item(i, text=f"{'x' * 100}") for i in range(10)]
lo, hi = _expand_outward(items, 5, max_items=10, max_chars=500)
lo, hi = _expand_outward(items, 5, max_chars=500)
assert lo <= 5
assert hi >= 5
total = sum(
@ -76,16 +77,9 @@ class TestExpandOutward:
# Should be around 500 chars (may overshoot by one item)
assert total >= 400
def test_respects_max_items(self):
items = [_item(i, text="x") for i in range(100)]
lo, hi = _expand_outward(items, 50, max_items=5, max_chars=999999)
count = hi - lo + 1
# May overshoot by 1-2 items due to alternating expansion
assert count <= 7
def test_respects_max_chars(self):
items = [_item(i, text=f"{'x' * 200}") for i in range(20)]
lo, hi = _expand_outward(items, 10, max_items=999, max_chars=500)
lo, hi = _expand_outward(items, 10, max_chars=500)
total = sum(
len(items[i].text)
for i in range(lo, hi + 1)
@ -96,12 +90,12 @@ class TestExpandOutward:
def test_center_at_start(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 0, max_items=5, max_chars=999999)
lo, hi = _expand_outward(items, 0, max_chars=999999)
assert lo == 0
def test_center_at_end(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 9, max_items=5, max_chars=999999)
lo, hi = _expand_outward(items, 9, max_chars=999999)
assert hi == 9
def test_skip_noise_excludes_from_char_count(self):
@ -113,7 +107,7 @@ class TestExpandOutward:
_item(4, label="footnote", text="f" * 5000),
_item(5, text="d" * 100),
]
lo, hi = _expand_outward(items, 2, max_items=10, max_chars=500, skip_noise=True)
lo, hi = _expand_outward(items, 2, max_chars=500, skip_noise=True)
# Footnotes (5000 chars each) should NOT count toward budget
# So we should expand past them
assert lo <= 0
@ -125,11 +119,17 @@ class TestExpandOutward:
_item(1, label="document_index", text="x" * 10000),
_item(2, text="b" * 200),
]
lo, hi = _expand_outward(items, 1, max_items=10, max_chars=500, skip_noise=True)
lo, hi = _expand_outward(items, 1, max_chars=500, skip_noise=True)
# Center is noise, should start at 0 chars and expand outward
assert lo == 0
assert hi == 2
def test_respects_bounds(self):
items = [_item(i, text="x" * 100) for i in range(20)]
lo, hi = _expand_outward(items, 10, max_chars=999999, lo_bound=8, hi_bound=12)
assert lo == 8
assert hi == 12
class TestFindExpansionRange:
def _structured_items(self):
@ -146,36 +146,47 @@ class TestFindExpansionRange:
def test_structured_returns_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
# Should return the Introduction section (items 0-3)
assert lo == 0
assert hi == 3
def test_structured_different_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {5}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {5}, has_sections=True, max_chars=5000)
# Should return the Methods section (items 4-6)
assert lo == 4
assert hi == 6
def test_structured_large_section_falls_back_to_outward(self):
def test_structured_large_section_bounded_by_section(self):
items = [
_item(0, label="section_header", text="Big Section"),
] + [_item(i, text="x" * 1000) for i in range(1, 20)]
# Section has 19 * 1000 = 19000 chars, way over 5000 budget
lo, hi = _find_expansion_range(
items, {10}, has_sections=True, max_items=50, max_chars=5000
)
# Should NOT return the full section
lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
# Should NOT return the full section, but should stay within it
total = sum(
len(items[i].text) for i in range(lo, hi + 1) if items[i].position >= lo
)
assert total < 10000
def test_structured_section_with_many_items_returned_whole(self):
"""A section that fits in char budget is returned even with many items."""
items = (
[
_item(0, label="section_header", text="Section"),
]
+ [_item(i, text="x" * 200) for i in range(1, 20)]
+ [
_item(20, label="section_header", text="Next"),
]
)
# Section has 19 * 200 = 3800 chars + header, under 5000 and over min_useful
lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
# Should return entire section despite 20 items
assert lo == 0
assert hi == 19
def test_structured_small_section_expands_outward(self):
items = [
_item(0, label="title", text="Paper Title"),
@ -186,26 +197,20 @@ class TestFindExpansionRange:
_item(5, text="Intro content. " * 50),
]
# Title section (items 0-1) is tiny (~25 chars) < 20% of 5000
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
# Should expand past the title section into the abstract
assert hi >= 3
def test_unstructured_expands_outward(self):
items = [_item(i, text=f"Paragraph {i}. " * 10) for i in range(10)]
lo, hi = _find_expansion_range(
items, {5}, has_sections=False, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {5}, has_sections=False, max_chars=5000)
assert lo < 5
assert hi > 5
def test_multiple_matched_positions_uses_center(self):
items = [_item(i, text="x" * 100) for i in range(20)]
# Match at positions 3 and 7, center should be index for position 5 (median)
lo, hi = _find_expansion_range(
items, {3, 7}, has_sections=False, max_items=5, max_chars=999999
)
# Use a char budget that forces partial expansion so center matters
lo, hi = _find_expansion_range(items, {3, 7}, has_sections=False, max_chars=500)
center = (lo + hi) // 2
# Center should be around position 5
assert 3 <= center <= 7
@ -219,9 +224,7 @@ class TestFindExpansionRange:
]
# Section non-noise chars: ~260 chars (items 0,1,3). Under 5000 budget.
# The footnote's 10000 chars should NOT count.
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
# Should return full section (it fits in budget excluding noise)
assert lo == 0
assert hi == 3
@ -233,9 +236,7 @@ class TestFindExpansionRange:
_item(2, label="section_header", text="First Section"),
_item(3, text="Section content."),
]
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
# Match is in preamble section (items 0-1), which is small
# Should expand outward into the first section
assert hi >= 2
@ -262,7 +263,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/999999"],
)
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [result], 10, 5000
rag.document_item_repository, doc.id, [result], 5000
)
assert len(expanded) == 1
assert expanded[0].content == "original"
@ -312,7 +313,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
# The TOC section's only non-header item is document_index (noise).
@ -376,7 +377,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
# Expansion produces "Steps\n\nClick\n\n+\n\nAdd a New Service" = 38 chars

28
tests/test_context_enhancement.py
View file

@ -80,7 +80,6 @@ def small_chunk_config() -> AppConfig:
"""Config with small chunk size to force splitting."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 25
config.search.max_context_chars = 10000
return config
@ -305,33 +304,6 @@ async def test_format_for_agent_output(temp_db_path, small_chunk_config):
assert "Content:" in formatted
@pytest.mark.vcr()
async def test_max_items_limit_caps_expansion(temp_db_path):
"""Expansion should respect max_context_items limit."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 2 # Very restrictive
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Limit Test")
assert doc.id is not None
results = await client.search("grapes", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
# With max_items=2, expansion should be limited
content = expanded[0].content.lower()
item_count = sum(
1 for item in ["apples", "bananas", "oranges", "grapes"] if item in content
)
# Should have at most 2 items (the limit)
assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
async def test_expand_context_single_item_document(temp_db_path):
"""Test expand_context with a single-item document."""
from haiku.rag.store.models.document import Document

48
tests/test_skill_generator.py
View file

@ -23,9 +23,8 @@ class TestAvailableTools:
"list_documents",
"get_document",
"search",
"ask",
"research",
"analyze",
"execute_code",
"cite",
}
@ -43,7 +42,7 @@ class TestValidateTools:
validate_tools(["search"])
def test_valid_multiple_tools(self):
validate_tools(["list_documents", "get_document", "search", "ask"])
validate_tools(["list_documents", "get_document", "search", "cite"])
def test_valid_all_tools(self):
validate_tools(list(AVAILABLE_TOOLS))
@ -97,7 +96,7 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["list_documents", "get_document", "search", "ask"],
tool_names=["list_documents", "get_document", "search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assert result.is_dir()
@ -112,7 +111,7 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="my-recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
assert result == tmp_path / "my-recipes-skill"
pkg = result / "my_recipes_skill"
@ -129,11 +128,11 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert '["search", "ask"]' in content
assert '["search", "cite"]' in content
def test_create_skill_tools_called(self, tmp_path):
render_templates(
@ -151,12 +150,12 @@ class TestRenderTemplates:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
assert '"search"' in content
assert '"ask"' in content
assert '"cite"' in content
def test_pyproject_toml(self, tmp_path):
render_templates(
@ -184,24 +183,23 @@ class TestRenderTemplates:
)
skill_md = tmp_path / "docs-skill" / "docs_skill" / "SKILL.md"
content = skill_md.read_text()
assert "**search**" in content
assert "**ask**" not in content
assert "**list_documents**" not in content
assert "**research**" not in content
assert "**analyze**" not in content
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", "ask", "analyze"],
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 "**ask**" in content
assert "**analyze**" in content
assert "search" in content
assert "execute_code" in content
assert "cite" in content
def test_custom_preamble(self, tmp_path):
render_templates(
@ -226,23 +224,23 @@ class TestRenderTemplates:
content = init.read_text()
assert 'state_namespace="recipes"' in content
def test_analyze_state_fields(self, tmp_path):
def test_execute_code_state_fields(self, tmp_path):
render_templates(
output_dir=tmp_path,
name="docs",
description="A docs skill.",
tool_names=["search", "analyze"],
tool_names=["search", "execute_code"],
)
init = tmp_path / "docs-skill" / "docs_skill" / "__init__.py"
content = init.read_text()
assert "analyses" in content
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", "ask", "analyze"],
tool_names=["search", "execute_code", "cite"],
)
init = tmp_path / "recipes-skill" / "recipes_skill" / "__init__.py"
content = init.read_text()
@ -329,7 +327,7 @@ class TestGenerateSkill:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
)
assert result == tmp_path / "recipes-skill"
assets = result / "recipes_skill" / "assets"
@ -479,7 +477,7 @@ class TestGenerateSkillRemote:
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
tool_names=["search", "cite"],
config_path=config_file,
)
assets = result / "recipes_skill" / "assets"

42
tests/tools/test_analysis.py
View file

@ -1,42 +0,0 @@
import pytest
from haiku.rag.tools.analysis import create_analysis_toolset
class TestAnalysisToolset:
"""Tests for create_analysis_toolset."""
def test_create_analysis_toolset_returns_function_toolset(self, analysis_config):
"""create_analysis_toolset returns a FunctionToolset."""
from pydantic_ai import FunctionToolset
toolset = create_analysis_toolset(analysis_config)
assert isinstance(toolset, FunctionToolset)
def test_analysis_toolset_has_analyze_tool(self, analysis_config):
"""The toolset includes an 'analyze' tool."""
toolset = create_analysis_toolset(analysis_config)
assert "analyze" in toolset.tools
def test_analysis_toolset_custom_tool_name(self, analysis_config):
"""Toolset supports custom tool name."""
toolset = create_analysis_toolset(analysis_config, tool_name="run_code")
assert "run_code" in toolset.tools
assert "analyze" not in toolset.tools
@pytest.fixture
async def analysis_client(temp_db_path):
"""Create a HaikuRAG client for analysis tests."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
yield rag
@pytest.fixture
def analysis_config():
"""Default AppConfig for analysis tests."""
from haiku.rag.config import Config
return Config

46
tests/tools/test_filters.py
View file

@ -1,21 +1,16 @@
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
)
from haiku.rag.tools.filters import _build_document_filter, build_multi_document_filter
def test_build_document_filter_simple():
"""Test build_document_filter with simple name."""
result = build_document_filter("mytest")
"""Test _build_document_filter with simple name."""
result = _build_document_filter("mytest")
assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
assert "LOWER(title) LIKE LOWER('%mytest%')" in result
def test_build_document_filter_with_spaces():
"""Test build_document_filter handles spaces correctly."""
result = build_document_filter("TB MED 593")
# Should include both the original (with spaces) and without spaces
"""Test _build_document_filter handles spaces correctly."""
result = _build_document_filter("TB MED 593")
assert "LOWER(uri) LIKE LOWER('%TB MED 593%')" in result
assert "LOWER(uri) LIKE LOWER('%TBMED593%')" in result
assert "LOWER(title) LIKE LOWER('%TB MED 593%')" in result
@ -23,9 +18,8 @@ def test_build_document_filter_with_spaces():
def test_build_document_filter_escapes_quotes():
"""Test build_document_filter escapes single quotes."""
result = build_document_filter("O'Reilly")
# Single quotes should be doubled for SQL escaping
"""Test _build_document_filter escapes single quotes."""
result = _build_document_filter("O'Reilly")
assert "O''Reilly" in result
@ -41,7 +35,6 @@ def test_build_multi_document_filter_single():
assert result is not None
assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
assert "LOWER(title) LIKE LOWER('%mytest%')" in result
# Single document should not have extra wrapping parentheses
assert " OR (" not in result
@ -49,31 +42,6 @@ def test_build_multi_document_filter_multiple():
"""Test build_multi_document_filter with multiple documents."""
result = build_multi_document_filter(["doc1", "doc2"])
assert result is not None
# Should have OR-combined filters
assert "doc1" in result
assert "doc2" in result
assert " OR (" in result
def test_combine_filters_both_none():
"""Test combine_filters with both None."""
result = combine_filters(None, None)
assert result is None
def test_combine_filters_first_only():
"""Test combine_filters with only first filter."""
result = combine_filters("uri = 'test'", None)
assert result == "uri = 'test'"
def test_combine_filters_second_only():
"""Test combine_filters with only second filter."""
result = combine_filters(None, "title = 'doc'")
assert result == "title = 'doc'"
def test_combine_filters_both():
"""Test combine_filters combines with AND."""
result = combine_filters("uri = 'test'", "title = 'doc'")
assert result == "(uri = 'test') AND (title = 'doc')"

17
tests/tools/test_models.py
View file

@ -1,17 +0,0 @@
from haiku.rag.tools.analysis import AnalysisResult
def test_analysis_result_defaults():
"""Test AnalysisResult has sensible defaults."""
result = AnalysisResult(answer="The result is 42")
assert result.code_executed is True
def test_analysis_result_with_values():
"""Test AnalysisResult with explicit values."""
result = AnalysisResult(
answer="The result is 42",
code_executed=True,
)
assert result.answer == "The result is 42"
assert result.code_executed is True

95
tests/tools/test_qa.py
View file

@ -1,95 +0,0 @@
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
class TestQAHistoryEntry:
"""Tests for QAHistoryEntry model."""
def test_defaults(self):
"""QAHistoryEntry has sensible defaults."""
entry = QAHistoryEntry(question="What is X?", answer="X is Y.")
assert entry.confidence == 0.9
assert entry.citations == []
assert entry.question_embedding is None
def test_sources_property(self):
"""sources returns unique document titles."""
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="doc1.md",
document_title="Document One",
content="Content 1",
),
Citation(
document_id="d1",
chunk_id="c2",
document_uri="doc1.md",
document_title="Document One",
content="Content 2",
),
Citation(
document_id="d2",
chunk_id="c3",
document_uri="doc2.md",
document_title="Document Two",
content="Content 3",
),
]
entry = QAHistoryEntry(question="Q", answer="A", citations=citations)
sources = entry.sources
assert len(sources) == 2
assert "Document One" in sources
assert "Document Two" in sources
def test_sources_uses_uri_as_fallback(self):
"""sources uses uri when title is None."""
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test.md",
document_title=None,
content="Content",
),
]
entry = QAHistoryEntry(question="Q", answer="A", citations=citations)
assert entry.sources == ["test.md"]
def test_to_search_answer(self):
"""to_search_answer converts to SearchAnswer."""
citation = Citation(
document_id="d1",
chunk_id="c1",
document_uri="doc1.md",
document_title="Doc",
content="Content",
)
entry = QAHistoryEntry(
question="What is X?",
answer="X is Y.",
confidence=0.85,
citations=[citation],
)
sa = entry.to_search_answer()
assert sa.query == "What is X?"
assert sa.answer == "X is Y."
assert sa.confidence == 0.85
assert sa.cited_chunks == ["c1"]
assert len(sa.citations) == 1
def test_question_embedding_excluded_from_serialization(self):
"""question_embedding is excluded from model_dump."""
entry = QAHistoryEntry(
question="Q",
answer="A",
question_embedding=[0.1, 0.2],
)
data = entry.model_dump()
assert "question_embedding" not in data
def test_prior_answer_relevance_threshold():
"""PRIOR_ANSWER_RELEVANCE_THRESHOLD is a sensible value."""
assert 0 < PRIOR_ANSWER_RELEVANCE_THRESHOLD < 1