claude-sql 0.4.0__py3-none-any.whl

claude_sql/embed_worker.py

@@ -0,0 +1,482 @@
+"""Cohere Embed v4 backfill worker for claude-sql.
+
+Discovers messages with no embedding yet, invokes ``cohere.embed-v4:0`` on
+Amazon Bedrock in parallel batches (up to 96 texts per call), and appends the
+resulting vectors to a parquet file keyed by message ``uuid``.
+
+The worker converts the int8 response to float on insert because DuckDB's VSS
+HNSW index requires ``FLOAT[]`` columns (storage loss of ~4x is accepted in
+v1 — see research notes).
+
+Tenacity retries on transient Bedrock errors via the loguru-native
+``loguru_before_sleep`` helper from ``logging_setup`` — keeps every module
+in claude-sql on a single logger (no stdlib ``logging`` imports).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import Any
+
+import boto3
+import duckdb
+import polars as pl
+from botocore.config import Config as BotoConfig
+from botocore.exceptions import (
+    ClientError,
+    ConnectionError as BotoConnectionError,
+    EndpointConnectionError,
+    ReadTimeoutError,
+    SSLError,
+)
+from loguru import logger
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+)
+
+from claude_sql.config import Settings
+from claude_sql.logging_setup import loguru_before_sleep
+from claude_sql.parquet_shards import iter_part_files, write_part
+
+#: Conservative per-text character cap before sending to Bedrock.  The real
+#: model limit is 128K tokens per text; this cap keeps total payload below
+#: the Bedrock 20 MB body ceiling even with a full batch of 96 large texts.
+MAX_CHARS_PER_TEXT = 50_000
+
+#: Bedrock error codes that tenacity should retry.
+_RETRY_CODES: set[str] = {
+    "ThrottlingException",
+    "ServiceUnavailableException",
+    "ModelTimeoutException",
+    "ModelErrorException",
+}
+
+
+def _is_retryable(exc: BaseException) -> bool:
+    """Return True if ``exc`` is a Bedrock error worth retrying.
+
+    Two buckets:
+    * ``ClientError`` with a code in :data:`_RETRY_CODES` — service-level
+      throttling and transient model failures.
+    * Network-layer errors (SSL, connection, endpoint, read-timeout) that
+      surface when long-running batches hit flaky TCP connections.
+    """
+    if isinstance(exc, SSLError | BotoConnectionError | EndpointConnectionError | ReadTimeoutError):
+        return True
+    if not isinstance(exc, ClientError):
+        return False
+    code = exc.response.get("Error", {}).get("Code")
+    return code in _RETRY_CODES
+
+
+def discover_unembedded(
+    con: duckdb.DuckDBPyConnection,
+    *,
+    embeddings_parquet: Path,
+    since_days: int | None = None,
+    limit: int | None = None,
+) -> list[tuple[str, str]]:
+    """Return ``(uuid, text)`` pairs that have no embedding yet.
+
+    Parameters
+    ----------
+    con
+        An open DuckDB connection with the ``messages_text`` view registered.
+    embeddings_parquet
+        Path to the parquet of already-embedded rows; may not exist yet.
+    since_days
+        If given, only include messages with ``ts >= now() - since_days``.
+    limit
+        Optional row cap.
+
+    Returns
+    -------
+    list of (uuid, text) tuples
+        Messages needing embedding, in DuckDB's scan order.
+    """
+    # Treat missing / truncated parquet files (or empty shard directories) as
+    # "no embeddings yet" so an aborted previous run doesn't lock discovery
+    # into "skip all" via a corrupt index.  Sharded directories are scanned
+    # via the part-file glob; legacy single files keep their original path.
+    parts = iter_part_files(embeddings_parquet)
+    parts = [p for p in parts if p.stat().st_size > 16]
+    if parts:
+        # CREATE VIEW doesn't accept prepared parameters in DuckDB; escape
+        # each part path inline as a SQL string literal and pass them all
+        # to ``read_parquet`` as a list.
+        path_literals = ", ".join(f"'{str(p).replace(chr(39), chr(39) * 2)}'" for p in parts)
+        con.execute(
+            "CREATE OR REPLACE TEMP VIEW _embedded AS "
+            f"SELECT uuid FROM read_parquet([{path_literals}]);"
+        )
+        # mt.uuid is typed UUID; parquet uuid column is VARCHAR. Cast to match.
+        anti = "AND CAST(mt.uuid AS VARCHAR) NOT IN (SELECT uuid FROM _embedded)"
+    else:
+        anti = ""
+
+    where = ["mt.text_content IS NOT NULL", "length(mt.text_content) > 0"]
+    if since_days is not None:
+        # DuckDB refuses to prepare an INTERVAL parameter; inline the coerced int.
+        where.append(f"mt.ts >= current_timestamp - INTERVAL {int(since_days)} DAY")
+
+    sql = (
+        f"SELECT mt.uuid, mt.text_content FROM messages_text mt WHERE {' AND '.join(where)} {anti}"
+    )
+    if limit is not None:
+        sql += f"\nLIMIT {int(limit)}"
+
+    rows = con.execute(sql).fetchall()
+    # DuckDB returns UUIDs as uuid.UUID objects; polars wants str for pl.Utf8.
+    return [(str(r[0]), r[1]) for r in rows]
+
+
+def _build_bedrock_client(settings: Settings) -> Any:
+    """Construct a boto3 ``bedrock-runtime`` client from settings.
+
+    Parameters
+    ----------
+    settings
+        Application settings providing the target AWS region.
+
+    Returns
+    -------
+    botocore client
+        A low-level ``bedrock-runtime`` client.
+    """
+    # Disable botocore's internal retry layer so tenacity sees throttling
+    # immediately — otherwise botocore silently absorbs 4 retries and our
+    # retry policy never kicks in.  Also bump read_timeout for large batches.
+    boto_cfg = BotoConfig(
+        region_name=settings.region,
+        retries={"max_attempts": 0, "mode": "standard"},
+        read_timeout=60,
+        connect_timeout=10,
+    )
+    return boto3.client("bedrock-runtime", config=boto_cfg)
+
+
+@retry(
+    # Cohere Embed v4 on Bedrock has a strict TPM bucket that replenishes over
+    # tens of seconds; wait up to 60s between attempts and try up to 10 times
+    # before surfacing the ThrottlingException.
+    stop=stop_after_attempt(10),
+    wait=wait_exponential(multiplier=2, min=2, max=60),
+    retry=retry_if_exception(_is_retryable),
+    before_sleep=loguru_before_sleep("WARNING"),
+    reraise=True,
+)
+def _invoke_bedrock_sync(
+    client: Any,
+    model_id: str,
+    texts: list[str],
+    *,
+    input_type: str,
+    output_dimension: int,
+    embedding_type: str,
+) -> list[list[int]] | list[list[float]]:
+    """Make one synchronous ``invoke_model`` call and return the vectors.
+
+    Parameters
+    ----------
+    client
+        A boto3 ``bedrock-runtime`` client.
+    model_id
+        Cohere Embed v4 model ID (direct or CRIS profile).
+    texts
+        Up to 96 strings; each is clipped to ``MAX_CHARS_PER_TEXT``.
+    input_type
+        Either ``"search_document"`` (corpus) or ``"search_query"``.
+    output_dimension
+        Target Matryoshka dimension: 256, 512, 1024, or 1536.
+    embedding_type
+        One of ``"int8"``, ``"float"``, ``"uint8"``, ``"binary"``, ``"ubinary"``.
+
+    Returns
+    -------
+    list of list of int or float
+        Flat list of vectors matching the order of ``texts``.
+    """
+    body = json.dumps(
+        {
+            "texts": [t[:MAX_CHARS_PER_TEXT] for t in texts],
+            "input_type": input_type,
+            "output_dimension": output_dimension,
+            "embedding_types": [embedding_type],
+            "truncate": "RIGHT",
+        }
+    )
+    resp = client.invoke_model(
+        modelId=model_id,
+        body=body,
+        contentType="application/json",
+        accept="application/json",
+    )
+    payload = json.loads(resp["body"].read())
+    return payload["embeddings"][embedding_type]
+
+
+async def _embed_one_batch(
+    client: Any,
+    texts: list[str],
+    model_id: str,
+    *,
+    input_type: str,
+    output_dimension: int,
+    embedding_type: str,
+    sem: asyncio.Semaphore,
+) -> list[list[int]] | list[list[float]]:
+    """Embed a single batch under a concurrency-limiting semaphore."""
+    async with sem:
+        return await asyncio.to_thread(
+            _invoke_bedrock_sync,
+            client,
+            model_id,
+            texts,
+            input_type=input_type,
+            output_dimension=output_dimension,
+            embedding_type=embedding_type,
+        )
+
+
+async def embed_documents_async(
+    texts: list[str],
+    *,
+    settings: Settings,
+) -> list[list[float]]:
+    """Embed corpus documents in parallel and return ``float`` vectors.
+
+    Uses ``input_type="search_document"``.  The Bedrock response (int8 or
+    other type per ``settings.embedding_type``) is cast to ``float`` so the
+    downstream parquet / DuckDB ``FLOAT[]`` column is directly consumable by
+    the VSS HNSW index.
+
+    Parameters
+    ----------
+    texts
+        Full corpus to embed; this function handles batching and concurrency.
+    settings
+        Application settings (model, batch size, concurrency, output dim).
+
+    Returns
+    -------
+    list of list of float
+        One vector per input text, same order.
+    """
+    if not texts:
+        return []
+
+    client = _build_bedrock_client(settings)
+    batch_size = settings.batch_size
+    batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]
+    sem = asyncio.Semaphore(settings.embed_concurrency)
+
+    logger.info(
+        "Embedding {} texts in {} batches (batch_size={}, concurrency={}, model={})",
+        len(texts),
+        len(batches),
+        batch_size,
+        settings.embed_concurrency,
+        settings.active_model_id,
+    )
+
+    t0 = time.monotonic()
+    coros = [
+        _embed_one_batch(
+            client,
+            batch,
+            settings.active_model_id,
+            input_type="search_document",
+            output_dimension=settings.output_dimension,
+            embedding_type=settings.embedding_type,
+            sem=sem,
+        )
+        for batch in batches
+    ]
+    results = await asyncio.gather(*coros)
+    elapsed = time.monotonic() - t0
+
+    vectors: list[list[float]] = [
+        [float(x) for x in v] for batch_vecs in results for v in batch_vecs
+    ]
+    logger.info(
+        "Embedded {} vectors across {} batches in {:.2f}s ({:.1f} vec/s)",
+        len(vectors),
+        len(batches),
+        elapsed,
+        len(vectors) / elapsed if elapsed > 0 else 0.0,
+    )
+    return vectors
+
+
+def embed_query(text: str, *, settings: Settings) -> list[float]:
+    """Embed a single query string for HNSW nearest-neighbor search.
+
+    Uses ``input_type="search_query"`` and forces ``embedding_type="float"``
+    regardless of ``settings.embedding_type`` because HNSW distance math
+    needs float vectors.
+
+    Parameters
+    ----------
+    text
+        The user's natural-language query.
+    settings
+        Application settings (model, output dim, region).
+
+    Returns
+    -------
+    list of float
+        A single vector of length ``settings.output_dimension``.
+    """
+    client = _build_bedrock_client(settings)
+    vectors = _invoke_bedrock_sync(
+        client,
+        settings.active_model_id,
+        [text],
+        input_type="search_query",
+        output_dimension=settings.output_dimension,
+        embedding_type="float",
+    )
+    return [float(x) for x in vectors[0]]
+
+
+async def run_backfill(
+    *,
+    con: duckdb.DuckDBPyConnection,
+    settings: Settings,
+    since_days: int | None = None,
+    limit: int | None = None,
+    dry_run: bool = False,
+) -> int | dict[str, Any]:
+    """Discover unembedded messages, embed them, and append to parquet.
+
+    Parameters
+    ----------
+    con
+        An open DuckDB connection with ``messages_text`` registered.
+    settings
+        Application settings (model, batch size, concurrency, parquet path).
+    since_days
+        If given, only consider messages newer than this many days.
+    limit
+        Optional cap on number of messages to embed this run.
+    dry_run
+        If true, log the plan and return a plan dict without calling Bedrock.
+
+    Returns
+    -------
+    int | dict
+        Under ``dry_run=True``, a plan dict with ``{pipeline, candidates,
+        batches, batch_size, concurrency, model, since_days, limit}``.
+        Otherwise, count of newly embedded rows (0 when nothing is pending).
+    """
+    pending = discover_unembedded(
+        con,
+        embeddings_parquet=settings.embeddings_parquet_path,
+        since_days=since_days,
+        limit=limit,
+    )
+    if not pending:
+        logger.info("No unembedded messages found - nothing to do")
+        if dry_run:
+            return {
+                "pipeline": "embed",
+                "candidates": 0,
+                "batches": 0,
+                "batch_size": settings.batch_size,
+                "concurrency": settings.embed_concurrency,
+                "model": settings.active_model_id,
+                "since_days": since_days,
+                "limit": limit,
+                "dry_run": True,
+            }
+        return 0
+
+    n_batches = (len(pending) + settings.batch_size - 1) // settings.batch_size
+    logger.info(
+        "Backfill plan: {} messages, {} batches, concurrency={}, model={}",
+        len(pending),
+        n_batches,
+        settings.embed_concurrency,
+        settings.active_model_id,
+    )
+    if dry_run:
+        logger.info("dry_run=True - skipping Bedrock calls")
+        return {
+            "pipeline": "embed",
+            "candidates": len(pending),
+            "batches": n_batches,
+            "batch_size": settings.batch_size,
+            "concurrency": settings.embed_concurrency,
+            "model": settings.active_model_id,
+            "since_days": since_days,
+            "limit": limit,
+            "dry_run": True,
+        }
+
+    # Checkpoint every N messages so a throttling-induced timeout doesn't
+    # discard work already embedded.  chunk must be a multiple of batch_size.
+    chunk_size = max(settings.batch_size * 4, 256)
+    path = settings.embeddings_parquet_path
+    total_t0 = time.monotonic()
+    written = 0
+    for i in range(0, len(pending), chunk_size):
+        slice_ = pending[i : i + chunk_size]
+        logger.info(
+            "Chunk {}/{}: embedding {} messages",
+            i // chunk_size + 1,
+            (len(pending) + chunk_size - 1) // chunk_size,
+            len(slice_),
+        )
+        t0 = time.monotonic()
+        texts = [p[1] for p in slice_]
+        vectors = await embed_documents_async(texts, settings=settings)
+        elapsed = time.monotonic() - t0
+        logger.info(
+            "Chunk done in {:.1f}s ({:.1f} vec/s)",
+            elapsed,
+            len(vectors) / elapsed if elapsed > 0 else 0.0,
+        )
+
+        now = datetime.now(UTC)
+        # Polars infers nested list[float] as Object when the batch is small or
+        # when rows are handed in as Python lists; force a fixed-size Array so
+        # write_parquet succeeds and DuckDB VSS sees FLOAT[dim] on read.
+        df = pl.DataFrame(
+            {
+                "uuid": [p[0] for p in slice_],
+                "model": [settings.active_model_id] * len(slice_),
+                "dim": [settings.output_dimension] * len(slice_),
+                "embedding": vectors,
+                "embedded_at": [now] * len(slice_),
+            },
+            schema={
+                "uuid": pl.Utf8,
+                "model": pl.Utf8,
+                "dim": pl.UInt16,
+                "embedding": pl.Array(pl.Float32, settings.output_dimension),
+                "embedded_at": pl.Datetime("us", "UTC"),
+            },
+        )
+        # Sharded write: drop a fresh ``part-<ts_ns>.parquet`` into the
+        # embeddings directory.  Legacy single-file caches still rewrite the
+        # whole file (handled inside ``write_part``) so existing installs
+        # keep working until they're migrated.
+        written_path = write_part(path, df)
+        written += len(slice_)
+        logger.info("Checkpoint: {} rows -> {}", len(df), written_path)
+
+    total_elapsed = time.monotonic() - total_t0
+    logger.info(
+        "Backfill complete: {} embeddings in {:.1f}s ({:.1f} vec/s overall)",
+        written,
+        total_elapsed,
+        written / total_elapsed if total_elapsed > 0 else 0.0,
+    )
+    return written

claude_sql/freeze.py

@@ -0,0 +1,189 @@
+"""Pre-registration: freeze + replay study manifests.
+
+The ``freeze`` subcommand hashes the full study spec (rubric YAML,
+judge panel, commit SHA, embed model, session-scoping rule) into a
+deterministic manifest SHA and writes it under ``~/.claude/studies/<sha>/``.
+
+The ``replay`` subcommand reads a manifest by SHA and rebuilds a
+``Study`` object so downstream workers (``judge``, ``ungrounded-claim``,
+``kappa``) can execute with the exact locked parameters.
+
+This is the IRR study's audit trail.  Every parquet the workers write
+carries the manifest SHA in a ``freeze_sha`` column.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import subprocess
+from dataclasses import asdict, dataclass, field
+from datetime import UTC, datetime
+from pathlib import Path
+
+from claude_sql import judges as judge_catalog
+
+
+@dataclass(frozen=True)
+class SessionScope:
+    """Session-scoping rule for the study."""
+
+    min_turns: int = 10
+    max_turns: int = 40
+    max_interrupt_minutes: int = 15
+    kind: str = "mechanical"
+
+
+@dataclass(frozen=True)
+class Study:
+    """A pre-registered IRR study specification."""
+
+    rubric_path: str
+    rubric_content_hash: str
+    panel_shortnames: tuple[str, ...]
+    panel_model_ids: tuple[str, ...]
+    embed_model_id: str
+    commit_sha: str
+    session_scope: SessionScope
+    seed: int
+    created_at_utc: str = field(
+        default_factory=lambda: datetime.now(UTC).isoformat(timespec="seconds")
+    )
+
+    @property
+    def manifest_sha(self) -> str:
+        """Deterministic SHA256 over the parameterised fields (excludes created_at)."""
+        payload = {
+            "rubric_content_hash": self.rubric_content_hash,
+            "panel_model_ids": list(self.panel_model_ids),
+            "embed_model_id": self.embed_model_id,
+            "commit_sha": self.commit_sha,
+            "session_scope": asdict(self.session_scope),
+            "seed": self.seed,
+        }
+        return hashlib.sha256(json.dumps(payload, sort_keys=True).encode("utf-8")).hexdigest()[:16]
+
+    def to_dict(self) -> dict[str, object]:
+        d = asdict(self)
+        d["manifest_sha"] = self.manifest_sha
+        return d
+
+
+def _read_rubric(rubric_path: Path) -> tuple[str, str]:
+    """Read rubric YAML/JSON from disk and return (content, SHA256 hash).
+
+    Supports YAML (loaded as text, no parsing) and JSON; we hash the
+    raw bytes so whitespace differences produce different manifests,
+    which is what we want — even a reformatted rubric is a new study.
+    """
+    if not rubric_path.exists():
+        raise FileNotFoundError(f"rubric not found: {rubric_path}")
+    content = rubric_path.read_text(encoding="utf-8")
+    h = hashlib.sha256(content.encode("utf-8")).hexdigest()
+    return content, h
+
+
+def _git_commit_sha(repo: Path) -> str:
+    """Return the current ``git rev-parse HEAD`` for ``repo``, or ``"<dirty>"``."""
+    try:
+        out = subprocess.run(
+            ["git", "-C", str(repo), "rev-parse", "HEAD"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        sha = out.stdout.strip()
+        # Flag dirty working trees so the commit hash doesn't over-claim reproducibility
+        dirty = subprocess.run(
+            ["git", "-C", str(repo), "status", "--porcelain"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return sha if not dirty.stdout.strip() else f"{sha}-dirty"
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        return "<no-git>"
+
+
+def _studies_root() -> Path:
+    root = Path(os.path.expanduser("~/.claude/studies"))
+    root.mkdir(parents=True, exist_ok=True)
+    return root
+
+
+def freeze(
+    rubric_path: Path,
+    panel_shortnames: tuple[str, ...],
+    embed_model_id: str = "global.cohere.embed-v4:0",
+    session_scope: SessionScope | None = None,
+    seed: int = 42,
+    repo: Path | None = None,
+) -> Study:
+    """Create and persist a study manifest.
+
+    Writes ``~/.claude/studies/<sha>/manifest.json`` plus a copy of the
+    rubric so the manifest is self-contained.  Returns the ``Study``.
+    """
+    rubric_content, rubric_hash = _read_rubric(rubric_path)
+    judges_resolved = judge_catalog.panel(list(panel_shortnames))
+    panel_model_ids = tuple(j.model_id for j in judges_resolved)
+    scope = session_scope or SessionScope()
+    study = Study(
+        rubric_path=str(rubric_path),
+        rubric_content_hash=rubric_hash,
+        panel_shortnames=tuple(panel_shortnames),
+        panel_model_ids=panel_model_ids,
+        embed_model_id=embed_model_id,
+        commit_sha=_git_commit_sha(repo or Path.cwd()),
+        session_scope=scope,
+        seed=seed,
+    )
+    study_dir = _studies_root() / study.manifest_sha
+    study_dir.mkdir(parents=True, exist_ok=True)
+    (study_dir / "manifest.json").write_text(
+        json.dumps(study.to_dict(), indent=2, sort_keys=True), encoding="utf-8"
+    )
+    (study_dir / "rubric.yaml").write_text(rubric_content, encoding="utf-8")
+    return study
+
+
+def replay(manifest_sha: str) -> Study:
+    """Load a previously-frozen study by its manifest SHA."""
+    study_dir = _studies_root() / manifest_sha
+    manifest_path = study_dir / "manifest.json"
+    if not manifest_path.exists():
+        raise FileNotFoundError(f"no manifest at {manifest_path}")
+    d = json.loads(manifest_path.read_text(encoding="utf-8"))
+    scope = SessionScope(**d["session_scope"])
+    return Study(
+        rubric_path=d["rubric_path"],
+        rubric_content_hash=d["rubric_content_hash"],
+        panel_shortnames=tuple(d["panel_shortnames"]),
+        panel_model_ids=tuple(d["panel_model_ids"]),
+        embed_model_id=d["embed_model_id"],
+        commit_sha=d["commit_sha"],
+        session_scope=scope,
+        seed=d["seed"],
+        created_at_utc=d["created_at_utc"],
+    )
+
+
+def list_studies() -> list[dict[str, object]]:
+    """Return a summary of every frozen study under ``~/.claude/studies/``."""
+    root = _studies_root()
+    out: list[dict[str, object]] = []
+    for d in sorted(root.iterdir()):
+        mf = d / "manifest.json"
+        if not mf.exists():
+            continue
+        payload = json.loads(mf.read_text(encoding="utf-8"))
+        out.append(
+            {
+                "manifest_sha": payload["manifest_sha"],
+                "created_at_utc": payload["created_at_utc"],
+                "commit_sha": payload["commit_sha"],
+                "n_judges": len(payload["panel_shortnames"]),
+            }
+        )
+    return out