atif-sql 0.2.0__py3-none-any.whl

atif_duck/infrastructure/lake_sessions.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Read one corpus's sessions out of the published lake, a batch at a time.
+
+The analytics read path. atif-cli wraps :class:`LakeSessionReader` in
+atif-analytics' ``SessionSource`` port (the two packages may not import each
+other, so the rows cross as plain tuples). The statements are the constants in
+:mod:`atif_duck.domain.lake_sessions`.
+
+The connection is private to the reader and runs only those statements, so
+it is not sandboxed the way ``query``'s is. It attaches the published reader
+catalog ``READ_ONLY`` through :func:`~atif_duck.infrastructure.lake.attach_lake_for_query`,
+so it never waits on the writer, sees the last committed state, and refuses
+a lake whose schema is stale or that holds another corpus under this name.
+
+A read that fails after the attach raises :class:`LakeReadError`. The one
+way that happens in practice: a ``lake rebuild`` swaps a new lake into place
+while a long ``analyze`` still holds the old attach, and the old data files
+it names are gone.
+"""
+
+from __future__ import annotations
+
+import shutil
+import tempfile
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from atif_duck.domain.lake_sessions import (
+    EDGE_UUIDS_SQL,
+    ERROR_STEPS_SQL,
+    LAST_EDGE_TS_SQL,
+    SESSION_FRESHNESS_SQL,
+    TOOL_CALLS_SQL,
+    TOOL_RESULTS_SQL,
+    TURNS_SQL,
+)
+from atif_duck.domain.sql_literal import sql_literal
+from atif_duck.infrastructure.lake import (
+    DEFAULT_WRITER_MEMORY_BYTES,
+    LakeReader,
+    LakeUnavailable,
+    attach_lake_for_query,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import duckdb
+
+    from atif_duck.infrastructure.lake import LakeLayout
+
+#: The reader's DuckDB memory cap (lowered further by the caller's). A batch
+#: of sessions' tool results is the largest thing it holds.
+DEFAULT_READER_MEMORY_BYTES: int = DEFAULT_WRITER_MEMORY_BYTES
+
+#: DuckDB threads for the reader's connection.
+_READER_THREADS: int = 4
+
+
+class LakeReadError(RuntimeError):
+    """A read on an open :class:`LakeSessionReader` failed (the lake moved under it)."""
+
+
+@dataclass(frozen=True, slots=True)
+class SessionFreshness:
+    """What a session's ``meta.json`` said when the lake loaded it."""
+
+    materialized_at: str | None
+    source_mtime_ns: int | None
+
+
+def _utc(value: datetime | None) -> datetime | None:
+    """The lake's naive UTC timestamp as an aware one (what the file path yields)."""
+    return None if value is None else value.replace(tzinfo=UTC)
+
+
+class LakeSessionReader:
+    """One corpus's session rows from the published lake."""
+
+    def __init__(self, con: duckdb.DuckDBPyConnection, corpus: str, spill: Path) -> None:
+        self._con = con
+        self._corpus = corpus
+        self._spill = spill
+
+    @classmethod
+    def open(
+        cls,
+        layout: LakeLayout,
+        corpus_root: Path,
+        *,
+        memory_limit_bytes: int | None = None,
+    ) -> LakeSessionReader | LakeUnavailable:
+        """Attach the lake for ``corpus_root``'s corpus, or say why not (nothing is installed)."""
+        import duckdb
+
+        spill = Path(tempfile.mkdtemp(prefix="atif-sql-lake-read-"))
+        con = duckdb.connect()
+        try:
+            cap = min(
+                memory_limit_bytes or DEFAULT_READER_MEMORY_BYTES, DEFAULT_READER_MEMORY_BYTES
+            )
+            con.execute(f"SET memory_limit='{int(cap)}B'")
+            con.execute(f"SET threads={int(_READER_THREADS)}")
+            con.execute(f"SET temp_directory={sql_literal(str(spill))}")
+            con.execute("SET autoinstall_known_extensions=false")
+            con.execute("SET autoload_known_extensions=false")
+            attached = attach_lake_for_query(
+                con, layout, corpus_root=corpus_root, all_corpora=False
+            )
+        except BaseException:
+            con.close()
+            shutil.rmtree(spill, ignore_errors=True)
+            raise
+        if not isinstance(attached, LakeReader) or attached.corpus is None:
+            con.close()
+            shutil.rmtree(spill, ignore_errors=True)
+            if isinstance(attached, LakeUnavailable):
+                return attached
+            return LakeUnavailable("the lake attached without a corpus scope")
+        return cls(con, attached.corpus, spill)
+
+    @property
+    def corpus(self) -> str:
+        """The corpus name the reads are scoped to."""
+        return self._corpus
+
+    def close(self) -> None:
+        """Close the connection and remove its spill directory."""
+        self._con.close()
+        shutil.rmtree(self._spill, ignore_errors=True)
+
+    def _fetch(self, statement: str, params: list[Any]) -> list[tuple[Any, ...]]:
+        import duckdb
+
+        try:
+            return self._con.execute(statement, params).fetchall()
+        except duckdb.Error as exc:
+            raise LakeReadError(str(exc)) from exc
+
+    def _batch(self, statement: str, session_ids: Sequence[str]) -> list[tuple[Any, ...]]:
+        if not session_ids:
+            return []
+        return self._fetch(statement, [self._corpus, list(session_ids)])
+
+    def freshness(self) -> dict[str, SessionFreshness]:
+        """``{session_id: SessionFreshness}`` for every session the lake holds for the corpus."""
+        rows = self._fetch(SESSION_FRESHNESS_SQL, [self._corpus])
+        return {
+            str(sid): SessionFreshness(
+                materialized_at=None if materialized is None else str(materialized),
+                source_mtime_ns=None if mtime is None else int(mtime),
+            )
+            for sid, materialized, mtime in rows
+        }
+
+    def last_edge_ts(self) -> dict[str, datetime | None]:
+        """``{session_id: newest record ts}`` (aware, UTC) over the corpus's edges."""
+        rows = self._fetch(LAST_EDGE_TS_SQL, [self._corpus])
+        return {str(sid): _utc(ts) for sid, ts in rows}
+
+    def turns(self, session_ids: Sequence[str]) -> list[tuple[Any, ...]]:
+        """``(session_id, step_id, ts, source, message, is_sidechain, is_compact_summary, source_uuids)``.
+
+        In session, then step order; ``source_uuids`` is JSON text.
+        """
+        return self._batch(TURNS_SQL, session_ids)
+
+    def error_steps(self, session_ids: Sequence[str]) -> set[tuple[str, int]]:
+        """``(session_id, step_id)`` of every step with an error result."""
+        return {(str(sid), int(step)) for sid, step in self._batch(ERROR_STEPS_SQL, session_ids)}
+
+    def tool_calls(self, session_ids: Sequence[str]) -> list[tuple[Any, ...]]:
+        """``(session_id, step_id, tool_name, tool_input)`` in source order (input as JSON text)."""
+        return self._batch(TOOL_CALLS_SQL, session_ids)
+
+    def tool_results(self, session_ids: Sequence[str]) -> list[tuple[Any, ...]]:
+        """``(session_id, step_id, tool_use_id, content, is_error)`` in source order (content as JSON text)."""
+        return self._batch(TOOL_RESULTS_SQL, session_ids)
+
+    def edge_uuids(self, session_id: str) -> set[str]:
+        """The session's non-empty raw-record uuids."""
+        rows = self._fetch(EDGE_UUIDS_SQL, [self._corpus, session_id])
+        return {str(uuid) for (uuid,) in rows}
+
+
+__all__ = [
+    "DEFAULT_READER_MEMORY_BYTES",
+    "LakeReadError",
+    "LakeSessionReader",
+    "SessionFreshness",
+]

