Spaces:
Running on Zero
Running on Zero
| """Restore a pinned, hash-checked RAG snapshot into the Space's writable disk.""" | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| from concurrent.futures import ThreadPoolExecutor | |
| from pathlib import Path | |
| from typing import TYPE_CHECKING | |
| from huggingface_hub import hf_hub_download | |
| if TYPE_CHECKING: | |
| from shared.config import PrecedentConfig | |
| def _inside(root: Path, relative: str) -> Path: | |
| path = (root / relative).resolve() | |
| if not path.is_relative_to(root.resolve()) or Path(relative).is_absolute(): | |
| raise ValueError("Artifact path escapes the runtime directory") | |
| return path | |
| def _sha256(path: Path) -> str: | |
| digest = hashlib.sha256() | |
| with path.open("rb") as handle: | |
| while block := handle.read(4 * 1024 * 1024): | |
| digest.update(block) | |
| return digest.hexdigest() | |
| def _validate_manifest(manifest: dict, root: Path) -> list[dict]: | |
| """Reject incomplete inventories and unsafe downloads before restoring files.""" | |
| schema = manifest.get("schema") | |
| if type(schema) is not int or schema not in {1, 2} or not manifest.get("files"): | |
| raise ValueError("Unsupported or empty RAG artifact manifest") | |
| paths = set() | |
| for item in manifest["files"]: | |
| path = _inside(root, item["path"]) | |
| if path in paths: | |
| raise ValueError("Duplicate artifact paths are not allowed") | |
| paths.add(path) | |
| if ( | |
| type(item.get("size")) is not int | |
| or item["size"] < 0 | |
| or not re.fullmatch(r"[a-f0-9]{64}", str(item.get("sha256", ""))) | |
| ): | |
| raise ValueError("Every artifact needs a size and SHA-256 hash") | |
| if "mtime_ns" in item and (type(item["mtime_ns"]) is not int or item["mtime_ns"] < 0): | |
| raise ValueError("Invalid artifact timestamp") | |
| if "repo_id" in item: | |
| download_root = _inside(root, item["download_root"]) | |
| downloaded = _inside(download_root, item["repo_path"]) | |
| if downloaded != path or not re.fullmatch(r"[a-f0-9]{40}", item["revision"]): | |
| raise ValueError("Source archives need consistent paths and pinned revisions") | |
| _inside(root, manifest["source_root"]) | |
| items = ( | |
| manifest.get("snapshots") | |
| if schema == 2 | |
| else [ | |
| { | |
| key: manifest[key] | |
| for key in ("corpus_path", "index_path", "collection", "documents", "chunks") | |
| } | |
| ] | |
| ) | |
| if not isinstance(items, list) or not 1 <= len(items) <= 4: | |
| raise ValueError("Release must contain one primary and at most three supplements") | |
| names = set() | |
| for index, item in enumerate(items): | |
| name = item.get("name", "primary") | |
| if ( | |
| not re.fullmatch(r"[A-Za-z0-9._-]{1,64}", name) | |
| or name in names | |
| or (index == 0 and name != "primary") | |
| ): | |
| raise ValueError("Invalid or duplicate snapshot names") | |
| names.add(name) | |
| collection = item["collection"] | |
| if not re.fullmatch(r"[A-Za-z0-9._-]{1,128}", collection) or collection in {".", ".."}: | |
| raise ValueError("Invalid vector collection name") | |
| if ( | |
| any(type(item.get(key)) is not int or item[key] < 1 for key in ("documents", "chunks")) | |
| or item["chunks"] < item["documents"] | |
| ): | |
| raise ValueError("Invalid snapshot document/passage counts") | |
| corpus, storage = _inside(root, item["corpus_path"]), _inside(root, item["index_path"]) | |
| required = { | |
| corpus, | |
| corpus.with_suffix(".jsonl.sqlite"), | |
| storage / "chroma.sqlite3", | |
| storage / f"{collection}.lexical.sqlite", | |
| storage / f"{collection}.lexical.dense-build.json", | |
| } | |
| if schema == 2: | |
| required.add(storage / f"{collection}.registry.sqlite") | |
| if not all( | |
| re.fullmatch(r"[a-f0-9]{64}", str(item.get(key, ""))) | |
| for key in ("corpus_sha256", "registry_snapshot_id") | |
| ): | |
| raise ValueError("V2 snapshots require registry and corpus identities") | |
| declared = ( | |
| next( | |
| entry["sha256"] | |
| for entry in manifest["files"] | |
| if _inside(root, entry["path"]) == corpus | |
| ) | |
| if corpus in paths | |
| else None | |
| ) | |
| if declared != item["corpus_sha256"]: | |
| raise ValueError("Registry corpus identity differs from artifact hash") | |
| if not required <= paths: | |
| raise ValueError("Release is missing a corpus, index or required sidecar") | |
| if schema == 2: | |
| if not re.fullmatch(r"[a-f0-9]{64}", str(manifest.get("coverage_snapshot_id", ""))): | |
| raise ValueError("V2 release requires a combined coverage identity") | |
| if any( | |
| manifest.get(key) != items[0][key] | |
| for key in ("corpus_path", "index_path", "collection") | |
| ): | |
| raise ValueError("Primary snapshot configuration is inconsistent") | |
| if any( | |
| manifest.get(key) != sum(item[key] for item in items) for key in ("documents", "chunks") | |
| ): | |
| raise ValueError("Combined snapshot counts are inconsistent") | |
| return items | |
| def prepare_artifacts(*, manifest_path: Path | None = None) -> dict: | |
| repo = os.environ["RAG_ARTIFACT_REPO"] | |
| revision = os.environ["RAG_ARTIFACT_REVISION"] | |
| if not re.fullmatch(r"[a-f0-9]{40}", revision): | |
| raise ValueError("RAG_ARTIFACT_REVISION must be an immutable commit SHA") | |
| root = Path(os.environ.get("RAG_RUNTIME_DIR", "/tmp/precedentai-rag")).resolve() | |
| root.mkdir(parents=True, exist_ok=True) | |
| # Bundle the expected hashes with the release. The platform mounts the | |
| # private dataset using the Space owner's identity; no expiring OAuth | |
| # access token needs to be copied into the runtime. | |
| manifest = json.loads( | |
| (manifest_path or Path(__file__).with_name("rag-manifest.json")).read_text() | |
| ) | |
| if manifest.get("artifact_repo") != repo or manifest.get("artifact_revision") != revision: | |
| raise ValueError("The bundled manifest does not match the configured RAG release") | |
| mounted = Path(os.environ.get("RAG_ARTIFACT_MOUNT", "/mnt/rag")).resolve() | |
| items = _validate_manifest(manifest, root) | |
| missing_bytes = sum( | |
| item["size"] | |
| for item in manifest["files"] | |
| if not _inside(root, item["path"]).is_file() | |
| or _inside(root, item["path"]).stat().st_size != item["size"] | |
| ) | |
| if shutil.disk_usage(root).free < missing_bytes + 4 * 1024**3: | |
| raise RuntimeError("Insufficient Space disk for the full corpus and model") | |
| def restore(item: dict) -> None: | |
| target = _inside(root, item["path"]) | |
| download_root = _inside(root, item.get("download_root", ".")) | |
| source_repo = item.get("repo_id", repo) | |
| source_revision = item.get("revision", revision) | |
| if not re.fullmatch(r"[a-f0-9]{40}", source_revision): | |
| raise ValueError("Every source archive must have a pinned revision") | |
| valid = ( | |
| target.is_file() | |
| and target.stat().st_size == item["size"] | |
| and _sha256(target) == item["sha256"] | |
| ) | |
| if not valid: | |
| if shutil.disk_usage(root).free < item["size"] + 4 * 1024**3: | |
| raise RuntimeError("Insufficient Space disk to repair a RAG artifact") | |
| if "repo_id" not in item: | |
| source = _inside(mounted, item["path"]) | |
| target.parent.mkdir(parents=True, exist_ok=True) | |
| # Chroma and SQLite need writable files. Use a temporary file | |
| # so a interrupted restore never leaves an accepted partial copy. | |
| partial = target.with_name(target.name + ".restoring") | |
| shutil.copyfile(source, partial) | |
| if partial.stat().st_size != item["size"] or _sha256(partial) != item["sha256"]: | |
| partial.unlink() | |
| raise ValueError(f"Mounted RAG artifact failed verification: {item['path']}") | |
| partial.replace(target) | |
| else: | |
| downloaded = Path( | |
| hf_hub_download( | |
| source_repo, | |
| item.get("repo_path", item["path"]), | |
| repo_type="dataset", | |
| revision=source_revision, | |
| local_dir=download_root, | |
| token=False, | |
| force_download=target.exists(), | |
| ) | |
| ).resolve() | |
| if downloaded != target: | |
| raise ValueError("Downloaded artifact does not match its declared path") | |
| if target.stat().st_size != item["size"] or _sha256(target) != item["sha256"]: | |
| raise ValueError(f"RAG artifact failed verification: {item['path']}") | |
| # Preserve corpus/parent freshness; the dense and lexical provenance | |
| # fingerprints remain the original matched pair rather than rewritten. | |
| if "mtime_ns" in item: | |
| os.utime(target, ns=(item["mtime_ns"], item["mtime_ns"])) | |
| print(f"Verified RAG artifact: {item['path']}", flush=True) | |
| with ThreadPoolExecutor(max_workers=3) as pool: | |
| list(pool.map(restore, manifest["files"])) | |
| os.environ["LEGAL_AI_CORPUS_PATH"] = str(_inside(root, manifest["corpus_path"])) | |
| os.environ["LEGAL_AI_CHROMA_DIR"] = str(_inside(root, manifest["index_path"])) | |
| os.environ["LEGAL_AI_SOURCE_ARCHIVE_ROOT"] = str(_inside(root, manifest["source_root"])) | |
| os.environ["LEGAL_AI_CHROMA_COLLECTION"] = manifest["collection"] | |
| os.environ["LEGAL_AI_RETRIEVAL_BACKEND"] = "fused" | |
| os.environ["LEGAL_AI_ENABLE_V2_DISCOVERY"] = "true" if manifest["schema"] == 2 else "false" | |
| os.environ["LEGAL_AI_EMBEDDING_MODEL"] = manifest.get( | |
| "embedding_model", "sentence-transformers/LaBSE" | |
| ) | |
| os.environ["LEGAL_AI_SUPPLEMENTS"] = json.dumps( | |
| [ | |
| { | |
| "name": item["name"], | |
| "corpus_path": str(_inside(root, item["corpus_path"])), | |
| "index_dir": str(_inside(root, item["index_path"])), | |
| "collection": item["collection"], | |
| } | |
| for item in items[1:] | |
| ] | |
| ) | |
| return manifest | |
| def verify_catalog(config: PrecedentConfig, manifest: dict) -> None: | |
| """Require every vector collection and v2 inventory to match the release.""" | |
| from tools.corpus_catalog import catalog_coverage, snapshots | |
| from tools.document_registry import coverage_manifest | |
| from tools.semantic_retrieve import _require_chroma | |
| root = Path(os.environ.get("RAG_RUNTIME_DIR", "/tmp/precedentai-rag")) | |
| items = _validate_manifest(manifest, root) | |
| configured = snapshots(config) | |
| if len(configured) != len(items): | |
| raise RuntimeError("Configured catalog does not match the release") | |
| for snapshot, expected in zip(configured, items): | |
| if ( | |
| snapshot.corpus_path.resolve() != _inside(root, expected["corpus_path"]) | |
| or snapshot.index_dir.resolve() != _inside(root, expected["index_path"]) | |
| or snapshot.collection != expected["collection"] | |
| ): | |
| raise RuntimeError("Configured snapshot paths differ from the release") | |
| collection = ( | |
| _require_chroma() | |
| .PersistentClient(path=str(snapshot.index_dir)) | |
| .get_collection(snapshot.collection) | |
| ) | |
| if collection.count() != expected["chunks"]: | |
| raise RuntimeError("A deployed vector collection is incomplete") | |
| if (collection.metadata or {}).get("embedding_model") != config.embedding_model: | |
| raise RuntimeError("A deployed vector encoder differs from the query encoder") | |
| if manifest["schema"] == 2: | |
| coverage = coverage_manifest(snapshot.index_dir, snapshot.collection) | |
| if ( | |
| coverage["snapshot_id"] != expected["registry_snapshot_id"] | |
| or coverage["documents"] != expected["documents"] | |
| ): | |
| raise RuntimeError("A deployed registry differs from the release") | |
| if ( | |
| manifest["schema"] == 2 | |
| and catalog_coverage(config)["snapshot_id"] != manifest["coverage_snapshot_id"] | |
| ): | |
| raise RuntimeError("Deployed coverage differs from the combined release") | |