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

rename RLM agent to analysis throughout the codebase

Browse source
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
Show stats Download patch file Download diff file Expand all files Collapse all files

8
CHANGELOG.md
View file

@ -21,6 +21,14 @@
- **`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 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

8
README.md
View file

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

22
docs/agents/rlm.md → docs/agents/analysis.md
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?"
- **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
# 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)
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
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
@ -35,18 +35,18 @@ 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'?")
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.rlm(
result = await client.analyze(
"What is the total revenue?",
filter="title LIKE '%Financial%'"
)
# Pre-load specific documents
result = await client.rlm(
result = await client.analyze(
"Compare the conclusions",
documents=["Report A", "Report B"]
)
@ -89,7 +89,7 @@ The `filter` parameter restricts what documents the agent can access. Unlike too
```python
# Agent can only see documents with "confidential" in the URI
result = await client.rlm(
result = await client.analyze(
"Summarize all findings",
filter="uri LIKE '%confidential%'"
)
@ -99,10 +99,10 @@ This is useful for scoping to specific document sets, enforcing access control,
## Configuration
RLM settings can be configured in `haiku.rag.yaml`:
Analysis settings can be configured in `haiku.rag.yaml`:
```yaml
rlm:
analysis:
model:
provider: anthropic
name: claude-sonnet-4-20250514

2
docs/agents/index.md
View file

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

6
docs/architecture.md
View file

@ -26,7 +26,7 @@ flowchart TB
QA[QA Agent]
Skill[RAG Skill]
Research[Research Graph]
RLM[RLM Agent]
Analysis[Analysis Agent]
end
subgraph Apps["Applications"]
@ -123,7 +123,7 @@ flowchart TB
Eval -->|Done| Synthesize[Synthesize]
end
subgraph RLM["RLM Agent"]
subgraph AnalysisAgent["Analysis Agent"]
Q4[Question] --> Code[Write Code]
Code --> Execute[Execute]
Execute --> Examine[Examine Results]
@ -151,7 +151,7 @@ flowchart TB
- Prior answers let the planner skip redundant searches
- 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
- Executes in sandboxed environment

10
docs/cli.md
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.
## 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 +245,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

2
docs/configuration/providers.md
View file

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

8
docs/configuration/qa-research.md
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.
## 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 +76,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.

4
docs/index.md
View file

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

2
docs/mcp.md
View file

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

12
docs/python.md
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 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

18
docs/skills/rlm.md → docs/skills/analysis.md
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?)`
```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)
```
@ -29,10 +29,10 @@ skill = create_skill(db_path=db_path, config=config)
## State
The skill manages an `RLMState` under the `"rlm"` namespace:
The skill manages an `AnalysisState` under the `"analysis"` namespace:
```python
class RLMState(BaseModel):
class AnalysisState(BaseModel):
document_filter: str | None = None
analyses: list[AnalysisEntry] = []
@ -51,14 +51,14 @@ 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.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)
rlm = create_rlm_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, rlm])
analysis = create_analysis_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, analysis])
agent = Agent(
"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.

6
docs/skills/index.md
View file

@ -7,7 +7,7 @@ haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggoz
| Skill | Description |
|-------|-------------|
| [`rag`](rag.md) | Search, retrieve, and answer questions from the knowledge base |
| [`rag-rlm`](rlm.md) | Computational analysis via code execution |
| [`rag-analysis`](analysis.md) | Computational analysis via code execution |
## Discovery
@ -16,7 +16,7 @@ Skills are registered as Python entrypoints under `haiku.skills`. They are disco
```bash
haiku-skills list --use-entrypoints
# rag — Search, retrieve and analyze documents using RAG.
# rag-rlm — Analyze documents using code execution in a sandboxed interpreter.
# rag-analysis — Analyze documents using code execution in a sandboxed interpreter.
```
## Usage
@ -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.

16
haiku_rag_slim/haiku/rag/agents/analysis/__init__.py Normal file
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",
]

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

@ -0,0 +1,59 @@
from pydantic_ai import Agent, RunContext
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import 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

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

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

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

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

2
haiku_rag_slim/haiku/rag/agents/rlm/prompts.py → haiku_rag_slim/haiku/rag/agents/analysis/prompts.py
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.

12
haiku_rag_slim/haiku/rag/agents/rlm/sandbox.py → haiku_rag_slim/haiku/rag/agents/analysis/sandbox.py
View file

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

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

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

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

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

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

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

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

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

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

@ -368,8 +368,8 @@ def ask( # pragma: no cover
)
@_cli.command("rlm", help="Answer questions using code execution (RLM agent)")
def rlm( # pragma: no cover
@_cli.command("analyze", help="Answer questions using code execution (analysis agent)")
def analyze( # pragma: no cover
question: str = typer.Argument(
help="The question to answer",
),
@ -393,7 +393,7 @@ def rlm( # pragma: no cover
):
app = create_app(db)
asyncio.run(
app.rlm(
app.analyze(
question=question,
document=document,
filter=filter,

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

@ -30,11 +30,11 @@ from haiku.rag.utils import escape_sql_string
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.analysis.models import AnalysisResult
from haiku.rag.agents.research.models import (
Citation,
ResearchReport,
)
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__)
@ -1192,17 +1192,17 @@ class HaikuRAG:
return await graph.run(state=state, deps=deps)
async def rlm(
async def analyze(
self,
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> "RLMResult":
"""Answer a question using the RLM agent with code execution.
) -> "AnalysisResult":
"""Answer a question using the analysis agent with code execution.
The RLM (Recursive Language Model) agent can write and execute Python
code in a sandboxed environment to solve problems that require
computation, aggregation, or complex traversal across documents.
The analysis agent can write and execute Python code in a sandboxed
environment to solve problems that require computation, aggregation,
or complex traversal across documents.
Args:
question: The question to answer.
@ -1210,16 +1210,16 @@ class HaikuRAG:
filter: SQL WHERE clause to filter documents during searches.
Returns:
RLMResult with the answer and the final consolidated program.
AnalysisResult with the answer and the final consolidated program.
"""
from haiku.rag.agents.rlm import (
RLMContext,
RLMDeps,
from haiku.rag.agents.analysis import (
AnalysisContext,
AnalysisDeps,
Sandbox,
create_rlm_agent,
create_analysis_agent,
)
context = RLMContext(filter=filter)
context = AnalysisContext(filter=filter)
if documents:
loaded_docs = []
@ -1234,12 +1234,12 @@ class HaikuRAG:
config=self._config,
context=context,
)
deps = RLMDeps(
deps = AnalysisDeps(
sandbox=sandbox,
context=context,
)
agent = create_rlm_agent(self._config)
agent = create_analysis_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output

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

@ -96,7 +96,7 @@ class ResearchConfig(BaseModel):
max_concurrency: int = 1
class RLMConfig(BaseModel):
class AnalysisConfig(BaseModel):
model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
@ -219,7 +219,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)

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

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

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

@ -216,7 +216,7 @@ async def skill_analyze(
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)
result = await rag.analyze(question, documents=documents, filter=filter)
output = result.answer
if result.program:
output += f"\n\nProgram:\n{result.program}"

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

@ -10,15 +10,15 @@ 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] = []
STATE_TYPE = RLMState
STATE_NAMESPACE = "rlm"
STATE_TYPE = AnalysisState
STATE_NAMESPACE = "analysis"
_skill_path = Path(__file__).parent / "rag-rlm"
_skill_path = Path(__file__).parent / "rag-analysis"
@cache
@ -45,7 +45,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 +67,7 @@ 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, ["analyze"])
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()

4
haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md → haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
View file

@ -1,5 +1,5 @@
---
name: rag-rlm
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,
@ -8,6 +8,6 @@ description: >
"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.

2
haiku_rag_slim/pyproject.toml
View file

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

4
mkdocs.yml
View file

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

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

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

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

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

@ -3,26 +3,26 @@ from pathlib import Path
import pytest
from pydantic_ai import Agent
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.agents.analysis.agent import create_analysis_agent
from haiku.rag.agents.analysis.dependencies import AnalysisDeps
from haiku.rag.agents.analysis.models import AnalysisResult, CodeExecution
from haiku.rag.config import AppConfig, Config
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_rlm")
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_analysis")
class TestCreateRLMAgent:
class TestCreateAnalysisAgent:
def test_creates_agent(self):
agent = create_rlm_agent(Config)
agent = create_analysis_agent(Config)
assert isinstance(agent, Agent)
assert agent.deps_type is RLMDeps
assert agent.output_type is RLMResult
assert agent.deps_type is AnalysisDeps
assert agent.output_type is AnalysisResult
def test_agent_has_execute_code_tool(self):
agent = create_rlm_agent(Config)
agent = create_analysis_agent(Config)
tool_names = list(agent._function_toolset.tools.keys())
assert "execute_code" in tool_names
@ -42,13 +42,13 @@ class TestCodeExecutionModel:
assert execution.success is True
class TestClientRLMIntegration:
"""Integration tests for client.rlm() method."""
class TestClientAnalysisIntegration:
"""Integration tests for client.analyze() method."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_count_documents(self, allow_model_requests, temp_db_path):
"""Test RLM agent can count documents.
async def test_analyze_count_documents(self, allow_model_requests, temp_db_path):
"""Test analysis agent can count documents.
Agent program:
docs = list_documents(limit=1000)
@ -63,14 +63,14 @@ class TestClientRLMIntegration:
await client.create_document("Second document about dogs.", title="Doc 2")
await client.create_document("Third document about birds.", title="Doc 3")
result = await client.rlm("How many documents are in the database?")
result = await client.analyze("How many documents are in the database?")
assert "3" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_aggregation(self, allow_model_requests, temp_db_path):
"""Test RLM agent can perform aggregation across documents.
async def test_analyze_aggregation(self, allow_model_requests, temp_db_path):
"""Test analysis agent can perform aggregation across documents.
Agent program:
import re
@ -103,7 +103,7 @@ class TestClientRLMIntegration:
"Sales report Q3: Revenue was $200,000.", title="Q3 Report"
)
result = await client.rlm(
result = await client.analyze(
"What is the total revenue across all quarterly reports?"
)
@ -111,8 +111,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_with_filter(self, allow_model_requests, temp_db_path):
"""Test RLM agent respects filter parameter.
async def test_analyze_with_filter(self, allow_model_requests, temp_db_path):
"""Test analysis agent respects filter parameter.
Agent program:
docs = list_documents(limit=1000)
@ -130,7 +130,7 @@ class TestClientRLMIntegration:
await client.create_document("Dog document.", title="Dogs")
await client.create_document("Bird document.", title="Birds")
result = await client.rlm(
result = await client.analyze(
"How many documents are available?",
filter="title = 'Cats'",
)
@ -139,8 +139,10 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_search_and_get_chunk(self, allow_model_requests, temp_db_path):
"""Test RLM agent can search and use get_chunk for citations.
async def test_analyze_search_and_get_chunk(
self, allow_model_requests, temp_db_path
):
"""Test analysis agent can search and use get_chunk for citations.
Agent program:
results = search("content", limit=5)
@ -158,7 +160,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 +169,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 +211,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 +222,8 @@ class TestClientRLMIntegration:
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_rlm_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test RLM agent can use search() to find content and extract information.
async def test_analyze_search_and_extract(self, allow_model_requests, temp_db_path):
"""Test analysis agent can use search() to find content and extract information.
Agent program:
results = search("document element types", limit=20)
@ -242,7 +244,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 +279,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 +305,7 @@ class TestClientRLMIntegration:
title="Mission Statement",
)
result = await client.rlm(
result = await client.analyze(
"Using the pre-loaded documents variable, "
"tell me when was the company founded and what is their mission?",
documents=["Company History", "Mission Statement"],

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

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

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

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

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

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

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

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

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

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

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

106
tests/skills/test_rlm.py → tests/skills/test_analysis.py
View file

@ -1,12 +1,12 @@
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.config.models import AppConfig
from haiku.rag.skills.rlm import (
from haiku.rag.skills.analysis import (
STATE_NAMESPACE,
STATE_TYPE,
RLMState,
AnalysisState,
instructions,
skill_metadata,
state_metadata,
@ -16,24 +16,24 @@ 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
class TestAnalysisModuleAPI:
def test_state_type_is_analysis_state(self):
assert STATE_TYPE is AnalysisState
def test_state_namespace(self):
assert STATE_NAMESPACE == "rlm"
assert STATE_NAMESPACE == "analysis"
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()
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-rlm"
assert result.name == "rag-analysis"
def test_instructions_returns_string(self):
result = instructions()
@ -41,7 +41,7 @@ class TestRLMModuleAPI:
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
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
@ -50,31 +50,31 @@ class TestRLMModuleAPI:
assert skill.instructions == instructions()
class TestRLMSkillCreation:
class TestAnalysisSkillCreation:
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)
assert skill.metadata.name == "rag-rlm"
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.rlm import create_skill
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 == {"analyze"}
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)
assert skill._state_type is RLMState
assert skill._state_namespace == "rlm"
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.rlm import create_skill
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
@ -86,22 +86,22 @@ class TestRLMSkillCreation:
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
from haiku.rag.skills.analysis import 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):
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)
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
from haiku.rag.skills.analysis import create_skill, instructions
config = AppConfig(
prompts=PromptsConfig(
@ -120,12 +120,12 @@ class TestDomainPreambleInRLMSkillInstructions:
class TestAnalyzeTool:
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(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
"analyze",
AsyncMock(return_value=AnalysisResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
@ -137,17 +137,17 @@ class TestAnalyzeTool:
assert "print(42)" in result
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(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
"analyze",
AsyncMock(return_value=AnalysisResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
state = RLMState()
state = AnalysisState()
ctx = _make_ctx(state)
await analyze(ctx, question="How many documents?")
assert len(state.analyses) == 1
@ -155,42 +155,20 @@ class TestAnalyzeTool:
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
async def test_analyze_with_document_filter_in_state(self, rag_db, monkeypatch):
from haiku.rag.skills.analysis import AnalysisState, create_skill
captured_kwargs = {}
async def mock_rlm(self, question, **kwargs):
async def mock_analyze(self, question, **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)
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'")
state = AnalysisState(document_filter="title = 'AI Overview'")
ctx = _make_ctx(state)
await analyze(
ctx,
@ -203,15 +181,15 @@ class TestAnalyzeTool:
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
from haiku.rag.skills.analysis import create_skill
captured_kwargs = {}
async def mock_rlm(self, question, **kwargs):
async def mock_analyze(self, question, **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)
analyze = _get_tool(skill, "analyze")