From ecad4c28b4c675c6dc3f984f8cb113c0c0dbc490 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 17 Sep 2025 13:06:53 +0300 Subject: [PATCH] Use ToolOutput in order to specify 3 retries --- src/haiku/rag/research/base.py | 5 ++--- src/haiku/rag/research/prompts.py | 2 ++ 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index 517bf26d..11c3c64a 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -1,11 +1,10 @@ -from __future__ import annotations - from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any from pydantic import BaseModel, Field from pydantic_ai import Agent from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.output import ToolOutput from pydantic_ai.providers.ollama import OllamaProvider from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai.run import AgentRunResult @@ -37,7 +36,7 @@ class BaseResearchAgent[T](ABC): self._agent = Agent( model=model_obj, deps_type=ResearchDependencies, - output_type=self.output_type, + output_type=ToolOutput(self.output_type, max_retries=3), system_prompt=self.get_system_prompt(), ) diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index ca540f0f..4bcec6a6 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -38,6 +38,8 @@ Tool usage: - The tool returns XML containing only a list of snippets, where each snippet has the verbatim `text`, a `score` indicating relevance, and the `document_uri` it came from. +- You may call the tool multiple times to refine or broaden context, but do not + exceed 3 total tool calls per question. Prefer precision over volume. - Use scores to prioritize evidence, but include only the minimal subset of snippet texts (verbatim) in SearchAnswer.context. - Set SearchAnswer.sources to the matching document_uris for the snippets you