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 re
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"
EMBEDDING_BATCH_SIZE = 128

# Chunking par paragraphe
PARA_TARGET = 1200
PARA_MAX = 2000
PARA_MIN = 200

_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)

_fallback_splitter = RecursiveCharacterTextSplitter(
    chunk_size=PARA_TARGET,
    chunk_overlap=100,
    separators=["\n", ". ", " ", ""],
)

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

def clean_text(text: str) -> str:
    text = re.sub(r"-\n", "", text)
    text = re.sub(r"(?<!\n)\n(?!\n)", " ", text)
    text = re.sub(r"[ \t]+", " ", text)
    return text.strip()

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 = clean_text(page.extract_text() or "")
        if text:
            pages_text.append(f"[Page {page_num}]\n{text}")
    return "\n\n".join(pages_text)

def chunk_text(text: str, source: str) -> list[dict]:
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]

    chunks, buffer = [], ""
    for para in paragraphs:
        if len(para) > PARA_MAX:
            if buffer:
                chunks.append(buffer); buffer = ""
            chunks.extend(_fallback_splitter.split_text(para))
            continue
        if len(buffer) + len(para) + 2 <= PARA_TARGET:
            buffer = (buffer + "\n\n" + para).strip()
        else:
            if buffer:
                chunks.append(buffer)
            buffer = para
    if buffer:
        chunks.append(buffer)

    merged = []
    for c in chunks:
        if merged and len(c) < PARA_MIN:
            merged[-1] = merged[-1] + "\n\n" + c
        else:
            merged.append(c)

    return [{"text": c, "source": source, "chunk_index": i} for i, c in enumerate(merged)]


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()