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

google/adk/sessions/vertex_ai_session_service.py

@@ -16,13 +16,12 @@ from __future__ import annotations
 import asyncio
 import logging
 import re
-import time
 from typing import Any
+from typing import Dict
 from typing import Optional
 import urllib.parse
 
 from dateutil import parser
-from google.genai import types
 from typing_extensions import override
 
 from google import genai
@@ -40,18 +39,27 @@ logger = logging.getLogger('google_adk.' + __name__)
 
 
 class VertexAiSessionService(BaseSessionService):
-  """Connects to the managed Vertex AI Session Service."""
+  """Connects to the Vertex AI Agent Engine Session Service using GenAI API client.
+
+  https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/sessions/overview
+  """
 
   def __init__(
       self,
-      project: str = None,
-      location: str = None,
+      project: Optional[str] = None,
+      location: Optional[str] = None,
+      agent_engine_id: Optional[str] = None,
   ):
-    self.project = project
-    self.location = location
+    """Initializes the VertexAiSessionService.
 
-    client = genai.Client(vertexai=True, project=project, location=location)
-    self.api_client = client._api_client
+    Args:
+      project: The project id of the project to use.
+      location: The location of the project to use.
+      agent_engine_id: The resource ID of the agent engine to use.
+    """
+    self._project = project
+    self._location = location
+    self._agent_engine_id = agent_engine_id
 
   @override
   async def create_session(
@@ -67,14 +75,13 @@ class VertexAiSessionService(BaseSessionService):
           'User-provided Session id is not supported for'
           ' VertexAISessionService.'
       )
-
-    reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+    reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    api_client = self._get_api_client()
 
     session_json_dict = {'user_id': user_id}
     if state:
       session_json_dict['session_state'] = state
 
-    api_client = _get_api_client(self.project, self.location)
     api_response = await api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions',
@@ -86,6 +93,7 @@ class VertexAiSessionService(BaseSessionService):
     operation_id = api_response['name'].split('/')[-1]
 
     max_retry_attempt = 5
