Merge pull request #22896 from dannon/chatgxy-agent-retries-config

[26.1] Make GalaxyAI agent retry budget configurable (fix "Exceeded maximum output retries")

doc/source/admin/galaxy_options.rst

@@ -5671,8 +5671,8 @@
 :Description:
     Configuration for AI inference services used by agents and
     visualization plugins. Supports per-agent or per-plugin model,
-    temperature, max_tokens, api_key, api_base_url, and enabled
-    settings. Valid keys include agent types (e.g. router,
+    temperature, max_tokens, retries, api_key, api_base_url, and
+    enabled settings. Valid keys include agent types (e.g. router,
     error_analysis) and plugin names (e.g. jupyterlite). Agents and
     plugins inherit from 'default' configuration, which itself falls
     back to global ai_model/ai_api_key settings. All agents are
@@ -5684,7 +5684,14 @@
     test/integration/static_agents.yml } Per-agent or default-block
     ``structured_output_override: true|false`` beats the model
     capability table -- see ``agent_model_capabilities_file`` for the
-    table's location and contents.
+    table's location and contents. Per-agent or default-block
+    ``retries`` sets the pydantic-ai retry budget (tool calls and
+    output validation); it defaults to 3. Raise it if a model
+    intermittently fails to produce conforming output ("Exceeded
+    maximum output retries"). custom_tool's producer keeps a budget of
+    0 because it runs its own reflection loop; a shared ``default``
+    block does not change that -- set ``custom_tool.retries``
+    explicitly to override it.
 :Default: ``None``
 :Type: any
 
@@ -6267,3 +6274,6 @@
     for user defined tools.
 :Default: ``false``
 :Type: bool
+
+
+

lib/galaxy/agents/base.py

@@ -29,6 +29,7 @@ from typing import (
 
 import yaml
 
+from galaxy.exceptions import ConfigurationError
 from galaxy.managers.context import ProvidesUserContext
 from galaxy.model import User
 from galaxy.schema.agents import (
@@ -377,6 +378,11 @@ class BaseGalaxyAgent(ABC):
     # backend we currently support (smallest is Qwen3-32B at 32k context).
     DEFAULT_MAX_TOKENS = 8192
 
+    # Retry budget passed to Agent(retries=...) (tool calls and output validation).
+    # pydantic-ai defaults to 1; 3 gives a flaky model a couple more chances to
+    # produce conforming output before the run fails.
+    DEFAULT_AGENT_RETRIES = 3
+
     def __init__(self, deps: GalaxyAgentDependencies):
         self.deps = deps
 
@@ -812,6 +818,19 @@ class BaseGalaxyAgent(ABC):
                 return self.deps.config.ai_api_base_url
         return default
 
+    def _get_agent_specific_config(self, key: str, default: Any = None) -> Any:
+        """Read a value only from this agent's own ``inference_services`` block.
+
+        Unlike :meth:`_get_agent_config`, this skips the shared ``default`` block so a
+        caller-pinned builtin is overridden only by an explicit per-agent entry.
+        """
+        inference_config = getattr(self.deps.config, "inference_services", {})
+        if isinstance(inference_config, dict):
+            agent_specific = inference_config.get(self.agent_type, {})
+            if isinstance(agent_specific, dict) and key in agent_specific:
+                return agent_specific[key]
+        return default
+
     def _get_model_name(self) -> str:
         return self._get_agent_config("model", "gpt-4o-mini")
 
@@ -856,6 +875,30 @@ class BaseGalaxyAgent(ABC):
     def _get_max_tokens(self) -> int:
         return self._get_agent_config("max_tokens", self.DEFAULT_MAX_TOKENS)
 
+    def _get_retries(self, default: Optional[int] = None) -> int:
+        """Retry budget for the agent's pydantic-ai ``Agent(retries=...)``.
+
+        With no ``default``, the budget resolves per-agent > ``default`` block >
+        builtin (:attr:`DEFAULT_AGENT_RETRIES`). A caller-pinned ``default`` (e.g.
+        custom_tool's producer keeps 0 so its own reflection loop owns the retry) is
+        a correctness requirement, not a tunable: only an explicit per-agent
+        ``retries`` overrides it -- a shared ``default`` block must not silently
+        re-enable pydantic-ai retries there.
+        """
+        if default is None:
+            raw = self._get_agent_config("retries", self.DEFAULT_AGENT_RETRIES)
+        else:
+            raw = self._get_agent_specific_config("retries", default)
+        try:
+            retries = int(raw)
+        except (TypeError, ValueError):
+            retries = None
+        if retries is None or retries < 0:
+            raise ConfigurationError(
+                f"inference_services 'retries' for agent '{self.agent_type}' must be a non-negative integer, got {raw!r}"
+            )
+        return retries
+
     async def _call_agent_from_tool(
         self,
         agent_type: str,
@@ -906,6 +949,7 @@ class SimpleGalaxyAgent(BaseGalaxyAgent):
             self._get_model(),
             deps_type=GalaxyAgentDependencies,
             system_prompt=self.get_system_prompt(),
+            retries=self._get_retries(),
         )
 
     def _format_response(self, result: Any, query: str, context: dict[str, Any]) -> AgentResponse:

lib/galaxy/agents/custom_tool.py

@@ -111,15 +111,16 @@ class CustomToolAgent(BaseGalaxyAgent):
     def _create_agent(self) -> Agent[GalaxyAgentDependencies, Any]:
         """Create agent with UserToolSource as the output type.
 
-        Sets retries=0 because the agent's explicit reflection loop
-        owns the validation retry (to provide a better prompt).
+        Defaults retries to 0 because the agent's explicit reflection loop owns the
+        validation retry (to provide a better prompt); operators can still override
+        via inference_services.
         """
         return Agent(
             self._get_model(),
             deps_type=GalaxyAgentDependencies,
             output_type=UserToolSource,
             system_prompt=self.get_system_prompt(),
-            retries=0,
+            retries=self._get_retries(default=0),
         )
 
     def get_system_prompt(self) -> str:
@@ -139,6 +140,7 @@ class CustomToolAgent(BaseGalaxyAgent):
                 deps_type=GalaxyAgentDependencies,
                 output_type=CritiqueReport,
                 system_prompt=self._get_critic_system_prompt(),
+                retries=self._get_retries(),
             )
         return self._critic_agent
 

lib/galaxy/agents/error_analysis.py

@@ -57,12 +57,14 @@ class ErrorAnalysisAgent(BaseGalaxyAgent):
                 deps_type=GalaxyAgentDependencies,
                 output_type=ErrorAnalysisResult,
                 system_prompt=self.get_system_prompt(),
+                retries=self._get_retries(),
             )
         else:
             agent = Agent(
                 self._get_model(),
                 deps_type=GalaxyAgentDependencies,
                 system_prompt=self._get_simple_system_prompt(),
+                retries=self._get_retries(),
             )
 
         return agent

lib/galaxy/agents/gtn_training.py

@@ -101,6 +101,7 @@ class GTNTrainingAgent(BaseGalaxyAgent):
                 self._get_model(),
                 deps_type=GalaxyAgentDependencies,
                 system_prompt=self._get_simple_system_prompt(),
+                retries=self._get_retries(),
             )
 
         agent = Agent(
@@ -110,8 +111,9 @@ class GTNTrainingAgent(BaseGalaxyAgent):
             system_prompt=self.get_system_prompt(),
             # gpt-oss occasionally emits prose instead of a valid GTNSearchResponse;
             # the pydantic-ai default of 1 output retry turns that into a hard error.
-            # Allow a couple more attempts so an occasional malformed output recovers.
-            retries=3,
+            # Configurable via inference_services, defaulting to 3 so an occasional
+            # malformed output recovers.
+            retries=self._get_retries(),
         )
 
         @agent.tool

