Spaces:
Runtime error
Runtime error
| from sqlalchemy.orm import Session | |
| from typing import List, Optional | |
| import os | |
| import shutil | |
| from app.database.models import Document | |
| from app.services.rag_service import rag_service | |
| from app.utils.helpers import generate_id | |
| from app.utils.validators import validate_file_type, validate_file_size | |
| class DocumentService: | |
| """Service for document management operations.""" | |
| def upload_document( | |
| db: Session, | |
| file, | |
| user_id: str, | |
| filename: str | |
| ) -> Document: | |
| """ | |
| Upload and process a document. | |
| Args: | |
| db: Database session | |
| file: File object | |
| user_id: User ID | |
| filename: Original filename | |
| Returns: | |
| Created Document object | |
| Raises: | |
| ValueError: If file validation fails | |
| """ | |
| # Validate file type (PDF only for now) | |
| if not validate_file_type(filename, ['pdf', 'txt', 'docx']): | |
| raise ValueError("Invalid file type. Only PDF, TXT, and DOCX files are allowed.") | |
| # Get file size | |
| file.seek(0, 2) # Seek to end | |
| file_size = file.tell() | |
| file.seek(0) # Reset to beginning | |
| # Validate file size (10MB limit) | |
| if not validate_file_size(file_size, max_size_mb=10): | |
| raise ValueError("File size exceeds 10MB limit.") | |
| # Generate document ID | |
| doc_id = generate_id() | |
| # Determine file type | |
| file_extension = filename.rsplit('.', 1)[1].lower() if '.' in filename else 'unknown' | |
| # Create upload directory if it doesn't exist | |
| upload_dir = os.path.join(os.path.dirname(__file__), "..", "..", "data", "uploads") | |
| os.makedirs(upload_dir, exist_ok=True) | |
| # Save file | |
| file_path = os.path.join(upload_dir, f"{doc_id}_{filename}") | |
| with open(file_path, "wb") as buffer: | |
| shutil.copyfileobj(file, buffer) | |
| # Create document record | |
| document = Document( | |
| id=doc_id, | |
| user_id=user_id, | |
| filename=filename, | |
| file_path=file_path, | |
| file_type=file_extension, | |
| file_size=file_size | |
| ) | |
| db.add(document) | |
| db.commit() | |
| db.refresh(document) | |
| return document | |
| def process_document_content( | |
| db: Session, | |
| document_id: str, | |
| content: str | |
| ) -> int: | |
| """ | |
| Process document content for RAG. | |
| Args: | |
| db: Database session | |
| document_id: Document ID | |
| content: Extracted text content | |
| Returns: | |
| Number of chunks created | |
| """ | |
| document = db.query(Document).filter(Document.id == document_id).first() | |
| if not document: | |
| raise ValueError("Document not found") | |
| # Process with RAG service | |
| num_chunks = rag_service.process_document( | |
| document_id=document.id, | |
| filename=document.filename, | |
| content=content, | |
| user_id=document.user_id | |
| ) | |
| return num_chunks | |
| def get_user_documents(db: Session, user_id: str) -> List[Document]: | |
| """ | |
| Get all documents for a user. | |
| Args: | |
| db: Database session | |
| user_id: User ID | |
| Returns: | |
| List of Document objects | |
| """ | |
| return db.query(Document).filter( | |
| Document.user_id == user_id | |
| ).order_by(Document.created_at.desc()).all() | |
| def delete_document(db: Session, document_id: str) -> bool: | |
| """ | |
| Delete a document and its chunks. | |
| Args: | |
| db: Database session | |
| document_id: Document ID | |
| Returns: | |
| True if deleted, False if not found | |
| """ | |
| document = db.query(Document).filter(Document.id == document_id).first() | |
| if not document: | |
| return False | |
| # Delete file from filesystem | |
| if os.path.exists(document.file_path): | |
| os.remove(document.file_path) | |
| # Delete chunks from vector database | |
| rag_service.delete_document_chunks(document_id) | |
| # Delete database record | |
| db.delete(document) | |
| db.commit() | |
| return True | |
| # Global document service instance | |
| document_service = DocumentService() | |