atif-sql 0.2.0__py3-none-any.whl

atif_analytics/infrastructure/parquet_cache.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Sharded parquet I/O helpers for the pipeline-append caches.
+
+Each pipeline chunk writes ``<dir>/part-<ts_ns>.parquet`` and readers glob
+the directory, so append cost is proportional to the chunk size rather than
+the cache size. The legacy single-file branch is kept for tool-side
+inspection of ad-hoc ``*.parquet`` targets (tests use it).
+
+Public API: :func:`is_sharded_dir`, :func:`write_part`, :func:`read_all`,
+:func:`iter_part_files`, :func:`count_rows`. (``replace_sessions``, the
+per-session shard rewrite, left with the trajectory pipeline, its one
+caller, on 2026-09-27.)
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+
+import polars as pl
+from loguru import logger
+
+#: Glob pattern for shard part files within a sharded cache directory.
+PART_GLOB: str = "part-*.parquet"
+
+#: Smallest byte count a readable parquet file can have. The format's own
+#: framing is ``PAR1`` + footer + 4-byte footer length + ``PAR1``, so anything
+#: below this is a file a crashed or half-flushed writer left behind, not a
+#: parquet with zero rows. Reading one raises inside polars, so the size check
+#: is what keeps a torn artifact from failing a pipeline that could skip it.
+MIN_PARQUET_BYTES: int = 16
+
+
+def is_sharded_dir(path: Path) -> bool:
+    """Return True iff ``path`` is (or should be) treated as a sharded cache directory.
+
+    Two cases qualify: ``path`` exists and is a directory, or ``path`` does
+    not exist yet and has no ``.parquet`` suffix (new caches default to the
+    directory layout; a missing ``x.parquet`` stays legacy single-file).
+    """
+    if path.exists():
+        return path.is_dir()
+    return path.suffix != ".parquet"
+
+
+def _is_whole_parquet(part: Path) -> bool:
+    """False for a file too small to carry parquet framing; warns when it drops one.
+
+    A shard at or below :data:`MIN_PARQUET_BYTES` is what a crashed or
+    half-flushed writer leaves behind, and :func:`write_part` drops new
+    shards straight into the directory readers glob, so the torn file is
+    live from the moment it appears.
+    """
+    try:
+        size = part.stat().st_size
+    except OSError as exc:
+        logger.warning("parquet cache: cannot stat {} ({}); skipping", part, exc)
+        return False
+    if size <= MIN_PARQUET_BYTES:
+        logger.warning(
+            "parquet cache: skipping torn shard {} ({} bytes, minimum {})",
+            part,
+            size,
+            MIN_PARQUET_BYTES,
+        )
+        return False
+    return True
+
+
+def iter_part_files(target: Path) -> list[Path]:
+    """Return a sorted list of whole parquet files backing ``target``.
+
+    For a sharded directory: every ``part-*.parquet`` under it, sorted by
+    name (timestamp-keyed, so the order is also chronological). For a
+    legacy single-file path: ``[target]`` if it exists, else ``[]``.
+
+    Torn files are dropped here rather than at each call site, because this
+    is the one choke point every reader shares — :func:`read_all` and
+    :func:`count_rows` — so
+    :data:`MIN_PARQUET_BYTES` governs the sharded read path and not only the
+    legacy rewrite in :func:`write_part`.
+    """
+    if not target.exists():
+        return []
+    candidates = sorted(target.glob(PART_GLOB)) if target.is_dir() else [target]
+    return [part for part in candidates if _is_whole_parquet(part)]
+
+
+def write_part(target: Path, df: pl.DataFrame) -> Path:
+    """Write ``df`` as a new shard (or rewrite the legacy single file).
+
+    Sharded branch: drop a brand-new ``part-<ns>.parquet`` into the
+    directory — no read-then-rewrite. Legacy branch: concat onto the
+    existing file and rewrite (historical behavior).
+    """
+    if is_sharded_dir(target):
+        target.mkdir(parents=True, exist_ok=True)
+        # Nanosecond timestamps are sortable, monotonic on Linux, and avoid
+        # filename collisions when two part-writes land in the same
+        # millisecond under concurrency.
+        part_path = target / f"part-{time.time_ns()}.parquet"
+        df.write_parquet(part_path)
+        return part_path
+
+    target.parent.mkdir(parents=True, exist_ok=True)
+    if target.exists() and target.stat().st_size > MIN_PARQUET_BYTES:
+        existing = pl.read_parquet(target)
+        df = pl.concat([existing, df], how="diagonal_relaxed")
+    df.write_parquet(target)
+    return target
+
+
+def read_all(
+    target: Path,
+    *,
+    columns: list[str] | None = None,
+) -> pl.DataFrame | None:
+    """Return the union of all part files (or the legacy single file).
+
+    Returns ``None`` when the cache is empty or missing. ``columns``
+    projects the read down via parquet column pushdown — callers that only
+    need a key column (the anti-join sets) never decode the wide text
+    columns.
+    """
+    parts = iter_part_files(target)
+    if not parts:
+        return None
+    return pl.read_parquet([str(p) for p in parts], columns=columns)
+
+
+def count_rows(target: Path) -> int:
+    """Total row count across every part file (parquet footers only)."""
+    parts = iter_part_files(target)
+    if not parts:
+        return 0
+    import pyarrow.parquet as pq
+
+    total = 0
+    for p in parts:
+        total += int(pq.ParquetFile(str(p)).metadata.num_rows)
+    return total
+
+
+class ParquetCache:
+    """The cache port over one sharded-parquet target.
+
+    Binds the target in the constructor; methods delegate to the
+    module-level functions so tests can monkeypatch them.
+    """
+
+    def __init__(self, target: Path) -> None:
+        self._target = target
+
+    @property
+    def target(self) -> Path:
+        """The bound cache directory / legacy file path."""
+        return self._target
+
+    def write_part(self, df: pl.DataFrame) -> Path:
+        """Append ``df`` as a new shard; return the written path."""
+        return write_part(self._target, df)
+
+    def read_all(self, *, columns: list[str] | None = None) -> pl.DataFrame | None:
+        """Union of all parts (``None`` when empty)."""
+        return read_all(self._target, columns=columns)
+
+    def count_rows(self) -> int:
+        """Total row count (footer reads only)."""
+        return count_rows(self._target)
+
+    def iter_part_files(self) -> list[Path]:
+        """Sorted list of backing parquet files."""
+        return iter_part_files(self._target)
+
+
+__all__ = [
+    "MIN_PARQUET_BYTES",
+    "PART_GLOB",
+    "ParquetCache",
+    "count_rows",
+    "is_sharded_dir",
+    "iter_part_files",
+    "read_all",
+    "write_part",
+]

