from fastapi import ( APIRouter, UploadFile, File, Depends, HTTPException, ) from sqlalchemy.orm import Session import shutil import os from DocPilot.backend.app.services.ingestion import ( process_document, TextExtractionError, ) from pilotcore.retrieval.vector_store import ( reset_vector_store, rebuild_index_without_document, ) from DocPilot.backend.app.core.dependencies import ( get_current_user, ) from DocPilot.backend.app.db.session import get_db from DocPilot.backend.app.models.document import Document from DocPilot.backend.app.schemas.document import ( DocumentResponse, ) router = APIRouter() @router.post("/upload") async def upload_document( file: UploadFile = File(...), current_user=Depends(get_current_user), db: Session = Depends(get_db), ): document_count = ( db.query(Document).filter(Document.owner_id == current_user.id).count() ) if current_user.plan == "free" and document_count >= 3: raise HTTPException( status_code=403, detail="Free plan upload limit reached.", ) os.makedirs( "temp", exist_ok=True, ) allowed_extensions = [ ".pdf", ".docx", ".pptx", ".txt", ".md", ".csv", ".xlsx", ".png", ".jpg", ".jpeg", ".webp", # treated as plain-text/code-like content ".py", ".js", ".jsx", ".ts", ".tsx", ".java", ".cpp", ".c", ".h", ".go", ".rs", ".json", ".yaml", ".yml", ".sql", ".css", ".html", ] file_ext = os.path.splitext(file.filename)[1].lower() if file_ext not in allowed_extensions: raise HTTPException( status_code=400, detail="Unsupported file type.", ) file_path = f"temp/{file.filename}" with open( file_path, "wb", ) as buffer: shutil.copyfileobj( file.file, buffer, ) document = Document( owner_id=current_user.id, filename=file.filename, filepath=file_path, file_size=os.path.getsize(file_path), ) db.add(document) db.commit() db.refresh(document) try: process_document( file_path, current_user.id, document.id, mime_type=file.content_type, ) except TextExtractionError as exc: db.delete(document) db.commit() raise HTTPException( status_code=422, detail=str(exc), ) from exc return { "message": "Document uploaded", "document_id": document.id, } @router.get( "/", response_model=list[DocumentResponse], ) def get_documents( db: Session = Depends(get_db), current_user=Depends(get_current_user), ): documents = db.query(Document).filter(Document.owner_id == current_user.id).all() return documents @router.delete("/reset") def reset_documents( current_user=Depends(get_current_user), db: Session = Depends(get_db), ): reset_vector_store(current_user.id) db.query(Document).filter(Document.owner_id == current_user.id).delete() db.commit() return {"message": "Vector store and documents cleared."} @router.delete("/{document_id}") def delete_document( document_id: int, db: Session = Depends(get_db), current_user=Depends(get_current_user), ): document = ( db.query(Document) .filter( Document.id == document_id, Document.owner_id == current_user.id, ) .first() ) if not document: raise HTTPException( status_code=404, detail="Document not found", ) if os.path.exists(document.filepath): os.remove(document.filepath) rebuild_index_without_document( current_user.id, document.id, ) db.delete(document) db.commit() return { "message": "Document deleted", }