Merge pull request galaxyproject#22011 from jmchilton/agent_plumbing

Refactor Agent Registry Wiring and Fix Config Options

Author: mvdbeek

lib/galaxy/agents/__init__.py

@@ -13,7 +13,10 @@
 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 @@
     "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)

lib/galaxy/agents/base.py

@@ -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 @@ async def _call_agent_from_tool(
             )
         """
         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

lib/galaxy/agents/orchestrator.py

@@ -191,15 +191,13 @@ async def _execute_sequential(
         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 @@ async def _execute_parallel(
         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

lib/galaxy/agents/registry.py

@@ -121,20 +121,46 @@ def list_agent_info(self) -> list[dict]:
         return [self.get_agent_info(agent_type) for agent_type in self._agents.keys()]
 
 
-# Global registry instance
-_global_registry = AgentRegistry()
-
-
-def get_global_registry() -> AgentRegistry:
-    """Get the global agent registry instance."""
-    return _global_registry
-
-
-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 get_agent(agent_type: str, deps: GalaxyAgentDependencies) -> BaseGalaxyAgent:
-    """Create an agent from the global registry."""
-    return _global_registry.get_agent(agent_type, deps)
+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
+
+    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_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()
+
+    # 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

lib/galaxy/app.py

@@ -25,6 +25,10 @@
     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.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.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 @@ def __init__(self, configure_logging=True, use_converters=True, use_display_appl
         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 @@ def __init__(self, **kwargs) -> None:
         # 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)
         )

lib/galaxy/config/__init__.py

@@ -1141,7 +1141,6 @@ def _process_config(self, kwargs: dict[str, Any]) -> None:
         # 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 @@ def _load_chat_prompts(self):
         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

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

@@ -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

lib/galaxy/config/schemas/config_schema.yml

@@ -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

lib/galaxy/dependencies/__init__.py

@@ -306,11 +306,6 @@ def check_tensorflow(self):
     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"

lib/galaxy/managers/agents.py

@@ -6,30 +6,15 @@
     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 @@ def __init__(
         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 @@ def create_dependencies(self, trans: ProvidesUserContext, user: User) -> GalaxyA
             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 @@ async def execute_agent(
 
         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 @@ async def route_and_execute(
             # 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)