src/ingest.py

"""Document ingestion pipeline: files -> chunks -> embeddings -> ChromaDB."""

import hashlib
import os
from pathlib import Path

import chromadb
from chromadb.utils import embedding_functions
# from sentence_transformers import SentenceTransformer
from tqdm import tqdm

from src.readers import read_file, READERS
from src.config_loader import load_config, get_api_key

MAX_CHUNK_CHARS = 6000


def split_text_recursive(text: str, chunk_size: int, chunk_overlap: int,
                         _sep_index: int = 0) -> list[str]:
    """Split text into overlapping chunks using recursive separators."""
    separators = ["\n\n", "\n", ". ", " "]

    if len(text) <= chunk_size:
        return [text] if text.strip() else []

    if _sep_index >= len(separators):
        # Last resort: character split
        step = max(1, chunk_size - chunk_overlap)
        chunks = []
        for i in range(0, len(text), step):
            chunk = text[i:i + chunk_size]
            if chunk.strip():
                chunks.append(chunk.strip())
        return chunks

    sep = separators[_sep_index]

    if sep not in text:
        return split_text_recursive(text, chunk_size, chunk_overlap, _sep_index + 1)

    parts = text.split(sep)
    merged = []
    current = ""
    for part in parts:
        candidate = current + sep + part if current else part
        if len(candidate) > chunk_size and current:
            merged.append(current)
            overlap_text = current[-chunk_overlap:] if chunk_overlap else ""
            current = overlap_text + sep + part if overlap_text else part
        else:
            current = candidate
    if current.strip():
        merged.append(current)

    # Recursively split any chunks that are still too large
    result = []
    for chunk in merged:
        if len(chunk) > chunk_size:
            result.extend(split_text_recursive(chunk, chunk_size, chunk_overlap,
                                               _sep_index + 1))
        elif chunk.strip():
            result.append(chunk.strip())
    return result


def chunk_documents(pages: list[dict], source_name: str, dataset_name: str,
                    chunk_size: int = 1000, chunk_overlap: int = 100) -> list[dict]:
    """Split pages into smaller chunks with metadata."""
    chunks = []
    global_chunk_index = 0
    for page_info in pages:
        text = page_info.get("text", "")
        if not text.strip():
            continue
        if len(text) > MAX_CHUNK_CHARS:
            splits = split_text_recursive(text, MAX_CHUNK_CHARS, chunk_overlap)
        else:
            splits = split_text_recursive(text, chunk_size, chunk_overlap)

        for split in splits:
            if len(split) > MAX_CHUNK_CHARS:
                split = split[:MAX_CHUNK_CHARS]
            chunks.append({
                "text": split,
                "metadata": {
                    "source": source_name,
                    "dataset": dataset_name,
                    "page": str(page_info.get("page", "?")),
                    "chunk_index": global_chunk_index,
                },
            })
            global_chunk_index += 1
    return chunks


def discover_files(documents_dir: str) -> list[tuple[Path, str]]:
    """Recursively discover all supported files.

    Returns list of (file_path, dataset_name) tuples.
    Dataset name is the immediate subfolder name, or "general" for root files.
    """
    root = Path(documents_dir)
    supported = set(READERS.keys())
    files = []

    for file_path in sorted(root.rglob("*")):
        if not file_path.is_file():
            continue
        if file_path.suffix.lower() not in supported:
            continue
        if file_path.name.startswith(".") or file_path.name.startswith("~$"):
            continue

        rel = file_path.relative_to(root)
        dataset_name = rel.parts[0] if len(rel.parts) > 1 else "general"
        files.append((file_path, dataset_name))

    return files


