diff --git a/CHANGELOG.md b/CHANGELOG.md index b97904b0..2fb4ac16 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index 7f7bbded..58888711 100644 --- a/docs/configuration/providers.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/agents/qa/agent.py b/haiku_rag_slim/haiku/rag/agents/qa/agent.py index 11de13ce..dc69c4c1 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/agent.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index d6996dcb..3d6b8283 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/agents/rlm/agent.py b/haiku_rag_slim/haiku/rag/agents/rlm/agent.py index 19ab3c4a..69c3ef0d 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/agent.py +++ b/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, ) diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index a07c857f..7bf1cc30 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/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): diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 0c89f49c..b5aecff2 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/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) diff --git a/tests/agents/rlm/test_agent.py b/tests/agents/rlm/test_agent.py index ba1cdf6c..aec486d4 100644 --- a/tests/agents/rlm/test_agent.py +++ b/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()) diff --git a/tests/test_utils.py b/tests/test_utils.py index fd3ad83d..33e039a5 100644 --- a/tests/test_utils.py +++ b/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")