gmakstutis/openrag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add Mike's UI enhancements for picker

Browse source
This commit is contained in:
Eric Hare 2025-09-04 09:27:18 -07:00
parent 64edbd8eed
commit 2dfc8faaac
No known key found for this signature in database
GPG key ID: A73DF73724270AB7
6 changed files with 438 additions and 154 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

117
frontend/src/app/connectors/GoogleDrivePicker.tsx
View file

@ -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<void>((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<void>((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 (
<div className="flex flex-col gap-2">
<div className="flex items-center gap-2">
<Button variant="secondary" size="sm" onClick={openPicker} disabled={loading}>
{loading ? "Loading Picker…" : buttonLabel}
</Button>
{(filesCount > 0 || foldersCount > 0) && (
<Badge variant="outline">{filesCount} file(s), {foldersCount} folder(s) selected</Badge>
)}
</div>
{(filesCount > 0 || foldersCount > 0) && (
<div className="flex flex-wrap gap-1 max-h-20 overflow-auto">
{value!.files.slice(0, 6).map((id) => <Badge key={id} variant="secondary">file:{id}</Badge>)}
{filesCount > 6 && <Badge>+{filesCount - 6} more</Badge>}
{value!.folders.slice(0, 6).map((id) => <Badge key={id} variant="secondary">folder:{id}</Badge>)}
{foldersCount > 6 && <Badge>+{foldersCount - 6} more</Badge>}
</div>
)}
</div>
)
}

109
frontend/src/app/settings/page.tsx
View file

@ -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<number>(10)
const [syncAllFiles, setSyncAllFiles] = useState<boolean>(false)
const [driveSelection, setDriveSelection] = useState<DriveSelection>({ 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() {
<CardContent className="flex-1 flex flex-col justify-end space-y-4">
{connector.status === "connected" ? (
<div className="space-y-3">
<div className="flex justify-center">
<GoogleDrivePicker value={driveSelection} onChange={setDriveSelection} />
</div>
{/* Google Drive file picker */}
{connector.type === "google_drive" && (
<GoogleDrivePicker
onFileSelected={(files) => handleFileSelection(connector.id, files)}
selectedFiles={selectedFiles[connector.id] || []}
isAuthenticated={connector.status === "connected"}
accessToken={connectorAccessTokens[connector.type]}
/>
)}
<Button
onClick={() => handleSync(connector)}
disabled={isSyncing === connector.id}
@ -453,7 +512,7 @@ function KnowledgeSourcesPage() {
) : (
<>
<RefreshCw className="mr-2 h-4 w-4" />
Sync Files
Sync Now
</>
)}
</Button>

269
frontend/src/components/google-drive-picker.tsx Normal file
View file

@ -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 <Folder className="h-4 w-4" />
}
return <FileText className="h-4 w-4" />
}
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 (
<div className="text-sm text-muted-foreground p-4 bg-muted/20 rounded-md">
Please connect to Google Drive first to select specific files.
</div>
)
}
return (
<div className="space-y-4">
<div className="flex items-center justify-between">
<div>
<h4 className="text-sm font-medium">File Selection</h4>
<p className="text-xs text-muted-foreground">
Choose specific files to sync instead of syncing everything
</p>
</div>
<Button
onClick={openPicker}
disabled={!isPickerLoaded || isPickerOpen || !accessToken}
size="sm"
variant="outline"
>
{isPickerOpen ? 'Opening Picker...' : 'Select Files'}
</Button>
</div>
{selectedFiles.length > 0 && (
<div className="space-y-2">
<p className="text-xs text-muted-foreground">
Selected files ({selectedFiles.length}):
</p>
<div className="max-h-32 overflow-y-auto space-y-1">
{selectedFiles.map((file) => (
<div
key={file.id}
className="flex items-center justify-between p-2 bg-muted/30 rounded-md text-xs"
>
<div className="flex items-center gap-2 flex-1 min-w-0">
{getFileIcon(file.mimeType)}
<span className="truncate font-medium">{file.name}</span>
<Badge variant="secondary" className="text-xs px-1 py-0.5 h-auto">
{getMimeTypeLabel(file.mimeType)}
</Badge>
</div>
<Button
onClick={() => removeFile(file.id)}
size="sm"
variant="ghost"
className="h-6 w-6 p-0"
>
<X className="h-3 w-3" />
</Button>
</div>
))}
</div>
<Button
onClick={() => onFileSelected([])}
size="sm"
variant="ghost"
className="text-xs h-6"
>
Clear all
</Button>
</div>
)}
</div>
)
}

44
src/api/connectors.py
View file

@ -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)

31
src/connectors/google_drive/connector.py
View file

@ -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:

22
src/main.py
View file

@ -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(