mirror of
https://github.com/galaxyproject/galaxy.git
synced 2026-10-31 02:28:09 +08:00
Merge pull request #22011 from jmchilton/agent_plumbing
Refactor Agent Registry Wiring and Fix Config Options
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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)
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(),
|
||||
)
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user