Add configurable structured output mode (tool vs native)

CHANGELOG.md

@@ -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()`
+- **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
 
 ### Changed
 

docs/configuration/providers.md

@@ -29,6 +29,7 @@ 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
 
@@ -66,6 +67,23 @@ 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.

haiku_rag_slim/haiku/rag/agents/qa/agent.py

@@ -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
@@ -60,7 +59,7 @@ class QuestionAnswerAgent:
         agent: Agent[_QARunDeps, RawSearchAnswer] = Agent(  # ty: ignore[invalid-assignment]
             model=get_model(self._model_config, self._config),
             deps_type=_QARunDeps,
-            output_type=ToolOutput(RawSearchAnswer, max_retries=3),
+            output_type=structured_output_type(RawSearchAnswer, self._model_config),
             instructions=self._system_prompt,
             toolsets=[search_toolset],
             retries=3,

haiku_rag_slim/haiku/rag/agents/research/graph.py

@@ -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:
@@ -68,7 +67,7 @@ async def _iterative_plan_logic(
 
     plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent(  # type: ignore[assignment]
         model=get_model(model_config, config),
-        output_type=ToolOutput(IterativePlanResult, max_retries=3),
+        output_type=structured_output_type(IterativePlanResult, model_config),
         instructions=effective_prompt,
         retries=3,
         deps_type=ResearchDependencies,
@@ -117,7 +116,7 @@ async def _search_one_step_logic(
     async with deps.semaphore:
         agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
-            output_type=ToolOutput(RawSearchAnswer, max_retries=3),
+            output_type=structured_output_type(RawSearchAnswer, model_config),
             instructions=search_prompt,
             retries=3,
             deps_type=ResearchDependencies,
@@ -218,7 +217,7 @@ def build_research_graph(
 
         agent: Agent[ResearchDependencies, ResearchReport] = Agent(  # type: ignore[assignment]
             model=get_model(model_config, config),
-            output_type=ToolOutput(ResearchReport, max_retries=3),
+            output_type=structured_output_type(ResearchReport, model_config),
             instructions=synthesis_prompt,
             retries=3,
             deps_type=ResearchDependencies,

haiku_rag_slim/haiku/rag/agents/rlm/agent.py

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

haiku_rag_slim/haiku/rag/config/models.py

@@ -25,6 +25,7 @@ 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):

haiku_rag_slim/haiku/rag/utils.py

@@ -306,6 +306,19 @@ def get_model(
         return f"{provider}:{model}"
 
 
+def structured_output_type(
+    result_type: type,
+    model_config: "ModelConfig",
+    max_retries: int = 3,
+) -> Any:
+    """Return a ToolOutput or NativeOutput wrapper based on model config."""
+    from pydantic_ai.output import NativeOutput, ToolOutput
+
+    if model_config.structured_output == "native":
+        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)

tests/agents/rlm/test_agent.py

@@ -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
@@ -17,12 +16,24 @@ def vcr_cassette_dir():
 
 class TestCreateRLMAgent:
     def test_creates_agent_with_correct_types(self):
+        from pydantic_ai.output import ToolOutput
+
         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_agent_has_execute_code_tool(self):
         agent = create_rlm_agent(Config)
         tool_names = list(agent._function_toolset.tools.keys())

tests/test_utils.py

@@ -139,6 +139,28 @@ 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)
+    assert isinstance(result, NativeOutput)
+    assert result.outputs is str
+
+
 def test_get_model_ollama():
     """Test get_model returns OpenAIChatModel for Ollama."""
     model_config = ModelConfig(provider="ollama", name="llama3")