Merge pull request #342 from ggozad/feat/analysis

Refactor analysis sandbox with document VFS and expanded search

CHANGELOG.md

@@ -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

README.md

@@ -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
+# Analyze — complex analytical tasks via code execution
-haiku-rag rlm "How many documents mention transformers?"
+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

app/frontend/components/Chat.tsx

@@ -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":
+			case "execute_code":
-				return "Analyze";
+				return "Code";
-			case "research":
+			case "cite":
-				return "Research";
+				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":
+			case "execute_code": {
-				return <span className="tool-query">{args.question as string}</span>;
+				const code = args.code as string | undefined;
-			case "research":
+				return (
-				return <span className="tool-query">{args.question as string}</span>;
+					<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) {
+						if (latestCitations.length > 0) {
-							const citations = citationsHistory[citIdx];
+							result.push(
-							if (citations?.length) {
+								<CitationBlock
-								result.push(
+									key={`citations-${i}`}
-									<CitationBlock
+									citations={latestCitations}
-										key={`citations-${citIdx}`}
+								/>,
-										citations={citations}
+							);
-									/>,
-								);
-							}
-							citIdx++;
 						}
 						seenToolCalls = false;
 					}

app/frontend/lib/sessionStorage.ts

@@ -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[];
+	citation_index: Record<string, Citation>;
-	qa_history: QAHistoryEntry[];
+	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 getLatestCitations(state: RAGState): Citation[] {
-export function deriveCitationsHistory(state: RAGState): Citation[][] {
+	const turns = state.citations;
-	return state.qa_history
+	if (turns.length === 0) return [];
-		.filter((entry) => entry.citations?.length > 0)
+	const latestIds = turns[turns.length - 1];
-		.map((entry) => entry.citations);
+	return latestIds
+		.map((id) => state.citation_index[id])
+		.filter((c): c is Citation => c !== undefined);
 }
 
 export function getAllSessions(): StoredSession[] {

docs/agents/analysis.md

@@ -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
+```

docs/agents/index.md

@@ -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.
 

docs/agents/rlm.md

@@ -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
-```

docs/apps.md

@@ -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
 

docs/architecture.md

@@ -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.

docs/cli.md

@@ -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."

docs/configuration/index.md

@@ -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

docs/configuration/providers.md

@@ -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

docs/configuration/qa-research.md

@@ -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.

docs/configuration/storage.md

@@ -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)

docs/index.md

@@ -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

docs/mcp.md

@@ -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)

docs/python.md

@@ -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
 

docs/skills/analysis.md

@@ -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.

docs/skills/index.md

@@ -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.

docs/skills/rag.md

@@ -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 |
+| `cite(chunk_ids)` | Register chunk IDs as citations for the current answer |
-| `research(question)` | Deep multi-agent research producing comprehensive reports |
 
 ## State
 
@@ -31,17 +30,13 @@ The skill manages a `RAGState` under the `"rag"` namespace:
 
 ```python
 class RAGState(BaseModel):
-    citations: list[Citation] = []
+    citation_index: dict[str, Citation] = {}
-    qa_history: list[QAHistoryEntry] = []
+    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.
+- **citation_index** — All citations indexed by chunk ID (deduplicated across turns).
-- **qa_history** — Questions and answers from `ask` calls. Prior Q&A is used as context for follow-up questions when embeddings are similar.
+- **citations** — Per-turn lists of chunk IDs registered via the `cite` tool.
-- **document_filter** — SQL WHERE clause applied to `search`, `list_documents`, `ask`, and `research` calls. Set this to scope queries to specific documents.
+- **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.

docs/skills/rlm.md

@@ -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.

docs/tools.md

@@ -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. Matches against both `uri` and `title`, case-insensitive.
-- **`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`.

docs/tuning.md

@@ -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
 

evaluations/evaluations/benchmark.py

@@ -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

haiku_rag_slim/haiku/rag/agents/analysis/__init__.py

@@ -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",
+]

haiku_rag_slim/haiku/rag/agents/analysis/agent.py

@@ -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

haiku_rag_slim/haiku/rag/agents/analysis/dependencies.py

@@ -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)

haiku_rag_slim/haiku/rag/agents/analysis/models.py

@@ -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):
+class RawAnalysisResult(BaseModel):
-    """Result from RLM agent execution."""
+    """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)

haiku_rag_slim/haiku/rag/agents/analysis/prompts.py

@@ -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."""

haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py

@@ -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)

haiku_rag_slim/haiku/rag/agents/rlm/__init__.py

@@ -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",
-]

haiku_rag_slim/haiku/rag/agents/rlm/agent.py

@@ -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

haiku_rag_slim/haiku/rag/agents/rlm/dependencies.py

@@ -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)

haiku_rag_slim/haiku/rag/agents/rlm/prompts.py

@@ -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."""

haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py

@@ -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)

haiku_rag_slim/haiku/rag/app.py

@@ -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"))

haiku_rag_slim/haiku/rag/chat/__init__.py

@@ -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, parse_model_option
-    from haiku.rag.utils import get_model
+    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),

haiku_rag_slim/haiku/rag/chat/app.py

@@ -30,6 +30,7 @@ from textual.worker import Worker
 from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import get_config
+from haiku.rag.skills.analysis import AnalysisState
 from haiku.rag.skills.rag import RAGState, get_agent_preamble
 from haiku.skills.agent import (
     SkillToolset,
@@ -51,6 +52,7 @@ if TYPE_CHECKING:
 
 
 RAG_STATE_NAMESPACE = "rag"
+ANALYSIS_STATE_NAMESPACE = "analysis"
 
 
 class ChatApp(App):
@@ -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_and_programs(chat_history)
-                        await self._show_citations(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:
+    async def _show_citations_and_programs(self, chat_history: "ChatHistory") -> None:
-        """Show citations from the RAG state after an agent response."""
+        """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)
+        citations = []
-        if rag_state is None:
+        for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
-            return
+            state = self._toolset.get_namespace(namespace)
-        citations = getattr(rag_state, "citations", [])
+            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(
+                rag_state.document_filter = doc_filter
-                    self._document_filter
+            analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
-                )
+            if isinstance(analysis_state, AnalysisState):
-                self._state = self._toolset.build_state_snapshot()
+                analysis_state.document_filter = doc_filter
+            self._state = self._toolset.build_state_snapshot()

haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py

@@ -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:

haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py

@@ -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()

haiku_rag_slim/haiku/rag/cli.py

@@ -368,8 +368,8 @@ def ask(  # pragma: no cover
     )
 
 
-@_cli.command("rlm", help="Answer questions using code execution (RLM agent)")
+@_cli.command("analyze", help="Answer questions using code execution (analysis agent)")
-def rlm(  # pragma: no cover
+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,
     )
 
 

haiku_rag_slim/haiku/rag/client.py

@@ -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":
+    ) -> "AnalysisResult":
-        """Answer a question using the RLM agent with code execution.
+        """Answer a question using the analysis agent with code execution.
 
-        The RLM (Recursive Language Model) agent can write and execute Python
+        The analysis agent can write and execute Python code in a sandboxed
-        code in a sandboxed environment to solve problems that require
+        environment to solve problems that require computation, aggregation,
-        computation, aggregation, or complex traversal across documents.
+        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 (
+        from haiku.rag.agents.analysis import (
-            RLMContext,
+            AnalysisContext,
-            RLMDeps,
+            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
+        Expands the chunk's context to find the full section, then resolves
-        from chunk metadata, and renders all pages that contain bounding boxes with
+        bounding boxes from all items in the expanded range. This ensures
-        yellow/orange highlight overlays.
+        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 []
 

haiku_rag_slim/haiku/rag/config/models.py

@@ -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)

haiku_rag_slim/haiku/rag/context.py

@@ -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
+  3. If the section fits within the char budget, include it entirely
-  4. If the section exceeds the budget, OR the section is too small (under
+  4. If the section exceeds the char budget, expand item-by-item from the
-     20% of max_context_chars), expand item-by-item from the match center
+     match center outward, bounded by section edges
-     outward, skipping noise labels. This lets small sections (e.g., a
+  5. If the section is too small (under 20% of max_context_chars), expand
-     title+authors area) grow into the next section's content.
+     item-by-item crossing into adjacent sections until the budget is filled.
-  5. Merge overlapping ranges from multiple results in the same document
+     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.
+    if sec_chars > max_chars:
-    # For too-large sections this stays within budget.
+        # Section too large — expand outward bounded by section edges
-    # For too-small sections (e.g., title+authors) this naturally grows
+        return _expand_outward(
-    # into adjacent sections until the budget is filled.
+            items,
-    return _expand_outward(items, center_idx, max_items, max_chars, skip_noise=True)
+            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(
+        lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
-            window_items, matched, has_sections, max_items, max_chars
-        )
         ranges.append((lo, hi, result))
 
     merged = _merge_ranges(ranges)

haiku_rag_slim/haiku/rag/mcp.py

@@ -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

haiku_rag_slim/haiku/rag/skill_generator/__init__.py

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

haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2

@@ -7,48 +7,58 @@ description: {{ description }}
 
 {{ preamble }}
 
-## How to decide which tool to use
+## Tools
-{% if "ask" in tool_names %}
+{% 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 %}
+{% if "execute_code" 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.
+
+### 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 %}
+{% if "cite" 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 %}
+### cite
-{% if "research" in tool_names %}
+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.
-- **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.
 {% 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:
+1. Call `search` with relevant keywords from the question
-{% if "ask" in tool_names %}
+2. Review results — they are ordered by relevance (rank 1 = best match)
-- Use **ask** if the question is factual
+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
-- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
+- 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?"
 {% endif %}

haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2

@@ -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 %}
+{% if "execute_code" in tool_names %}
-from haiku.rag.skills._tools import ResearchEntry
+from haiku.rag.skills._tools import CodeExecutionEntry
-{% endif %}
-{% if "analyze" in tool_names %}
-from haiku.rag.skills._tools import AnalysisEntry
 {% 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 %}
+{% if "cite" in tool_names %}
-    citations: list[Citation] = Field(default_factory=list)
+    citation_index: dict[str, Citation] = Field(default_factory=dict)
-{% endif %}
+    citations: list[list[str]] = Field(default_factory=list)
-{% if "ask" in tool_names %}
-    qa_history: list[QAHistoryEntry] = 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 %}
+{% if "execute_code" in tool_names %}
-    documents: list[DocumentInfo] = Field(default_factory=list)
+    executions: list[CodeExecutionEntry] = 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)
 {% endif %}
 
 

haiku_rag_slim/haiku/rag/skills/_tools.py

@@ -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):
+class CodeExecutionEntry(BaseModel):
-    question: str
+    code: str
-    title: str
+    stdout: str
-    executive_summary: str
+    stderr: str = ""
-
+    success: bool = True
-
-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
 
 
 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.
+            """List all documents in the knowledge base."""
-
-            Args:
-                limit: Maximum number of documents to return.
-                offset: Number of documents to skip.
-            """
             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)
+            return 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
 
         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:
+        async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
-            """Ask a question and get an answer with citations from the knowledge base.
+            """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(
+            doc_filter = state.document_filter if state else None
-                db_path,
+            context = AnalysisContext(filter=doc_filter)
-                config,
+            sandbox = Sandbox(db_path=db_path, config=config, context=context)
-                question,
+            result = await sandbox.execute(code)
-                qa_history=state.qa_history if state else None,
+
-                document_filter=state.document_filter if state else None,
+            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
+                state.executions.append(
-                for citation in citations:
+                    CodeExecutionEntry(
-                    citation.index = next_index
+                        code=code,
-                    next_index += 1
+                        stdout=result.stdout,
-                state.citations.extend(citations)
+                        stderr=result.stderr,
-                state.qa_history.append(
+                        success=result.success,
-                    QAHistoryEntry(
-                        question=question, answer=answer, citations=citations
                     )
                 )
 
+            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["cite"] = cite
-
-        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
 
     return tools

haiku_rag_slim/haiku/rag/skills/analysis.py

@@ -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_TYPE = AnalysisState
-STATE_NAMESPACE = "rlm"
+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()

haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md

@@ -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

haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md

@@ -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.

haiku_rag_slim/haiku/rag/skills/rag.py

@@ -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)
+    citation_index: dict[str, Citation] = Field(default_factory=dict)
-    qa_history: list[QAHistoryEntry] = Field(default_factory=list)
+    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

haiku_rag_slim/haiku/rag/skills/rag/SKILL.md

@@ -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").
+Each result includes:
-- **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.
+- `chunk_id` in brackets and rank position (rank 1 = most relevant)
-- **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.
+- Source: document title and section hierarchy
-- **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.
+- Type: content type (paragraph, table, code, list_item)
-- **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.
+- 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:
+### get_document
-- Use **ask** if the question is factual
+Retrieve a document by ID, title, or URI. Partial matches work. Use when the user wants the full content of a specific document.
-- Report that the knowledge base doesn't contain relevant information
+
+### 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
-- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
+- Then search for the topic
 
-Examples:
+## When search returns irrelevant results
-- "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?"
+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

haiku_rag_slim/haiku/rag/store/repositories/document.py

@@ -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:

haiku_rag_slim/haiku/rag/store/repositories/document_item.py

@@ -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]:

haiku_rag_slim/haiku/rag/tools/__init__.py

@@ -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 (
+from haiku.rag.tools.filters import build_multi_document_filter
-    build_document_filter,
-    build_multi_document_filter,
-    combine_filters,
-)
-from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
 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",
 ]

haiku_rag_slim/haiku/rag/tools/analysis.py

@@ -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

haiku_rag_slim/haiku/rag/tools/filters.py

@@ -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]})"

haiku_rag_slim/haiku/rag/tools/qa.py

@@ -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,
-        )

haiku_rag_slim/pyproject.toml

@@ -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"

mkdocs.yml

@@ -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

tests/agents/analysis/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 
-from haiku.rag.agents.rlm.dependencies import RLMContext
+from haiku.rag.agents.analysis.dependencies import AnalysisContext
-from haiku.rag.agents.rlm.sandbox import Sandbox
+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()
+    async with HaikuRAG(temp_db_path, create=True):
-    context = RLMContext()
+        config = AppConfig()
-    return Sandbox(client=empty_client, config=config, context=context)
+        context = AnalysisContext()
+        return Sandbox(db_path=temp_db_path, config=config, context=context)

tests/agents/analysis/test_agent.py

@@ -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.analysis.agent import create_analysis_agent
-from haiku.rag.agents.rlm.dependencies import RLMDeps
+from haiku.rag.agents.analysis.dependencies import AnalysisDeps
-from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
+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.deps_type is AnalysisDeps
-        assert agent.output_type is RLMResult
+        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:
+class TestClientAnalysisIntegration:
-    """Integration tests for client.rlm() method."""
+    """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):
+    async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
-        """Test RLM agent can count documents.
+        """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):
+    async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
-        """Test RLM agent can perform aggregation across documents.
+        """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):
+    async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
-        """Test RLM agent respects filter parameter.
+        """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):
+    async def test_analyze_search_and_identify_source(
-        """Test RLM agent can search and use get_chunk for citations.
+        self, allow_model_requests, temp_db_path
-
+    ):
-        Agent program:
+        """Test analysis agent can search and identify source documents."""
-            results = search("content", limit=5)
-            for r in results:
-                chunk = get_chunk(r['chunk_id'])
-                print(chunk['document_title'], chunk['chunk_id'])
-        """
         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):
+    async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
-        """Test RLM agent can use search() to find content and extract information.
+        """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"],

tests/agents/analysis/test_models.py

@@ -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)"

tests/agents/analysis/test_sandbox.py

@@ -2,8 +2,8 @@ from pathlib import Path
 
 import pytest
 
-from haiku.rag.agents.rlm.dependencies import RLMContext
+from haiku.rag.agents.analysis.dependencies import AnalysisContext
-from haiku.rag.agents.rlm.sandbox import Sandbox, SandboxResult
+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:
+class TestSandboxListDocuments:
-    """Test haiku.rag functions in sandbox."""
+    """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()
+            context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            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()
+            context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            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):
+    async def test_search_returns_doc_item_refs_and_labels(self, temp_db_path):
-        """Test get_document function."""
+        """Search results include doc_item_refs and labels."""
-        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."""
         config = AppConfig()
         async with HaikuRAG(temp_db_path, create=True) as client:
             await client.create_document(
-                content="Content about foxes and dogs.",
+                content="The quick brown fox jumps over the lazy dog.",
-                uri="test://doc",
+                uri="test://animals",
-                title="Fox Document",
+                title="Animals",
             )
 
-            context = RLMContext()
+            context = AnalysisContext()
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
-            # First search to get a chunk_id
             result = await sb.execute(
-                "results = await search('foxes', limit=1)\n"
+                "results = await search('fox', limit=1)\n"
-                "chunk_id = results[0]['chunk_id']\n"
+                "r = results[0]\n"
-                "chunk = await get_chunk(chunk_id)\n"
+                "print('doc_item_refs' in r)\n"
-                "print(chunk['document_title'])\n"
+                "print('labels' in r)\n"
-                "print('content' in chunk)"
+                "print(type(r['doc_item_refs']).__name__)\n"
+                "print(type(r['labels']).__name__)"
             )
             assert result.success
-            assert "Fox Document" in result.stdout
+            assert "True\nTrue" in result.stdout
-            assert "True" in result.stdout
+            assert "list\nlist" in result.stdout
 
     @pytest.mark.asyncio
-    async def test_get_chunk_not_found(self, sandbox):
+    @pytest.mark.vcr()
-        """Test get_chunk returns None for missing chunk."""
+    async def test_search_returns_expanded_content(self, temp_db_path):
-        result = await sandbox.execute(
+        """search() returns context-expanded results."""
-            "chunk = await get_chunk('nonexistent-id')\nprint(chunk is None)"
+        config = AppConfig()
-        )
+        async with HaikuRAG(temp_db_path, create=True) as client:
-        assert result.success
+            await client.create_document(
-        assert "True" in result.stdout
+                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()
+        async with HaikuRAG(temp_db_path, create=True):
-        config.rlm.max_output_chars = 20
+            config = AppConfig()
-        context = RLMContext()
+            config.analysis.max_output_chars = 20
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            context = AnalysisContext()
-        result = await sb.execute("print('a' * 100)\nx = 1/0")
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
-        assert not result.success
+            result = await sb.execute("print('a' * 100)\nx = 1/0")
-        assert "ZeroDivisionError" in result.stderr
+            assert not result.success
-        assert result.stdout.endswith("... (output truncated)")
+            assert "ZeroDivisionError" in result.stderr
-        assert len(result.stdout) < 100
+            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()
+        async with HaikuRAG(temp_db_path, create=True):
-        config.rlm.max_output_chars = 20
+            config = AppConfig()
-        context = RLMContext()
+            config.analysis.max_output_chars = 20
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            context = AnalysisContext()
-        result = await sb.execute("print('b' * 100)")
+            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)")
+        # /documents dir may or may not exist when empty, both are valid
-        assert len(result.stdout) < 100
+        # The key is it doesn't error
-
-
-class TestSandboxContextFilter:
-    """Test context filter is applied."""
 
     @pytest.mark.asyncio
     @pytest.mark.vcr()
-    async def test_filter_applied_to_list_documents(self, temp_db_path):
+    async def test_iterdir_discovers_documents(self, temp_db_path):
-        """Test that context filter is passed to list_documents."""
+        """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://%'")
+            context = AnalysisContext(filter="uri LIKE 'public://%'")
-            sb = Sandbox(client=client, config=config, context=context)
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
             result = await sb.execute(
-                "docs = await list_documents()\n"
+                "from pathlib import Path\n"
-                "print(len(docs))\n"
+                "import json\n"
-                "if docs:\n"
+                "dirs = list(Path('/documents').iterdir())\n"
-                "    print(docs[0]['title'])"
+                "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()
+        async with HaikuRAG(temp_db_path, create=True):
-        docs = [
+            config = AppConfig()
-            Document(id="1", content="Content A", title="Doc A", uri="a://1"),
+            docs = [
-            Document(id="2", content="Content B", title="Doc B", uri="b://2"),
+                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)
+            context = AnalysisContext(documents=docs)
-        result = await sb.execute(
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
-            "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)
             result = await sb.execute(
-                f"doc = await get_docling_document('{doc.id}')\n"
+                "print(len(documents))\n"
-                "print(type(doc).__name__)\n"
+                "print(documents[0]['title'])\n"
-                "print(doc['name'])\n"
+                "print(documents[1]['title'])"
-                "print('texts' in doc)"
             )
             assert result.success
-            assert "dict" in result.stdout
+            assert "2" in result.stdout
-            assert "True" 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()
+        async with HaikuRAG(temp_db_path, create=True):
-        context = RLMContext()
+            config = AppConfig()
-        sb = Sandbox(client=empty_client, config=config, context=context)
+            context = AnalysisContext()
-        result = await sb.execute(
+            sb = Sandbox(db_path=temp_db_path, config=config, context=context)
-            "answer = await llm('What is 2 + 2? Reply with just the number.')\n"
+            result = await sb.execute(
-            "print(answer)"
+                "answer = await llm('What is 2 + 2? Reply with just the number.')\n"
-        )
+                "print(answer)"
-        assert result.success
+            )
-        assert "4" in result.stdout
+            assert result.success
+            assert "4" in result.stdout

tests/chat/test_chat_app.py

@@ -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
 

tests/skills/test_analysis.py

@@ -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

tests/skills/test_rag.py

@@ -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 == {
+        assert tool_names == {"search", "list_documents", "get_document", "cite"}
-            "search",
-            "list_documents",
-            "get_document",
-            "ask",
-            "research",
-        }
 
     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:
+class TestCiteTool:
-    async def test_ask_returns_answer_with_citations(self, rag_db, monkeypatch):
+    async def test_cite_registers_citations(self, rag_db):
-        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):
         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 len(state.citations[0]) == 2
-        assert isinstance(state.qa_history[0], QAHistoryEntry)
+        assert all(cid in state.citation_index for cid in chunk_ids)
-        assert state.qa_history[0].question == "What is AI?"
 
-    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")
+        await search(ctx, query="artificial intelligence")
-        assert state.citations[0].index == 1
+        chunk_ids = [
-        assert state.citations[1].index == 2
+            sr.chunk_id
+            for results in state.searches.values()
+            for sr in results
+            if sr.chunk_id
+        ][:1]
 
-        await ask(ctx, question="Second question")
+        await cite(ctx, chunk_ids=chunk_ids)
-        assert state.citations[2].index == 3
+        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):
+    async def test_cite_without_state(self, rag_db):
-        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):
         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")
+        cite = _get_tool(skill, "cite")
-        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")
         ctx = _make_ctx(state=None)
-        result = await research(ctx, question="What is AI?")
+        result = await cite(ctx, chunk_ids=["nonexistent"])
-        assert isinstance(result, str)
+        assert "No state" in result

tests/skills/test_rlm.py

@@ -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'"

tests/test_context.py

@@ -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 len(merged) == 2
-        assert merged[0] == (0, 10, [r1, r2])
+        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(
+        lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
-            items, {1}, has_sections=True, max_items=20, 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(
+        lo, hi = _find_expansion_range(items, {5}, has_sections=True, max_chars=5000)
-            items, {5}, has_sections=True, max_items=20, 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(
+        lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
-            items, {10}, has_sections=True, max_items=50, max_chars=5000
+        # Should NOT return the full section, but should stay within it
-        )
-        # Should NOT return the full section
         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(
+        lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
-            items, {0}, has_sections=True, max_items=20, 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(
+        lo, hi = _find_expansion_range(items, {5}, has_sections=False, max_chars=5000)
-            items, {5}, has_sections=False, max_items=20, 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)
+        # Use a char budget that forces partial expansion so center matters
-        lo, hi = _find_expansion_range(
+        lo, hi = _find_expansion_range(items, {3, 7}, has_sections=False, max_chars=500)
-            items, {3, 7}, has_sections=False, max_items=5, max_chars=999999
-        )
         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(
+        lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
-            items, {1}, has_sections=True, max_items=20, 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(
+        lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
-            items, {0}, has_sections=True, max_items=20, 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

tests/test_context_enhancement.py

@@ -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

tests/test_skill_generator.py

@@ -23,9 +23,8 @@ class TestAvailableTools:
             "list_documents",
             "get_document",
             "search",
-            "ask",
+            "execute_code",
-            "research",
+            "cite",
-            "analyze",
         }
 
 
@@ -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 "### search" in content
-        assert "**ask**" not in content
+        assert "### cite" not in content
-        assert "**list_documents**" not in content
+        assert "### list_documents" not in content
-        assert "**research**" not in content
+        assert "### execute_code" not in content
-        assert "**analyze**" 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 "search" in content
-        assert "**ask**" in content
+        assert "execute_code" in content
-        assert "**analyze**" 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"

tests/tools/test_analysis.py

@@ -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

tests/tools/test_filters.py

@@ -1,21 +1,16 @@
-from haiku.rag.tools.filters import (
+from haiku.rag.tools.filters import _build_document_filter, build_multi_document_filter
-    build_document_filter,
-    build_multi_document_filter,
-    combine_filters,
-)
 
 
 def test_build_document_filter_simple():
-    """Test build_document_filter with simple name."""
+    """Test _build_document_filter with simple name."""
-    result = build_document_filter("mytest")
+    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."""
+    """Test _build_document_filter handles spaces correctly."""
-    result = build_document_filter("TB MED 593")
+    result = _build_document_filter("TB MED 593")
-    # Should include both the original (with spaces) and without spaces
     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."""
+    """Test _build_document_filter escapes single quotes."""
-    result = build_document_filter("O'Reilly")
+    result = _build_document_filter("O'Reilly")
-    # Single quotes should be doubled for SQL escaping
     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')"

tests/tools/test_models.py

@@ -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

tests/tools/test_qa.py

@@ -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