Spaces:
Sleeping
Sleeping
| 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() | |