def get_chroma_collection(cfg: dict):
    """Get or create the ChromaDB collection with embeddings."""
    db_path = cfg.get("paths", {}).get("vector_db", "chroma_db")
    if not os.path.isabs(db_path):
        project_root = Path(__file__).resolve().parent.parent
        db_path = os.path.join(str(project_root), db_path)

    client = chromadb.PersistentClient(path=db_path)

    embed_provider = cfg.get("embeddings", {}).get("provider", "local")
    if embed_provider == "openai":
        api_key = get_api_key(cfg, "openai")
        model = cfg.get("embeddings", {}).get("openai_model", "text-embedding-3-small")
        ef = embedding_functions.OpenAIEmbeddingFunction(
            api_key=api_key, model_name=model,
        )
    elif embed_provider == "meta-llama":
        api_key = get_api_key(cfg, "meta-llama")
        model = cfg.get("embeddings", {}).get("emb_model", "sentence-transformers/all-mpnet-base-v2")
        ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model)
    else:
        ef = embedding_functions.DefaultEmbeddingFunction()

    collection = client.get_or_create_collection(
        name="knowledge_base",
        embedding_function=ef,
        metadata={"hnsw:space": "cosine"},
    )
    return collection


def ingest_documents(cfg: dict = None, documents_dir: str = None) -> int:
    """Ingest all supported files from knowledge_base/ into ChromaDB.

    Returns the number of chunks ingested.
    """
    if cfg is None:
        cfg = load_config()

    if documents_dir is None:
        documents_dir = cfg.get("paths", {}).get("knowledge_base", "knowledge_base")
        if not os.path.isabs(documents_dir):
            project_root = Path(__file__).resolve().parent.parent
            documents_dir = os.path.join(str(project_root), documents_dir)

    files = discover_files(documents_dir)
    if not files:
        print(f"No supported files found in {documents_dir}")
        print(f"Supported types: {', '.join(sorted(READERS.keys()))}")
        # Clear stale data from previous ingestions
        try:
            collection = get_chroma_collection(cfg)
            existing = collection.count()
            if existing > 0:
                print(f"Clearing {existing} stale chunks from previous ingestion...")
                all_ids = collection.get().get("ids", [])
                if all_ids:
                    for ci in range(0, len(all_ids), 5000):
                        collection.delete(ids=all_ids[ci:ci + 5000])
        except Exception as e:
            print(f"Warning: Could not clear vector DB: {e}")
        # Clear SQL database
        sql_enabled = cfg.get("sql", {}).get("enabled", True)
        if sql_enabled:
            try:
                sql_db_dir = cfg.get("paths", {}).get("sql_db", "sql_db")
                if not os.path.isabs(sql_db_dir):
                    project_root = Path(__file__).resolve().parent.parent
                    sql_db_dir = os.path.join(str(project_root), sql_db_dir)
                db_file = os.path.join(sql_db_dir, "knowledge_base.db")
                schema_file = os.path.join(sql_db_dir, "sql_schemas.json")
                if os.path.exists(db_file):
                    os.remove(db_file)
                    print("Cleared stale SQL database.")
                if os.path.exists(schema_file):
                    os.remove(schema_file)
            except Exception as e:
                print(f"Warning: Could not clear SQL database: {e}")
        # Clear KB meta
        try:
            db_path = cfg.get("paths", {}).get("vector_db", "chroma_db")
            if not os.path.isabs(db_path):
                project_root = Path(__file__).resolve().parent.parent
                db_path = os.path.join(str(project_root), db_path)
            cleared_any = False
            for meta_name in ("kb_meta.txt", "kb_meta_brief.txt"):
                meta_file = os.path.join(db_path, meta_name)
                if os.path.exists(meta_file):
                    os.remove(meta_file)
                    cleared_any = True
            if cleared_any:
                print("Cleared stale KB overview.")
        except Exception as e:
            print(f"Warning: Could not clear KB meta: {e}")
        return 0

    # Summarize
    datasets = {}
    for f, ds in files:
        datasets.setdefault(ds, []).append(f)
    print(f"Found {len(files)} files across {len(datasets)} dataset(s):\n")
    for ds_name, ds_files in sorted(datasets.items()):
        exts = [f.suffix for f in ds_files]
        print(f"  {ds_name}: {len(ds_files)} files ({', '.join(sorted(set(exts)))})")
    print()

    chunk_size = cfg.get("retrieval", {}).get("chunk_size", 1000)
    chunk_overlap = cfg.get("retrieval", {}).get("chunk_overlap", 100)

    # Guard: overlap >= 80% of chunk_size would cause near-duplicate chunks
    max_overlap = int(chunk_size * 0.8)
    if chunk_overlap >= max_overlap:
        print(f"Warning: chunk_overlap ({chunk_overlap}) >= 80% of chunk_size ({chunk_size}). Capping at {max_overlap}.")
        chunk_overlap = max_overlap

    collection = get_chroma_collection(cfg)

    # Defer clearing until first successful chunk generation (prevents data loss
    # if all file reads fail after clearing)
    needs_clear = True
    total_chunks = 0
    for file_path, dataset_name in files:
        rel_path = file_path.relative_to(documents_dir)
        print(f"Processing: {rel_path}")

        try:
            pages = read_file(str(file_path))
        except Exception as e:
            print(f"  Error reading {file_path.name}: {e}")
            continue

        if not pages:
            print(f"  No text extracted, skipping.")
            continue

        source_name = file_path.relative_to(Path(documents_dir)).as_posix()
        chunks = chunk_documents(pages, source_name, dataset_name, chunk_size, chunk_overlap)
        print(f"  -> {len(chunks)} chunks")

        batch_size = 5000
        safe_id_prefix = hashlib.md5(source_name.encode()).hexdigest()[:12]
        file_chunks_added = 0
        for i in range(0, len(chunks), batch_size):
            batch = chunks[i:i + batch_size]
            ids = [f"{safe_id_prefix}_{i + j}" for j in range(len(batch))]
            documents = [c.get("text", "") for c in batch]
            metadatas = [c.get("metadata", {}) for c in batch]

            # emb_model = SentenceTransformer("all-MiniLM-L6-v2")
            # # print("computing embeddings")
            # precomputed_embeddings = emb_model.encode(documents, show_progress_bar=True).tolist()

            # print("adding docs")
            try:
                # Clear existing data only on first successful add
                if needs_clear:
                    existing = collection.count()
                    if existing > 0:
                        print(f"Clearing {existing} existing chunks...\n")
                        all_ids = collection.get().get("ids", [])
                        if all_ids:
                            for ci in range(0, len(all_ids), 5000):
                                collection.delete(ids=all_ids[ci:ci + 5000])
                    needs_clear = False
                collection.add(ids=ids, documents=documents, metadatas=metadatas)
                file_chunks_added += len(batch)
            except Exception as e:
                print(f"  Warning: Failed to embed batch for {source_name}: {e}")

        total_chunks += file_chunks_added

    # ── SQL ingestion for tabular files ──────────────────────────────
    sql_enabled = cfg.get("sql", {}).get("enabled", True)
    if sql_enabled:
        from src.sql_ingest import ingest_to_sql, SQL_EXTENSIONS
        tabular_count = sum(1 for f, _ in files if f.suffix.lower() in SQL_EXTENSIONS)
        if tabular_count > 0:
            print(f"\nIngesting {tabular_count} tabular file(s) into SQLite...")
            try:
                schema = ingest_to_sql(files, documents_dir, cfg)
                print(f"SQL ingestion complete: {len(schema)} table(s).")
            except Exception as e:
                print(f"SQL ingestion error (non-fatal): {e}")

    # ── KB meta overview (LLM-generated) ────────────────────────────
    _db_path = cfg.get("paths", {}).get("vector_db", "chroma_db")
    if not os.path.isabs(_db_path):
        _db_path = os.path.join(str(Path(__file__).resolve().parent.parent), _db_path)

    print("\nGenerating knowledge base overview...")
    try:
        from src.kb_meta import build_and_store_overview
        # build_and_store_overview overwrites meta files via write_text();
        # no pre-deletion needed — old files stay intact if generation fails.
        overview = build_and_store_overview(collection, cfg)
        if overview:
            print("KB overview generated and indexed.")
        else:
            print("KB overview: nothing to summarize.")
    except Exception as e:
        print(f"KB overview generation failed (non-fatal): {e}")

    print(f"\nIngestion complete: {total_chunks} chunks from {len(files)} files.")
    return total_chunks