From 034ee27daf7a6656d70c1fa2a9b051a7d6724c29 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 3 Mar 2026 14:18:28 +0200 Subject: [PATCH] Use ToolOutput for structured output. --- haiku_rag_slim/haiku/rag/agents/qa/agent.py | 4 ++-- haiku_rag_slim/haiku/rag/agents/research/graph.py | 6 ++---- haiku_rag_slim/haiku/rag/agents/rlm/agent.py | 3 ++- tests/agents/rlm/test_agent.py | 4 +++- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/qa/agent.py b/haiku_rag_slim/haiku/rag/agents/qa/agent.py index 9e5ea2c0..11de13ce 100644 --- a/haiku_rag_slim/haiku/rag/agents/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/qa/agent.py @@ -1,6 +1,7 @@ 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 ( @@ -59,8 +60,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=RawSearchAnswer, - output_retries=3, + output_type=ToolOutput(RawSearchAnswer, max_retries=3), 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 e445622e..d6996dcb 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/haiku_rag_slim/haiku/rag/agents/research/graph.py @@ -68,10 +68,9 @@ async def _iterative_plan_logic( plan_agent: Agent[ResearchDependencies, IterativePlanResult] = Agent( # type: ignore[assignment] model=get_model(model_config, config), - output_type=IterativePlanResult, + output_type=ToolOutput(IterativePlanResult, max_retries=3), instructions=effective_prompt, retries=3, - output_retries=3, deps_type=ResearchDependencies, ) @@ -219,10 +218,9 @@ def build_research_graph( agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment] model=get_model(model_config, config), - output_type=ResearchReport, + output_type=ToolOutput(ResearchReport, max_retries=3), instructions=synthesis_prompt, retries=3, - output_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 4c234832..19ab3c4a 100644 --- a/haiku_rag_slim/haiku/rag/agents/rlm/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/rlm/agent.py @@ -1,4 +1,5 @@ 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 @@ -25,7 +26,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=RLMResult, + output_type=ToolOutput(RLMResult, max_retries=3), instructions=RLM_SYSTEM_PROMPT, retries=3, ) diff --git a/tests/agents/rlm/test_agent.py b/tests/agents/rlm/test_agent.py index f78af4e1..ba1cdf6c 100644 --- a/tests/agents/rlm/test_agent.py +++ b/tests/agents/rlm/test_agent.py @@ -2,6 +2,7 @@ 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 @@ -19,7 +20,8 @@ class TestCreateRLMAgent: agent = create_rlm_agent(Config) assert isinstance(agent, Agent) assert agent.deps_type is RLMDeps - assert agent.output_type is RLMResult + 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)