UNPKG

mcp-gdrive-enhanced-markov

Version:

MCP server for Google Drive and Sheets with write capabilities, AI-powered PDF analysis, SQLite document indexing, and advanced directory exploration tools

214 lines (213 loc) 8.43 kB
import { google } from 'googleapis'; import { getValidCredentials } from '../auth.js'; import { db } from './database.js'; import { aiParser } from './ai_parser.js'; export class DriveIndexer { progressCallback; constructor(progressCallback) { this.progressCallback = progressCallback; } updateProgress(progress) { if (this.progressCallback) { this.progressCallback(progress); } } async indexFile(fileId) { try { const authClient = await getValidCredentials(); const drive = google.drive({ version: 'v3', auth: authClient }); // Get file metadata const metadataResponse = await drive.files.get({ fileId, fields: 'id,name,mimeType,modifiedTime,size,parents' }); const file = metadataResponse.data; // Skip non-PDF files for now if (file.mimeType !== 'application/pdf') { console.log(`Skipping non-PDF file: ${file.name} (${file.mimeType})`); return -1; } // Get file path by traversing parents const filePath = await this.getFilePath(drive, fileId, file.name || 'Unknown'); // Download file content const response = await drive.files.get({ fileId, alt: 'media' }, { responseType: 'arraybuffer' }); const fileContent = Buffer.from(response.data); // Check file size limit (20MB) const MAX_SIZE = 20 * 1024 * 1024; if (fileContent.length > MAX_SIZE) { throw new Error(`File too large: ${fileContent.length} bytes (max: ${MAX_SIZE} bytes)`); } // Parse with AI const aiResult = await aiParser.parseWithRetry(fileContent, file.name || 'Unknown', file.mimeType || 'application/pdf'); // Prepare document record const doc = { drive_id: file.id, drive_path: filePath, name: file.name || 'Unknown', mime_type: file.mimeType || 'application/pdf', doc_type: aiResult.doc_type, synopsis: aiResult.synopsis, ai_json: aiResult.fields, ai_model: 'gemini-2.5-flash-preview-05-20', ai_ts_utc: new Date().toISOString(), modified_time: file.modifiedTime || undefined, size: file.size ? parseInt(file.size) : undefined }; // Upsert into database const docId = db.upsertDocument(doc); console.log(`Indexed document: ${file.name} (ID: ${docId})`); return docId; } catch (error) { console.error(`Error indexing file ${fileId}:`, error); throw error; } } async indexFolder(folderId, recurse = true) { const authClient = await getValidCredentials(); const drive = google.drive({ version: 'v3', auth: authClient }); const result = { indexed: 0, skipped: 0, errors: 0, errorDetails: [] }; const progress = { total: 0, processed: 0, errors: 0 }; // First, count total files await this.countFiles(drive, folderId, recurse, progress); this.updateProgress(progress); // Then process files await this.processFolder(drive, folderId, recurse, result, progress); return result; } async countFiles(drive, folderId, recurse, progress) { let pageToken; do { const response = await drive.files.list({ q: `'${folderId}' in parents and trashed = false`, fields: 'nextPageToken, files(id, mimeType)', pageSize: 100, pageToken }); for (const file of response.data.files || []) { if (file.mimeType === 'application/pdf') { progress.total++; } else if (recurse && file.mimeType === 'application/vnd.google-apps.folder') { await this.countFiles(drive, file.id, recurse, progress); } } pageToken = response.data.nextPageToken; } while (pageToken); } async processFolder(drive, folderId, recurse, result, progress) { let pageToken; do { const response = await drive.files.list({ q: `'${folderId}' in parents and trashed = false`, fields: 'nextPageToken, files(id, name, mimeType)', pageSize: 20, // Smaller batch size for rate limiting pageToken }); for (const file of response.data.files || []) { if (file.mimeType === 'application/pdf') { progress.currentFile = file.name || 'Unknown'; this.updateProgress(progress); try { await this.indexFile(file.id); result.indexed++; progress.processed++; } catch (error) { result.errors++; progress.errors++; result.errorDetails.push({ fileId: file.id, fileName: file.name || 'Unknown', error: error instanceof Error ? error.message : 'Unknown error' }); } this.updateProgress(progress); // Rate limiting - wait 100ms between files await new Promise(resolve => setTimeout(resolve, 100)); } else if (file.mimeType === 'application/vnd.google-apps.folder') { if (recurse) { await this.processFolder(drive, file.id, recurse, result, progress); } } else { result.skipped++; } } pageToken = response.data.nextPageToken; } while (pageToken); } async reindexChanged(since) { const checkDate = since || new Date(Date.now() - 24 * 60 * 60 * 1000); // Default: last 24 hours const documents = db.getDocumentsNeedingReindex(checkDate); const result = { indexed: 0, skipped: 0, errors: 0, errorDetails: [] }; const progress = { total: documents.length, processed: 0, errors: 0 }; this.updateProgress(progress); for (const doc of documents) { progress.currentFile = doc.name; this.updateProgress(progress); try { await this.indexFile(doc.drive_id); result.indexed++; progress.processed++; } catch (error) { result.errors++; progress.errors++; result.errorDetails.push({ fileId: doc.drive_id, fileName: doc.name, error: error instanceof Error ? error.message : 'Unknown error' }); } this.updateProgress(progress); } return result; } async getFilePath(drive, fileId, fileName) { try { const file = await drive.files.get({ fileId, fields: 'parents' }); if (!file.data.parents || file.data.parents.length === 0) { return fileName; } const parentId = file.data.parents[0]; const parent = await drive.files.get({ fileId: parentId, fields: 'name,parents' }); if (parent.data.name === 'My Drive' || !parent.data.parents) { return fileName; } const parentPath = await this.getFilePath(drive, parentId, parent.data.name); return `${parentPath}/${fileName}`; } catch (error) { console.error('Error getting file path:', error); return fileName; } } } // Export singleton instance export const indexer = new DriveIndexer();