Auto-detect structured output mode from model profile. Remove the structured_output config field from ModelConfig.

This commit is contained in:
Yiorgis Gozadinos 2026-03-04 14:51:09 +02:00
parent fdd7c21757
commit 28e90e5a87
No known key found for this signature in database
19 changed files with 3113 additions and 7595 deletions

View file

@ -4,7 +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()`
- **Configurable structured output mode**: New `structured_output` setting on model config (`"tool"` or `"native"`). Models like OpenAI and Anthropic can use native JSON schema enforcement instead of the default tool-call approach
- **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

View file

@ -29,7 +29,6 @@ qa:
- **max_tokens**: Maximum tokens in response
- **enable_thinking**: Control reasoning behavior (see below)
- **base_url**: Custom endpoint for OpenAI-compatible servers (vLLM, LM Studio, etc.)
- **structured_output**: How the model returns structured data — `"tool"` (default) or `"native"` (see below)
### Thinking Control
@ -67,23 +66,6 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- Disable for simple queries, RAG workflows, speed-critical applications
- Enable for complex reasoning, mathematical problems, research tasks
### Structured Output Mode
The `structured_output` setting controls how the model returns structured data (JSON responses for QA citations, research reports, etc.).
```yaml
qa:
model:
provider: openai
name: gpt-4o
structured_output: native # Use model's native JSON schema enforcement
```
**Values:**
- `"tool"` (default): Uses a fake tool call to extract structured output. Works with all providers that support tool calling.
- `"native"`: Uses the model's native JSON schema enforcement. Can be more reliable for models that support it well (OpenAI, Anthropic), but may not work with all models.
## Embedding Providers
Embedding models require three settings: `provider`, `name`, and `vector_dim`. Optionally, use `base_url` for OpenAI-compatible servers.

View file

@ -56,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=structured_output_type(RawSearchAnswer, self._model_config),
output_type=structured_output_type(RawSearchAnswer, model),
instructions=self._system_prompt,
toolsets=[search_toolset],
retries=3,

View file

@ -65,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=structured_output_type(IterativePlanResult, model_config),
model=model,
output_type=structured_output_type(IterativePlanResult, model),
instructions=effective_prompt,
retries=3,
deps_type=ResearchDependencies,
@ -114,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=structured_output_type(RawSearchAnswer, model_config),
model=model,
output_type=structured_output_type(RawSearchAnswer, model),
instructions=search_prompt,
retries=3,
deps_type=ResearchDependencies,
@ -215,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=structured_output_type(ResearchReport, model_config),
model=model,
output_type=structured_output_type(ResearchReport, model),
instructions=synthesis_prompt,
retries=3,
deps_type=ResearchDependencies,

View file

@ -25,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=structured_output_type(RLMResult, config.rlm.model),
output_type=structured_output_type(RLMResult, model),
instructions=RLM_SYSTEM_PROMPT,
retries=3,
)

View file

@ -25,7 +25,6 @@ class ModelConfig(BaseModel):
enable_thinking: bool | None = None
temperature: float | None = None
max_tokens: int | None = None
structured_output: Literal["tool", "native"] = "tool"
class EmbeddingModelConfig(BaseModel):

View file

@ -308,13 +308,14 @@ def get_model(
def structured_output_type(
result_type: type,
model_config: "ModelConfig",
model: Any,
max_retries: int = 3,
) -> Any:
"""Return a ToolOutput or NativeOutput wrapper based on model config."""
"""Return a NativeOutput or ToolOutput wrapper based on model capability."""
from pydantic_ai.models import Model
from pydantic_ai.output import NativeOutput, ToolOutput
if model_config.structured_output == "native":
if isinstance(model, Model) and model.profile.supports_json_schema_output:
return NativeOutput(result_type)
return ToolOutput(result_type, max_retries=max_retries)

View file

@ -15,25 +15,25 @@ def vcr_cassette_dir():
class TestCreateRLMAgent:
def test_creates_agent_with_correct_types(self):
from pydantic_ai.output import ToolOutput
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, ToolOutput)
assert agent.output_type.output is RLMResult
def test_creates_agent_with_native_output(self):
from pydantic_ai.output import NativeOutput
config = AppConfig()
config.rlm.model.structured_output = "native"
agent = create_rlm_agent(config)
assert isinstance(agent, Agent)
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
def test_agent_has_execute_code_tool(self):
agent = create_rlm_agent(Config)
tool_names = list(agent._function_toolset.tools.keys())

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

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

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

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

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -139,28 +139,38 @@ Emoji test: 🚀 ✅ 📝"""
assert "🚀" in result_markdown
def test_structured_output_type_default():
from pydantic_ai.output import ToolOutput
from haiku.rag.utils import structured_output_type
mc = ModelConfig(provider="openai", name="gpt-4o")
result = structured_output_type(str, mc)
assert isinstance(result, ToolOutput)
assert result.output is str
def test_structured_output_type_native():
from pydantic_ai.output import NativeOutput
from haiku.rag.utils import structured_output_type
mc = ModelConfig(provider="openai", name="gpt-4o", structured_output="native")
result = structured_output_type(str, mc)
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")