Spaces:
Running
Running
| # src/utils/drive_document_processor.py | |
| from pathlib import Path | |
| from typing import Dict, List, Any, Tuple | |
| import logging | |
| from fastapi import HTTPException | |
| from src.utils.google_drive_service import GoogleDriveService | |
| from src.utils.document_processor import DocumentProcessor | |
| from src.vectorstores.chroma_vectorstore import ChromaVectorStore | |
| from src.utils.logger import logger | |
| from src.db.mongodb_store import MongoDBStore | |
| class DriveDocumentProcessor: | |
| def __init__( | |
| self, | |
| google_service_account_path: str, | |
| folder_id: str, | |
| temp_dir: str, | |
| doc_processor: DocumentProcessor, | |
| mongodb: MongoDBStore # Add MongoDB | |
| ): | |
| """ | |
| Initialize Drive Document Processor | |
| Args: | |
| google_service_account_path (str): Path to Google service account credentials | |
| folder_id (str): Google Drive folder ID to process | |
| temp_dir (str): Directory for temporary files | |
| doc_processor (DocumentProcessor): Instance of DocumentProcessor | |
| """ | |
| self.google_drive_service = GoogleDriveService( | |
| google_service_account_path) | |
| self.folder_id = folder_id | |
| self.temp_dir = Path(temp_dir) | |
| self.doc_processor = doc_processor | |
| self.mongodb = mongodb # Store MongoDB instance | |
| # Create temp directory if it doesn't exist | |
| self.temp_dir.mkdir(parents=True, exist_ok=True) | |
| # Define supported MIME types | |
| self.supported_mime_types = { | |
| # Google Docs | |
| 'application/vnd.google-apps.document': '.docx', | |
| # Microsoft Word Documents | |
| 'application/vnd.openxmlformats-officedocument.wordprocessingml.document': '.docx', | |
| 'application/msword': '.doc', | |
| # Microsoft Excel Documents | |
| 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': '.xlsx', | |
| 'application/vnd.ms-excel': '.xls', | |
| # Text Documents | |
| 'text/plain': '.txt', | |
| 'text/csv': '.csv', | |
| 'text/markdown': '.md', | |
| 'text/html': '.html', | |
| 'text/xml': '.xml', | |
| 'application/json': '.json', | |
| 'application/rtf': '.rtf', | |
| # PDF Documents | |
| 'application/pdf': '.pdf' | |
| } | |
| self.google_docs_export_types = { | |
| 'application/vnd.google-apps.document': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document' | |
| } | |
| async def _cleanup_orphaned_documents( | |
| self, | |
| drive_files: List[Dict[str, Any]], | |
| vector_store: ChromaVectorStore | |
| ) -> Dict[str, Any]: | |
| """ | |
| Clean up documents that exist in MongoDB but not in Google Drive | |
| Args: | |
| drive_files (List[Dict[str, Any]]): List of files from Google Drive | |
| vector_store (ChromaVectorStore): Vector store instance | |
| Returns: | |
| Dict[str, Any]: Cleanup statistics | |
| """ | |
| try: | |
| # Get all documents from MongoDB | |
| mongo_docs = await self.mongodb.get_all_documents() | |
| # Create set of Google Drive file IDs | |
| drive_file_ids = {file['id'] for file in drive_files} | |
| deleted_count = 0 | |
| failed_deletions = [] | |
| # Check each MongoDB document | |
| for doc in mongo_docs: | |
| # Only process Google Drive documents | |
| if doc.get('source') != 'google_drive': | |
| continue | |
| doc_id = doc.get('document_id') | |
| if not doc_id or doc_id not in drive_file_ids: | |
| try: | |
| # Delete from MongoDB | |
| await self.mongodb.delete_document(doc_id) | |
| # Delete from vector store | |
| vector_store.delete_document(doc_id) | |
| deleted_count += 1 | |
| except Exception as e: | |
| logger.error( | |
| f"Error deleting orphaned document {doc_id}: {str(e)}") | |
| failed_deletions.append({ | |
| 'document_id': doc_id, | |
| 'error': str(e) | |
| }) | |
| return { | |
| 'orphaned_documents_deleted': deleted_count, | |
| 'failed_deletions': failed_deletions | |
| } | |
| except Exception as e: | |
| logger.error(f"Error in cleanup_orphaned_documents: {str(e)}") | |
| raise | |
| async def process_documents( | |
| self, | |
| vector_store: ChromaVectorStore, | |
| # New parameter with default True for backward compatibility | |
| include_subfolders: bool = True | |
| ) -> Dict[str, Any]: | |
| """ | |
| Process all documents in the specified Drive folder | |
| Args: | |
| vector_store (ChromaVectorStore): Vector store instance | |
| include_subfolders (bool): Whether to process documents in subfolders | |
| Returns: | |
| Dict[str, Any]: Processing results | |
| """ | |
| try: | |
| # Get documents from folder | |
| files = self.google_drive_service.get_folder_contents( | |
| self.folder_id, | |
| include_subfolders=include_subfolders | |
| ) | |
| # Clean up orphaned documents first | |
| cleanup_results = await self._cleanup_orphaned_documents(files, vector_store) | |
| processed_files = [] | |
| skipped_files = [] | |
| errors = [] | |
| for file in files: | |
| # Skip if it's a folder | |
| if file.get('mimeType') == 'application/vnd.google-apps.folder': | |
| continue | |
| # Get file path (including folder structure if available) | |
| file_path = self._get_file_path(file) | |
| file['display_path'] = file_path | |
| result = await self._process_single_file(file, vector_store) | |
| if result['status'] == 'processed': | |
| processed_files.append(result['data']) | |
| elif result['status'] == 'skipped': | |
| skipped_files.append(result['data']) | |
| else: # status == 'error' | |
| errors.append(result['data']) | |
| # Clean up temporary directory if empty | |
| self._cleanup_temp_dir() | |
| return { | |
| "status": "completed", | |
| "processed_files": { | |
| "count": len(processed_files), | |
| "details": processed_files | |
| }, | |
| "skipped_files": { | |
| "count": len(skipped_files), | |
| "details": skipped_files | |
| }, | |
| "errors": { | |
| "count": len(errors), | |
| "details": errors | |
| } | |
| } | |
| except Exception as e: | |
| logger.error(f"Error processing Drive documents: {str(e)}") | |
| raise HTTPException( | |
| status_code=500, | |
| detail=f"Failed to process drive documents: {str(e)}" | |
| ) | |
| def _get_file_path(self, file: Dict[str, Any]) -> str: | |
| """ | |
| Get the full path for a file including its folder structure | |
| Args: | |
| file (Dict[str, Any]): File metadata | |
| Returns: | |
| str: Display path of the file | |
| """ | |
| path_parts = [file['name']] | |
| # Add folder path if available (new structure) | |
| if folder_path := file.get('folder_path', []): | |
| for folder in reversed(folder_path): | |
| path_parts.insert(0, folder['name']) | |
| return '/'.join(path_parts) | |
| async def _process_single_file( | |
| self, | |
| file: Dict[str, Any], | |
| vector_store: ChromaVectorStore | |
| ) -> Dict[str, Any]: | |
| """Process a single Drive file""" | |
| mime_type = file.get('mimeType', '') | |
| # Skip if mime type not supported | |
| if mime_type not in self.supported_mime_types: | |
| return { | |
| 'status': 'skipped', | |
| 'data': { | |
| 'name': file['name'], | |
| 'path': file.get('display_path', file['name']), | |
| 'reason': f'Unsupported mime type: {mime_type}' | |
| } | |
| } | |
| try: | |
| document_id = file['id'] | |
| modified_time = file.get('modifiedTime', 'N/A') | |
| # Check if document should be processed | |
| if self.save_document(document_id, vector_store, modified_time): | |
| # Download and process file | |
| temp_file_path = await self._download_and_save_file( | |
| file['id'], | |
| mime_type | |
| ) | |
| try: | |
| # Process document | |
| processed_doc = await self.doc_processor.process_document( | |
| str(temp_file_path) | |
| ) | |
| # Add to vector store with path information | |
| self._add_to_vector_store( | |
| processed_doc['chunks'], | |
| file, | |
| mime_type, | |
| vector_store | |
| ) | |
| # Add MongoDB storage - Store Google Drive URL | |
| await self.mongodb.store_document( | |
| document_id=document_id, | |
| filename=file['name'], | |
| content_type=mime_type, | |
| file_size=0, # Not needed for drive documents | |
| url_path=f"https://drive.google.com/file/d/{document_id}/view", | |
| source="google_drive" | |
| ) | |
| return { | |
| 'status': 'processed', | |
| 'data': { | |
| 'name': file['name'], | |
| 'path': file.get('display_path', file['name']), | |
| 'id': file['id'], | |
| 'chunks_processed': len(processed_doc['chunks']) | |
| } | |
| } | |
| finally: | |
| # Clean up temporary file | |
| if temp_file_path.exists(): | |
| temp_file_path.unlink() | |
| else: | |
| return { | |
| 'status': 'skipped', | |
| 'data': { | |
| 'name': file['name'], | |
| 'path': file.get('display_path', file['name']), | |
| 'reason': 'Document already exists in the memory.' | |
| } | |
| } | |
| except Exception as e: | |
| logger.error(f"Error processing file {file['name']}: {str(e)}") | |
| return { | |
| 'status': 'error', | |
| 'data': { | |
| 'file_name': file['name'], | |
| 'path': file.get('display_path', file['name']), | |
| 'error': str(e) | |
| } | |
| } | |
| def _add_to_vector_store( | |
| self, | |
| chunks: List[str], | |
| file: Dict[str, Any], | |
| mime_type: str, | |
| vector_store: ChromaVectorStore | |
| ) -> None: | |
| """Add processed chunks to vector store with path information""" | |
| chunk_metadatas = [] | |
| chunk_ids = [] | |
| modified_time = file.get('modifiedTime', 'N/A') | |
| file_path = file.get('display_path', file['name']) | |
| for i, chunk in enumerate(chunks): | |
| chunk_id = f"{file['id']}-chunk-{i}" | |
| chunk_ids.append(chunk_id) | |
| chunk_metadatas.append({ | |
| "source": file_path, # Use full path instead of just name | |
| "document_id": file['id'], | |
| "chunk_index": i, | |
| "mime_type": mime_type, | |
| "modified_time": modified_time, | |
| "total_chunks": len(chunks), | |
| "file_type": self.supported_mime_types[mime_type], | |
| "is_google_doc": mime_type.startswith('application/vnd.google-apps') | |
| }) | |
| vector_store.add_documents( | |
| documents=chunks, | |
| metadatas=chunk_metadatas, | |
| ids=chunk_ids | |
| ) | |
| async def _download_and_save_file( | |
| self, | |
| file_id: str, | |
| mime_type: str | |
| ) -> Path: | |
| """Download and save file to temporary location""" | |
| extension = self.supported_mime_types[mime_type] | |
| temp_file_path = self.temp_dir / f"{file_id}{extension}" | |
| if mime_type in self.google_docs_export_types: | |
| # Download Google Doc in the specified export format | |
| content = self.google_drive_service.export_file( | |
| file_id, | |
| self.google_docs_export_types[mime_type] | |
| ) | |
| else: | |
| # Download regular file | |
| content = self.google_drive_service.download_file(file_id) | |
| with open(temp_file_path, 'wb') as f: | |
| if isinstance(content, str): | |
| f.write(content.encode('utf-8')) | |
| else: | |
| f.write(content) | |
| return temp_file_path | |
| def save_document( | |
| self, | |
| document_id: str, | |
| vector_store: ChromaVectorStore, | |
| modified_date: str | |
| ) -> bool: | |
| """ | |
| Check if document needs to be processed based on modification date | |
| Args: | |
| document_id (str): ID of the document to check | |
| vector_store (ChromaVectorStore): Vector store instance | |
| modified_date (str): Modified date to compare against | |
| Returns: | |
| bool: True if document should be processed, False otherwise | |
| """ | |
| try: | |
| # Retrieve all chunks for the given document_id | |
| chunks = vector_store.get_document_chunks(document_id) | |
| if not chunks: | |
| # Document doesn't exist in vector store | |
| return True | |
| # Check the modified_time of the first chunk | |
| first_chunk_metadata = chunks[0].get("metadata", {}) | |
| if first_chunk_metadata.get("modified_time") != modified_date: | |
| # If modified_time doesn't match, delete existing chunks | |
| vector_store.delete_document(document_id) | |
| logger.info( | |
| f"Document {document_id} has been modified, will reprocess") | |
| return True | |
| logger.info(f"Document {document_id} is up to date, skipping") | |
| return False | |
| except Exception as e: | |
| logger.error(f"Error checking document status: {str(e)}") | |
| # In case of error, process the document to be safe | |
| return True | |
| def _cleanup_temp_dir(self) -> None: | |
| """Clean up temporary directory if empty""" | |
| try: | |
| if self.temp_dir.exists() and not any(self.temp_dir.iterdir()): | |
| self.temp_dir.rmdir() | |
| except Exception as e: | |
| logger.error(f"Error cleaning up temp directory: {str(e)}") | |
| # Don't raise the error as this is a cleanup operation | |