From c4e03eef6a2b23e9bde6a8136a708f0b33f7665c Mon Sep 17 00:00:00 2001 From: darshankparmar Date: Fri, 15 Aug 2025 22:30:28 +0530 Subject: [PATCH] feat: Add Google Drive toolkit with comprehensive file operations - Implement GoogleDriveTools class with OAuth2 authentication - Support file operations: list, search, upload, download, create folder, delete - Include pagination support for large file collections - Add comprehensive unit tests with mocked Google Drive API - Provide cookbook example demonstrating agent-driven Google Drive usage - Support both credentials.json and environment variable authentication - Handle soft/hard delete operations with proper error handling --- cookbook/tools/googledrive_tools.py | 35 ++ libs/agno/agno/tools/googledrive.py | 301 ++++++++++++++++++ .../unit/tools/test_googledrive_tools.py | 102 ++++++ 3 files changed, 438 insertions(+) create mode 100644 cookbook/tools/googledrive_tools.py create mode 100644 libs/agno/agno/tools/googledrive.py create mode 100644 libs/agno/tests/unit/tools/test_googledrive_tools.py diff --git a/cookbook/tools/googledrive_tools.py b/cookbook/tools/googledrive_tools.py new file mode 100644 index 0000000000..8f915133f6 --- /dev/null +++ b/cookbook/tools/googledrive_tools.py @@ -0,0 +1,35 @@ +""" +Google Drive Toolkit example for Agno (Agent-driven). + +Demonstrates: +1. Agent lists files whose name starts with "report" (max 10). + +Requires: +- credentials.json or env vars for OAuth +- Authorised redirect URI matching oauth_port in Google Cloud Console +Example: http://localhost:8080/flowName=GeneralOAuthFlow +""" + +from agno.agent import Agent +from agno.tools.googledrive import GoogleDriveTools + +# --- Initialize Google Drive Tools --- +google_drive_tools = GoogleDriveTools( + oauth_port=8080 # Change if needed +) + +# Create Agent with debug and monitoring +agent = Agent( + tools=[google_drive_tools], + show_tool_calls=True, + markdown=True, + instructions=[ + "You help users interact with Google Drive using the Google Drive API.", + "You can list, search, download, upload, and delete files as needed.", + ], +) + +# Run all tasks through Agent +agent.print_response(""" +List up to 10 files in my Google Drive whose name starts with 'report' +""") diff --git a/libs/agno/agno/tools/googledrive.py b/libs/agno/agno/tools/googledrive.py new file mode 100644 index 0000000000..82a0bd61ee --- /dev/null +++ b/libs/agno/agno/tools/googledrive.py @@ -0,0 +1,301 @@ +""" +Google Drive Toolkit for Agno + +Provides tools to interact with Google Drive: +- List files with pagination +- Search files +- Get file info +- Upload files +- Download files +- Create folders +- Delete files + +Authentication: +--------------- +Requires either a `credentials.json` file or environment variables: +- GOOGLE_CLIENT_ID +- GOOGLE_CLIENT_SECRET +- GOOGLE_PROJECT_ID +- GOOGLE_REDIRECT_URI (default: http://localhost) + +Google API References: +- https://developers.google.com/drive/api/v3/about-sdk +""" + +import os +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +from agno.tools import Toolkit +from agno.utils.log import logger + +try: + from google.auth.transport.requests import Request + from google.oauth2.credentials import Credentials + from google_auth_oauthlib.flow import InstalledAppFlow + from googleapiclient.discovery import Resource, build + from googleapiclient.errors import HttpError +except ImportError: + raise ImportError( + "`google-api-python-client` `google-auth-httplib2` `google-auth-oauthlib` not installed.\n" + "Install with: pip install google-api-python-client google-auth-httplib2 google-auth-oauthlib" + ) + + +class GoogleDriveTools(Toolkit): + """A toolkit for interacting with Google Drive files and folders.""" + + DEFAULT_SCOPES = { + "read": "https://www.googleapis.com/auth/drive.readonly", + "write": "https://www.googleapis.com/auth/drive.file", + "full": "https://www.googleapis.com/auth/drive", + "metadata": "https://www.googleapis.com/auth/drive.metadata.readonly", + } + + service: Optional[Resource] + + def __init__( + self, + creds_path: Optional[str] = None, + token_path: Optional[str] = None, + scopes: Optional[List[str]] = None, + oauth_port: int = 0, + **kwargs, + ): + """Initialize GoogleDriveTools.""" + self.creds = None + self.credentials_path = creds_path + self.token_path = token_path + self.oauth_port = oauth_port + self.scopes = scopes or [ + self.DEFAULT_SCOPES["metadata"], + self.DEFAULT_SCOPES["read"], + self.DEFAULT_SCOPES["write"], + ] + self.service: Optional[Resource] = None + + tools: List[Callable[..., Any]] = [ + self.list_files, + self.get_file_info, + self.search_files, + self.upload_file, + self.download_file, + self.create_folder, + self.delete_file, + ] + super().__init__(name="google_drive_tools", tools=tools, **kwargs) + + def _auth(self): + """Authenticate with Google Drive API.""" + token_file = Path(self.token_path or "token.json") + creds_file = Path(self.credentials_path or "credentials.json") + + if token_file.exists(): + self.creds = Credentials.from_authorized_user_file(str(token_file), self.scopes) + + if not self.creds or not self.creds.valid: + if self.creds and self.creds.expired and self.creds.refresh_token: + logger.info("Refreshing Google Drive credentials...") + self.creds.refresh(Request()) + else: + client_config = { + "installed": { + "client_id": os.getenv("GOOGLE_CLIENT_ID"), + "client_secret": os.getenv("GOOGLE_CLIENT_SECRET"), + "project_id": os.getenv("GOOGLE_PROJECT_ID"), + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "redirect_uris": [os.getenv("GOOGLE_REDIRECT_URI", "http://localhost")], + } + } + if creds_file.exists(): + flow = InstalledAppFlow.from_client_secrets_file(str(creds_file), self.scopes) + else: + flow = InstalledAppFlow.from_client_config(client_config, self.scopes) + self.creds = flow.run_local_server(port=self.oauth_port) + + if self.creds: + token_file.write_text(self.creds.to_json()) + + self.service = build("drive", "v3", credentials=self.creds) + + def _ensure_service(self): + if not self.service: + self._auth() + + # --------------------------- + # Core Operations + # --------------------------- + + def list_files( + self, folder_id: Optional[str] = None, page_size: int = 10, page_token: Optional[str] = None + ) -> Dict[str, Any]: + """List files/folders with optional pagination.""" + self._ensure_service() + + assert self.service is not None + + try: + query = None + if folder_id is None: + query = "'root' in parents and trashed=false" + else: + query = f"'{folder_id}' in parents and trashed=false" + + results = ( + self.service.files() + .list( + q=query, + pageSize=page_size, + pageToken=page_token, + fields="nextPageToken, files(id,name,mimeType,size,createdTime,modifiedTime,parents,webViewLink)", + orderBy="name", + ) + .execute() + ) + + return { + "files": results.get("files", []), + "nextPageToken": results.get("nextPageToken"), + } + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def search_files(self, query: str, page_size: int = 10, page_token: Optional[str] = None) -> Dict[str, Any]: + """Search for files by query string.""" + self._ensure_service() + + assert self.service is not None + + try: + results = ( + self.service.files() + .list( + q=f"{query} and trashed=false", + pageSize=page_size, + pageToken=page_token, + fields="nextPageToken, files(id,name,mimeType,size,createdTime,modifiedTime,parents,webViewLink)", + orderBy="modifiedTime desc", + ) + .execute() + ) + + return { + "files": results.get("files", []), + "nextPageToken": results.get("nextPageToken"), + } + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def get_file_info(self, file_id: str) -> Dict[str, Any]: + """Get detailed information for a file.""" + self._ensure_service() + + assert self.service is not None + + try: + return ( + self.service.files() + .get( + fileId=file_id, + fields="id,name,mimeType,size,createdTime,modifiedTime,parents,webViewLink,description,owners,permissions,capabilities", + ) + .execute() + ) + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def upload_file(self, local_path: str, folder_id: Optional[str] = None) -> Dict[str, Any]: + """Upload a local file to Google Drive.""" + from googleapiclient.http import MediaFileUpload + + self._ensure_service() + + assert self.service is not None + + try: + file_metadata: Dict[str, Any] = {"name": Path(local_path).name} + if folder_id: + file_metadata["parents"] = [folder_id] + + media = MediaFileUpload(local_path, resumable=True) + return ( + self.service.files() + .create(body=file_metadata, media_body=media, fields="id, name, webViewLink") + .execute() + ) + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def download_file(self, file_id: str, destination_path: str) -> Dict[str, Any]: + """Download a file from Google Drive to local path.""" + import io + + from googleapiclient.http import MediaIoBaseDownload + + self._ensure_service() + + assert self.service is not None + + try: + request = self.service.files().get_media(fileId=file_id) + fh = io.FileIO(destination_path, "wb") + downloader = MediaIoBaseDownload(fh, request) + done = False + while not done: + status, done = downloader.next_chunk() + logger.info(f"Download progress: {int(status.progress() * 100)}%") + return {"message": f"File downloaded to {destination_path}"} + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def create_folder(self, name: str, parent_id: Optional[str] = None) -> Dict[str, Any]: + """Create a folder in Google Drive.""" + self._ensure_service() + + assert self.service is not None + + try: + file_metadata: Dict[str, Any] = {"name": name, "mimeType": "application/vnd.google-apps.folder"} + if parent_id: + file_metadata["parents"] = [parent_id] + + return self.service.files().create(body=file_metadata, fields="id, name, webViewLink").execute() + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} + + def delete_file(self, file_id: str, permanent: bool = False) -> Dict[str, Any]: + """ + Delete a file from Google Drive. + + Args: + file_id (str): The ID of the file to delete. + permanent (bool): If True, permanently deletes the file (bypasses trash). + If False (default), moves the file to Trash (soft delete). + + Returns: + dict: Message or error. + """ + self._ensure_service() + + assert self.service is not None + + try: + if permanent: + # Hard delete – file is gone forever + self.service.files().delete(fileId=file_id).execute() + return {"message": f"File {file_id} permanently deleted"} + else: + # Soft delete – move to trash + self.service.files().update(fileId=file_id, body={"trashed": True}).execute() + return {"message": f"File {file_id} moved to Trash"} + except HttpError as e: + logger.error(f"Drive API error: {e}") + return {"error": str(e)} diff --git a/libs/agno/tests/unit/tools/test_googledrive_tools.py b/libs/agno/tests/unit/tools/test_googledrive_tools.py new file mode 100644 index 0000000000..513d085373 --- /dev/null +++ b/libs/agno/tests/unit/tools/test_googledrive_tools.py @@ -0,0 +1,102 @@ +"""Unit tests for GoogleDriveTools class.""" + +from unittest.mock import MagicMock, Mock, patch + +import pytest +from google.oauth2.credentials import Credentials + +from agno.tools.googledrive import GoogleDriveTools + + +@pytest.fixture +def mock_credentials(): + """Mock Google OAuth2 credentials.""" + mock_creds = Mock(spec=Credentials) + mock_creds.valid = True + mock_creds.expired = False + return mock_creds + + +@pytest.fixture +def mock_drive_service(): + """Mock Google Drive API service.""" + return MagicMock() + + +@pytest.fixture +def drive_tools(mock_credentials, mock_drive_service): + """Create GoogleDriveTools instance with mocked service.""" + with patch("agno.tools.googledrive.build") as mock_build: + mock_build.return_value = mock_drive_service + tools = GoogleDriveTools(creds_path=None, token_path=None) + tools.creds = mock_credentials # inject mock creds + tools.service = mock_drive_service + return tools + + +def test_list_files(drive_tools, mock_drive_service): + """Test listing files.""" + mock_drive_service.files().list().execute.return_value = {"files": [{"id": "1", "name": "test.txt"}]} + result = drive_tools.list_files() + assert result["files"][0]["name"] == "test.txt" + + +def test_search_files(drive_tools, mock_drive_service): + """Test searching files.""" + mock_drive_service.files().list().execute.return_value = {"files": [{"id": "1", "name": "report.doc"}]} + result = drive_tools.search_files(query="name contains 'report'") + assert any("report" in f["name"] for f in result["files"]) + + +def test_get_file_info(drive_tools, mock_drive_service): + """Test retrieving file info.""" + mock_drive_service.files().get().execute.return_value = {"id": "1", "name": "info.txt"} + result = drive_tools.get_file_info(file_id="1") + assert result["name"] == "info.txt" + + +def test_upload_file(drive_tools, mock_drive_service, tmp_path): + """Test uploading a file.""" + # Create temp file + file_path = tmp_path / "upload.txt" + file_path.write_text("content") + + mock_drive_service.files().create().execute.return_value = {"id": "123", "name": "upload.txt"} + result = drive_tools.upload_file(local_path=str(file_path)) + assert result["name"] == "upload.txt" + + +def test_download_file(drive_tools, mock_drive_service, tmp_path): + """Test downloading a file.""" + mock_request = MagicMock() + mock_drive_service.files().get_media.return_value = mock_request + + # Patch MediaIoBaseDownload to simulate progress + with patch("googleapiclient.http.MediaIoBaseDownload") as mock_downloader_cls: + mock_downloader = MagicMock() + mock_downloader.next_chunk.side_effect = [(MagicMock(progress=lambda: 1.0), True)] + mock_downloader_cls.return_value = mock_downloader + + result = drive_tools.download_file(file_id="1", destination_path=str(tmp_path / "dl.txt")) + assert "downloaded" in result["message"] + + +def test_create_folder(drive_tools, mock_drive_service): + """Test creating a folder.""" + mock_drive_service.files().create().execute.return_value = {"id": "f123", "name": "NewFolder"} + result = drive_tools.create_folder("NewFolder") + assert result["name"] == "NewFolder" + + +def test_delete_file_soft(drive_tools, mock_drive_service): + """Test moving file to trash.""" + result = drive_tools.delete_file(file_id="abc123", permanent=False) + mock_drive_service.files().update.assert_called_once() + assert "Trash" in result["message"] + + +def test_delete_file_permanent(drive_tools, mock_drive_service): + """Test hard delete of file.""" + result = drive_tools.delete_file(file_id="abc123", permanent=True) + mock_drive_service.files().delete.assert_called_once() + assert "permanently" in result["message"]