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

Return consolidated program from RLM agent instead of execution history

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-06 11:07:46 +01:00
parent e882831afb
commit f8ec511250
No known key found for this signature in database
17 changed files with 4464 additions and 5308 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

200
TOOLS_REFACTORING_PLAN.md
View file

@ -1,200 +0,0 @@
# Tools Extraction Refactoring Plan
## Goal
Extract tools from haiku.rag agents into a reusable `tools/` module, enabling users to create pydantic-ai agents outside haiku.rag and compose toolsets as needed.
## Target API
```python
from pydantic_ai import Agent
from haiku.rag import HaikuRAG
from haiku.rag.tools import ToolContext, create_search_toolset, create_document_toolset
async with HaikuRAG(db_path) as client:
context = ToolContext()
search_tools = create_search_toolset(client, config, context)
doc_tools = create_document_toolset(client, config, context)
agent = Agent(
'anthropic:claude-sonnet',
toolsets=[search_tools, doc_tools]
)
result = await agent.run("Find documents about X")
# Access accumulated state after run
search_state = context.get("haiku.rag.search")
for result in search_state.results:
print(f"{result.document_title}")
```
## Design Principles
1. **ToolContext is a pure generic container** - No special-cased fields. Toolsets register their own Pydantic model state under namespaces.
2. **Shared state via same namespace** - Multiple toolsets can share state (e.g., citations, filters) by registering under the same namespace.
3. **App manages identity** - ToolContext has no session/user identity. The app layer manages `session_id -> ToolContext` mapping.
4. **Toolsets are stateless factories** - `create_*_toolset()` returns a `FunctionToolset`. State lives in the context they're given.
## ToolContext Design
```python
class ToolContext(BaseModel):
"""Generic state container for toolsets.
Toolsets register Pydantic model state under namespaces.
Multiple toolsets can share state via the same namespace.
"""
_namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
def register(self, namespace: str, state: BaseModel) -> None: ...
def get(self, namespace: str) -> BaseModel | None: ...
def get_or_create(self, namespace: str, factory: Callable[[], T]) -> T: ...
def clear_namespace(self, namespace: str) -> None: ...
def clear_all(self) -> None: ...
def dump_namespaces(self) -> dict[str, dict[str, Any]]: ...
def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T: ...
```
## Toolset State Examples
Each toolset defines its own state model:
```python
# Search toolset state
class SearchState(BaseModel):
results: list[SearchResult] = []
filter: str | None = None
SEARCH_NAMESPACE = "haiku.rag.search"
# QA toolset state
class QAState(BaseModel):
history: list[QAResult] = []
QA_NAMESPACE = "haiku.rag.qa"
# Shared citation state (used by multiple toolsets)
class CitationState(BaseModel):
registry: dict[str, int] = {}
def get_or_assign_index(self, chunk_id: str) -> int:
if chunk_id in self.registry:
return self.registry[chunk_id]
new_index = len(self.registry) + 1
self.registry[chunk_id] = new_index
return new_index
CITATION_NAMESPACE = "haiku.rag.citations"
```
## Multi-User/Session Management
App layer manages context routing:
```python
# App maintains context per session
contexts: dict[str, ToolContext] = {}
def get_context(session_id: str) -> ToolContext:
if session_id not in contexts:
contexts[session_id] = ToolContext()
return contexts[session_id]
# When running agent
context = get_context(user_session_id)
toolsets = [create_search_toolset(client, config, context)]
await agent.run(prompt, toolsets=toolsets)
```
## New Module Structure
```
haiku_rag_slim/haiku/rag/
├── tools/ # NEW
│ ├── __init__.py # Public exports
│ ├── context.py # ToolContext (generic state container)
│ ├── models.py # QAResult, AnalysisResult
│ ├── filters.py # build_document_filter, combine_filters
│ ├── search.py # create_search_toolset()
│ ├── document.py # create_document_toolset()
│ ├── qa.py # create_qa_toolset()
│ └── analysis.py # create_analysis_toolset()
├── agents/ # REFACTORED to use tools/
```
## Implementation Chunks
### Chunk 1: Create tools module foundation ✅ DONE
- Created `tools/__init__.py`, `tools/context.py`, `tools/models.py`, `tools/filters.py`
- Created `ToolContext` as generic namespace-based Pydantic model
- Moved filter utilities from `agents/chat/state.py` to `tools/filters.py`
- Created result models (`QAResult`, `AnalysisResult`)
- Added tests for ToolContext and filters
### Chunk 2: Create SearchToolset ✅ DONE
- Created `tools/search.py` with `create_search_toolset()`
- Defined `SearchState` model for accumulating search results
- Core search logic: `client.search()` → `client.expand_context()` → `format_for_agent()`
- Results accumulated in `SearchState` under `SEARCH_NAMESPACE`
- Added 13 tests for SearchToolset
### Chunk 3: Refactor QA Agent to use SearchToolset ✅ DONE
- Updated `agents/qa/agent.py` to use `create_search_toolset()`
- Added `base_filter` and `tool_name` parameters to `create_search_toolset()`
- QA agent now uses ToolContext + SearchState for result accumulation
- Public interface (`answer(question, filter)`) unchanged
- All 5 QA tests pass
### Chunk 4: Create DocumentToolset ✅ DONE
- Created `tools/document.py` with `create_document_toolset()`
- Defined `DocumentState`, `DocumentInfo`, `DocumentListResponse` models
- Extracted `list_documents`, `get_document`, `summarize_document` tools
- Moved `find_document` helper (now public)
- Added 13 tests
### Chunk 5: Create QAToolset ✅ DONE
- Created `tools/qa.py` with `create_qa_toolset()`
- Defined `QAState` model (tracks QA history)
- Runs research graph, returns structured `QAResult`
- Supports `base_filter`, `tool_name`, `session_context`, `prior_answers` params
- Added 7 tests
### Chunk 6: Create AnalysisToolset ✅ DONE
- Created `tools/analysis.py` with `create_analysis_toolset()`
- Defined `AnalysisState` model (tracks CodeExecution history)
- Extracted `analyze` tool (RLM delegation with filter support)
- Fixed circular import by using direct submodule imports
- Added 6 tests
### Chunk 7: Refactor Chat Agent ✅ DONE
- Removed `analyze` tool from chat agent (kept hardcoded, not composing toolsets)
- Reverted system prompt to pre-analyze version
- Removed `test_analyze_tool` test and cassette file
- All 47 chat agent tests pass
### Chunk 8: Refactor Research Graph
- Update `_search_one_step_logic` to use search toolset
- Verify research tests pass
### Chunk 9: Public API and Documentation
- Export from `haiku.rag.tools` and `haiku.rag`
- Update CLAUDE.md
- Add usage examples
## Verification
- Run `pytest` after each chunk
- Run `ty check` and `ruff check`
- Test with existing agents (QA, Chat, Research)
- Test with external agent using new toolsets
## Critical Files
- `haiku_rag_slim/haiku/rag/agents/chat/agent.py` - largest tool collection
- `haiku_rag_slim/haiku/rag/agents/qa/agent.py` - simplest, good starting point
- `haiku_rag_slim/haiku/rag/agents/chat/state.py` - filter utilities (now moved)
- `haiku_rag_slim/haiku/rag/agents/research/graph.py` - search tool inside step
- `haiku_rag_slim/haiku/rag/store/models/chunk.py` - SearchResult.format_for_agent()

