Merge branch 'main' into custom-session-id-with-agent-engine

Author: guillaumeblaquiere

src/google/adk/auth/auth_provider_registry.py

@@ -42,13 +42,18 @@ def register(
     """
     self._providers[auth_scheme_type] = provider_instance
 
-  def get_provider(self, auth_scheme: AuthScheme) -> BaseAuthProvider | None:
+  def get_provider(
+      self, auth_scheme: AuthScheme | type[AuthScheme]
+  ) -> BaseAuthProvider | None:
     """Get the provider instance for an auth scheme.
 
     Args:
-        auth_scheme: The auth scheme to get provider for.
+        auth_scheme: The auth scheme or the auth scheme type to get the provider
+            for.
 
     Returns:
         The provider instance if registered, None otherwise.
     """
+    if isinstance(auth_scheme, type):
+      return self._providers.get(auth_scheme)
     return self._providers.get(type(auth_scheme))

src/google/adk/auth/auth_schemes.py

@@ -27,6 +27,7 @@
 from pydantic import Field
 
 from ..utils.feature_decorator import experimental
+from .auth_credential import BaseModelWithConfig
 
 
 class OpenIdConnectWithConfig(SecurityBase):
@@ -42,8 +43,20 @@ class OpenIdConnectWithConfig(SecurityBase):
   scopes: Optional[List[str]] = None
 
 
-# AuthSchemes contains SecuritySchemes from OpenAPI 3.0 and an extra flattened OpenIdConnectWithConfig.
-AuthScheme = Union[SecurityScheme, OpenIdConnectWithConfig]
+class CustomAuthScheme(BaseModelWithConfig):
+  """A flexible model for custom authentication schemes.
+
+  The subclasses must define a `default` for the `type_` field, if using OAuth2
+  user consent flow, to ensure correct rehydration.
+  """
+
+  type_: str = Field(alias="type")
+
+
+# AuthSchemes contains SecuritySchemes from OpenAPI 3.0, an extra flattened
+# OpenIdConnectWithConfig, and supports external schemes that subclasses
+# CustomAuthScheme
+AuthScheme = Union[SecurityScheme, OpenIdConnectWithConfig, CustomAuthScheme]
 
 
 class OAuthGrantType(str, Enum):

src/google/adk/auth/auth_tool.py

@@ -106,8 +106,11 @@ def get_credential_key(self):
     if auth_scheme.model_extra:
       auth_scheme = auth_scheme.model_copy(deep=True)
       auth_scheme.model_extra.clear()
+
+    type_ = auth_scheme.type_
+    type_name = type_.name if type_ and hasattr(type_, "name") else str(type_)
     scheme_name = (
-        f"{auth_scheme.type_.name}_{_stable_model_digest(auth_scheme)}"
+        f"{type_name}_{_stable_model_digest(auth_scheme)}"
         if auth_scheme
         else ""
     )

src/google/adk/auth/base_auth_provider.py

@@ -16,6 +16,10 @@
 
 from abc import ABC
 from abc import abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+  from .auth_schemes import AuthScheme
 
 from ..agents.callback_context import CallbackContext
 from ..features import experimental
@@ -28,6 +32,15 @@
 class BaseAuthProvider(ABC):
   """Abstract base class for custom authentication providers."""
 
+  @property
+  def supported_auth_schemes(self) -> tuple[type[AuthScheme], ...]:
+    """The AuthScheme types supported by this provider.
+
+    Subclasses can override this to return a tuple of scheme types, enabling
+    1-parameter registration.
+    """
+    return ()
+
   @abstractmethod
   async def get_auth_credential(
       self, auth_config: AuthConfig, context: CallbackContext

src/google/adk/auth/credential_manager.py

@@ -14,7 +14,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 import logging
+import threading
 from typing import Optional
 
 from fastapi.openapi.models import OAuth2
@@ -25,10 +27,13 @@
 from .auth_credential import AuthCredential
 from .auth_credential import AuthCredentialTypes
 from .auth_provider_registry import AuthProviderRegistry
+from .auth_schemes import AuthScheme
 from .auth_schemes import AuthSchemeType
+from .auth_schemes import CustomAuthScheme
 from .auth_schemes import ExtendedOAuth2
 from .auth_schemes import OpenIdConnectWithConfig
 from .auth_tool import AuthConfig
+from .base_auth_provider import BaseAuthProvider
 from .exchanger.base_credential_exchanger import BaseCredentialExchanger
 from .exchanger.base_credential_exchanger import ExchangeResult
 from .exchanger.credential_exchanger_registry import CredentialExchangerRegistry
@@ -38,6 +43,25 @@
 logger = logging.getLogger("google_adk." + __name__)
 
 
+def _rehydrate_custom_scheme(
+    scheme: CustomAuthScheme, supported_schemes: Sequence[type[AuthScheme]]
+) -> CustomAuthScheme:
+  """Rehydrate a CustomAuthScheme into one of the given supported_schemes."""
+  incoming_type = scheme.type_
+  for scheme_class in supported_schemes:
+    type_field = scheme_class.model_fields.get("type_")
+    # Custom AuthScheme classes must define a `default` for their `type_` field
+    # to be rehydrated correctly.
+    if type_field and type_field.default == incoming_type:
+      data = scheme.model_dump(by_alias=True)
+      if scheme.model_extra:
+        data.update(scheme.model_extra)
+      return scheme_class.model_validate(data)
+  raise ValueError(
+      f"Cannot rehydrate: no registered scheme matches type '{incoming_type}'"
+  )
+
+
 @experimental
 class CredentialManager:
   """Manages authentication credentials through a structured workflow.
@@ -77,12 +101,32 @@ class CredentialManager:
       ```
   """
 
+  _auth_provider_registry = AuthProviderRegistry()
+  _registry_lock = threading.Lock()
+
+  @classmethod
+  def register_auth_provider(cls, provider: BaseAuthProvider) -> None:
+    """Public API for developers to register custom auth providers."""
+    with cls._registry_lock:
+      for scheme_type in provider.supported_auth_schemes:
+        existing_provider = cls._auth_provider_registry.get_provider(
+            scheme_type
+        )
+        if existing_provider is not None:
+          if existing_provider is not provider:
+            logger.warning(
+                "An auth provider is already registered for scheme %s. "
+                "Ignoring the new provider.",
+                scheme_type,
+            )
+            continue
+        cls._auth_provider_registry.register(scheme_type, provider)
+
   def __init__(
       self,
       auth_config: AuthConfig,
   ):
     self._auth_config = auth_config
-    self._auth_provider_registry = AuthProviderRegistry()
     self._exchanger_registry = CredentialExchangerRegistry()
     self._refresher_registry = CredentialRefresherRegistry()
     self._discovery_manager = OAuth2DiscoveryManager()
@@ -139,6 +183,20 @@ async def get_auth_credential(
   ) -> Optional[AuthCredential]:
     """Load and prepare authentication credential through a structured workflow."""
 
+    # Pydantic may have deserialized an unknown scheme into a generic
+    # CustomAuthScheme. If so, rehydrate it first into a specific subclass.
+    # Note: Custom authentication scheme classes must have been imported into
+    # the Python runtime before get_auth_credential is called for their
+    # subclasses to be registered. This is fine as developer will anyway import
+    # them while registering the auth providers.
+    # Note: `__subclasses__()` only returns immediate subclasses, if there is a
+    # subclass of a subclass of CustomAuthScheme then it will not be returned.
+    # pylint: disable=unidiomatic-typecheck Needs exact class matching.
+    if type(self._auth_config.auth_scheme) is CustomAuthScheme:
+      self._auth_config.auth_scheme = _rehydrate_custom_scheme(
+          self._auth_config.auth_scheme,
+          CustomAuthScheme.__subclasses__(),
+      )
     # First, check if a registered auth provider is available before attempting
     # to retrieve tokens natively.
     provider = self._auth_provider_registry.get_provider(

src/google/adk/cli/cli_tools_click.py

@@ -1064,6 +1064,7 @@ def cli_optimize(
     from .cli_eval import _collect_eval_results
     from .cli_eval import _collect_inferences
     from .cli_eval import get_root_agent
+
   except ModuleNotFoundError as mnf:
     raise click.ClickException(MISSING_EVAL_DEPENDENCIES_MESSAGE) from mnf
 
@@ -1199,6 +1200,7 @@ def cli_add_eval_case(
     from ..evaluation.eval_case import EvalCase
     from ..evaluation.eval_case import SessionInput
     from .cli_eval import get_eval_sets_manager
+
   except ModuleNotFoundError as mnf:
     raise click.ClickException(MISSING_EVAL_DEPENDENCIES_MESSAGE) from mnf
 
@@ -1247,6 +1249,127 @@ def cli_add_eval_case(
     raise click.ClickException(f"Failed to add eval case(s): {e}") from e
 
 
+@eval_set.command("generate_eval_cases", cls=HelpfulCommand)
+@click.argument(
+    "agent_module_file_path",
+    type=click.Path(
+        exists=True, dir_okay=True, file_okay=False, resolve_path=True
+    ),
+)
+@click.argument("eval_set_id", type=str, required=True)
+@click.option(
+    "--user_simulation_config_file",
+    type=click.Path(
+        exists=True, dir_okay=False, file_okay=True, resolve_path=True
+    ),
+    help=(
+        "A path to file containing JSON serialized "
+        "UserScenarioGenerationConfig dict."
+    ),
+    required=True,
+)
+@eval_options()
+def cli_generate_eval_cases(
+    agent_module_file_path: str,
+    eval_set_id: str,
+    user_simulation_config_file: str,
+    eval_storage_uri: Optional[str] = None,
+    log_level: str = "INFO",
+):
+  """Generates eval cases dynamically and adds them to the given eval set.
+
+  Uses Vertex AI Eval SDK to generate conversation scenarios based on an
+  Agent's info and definitions. It will automatically create the empty eval_set
+  if it has not been created in advance.
+
+  Args:
+    agent_module_file_path: The path to the agent module file.
+    eval_set_id: The id of the eval set to generate cases for.
+    user_simulation_config_file: The path to the user simulation config file.
+    eval_storage_uri: The eval storage uri.
+    log_level: The log level.
+  """
+  logs.setup_adk_logger(getattr(logging, log_level.upper()))
+  try:
+    from ..evaluation._vertex_ai_scenario_generation_facade import ScenarioGenerator
+    from ..evaluation.conversation_scenarios import ConversationGenerationConfig
+    from ..evaluation.eval_case import EvalCase
+    from ..evaluation.eval_case import SessionInput
+    from .cli_eval import get_eval_sets_manager
+    from .cli_eval import get_root_agent
+    from .utils.state import create_empty_state
+
+  except ModuleNotFoundError as mnf:
+    raise click.ClickException(MISSING_EVAL_DEPENDENCIES_MESSAGE) from mnf
+
+  app_name = os.path.basename(agent_module_file_path)
+  agents_dir = os.path.dirname(agent_module_file_path)
+
+  try:
+    eval_sets_manager = get_eval_sets_manager(eval_storage_uri, agents_dir)
+    root_agent = get_root_agent(agent_module_file_path)
+
+    # Try to create if it doesn't already exist.
+    if (
+        eval_sets_manager.get_eval_set(
+            app_name=app_name, eval_set_id=eval_set_id
+        )
+        is None
+    ):
+      eval_sets_manager.create_eval_set(
+          app_name=app_name, eval_set_id=eval_set_id
+      )
+      click.echo(f"Eval set '{eval_set_id}' created for app '{app_name}'.")
+    else:
+      click.echo(f"Eval set '{eval_set_id}' already exists.")
+
+    with open(user_simulation_config_file, "r") as f:
+      config = ConversationGenerationConfig.model_validate_json(f.read())
+
+    generator = ScenarioGenerator()
+    click.echo("Generating scenarios utilizing Vertex AI Eval SDK...")
+    scenarios = generator.generate_scenarios(root_agent, config)
+
+    # TODO(pthodoroff): Expose initial session state when simulation library
+    # supports it.
+    initial_session_state = create_empty_state(root_agent)
+
+    session_input = SessionInput(
+        app_name=app_name, user_id="test_user_id", state=initial_session_state
+    )
+
+    for scenario in scenarios:
+      scenario_str = json.dumps(scenario.model_dump(), sort_keys=True)
+      eval_id = hashlib.sha256(scenario_str.encode("utf-8")).hexdigest()[:8]
+      eval_case = EvalCase(
+          eval_id=eval_id,
+          conversation_scenario=scenario,
+          session_input=session_input,
+          creation_timestamp=datetime.now().timestamp(),
+      )
+
+      if (
+          eval_sets_manager.get_eval_case(
+              app_name=app_name, eval_set_id=eval_set_id, eval_case_id=eval_id
+          )
+          is None
+      ):
+        eval_sets_manager.add_eval_case(
+            app_name=app_name, eval_set_id=eval_set_id, eval_case=eval_case
+        )
+        click.echo(
+            f"Eval case '{eval_case.eval_id}' added to eval set"
+            f" '{eval_set_id}'."
+        )
+      else:
+        click.echo(
+            f"Eval case '{eval_case.eval_id}' already exists in eval set"
+            f" '{eval_set_id}', skipped adding."
+        )
+  except Exception as e:
+    raise click.ClickException(f"Failed to generate eval case(s): {e}") from e
+
+
 def web_options():
   """Decorator to add web UI options to click commands."""