Auto-detect structured output mode from model profile. Remove the structured_output config field from ModelConfig.
This commit is contained in:
parent
fdd7c21757
commit
28e90e5a87
19 changed files with 3113 additions and 7595 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
File diff suppressed because it is too large
Load diff
|
|
@ -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
|
|
@ -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
|
|
@ -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")
|
||||
|
|
|
|||
Loading…
Reference in a new issue