Vizz17/context-aware-rag

1"""PDF upload & ingestion endpoint."""2 3from __future__ import annotations4 5import shutil6from datetime import datetime, timezone7from pathlib import Path8 9from fastapi import APIRouter, File, HTTPException, UploadFile, Depends10 11from app.core.config import settings12from app.api.auth import get_current_user13from app.models.schemas import DocumentInfo, DocumentListResponse, UploadResponse14from app.services.chunker import chunk_pages15from app.services.embedder import embed_texts16from app.services.pdf_parser import generate_doc_id, parse_pdf17from app.services.vector_store import get_vector_store18 19router = APIRouter()20 21 22@router.post("/upload", response_model=UploadResponse)23def upload_pdf(24    file: UploadFile = File(...),25    user_id: str = Depends(get_current_user)26):27    """Upload a PDF, parse it, chunk it, embed it, and index it."""28    if not file.filename or not file.filename.lower().endswith(".pdf"):29        raise HTTPException(status_code=400, detail="Only PDF files are accepted.")30 31    # Save uploaded file32    upload_dir = Path(settings.upload_dir)33    upload_dir.mkdir(parents=True, exist_ok=True)34    filepath = upload_dir / file.filename35 36    with open(filepath, "wb") as f:37        shutil.copyfileobj(file.file, f)38 39    try:40        # 1. Parse PDF41        pages = parse_pdf(str(filepath))42        if not pages:43            raise HTTPException(status_code=400, detail="Could not extract text from PDF.")44 45        # 2. Chunk46        chunks = chunk_pages(pages)47 48        if not chunks:49            raise HTTPException(50                status_code=400,51                detail=(52                    "No text could be extracted from this PDF. "53                    "It may be a scanned/image-based document. "54                    "Try uploading a text-based PDF (exported from Word/Google Docs) for best results."55                ),56            )57 58        # 3. Embed59        texts = [c.text for c in chunks]60        embeddings = embed_texts(texts)61 62        # 4. Index63        store = get_vector_store()64        store.add_chunks(chunks, embeddings, user_id=user_id)65 66        doc_id = generate_doc_id(str(filepath))67        return UploadResponse(68            doc_id=doc_id,69            filename=file.filename,70            page_count=len(pages),71            chunk_count=len(chunks),72        )73 74    except HTTPException:75        raise76    except Exception as e:77        raise HTTPException(status_code=500, detail=f"Ingestion failed: {e}")78 79 80@router.get("/documents", response_model=DocumentListResponse)81def list_documents(user_id: str = Depends(get_current_user)):82    """List all indexed documents."""83    store = get_vector_store()84    docs = store.list_documents(user_id=user_id)85    return DocumentListResponse(86        documents=[87            DocumentInfo(88                doc_id=d["doc_id"],89                filename=d["filename"],90                page_count=d["page_count"],91                chunk_count=d["chunk_count"],92                indexed_at=datetime.now(timezone.utc),93            )94            for d in docs95        ],96        total=len(docs),97    )98 99 100@router.delete("/documents/{doc_id}")101def delete_document(doc_id: str, user_id: str = Depends(get_current_user)):102    """Remove a document from the vector store."""103    store = get_vector_store()104    deleted = store.delete_document(doc_id, user_id=user_id)105    if deleted == 0:106        raise HTTPException(status_code=404, detail="Document not found.")107    return {"message": f"Deleted {deleted} chunks for document {doc_id}"}108