ingest_train_data.py

"""
Ingest Training Data Script
============================

Standalone script to convert all documents in ./train_data into embeddings
and store them in the ChromaDB vector store.

Usage:
    python ingest_train_data.py

Requires:
    - AZURE_API_KEY environment variable set
    - data/config.json (or root config.json as fallback) with endpoint URLs
    - Documents (PDF/TXT) in ./train_data/<category>/
"""

import os
import sys
import json
import logging
import time
from pathlib import Path

os.environ.setdefault("ANONYMIZED_TELEMETRY", "False")

import requests as http_requests
import chromadb
from chromadb.config import Settings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pypdf import PdfReader

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

PROJECT_ROOT = Path(__file__).parent
DATA_DIR = Path("/data") if Path("/data").is_dir() else PROJECT_ROOT / "data"
TRAIN_DATA_DIR = Path("/train_data") if Path("/train_data").is_dir() else PROJECT_ROOT / "train_data"

CHROMA_PERSIST_DIR = str(DATA_DIR / "chroma_db")
COLLECTION_NAME = "rag_documents"
CHUNK_SIZE = 512
CHUNK_OVERLAP = 50
EMBEDDING_BATCH_SIZE = 16

_CONFIG_PATH = DATA_DIR / "config.json"
if not _CONFIG_PATH.exists():
    _CONFIG_PATH = PROJECT_ROOT / "config.json"
    logger.warning(f"No config.json in {DATA_DIR} — using root config.json as fallback.")

with open(_CONFIG_PATH, encoding="utf-8") as _f:
    _config = json.load(_f)

EMBEDDING_ENDPOINT_URL = _config["embedding"]["endpoint_url"]
EMBEDDING_MODEL_NAME = _config["embedding"]["model"]

AZURE_API_KEY = os.environ.get("AZURE_API_KEY")
if not AZURE_API_KEY:
    logger.error("AZURE_API_KEY is not set. Export it before running this script.")
    sys.exit(1)

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def extract_text_from_pdf(pdf_path: Path) -> str:
    reader = PdfReader(str(pdf_path))
    pages_text = []
    for page_num, page in enumerate(reader.pages, start=1):
        text = page.extract_text()
        if text and text.strip():
            pages_text.append(f"[Page {page_num}]\n{text.strip()}")
    return "\n\n".join(pages_text)


def chunk_text(text: str, source: str) -> list[dict]:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    chunks = splitter.split_text(text)
    return [{"text": chunk, "source": source, "chunk_index": i} for i, chunk in enumerate(chunks)]


def generate_embeddings(texts: list[str]) -> list[list[float]]:
    headers = {"api-key": AZURE_API_KEY, "Content-Type": "application/json"}
    payload = {"input": texts, "model": EMBEDDING_MODEL_NAME}
    resp = http_requests.post(EMBEDDING_ENDPOINT_URL, headers=headers, json=payload, timeout=120)
    resp.raise_for_status()
    data = resp.json()
    return [item["embedding"] for item in data["data"]]


def generate_embeddings_batched(texts: list[str]) -> list[list[float]]:
    all_embeddings = []
    for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
        batch = texts[i:i + EMBEDDING_BATCH_SIZE]
        logger.info(f"  Embedding batch {i // EMBEDDING_BATCH_SIZE + 1} ({len(batch)} chunks)...")
        embeddings = generate_embeddings(batch)
        all_embeddings.extend(embeddings)
        time.sleep(0.5)  # Rate-limit courtesy
    return all_embeddings


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    if not TRAIN_DATA_DIR.exists():
        logger.error(f"Training data directory not found: {TRAIN_DATA_DIR}")
        sys.exit(1)

    logger.info(f"Training data directory: {TRAIN_DATA_DIR}")
    logger.info(f"ChromaDB path: {CHROMA_PERSIST_DIR}")
    logger.info(f"Embedding model: {EMBEDDING_MODEL_NAME}")

    chroma_client = chromadb.PersistentClient(
        path=CHROMA_PERSIST_DIR,
        settings=Settings(anonymized_telemetry=False),
    )
    collection = chroma_client.get_or_create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )
    logger.info(f"ChromaDB collection '{COLLECTION_NAME}' — existing documents: {collection.count()}")

    # Gather all PDF and TXT files recursively
    files = sorted(
        list(TRAIN_DATA_DIR.rglob("*.pdf")) + list(TRAIN_DATA_DIR.rglob("*.txt"))
    )
    logger.info(f"Found {len(files)} files to process.")

    total_chunks_added = 0

    for idx, file_path in enumerate(files, start=1):
        relative = file_path.relative_to(TRAIN_DATA_DIR)
        # Use category/filename as the source identifier
        source = str(relative)
        logger.info(f"[{idx}/{len(files)}] Processing: {source}")

        try:
            if file_path.suffix.lower() == ".pdf":
                text = extract_text_from_pdf(file_path)
            else:
                text = file_path.read_text(encoding="utf-8")
        except Exception as e:
            logger.warning(f"  Skipped (read error): {e}")
            continue

        if not text.strip():
            logger.warning(f"  Skipped (no extractable text)")
            continue

        chunks = chunk_text(text, source=source)
        if not chunks:
            logger.warning(f"  Skipped (no chunks produced)")
            continue

        logger.info(f"  {len(chunks)} chunks extracted")

        texts = [c["text"] for c in chunks]
        try:
            embeddings = generate_embeddings_batched(texts)
        except http_requests.exceptions.HTTPError as e:
            logger.error(f"  Embedding failed: {e}")
            continue
        except Exception as e:
            logger.error(f"  Unexpected embedding error: {e}")
            continue

        existing_count = collection.count()
        ids = [f"doc_{existing_count + i}" for i in range(len(chunks))]
        metadatas = [{"source": c["source"], "chunk_index": c["chunk_index"]} for c in chunks]

        collection.add(
            ids=ids,
            embeddings=embeddings,
            documents=texts,
            metadatas=metadatas,
        )
        total_chunks_added += len(chunks)
        logger.info(f"  Added {len(chunks)} chunks (total in store: {collection.count()})")

    logger.info("=" * 60)
    logger.info(f"Ingestion complete. Chunks added this run: {total_chunks_added}")
    logger.info(f"Total documents in vector store: {collection.count()}")


if __name__ == "__main__":
    main()