Break apart RLM analysis as a separate skill
This commit is contained in:
parent
37f28ea8de
commit
a1c09e1a47
7 changed files with 227 additions and 72 deletions
33
haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md
Normal file
33
haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md
Normal file
|
|
@ -0,0 +1,33 @@
|
|||
---
|
||||
name: rag-rlm
|
||||
description: Analyze documents using code execution in a Docker sandbox.
|
||||
---
|
||||
|
||||
# RLM (Reflexion Language Model) Analysis
|
||||
|
||||
You have access to a computational analysis tool that can write and execute Python code against the knowledge base.
|
||||
|
||||
## When to use `analyze`
|
||||
|
||||
Use the `analyze` tool for questions that require:
|
||||
|
||||
- **Computation** — counting, aggregation, averages, statistics
|
||||
- **Data traversal** — iterating over documents, comparing tables, extracting structured data
|
||||
- **Code execution** — any question best answered by writing and running Python code
|
||||
- **Complex reasoning** — multi-step analysis that goes beyond simple search or Q&A
|
||||
|
||||
Examples:
|
||||
- "How many pages are in this document?"
|
||||
- "Compare the results in table 3 across all documents"
|
||||
- "Calculate the average word count per document"
|
||||
- "Write code to extract all email addresses"
|
||||
|
||||
## Requirements
|
||||
|
||||
The analyze tool requires Docker to be running, as code execution happens in an isolated Docker sandbox.
|
||||
|
||||
## Parameters
|
||||
|
||||
- `question` (required) — The analytical question to answer
|
||||
- `document` — Optional document ID or title to pre-load for analysis
|
||||
- `filter` — Optional SQL WHERE clause to filter documents
|
||||
|
|
@ -267,38 +267,6 @@ def create_skill(
|
|||
|
||||
return answer
|
||||
|
||||
async def analyze(
|
||||
ctx: RunContext[SkillRunDeps],
|
||||
question: str,
|
||||
document: str | None = None,
|
||||
filter: str | None = None,
|
||||
) -> str:
|
||||
"""Answer complex analytical questions using code execution.
|
||||
|
||||
Use this for questions requiring computation, aggregation, or
|
||||
data traversal across documents.
|
||||
|
||||
Args:
|
||||
question: The question to answer.
|
||||
document: Optional document ID or title to pre-load for analysis.
|
||||
filter: Optional SQL WHERE clause to filter documents.
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
||||
documents = [document] if document else None
|
||||
result = await rag.rlm(question, documents=documents, filter=filter)
|
||||
output = result.answer
|
||||
if result.program:
|
||||
output += f"\n\nProgram:\n{result.program}"
|
||||
|
||||
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
|
||||
ctx.deps.state.qa_history.append(
|
||||
QAHistoryEntry(question=question, answer=output)
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
|
||||
"""Conduct deep multi-agent research on a question.
|
||||
|
||||
|
|
@ -368,7 +336,6 @@ def create_skill(
|
|||
list_documents,
|
||||
get_document,
|
||||
ask,
|
||||
analyze,
|
||||
research,
|
||||
],
|
||||
state_type=RAGState,
|
||||
|
|
|
|||
|
|
@ -14,13 +14,11 @@ Use your tools to search and answer questions. Never make up information — alw
|
|||
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
|
||||
- **search** — Use when the user wants to find relevant passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns matching chunks with metadata.
|
||||
- **ask** — Use for questions about topics in the knowledge base (e.g., "what is DocLayNet?", "explain the methodology"). Returns an answer with citations. Always include the citations in your response.
|
||||
- **analyze** — Use for any question that involves code, computation, counting, aggregation, comparison, or complex reasoning (e.g., "how many pages?", "compare the results in table 3", "write code to find the longest word", "calculate the average"). The analyze tool can write and execute Python code with full access to the knowledge base. **When in doubt between search and analyze, prefer analyze** — it can search internally and also compute over results.
|
||||
- **research** — Deep multi-agent research that produces comprehensive reports. **Only use when the user explicitly requests deep research** (e.g., "do a deep research on X", "research this topic thoroughly"). Never call this tool on your own — it is slow and expensive.
|
||||
|
||||
## When search returns irrelevant results
|
||||
|
||||
If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
|
||||
- Use **analyze** if the question involves computation or code
|
||||
- Use **ask** if the question is factual
|
||||
- Report that the knowledge base doesn't contain relevant information
|
||||
|
||||
|
|
|
|||
95
haiku_rag_slim/haiku/rag/skills/rlm.py
Normal file
95
haiku_rag_slim/haiku/rag/skills/rlm.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import RunContext
|
||||
|
||||
from haiku.skills.models import Skill, SkillSource
|
||||
from haiku.skills.parser import parse_skill_md
|
||||
from haiku.skills.state import SkillRunDeps
|
||||
|
||||
|
||||
class AnalysisEntry(BaseModel):
|
||||
question: str
|
||||
answer: str
|
||||
program: str | None = None
|
||||
|
||||
|
||||
class RLMState(BaseModel):
|
||||
analyses: list[AnalysisEntry] = []
|
||||
|
||||
|
||||
def create_skill(
|
||||
db_path: Path | None = None,
|
||||
config: Any = None,
|
||||
) -> Skill:
|
||||
"""Create an RLM analysis skill for computational document analysis.
|
||||
|
||||
Args:
|
||||
db_path: Path to the LanceDB database. Resolved from:
|
||||
1. This argument
|
||||
2. HAIKU_RAG_DB environment variable
|
||||
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
|
||||
config: haiku.rag AppConfig instance. If None, uses get_config().
|
||||
"""
|
||||
from haiku.rag.config import get_config
|
||||
|
||||
if config is None:
|
||||
config = get_config()
|
||||
|
||||
if db_path is None:
|
||||
env_db = os.environ.get("HAIKU_RAG_DB")
|
||||
if env_db:
|
||||
db_path = Path(env_db).expanduser()
|
||||
else:
|
||||
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
path = Path(__file__).parent / "rag-rlm"
|
||||
metadata, instructions = parse_skill_md(path / "SKILL.md")
|
||||
|
||||
async def analyze(
|
||||
ctx: RunContext[SkillRunDeps],
|
||||
question: str,
|
||||
document: str | None = None,
|
||||
filter: str | None = None,
|
||||
) -> str:
|
||||
"""Answer complex analytical questions using code execution.
|
||||
|
||||
Use this for questions requiring computation, aggregation, or
|
||||
data traversal across documents.
|
||||
|
||||
Args:
|
||||
question: The question to answer.
|
||||
document: Optional document ID or title to pre-load for analysis.
|
||||
filter: Optional SQL WHERE clause to filter documents.
|
||||
"""
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
|
||||
documents = [document] if document else None
|
||||
result = await rag.rlm(question, documents=documents, filter=filter)
|
||||
output = result.answer
|
||||
if result.program:
|
||||
output += f"\n\nProgram:\n{result.program}"
|
||||
|
||||
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RLMState):
|
||||
ctx.deps.state.analyses.append(
|
||||
AnalysisEntry(
|
||||
question=question,
|
||||
answer=result.answer,
|
||||
program=result.program,
|
||||
)
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
return Skill(
|
||||
metadata=metadata,
|
||||
source=SkillSource.ENTRYPOINT,
|
||||
path=path,
|
||||
instructions=instructions,
|
||||
tools=[analyze],
|
||||
state_type=RLMState,
|
||||
state_namespace="rlm",
|
||||
)
|
||||
|
|
@ -60,6 +60,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"
|
||||
|
||||
[project.scripts]
|
||||
haiku-rag = "haiku.rag.cli:cli"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,6 @@
|
|||
from unittest.mock import AsyncMock
|
||||
|
||||
from haiku.rag.agents.research.models import Citation, ResearchReport
|
||||
from haiku.rag.agents.rlm.models import RLMResult
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.store.models.chunk import SearchResult
|
||||
from haiku.rag.tools.document import DocumentInfo
|
||||
|
|
@ -29,7 +28,6 @@ class TestRAGSkillCreation:
|
|||
"list_documents",
|
||||
"get_document",
|
||||
"ask",
|
||||
"analyze",
|
||||
"research",
|
||||
}
|
||||
|
||||
|
|
@ -413,41 +411,6 @@ class TestAskTool:
|
|||
assert captured_questions[0] == "Explain quantum computing"
|
||||
|
||||
|
||||
class TestAnalyzeTool:
|
||||
async def test_analyze_returns_result(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
|
||||
monkeypatch.setattr(
|
||||
HaikuRAG,
|
||||
"rlm",
|
||||
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
|
||||
)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
analyze = _get_tool(skill, "analyze")
|
||||
ctx = _make_ctx()
|
||||
result = await analyze(ctx, question="How many documents?")
|
||||
assert isinstance(result, str)
|
||||
assert "42" in result
|
||||
|
||||
async def test_analyze_updates_state(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rag import RAGState, create_skill
|
||||
|
||||
monkeypatch.setattr(
|
||||
HaikuRAG,
|
||||
"rlm",
|
||||
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
|
||||
)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
analyze = _get_tool(skill, "analyze")
|
||||
state = RAGState()
|
||||
ctx = _make_ctx(state)
|
||||
await analyze(ctx, question="How many documents?")
|
||||
assert len(state.qa_history) == 1
|
||||
assert state.qa_history[0].question == "How many documents?"
|
||||
|
||||
|
||||
class TestResearchTool:
|
||||
async def test_research_returns_report(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rag import create_skill
|
||||
|
|
|
|||
98
tests/skills/test_rlm.py
Normal file
98
tests/skills/test_rlm.py
Normal file
|
|
@ -0,0 +1,98 @@
|
|||
from unittest.mock import AsyncMock
|
||||
|
||||
from haiku.rag.agents.rlm.models import RLMResult
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
from .conftest import _get_tool, _make_ctx
|
||||
|
||||
|
||||
class TestRLMSkillCreation:
|
||||
def test_create_skill_returns_valid_skill(self, temp_db_path):
|
||||
from haiku.rag.skills.rlm import create_skill
|
||||
|
||||
skill = create_skill(db_path=temp_db_path)
|
||||
assert skill.metadata.name == "rag-rlm"
|
||||
assert skill.metadata.description
|
||||
assert skill.instructions
|
||||
|
||||
def test_create_skill_has_expected_tools(self, temp_db_path):
|
||||
from haiku.rag.skills.rlm import create_skill
|
||||
|
||||
skill = create_skill(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, temp_db_path):
|
||||
from haiku.rag.skills.rlm import RLMState, create_skill
|
||||
|
||||
skill = create_skill(db_path=temp_db_path)
|
||||
assert skill._state_type is RLMState
|
||||
assert skill._state_namespace == "rlm"
|
||||
|
||||
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
|
||||
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
|
||||
from haiku.rag.skills.rlm import create_skill
|
||||
|
||||
skill = create_skill()
|
||||
assert skill.metadata.name == "rag-rlm"
|
||||
|
||||
|
||||
class TestAnalyzeTool:
|
||||
async def test_analyze_returns_result(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rlm import create_skill
|
||||
|
||||
monkeypatch.setattr(
|
||||
HaikuRAG,
|
||||
"rlm",
|
||||
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
|
||||
)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
analyze = _get_tool(skill, "analyze")
|
||||
ctx = _make_ctx()
|
||||
result = await analyze(ctx, question="How many documents?")
|
||||
assert isinstance(result, str)
|
||||
assert "42" in result
|
||||
assert "print(42)" in result
|
||||
|
||||
async def test_analyze_updates_state(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rlm import RLMState, create_skill
|
||||
|
||||
monkeypatch.setattr(
|
||||
HaikuRAG,
|
||||
"rlm",
|
||||
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
|
||||
)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
analyze = _get_tool(skill, "analyze")
|
||||
state = RLMState()
|
||||
ctx = _make_ctx(state)
|
||||
await analyze(ctx, question="How many documents?")
|
||||
assert len(state.analyses) == 1
|
||||
assert state.analyses[0].question == "How many documents?"
|
||||
assert state.analyses[0].answer == "42"
|
||||
assert state.analyses[0].program == "print(42)"
|
||||
|
||||
async def test_analyze_with_document_and_filter(self, rag_db, monkeypatch):
|
||||
from haiku.rag.skills.rlm import create_skill
|
||||
|
||||
captured_kwargs = {}
|
||||
|
||||
async def mock_rlm(self, question, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return RLMResult(answer="Result", program="code()")
|
||||
|
||||
monkeypatch.setattr(HaikuRAG, "rlm", mock_rlm)
|
||||
|
||||
skill = create_skill(db_path=rag_db)
|
||||
analyze = _get_tool(skill, "analyze")
|
||||
ctx = _make_ctx()
|
||||
await analyze(
|
||||
ctx,
|
||||
question="Count pages",
|
||||
document="AI Overview",
|
||||
filter="title = 'AI Overview'",
|
||||
)
|
||||
assert captured_kwargs.get("documents") == ["AI Overview"]
|
||||
assert captured_kwargs.get("filter") == "title = 'AI Overview'"
|
||||
Loading…
Reference in a new issue