From deb7ecd706e4021305ddcb1861bb84db0a0875ef Mon Sep 17 00:00:00 2001 From: Dannon Baker Date: Wed, 10 Jun 2026 18:45:44 -0400 Subject: [PATCH] Make agent retry budget configurable via inference_services ChatGXY's router was hard-failing with 'Exceeded maximum output retries (1)': its pydantic-ai Agent used the default retry budget of 1, and the router's process() only caught (OSError, ValueError), so an UnexpectedModelBehavior sailed past the graceful _handle_fallback and surfaced to the user as a generic 500. This adds a per-agent 'retries' knob to inference_services, resolved through the same _get_agent_config chain as temperature/max_tokens and wired into every agent's Agent(retries=...). The default goes from pydantic-ai's 1 to 3 so a transient bout of non-conforming model output recovers instead of failing the run. custom_tool's producer keeps its retries=0 reflection-loop default but can now be overridden. The router also catches UnexpectedModelBehavior so a genuinely exhausted budget degrades to the existing fallback message instead of a 500. Adds unit tests for the resolution chain, the config->Agent wiring, and the router's graceful fallback (forced via an empty ModelResponse). --- doc/source/admin/galaxy_options.rst | 14 ++++-- lib/galaxy/agents/base.py | 14 ++++++ lib/galaxy/agents/custom_tool.py | 8 ++-- lib/galaxy/agents/error_analysis.py | 2 + lib/galaxy/agents/gtn_training.py | 6 ++- lib/galaxy/agents/history.py | 1 + lib/galaxy/agents/orchestrator.py | 2 + lib/galaxy/agents/page_assistant.py | 1 + lib/galaxy/agents/router.py | 5 +- lib/galaxy/agents/tools.py | 2 + lib/galaxy/config/sample/galaxy.yml.sample | 22 +++++---- lib/galaxy/config/schemas/config_schema.yml | 7 ++- test/unit/app/test_agents.py | 52 +++++++++++++++++++++ 13 files changed, 118 insertions(+), 18 deletions(-) diff --git a/doc/source/admin/galaxy_options.rst b/doc/source/admin/galaxy_options.rst index dbb525bed4b..b64062ec8cd 100644 --- a/doc/source/admin/galaxy_options.rst +++ b/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,12 @@ 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 defaults to 0 + since it runs its own reflection loop. :Default: ``None`` :Type: any @@ -6267,3 +6272,6 @@ for user defined tools. :Default: ``false`` :Type: bool + + + diff --git a/lib/galaxy/agents/base.py b/lib/galaxy/agents/base.py index 3cd42d0ba4e..fa38b184fcb 100644 --- a/lib/galaxy/agents/base.py +++ b/lib/galaxy/agents/base.py @@ -377,6 +377,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 @@ -856,6 +861,14 @@ 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=...)``. + + ``default`` lets a caller override the builtin (e.g. custom_tool's producer + keeps 0 so its own reflection loop owns the retry). + """ + return int(self._get_agent_config("retries", self.DEFAULT_AGENT_RETRIES if default is None else default)) + async def _call_agent_from_tool( self, agent_type: str, @@ -906,6 +919,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: diff --git a/lib/galaxy/agents/custom_tool.py b/lib/galaxy/agents/custom_tool.py index 4c07b8afe60..f80a743cb49 100644 --- a/lib/galaxy/agents/custom_tool.py +++ b/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 diff --git a/lib/galaxy/agents/error_analysis.py b/lib/galaxy/agents/error_analysis.py index f56fcb1719d..4aea0134422 100644 --- a/lib/galaxy/agents/error_analysis.py +++ b/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 diff --git a/lib/galaxy/agents/gtn_training.py b/lib/galaxy/agents/gtn_training.py index 50229ef8586..0e533d8ab86 100644 --- a/lib/galaxy/agents/gtn_training.py +++ b/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 diff --git a/lib/galaxy/agents/history.py b/lib/galaxy/agents/history.py index b500d90acfa..9ca0d9bdd0d 100644 --- a/lib/galaxy/agents/history.py +++ b/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 diff --git a/lib/galaxy/agents/orchestrator.py b/lib/galaxy/agents/orchestrator.py index f71658a8f9f..2a7b550c34a 100644 --- a/lib/galaxy/agents/orchestrator.py +++ b/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 diff --git a/lib/galaxy/agents/page_assistant.py b/lib/galaxy/agents/page_assistant.py index 375039528b7..567257250d1 100644 --- a/lib/galaxy/agents/page_assistant.py +++ b/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 diff --git a/lib/galaxy/agents/router.py b/lib/galaxy/agents/router.py index 245337d38dc..981e7c453e4 100644 --- a/lib/galaxy/agents/router.py +++ b/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) @@ -578,7 +581,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)) diff --git a/lib/galaxy/agents/tools.py b/lib/galaxy/agents/tools.py index 4140893e6c4..1539dfc71a3 100644 --- a/lib/galaxy/agents/tools.py +++ b/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 diff --git a/lib/galaxy/config/sample/galaxy.yml.sample b/lib/galaxy/config/sample/galaxy.yml.sample index c94c2358253..5e470624224 100644 --- a/lib/galaxy/config/sample/galaxy.yml.sample +++ b/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,12 @@ 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 defaults to 0 since it runs its + # own reflection loop. #inference_services: null # YAML file with capability hints for agent inference models. Maps @@ -3365,3 +3370,4 @@ galaxy: # Enable beta tool formats (yaml, cwl, ...) which is a prerequisite # for user defined tools. #enable_beta_tool_formats: false + diff --git a/lib/galaxy/config/schemas/config_schema.yml b/lib/galaxy/config/schemas/config_schema.yml index f737f1072c1..9f05a1b0b32 100644 --- a/lib/galaxy/config/schemas/config_schema.yml +++ b/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,11 @@ 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 defaults to 0 since it runs its own + reflection loop. agent_model_capabilities_file: type: str diff --git a/test/unit/app/test_agents.py b/test/unit/app/test_agents.py index b808c16a301..427fb6d8a51 100644 --- a/test/unit/app/test_agents.py +++ b/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 ( @@ -158,6 +159,57 @@ 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 + + @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)