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

Merge pull request #300 from ggozad/feat/support-native-structured-output

Browse source 
Automatically set structured output mode (tool vs native)
This commit is contained in:
Yiorgis Gozadinos 2026-03-04 15:01:47 +02:00 committed by GitHub
commit d08d32dd5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 3136 additions and 7555 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

@ -4,6 +4,7 @@
### Added
- **Module-level skill introspection API**: `STATE_TYPE`, `STATE_NAMESPACE`, `skill_metadata()`, `instructions()`, and `state_metadata()` on `haiku.rag.skills.rag` and `haiku.rag.skills.rlm` — allows introspecting skill configuration without calling `create_skill()`
- **Automatic structured output detection**: Native JSON schema output is used automatically when the model supports it, with tool-call fallback otherwise. No configuration needed.
### Changed

8
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -1,7 +1,6 @@
from dataclasses import dataclass
from pydantic_ai import Agent
from pydantic_ai.output import ToolOutput
from haiku.rag.agents.qa.prompts import QA_SYSTEM_PROMPT
from haiku.rag.agents.research.models import (
@ -14,7 +13,7 @@ from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.utils import get_model
from haiku.rag.utils import get_model, structured_output_type
@dataclass
@ -57,10 +56,11 @@ class QuestionAnswerAgent:
# Agent created per-call: toolset varies with filter, and Agent
# construction is pure Python (no IO).
model = get_model(self._model_config, self._config)
agent: Agent[_QARunDeps, RawSearchAnswer] = Agent( # ty: ignore[invalid-assignment]
model=get_model(self._model_config, self._config),
model=model,
deps_type=_QARunDeps,
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
output_type=structured_output_type(RawSearchAnswer, model),
instructions=self._system_prompt,
toolsets=[search_toolset],
retries=3,

18
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -1,7 +1,6 @@
import asyncio
from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
@ -20,7 +19,7 @@ from haiku.rag.agents.research.prompts import (
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.utils import build_prompt, get_model
from haiku.rag.utils import build_prompt, get_model, structured_output_type
def format_context_for_prompt(context: ResearchContext) -> str:
@ -66,9 +65,10 @@ async def _iterative_plan_logic(
else:
effective_prompt = build_prompt(ITERATIVE_PLAN_PROMPT, config)
model = get_model(model_config, config)
plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ToolOutput(IterativePlanResult, max_retries=3),
model=model,
output_type=structured_output_type(IterativePlanResult, model),
instructions=effective_prompt,
retries=3,
deps_type=ResearchDependencies,
@ -115,9 +115,10 @@ async def _search_one_step_logic(
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
model = get_model(model_config, config)
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
model=model,
output_type=structured_output_type(RawSearchAnswer, model),
instructions=search_prompt,
retries=3,
deps_type=ResearchDependencies,
@ -216,9 +217,10 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
model = get_model(model_config, config)
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ToolOutput(ResearchReport, max_retries=3),
model=model,
output_type=structured_output_type(ResearchReport, model),
instructions=synthesis_prompt,
retries=3,
deps_type=ResearchDependencies,

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

@ -1,11 +1,10 @@
from pydantic_ai import Agent, RunContext
from pydantic_ai.output import ToolOutput
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
from haiku.rag.utils import get_model, structured_output_type
def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
@ -26,7 +25,7 @@ def create_rlm_agent(config: AppConfig) -> Agent[RLMDeps, RLMResult]:
agent: Agent[RLMDeps, RLMResult] = Agent( # type: ignore[invalid-assignment]
model,
deps_type=RLMDeps,
output_type=ToolOutput(RLMResult, max_retries=3),
output_type=structured_output_type(RLMResult, model),
instructions=RLM_SYSTEM_PROMPT,
retries=3,
)

14
haiku_rag_slim/haiku/rag/utils.py
View file

@ -306,6 +306,20 @@ def get_model(
return f"{provider}:{model}"
def structured_output_type(
result_type: type,
model: Any,
max_retries: int = 3,
) -> Any:
"""Return a NativeOutput or ToolOutput wrapper based on model capability."""
from pydantic_ai.models import Model
from pydantic_ai.output import NativeOutput, ToolOutput
if isinstance(model, Model) and model.profile.supports_json_schema_output:
return NativeOutput(result_type)
return ToolOutput(result_type, max_retries=max_retries)
def format_bytes(num_bytes: int) -> str:
"""Format bytes as human-readable string."""
size = float(num_bytes)

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

@ -2,7 +2,6 @@ from pathlib import Path
import pytest
from pydantic_ai import Agent
from pydantic_ai.output import ToolOutput
from haiku.rag.agents.rlm.agent import create_rlm_agent
from haiku.rag.agents.rlm.dependencies import RLMDeps
@ -16,10 +15,22 @@ def vcr_cassette_dir():
class TestCreateRLMAgent:
def test_creates_agent_with_correct_types(self):
def test_creates_agent_native_output_when_supported(self):
from pydantic_ai.output import NativeOutput
agent = create_rlm_agent(Config)
assert isinstance(agent, Agent)
assert agent.deps_type is RLMDeps
assert isinstance(agent.output_type, NativeOutput)
assert agent.output_type.outputs is RLMResult
def test_creates_agent_tool_output_when_not_supported(self):
from pydantic_ai.output import ToolOutput
config = AppConfig()
config.rlm.model.name = "qwen3"
agent = create_rlm_agent(config)
assert isinstance(agent, Agent)
assert isinstance(agent.output_type, ToolOutput)
assert agent.output_type.output is RLMResult

232
tests/cassettes/test_qa/test_qa_ollama.yaml
View file

File diff suppressed because one or more lines are too long

542
tests/cassettes/test_research_graph/test_graph_end_to_end.yaml
View file

File diff suppressed because one or more lines are too long

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

File diff suppressed because it is too large Load diff

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

@ -128,7 +128,7 @@ interactions:
connection:
- keep-alive
content-length:
- '7704'
- '7721'
content-type:
- application/json
host:
@ -276,6 +276,25 @@ interactions:
role: user
model: gpt-oss
reasoning_effort: low
response_format:
json_schema:
description: Result from RLM agent execution.
name: RLMResult
schema:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: json_schema
stream: false
tool_choice: auto
tools:
@ -302,29 +321,11 @@ interactions:
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
- '516'
- '522'
content-type:
- application/json
parsed_body:
@ -333,24 +334,24 @@ interactions:
index: 0
message:
content: ''
reasoning: Need list_documents.
reasoning: We need to list documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_oyaoz18v
id: call_cthflnpr
index: 0
type: function
created: 1772549310
id: chatcmpl-325
created: 1772626945
id: chatcmpl-979
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 43
prompt_tokens: 1730
total_tokens: 1773
completion_tokens: 45
prompt_tokens: 1686
total_tokens: 1731
status:
code: 200
message: OK
@ -363,7 +364,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8140'
- '8163'
content-type:
- application/json
host:
@ -510,19 +511,38 @@ interactions:
- content: How many documents are in the database?
role: user
- content: null
reasoning: Need list_documents.
reasoning: We need to list documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_oyaoz18v
id: call_cthflnpr
type: function
- content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
role: tool
tool_call_id: call_oyaoz18v
tool_call_id: call_cthflnpr
model: gpt-oss
reasoning_effort: low
response_format:
json_schema:
description: Result from RLM agent execution.
name: RLMResult
schema:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: json_schema
stream: false
tool_choice: auto
tools:
@ -549,29 +569,11 @@ interactions:
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
- '416'
- '459'
content-type:
- application/json
parsed_body:
@ -579,272 +581,18 @@ interactions:
- finish_reason: stop
index: 0
message:
content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(f''There
are {len(docs)} documents in the database.'')" }'
role: assistant
created: 1772549311
id: chatcmpl-670
created: 1772626947
id: chatcmpl-490
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 38
prompt_tokens: 1815
total_tokens: 1853
status:
code: 200
message: OK
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '8432'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a Recursive Language Model (RLM) 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.
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
- results = await search("query") ✓ CORRECT
- import search ✗ WRONG - will fail
- results = search("query") ✗ WRONG - must use await
## Available Functions
### await 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
### await 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
### await get_document(id_or_title) -> str | None
Get the full text content of a document by ID, title, or URI.
Returns the document content as a string, or None if not found.
### await get_chunk(chunk_id) -> dict | None
Get a specific chunk by its ID (from search results).
Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
Use this to retrieve full chunk details and metadata for citation.
### await get_docling_document(document_id) -> dict | None
Get the full document structure as a dict (DoclingDocument format).
Use `list_documents()` or search results to get document IDs first.
- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
- `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
- `pictures`: list of figures/images with metadata
- `pages`: page dimensions and metadata
### await regex_findall(pattern, text) -> list[str]
Find all non-overlapping matches of a regular expression pattern in text.
### await regex_sub(pattern, repl, text) -> str
Replace all occurrences of a regular expression pattern with a replacement string.
### await regex_search(pattern, text) -> dict | None
Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
### await regex_split(pattern, text) -> list[str]
Split text by a regular expression pattern.
### await llm(prompt) -> str
Call an LLM directly with the given prompt. Returns the response as a string.
Use this for classification, summarization, extraction, or any task where you
already have the content and just need LLM reasoning.
## Pre-loaded Documents Variable
If documents were pre-loaded for this session, a `documents` variable is available:
```python
# documents is a list of dicts with keys: id, title, uri, content
for doc in documents:
print(doc['title'], len(doc['content']))
```
Check if it exists with: `try: documents ... except NameError: ...`
## Available Python Features
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
## Strategy Guide
1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
## Example Patterns
### Counting documents matching a condition
```python
docs = await list_documents(limit=100)
count = 0
for doc in docs:
content = await get_document(doc['id'])
if content and 'keyword' in content.lower():
count += 1
print(f"Found in: {doc['title']}")
print(f"Total: {count}")
```
### Extracting data with regex
```python
numbers = []
results = await search("financial data", limit=20)
for r in results:
amounts = await regex_findall(r'\$([\d,]+)', r['content'])
for a in amounts:
numbers.append(int(a.replace(',', '')))
if numbers:
print(f"Average: {sum(numbers) / len(numbers)}")
```
### Extracting tables from a document
```python
docs = await list_documents(limit=10)
for d in docs:
doc = await get_docling_document(d['id'])
if doc:
tables = doc.get('tables', [])
if tables:
print(f"{d['title']}: {len(tables)} table(s)")
for i, table in enumerate(tables):
grid = table.get('data', {}).get('grid', [])
for row in grid:
cells = [cell.get('text', '') for cell in row]
print(f" Table {i}: {cells}")
```
## Output Format
Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
- `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": "...", "program": "..."}
You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
- content: How many documents are in the database?
role: user
- content: null
reasoning: Need list_documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_oyaoz18v
type: function
- content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}'
role: tool
tool_call_id: call_oyaoz18v
- content: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
role: assistant
- content: |-
Validation feedback:
Please include your response in a tool call.
Fix the errors and try again.
role: user
model: gpt-oss
reasoning_effort: low
stream: false
tool_choice: auto
tools:
- function:
description: |-
<summary>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.</summary>
<returns>
<description>Structured result with success status, stdout, and stderr.</description>
</returns>
name: execute_code
parameters:
additionalProperties: false
properties:
code:
description: Python code to execute.
type: string
required:
- code
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
- '610'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: tool_calls
index: 0
message:
content: ''
reasoning: Need to return via tool call? We should use final_result.
role: assistant
tool_calls:
- function:
arguments: '{"answer":"There are 3 documents in the database.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: final_result
id: call_46d25765
index: 0
type: function
created: 1772549312
id: chatcmpl-412
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 63
prompt_tokens: 1865
total_tokens: 1928
completion_tokens: 49
prompt_tokens: 1773
total_tokens: 1822
status:
code: 200
message: OK

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

File diff suppressed because one or more lines are too long

300
tests/cassettes/test_rlm/TestClientRLMIntegration.test_rlm_search_and_get_chunk.yaml
View file

File diff suppressed because one or more lines are too long

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

File diff suppressed because one or more lines are too long

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

@ -128,7 +128,7 @@ interactions:
connection:
- keep-alive
content-length:
- '7698'
- '7715'
content-type:
- application/json
host:
@ -276,6 +276,25 @@ interactions:
role: user
model: gpt-oss
reasoning_effort: low
response_format:
json_schema:
description: Result from RLM agent execution.
name: RLMResult
schema:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: json_schema
stream: false
tool_choice: auto
tools:
@ -302,29 +321,11 @@ interactions:
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
- '519'
- '522'
content-type:
- application/json
parsed_body:
@ -337,20 +338,20 @@ interactions:
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))\n"}'
name: execute_code
id: call_qqoyb2of
id: call_jsd3hga7
index: 0
type: function
created: 1772548188
id: chatcmpl-356
created: 1772626976
id: chatcmpl-987
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 44
prompt_tokens: 1728
total_tokens: 1772
completion_tokens: 46
prompt_tokens: 1684
total_tokens: 1730
status:
code: 200
message: OK
@ -363,7 +364,7 @@ interactions:
connection:
- keep-alive
content-length:
- '8137'
- '8160'
content-type:
- application/json
host:
@ -514,15 +515,34 @@ interactions:
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))\n"}'
name: execute_code
id: call_qqoyb2of
id: call_jsd3hga7
type: function
- content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"1\n","stderr":"","success":true}'
- content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))\n","stdout":"1\n","stderr":"","success":true}'
role: tool
tool_call_id: call_qqoyb2of
tool_call_id: call_jsd3hga7
model: gpt-oss
reasoning_effort: low
response_format:
json_schema:
description: Result from RLM agent execution.
name: RLMResult
schema:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
program:
description: The final consolidated program
type: string
required:
- answer
- program
type: object
strict: true
type: json_schema
stream: false
tool_choice: auto
tools:
@ -549,24 +569,6 @@ interactions:
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
@ -581,270 +583,15 @@ interactions:
message:
content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
role: assistant
created: 1772548189
id: chatcmpl-173
created: 1772626977
id: chatcmpl-959
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 40
prompt_tokens: 1814
total_tokens: 1854
status:
code: 200
message: OK
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '8443'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a Recursive Language Model (RLM) 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.
Inside execute_code, these functions are ALREADY available in the namespace. Do NOT import them - just call them with `await`:
- results = await search("query") ✓ CORRECT
- import search ✗ WRONG - will fail
- results = search("query") ✗ WRONG - must use await
## Available Functions
### await 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
### await 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
### await get_document(id_or_title) -> str | None
Get the full text content of a document by ID, title, or URI.
Returns the document content as a string, or None if not found.
### await get_chunk(chunk_id) -> dict | None
Get a specific chunk by its ID (from search results).
Returns dict with keys: chunk_id, content, document_id, document_title, headings, page_numbers, labels
Use this to retrieve full chunk details and metadata for citation.
### await get_docling_document(document_id) -> dict | None
Get the full document structure as a dict (DoclingDocument format).
Use `list_documents()` or search results to get document IDs first.
- `texts`: list of text items, each with `text`, `label` (e.g. "title", "text", "section_header", "list_item"), and `prov` (provenance with page/bounding box)
- `tables`: list of tables, each with `data` containing `grid` (list of rows, each row a list of cells with `text`), `num_rows`, `num_cols`
- `pictures`: list of figures/images with metadata
- `pages`: page dimensions and metadata
### await regex_findall(pattern, text) -> list[str]
Find all non-overlapping matches of a regular expression pattern in text.
### await regex_sub(pattern, repl, text) -> str
Replace all occurrences of a regular expression pattern with a replacement string.
### await regex_search(pattern, text) -> dict | None
Search for the first match of a pattern. Returns a dict with keys: group, groups, start, end — or None if no match.
### await regex_split(pattern, text) -> list[str]
Split text by a regular expression pattern.
### await llm(prompt) -> str
Call an LLM directly with the given prompt. Returns the response as a string.
Use this for classification, summarization, extraction, or any task where you
already have the content and just need LLM reasoning.
## Pre-loaded Documents Variable
If documents were pre-loaded for this session, a `documents` variable is available:
```python
# documents is a list of dicts with keys: id, title, uri, content
for doc in documents:
print(doc['title'], len(doc['content']))
```
Check if it exists with: `try: documents ... except NameError: ...`
## Available Python Features
The interpreter supports: variables, arithmetic, strings, f-strings, lists, dicts, tuples, sets, loops, conditionals, comprehensions, functions, async/await, `map()`, `sorted()`/`.sort(key=...)`, try/except, and the `json` module.
Not supported: imports (other than `json`), class definitions, generators/yield, match statements, decorators, `with` statements.
For pattern matching or text extraction, use the `regex_*` functions, string methods (`str.split`, `str.find`, `str.startswith`, `in` operator), or the `llm()` function.
## Strategy Guide
1. **Explore First**: Start by listing documents or searching to understand what's available. Document `title` is often None — use `uri` or `id` to identify documents instead.
2. **If get_document returns None**: Use `await list_documents()` to see available documents (check `uri` and `id`), or `await search()` to find relevant content.
3. **Iterative Refinement**: Run code, examine results, adjust your approach based on what you find.
4. **Use llm() for Classification/Extraction**: When you need to classify, summarize, or extract structured data from content you already have, use `await llm()`.
5. **Cite Your Sources**: Use get_chunk() to retrieve chunk metadata for citations. Track which documents/chunks informed your answer.
## Example Patterns
### Counting documents matching a condition
```python
docs = await list_documents(limit=100)
count = 0
for doc in docs:
content = await get_document(doc['id'])
if content and 'keyword' in content.lower():
count += 1
print(f"Found in: {doc['title']}")
print(f"Total: {count}")
```
### Extracting data with regex
```python
numbers = []
results = await search("financial data", limit=20)
for r in results:
amounts = await regex_findall(r'\$([\d,]+)', r['content'])
for a in amounts:
numbers.append(int(a.replace(',', '')))
if numbers:
print(f"Average: {sum(numbers) / len(numbers)}")
```
### Extracting tables from a document
```python
docs = await list_documents(limit=10)
for d in docs:
doc = await get_docling_document(d['id'])
if doc:
tables = doc.get('tables', [])
if tables:
print(f"{d['title']}: {len(tables)} table(s)")
for i, table in enumerate(tables):
grid = table.get('data', {}).get('grid', [])
for row in grid:
cells = [cell.get('text', '') for cell in row]
print(f" Table {i}: {cells}")
```
## Output Format
Your final response MUST be valid JSON matching this exact schema:
```json
{"answer": "Your complete answer here as a string", "program": "Your final consolidated program here as a string"}
```
- `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": "...", "program": "..."}
You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first.
role: system
- content: How many documents are available?
role: user
- content: null
reasoning: Need to list documents.
role: assistant
tool_calls:
- function:
arguments: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code
id: call_qqoyb2of
type: function
- content: '{"code":"docs = await list_documents(limit=1000)\nprint(len(docs))","stdout":"1\n","stderr":"","success":true}'
role: tool
tool_call_id: call_qqoyb2of
- content: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
role: assistant
- content: |-
Validation feedback:
Please include your response in a tool call.
Fix the errors and try again.
role: user
model: gpt-oss
reasoning_effort: low
stream: false
tool_choice: auto
tools:
- function:
description: |-
<summary>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.</summary>
<returns>
<description>Structured result with success status, stdout, and stderr.</description>
</returns>
name: execute_code
parameters:
additionalProperties: false
properties:
code:
description: Python code to execute.
type: string
required:
- code
type: object
strict: true
type: function
- function:
description: Result from RLM agent execution.
name: final_result
parameters:
additionalProperties: false
properties:
answer:
description: The answer to the user's question
type: string
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:
- '620'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: tool_calls
index: 0
message:
content: ''
reasoning: Need to output JSON in a tool call. Use final_result.
role: assistant
tool_calls:
- function:
arguments: '{"answer":"There is 1 document available in the knowledge base.","program":"docs = await list_documents(limit=1000)\nprint(len(docs))"}'
name: final_result
id: call_hli4bq9m
index: 0
type: function
created: 1772548191
id: chatcmpl-603
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 65
prompt_tokens: 1865
total_tokens: 1930
completion_tokens: 37
prompt_tokens: 1774
total_tokens: 1811
status:
code: 200
message: OK

725
tests/cassettes/test_search_filter/test_research_graph_uses_search_filter.yaml
View file

File diff suppressed because one or more lines are too long

1301
tests/cassettes/test_search_filter/test_search_filter_none_searches_all.yaml
View file

File diff suppressed because one or more lines are too long

32
tests/test_utils.py
View file

@ -139,6 +139,38 @@ Emoji test: 🚀 ✅ 📝"""
assert "🚀" in result_markdown
def test_structured_output_type_native():
from pydantic_ai.output import NativeOutput
from haiku.rag.utils import structured_output_type
model = get_model(ModelConfig(provider="openai", name="gpt-4o"))
result = structured_output_type(str, model)
assert isinstance(result, NativeOutput)
assert result.outputs is str
def test_structured_output_type_tool_fallback():
from pydantic_ai.output import ToolOutput
from haiku.rag.utils import structured_output_type
model = get_model(ModelConfig(provider="ollama", name="qwen3"))
result = structured_output_type(str, model)
assert isinstance(result, ToolOutput)
assert result.output is str
def test_structured_output_type_string_model():
from pydantic_ai.output import ToolOutput
from haiku.rag.utils import structured_output_type
result = structured_output_type(str, "unknown:model")
assert isinstance(result, ToolOutput)
assert result.output is str
def test_get_model_ollama():
"""Test get_model returns OpenAIChatModel for Ollama."""
model_config = ModelConfig(provider="ollama", name="llama3")