lib/galaxy/agents/history.py

@@ -39,6 +39,7 @@ class HistoryAgent(BaseGalaxyAgent):
             self._get_model(),
             deps_type=GalaxyAgentDependencies,
             system_prompt=self.get_system_prompt(),
+            retries=self._get_retries(),
         )
 
         @agent.tool

lib/galaxy/agents/orchestrator.py

@@ -61,12 +61,14 @@ class WorkflowOrchestratorAgent(BaseGalaxyAgent):
                 deps_type=GalaxyAgentDependencies,
                 output_type=AgentPlan,
                 system_prompt=self.get_system_prompt(),
+                retries=self._get_retries(),
             )
         else:
             agent = Agent(
                 self._get_model(),
                 deps_type=GalaxyAgentDependencies,
                 system_prompt=self._get_simple_system_prompt(),
+                retries=self._get_retries(),
             )
 
         return agent

lib/galaxy/agents/page_assistant.py

@@ -209,6 +209,7 @@ class PageAssistantAgent(BaseGalaxyAgent):
                 ),
                 str,  # Conversational response — questions, clarifications, no document content.
             ],
+            retries=self._get_retries(),
         )
 
         # Dynamic system prompt — reads self.page_content at call time

lib/galaxy/agents/router.py

@@ -27,6 +27,7 @@ from pydantic import ValidationError
 from pydantic_ai import (
     Agent,
     RunContext,
+    UnexpectedModelBehavior,
 )
 
 from galaxy.agents.operations import AgentOperationsManager
