atif-sql 0.2.0__py3-none-any.whl

atif_duck/domain/columnar.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""The typed columnar artifact contract: names, schema version, column shapes.
+
+Beside the JSON artifacts in ``<corpus_root>/sessions/<id>/``, a session
+may carry five Parquet files that hold what the transcript-derived views need,
+already typed, so a query never parses ``trajectory.json``:
+
+* ``session.parquet``      one row: the trajectory's top-level members minus
+  ``steps`` (what the ``sessions`` view reads), keyed by ``session_id_path``.
+* ``steps.parquet``        one row per ATIF step, exactly the ``steps`` view's
+  columns.
+* ``tool_calls.parquet``   one row per tool call, exactly the ``tool_calls``
+  view's columns.
+* ``tool_results.parquet`` one row per observation result, exactly the
+  ``tool_results`` view's columns.
+* ``session_events.parquet`` one row per kept non-message record, exactly the
+  ``session_events`` view's columns (schema 2 onward). Built from the staged
+  ``session_events.jsonl``, not from the trajectory.
+
+atif-duck owns this shape (the views define what the columns mean) and writes
+it through :class:`atif_duck.infrastructure.columnar.ColumnarArtifactProducer`;
+atif-corpus calls that producer through its ``ArtifactProducer`` port without
+knowing these names. The registry reads a session from these files only when
+``meta.json`` carries ``columnar_schema`` equal to
+:data:`COLUMNAR_SCHEMA_VERSION` and all five files are present; otherwise it
+reads the session's ``trajectory.json`` as before, so a corpus materialized
+before this contract existed (or by a build with a different schema version)
+still queries correctly.
+
+Pure constants and path arithmetic: no duckdb import, no filesystem access.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from atif_duck.domain import artifacts
+from atif_duck.domain.catalog import VIEW_SCHEMA
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+#: Bump when any column, type, or filename below changes. A session whose
+#: ``meta.json`` names another version is read from ``trajectory.json``.
+#:
+#: 2: ``session_events.parquet`` added; ``steps`` gained ``agent_id`` /
+#: ``images``; ``tool_results`` gained ``is_error`` / ``exit_code`` /
+#: ``interrupted`` / ``images``. One bump for both changes: they ship together,
+#: and no corpus was ever stamped with only one of them, so a single version
+#: separates every schema-1 session from every schema-2 one.
+COLUMNAR_SCHEMA_VERSION: int = 2
+
+#: The ``meta.json`` key that records which columnar schema a session's
+#: parquet files were written against. Absent on sessions materialized
+#: before the columnar artifacts existed, which is a legal, JSON-read state.
+META_COLUMNAR_KEY: str = "columnar_schema"
+
+SESSION_PARQUET: str = "session.parquet"
+STEPS_PARQUET: str = "steps.parquet"
+TOOL_CALLS_PARQUET: str = "tool_calls.parquet"
+TOOL_RESULTS_PARQUET: str = "tool_results.parquet"
+#: Added in schema 2: the ``session_events`` view's rows, built from the staged
+#: ``session_events.jsonl`` the corpus writer puts beside ``edges.jsonl``.
+SESSION_EVENTS_PARQUET: str = "session_events.parquet"
+
+#: Every columnar filename, in the order the producer writes them.
+COLUMNAR_FILENAMES: tuple[str, ...] = (
+    SESSION_PARQUET,
+    STEPS_PARQUET,
+    TOOL_CALLS_PARQUET,
+    TOOL_RESULTS_PARQUET,
+    SESSION_EVENTS_PARQUET,
+)
+
+#: The JSON artifact the ``session_events`` rows come from (atif-corpus writes
+#: it, compressed on a current corpus: see :mod:`atif_duck.domain.artifacts`).
+SESSION_EVENTS_JSONL: str = artifacts.SESSION_EVENTS_JSONL
+
+#: ``session.parquet`` columns. ``session_id_path`` is the session directory
+#: name, the canonical key every view uses as ``session_id``; the rest are the
+#: trajectory's top-level members with the types the JSON reader projects them
+#: to (``atif_duck.infrastructure.registry._TRAJECTORY_COLUMNS``), minus
+#: ``steps``.
+SESSION_COLUMNS: tuple[tuple[str, str], ...] = (
+    ("session_id_path", "VARCHAR"),
+    ("schema_version", "VARCHAR"),
+    ("session_id", "VARCHAR"),
+    ("trajectory_id", "VARCHAR"),
+    ("agent", "STRUCT(name VARCHAR, version VARCHAR, model_name VARCHAR, extra JSON)"),
+    (
+        "final_metrics",
+        (
+            "STRUCT(total_prompt_tokens BIGINT, total_completion_tokens BIGINT, "
+            "total_cached_tokens BIGINT, total_cost_usd DOUBLE, total_steps BIGINT, extra JSON)"
+        ),
+    ),
+    ("extra", "JSON"),
+)
+
+#: Column shape per columnar file. The three step-derived files carry exactly
+#: their view's catalog columns, so the catalog stays the one place a column
+#: is declared and the schema drift test covers these files too.
+COLUMNAR_SCHEMAS: dict[str, tuple[tuple[str, str], ...]] = {
+    SESSION_PARQUET: SESSION_COLUMNS,
+    STEPS_PARQUET: VIEW_SCHEMA["steps"],
+    TOOL_CALLS_PARQUET: VIEW_SCHEMA["tool_calls"],
+    TOOL_RESULTS_PARQUET: VIEW_SCHEMA["tool_results"],
+    SESSION_EVENTS_PARQUET: VIEW_SCHEMA["session_events"],
+}
+
+#: Smallest byte count a readable parquet file can have (``PAR1`` + footer +
+#: footer length + ``PAR1``). A shorter file is what a crashed writer leaves
+#: behind; the registry treats it as absent and reads the session from JSON.
+MIN_PARQUET_BYTES: int = 16
+
+
+def columnar_paths(session_dir: Path) -> tuple[Path, ...]:
+    """The five columnar artifact paths inside one session directory."""
+    return tuple(session_dir / name for name in COLUMNAR_FILENAMES)
+
+
+def is_current_columnar_schema(value: object) -> bool:
+    """True when a ``meta.json`` ``columnar_schema`` value names this build's schema.
+
+    Anything but the exact integer (absent key, ``None``, a string, another
+    version) answers ``False`` and sends the session down the JSON path.
+    """
+    return (
+        isinstance(value, int) and not isinstance(value, bool) and value == COLUMNAR_SCHEMA_VERSION
+    )
+
+
+__all__ = [
+    "COLUMNAR_FILENAMES",
+    "COLUMNAR_SCHEMAS",
+    "COLUMNAR_SCHEMA_VERSION",
+    "META_COLUMNAR_KEY",
+    "MIN_PARQUET_BYTES",
+    "SESSION_COLUMNS",
+    "SESSION_EVENTS_JSONL",
+    "SESSION_EVENTS_PARQUET",
+    "SESSION_PARQUET",
+    "STEPS_PARQUET",
+    "TOOL_CALLS_PARQUET",
+    "TOOL_RESULTS_PARQUET",
+    "columnar_paths",
+    "is_current_columnar_schema",
+]