atif_duck/infrastructure/lake_settings.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Where the lake lives and how it is written, from ``ATIF_SQL_*`` env vars.
+
+Reading process env is I/O, so this lives in ``infrastructure``. Every default
+factory reads the environment at call time, so a test that re-points ``HOME``
+sees the new value without a module reload.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+from atif_duck.infrastructure.lake import (
+    DEFAULT_LOAD_BATCH_SIZE,
+    DEFAULT_LOCK_TIMEOUT_SECONDS,
+    DEFAULT_STAGE_WORKERS,
+    DEFAULT_SYNC_BATCH_SIZE,
+)
+
+
+def _default_lake_root() -> Path:
+    """``~/.atif-sql/lake``: one lake for every corpus."""
+    return Path.home() / ".atif-sql" / "lake"
+
+
+def _default_corpus_base() -> Path:
+    """``~/.atif-sql/corpus``: where the default corpus roots live (``lake rebuild`` scans it)."""
+    return Path.home() / ".atif-sql" / "corpus"
+
+
+class LakeSettings(BaseSettings):
+    """Env-driven settings for the DuckLake every corpus is queried through."""
+
+    model_config = SettingsConfigDict(env_prefix="ATIF_SQL_", env_file=".env", extra="ignore")
+
+    #: The lake root (``ATIF_SQL_LAKE_ROOT``): ``catalog.duckdb``, its
+    #: published reader copy, and ``data/``.
+    lake_root: Path = Field(default_factory=_default_lake_root)
+    #: The directory ``lake rebuild`` looks in for corpora when none is named
+    #: (``ATIF_SQL_CORPUS_BASE``); every child holding ``sessions/`` is one.
+    corpus_base: Path = Field(default_factory=_default_corpus_base)
+    #: Sessions per lake transaction when materialize syncs what it published
+    #: (``ATIF_SQL_LAKE_SYNC_BATCH_SIZE``).
+    lake_sync_batch_size: int = Field(default=DEFAULT_SYNC_BATCH_SIZE, ge=1)
+    #: Sessions per lake transaction when a whole corpus is loaded
+    #: (``ATIF_SQL_LAKE_LOAD_BATCH_SIZE``).
+    lake_load_batch_size: int = Field(default=DEFAULT_LOAD_BATCH_SIZE, ge=1)
+    #: Processes that stage a batch's parquet from its compressed trajectories
+    #: while the lake loads it (``ATIF_SQL_LAKE_STAGE_WORKERS``); 1 stages in
+    #: the loading process.
+    lake_stage_workers: int = Field(default=DEFAULT_STAGE_WORKERS, ge=1)
+    #: Seconds a writer waits for another writer (``ATIF_SQL_LAKE_LOCK_TIMEOUT_SECONDS``).
+    lake_lock_timeout_seconds: float = Field(default=DEFAULT_LOCK_TIMEOUT_SECONDS, gt=0)
+    #: ``lake compact`` expires snapshots older than this many days
+    #: (``ATIF_SQL_LAKE_EXPIRE_DAYS``).
+    lake_expire_days: int = Field(default=30, ge=0)
+
+
+__all__ = ["LakeSettings"]

