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

google/adk/tools/mcp_tool/mcp_toolset.py

@@ -22,13 +22,16 @@ from typing import TextIO
 from typing import Union
 
 from ...agents.readonly_context import ReadonlyContext
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_schemes import AuthScheme
 from ..base_tool import BaseTool
 from ..base_toolset import BaseToolset
 from ..base_toolset import ToolPredicate
 from .mcp_session_manager import MCPSessionManager
 from .mcp_session_manager import retry_on_closed_resource
-from .mcp_session_manager import SseServerParams
-from .mcp_session_manager import StreamableHTTPServerParams
+from .mcp_session_manager import SseConnectionParams
+from .mcp_session_manager import StdioConnectionParams
+from .mcp_session_manager import StreamableHTTPConnectionParams
 
 # Attempt to import MCP Tool from the MCP library, and hints user to upgrade
 # their Python version to 3.10 if it fails.
@@ -85,23 +88,34 @@ class MCPToolset(BaseToolset):
   def __init__(
       self,
       *,
-      connection_params: (
-          StdioServerParameters | SseServerParams | StreamableHTTPServerParams
-      ),
+      connection_params: Union[
+          StdioServerParameters,
+          StdioConnectionParams,
+          SseConnectionParams,
+          StreamableHTTPConnectionParams,
+      ],
       tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
       errlog: TextIO = sys.stderr,
+      auth_scheme: Optional[AuthScheme] = None,
+      auth_credential: Optional[AuthCredential] = None,
   ):
     """Initializes the MCPToolset.
 
     Args:
       connection_params: The connection parameters to the MCP server. Can be:
-        `StdioServerParameters` for using local mcp server (e.g. using `npx` or
-        `python3`); or `SseServerParams` for a local/remote SSE server; or
-        `StreamableHTTPServerParams` for local/remote Streamable http server.
-      tool_filter: Optional filter to select specific tools. Can be either:
-        - A list of tool names to include
-        - A ToolPredicate function for custom filtering logic
+        `StdioConnectionParams` for using local mcp server (e.g. using `npx` or
+        `python3`); or `SseConnectionParams` for a local/remote SSE server; or
+        `StreamableHTTPConnectionParams` for local/remote Streamable http
+        server. Note, `StdioServerParameters` is also supported for using local
+        mcp server (e.g. using `npx` or `python3` ), but it does not support
+        timeout, and we recommend to use `StdioConnectionParams` instead when
+        timeout is needed.
+      tool_filter: Optional filter to select specific tools. Can be either: - A
+        list of tool names to include - A ToolPredicate function for custom
+        filtering logic
       errlog: TextIO stream for error logging.
+      auth_scheme: The auth scheme of the tool for tool calling
+      auth_credential: The auth credential of the tool for tool calling
     """
     super().__init__(tool_filter=tool_filter)
 
@@ -116,10 +130,10 @@ class MCPToolset(BaseToolset):
         connection_params=self._connection_params,
         errlog=self._errlog,
     )
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
 
