class MyGoogleDriveToolkit(GoogleDriveToolkit):
def __init__(
self,
*,
credentials,
download_folder,
):
super().__init__(credentials=credentials, download_folder=download_folder)
# Define a custom tool
@tool(description="List documents in a folder")
def list_docs(
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]:
doc_mime_types = [
"application/vnd.google-apps.document", # Google Docs
"application/pdf", # PDFs
"application/msword", # MS Word
"application/vnd.openxmlformats-officedocument.wordprocessingml.document", # DOCX
]
mime_type_filter = " or ".join(f"mimeType='{mime}'" for mime in doc_mime_types)
query = f"({mime_type_filter}) and trashed=false"
if folder_id:
query = f"'{folder_id}' in parents and {query}"
kwargs = {
"pageSize": page_size,
"fields": "nextPageToken, files(id, name, mimeType)",
"q": query, # Apply filtering in the query itself
}
response = self.service.files().list(**kwargs).execute()
result = response.get("files", [])
if not isinstance(result, list):
raise ValueError(f"Expected a list of files, but got {result}")
return [GoogleFileInfo(**file_info) for file_info in result]
# Remove tool which you don't want to use
self.remove_tool("list_drive_files_and_folders")
# Add your custom tool
self.set_tool(list_docs)