atif_duck/infrastructure/lake_steps.py

@@ -0,0 +1,225 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""One corpus's step texts and keys, read from the lake for the embedding store.
+
+atif-embed discovers what to embed from these reads (atif-cli adapts
+:class:`LakeSteps` to atif-embed's ``LakeStepsPort``; the two packages may not
+import each other). Everything here reads the published catalog copy
+``READ_ONLY`` through :func:`atif_duck.infrastructure.lake.attach_lake_for_query`,
+so it never waits on a writer and falls back for exactly the reasons ``query``
+does (no lake, a stale one, a corpus it doesn't hold).
+
+The rows are the ``steps`` table's, in the per-session reader's order
+(session id, then step id), keyed like the store: a step's primary uuid is
+the first entry of ``source_uuids``, and its text is the flattened
+``message`` the lake already stores.
+
+Change reads use DuckLake's ``ducklake_table_changes``: the sink replaces a
+session by deleting and re-inserting its rows, so a session written since a
+snapshot shows up as deletes of its old rows and inserts of its new ones.
+File merges and rewrites (``lake compact``) change no row and show up as
+nothing. Expired snapshots can't be read across, which
+:meth:`LakeSteps.can_read_changes_since` reports.
+"""
+
+from __future__ import annotations
+
+import shutil
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING, Self
+
+from atif_duck.domain.lake import CORPORA_TABLE, LAKE_ALIAS, SESSION_KEY
+from atif_duck.domain.raw_readers import CORPUS_COLUMN
+from atif_duck.domain.sql_literal import SqlFragment, sql_literal
+from atif_duck.infrastructure.lake import (
+    DEFAULT_WRITER_MEMORY_BYTES,
+    LakeReader,
+    LakeUnavailable,
+    attach_lake_for_query,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from types import TracebackType
+
+    import duckdb
+
+    from atif_duck.infrastructure.lake import LakeLayout
+
+#: The lake table the embeddable rows live in.
+STEPS_TABLE: str = "steps"
+
+#: A step's primary uuid: the key the embedding store uses.
+_PRIMARY_UUID: str = "json_extract_string(source_uuids, '$[0]')"
+
+#: One corpus's steps with the key computed once (param: corpus).
+_KEYED_STEPS: str = (
+    f"SELECT {SESSION_KEY} AS session_id, step_id, {_PRIMARY_UUID} AS uuid, message "  # noqa: S608  # nosec B608 - module constants only
+    f"FROM {LAKE_ALIAS}.{STEPS_TABLE} WHERE {CORPUS_COLUMN} = ?"
+)
+
+#: What makes a row embeddable (param: the minimum text length).
+_QUALIFIES: str = "uuid IS NOT NULL AND message IS NOT NULL AND length(message) >= ?"
+
+#: Every qualifying row of one corpus, in corpus order (params: corpus, min chars).
+STEP_TEXTS_SQL: SqlFragment = SqlFragment(
+    f"SELECT uuid, message FROM ({_KEYED_STEPS}) WHERE {_QUALIFIES} "  # noqa: S608  # nosec B608 - module constants only
+    "ORDER BY session_id, step_id"
+)
+
+#: The qualifying rows of every uuid a step row inserted or deleted after a
+#: snapshot touched, in corpus order (params: first snapshot, last snapshot,
+#: corpus, corpus, min chars).
+CHANGED_STEP_TEXTS_SQL: SqlFragment = SqlFragment(
+    f"WITH changed AS (SELECT DISTINCT {_PRIMARY_UUID} AS uuid "  # noqa: S608  # nosec B608 - module constants only
+    f"FROM ducklake_table_changes({sql_literal(LAKE_ALIAS)}, 'main', {sql_literal(STEPS_TABLE)}, ?, ?) "
+    f"WHERE {CORPUS_COLUMN} = ?) "
+    f"SELECT uuid, message FROM ({_KEYED_STEPS}) WHERE {_QUALIFIES} "
+    "AND uuid IN (SELECT uuid FROM changed) ORDER BY session_id, step_id"
+)
+
+#: Every step's primary uuid in one corpus, text or none (param: corpus).
+STEP_KEYS_SQL: SqlFragment = SqlFragment(
+    f"SELECT DISTINCT uuid FROM ({_KEYED_STEPS}) WHERE uuid IS NOT NULL"  # noqa: S608  # nosec B608 - module constants only
+)
+
+#: When the corpus was last loaded whole (param: corpus). A rebuild and a
+#: first load both reset it, which is what makes it the lineage.
+REGISTERED_AT_SQL: SqlFragment = SqlFragment(
+    f"SELECT registered_at FROM {LAKE_ALIAS}.{CORPORA_TABLE} WHERE {CORPUS_COLUMN} = ?"  # noqa: S608  # nosec B608 - module constants only
+)
+
+#: The oldest and newest snapshot the attached catalog still holds.
+SNAPSHOT_RANGE_SQL: SqlFragment = SqlFragment(
+    f"SELECT min(snapshot_id), max(snapshot_id) FROM ducklake_snapshots({sql_literal(LAKE_ALIAS)})"  # noqa: S608  # nosec B608 - module constants only
+)
+
+#: Rows pulled off a result per ``fetchmany``.
+_FETCH_PAGE_ROWS: int = 2048
+
+#: Threads for a step read. The reads are one scan and one sort; more
+#: threads only raise the peak.
+_READ_THREADS: int = 4
+
+
+@dataclass(slots=True)
+class LakeSteps:
+    """One corpus's steps at the published catalog's current snapshot. Close it when done."""
+
+    con: duckdb.DuckDBPyConnection
+    corpus: str
+    #: The corpus's ``registered_at``: changes when its rows are loaded from scratch.
+    lineage: str
+    snapshot_id: int
+    oldest_snapshot_id: int
+    spill_dir: Path
+
+    def can_read_changes_since(self, snapshot_id: int) -> bool:
+        """True when every snapshot after ``snapshot_id`` is still in the catalog."""
+        return self.oldest_snapshot_id <= snapshot_id + 1 and snapshot_id <= self.snapshot_id
+
+    def step_texts(
+        self, *, min_chars: int, since_snapshot: int | None
+    ) -> Iterator[tuple[str, str]]:
+        """``(uuid, text)`` rows in corpus order; with ``since_snapshot``, only the changed uuids'."""
+        if since_snapshot is None:
+            result = self.con.execute(STEP_TEXTS_SQL, [self.corpus, int(min_chars)])
+        elif since_snapshot >= self.snapshot_id:
+            return
+        else:
+            result = self.con.execute(
+                CHANGED_STEP_TEXTS_SQL,
+                [
+                    int(since_snapshot) + 1,
+                    self.snapshot_id,
+                    self.corpus,
+                    self.corpus,
+                    int(min_chars),
+                ],
+            )
+        while page := result.fetchmany(_FETCH_PAGE_ROWS):
+            for uuid, text in page:
+                yield str(uuid), str(text)
+
+    def step_keys(self) -> frozenset[str]:
+        """Every step's primary uuid in the corpus."""
+        return frozenset(
+            str(row[0]) for row in self.con.execute(STEP_KEYS_SQL, [self.corpus]).fetchall()
+        )
+
+    def close(self) -> None:
+        """Close the connection and remove its spill directory."""
+        self.con.close()
+        shutil.rmtree(self.spill_dir, ignore_errors=True)
+
+    def __enter__(self) -> Self:
+        """Use as a context manager: the connection closes on exit."""
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Close on the way out, whatever happened."""
+        self.close()
+
+
+def open_lake_steps(
+    layout: LakeLayout, corpus_root: Path, *, memory_limit_bytes: int | None = None
+) -> LakeSteps | LakeUnavailable:
+    """Attach the published lake for ``corpus_root``'s steps, or say why not.
+
+    DuckDB runs capped at ``memory_limit_bytes`` (at most the writer's cap),
+    spilling to a private temporary directory the returned object removes on
+    close.
+    """
+    import duckdb
+
+    spill = Path(tempfile.mkdtemp(prefix="atif-sql-lake-steps-"))
+    con = duckdb.connect()
+    try:
+        cap = min(memory_limit_bytes or DEFAULT_WRITER_MEMORY_BYTES, DEFAULT_WRITER_MEMORY_BYTES)
+        con.execute(f"SET memory_limit='{int(cap)}B'")
+        con.execute(f"SET threads={int(_READ_THREADS)}")
+        con.execute(f"SET temp_directory={sql_literal(str(spill))}")
+        con.execute("SET autoinstall_known_extensions=false")
+        con.execute("SET autoload_known_extensions=false")
+        attached = attach_lake_for_query(con, layout, corpus_root=corpus_root, all_corpora=False)
+        if not isinstance(attached, LakeReader) or attached.corpus is None:
+            con.close()
+            shutil.rmtree(spill, ignore_errors=True)
+            return (
+                attached
+                if isinstance(attached, LakeUnavailable)
+                else LakeUnavailable("the lake did not scope to one corpus")
+            )
+        registered = con.execute(REGISTERED_AT_SQL, [attached.corpus]).fetchone()
+        snapshots = con.execute(SNAPSHOT_RANGE_SQL).fetchone()
+    except BaseException:
+        con.close()
+        shutil.rmtree(spill, ignore_errors=True)
+        raise
+    return LakeSteps(
+        con=con,
+        corpus=attached.corpus,
+        lineage=str(registered[0]) if registered else "",
+        snapshot_id=int(snapshots[1]) if snapshots and snapshots[1] is not None else 0,
+        oldest_snapshot_id=int(snapshots[0]) if snapshots and snapshots[0] is not None else 0,
+        spill_dir=spill,
+    )
+
+
+__all__ = [
+    "CHANGED_STEP_TEXTS_SQL",
+    "REGISTERED_AT_SQL",
+    "SNAPSHOT_RANGE_SQL",
+    "STEPS_TABLE",
+    "STEP_KEYS_SQL",
+    "STEP_TEXTS_SQL",
+    "LakeSteps",
+    "open_lake_steps",
+]