mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-07-18 23:11:12 +02:00
feat: enhance Composio Google Drive integration with folder and file selection
- Added a new endpoint to list folders and files in a user's Composio Google Drive, supporting hierarchical structure. - Implemented UI components for selecting specific folders and files to index, improving user control over indexing options. - Introduced indexing options for maximum files per folder and inclusion of subfolders, allowing for customizable indexing behavior. - Enhanced error handling and logging for Composio Drive operations, ensuring better visibility into issues during file retrieval and indexing. - Updated the Composio configuration component to reflect new selection capabilities and indexing options.
This commit is contained in:
Anish Sarkar
parent
e6a4ac7c9c
commit
7ec7ed5c3b
11 changed files with 1069 additions and 24 deletions
|
|
@ -8,6 +8,7 @@ Endpoints:
|
|||
- GET /composio/toolkits - List available Composio toolkits
|
||||
- GET /auth/composio/connector/add - Initiate OAuth for a specific toolkit
|
||||
- GET /auth/composio/connector/callback - Handle OAuth callback
|
||||
- GET /connectors/{connector_id}/composio-drive/folders - List folders/files for Composio Google Drive
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -369,3 +370,124 @@ async def composio_callback(
|
|||
raise HTTPException(
|
||||
status_code=500, detail=f"Failed to complete Composio OAuth: {e!s}"
|
||||
) from e
|
||||
|
||||
|
||||
@router.get("/connectors/{connector_id}/composio-drive/folders")
|
||||
async def list_composio_drive_folders(
|
||||
connector_id: int,
|
||||
parent_id: str | None = None,
|
||||
session: AsyncSession = Depends(get_async_session),
|
||||
user: User = Depends(current_active_user),
|
||||
):
|
||||
"""
|
||||
List folders AND files in user's Google Drive via Composio with hierarchical support.
|
||||
|
||||
This is called at index time from the manage connector page to display
|
||||
the complete file system (folders and files). Only folders are selectable.
|
||||
|
||||
Args:
|
||||
connector_id: ID of the Composio Google Drive connector
|
||||
parent_id: Optional parent folder ID to list contents (None for root)
|
||||
|
||||
Returns:
|
||||
JSON with list of items: {
|
||||
"items": [
|
||||
{"id": str, "name": str, "mimeType": str, "isFolder": bool, ...},
|
||||
...
|
||||
]
|
||||
}
|
||||
"""
|
||||
if not ComposioService.is_enabled():
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Composio integration is not enabled.",
|
||||
)
|
||||
|
||||
try:
|
||||
# Get connector and verify ownership
|
||||
result = await session.execute(
|
||||
select(SearchSourceConnector).filter(
|
||||
SearchSourceConnector.id == connector_id,
|
||||
SearchSourceConnector.user_id == user.id,
|
||||
SearchSourceConnector.connector_type
|
||||
== SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
|
||||
)
|
||||
)
|
||||
connector = result.scalars().first()
|
||||
|
||||
if not connector:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Composio Google Drive connector not found or access denied",
|
||||
)
|
||||
|
||||
# Get Composio connected account ID from config
|
||||
composio_connected_account_id = connector.config.get("composio_connected_account_id")
|
||||
if not composio_connected_account_id:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Composio connected account not found. Please reconnect the connector.",
|
||||
)
|
||||
|
||||
# Initialize Composio service and fetch files
|
||||
service = ComposioService()
|
||||
entity_id = f"surfsense_{user.id}"
|
||||
|
||||
# Fetch files/folders from Composio Google Drive
|
||||
files, next_token, error = await service.get_drive_files(
|
||||
connected_account_id=composio_connected_account_id,
|
||||
entity_id=entity_id,
|
||||
folder_id=parent_id,
|
||||
page_size=100,
|
||||
)
|
||||
|
||||
if error:
|
||||
logger.error(f"Failed to list Composio Drive files: {error}")
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Failed to list folder contents: {error}"
|
||||
)
|
||||
|
||||
# Transform files to match the expected format with isFolder field
|
||||
items = []
|
||||
for file_info in files:
|
||||
file_id = file_info.get("id", "") or file_info.get("fileId", "")
|
||||
file_name = file_info.get("name", "") or file_info.get("fileName", "") or "Untitled"
|
||||
mime_type = file_info.get("mimeType", "") or file_info.get("mime_type", "")
|
||||
|
||||
if not file_id:
|
||||
continue
|
||||
|
||||
is_folder = mime_type == "application/vnd.google-apps.folder"
|
||||
|
||||
items.append({
|
||||
"id": file_id,
|
||||
"name": file_name,
|
||||
"mimeType": mime_type,
|
||||
"isFolder": is_folder,
|
||||
"parents": file_info.get("parents", []),
|
||||
"size": file_info.get("size"),
|
||||
"iconLink": file_info.get("iconLink"),
|
||||
})
|
||||
|
||||
# Sort: folders first, then files, both alphabetically
|
||||
folders = sorted([item for item in items if item["isFolder"]], key=lambda x: x["name"].lower())
|
||||
files_list = sorted([item for item in items if not item["isFolder"]], key=lambda x: x["name"].lower())
|
||||
items = folders + files_list
|
||||
|
||||
folder_count = len(folders)
|
||||
file_count = len(files_list)
|
||||
|
||||
logger.info(
|
||||
f"✅ Listed {len(items)} total items ({folder_count} folders, {file_count} files) for Composio connector {connector_id}"
|
||||
+ (f" in folder {parent_id}" if parent_id else " in ROOT")
|
||||
)
|
||||
|
||||
return {"items": items}
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Error listing Composio Drive contents: {e!s}", exc_info=True)
|
||||
raise HTTPException(
|
||||
status_code=500, detail=f"Failed to list Drive contents: {e!s}"
|
||||
) from e
|
||||
|
|
|
|||
|
|
@ -897,8 +897,46 @@ async def index_connector_content(
|
|||
)
|
||||
response_message = "Web page indexing started in the background."
|
||||
|
||||
elif connector.connector_type == SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR:
|
||||
from app.tasks.celery_tasks.connector_tasks import (
|
||||
index_composio_connector_task,
|
||||
)
|
||||
|
||||
# For Composio Google Drive, if drive_items is provided, update connector config
|
||||
# This allows the UI to pass folder/file selection like the regular Google Drive connector
|
||||
if drive_items and drive_items.has_items():
|
||||
# Update connector config with the selected folders/files
|
||||
config = connector.config or {}
|
||||
config["selected_folders"] = [{"id": f.id, "name": f.name} for f in drive_items.folders]
|
||||
config["selected_files"] = [{"id": f.id, "name": f.name} for f in drive_items.files]
|
||||
if drive_items.indexing_options:
|
||||
config["indexing_options"] = {
|
||||
"max_files_per_folder": drive_items.indexing_options.max_files_per_folder,
|
||||
"incremental_sync": drive_items.indexing_options.incremental_sync,
|
||||
"include_subfolders": drive_items.indexing_options.include_subfolders,
|
||||
}
|
||||
connector.config = config
|
||||
from sqlalchemy.orm.attributes import flag_modified
|
||||
flag_modified(connector, "config")
|
||||
await session.commit()
|
||||
await session.refresh(connector)
|
||||
|
||||
logger.info(
|
||||
f"Triggering Composio Google Drive indexing for connector {connector_id} into search space {search_space_id}, "
|
||||
f"folders: {len(drive_items.folders)}, files: {len(drive_items.files)}"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
f"Triggering Composio Google Drive indexing for connector {connector_id} into search space {search_space_id} "
|
||||
f"using existing config (from {indexing_from} to {indexing_to})"
|
||||
)
|
||||
|
||||
index_composio_connector_task.delay(
|
||||
connector_id, search_space_id, str(user.id), indexing_from, indexing_to
|
||||
)
|
||||
response_message = "Composio Google Drive indexing started in the background."
|
||||
|
||||
elif connector.connector_type in [
|
||||
SearchSourceConnectorType.COMPOSIO_GOOGLE_DRIVE_CONNECTOR,
|
||||
SearchSourceConnectorType.COMPOSIO_GMAIL_CONNECTOR,
|
||||
SearchSourceConnectorType.COMPOSIO_GOOGLE_CALENDAR_CONNECTOR,
|
||||
]:
|
||||
|
|
|
|||
|
|
@ -397,7 +397,11 @@ class ComposioService:
|
|||
"page_size": min(page_size, 100),
|
||||
}
|
||||
if folder_id:
|
||||
params["folder_id"] = folder_id
|
||||
# List contents of a specific folder (exclude shortcuts - we don't have access to them)
|
||||
params["q"] = f"'{folder_id}' in parents and trashed = false and mimeType != 'application/vnd.google-apps.shortcut'"
|
||||
else:
|
||||
# List root-level items only (My Drive root), exclude shortcuts
|
||||
params["q"] = "'root' in parents and trashed = false and mimeType != 'application/vnd.google-apps.shortcut'"
|
||||
if page_token:
|
||||
params["page_token"] = page_token
|
||||
|
||||
|
|
|
|||
|
|
@ -252,37 +252,123 @@ async def _index_composio_google_drive(
|
|||
update_last_indexed: bool = True,
|
||||
max_items: int = 1000,
|
||||
) -> tuple[int, str]:
|
||||
"""Index Google Drive files via Composio."""
|
||||
"""Index Google Drive files via Composio.
|
||||
|
||||
Supports folder/file selection via connector config:
|
||||
- selected_folders: List of {id, name} for folders to index
|
||||
- selected_files: List of {id, name} for individual files to index
|
||||
- indexing_options: {max_files_per_folder, incremental_sync, include_subfolders}
|
||||
"""
|
||||
try:
|
||||
composio_connector = ComposioConnector(session, connector_id)
|
||||
connector_config = await composio_connector.get_config()
|
||||
|
||||
# Get folder/file selection configuration
|
||||
selected_folders = connector_config.get("selected_folders", [])
|
||||
selected_files = connector_config.get("selected_files", [])
|
||||
indexing_options = connector_config.get("indexing_options", {})
|
||||
|
||||
max_files_per_folder = indexing_options.get("max_files_per_folder", 100)
|
||||
include_subfolders = indexing_options.get("include_subfolders", True)
|
||||
|
||||
await task_logger.log_task_progress(
|
||||
log_entry,
|
||||
f"Fetching Google Drive files via Composio for connector {connector_id}",
|
||||
{"stage": "fetching_files"},
|
||||
{"stage": "fetching_files", "selected_folders": len(selected_folders), "selected_files": len(selected_files)},
|
||||
)
|
||||
|
||||
# Fetch files
|
||||
all_files = []
|
||||
page_token = None
|
||||
|
||||
while len(all_files) < max_items:
|
||||
files, next_token, error = await composio_connector.list_drive_files(
|
||||
page_token=page_token,
|
||||
page_size=min(100, max_items - len(all_files)),
|
||||
)
|
||||
# If specific folders/files are selected, fetch from those
|
||||
if selected_folders or selected_files:
|
||||
# Fetch files from selected folders
|
||||
for folder in selected_folders:
|
||||
folder_id = folder.get("id")
|
||||
folder_name = folder.get("name", "Unknown")
|
||||
|
||||
if not folder_id:
|
||||
continue
|
||||
|
||||
# Handle special case for "root" folder
|
||||
actual_folder_id = None if folder_id == "root" else folder_id
|
||||
|
||||
logger.info(f"Fetching files from folder: {folder_name} ({folder_id})")
|
||||
|
||||
# Fetch files from this folder
|
||||
folder_files = []
|
||||
page_token = None
|
||||
|
||||
while len(folder_files) < max_files_per_folder:
|
||||
files, next_token, error = await composio_connector.list_drive_files(
|
||||
folder_id=actual_folder_id,
|
||||
page_token=page_token,
|
||||
page_size=min(100, max_files_per_folder - len(folder_files)),
|
||||
)
|
||||
|
||||
if error:
|
||||
await task_logger.log_task_failure(
|
||||
log_entry, f"Failed to fetch Drive files: {error}", {}
|
||||
if error:
|
||||
logger.warning(f"Failed to fetch files from folder {folder_name}: {error}")
|
||||
break
|
||||
|
||||
# Process files
|
||||
for file_info in files:
|
||||
mime_type = file_info.get("mimeType", "") or file_info.get("mime_type", "")
|
||||
|
||||
# If it's a folder and include_subfolders is enabled, recursively fetch
|
||||
if mime_type == "application/vnd.google-apps.folder":
|
||||
if include_subfolders:
|
||||
# Add subfolder files recursively
|
||||
subfolder_files = await _fetch_folder_files_recursively(
|
||||
composio_connector,
|
||||
file_info.get("id"),
|
||||
max_files=max_files_per_folder,
|
||||
current_count=len(folder_files),
|
||||
)
|
||||
folder_files.extend(subfolder_files)
|
||||
else:
|
||||
folder_files.append(file_info)
|
||||
|
||||
if not next_token:
|
||||
break
|
||||
page_token = next_token
|
||||
|
||||
all_files.extend(folder_files[:max_files_per_folder])
|
||||
logger.info(f"Found {len(folder_files)} files in folder {folder_name}")
|
||||
|
||||
# Add specifically selected files
|
||||
for selected_file in selected_files:
|
||||
file_id = selected_file.get("id")
|
||||
file_name = selected_file.get("name", "Unknown")
|
||||
|
||||
if not file_id:
|
||||
continue
|
||||
|
||||
# Add file info (we'll fetch content later during indexing)
|
||||
all_files.append({
|
||||
"id": file_id,
|
||||
"name": file_name,
|
||||
"mimeType": "", # Will be determined later
|
||||
})
|
||||
else:
|
||||
# No selection specified - fetch all files (original behavior)
|
||||
page_token = None
|
||||
|
||||
while len(all_files) < max_items:
|
||||
files, next_token, error = await composio_connector.list_drive_files(
|
||||
page_token=page_token,
|
||||
page_size=min(100, max_items - len(all_files)),
|
||||
)
|
||||
return 0, f"Failed to fetch Drive files: {error}"
|
||||
|
||||
all_files.extend(files)
|
||||
if error:
|
||||
await task_logger.log_task_failure(
|
||||
log_entry, f"Failed to fetch Drive files: {error}", {}
|
||||
)
|
||||
return 0, f"Failed to fetch Drive files: {error}"
|
||||
|
||||
if not next_token:
|
||||
break
|
||||
page_token = next_token
|
||||
all_files.extend(files)
|
||||
|
||||
if not next_token:
|
||||
break
|
||||
page_token = next_token
|
||||
|
||||
if not all_files:
|
||||
success_msg = "No Google Drive files found"
|
||||
|
|
@ -479,6 +565,81 @@ async def _index_composio_google_drive(
|
|||
return 0, f"Failed to index Google Drive via Composio: {e!s}"
|
||||
|
||||
|
||||
async def _fetch_folder_files_recursively(
|
||||
composio_connector: ComposioConnector,
|
||||
folder_id: str,
|
||||
max_files: int = 100,
|
||||
current_count: int = 0,
|
||||
depth: int = 0,
|
||||
max_depth: int = 10,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Recursively fetch files from a Google Drive folder via Composio.
|
||||
|
||||
Args:
|
||||
composio_connector: The Composio connector instance
|
||||
folder_id: Google Drive folder ID
|
||||
max_files: Maximum number of files to fetch
|
||||
current_count: Current number of files already fetched
|
||||
depth: Current recursion depth
|
||||
max_depth: Maximum recursion depth to prevent infinite loops
|
||||
|
||||
Returns:
|
||||
List of file info dictionaries
|
||||
"""
|
||||
if depth >= max_depth:
|
||||
logger.warning(f"Max recursion depth reached for folder {folder_id}")
|
||||
return []
|
||||
|
||||
if current_count >= max_files:
|
||||
return []
|
||||
|
||||
all_files = []
|
||||
page_token = None
|
||||
|
||||
try:
|
||||
while len(all_files) + current_count < max_files:
|
||||
files, next_token, error = await composio_connector.list_drive_files(
|
||||
folder_id=folder_id,
|
||||
page_token=page_token,
|
||||
page_size=min(100, max_files - len(all_files) - current_count),
|
||||
)
|
||||
|
||||
if error:
|
||||
logger.warning(f"Error fetching files from subfolder {folder_id}: {error}")
|
||||
break
|
||||
|
||||
for file_info in files:
|
||||
mime_type = file_info.get("mimeType", "") or file_info.get("mime_type", "")
|
||||
|
||||
if mime_type == "application/vnd.google-apps.folder":
|
||||
# Recursively fetch from subfolders
|
||||
subfolder_files = await _fetch_folder_files_recursively(
|
||||
composio_connector,
|
||||
file_info.get("id"),
|
||||
max_files=max_files,
|
||||
current_count=current_count + len(all_files),
|
||||
depth=depth + 1,
|
||||
max_depth=max_depth,
|
||||
)
|
||||
all_files.extend(subfolder_files)
|
||||
else:
|
||||
all_files.append(file_info)
|
||||
|
||||
if len(all_files) + current_count >= max_files:
|
||||
break
|
||||
|
||||
if not next_token:
|
||||
break
|
||||
page_token = next_token
|
||||
|
||||
return all_files[:max_files - current_count]
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error in recursive folder fetch: {e!s}")
|
||||
return all_files
|
||||
|
||||
|
||||
async def _process_gmail_message_batch(
|
||||
session: AsyncSession,
|
||||
messages: list[dict[str, Any]],
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue