spb/ai-atlas

1"""Raw source archive — content-addressed, gzip-compressed, deduplicated. Historical snapshots are never discarded.23Layout (under AIA_DATA_DIR):4    raw/<sha256[0:2]>/<sha256[2:4]>/<sha256>.gz      original bytes (HTML, JSON, XML, PDF…)5    text/<sha256[0:2]>/<sha256[2:4]>/<sha256>.txt.gz cleaned text6"""7from __future__ import annotations89import gzip10import hashlib11from pathlib import Path1213from aiatlas.config import settings141516def _path(root: Path, digest: str, suffix: str) -> Path:17    return root / digest[:2] / digest[2:4] / f"{digest}{suffix}"181920def store_raw(content: bytes, *, sha256: str | None = None) -> str:21    """Write bytes once (dedupe by hash); return the path relative to raw_dir."""22    digest = sha256 or hashlib.sha256(content).hexdigest()23    path = _path(settings.raw_dir, digest, ".gz")24    if not path.exists():25        path.parent.mkdir(parents=True, exist_ok=True)26        tmp = path.with_suffix(".gz.tmp")27        with gzip.open(tmp, "wb", compresslevel=6) as fh:28            fh.write(content)29        tmp.replace(path)30    return str(path.relative_to(settings.raw_dir))313233def store_text(text: str) -> tuple[str, str]:34    """Write cleaned text; return (relative path, sha256 of the text)."""35    data = text.encode("utf-8")36    digest = hashlib.sha256(data).hexdigest()37    path = _path(settings.text_dir, digest, ".txt.gz")38    if not path.exists():39        path.parent.mkdir(parents=True, exist_ok=True)40        tmp = path.with_suffix(".gz.tmp")41        with gzip.open(tmp, "wb", compresslevel=6) as fh:42            fh.write(data)43        tmp.replace(path)44    return str(path.relative_to(settings.text_dir)), digest454647def load_raw(rel_path: str) -> bytes:48    with gzip.open(settings.raw_dir / rel_path, "rb") as fh:49        return fh.read()505152def load_text(rel_path: str) -> str:53    with gzip.open(settings.text_dir / rel_path, "rb") as fh:54        return fh.read().decode("utf-8")555657def archive_size() -> dict[str, int]:58    out: dict[str, int] = {}59    for name, root in (("raw", settings.raw_dir), ("text", settings.text_dir)):60        total = 061        count = 062        if root.exists():63            for p in root.rglob("*.gz"):64                try:65                    total += p.stat().st_size66                    count += 167                except OSError:68                    pass69        out[f"{name}_bytes"] = total70        out[f"{name}_files"] = count71    return out727374__all__ = ["archive_size", "load_raw", "load_text", "store_raw", "store_text"]75