diff --git a/lib/galaxy/agents/__init__.py b/lib/galaxy/agents/__init__.py index 83ef2267711..a27ddb7bdf3 100644 --- a/lib/galaxy/agents/__init__.py +++ b/lib/galaxy/agents/__init__.py @@ -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) diff --git a/lib/galaxy/agents/base.py b/lib/galaxy/agents/base.py index 4324ca5965a..38bdb5832d6 100644 --- a/lib/galaxy/agents/base.py +++ b/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 @@ 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 diff --git a/lib/galaxy/agents/orchestrator.py b/lib/galaxy/agents/orchestrator.py index 3416bef2ecb..1891e675a87 100644 --- a/lib/galaxy/agents/orchestrator.py +++ b/lib/galaxy/agents/orchestrator.py @@ -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 diff --git a/lib/galaxy/agents/registry.py b/lib/galaxy/agents/registry.py index ea165c9b413..12644f14bfe 100644 --- a/lib/galaxy/agents/registry.py +++ b/lib/galaxy/agents/registry.py @@ -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 diff --git a/lib/galaxy/app.py b/lib/galaxy/app.py index 741f25c8294..a31bf45b048 100644 --- a/lib/galaxy/app.py +++ b/lib/galaxy/app.py @@ -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) ) diff --git a/lib/galaxy/config/__init__.py b/lib/galaxy/config/__init__.py index 10dccc0cab3..88269e13962 100644 --- a/lib/galaxy/config/__init__.py +++ b/lib/galaxy/config/__init__.py @@ -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 diff --git a/lib/galaxy/config/sample/galaxy.yml.sample b/lib/galaxy/config/sample/galaxy.yml.sample index 8cfa0d0cfae..c0cf9e47245 100644 --- a/lib/galaxy/config/sample/galaxy.yml.sample +++ b/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 diff --git a/lib/galaxy/config/schemas/config_schema.yml b/lib/galaxy/config/schemas/config_schema.yml index d260b8de9e4..f0a4ac4d1dd 100644 --- a/lib/galaxy/config/schemas/config_schema.yml +++ b/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 diff --git a/lib/galaxy/dependencies/__init__.py b/lib/galaxy/dependencies/__init__.py index a870f9bc243..bca7f20ff8d 100644 --- a/lib/galaxy/dependencies/__init__.py +++ b/lib/galaxy/dependencies/__init__.py @@ -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" diff --git a/lib/galaxy/managers/agents.py b/lib/galaxy/managers/agents.py index 632d3c58fd6..4a09fa656e5 100644 --- a/lib/galaxy/managers/agents.py +++ b/lib/galaxy/managers/agents.py @@ -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) diff --git a/lib/galaxy/webapps/galaxy/api/agents.py b/lib/galaxy/webapps/galaxy/api/agents.py index 9a8974d9ccc..ea9c842b964 100644 --- a/lib/galaxy/webapps/galaxy/api/agents.py +++ b/lib/galaxy/webapps/galaxy/api/agents.py @@ -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: diff --git a/test/integration/test_agents.py b/test/integration/test_agents.py index 2868095ce51..3868bc09337 100644 --- a/test/integration/test_agents.py +++ b/test/integration/test_agents.py @@ -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(), ) diff --git a/test/unit/app/test_agents.py b/test/unit/app/test_agents.py index f3e63a45c45..8bcb396005b 100644 --- a/test/unit/app/test_agents.py +++ b/test/unit/app/test_agents.py @@ -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, )