backend/rag/parser.py

import tempfile
import os

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS

from core.config import CHUNK_SIZE, CHUNK_OVERLAP
from core.llm import embeddings
from rag.vector_store import save_db, set_db

SUPPORTED_EXTENSIONS = {".pdf", ".docx", ".pptx", ".txt", ".md"}


def _load_documents(path: str, ext: str):
    """Dispatch to the right loader based on file extension."""

    if ext == ".pdf":
        from langchain_community.document_loaders import PyMuPDFLoader
        loader = PyMuPDFLoader(path)

    elif ext == ".docx":
        from langchain_community.document_loaders import Docx2txtLoader
        loader = Docx2txtLoader(path)

    elif ext == ".pptx":
        from pptx import Presentation
        from langchain_core.documents import Document as LCDoc
        prs = Presentation(path)
        raw_docs = []
        for slide_num, slide in enumerate(prs.slides, start=1):
            texts = []
            for shape in slide.shapes:
                if shape.has_text_frame:
                    for para in shape.text_frame.paragraphs:
                        line = " ".join(run.text for run in para.runs).strip()
                        if line:
                            texts.append(line)
            if texts:
                raw_docs.append(LCDoc(
                    page_content="\n".join(texts),
                    metadata={"source": path, "slide": slide_num}
                ))
        return raw_docs

    elif ext in (".txt", ".md"):
        from langchain_community.document_loaders import TextLoader
        loader = TextLoader(path, encoding="utf-8")

    else:
        raise ValueError(f"Unsupported file type: '{ext}'.")

    return loader.load()


async def parse_document(file):
    ext = os.path.splitext(file.filename or "")[-1].lower()

    if ext not in SUPPORTED_EXTENSIONS:
        raise ValueError(
            f"Unsupported file type '{ext}'. "
            f"Please upload one of: {', '.join(sorted(SUPPORTED_EXTENSIONS))}"
        )

    with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
        tmp.write(await file.read())
        path = tmp.name

    try:
        raw_docs = _load_documents(path, ext)

        splitter = RecursiveCharacterTextSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP,
        )

        docs = splitter.split_documents(raw_docs)

        db = FAISS.from_documents(docs, embeddings)
        set_db(db)
        save_db(db)

        # Log chunk count and embedding dimension
        dim = len(embeddings.embed_query("test"))
        print(f"[Parser] File: {file.filename} | Chunks: {len(docs)} | Embedding dim: {dim}")

    finally:
        os.remove(path)

    return True


# Backward-compatible alias so existing imports still work
parse_pdf = parse_document