google-adk 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

google/adk/auth/exchanger/service_account_credential_exchanger.py

@@ -0,0 +1,104 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Credential fetcher for Google Service Account."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import google.auth
+from google.auth.transport.requests import Request
+from google.oauth2 import service_account
+from typing_extensions import override
+
+from ...utils.feature_decorator import experimental
+from ..auth_credential import AuthCredential
+from ..auth_credential import AuthCredentialTypes
+from ..auth_schemes import AuthScheme
+from .base_credential_exchanger import BaseCredentialExchanger
+
+
+@experimental
+class ServiceAccountCredentialExchanger(BaseCredentialExchanger):
+  """Exchanges Google Service Account credentials for an access token.
+
+  Uses the default service credential if `use_default_credential = True`.
+  Otherwise, uses the service account credential provided in the auth
+  credential.
+  """
+
+  @override
+  async def exchange(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Exchanges the service account auth credential for an access token.
+
+    If the AuthCredential contains a service account credential, it will be used
+    to exchange for an access token. Otherwise, if use_default_credential is True,
+    the default application credential will be used for exchanging an access token.
+
+    Args:
+        auth_scheme: The authentication scheme.
+        auth_credential: The credential to exchange.
+
+    Returns:
+        An AuthCredential in OAUTH2 format, containing the exchanged credential JSON.
+
+    Raises:
+        ValueError: If service account credentials are missing or invalid.
+        Exception: If credential exchange or refresh fails.
+    """
+    if auth_credential is None:
+      raise ValueError("Credential cannot be None.")
+
+    if auth_credential.auth_type != AuthCredentialTypes.SERVICE_ACCOUNT:
+      raise ValueError("Credential is not a service account credential.")
+
+    if auth_credential.service_account is None:
+      raise ValueError(
+          "Service account credentials are missing. Please provide them."
+      )
+
+    if (
+        auth_credential.service_account.service_account_credential is None
+        and not auth_credential.service_account.use_default_credential
+    ):
+      raise ValueError(
+          "Service account credentials are invalid. Please set the"
+          " service_account_credential field or set `use_default_credential ="
+          " True` to use application default credential in a hosted service"
+          " like Google Cloud Run."
+      )
+
+    try:
+      if auth_credential.service_account.use_default_credential:
+        credentials, _ = google.auth.default()
+      else:
+        config = auth_credential.service_account
+        credentials = service_account.Credentials.from_service_account_info(
+            config.service_account_credential.model_dump(), scopes=config.scopes
+        )
+
+      # Refresh credentials to ensure we have a valid access token
+      credentials.refresh(Request())
+
+      return AuthCredential(
+          auth_type=AuthCredentialTypes.OAUTH2,
+          google_oauth2_json=credentials.to_json(),
+      )
+    except Exception as e:
+      raise ValueError(f"Failed to exchange service account token: {e}") from e

google/adk/auth/oauth2_credential_util.py

@@ -0,0 +1,107 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+from typing import Tuple
+
+from fastapi.openapi.models import OAuth2
+
+from ..utils.feature_decorator import experimental
+from .auth_credential import AuthCredential
+from .auth_schemes import AuthScheme
+from .auth_schemes import OpenIdConnectWithConfig
+
+try:
+  from authlib.integrations.requests_client import OAuth2Session
+  from authlib.oauth2.rfc6749 import OAuth2Token
+
+  AUTHLIB_AVIALABLE = True
+except ImportError:
+  AUTHLIB_AVIALABLE = False
+
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+def create_oauth2_session(
+    auth_scheme: AuthScheme,
+    auth_credential: AuthCredential,
+) -> Tuple[Optional[OAuth2Session], Optional[str]]:
+  """Create an OAuth2 session for token operations.
+
+  Args:
+      auth_scheme: The authentication scheme configuration.
+      auth_credential: The authentication credential.
+
+  Returns:
+      Tuple of (OAuth2Session, token_endpoint) or (None, None) if cannot create session.
+  """
+  if isinstance(auth_scheme, OpenIdConnectWithConfig):
+    if not hasattr(auth_scheme, "token_endpoint"):
+      return None, None
+    token_endpoint = auth_scheme.token_endpoint
+    scopes = auth_scheme.scopes
+  elif isinstance(auth_scheme, OAuth2):
+    if (
+        not auth_scheme.flows.authorizationCode
+        or not auth_scheme.flows.authorizationCode.tokenUrl
+    ):
+      return None, None
+    token_endpoint = auth_scheme.flows.authorizationCode.tokenUrl
+    scopes = list(auth_scheme.flows.authorizationCode.scopes.keys())
+  else:
+    return None, None
+
+  if (
+      not auth_credential
+      or not auth_credential.oauth2
+      or not auth_credential.oauth2.client_id
+      or not auth_credential.oauth2.client_secret
+  ):
+    return None, None
+
+  return (
+      OAuth2Session(
+          auth_credential.oauth2.client_id,
+          auth_credential.oauth2.client_secret,
+          scope=" ".join(scopes),
+          redirect_uri=auth_credential.oauth2.redirect_uri,
+          state=auth_credential.oauth2.state,
+      ),
+      token_endpoint,
+  )
+
+
+@experimental
+def update_credential_with_tokens(
+    auth_credential: AuthCredential, tokens: OAuth2Token
+) -> None:
+  """Update the credential with new tokens.
+
+  Args:
+      auth_credential: The authentication credential to update.
+      tokens: The OAuth2Token object containing new token information.
+  """
+  auth_credential.oauth2.access_token = tokens.get("access_token")
+  auth_credential.oauth2.refresh_token = tokens.get("refresh_token")
+  auth_credential.oauth2.expires_at = (
+      int(tokens.get("expires_at")) if tokens.get("expires_at") else None
+  )
+  auth_credential.oauth2.expires_in = (
+      int(tokens.get("expires_in")) if tokens.get("expires_in") else None
+  )

google/adk/auth/refresher/__init__.py

@@ -0,0 +1,21 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Credential refresher module."""
+
+from .base_credential_refresher import BaseCredentialRefresher
+
+__all__ = [
+    "BaseCredentialRefresher",
+]

google/adk/auth/refresher/base_credential_refresher.py

@@ -0,0 +1,74 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base credential refresher interface."""
+
+from __future__ import annotations
+
+import abc
+from typing import Optional
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.utils.feature_decorator import experimental
+
+
+class CredentialRefresherError(Exception):
+  """Base exception for credential refresh errors."""
+
+
+@experimental
+class BaseCredentialRefresher(abc.ABC):
+  """Base interface for credential refreshers.
+
+  Credential refreshers are responsible for checking if a credential is expired
+  or needs to be refreshed, and for refreshing it if necessary.
+  """
+
+  @abc.abstractmethod
+  async def is_refresh_needed(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> bool:
+    """Checks if a credential needs to be refreshed.
+
+    Args:
+        auth_credential: The credential to check.
+        auth_scheme: The authentication scheme (optional, some refreshers don't need it).
+
+    Returns:
+        True if the credential needs to be refreshed, False otherwise.
+    """
+    pass
+
+  @abc.abstractmethod
+  async def refresh(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Refreshes a credential if needed.
+
+    Args:
+        auth_credential: The credential to refresh.
+        auth_scheme: The authentication scheme (optional, some refreshers don't need it).
+
+    Returns:
+        The refreshed credential.
+
+    Raises:
+        CredentialRefresherError: If credential refresh fails.
+    """
+    pass

google/adk/auth/refresher/credential_refresher_registry.py

@@ -0,0 +1,59 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Credential refresher registry."""
+
+from __future__ import annotations
+
+from typing import Dict
+from typing import Optional
+
+from google.adk.auth.auth_credential import AuthCredentialTypes
+from google.adk.utils.feature_decorator import experimental
+
+from .base_credential_refresher import BaseCredentialRefresher
+
+
+@experimental
+class CredentialRefresherRegistry:
+  """Registry for credential refresher instances."""
+
+  def __init__(self):
+    self._refreshers: Dict[AuthCredentialTypes, BaseCredentialRefresher] = {}
+
+  def register(
+      self,
+      credential_type: AuthCredentialTypes,
+      refresher_instance: BaseCredentialRefresher,
+  ) -> None:
+    """Register a refresher instance for a credential type.
+
+    Args:
+        credential_type: The credential type to register for.
+        refresher_instance: The refresher instance to register.
+    """
+    self._refreshers[credential_type] = refresher_instance
+
+  def get_refresher(
+      self, credential_type: AuthCredentialTypes
+  ) -> Optional[BaseCredentialRefresher]:
+    """Get the refresher instance for a credential type.
+
+    Args:
+        credential_type: The credential type to get refresher for.
+
+    Returns:
+        The refresher instance if registered, None otherwise.
+    """
+    return self._refreshers.get(credential_type)

google/adk/auth/refresher/oauth2_credential_refresher.py

@@ -0,0 +1,154 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""OAuth2 credential refresher implementation."""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import Optional
+
+from google.adk.auth.auth_credential import AuthCredential
+from google.adk.auth.auth_schemes import AuthScheme
+from google.adk.auth.oauth2_credential_util import create_oauth2_session
+from google.adk.auth.oauth2_credential_util import update_credential_with_tokens
+from google.adk.utils.feature_decorator import experimental
+from google.auth.transport.requests import Request
+from google.oauth2.credentials import Credentials
+from typing_extensions import override
+
+from .base_credential_refresher import BaseCredentialRefresher
+
+try:
+  from authlib.oauth2.rfc6749 import OAuth2Token
+
+  AUTHLIB_AVIALABLE = True
+except ImportError:
+  AUTHLIB_AVIALABLE = False
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class OAuth2CredentialRefresher(BaseCredentialRefresher):
+  """Refreshes OAuth2 credentials including Google OAuth2 JSON credentials."""
+
+  @override
+  async def is_refresh_needed(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> bool:
+    """Check if the OAuth2 credential needs to be refreshed.
+
+    Args:
+        auth_credential: The OAuth2 credential to check.
+        auth_scheme: The OAuth2 authentication scheme (optional for Google OAuth2 JSON).
+
+    Returns:
+        True if the credential needs to be refreshed, False otherwise.
+    """
+    # Handle Google OAuth2 credentials (from service account exchange)
+    if auth_credential.google_oauth2_json:
+      try:
+        google_credential = Credentials.from_authorized_user_info(
+            json.loads(auth_credential.google_oauth2_json)
+        )
+        return google_credential.expired and bool(
+            google_credential.refresh_token
+        )
+      except Exception as e:
+        logger.warning("Failed to parse Google OAuth2 JSON credential: %s", e)
+        return False
+
+    # Handle regular OAuth2 credentials
+    elif auth_credential.oauth2 and auth_scheme:
+      if not AUTHLIB_AVIALABLE:
+        return False
+
+      if not auth_credential.oauth2:
+        return False
+
+      return OAuth2Token({
+          "expires_at": auth_credential.oauth2.expires_at,
+          "expires_in": auth_credential.oauth2.expires_in,
+      }).is_expired()
+
+    return False
+
+  @override
+  async def refresh(
+      self,
+      auth_credential: AuthCredential,
+      auth_scheme: Optional[AuthScheme] = None,
+  ) -> AuthCredential:
+    """Refresh the OAuth2 credential.
+    If refresh failed, return the original credential.
+
+    Args:
+        auth_credential: The OAuth2 credential to refresh.
+        auth_scheme: The OAuth2 authentication scheme (optional for Google OAuth2 JSON).
+
+    Returns:
+        The refreshed credential.
+
+    """
+    # Handle Google OAuth2 credentials (from service account exchange)
+    if auth_credential.google_oauth2_json:
+      try:
+        google_credential = Credentials.from_authorized_user_info(
+            json.loads(auth_credential.google_oauth2_json)
+        )
+        if google_credential.expired and google_credential.refresh_token:
+          google_credential.refresh(Request())
+          auth_credential.google_oauth2_json = google_credential.to_json()
+          logger.info("Successfully refreshed Google OAuth2 JSON credential")
+      except Exception as e:
+        # TODO reconsider whether we should raise error when refresh failed.
+        logger.error("Failed to refresh Google OAuth2 JSON credential: %s", e)
+
+    # Handle regular OAuth2 credentials
+    elif auth_credential.oauth2 and auth_scheme:
+      if not AUTHLIB_AVIALABLE:
+        return auth_credential
+
+      if not auth_credential.oauth2:
+        return auth_credential
+
+      if OAuth2Token({
+          "expires_at": auth_credential.oauth2.expires_at,
+          "expires_in": auth_credential.oauth2.expires_in,
+      }).is_expired():
+        client, token_endpoint = create_oauth2_session(
+            auth_scheme, auth_credential
+        )
+        if not client:
+          logger.warning("Could not create OAuth2 session for token refresh")
+          return auth_credential
+
+        try:
+          tokens = client.refresh_token(
+              url=token_endpoint,
+              refresh_token=auth_credential.oauth2.refresh_token,
+          )
+          update_credential_with_tokens(auth_credential, tokens)
+          logger.debug("Successfully refreshed OAuth2 tokens")
+        except Exception as e:
+          # TODO reconsider whether we should raise error when refresh failed.
+          logger.error("Failed to refresh OAuth2 tokens: %s", e)
+          # Return original credential on failure
+          return auth_credential
+
+    return auth_credential

google/adk/cli/agent_graph.py

@@ -64,11 +64,11 @@ async def build_graph(
     if isinstance(tool_or_agent, BaseAgent):
       # Added Workflow Agent checks for different agent types
       if isinstance(tool_or_agent, SequentialAgent):
-        return tool_or_agent.name + f' (Sequential Agent)'
+        return tool_or_agent.name + ' (Sequential Agent)'
       elif isinstance(tool_or_agent, LoopAgent):
-        return tool_or_agent.name + f' (Loop Agent)'
+        return tool_or_agent.name + ' (Loop Agent)'
       elif isinstance(tool_or_agent, ParallelAgent):
-        return tool_or_agent.name + f' (Parallel Agent)'
+        return tool_or_agent.name + ' (Parallel Agent)'
       else:
         return tool_or_agent.name
     elif isinstance(tool_or_agent, BaseTool):
@@ -144,49 +144,53 @@ async def build_graph(
       )
       return False
 
-  def build_cluster(child: graphviz.Digraph, agent: BaseAgent, name: str):
+  async def build_cluster(child: graphviz.Digraph, agent: BaseAgent, name: str):
     if isinstance(agent, LoopAgent):
       # Draw the edge from the parent agent to the first sub-agent
-      draw_edge(parent_agent.name, agent.sub_agents[0].name)
+      if parent_agent:
+        draw_edge(parent_agent.name, agent.sub_agents[0].name)
       length = len(agent.sub_agents)
-      currLength = 0
+      curr_length = 0
       # Draw the edges between the sub-agents
       for sub_agent_int_sequential in agent.sub_agents:
-        build_graph(child, sub_agent_int_sequential, highlight_pairs)
+        await build_graph(child, sub_agent_int_sequential, highlight_pairs)
         # Draw the edge between the current sub-agent and the next one
         # If it's the last sub-agent, draw an edge to the first one to indicating a loop
         draw_edge(
-            agent.sub_agents[currLength].name,
+            agent.sub_agents[curr_length].name,
             agent.sub_agents[
-                0 if currLength == length - 1 else currLength + 1
+                0 if curr_length == length - 1 else curr_length + 1
             ].name,
         )
-        currLength += 1
+        curr_length += 1
     elif isinstance(agent, SequentialAgent):
       # Draw the edge from the parent agent to the first sub-agent
-      draw_edge(parent_agent.name, agent.sub_agents[0].name)
+      if parent_agent:
+        draw_edge(parent_agent.name, agent.sub_agents[0].name)
       length = len(agent.sub_agents)
-      currLength = 0
+      curr_length = 0
 
       # Draw the edges between the sub-agents
       for sub_agent_int_sequential in agent.sub_agents:
-        build_graph(child, sub_agent_int_sequential, highlight_pairs)
+        await build_graph(child, sub_agent_int_sequential, highlight_pairs)
         # Draw the edge between the current sub-agent and the next one
         # If it's the last sub-agent, don't draw an edge to avoid a loop
-        draw_edge(
-            agent.sub_agents[currLength].name,
-            agent.sub_agents[currLength + 1].name,
-        ) if currLength != length - 1 else None
-        currLength += 1
+        if curr_length != length - 1:
+          draw_edge(
+              agent.sub_agents[curr_length].name,
+              agent.sub_agents[curr_length + 1].name,
+          )
+        curr_length += 1
 
     elif isinstance(agent, ParallelAgent):
       # Draw the edge from the parent agent to every sub-agent
       for sub_agent in agent.sub_agents:
-        build_graph(child, sub_agent, highlight_pairs)
-        draw_edge(parent_agent.name, sub_agent.name)
+        await build_graph(child, sub_agent, highlight_pairs)
+        if parent_agent:
+          draw_edge(parent_agent.name, sub_agent.name)
     else:
       for sub_agent in agent.sub_agents:
-        build_graph(child, sub_agent, highlight_pairs)
+        await build_graph(child, sub_agent, highlight_pairs)
         draw_edge(agent.name, sub_agent.name)
 
     child.attr(
@@ -196,21 +200,20 @@ async def build_graph(
         fontcolor=light_gray,
     )
 
-  def draw_node(tool_or_agent: Union[BaseAgent, BaseTool]):
+  async def draw_node(tool_or_agent: Union[BaseAgent, BaseTool]):
     name = get_node_name(tool_or_agent)
     shape = get_node_shape(tool_or_agent)
     caption = get_node_caption(tool_or_agent)
-    asCluster = should_build_agent_cluster(tool_or_agent)
-    child = None
+    as_cluster = should_build_agent_cluster(tool_or_agent)
     if highlight_pairs:
       for highlight_tuple in highlight_pairs:
         if name in highlight_tuple:
           # if in highlight, draw highlight node
-          if asCluster:
+          if as_cluster:
             cluster = graphviz.Digraph(
                 name='cluster_' + name
             )  # adding "cluster_" to the name makes the graph render as a cluster subgraph
-            build_cluster(cluster, agent, name)
+            await build_cluster(cluster, agent, name)
             graph.subgraph(cluster)
           else:
             graph.node(
@@ -224,12 +227,12 @@ async def build_graph(
             )
           return
     # if not in highlight, draw non-highlight node
-    if asCluster:
+    if as_cluster:
 
       cluster = graphviz.Digraph(
           name='cluster_' + name
       )  # adding "cluster_" to the name makes the graph render as a cluster subgraph
-      build_cluster(cluster, agent, name)
+      await build_cluster(cluster, agent, name)
       graph.subgraph(cluster)
 
     else:
@@ -264,10 +267,9 @@ async def build_graph(
     else:
       graph.edge(from_name, to_name, arrowhead='none', color=light_gray)
 
-  draw_node(agent)
+  await draw_node(agent)
   for sub_agent in agent.sub_agents:
-
-    build_graph(graph, sub_agent, highlight_pairs, agent)
+    await build_graph(graph, sub_agent, highlight_pairs, agent)
     if not should_build_agent_cluster(
         sub_agent
     ) and not should_build_agent_cluster(
@@ -276,7 +278,7 @@ async def build_graph(
       draw_edge(agent.name, sub_agent.name)
   if isinstance(agent, LlmAgent):
     for tool in await agent.canonical_tools():
-      draw_node(tool)
+      await draw_node(tool)
       draw_edge(agent.name, get_node_name(tool))