index_documents_to_chroma.py

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any

import chromadb
from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from pypdf import PdfReader

from hf_text_embeddings import HFTextEmbeddings


DEFAULT_PERSIST_DIR = "./chroma_db_docs"
DEFAULT_COLLECTION = "document_context"
DEFAULT_DOCUMENTS_DIR = "documentos"
PDF_EXTENSIONS = {".pdf"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Index local documents into Chroma for RAG.")
    parser.add_argument("--documents-dir", default=DEFAULT_DOCUMENTS_DIR)
    parser.add_argument("--files", help="Comma-separated document paths. Defaults to all PDFs under documents-dir.")
    parser.add_argument("--persist-dir", default=DEFAULT_PERSIST_DIR)
    parser.add_argument("--collection", default=DEFAULT_COLLECTION)
    parser.add_argument("--chunk-size", type=int, default=1200)
    parser.add_argument("--chunk-overlap", type=int, default=180)
    parser.add_argument("--reset-collection", action="store_true")
    return parser.parse_args()


def iter_document_paths(documents_dir: str, files: str | None) -> list[Path]:
    if files:
        return [Path(item.strip()) for item in files.split(",") if item.strip()]

    root = Path(documents_dir)
    if not root.exists():
        return []

    return [
        path
        for path in sorted(root.rglob("*"))
        if path.is_file() and path.suffix.lower() in PDF_EXTENSIONS
    ]


def clean_pdf_text(text: str) -> str:
    lines = [line.strip() for line in text.splitlines()]
    return "\n".join(line for line in lines if line)


def metadata_path_for(path: Path) -> Path:
    return path.with_suffix(".metadata.json")


def load_document_metadata(path: Path) -> dict[str, Any]:
    metadata_path = metadata_path_for(path)
    if not metadata_path.exists():
        return {}
    payload = json.loads(metadata_path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError(f"Invalid metadata object in {metadata_path}.")
    return payload


def load_pdf(path: Path) -> list[Document]:
    reader = PdfReader(str(path))
    docs: list[Document] = []
    document_metadata = load_document_metadata(path)
    for page_index, page in enumerate(reader.pages, start=1):
        text = clean_pdf_text(page.extract_text() or "")
        if not text:
            continue
        docs.append(
            Document(
                page_content=text,
                metadata={
                    "source_type": "document",
                    "source": str(path),
                    "filename": path.name,
                    "page": page_index,
                    "document_title": path.stem,
                    **document_metadata,
                },
            )
        )
    return docs


def stable_id(doc: Document, chunk_index: int) -> str:
    source = doc.metadata.get("source", "")
    page = doc.metadata.get("page", "")
    digest = hashlib.sha1(f"{source}|{page}|{chunk_index}|{doc.page_content}".encode("utf-8")).hexdigest()
    return f"document:{Path(str(source)).stem}:p{page}:{digest[:16]}"


def main() -> None:
    load_dotenv()
    args = parse_args()

    paths = iter_document_paths(args.documents_dir, args.files)
    if not paths:
        raise RuntimeError("No PDF documents found to index.")

    if args.reset_collection:
        client = chromadb.PersistentClient(path=args.persist_dir)
        try:
            client.delete_collection(args.collection)
            print(f"Deleted existing Chroma collection: {args.collection}")
        except Exception:
            print(f"Chroma collection did not exist yet: {args.collection}")

    embeddings = HFTextEmbeddings()
    vectorstore = Chroma(
        collection_name=args.collection,
        embedding_function=embeddings,
        persist_directory=args.persist_dir,
    )
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=args.chunk_size,
        chunk_overlap=args.chunk_overlap,
    )

    total_pages = 0
    total_chunks = 0
    for path in paths:
        docs = load_pdf(path)
        splits = splitter.split_documents(docs)
        ids = [stable_id(doc, index) for index, doc in enumerate(splits)]
        try:
            vectorstore.delete(ids=ids)
        except Exception:
            pass
        vectorstore.add_documents(splits, ids=ids)
        total_pages += len(docs)
        total_chunks += len(splits)
        print(f"Indexed {path}: {len(docs)} pages, {len(splits)} chunks")

    print(f"Done. Added {total_chunks} chunks from {total_pages} pages into {args.persist_dir} / {args.collection}.")


if __name__ == "__main__":
    main()