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

488 lines (486 loc) 22.7 kB
import PQueue from 'p-queue'; import { google } from 'googleapis'; // Use built-in fetch instead of node-fetch import { getValidCredentials } from '../auth.js'; import { dbV2 } from './database_v2.js'; import { DocumentAnalysisSchema, zodToJsonSchema } from './document_schemas.js'; const OPENROUTER_API_URL = 'https://openrouter.ai/api/v1/chat/completions'; export class ParallelIndexer { driveQueue; aiQueue; progressCallback; progress; constructor(progressCallback) { // Google Drive API rate limits // 20,000 queries per 100 seconds per user // We'll be conservative and do 100 requests per 5 seconds this.driveQueue = new PQueue({ concurrency: 50, interval: 5000, intervalCap: 50 }); // OpenRouter rate limits (adjust based on your tier) // Free tier typically allows 20 requests per minute // We'll use 15 to be safe this.aiQueue = new PQueue({ concurrency: 10, interval: 60000, intervalCap: 15 }); this.progressCallback = progressCallback; this.progress = { total: 0, processed: 0, succeeded: 0, failed: 0, skipped: 0, errors: [], startTime: Date.now() }; } updateProgress(update) { this.progress = { ...this.progress, ...update }; // Calculate estimated time remaining if (this.progress.processed > 0) { const elapsedSeconds = (Date.now() - this.progress.startTime) / 1000; const rate = this.progress.processed / elapsedSeconds; const remaining = this.progress.total - this.progress.processed; this.progress.estimatedTimeRemaining = remaining / rate; } if (this.progressCallback) { this.progressCallback(this.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,size,createdTime,modifiedTime,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; } this.updateProgress({ currentFile: file.name || undefined }); // Check if already indexed and up to date const existing = await dbV2.getDocumentByDriveId(file.id); if (existing && existing.modified_time === file.modifiedTime) { console.log(`Skipping up-to-date file: ${file.name}`); this.updateProgress({ processed: this.progress.processed + 1, skipped: this.progress.skipped + 1 }); return; } // 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 / 1024 / 1024).toFixed(2)}MB (max: 20MB)`); } // Get file path const filePath = await this.getFilePath(drive, file.id, file.name); // Analyze with AI using rate-limited queue const startTime = Date.now(); const analysis = await this.aiQueue.add(() => this.analyzeDocument(fileContent, file)); if (!analysis) { throw new Error('Failed to analyze document'); } const processingTime = Date.now() - startTime; // Save to database try { // Insert or update base document const docId = await dbV2.insertDocument({ drive_id: file.id, drive_path: filePath, name: file.name, mime_type: file.mimeType, size: file.size ? parseInt(file.size) : undefined, created_time: file.createdTime || undefined, modified_time: file.modifiedTime || undefined, ai_model: 'gemini-2.5-flash-preview-05-20', ai_processing_time_ms: processingTime, doc_type: analysis.doc_type, synopsis: analysis.synopsis, confidence_score: analysis.confidence_score }); // Insert type-specific data switch (analysis.doc_type) { case 'contract': if ('contract_data' in analysis) { const contractId = await dbV2.insertContract(docId, { contract_number: analysis.contract_data.contract_number || undefined, contract_type: analysis.contract_data.contract_type || undefined, effective_date: analysis.contract_data.effective_date || undefined, expiration_date: analysis.contract_data.expiration_date || undefined, auto_renewal: analysis.contract_data.auto_renewal ?? undefined, notice_period_days: analysis.contract_data.notice_period_days ?? undefined, total_value: analysis.contract_data.total_value ?? undefined, currency: analysis.contract_data.currency || undefined, payment_terms: analysis.contract_data.payment_terms || undefined, governing_law: analysis.contract_data.governing_law || undefined }); if (analysis.contract_data.parties.length > 0) { await dbV2.insertContractParties(contractId, analysis.contract_data.parties.map(p => ({ party_name: p.name, party_role: p.role }))); } } break; case 'invoice': if ('invoice_data' in analysis) { const invoiceId = await dbV2.insertInvoice(docId, { invoice_number: analysis.invoice_data.invoice_number || undefined, invoice_date: analysis.invoice_data.invoice_date || undefined, due_date: analysis.invoice_data.due_date || undefined, supplier_name: analysis.invoice_data.supplier.name || undefined, supplier_address: analysis.invoice_data.supplier.address || undefined, customer_name: analysis.invoice_data.customer.name || undefined, customer_address: analysis.invoice_data.customer.address || undefined, subtotal: analysis.invoice_data.subtotal ?? undefined, tax_amount: analysis.invoice_data.tax_amount ?? undefined, total_amount: analysis.invoice_data.total_amount ?? undefined, currency: analysis.invoice_data.currency || undefined, payment_status: analysis.invoice_data.payment_status || undefined, payment_method: analysis.invoice_data.payment_method || undefined }); if (analysis.invoice_data.line_items.length > 0) { await dbV2.insertInvoiceItems(invoiceId, analysis.invoice_data.line_items); } } break; case 'proposal': if ('proposal_data' in analysis) { const proposalId = await dbV2.insertProposal(docId, { proposal_number: analysis.proposal_data.proposal_number || undefined, client_name: analysis.proposal_data.client_name || undefined, project_name: analysis.proposal_data.project_name || undefined, proposal_date: analysis.proposal_data.proposal_date || undefined, valid_until: analysis.proposal_data.valid_until || undefined, total_value: analysis.proposal_data.total_value ?? undefined, currency: analysis.proposal_data.currency || undefined, timeline_days: analysis.proposal_data.timeline_days ?? undefined, status: analysis.proposal_data.status || undefined }); // Note: Proposal deliverables could be saved here if needed } break; // Add other document types as needed } } catch (dbError) { console.error('Database error:', dbError); throw dbError; } console.log(`Successfully indexed: ${file.name} as ${analysis.doc_type}`); this.updateProgress({ processed: this.progress.processed + 1, succeeded: this.progress.succeeded + 1 }); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error'; console.error(`Error indexing file ${fileId}:`, errorMessage); this.updateProgress({ processed: this.progress.processed + 1, failed: this.progress.failed + 1, errors: [...this.progress.errors, { fileId, fileName: this.progress.currentFile || 'Unknown', error: errorMessage }] }); // Update queue item status if applicable throw error; } } async indexFolder(folderId, recursive = true) { const startTime = Date.now(); // Reset progress this.progress = { total: 0, processed: 0, succeeded: 0, failed: 0, skipped: 0, errors: [], startTime }; try { const authClient = await getValidCredentials(); const drive = google.drive({ version: 'v3', auth: authClient }); // First, collect all files to index const files = await this.collectFiles(drive, folderId, recursive); this.updateProgress({ total: files.length }); console.log(`Found ${files.length} files to index`); // Add files to processing queue in the database const queueItems = files.map(file => ({ drive_id: file.id, operation: 'index', priority: 5 })); await dbV2.addToProcessingQueue(queueItems); // Process files in parallel with rate limiting await Promise.all(files.map(file => this.driveQueue.add(async () => { try { await this.indexFile(file.id); } catch (error) { // Error is already logged and recorded } }))); const duration = (Date.now() - startTime) / 1000; return { total_files: this.progress.total, indexed: this.progress.succeeded, skipped: this.progress.skipped, errors: this.progress.failed, duration_seconds: duration, files_per_second: this.progress.processed / duration, error_details: this.progress.errors }; } catch (error) { console.error('Error in indexFolder:', error); throw error; } } async collectFiles(drive, folderId, recursive) { const files = []; let pageToken; do { const response = await drive.files.list({ q: `'${folderId}' in parents and trashed = false`, fields: 'nextPageToken, files(id, name, mimeType, size, createdTime, modifiedTime)', pageSize: 1000, pageToken }); for (const file of response.data.files || []) { if (file.mimeType === 'application/pdf') { files.push(file); } else if (recursive && file.mimeType === 'application/vnd.google-apps.folder') { // Recursively collect files from subfolders const subfolderFiles = await this.collectFiles(drive, file.id, recursive); files.push(...subfolderFiles); } } pageToken = response.data.nextPageToken; } while (pageToken); return files; } 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; } } async analyzeDocument(content, file) { const systemPrompt = `You are a document analyzer. Analyze the provided document and extract structured data according to the JSON schema provided. Key instructions: 1. Identify the document type (contract, invoice, proposal, project_plan, report, or unknown) 2. Extract all relevant fields based on the document type 3. Provide a confidence score (0-1) indicating how confident you are about the classification 4. Write a clear 1-2 sentence synopsis describing the document's purpose and key parties/details 5. For dates, use ISO format (YYYY-MM-DD) 6. For currency amounts, extract the numeric value only 7. For currency codes, use standard 3-letter ISO codes (USD, EUR, GBP, etc.) 8. If a field cannot be determined, use null 9. Be thorough in extracting all line items, parties, and other detailed information`; try { const response = await fetch(OPENROUTER_API_URL, { method: 'POST', headers: { 'Authorization': `Bearer ${process.env.OPENROUTER_API_KEY}`, 'Content-Type': 'application/json', 'HTTP-Referer': 'https://github.com/markov/google-drive-mcp', 'X-Title': 'Google Drive MCP Document Indexer' }, body: JSON.stringify({ model: 'google/gemini-2.5-flash-preview-05-20', messages: [ { role: 'system', content: systemPrompt }, { role: 'user', content: [ { type: 'text', text: `Analyze this document named "${file.name}" and extract structured information according to the schema.` }, { type: 'file', file: { filename: file.name || 'document.pdf', file_data: `data:${file.mimeType};base64,${content.toString('base64')}` } } ] } ], temperature: 0.7, plugins: [ { id: "file-parser", pdf: { engine: "native" } } ], response_format: { type: 'json_schema', json_schema: { name: 'document_analysis', schema: zodToJsonSchema(DocumentAnalysisSchema), strict: true } } }) }); if (!response.ok) { const errorData = await response.text(); console.error('OpenRouter API error:', errorData); // Check if it's because structured outputs aren't supported if (errorData.includes('json_schema') || errorData.includes('response_format')) { // Fallback to regular JSON mode return await this.analyzeDocumentFallback(content, file); } throw new Error(`OpenRouter API error: ${response.status}`); } const data = await response.json(); if (!data.choices || !data.choices[0] || !data.choices[0].message) { throw new Error('Invalid response from OpenRouter API'); } const content_text = data.choices[0].message.content; const parsed = JSON.parse(content_text); // Validate against schema const validated = DocumentAnalysisSchema.parse(parsed); return validated; } catch (error) { console.error('Error analyzing document:', error); // Try fallback method return await this.analyzeDocumentFallback(content, file); } } async analyzeDocumentFallback(content, file) { // Fallback prompt without structured output const prompt = `Analyze this document and return a JSON response with this structure: { "doc_type": "contract" | "invoice" | "proposal" | "project_plan" | "report" | "unknown", "synopsis": "1-2 sentence summary", "confidence_score": 0.0-1.0, // Then add the appropriate data field based on doc_type: "contract_data": { ... } // if contract "invoice_data": { ... } // if invoice "proposal_data": { ... } // if proposal "project_data": { ... } // if project_plan "report_data": { ... } // if report "extracted_data": { ... } // if unknown }`; try { const response = await fetch(OPENROUTER_API_URL, { method: 'POST', headers: { 'Authorization': `Bearer ${process.env.OPENROUTER_API_KEY}`, 'Content-Type': 'application/json', }, body: JSON.stringify({ model: 'google/gemini-2.5-flash-preview-05-20', messages: [ { role: 'user', content: [ { type: 'text', text: prompt }, { type: 'file', file: { filename: file.name || 'document.pdf', file_data: `data:${file.mimeType};base64,${content.toString('base64')}` } } ] } ], temperature: 0.7, plugins: [ { id: "file-parser", pdf: { engine: "native" } } ], response_format: { type: 'json_object' } }) }); if (!response.ok) { throw new Error(`OpenRouter API error: ${response.status}`); } const data = await response.json(); const parsed = JSON.parse(data.choices[0].message.content); // Basic validation and return return DocumentAnalysisSchema.parse(parsed); } catch (error) { console.error('Fallback analysis also failed:', error); return null; } } async processQueue() { let item = await dbV2.getNextQueueItem(); while (item) { try { await dbV2.updateQueueItemStatus(item.id, 'processing'); switch (item.operation) { case 'index': case 'reindex': await this.indexFile(item.drive_id); break; case 'delete': // Handle deletion if needed break; } await dbV2.updateQueueItemStatus(item.id, 'completed'); } catch (error) { await dbV2.updateQueueItemStatus(item.id, 'failed', error instanceof Error ? error.message : 'Unknown error'); } item = await dbV2.getNextQueueItem(); } } } // Export singleton instance for convenience export const parallelIndexer = new ParallelIndexer();