@@ -65,6 +66,7 @@ class QueryRouterAgent(BaseGalaxyAgent):
                 self._get_model(),
                 deps_type=GalaxyAgentDependencies,
                 system_prompt=self._get_simple_system_prompt(),
+                retries=self._get_retries(),
             )
 
         error_handoff = self._create_error_analysis_handoff()
@@ -91,6 +93,7 @@ class QueryRouterAgent(BaseGalaxyAgent):
                 str,  # Default: answer directly
             ],
             system_prompt=self.get_system_prompt(),
+            retries=self._get_retries(),
         )
 
         self._register_fast_path_tools(agent)
@@ -579,7 +582,7 @@ class QueryRouterAgent(BaseGalaxyAgent):
                 query=query,
             )
 
-        except (OSError, ValueError) as e:
+        except (UnexpectedModelBehavior, OSError, ValueError) as e:
             log.warning(f"Router agent error, using fallback: {e}")
             return self._handle_fallback(query, context, str(e))
 

lib/galaxy/agents/tools.py

@@ -78,12 +78,14 @@ class ToolRecommendationAgent(BaseGalaxyAgent):
                 deps_type=GalaxyAgentDependencies,
                 output_type=SimplifiedToolRecommendationResult,
                 system_prompt=self.get_system_prompt(),
+                retries=self._get_retries(),
             )
         else:
             agent = Agent(
                 self._get_model(),
                 deps_type=GalaxyAgentDependencies,
                 system_prompt=self._get_simple_system_prompt(),
+                retries=self._get_retries(),
             )
 
         @agent.tool

lib/galaxy/config/sample/galaxy.yml.sample

@@ -1,21 +1,21 @@
 # Galaxy is configured by default to be usable in a single-user development
 # environment.  To tune the application for a multi-user production
 # environment, see the documentation at:
-#
+# 
 #  https://docs.galaxyproject.org/en/master/admin/production.html
-#
+# 
 # Throughout this sample configuration file, except where stated otherwise,
 # uncommented values override the default if left unset, whereas commented
 # values are set to the default value.  Relative paths are relative to the root
 # Galaxy directory.
-#
+# 
 # Examples of many of these options are explained in more detail in the Galaxy
 # Community Hub.
-#
+# 
 #   https://galaxyproject.org/admin/config
-#
+# 
 # Config hackers are encouraged to check there before asking for help.
-#
+# 
 # Configuration for Gravity process manager.
 # ``uwsgi:`` section will be ignored if Galaxy is started via Gravity commands (e.g ``./run.sh``, ``galaxy`` or ``galaxyctl``).
 gravity:
@@ -3067,7 +3067,7 @@ galaxy:
 
   # Configuration for AI inference services used by agents and
   # visualization plugins. Supports per-agent or per-plugin model,
