RAG_helper/routers/documents.py

import logging
from datetime import datetime, timezone

from fastapi import APIRouter, File, Form, HTTPException, UploadFile

from models.responses import (
    ChunkDetail,
    ChunkPreview,
    DocumentChunksResponse,
    DocumentDeleteResponse,
    DocumentInfo,
    DocumentListResponse,
    DocumentUploadResponse,
)
from services.document_processor import process_document

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/documents", tags=["documents"])

ALLOWED_EXTENSIONS = {".pdf", ".docx", ".doc", ".txt", ".md"}
MAX_FILE_SIZE = 50 * 1024 * 1024  # 50 MB


@router.post("/upload", response_model=DocumentUploadResponse)
async def upload_document(
    file: UploadFile = File(...),
    document_name: str | None = Form(None),
):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    filename = file.filename or "unknown"
    ext = "." + filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
    if ext not in ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file format: {ext}. Allowed: {', '.join(ALLOWED_EXTENSIONS)}",
        )

    file_bytes = await file.read()
    if len(file_bytes) > MAX_FILE_SIZE:
        raise HTTPException(status_code=400, detail="File too large (max 50 MB)")
    if len(file_bytes) == 0:
        raise HTTPException(status_code=400, detail="Empty file")

    display_name = document_name or filename
    try:
        document_id, sections, chunks = process_document(file_bytes, filename)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        logger.exception("Failed to process document: %s", filename)
        raise HTTPException(status_code=500, detail=f"Error processing document: {e}")

    if not chunks:
        raise HTTPException(status_code=400, detail="No content could be extracted from the document")

    file_type = ext.lstrip(".")
    chunks_count = vectorstore_service.add_document(
        document_id=document_id,
        document_name=display_name,
        file_type=file_type,
        chunks=[
            {
                "text": c.text,
                "section": c.section,
                "page_number": c.page_number,
                "chunk_index": c.chunk_index,
            }
            for c in chunks
        ],
    )

    chunks_prev = [
        ChunkPreview(
            index=c.chunk_index,
            section=c.section,
            page_number=c.page_number,
            text_preview=c.text[:300],
            char_length=len(c.text),
        )
        for c in chunks[:3]
    ]

    return DocumentUploadResponse(
        document_id=document_id,
        name=display_name,
        chunks_count=chunks_count,
        status="indexed",
        created_at=datetime.now(timezone.utc).isoformat(),
        chunks_preview=chunks_prev,
    )


@router.get("", response_model=DocumentListResponse)
async def list_documents():
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    docs = vectorstore_service.list_documents()
    return DocumentListResponse(
        documents=[DocumentInfo(**d) for d in docs],
        total=len(docs),
    )


@router.get("/{document_id}/chunks", response_model=DocumentChunksResponse)
async def get_document_chunks(document_id: str):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    raw_chunks = vectorstore_service.get_document_chunks(document_id)
    if not raw_chunks:
        raise HTTPException(status_code=404, detail="Document not found")

    meta0 = raw_chunks[0]["metadata"]
    chunks = [
        ChunkDetail(
            index=c["metadata"].get("chunk_index", 0),
            section=c["metadata"].get("section", ""),
            page_number=c["metadata"].get("page_number", 0),
            text=c["text"],
            char_length=len(c["text"]),
        )
        for c in raw_chunks
    ]

    return DocumentChunksResponse(
        document_id=document_id,
        name=meta0.get("document_name", ""),
        file_type=meta0.get("file_type", ""),
        chunks_count=len(chunks),
        chunks=chunks,
    )


@router.delete("/{document_id}", response_model=DocumentDeleteResponse)
async def delete_document(document_id: str):
    from main import vectorstore_service

    if vectorstore_service is None:
        raise HTTPException(status_code=503, detail="Service not ready")

    deleted = vectorstore_service.delete_document(document_id)
    if deleted == 0:
        raise HTTPException(status_code=404, detail="Document not found")

    return DocumentDeleteResponse(ok=True, deleted_chunks=deleted)