diff --git a/src/google/adk/cli/cli_deploy.py b/src/google/adk/cli/cli_deploy.py index d6394ff2..dcec5186 100644 --- a/src/google/adk/cli/cli_deploy.py +++ b/src/google/adk/cli/cli_deploy.py @@ -14,10 +14,13 @@ from __future__ import annotations from datetime import datetime +import importlib import json import os import shutil import subprocess +import sys +import traceback from typing import Final from typing import Optional import warnings @@ -465,6 +468,122 @@ def _validate_gcloud_extra_args( ) +def _validate_agent_import( + agent_src_path: str, + adk_app_object: str, + is_config_agent: bool, +) -> None: + """Validates that the agent module can be imported successfully. + + This pre-deployment validation catches common issues like missing + dependencies or import errors in custom BaseLlm implementations before + the agent is deployed to Agent Engine. This provides clearer error + messages and prevents deployments that would fail at runtime. + + Args: + agent_src_path: Path to the staged agent source code. + adk_app_object: The Python object name to import ('root_agent' or 'app'). + is_config_agent: Whether this is a config-based agent. + + Raises: + click.ClickException: If the agent module cannot be imported. + """ + if is_config_agent: + # Config agents are loaded from YAML, skip Python import validation + return + + agent_module_path = os.path.join(agent_src_path, 'agent.py') + if not os.path.exists(agent_module_path): + raise click.ClickException( + f'Agent module not found at {agent_module_path}. ' + 'Please ensure your agent folder contains an agent.py file.' + ) + + # Add the parent directory to sys.path temporarily for import resolution + parent_dir = os.path.dirname(agent_src_path) + module_name = os.path.basename(agent_src_path) + + original_sys_path = sys.path.copy() + original_sys_modules_keys = set(sys.modules.keys()) + try: + # Add parent directory to path so imports work correctly + if parent_dir not in sys.path: + sys.path.insert(0, parent_dir) + try: + module = importlib.import_module(f'{module_name}.agent') + except ImportError as e: + error_msg = str(e) + tb = traceback.format_exc() + + # Check for common issues + if 'BaseLlm' in tb or 'base_llm' in tb.lower(): + raise click.ClickException( + 'Failed to import agent module due to a BaseLlm-related error:\n' + f'{error_msg}\n\n' + 'This error often occurs when deploying agents with custom LLM ' + 'implementations. Please ensure:\n' + '1. All custom LLM classes are defined in files within your agent ' + 'folder\n' + '2. All required dependencies are listed in requirements.txt\n' + '3. Import paths use relative imports (e.g., "from .my_llm import ' + 'MyLlm")\n' + '4. Your custom BaseLlm class and its dependencies are installed\n' + '\n' + 'If this failure is expected (e.g., missing local dependencies), ' + 'disable agent import validation by omitting ' + '--validate-agent-import (default) or passing ' + '--skip-agent-import-validation (or --no-validate-agent-import).' + ) from e + else: + raise click.ClickException( + f'Failed to import agent module:\n{error_msg}\n\n' + 'Please ensure all dependencies are listed in requirements.txt ' + 'and all imports are resolvable.\n\n' + f'Full traceback:\n{tb}\n\n' + 'If this failure is expected (e.g., missing local dependencies), ' + 'disable agent import validation by omitting ' + '--validate-agent-import (default) or passing ' + '--skip-agent-import-validation (or --no-validate-agent-import).' + ) from e + except Exception as e: + tb = traceback.format_exc() + raise click.ClickException( + f'Error while loading agent module:\n{e}\n\n' + 'Please check your agent code for errors.\n\n' + f'Full traceback:\n{tb}\n\n' + 'If this failure is expected (e.g., missing local dependencies), ' + 'disable agent import validation by omitting ' + '--validate-agent-import (default) or passing ' + '--skip-agent-import-validation (or --no-validate-agent-import).' + ) from e + + # Check that the expected object exists + if not hasattr(module, adk_app_object): + available_attrs = [ + attr for attr in dir(module) if not attr.startswith('_') + ] + raise click.ClickException( + f"Agent module does not export '{adk_app_object}'. " + f'Available exports: {available_attrs}\n\n' + 'Please ensure your agent.py exports either "root_agent" or "app".' + ) + + click.echo( + 'Agent module validation successful: ' + f'found "{adk_app_object}" in agent.py' + ) + + finally: + # Restore original sys.path + sys.path[:] = original_sys_path + # Clean up modules introduced by validation. + for key in list(sys.modules.keys()): + if key in original_sys_modules_keys: + continue + if key == module_name or key.startswith(f'{module_name}.'): + sys.modules.pop(key, None) + + def _get_service_option_by_adk_version( adk_version: str, session_uri: Optional[str], @@ -702,6 +821,7 @@ def to_agent_engine( requirements_file: Optional[str] = None, env_file: Optional[str] = None, agent_engine_config_file: Optional[str] = None, + skip_agent_import_validation: bool = True, ): """Deploys an agent to Vertex AI Agent Engine. @@ -761,6 +881,11 @@ def to_agent_engine( agent_engine_config_file (str): The filepath to the agent engine config file to use. If not specified, the `.agent_engine_config.json` file in the `agent_folder` will be used. + skip_agent_import_validation (bool): Optional. Default is True. If True, + skip the + pre-deployment import validation of `agent.py`. This can be useful when + the local environment does not have the same dependencies as the + deployment environment. """ app_name = os.path.basename(agent_folder) display_name = display_name or app_name @@ -953,6 +1078,11 @@ def to_agent_engine( click.echo(f'Config agent detected: {config_root_agent_file}') is_config_agent = True + # Validate that the agent module can be imported before deployment. + if not skip_agent_import_validation: + click.echo('Validating agent module...') + _validate_agent_import(agent_src_path, adk_app_object, is_config_agent) + adk_app_file = os.path.join(temp_folder, f'{adk_app}.py') if adk_app_object == 'root_agent': adk_app_type = 'agent' diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index c4f1d405..c686211e 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -1867,6 +1867,25 @@ def cli_migrate_session( " directory, if any.)" ), ) +@click.option( + "--validate-agent-import/--no-validate-agent-import", + default=False, + help=( + "Optional. Validate that the agent module can be imported before" + " deployment. This requires your local environment to have the same" + " dependencies as the deployment environment. (default: disabled)" + ), +) +@click.option( + "--skip-agent-import-validation", + "skip_agent_import_validation_alias", + is_flag=True, + default=False, + help=( + "Optional. Skip pre-deployment import validation of `agent.py`. This is" + " the default; use --validate-agent-import to enable validation." + ), +) @click.argument( "agent", type=click.Path( @@ -1891,6 +1910,8 @@ def cli_deploy_agent_engine( requirements_file: str, absolutize_imports: bool, agent_engine_config_file: str, + validate_agent_import: bool = False, + skip_agent_import_validation_alias: bool = False, ): """Deploys an agent to Agent Engine. @@ -1905,6 +1926,11 @@ def cli_deploy_agent_engine( """ logging.getLogger("vertexai_genai.agentengines").setLevel(logging.INFO) try: + if validate_agent_import and skip_agent_import_validation_alias: + raise click.UsageError( + "Do not pass both --validate-agent-import and" + " --skip-agent-import-validation." + ) cli_deploy.to_agent_engine( agent_folder=agent, project=project, @@ -1922,6 +1948,7 @@ def cli_deploy_agent_engine( requirements_file=requirements_file, absolutize_imports=absolutize_imports, agent_engine_config_file=agent_engine_config_file, + skip_agent_import_validation=not validate_agent_import, ) except Exception as e: click.secho(f"Deploy failed: {e}", fg="red", err=True) diff --git a/tests/unittests/cli/utils/test_cli_deploy.py b/tests/unittests/cli/utils/test_cli_deploy.py index 36484da1..e1829d2f 100644 --- a/tests/unittests/cli/utils/test_cli_deploy.py +++ b/tests/unittests/cli/utils/test_cli_deploy.py @@ -14,7 +14,6 @@ """Tests for utilities in cli_deploy.""" - from __future__ import annotations import importlib @@ -83,7 +82,9 @@ def agent_dir(tmp_path: Path) -> Callable[[bool, bool], Path]: def _factory(include_requirements: bool, include_env: bool) -> Path: base = tmp_path / "agent" base.mkdir() - (base / "agent.py").write_text("# dummy agent") + (base / "agent.py").write_text( + "# dummy agent\nroot_agent = 'dummy_agent'\n" + ) (base / "__init__.py").touch() if include_requirements: (base / "requirements.txt").write_text("pytest\n") @@ -305,6 +306,110 @@ def test_to_agent_engine_happy_path( assert str(rmtree_recorder.get_last_call_args()[0]) == str(tmp_dir) +def test_to_agent_engine_skips_agent_import_validation_by_default( + monkeypatch: pytest.MonkeyPatch, + agent_dir: Callable[[bool, bool], Path], +) -> None: + """It should skip agent.py import validation by default.""" + validate_recorder = _Recorder() + + def _validate_agent_import(*args: Any, **kwargs: Any) -> None: + validate_recorder(*args, **kwargs) + raise AssertionError("_validate_agent_import should not be called") + + monkeypatch.setattr( + cli_deploy, "_validate_agent_import", _validate_agent_import + ) + + fake_vertexai = types.ModuleType("vertexai") + + class _FakeAgentEngines: + + def create(self, *, config: Dict[str, Any]) -> Any: + del config + return types.SimpleNamespace( + api_resource=types.SimpleNamespace( + name="projects/p/locations/l/reasoningEngines/e" + ) + ) + + class _FakeVertexClient: + + def __init__(self, *args: Any, **kwargs: Any) -> None: + del args + del kwargs + self.agent_engines = _FakeAgentEngines() + + fake_vertexai.Client = _FakeVertexClient + monkeypatch.setitem(sys.modules, "vertexai", fake_vertexai) + + src_dir = agent_dir(False, False) + cli_deploy.to_agent_engine( + agent_folder=str(src_dir), + temp_folder="tmp", + adk_app="my_adk_app", + trace_to_cloud=True, + project="my-gcp-project", + region="us-central1", + display_name="My Test Agent", + description="A test agent.", + ) + + assert validate_recorder.calls == [] + + +def test_to_agent_engine_validates_agent_import_when_enabled( + monkeypatch: pytest.MonkeyPatch, + agent_dir: Callable[[bool, bool], Path], +) -> None: + """It should run agent.py import validation when enabled.""" + validate_recorder = _Recorder() + + def _validate_agent_import(*args: Any, **kwargs: Any) -> None: + validate_recorder(*args, **kwargs) + + monkeypatch.setattr( + cli_deploy, "_validate_agent_import", _validate_agent_import + ) + + fake_vertexai = types.ModuleType("vertexai") + + class _FakeAgentEngines: + + def create(self, *, config: Dict[str, Any]) -> Any: + del config + return types.SimpleNamespace( + api_resource=types.SimpleNamespace( + name="projects/p/locations/l/reasoningEngines/e" + ) + ) + + class _FakeVertexClient: + + def __init__(self, *args: Any, **kwargs: Any) -> None: + del args + del kwargs + self.agent_engines = _FakeAgentEngines() + + fake_vertexai.Client = _FakeVertexClient + monkeypatch.setitem(sys.modules, "vertexai", fake_vertexai) + + src_dir = agent_dir(False, False) + cli_deploy.to_agent_engine( + agent_folder=str(src_dir), + temp_folder="tmp", + adk_app="my_adk_app", + trace_to_cloud=True, + project="my-gcp-project", + region="us-central1", + display_name="My Test Agent", + description="A test agent.", + skip_agent_import_validation=False, + ) + + assert len(validate_recorder.calls) == 1 + + @pytest.mark.parametrize("include_requirements", [True, False]) def test_to_gke_happy_path( monkeypatch: pytest.MonkeyPatch, @@ -407,3 +512,145 @@ def test_to_gke_happy_path( # 4. Verify cleanup assert str(rmtree_recorder.get_last_call_args()[0]) == str(tmp_path) + + +# _validate_agent_import tests +class TestValidateAgentImport: + """Tests for the _validate_agent_import function.""" + + def test_skips_config_agents(self, tmp_path: Path) -> None: + """Config agents should skip validation.""" + # This should not raise even with no agent.py file + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=True + ) + + def test_raises_on_missing_agent_module(self, tmp_path: Path) -> None: + """Should raise when agent.py is missing.""" + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Agent module not found" in str(exc_info.value) + + def test_raises_on_missing_export(self, tmp_path: Path) -> None: + """Should raise when the expected export is missing.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("some_other_var = 'hello'\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "does not export 'root_agent'" in str(exc_info.value) + assert "some_other_var" in str(exc_info.value) + + def test_success_with_root_agent_export(self, tmp_path: Path) -> None: + """Should succeed when root_agent is exported.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + # Should not raise + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + def test_success_with_app_export(self, tmp_path: Path) -> None: + """Should succeed when app is exported.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("app = 'my_app'\n") + (tmp_path / "__init__.py").touch() + + # Should not raise + cli_deploy._validate_agent_import( + str(tmp_path), "app", is_config_agent=False + ) + + def test_success_with_relative_imports(self, tmp_path: Path) -> None: + """Should succeed when agent.py uses relative imports.""" + (tmp_path / "helper.py").write_text("VALUE = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + (tmp_path / "agent.py").write_text( + "from .helper import VALUE\n\nroot_agent = VALUE\n" + ) + + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + def test_raises_on_import_error(self, tmp_path: Path) -> None: + """Should raise with helpful message on ImportError.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("from nonexistent_module import something\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Failed to import agent module" in str(exc_info.value) + assert "nonexistent_module" in str(exc_info.value) + + def test_raises_on_basellm_import_error(self, tmp_path: Path) -> None: + """Should provide specific guidance for BaseLlm import errors.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text( + "from google.adk.models.base_llm import NonexistentBaseLlm\n" + ) + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "BaseLlm-related error" in str(exc_info.value) + assert "custom LLM" in str(exc_info.value) + + def test_raises_on_syntax_error(self, tmp_path: Path) -> None: + """Should raise on syntax errors in agent.py.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("def invalid syntax here:\n") + (tmp_path / "__init__.py").touch() + + with pytest.raises(click.ClickException) as exc_info: + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + assert "Error while loading agent module" in str(exc_info.value) + + def test_cleans_up_sys_modules(self, tmp_path: Path) -> None: + """Should clean up sys.modules after validation.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + module_name = tmp_path.name + agent_module_key = f"{module_name}.agent" + + # Ensure module is not in sys.modules before + assert module_name not in sys.modules + assert agent_module_key not in sys.modules + + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + # Ensure module is cleaned up after + assert module_name not in sys.modules + assert agent_module_key not in sys.modules + + def test_restores_sys_path(self, tmp_path: Path) -> None: + """Should restore sys.path after validation.""" + agent_file = tmp_path / "agent.py" + agent_file.write_text("root_agent = 'my_agent'\n") + (tmp_path / "__init__.py").touch() + + original_path = sys.path.copy() + + cli_deploy._validate_agent_import( + str(tmp_path), "root_agent", is_config_agent=False + ) + + assert sys.path == original_path