Use ToolOutput in order to specify 3 retries

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

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