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

Integrate with client, cli, app, mcp

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-29 18:12:11 +02:00
parent 75de81accf
commit b68b2393e9
No known key found for this signature in database
15 changed files with 262 additions and 71 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,12 +1,11 @@
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext, RLMDeps
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 REPLEnvironment, REPLResult
__all__ = [
"CodeExecution",
"RLMConfig",
"RLMContext",
"RLMDeps",
"RLMResult",

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

@ -16,7 +16,7 @@ def _get_or_create_repl(ctx) -> REPLEnvironment:
if key not in _repl_cache:
_repl_cache[key] = REPLEnvironment(
client=ctx.deps.client,
config=ctx.deps.rlm_config,
config=ctx.deps.config.rlm,
context=ctx.deps.context,
)
return _repl_cache[key]
@ -35,7 +35,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
Returns:
A pydantic-ai Agent configured for RLM execution.
"""
model = get_model(config.qa.model, config)
model = get_model(config.rlm.model, config)
agent: Agent[RLMDeps, RLMResult] = Agent( # type: ignore[invalid-assignment]
model,

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

@ -1,8 +1,6 @@
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
from pydantic import BaseModel
from haiku.rag.store.models import Document, SearchResult
if TYPE_CHECKING:
@ -11,19 +9,12 @@ if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
class RLMConfig(BaseModel):
"""Configuration for RLM agent sandbox execution."""
code_timeout: float = 60.0
max_output_chars: int = 50_000
max_tool_calls: int = 20
@dataclass
class RLMContext:
"""Mutable context accumulating data during RLM execution."""
documents: list[Document] | None = None
filter: str | None = None
search_results: list[SearchResult] = field(default_factory=list)
code_executions: "list[CodeExecution]" = field(default_factory=list)
@ -34,5 +25,4 @@ class RLMDeps:
client: "HaikuRAG"
config: "AppConfig"
rlm_config: RLMConfig = field(default_factory=RLMConfig)
context: RLMContext = field(default_factory=RLMContext)

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

@ -1,7 +1,5 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
class CodeExecution(BaseModel):
"""Result of executing a code block in the RLM sandbox."""
@ -16,10 +14,6 @@ class RLMResult(BaseModel):
"""Result from RLM agent execution."""
answer: str = Field(description="The answer to the user's question")
citations: list[Citation] = Field(
default_factory=list,
description="Citations for sources used in the answer",
)
code_executions: list[CodeExecution] = Field(
default_factory=list,
description="History of code executions during the RLM session",

35
haiku_rag_slim/haiku/rag/agents/rlm/prompts.py
View file

@ -1,14 +1,20 @@
RLM_SYSTEM_PROMPT = """You are a Recursive Language Model (RLM) agent that solves complex research questions by writing and executing Python code.
You have access to a sandboxed Python environment with these haiku.rag functions:
IMPORTANT: You MUST use the `execute_code` tool to run Python code. The functions described below are ONLY available inside the execute_code tool - you cannot access them any other way. Always execute code to answer questions; do not just describe what code would do.
CRITICAL: Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just use them directly:
- search("query") ✓ CORRECT
- from haiku.rag import search ✗ WRONG - will fail
You have access to a sandboxed Python environment with these haiku.rag functions (use them directly, no imports needed):
## Available Functions
### search(query, limit=10, filter=None) -> list[dict]
### search(query, limit=10) -> list[dict]
Search the knowledge base using hybrid search (vector + full-text).
Returns list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings
### list_documents(limit=10, offset=0, filter=None) -> list[dict]
### list_documents(limit=10, offset=0) -> list[dict]
List available documents in the knowledge base.
Returns list of dicts with keys: id, title, uri, created_at
@ -21,7 +27,7 @@ Get the structured DoclingDocument object for advanced analysis.
Returns a DoclingDocument object, or None if not found.
See "DoclingDocument API" section below for how to use it.
### ask(question, filter=None) -> str
### ask(question) -> str
Ask a question using the QA agent with RAG. Returns the answer as a string.
Use this for semantic analysis that benefits from LLM reasoning.
@ -30,12 +36,13 @@ You can import: json, re, collections, math, statistics, itertools, functools, d
## Strategy Guide
1. **Explore First**: Start by listing documents or searching to understand what's available.
2. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
3. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with.
4. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections.
5. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function.
6. **Cite Your Sources**: Track which documents/chunks informed your answer for citation.
1. **Explore First**: Start by listing documents or searching to understand what's available. Document names may differ from filenames (e.g., "tbmed593.pdf" might be stored as "TB MED 593" or similar).
2. **If get_document returns None**: Use `list_documents()` to see actual document titles, or `search()` to find relevant content.
3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
4. **Use print() Liberally**: The REPL captures stdout - print intermediate results to see what you're working with.
5. **Aggregate with Code**: For counting, averaging, or comparing across documents, write loops and use collections.
6. **Use ask() for Semantic Analysis**: When you need to understand meaning or interpret content, use the ask() function.
7. **Cite Your Sources**: Track which documents/chunks informed your answer for citation.
## DoclingDocument API
@ -114,6 +121,12 @@ summary = ask("What are the main machine learning approaches discussed?")
print(summary)
```
## Workflow
1. **ALWAYS start by using execute_code** to explore the knowledge base
2. Run multiple code blocks as needed to gather information
3. After collecting data, provide your final answer
## Output Format
After executing code and gathering information, provide:
@ -121,4 +134,4 @@ After executing code and gathering information, provide:
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
Remember: You're solving problems that require computation, aggregation, or complex traversal - things traditional RAG can't do well. Write code to do the heavy lifting."""
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

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

@ -6,7 +6,8 @@ import traceback
from io import StringIO
from typing import TYPE_CHECKING, Any
from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.config.models import RLMConfig
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
@ -151,11 +152,11 @@ class REPLEnvironment:
def _make_search(self):
"""Create sync search function that bridges to async client."""
def search(
query: str, limit: int = 10, filter: str | None = None
) -> list[dict]:
def search(query: str, limit: int = 10) -> list[dict]:
async def _search():
return await self.client.search(query, limit=limit, filter=filter)
return await self.client.search(
query, limit=limit, filter=self.context.filter
)
results = self._run_async_from_thread(_search())
self.context.search_results.extend(results)
@ -178,12 +179,10 @@ class REPLEnvironment:
def _make_list_documents(self):
"""Create sync list_documents function."""
def list_documents(
limit: int = 10, offset: int = 0, filter: str | None = None
) -> list[dict]:
def list_documents(limit: int = 10, offset: int = 0) -> list[dict]:
async def _list():
return await self.client.list_documents(
limit=limit, offset=offset, filter=filter
limit=limit, offset=offset, filter=self.context.filter
)
docs = self._run_async_from_thread(_list())
@ -250,9 +249,11 @@ class REPLEnvironment:
def _make_ask(self):
"""Create sync ask function that uses QA agent."""
def ask(question: str, filter: str | None = None) -> str:
def ask(question: str) -> str:
async def _ask():
answer, citations = await self.client.ask(question, filter=filter)
answer, citations = await self.client.ask(
question, filter=self.context.filter
)
for c in citations:
for sr in self.context.search_results:
if sr.chunk_id == c.chunk_id:

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

@ -432,6 +432,37 @@ class HaikuRAGApp:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
async def rlm(
self,
question: str,
document: str | None = None,
filter: str | None = None,
):
"""Answer a question using the RLM agent with code execution.
Args:
question: The question to answer
document: Optional document ID or title to pre-load
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
documents = [document] if document else None
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[dim]Running RLM agent with code execution...[/dim]")
self.console.print()
answer = await self.client.rlm(question, documents=documents, filter=filter)
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
async def research(
self,
question: str,

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

@ -364,6 +364,39 @@ def ask(
)
@_cli.command("rlm", help="Answer questions using code execution (RLM agent)")
def rlm(
question: str = typer.Argument(
help="The question to answer",
),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
document: str | None = typer.Option(
None,
"--document",
"-d",
help="Document ID or title to pre-load for analysis",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
),
):
app = create_app(db)
asyncio.run(
app.rlm(
question=question,
document=document,
filter=filter,
)
)
@_cli.command("research", help="Run multi-agent research and output a concise report")
def research(
question: str = typer.Argument(..., help="The research question to investigate"),

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

@ -1293,6 +1293,53 @@ class HaikuRAG:
qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
return await qa_agent.answer(question, filter=filter)
async def rlm(
self,
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> str:
"""Answer a question using the RLM agent with code execution.
The RLM (Recursive Language Model) agent can write and execute Python
code in a sandboxed environment to solve problems that require
computation, aggregation, or complex traversal across documents.
Args:
question: The question to answer.
documents: Optional list of document IDs or titles to pre-load.
filter: SQL WHERE clause to filter documents during searches.
Returns:
The answer as a string.
"""
from haiku.rag.agents.rlm import RLMContext, RLMDeps, create_rlm_agent
context = RLMContext(filter=filter)
if documents:
loaded_docs = []
for doc_ref in documents:
doc = await self.get_document_by_id(doc_ref)
if not doc:
docs = await self.list_documents(filter=f"title = '{doc_ref}'")
if docs and docs[0].id:
doc = await self.get_document_by_id(docs[0].id)
if doc:
loaded_docs.append(doc)
context.documents = loaded_docs if loaded_docs else None
deps = RLMDeps(
client=self,
config=self._config,
context=context,
)
agent = create_rlm_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output.answer
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.

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

@ -94,6 +94,19 @@ class ResearchConfig(BaseModel):
max_concurrency: int = 1
class RLMConfig(BaseModel):
model: ModelConfig = Field(
default_factory=lambda: ModelConfig(
provider="ollama",
name="gpt-oss",
enable_thinking=False,
)
)
code_timeout: float = 60.0
max_output_chars: int = 50_000
max_tool_calls: int = 20
class PictureDescriptionConfig(BaseModel):
"""Configuration for VLM-based picture description."""
@ -194,6 +207,7 @@ class AppConfig(BaseModel):
reranking: RerankingConfig = Field(default_factory=RerankingConfig)
qa: QAConfig = Field(default_factory=QAConfig)
research: ResearchConfig = Field(default_factory=ResearchConfig)
rlm: RLMConfig = Field(default_factory=RLMConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)

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

@ -245,4 +245,31 @@ def create_mcp_server(
except Exception:
return None
@mcp.tool()
async def rlm_question(
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex questions using code execution (RLM agent).
Use this for questions requiring computation, aggregation, or
complex traversal across documents. The agent can write Python
code to search, analyze, and compute answers.
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.
Returns:
The answer as a string.
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
documents = [document] if document else None
return await rag.rlm(question, documents=documents, filter=filter)
except Exception as e:
return f"Error running RLM agent: {e!s}"
return mcp

3
tests/agents/rlm/conftest.py
View file

@ -1,8 +1,9 @@
import pytest
from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
@pytest.fixture

8
tests/agents/rlm/test_agent.py
View file

@ -2,7 +2,7 @@ import pytest
from pydantic_ai import Agent
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMConfig, RLMContext, RLMDeps
from haiku.rag.agents.rlm.dependencies import RLMContext, RLMDeps
from haiku.rag.agents.rlm.models import CodeExecution, RLMResult
from haiku.rag.config import Config
@ -26,12 +26,10 @@ class TestExecuteCodeTool:
"""Test that execute_code tool produces structured CodeExecution output."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
config = RLMConfig()
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
rlm_config=config,
context=context,
)
@ -51,12 +49,10 @@ class TestExecuteCodeTool:
"""Test that code executions are tracked as CodeExecution objects in RLMContext."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
config = RLMConfig()
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
rlm_config=config,
context=context,
)
@ -91,12 +87,10 @@ class TestExecuteCodeTool:
"""Test that failed executions are properly captured."""
from haiku.rag.agents.rlm.agent import _get_or_create_repl
config = RLMConfig()
context = RLMContext()
deps = RLMDeps(
client=empty_client,
config=Config,
rlm_config=config,
context=context,
)

20
tests/agents/rlm/test_models.py
View file

@ -29,7 +29,6 @@ class TestRLMResult:
def test_create_result_with_answer_only(self):
result = RLMResult(answer="The answer is 42")
assert result.answer == "The answer is 42"
assert result.citations == []
assert result.code_executions == []
def test_create_result_with_code_executions(self):
@ -53,22 +52,3 @@ class TestRLMResult:
)
assert len(result.code_executions) == 2
assert result.code_executions[1].stdout == "2\n"
def test_create_result_with_citations(self):
from haiku.rag.agents.research.models import Citation
citations = [
Citation(
document_id="doc1",
chunk_id="chunk1",
document_uri="file://test.pdf",
document_title="Test Doc",
content="Some content",
)
]
result = RLMResult(
answer="Found in Test Doc",
citations=citations,
)
assert len(result.citations) == 1
assert result.citations[0].document_title == "Test Doc"

67
tests/agents/rlm/test_sandbox.py
View file

@ -445,6 +445,73 @@ class TestSandboxExecution:
) # Allow some margin for truncation message
class TestContextFilter:
"""Test that context filter is applied to all searches."""
@pytest.mark.asyncio
async def test_context_filter_applied_to_search(self, temp_db_path):
"""Search applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="uri LIKE '%medical%'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.search = AsyncMock(return_value=[])
await repl.execute_async("search('test query')")
client.search.assert_called_once_with(
"test query", limit=10, filter="uri LIKE '%medical%'"
)
@pytest.mark.asyncio
async def test_context_filter_applied_to_list_documents(self, temp_db_path):
"""list_documents applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="title = 'Report'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.list_documents = AsyncMock(return_value=[])
await repl.execute_async("list_documents()")
client.list_documents.assert_called_once_with(
limit=10, offset=0, filter="title = 'Report'"
)
@pytest.mark.asyncio
async def test_context_filter_applied_to_ask(self, temp_db_path):
"""ask applies context filter automatically."""
from unittest.mock import AsyncMock
from haiku.rag.agents.rlm.dependencies import RLMContext
from haiku.rag.agents.rlm.sandbox import REPLEnvironment
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import RLMConfig
async with HaikuRAG(temp_db_path, create=True) as client:
context = RLMContext(filter="metadata->>'category' = 'finance'")
repl = REPLEnvironment(client=client, config=RLMConfig(), context=context)
client.ask = AsyncMock(return_value=("Answer", []))
await repl.execute_async("ask('What is the revenue?')")
client.ask.assert_called_once_with(
"What is the revenue?", filter="metadata->>'category' = 'finance'"
)
class TestSecurityEscapes:
"""Test that common security escape attempts are blocked."""