ingest.py

# -*- coding: utf-8 -*-
"""ingest

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1jCKxVl3ahePkWWs__LZjCIwbiQzfEl1c
"""

# ingest.py

from pathlib import Path

from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
DB_DIR = BASE_DIR / "db"


def load_documents():
    """
    Carga todos los PDFs de la carpeta ./data
    """
    if not DATA_DIR.exists():
        raise FileNotFoundError(f"La carpeta {DATA_DIR} no existe. Creala y poné ahí tus PDFs.")

    loader = DirectoryLoader(
        str(DATA_DIR),
        glob="*.pdf",
        loader_cls=PyPDFLoader,
    )
    docs = loader.load()
    print(f"[INGEST] Documentos cargados: {len(docs)}")
    return docs


def split_documents(documents, chunk_size=800, chunk_overlap=200):
    """
    Divide los documentos en chunks de texto para RAG.
    """
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_documents(documents)
    print(f"[INGEST] Chunks generados: {len(chunks)}")
    return chunks


def get_embeddings():
    """
    Configura el modelo de embeddings de Hugging Face.
    all-MiniLM-L6-v2 es un estándar para RAG y es liviano.
    """
    model_name = "sentence-transformers/all-MiniLM-L6-v2"
    embeddings = HuggingFaceEmbeddings(model_name=model_name)
    return embeddings


def build_vectorstore(chunks):
    """
    Crea/reescribe la base ChromaDB en ./db
    """
    embeddings = get_embeddings()
    DB_DIR.mkdir(exist_ok=True)

    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embeddings,
        persist_directory=str(DB_DIR),
    )

    vectorstore.persist()
    print(f"[INGEST] Vectorstore creado y persistido en: {DB_DIR.resolve()}")


def main():
    print("[INGEST] Iniciando proceso de ingesta...")
    docs = load_documents()
    chunks = split_documents(docs)

    print("\n[INGEST] Ejemplo de chunk:")
    if chunks:
        sample = chunks[0]
        print(f"  - source: {sample.metadata.get('source')}")
        print(f"  - page:   {sample.metadata.get('page')}")
        print(f"  - texto:  {sample.page_content[:300]}...")

    build_vectorstore(chunks)
    print("[INGEST] Proceso completado ✔")


if __name__ == "__main__":
    main()