assemble_chroma.py

from __future__ import annotations

import hashlib
import tarfile
from pathlib import Path


ROOT = Path(__file__).resolve().parent
CHROMA_DIR = ROOT / "chroma_db_docs"
PARTS_DIR = ROOT / "chroma_package"
ARCHIVE_PATH = ROOT / "chroma_db_docs.tar"
DATABASE_PATH = CHROMA_DIR / "chroma.sqlite3"
EXPECTED_ARCHIVE_SHA256 = "656a7155238ba2b6ec99baad3fc22698a4765696ba5cad613a6469ab8dba8add"
EXPECTED_DATABASE_SHA256 = "55b525aca0298422c27393485f9feb146ceb2119fb60f382c81ec44e2b670db6"


def main() -> None:
    parts = sorted(PARTS_DIR.glob("chroma.tar.part*"))
    if not parts:
        raise FileNotFoundError(f"No packaged Chroma archive parts found in {PARTS_DIR}")

    digest = hashlib.sha256()
    with ARCHIVE_PATH.open("wb") as output:
        for part in parts:
            payload = part.read_bytes()
            output.write(payload)
            digest.update(payload)

    actual_sha256 = digest.hexdigest()
    if actual_sha256 != EXPECTED_ARCHIVE_SHA256:
        ARCHIVE_PATH.unlink(missing_ok=True)
        raise RuntimeError(
            f"Invalid assembled Chroma archive: expected {EXPECTED_ARCHIVE_SHA256}, got {actual_sha256}"
        )

    with tarfile.open(ARCHIVE_PATH) as archive:
        archive.extractall(ROOT, filter="data")
    ARCHIVE_PATH.unlink()

    database_sha256 = hashlib.sha256(DATABASE_PATH.read_bytes()).hexdigest()
    if database_sha256 != EXPECTED_DATABASE_SHA256:
        raise RuntimeError(f"Invalid extracted Chroma database: {database_sha256}")
    print(f"Assembled and extracted packaged Chroma database from {len(parts)} parts.")


if __name__ == "__main__":
    main()