diff --git a/frontend/src/app/connectors/GoogleDrivePicker.tsx b/frontend/src/app/connectors/GoogleDrivePicker.tsx deleted file mode 100644 index 7723ca1e..00000000 --- a/frontend/src/app/connectors/GoogleDrivePicker.tsx +++ /dev/null @@ -1,117 +0,0 @@ -"use client" - -import { useCallback, useState } from "react" -import { Button } from "@/components/ui/button" -import { Badge } from "@/components/ui/badge" - -// declare globals to silence TS -declare global { - interface Window { google?: any; gapi?: any } -} - -const loadScript = (src: string) => - new Promise((resolve, reject) => { - if (document.querySelector(`script[src="${src}"]`)) return resolve() - const s = document.createElement("script") - s.src = src - s.async = true - s.onload = () => resolve() - s.onerror = () => reject(new Error(`Failed to load ${src}`)) - document.head.appendChild(s) - }) - -export type DriveSelection = { files: string[]; folders: string[] } - -export function GoogleDrivePicker({ - value, - onChange, - buttonLabel = "Choose in Drive", -}: { - value?: DriveSelection - onChange: (sel: DriveSelection) => void - buttonLabel?: string -}) { - const [loading, setLoading] = useState(false) - - const ensureGoogleApis = useCallback(async () => { - await loadScript("https://accounts.google.com/gsi/client") - await loadScript("https://apis.google.com/js/api.js") - await new Promise((res) => window.gapi?.load("picker", () => res())) - }, []) - - const openPicker = useCallback(async () => { - const clientId = process.env.NEXT_PUBLIC_GOOGLE_CLIENT_ID - const apiKey = process.env.NEXT_PUBLIC_GOOGLE_API_KEY - if (!clientId || !apiKey) { - alert("Google Picker requires NEXT_PUBLIC_GOOGLE_CLIENT_ID and NEXT_PUBLIC_GOOGLE_API_KEY") - return - } - try { - setLoading(true) - await ensureGoogleApis() - const tokenClient = window.google.accounts.oauth2.initTokenClient({ - client_id: clientId, - scope: "https://www.googleapis.com/auth/drive.readonly https://www.googleapis.com/auth/drive.metadata.readonly", - callback: (tokenResp: any) => { - const viewDocs = new window.google.picker.DocsView() - .setIncludeFolders(true) - .setSelectFolderEnabled(true) - - console.log("Picker using clientId:", clientId, "apiKey:", apiKey) - - const picker = new window.google.picker.PickerBuilder() - .enableFeature(window.google.picker.Feature.MULTISELECT_ENABLED) - .setOAuthToken(tokenResp.access_token) - .setDeveloperKey(apiKey) - .addView(viewDocs) - .setCallback((data: any) => { - if (data.action === window.google.picker.Action.PICKED) { - const pickedFiles: string[] = [] - const pickedFolders: string[] = [] - for (const doc of data.docs || []) { - const id = doc.id - const isFolder = doc?.type === "folder" || doc?.mimeType === "application/vnd.google-apps.folder" - if (isFolder) pickedFolders.push(id) - else pickedFiles.push(id) - } - onChange({ files: pickedFiles, folders: pickedFolders }) - } - }) - .build() - picker.setVisible(true) - }, - }) - tokenClient.requestAccessToken() - } catch (e) { - console.error("Drive Picker error", e) - alert("Failed to open Google Drive Picker. See console.") - } finally { - setLoading(false) - } - }, [ensureGoogleApis, onChange]) - - const filesCount = value?.files?.length ?? 0 - const foldersCount = value?.folders?.length ?? 0 - - return ( -
-
- - {(filesCount > 0 || foldersCount > 0) && ( - {filesCount} file(s), {foldersCount} folder(s) selected - )} -
- - {(filesCount > 0 || foldersCount > 0) && ( -
- {value!.files.slice(0, 6).map((id) => file:{id})} - {filesCount > 6 && +{filesCount - 6} more} - {value!.folders.slice(0, 6).map((id) => folder:{id})} - {foldersCount > 6 && +{foldersCount - 6} more} -
- )} -
- ) -} diff --git a/frontend/src/app/settings/page.tsx b/frontend/src/app/settings/page.tsx index cbc17449..a1eea2a7 100644 --- a/frontend/src/app/settings/page.tsx +++ b/frontend/src/app/settings/page.tsx @@ -12,9 +12,17 @@ import { Loader2, PlugZap, RefreshCw } from "lucide-react" import { ProtectedRoute } from "@/components/protected-route" import { useTask } from "@/contexts/task-context" import { useAuth } from "@/contexts/auth-context" -import { GoogleDrivePicker, type DriveSelection } from "../connectors/GoogleDrivePicker" +import { GoogleDrivePicker } from "@/components/google-drive-picker" +interface GoogleDriveFile { + id: string + name: string + mimeType: string + webViewLink?: string + iconLink?: string +} + interface Connector { id: string name: string @@ -24,6 +32,7 @@ interface Connector { type: string connectionId?: string access_token?: string + selectedFiles?: GoogleDriveFile[] } interface SyncResult { @@ -54,7 +63,8 @@ function KnowledgeSourcesPage() { const [syncResults, setSyncResults] = useState<{[key: string]: SyncResult | null}>({}) const [maxFiles, setMaxFiles] = useState(10) const [syncAllFiles, setSyncAllFiles] = useState(false) - const [driveSelection, setDriveSelection] = useState({ files: [], folders: [] }) + const [selectedFiles, setSelectedFiles] = useState<{[connectorId: string]: GoogleDriveFile[]}>({}) + const [connectorAccessTokens, setConnectorAccessTokens] = useState<{[connectorId: string]: string}>({}) // Settings state // Note: backend internal Langflow URL is not needed on the frontend @@ -145,6 +155,24 @@ function KnowledgeSourcesPage() { const activeConnection = connections.find((conn: Connection) => conn.is_active) const isConnected = activeConnection !== undefined + // For Google Drive, try to get access token for the picker + if (connectorType === 'google_drive' && activeConnection) { + try { + const tokenResponse = await fetch(`/api/connectors/${connectorType}/token?connection_id=${activeConnection.connection_id}`) + if (tokenResponse.ok) { + const tokenData = await tokenResponse.json() + if (tokenData.access_token) { + setConnectorAccessTokens(prev => ({ + ...prev, + [connectorType]: tokenData.access_token + })) + } + } + } catch (e) { + console.log('Could not fetch access token for Google Drive picker:', e) + } + } + setConnectors(prev => prev.map(c => c.type === connectorType ? { @@ -210,47 +238,71 @@ function KnowledgeSourcesPage() { } } + const handleFileSelection = (connectorId: string, files: GoogleDriveFile[]) => { + setSelectedFiles(prev => ({ + ...prev, + [connectorId]: files + })) + + // Update the connector with selected files + setConnectors(prev => prev.map(c => + c.id === connectorId + ? { ...c, selectedFiles: files } + : c + )) + } + const handleSync = async (connector: Connector) => { if (!connector.connectionId) return - + setIsSyncing(connector.id) setSyncResults(prev => ({ ...prev, [connector.id]: null })) - + try { - const body: any = { + const syncBody: { + connection_id: string; + max_files?: number; + selected_files?: string[]; + } = { connection_id: connector.connectionId, - max_files: syncAllFiles ? 0 : (maxFiles || undefined), + max_files: syncAllFiles ? 0 : (maxFiles || undefined) } - - if (connector.type === "google-drive") { - body.file_ids = driveSelection.files - body.folder_ids = driveSelection.folders - body.recursive = true // or expose a checkbox if you want + + // Add selected files for Google Drive + if (connector.type === "google_drive" && selectedFiles[connector.id]?.length > 0) { + syncBody.selected_files = selectedFiles[connector.id].map(file => file.id) } - + const response = await fetch(`/api/connectors/${connector.type}/sync`, { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(body), + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(syncBody), }) - + const result = await response.json() + if (response.status === 201) { const taskId = result.task_id if (taskId) { addTask(taskId) - setSyncResults(prev => ({ - ...prev, - [connector.id]: { processed: 0, total: result.total_files || 0 } + setSyncResults(prev => ({ + ...prev, + [connector.id]: { + processed: 0, + total: result.total_files || 0 + } })) } } else if (response.ok) { setSyncResults(prev => ({ ...prev, [connector.id]: result })) + // Note: Stats will auto-refresh via task completion watcher for async syncs } else { - console.error("Sync failed:", result.error) + console.error('Sync failed:', result.error) } } catch (error) { - console.error("Sync error:", error) + console.error('Sync error:', error) } finally { setIsSyncing(null) } @@ -436,9 +488,16 @@ function KnowledgeSourcesPage() { {connector.status === "connected" ? (
-
- -
+ {/* Google Drive file picker */} + {connector.type === "google_drive" && ( + handleFileSelection(connector.id, files)} + selectedFiles={selectedFiles[connector.id] || []} + isAuthenticated={connector.status === "connected"} + accessToken={connectorAccessTokens[connector.type]} + /> + )} + diff --git a/frontend/src/components/google-drive-picker.tsx b/frontend/src/components/google-drive-picker.tsx new file mode 100644 index 00000000..e3a9555b --- /dev/null +++ b/frontend/src/components/google-drive-picker.tsx @@ -0,0 +1,269 @@ +"use client" + +import { useState, useEffect } from "react" +import { Button } from "@/components/ui/button" +import { Badge } from "@/components/ui/badge" +import { FileText, Folder, X } from "lucide-react" + +interface GoogleDrivePickerProps { + onFileSelected: (files: GoogleDriveFile[]) => void + selectedFiles?: GoogleDriveFile[] + isAuthenticated: boolean + accessToken?: string +} + +interface GoogleDriveFile { + id: string + name: string + mimeType: string + webViewLink?: string + iconLink?: string +} + +interface GoogleAPI { + load: (api: string, options: { callback: () => void; onerror?: () => void }) => void +} + +interface GooglePickerData { + action: string + docs: GooglePickerDocument[] +} + +interface GooglePickerDocument { + [key: string]: string +} + +declare global { + interface Window { + gapi: GoogleAPI + google: { + picker: { + api: { + load: (callback: () => void) => void + } + PickerBuilder: new () => GooglePickerBuilder + ViewId: { + DOCS: string + FOLDERS: string + DOCS_IMAGES_AND_VIDEOS: string + DOCUMENTS: string + PRESENTATIONS: string + SPREADSHEETS: string + } + Feature: { + MULTISELECT_ENABLED: string + NAV_HIDDEN: string + SIMPLE_UPLOAD_ENABLED: string + } + Action: { + PICKED: string + CANCEL: string + } + Document: { + ID: string + NAME: string + MIME_TYPE: string + URL: string + ICON_URL: string + } + } + } + } +} + +interface GooglePickerBuilder { + addView: (view: string) => GooglePickerBuilder + setOAuthToken: (token: string) => GooglePickerBuilder + setCallback: (callback: (data: GooglePickerData) => void) => GooglePickerBuilder + enableFeature: (feature: string) => GooglePickerBuilder + setTitle: (title: string) => GooglePickerBuilder + build: () => GooglePicker +} + +interface GooglePicker { + setVisible: (visible: boolean) => void +} + +export function GoogleDrivePicker({ + onFileSelected, + selectedFiles = [], + isAuthenticated, + accessToken +}: GoogleDrivePickerProps) { + const [isPickerLoaded, setIsPickerLoaded] = useState(false) + const [isPickerOpen, setIsPickerOpen] = useState(false) + + useEffect(() => { + const loadPickerApi = () => { + if (typeof window !== 'undefined' && window.gapi) { + window.gapi.load('picker', { + callback: () => { + setIsPickerLoaded(true) + }, + onerror: () => { + console.error('Failed to load Google Picker API') + } + }) + } + } + + // Load Google API script if not already loaded + if (typeof window !== 'undefined') { + if (!window.gapi) { + const script = document.createElement('script') + script.src = 'https://apis.google.com/js/api.js' + script.async = true + script.defer = true + script.onload = loadPickerApi + script.onerror = () => { + console.error('Failed to load Google API script') + } + document.head.appendChild(script) + + return () => { + if (document.head.contains(script)) { + document.head.removeChild(script) + } + } + } else { + loadPickerApi() + } + } + }, []) + + const openPicker = () => { + if (!isPickerLoaded || !accessToken || !window.google?.picker) { + return + } + + try { + setIsPickerOpen(true) + + const picker = new window.google.picker.PickerBuilder() + .addView(window.google.picker.ViewId.DOCS) + .addView(window.google.picker.ViewId.FOLDERS) + .setOAuthToken(accessToken) + .enableFeature(window.google.picker.Feature.MULTISELECT_ENABLED) + .setTitle('Select files from Google Drive') + .setCallback(pickerCallback) + .build() + + picker.setVisible(true) + } catch (error) { + console.error('Error creating picker:', error) + setIsPickerOpen(false) + } + } + + const pickerCallback = (data: GooglePickerData) => { + if (data.action === window.google.picker.Action.PICKED) { + const files: GoogleDriveFile[] = data.docs.map((doc: GooglePickerDocument) => ({ + id: doc[window.google.picker.Document.ID], + name: doc[window.google.picker.Document.NAME], + mimeType: doc[window.google.picker.Document.MIME_TYPE], + webViewLink: doc[window.google.picker.Document.URL], + iconLink: doc[window.google.picker.Document.ICON_URL] + })) + + onFileSelected(files) + } + + setIsPickerOpen(false) + } + + const removeFile = (fileId: string) => { + const updatedFiles = selectedFiles.filter(file => file.id !== fileId) + onFileSelected(updatedFiles) + } + + const getFileIcon = (mimeType: string) => { + if (mimeType.includes('folder')) { + return + } + return + } + + const getMimeTypeLabel = (mimeType: string) => { + const typeMap: { [key: string]: string } = { + 'application/vnd.google-apps.document': 'Google Doc', + 'application/vnd.google-apps.spreadsheet': 'Google Sheet', + 'application/vnd.google-apps.presentation': 'Google Slides', + 'application/vnd.google-apps.folder': 'Folder', + 'application/pdf': 'PDF', + 'text/plain': 'Text', + 'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'Word Doc', + 'application/vnd.openxmlformats-officedocument.presentationml.presentation': 'PowerPoint' + } + + return typeMap[mimeType] || 'Document' + } + + if (!isAuthenticated) { + return ( +
+ Please connect to Google Drive first to select specific files. +
+ ) + } + + return ( +
+
+
+

File Selection

+

+ Choose specific files to sync instead of syncing everything +

+
+ +
+ + {selectedFiles.length > 0 && ( +
+

+ Selected files ({selectedFiles.length}): +

+
+ {selectedFiles.map((file) => ( +
+
+ {getFileIcon(file.mimeType)} + {file.name} + + {getMimeTypeLabel(file.mimeType)} + +
+ +
+ ))} +
+ +
+ )} +
+ ) +} diff --git a/src/api/connectors.py b/src/api/connectors.py index 87f21b4b..43b9400d 100644 --- a/src/api/connectors.py +++ b/src/api/connectors.py @@ -111,6 +111,8 @@ async def connector_status(request: Request, connector_service, session_manager) async def connector_webhook(request: Request, connector_service, session_manager): """Handle webhook notifications from any connector type""" connector_type = request.path_params.get("connector_type") + if connector_type is None: + connector_type = "unknown" # Handle webhook validation (connector-specific) temp_config = {"token_file": "temp.json"} @@ -118,7 +120,7 @@ async def connector_webhook(request: Request, connector_service, session_manager temp_connection = ConnectionConfig( connection_id="temp", - connector_type=connector_type, + connector_type=str(connector_type), name="temp", config=temp_config, ) @@ -186,7 +188,6 @@ async def connector_webhook(request: Request, connector_service, session_manager ) # Process webhook for the specific connection - results = [] try: # Get the connector instance connector = await connector_service._get_connector(connection.connection_id) @@ -272,3 +273,42 @@ async def connector_webhook(request: Request, connector_service, session_manager return JSONResponse( {"error": f"Webhook processing failed: {str(e)}"}, status_code=500 ) + +async def connector_token(request: Request, connector_service, session_manager): + """Get access token for connector API calls (e.g., Google Picker)""" + connector_type = request.path_params.get("connector_type") + connection_id = request.query_params.get("connection_id") + + if not connection_id: + return JSONResponse({"error": "connection_id is required"}, status_code=400) + + user = request.state.user + + try: + # Get the connection and verify it belongs to the user + connection = await connector_service.connection_manager.get_connection(connection_id) + if not connection or connection.user_id != user.user_id: + return JSONResponse({"error": "Connection not found"}, status_code=404) + + # Get the connector instance + connector = await connector_service._get_connector(connection_id) + if not connector: + return JSONResponse({"error": "Connector not available"}, status_code=404) + + # For Google Drive, get the access token + if connector_type == "google_drive" and hasattr(connector, 'oauth'): + await connector.oauth.load_credentials() + if connector.oauth.creds and connector.oauth.creds.valid: + return JSONResponse({ + "access_token": connector.oauth.creds.token, + "expires_in": (connector.oauth.creds.expiry.timestamp() - + __import__('time').time()) if connector.oauth.creds.expiry else None + }) + else: + return JSONResponse({"error": "Invalid or expired credentials"}, status_code=401) + + return JSONResponse({"error": "Token not available for this connector type"}, status_code=400) + + except Exception as e: + print(f"Error getting connector token: {e}") + return JSONResponse({"error": str(e)}, status_code=500) diff --git a/src/connectors/google_drive/connector.py b/src/connectors/google_drive/connector.py index 5a8099e0..b392f8cc 100644 --- a/src/connectors/google_drive/connector.py +++ b/src/connectors/google_drive/connector.py @@ -69,6 +69,10 @@ class GoogleDriveConnector(BaseConnector): CLIENT_ID_ENV_VAR: str = "GOOGLE_OAUTH_CLIENT_ID" CLIENT_SECRET_ENV_VAR: str = "GOOGLE_OAUTH_CLIENT_SECRET" + # Supported alias keys coming from various frontends / pickers + _FILE_ID_ALIASES = ("file_ids", "selected_file_ids", "selected_files") + _FOLDER_ID_ALIASES = ("folder_ids", "selected_folder_ids", "selected_folders") + def log(self, message: str) -> None: print(message) @@ -106,12 +110,24 @@ class GoogleDriveConnector(BaseConnector): f"Provide config['client_secret'] or set {self.CLIENT_SECRET_ENV_VAR}." ) + # Normalize incoming IDs from any of the supported alias keys + def _first_present_list(cfg: Dict[str, Any], keys: Iterable[str]) -> Optional[List[str]]: + for k in keys: + v = cfg.get(k) + if v: # accept non-empty list + return list(v) + return None + + normalized_file_ids = _first_present_list(config, self._FILE_ID_ALIASES) + normalized_folder_ids = _first_present_list(config, self._FOLDER_ID_ALIASES) + self.cfg = GoogleDriveConfig( client_id=client_id, client_secret=client_secret, token_file=token_file, - file_ids=config.get("file_ids") or config.get("selected_file_ids"), - folder_ids=config.get("folder_ids") or config.get("selected_folder_ids"), + # Accept "selected_files" and "selected_folders" used by the Drive Picker flow + file_ids=normalized_file_ids, + folder_ids=normalized_folder_ids, recursive=bool(config.get("recursive", True)), drive_id=config.get("drive_id"), corpora=config.get("corpora"), @@ -417,7 +433,11 @@ class GoogleDriveConnector(BaseConnector): self.log(f"GoogleDriveConnector.authenticate failed: {e}") return False - async def list_files(self, page_token: Optional[str] = None, **kwargs) -> Dict[str, Any]: + async def list_files( + self, + page_token: Optional[str] = None, + **kwargs + ) -> Dict[str, Any]: """ List files in the currently selected scope (file_ids/folder_ids/recursive). Returns a dict with 'files' and 'next_page_token'. @@ -429,6 +449,11 @@ class GoogleDriveConnector(BaseConnector): try: items = self._iter_selected_items() + # Optionally honor a request-scoped max_files (e.g., from your API payload) + max_files = kwargs.get("max_files") + if isinstance(max_files, int) and max_files > 0: + items = items[:max_files] + # Simplest: ignore page_token and just dump all # If you want real pagination, slice items here if page_token: diff --git a/src/main.py b/src/main.py index 611d9885..ebe89149 100644 --- a/src/main.py +++ b/src/main.py @@ -6,18 +6,12 @@ import subprocess from functools import partial from starlette.applications import Starlette from starlette.routing import Route - -# Set multiprocessing start method to 'spawn' for CUDA compatibility -multiprocessing.set_start_method("spawn", force=True) - -# Create process pool FIRST, before any torch/CUDA imports from utils.process_pool import process_pool import torch # Configuration and setup from config.settings import clients, INDEX_NAME, INDEX_BODY, SESSION_SECRET -from utils.gpu_detection import detect_gpu_devices # Services from services.document_service import DocumentService @@ -46,6 +40,9 @@ from api import ( settings, ) +# Set multiprocessing start method to 'spawn' for CUDA compatibility +multiprocessing.set_start_method("spawn", force=True) + print("CUDA available:", torch.cuda.is_available()) print("CUDA version PyTorch was built with:", torch.version.cuda) @@ -240,7 +237,7 @@ async def initialize_services(): except Exception as e: print(f"[WARNING] Failed to load persisted connections on startup: {e}") else: - print(f"[CONNECTORS] Skipping connection loading in no-auth mode") + print("[CONNECTORS] Skipping connection loading in no-auth mode") return { "document_service": document_service, @@ -586,6 +583,17 @@ async def create_app(): ), methods=["GET"], ), + Route( + "/connectors/{connector_type}/token", + require_auth(services["session_manager"])( + partial( + connectors.connector_token, + connector_service=services["connector_service"], + session_manager=services["session_manager"], + ) + ), + methods=["GET"], + ), Route( "/connectors/{connector_type}/webhook", partial(