danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use ToolOutput for structured output.

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-03-03 14:18:28 +02:00
parent 1eca0fc308
commit 034ee27daf
No known key found for this signature in database
4 changed files with 9 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -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,

6
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -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,
)

3
haiku_rag_slim/haiku/rag/agents/rlm/agent.py
View file

@ -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,
)

4
tests/agents/rlm/test_agent.py
View file

@ -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)