atif_duck/domain/embedding_guard.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Fail-loud embedding provider/dimension guard for the VSS bind path.
+
+DELIBERATE TWIN of ``atif_embed.domain.embedding_guard``: atif-duck may never
+import atif-embed (import-linter independence contract — only atif-cli
+composes), and this is a tiny pure rule (string/int comparison, no I/O), so
+each package carries its own copy rather than growing a shared-kernel package
+for one function. ``test_guard_twin_pin.py`` is what keeps the copies honest.
+The two copies are coupled through the CONTRACT's store-stamp semantics; a
+behavioral change in one must be ported to the other.
+
+The read/bind path (``register_vss``) reads the Lance store's stamped
+``(model, dim)`` back and calls :func:`ensure_store_matches` BEFORE binding
+the ``message_embeddings`` view — a store written by a different provider
+raises rather than serving numerically valid but garbage cosine scores.
+"""
+
+from __future__ import annotations
+
+#: Recovery instruction appended to the mismatch message. The store lives at
+#: ``<corpus_root>/embeddings_lance`` by default, NOT under a fixed home
+#: directory — naming a path the store isn't at makes the operator delete
+#: nothing and stay broken. Kept identical to atif-embed's twin copy.
+RECOVERY_HINT = (
+    "Re-embed under the new provider: rm -rf the Lance store directory "
+    "(ATIF_SQL_LANCE_URI, or <corpus-root>/embeddings_lance by default) "
+    "and re-run `atif-sql embed --all --no-dry-run` (a bare `atif-sql embed` "
+    "exits 64: a real run needs an explicit scope)."
+)
+
+
+class EmbeddingProviderMismatch(Exception):  # noqa: N818 — names a store state, not an "*Error"
+    """Raised on a stamped-vs-active ``(model, dim)`` mismatch.
+
+    The Lance store records the identity of the embedder that wrote it, and it
+    differs from the identity of the embedder now asking to read or extend it.
+
+    The name carries no ``Error`` suffix because it names a STATE OF THE STORE
+    that is true before anything is attempted, not a fault that occurred while
+    doing something. Nothing has gone wrong at runtime: two consistent stores
+    simply cannot be the same store.
+
+    Different embedding models produce vectors in incompatible spaces, so
+    querying across a provider switch yields numerically valid but
+    semantically garbage cosine scores. This error is terminal: the store
+    must be dropped and re-embedded under the new provider.
+    """
+
+
+def ensure_store_matches(
+    *,
+    stored_model: str | None,
+    stored_dim: int | None,
+    expected_model: str,
+    expected_dim: int | None,
+) -> None:
+    """Fail loud if a store's stamped identity differs from the active embedder.
+
+    ``None`` for either stored value means the store is empty (fresh
+    install) and any provider may claim it, so the check is a no-op.
+    ``model_id`` is the primary identity; ``expected_dim`` is also checked
+    when supplied (Cohere's single model id can emit different Matryoshka
+    widths). ``expected_dim=None`` trusts ``model_id`` alone.
+    """
+    if stored_model is None or stored_dim is None:
+        return
+    model_ok = stored_model == expected_model
+    dim_ok = expected_dim is None or stored_dim == expected_dim
+    if model_ok and dim_ok:
+        return
+    # The message stays INSIDE the raise: `test_guard_twin_pin.py` reads both
+    # twin modules as source text and requires `{RECOVERY_HINT}` to appear after
+    # the `raise EmbeddingProviderMismatch` keyword, which is what proves the
+    # shared constant reaches the operator rather than only sitting in the
+    # module. Hoisting the message to a local satisfies EM102 / TRY003 and
+    # defeats that pin, so both are suppressed here instead.
+    raise EmbeddingProviderMismatch(  # noqa: TRY003
+        "Embedding store was written by a different provider/model. "  # noqa: EM102
+        f"stored=(model={stored_model!r}, dim={stored_dim}) "
+        f"active=(model={expected_model!r}, dim={expected_dim}). "
+        "Vectors from different models live in incompatible spaces (even at "
+        "matching dimensions), so mixing them silently corrupts kNN search. "
+        f"{RECOVERY_HINT}"
+    )
+
+
+__all__ = ["RECOVERY_HINT", "EmbeddingProviderMismatch", "ensure_store_matches"]

atif_duck/domain/examples.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+
+"""Derived, test-executed example queries for the atif-duck SQL surface.
+
+Operator's principle: DON'T hardcode example strings — they rot the moment a
+macro signature or view changes. Every example here is DERIVED from the
+static catalogs (:data:`~atif_duck.domain.catalog.VIEW_NAMES`,
+:data:`~atif_duck.domain.catalog.ANALYTICS_VIEW_NAMES`,
+:data:`~atif_duck.domain.catalog.MACRO_SIGNATURES`,
+:data:`~atif_duck.domain.catalog.ANALYTICS_MACRO_SIGNATURES`) plus one small
+exemplar table mapping parameter NAMES to literal SQL argument text
+(:data:`ARG_EXEMPLARS`). The catalogs are already drift-pinned against the
+real DDL by the tests in ``tests/test_duck_views.py`` /
+``tests/test_analytics_views.py``, so the examples inherit that freshness.
+
+THE TEST IS THE CONTRACT: ``tests/test_examples.py`` EXECUTES every example
+:func:`build_examples` emits against the fixture corpus (analytics parquets
+registered, tiny Lance store attached). An example that stops running stops
+shipping — the CLI surface (``atif-sql examples``) can honestly claim
+"test-executed against this version".
+
+The only hand-maintained prose is
+:data:`~atif_duck.domain.catalog.DESCRIPTIONS` in ``catalog.py``; a coverage
+drift test forces one description per catalog object.
+
+Pure domain module: no duckdb import, no I/O — safe on the CLI's lean
+import path (``schema`` / ``examples`` / ``--help``).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal, cast, get_args
+
+from atif_duck.domain.catalog import (
+    ANALYTICS_MACRO_SIGNATURES,
+    ANALYTICS_VIEW_NAMES,
+    ANALYTICS_VIEW_SCHEMA,
+    DESCRIPTIONS,
+    MACRO_SIGNATURES,
+    TABLE_MACRO_NAMES,
+    VIEW_NAMES,
+    VIEW_SCHEMA,
+)
+
+#: What an example needs registered before it runs: the base corpus views
+#: (``core``), the analytics parquet views (``analytics``), or the
+#: Lance-backed embeddings surface (``vss``).
+Requires = Literal["core", "analytics", "vss"]
+
+#: Structural kind, derived from which catalog the object lives in and
+#: whether its DDL is ``AS TABLE`` (see
+#: :data:`~atif_duck.domain.catalog.TABLE_MACRO_NAMES`).
+Category = Literal["view", "table-macro", "scalar-macro"]
+
+REQUIRES_VALUES: tuple[str, ...] = cast("tuple[str, ...]", get_args(Requires))
+CATEGORY_VALUES: tuple[str, ...] = cast("tuple[str, ...]", get_args(Category))
+
+#: Exemplar SQL argument text per macro parameter NAME. The single seam
+#: between the derived examples and real macro signatures: a new parameter
+#: name lands here (one line) or ``build_examples`` fails loud — and the
+#: executing test in ``tests/test_examples.py`` fails with it.
+#:
+#: * ``sid`` — a real session id, picked from the corpus itself so the
+#:   example runs verbatim on ANY materialized corpus.
+#: * ``query_vec`` — a stored embedding L2-NORMALIZED into a unit-norm probe,
+#:   matching what ``embed_query`` returns, so the example runs on any corpus
+#:   without fabricating a dim-wide literal. The normalization is not
+#:   cosmetic: stored vectors are un-normalized, and a RAW stored vector is
+#:   its own nearest neighbor under every metric, so an un-normalized probe
+#:   would make this example score sim == 1.0 no matter which metric the macro
+#:   ranks by — it could not tell a correct macro from a broken one. For TEXT
+#:   queries the flow is two-step: embed the query text first (``atif-sql
+#:   search`` wraps both steps); see
+#:   :data:`~atif_duck.domain.catalog.DESCRIPTIONS`.
+ARG_EXEMPLARS: dict[str, str] = {
+    "sid": "(SELECT session_id FROM sessions LIMIT 1)",
+    "since_days": "30",
+    "last_n_days": "30",
+    "k": "10",
+    "n": "10",
+    "label_name": "'correction'",
+    "signal_name": "'correction'",
+    "src": "'user'",
+    "msg": "'Stop hook feedback: the claim has no evidence'",
+    "interval_text": "'7 days'",
+    "query_vec": (
+        "(SELECT list_transform("
+        "me.embedding, v -> (v / sqrt(list_sum("
+        "list_transform(me.embedding, x -> x::DOUBLE * x))))::FLOAT"
+        ") FROM message_embeddings me LIMIT 1)"
+    ),
+}
+
+#: Catalog objects DELIBERATELY without an example, mapped to the reason.
+#: The coverage drift test treats these as documented exclusions; everything
+#: else in the catalogs MUST yield an example. Currently empty — every view
+#: and macro is exampled — but the seam stays so a future exclusion is a
+#: reviewed one-liner instead of a silent coverage hole.
+EXCLUSIONS: dict[str, str] = {}
+
+
+@dataclass(frozen=True, slots=True)
+class ExampleQuery:
+    """One runnable example query, derived from the catalog.
+
+    Attributes
+    ----------
+    name
+        The catalog object (view or macro) the example demonstrates.
+    sql
+        A complete, runnable statement for ``atif-sql query``.
+    description
+        Agent-facing prose from :data:`~atif_duck.domain.catalog.DESCRIPTIONS`.
+    requires
+        ``core`` | ``analytics`` | ``vss`` — what must be registered/populated.
+    category
+        ``view`` | ``table-macro`` | ``scalar-macro``.
+    """
+
+    name: str
+    sql: str
+    description: str
+    requires: Requires
+    category: Category
+
+
+#: Catalog objects that need the embeddings store (``atif-sql embed``).
+_VSS_OBJECTS: frozenset[str] = frozenset({"message_embeddings", "semantic_search"})
+
+
+def object_requires(name: str) -> Requires:
+    """What must be registered/populated before catalog object ``name`` binds.
+
+    ``analytics`` for the parquet-gated analytics views and macros, ``vss``
+    for the embeddings surface, ``core`` for everything the corpus alone
+    registers (the authorship views included). ``atif-sql schema`` prints
+    this beside every view and macro, and each derived example carries it.
+    """
+    if name in ANALYTICS_VIEW_SCHEMA or name in ANALYTICS_MACRO_SIGNATURES:
+        return "analytics"
+    return "vss" if name in _VSS_OBJECTS else "core"
+
+
+def _view_sql(name: str) -> str:
+    """Canonical top-N SELECT for one view.
+
+    When the static :data:`~atif_duck.domain.catalog.VIEW_SCHEMA` (or
+    :data:`~atif_duck.domain.catalog.ANALYTICS_VIEW_SCHEMA`) declares a
+    timestamp column, the example orders by the FIRST one descending
+    (newest-first is the agent's default question); otherwise a plain LIMIT.
+    """
+    schema = VIEW_SCHEMA.get(name) or ANALYTICS_VIEW_SCHEMA.get(name, ())
+    ts_col = next((col for col, typ in schema if typ.startswith("TIMESTAMP")), None)
+    if ts_col is not None:
+        # The view name and its timestamp column both come from the static catalog.
+        return f"SELECT * FROM {name} ORDER BY {ts_col} DESC LIMIT 10"  # noqa: S608  # nosec B608 - view name and column are VIEW_SCHEMA constants
+    return f"SELECT * FROM {name} LIMIT 10"  # noqa: S608  # nosec B608 - view name is a VIEW_SCHEMA constant
+
+
+def _macro_sql(name: str, params: tuple[str, ...]) -> str:
+    """Runnable invocation for one macro, args filled from the exemplars."""
+    try:
+        args = ", ".join(ARG_EXEMPLARS[param] for param in params)
+    except KeyError as exc:
+        msg = (
+            f"macro {name!r} has parameter {exc.args[0]!r} with no entry in "
+            "ARG_EXEMPLARS — add an exemplar literal so build_examples() can "
+            "derive a runnable example"
+        )
+        raise KeyError(msg) from exc
+    if name in TABLE_MACRO_NAMES:
+        # The macro name comes from MACRO_SIGNATURES and args from ARG_EXEMPLARS.
+        return f"SELECT * FROM {name}({args}) LIMIT 10"  # noqa: S608  # nosec B608 - macro name from MACRO_SIGNATURES, args from ARG_EXEMPLARS
+    return f"SELECT {name}({args}) AS {name}"
+
+
+def _description(name: str) -> str:
+    """DESCRIPTIONS lookup that fails loud on a missing catalog object."""
+    try:
+        return DESCRIPTIONS[name]
+    except KeyError as exc:
+        msg = (
+            f"catalog object {name!r} has no entry in "
+            "atif_duck.domain.catalog.DESCRIPTIONS — add one (the drift test "
+            "in tests/test_examples.py pins this coverage)"
+        )
+        raise KeyError(msg) from exc
+
+
+def build_examples() -> tuple[ExampleQuery, ...]:
+    """Derive the full example inventory from the catalogs.
+
+    Order: core views, VSS view, analytics views, then core macros, VSS
+    macro, analytics macros — stable, so the CLI listing and the JSON array
+    are deterministic across runs.
+    """
+    examples: list[ExampleQuery] = []
+
+    for view in VIEW_NAMES:
+        if view in EXCLUSIONS:
+            continue
+        examples.append(
+            ExampleQuery(
+                name=view,
+                sql=_view_sql(view),
+                description=_description(view),
+                requires=object_requires(view),
+                category="view",
+            )
+        )
+
+    for view in ANALYTICS_VIEW_NAMES:
+        if view in EXCLUSIONS:
+            continue
+        examples.append(
+            ExampleQuery(
+                name=view,
+                sql=_view_sql(view),
+                description=_description(view),
+                requires=object_requires(view),
+                category="view",
+            )
+        )
+
+    for macro, params in MACRO_SIGNATURES.items():
+        if macro in EXCLUSIONS:
+            continue
+        examples.append(
+            ExampleQuery(
+                name=macro,
+                sql=_macro_sql(macro, params),
+                description=_description(macro),
+                requires=object_requires(macro),
+                category="table-macro" if macro in TABLE_MACRO_NAMES else "scalar-macro",
+            )
+        )
+
+    for macro, params in ANALYTICS_MACRO_SIGNATURES.items():
+        if macro in EXCLUSIONS:
+            continue
+        examples.append(
+            ExampleQuery(
+                name=macro,
+                sql=_macro_sql(macro, params),
+                description=_description(macro),
+                requires=object_requires(macro),
+                category="table-macro" if macro in TABLE_MACRO_NAMES else "scalar-macro",
+            )
+        )
+
+    return tuple(examples)
+
+
+__all__ = [
+    "ARG_EXEMPLARS",
+    "CATEGORY_VALUES",
+    "EXCLUSIONS",
+    "REQUIRES_VALUES",
+    "Category",
+    "ExampleQuery",
+    "Requires",
+    "build_examples",
+    "object_requires",
+]