ingest.py

"""Pipeline de ingesta para el sistema RAG de EcoMarket.

Carga los tres documentos de la base de conocimiento (política Markdown,
catálogo de pedidos CSV y FAQ JSON), aplica una estrategia de chunking
distinta a cada formato, los embebe con `bge-m3` vía Ollama y los persiste
en una colección de ChromaDB.
"""
from __future__ import annotations

import json
import shutil
from pathlib import Path

from langchain_chroma import Chroma
from langchain_community.document_loaders import CSVLoader, TextLoader
from langchain_core.documents import Document
from langchain_ollama import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter

PROJECT_ROOT = Path(__file__).resolve().parent.parent
DATA_DIR = PROJECT_ROOT / "data"
CHROMA_DIR = PROJECT_ROOT / "chroma_db"

EMBEDDING_MODEL = "bge-m3"
COLLECTION_NAME = "ecomarket_knowledge"


def load_policy_chunks() -> list[Document]:
    """Carga la política de devoluciones (Markdown) y la fragmenta de forma recursiva.

    Estrategia: separadores jerárquicos que respetan primero los encabezados
    (`##`, `###`), luego párrafos. chunk_size=600 caracteres con overlap=100
    preserva contexto suficiente para que cada fragmento sea autoexplicativo.
    """
    loader = TextLoader(str(DATA_DIR / "politica_devoluciones.md"), encoding="utf-8")
    raw_docs = loader.load()
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=600,
        chunk_overlap=100,
        separators=["\n## ", "\n### ", "\n\n", "\n", " ", ""],
    )
    chunks = splitter.split_documents(raw_docs)
    for chunk in chunks:
        chunk.metadata["fuente"] = "politica_devoluciones"
        chunk.metadata["tipo"] = "politica"
    return chunks


def load_orders_chunks() -> list[Document]:
    """Carga el catálogo de pedidos. Cada fila CSV se convierte en un chunk.

    No re-segmentamos: cada pedido es atómico. Mezclar dos pedidos en un
    mismo chunk introduciría ambigüedad ("¿de qué pedido es esa fecha?").
    """
    loader = CSVLoader(
        file_path=str(DATA_DIR / "catalogo_pedidos.csv"),
        encoding="utf-8",
        csv_args={"delimiter": ","},
    )
    docs = loader.load()
    for doc in docs:
        doc.metadata["fuente"] = "catalogo_pedidos"
        doc.metadata["tipo"] = "pedido"
    return docs


def load_faq_chunks() -> list[Document]:
    """Carga el FAQ JSON: cada par {pregunta, respuesta} se convierte en un chunk."""
    with (DATA_DIR / "faq_ecomarket.json").open(encoding="utf-8") as f:
        faqs = json.load(f)
    chunks: list[Document] = []
    for item in faqs:
        text = f"P: {item['pregunta']}\nR: {item['respuesta']}"
        chunks.append(
            Document(
                page_content=text,
                metadata={
                    "fuente": "faq_ecomarket",
                    "tipo": "faq",
                    "categoria": item.get("categoria", "general"),
                },
            )
        )
    return chunks


def build_vector_store() -> None:
    """Construye (o reconstruye) la base vectorial persistente."""
    if CHROMA_DIR.exists():
        print(f"Eliminando colección existente en {CHROMA_DIR}...")
        shutil.rmtree(CHROMA_DIR)

    print("Cargando documentos...")
    all_chunks: list[Document] = []
    all_chunks.extend(load_policy_chunks())
    all_chunks.extend(load_orders_chunks())
    all_chunks.extend(load_faq_chunks())
    print(f"Total de chunks generados: {len(all_chunks)}")

    print(f"Inicializando embeddings con '{EMBEDDING_MODEL}' vía Ollama...")
    embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)

    print(f"Indexando en ChromaDB ({CHROMA_DIR})...")
    Chroma.from_documents(
        documents=all_chunks,
        embedding=embeddings,
        collection_name=COLLECTION_NAME,
        persist_directory=str(CHROMA_DIR),
    )
    print("Indexación completa. La base vectorial está lista para consultas.")


if __name__ == "__main__":
    build_vector_store()