pipeline.py

"""
Ingestion pipeline: load raw documents, chunk them, build/load vector and
keyword stores.

Phase 2: curated product summaries from knowledge/product_summary.py are now
included in the ingest so the system can answer from clean, hand-authored text.

Phase 3: the index manifest detects when source files change and automatically
triggers a rebuild so the Chroma vector index and BM25 are always in sync.
"""
import logging

from langchain_community.vectorstores import Chroma
from langchain_community.retrievers import BM25Retriever

from config.settings import (
    CHROMA_DB_PATH,
    RAW_PDF_PATH,
    RAW_EXCEL_PATH,
    RAW_TEXT_PATH,
    RAW_IMAGE_PATH,
)
from config.runtime import is_keyword_enabled, is_force_rebuild, get_retriever_k
from llm.providers import get_embeddings
from processors.pdf_processor import process_pdfs
from processors.excel_processor import process_excel
from processors.text_processor import process_texts
from processors.image_processor import process_images
from processors.document_processor import process_documents
from knowledge.product_summary import create_product_summaries
from ingestion.manifest import needs_rebuild, save_manifest

logger = logging.getLogger(__name__)

_RAW_DIRS = [RAW_PDF_PATH, RAW_EXCEL_PATH, RAW_TEXT_PATH, RAW_IMAGE_PATH]


# ---------------------------------------------------------------------------
# Document collection
# ---------------------------------------------------------------------------

def ingest_all_documents():
    """
    Load raw documents from all sources including curated product summaries.
    Returns a flat list of Document objects (not yet chunked).
    """
    all_docs = []
    all_docs.extend(process_pdfs())
    all_docs.extend(process_excel())
    all_docs.extend(process_texts())
    all_docs.extend(process_images())

    # Phase 2: ingest curated summaries alongside raw files
    try:
        summaries = create_product_summaries()
        all_docs.extend(summaries)
        logger.info("Ingested %d curated product summary documents.", len(summaries))
    except Exception as e:
        logger.warning("Could not load product summaries: %s", e)

    print(f"Ingestion complete: {len(all_docs)} raw documents loaded")
    return all_docs


# ---------------------------------------------------------------------------
# Keyword retriever
# ---------------------------------------------------------------------------

def create_keyword_retriever(chunks):
    if not is_keyword_enabled():
        return None
    if not chunks:
        return None
    retriever = BM25Retriever.from_documents(chunks)
    retriever.k = get_retriever_k()
    print("[OK] Keyword index built (BM25)")
    return retriever


# ---------------------------------------------------------------------------
# Database build / load
# ---------------------------------------------------------------------------

def build_database():
    """Process all documents, embed them, and persist to Chroma."""
    print("=== Processing documents ===")

    raw_docs = ingest_all_documents()
    all_chunks = process_documents(raw_docs)

    if not all_chunks:
        print("[WARN] No documents found! Place files in the data/raw/ folder.")
        return None, []

    print(f"\nTotal: {len(all_chunks)} chunks")
    print("=== Building vector database ===")

    embeddings = get_embeddings()
    vectorstore = Chroma.from_documents(
        documents=all_chunks,
        embedding=embeddings,
        persist_directory=CHROMA_DB_PATH,
    )

    # Phase 3: update manifest so next run knows the index is fresh
    save_manifest(CHROMA_DB_PATH, _RAW_DIRS)

    print("[OK] Vector database built successfully!")
    return vectorstore, all_chunks


def load_or_build_database():
    """
    Load the persisted Chroma DB when it exists and source data has not
    changed; otherwise build from scratch.

    Phase 3: the manifest hash check replaces the previous pattern where
    BM25 was always rebuilt from raw files on every startup.  When the index
    is still valid, BM25 chunks are rebuilt from the same ingest so both
    indexes share identical documents.
    """
    import os
    embeddings = get_embeddings()
    force_rebuild = is_force_rebuild()

    # Determine whether raw data has changed since the last build
    index_exists = os.path.isdir(CHROMA_DB_PATH) and os.listdir(CHROMA_DB_PATH)
    data_changed = needs_rebuild(CHROMA_DB_PATH, _RAW_DIRS) if index_exists else True

    if force_rebuild or data_changed:
        if force_rebuild:
            print("=== Force-rebuilding vector database ===")
        else:
            print("=== Source data changed — rebuilding vector database ===")
        return build_database()

    print("=== Loading existing vector database ===")
    vectorstore = Chroma(
        persist_directory=CHROMA_DB_PATH,
        embedding_function=embeddings,
    )
    print("[OK] Existing vector database loaded")

    keyword_chunks = []
    if is_keyword_enabled():
        raw_docs = ingest_all_documents()
        keyword_chunks = process_documents(raw_docs)

    return vectorstore, keyword_chunks