diff --git a/src/google/adk/tools/google_api_tool/__init__.py b/src/google/adk/tools/google_api_tool/__init__.py index 736d6d7..0abbf8f 100644 --- a/src/google/adk/tools/google_api_tool/__init__.py +++ b/src/google/adk/tools/google_api_tool/__init__.py @@ -12,75 +12,24 @@ # See the License for the specific language governing permissions and # limitations under the License. __all__ = [ - 'bigquery_toolset', - 'calendar_toolset', - 'gmail_toolset', - 'youtube_toolset', - 'slides_toolset', - 'sheets_toolset', - 'docs_toolset', + 'BigQueryToolset', + 'CalendarToolset', + 'GmailToolset', + 'YoutubeToolset', + 'SlidesToolset', + 'SheetsToolset', + 'DocsToolset', + 'GoogleApiToolset', + 'GoogleApiTool', ] -# Nothing is imported here automatically -# Each tool set will only be imported when accessed -_bigquery_toolset = None -_calendar_toolset = None -_gmail_toolset = None -_youtube_toolset = None -_slides_toolset = None -_sheets_toolset = None -_docs_toolset = None - - -def __getattr__(name): - global _bigquery_toolset, _calendar_toolset, _gmail_toolset, _youtube_toolset, _slides_toolset, _sheets_toolset, _docs_toolset - - if name == 'bigquery_toolset': - if _bigquery_toolset is None: - from .google_api_toolsets import bigquery_toolset as bigquery - - _bigquery_toolset = bigquery - return _bigquery_toolset - - if name == 'calendar_toolset': - if _calendar_toolset is None: - from .google_api_toolsets import calendar_toolset as calendar - - _calendar_toolset = calendar - return _calendar_toolset - - if name == 'gmail_toolset': - if _gmail_toolset is None: - from .google_api_toolsets import gmail_toolset as gmail - - _gmail_toolset = gmail - return _gmail_toolset - - if name == 'youtube_toolset': - if _youtube_toolset is None: - from .google_api_toolsets import youtube_toolset as youtube - - _youtube_toolset = youtube - return _youtube_toolset - - if name == 'slides_toolset': - if _slides_toolset is None: - from .google_api_toolsets import slides_toolset as slides - - _slides_toolset = slides - return _slides_toolset - - if name == 'sheets_toolset': - if _sheets_toolset is None: - from .google_api_toolsets import sheets_toolset as sheets - - _sheets_toolset = sheets - return _sheets_toolset - - if name == 'docs_toolset': - if _docs_toolset is None: - from .google_api_toolsets import docs_toolset as docs - - _docs_toolset = docs - return _docs_toolset +from .google_api_tool import GoogleApiTool +from .google_api_toolset import GoogleApiToolset +from .google_api_toolsets import BigQueryToolset +from .google_api_toolsets import CalendarToolset +from .google_api_toolsets import DocsToolset +from .google_api_toolsets import GmailToolset +from .google_api_toolsets import SheetsToolset +from .google_api_toolsets import SlidesToolset +from .google_api_toolsets import YoutubeToolset diff --git a/src/google/adk/tools/google_api_tool/google_api_tool.py b/src/google/adk/tools/google_api_tool/google_api_tool.py index 84e16c3..33cd1d8 100644 --- a/src/google/adk/tools/google_api_tool/google_api_tool.py +++ b/src/google/adk/tools/google_api_tool/google_api_tool.py @@ -29,13 +29,19 @@ from ..tool_context import ToolContext class GoogleApiTool(BaseTool): - def __init__(self, rest_api_tool: RestApiTool): + def __init__( + self, + rest_api_tool: RestApiTool, + client_id: Optional[str] = None, + client_secret: Optional[str] = None, + ): super().__init__( name=rest_api_tool.name, description=rest_api_tool.description, is_long_running=rest_api_tool.is_long_running, ) self._rest_api_tool = rest_api_tool + self.configure_auth(client_id, client_secret) @override def _get_declaration(self) -> FunctionDeclaration: diff --git a/src/google/adk/tools/google_api_tool/google_api_toolset.py b/src/google/adk/tools/google_api_tool/google_api_toolset.py index 5a92651..7131ebc 100644 --- a/src/google/adk/tools/google_api_tool/google_api_toolset.py +++ b/src/google/adk/tools/google_api_tool/google_api_toolset.py @@ -36,22 +36,39 @@ from .googleapi_to_openapi_converter import GoogleApiToOpenApiConverter class GoogleApiToolset(BaseToolset): """Google API Toolset contains tools for interacting with Google APIs. - Usually one toolsets will contains tools only replated to one Google API, e.g. + Usually one toolsets will contains tools only related to one Google API, e.g. Google Bigquery API toolset will contains tools only related to Google Bigquery API, like list dataset tool, list table tool etc. """ def __init__( self, - openapi_toolset: OpenAPIToolset, + api_name: str, + api_version: str, client_id: Optional[str] = None, client_secret: Optional[str] = None, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, ): - self._openapi_toolset = openapi_toolset - self.tool_filter = tool_filter + self.api_name = api_name + self.api_version = api_version self._client_id = client_id self._client_secret = client_secret + self._openapi_toolset = self._load_toolset_with_oidc_auth() + self.tool_filter = tool_filter + + def _is_tool_selected( + self, tool: GoogleApiTool, readonly_context: ReadonlyContext + ) -> bool: + if not self.tool_filter: + return True + + if isinstance(self.tool_filter, ToolPredicate): + return self.tool_filter(tool, readonly_context) + + if isinstance(self.tool_filter, list): + return tool.name in self.tool_filter + + return False @override async def get_tools( @@ -60,44 +77,26 @@ class GoogleApiToolset(BaseToolset): """Get all tools in the toolset.""" tools = [] - for tool in await self._openapi_toolset.get_tools(readonly_context): - if self.tool_filter and ( - isinstance(self.tool_filter, ToolPredicate) - and not self.tool_filter(tool, readonly_context) - or isinstance(self.tool_filter, list) - and tool.name not in self.tool_filter - ): - continue - google_api_tool = GoogleApiTool(tool) - google_api_tool.configure_auth(self._client_id, self._client_secret) - tools.append(google_api_tool) - - return tools + return [ + GoogleApiTool(tool, self._client_id, self._client_secret) + for tool in await self._openapi_toolset.get_tools(readonly_context) + if self._is_tool_selected(tool, readonly_context) + ] def set_tool_filter(self, tool_filter: Union[ToolPredicate, List[str]]): self.tool_filter = tool_filter - @staticmethod - def _load_toolset_with_oidc_auth( - spec_file: Optional[str] = None, - spec_dict: Optional[dict[str, Any]] = None, - scopes: Optional[list[str]] = None, - ) -> OpenAPIToolset: - spec_str = None - if spec_file: - # Get the frame of the caller - caller_frame = inspect.stack()[1] - # Get the filename of the caller - caller_filename = caller_frame.filename - # Get the directory of the caller - caller_dir = os.path.dirname(os.path.abspath(caller_filename)) - # Join the directory path with the filename - yaml_path = os.path.join(caller_dir, spec_file) - with open(yaml_path, 'r', encoding='utf-8') as file: - spec_str = file.read() - toolset = OpenAPIToolset( + def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset: + spec_dict = GoogleApiToOpenApiConverter( + self.api_name, self.api_version + ).convert() + scope = list( + spec_dict['components']['securitySchemes']['oauth2']['flows'][ + 'authorizationCode' + ]['scopes'].keys() + )[0] + return OpenAPIToolset( spec_dict=spec_dict, - spec_str=spec_str, spec_str_type='yaml', auth_scheme=OpenIdConnectWithConfig( authorization_endpoint=( @@ -113,31 +112,14 @@ class GoogleApiToolset(BaseToolset): 'client_secret_basic', ], grant_types_supported=['authorization_code'], - scopes=scopes, + scopes=[scope], ), ) - return toolset def configure_auth(self, client_id: str, client_secret: str): self._client_id = client_id self._client_secret = client_secret - @classmethod - def load_toolset( - cls: Type[GoogleApiToolset], - api_name: str, - api_version: str, - ) -> GoogleApiToolset: - spec_dict = GoogleApiToOpenApiConverter(api_name, api_version).convert() - scope = list( - spec_dict['components']['securitySchemes']['oauth2']['flows'][ - 'authorizationCode' - ]['scopes'].keys() - )[0] - return cls( - cls._load_toolset_with_oidc_auth(spec_dict=spec_dict, scopes=[scope]) - ) - @override async def close(self): if self._openapi_toolset: diff --git a/src/google/adk/tools/google_api_tool/google_api_toolsets.py b/src/google/adk/tools/google_api_tool/google_api_toolsets.py index 4e56f50..4d61f4f 100644 --- a/src/google/adk/tools/google_api_tool/google_api_toolsets.py +++ b/src/google/adk/tools/google_api_tool/google_api_toolsets.py @@ -14,98 +14,88 @@ import logging +from typing import List, Optional, Union + +from google.adk.tools.base_toolset import ToolPredicate from .google_api_toolset import GoogleApiToolset + logger = logging.getLogger(__name__) -_bigquery_toolset = None -_calendar_toolset = None -_gmail_toolset = None -_youtube_toolset = None -_slides_toolset = None -_sheets_toolset = None -_docs_toolset = None + +class BigQueryToolset(GoogleApiToolset): + + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("bigquery", "v2", client_id, client_secret, tool_filter) -def __getattr__(name): - """This method dynamically loads and returns GoogleApiToolSet instances for +class CalendarToolset(GoogleApiToolset): - various Google APIs. It uses a lazy loading approach, initializing each - tool set only when it is first requested. This avoids unnecessary loading - of tool sets that are not used in a given session. + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("calendar", "v3", client_id, client_secret, tool_filter) - Args: - name (str): The name of the tool set to retrieve (e.g., - "bigquery_toolset"). - Returns: - GoogleApiToolSet: The requested tool set instance. +class GmailToolset(GoogleApiToolset): - Raises: - AttributeError: If the requested tool set name is not recognized. - """ - global _bigquery_toolset, _calendar_toolset, _gmail_toolset, _youtube_toolset, _slides_toolset, _sheets_toolset, _docs_toolset + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("gmail", "v1", client_id, client_secret, tool_filter) - if name == "bigquery_toolset": - if _bigquery_toolset is None: - _bigquery_toolset = GoogleApiToolset.load_toolset( - api_name="bigquery", - api_version="v2", - ) - return _bigquery_toolset +class YoutubeToolset(GoogleApiToolset): - if name == "calendar_toolset": - if _calendar_toolset is None: - _calendar_toolset = GoogleApiToolset.load_toolset( - api_name="calendar", - api_version="v3", - ) + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("youtube", "v3", client_id, client_secret, tool_filter) - return _calendar_toolset - if name == "gmail_toolset": - if _gmail_toolset is None: - _gmail_toolset = GoogleApiToolset.load_toolset( - api_name="gmail", - api_version="v1", - ) +class SlidesToolset(GoogleApiToolset): - return _gmail_toolset + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("slides", "v1", client_id, client_secret, tool_filter) - if name == "youtube_toolset": - if _youtube_toolset is None: - _youtube_toolset = GoogleApiToolset.load_toolset( - api_name="youtube", - api_version="v3", - ) - return _youtube_toolset +class SheetsToolset(GoogleApiToolset): - if name == "slides_toolset": - if _slides_toolset is None: - _slides_toolset = GoogleApiToolset.load_toolset( - api_name="slides", - api_version="v1", - ) + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("sheets", "v4", client_id, client_secret, tool_filter) - return _slides_toolset - if name == "sheets_toolset": - if _sheets_toolset is None: - _sheets_toolset = GoogleApiToolset.load_toolset( - api_name="sheets", - api_version="v4", - ) +class DocsToolset(GoogleApiToolset): - return _sheets_toolset - - if name == "docs_toolset": - if _docs_toolset is None: - _docs_toolset = GoogleApiToolset.load_toolset( - api_name="docs", - api_version="v1", - ) - - return _docs_toolset + def __init__( + self, + client_id: str = None, + client_secret: str = None, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + ): + super().__init__("docs", "v1", client_id, client_secret, tool_filter)