89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
from pathlib import Path
|
|
from typing import Annotated, Literal, Optional, Union
|
|
|
|
from .....doc_utils import export_module
|
|
from .....import_utils import optional_import_block
|
|
from .... import Toolkit, tool
|
|
from ..model import GoogleFileInfo
|
|
from ..toolkit_protocol import GoogleToolkitProtocol
|
|
from .drive_functions import download_file, list_files_and_folders
|
|
|
|
with optional_import_block():
|
|
from google.oauth2.credentials import Credentials
|
|
from googleapiclient.discovery import build
|
|
|
|
__all__ = [
|
|
"GoogleDriveToolkit",
|
|
]
|
|
|
|
|
|
@export_module("autogen.tools.experimental.google.drive")
|
|
class GoogleDriveToolkit(Toolkit, GoogleToolkitProtocol):
|
|
"""A tool map for Google Drive."""
|
|
|
|
def __init__( # type: ignore[no-any-unimported]
|
|
self,
|
|
*,
|
|
credentials: "Credentials",
|
|
download_folder: Union[Path, str],
|
|
exclude: Optional[list[Literal["list_drive_files_and_folders", "download_file_from_drive"]]] = None,
|
|
api_version: str = "v3",
|
|
) -> None:
|
|
"""Initialize the Google Drive tool map.
|
|
|
|
Args:
|
|
credentials: The Google OAuth2 credentials.
|
|
download_folder: The folder to download files to.
|
|
exclude: The tool names to exclude.
|
|
api_version: The Google Drive API version to use."
|
|
"""
|
|
self.service = build(serviceName="drive", version=api_version, credentials=credentials)
|
|
|
|
if isinstance(download_folder, str):
|
|
download_folder = Path(download_folder)
|
|
download_folder.mkdir(parents=True, exist_ok=True)
|
|
|
|
@tool(description="List files and folders in a Google Drive")
|
|
def list_drive_files_and_folders(
|
|
page_size: Annotated[int, "The number of files to list per page."] = 10,
|
|
folder_id: Annotated[
|
|
Optional[str],
|
|
"The ID of the folder to list files from. If not provided, lists all files in the root folder.",
|
|
] = None,
|
|
) -> list[GoogleFileInfo]:
|
|
return list_files_and_folders(service=self.service, page_size=page_size, folder_id=folder_id)
|
|
|
|
@tool(description="download a file from Google Drive")
|
|
def download_file_from_drive(
|
|
file_info: Annotated[GoogleFileInfo, "The file info to download."],
|
|
subfolder_path: Annotated[
|
|
Optional[str],
|
|
"The subfolder path to save the file in. If not provided, saves in the main download folder.",
|
|
] = None,
|
|
) -> str:
|
|
return download_file(
|
|
service=self.service,
|
|
file_id=file_info.id,
|
|
file_name=file_info.name,
|
|
mime_type=file_info.mime_type,
|
|
download_folder=download_folder,
|
|
subfolder_path=subfolder_path,
|
|
)
|
|
|
|
if exclude is None:
|
|
exclude = []
|
|
|
|
tools = [tool for tool in [list_drive_files_and_folders, download_file_from_drive] if tool.name not in exclude]
|
|
super().__init__(tools=tools)
|
|
|
|
@classmethod
|
|
def recommended_scopes(cls) -> list[str]:
|
|
"""Return the recommended scopes manatory for using tools from this tool map."""
|
|
return [
|
|
"https://www.googleapis.com/auth/drive.readonly",
|
|
]
|