-  # temperature, max_tokens, api_key, api_base_url, and enabled
+  # temperature, max_tokens, retries, api_key, api_base_url, and enabled
   # settings. Valid keys include agent types (e.g. router,
   # error_analysis) and plugin names (e.g. jupyterlite). Agents and
   # plugins inherit from 'default' configuration, which itself falls
@@ -3080,7 +3080,14 @@ galaxy:
   # test/integration/static_agents.yml } Per-agent or default-block
   # ``structured_output_override: true|false`` beats the model
   # capability table -- see ``agent_model_capabilities_file`` for the
-  # table's location and contents.
+  # table's location and contents. Per-agent or default-block
+  # ``retries`` sets the pydantic-ai retry budget (tool calls and output
+  # validation); it defaults to 3. Raise it if a model intermittently
+  # fails to produce conforming output ("Exceeded maximum output
+  # retries"). custom_tool's producer keeps a budget of 0 because it
+  # runs its own reflection loop; a shared ``default`` block does not
+  # change that -- set ``custom_tool.retries`` explicitly to override
+  # it.
   #inference_services: null
 
   # YAML file with capability hints for agent inference models. Maps
@@ -3365,3 +3372,4 @@ galaxy:
   # Enable beta tool formats (yaml, cwl, ...) which is a prerequisite
   # for user defined tools.
   #enable_beta_tool_formats: false
+

lib/galaxy/config/schemas/config_schema.yml

@@ -4197,7 +4197,7 @@ mapping:
         required: false
         desc: |
           Configuration for AI inference services used by agents and visualization plugins.
-          Supports per-agent or per-plugin model, temperature, max_tokens, api_key, api_base_url, and enabled settings.
+          Supports per-agent or per-plugin model, temperature, max_tokens, retries, api_key, api_base_url, and enabled settings.
           Valid keys include agent types (e.g. router, error_analysis) and plugin names (e.g. jupyterlite).
           Agents and plugins inherit from 'default' configuration, which itself falls back to global ai_model/ai_api_key settings.
           All agents are enabled by default.
@@ -4208,6 +4208,12 @@ mapping:
           Per-agent or default-block ``structured_output_override: true|false``
           beats the model capability table -- see ``agent_model_capabilities_file``
           for the table's location and contents.