+    lro_response = None
     while max_retry_attempt >= 0:
       lro_response = await api_client.async_request(
           http_method='GET',
@@ -99,6 +107,11 @@ class VertexAiSessionService(BaseSessionService):
       await asyncio.sleep(1)
       max_retry_attempt -= 1
 
+    if lro_response is None or not lro_response.get('done', None):
+      raise TimeoutError(
+          f'Timeout waiting for operation {operation_id} to complete.'
+      )
+
     # Get session resource
     get_session_api_response = await api_client.async_request(
         http_method='GET',
@@ -127,10 +140,10 @@ class VertexAiSessionService(BaseSessionService):
       session_id: str,
       config: Optional[GetSessionConfig] = None,
   ) -> Optional[Session]:
-    reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+    reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    api_client = self._get_api_client()
 
     # Get session resource
-    api_client = _get_api_client(self.project, self.location)
     get_session_api_response = await api_client.async_request(
         http_method='GET',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
@@ -200,14 +213,14 @@ class VertexAiSessionService(BaseSessionService):
   async def list_sessions(
       self, *, app_name: str, user_id: str
   ) -> ListSessionsResponse:
-    reasoning_engine_id = _parse_reasoning_engine_id(app_name)
+    reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    api_client = self._get_api_client()
 
     path = f'reasoningEngines/{reasoning_engine_id}/sessions'
     if user_id:
       parsed_user_id = urllib.parse.quote(f'''"{user_id}"''', safe='')
       path = path + f'?filter=user_id={parsed_user_id}'
 
-    api_client = _get_api_client(self.project, self.location)
     api_response = await api_client.async_request(
         http_method='GET',
         path=path,
@@ -233,21 +246,26 @@ class VertexAiSessionService(BaseSessionService):
   async def delete_session(
       self, *, app_name: str, user_id: str, session_id: str
   ) -> None:
-    reasoning_engine_id = _parse_reasoning_engine_id(app_name)
-    api_client = _get_api_client(self.project, self.location)
-    await api_client.async_request(
-        http_method='DELETE',
-        path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
-        request_dict={},
-    )
+    reasoning_engine_id = self._get_reasoning_engine_id(app_name)
+    api_client = self._get_api_client()
+
+    try:
+      await api_client.async_request(
+          http_method='DELETE',
+          path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session_id}',
+          request_dict={},
+      )
+    except Exception as e:
+      logger.error(f'Error deleting session {session_id}: {e}')
+      raise e
 
   @override
   async def append_event(self, session: Session, event: Event) -> Event:
     # Update the in-memory session.
     await super().append_event(session=session, event=event)
 
-    reasoning_engine_id = _parse_reasoning_engine_id(session.app_name)
-    api_client = _get_api_client(self.project, self.location)
+    reasoning_engine_id = self._get_reasoning_engine_id(session.app_name)
+    api_client = self._get_api_client()
     await api_client.async_request(
         http_method='POST',
         path=f'reasoningEngines/{reasoning_engine_id}/sessions/{session.id}:appendEvent',
@@ -255,18 +273,37 @@ class VertexAiSessionService(BaseSessionService):
     )
     return event
 
+  def _get_reasoning_engine_id(self, app_name: str):
+    if self._agent_engine_id:
+      return self._agent_engine_id
 
-def _get_api_client(project: str, location: str):
-  """Instantiates an API client for the given project and location.
+    if app_name.isdigit():
+      return app_name
 
-  It needs to be instantiated inside each request so that the event loop
-  management.
-  """
-  client = genai.Client(vertexai=True, project=project, location=location)
-  return client._api_client
+    pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
+    match = re.fullmatch(pattern, app_name)
+
+    if not bool(match):
+      raise ValueError(
+          f'App name {app_name} is not valid. It should either be the full'
+          ' ReasoningEngine resource name, or the reasoning engine id.'
+      )
+
+    return match.groups()[-1]
+
+  def _get_api_client(self):
+    """Instantiates an API client for the given project and location.
+
+    It needs to be instantiated inside each request so that the event loop
+    management can be properly propagated.
+    """
+    client = genai.Client(
+        vertexai=True, project=self._project, location=self._location
+    )
+    return client._api_client
 
 
-def _convert_event_to_json(event: Event):
+def _convert_event_to_json(event: Event) -> Dict[str, Any]:
   metadata_json = {
       'partial': event.partial,
       'turn_complete': event.turn_complete,
@@ -318,7 +355,7 @@ def _convert_event_to_json(event: Event):
   return event_json
 
 
-def _from_api_event(api_event: dict) -> Event:
+def _from_api_event(api_event: Dict[str, Any]) -> Event:
   event_actions = EventActions()
   if api_event.get('actions', None):
     event_actions = EventActions(
@@ -359,19 +396,3 @@ def _from_api_event(api_event: dict) -> Event:
     )
 
   return event
-
-
-def _parse_reasoning_engine_id(app_name: str):
-  if app_name.isdigit():
-    return app_name
-
-  pattern = r'^projects/([a-zA-Z0-9-_]+)/locations/([a-zA-Z0-9-_]+)/reasoningEngines/(\d+)$'
-  match = re.fullmatch(pattern, app_name)
-
-  if not bool(match):
-    raise ValueError(
-        f'App name {app_name} is not valid. It should either be the full'
-        ' ReasoningEngine resource name, or the reasoning engine id.'
-    )
-
-  return match.groups()[-1]

google/adk/tools/__init__.py

@@ -27,6 +27,7 @@ from .long_running_tool import LongRunningFunctionTool
 from .preload_memory_tool import preload_memory_tool as preload_memory
 from .tool_context import ToolContext
 from .transfer_to_agent_tool import transfer_to_agent
+from .url_context_tool import url_context
 from .vertex_ai_search_tool import VertexAiSearchTool
 
 __all__ = [
@@ -34,6 +35,7 @@ __all__ = [
     'AuthToolArguments',
     'BaseTool',
     'google_search',
+    'url_context',
     'VertexAiSearchTool',
     'ExampleTool',
     'exit_loop',

google/adk/tools/_automatic_function_calling_util.py

@@ -230,6 +230,7 @@ def build_function_declaration(
       )
       new_func.__signature__ = new_sig
       new_func.__doc__ = func.__doc__
+      new_func.__annotations__ = func.__annotations__
 
   return (
       from_function_with_options(func, variant)

google/adk/tools/_forwarding_artifact_service.py

@@ -0,0 +1,96 @@
+# 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
+
+from typing import Optional
+from typing import TYPE_CHECKING
+
+from google.genai import types
+from typing_extensions import override
+
+from ..artifacts.base_artifact_service import BaseArtifactService
+
+if TYPE_CHECKING:
+  from .tool_context import ToolContext
+
+
+class ForwardingArtifactService(BaseArtifactService):
+  """Artifact service that forwards to the parent tool context."""
+
+  def __init__(self, tool_context: ToolContext):
+    self.tool_context = tool_context
+    self._invocation_context = tool_context._invocation_context
+
+  @override
+  async def save_artifact(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      filename: str,
+      artifact: types.Part,
+  ) -> int:
+    return await self.tool_context.save_artifact(
+        filename=filename, artifact=artifact
+    )
+
+  @override
+  async def load_artifact(
+      self,
+      *,
+      app_name: str,
+      user_id: str,
+      session_id: str,
+      filename: str,
+      version: Optional[int] = None,
+  ) -> Optional[types.Part]:
+    return await self.tool_context.load_artifact(
+        filename=filename, version=version
+    )
+
+  @override
+  async def list_artifact_keys(
+      self, *, app_name: str, user_id: str, session_id: str
+  ) -> list[str]:
+    return await self.tool_context.list_artifacts()
+
+  @override
+  async def delete_artifact(
+      self, *, app_name: str, user_id: str, session_id: str, filename: str
+  ) -> None:
+    del app_name, user_id, session_id
+    if self._invocation_context.artifact_service is None:
+      raise ValueError("Artifact service is not initialized.")
+    await self._invocation_context.artifact_service.delete_artifact(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        session_id=self._invocation_context.session.id,
+        filename=filename,
+    )
+
+  @override
+  async def list_versions(
+      self, *, app_name: str, user_id: str, session_id: str, filename: str
+  ) -> list[int]:
+    del app_name, user_id, session_id
+    if self._invocation_context.artifact_service is None:
+      raise ValueError("Artifact service is not initialized.")
+    return await self._invocation_context.artifact_service.list_versions(
+        app_name=self._invocation_context.app_name,
+        user_id=self._invocation_context.user_id,
+        session_id=self._invocation_context.session.id,
+        filename=filename,
+    )

google/adk/tools/_function_parameter_parse_util.py

@@ -37,6 +37,7 @@ _py_builtin_type_to_schema_type = {
     bool: types.Type.BOOLEAN,
     list: types.Type.ARRAY,
     dict: types.Type.OBJECT,
+    None: types.Type.NULL,
 }
 
 logger = logging.getLogger('google_adk.' + __name__)

google/adk/tools/agent_tool.py

@@ -25,6 +25,7 @@ from . import _automatic_function_calling_util
 from ..memory.in_memory_memory_service import InMemoryMemoryService
 from ..runners import Runner
 from ..sessions.in_memory_session_service import InMemorySessionService
+from ._forwarding_artifact_service import ForwardingArtifactService
 from .base_tool import BaseTool
 from .tool_context import ToolContext
 
@@ -96,17 +97,6 @@ class AgentTool(BaseTool):
 
     if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
       input_value = self.agent.input_schema.model_validate(args)
-    else:
-      input_value = args['request']
-
-    if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
-      if isinstance(input_value, dict):
-        input_value = self.agent.input_schema.model_validate(input_value)
-      if not isinstance(input_value, self.agent.input_schema):
-        raise ValueError(
-            f'Input value {input_value} is not of type'
-            f' `{self.agent.input_schema}`.'
-        )
       content = types.Content(
           role='user',
           parts=[
@@ -118,14 +108,12 @@ class AgentTool(BaseTool):
     else:
       content = types.Content(
           role='user',
-          parts=[types.Part.from_text(text=input_value)],
+          parts=[types.Part.from_text(text=args['request'])],
       )
     runner = Runner(
         app_name=self.agent.name,
         agent=self.agent,
-        # TODO(kech): Remove the access to the invocation context.
-        #   It seems we don't need re-use artifact_service if we forward below.
-        artifact_service=tool_context._invocation_context.artifact_service,
+        artifact_service=ForwardingArtifactService(tool_context),
         session_service=InMemorySessionService(),
         memory_service=InMemoryMemoryService(),
     )
@@ -144,35 +132,13 @@ class AgentTool(BaseTool):
         tool_context.state.update(event.actions.state_delta)
       last_event = event
 
-    if runner.artifact_service:
-      # Forward all artifacts to parent session.
-      artifact_names = await runner.artifact_service.list_artifact_keys(
-          app_name=session.app_name,
-          user_id=session.user_id,
-          session_id=session.id,
-      )
-      for artifact_name in artifact_names:
-        if artifact := await runner.artifact_service.load_artifact(
-            app_name=session.app_name,
-            user_id=session.user_id,
-            session_id=session.id,
-            filename=artifact_name,
-        ):
-          await tool_context.save_artifact(
-              filename=artifact_name, artifact=artifact
-          )
-
     if not last_event or not last_event.content or not last_event.content.parts:
       return ''
+    merged_text = '\n'.join(p.text for p in last_event.content.parts if p.text)
     if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
-      merged_text = '\n'.join(
-          [p.text for p in last_event.content.parts if p.text]
-      )
       tool_result = self.agent.output_schema.model_validate_json(
           merged_text
       ).model_dump(exclude_none=True)
     else:
-      tool_result = '\n'.join(
-          [p.text for p in last_event.content.parts if p.text]
-      )
+      tool_result = merged_text
     return tool_result

google/adk/tools/application_integration_tool/integration_connector_tool.py

@@ -150,7 +150,7 @@ class IntegrationConnectorTool(BaseTool):
     tool_auth_handler = ToolAuthHandler.from_tool_context(
         tool_context, self._auth_scheme, self._auth_credential
     )
-    auth_result = tool_auth_handler.prepare_auth_credentials()
+    auth_result = await tool_auth_handler.prepare_auth_credentials()
 
     if auth_result.state == 'pending':
       return {
@@ -178,7 +178,7 @@ class IntegrationConnectorTool(BaseTool):
     args['operation'] = self._operation
     args['action'] = self._action
     logger.info('Running tool: %s with args: %s', self.name, args)
-    return self._rest_api_tool.call(args=args, tool_context=tool_context)
+    return await self._rest_api_tool.call(args=args, tool_context=tool_context)
 
   def __str__(self):
     return (

google/adk/tools/authenticated_function_tool.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 inspect
+import logging
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Optional
+from typing import Union
+
+from typing_extensions import override
+
+from ..auth.auth_credential import AuthCredential
+from ..auth.auth_tool import AuthConfig
+from ..auth.credential_manager import CredentialManager
+from ..utils.feature_decorator import experimental
+from .function_tool import FunctionTool
+from .tool_context import ToolContext
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class AuthenticatedFunctionTool(FunctionTool):
+  """A FunctionTool that handles authentication before the actual tool logic
+  gets called. Functions can accept a special `credential` argument which is the
+  credential ready for use.(Experimental)
+  """
+
+  def __init__(
+      self,
+      *,
+      func: Callable[..., Any],
+      auth_config: AuthConfig = None,
+      response_for_auth_required: Optional[Union[dict[str, Any], str]] = None,
+  ):
+    """Initializes the AuthenticatedFunctionTool.
+
+    Args:
+        func: The function to be called.
+        auth_config: The authentication configuration.
+        response_for_auth_required: The response to return when the tool is
+          requesting auth credential from the client. There could be two case,
+          the tool doesn't configure any credentials
+          (auth_config.raw_auth_credential is missing) or the credentials
+          configured is not enough to authenticate the tool (e.g. an OAuth
+          client id and client secrect is configured.) and needs client input
+          (e.g. client need to involve the end user in an oauth flow and get
+          back the oauth response.)
+    """
+    super().__init__(func=func)
+    self._ignore_params.append("credential")
+
+    if auth_config and auth_config.auth_scheme:
+      self._credentials_manager = CredentialManager(auth_config=auth_config)
+    else:
+      logger.warning(
+          "auth_config or auth_config.auth_scheme is missing. Will skip"
+          " authentication.Using FunctionTool instead if authentication is not"
+          " required."
+      )
+      self._credentials_manager = None
+    self._response_for_auth_required = response_for_auth_required
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    credential = None
+    if self._credentials_manager:
+      credential = await self._credentials_manager.get_auth_credential(
+          tool_context
+      )
+      if not credential:
+        await self._credentials_manager.request_credential(tool_context)
+        return self._response_for_auth_required or "Pending User Authorization."
+
+    return await self._run_async_impl(
+        args=args, tool_context=tool_context, credential=credential
+    )
+
+  async def _run_async_impl(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      credential: AuthCredential,
+  ) -> Any:
+    args_to_call = args.copy()
+    signature = inspect.signature(self.func)
+    if "credential" in signature.parameters:
+      args_to_call["credential"] = credential
+    return await super().run_async(args=args_to_call, tool_context=tool_context)

google/adk/tools/base_authenticated_tool.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
+
+from abc import abstractmethod
+import logging
+from typing import Any
+from typing import Optional
+from typing import Union
+
+from typing_extensions import override
+
+from ..auth.auth_credential import AuthCredential
+from ..auth.auth_tool import AuthConfig
+from ..auth.credential_manager import CredentialManager
+from ..utils.feature_decorator import experimental
+from .base_tool import BaseTool
+from .tool_context import ToolContext
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+@experimental
+class BaseAuthenticatedTool(BaseTool):
+  """A base tool class that handles authentication before the actual tool logic
+  gets called. Functions can accept a special `credential` argument which is the
+  credential ready for use.(Experimental)
+  """
+
+  def __init__(
+      self,
+      *,
+      name,
+      description,
+      auth_config: AuthConfig = None,
+      response_for_auth_required: Optional[Union[dict[str, Any], str]] = None,
+  ):
+    """
+    Args:
+      name: The name of the tool.
+      description: The description of the tool.
+      auth_config: The auth configuration of the tool.
+      response_for_auth_required: The response to return when the tool is
+          requesting auth credential from the client. There could be two case,
+          the tool doesn't configure any credentials
+          (auth_config.raw_auth_credential is missing) or the credentials
+          configured is not enough to authenticate the tool (e.g. an OAuth
+          client id and client secrect is configured.) and needs client input
+          (e.g. client need to involve the end user in an oauth flow and get
+          back the oauth response.)
+    """
+    super().__init__(
+        name=name,
+        description=description,
+    )
+
+    if auth_config and auth_config.auth_scheme:
+      self._credentials_manager = CredentialManager(auth_config=auth_config)
+    else:
+      logger.warning(
+          "auth_config or auth_config.auth_scheme is missing. Will skip"
+          " authentication.Using FunctionTool instead if authentication is not"
+          " required."
+      )
+      self._credentials_manager = None
+    self._response_for_auth_required = response_for_auth_required
+
+  @override
+  async def run_async(
+      self, *, args: dict[str, Any], tool_context: ToolContext
+  ) -> Any:
+    credential = None
+    if self._credentials_manager:
+      credential = await self._credentials_manager.get_auth_credential(
+          tool_context
+      )
+      if not credential:
+        await self._credentials_manager.request_credential(tool_context)
+        return self._response_for_auth_required or "Pending User Authorization."
+
+    return await self._run_async_impl(
+        args=args,
+        tool_context=tool_context,
+        credential=credential,
+    )
+
+  @abstractmethod
+  async def _run_async_impl(
+      self,
+      *,
+      args: dict[str, Any],
+      tool_context: ToolContext,
+      credential: AuthCredential,
+  ) -> Any:
+    pass