Merge pull request #22011 from jmchilton/agent_plumbing

Refactor Agent Registry Wiring and Fix Config Options
This commit is contained in:
Marius van den Beek
2026-03-19 10:19:27 +01:00
committed by GitHub
13 changed files with 170 additions and 165 deletions
+5 -11
View File
@@ -13,7 +13,10 @@ from .base import (
from .custom_tool import CustomToolAgent
from .error_analysis import ErrorAnalysisAgent
from .orchestrator import WorkflowOrchestratorAgent
from .registry import AgentRegistry
from .registry import (
AgentRegistry,
build_default_registry,
)
from .router import QueryRouterAgent
from .tools import ToolRecommendationAgent
@@ -22,19 +25,10 @@ __all__ = [
"BaseGalaxyAgent",
"GalaxyAgentDependencies",
"AgentRegistry",
"build_default_registry",
"QueryRouterAgent",
"ErrorAnalysisAgent",
"CustomToolAgent",
"WorkflowOrchestratorAgent",
"ToolRecommendationAgent",
]
# Global agent registry instance
agent_registry = AgentRegistry()
# Register default agents
agent_registry.register(AgentType.ROUTER, QueryRouterAgent)
agent_registry.register(AgentType.ERROR_ANALYSIS, ErrorAnalysisAgent)
agent_registry.register(AgentType.CUSTOM_TOOL, CustomToolAgent)
agent_registry.register(AgentType.ORCHESTRATOR, WorkflowOrchestratorAgent)
agent_registry.register(AgentType.TOOL_RECOMMENDATION, ToolRecommendationAgent)
+2 -6
View File
@@ -222,13 +222,13 @@ class GalaxyAgentDependencies:
trans: ProvidesUserContext
user: User
config: "GalaxyAppConfiguration"
# Callable to get agent instances, avoids circular import in base.py
get_agent: Callable[[str, "GalaxyAgentDependencies"], "BaseGalaxyAgent"]
job_manager: Optional["JobManager"] = None
dataset_manager: Optional["DatasetManager"] = None
workflow_manager: Optional["WorkflowsManager"] = None
tool_cache: Optional["ToolCache"] = None
toolbox: Optional["ToolBox"] = None
# Callable to get agent instances, avoids circular import in base.py
get_agent: Optional[Callable[[str, "GalaxyAgentDependencies"], "BaseGalaxyAgent"]] = None
# Optional factory for creating model instances (useful for testing)
model_factory: Optional[Callable[[], Any]] = None
@@ -720,10 +720,6 @@ class BaseGalaxyAgent(ABC):
)
"""
try:
if ctx.deps.get_agent is None:
raise RuntimeError("get_agent not configured in dependencies")
# Get the target agent using the injected callable
target_agent = ctx.deps.get_agent(agent_type, ctx.deps)
# Prepare query with context if available
+2 -6
View File
@@ -191,15 +191,13 @@ class WorkflowOrchestratorAgent(BaseGalaxyAgent):
self, agents: list[str], query: str, context: Optional[dict[str, Any]] = None
) -> dict[str, AgentResponse]:
"""Execute agents sequentially with timeout protection."""
from galaxy.agents import agent_registry
responses = {}
current_query = query
timeout = self._get_agent_timeout()
for agent_name in agents:
try:
agent = agent_registry.get_agent(agent_name, self.deps)
agent = self.deps.get_agent(agent_name, self.deps)
# Execute with timeout protection
response = await asyncio.wait_for(agent.process(current_query, context or {}), timeout=timeout)
responses[agent_name] = response
@@ -225,13 +223,11 @@ class WorkflowOrchestratorAgent(BaseGalaxyAgent):
self, agents: list[str], query: str, context: Optional[dict[str, Any]] = None
) -> dict[str, AgentResponse]:
"""Execute agents in parallel with timeout protection."""
from galaxy.agents import agent_registry
timeout = self._get_agent_timeout()
async def call_agent(agent_name: str):
try:
agent = agent_registry.get_agent(agent_name, self.deps)
agent = self.deps.get_agent(agent_name, self.deps)
# Execute with timeout protection
response = await asyncio.wait_for(agent.process(query, context or {}), timeout=timeout)
return agent_name, response
+37 -11
View File
@@ -121,20 +121,46 @@ class AgentRegistry:
return [self.get_agent_info(agent_type) for agent_type in self._agents.keys()]
# Global registry instance
_global_registry = AgentRegistry()
def build_default_registry(config=None) -> AgentRegistry:
"""Create an AgentRegistry with all default Galaxy agents.
Args:
config: Optional app config. When provided, agents with
``enabled: false`` in ``inference_services`` are skipped.
The router agent is always registered regardless of config.
"""
from .base import AgentType
from .custom_tool import CustomToolAgent
from .error_analysis import ErrorAnalysisAgent
from .orchestrator import WorkflowOrchestratorAgent
from .router import QueryRouterAgent
from .tools import ToolRecommendationAgent
def get_global_registry() -> AgentRegistry:
"""Get the global agent registry instance."""
return _global_registry
inference_config: dict = {}
if config is not None:
inference_config = getattr(config, "inference_services", {}) or {}
def _is_enabled(agent_type: str) -> bool:
agent_cfg = inference_config.get(agent_type, {})
if isinstance(agent_cfg, dict):
return agent_cfg.get("enabled", True)
return True
def register_agent(agent_type: str, agent_class: type[BaseGalaxyAgent], metadata: Optional[dict] = None):
"""Register an agent in the global registry."""
_global_registry.register(agent_type, agent_class, metadata)
def _register_if_enabled(registry: AgentRegistry, agent_type: str, agent_class: type[BaseGalaxyAgent]):
if _is_enabled(agent_type):
registry.register(agent_type, agent_class)
else:
log.info(f"Agent '{agent_type}' disabled by configuration, skipping registration")
registry = AgentRegistry()
def get_agent(agent_type: str, deps: GalaxyAgentDependencies) -> BaseGalaxyAgent:
"""Create an agent from the global registry."""
return _global_registry.get_agent(agent_type, deps)
# Router is always registered
if not _is_enabled(AgentType.ROUTER):
log.warning("Router agent cannot be disabled — ignoring enabled: false")
registry.register(AgentType.ROUTER, QueryRouterAgent)
_register_if_enabled(registry, AgentType.ERROR_ANALYSIS, ErrorAnalysisAgent)
_register_if_enabled(registry, AgentType.CUSTOM_TOOL, CustomToolAgent)
_register_if_enabled(registry, AgentType.ORCHESTRATOR, WorkflowOrchestratorAgent)
_register_if_enabled(registry, AgentType.TOOL_RECOMMENDATION, ToolRecommendationAgent)
return registry
+15 -1
View File
@@ -25,6 +25,10 @@ from galaxy import (
jobs,
tools,
)
from galaxy.agents.registry import (
AgentRegistry,
build_default_registry,
)
from galaxy.carbon_emissions import get_carbon_intensity_entry
from galaxy.celery.base_task import (
GalaxyTaskBeforeStart,
@@ -44,6 +48,7 @@ from galaxy.files.plugins import FileSourcePluginLoader
from galaxy.files.templates import ConfiguredFileSourceTemplates
from galaxy.job_metrics import JobMetrics
from galaxy.jobs.manager import JobManager
from galaxy.managers.agents import AgentService
from galaxy.managers.api_keys import ApiKeyManager
from galaxy.managers.citations import CitationsManager
from galaxy.managers.collections import DatasetCollectionManager
@@ -57,7 +62,10 @@ from galaxy.managers.folders import FolderManager
from galaxy.managers.hdas import HDAManager
from galaxy.managers.histories import HistoryManager
from galaxy.managers.interactivetool import InteractiveToolManager
from galaxy.managers.jobs import JobSearch
from galaxy.managers.jobs import (
JobManager as JobQueryManager,
JobSearch,
)
from galaxy.managers.libraries import LibraryManager
from galaxy.managers.library_datasets import LibraryDatasetsManager
from galaxy.managers.notification import NotificationManager
@@ -641,6 +649,7 @@ class GalaxyManagerApplication(MinimalManagerApp, MinimalGalaxyApplication):
self.library_datasets_manager = self._register_singleton(LibraryDatasetsManager)
self.role_manager = self._register_singleton(RoleManager)
self.job_manager = self._register_singleton(JobManager)
self.notification_manager = self._register_singleton(NotificationManager)
self.interactivetool_manager = InteractiveToolManager(self)
@@ -765,6 +774,11 @@ class UniverseApplication(StructuredApp, GalaxyManagerApplication, InstallationT
# want to and we'll allow postfork to bind and start it.
self.queue_worker = self._register_singleton(GalaxyQueueWorker, GalaxyQueueWorker(self))
# AI agent registry and service
agent_registry = build_default_registry(self.config)
self._register_singleton(AgentRegistry, agent_registry)
self._register_singleton(AgentService, AgentService(self.config, JobQueryManager(self), agent_registry))
self.dependency_resolvers_view = self._register_singleton(
DependencyResolversView, DependencyResolversView(self)
)
-45
View File
@@ -1141,7 +1141,6 @@ class GalaxyAppConfiguration(BaseAppConfiguration, CommonConfigurationMixin):
# load in the chat_prompts if AI is configured (old key & base URL, or inference_services)
if self.ai_api_key or self.ai_api_base_url or getattr(self, "inference_services", None):
self._load_chat_prompts()
self._load_agent_config()
self.pretty_datetime_format = expand_pretty_datetime_format(self.pretty_datetime_format)
try:
@@ -1331,50 +1330,6 @@ class GalaxyAppConfiguration(BaseAppConfiguration, CommonConfigurationMixin):
else:
log.warning(f"Chat prompts file not found at {chat_prompts_path}")
def _load_agent_config(self):
"""Load agent configuration with defaults."""
# Set default agent configuration if not present
if not hasattr(self, "agents"):
self.agents: dict = {}
# Default agent configurations
default_agents = {
"router": {
"enabled": True,
"model": "openai:gpt-4",
"temperature": 0.3,
"max_tokens": 1000,
},
"error_analysis": {
"enabled": True,
"model": "openai:gpt-4",
"temperature": 0.2,
"max_tokens": 2000,
},
"dataset_analyzer": {
"enabled": False, # Beta feature
"model": "openai:gpt-4",
"temperature": 0.3,
"max_tokens": 1500,
},
"custom_tool": {
"enabled": True,
"model": "openai:gpt-4",
"temperature": 0.4,
"max_tokens": 2000,
},
}
# Merge with any existing config
for agent_type, default_config in default_agents.items():
if agent_type not in self.agents:
self.agents[agent_type] = default_config
else:
# Fill in missing keys with defaults
for key, value in default_config.items():
if key not in self.agents[agent_type]:
self.agents[agent_type][key] = value
def _process_celery_config(self):
if self.celery_conf and self.celery_conf.get("result_backend") is None:
# If the result_backend is not set, use a SQLite database in the data directory
+6 -4
View File
@@ -2993,10 +2993,12 @@ galaxy:
#ai_model: gpt-4o
# Configuration for AI inference services used by agents. Supports
# per-agent model, temperature, and token settings. Agents inherit
# from 'default' configuration, which itself falls back to global
# ai_model/ai_api_key settings. Example: inference_services: {
# default: { model: gpt-4o-mini, temperature: 0.7 } }
# per-agent model, temperature, max_tokens, api_key, api_base_url,
# and enabled settings. Agents inherit from 'default' configuration,
# which itself falls back to global ai_model/ai_api_key settings. All
# agents are enabled by default. Example: inference_services: {
# default: { model: gpt-4o-mini, temperature: 0.7 }, custom_tool: {
# enabled: false } }
#inference_services: null
# Allow the display of tool recommendations in workflow editor and
+3 -2
View File
@@ -4086,10 +4086,11 @@ mapping:
required: false
desc: |
Configuration for AI inference services used by agents and visualization plugins.
Supports per-agent or per-plugin model, temperature, and token settings.
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, 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.
Example: inference_services: { default: { model: gpt-4o-mini }, jupyterlite: { model: gpt-4o } }
All agents are enabled by default.
Example: inference_services: { default: { model: gpt-4o-mini, temperature: 0.7 }, custom_tool: { enabled: false }, jupyterlite: { model: gpt-4o } }
enable_tool_recommendations:
type: bool
-5
View File
@@ -306,11 +306,6 @@ class ConditionalDependencies:
def check_openai(self):
return self.config.get("openai_api_key", None) is not None
def check_pydantic_ai(self):
return (
self.config.get("ai_api_key", None) is not None or self.config.get("inference_services", None) is not None
)
def check_weasyprint(self):
# See notes in ./conditional-requirements.txt for more information.
return os.environ.get("GALAXY_DEPENDENCIES_INSTALL_WEASYPRINT") == "1"
+13 -23
View File
@@ -6,30 +6,15 @@ from typing import (
Optional,
)
from galaxy.agents import GalaxyAgentDependencies
from galaxy.agents.registry import AgentRegistry
from galaxy.agents.router import QueryRouterAgent
from galaxy.config import GalaxyAppConfiguration
from galaxy.exceptions import ConfigurationError
from galaxy.managers.context import ProvidesUserContext
from galaxy.managers.jobs import JobManager
from galaxy.model import User
from galaxy.schema.agents import AgentResponse
# Import agent system (pydantic_ai is optional)
try:
from galaxy.agents import (
agent_registry,
GalaxyAgentDependencies,
)
from galaxy.agents.error_analysis import ErrorAnalysisAgent
from galaxy.agents.router import QueryRouterAgent
HAS_AGENTS = True
except ImportError:
HAS_AGENTS = False
agent_registry = None # type: ignore[assignment,misc,unused-ignore]
GalaxyAgentDependencies = None # type: ignore[assignment,misc,unused-ignore]
QueryRouterAgent = None # type: ignore[assignment,misc,unused-ignore]
ErrorAnalysisAgent = None # type: ignore[assignment,misc,unused-ignore]
log = logging.getLogger(__name__)
@@ -40,12 +25,11 @@ class AgentService:
self,
config: GalaxyAppConfiguration,
job_manager: JobManager,
registry: AgentRegistry,
):
if not HAS_AGENTS:
raise ConfigurationError("Agent system is not available")
self.config = config
self.job_manager = job_manager
self.registry = registry
def create_dependencies(self, trans: ProvidesUserContext, user: User) -> GalaxyAgentDependencies:
"""Create agent dependencies for dependency injection."""
@@ -56,7 +40,7 @@ class AgentService:
config=self.config,
job_manager=self.job_manager,
toolbox=toolbox,
get_agent=agent_registry.get_agent,
get_agent=self.registry.get_agent,
)
async def execute_agent(
@@ -75,7 +59,7 @@ class AgentService:
try:
log.info(f"Executing {agent_type} agent for query: '{query[:100]}...'")
agent = agent_registry.get_agent(agent_type, deps)
agent = self.registry.get_agent(agent_type, deps)
response = await agent.process(query, context)
return AgentResponse(
@@ -134,3 +118,9 @@ class AgentService:
# Explicit agent request - execute directly
log.info(f"User explicitly requested agent: {agent_type}")
return await self.execute_agent(agent_type, query, trans, user, context)
def list_agents(self) -> list[str]:
return self.registry.list_agents()
def get_agent_info(self, agent_type: str) -> dict:
return self.registry.get_agent_info(agent_type)
+16 -22
View File
@@ -30,15 +30,6 @@ from galaxy.webapps.galaxy.api import (
Router,
)
# Import agent system
try:
from galaxy.agents import agent_registry
HAS_AGENTS = True
except ImportError:
HAS_AGENTS = False
agent_registry = None # type: ignore[assignment,unused-ignore]
log = logging.getLogger(__name__)
router = Router(tags=["ai"])
@@ -62,26 +53,32 @@ class AgentAPI:
user: User = DependsOnUser,
) -> AgentListResponse:
"""List available AI agents."""
if not HAS_AGENTS:
raise ConfigurationError("Agent system is not available")
config = trans.app.config
inference_config = getattr(config, "inference_services", {}) or {}
agents = []
for agent_type in agent_registry.list_agents():
agent_info = agent_registry.get_agent_info(agent_type)
for agent_type in self.agent_service.list_agents():
agent_info = self.agent_service.get_agent_info(agent_type)
agent_config = inference_config.get(agent_type, {})
# Check if agent is enabled in config
agent_config = getattr(config, "agents", {}).get(agent_type, {})
enabled = agent_config.get("enabled", True)
# Resolve model: agent-specific -> default -> global
model = None
if isinstance(agent_config, dict):
model = agent_config.get("model")
if model is None:
default_config = inference_config.get("default", {})
if isinstance(default_config, dict):
model = default_config.get("model")
if model is None:
model = getattr(config, "ai_model", None)
agents.append(
AvailableAgent(
agent_type=agent_type,
name=agent_info["class_name"].replace("Agent", "").replace("_", " ").title(),
description=agent_info.get("description", "No description available"),
enabled=enabled,
model=agent_config.get("model"),
enabled=True,
model=model,
specialties=self._get_agent_specialties(agent_type),
)
)
@@ -102,9 +99,6 @@ class AgentAPI:
"""
log.warning("DEPRECATED: /api/ai/agents/query is deprecated. Use /api/chat instead.")
if not HAS_AGENTS:
raise ConfigurationError("Agent system is not available")
start_time = time.time()
try:
+6 -5
View File
@@ -34,11 +34,9 @@ from unittest.mock import (
patch,
)
from galaxy.agents import (
agent_registry,
GalaxyAgentDependencies,
)
from galaxy.agents import GalaxyAgentDependencies
from galaxy.agents.error_analysis import ErrorAnalysisResult
from galaxy.agents.registry import build_default_registry
from galaxy.tool_util_models import UserToolSource
from galaxy.util.unittest_utils import pytestmark_live_llm
from galaxy_test.base.populators import (
@@ -72,6 +70,9 @@ class AgentIntegrationTestCase(IntegrationTestCase):
config["ai_model"] = ai_model
_registry = build_default_registry()
def _create_deps_with_mock_model(self, trans, user):
"""Replacement for AgentService.create_dependencies that injects a mock model_factory."""
toolbox = trans.app.toolbox if hasattr(trans, "app") and hasattr(trans.app, "toolbox") else None
@@ -81,7 +82,7 @@ def _create_deps_with_mock_model(self, trans, user):
config=self.config,
job_manager=self.job_manager,
toolbox=toolbox,
get_agent=agent_registry.get_agent,
get_agent=_registry.get_agent,
model_factory=lambda: MagicMock(),
)
+65 -24
View File
@@ -33,12 +33,14 @@ from pydantic_ai import Agent
from pydantic_ai.models.test import TestModel
from galaxy.agents import (
agent_registry,
CustomToolAgent,
ErrorAnalysisAgent,
GalaxyAgentDependencies,
QueryRouterAgent,
)
from galaxy.agents.registry import build_default_registry
agent_registry = build_default_registry()
from galaxy.agents.error_analysis import ErrorAnalysisResult
from galaxy.agents.orchestrator import (
AgentPlan,
@@ -71,6 +73,7 @@ class TestAgentUnitMocked:
trans=self.mock_trans,
user=self.mock_user,
config=self.mock_config,
get_agent=agent_registry.get_agent,
job_manager=None,
)
@@ -160,6 +163,49 @@ class TestAgentUnitMocked:
assert "structured output" in response.content.lower()
assert response.confidence.value == "low"
def test_build_default_registry(self):
"""Test that build_default_registry creates a fully populated registry."""
registry = build_default_registry()
assert registry.is_registered("router")
assert registry.is_registered("error_analysis")
assert registry.is_registered("custom_tool")
assert registry.is_registered("orchestrator")
assert registry.is_registered("tool_recommendation")
assert len(registry.list_agents()) == 5
def test_disabled_agent_not_registered(self):
"""Disabled agent should not be in registry."""
config = mock.Mock()
config.inference_services = {
"custom_tool": {"enabled": False},
"error_analysis": {"enabled": True},
}
registry = build_default_registry(config)
assert not registry.is_registered("custom_tool")
assert registry.is_registered("error_analysis")
# Router always registered even if disabled in config
assert registry.is_registered("router")
def test_router_always_registered(self):
"""Router should always be registered even if config says disabled."""
config = mock.Mock()
config.inference_services = {"router": {"enabled": False}}
registry = build_default_registry(config)
assert registry.is_registered("router")
def test_build_registry_no_config_registers_all(self):
"""Without config, all agents registered (backwards compat)."""
registry = build_default_registry()
assert len(registry.list_agents()) == 5
def test_disabled_agent_registry_get_agent_raises(self):
"""Registry.get_agent for a disabled agent gives 'Unknown agent type' error."""
config = mock.Mock()
config.inference_services = {"custom_tool": {"enabled": False}}
registry = build_default_registry(config)
with pytest.raises(ValueError, match="Unknown agent type"):
registry.get_agent("custom_tool", self.deps)
def test_agent_registry(self):
"""Test that all required agents are registered."""
required_agents = [
@@ -277,21 +323,20 @@ class TestAgentUnitMocked:
reasoning="Single error analysis needed",
)
# Mock the actual agent call to avoid running it
with patch("galaxy.agents.agent_registry.get_agent") as mock_get_agent:
mock_error_agent = AsyncMock()
mock_error_agent.process.return_value = MagicMock(
content="The job failed due to memory limits.",
agent_type="error_analysis",
)
mock_get_agent.return_value = mock_error_agent
# Mock the deps.get_agent callback to avoid running real agents
mock_error_agent = AsyncMock()
mock_error_agent.process.return_value = MagicMock(
content="The job failed due to memory limits.",
agent_type="error_analysis",
)
self.deps.get_agent = MagicMock(return_value=mock_error_agent)
response = await agent.process("Why did my job fail?")
response = await agent.process("Why did my job fail?")
# Should not orchestrate, just return single agent response
assert response.agent_type == "orchestrator"
assert response.metadata.get("agents_used") == ["error_analysis"]
assert "memory limits" in response.content
# Should not orchestrate, just return single agent response
assert response.agent_type == "orchestrator"
assert response.metadata.get("agents_used") == ["error_analysis"]
assert "memory limits" in response.content
@pytest.mark.asyncio
async def test_workflow_orchestrator_sequential_execution(self):
@@ -306,10 +351,7 @@ class TestAgentUnitMocked:
)
# Mock each agent call in the sequential workflow
with (
patch.object(agent, "_get_agent_plan") as mock_get_plan,
patch("galaxy.agents.agent_registry.get_agent") as mock_get_agent,
):
with patch.object(agent, "_get_agent_plan") as mock_get_plan:
mock_get_plan.return_value = complex_plan
# Mock individual agent responses
@@ -332,7 +374,7 @@ class TestAgentUnitMocked:
else:
raise ValueError(f"Unexpected agent type: {agent_type}")
mock_get_agent.side_effect = get_agent_side_effect
self.deps.get_agent = MagicMock(side_effect=get_agent_side_effect)
response = await agent.process("My tool failed with memory error, help me create a fixed version")
@@ -358,10 +400,7 @@ class TestAgentUnitMocked:
reasoning="Independent tasks can run in parallel",
)
with (
patch.object(agent, "_get_agent_plan") as mock_get_plan,
patch("galaxy.agents.agent_registry.get_agent") as mock_get_agent,
):
with patch.object(agent, "_get_agent_plan") as mock_get_plan:
mock_get_plan.return_value = parallel_plan
# Mock agent responses
@@ -383,7 +422,7 @@ class TestAgentUnitMocked:
else:
raise ValueError(f"Unexpected agent type: {agent_type}")
mock_get_agent.side_effect = get_agent_side_effect
self.deps.get_agent = MagicMock(side_effect=get_agent_side_effect)
response = await agent.process("Help with my error and create a custom tool")
@@ -438,6 +477,7 @@ class TestAgentUnitLiveLLM:
trans=self.mock_trans,
user=self.mock_user,
config=self.mock_config,
get_agent=agent_registry.get_agent,
job_manager=None,
)
@@ -534,6 +574,7 @@ class TestAgentConsistencyLiveLLM:
trans=mock_trans,
user=mock_user,
config=mock_config,
get_agent=agent_registry.get_agent,
job_manager=None,
)