+          Per-agent or default-block ``retries`` sets the pydantic-ai retry budget
+          (tool calls and output validation); it defaults to 3. Raise it if a model
+          intermittently fails to produce conforming output ("Exceeded maximum output
+          retries"). custom_tool's producer keeps a budget of 0 because it runs its
+          own reflection loop; a shared ``default`` block does not change that -- set
+          ``custom_tool.retries`` explicitly to override it.
 
       agent_model_capabilities_file:
         type: str

packages/app/setup.cfg

@@ -49,7 +49,7 @@ install_requires =
     pebble
     pulsar-galaxy-lib>=0.15.0.dev0
     pydantic>=2.7.4
-    pydantic-ai>=0.1.16
+    pydantic-ai>=1.99.0
     pysam>=0.21
     python-slugify
     PyJWT

test/unit/app/test_agents.py

@@ -46,6 +46,7 @@ from pydantic_ai.messages import (
     TextPart,
     UserPromptPart,
 )
+from pydantic_ai.models.function import FunctionModel
 from pydantic_ai.models.test import TestModel
 
 from galaxy.agents import (
@@ -83,6 +84,7 @@ from galaxy.agents.page_assistant import (
     FullReplacementEdit,
     SectionPatchEdit,
 )
+from galaxy.exceptions import ConfigurationError
 from galaxy.schema.agents import ConfidenceLevel
 from galaxy.tool_util_models import UserToolSource
 from galaxy.util.unittest_utils import pytestmark_live_llm
@@ -158,6 +160,88 @@ class TestAgentUnitMocked:
         assert router_agent._get_agent_config("temperature", 0.5) == 0.5
         assert router_agent._get_agent_config("max_tokens", 1500) == 1500
 
+    def test_get_retries_resolution(self):
+        # Unset -> builtin default (bumped above pydantic-ai's default of 1).
+        self.mock_config.inference_services = None
+        router = QueryRouterAgent(self.deps)
+        assert router._get_retries() == agents_base.BaseGalaxyAgent.DEFAULT_AGENT_RETRIES == 3
+        # A caller-supplied default wins over the builtin when the key is unset
+        # (custom_tool's producer relies on this to keep its reflection-loop 0).
+        assert router._get_retries(default=0) == 0
+
+        # default entry applies to every agent; a per-agent entry overrides it.
+        # String values from YAML are coerced to int.
+        self.mock_config.inference_services = {
+            "default": {"retries": "4"},
+            "router": {"retries": 2},
+        }
+        assert QueryRouterAgent(self.deps)._get_retries() == 2
+        error_retries = ErrorAnalysisAgent(self.deps)._get_retries()
+        assert error_retries == 4 and isinstance(error_retries, int)
+
+    def test_configured_retries_wired_into_agent(self):
+        # The configured budget reaches the constructed pydantic-ai Agent
+        # (both tool and output budgets, via Agent(retries=N)).
+        self.mock_config.inference_services = {"default": {"retries": 4}}
+        router = QueryRouterAgent(self.deps)
+        assert router.agent._max_output_retries == 4
+        assert router.agent._max_tool_retries == 4
+
+        # custom_tool's producer keeps its reflection-loop default of 0 when unconfigured.
+        self.mock_config.inference_services = None
+        producer = CustomToolAgent(self.deps)
+        assert producer.agent._max_output_retries == 0
+
+    def test_producer_retries_default_ignores_default_block(self):
+        # custom_tool's producer pins retries=0 so its own reflection loop owns
+        # the retry. A shared `default` block must NOT silently re-enable
+        # pydantic-ai retries there -- only an explicit per-agent entry may.
+        producer = CustomToolAgent(self.deps)
+
+        self.mock_config.inference_services = {"default": {"retries": 5}}
+        assert producer._get_retries(default=0) == 0
+        # The normal (critic) lookup still honors the default block.
+        assert producer._get_retries() == 5
+
+        # An explicit custom_tool entry still overrides the pinned builtin.
+        self.mock_config.inference_services = {"custom_tool": {"retries": 7}}
+        assert producer._get_retries(default=0) == 7
+
+    def test_invalid_retries_config_raises_configuration_error(self):
+        # Non-numeric, blank, or negative `retries` is operator misconfiguration;
+        # it must surface as a clear ConfigurationError rather than a bare
+        # TypeError/ValueError (which the manager mistakes for an unknown-agent
+        # fallback) or a silently-broken negative budget that fails every request.
+        router = QueryRouterAgent(self.deps)
+
+        for bad in ("three", None, -1):
+            self.mock_config.inference_services = {"default": {"retries": bad}}
+            with pytest.raises(ConfigurationError, match="retries"):
+                router._get_retries()
+
+        # 0 is valid (custom_tool's producer relies on it) and must not raise.
+        self.mock_config.inference_services = {"default": {"retries": 0}}
+        assert router._get_retries() == 0
+
+    @pytest.mark.asyncio
+    async def test_router_falls_back_on_output_retry_exhaustion(self):
+        # When the model never produces a valid structured output, pydantic-ai
+        # raises UnexpectedModelBehavior after exhausting the output-retry budget.
+        # The router must degrade to its graceful fallback rather than propagate.
+        self.mock_config.inference_services = None
+        router = QueryRouterAgent(self.deps)
+
+        # Empty response (not text): str is in the router's output_type union, so
+        # any text would succeed via the str branch and never trigger a retry.
+        def empty_response(messages, info):
+            return ModelResponse(parts=[])
+
+        with router.agent.override(model=FunctionModel(empty_response)):
+            response = await router.process("Summarize my history")
+
+        assert isinstance(response, AgentResponse)
+        assert "unavailable" in response.content.lower()
+
     @pytest.mark.asyncio
     async def test_custom_tool_agent_structured_output(self):
         # Test with a model that supports structured output (gpt-4o)