diff --git a/arborist/cli.py b/arborist/cli.py index ee1230f..abf9467 100644 --- a/arborist/cli.py +++ b/arborist/cli.py @@ -7510,12 +7510,27 @@ def _cmd_cloud_ask(args: argparse.Namespace) -> int: if getattr(args, "bucket_url", None): from arborist.wallet.bucket import ( MultiShardBucketCorpus, + MultiShardSidecarCorpus, load_bucket_manifest, ) - manifest = load_bucket_manifest(args.bucket_url) - client = MultiShardBucketCorpus( - manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024, - ) + manifest_url = args.bucket_url + # Convenience: allow passing the manifest path directly. + if not manifest_url.rstrip("/").endswith(".json"): + manifest_url = manifest_url # base URL → load_bucket_manifest appends path + manifest = load_bucket_manifest(manifest_url) + # Pick the right corpus class based on whether the manifest + # advertises per-shard sidecars. Mixed (some sidecar, some not) + # also uses the sidecar-aware path — it falls back to bucket- + # direct FTS5 for shards without one. + any_sidecar = any(sh.get("sidecar_url") for sh in manifest.shards) + if any_sidecar: + client = MultiShardSidecarCorpus( + manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024, + ) + else: + client = MultiShardBucketCorpus( + manifest, cache_bytes_per_shard=args.cache_mb * 1024 * 1024, + ) multi_shard = True else: if not getattr(args, "shard_url", None): diff --git a/arborist/wallet/bucket.py b/arborist/wallet/bucket.py index 9d83fe5..3dcfca7 100644 --- a/arborist/wallet/bucket.py +++ b/arborist/wallet/bucket.py @@ -542,15 +542,20 @@ class BucketManifest: def load_bucket_manifest(bucket_url: str, *, timeout_s: float = 30.0) -> BucketManifest: - """Fetch + parse clones/manifest.json from a bucket base URL. + """Fetch + parse a manifest from a bucket URL. - bucket_url is the path that resolves to the manifest's parent — - typically just the bucket root (e.g. https://bucket.example.com/). - The client appends ``clones/manifest.json`` and parses the response. + Accepts either: + - bucket root (e.g. https://bucket.example.com) → client appends + ``clones/manifest.json`` + - explicit manifest URL (e.g. .../clones/manifest-sidecar.json) → + used as-is, lets operators pin a non-default manifest """ import json as _json base = bucket_url.rstrip("/") - manifest_url = f"{base}/clones/manifest.json" + if base.endswith(".json"): + manifest_url = base + else: + manifest_url = f"{base}/clones/manifest.json" req = urllib.request.Request(manifest_url) with urllib.request.urlopen(req, timeout=timeout_s) as resp: body = _json.loads(resp.read().decode("utf-8")) @@ -667,3 +672,240 @@ class MultiShardBucketCorpus: "bytes_fetched": (s.get("cache") or {}).get("bytes_fetched"), }) return merged + + +# --------------------------------------------------------------------------- +# Sidecar-aware corpus: download per-shard sidecar.bin once, query locally. +# --------------------------------------------------------------------------- + + +def _sidecar_cache_path(url: str, *, base_dir: str | None = None) -> str: + """Stable local cache path for a sidecar URL. + + Hashes the URL → places under ~/.arborist/sidecar-cache/.bin so + re-runs reuse the download. base_dir overrides for tests. + """ + import hashlib + if base_dir is None: + base_dir = os.path.join( + os.path.expanduser("~"), ".arborist", "sidecar-cache" + ) + os.makedirs(base_dir, exist_ok=True) + h = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16] + return os.path.join(base_dir, f"{h}.bin") + + +def _download_sidecar(url: str, dest: str, *, timeout_s: float = 300.0) -> int: + """Stream a sidecar URL into `dest`. Returns bytes downloaded. + + Writes to dest.part then renames atomically — interrupted downloads + don't corrupt the cache. The cache file is content-addressed by URL + hash; if you change the bucket-side artifact you must change the + URL (or clear ~/.arborist/sidecar-cache/) — sidecars are not + Merkle-bound, so the URL is the identity.""" + import shutil + tmp = dest + ".part" + req = urllib.request.Request(url) + n = 0 + with urllib.request.urlopen(req, timeout=timeout_s) as resp, open(tmp, "wb") as out: + while True: + chunk = resp.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + n += len(chunk) + shutil.move(tmp, dest) + return n + + +class SidecarShardClient: + """One shard's view: a local sidecar (download-then-query-locally) + plus a BucketClient against the shard .db for chunk content reads. + + fts_search uses the sidecar (sub-second BM25); chunks_for_doc and + SQL go through the bucket VFS so chunk content fetches benefit + from the same on-disk page cache.""" + + def __init__( + self, + shard_url: str, + sidecar_url: str, + blob_base: str, + *, + cache_bytes: int = DEFAULT_CACHE_BYTES, + sidecar_cache_dir: str | None = None, + ): + from arborist.wallet.sidecar import SidecarReader + + self.shard_url = shard_url + self.sidecar_url = sidecar_url + # 1. Download or reuse cached sidecar. + cache_path = _sidecar_cache_path(sidecar_url, base_dir=sidecar_cache_dir) + if not os.path.exists(cache_path): + _download_sidecar(sidecar_url, cache_path) + self._sidecar_path = cache_path + self.sidecar = SidecarReader(cache_path) + # 2. Bucket client for chunk content reads. + self.bucket = BucketClient( + BucketEndpoint(shard_url=shard_url, blob_base=blob_base), + cache_bytes=cache_bytes, + ) + + def close(self): + self.bucket.close() + + def fts_search(self, query: str, *, limit: int = 8, raw: bool = False) -> list[dict]: + """Sidecar-driven FTS. ``raw`` is accepted for API parity with + BucketClient.fts_search but ignored — sidecar tokenizes via + tokenize_text (matches the OR-sanitizer behavior).""" + hits = self.sidecar.search(query, limit=limit, mode="or") + return [ + { + "document_root": h.document_root, + "document_uri": h.document_uri, + "title": h.title, + "score": h.score, + } + for h in hits + ] + + @property + def conn(self): + """apsw connection for SQL pulls (chunk content). Goes through + the bucket VFS — chunk reads are cached.""" + return self.bucket.conn + + +class MultiShardSidecarCorpus: + """Manifest-driven corpus where each shard MAY have a sidecar. + + Shards with a `sidecar_url` use SidecarShardClient (download-then- + local FTS). Shards without one fall back to BucketClient (in-place + bucket-direct FTS5 — slow on multi-GB shards, but works for small + shards or when no sidecar is published yet). + + fts_search runs per-shard in parallel, merges by score, picks top-K. + Mixed-mode results are score-comparable in practice: sidecar BM25 + is on a similar magnitude (~10-30) as FTS5's BM25 (~-15 to -1) but + sign-flipped (sidecar is positive, FTS5 stored as negative). We + normalize sidecar to negative before merging so a single ORDER BY + DESC works across both. + """ + + def __init__( + self, + manifest: "BucketManifest", + *, + cache_bytes_per_shard: int = DEFAULT_CACHE_BYTES, + sidecar_cache_dir: str | None = None, + ): + self.manifest = manifest + self.shard_clients = [] + self._has_sidecar: list[bool] = [] + for sh in manifest.shards: + if sh.get("sidecar_url"): + client = SidecarShardClient( + shard_url=sh["url"], + sidecar_url=sh["sidecar_url"], + blob_base=manifest.blob_base, + cache_bytes=cache_bytes_per_shard, + sidecar_cache_dir=sidecar_cache_dir, + ) + self._has_sidecar.append(True) + else: + client = BucketClient( + BucketEndpoint(shard_url=sh["url"], blob_base=manifest.blob_base), + cache_bytes=cache_bytes_per_shard, + ) + self._has_sidecar.append(False) + self.shard_clients.append(client) + self._by_url = {sh["url"]: c for sh, c in zip(manifest.shards, self.shard_clients)} + + def close(self): + for c in self.shard_clients: + c.close() + + def fts_search( + self, query: str, *, limit: int = 8, raw: bool = False, + rrf_k: int = 60, + ) -> list[dict]: + """Run each shard's FTS in parallel, merge via reciprocal-rank-fusion. + + BM25 scores aren't comparable across heterogeneous backends + (sidecar BM25 is positive in the [5, 50] range; FTS5 BM25 is + negative in [-20, 0]). RRF sidesteps the scale mismatch + entirely: each shard ranks its own hits, the global rank is + the sum of reciprocal ranks across shards. Standard k=60 is + TREC-recommended. + + score_rrf(d) = sum over shards of 1 / (rrf_k + rank(d, shard)) + """ + from concurrent.futures import ThreadPoolExecutor + + def _one(item): + sh_url, sh_client, _ = item + try: + # Pull per-shard top-K (oversample so RRF has more to + # work with — a doc that's barely in top-K on one + # shard but #1 on another should still bubble up). + return sh_url, sh_client.fts_search(query, limit=limit * 4, raw=raw) + except Exception: + return sh_url, [] + + triples = [ + (sh["url"], self._by_url[sh["url"]], self._has_sidecar[i]) + for i, sh in enumerate(self.manifest.shards) + ] + per_shard: list[tuple[str, list[dict]]] = [] + with ThreadPoolExecutor(max_workers=max(1, len(triples))) as ex: + for sh_url, hits in ex.map(_one, triples): + per_shard.append((sh_url, hits)) + + # RRF merge: keyed by document_root so duplicates across shards + # (rare but possible — e.g. a doc replicated for redundancy) + # combine. + rrf_score: dict[str, float] = {} + hit_by_root: dict[str, dict] = {} + for sh_url, hits in per_shard: + for rank, h in enumerate(hits, start=1): + droot = h["document_root"] + rrf_score[droot] = ( + rrf_score.get(droot, 0.0) + 1.0 / (rrf_k + rank) + ) + if droot not in hit_by_root: + h["_shard_url"] = sh_url + hit_by_root[droot] = h + # Surface rrf score so callers see "this doc was strong across N shards." + for droot, h in hit_by_root.items(): + h["rrf_score"] = rrf_score[droot] + ranked = sorted( + hit_by_root.values(), + key=lambda h: -h["rrf_score"], + ) + return ranked[:limit] + + def conn_for_shard(self, shard_url: str): + return self._by_url[shard_url].conn + + def stats(self) -> dict: + merged = { + "total_http_requests": 0, + "total_bytes_fetched": 0, + "per_shard": [], + } + for sh, c, has_sc in zip( + self.manifest.shards, self.shard_clients, self._has_sidecar, + ): + bucket = c.bucket if has_sc else c # type: ignore[attr-defined] + s = bucket.stats() + merged["total_http_requests"] += s.get("http_requests") or 0 + merged["total_bytes_fetched"] += ( + (s.get("cache") or {}).get("bytes_fetched") or 0 + ) + merged["per_shard"].append({ + "label": sh.get("label", sh["url"].rsplit("/", 1)[-1]), + "sidecar": has_sc, + "http_requests": s.get("http_requests"), + "bytes_fetched": (s.get("cache") or {}).get("bytes_fetched"), + }) + return merged