-    self._session = None
-
-  @retry_on_closed_resource("_reinitialize_session")
+  @retry_on_closed_resource
   async def get_tools(
       self,
       readonly_context: Optional[ReadonlyContext] = None,
@@ -134,11 +148,10 @@ class MCPToolset(BaseToolset):
         List[BaseTool]: A list of tools available under the specified context.
     """
     # Get session from session manager
-    if not self._session:
-      self._session = await self._mcp_session_manager.create_session()
+    session = await self._mcp_session_manager.create_session()
 
     # Fetch available tools from the MCP server
-    tools_response: ListToolsResult = await self._session.list_tools()
+    tools_response: ListToolsResult = await session.list_tools()
 
     # Apply filtering based on context and tool_filter
     tools = []
@@ -146,20 +159,14 @@ class MCPToolset(BaseToolset):
       mcp_tool = MCPTool(
           mcp_tool=tool,
           mcp_session_manager=self._mcp_session_manager,
+          auth_scheme=self._auth_scheme,
+          auth_credential=self._auth_credential,
       )
 
       if self._is_tool_selected(mcp_tool, readonly_context):
         tools.append(mcp_tool)
     return tools
 
-  async def _reinitialize_session(self):
-    """Reinitializes the session when connection is lost."""
-    # Close the old session and clear cache
-    await self._mcp_session_manager.close()
-    self._session = await self._mcp_session_manager.create_session()
-
-    # Tools will be reloaded on next get_tools call
-
   async def close(self) -> None:
     """Performs cleanup and releases resources held by the toolset.
 
@@ -172,7 +179,3 @@ class MCPToolset(BaseToolset):
     except Exception as e:
       # Log the error but don't re-raise to avoid blocking shutdown
       print(f"Warning: Error during MCPToolset cleanup: {e}", file=self._errlog)
-    finally:
-      # Clear cached tools
-      self._tools_cache = None
-      self._tools_loaded = False

google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py

@@ -345,9 +345,9 @@ class RestApiTool(BaseTool):
   async def run_async(
       self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
   ) -> Dict[str, Any]:
-    return self.call(args=args, tool_context=tool_context)
+    return await self.call(args=args, tool_context=tool_context)
 
-  def call(
+  async def call(
       self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
   ) -> Dict[str, Any]:
     """Executes the REST API call.
@@ -364,7 +364,7 @@ class RestApiTool(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()
     auth_state, auth_scheme, auth_credential = (
         auth_result.state,
         auth_result.auth_scheme,

google/adk/tools/openapi_tool/openapi_spec_parser/tool_auth_handler.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
 
 import logging
 from typing import Literal
 from typing import Optional
 
-from fastapi.encoders import jsonable_encoder
 from pydantic import BaseModel
 
 from ....auth.auth_credential import AuthCredential
@@ -25,6 +25,7 @@ from ....auth.auth_credential import AuthCredentialTypes
 from ....auth.auth_schemes import AuthScheme
 from ....auth.auth_schemes import AuthSchemeType
 from ....auth.auth_tool import AuthConfig
+from ....auth.refresher.oauth2_credential_refresher import OAuth2CredentialRefresher
 from ...tool_context import ToolContext
 from ..auth.credential_exchangers.auto_auth_credential_exchanger import AutoAuthCredentialExchanger
 from ..auth.credential_exchangers.base_credential_exchanger import AuthCredentialMissingError
@@ -95,10 +96,9 @@ class ToolContextCredentialStore:
       auth_credential: Optional[AuthCredential],
   ):
     if self.tool_context:
-      serializable_credential = jsonable_encoder(
-          auth_credential, exclude_none=True
+      self.tool_context.state[key] = auth_credential.model_dump(
+          exclude_none=True
       )
-      self.tool_context.state[key] = serializable_credential
 
   def remove_credential(self, key: str):
     del self.tool_context.state[key]
@@ -146,20 +146,22 @@ class ToolAuthHandler:
         credential_store,
     )
 
-  def _handle_existing_credential(
+  async def _get_existing_credential(
       self,
-  ) -> Optional[AuthPreparationResult]:
+  ) -> Optional[AuthCredential]:
     """Checks for and returns an existing, exchanged credential."""
     if self.credential_store:
       existing_credential = self.credential_store.get_credential(
           self.auth_scheme, self.auth_credential
       )
       if existing_credential:
-        return AuthPreparationResult(
-            state="done",
-            auth_scheme=self.auth_scheme,
-            auth_credential=existing_credential,
-        )
+        if existing_credential.oauth2:
+          refresher = OAuth2CredentialRefresher()
+          if await refresher.is_refresh_needed(existing_credential):
+            existing_credential = await refresher.refresh(
+                existing_credential, self.auth_scheme
+            )
+        return existing_credential
     return None
 
   def _exchange_credential(
@@ -223,7 +225,17 @@ class ToolAuthHandler:
         )
     )
 
-  def prepare_auth_credentials(
+  def _external_exchange_required(self, credential) -> bool:
+    return (
+        credential.auth_type
+        in (
+            AuthCredentialTypes.OAUTH2,
+            AuthCredentialTypes.OPEN_ID_CONNECT,
+        )
+        and not credential.oauth2.access_token
+    )
+
+  async def prepare_auth_credentials(
       self,
   ) -> AuthPreparationResult:
     """Prepares authentication credentials, handling exchange and user interaction."""
@@ -233,31 +245,41 @@ class ToolAuthHandler:
       return AuthPreparationResult(state="done")
 
     # Check for existing credential.
-    existing_result = self._handle_existing_credential()
-    if existing_result:
-      return existing_result
+    existing_credential = await self._get_existing_credential()
 
+    credential = existing_credential or self.auth_credential
     # fetch credential from adk framework
     # Some auth scheme like OAuth2 AuthCode & OpenIDConnect may require
     # multi-step exchange:
     # client_id , client_secret -> auth_uri -> auth_code -> access_token
-    # -> bearer token
     # adk framework supports exchange access_token already
-    fetched_credential = self._get_auth_response() or self.auth_credential
-
-    exchanged_credential = self._exchange_credential(fetched_credential)
+    # for other credential, adk can also get back the credential directly
+    if not credential or self._external_exchange_required(credential):
+      credential = self._get_auth_response()
+      # store fetched credential
+      if credential:
+        self._store_credential(credential)
+      else:
+        self._request_credential()
+        return AuthPreparationResult(
+            state="pending",
+            auth_scheme=self.auth_scheme,
+            auth_credential=self.auth_credential,
+        )
 
-    if exchanged_credential:
-      self._store_credential(exchanged_credential)
-      return AuthPreparationResult(
-          state="done",
-          auth_scheme=self.auth_scheme,
-          auth_credential=exchanged_credential,
-      )
-    else:
-      self._request_credential()
-      return AuthPreparationResult(
-          state="pending",
-          auth_scheme=self.auth_scheme,
-          auth_credential=self.auth_credential,
-      )
+    # here exchangers are doing two different thing:
+    # for service account the exchanger is doing actualy token exchange
+    # while for oauth2 it's actually doing the credentail conversion
+    # from OAuth2 credential to HTTP credentails for setting credential in
+    # http header
+    # TODO cleanup the logic:
+    # 1. service account token exchanger should happen before we store them in
+    #    the token store
+    # 2. blow line should only do credential conversion
+
+    exchanged_credential = self._exchange_credential(credential)
+    return AuthPreparationResult(
+        state="done",
+        auth_scheme=self.auth_scheme,
+        auth_credential=exchanged_credential,
+    )

google/adk/tools/retrieval/files_retrieval.py

@@ -14,11 +14,17 @@
 
 """Provides data for the agent."""
 
+from __future__ import annotations
+
+import logging
+
 from llama_index.core import SimpleDirectoryReader
 from llama_index.core import VectorStoreIndex
 
 from .llama_index_retrieval import LlamaIndexRetrieval
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
 class FilesRetrieval(LlamaIndexRetrieval):
 
@@ -26,7 +32,7 @@ class FilesRetrieval(LlamaIndexRetrieval):
 
     self.input_dir = input_dir
 
-    print(f'Loading data from {input_dir}')
+    logger.info("Loading data from %s", input_dir)
     retriever = VectorStoreIndex.from_documents(
         SimpleDirectoryReader(input_dir).load_data()
     ).as_retriever()

google/adk/tools/url_context_tool.py

@@ -0,0 +1,61 @@
+# 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 TYPE_CHECKING
+
+from google.genai import types
+from typing_extensions import override
+
+from .base_tool import BaseTool
+from .tool_context import ToolContext
+
+if TYPE_CHECKING:
+  from ..models import LlmRequest
+
+
+class UrlContextTool(BaseTool):
+  """A built-in tool that is automatically invoked by Gemini 2 models to retrieve content from the URLs and use that content to inform and shape its response.
+
+  This tool operates internally within the model and does not require or perform
+  local code execution.
+  """
+
+  def __init__(self):
+    # Name and description are not used because this is a model built-in tool.
+    super().__init__(name='url_context', description='url_context')
+
+  @override
+  async def process_llm_request(
+      self,
+      *,
+      tool_context: ToolContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    llm_request.config = llm_request.config or types.GenerateContentConfig()
+    llm_request.config.tools = llm_request.config.tools or []
+    if llm_request.model and 'gemini-1' in llm_request.model:
+      raise ValueError('Url context tool can not be used in Gemini 1.x.')
+    elif llm_request.model and 'gemini-2' in llm_request.model:
+      llm_request.config.tools.append(
+          types.Tool(url_context=types.UrlContext())
+      )
+    else:
+      raise ValueError(
+          f'Url context tool is not supported for model {llm_request.model}'
+      )
+
+
+url_context = UrlContextTool()

google/adk/tools/vertex_ai_search_tool.py

@@ -39,6 +39,9 @@ class VertexAiSearchTool(BaseTool):
       self,
       *,
       data_store_id: Optional[str] = None,
+      data_store_specs: Optional[
+          list[types.VertexAISearchDataStoreSpec]
+      ] = None,
       search_engine_id: Optional[str] = None,
       filter: Optional[str] = None,
       max_results: Optional[int] = None,
@@ -49,6 +52,8 @@ class VertexAiSearchTool(BaseTool):
       data_store_id: The Vertex AI search data store resource ID in the format
         of
         "projects/{project}/locations/{location}/collections/{collection}/dataStores/{dataStore}".
+      data_store_specs: Specifications that define the specific DataStores to be
+        searched. It should only be set if engine is used.
       search_engine_id: The Vertex AI search engine resource ID in the format of
         "projects/{project}/locations/{location}/collections/{collection}/engines/{engine}".
 
@@ -64,7 +69,12 @@ class VertexAiSearchTool(BaseTool):
       raise ValueError(
           'Either data_store_id or search_engine_id must be specified.'
       )
+    if data_store_specs is not None and search_engine_id is None:
+      raise ValueError(
+          'search_engine_id must be specified if data_store_specs is specified.'
+      )
     self.data_store_id = data_store_id
+    self.data_store_specs = data_store_specs
     self.search_engine_id = search_engine_id
     self.filter = filter
     self.max_results = max_results
@@ -76,8 +86,8 @@ class VertexAiSearchTool(BaseTool):
       tool_context: ToolContext,
       llm_request: LlmRequest,
   ) -> None:
-    if llm_request.model and llm_request.model.startswith('gemini-'):
-      if llm_request.model.startswith('gemini-1') and llm_request.config.tools:
+    if llm_request.model and 'gemini-' in llm_request.model:
+      if 'gemini-1' in llm_request.model and llm_request.config.tools:
         raise ValueError(
             'Vertex AI search tool can not be used with other tools in Gemini'
             ' 1.x.'
@@ -89,6 +99,7 @@ class VertexAiSearchTool(BaseTool):
               retrieval=types.Retrieval(
                   vertex_ai_search=types.VertexAISearch(
                       datastore=self.data_store_id,
+                      data_store_specs=self.data_store_specs,
                       engine=self.search_engine_id,
                       filter=self.filter,
                       max_results=self.max_results,

google/adk/utils/feature_decorator.py

@@ -0,0 +1,175 @@
+# 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 functools
+import os
+from typing import Callable
+from typing import cast
+from typing import Optional
+from typing import TypeVar
+from typing import Union
+import warnings
+
+from dotenv import load_dotenv
+
+T = TypeVar("T", bound=Union[Callable, type])
+
+
+def _make_feature_decorator(
+    *,
+    label: str,
+    default_message: str,
+    block_usage: bool = False,
+    bypass_env_var: Optional[str] = None,
+) -> Callable:
+  def decorator_factory(message_or_obj=None):
+    # Case 1: Used as @decorator without parentheses
+    # message_or_obj is the decorated class/function
+    if message_or_obj is not None and (
+        isinstance(message_or_obj, type) or callable(message_or_obj)
+    ):
+      return _create_decorator(
+          default_message, label, block_usage, bypass_env_var
+      )(message_or_obj)
+
+    # Case 2: Used as @decorator() with or without message
+    # message_or_obj is either None or a string message
+    message = (
+        message_or_obj if isinstance(message_or_obj, str) else default_message
+    )
+    return _create_decorator(message, label, block_usage, bypass_env_var)
+
+  return decorator_factory
+
+
+def _create_decorator(
+    message: str, label: str, block_usage: bool, bypass_env_var: Optional[str]
+) -> Callable[[T], T]:
+  def decorator(obj: T) -> T:
+    obj_name = getattr(obj, "__name__", type(obj).__name__)
+    msg = f"[{label.upper()}] {obj_name}: {message}"
+
+    if isinstance(obj, type):  # decorating a class
+      orig_init = obj.__init__
+
+      @functools.wraps(orig_init)
+      def new_init(self, *args, **kwargs):
+        # Load .env file if dotenv is available
+        load_dotenv()
+
+        # Check if usage should be bypassed via environment variable at call time
+        should_bypass = (
+            bypass_env_var is not None
+            and os.environ.get(bypass_env_var, "").lower() == "true"
+        )
+
+        if should_bypass:
+          # Bypass completely - no warning, no error
+          pass
+        elif block_usage:
+          raise RuntimeError(msg)
+        else:
+          warnings.warn(msg, category=UserWarning, stacklevel=2)
+        return orig_init(self, *args, **kwargs)
+
+      obj.__init__ = new_init  # type: ignore[attr-defined]
+      return cast(T, obj)
+
+    elif callable(obj):  # decorating a function or method
+
+      @functools.wraps(obj)
+      def wrapper(*args, **kwargs):
+        # Load .env file if dotenv is available
+        load_dotenv()
+
+        # Check if usage should be bypassed via environment variable at call time
+        should_bypass = (
+            bypass_env_var is not None
+            and os.environ.get(bypass_env_var, "").lower() == "true"
+        )
+
+        if should_bypass:
+          # Bypass completely - no warning, no error
+          pass
+        elif block_usage:
+          raise RuntimeError(msg)
+        else:
+          warnings.warn(msg, category=UserWarning, stacklevel=2)
+        return obj(*args, **kwargs)
+
+      return cast(T, wrapper)
+
+    else:
+      raise TypeError(
+          f"@{label} can only be applied to classes or callable objects"
+      )
+
+  return decorator
+
+
+working_in_progress = _make_feature_decorator(
+    label="WIP",
+    default_message=(
+        "This feature is a work in progress and is not working completely. ADK"
+        " users are not supposed to use it."
+    ),
+    block_usage=True,
+    bypass_env_var="ADK_ALLOW_WIP_FEATURES",
+)
+"""Mark a class or function as a work in progress.
+
+By default, decorated functions/classes will raise RuntimeError when used.
+Set ADK_ALLOW_WIP_FEATURES=true environment variable to bypass this restriction.
+ADK users are not supposed to set this environment variable.
+
+Sample usage:
+
+```
+@working_in_progress("This feature is not ready for production use.")
+def my_wip_function():
+  pass
+```
+"""
+
+experimental = _make_feature_decorator(
+    label="EXPERIMENTAL",
+    default_message=(
+        "This feature is experimental and may change or be removed in future"
+        " versions without notice. It may introduce breaking changes at any"
+        " time."
+    ),
+)
+"""Mark a class or a function as an experimental feature.
+
+Sample usage:
+
+```
+# Use with default message
+@experimental
+class ExperimentalClass:
+  pass
+
+# Use with custom message
+@experimental("This API may have breaking change in the future.")
+class CustomExperimentalClass:
+  pass
+
+# Use with empty parentheses (same as default message)
+@experimental()
+def experimental_function():
+  pass
+```
+"""

google/adk/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 # version: major.minor.patch
-__version__ = "1.2.1"
+__version__ = "1.4.1"

{google_adk-1.2.1.dist-info → google_adk-1.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-adk
-Version: 1.2.1
+Version: 1.4.1
 Summary: Agent Development Kit
 Author-email: Google LLC <googleapis-packages@google.com>
 Requires-Python: >=3.9
@@ -19,6 +19,7 @@ Classifier: Operating System :: OS Independent
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: License :: OSI Approved :: Apache Software License
 License-File: LICENSE
+Requires-Dist: anyio>=4.9.0;python_version>='3.10'
 Requires-Dist: authlib>=1.5.1
 Requires-Dist: click>=8.1.8
 Requires-Dist: fastapi>=0.115.0
@@ -34,12 +35,17 @@ Requires-Dist: opentelemetry-api>=1.31.0
 Requires-Dist: opentelemetry-exporter-gcp-trace>=1.9.0
 Requires-Dist: opentelemetry-sdk>=1.31.0
 Requires-Dist: pydantic>=2.0, <3.0.0
+Requires-Dist: python-dateutil>=2.9.0.post0
 Requires-Dist: python-dotenv>=1.0.0
 Requires-Dist: PyYAML>=6.0.2
+Requires-Dist: requests>=2.32.4
 Requires-Dist: sqlalchemy>=2.0
-Requires-Dist: tzlocal>=5.3
+Requires-Dist: starlette>=0.46.2
 Requires-Dist: typing-extensions>=4.5, <5
+Requires-Dist: tzlocal>=5.3
 Requires-Dist: uvicorn>=0.34.0
+Requires-Dist: websockets>=15.0.1
+Requires-Dist: a2a-sdk>=0.2.7 ; extra == "a2a" and (python_version>='3.10')
 Requires-Dist: flit>=3.10.0 ; extra == "dev"
 Requires-Dist: isort>=6.0.0 ; extra == "dev"
 Requires-Dist: pyink>=24.10.0 ; extra == "dev"
@@ -76,6 +82,7 @@ Project-URL: changelog, https://github.com/google/adk-python/blob/main/CHANGELOG
 Project-URL: documentation, https://google.github.io/adk-docs/
 Project-URL: homepage, https://google.github.io/adk-docs/
 Project-URL: repository, https://github.com/google/adk-python
+Provides-Extra: a2a
 Provides-Extra: dev
 Provides-Extra: docs
 Provides-Extra: eval
@@ -87,6 +94,7 @@ Provides-Extra: test
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
 [![Python Unit Tests](https://github.com/google/adk-python/actions/workflows/python-unit-tests.yml/badge.svg)](https://github.com/google/adk-python/actions/workflows/python-unit-tests.yml)
 [![r/agentdevelopmentkit](https://img.shields.io/badge/Reddit-r%2Fagentdevelopmentkit-FF4500?style=flat&logo=reddit&logoColor=white)](https://www.reddit.com/r/agentdevelopmentkit/)
+[![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/google/adk-python)
 
 <html>
     <h2 align="center">