atif_analytics/infrastructure/settings.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Env-driven settings for the analytics pipelines.
+
+Pydantic v2 ``BaseSettings`` under the workspace ``ATIF_SQL_`` prefix
+(atif-corpus / atif-models precedent). Carries the corpus root and the
+per-pipeline knobs (friction char cutoff, batch size, budget ceilings,
+transcript caps).
+
+Composes :class:`atif_models.infrastructure.settings.LlmSettings` for the
+LLM family/size/region/concurrency selection rather than duplicating those
+fields.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from atif_analytics.domain.config import TranscriptCaps
+from atif_analytics.domain.layout import AnalyticsLayout
+
+if TYPE_CHECKING:
+    from atif_models.infrastructure.settings import LlmSettings
+
+
+def _default_corpus_root() -> Path:
+    """Mirror atif-corpus's default: ``~/.atif-sql/corpus/<slug of source root>``.
+
+    Re-implements the slug locally (sha256 of the resolved source root,
+    ``default`` for ``~/.claude``) because atif-analytics may not import
+    atif-corpus (independence contract). The parity test pins both against
+    the same expected strings.
+    """
+    import hashlib
+    import re
+
+    env_source = os.environ.get("ATIF_SQL_SOURCE_ROOT")
+    if env_source is not None:
+        source_root = Path(env_source)
+    else:
+        config_root = Path(os.environ.get("CLAUDE_CONFIG_DIR", "~/.claude")).expanduser()
+        source_root = config_root / "projects"
+    resolved = source_root.expanduser().resolve()
+    # corpus_slug parity: a bare ~/.claude source maps to the reserved
+    # "default" key; everything else is <sanitized-dirname>-<8-hex sha256>.
+    if resolved == Path("~/.claude").expanduser().resolve():
+        slug = "default"
+    else:
+        name = re.sub(r"[^a-z0-9]+", "-", resolved.name.lower()).strip("-")[:32]
+        digest = hashlib.sha256(str(resolved).encode("utf-8")).hexdigest()[:8]
+        slug = f"{name}-{digest}" if name else digest
+    return Path.home() / ".atif-sql" / "corpus" / slug
+
+
+class AnalyticsSettings(BaseSettings):
+    """Env-driven configuration for the analytics pipelines."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="ATIF_SQL_",
+        env_file=".env",
+        extra="ignore",
+    )
+
+    #: Materialized corpus root (the directory containing ``sessions/``).
+    corpus_root: Path = Field(default_factory=_default_corpus_root)
+
+    # --- LLM pipeline knobs ---
+    #: Char cutoff for friction candidates: longer user turns are almost
+    #: always genuine instructions.
+    friction_max_chars: int = 300
+    #: Hard per-run session ceiling for each LLM pipeline's candidate list
+    #: (newest-first, so fresh sessions win). The unattended nightly cron
+    #: relies on this: a backlogged corpus must never turn one tick into an
+    #: unbounded spend. Env: ``ATIF_SQL_LLM_MAX_SESSIONS_PER_RUN``.
+    llm_max_sessions_per_run: int = 50
+    #: Hard per-run dollar ceiling across ALL LLM pipelines, checked against
+    #: the UsageAccumulator's running actuals between stages — when crossed,
+    #: the remaining LLM stages are aborted cleanly (nothing stamped for
+    #: unstarted sessions). Env: ``ATIF_SQL_LLM_MAX_COST_USD_PER_RUN``.
+    llm_max_cost_usd_per_run: float = 25.0
+    #: Write chunking: rows land every ``max(batch_size * 4, 256)`` units.
+    batch_size: int = 96
+    #: Transcript caps: 800K per session, 50K per tool_result.
+    session_text_total_max_chars: int = 800_000
+    session_text_tool_result_max_chars: int = 50_000
+
+    # ------------------------------------------------------------------
+    # Derivations
+    # ------------------------------------------------------------------
+
+    def layout(self) -> AnalyticsLayout:
+        """The analytics artifact layout under :attr:`corpus_root`."""
+        return AnalyticsLayout(corpus_root=self.corpus_root)
+
+    def llm(self) -> LlmSettings:
+        """The atif-models LLM settings (family/sizes/region/concurrency)."""
+        from atif_models.infrastructure.settings import LlmSettings
+
+        return LlmSettings()
+
+    def transcript_caps(self) -> TranscriptCaps:
+        """Project the transcript char caps."""
+        return TranscriptCaps(
+            session_text_total_max_chars=self.session_text_total_max_chars,
+            session_text_tool_result_max_chars=self.session_text_tool_result_max_chars,
+        )
+
+
+__all__ = ["AnalyticsSettings"]

atif_analytics/infrastructure/sqlite_state/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""SQLite WAL state.db: per-corpus checkpoint + retry-queue tables."""

