Fix/agent engine custom basellm validation

src/google/adk/cli/cli_deploy.py

@@ -14,17 +14,23 @@
 from __future__ import annotations
 
 from datetime import datetime
+import importlib.util
 import json
+import logging
 import os
 import shutil
 import subprocess
+import sys
+import traceback
 from typing import Final
 from typing import Optional
 import warnings
 
 import click
 from packaging.version import parse
 
+logger = logging.getLogger('google_adk.' + __name__)
+
 _IS_WINDOWS = os.name == 'nt'
 _GCLOUD_CMD = 'gcloud.cmd' if _IS_WINDOWS else 'gcloud'
 _LOCAL_STORAGE_FLAG_MIN_VERSION: Final[str] = '1.21.0'
@@ -465,6 +471,120 @@ 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()
+  try:
+    # Add parent directory to path so imports work correctly
+    if parent_dir not in sys.path:
+      sys.path.insert(0, parent_dir)
+
+    # Load the agent module spec
+    spec = importlib.util.spec_from_file_location(
+        f'{module_name}.agent',
+        agent_module_path,
+        submodule_search_locations=[agent_src_path],
+    )
+    if spec is None or spec.loader is None:
+      raise click.ClickException(
+          f'Failed to load module spec from {agent_module_path}'
+      )
+
+    # Try to load the module
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[f'{module_name}.agent'] = module
+
+    try:
+      spec.loader.exec_module(module)
+    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 implementation is serializable'
+        ) 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}'
+        ) 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}'
+      ) 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 the module from sys.modules
+    sys.modules.pop(f'{module_name}.agent', None)
+
+
 def _get_service_option_by_adk_version(
     adk_version: str,
     session_uri: Optional[str],
@@ -953,6 +1073,10 @@ 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
+    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'

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")
@@ -407,3 +408,131 @@ def mock_subprocess_run(*args, **kwargs):
 
   # 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_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 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 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