"""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")