vectorstore.py

"""
Vector store management: PDF loading, chunking, embedding, and retrieval.
"""

import os
import tempfile
from typing import List

from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.schema import Document


def load_and_chunk_pdfs(uploaded_files) -> List[Document]:
    """
    Load PDF files from Streamlit UploadedFile objects and split into chunks.

    Args:
        uploaded_files: List of Streamlit UploadedFile objects.

    Returns:
        List of LangChain Document objects with metadata (source, page).
    """
    documents = []
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
    )

    for uploaded_file in uploaded_files:
        # PyPDFLoader requires a file path, so write to a temp file
        uploaded_file.seek(0)
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded_file.read())
            tmp_path = tmp.name

        try:
            loader = PyPDFLoader(tmp_path)
            pages = loader.load()

            # Overwrite source metadata with the original filename
            for page in pages:
                page.metadata["source"] = uploaded_file.name

            chunks = text_splitter.split_documents(pages)
            documents.extend(chunks)
        finally:
            os.unlink(tmp_path)

    return documents


def create_vectorstore(documents: List[Document]) -> FAISS:
    """
    Create an in-memory FAISS vector store from document chunks.

    Args:
        documents: List of LangChain Document objects.

    Returns:
        FAISS vector store with embedded documents.
    """
    embeddings = HuggingFaceEmbeddings(
        model_name="all-MiniLM-L6-v2",
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True},
    )
    vectorstore = FAISS.from_documents(documents, embeddings)
    return vectorstore


def get_retriever(vectorstore: FAISS, k: int = 4):
    """
    Return a retriever that fetches the top-k most relevant chunks.

    Args:
        vectorstore: Existing FAISS vector store.
        k: Number of chunks to retrieve per query.

    Returns:
        LangChain retriever object.
    """
    return vectorstore.as_retriever(
        search_type="similarity",
        search_kwargs={"k": k},
    )