atif_analytics/infrastructure/sqlite_state/checkpointer.py

@@ -0,0 +1,326 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Per-(session_id, pipeline) checkpoint backed by a SQLite WAL file.
+
+Tracks when each LLM pipeline last processed each session so re-runs skip
+sessions whose transcripts have not advanced. One row per
+``(session_id, pipeline)``; UPSERT via ``INSERT ... ON CONFLICT DO UPDATE``.
+
+Schema::
+
+    CREATE TABLE session_checkpoint (
+        session_id            TEXT NOT NULL,
+        pipeline              TEXT NOT NULL,
+        last_ts_processed     TEXT,
+        last_mtime_processed  TEXT,
+        completed_at          TEXT NOT NULL,
+        PRIMARY KEY (session_id, pipeline)
+    );
+
+All timestamps are stored as ISO-8601 UTC strings; they sort
+lexicographically. WAL journal mode lets multiple readers run concurrently
+with one writer; ``busy_timeout`` absorbs transient writer contention.
+The file lives at ``<corpus_root>/analytics/state.db`` — corpus-scoped, so
+one corpus's completions can never skip another's re-scoring.
+"""
+
+from __future__ import annotations
+
+import sqlite3
+import threading
+from collections.abc import Iterable
+from datetime import UTC, datetime
+from pathlib import Path
+
+#: The checkpointed LLM pipelines. ``trajectory`` left on 2026-09-27; its
+#: rows stay in an older ``state.db`` and nothing reads them.
+PIPELINE_NAMES: tuple[str, ...] = (
+    "classify",
+    "conflicts",
+    "user_friction",
+    "perceived",
+)
+
+_CREATE_TABLES_SQL = """
+CREATE TABLE IF NOT EXISTS session_checkpoint (
+    session_id            TEXT NOT NULL,
+    pipeline              TEXT NOT NULL,
+    last_ts_processed     TEXT,
+    last_mtime_processed  TEXT,
+    completed_at          TEXT NOT NULL,
+    PRIMARY KEY (session_id, pipeline)
+);
+"""
+
+# Starvation counter: one row per budget-skipped stage per run. A stage
+# that actually runs clears its rows, so the row count IS the consecutive
+# budget-skip streak (run_analyze escalates to ERROR at >= 3).
+_CREATE_BUDGET_SKIPS_SQL = """
+CREATE TABLE IF NOT EXISTS budget_skips (
+    pipeline TEXT NOT NULL,
+    run_at   TEXT NOT NULL
+);
+"""
+
+# Process-local set of paths whose schema has already been bootstrapped this
+# process. Skipping the redundant ``CREATE TABLE IF NOT EXISTS`` on every
+# ``connect`` call avoids racing the writer lock when concurrent threads
+# open the same file: on a cold start the WAL transition and the CREATE
+# both take the writer lock, and racing them raises "database is locked".
+_SCHEMA_BOOTSTRAPPED: set[str] = set()
+_SCHEMA_BOOTSTRAP_LOCK = threading.Lock()
+
+
+def to_iso(dt: datetime | None) -> str | None:
+    """ISO-8601 UTC string, or None."""
+    if dt is None:
+        return None
+    return dt.astimezone(UTC).isoformat()
+
+
+def _from_iso(s: str | None) -> datetime | None:
+    """Parse an ISO-8601 string back to a tz-aware UTC datetime."""
+    if s is None:
+        return None
+    dt = datetime.fromisoformat(s)
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=UTC)
+    return dt.astimezone(UTC)
+
+
+def connect(path: Path) -> sqlite3.Connection:
+    """Open the SQLite checkpoint DB and ensure tables exist.
+
+    Per-connection PRAGMAs run every open; the file-level WAL/synchronous
+    setup + DDL run once per (process, path) under a lock, because the WAL
+    transition and ``CREATE TABLE`` both grab the writer lock and racing
+    them on cold start raises ``database is locked``.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    con = sqlite3.connect(str(path), isolation_level=None, timeout=30.0)
+    con.execute("PRAGMA busy_timeout=10000")
+    con.execute("PRAGMA foreign_keys=ON")
+    key = str(path.resolve())
+    with _SCHEMA_BOOTSTRAP_LOCK:
+        if key not in _SCHEMA_BOOTSTRAPPED:
+            con.execute("PRAGMA journal_mode=WAL")
+            con.execute("PRAGMA synchronous=NORMAL")
+            con.execute(_CREATE_TABLES_SQL)
+            con.execute(_CREATE_BUDGET_SKIPS_SQL)
+            _SCHEMA_BOOTSTRAPPED.add(key)
+    # Writes wrap in implicit BEGIN/COMMIT from here on.
+    con.isolation_level = "DEFERRED"
+    return con
+
+
+def load_as_map(db_path: Path, pipeline: str) -> dict[str, tuple[datetime | None, datetime | None]]:
+    """Return ``{session_id: (last_ts, last_mtime)}`` for one pipeline.
+
+    Empty dict when the DB doesn't exist yet or the pipeline has no rows.
+    """
+    if not db_path.exists():
+        return {}
+    con = connect(db_path)
+    try:
+        rows = con.execute(
+            "SELECT session_id, last_ts_processed, last_mtime_processed "
+            "FROM session_checkpoint WHERE pipeline = ?",
+            [pipeline],
+        ).fetchall()
+    finally:
+        con.close()
+    return {
+        str(sid): (_from_iso(last_ts), _from_iso(last_mtime)) for sid, last_ts, last_mtime in rows
+    }
+
+
+def filter_unchanged(
+    candidates: Iterable[tuple[str, datetime | None, datetime | None]],
+    *,
+    pipeline: str,
+    checkpoint_db_path: Path,
+) -> tuple[list[str], int]:
+    """Drop sessions whose ``(last_ts, last_mtime)`` has not advanced.
+
+    ``candidates`` is an iterable of ``(session_id, current_last_ts,
+    current_last_mtime)``. Returns ``(pending_session_ids, skipped_count)``.
+
+    A session is skipped iff a checkpoint row exists for ``pipeline`` AND
+    both ``current_last_ts <= ckpt.last_ts`` AND ``current_last_mtime <=
+    ckpt.last_mtime``. Either bound moving forward invalidates the skip.
+    """
+    ckpt = load_as_map(checkpoint_db_path, pipeline)
+    pending: list[str] = []
+    skipped = 0
+    for sid, cur_ts, cur_mtime in candidates:
+        prev = ckpt.get(sid)
+        if prev is None:
+            pending.append(sid)
+            continue
+        prev_ts, prev_mtime = prev
+        if _stale_or_equal(cur_ts, prev_ts) and _stale_or_equal(cur_mtime, prev_mtime):
+            skipped += 1
+            continue
+        pending.append(sid)
+    return pending, skipped
+
+
+def _stale_or_equal(cur: datetime | None, prev: datetime | None) -> bool:
+    """True iff both are present and ``cur`` has not advanced past ``prev``.
+
+    None on either side returns False (treat as advanced).
+    """
+    if cur is None or prev is None:
+        return False
+    cur_aware = cur.astimezone(UTC) if cur.tzinfo else cur.replace(tzinfo=UTC)
+    prev_aware = prev.astimezone(UTC) if prev.tzinfo else prev.replace(tzinfo=UTC)
+    return cur_aware <= prev_aware
+
+
+def mark_completed(
+    db_path: Path,
+    *,
+    pipeline: str,
+    rows: Iterable[tuple[str, datetime | None, datetime | None]],
+) -> int:
+    """Upsert checkpoint rows for ``(session_id, pipeline)``.
+
+    Each row is ``(session_id, last_ts_processed, last_mtime_processed)``.
+    ``completed_at`` is stamped with ``datetime.now(UTC)``. Returns the
+    number of upserted rows; an empty ``rows`` leaves the DB untouched.
+    """
+    incoming = list(rows)
+    if not incoming:
+        return 0
+    now_iso = to_iso(datetime.now(UTC))
+    payload = [
+        (sid, pipeline, to_iso(last_ts), to_iso(last_mtime), now_iso)
+        for sid, last_ts, last_mtime in incoming
+    ]
+    con = connect(db_path)
+    try:
+        con.executemany(
+            "INSERT INTO session_checkpoint "
+            "(session_id, pipeline, last_ts_processed, last_mtime_processed, completed_at) "
+            "VALUES (?, ?, ?, ?, ?) "
+            "ON CONFLICT(session_id, pipeline) DO UPDATE SET "
+            "last_ts_processed = excluded.last_ts_processed, "
+            "last_mtime_processed = excluded.last_mtime_processed, "
+            "completed_at = excluded.completed_at",
+            payload,
+        )
+        con.commit()
+    finally:
+        con.close()
+    return len(incoming)
+
+
+def record_budget_skip(db_path: Path, pipeline: str, *, now: datetime | None = None) -> int:
+    """Append one budget-skip row for ``pipeline``; return the streak length.
+
+    The streak is the number of consecutive runs this stage was skipped
+    with ``budget_exhausted`` — a successful run clears the rows via
+    :func:`clear_budget_skips`, so a plain count is the streak.
+    """
+    con = connect(db_path)
+    try:
+        con.execute(
+            "INSERT INTO budget_skips (pipeline, run_at) VALUES (?, ?)",
+            [pipeline, to_iso(now or datetime.now(UTC))],
+        )
+        con.commit()
+        row = con.execute(
+            "SELECT count(*) FROM budget_skips WHERE pipeline = ?", [pipeline]
+        ).fetchone()
+    finally:
+        con.close()
+    return int(row[0]) if row else 1
+
+
+def clear_budget_skips(db_path: Path, pipeline: str) -> None:
+    """Reset ``pipeline``'s budget-skip streak (called when the stage runs)."""
+    con = connect(db_path)
+    try:
+        con.execute("DELETE FROM budget_skips WHERE pipeline = ?", [pipeline])
+        con.commit()
+    finally:
+        con.close()
+
+
+def budget_skip_count(db_path: Path, pipeline: str) -> int:
+    """Current consecutive budget-skip streak for ``pipeline`` (0 when absent)."""
+    if not db_path.exists():
+        return 0
+    con = connect(db_path)
+    try:
+        row = con.execute(
+            "SELECT count(*) FROM budget_skips WHERE pipeline = ?", [pipeline]
+        ).fetchone()
+    finally:
+        con.close()
+    return int(row[0]) if row else 0
+
+
+def count_rows(db_path: Path) -> int:
+    """Total number of checkpoint rows, or 0 when the DB is missing."""
+    if not db_path.exists():
+        return 0
+    con = connect(db_path)
+    try:
+        row = con.execute("SELECT count(*) FROM session_checkpoint").fetchone()
+    finally:
+        con.close()
+    return int(row[0]) if row else 0
+
+
+class SqliteCheckpoint:
+    """The checkpoint port over one ``state.db``."""
+
+    def __init__(self, db_path: Path) -> None:
+        self._db_path = db_path
+
+    @property
+    def db_path(self) -> Path:
+        """The bound sqlite file."""
+        return self._db_path
+
+    def load_as_map(self, pipeline: str) -> dict[str, tuple[datetime | None, datetime | None]]:
+        """``{session_id: (last_ts, last_mtime)}`` for one pipeline."""
+        return load_as_map(self._db_path, pipeline)
+
+    def filter_unchanged(
+        self,
+        candidates: Iterable[tuple[str, datetime | None, datetime | None]],
+        *,
+        pipeline: str,
+    ) -> tuple[list[str], int]:
+        """``(pending_session_ids, skipped_count)`` per the staleness rule."""
+        return filter_unchanged(candidates, pipeline=pipeline, checkpoint_db_path=self._db_path)
+
+    def mark_completed(
+        self,
+        *,
+        pipeline: str,
+        rows: Iterable[tuple[str, datetime | None, datetime | None]],
+    ) -> int:
+        """Upsert checkpoint rows; return the number upserted."""
+        return mark_completed(self._db_path, pipeline=pipeline, rows=rows)
+
+    def count_rows(self) -> int:
+        """Total checkpoint rows (``0`` when absent)."""
+        return count_rows(self._db_path)
+
+
+__all__ = [
+    "PIPELINE_NAMES",
+    "SqliteCheckpoint",
+    "budget_skip_count",
+    "clear_budget_skips",
+    "connect",
+    "count_rows",
+    "filter_unchanged",
+    "load_as_map",
+    "mark_completed",
+    "record_budget_skip",
+    "to_iso",
+]