8
docs/python.md
View file

@ -403,16 +403,18 @@ Answer complex analytical questions via code execution:
```python
# Aggregation across documents
answer = await client.rlm("Which quarter had the highest revenue?")
result = await client.rlm("Which quarter had the highest revenue?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
# Computation within a document set
answer = await client.rlm(
result = await client.rlm(
"What is the average deal size mentioned in these contracts?",
filter="uri LIKE '%contracts%'"
)
# Multi-document comparison
answer = await client.rlm(
result = await client.rlm(
"What changed between these two versions of the policy?",
documents=["Policy v1.0", "Policy v2.0"]
)

11
docs/rlm.md
View file

@ -35,17 +35,18 @@ from haiku.rag.client import HaikuRAG
async with HaikuRAG(path_to_db) as client:
# Basic question
answer = await client.rlm("How many documents mention 'security'?")
print(answer)
result = await client.rlm("How many documents mention 'security'?")
print(result.answer) # The answer
print(result.program) # The final consolidated program
# With filter (agent can only see filtered documents)
answer = await client.rlm(
result = await client.rlm(
"What is the total revenue?",
filter="title LIKE '%Financial%'"
)
# Pre-load specific documents
answer = await client.rlm(
result = await client.rlm(
"Compare the conclusions",
documents=["Report A", "Report B"]
)
@ -169,7 +170,7 @@ The `filter` parameter restricts what documents the agent can access. Unlike too
```python
# Agent can only see documents with "confidential" in the URI
answer = await client.rlm(
result = await client.rlm(
"Summarize all findings",
filter="uri LIKE '%confidential%'"
)

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

@ -14,7 +14,4 @@ class RLMResult(BaseModel):
"""Result from RLM agent execution."""
answer: str = Field(description="The answer to the user's question")
code_executions: list[CodeExecution] = Field(
default_factory=list,
description="History of code executions during the RLM session",
)
program: str = Field(description="The final consolidated program")

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

@ -43,7 +43,7 @@ for doc in documents:
Check if it exists with: `if 'documents' in dir(): ...`
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
You can import any Python standard library module.
## Strategy Guide
@ -72,7 +72,7 @@ When you call `get_docling_document(id_or_title)`, you get a DoclingDocument obj
### Text Item Properties
- `item.text` - The text content
- `item.label` - Type: TITLE, PARAGRAPH, SECTION_HEADER, LIST_ITEM, etc.
- `item.label` - Type: title, paragraph, section_header, list_item, etc. (lowercase enum values)
- `item.prov` - Provenance (page numbers, bounding boxes)
### Table Access
@ -85,7 +85,7 @@ When you call `get_docling_document(id_or_title)`, you get a DoclingDocument obj
doc = get_docling_document("My Document")
# Get all headings
headings = [t.text for t in doc.texts if "HEADER" in str(t.label)]
headings = [t.text for t in doc.texts if "header" in str(t.label)]
# Iterate with structure
for item, level in doc.iterate_items():
@ -142,14 +142,12 @@ print(sentiment)
CRITICAL: Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string"}
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
The `answer` field should contain:
1. A clear answer to the user's question
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

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

@ -16,6 +16,7 @@ from rich.progress import (
TextColumn,
TransferSpeedColumn,
)
from rich.syntax import Syntax
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
@ -458,10 +459,13 @@ class HaikuRAGApp:
self.console.print("[dim]Running RLM agent with code execution...[/dim]")
self.console.print()
answer = await self.client.rlm(question, documents=documents, filter=filter)
result = await self.client.rlm(question, documents=documents, filter=filter)
self.console.print("[bold yellow]Program:[/bold yellow]")
self.console.print(Syntax(result.program, "python"))
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
self.console.print(Markdown(result.answer))
async def research(
self,

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

@ -32,6 +32,7 @@ if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.research.models import Citation
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__)
@ -1325,7 +1326,7 @@ class HaikuRAG:
question: str,
documents: list[str] | None = None,
filter: str | None = None,
) -> str:
) -> "RLMResult":
"""Answer a question using the RLM agent with code execution.
The RLM (Recursive Language Model) agent can write and execute Python
@ -1338,7 +1339,7 @@ class HaikuRAG:
filter: SQL WHERE clause to filter documents during searches.
Returns:
The answer as a string.
RLMResult with the answer and the final consolidated program.
"""
from haiku.rag.agents.rlm import (
DockerSandbox,
@ -1371,7 +1372,7 @@ class HaikuRAG:
agent = create_rlm_agent(self._config)
result = await agent.run(question, deps=deps)
return result.output.answer
return result.output
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.

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

@ -268,7 +268,8 @@ def create_mcp_server(
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)
result = await rag.rlm(question, documents=documents, filter=filter)
return result.answer
except Exception as e:
return f"Error running RLM agent: {e!s}"

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

@ -65,9 +65,9 @@ class TestClientRLMIntegration:
await client.create_document("Second document about dogs.", title="Doc 2")
await client.create_document("Third document about birds.", title="Doc 3")
answer = await client.rlm("How many documents are in the database?")
result = await client.rlm("How many documents are in the database?")
assert "3" in answer
assert "3" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -107,11 +107,11 @@ class TestClientRLMIntegration:
"Sales report Q3: Revenue was $200,000.", title="Q3 Report"
)
answer = await client.rlm(
result = await client.rlm(
"What is the total revenue across all quarterly reports?"
)
assert "450" in answer or "450,000" in answer
assert "450" in result.answer or "450,000" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -136,12 +136,12 @@ class TestClientRLMIntegration:
await client.create_document("Dog document.", title="Dogs")
await client.create_document("Bird document.", title="Birds")
answer = await client.rlm(
result = await client.rlm(
"How many documents are available?",
filter="title = 'Cats'",
)
assert "1" in answer
assert "1" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -169,13 +169,13 @@ class TestClientRLMIntegration:
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path)
answer = await client.rlm(
result = await client.rlm(
"How many tables are in the document? "
"Also tell me how many pictures/figures it contains."
)
# The doclaynet.pdf has 1 table and 1 picture
assert "1" in answer
assert "1" in result.answer
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -221,14 +221,14 @@ class TestClientRLMIntegration:
title="Q3 Update",
)
answer = await client.rlm(
result = await client.rlm(
"Analyze the sentiment of each quarterly update. "
"How many quarters were positive, negative, and mixed?"
)
# Should identify: Q1=positive, Q2=negative, Q3=mixed
assert "positive" in answer.lower()
assert "negative" in answer.lower()
assert "positive" in result.answer.lower()
assert "negative" in result.answer.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
@ -257,7 +257,7 @@ class TestClientRLMIntegration:
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
await client.create_document_from_source(pdf_path)
answer = await client.rlm(
result = await client.rlm(
"Search for content about document element types or labels. "
"What are all the different document element types mentioned? "
"List them all."
@ -265,7 +265,7 @@ class TestClientRLMIntegration:
# The doclaynet.pdf defines exactly 11 class labels for document elements
# Normalize Unicode hyphens (U+2011 non-breaking hyphen) to regular hyphens
answer_lower = answer.lower().replace("\u2011", "-")
answer_lower = result.answer.lower().replace("\u2011", "-")
expected_labels = [
"caption",
"footnote",
@ -318,11 +318,14 @@ class TestClientRLMIntegration:
title="Mission Statement",
)
answer = await client.rlm(
result = await client.rlm(
"Using the pre-loaded documents variable, "
"tell me when was the company founded and what is their mission?",
documents=["Company History", "Mission Statement"],
)
assert "1985" in answer
assert "accessible" in answer.lower() or "technology" in answer.lower()
assert "1985" in result.answer
assert (
"accessible" in result.answer.lower()
or "technology" in result.answer.lower()
)

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

@ -26,29 +26,7 @@ class TestCodeExecution:
class TestRLMResult:
def test_create_result_with_answer_only(self):
result = RLMResult(answer="The answer is 42")
def test_create_result(self):
result = RLMResult(answer="The answer is 42", program="print(42)")
assert result.answer == "The answer is 42"
assert result.code_executions == []
def test_create_result_with_code_executions(self):
executions = [
CodeExecution(
code="x = 1 + 1",
stdout="",
stderr="",
success=True,
),
CodeExecution(
code="print(x)",
stdout="2\n",
stderr="",
success=True,
),
]
result = RLMResult(
answer="x equals 2",
code_executions=executions,
)
assert len(result.code_executions) == 2
assert result.code_executions[1].stdout == "2\n"
assert result.program == "print(42)"

2075
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_aggregation.yaml
View file

File diff suppressed because it is too large Load diff

163
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_count_documents.yaml
View file

@ -128,7 +128,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8280'
- '7774'
content-type:
- application/json
host:
@ -182,7 +182,7 @@ interactions:
Check if it exists with: `if 'documents' in dir(): ...`
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
You can import any Python standard library module.
## Strategy Guide
@ -211,7 +211,7 @@ interactions:
### Text Item Properties
- `item.text` - The text content
- `item.label` - Type: TITLE, PARAGRAPH, SECTION_HEADER, LIST_ITEM, etc.
- `item.label` - Type: title, paragraph, section_header, list_item, etc. (lowercase enum values)
- `item.prov` - Provenance (page numbers, bounding boxes)
### Table Access
@ -224,7 +224,7 @@ interactions:
doc = get_docling_document("My Document")
# Get all headings
headings = [t.text for t in doc.texts if "HEADER" in str(t.label)]
headings = [t.text for t in doc.texts if "header" in str(t.label)]
# Iterate with structure
for item, level in doc.iterate_items():
@ -281,15 +281,13 @@ interactions:
CRITICAL: Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string"}
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
The `answer` field should contain:
1. A clear answer to the user's question
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
@ -302,12 +300,11 @@ interactions:
tools:
- function:
description: |-
<summary>Execute Python code in the sandboxed environment.
<summary>Execute Python code in a Docker-sandboxed environment.
The code has access to haiku.rag functions (search, list_documents,
get_document, get_docling_document, ask) and safe standard library
modules (json, re, collections, math, statistics, itertools,
functools, datetime, typing).
get_document, get_docling_document, llm) and any Python standard
library module.
Use print() to output results.</summary>
<returns>
@ -329,48 +326,25 @@ interactions:
description: Result from RLM agent execution.
name: final_result
parameters:
$defs:
CodeExecution:
additionalProperties: false
description: Result of executing a code block in the RLM sandbox.
properties:
code:
description: The Python code that was executed
type: string
stderr:
description: Standard error captured during execution
type: string
stdout:
description: Standard output captured during execution
type: string
success:
description: Whether execution completed without error
type: boolean
required:
- code
- stdout
- stderr
- success
type: object
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
code_executions:
description: History of code executions during the RLM session
items:
$ref: '#/$defs/CodeExecution'
type: array
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '513'
- '523'
content-type:
- application/json
parsed_body:
@ -379,24 +353,24 @@ interactions:
index: 0
message:
content: ''
reasoning: Need to list_documents.
reasoning: We need to list documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
arguments: '{"code":"# list documents\nimport json\nprint(list_documents())\n"}'
name: execute_code
id: call_3ltyid1i
id: call_d8xhmimu
index: 0
type: function
created: 1770322488
id: chatcmpl-467
created: 1770373335
id: chatcmpl-184
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 43
prompt_tokens: 1752
total_tokens: 1795
completion_tokens: 42
prompt_tokens: 1747
total_tokens: 1789
status:
code: 200
message: OK
@ -409,7 +383,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8709'
- '8588'
content-type:
- application/json
host:
@ -463,7 +437,7 @@ interactions:
Check if it exists with: `if 'documents' in dir(): ...`
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
You can import any Python standard library module.
## Strategy Guide
@ -492,7 +466,7 @@ interactions:
### Text Item Properties
- `item.text` - The text content
- `item.label` - Type: TITLE, PARAGRAPH, SECTION_HEADER, LIST_ITEM, etc.
- `item.label` - Type: title, paragraph, section_header, list_item, etc. (lowercase enum values)
- `item.prov` - Provenance (page numbers, bounding boxes)
### Table Access
@ -505,7 +479,7 @@ interactions:
doc = get_docling_document("My Document")
# Get all headings
headings = [t.text for t in doc.texts if "HEADER" in str(t.label)]
headings = [t.text for t in doc.texts if "header" in str(t.label)]
# Iterate with structure
for item, level in doc.iterate_items():
@ -562,15 +536,13 @@ interactions:
CRITICAL: Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string"}
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
The `answer` field should contain:
1. A clear answer to the user's question
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
@ -578,18 +550,21 @@ interactions:
role: user
- content: |-
<think>
Need to list_documents.
We need to list documents.
</think>
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
arguments: '{"code":"# list documents\nimport json\nprint(list_documents())\n"}'
name: execute_code
id: call_3ltyid1i
id: call_d8xhmimu
type: function
- content: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
- content: '{"code":"# list documents\nimport json\nprint(list_documents())\n","stdout":"[{''id'': ''b73f8a17-4328-475c-84db-3d81ce52adce'',
''title'': ''Doc 1'', ''uri'': None, ''created_at'': ''2026-02-06 11:21:55.843558''}, {''id'': ''accb877b-f04e-4bf2-ba4c-2d90339fa875'',
''title'': ''Doc 2'', ''uri'': None, ''created_at'': ''2026-02-06 11:21:57.397026''}, {''id'': ''afdb966f-5e9d-4759-a08f-28eb5108c80f'',
''title'': ''Doc 3'', ''uri'': None, ''created_at'': ''2026-02-06 11:21:58.988378''}]\n","stderr":"","success":true}'
role: tool
tool_call_id: call_3ltyid1i
tool_call_id: call_d8xhmimu
model: gpt-oss
reasoning_effort: low
stream: false
@ -597,12 +572,11 @@ interactions:
tools:
- function:
description: |-
<summary>Execute Python code in the sandboxed environment.
<summary>Execute Python code in a Docker-sandboxed environment.
The code has access to haiku.rag functions (search, list_documents,
get_document, get_docling_document, ask) and safe standard library
modules (json, re, collections, math, statistics, itertools,
functools, datetime, typing).
get_document, get_docling_document, llm) and any Python standard
library module.
Use print() to output results.</summary>
<returns>
@ -624,48 +598,25 @@ interactions:
description: Result from RLM agent execution.
name: final_result
parameters:
$defs:
CodeExecution:
additionalProperties: false
description: Result of executing a code block in the RLM sandbox.
properties:
code:
description: The Python code that was executed
type: string
stderr:
description: Standard error captured during execution
type: string
stdout:
description: Standard output captured during execution
type: string
success:
description: Whether execution completed without error
type: boolean
required:
- code
- stdout
- stderr
- success
type: object
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
code_executions:
description: History of code executions during the RLM session
items:
$ref: '#/$defs/CodeExecution'
type: array
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '341'
- '523'
content-type:
- application/json
parsed_body:
@ -673,17 +624,19 @@ interactions:
- finish_reason: stop
index: 0
message:
content: '{"answer":"There are 3 documents in the database."}'
content: '{"answer":"There are 3 documents in the database.","program":"# List and count documents\nimport json\n\ndocs
= list_documents()\nprint(f\"Number of documents: {len(docs)}\")\n"}'
reasoning: Count is 3. Provide answer.
role: assistant
created: 1770322490
id: chatcmpl-796
created: 1770373336
id: chatcmpl-441
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 20
prompt_tokens: 1842
total_tokens: 1862
completion_tokens: 68
prompt_tokens: 2019
total_tokens: 2087
status:
code: 200
message: OK

1801
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_docling_document_structure.yaml
View file

File diff suppressed because it is too large Load diff

3103
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_search_and_extract.yaml
View file

File diff suppressed because it is too large Load diff

1153
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_semantic_analysis_with_llm.yaml
View file

File diff suppressed because it is too large Load diff

158
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_with_filter.yaml
View file

@ -128,7 +128,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8274'
- '7768'
content-type:
- application/json
host:
@ -182,7 +182,7 @@ interactions:
Check if it exists with: `if 'documents' in dir(): ...`
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
You can import any Python standard library module.
## Strategy Guide
@ -211,7 +211,7 @@ interactions:
### Text Item Properties
- `item.text` - The text content
- `item.label` - Type: TITLE, PARAGRAPH, SECTION_HEADER, LIST_ITEM, etc.
- `item.label` - Type: title, paragraph, section_header, list_item, etc. (lowercase enum values)
- `item.prov` - Provenance (page numbers, bounding boxes)
### Table Access
@ -224,7 +224,7 @@ interactions:
doc = get_docling_document("My Document")
# Get all headings
headings = [t.text for t in doc.texts if "HEADER" in str(t.label)]
headings = [t.text for t in doc.texts if "header" in str(t.label)]
# Iterate with structure
for item, level in doc.iterate_items():
@ -281,15 +281,13 @@ interactions:
CRITICAL: Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string"}
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
The `answer` field should contain:
1. A clear answer to the user's question
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
@ -302,12 +300,11 @@ interactions:
tools:
- function:
description: |-
<summary>Execute Python code in the sandboxed environment.
<summary>Execute Python code in a Docker-sandboxed environment.
The code has access to haiku.rag functions (search, list_documents,
get_document, get_docling_document, ask) and safe standard library
modules (json, re, collections, math, statistics, itertools,
functools, datetime, typing).
get_document, get_docling_document, llm) and any Python standard
library module.
Use print() to output results.</summary>
<returns>
@ -329,48 +326,25 @@ interactions:
description: Result from RLM agent execution.
name: final_result
parameters:
$defs:
CodeExecution:
additionalProperties: false
description: Result of executing a code block in the RLM sandbox.
properties:
code:
description: The Python code that was executed
type: string
stderr:
description: Standard error captured during execution
type: string
stdout:
description: Standard output captured during execution
type: string
success:
description: Whether execution completed without error
type: boolean
required:
- code
- stdout
- stderr
- success
type: object
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
code_executions:
description: History of code executions during the RLM session
items:
$ref: '#/$defs/CodeExecution'
type: array
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '490'
- '517'
content-type:
- application/json
parsed_body:
@ -379,24 +353,24 @@ interactions:
index: 0
message:
content: ''
reasoning: We need to list documents.
reasoning: Need to get list_documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"print(len(list_documents()))"}'
arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_wvyfhnmo
id: call_ly3bn3y0
index: 0
type: function
created: 1770322511
id: chatcmpl-2
created: 1770373376
id: chatcmpl-196
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 34
prompt_tokens: 1750
total_tokens: 1784
completion_tokens: 44
prompt_tokens: 1745
total_tokens: 1789
status:
code: 200
message: OK
@ -409,7 +383,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8658'
- '8201'
content-type:
- application/json
host:
@ -463,7 +437,7 @@ interactions:
Check if it exists with: `if 'documents' in dir(): ...`
## Standard Library Modules
You can import: json, re, collections, math, statistics, itertools, functools, datetime, typing
You can import any Python standard library module.
## Strategy Guide
@ -492,7 +466,7 @@ interactions:
### Text Item Properties
- `item.text` - The text content
- `item.label` - Type: TITLE, PARAGRAPH, SECTION_HEADER, LIST_ITEM, etc.
- `item.label` - Type: title, paragraph, section_header, list_item, etc. (lowercase enum values)
- `item.prov` - Provenance (page numbers, bounding boxes)
### Table Access
@ -505,7 +479,7 @@ interactions:
doc = get_docling_document("My Document")
# Get all headings
headings = [t.text for t in doc.texts if "HEADER" in str(t.label)]
headings = [t.text for t in doc.texts if "header" in str(t.label)]
# Iterate with structure
for item, level in doc.iterate_items():
@ -562,15 +536,13 @@ interactions:
CRITICAL: Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string"}
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
The `answer` field should contain:
1. A clear answer to the user's question
2. Key findings from your analysis
3. References to specific documents/chunks that informed your answer
- `answer`: A clear answer to the user's question with key findings and references to specific documents/chunks.
- `program`: A single, self-contained Python program that produces the answer. Consolidate your exploratory code executions into one clean script.
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "...", "program": "..."}
CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
@ -578,18 +550,18 @@ interactions:
role: user
- content: |-
<think>
We need to list documents.
Need to get list_documents.
</think>
role: assistant
tool_calls:
- function:
arguments: '{"code":"print(len(list_documents()))"}'
arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_wvyfhnmo
id: call_ly3bn3y0
type: function
- content: '{"code":"print(len(list_documents()))","stdout":"1\n","stderr":"","success":true}'
- content: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))","stdout":"1\n","stderr":"","success":true}'
role: tool
tool_call_id: call_wvyfhnmo
tool_call_id: call_ly3bn3y0
model: gpt-oss
reasoning_effort: low
stream: false
@ -597,12 +569,11 @@ interactions:
tools:
- function:
description: |-
<summary>Execute Python code in the sandboxed environment.
<summary>Execute Python code in a Docker-sandboxed environment.
The code has access to haiku.rag functions (search, list_documents,
get_document, get_docling_document, ask) and safe standard library
modules (json, re, collections, math, statistics, itertools,
functools, datetime, typing).
get_document, get_docling_document, llm) and any Python standard
library module.
Use print() to output results.</summary>
<returns>
@ -624,48 +595,25 @@ interactions:
description: Result from RLM agent execution.
name: final_result
parameters:
$defs:
CodeExecution:
additionalProperties: false
description: Result of executing a code block in the RLM sandbox.
properties:
code:
description: The Python code that was executed
type: string
stderr:
description: Standard error captured during execution
type: string
stdout:
description: Standard output captured during execution
type: string
success:
description: Whether execution completed without error
type: boolean
required:
- code
- stdout
- stderr
- success
type: object
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
code_executions:
description: History of code executions during the RLM session
items:
$ref: '#/$defs/CodeExecution'
type: array
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '333'
- '424'
content-type:
- application/json
parsed_body:
@ -673,17 +621,17 @@ interactions:
- finish_reason: stop
index: 0
message:
content: '{"answer":"There is 1 document available."}'
content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
role: assistant
created: 1770322513
id: chatcmpl-541
created: 1770373377
id: chatcmpl-195
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 18
prompt_tokens: 1821
total_tokens: 1839
completion_tokens: 39
prompt_tokens: 1836
total_tokens: 1875
status:
code: 200
message: OK

998
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_with_preloaded_documents.yaml
View file

File diff suppressed because it is too large Load diff