ingest.py

"""
Document Ingestion Pipeline
Încarcă documente (.md, .txt, .pdf), le sparge în chunk-uri, și le indexează în ChromaDB.
"""

import os
from pathlib import Path

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
    TextLoader,
    UnstructuredMarkdownLoader,
    PyPDFLoader,
    DirectoryLoader,
)
import chromadb
from sentence_transformers import SentenceTransformer

from config import (
    DOCS_DIR,
    CHROMA_DIR,
    EMBEDDING_MODEL,
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    COLLECTION_NAME,
)


LOADER_MAP = {
    ".txt": (TextLoader, {"encoding": "utf-8"}),
    ".md": (TextLoader, {"encoding": "utf-8"}),
    ".pdf": (PyPDFLoader, {}),
}


def load_documents(docs_dir: Path | str = DOCS_DIR) -> list:
    """Încarcă toate documentele suportate dintr-un director."""
    docs_dir = Path(docs_dir)
    if not docs_dir.exists():
        raise FileNotFoundError(f"Directorul {docs_dir} nu există. Creează-l și adaugă documente.")

    documents = []
    for file_path in sorted(docs_dir.rglob("*")):
        ext = file_path.suffix.lower()
        if ext not in LOADER_MAP:
            continue
        loader_cls, loader_kwargs = LOADER_MAP[ext]
        try:
            loader = loader_cls(str(file_path), **loader_kwargs)
            loaded = loader.load()
            for doc in loaded:
                doc.metadata["source"] = str(file_path.relative_to(docs_dir))
                doc.metadata["filename"] = file_path.name
            documents.extend(loaded)
            print(f"  ✓ Încărcat: {file_path.name} ({len(loaded)} pagini/secțiuni)")
        except Exception as e:
            print(f"  ✗ Eroare la {file_path.name}: {e}")

    print(f"\nTotal documente încărcate: {len(documents)}")
    return documents


def chunk_documents(documents: list) -> list:
    """Sparge documentele în chunk-uri mici cu overlap."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE,
        chunk_overlap=CHUNK_OVERLAP,
        separators=["\n\n", "\n", ". ", " ", ""],
        length_function=len,
    )
    chunks = splitter.split_documents(documents)
    print(f"Documente sparte în {len(chunks)} chunk-uri (size={CHUNK_SIZE}, overlap={CHUNK_OVERLAP})")
    return chunks


def index_chunks(chunks: list) -> chromadb.Collection:
    """Creează embeddings și indexează chunk-urile în ChromaDB."""
    print(f"\nÎncărcăm modelul de embedding: {EMBEDDING_MODEL} ...")
    model = SentenceTransformer(EMBEDDING_MODEL)

    texts = [chunk.page_content for chunk in chunks]
    metadatas = [chunk.metadata for chunk in chunks]
    ids = [f"chunk_{i}" for i in range(len(chunks))]

    print(f"Generăm embeddings pentru {len(texts)} chunk-uri ...")
    embeddings = model.encode(texts, show_progress_bar=True, batch_size=64)
    embeddings_list = embeddings.tolist()

    print(f"Salvăm în ChromaDB ({CHROMA_DIR}) ...")
    CHROMA_DIR.mkdir(parents=True, exist_ok=True)
    client = chromadb.PersistentClient(path=str(CHROMA_DIR))

    # Șterge colecția veche dacă există (re-indexare completă)
    try:
        client.delete_collection(COLLECTION_NAME)
    except ValueError:
        pass

    collection = client.create_collection(
        name=COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )

    # ChromaDB are limită de batch, trimitem în loturi de 5000
    batch_size = 5000
    for i in range(0, len(ids), batch_size):
        end = min(i + batch_size, len(ids))
        collection.add(
            ids=ids[i:end],
            embeddings=embeddings_list[i:end],
            documents=texts[i:end],
            metadatas=metadatas[i:end],
        )

    print(f"Indexare completă! {collection.count()} chunk-uri în colecția '{COLLECTION_NAME}'")
    return collection


def run_ingestion(docs_dir: Path | str = DOCS_DIR):
    """Pipeline complet de ingestie: load → chunk → index."""
    print("=" * 60)
    print("  RAG INGESTION PIPELINE")
    print("=" * 60)

    print(f"\n[1/3] Încărcăm documentele din: {docs_dir}")
    documents = load_documents(docs_dir)
    if not documents:
        print("Nu s-au găsit documente! Adaugă fișiere .md, .txt sau .pdf în directorul docs/")
        return

    print(f"\n[2/3] Chunking ...")
    chunks = chunk_documents(documents)

    print(f"\n[3/3] Embedding + Indexare în ChromaDB ...")
    collection = index_chunks(chunks)

    print("\n" + "=" * 60)
    print(f"  FINALIZAT - {collection.count()} chunk-uri indexate")
    print("=" * 60)


if __name__ == "__main__":
    run_ingestion()