rename RLM agent to analysis throughout the codebase

This commit is contained in:
Yiorgis Gozadinos 2026-04-09 15:34:21 +03:00
parent 499a843a43
commit d2b3ba1b59
No known key found for this signature in database
45 changed files with 306 additions and 314 deletions

View file

@ -21,6 +21,14 @@
- **`max_searches` default**: Raised from 3 to 5 — faster expansion makes additional searches inexpensive - **`max_searches` default**: Raised from 3 to 5 — faster expansion makes additional searches inexpensive
- **Improved QA prompt**: Stronger instruction to refuse answering from tangentially related content - **Improved QA prompt**: Stronger instruction to refuse answering from tangentially related content
- **Improved judge prompt**: Asymmetric evaluation — generated answers that are more comprehensive than expected are not penalized - **Improved judge prompt**: Asymmetric evaluation — generated answers that are more comprehensive than expected are not penalized
- **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: `rag-rlm``rag-analysis`, `skills/rlm.py``skills/analysis.py`
- State namespace: `"rlm"``"analysis"`
### Removed ### Removed

View file

@ -11,7 +11,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Question answering** — QA agents with citations (page numbers, section headings) - **Question answering** — QA agents with citations (page numbers, section headings)
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM - **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize - **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 - **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 - **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 - **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 # Research mode — iterative planning and search
haiku-rag research "What are the limitations of the approach?" 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 # Interactive chat — multi-turn conversations with memory
haiku-rag chat 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 - [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs - [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA and research agents - [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 - [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 - [Server](https://ggozad.github.io/haiku.rag/server/) - File monitoring and MCP
- [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration - [MCP](https://ggozad.github.io/haiku.rag/mcp/) - Model Context Protocol integration

View file

@ -1,6 +1,6 @@
# RLM Agent (Recursive Language Model) # Analysis Agent
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: 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?" - **Aggregation**: "How many documents mention security vulnerabilities?"
- **Computation**: "What's the average revenue across all quarterly reports?" - **Computation**: "What's the average revenue across all quarterly reports?"
@ -19,13 +19,13 @@ The RLM agent enables complex analytical tasks by writing and executing Python c
```bash ```bash
# Basic usage # Basic usage
haiku-rag rlm "How many documents are in the database?" haiku-rag analyze "How many documents are in the database?"
# With document filter (restricts what the agent can access) # With document filter (restricts what the agent can access)
haiku-rag rlm "Summarize the key points" --filter "uri LIKE '%report%'" haiku-rag analyze "Summarize the key points" --filter "uri LIKE '%report%'"
# Pre-load specific documents # Pre-load specific documents
haiku-rag rlm "Compare these two reports" --document "Q1 Report" --document "Q2 Report" haiku-rag analyze "Compare these two reports" --document "Q1 Report" --document "Q2 Report"
``` ```
## Python Usage ## Python Usage
@ -35,18 +35,18 @@ from haiku.rag.client import HaikuRAG
async with HaikuRAG(path_to_db) as client: async with HaikuRAG(path_to_db) as client:
# Basic question # Basic question
result = await client.rlm("How many documents mention 'security'?") result = await client.analyze("How many documents mention 'security'?")
print(result.answer) # The answer print(result.answer) # The answer
print(result.program) # The final consolidated program print(result.program) # The final consolidated program
# With filter (agent can only see filtered documents) # With filter (agent can only see filtered documents)
result = await client.rlm( result = await client.analyze(
"What is the total revenue?", "What is the total revenue?",
filter="title LIKE '%Financial%'" filter="title LIKE '%Financial%'"
) )
# Pre-load specific documents # Pre-load specific documents
result = await client.rlm( result = await client.analyze(
"Compare the conclusions", "Compare the conclusions",
documents=["Report A", "Report B"] documents=["Report A", "Report B"]
) )
@ -89,7 +89,7 @@ The `filter` parameter restricts what documents the agent can access. Unlike too
```python ```python
# Agent can only see documents with "confidential" in the URI # Agent can only see documents with "confidential" in the URI
result = await client.rlm( result = await client.analyze(
"Summarize all findings", "Summarize all findings",
filter="uri LIKE '%confidential%'" filter="uri LIKE '%confidential%'"
) )
@ -99,10 +99,10 @@ This is useful for scoping to specific document sets, enforcing access control,
## Configuration ## Configuration
RLM settings can be configured in `haiku.rag.yaml`: Analysis settings can be configured in `haiku.rag.yaml`:
```yaml ```yaml
rlm: analysis:
model: model:
provider: anthropic provider: anthropic
name: claude-sonnet-4-20250514 name: claude-sonnet-4-20250514

View file

@ -4,7 +4,7 @@ Three agentic flows are provided by haiku.rag:
- **Simple QA Agent** — a focused question answering agent - **Simple QA Agent** — a focused question answering agent
- **Research Graph** — a multi-step research workflow with question decomposition - **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. 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.

View file

@ -26,7 +26,7 @@ flowchart TB
QA[QA Agent] QA[QA Agent]
Skill[RAG Skill] Skill[RAG Skill]
Research[Research Graph] Research[Research Graph]
RLM[RLM Agent] Analysis[Analysis Agent]
end end
subgraph Apps["Applications"] subgraph Apps["Applications"]
@ -123,7 +123,7 @@ flowchart TB
Eval -->|Done| Synthesize[Synthesize] Eval -->|Done| Synthesize[Synthesize]
end end
subgraph RLM["RLM Agent"] subgraph AnalysisAgent["Analysis Agent"]
Q4[Question] --> Code[Write Code] Q4[Question] --> Code[Write Code]
Code --> Execute[Execute] Code --> Execute[Execute]
Execute --> Examine[Examine Results] Execute --> Examine[Examine Results]
@ -151,7 +151,7 @@ flowchart TB
- Prior answers let the planner skip redundant searches - Prior answers let the planner skip redundant searches
- Synthesizes structured report - Synthesizes structured report
**RLM Agent** - Complex analytical tasks via code execution: **Analysis Agent** - Complex analytical tasks via code execution:
- Writes Python code to explore the knowledge base - Writes Python code to explore the knowledge base
- Executes in sandboxed environment - Executes in sandboxed environment

View file

@ -220,24 +220,24 @@ Flags:
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section. 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: Answer complex analytical questions via code execution:
```bash ```bash
haiku-rag rlm "How many documents mention security?" haiku-rag analyze "How many documents mention security?"
``` ```
Filter to specific documents: Filter to specific documents:
```bash ```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: Pre-load specific documents for comparison:
```bash ```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: Flags:
@ -245,7 +245,7 @@ Flags:
- `--filter` / `-f`: SQL WHERE clause to restrict document access - `--filter` / `-f`: SQL WHERE clause to restrict document access
- `--document` / `-d`: Pre-load a document by title or ID (can repeat) - `--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 ## Create Skill

View file

@ -22,7 +22,7 @@ qa:
**Available options:** **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 - Lower (0.0-0.3): Deterministic, focused responses
- Medium (0.4-0.7): Balanced - Medium (0.4-0.7): Balanced
- Higher (0.8-1.0+): Creative, varied responses - Higher (0.8-1.0+): Creative, varied responses

View file

@ -58,12 +58,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. 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 ```yaml
rlm: analysis:
model: model:
provider: anthropic provider: anthropic
name: claude-sonnet-4-20250514 name: claude-sonnet-4-20250514
@ -76,4 +76,4 @@ rlm:
- **code_timeout**: Maximum seconds for each code execution (default: 60) - **code_timeout**: Maximum seconds for each code execution (default: 60)
- **max_output_chars**: Truncate code output after this many characters (default: 50000) - **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.

View file

@ -8,7 +8,7 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
- **Question answering** — QA agents with citations (page numbers, section headings) - **Question answering** — QA agents with citations (page numbers, section headings)
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM - **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Research agents** — Multi-agent workflows via pydantic-graph: plan, search, evaluate, synthesize - **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 - **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 - **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 - **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM. QA/Research: any model supported by Pydantic AI
@ -65,7 +65,7 @@ haiku-rag chat # Interactive conversation mode
- [Python](python.md) - Python API reference - [Python](python.md) - Python API reference
- [Custom Pipelines](custom-pipelines.md) - Build custom processing workflows - [Custom Pipelines](custom-pipelines.md) - Build custom processing workflows
- [Agents](agents/index.md) - QA, chat, and research agents - [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 - [Applications](apps.md) - Chat TUI, web app, and inspector
- [Server](server.md) - File monitoring and server mode - [Server](server.md) - File monitoring and server mode
- [MCP](mcp.md) - Model Context Protocol integration - [MCP](mcp.md) - Model Context Protocol integration

View file

@ -50,7 +50,7 @@ The MCP server exposes `haiku.rag` as MCP tools for compatible MCP clients like
- `question` (required): The research question - `question` (required): The research question
- Returns a structured research report with findings, conclusions, and sources - 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 - `question` (required): The question to answer
- `filter` (optional): SQL WHERE clause to restrict document access - `filter` (optional): SQL WHERE clause to restrict document access
- `document` (optional): Document title/ID to pre-load (can repeat) - `document` (optional): Document title/ID to pre-load (can repeat)

View file

@ -420,32 +420,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 multiagent research workflow. See also: [Agents](agents/index.md) for details on the QA agent and the multiagent research workflow.
## RLM (Recursive Language Model) ## Analysis
Answer complex analytical questions via code execution: Answer complex analytical questions via code execution:
```python ```python
# Aggregation across documents # 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.answer) # The answer
print(result.program) # The final consolidated program print(result.program) # The final consolidated program
# Computation within a document set # Computation within a document set
result = await client.rlm( result = await client.analyze(
"What is the average deal size mentioned in these contracts?", "What is the average deal size mentioned in these contracts?",
filter="uri LIKE '%contracts%'" filter="uri LIKE '%contracts%'"
) )
# Multi-document comparison # Multi-document comparison
result = await client.rlm( result = await client.analyze(
"What changed between these two versions of the policy?", "What changed between these two versions of the policy?",
documents=["Policy v1.0", "Policy v2.0"] 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 ## Building Custom Agents

View file

@ -1,11 +1,11 @@
# RLM Skill # Analysis 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. 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?)` ## `create_skill(db_path?, config?)`
```python ```python
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=db_path, config=config) skill = create_skill(db_path=db_path, config=config)
``` ```
@ -29,10 +29,10 @@ skill = create_skill(db_path=db_path, config=config)
## State ## State
The skill manages an `RLMState` under the `"rlm"` namespace: The skill manages an `AnalysisState` under the `"analysis"` namespace:
```python ```python
class RLMState(BaseModel): class AnalysisState(BaseModel):
document_filter: str | None = None document_filter: str | None = None
analyses: list[AnalysisEntry] = [] analyses: list[AnalysisEntry] = []
@ -51,14 +51,14 @@ Combine both skills to give the agent full RAG + analysis capabilities:
```python ```python
from haiku.rag.skills.rag import create_skill as create_rag_skill 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.rag.skills.analysis import create_skill as create_analysis_skill
from haiku.skills.agent import SkillToolset from haiku.skills.agent import SkillToolset
from haiku.skills.prompts import build_system_prompt from haiku.skills.prompts import build_system_prompt
from pydantic_ai import Agent from pydantic_ai import Agent
rag = create_rag_skill(db_path=db_path) rag = create_rag_skill(db_path=db_path)
rlm = create_rlm_skill(db_path=db_path) analysis = create_analysis_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, rlm]) toolset = SkillToolset(skills=[rag, analysis])
agent = Agent( agent = Agent(
"openai:gpt-4o", "openai:gpt-4o",
@ -67,4 +67,4 @@ agent = Agent(
) )
``` ```
See the [RLM Agent](../agents/rlm.md) documentation for details on how the underlying agent works. See the [Analysis Agent](../agents/analysis.md) documentation for details on how the underlying agent works.

View file

@ -7,7 +7,7 @@ haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggoz
| Skill | Description | | Skill | Description |
|-------|-------------| |-------|-------------|
| [`rag`](rag.md) | Search, retrieve, and answer questions from the knowledge base | | [`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 ## Discovery
@ -16,7 +16,7 @@ Skills are registered as Python entrypoints under `haiku.skills`. They are disco
```bash ```bash
haiku-skills list --use-entrypoints haiku-skills list --use-entrypoints
# rag — Search, retrieve and analyze documents using RAG. # 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 ## Usage
@ -88,7 +88,7 @@ Each skill manages its own state under a dedicated namespace. State is automatic
```python ```python
rag_state = toolset.get_namespace("rag") 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. See the individual skill pages for state model details.

View file

@ -0,0 +1,16 @@
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
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",
"AnalysisResult",
"CodeExecution",
"Sandbox",
"SandboxResult",
"create_analysis_agent",
]

View file

@ -0,0 +1,59 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
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, AnalysisResult]:
"""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, AnalysisResult] = Agent( # type: ignore[assignment] # ty: ignore[invalid-assignment]
model,
deps_type=AnalysisDeps,
output_type=AnalysisResult,
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 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

View file

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

View file

@ -2,7 +2,7 @@ from pydantic import BaseModel, Field
class CodeExecution(BaseModel): 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") code: str = Field(description="The Python code that was executed")
stdout: str = Field(description="Standard output captured during execution") stdout: str = Field(description="Standard output captured during execution")
@ -10,8 +10,8 @@ class CodeExecution(BaseModel):
success: bool = Field(description="Whether execution completed without error") success: bool = Field(description="Whether execution completed without error")
class RLMResult(BaseModel): class AnalysisResult(BaseModel):
"""Result from RLM agent execution.""" """Result from analysis agent execution."""
answer: str = Field(description="The answer to the user's question") answer: str = Field(description="The answer to the user's question")
program: str = Field(description="The final consolidated program") program: str = Field(description="The final consolidated program")

View file

@ -1,4 +1,4 @@
RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code. 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 described below are ONLY available inside execute_code. Always execute code to answer questions; do not just describe what code would do. 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.

View file

@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Literal
import pydantic_monty import pydantic_monty
from haiku.rag.agents.rlm.dependencies import RLMContext from haiku.rag.agents.analysis.dependencies import AnalysisContext
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.store.compression import decompress_json from haiku.rag.store.compression import decompress_json
@ -34,13 +34,13 @@ class Sandbox:
_client: "HaikuRAG" _client: "HaikuRAG"
_config: AppConfig _config: AppConfig
_context: RLMContext _context: AnalysisContext
def __init__( def __init__(
self, self,
client: "HaikuRAG", client: "HaikuRAG",
config: AppConfig, config: AppConfig,
context: RLMContext, context: AnalysisContext,
): ):
self._client = client self._client = client
self._config = config self._config = config
@ -122,7 +122,7 @@ class Sandbox:
from haiku.rag.utils import get_model from haiku.rag.utils import get_model
model = get_model(config.rlm.model, config) model = get_model(config.analysis.model, config)
agent: Agent[None, str] = Agent(model, output_type=str) agent: Agent[None, str] = Agent(model, output_type=str)
result = await agent.run(prompt) result = await agent.run(prompt)
return result.output return result.output
@ -172,9 +172,9 @@ class Sandbox:
def print_callback(_stream: Literal["stdout"], text: str) -> None: def print_callback(_stream: Literal["stdout"], text: str) -> None:
stdout_lines.append(text) stdout_lines.append(text)
max_chars = self._config.rlm.max_output_chars max_chars = self._config.analysis.max_output_chars
limits: pydantic_monty.ResourceLimits = { limits: pydantic_monty.ResourceLimits = {
"max_duration_secs": self._config.rlm.code_timeout, "max_duration_secs": self._config.analysis.code_timeout,
} }
try: try:

View file

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

View file

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

View file

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

View file

@ -446,13 +446,13 @@ class HaikuRAGApp: # pragma: no cover
for renderable in format_citations_rich(citations): for renderable in format_citations_rich(citations):
self.console.print(renderable) self.console.print(renderable)
async def rlm( async def analyze(
self, self,
question: str, question: str,
document: str | None = None, document: str | None = None,
filter: 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: Args:
question: The question to answer 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(f"[bold blue]Question:[/bold blue] {question}")
self.console.print() 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() 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("[bold yellow]Program:[/bold yellow]")
self.console.print(Syntax(result.program, "python")) self.console.print(Syntax(result.program, "python"))

View file

@ -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( question: str = typer.Argument(
help="The question to answer", help="The question to answer",
), ),
@ -393,7 +393,7 @@ def rlm( # pragma: no cover
): ):
app = create_app(db) app = create_app(db)
asyncio.run( asyncio.run(
app.rlm( app.analyze(
question=question, question=question,
document=document, document=document,
filter=filter, filter=filter,

View file

@ -30,11 +30,11 @@ from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING: if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import ( from haiku.rag.agents.research.models import (
Citation, Citation,
ResearchReport, ResearchReport,
) )
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -1192,17 +1192,17 @@ class HaikuRAG:
return await graph.run(state=state, deps=deps) return await graph.run(state=state, deps=deps)
async def rlm( async def analyze(
self, self,
question: str, question: str,
documents: list[str] | None = None, documents: list[str] | None = None,
filter: 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: Args:
question: The question to answer. question: The question to answer.
@ -1210,16 +1210,16 @@ class HaikuRAG:
filter: SQL WHERE clause to filter documents during searches. filter: SQL WHERE clause to filter documents during searches.
Returns: 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, Sandbox,
create_rlm_agent, create_analysis_agent,
) )
context = RLMContext(filter=filter) context = AnalysisContext(filter=filter)
if documents: if documents:
loaded_docs = [] loaded_docs = []
@ -1234,12 +1234,12 @@ class HaikuRAG:
config=self._config, config=self._config,
context=context, context=context,
) )
deps = RLMDeps( deps = AnalysisDeps(
sandbox=sandbox, sandbox=sandbox,
context=context, context=context,
) )
agent = create_rlm_agent(self._config) agent = create_analysis_agent(self._config)
result = await agent.run(question, deps=deps) result = await agent.run(question, deps=deps)
return result.output return result.output

View file

@ -96,7 +96,7 @@ class ResearchConfig(BaseModel):
max_concurrency: int = 1 max_concurrency: int = 1
class RLMConfig(BaseModel): class AnalysisConfig(BaseModel):
model: ModelConfig = Field( model: ModelConfig = Field(
default_factory=lambda: ModelConfig( default_factory=lambda: ModelConfig(
provider="ollama", provider="ollama",
@ -219,7 +219,7 @@ class AppConfig(BaseModel):
reranking: RerankingConfig = Field(default_factory=RerankingConfig) reranking: RerankingConfig = Field(default_factory=RerankingConfig)
qa: QAConfig = Field(default_factory=QAConfig) qa: QAConfig = Field(default_factory=QAConfig)
research: ResearchConfig = Field(default_factory=ResearchConfig) research: ResearchConfig = Field(default_factory=ResearchConfig)
rlm: RLMConfig = Field(default_factory=RLMConfig) analysis: AnalysisConfig = Field(default_factory=AnalysisConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig) processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig) search: SearchConfig = Field(default_factory=SearchConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig) providers: ProvidersConfig = Field(default_factory=ProvidersConfig)

View file

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

View file

@ -216,7 +216,7 @@ async def skill_analyze(
async with HaikuRAG(db_path, config=config, read_only=True) as rag: async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = [document] if document else None documents = [document] if document else None
result = await rag.rlm(question, documents=documents, filter=filter) result = await rag.analyze(question, documents=documents, filter=filter)
output = result.answer output = result.answer
if result.program: if result.program:
output += f"\n\nProgram:\n{result.program}" output += f"\n\nProgram:\n{result.program}"

View file

@ -10,15 +10,15 @@ from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md from haiku.skills.parser import parse_skill_md
class RLMState(BaseModel): class AnalysisState(BaseModel):
document_filter: str | None = None document_filter: str | None = None
analyses: list[AnalysisEntry] = [] analyses: list[AnalysisEntry] = []
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 @cache
@ -45,7 +45,7 @@ def create_skill(
db_path: Path | None = None, db_path: Path | None = None,
config: AppConfig | None = None, config: AppConfig | None = None,
) -> Skill: ) -> Skill:
"""Create an RLM analysis skill for computational document analysis. """Create an analysis skill for computational document analysis.
Args: Args:
db_path: Path to the LanceDB database. Resolved from: db_path: Path to the LanceDB database. Resolved from:
@ -67,7 +67,7 @@ def create_skill(
else: else:
db_path = config.storage.data_dir / "haiku.rag.lancedb" 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, ["analyze"])
extras = create_skill_extras(db_path, config) extras = create_skill_extras(db_path, config)
skill_instructions = instructions() skill_instructions = instructions()

View file

@ -1,5 +1,5 @@
--- ---
name: rag-rlm name: rag-analysis
description: > description: >
Computational analysis of the knowledge base via code execution in a sandboxed Python interpreter. Computational analysis of the knowledge base via code execution in a sandboxed Python interpreter.
Use for questions requiring counting, aggregation, statistics, data traversal, Use for questions requiring counting, aggregation, statistics, data traversal,
@ -8,6 +8,6 @@ description: >
"calculate average word count", "extract all email addresses". "calculate average word count", "extract all email addresses".
--- ---
# RLM Analysis # Analysis
Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in a sandboxed Python interpreter. Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in a sandboxed Python interpreter.

View file

@ -62,7 +62,7 @@ vertexai = ["pydantic-ai-slim[vertexai]"]
[project.entry-points."haiku.skills"] [project.entry-points."haiku.skills"]
rag = "haiku.rag.skills.rag:create_skill" 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] [project.scripts]
haiku-rag = "haiku.rag.cli:cli" haiku-rag = "haiku.rag.cli:cli"

View file

@ -73,11 +73,11 @@ nav:
- Tuning: tuning.md - Tuning: tuning.md
- Agents: - Agents:
- agents/index.md - agents/index.md
- RLM Agent: agents/rlm.md - Analysis Agent: agents/analysis.md
- Skills: - Skills:
- skills/index.md - skills/index.md
- RAG: skills/rag.md - RAG: skills/rag.md
- RLM: skills/rlm.md - Analysis: skills/analysis.md
- Toolsets: tools.md - Toolsets: tools.md
- Applications: apps.md - Applications: apps.md
- Server: server.md - Server: server.md

View file

@ -1,7 +1,7 @@
import pytest 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.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
@ -17,5 +17,5 @@ async def empty_client(temp_db_path):
async def sandbox(empty_client): async def sandbox(empty_client):
"""Create a Monty sandbox for testing.""" """Create a Monty sandbox for testing."""
config = AppConfig() config = AppConfig()
context = RLMContext() context = AnalysisContext()
return Sandbox(client=empty_client, config=config, context=context) return Sandbox(client=empty_client, config=config, context=context)

View file

@ -3,26 +3,26 @@ from pathlib import Path
import pytest import pytest
from pydantic_ai import Agent 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 AnalysisResult, CodeExecution
from haiku.rag.config import AppConfig, Config from haiku.rag.config import AppConfig, Config
@pytest.fixture(scope="module") @pytest.fixture(scope="module")
def vcr_cassette_dir(): 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): def test_creates_agent(self):
agent = create_rlm_agent(Config) agent = create_analysis_agent(Config)
assert isinstance(agent, Agent) 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 AnalysisResult
def test_agent_has_execute_code_tool(self): 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()) tool_names = list(agent._function_toolset.tools.keys())
assert "execute_code" in tool_names assert "execute_code" in tool_names
@ -42,13 +42,13 @@ class TestCodeExecutionModel:
assert execution.success is True 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.asyncio
@pytest.mark.vcr() @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: Agent program:
docs = list_documents(limit=1000) 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("Second document about dogs.", title="Doc 2")
await client.create_document("Third document about birds.", title="Doc 3") 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 assert "3" in result.answer
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @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: Agent program:
import re import re
@ -103,7 +103,7 @@ class TestClientRLMIntegration:
"Sales report Q3: Revenue was $200,000.", title="Q3 Report" "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?" "What is the total revenue across all quarterly reports?"
) )
@ -111,8 +111,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @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: Agent program:
docs = list_documents(limit=1000) docs = list_documents(limit=1000)
@ -130,7 +130,7 @@ class TestClientRLMIntegration:
await client.create_document("Dog document.", title="Dogs") await client.create_document("Dog document.", title="Dogs")
await client.create_document("Bird document.", title="Birds") await client.create_document("Bird document.", title="Birds")
result = await client.rlm( result = await client.analyze(
"How many documents are available?", "How many documents are available?",
filter="title = 'Cats'", filter="title = 'Cats'",
) )
@ -139,8 +139,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @pytest.mark.vcr()
async def test_rlm_search_and_get_chunk(self, allow_model_requests, temp_db_path): async def test_analyze_search_and_get_chunk(
"""Test RLM agent can search and use get_chunk for citations. self, allow_model_requests, temp_db_path
):
"""Test analysis agent can search and use get_chunk for citations.
Agent program: Agent program:
results = search("content", limit=5) results = search("content", limit=5)
@ -158,7 +160,7 @@ class TestClientRLMIntegration:
title="Animal Facts", title="Animal Facts",
) )
result = await client.rlm( result = await client.analyze(
"Search for content about animals and tell me " "Search for content about animals and tell me "
"which document it came from." "which document it came from."
) )
@ -167,10 +169,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @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 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: Agent program:
docs = list_documents(limit=100) docs = list_documents(limit=100)
@ -209,7 +211,7 @@ class TestClientRLMIntegration:
title="Q3 Update", title="Q3 Update",
) )
result = await client.rlm( result = await client.analyze(
"Analyze the sentiment of each quarterly update. " "Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?" "How many quarters were positive, negative, and mixed?"
) )
@ -220,8 +222,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @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: Agent program:
results = search("document element types", limit=20) results = search("document element types", limit=20)
@ -242,7 +244,7 @@ class TestClientRLMIntegration:
async with HaikuRAG(temp_db_path, config=config, create=True) as client: async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path) 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. " "Search for content about document element types or labels. "
"What are all the different document element types mentioned? " "What are all the different document element types mentioned? "
"List them all." "List them all."
@ -277,10 +279,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.vcr() @pytest.mark.vcr()
async def test_rlm_with_preloaded_documents( async def test_analyze_with_preloaded_documents(
self, allow_model_requests, temp_db_path 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: Agent program:
if 'documents' in dir(): if 'documents' in dir():
@ -303,7 +305,7 @@ class TestClientRLMIntegration:
title="Mission Statement", title="Mission Statement",
) )
result = await client.rlm( result = await client.analyze(
"Using the pre-loaded documents variable, " "Using the pre-loaded documents variable, "
"tell me when was the company founded and what is their mission?", "tell me when was the company founded and what is their mission?",
documents=["Company History", "Mission Statement"], documents=["Company History", "Mission Statement"],

View file

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

View file

@ -2,8 +2,8 @@ from pathlib import Path
import pytest 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.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import Document from haiku.rag.store.models import Document
@ -98,7 +98,7 @@ class TestSandboxHaikuRAG:
title="Test Document", title="Test Document",
) )
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context) sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
"docs = await list_documents()\n" "docs = await list_documents()\n"
@ -121,7 +121,7 @@ class TestSandboxHaikuRAG:
title="Animals", title="Animals",
) )
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context) sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
"results = await search('fox', limit=5)\n" "results = await search('fox', limit=5)\n"
@ -144,7 +144,7 @@ class TestSandboxHaikuRAG:
title="Fox Document", title="Fox Document",
) )
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context) sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
f"content = await get_document('{doc.id}')\n" f"content = await get_document('{doc.id}')\n"
@ -174,7 +174,7 @@ class TestSandboxHaikuRAG:
title="Fox Document", title="Fox Document",
) )
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context) sb = Sandbox(client=client, config=config, context=context)
# First search to get a chunk_id # First search to get a chunk_id
result = await sb.execute( result = await sb.execute(
@ -257,8 +257,8 @@ class TestSandboxOutputTruncation:
async def test_truncate_stdout_on_runtime_error(self, empty_client): async def test_truncate_stdout_on_runtime_error(self, empty_client):
"""Test stdout is truncated when a runtime error occurs after large output.""" """Test stdout is truncated when a runtime error occurs after large output."""
config = AppConfig() config = AppConfig()
config.rlm.max_output_chars = 20 config.analysis.max_output_chars = 20
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=empty_client, config=config, context=context) sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('a' * 100)\nx = 1/0") result = await sb.execute("print('a' * 100)\nx = 1/0")
assert not result.success assert not result.success
@ -270,8 +270,8 @@ class TestSandboxOutputTruncation:
async def test_truncate_successful_output(self, empty_client): async def test_truncate_successful_output(self, empty_client):
"""Test output is truncated on successful execution with large output.""" """Test output is truncated on successful execution with large output."""
config = AppConfig() config = AppConfig()
config.rlm.max_output_chars = 20 config.analysis.max_output_chars = 20
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=empty_client, config=config, context=context) sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute("print('b' * 100)") result = await sb.execute("print('b' * 100)")
assert result.success assert result.success
@ -299,7 +299,7 @@ class TestSandboxContextFilter:
title="Private Doc", 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(client=client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
"docs = await list_documents()\n" "docs = await list_documents()\n"
@ -331,7 +331,7 @@ class TestSandboxPreloadedDocuments:
Document(id="1", content="Content A", title="Doc A", uri="a://1"), Document(id="1", content="Content A", title="Doc A", uri="a://1"),
Document(id="2", content="Content B", title="Doc B", uri="b://2"), Document(id="2", content="Content B", title="Doc B", uri="b://2"),
] ]
context = RLMContext(documents=docs) context = AnalysisContext(documents=docs)
sb = Sandbox(client=empty_client, config=config, context=context) sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
"print(len(documents))\n" "print(len(documents))\n"
@ -368,7 +368,7 @@ class TestSandboxDoclingDocument:
title="Docling Doc", title="Docling Doc",
) )
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=client, config=config, context=context) sb = Sandbox(client=client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
f"doc = await get_docling_document('{doc.id}')\n" f"doc = await get_docling_document('{doc.id}')\n"
@ -389,7 +389,7 @@ class TestSandboxLLM:
async def test_llm_function(self, allow_model_requests, empty_client): async def test_llm_function(self, allow_model_requests, empty_client):
"""Test llm() calls the model and returns a string.""" """Test llm() calls the model and returns a string."""
config = AppConfig() config = AppConfig()
context = RLMContext() context = AnalysisContext()
sb = Sandbox(client=empty_client, config=config, context=context) sb = Sandbox(client=empty_client, config=config, context=context)
result = await sb.execute( result = await sb.execute(
"answer = await llm('What is 2 + 2? Reply with just the number.')\n" "answer = await llm('What is 2 + 2? Reply with just the number.')\n"

View file

@ -1,12 +1,12 @@
from unittest.mock import AsyncMock from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.models import RLMResult from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rlm import ( from haiku.rag.skills.analysis import (
STATE_NAMESPACE, STATE_NAMESPACE,
STATE_TYPE, STATE_TYPE,
RLMState, AnalysisState,
instructions, instructions,
skill_metadata, skill_metadata,
state_metadata, state_metadata,
@ -16,24 +16,24 @@ from haiku.skills.models import SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx from .conftest import _get_tool, _make_ctx
class TestRLMModuleAPI: class TestAnalysisModuleAPI:
def test_state_type_is_rlm_state(self): def test_state_type_is_analysis_state(self):
assert STATE_TYPE is RLMState assert STATE_TYPE is AnalysisState
def test_state_namespace(self): def test_state_namespace(self):
assert STATE_NAMESPACE == "rlm" assert STATE_NAMESPACE == "analysis"
def test_state_metadata_returns_state_metadata(self): def test_state_metadata_returns_state_metadata(self):
result = state_metadata() result = state_metadata()
assert isinstance(result, StateMetadata) assert isinstance(result, StateMetadata)
assert result.namespace == "rlm" assert result.namespace == "analysis"
assert result.type is RLMState assert result.type is AnalysisState
assert result.schema == RLMState.model_json_schema() assert result.schema == AnalysisState.model_json_schema()
def test_skill_metadata_returns_skill_metadata(self): def test_skill_metadata_returns_skill_metadata(self):
result = skill_metadata() result = skill_metadata()
assert isinstance(result, SkillMetadata) assert isinstance(result, SkillMetadata)
assert result.name == "rag-rlm" assert result.name == "rag-analysis"
def test_instructions_returns_string(self): def test_instructions_returns_string(self):
result = instructions() result = instructions()
@ -41,7 +41,7 @@ class TestRLMModuleAPI:
assert len(result) > 0 assert len(result) > 0
def test_constants_match_create_skill(self, test_app_config, temp_db_path): def test_constants_match_create_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.state_type is STATE_TYPE assert skill.state_type is STATE_TYPE
@ -50,31 +50,31 @@ class TestRLMModuleAPI:
assert skill.instructions == instructions() assert skill.instructions == instructions()
class TestRLMSkillCreation: class TestAnalysisSkillCreation:
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path): def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag-rlm" assert skill.metadata.name == "rag-analysis"
assert skill.metadata.description assert skill.metadata.description
assert skill.instructions assert skill.instructions
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path): def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path) 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)} tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {"analyze"} assert tool_names == {"analyze"}
def test_create_skill_has_state(self, test_app_config, temp_db_path): def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import RLMState, create_skill from haiku.rag.skills.analysis import AnalysisState, create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is RLMState assert skill._state_type is AnalysisState
assert skill._state_namespace == "rlm" assert skill._state_namespace == "analysis"
def test_create_skill_has_extras(self, test_app_config, temp_db_path): def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.extras["config"] is test_app_config assert skill.extras["config"] is test_app_config
@ -86,22 +86,22 @@ class TestRLMSkillCreation:
def test_create_skill_from_env(self, monkeypatch, temp_db_path): def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path)) monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
skill = create_skill() skill = create_skill()
assert skill.metadata.name == "rag-rlm" assert skill.metadata.name == "rag-analysis"
class TestDomainPreambleInRLMSkillInstructions: class TestDomainPreambleInAnalysisSkillInstructions:
def test_create_skill_without_domain_preamble(self, test_app_config, temp_db_path): def test_create_skill_without_domain_preamble(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill, instructions from haiku.rag.skills.analysis import create_skill, instructions
skill = create_skill(config=test_app_config, db_path=temp_db_path) skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.instructions == instructions() assert skill.instructions == instructions()
def test_create_skill_with_domain_preamble(self, temp_db_path): def test_create_skill_with_domain_preamble(self, temp_db_path):
from haiku.rag.config.models import PromptsConfig from haiku.rag.config.models import PromptsConfig
from haiku.rag.skills.rlm import create_skill, instructions from haiku.rag.skills.analysis import create_skill, instructions
config = AppConfig( config = AppConfig(
prompts=PromptsConfig( prompts=PromptsConfig(
@ -120,12 +120,12 @@ class TestDomainPreambleInRLMSkillInstructions:
class TestAnalyzeTool: class TestAnalyzeTool:
async def test_analyze_returns_result(self, rag_db, monkeypatch): async def test_analyze_returns_result(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
monkeypatch.setattr( monkeypatch.setattr(
HaikuRAG, HaikuRAG,
"rlm", "analyze",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")), AsyncMock(return_value=AnalysisResult(answer="42", program="print(42)")),
) )
skill = create_skill(db_path=rag_db) skill = create_skill(db_path=rag_db)
@ -137,17 +137,17 @@ class TestAnalyzeTool:
assert "print(42)" in result assert "print(42)" in result
async def test_analyze_updates_state(self, rag_db, monkeypatch): async def test_analyze_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import RLMState, create_skill from haiku.rag.skills.analysis import AnalysisState, create_skill
monkeypatch.setattr( monkeypatch.setattr(
HaikuRAG, HaikuRAG,
"rlm", "analyze",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")), AsyncMock(return_value=AnalysisResult(answer="42", program="print(42)")),
) )
skill = create_skill(db_path=rag_db) skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze") analyze = _get_tool(skill, "analyze")
state = RLMState() state = AnalysisState()
ctx = _make_ctx(state) ctx = _make_ctx(state)
await analyze(ctx, question="How many documents?") await analyze(ctx, question="How many documents?")
assert len(state.analyses) == 1 assert len(state.analyses) == 1
@ -155,42 +155,20 @@ class TestAnalyzeTool:
assert state.analyses[0].answer == "42" assert state.analyses[0].answer == "42"
assert state.analyses[0].program == "print(42)" assert state.analyses[0].program == "print(42)"
async def test_analyze_applies_document_filter_from_state( async def test_analyze_with_document_filter_in_state(self, rag_db, monkeypatch):
self, rag_db, monkeypatch from haiku.rag.skills.analysis import AnalysisState, create_skill
):
from haiku.rag.skills.rlm import RLMState, create_skill
captured_kwargs = {} captured_kwargs = {}
async def mock_rlm(self, question, **kwargs): async def mock_analyze(self, question, **kwargs):
captured_kwargs.update(kwargs) captured_kwargs.update(kwargs)
return RLMResult(answer="42", program="print(42)") return AnalysisResult(answer="Result", program="code()")
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm) monkeypatch.setattr(HaikuRAG, "analyze", mock_analyze)
skill = create_skill(db_path=rag_db) skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze") analyze = _get_tool(skill, "analyze")
state = RLMState(document_filter="title = 'AI Overview'") state = AnalysisState(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) ctx = _make_ctx(state)
await analyze( await analyze(
ctx, ctx,
@ -203,15 +181,15 @@ class TestAnalyzeTool:
assert "uri LIKE '%test%'" in result_filter assert "uri LIKE '%test%'" in result_filter
async def test_analyze_with_document_and_filter(self, rag_db, monkeypatch): async def test_analyze_with_document_and_filter(self, rag_db, monkeypatch):
from haiku.rag.skills.rlm import create_skill from haiku.rag.skills.analysis import create_skill
captured_kwargs = {} captured_kwargs = {}
async def mock_rlm(self, question, **kwargs): async def mock_analyze(self, question, **kwargs):
captured_kwargs.update(kwargs) captured_kwargs.update(kwargs)
return RLMResult(answer="Result", program="code()") return AnalysisResult(answer="Result", program="code()")
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm) monkeypatch.setattr(HaikuRAG, "analyze", mock_analyze)
skill = create_skill(db_path=rag_db) skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze") analyze = _get_tool(skill, "analyze")