document_utils.py

from langchain_core.embeddings import Embeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.schema.document import Document
from typing import List


def split_documents(documents: List[Document]) -> List[Document]:
    """
    Splits a list of documents into smaller chunks using a recursive character text splitter.

    Args:
        documents (List[Document]): A list of Document objects to be split.

    Returns:
        List[Document]: A list of smaller Document objects after splitting.
    """
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
        is_separator_regex=False
    )
    return text_splitter.split_documents(documents)


def get_embedding_function(model_choice: str = "mini") -> Embeddings:
    """
    Returns a Hugging Face embedding function with model selection.

    Args:
        model_choice: One of "mini", "base", "large", or "multilingual"

    Returns:
        Embeddings: A HuggingFace embeddings instance.
    """
    from langchain_community.embeddings import HuggingFaceEmbeddings

    model_map = {
        # Fast & efficient (384 dims, ~80MB)
        "mini": "sentence-transformers/all-MiniLM-L6-v2",

        # Better quality (768 dims, ~420MB)
        "base": "sentence-transformers/all-mpnet-base-v2",

        # High quality (1024 dims, ~1.3GB)
        "large": "sentence-transformers/all-roberta-large-v1",

        # Multilingual support (768 dims, ~1.2GB)
        "multilingual": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"
    }

    model_name = model_map.get(model_choice, model_map["mini"])

    embeddings = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs={'device': 'cpu'},
        encode_kwargs={'normalize_embeddings': True}
    )

    return embeddings