rename RLM agent to analysis throughout the codebase
This commit is contained in:
parent
499a843a43
commit
d2b3ba1b59
45 changed files with 306 additions and 314 deletions
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
@ -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.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
10
docs/cli.md
10
docs/cli.md
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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.
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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.
|
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:
|
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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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.
|
||||||
|
|
@ -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.
|
||||||
|
|
|
||||||
16
haiku_rag_slim/haiku/rag/agents/analysis/__init__.py
Normal file
16
haiku_rag_slim/haiku/rag/agents/analysis/__init__.py
Normal 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
59
haiku_rag_slim/haiku/rag/agents/analysis/agent.py
Normal 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
23
haiku_rag_slim/haiku/rag/agents/analysis/dependencies.py
Normal 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)
|
||||||
|
|
@ -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")
|
||||||
|
|
@ -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.
|
||||||
|
|
||||||
|
|
@ -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:
|
||||||
|
|
@ -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",
|
|
||||||
]
|
|
||||||
|
|
@ -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
|
|
||||||
|
|
@ -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)
|
|
||||||
|
|
@ -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"))
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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}"
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
@ -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.
|
||||||
|
|
@ -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"
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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)
|
||||||
|
|
@ -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"],
|
||||||
|
|
@ -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)"
|
||||||
|
|
@ -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"
|
||||||
|
|
@ -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")
|
||||||
Loading…
Reference in a new issue