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'
@@ -99,15 +105,33 @@ def _ensure_agent_engine_dependency(requirements_txt_path: str) -> None:
 """
 
 _AGENT_ENGINE_APP_TEMPLATE: Final[str] = """
+import logging
 import os
+import sys
+import cloudpickle
 import vertexai
 from vertexai.agent_engines import AdkApp
 
+_logger = logging.getLogger("google_adk." + __name__)
+
 if {is_config_agent}:
   from google.adk.agents import config_agent_utils
   root_agent = config_agent_utils.from_config("{agent_folder}/root_agent.yaml")
 else:
   from .agent import {adk_app_object}
+  # Register the agent module for pickle-by-value serialization.
+  # This ensures custom BaseLlm implementations are serialized with their
+  # full class definition instead of just the import path, which fixes
+  # the "query method not found" error when using custom LLM clients.
+  from . import agent as _agent_module
+  cloudpickle.register_pickle_by_value(_agent_module)
+  # Also register any submodules that contain custom classes
+  for name, module in list(sys.modules.items()):
+    if module is not None and name.startswith(_agent_module.__name__.rsplit('.', 1)[0] + '.'):
+      try:
+        cloudpickle.register_pickle_by_value(module)
+      except Exception as e:
+        _logger.debug("Failed to register module %s for pickle-by-value: %s", name, e)
 
 if {express_mode}: # Whether or not to use Express Mode
   vertexai.init(api_key=os.environ.get("GOOGLE_API_KEY"))
@@ -464,6 +488,218 @@ def _validate_gcloud_extra_args(
       )
 
 
+def _validate_agent_object(exported_obj: object, adk_app_object: str) -> None:
+  """Validates that the exported agent/app object is properly configured.
+
+  This function performs deeper validation beyond just checking that the
+  object exists. It verifies that agents with custom BaseLlm implementations
+  have properly configured models that will work at Agent Engine runtime.
+
+  Args:
+    exported_obj: The exported root_agent or app object.
+    adk_app_object: The name of the exported object ('root_agent' or 'app').
+
+  Raises:
+    click.ClickException: If the agent object is not properly configured.
+  """
+  # Import here to avoid circular imports
+  try:
+    from google.adk.agents import BaseAgent
+    from google.adk.models import BaseLlm
+  except ImportError as e:
+    # If we can't import these, skip validation. This can happen in partial
+    # ADK environments or when optional dependencies are not installed.
+    logger.debug(
+        'Skipping agent object validation: could not import ADK classes: %s', e
+    )
+    return
+
+  # For 'app' exports (AdkApp instances), we don't validate the internal agent
+  if adk_app_object == 'app':
+    return
+
+  # For 'root_agent' exports, validate it's an agent with a valid model
+  if not isinstance(exported_obj, BaseAgent):
+    click.secho(
+        f'Warning: {adk_app_object} is not a BaseAgent instance. '
+        'Skipping model validation.',
+        fg='yellow',
+    )
+    return
+
+  # Check if the agent has a model attribute
+  model = getattr(exported_obj, 'model', None)
+  if model is None:
+    # Some agents might not have a model (e.g., workflow agents)
+    return
+
+  # If the model is a string, it will be resolved by LLMRegistry at runtime
+  if isinstance(model, str):
+    return
+
+  # If the model is a BaseLlm instance, validate it
+  if isinstance(model, BaseLlm):
+    model_class = type(model)
+    model_module = model_class.__module__
+
+    # Check if this is a custom BaseLlm (not from google.adk.models)
+    if not model_module.startswith('google.adk.models'):
+      click.echo(
+          f'Detected custom BaseLlm implementation: {model_class.__name__} '
+          f'from {model_module}'
+      )
+
+      # Validate that the custom model can be pickled (required for Agent Engine)
+      try:
+        import cloudpickle
+
+        cloudpickle.dumps(model)
+        click.echo(
+            f'Custom model {model_class.__name__} passed serialization check'
+        )
+      except Exception as e:
+        raise click.ClickException(
+            f'Custom BaseLlm implementation {model_class.__name__} cannot be '
+            f'serialized:\n{e}\n\n'
+            'Agent Engine requires all custom LLM implementations to be '
+            'serializable. Please ensure:\n'
+            '1. Your custom BaseLlm class does not have non-serializable '
+            'attributes (file handles, connections, etc.)\n'
+            "2. All fields are JSON-serializable or use Pydantic's "
+            'ConfigDict(arbitrary_types_allowed=True)\n'
+            '3. Consider implementing __getstate__ and __setstate__ methods '
+            'for custom serialization logic'
+        ) from e
+
+      # Check if the model class is defined in a file within the agent folder
+      # by verifying the module can be imported with relative imports
+      if '.' not in model_module or model_module.count('.') < 2:
+        click.secho(
+            f'Warning: Custom model {model_class.__name__} is defined in '
+            f'{model_module}. For Agent Engine deployment, ensure this module '
+            'is within your agent folder and uses relative imports.',
+            fg='yellow',
+        )
+
+
+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".'
+      )
+
+    # Validate that the exported object is properly configured for Agent Engine
+    exported_obj = getattr(module, adk_app_object)
+    _validate_agent_object(exported_obj, adk_app_object)
+
+    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],
@@ -952,6 +1188,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'