atif-sql 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (164) hide show
  1. atif_analytics/__init__.py +9 -0
  2. atif_analytics/application/__init__.py +3 -0
  3. atif_analytics/application/analyze.py +158 -0
  4. atif_analytics/application/prompts.py +750 -0
  5. atif_analytics/application/use_cases/__init__.py +3 -0
  6. atif_analytics/application/use_cases/_shared.py +264 -0
  7. atif_analytics/application/use_cases/classify.py +430 -0
  8. atif_analytics/application/use_cases/conflicts.py +448 -0
  9. atif_analytics/application/use_cases/friction.py +634 -0
  10. atif_analytics/application/use_cases/perceived.py +451 -0
  11. atif_analytics/domain/__init__.py +3 -0
  12. atif_analytics/domain/authorship.py +166 -0
  13. atif_analytics/domain/config.py +30 -0
  14. atif_analytics/domain/costs.py +50 -0
  15. atif_analytics/domain/friction.py +117 -0
  16. atif_analytics/domain/layout.py +94 -0
  17. atif_analytics/domain/models.py +385 -0
  18. atif_analytics/domain/ports.py +100 -0
  19. atif_analytics/domain/transcript.py +359 -0
  20. atif_analytics/infrastructure/__init__.py +3 -0
  21. atif_analytics/infrastructure/corpus_reader.py +546 -0
  22. atif_analytics/infrastructure/parquet_cache.py +186 -0
  23. atif_analytics/infrastructure/settings.py +114 -0
  24. atif_analytics/infrastructure/sqlite_state/__init__.py +3 -0
  25. atif_analytics/infrastructure/sqlite_state/checkpointer.py +326 -0
  26. atif_analytics/infrastructure/sqlite_state/retry_queue.py +296 -0
  27. atif_analytics/py.typed +0 -0
  28. atif_cli/__init__.py +3 -0
  29. atif_cli/__main__.py +8 -0
  30. atif_cli/app.py +2459 -0
  31. atif_cli/converter_adapter.py +151 -0
  32. atif_cli/corpus.py +409 -0
  33. atif_cli/cron.py +311 -0
  34. atif_cli/duck_errors.py +97 -0
  35. atif_cli/embed_lake.py +101 -0
  36. atif_cli/errors.py +89 -0
  37. atif_cli/lake.py +373 -0
  38. atif_cli/lake_sessions.py +270 -0
  39. atif_cli/output.py +297 -0
  40. atif_cli/py.typed +0 -0
  41. atif_converter/__init__.py +3 -0
  42. atif_converter/application/__init__.py +3 -0
  43. atif_converter/application/convert_and_audit.py +294 -0
  44. atif_converter/application/convert_codex.py +210 -0
  45. atif_converter/domain/__init__.py +9 -0
  46. atif_converter/domain/agents.py +40 -0
  47. atif_converter/domain/atif/__init__.py +55 -0
  48. atif_converter/domain/atif/agent.py +37 -0
  49. atif_converter/domain/atif/content.py +161 -0
  50. atif_converter/domain/atif/final_metrics.py +45 -0
  51. atif_converter/domain/atif/metrics.py +49 -0
  52. atif_converter/domain/atif/observation.py +21 -0
  53. atif_converter/domain/atif/observation_result.py +47 -0
  54. atif_converter/domain/atif/step.py +144 -0
  55. atif_converter/domain/atif/subagent_trajectory_ref.py +90 -0
  56. atif_converter/domain/atif/tool_call.py +36 -0
  57. atif_converter/domain/atif/trajectory.py +211 -0
  58. atif_converter/domain/atif/trajectory_validator.py +299 -0
  59. atif_converter/domain/blobs.py +510 -0
  60. atif_converter/domain/claude_code_conversion.py +1077 -0
  61. atif_converter/domain/codex_conversion.py +1036 -0
  62. atif_converter/domain/codex_edges.py +192 -0
  63. atif_converter/domain/codex_enrichment.py +419 -0
  64. atif_converter/domain/codex_fidelity.py +124 -0
  65. atif_converter/domain/codex_nested_calls.py +165 -0
  66. atif_converter/domain/edges.py +113 -0
  67. atif_converter/domain/enrichment.py +412 -0
  68. atif_converter/domain/errors.py +60 -0
  69. atif_converter/domain/fidelity.py +171 -0
  70. atif_converter/domain/model_prices.json +8538 -0
  71. atif_converter/domain/pricing.py +790 -0
  72. atif_converter/domain/result_signals.py +774 -0
  73. atif_converter/domain/schema_version.py +90 -0
  74. atif_converter/domain/session_events.py +408 -0
  75. atif_converter/infrastructure/__init__.py +3 -0
  76. atif_converter/infrastructure/census.py +89 -0
  77. atif_converter/infrastructure/claude_code_converter.py +127 -0
  78. atif_converter/infrastructure/codex_adapter.py +132 -0
  79. atif_converter/infrastructure/codex_census.py +100 -0
  80. atif_converter/infrastructure/codex_converter.py +82 -0
  81. atif_converter/infrastructure/harbor_adapter.py +187 -0
  82. atif_converter/infrastructure/raw_records.py +426 -0
  83. atif_converter/infrastructure/source_archive.py +154 -0
  84. atif_converter/py.typed +0 -0
  85. atif_corpus/__init__.py +13 -0
  86. atif_corpus/application/__init__.py +7 -0
  87. atif_corpus/application/materialize.py +1826 -0
  88. atif_corpus/domain/__init__.py +9 -0
  89. atif_corpus/domain/agents.py +33 -0
  90. atif_corpus/domain/generation.py +34 -0
  91. atif_corpus/domain/layout.py +168 -0
  92. atif_corpus/domain/ports.py +175 -0
  93. atif_corpus/domain/session_id.py +68 -0
  94. atif_corpus/domain/sessions.py +228 -0
  95. atif_corpus/domain/slug.py +72 -0
  96. atif_corpus/domain/source_layout.py +143 -0
  97. atif_corpus/domain/watermark.py +76 -0
  98. atif_corpus/infrastructure/__init__.py +7 -0
  99. atif_corpus/infrastructure/atomic.py +266 -0
  100. atif_corpus/infrastructure/compress_artifacts.py +161 -0
  101. atif_corpus/infrastructure/fake_converter.py +174 -0
  102. atif_corpus/infrastructure/scanner.py +316 -0
  103. atif_corpus/infrastructure/settings.py +130 -0
  104. atif_corpus/infrastructure/source_archive.py +202 -0
  105. atif_corpus/py.typed +0 -0
  106. atif_duck/__init__.py +43 -0
  107. atif_duck/domain/__init__.py +32 -0
  108. atif_duck/domain/artifacts.py +66 -0
  109. atif_duck/domain/authorship.py +165 -0
  110. atif_duck/domain/catalog.py +626 -0
  111. atif_duck/domain/columnar.py +149 -0
  112. atif_duck/domain/embedding_guard.py +88 -0
  113. atif_duck/domain/examples.py +265 -0
  114. atif_duck/domain/lake.py +427 -0
  115. atif_duck/domain/lake_sessions.py +97 -0
  116. atif_duck/domain/raw_readers.py +140 -0
  117. atif_duck/domain/session_id.py +69 -0
  118. atif_duck/domain/sql_literal.py +50 -0
  119. atif_duck/infrastructure/__init__.py +34 -0
  120. atif_duck/infrastructure/analytics.py +469 -0
  121. atif_duck/infrastructure/authorship.py +139 -0
  122. atif_duck/infrastructure/columnar.py +641 -0
  123. atif_duck/infrastructure/lake.py +1383 -0
  124. atif_duck/infrastructure/lake_sessions.py +194 -0
  125. atif_duck/infrastructure/lake_settings.py +63 -0
  126. atif_duck/infrastructure/lake_steps.py +225 -0
  127. atif_duck/infrastructure/projections.py +221 -0
  128. atif_duck/infrastructure/registry.py +2194 -0
  129. atif_duck/infrastructure/stored_artifacts.py +115 -0
  130. atif_duck/py.typed +0 -0
  131. atif_embed/__init__.py +12 -0
  132. atif_embed/application/__init__.py +3 -0
  133. atif_embed/application/embed.py +293 -0
  134. atif_embed/application/prune.py +80 -0
  135. atif_embed/domain/__init__.py +3 -0
  136. atif_embed/domain/discovery.py +70 -0
  137. atif_embed/domain/embedding_guard.py +83 -0
  138. atif_embed/domain/errors.py +94 -0
  139. atif_embed/domain/ports.py +223 -0
  140. atif_embed/domain/sql_literal.py +49 -0
  141. atif_embed/domain/text_stamp.py +71 -0
  142. atif_embed/infrastructure/__init__.py +3 -0
  143. atif_embed/infrastructure/cohere_bedrock.py +433 -0
  144. atif_embed/infrastructure/corpus_text_rows.py +273 -0
  145. atif_embed/infrastructure/gemma_local.py +355 -0
  146. atif_embed/infrastructure/lake_text_rows.py +203 -0
  147. atif_embed/infrastructure/lance_store.py +440 -0
  148. atif_embed/infrastructure/providers.py +30 -0
  149. atif_embed/infrastructure/settings.py +154 -0
  150. atif_embed/py.typed +0 -0
  151. atif_models/__init__.py +14 -0
  152. atif_models/domain/__init__.py +9 -0
  153. atif_models/domain/ports.py +143 -0
  154. atif_models/domain/registry.py +151 -0
  155. atif_models/domain/schema.py +109 -0
  156. atif_models/infrastructure/__init__.py +8 -0
  157. atif_models/infrastructure/openai_bedrock.py +314 -0
  158. atif_models/infrastructure/settings.py +85 -0
  159. atif_models/py.typed +0 -0
  160. atif_sql-0.2.0.dist-info/METADATA +272 -0
  161. atif_sql-0.2.0.dist-info/RECORD +164 -0
  162. atif_sql-0.2.0.dist-info/WHEEL +4 -0
  163. atif_sql-0.2.0.dist-info/entry_points.txt +2 -0
  164. atif_sql-0.2.0.dist-info/licenses/LICENSE +202 -0
atif_cli/app.py ADDED
@@ -0,0 +1,2459 @@
1
+ # SPDX-License-Identifier: Apache-2.0
2
+
3
+ """The atif-sql CLI (cyclopts) — the workspace's composition root.
4
+
5
+ The ONLY package that may import another workspace member, and it imports
6
+ five: atif-converter, atif-corpus, atif-duck, atif-embed, and atif-analytics.
7
+ Those five may never import each other, so every cross-package seam (the
8
+ ConverterPort adapter, the clock, version pins, the DuckDB connection) is
9
+ wired here, per docs/CONTRACT.md §CLI.
10
+
11
+ The twelve registered commands — nine ``@app.command`` functions plus the
12
+ ``cron``, ``lake`` and ``corpus`` sub-apps:
13
+
14
+ * ``convert`` one-shot convert+audit for a single session JSONL
15
+ * ``materialize`` sync the materialized corpus (RealConverter behind the port)
16
+ * ``status`` corpus freshness — read-only, fast
17
+ * ``query`` SQL over the atif-duck views/macros with classified errors
18
+ * ``analyze`` run the analytics pipelines — CALLS BEDROCK, spends money
19
+ * ``embed`` backfill step embeddings into LanceDB — CALLS BEDROCK
20
+ * ``search`` embed one query string, then kNN over the store — CALLS BEDROCK
21
+ * ``examples`` derived, test-executed example queries (also ``query --examples``)
22
+ * ``schema`` static catalog dump, no DuckDB bind, <50ms
23
+ * ``cron`` sub-app: ``install`` (prints, never writes) and ``status``
24
+ * ``lake`` sub-app: ``rebuild``, ``verify``, ``status``, ``compact`` for
25
+ the DuckLake every corpus is queried through
26
+ * ``corpus`` sub-app: ``slim`` converts a corpus to the compressed layout
27
+
28
+ A command that reaches Bedrock is marked above because its spend is not
29
+ recoverable. Each one guards itself differently: ``analyze`` defaults to a DRY
30
+ RUN that only plans and estimates, a bare ``embed`` exits 64 rather than
31
+ choosing a scope for the caller, and ``search`` embeds exactly one query
32
+ string per invocation.
33
+
34
+ Two agents
35
+ ----------
36
+ ``convert``, ``materialize`` and ``status`` take ``--agent claude-code|codex``.
37
+ The flag picks three things at once and they must move together: which
38
+ converter reads a transcript, which discovery layout finds one, and which
39
+ source root and corpus slug the settings default to. One corpus root therefore
40
+ holds exactly one agent's sessions, which is what lets everything downstream —
41
+ the DuckDB views, the analytics pipelines, the embedding store — stay unaware
42
+ that a second agent exists.
43
+
44
+ Agent-friendly defaults
45
+ -----------------------
46
+ * ``--format auto`` emits a human table on a TTY and JSON on a pipe.
47
+ * DuckDB errors classify into parse/catalog/runtime -> exit 64/65/70 with a
48
+ JSON error envelope on non-TTY (see :mod:`atif_cli.errors`).
49
+
50
+ Heavy imports (duckdb, the ATIF models via atif_converter, pydantic via atif_corpus)
51
+ are DEFERRED into the command bodies that use them so the fast path
52
+ (``schema`` / ``--help`` / ``--version``) stays on a lean import graph —
53
+ pinned by the fresh-interpreter lean-import test in ``tests/test_lean_import``.
54
+ """
55
+
56
+ from __future__ import annotations
57
+
58
+ import json
59
+ import os
60
+ import re
61
+ import shutil
62
+ import sys
63
+ import tempfile
64
+ from contextlib import contextmanager
65
+ from dataclasses import dataclass
66
+ from pathlib import Path
67
+ from typing import TYPE_CHECKING, Annotated, Any
68
+
69
+ import cyclopts
70
+
71
+ from atif_cli.corpus import corpus_app
72
+ from atif_cli.cron import cron_app
73
+ from atif_cli.errors import EXIT_CODES, ClassifiedError
74
+ from atif_cli.lake import lake_app
75
+ from atif_cli.output import (
76
+ OutputFormat,
77
+ emit_cursor,
78
+ emit_error,
79
+ emit_json,
80
+ emit_rows,
81
+ resolve_format,
82
+ )
83
+
84
+ if TYPE_CHECKING:
85
+ from collections.abc import Generator, Sequence
86
+
87
+ from atif_corpus.application.materialize import MaterializationReport
88
+ from atif_embed.infrastructure.settings import EmbedSettings
89
+
90
+ app = cyclopts.App(
91
+ name="atif-sql",
92
+ help="ATIF-native analytics over agent trajectories.",
93
+ # The DISTRIBUTION, which is not the module. cyclopts resolves a default
94
+ # `--version` by looking the calling module's name up in the installed metadata;
95
+ # the module is `atif_cli` and the distribution is `atif-sql`, so the default
96
+ # lookup raises PackageNotFoundError and cyclopts answers `0.0.0`. A callable
97
+ # also keeps `importlib.metadata` out of the import path until `--version` is
98
+ # actually asked for, which is what the lean-import test pins.
99
+ version=lambda: _version_of("atif-sql"),
100
+ )
101
+
102
+ # The `cron` subcommand group (install / status) — lean stdlib+cyclopts module,
103
+ # safe to register eagerly (see atif_cli.cron's module docstring).
104
+ app.command(cron_app)
105
+ # The `lake` subcommand group (rebuild / verify / status / compact): lean at
106
+ # import like `cron`, every heavy import deferred into its command bodies.
107
+ app.command(lake_app)
108
+ # The `corpus` subcommand group (slim): lean at import, like `lake`.
109
+ app.command(corpus_app)
110
+
111
+
112
+ # ---------------------------------------------------------------------------
113
+ # Shared plumbing
114
+ # ---------------------------------------------------------------------------
115
+
116
+
117
+ def _now_iso() -> str:
118
+ """ISO-8601 UTC instant — the CLI layer owns the wall clock."""
119
+ from datetime import UTC, datetime
120
+
121
+ return datetime.now(tz=UTC).isoformat()
122
+
123
+
124
+ def _version_of(distribution: str) -> str:
125
+ """Installed version of ``distribution``, or ``"unknown"`` off-venv."""
126
+ import importlib.metadata
127
+
128
+ try:
129
+ return importlib.metadata.version(distribution)
130
+ except importlib.metadata.PackageNotFoundError:
131
+ return "unknown"
132
+
133
+
134
+ def _sql_str(value: str) -> str:
135
+ """Escape a Python string as a single-quoted SQL literal."""
136
+ escaped = value.replace("'", "''")
137
+ return f"'{escaped}'"
138
+
139
+
140
+ #: Env override for the DuckDB memory cap ``query`` runs under: any size
141
+ #: literal DuckDB's ``SET memory_limit`` accepts (``6GB``, ``512MiB``, ``2000000000B``).
142
+ QUERY_MEMORY_LIMIT_ENV = "ATIF_SQL_QUERY_MEMORY_LIMIT"
143
+
144
+ #: Env override for the DuckDB thread count ``query`` runs under.
145
+ QUERY_THREADS_ENV = "ATIF_SQL_QUERY_THREADS"
146
+
147
+ #: Set to ``1`` to let ``query``, ``search`` and ``analyze`` run as uid 0.
148
+ ALLOW_ROOT_ENV = "ATIF_SQL_ALLOW_ROOT"
149
+
150
+ #: Target for ``query``'s DuckDB memory cap when the host can afford it: the
151
+ #: base views (``tool_calls``, ``tool_rank``) need several GiB on a multi-GB
152
+ #: corpus, and a tighter cap turns working queries into OutOfMemory. A
153
+ #: target, not a floor: it never exceeds what the host has (the old 8 GiB
154
+ #: floor set a "limit" above physical RAM on an 8 GiB guest).
155
+ _QUERY_MEMORY_TARGET_BYTES = 8 * 1024**3
156
+
157
+ #: Never cap below this. DuckDB needs room to open the readers at all, and a
158
+ #: host with less available than this fails to register whatever the cap says.
159
+ _QUERY_MEMORY_MIN_BYTES = 512 * 1024**2
160
+
161
+ #: Cap budgeted per DuckDB thread when deriving the thread count. The JSON
162
+ #: readers reserve about twice their ``maximum_object_size`` per thread, and
163
+ #: that bound is now sized from the largest file present (the largest
164
+ #: trajectory.json seen is 436 MB), so 2 GiB a thread keeps a corpus on the
165
+ #: JSON path registering under the cap.
166
+ _QUERY_BYTES_PER_THREAD = 2 * 1024**3
167
+
168
+
169
+ #: Where this process's cgroup v2 membership is listed (``0::/<path>``).
170
+ _PROC_SELF_CGROUP = Path("/proc/self/cgroup")
171
+ #: The cgroup v2 unified hierarchy's mount point.
172
+ _CGROUP_ROOT = Path("/sys/fs/cgroup")
173
+ #: /proc/meminfo, read for ``MemAvailable``.
174
+ _PROC_MEMINFO = Path("/proc/meminfo")
175
+
176
+
177
+ def _read_cgroup_int(path: Path) -> int | None:
178
+ """An integer cgroup file (``memory.max``, ``memory.current``); ``None`` for ``max`` or unreadable."""
179
+ try:
180
+ text = path.read_text(encoding="ascii").strip()
181
+ except OSError:
182
+ return None
183
+ return int(text) if text.isdigit() else None
184
+
185
+
186
+ def _reclaimable_file_bytes(level: Path) -> int:
187
+ """Page cache the kernel can reclaim from this cgroup: ``active_file`` plus ``inactive_file``.
188
+
189
+ Read from ``memory.stat``; 0 when the file is missing or unreadable, which
190
+ leaves ``memory.current`` as the (conservative) usage.
191
+ """
192
+ try:
193
+ text = (level / "memory.stat").read_text(encoding="ascii")
194
+ except OSError:
195
+ return 0
196
+ total = 0
197
+ for line in text.splitlines():
198
+ key, _, value = line.partition(" ")
199
+ if key in {"active_file", "inactive_file"} and value.strip().isdigit():
200
+ total += int(value)
201
+ return total
202
+
203
+
204
+ def _cgroup_memory() -> tuple[int | None, int | None]:
205
+ """``(limit, headroom)`` bytes from this process's cgroup v2 chain; ``None`` when unlimited.
206
+
207
+ A limit can sit on any ancestor (``systemd-run --scope -p MemoryMax=``
208
+ sets it on the scope; a container runtime on a parent), so every level
209
+ from the process's own cgroup up to the root is read: the limit is the
210
+ smallest ``memory.max``, and the headroom the smallest ``memory.max`` minus
211
+ what that level holds and cannot give back. ``memory.current`` counts page
212
+ cache, which the kernel reclaims under pressure, so the level's
213
+ ``active_file`` and ``inactive_file`` from ``memory.stat`` are subtracted
214
+ first, the same way ``MemAvailable`` counts reclaimable cache as free on
215
+ the host. Without that, a slice full of cache reports almost no room and
216
+ every query and lake command drops to the floor. A host without cgroup v2
217
+ (macOS, a v1 host) has no ``0::`` line and answers ``(None, None)``.
218
+ """
219
+ try:
220
+ lines = _PROC_SELF_CGROUP.read_text(encoding="ascii").splitlines()
221
+ except OSError:
222
+ return None, None
223
+ relative = next((line[3:] for line in lines if line.startswith("0::")), None)
224
+ if relative is None:
225
+ return None, None
226
+ level = _CGROUP_ROOT / relative.lstrip("/")
227
+ limit: int | None = None
228
+ headroom: int | None = None
229
+ while True:
230
+ level_max = _read_cgroup_int(level / "memory.max")
231
+ if level_max is not None:
232
+ limit = level_max if limit is None else min(limit, level_max)
233
+ used = _read_cgroup_int(level / "memory.current") or 0
234
+ used = max(0, used - _reclaimable_file_bytes(level))
235
+ room = max(0, level_max - used)
236
+ headroom = room if headroom is None else min(headroom, room)
237
+ if level == _CGROUP_ROOT or _CGROUP_ROOT not in level.parents:
238
+ break
239
+ level = level.parent
240
+ return limit, headroom
241
+
242
+
243
+ def _host_memory() -> tuple[int, int]:
244
+ """``(physical, available)`` bytes this process can use.
245
+
246
+ Physical comes from ``os.sysconf``. Available is Linux's ``MemAvailable``
247
+ from ``/proc/meminfo`` (what a new allocation can really get, reclaimable
248
+ page cache included); where that file is absent (macOS) available is
249
+ taken as physical. A cgroup v2 ``memory.max`` below either one lowers it
250
+ (:func:`_cgroup_memory`): /proc/meminfo describes the host, and a query
251
+ capped by its cgroup that sized itself from the host would be OOM-killed
252
+ instead of spilling.
253
+ """
254
+ physical = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
255
+ available = physical
256
+ try:
257
+ with _PROC_MEMINFO.open(encoding="ascii") as handle:
258
+ for line in handle:
259
+ if line.startswith("MemAvailable:"):
260
+ available = int(line.split()[1]) * 1024
261
+ break
262
+ except (OSError, ValueError, IndexError):
263
+ pass
264
+ limit, headroom = _cgroup_memory()
265
+ if limit is not None:
266
+ physical = min(physical, limit)
267
+ if headroom is not None:
268
+ available = min(available, headroom)
269
+ return physical, min(physical, available)
270
+
271
+
272
+ def _query_memory_limit_bytes() -> int:
273
+ """Bytes to cap ``query``'s DuckDB heap at, derived from the host.
274
+
275
+ Half of physical RAM or :data:`_QUERY_MEMORY_TARGET_BYTES`, whichever is
276
+ larger, but never above 80% of physical RAM (DuckDB's own default) and
277
+ never above 80% of the memory available right now, floored at
278
+ :data:`_QUERY_MEMORY_MIN_BYTES`. On a 124 GiB host that is 62 GiB; on an
279
+ 8 GiB guest with 6 GiB free it is 4.8 GiB, where the previous rule said
280
+ 8 GiB and DuckDB's default said 6.4 GiB, neither of which the guest had.
281
+ """
282
+ physical, available = _host_memory()
283
+ target = max(int(physical * 0.5), _QUERY_MEMORY_TARGET_BYTES)
284
+ ceiling = min(int(physical * 0.8), int(available * 0.8))
285
+ return max(_QUERY_MEMORY_MIN_BYTES, min(target, ceiling))
286
+
287
+
288
+ def parse_size(text: str) -> int:
289
+ """Bytes for a DuckDB-style size literal; raises ``ValueError`` on anything else."""
290
+ return _parse_size(text)
291
+
292
+
293
+ def query_memory_limit_bytes() -> int:
294
+ """The host- and cgroup-derived DuckDB cap; the ``lake`` commands open their writers under it."""
295
+ return _query_memory_limit_bytes()
296
+
297
+
298
+ def _query_threads(memory_limit_bytes: int) -> int:
299
+ """DuckDB threads for ``query``: one per :data:`_QUERY_BYTES_PER_THREAD` of cap.
300
+
301
+ Capped at the CPUs this process may run on (``sched_getaffinity``, which
302
+ sees a container's cpuset where ``cpu_count`` does not) and floored at
303
+ one. DuckDB's own default is the core count, which is what made a 4 vCPU
304
+ guest reserve four readers' worth of memory it did not have.
305
+ """
306
+ affinity = getattr(os, "sched_getaffinity", None)
307
+ cpus = len(affinity(0)) if affinity is not None else (os.cpu_count() or 1)
308
+ return max(1, min(cpus, memory_limit_bytes // _QUERY_BYTES_PER_THREAD))
309
+
310
+
311
+ _SIZE_UNITS: dict[str, int] = {
312
+ "B": 1,
313
+ "KB": 10**3,
314
+ "MB": 10**6,
315
+ "GB": 10**9,
316
+ "TB": 10**12,
317
+ "KIB": 1024,
318
+ "MIB": 1024**2,
319
+ "GIB": 1024**3,
320
+ "TIB": 1024**4,
321
+ }
322
+
323
+ _SIZE_RE = re.compile(r"\s*(\d+(?:\.\d+)?)\s*([A-Za-z]*)\s*")
324
+
325
+
326
+ def _parse_size(text: str) -> int:
327
+ """Bytes for a DuckDB-style size literal (``6GB``, ``512 MiB``, ``123B``)."""
328
+ match = _SIZE_RE.fullmatch(text)
329
+ if match is None:
330
+ problem = f"{text!r} is not a size (expected e.g. 6GB, 512MiB, 2000000000B)"
331
+ raise ValueError(problem)
332
+ number, unit = match.group(1), (match.group(2) or "B").upper()
333
+ if unit not in _SIZE_UNITS:
334
+ problem = f"{text!r} has an unknown unit (use B, KB, MB, GB, TB, KiB, MiB, GiB, TiB)"
335
+ raise ValueError(problem)
336
+ return int(float(number) * _SIZE_UNITS[unit])
337
+
338
+
339
+ @dataclass(frozen=True, slots=True)
340
+ class QueryResources:
341
+ """What ``query`` hands DuckDB before it registers anything."""
342
+
343
+ memory_limit_bytes: int
344
+ threads: int
345
+
346
+
347
+ def _query_resources() -> QueryResources:
348
+ """Resolve the memory cap and thread count for one ``query`` process.
349
+
350
+ :data:`QUERY_MEMORY_LIMIT_ENV` and :data:`QUERY_THREADS_ENV` override the
351
+ host-derived values; the thread default follows whichever cap is in
352
+ force, so a caller who lowers the cap gets fewer threads for free.
353
+ Raises ``ValueError`` on a malformed override; the command turns that
354
+ into exit 64.
355
+ """
356
+ memory_env = os.environ.get(QUERY_MEMORY_LIMIT_ENV, "").strip()
357
+ threads_env = os.environ.get(QUERY_THREADS_ENV, "").strip()
358
+ if memory_env:
359
+ memory = _parse_size(memory_env)
360
+ if memory <= 0:
361
+ problem = f"{QUERY_MEMORY_LIMIT_ENV} must be a positive size"
362
+ raise ValueError(problem)
363
+ else:
364
+ memory = _query_memory_limit_bytes()
365
+ if threads_env:
366
+ threads = int(threads_env)
367
+ if threads < 1:
368
+ problem = f"{QUERY_THREADS_ENV} must be a positive integer"
369
+ raise ValueError(problem)
370
+ else:
371
+ threads = _query_threads(memory)
372
+ return QueryResources(memory_limit_bytes=memory, threads=threads)
373
+
374
+
375
+ def _configure_query_resources(con: Any, resources: QueryResources, spill_dir: Path) -> None:
376
+ """Size the connection to the host BEFORE registration.
377
+
378
+ Registration is the heaviest thing ``query`` does (the eager JSON readers
379
+ reserve memory per thread), so the cap and thread count must be in force
380
+ before it runs, not after, which is where they used to be applied. The
381
+ spill directory is set here for the same reason: a registration that
382
+ spills must spill into the private directory. Extension auto-install and
383
+ auto-load are switched off so nothing during registration can reach the
384
+ network; the lance extension is LOADed only when it is already installed
385
+ (:func:`atif_duck.infrastructure.registry.load_lance_extension`).
386
+ """
387
+ con.execute(f"SET threads={int(resources.threads)}")
388
+ con.execute(f"SET memory_limit='{int(resources.memory_limit_bytes)}B'")
389
+ con.execute(f"SET temp_directory={_sql_str(str(spill_dir))}")
390
+ con.execute("SET autoinstall_known_extensions=false")
391
+ con.execute("SET autoload_known_extensions=false")
392
+
393
+
394
+ @contextmanager
395
+ def _private_spill_dir() -> Generator[Path]:
396
+ """A per-process spill directory outside the corpus, gone when the process is.
397
+
398
+ ``mkdtemp`` creates it mode 0700 under the system temp dir. It is the
399
+ ONLY directory the sandbox grants, so it is where DuckDB spills a query
400
+ that exceeds ``memory_limit``; it used to be ``<corpus_root>/.duckdb_tmp``,
401
+ inside the tree ``materialize`` scans, and files written there by caller
402
+ SQL persisted between runs. Removed on every exit path, error included,
403
+ after the connection is closed.
404
+ """
405
+ path = Path(tempfile.mkdtemp(prefix="atif-sql-query-"))
406
+ try:
407
+ yield path
408
+ finally:
409
+ shutil.rmtree(path, ignore_errors=True)
410
+
411
+
412
+ def _refuse_root(command: str, fmt: OutputFormat) -> None:
413
+ """Exit 77 when running as uid 0, unless :data:`ALLOW_ROOT_ENV` is ``1``.
414
+
415
+ The corpus artifacts are written mode 0444, which is what stops a
416
+ ``COPY ... TO`` at a granted parquet from succeeding at the filesystem.
417
+ Root ignores file modes, so as uid 0 that protection is gone and caller
418
+ SQL could overwrite the corpus (measured: probe d2 of the MicroVM review
419
+ rewrote ``steps.parquet`` as a one-row file). The override exists for
420
+ containers that have no other user; it logs a warning so the run is on
421
+ record.
422
+ """
423
+ geteuid = getattr(os, "geteuid", None)
424
+ if geteuid is None or geteuid() != 0:
425
+ return
426
+ if os.environ.get(ALLOW_ROOT_ENV, "").strip() == "1":
427
+ from loguru import logger
428
+
429
+ logger.warning(
430
+ "atif-sql {} is running as root because {}=1; the 0444 file modes no longer "
431
+ "protect the corpus from the SQL this process runs",
432
+ command,
433
+ ALLOW_ROOT_ENV,
434
+ )
435
+ return
436
+ err = ClassifiedError(
437
+ kind="root_refused",
438
+ exit_code=EXIT_CODES["root_refused"],
439
+ message=f"atif-sql {command} refuses to run as root (uid 0)",
440
+ hint=f"run it as an unprivileged user, or set {ALLOW_ROOT_ENV}=1 to override "
441
+ "(a warning is logged; root can then overwrite corpus files from SQL)",
442
+ )
443
+ emit_error(err, fmt)
444
+ raise SystemExit(err.exit_code)
445
+
446
+
447
+ #: Statement kinds the query sandbox executes. Everything else is refused
448
+ #: before execution, by name, using DuckDB's own parser on the hardened
449
+ #: connection. The refused kinds are the ones that name a file or defer a
450
+ #: statement past this check: COPY and COPY_DATABASE (DuckDB's ``allowed_paths``
451
+ #: grants are read-write, so ``COPY ... TO <granted parquet> (USE_TMP_FILE
452
+ #: false)`` would overwrite corpus data, and root ignores the 0444 mode that
453
+ #: used to stop it), EXPORT, ATTACH, DETACH, LOAD, EXTENSION (INSTALL),
454
+ #: PREPARE and EXECUTE (a prepared COPY parses as PREPARE). An allowlist, so a
455
+ #: statement kind a future DuckDB adds is refused until someone reads what it
456
+ #: does.
457
+ _QUERY_STATEMENT_KINDS: frozenset[str] = frozenset(
458
+ {
459
+ "SELECT",
460
+ "EXPLAIN",
461
+ "SET",
462
+ "VARIABLE_SET",
463
+ "CREATE",
464
+ "CREATE_FUNC",
465
+ "DROP",
466
+ "ALTER",
467
+ "INSERT",
468
+ "UPDATE",
469
+ "DELETE",
470
+ "MERGE_INTO",
471
+ "CALL",
472
+ "PRAGMA",
473
+ "TRANSACTION",
474
+ "VACUUM",
475
+ "ANALYZE",
476
+ }
477
+ )
478
+
479
+
480
+ def _refused_statement_kinds(con: Any, sql: str) -> list[str]:
481
+ """Statement kinds in ``sql`` outside :data:`_QUERY_STATEMENT_KINDS`, in order, once each.
482
+
483
+ ``extract_statements`` is the parser ``execute`` uses, so there is no
484
+ second grammar to disagree with it, and it runs on the hardened
485
+ connection so a PRAGMA that reads a file while parsing meets the same
486
+ allowlist the query would. A parse error propagates as DuckDB's own
487
+ ``ParserException`` (exit 64), exactly as ``execute`` would have raised it.
488
+ """
489
+ refused: list[str] = []
490
+ for statement in con.extract_statements(sql):
491
+ kind = str(statement.type.name)
492
+ if kind not in _QUERY_STATEMENT_KINDS and kind not in refused:
493
+ refused.append(kind)
494
+ return refused
495
+
496
+
497
+ #: Config options caller SQL may still change after the sandbox locks. Only
498
+ #: ``TimeZone``: rendering timestamps in the reader's zone is an ordinary
499
+ #: analytics need and grants no filesystem or network reach. Everything else
500
+ #: (``threads``, ``errors_as_json``, ``memory_limit``, the allowlists) stays
501
+ #: frozen, because a widenable allowlist is not an allowlist.
502
+ _QUERY_UNLOCKED_CONFIGS: tuple[str, ...] = ("TimeZone",)
503
+
504
+
505
+ def _lazy_read_paths(*corpus_roots: Path) -> list[Path]:
506
+ """Analytics parquets a registered view reads LAZILY, at caller-query time.
507
+
508
+ ``register_raw`` materializes the JSON session artifacts into TEMP
509
+ TABLEs, and ``register_vss`` holds the Lance store open through an
510
+ ATTACH, so both survive the sandbox with no allowlist entry. The
511
+ analytics views are ``read_parquet`` over files on disk, bound but not
512
+ read until the caller selects from them — those paths are the ones that
513
+ must stay reachable. The other lazily-read set, the per-session columnar
514
+ parquets, is reported by ``register`` itself (``RawSources.lazy_read_paths``)
515
+ because only the registry knows which sessions it bound that way.
516
+ """
517
+ return sorted(
518
+ p for root in corpus_roots for p in (root / "analytics").rglob("*.parquet") if p.is_file()
519
+ )
520
+
521
+
522
+ def _harden_query_connection(
523
+ con: Any,
524
+ *,
525
+ corpus_root: Path,
526
+ spill_dir: Path,
527
+ columnar_paths: Sequence[Path] = (),
528
+ analytics_roots: Sequence[Path] = (),
529
+ ) -> None:
530
+ """Sandbox a fully-registered connection before it runs caller SQL.
531
+
532
+ ``query`` executes SQL an agent composed while reading third-party text
533
+ out of the transcript corpus, so the statement is untrusted: without this
534
+ it reaches ``read_text`` on any file the user can read, ``COPY`` to any
535
+ path, ``ATTACH`` of unrelated databases, and ``INSTALL httpfs`` for
536
+ network egress.
537
+
538
+ The memory cap, thread count and spill directory are NOT set here any
539
+ more: :func:`_configure_query_resources` applies them before
540
+ ``register`` runs, because registration is what needs them. This
541
+ function arms the allowlists and freezes the configuration.
542
+
543
+ Reach is granted at two granularities, and the split is what keeps
544
+ injected SQL from overwriting the corpus it was read from:
545
+
546
+ * ``allowed_directories`` holds ONLY ``spill_dir``, the per-process
547
+ private directory from :func:`_private_spill_dir`, which must accept
548
+ writes for a query that exceeds ``memory_limit`` to complete at all.
549
+ It lives outside the corpus and is removed when the process exits, so
550
+ nothing written there persists and nothing under the corpus root is
551
+ writable. (DuckDB grants a directory read-write, so a ``COPY`` into it
552
+ would succeed at this layer; the statement gate in :func:`query`
553
+ refuses COPY before it gets here.)
554
+ * ``allowed_paths`` holds the individual files the views read lazily:
555
+ the analytics parquets from :func:`_lazy_read_paths` and the
556
+ per-session columnar parquets in ``columnar_paths`` (what
557
+ ``register`` bound this connection to). A file grant lets
558
+ ``read_parquet`` open that exact path. Nothing else under the corpus
559
+ is named, so ``read_text`` of a ``trajectory.json`` is refused. DuckDB
560
+ 1.5.5 has no read-only grant (measured: a ``read_parquet`` relation
561
+ bound before ``enable_external_access=false`` is refused at query time
562
+ without a grant, and ``duckdb_settings()`` lists no write switch), so
563
+ each grant is read-write and ``COPY ... TO <granted parquet>
564
+ (USE_TMP_FILE false)`` would overwrite it. Two things close that: the
565
+ statement gate refuses COPY for any uid, and :func:`_refuse_root`
566
+ keeps the producer's 0444 modes meaningful by refusing uid 0.
567
+
568
+ Four ordering constraints, each one required:
569
+
570
+ * Both allowlists must be set BEFORE ``enable_external_access``. On their
571
+ own they block nothing; disabling external access is what arms them,
572
+ and DuckDB then refuses to widen either list.
573
+ * Those settings are rejected pre-boot, so they cannot move into
574
+ ``duckdb.connect(config=...)`` — and they must follow ``register`` so
575
+ the corpus globs and the Lance ATTACH still bind.
576
+ * ``allowed_configs`` must precede ``lock_configuration``, which is what
577
+ makes its exemptions mean anything.
578
+ * ``lock_configuration`` must come LAST (it freezes every option above,
579
+ including itself). Without it the memory cap is decorative: caller SQL
580
+ can just ``SET memory_limit`` back up.
581
+ """
582
+ con.execute(f"SET allowed_directories=[{_sql_str(str(spill_dir))}]")
583
+ lazy_paths = ", ".join(
584
+ _sql_str(str(path))
585
+ for path in (*columnar_paths, *_lazy_read_paths(*(analytics_roots or (corpus_root,))))
586
+ )
587
+ con.execute(f"SET allowed_paths=[{lazy_paths}]")
588
+ unlocked = ", ".join(_sql_str(name) for name in _QUERY_UNLOCKED_CONFIGS)
589
+ con.execute(f"SET allowed_configs=[{unlocked}]")
590
+ con.execute("SET enable_external_access=false")
591
+ con.execute("SET lock_configuration=true")
592
+
593
+
594
+ def _corpus_settings(
595
+ source_root: Path | None,
596
+ corpus_root: Path | None,
597
+ agent: str | None = None,
598
+ ) -> Any:
599
+ """Resolve CorpusSettings, letting explicit flags override env/defaults.
600
+
601
+ When ``--source-root`` is given without ``--corpus-root``, the corpus
602
+ root re-derives from the OVERRIDDEN source root's slug, so re-pointing
603
+ the source can never write into another source's corpus. An explicit
604
+ ``ATIF_SQL_CORPUS_ROOT`` in the env wins over that re-derivation, because
605
+ a pinned root is a deliberate statement about where artifacts belong.
606
+
607
+ ``agent`` is applied at CONSTRUCTION rather than copied in afterwards,
608
+ because it is what the settings' own defaults for both roots derive from:
609
+ a copy would leave a Codex run pointed at the Claude Code source root.
610
+ """
611
+ import os
612
+
613
+ from atif_corpus.domain.agents import AgentSource as CorpusAgentSource
614
+ from atif_corpus.domain.slug import corpus_slug
615
+ from atif_corpus.infrastructure.settings import CorpusSettings
616
+
617
+ settings = (
618
+ CorpusSettings(agent=_resolve_agent(agent, CorpusAgentSource))
619
+ if agent is not None
620
+ else CorpusSettings()
621
+ )
622
+ if source_root is not None:
623
+ resolved_corpus = corpus_root
624
+ if resolved_corpus is None and "ATIF_SQL_CORPUS_ROOT" not in os.environ:
625
+ resolved_corpus = Path.home() / ".atif-sql" / "corpus" / corpus_slug(source_root)
626
+ settings = settings.model_copy(
627
+ update={
628
+ "source_root": source_root,
629
+ **({"corpus_root": resolved_corpus} if resolved_corpus is not None else {}),
630
+ }
631
+ )
632
+ elif corpus_root is not None:
633
+ settings = settings.model_copy(update={"corpus_root": corpus_root})
634
+ return settings
635
+
636
+
637
+ def _env_agent() -> str:
638
+ """The ``ATIF_SQL_AGENT`` setting, or the default spelling.
639
+
640
+ ``materialize`` / ``status`` / ``query`` read this through
641
+ ``CorpusSettings``, which is env-aware for every field. ``convert`` takes no
642
+ corpus settings at all, so it would have ignored the variable and run the
643
+ Claude Code converter over a Codex rollout — which fails with "no
644
+ convertible events", blaming the transcript for the wrong adapter.
645
+ """
646
+ import os
647
+
648
+ return os.environ.get("ATIF_SQL_AGENT", "claude-code")
649
+
650
+
651
+ def _resolve_agent(value: str, agent_enum: Any) -> Any:
652
+ """Parse an ``--agent`` value, or exit 64 naming the accepted spellings.
653
+
654
+ cyclopts would coerce a StrEnum parameter itself, but the enum lives behind
655
+ a DEFERRED import (it comes from atif-converter, which builds the ATIF
656
+ models), and the lean-import test pins that atif_converter stays out of
657
+ the fast path. So the
658
+ flag is typed ``str`` at the signature and parsed here, inside the command
659
+ body, with the same exit code an unparseable path would get.
660
+ """
661
+ try:
662
+ return agent_enum(value)
663
+ except ValueError:
664
+ accepted = ", ".join(member.value for member in agent_enum)
665
+ print(f"error: unknown --agent {value!r} (expected one of: {accepted})", file=sys.stderr)
666
+ raise SystemExit(EXIT_CODES["invalid_input"]) from None
667
+
668
+
669
+ # ---------------------------------------------------------------------------
670
+ # convert
671
+ # ---------------------------------------------------------------------------
672
+
673
+
674
+ @app.command
675
+ def convert(
676
+ session_jsonl: Path,
677
+ *,
678
+ agent: str | None = None,
679
+ include_subagents: Annotated[bool, cyclopts.Parameter(negative="--no-subagents")] = True,
680
+ trajectory_out: Path | None = None,
681
+ ) -> None:
682
+ """Convert one agent transcript (Claude Code session or Codex rollout) to ATIF.
683
+
684
+ One-shot convert-and-audit (CONTRACT §CLI). With ``--trajectory-out``,
685
+ the ENRICHED trajectory is written there, ``edges.jsonl`` lands next to
686
+ it (always — the edges are part of the contract's conversion output),
687
+ and stdout carries the trajectory path plus the loss summary. Without
688
+ it, the trajectory JSON itself goes to stdout (edges are skipped: there
689
+ is no "next to the output" on a stream).
690
+
691
+ Parameters
692
+ ----------
693
+ session_jsonl
694
+ Path to the transcript: a Claude Code session
695
+ (``~/.claude/projects/<proj>/<session>.jsonl``) or a Codex rollout
696
+ (``~/.codex/sessions/<YYYY>/<MM>/<DD>/rollout-<ts>-<uuid>.jsonl``).
697
+ agent
698
+ Which agent wrote it: ``claude-code`` (default) or ``codex``. It
699
+ selects the harbor adapter, the record taxonomy and the fidelity
700
+ policy the loss report is written against.
701
+ include_subagents
702
+ Stage ``<session>/subagents/**.jsonl`` side-files alongside the main
703
+ chain. Claude Code only — a Codex rollout has no side-files, so the
704
+ flag is accepted and ignored there.
705
+ trajectory_out
706
+ Write the trajectory JSON here (plus ``edges.jsonl`` beside it)
707
+ instead of stdout.
708
+
709
+ Exit codes: 0 ok, 2 empty session, 64 invalid input or unknown agent,
710
+ 65 validation, 70 conversion.
711
+ """
712
+ from atif_converter.application.convert_and_audit import convert_and_audit
713
+ from atif_converter.application.convert_codex import convert_codex_and_audit
714
+ from atif_converter.domain.agents import AgentSource
715
+ from atif_converter.domain.errors import (
716
+ DomainError,
717
+ EmptySessionError,
718
+ InvalidSessionInput,
719
+ )
720
+
721
+ resolved_agent = _resolve_agent(agent or _env_agent(), AgentSource)
722
+ try:
723
+ if resolved_agent is AgentSource.CODEX:
724
+ result, report = convert_codex_and_audit(session_jsonl)
725
+ else:
726
+ result, report = convert_and_audit(session_jsonl, include_subagents=include_subagents)
727
+ except InvalidSessionInput as exc:
728
+ print(f"error: {exc}", file=sys.stderr)
729
+ raise SystemExit(EXIT_CODES["invalid_input"]) from exc
730
+ except EmptySessionError as exc:
731
+ print(f"error: {exc}", file=sys.stderr)
732
+ raise SystemExit(EXIT_CODES["empty_session"]) from exc
733
+ except DomainError as exc:
734
+ print(f"error: {exc}", file=sys.stderr)
735
+ raise SystemExit(EXIT_CODES["runtime_error"]) from exc
736
+
737
+ if trajectory_out is not None:
738
+ trajectory_out.parent.mkdir(parents=True, exist_ok=True)
739
+ trajectory_out.write_text(json.dumps(result.trajectory, indent=2, default=str) + "\n")
740
+ edges_out = trajectory_out.parent / "edges.jsonl"
741
+ edges_out.write_text("".join(f"{line}\n" for line in result.edges_lines))
742
+ print(f"trajectory: {trajectory_out}")
743
+ print(f"edges: {edges_out} ({len(result.edges_lines)} lines)")
744
+ else:
745
+ print(json.dumps(result.trajectory, indent=2, default=str))
746
+
747
+ print(
748
+ json.dumps(
749
+ {
750
+ "loss_report": report.to_json(),
751
+ "validation_errors": list(result.validation_errors),
752
+ },
753
+ indent=2,
754
+ )
755
+ )
756
+ if result.validation_errors:
757
+ raise SystemExit(EXIT_CODES["validation_error"])
758
+
759
+
760
+ # ---------------------------------------------------------------------------
761
+ # materialize
762
+ # ---------------------------------------------------------------------------
763
+
764
+
765
+ def _print_report(report: MaterializationReport, fmt: OutputFormat) -> None:
766
+ """Emit a MaterializationReport: human lines on TTY, JSON otherwise.
767
+
768
+ An unreadable session lands in no other counter — not materialized, not
769
+ up-to-date, not skipped-live, not failed, and never marked — so a pass
770
+ that could only stat nothing prints all zeroes and reads as an idle,
771
+ complete corpus unless ``unreadable`` is shown alongside them. A rejected
772
+ session (a transcript whose name fails the session id boundary) is the
773
+ same kind of silence and gets the same treatment; its name is printed
774
+ ``repr``-quoted because the whole point is that it carries odd characters.
775
+
776
+ ``removed`` counts sessions whose source vanished THIS pass. Nothing is
777
+ deleted any more: their artifacts are kept and marked
778
+ ``source_present: false``, and ``retained`` counts every session kept
779
+ without a source. ``empty`` counts transcripts with nothing to convert,
780
+ which are recorded rather than failed; ``from_archive`` counts
781
+ source-removed sessions re-converted from their raw source archive.
782
+ """
783
+ payload = {
784
+ "materialized": report.materialized_count,
785
+ "up_to_date": report.up_to_date_count,
786
+ "skipped_live": report.skipped_live_count,
787
+ "failed": report.failed_count,
788
+ "failures": [{"session_id": f.session_id, "error": f.error} for f in report.failures],
789
+ "empty": report.empty_count,
790
+ "empty_session_ids": list(report.empty_session_ids),
791
+ "removed": report.sessions_removed,
792
+ "removed_session_ids": list(report.removed_session_ids),
793
+ "retained": report.retained_count,
794
+ "from_archive": report.archive_count,
795
+ "from_archive_session_ids": list(report.archive_session_ids),
796
+ "unreadable": report.unreadable_count,
797
+ "unreadable_session_ids": list(report.unreadable_session_ids),
798
+ "rejected": report.rejected_count,
799
+ "rejected_session_ids": list(report.rejected_session_ids),
800
+ "total_seconds": round(report.total_seconds, 3),
801
+ "convert_seconds": round(report.convert_seconds, 3),
802
+ "workers": report.workers,
803
+ "artifact_seconds": round(report.artifact_seconds, 3),
804
+ "lake_synced": report.sink_synced_count,
805
+ "lake_pending": len(report.sink_pending_session_ids),
806
+ "lake_pending_session_ids": list(report.sink_pending_session_ids),
807
+ "lake_error": report.sink_error,
808
+ "lake_seconds": round(report.sink_seconds, 3),
809
+ }
810
+ if resolve_format(fmt) is OutputFormat.TABLE:
811
+ print(
812
+ f"materialized: {report.materialized_count} "
813
+ f"up-to-date: {report.up_to_date_count} "
814
+ f"skipped-live: {report.skipped_live_count} "
815
+ f"failed: {report.failed_count} "
816
+ f"empty: {report.empty_count} "
817
+ f"removed: {report.sessions_removed} "
818
+ f"retained: {report.retained_count} "
819
+ f"from-archive: {report.archive_count} "
820
+ f"unreadable: {report.unreadable_count} "
821
+ f"rejected: {report.rejected_count}"
822
+ )
823
+ print(
824
+ f"total: {report.total_seconds:.2f}s "
825
+ f"(convert: {report.convert_seconds:.2f}s summed over {report.workers} worker(s), "
826
+ f"lake: {report.sink_synced_count} synced in {report.sink_seconds:.2f}s)"
827
+ )
828
+ if report.sink_error is not None:
829
+ print(
830
+ f" LAKE PENDING {len(report.sink_pending_session_ids)} session(s): "
831
+ f"{report.sink_error}",
832
+ file=sys.stderr,
833
+ )
834
+ for failure in report.failures:
835
+ print(f" FAILED {failure.session_id}: {failure.error}", file=sys.stderr)
836
+ for session_id in report.unreadable_session_ids:
837
+ print(f" UNREADABLE {session_id}", file=sys.stderr)
838
+ for session_name in report.rejected_session_ids:
839
+ print(f" REJECTED {session_name!r}", file=sys.stderr)
840
+ else:
841
+ emit_json(payload, fmt)
842
+
843
+
844
+ def _materialize_worker_setup() -> None:
845
+ """Give a materialize pool worker the stderr sink :func:`main` gives the parent.
846
+
847
+ A spawned worker starts with loguru's default DEBUG sink, so without this
848
+ every debug line the converter or the atomic writer emits in a worker
849
+ would land on the CLI's stderr in a format the parent never uses. Reads
850
+ :data:`LOG_LEVEL_ENV` the same way, because the environment is what a
851
+ spawned child inherits. Module-level so the executor can pickle it.
852
+ """
853
+ from loguru import logger
854
+
855
+ logger.remove()
856
+ try:
857
+ logger.add(sys.stderr, level=_stderr_log_level())
858
+ except ValueError:
859
+ logger.add(sys.stderr, level=DEFAULT_LOG_LEVEL)
860
+
861
+
862
+ @app.command
863
+ def materialize(
864
+ *,
865
+ agent: str | None = None,
866
+ force: bool = False,
867
+ quiesce_seconds: int | None = None,
868
+ source_root: Path | None = None,
869
+ corpus_root: Path | None = None,
870
+ sessions: str | None = None,
871
+ workers: int | None = None,
872
+ lake: Annotated[bool, cyclopts.Parameter(negative="--no-lake")] = True,
873
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
874
+ ) -> None:
875
+ """Sync the materialized corpus with the raw transcript corpus.
876
+
877
+ Runs one scan → plan → convert → write pass through atif-corpus's
878
+ materialize use case, with atif-converter's real converter adapted
879
+ behind the ConverterPort (see :mod:`atif_cli.converter_adapter`). Each
880
+ session is written as ``trajectory.json.zst``, ``edges.jsonl.zst`` and
881
+ ``session_events.jsonl.zst`` (zstd), ``loss_report.json``, ``meta.json``
882
+ and its raw source archive; the lake holds the queryable rows. The CLI
883
+ owns the wall clock and the version pins stamped into ``meta.json``.
884
+
885
+ Parameters
886
+ ----------
887
+ agent
888
+ Which agent's transcripts to discover and convert: ``claude-code``
889
+ (default) or ``codex``. It selects the discovery layout, the converter,
890
+ and — unless a root is given explicitly — both default roots, so
891
+ ``--agent codex`` alone materializes ``<CODEX_HOME>/sessions`` into its
892
+ own corpus without touching the Claude Code one.
893
+ force
894
+ Re-materialize every quiescent session regardless of the watermark.
895
+ quiesce_seconds
896
+ Source-silence threshold; default from settings (contract: 300).
897
+ source_root
898
+ Override the raw transcript root (default: env ``ATIF_SQL_SOURCE_ROOT``,
899
+ or ``<CLAUDE_CONFIG_DIR>/projects`` / ``<CODEX_HOME>/sessions`` per
900
+ ``--agent``).
901
+ corpus_root
902
+ Override the materialized corpus root (default: env or
903
+ ``~/.atif-sql/corpus/<slug>``).
904
+ sessions
905
+ Comma-separated session-id filter — only these sessions are planned
906
+ this pass (contract-compatible extension; other sessions' watermark
907
+ entries are left untouched).
908
+ workers
909
+ Processes for the convert+write stage. Default from
910
+ ``ATIF_SQL_MATERIALIZE_WORKERS``, else ``min(8, cpu_count)``. ``1``
911
+ is the single-process reference path; above that a process pool
912
+ converts sessions in parallel and writes byte-identical artifacts.
913
+ Below ``1`` exits 64.
914
+ lake
915
+ After the sessions are swapped into place, replace their rows in the
916
+ DuckLake at ``ATIF_SQL_LAKE_ROOT`` (default ``~/.atif-sql/lake``), one
917
+ transaction per batch, in this process. Does nothing until ``atif-sql
918
+ lake rebuild`` has created the lake. A failed lake write leaves the
919
+ sessions in ``<corpus>/sink_pending.json`` and the next pass retries
920
+ them; the report's ``lake_pending`` counts them. ``--no-lake`` skips the
921
+ lake and leaves that file alone.
922
+ fmt
923
+ Report format; ``auto`` = human lines on TTY, JSON on a pipe.
924
+ """
925
+ from atif_cli.converter_adapter import RealConverter
926
+ from atif_converter.domain.atif import UPSTREAM_VERSION as ATIF_MODELS_VERSION
927
+ from atif_converter.domain.schema_version import CONVERTER_SCHEMA_VERSION
928
+ from atif_corpus.application.materialize import (
929
+ CorpusAgentMismatchError,
930
+ SuspiciousEmptyScanError,
931
+ materialize as materialize_use_case,
932
+ )
933
+ from atif_corpus.domain.source_layout import layout_for
934
+ from atif_duck.infrastructure.lake import DuckLakeSessionSink, LakeLayout
935
+ from atif_duck.infrastructure.lake_settings import LakeSettings
936
+
937
+ settings = _corpus_settings(source_root, corpus_root, agent)
938
+ lake_settings = LakeSettings()
939
+ source_layout = layout_for(settings.agent)
940
+ session_filter = (
941
+ [s for s in (part.strip() for part in sessions.split(",")) if s]
942
+ if sessions is not None
943
+ else None
944
+ )
945
+ worker_count = workers if workers is not None else settings.materialize_workers
946
+ if worker_count < 1:
947
+ emit_error(
948
+ ClassifiedError(
949
+ kind="invalid_input",
950
+ exit_code=EXIT_CODES["invalid_input"],
951
+ message=f"--workers must be >= 1, got {worker_count}",
952
+ hint="pass --workers 1 for the single-process path",
953
+ ),
954
+ fmt,
955
+ )
956
+ raise SystemExit(EXIT_CODES["invalid_input"])
957
+ try:
958
+ report = materialize_use_case(
959
+ source_root=settings.source_root,
960
+ corpus_root=settings.corpus_root,
961
+ converter=RealConverter(agent=settings.agent),
962
+ source_layout=source_layout,
963
+ materialized_at=_now_iso(),
964
+ # The harbor release the vendored ATIF models match; harbor itself
965
+ # isn't installed with atif-sql, so this is provenance, not a lookup.
966
+ harbor_version=ATIF_MODELS_VERSION,
967
+ # The release that did the converting, looked up under the ONE
968
+ # published distribution. `atif-converter` isn't a distribution in the
969
+ # bundled wheel, so looking it up answered "unknown" for every
970
+ # installed run. Provenance only: staleness keys on the schema below.
971
+ converter_version=_version_of("atif-sql"),
972
+ # The converter's own output version. A session recording another
973
+ # one re-converts, from its source archive if the source is gone.
974
+ converter_schema=CONVERTER_SCHEMA_VERSION,
975
+ # No per-session parquet any more, so no columnar schema to
976
+ # expect: a session written with them (before `corpus slim`) is
977
+ # as current as one written without, and the lake loads either.
978
+ quiesce_seconds=(
979
+ quiesce_seconds if quiesce_seconds is not None else settings.quiesce_seconds
980
+ ),
981
+ force=force,
982
+ session_ids=session_filter,
983
+ workers=worker_count,
984
+ worker_setup=_materialize_worker_setup,
985
+ session_sink=(
986
+ DuckLakeSessionSink(
987
+ LakeLayout(lake_settings.lake_root),
988
+ lock_timeout_seconds=lake_settings.lake_lock_timeout_seconds,
989
+ load_batch_size=lake_settings.lake_load_batch_size,
990
+ memory_limit_bytes=_query_memory_limit_bytes(),
991
+ stage_workers=lake_settings.lake_stage_workers,
992
+ )
993
+ if lake
994
+ else None
995
+ ),
996
+ sink_batch_size=lake_settings.lake_sync_batch_size,
997
+ )
998
+ except CorpusAgentMismatchError as exc:
999
+ # One corpus holds one agent. Nothing was removed and nothing written —
1000
+ # the same posture as a suspicious scan, and the same exit code, because
1001
+ # an unattended lane must branch on it rather than retry.
1002
+ emit_error(
1003
+ ClassifiedError(
1004
+ kind="terminal_state",
1005
+ exit_code=EXIT_CODES["terminal_state"],
1006
+ message=str(exc),
1007
+ hint="point --corpus-root / ATIF_SQL_CORPUS_ROOT at this agent's own corpus",
1008
+ ),
1009
+ fmt,
1010
+ )
1011
+ raise SystemExit(EXIT_CODES["terminal_state"]) from exc
1012
+ except SuspiciousEmptyScanError as exc:
1013
+ # The corpus's loudest data-loss tripwire, so it gets a code an
1014
+ # unattended lane can branch on. Uncaught it would exit 1, which a
1015
+ # caller cannot tell apart from a crash — and the two want opposite
1016
+ # responses: a crash is worth retrying, a wrong source_root is not.
1017
+ emit_error(
1018
+ ClassifiedError(
1019
+ kind="suspicious_scan",
1020
+ exit_code=EXIT_CODES["suspicious_scan"],
1021
+ message=str(exc),
1022
+ hint="check --source-root / ATIF_SQL_SOURCE_ROOT; nothing was removed",
1023
+ ),
1024
+ fmt,
1025
+ )
1026
+ raise SystemExit(EXIT_CODES["suspicious_scan"]) from exc
1027
+ _print_report(report, fmt)
1028
+
1029
+
1030
+ # ---------------------------------------------------------------------------
1031
+ # status
1032
+ # ---------------------------------------------------------------------------
1033
+
1034
+
1035
+ def _embed_settings(fmt: OutputFormat) -> EmbedSettings:
1036
+ """The embedding settings from the environment, or exit 64 naming the bad one.
1037
+
1038
+ A width the selected provider cannot emit (``ATIF_SQL_OUTPUT_DIMENSION=1024``
1039
+ with ``ATIF_SQL_EMBED_PROVIDER=gemma``) is malformed input, not a traceback.
1040
+ """
1041
+ from pydantic import ValidationError
1042
+
1043
+ from atif_embed.infrastructure.settings import EmbedSettings
1044
+
1045
+ try:
1046
+ return EmbedSettings()
1047
+ except ValidationError as exc:
1048
+ reasons = "; ".join(str(e["msg"]).removeprefix("Value error, ") for e in exc.errors())
1049
+ err = ClassifiedError(
1050
+ kind="invalid_input",
1051
+ exit_code=EXIT_CODES["invalid_input"],
1052
+ message=f"invalid embedding settings: {reasons}",
1053
+ hint="check the ATIF_SQL_EMBED_PROVIDER and ATIF_SQL_OUTPUT_DIMENSION variables",
1054
+ )
1055
+ emit_error(err, fmt)
1056
+ raise SystemExit(err.exit_code) from exc
1057
+
1058
+
1059
+ def _vector_surface(corpus_root: Path, fmt: OutputFormat) -> dict[str, Any]:
1060
+ """How ``query`` and ``search`` will see the embeddings store, without installing anything.
1061
+
1062
+ The extension check reads ``duckdb_extensions()`` on a throwaway
1063
+ connection (the local extension directory, a few milliseconds, no
1064
+ network); the store check is the directory ``embed`` writes, resolved the
1065
+ way ``query`` and ``search`` resolve it (``ATIF_SQL_LANCE_URI`` wins, then
1066
+ the selected provider's directory under the corpus root).
1067
+ """
1068
+ import duckdb
1069
+
1070
+ from atif_duck.infrastructure.registry import lance_extension_installed
1071
+
1072
+ embed_settings = _embed_settings(fmt)
1073
+ store = embed_settings.resolve_lance_uri(corpus_root)
1074
+ con = duckdb.connect()
1075
+ try:
1076
+ installed = lance_extension_installed(con)
1077
+ finally:
1078
+ con.close()
1079
+ store_present = store.is_dir()
1080
+ if store_present and installed:
1081
+ state, note = "ready", f"store at {store}"
1082
+ elif store_present:
1083
+ state = "extension_missing"
1084
+ note = (
1085
+ "store present but the lance DuckDB extension is not installed; "
1086
+ "run `atif-sql embed --install-extension`"
1087
+ )
1088
+ else:
1089
+ state, note = "no_store", "no embeddings store; run `atif-sql embed --all --no-dry-run`"
1090
+ return {
1091
+ "embed_provider": embed_settings.embed_provider,
1092
+ "embeddings_store": str(store),
1093
+ "lance_extension_installed": installed,
1094
+ "embeddings_store_present": store_present,
1095
+ "vector_search": state,
1096
+ "note": note,
1097
+ }
1098
+
1099
+
1100
+ def _lake_surface(corpus_root: Path) -> dict[str, Any]:
1101
+ """Whether ``query`` will read this corpus from the lake, and the lake's vitals.
1102
+
1103
+ ``state`` is ``ready`` (query reads the lake), ``absent``, ``stale``,
1104
+ ``extension_missing``, or ``corpus_missing`` (the lake does not hold this
1105
+ corpus yet; the next materialize loads it). ``pending`` counts this
1106
+ corpus's sessions waiting for a lake write.
1107
+ """
1108
+ from atif_corpus.application.materialize import read_sink_pending
1109
+ from atif_corpus.domain.layout import CorpusLayout
1110
+ from atif_duck.infrastructure.lake import LakeLayout, lake_status
1111
+ from atif_duck.infrastructure.lake_settings import LakeSettings
1112
+
1113
+ state = lake_status(LakeLayout(LakeSettings().lake_root)).as_dict()
1114
+ registered = {c["corpus"]: c for c in state["corpora"]}
1115
+ mine = registered.get(corpus_root.name)
1116
+ if not state["present"]:
1117
+ name, note = "absent", "no lake; run `atif-sql lake rebuild`"
1118
+ elif not state["extension_installed"]:
1119
+ name, note = "extension_missing", "the ducklake extension is not installed"
1120
+ elif not state["schema_current"]:
1121
+ name, note = "stale", f"schema {', '.join(state['stale']) or state['error']}"
1122
+ elif mine is None or Path(mine["corpus_root"]).resolve() != corpus_root.resolve():
1123
+ name, note = "corpus_missing", "the lake does not hold this corpus"
1124
+ else:
1125
+ name, note = "ready", f"{mine['sessions']} sessions, last write {state['last_write']}"
1126
+ return {
1127
+ "state": name,
1128
+ "note": note,
1129
+ "root": state["root"],
1130
+ "schema": state["schema"],
1131
+ "snapshots": state["snapshots"],
1132
+ "data_files": state["data_files"],
1133
+ "last_write": state["last_write"],
1134
+ "sessions": mine["sessions"] if mine else 0,
1135
+ "pending": len(read_sink_pending(CorpusLayout(corpus_root=corpus_root).sink_pending_path)),
1136
+ }
1137
+
1138
+
1139
+ def _dir_bytes(root: Path) -> int:
1140
+ """Total bytes of every file under ``root`` (0 if absent)."""
1141
+ if not root.is_dir():
1142
+ return 0
1143
+ return sum(p.stat().st_size for p in root.rglob("*") if p.is_file())
1144
+
1145
+
1146
+ @app.command
1147
+ def status(
1148
+ *,
1149
+ agent: str | None = None,
1150
+ source_root: Path | None = None,
1151
+ corpus_root: Path | None = None,
1152
+ quiesce_seconds: int | None = None,
1153
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
1154
+ ) -> None:
1155
+ """Report corpus freshness: watermark age, counts, bytes, staleness.
1156
+
1157
+ Read-only and fast: scans source mtimes and replays the same planning
1158
+ decision ``materialize`` would make (quiescence, watermark, and the
1159
+ recorded converter schema), without converting
1160
+ anything. The staleness summary is therefore exactly "what would a
1161
+ materialize pass do right now", and ``stale`` includes sessions a
1162
+ converter upgrade made stale (``generation_stale`` counts those alone).
1163
+
1164
+ ``retained`` counts sessions the corpus keeps whose source transcript is
1165
+ gone (``meta.source_present`` false, or about to be marked so), and
1166
+ ``from_archive`` how many of them the next pass would re-convert from
1167
+ their raw source archive. ``empty`` counts transcripts recorded as having
1168
+ nothing to convert.
1169
+
1170
+ ``--agent`` selects which corpus is reported, resolving the same way
1171
+ ``materialize`` resolves it, so the two commands always describe the same
1172
+ pair of roots. The JSON payload carries the resolved agent, which is what
1173
+ lets an unattended lane confirm it ticked the corpus it meant to.
1174
+
1175
+ ``layout`` says how many sessions are still stored the old way (plain
1176
+ ``trajectory.json`` / ``edges.jsonl`` / ``session_events.jsonl``, or the
1177
+ per-session parquet files materialize no longer writes) and how many
1178
+ bytes that is; ``atif-sql corpus slim`` converts them. ``query path``
1179
+ says how ``query`` reads this corpus when it can't use the lake:
1180
+ ``columnar`` when every complete session still carries its parquet,
1181
+ ``json`` when none does (every session materialized or slimmed since the
1182
+ parquet was dropped), ``mixed`` in between. It applies the same
1183
+ per-session predicate the registry does.
1184
+
1185
+ ``vector search`` says whether ``semantic_search`` and ``search`` can
1186
+ reach an embeddings store: ``ready`` (store present, lance extension
1187
+ installed), ``no_store`` (run ``atif-sql embed``), or
1188
+ ``extension_missing`` (a store exists but the DuckDB lance extension is
1189
+ not installed; ``query`` binds ``message_embeddings`` empty rather than
1190
+ downloading it, so run ``atif-sql embed --install-extension``). Read
1191
+ from DuckDB's local extension listing; never installs anything.
1192
+ """
1193
+ import time
1194
+
1195
+ from atif_cli.corpus import storage_layout
1196
+ from atif_converter.domain.schema_version import CONVERTER_SCHEMA_VERSION
1197
+ from atif_corpus.application.materialize import preview_pass, read_watermark
1198
+ from atif_corpus.domain.layout import CorpusLayout
1199
+ from atif_corpus.domain.source_layout import layout_for
1200
+ from atif_duck.domain.columnar import COLUMNAR_SCHEMA_VERSION
1201
+ from atif_duck.infrastructure.columnar import columnar_coverage
1202
+
1203
+ settings = _corpus_settings(source_root, corpus_root, agent)
1204
+ layout = CorpusLayout(corpus_root=settings.corpus_root)
1205
+ quiesce = quiesce_seconds if quiesce_seconds is not None else settings.quiesce_seconds
1206
+
1207
+ watermark = read_watermark(layout.watermark_path)
1208
+ # Planned as a default `materialize` would plan it.
1209
+ preview = preview_pass(
1210
+ source_root=settings.source_root,
1211
+ corpus_root=settings.corpus_root,
1212
+ converter_version=_version_of("atif-sql"),
1213
+ converter_schema=CONVERTER_SCHEMA_VERSION,
1214
+ quiesce_seconds=quiesce,
1215
+ source_layout=layout_for(settings.agent),
1216
+ )
1217
+ plan = preview.plan
1218
+ source_sessions = (*plan.to_materialize, *plan.up_to_date, *plan.skipped_live)
1219
+ # Read after the scan, like the pass's own "now".
1220
+ now_ns = time.time_ns()
1221
+ watermark_age_seconds: float | None = None
1222
+ if layout.watermark_path.exists():
1223
+ watermark_age_seconds = round((now_ns - layout.watermark_path.stat().st_mtime_ns) / 1e9, 1)
1224
+ materialized_dirs = (
1225
+ sorted(p.name for p in layout.sessions_dir.iterdir() if p.is_dir())
1226
+ if layout.sessions_dir.is_dir()
1227
+ else []
1228
+ )
1229
+
1230
+ corpus_bytes = _dir_bytes(settings.corpus_root)
1231
+ staleness = {
1232
+ "stale": len(plan.to_materialize),
1233
+ "up_to_date": len(plan.up_to_date),
1234
+ "live": len(plan.skipped_live),
1235
+ "generation_stale": len(preview.generation_stale_session_ids),
1236
+ "from_archive": len(preview.archive_session_ids),
1237
+ }
1238
+ coverage = columnar_coverage(settings.corpus_root)
1239
+ storage = storage_layout(settings.corpus_root)
1240
+ vector = _vector_surface(settings.corpus_root, fmt)
1241
+ lake_block = _lake_surface(settings.corpus_root)
1242
+ if resolve_format(fmt) is OutputFormat.TABLE:
1243
+ print(f"agent: {settings.agent.value}")
1244
+ print(f"source root: {settings.source_root}")
1245
+ print(f"corpus root: {settings.corpus_root}")
1246
+ age = "never materialized" if watermark_age_seconds is None else f"{watermark_age_seconds}s"
1247
+ print(f"watermark: {age} old ({len(watermark)} entries)")
1248
+ print(
1249
+ f"sessions: {len(source_sessions)} in source, {len(materialized_dirs)} materialized"
1250
+ )
1251
+ print(f"corpus bytes: {corpus_bytes:,}")
1252
+ print(
1253
+ f"staleness: {staleness['stale']} stale "
1254
+ f"({staleness['generation_stale']} from a converter or schema change), "
1255
+ f"{staleness['up_to_date']} up-to-date, {staleness['live']} live"
1256
+ )
1257
+ print(
1258
+ f"retained: {len(preview.retained_session_ids)} kept without a source "
1259
+ f"({staleness['from_archive']} to re-convert from archive), "
1260
+ f"{len(preview.empty_session_ids)} empty"
1261
+ )
1262
+ print(f"layout: {storage.summary}")
1263
+ print(
1264
+ f"query path: {coverage.query_path} without the lake "
1265
+ f"({coverage.columnar_sessions} of {coverage.total_sessions} complete sessions "
1266
+ "still carry per-session parquet)"
1267
+ )
1268
+ print(f"vector search: {vector['vector_search']} ({vector['note']})")
1269
+ print(f"lake: {lake_block['state']} ({lake_block['note']})")
1270
+ else:
1271
+ emit_json(
1272
+ {
1273
+ "agent": settings.agent.value,
1274
+ "source_root": str(settings.source_root),
1275
+ "corpus_root": str(settings.corpus_root),
1276
+ "watermark_age_seconds": watermark_age_seconds,
1277
+ "watermark_entries": len(watermark),
1278
+ "source_sessions": len(source_sessions),
1279
+ "materialized_sessions": len(materialized_dirs),
1280
+ "corpus_bytes": corpus_bytes,
1281
+ "staleness": staleness,
1282
+ "retained_sessions": len(preview.retained_session_ids),
1283
+ "empty_sessions": len(preview.empty_session_ids),
1284
+ "converter_schema": CONVERTER_SCHEMA_VERSION,
1285
+ "columnar_schema": COLUMNAR_SCHEMA_VERSION,
1286
+ "columnar_sessions": coverage.columnar_sessions,
1287
+ "json_sessions": coverage.json_sessions,
1288
+ "query_path": coverage.query_path,
1289
+ "layout": storage.as_dict(),
1290
+ "lance_extension_installed": vector["lance_extension_installed"],
1291
+ "embeddings_store_present": vector["embeddings_store_present"],
1292
+ "vector_search": vector["vector_search"],
1293
+ "lake": lake_block,
1294
+ },
1295
+ fmt,
1296
+ )
1297
+
1298
+
1299
+ # ---------------------------------------------------------------------------
1300
+ # query
1301
+ # ---------------------------------------------------------------------------
1302
+
1303
+
1304
+ @app.command
1305
+ def query(
1306
+ sql: str | None = None,
1307
+ /,
1308
+ *,
1309
+ examples_flag: Annotated[bool, cyclopts.Parameter(name="--examples")] = False,
1310
+ category: str | None = None,
1311
+ requires: str | None = None,
1312
+ agent: str | None = None,
1313
+ corpus_root: Path | None = None,
1314
+ all_corpora: bool = False,
1315
+ lake: Annotated[bool, cyclopts.Parameter(negative="--no-lake")] = True,
1316
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
1317
+ ) -> None:
1318
+ """Run one SQL statement against the atif-duck catalog and emit results.
1319
+
1320
+ Opens an in-memory DuckDB connection, registers the corpus views +
1321
+ macros via ``atif_duck.register``, executes, and emits: TTY -> plain
1322
+ table, pipe -> JSON array of row objects.
1323
+
1324
+ ``--examples`` short-circuits to the ``examples`` listing (same output,
1325
+ same ``--category`` / ``--requires`` filters) without opening DuckDB —
1326
+ the natural discovery path when an agent is already composing a query.
1327
+
1328
+ Where the rows come from
1329
+ ------------------------
1330
+ When the DuckLake at ``ATIF_SQL_LAKE_ROOT`` exists, its schema is current
1331
+ and it holds the requested corpus, every view reads the lake, scoped to
1332
+ that corpus (``--agent`` / ``--corpus-root`` pick it, as before), and
1333
+ nothing is loaded before the statement runs. ``--all-corpora`` scopes the
1334
+ views to every corpus the lake holds instead; ``sessions.corpus`` tells
1335
+ them apart (the analytics views stay the selected corpus's, and
1336
+ ``message_embeddings`` spans every corpus's store unless
1337
+ ``ATIF_SQL_LANCE_URI`` pins one). Otherwise, or with ``--no-lake``, ``query`` reads the
1338
+ per-session artifacts as it always has, after a one-line warning saying
1339
+ why. ``--all-corpora`` has no per-session fallback and exits 78
1340
+ (``lake_unavailable``) without a usable lake.
1341
+
1342
+ Sandbox
1343
+ -------
1344
+ Before anything is registered the connection is sized to the host
1345
+ (:func:`_query_resources`: a memory cap derived from available RAM and a
1346
+ thread count derived from that cap, overridable with
1347
+ ``ATIF_SQL_QUERY_MEMORY_LIMIT`` and ``ATIF_SQL_QUERY_THREADS``), its spill
1348
+ directory is a private ``mkdtemp`` (mode 0700) outside the corpus that is
1349
+ removed when the process exits, and extension auto-install and auto-load
1350
+ are off. The lance extension is loaded only when it is already installed;
1351
+ when a store exists and it is not, the vector surface binds empty with a
1352
+ warning and ``atif-sql status`` says so. Registration therefore reaches
1353
+ neither the network nor the corpus for writing, and it fits an 8 GiB
1354
+ guest, where it used to exit 70 out of memory before the cap applied.
1355
+
1356
+ The statement then runs against the hardened connection (see
1357
+ :func:`_harden_query_connection`). Reads reach the registered views and
1358
+ nothing else. ``read_text`` outside the corpus, ``ATTACH`` of unrelated
1359
+ databases and extension installs fail with exit 70 rather than reading
1360
+ credentials or reaching the network, and nothing under the corpus root
1361
+ is writable. Two layers keep caller SQL from writing the corpus:
1362
+
1363
+ * DuckDB's file grants are read-write, and ``COPY ... TO <granted
1364
+ parquet> (USE_TMP_FILE false)`` used to overwrite one, so every
1365
+ statement kind that names a file is refused BEFORE execution, for any
1366
+ uid, by DuckDB's own parser: COPY, EXPORT, ATTACH, DETACH, INSTALL,
1367
+ LOAD, and PREPARE/EXECUTE (which could defer one). Exit 70, kind
1368
+ ``sandbox_refused``.
1369
+ * A 0444 file mode does not bind root, so ``query`` refuses to run as
1370
+ uid 0 (exit 77) unless ``ATIF_SQL_ALLOW_ROOT=1`` is set, which logs a
1371
+ warning.
1372
+
1373
+ The embedding-store guard binds here exactly as it does for ``search``,
1374
+ so a store written by a different provider refuses to bind (exit 65)
1375
+ instead of scoring garbage.
1376
+
1377
+ One residual hole remains, over recomputable derived data rather than
1378
+ the transcript artifacts: ``DELETE``/``INSERT`` against
1379
+ ``lance_store.main.embeddings`` reach the ATTACHed store, which the
1380
+ filesystem allowlist does not cover. Re-run
1381
+ ``atif-sql embed --all --no-dry-run`` to rebuild.
1382
+
1383
+ What caller SQL can still see, and why that is accepted:
1384
+ ``duckdb_settings()`` and ``current_setting(...)`` return the sandbox's
1385
+ own configuration, including the corpus root, the spill directory, the
1386
+ memory cap and every granted parquet path, which names every session id.
1387
+ DuckDB cannot hide a setting from SQL, and the grants have to be per-file
1388
+ for the lazy reads to work at all, so the inventory is a property of the
1389
+ design. The caller is the local user, who can list the corpus directory
1390
+ and ``SELECT session_id FROM sessions`` anyway; ``query`` is not a
1391
+ privilege boundary (SECURITY.md) and must not be exposed to a caller who
1392
+ may not know the session ids.
1393
+
1394
+ ``lock_configuration`` freezes the connection's settings, so ``SET`` is
1395
+ refused for everything except ``TimeZone``, which stays open for
1396
+ timestamp rendering. ``SET threads`` and ``SET errors_as_json`` fail with
1397
+ exit 70 — deliberately: an unfreezable option is one injected SQL can
1398
+ turn back off.
1399
+
1400
+ Exit codes
1401
+ ----------
1402
+ * 64 parse_error malformed SQL (or no SQL and no --examples)
1403
+ * 64 invalid_input a malformed ``ATIF_SQL_QUERY_MEMORY_LIMIT`` or
1404
+ ``ATIF_SQL_QUERY_THREADS``
1405
+ * 65 catalog_error unknown view/macro/column (try ``atif-sql schema``)
1406
+ * 65 embedding_mismatch the Lance store was written by another provider
1407
+ * 70 sandbox_refused a statement kind the sandbox never runs (COPY,
1408
+ EXPORT, ATTACH, INSTALL, LOAD, PREPARE, ...)
1409
+ * 70 runtime_error everything else (an unmaterialized corpus, or SQL
1410
+ the sandbox refused)
1411
+ * 77 root_refused running as uid 0 without ``ATIF_SQL_ALLOW_ROOT=1``
1412
+
1413
+ Examples
1414
+ --------
1415
+ ::
1416
+
1417
+ atif-sql query 'SELECT count(*) FROM sessions'
1418
+ atif-sql query 'SELECT * FROM tool_rank(14)' --format json
1419
+ atif-sql query --examples --requires core
1420
+ atif-sql query --agent codex 'SELECT agent, count(*) FROM sessions GROUP BY 1'
1421
+
1422
+ ``--agent codex`` selects the Codex corpus root the way ``materialize``
1423
+ selects it, so the corpus a Codex materialize wrote is queryable without
1424
+ spelling its path. One corpus holds one agent; ``--corpus-root`` still wins.
1425
+ """
1426
+ if examples_flag:
1427
+ examples(category=category, requires=requires, fmt=fmt)
1428
+ return
1429
+ if sql is None:
1430
+ emit_error(
1431
+ ClassifiedError(
1432
+ kind="parse_error",
1433
+ exit_code=EXIT_CODES["parse_error"],
1434
+ message="no SQL given",
1435
+ hint="pass a statement (atif-sql query 'SELECT ...') or list "
1436
+ "tested templates with: atif-sql examples",
1437
+ ),
1438
+ fmt,
1439
+ )
1440
+ raise SystemExit(EXIT_CODES["parse_error"])
1441
+
1442
+ _refuse_root("query", fmt)
1443
+ try:
1444
+ resources = _query_resources()
1445
+ except ValueError as exc:
1446
+ err = ClassifiedError(
1447
+ kind="invalid_input",
1448
+ exit_code=EXIT_CODES["invalid_input"],
1449
+ message=str(exc),
1450
+ hint=f"{QUERY_MEMORY_LIMIT_ENV} takes a size such as 6GB or 512MiB; "
1451
+ f"{QUERY_THREADS_ENV} takes a positive integer",
1452
+ )
1453
+ emit_error(err, fmt)
1454
+ raise SystemExit(err.exit_code) from exc
1455
+
1456
+ import duckdb
1457
+
1458
+ from atif_cli.duck_errors import REGISTRATION_ERRORS, classify_registration_error
1459
+ from atif_duck.infrastructure.registry import analytics_roots, register
1460
+
1461
+ settings = _corpus_settings(None, corpus_root, agent)
1462
+ embed_settings = _embed_settings(fmt)
1463
+ expected_model, expected_dim = embed_settings.expected_embedding_identity()
1464
+ lance_uri = embed_settings.resolve_lance_uri(settings.corpus_root)
1465
+ if all_corpora and not lake:
1466
+ err = ClassifiedError(
1467
+ kind="invalid_input",
1468
+ exit_code=EXIT_CODES["invalid_input"],
1469
+ message="--all-corpora reads the lake, and --no-lake turns it off",
1470
+ hint="drop one of the two flags",
1471
+ )
1472
+ emit_error(err, fmt)
1473
+ raise SystemExit(err.exit_code)
1474
+
1475
+ with _private_spill_dir() as spill_dir:
1476
+ con = duckdb.connect()
1477
+ lake_reader = None
1478
+ try:
1479
+ try:
1480
+ _configure_query_resources(con, resources, spill_dir)
1481
+ lake_reader = (
1482
+ _attach_query_lake(con, settings.corpus_root, all_corpora=all_corpora, fmt=fmt)
1483
+ if lake
1484
+ else None
1485
+ )
1486
+ sources = register(
1487
+ con,
1488
+ settings.corpus_root,
1489
+ lance_uri=lance_uri,
1490
+ expected_model=expected_model,
1491
+ expected_dim=expected_dim,
1492
+ lake=lake_reader,
1493
+ lance_uris=_embedding_stores(lake_reader, embed_settings),
1494
+ )
1495
+ _harden_query_connection(
1496
+ con,
1497
+ corpus_root=settings.corpus_root,
1498
+ spill_dir=spill_dir,
1499
+ columnar_paths=sources.lazy_read_paths,
1500
+ analytics_roots=analytics_roots(settings.corpus_root, lake_reader),
1501
+ )
1502
+ refused = _refused_statement_kinds(con, sql)
1503
+ if refused:
1504
+ err = ClassifiedError(
1505
+ kind="sandbox_refused",
1506
+ exit_code=EXIT_CODES["sandbox_refused"],
1507
+ message=f"the query sandbox does not run {', '.join(refused)} statements",
1508
+ hint="query runs SELECT and in-memory DDL/DML only; COPY, EXPORT, "
1509
+ "ATTACH, DETACH, INSTALL, LOAD, PREPARE and EXECUTE are refused "
1510
+ "because they reach files or defer a statement past this check",
1511
+ )
1512
+ emit_error(err, fmt)
1513
+ raise SystemExit(err.exit_code)
1514
+ cursor = con.execute(sql)
1515
+ except REGISTRATION_ERRORS as exc:
1516
+ err = _per_session_memory_hint(classify_registration_error(exc), exc, lake_reader)
1517
+ emit_error(err, fmt)
1518
+ raise SystemExit(err.exit_code) from exc
1519
+ try:
1520
+ emit_cursor(cursor, fmt)
1521
+ except REGISTRATION_ERRORS as exc:
1522
+ err = _per_session_memory_hint(classify_registration_error(exc), exc, lake_reader)
1523
+ emit_error(err, fmt)
1524
+ raise SystemExit(err.exit_code) from exc
1525
+ finally:
1526
+ con.close()
1527
+
1528
+
1529
+ #: The hint a query that ran out of memory on the per-session path carries.
1530
+ PER_SESSION_MEMORY_HINT = (
1531
+ "this query read the per-session files because no usable lake holds the corpus, "
1532
+ "and that path holds every document it reads in memory; the lake answers it in "
1533
+ "bounded memory: run `atif-sql lake rebuild` (or `atif-sql materialize`), or raise "
1534
+ f"{QUERY_MEMORY_LIMIT_ENV}"
1535
+ )
1536
+
1537
+
1538
+ def _per_session_memory_hint(
1539
+ err: ClassifiedError, exc: Exception, lake_reader: object | None
1540
+ ) -> ClassifiedError:
1541
+ """Point an out-of-memory error on the per-session path at the lake.
1542
+
1543
+ The per-session reader loads each trajectory whole through DuckDB's JSON
1544
+ reader, so a full scan of a large corpus can exceed the query cap where
1545
+ the lake, reading typed parquet, doesn't.
1546
+ """
1547
+ import dataclasses
1548
+
1549
+ import duckdb
1550
+
1551
+ if lake_reader is not None or not isinstance(exc, duckdb.OutOfMemoryException):
1552
+ return err
1553
+ return dataclasses.replace(err, hint=PER_SESSION_MEMORY_HINT)
1554
+
1555
+
1556
+ def _embedding_stores(lake_reader: Any, embed_settings: Any) -> list[Path] | None:
1557
+ """Every lake corpus's embeddings store, for a connection scoped to all corpora.
1558
+
1559
+ ``None`` (bind the selected corpus's store, as before) unless the lake
1560
+ reader spans every corpus and the stores follow the per-corpus layout;
1561
+ ``ATIF_SQL_LANCE_URI`` names one store for all of them, so there is
1562
+ nothing to join.
1563
+ """
1564
+ if (
1565
+ lake_reader is None
1566
+ or lake_reader.corpus is not None
1567
+ or embed_settings.lance_uri is not None
1568
+ ):
1569
+ return None
1570
+ from atif_duck.infrastructure.lake import LakeLayout, registered_corpora
1571
+ from atif_duck.infrastructure.lake_settings import LakeSettings
1572
+
1573
+ corpora = registered_corpora(LakeLayout(LakeSettings().lake_root))
1574
+ stores: list[Path] = [
1575
+ embed_settings.resolve_lance_uri(corpora[name].root) for name in sorted(corpora)
1576
+ ]
1577
+ return stores or None
1578
+
1579
+
1580
+ def _attach_query_lake(
1581
+ con: Any, corpus_root: Path, *, all_corpora: bool, fmt: OutputFormat, command: str = "query"
1582
+ ) -> Any:
1583
+ """Attach the lake for ``query``, or warn once and return ``None`` for the per-session path.
1584
+
1585
+ With ``all_corpora`` there is no per-session equivalent, so an unusable
1586
+ lake exits 78 instead. A corpus whose last pass left sessions pending
1587
+ for the lake still reads the lake, after a warning naming how many.
1588
+ """
1589
+ from loguru import logger
1590
+
1591
+ from atif_corpus.application.materialize import read_sink_pending
1592
+ from atif_corpus.domain.layout import CorpusLayout
1593
+ from atif_duck.infrastructure.lake import LakeLayout, LakeReader, attach_lake_for_query
1594
+ from atif_duck.infrastructure.lake_settings import LakeSettings
1595
+
1596
+ attached = attach_lake_for_query(
1597
+ con,
1598
+ LakeLayout(LakeSettings().lake_root),
1599
+ corpus_root=corpus_root,
1600
+ all_corpora=all_corpora,
1601
+ )
1602
+ if not isinstance(attached, LakeReader):
1603
+ if all_corpora:
1604
+ err = ClassifiedError(
1605
+ kind="lake_unavailable",
1606
+ exit_code=EXIT_CODES["lake_unavailable"],
1607
+ message=f"--all-corpora needs the lake: {attached.reason}",
1608
+ hint="run `atif-sql lake rebuild`",
1609
+ )
1610
+ emit_error(err, fmt)
1611
+ raise SystemExit(err.exit_code)
1612
+ logger.warning(
1613
+ "{}: {}; reading the per-session artifacts instead", command, attached.reason
1614
+ )
1615
+ return None
1616
+ if attached.corpus is not None:
1617
+ pending = read_sink_pending(CorpusLayout(corpus_root=corpus_root).sink_pending_path)
1618
+ if pending:
1619
+ logger.warning(
1620
+ "{}: {} session(s) of this corpus are pending a lake write; "
1621
+ "their rows show the previous materialize until the next pass lands",
1622
+ command,
1623
+ len(pending),
1624
+ )
1625
+ return attached
1626
+
1627
+
1628
+ # ---------------------------------------------------------------------------
1629
+ # analyze
1630
+ # ---------------------------------------------------------------------------
1631
+
1632
+
1633
+ @app.command
1634
+ def analyze(
1635
+ *,
1636
+ since_days: int | None = 30,
1637
+ limit: int | None = None,
1638
+ max_sessions: Annotated[int | None, cyclopts.Parameter(name="--max-sessions")] = None,
1639
+ max_cost_usd: Annotated[float | None, cyclopts.Parameter(name="--max-cost-usd")] = None,
1640
+ no_dry_run: Annotated[bool, cyclopts.Parameter(name="--no-dry-run")] = False,
1641
+ llm_only: bool = False,
1642
+ skip_classify: bool = False,
1643
+ skip_conflicts: bool = False,
1644
+ skip_friction: bool = False,
1645
+ skip_perceived: bool = False,
1646
+ corpus_root: Path | None = None,
1647
+ lake: Annotated[bool, cyclopts.Parameter(negative="--no-lake")] = True,
1648
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
1649
+ ) -> None:
1650
+ """Run the four LLM analytics pipelines: classify, conflicts, friction, perceived.
1651
+
1652
+ Defaults to a DRY RUN (plan dicts + cost estimates, zero LLM spend);
1653
+ pass ``--no-dry-run`` to execute them. ``--skip-<stage>`` subtracts
1654
+ individual stages. The deterministic surfaces (``user_steps``,
1655
+ ``human_turns``, ``session_outcomes``) are views, so they need no run.
1656
+
1657
+ Session data comes from the lake when it holds the corpus (a session
1658
+ whose lake rows are not its current artifacts is read from its files),
1659
+ and from the per-session files otherwise, after one warning. The
1660
+ summary's ``session_source`` says which. Outputs still land under the
1661
+ corpus's ``analytics/`` directory either way.
1662
+
1663
+ Parameters
1664
+ ----------
1665
+ since_days
1666
+ Restrict the stages to sessions whose last step is within N days
1667
+ (default 30).
1668
+ limit
1669
+ Cap the number of sessions (newest-first) per stage.
1670
+ max_sessions
1671
+ Hard per-run session ceiling per LLM pipeline (newest-first wins);
1672
+ overrides ``ATIF_SQL_LLM_MAX_SESSIONS_PER_RUN`` (default 50).
1673
+ max_cost_usd
1674
+ Hard per-run dollar ceiling across all LLM pipelines, checked against
1675
+ running actual usage; overrides ``ATIF_SQL_LLM_MAX_COST_USD_PER_RUN``
1676
+ (default 25.0). When crossed, remaining LLM work aborts cleanly and
1677
+ the summary flags ``budget_exhausted``.
1678
+ no_dry_run
1679
+ Execute the LLM stages for real (costs money).
1680
+ llm_only
1681
+ Accepted and ignored: since the structural stages were cut
1682
+ (2026-09-27) every stage is an LLM stage. Kept so existing crontab
1683
+ lines and scripts keep parsing.
1684
+ skip_classify, skip_conflicts, skip_friction, skip_perceived
1685
+ Opt out of one stage.
1686
+ corpus_root
1687
+ Override the materialized corpus root.
1688
+ lake
1689
+ Read session data from the lake (default). ``--no-lake`` reads the
1690
+ per-session files.
1691
+ fmt
1692
+ Summary format; ``auto`` = JSON on a pipe.
1693
+ """
1694
+ _refuse_root("analyze", fmt)
1695
+ del llm_only # a no-op flag; see the docstring
1696
+
1697
+ from atif_analytics.application.analyze import run_analyze
1698
+ from atif_analytics.infrastructure.settings import AnalyticsSettings
1699
+
1700
+ settings = AnalyticsSettings()
1701
+ if corpus_root is None:
1702
+ # Reuse the corpus-root resolution the other commands share so
1703
+ # ATIF_SQL_SOURCE_ROOT/-CORPUS_ROOT behave identically here.
1704
+ settings = settings.model_copy(
1705
+ update={"corpus_root": _corpus_settings(None, None).corpus_root}
1706
+ )
1707
+ else:
1708
+ settings = settings.model_copy(update={"corpus_root": corpus_root})
1709
+ # Budget ceilings: explicit flags beat env/defaults so the crontab line
1710
+ # carries the spend cap visibly (CONTRACT-V2 §Cost guards).
1711
+ if max_sessions is not None:
1712
+ settings = settings.model_copy(update={"llm_max_sessions_per_run": max_sessions})
1713
+ if max_cost_usd is not None:
1714
+ settings = settings.model_copy(update={"llm_max_cost_usd_per_run": max_cost_usd})
1715
+
1716
+ source = _analyze_lake_source(settings.corpus_root) if lake else None
1717
+ try:
1718
+ summary = run_analyze(
1719
+ settings,
1720
+ since_days=since_days,
1721
+ limit=limit,
1722
+ dry_run=not no_dry_run,
1723
+ skip_classify=skip_classify,
1724
+ skip_conflicts=skip_conflicts,
1725
+ skip_friction=skip_friction,
1726
+ skip_perceived=skip_perceived,
1727
+ source=source,
1728
+ )
1729
+ finally:
1730
+ if source is not None:
1731
+ source.close()
1732
+ if source is not None:
1733
+ summary["sessions_read_from_files"] = len(source.from_files)
1734
+ summary["lake_read_failed"] = source.failed
1735
+ emit_json(summary, fmt)
1736
+
1737
+
1738
+ def _analyze_lake_source(corpus_root: Path) -> Any:
1739
+ """The lake's session source for ``corpus_root``, or ``None`` (after a warning) to read the files."""
1740
+ from loguru import logger
1741
+
1742
+ from atif_cli.lake_sessions import LakeSessionSource
1743
+ from atif_duck.infrastructure.lake import LakeLayout
1744
+ from atif_duck.infrastructure.lake_settings import LakeSettings
1745
+
1746
+ opened = LakeSessionSource.open(
1747
+ LakeLayout(LakeSettings().lake_root),
1748
+ corpus_root,
1749
+ memory_limit_bytes=query_memory_limit_bytes(),
1750
+ )
1751
+ if not isinstance(opened, LakeSessionSource):
1752
+ logger.warning("analyze: {}; reading the per-session artifacts instead", opened.reason)
1753
+ return None
1754
+ return opened
1755
+
1756
+
1757
+ # ---------------------------------------------------------------------------
1758
+ # embed
1759
+ # ---------------------------------------------------------------------------
1760
+
1761
+
1762
+ def _install_lance_extension(fmt: OutputFormat, *, quiet: bool = False) -> None:
1763
+ """``INSTALL lance; LOAD lance`` on a throwaway connection.
1764
+
1765
+ The one network reach outside Bedrock, so it lives behind ``embed``.
1766
+ With ``quiet`` (a real embed run) a failure is a warning rather than an
1767
+ exit: the backfill itself still runs, and ``status`` will keep saying
1768
+ the extension is missing until a later attempt succeeds. Without
1769
+ ``quiet`` (``--install-extension``) the result is emitted as JSON and a
1770
+ failure exits 70.
1771
+ """
1772
+ import duckdb
1773
+
1774
+ from atif_cli.duck_errors import classify_duckdb_error
1775
+ from atif_duck.infrastructure.registry import install_lance_extension
1776
+
1777
+ con = duckdb.connect()
1778
+ try:
1779
+ try:
1780
+ install_path = install_lance_extension(con)
1781
+ except duckdb.Error as exc:
1782
+ if quiet:
1783
+ from loguru import logger
1784
+
1785
+ logger.warning(
1786
+ "Could not install the lance DuckDB extension ({}); vector search stays "
1787
+ "unavailable until `atif-sql embed --install-extension` succeeds",
1788
+ str(exc).splitlines()[0],
1789
+ )
1790
+ return
1791
+ err = classify_duckdb_error(exc)
1792
+ emit_error(err, fmt)
1793
+ raise SystemExit(err.exit_code) from exc
1794
+ finally:
1795
+ con.close()
1796
+ if not quiet:
1797
+ emit_json({"extension": "lance", "installed": True, "install_path": install_path}, fmt)
1798
+
1799
+
1800
+ @app.command
1801
+ def embed(
1802
+ *,
1803
+ limit: int | None = None,
1804
+ all_steps: Annotated[bool, cyclopts.Parameter(name="--all")] = False,
1805
+ dry_run: bool | None = None,
1806
+ install_extension: Annotated[bool, cyclopts.Parameter(name="--install-extension")] = False,
1807
+ prune_orphans: Annotated[bool, cyclopts.Parameter(name="--prune-orphans")] = False,
1808
+ lake: Annotated[bool, cyclopts.Parameter(negative="--no-lake")] = True,
1809
+ corpus_root: Path | None = None,
1810
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
1811
+ ) -> None:
1812
+ """Embed unembedded corpus steps and append them to LanceDB.
1813
+
1814
+ Provider
1815
+ --------
1816
+ ``ATIF_SQL_EMBED_PROVIDER`` picks the embedder. ``cohere`` (the default)
1817
+ calls Cohere Embed v4 on Bedrock and writes ``<corpus_root>/embeddings_lance``.
1818
+ ``gemma`` runs EmbeddingGemma 2 on this machine, text-only, and writes
1819
+ ``<corpus_root>/embeddings_lance_gemma``; it needs the ``local`` extra and
1820
+ downloads about 1.5 GB into the Hugging Face cache on first use. Each
1821
+ store is stamped with its model and width, so ``search`` reads the store
1822
+ of the provider selected when it runs.
1823
+
1824
+ Discovery
1825
+ ---------
1826
+ When the DuckLake at ``ATIF_SQL_LAKE_ROOT`` holds this corpus (current
1827
+ schema, same directory), the steps to embed are read from its ``steps``
1828
+ table rather than from every session's ``trajectory.json``, and only the
1829
+ rows changed since the last complete run: the lake snapshot that run
1830
+ reached is kept in ``lake_watermark.json`` inside the store directory.
1831
+ A run cut short by ``--limit``, or one where any row failed to embed,
1832
+ leaves the watermark where it was. Without a usable lake, or with
1833
+ ``--no-lake``, discovery reads the per-session artifacts as before. The
1834
+ dry-run plan's ``discovery`` says which path ran (``corpus``,
1835
+ ``lake-full``, ``lake-incremental``, ``lake-unchanged``).
1836
+
1837
+ Orphans
1838
+ -------
1839
+ ``--prune-orphans`` counts the stored rows whose uuid is no step's
1840
+ primary uuid in the lake (sessions deleted before the corpus kept them,
1841
+ re-keyed steps). It's a dry run unless ``--no-dry-run`` is given too,
1842
+ and it never calls Bedrock. It needs the lake (exit 78 without one), and
1843
+ it refuses to delete while any session of the corpus is still pending a
1844
+ lake write (exit 78), since those sessions' rows would look orphaned.
1845
+ Output: ``{pipeline, lake_snapshot, stored, lake_keys, orphans, deleted,
1846
+ dry_run}``.
1847
+
1848
+ Extension
1849
+ ---------
1850
+ ``query`` and ``search`` read the store through DuckDB's lance extension
1851
+ and never install it themselves (installing is a 242 MB download from
1852
+ the extension repository, and ``query`` runs unattended). It is
1853
+ installed here, where the network is already a deliberate act: every
1854
+ REAL run installs it first, and ``--install-extension`` installs it and
1855
+ exits without touching Bedrock or the store (``{"extension": "lance",
1856
+ "installed": true, "install_path": ...}``). ``atif-sql status`` reports
1857
+ whether it is present.
1858
+
1859
+ Cost
1860
+ ----
1861
+ With Cohere, calls Bedrock (``global.cohere.embed-v4:0``) on every
1862
+ unembedded step (main + sidechain text >= 32 chars, keyed by the step's
1863
+ primary source uuid); with Gemma, the same steps embed on this machine
1864
+ for no charge. A REAL run requires an explicit scope: either ``--limit N``
1865
+ or ``--all`` — a bare ``atif-sql embed`` exits 64 with a hint, so an
1866
+ accidental full backfill (potentially every step ever recorded) can't
1867
+ happen from a mistyped command. ``--dry-run`` needs no scope: it spends
1868
+ nothing and previews the full plan.
1869
+
1870
+ Flags
1871
+ -----
1872
+ --limit N Cap the number of steps embedded this run.
1873
+ --all Explicitly embed EVERY unembedded step (full backfill).
1874
+ --dry-run Preview only; emit plan JSON, no embedding calls.
1875
+ --install-extension Install the lance DuckDB extension and exit (no Bedrock).
1876
+ --prune-orphans Count (and with --no-dry-run, delete) rows no lake step names.
1877
+ --no-lake Discover from the per-session artifacts even when a lake exists.
1878
+ --corpus-root Override the materialized corpus root.
1879
+
1880
+ Output
1881
+ ------
1882
+ Dry run: the plan JSON ``{pipeline, discovery, candidates, batches,
1883
+ batch_size, concurrency, provider, model, dim, store, limit, dry_run}``.
1884
+ Real run:
1885
+ ``{"pipeline": "embed", "rows_processed": N, "dry_run": false}``.
1886
+
1887
+ Exit codes: 0 success, 64 missing --limit/--all or an embedding setting
1888
+ the provider cannot take, 70 runtime (Bedrock / local model / DuckDB /
1889
+ Lance failure — transient, safe to retry), 78 terminal state (the store
1890
+ or its config requires operator action, such as the gemma provider without
1891
+ the ``local`` extra; retrying without intervention cannot succeed, so
1892
+ unattended lanes suppress retries on 78).
1893
+ """
1894
+ if install_extension:
1895
+ _install_lance_extension(fmt)
1896
+ return
1897
+ if prune_orphans:
1898
+ _prune_orphans(corpus_root, dry_run=dry_run is not False, fmt=fmt)
1899
+ return
1900
+
1901
+ import asyncio
1902
+
1903
+ from atif_embed.application.embed import run_backfill
1904
+ from atif_embed.domain.errors import DomainError
1905
+ from atif_embed.infrastructure.corpus_text_rows import DuckDbTextRows
1906
+
1907
+ dry_run = bool(dry_run)
1908
+ if not dry_run and limit is None and not all_steps:
1909
+ emit_error(
1910
+ ClassifiedError(
1911
+ kind="invalid_input",
1912
+ exit_code=EXIT_CODES["invalid_input"],
1913
+ message="a real embed run needs an explicit scope",
1914
+ hint="pass --limit N (bounded) or --all (full backfill), or preview with --dry-run",
1915
+ ),
1916
+ fmt,
1917
+ )
1918
+ raise SystemExit(EXIT_CODES["invalid_input"])
1919
+
1920
+ settings = _corpus_settings(None, corpus_root)
1921
+ embed_settings = _embed_settings(fmt)
1922
+ if not dry_run:
1923
+ _install_lance_extension(fmt, quiet=True)
1924
+ text_rows: Any = DuckDbTextRows()
1925
+ if lake:
1926
+ from atif_cli.embed_lake import lake_steps_port
1927
+ from atif_embed.infrastructure.lake_text_rows import LAKE_WATERMARK_FILE, LakeTextRows
1928
+
1929
+ text_rows = LakeTextRows(
1930
+ lake=lake_steps_port(),
1931
+ watermark_path=embed_settings.resolve_lance_uri(settings.corpus_root)
1932
+ / LAKE_WATERMARK_FILE,
1933
+ fallback=text_rows,
1934
+ )
1935
+ try:
1936
+ result = asyncio.run(
1937
+ run_backfill(
1938
+ corpus_root=settings.corpus_root,
1939
+ settings=embed_settings,
1940
+ text_rows=text_rows,
1941
+ limit=limit,
1942
+ dry_run=dry_run,
1943
+ )
1944
+ )
1945
+ except DomainError as exc:
1946
+ # Terminal vs transient decides the exit code: a terminal state
1947
+ # (provider mismatch, unhealable schema) needs an operator, and the
1948
+ # refresh lane suppresses retries on 78 instead of re-running an
1949
+ # identical failure every tick.
1950
+ kind = "terminal_state" if exc.terminal else "runtime_error"
1951
+ emit_error(
1952
+ ClassifiedError(
1953
+ kind=kind,
1954
+ exit_code=EXIT_CODES[kind],
1955
+ message=str(exc),
1956
+ ),
1957
+ fmt,
1958
+ )
1959
+ raise SystemExit(EXIT_CODES[kind]) from exc
1960
+ if isinstance(result, dict):
1961
+ emit_json(result, fmt)
1962
+ else:
1963
+ emit_json({"pipeline": "embed", "rows_processed": result, "dry_run": False}, fmt)
1964
+
1965
+
1966
+ def _prune_orphans(corpus_root: Path | None, *, dry_run: bool, fmt: OutputFormat) -> None:
1967
+ """``embed --prune-orphans``: count, and unless ``dry_run`` delete, the store's orphan rows."""
1968
+ from atif_cli.embed_lake import lake_steps_port
1969
+ from atif_corpus.application.materialize import read_sink_pending
1970
+ from atif_corpus.domain.layout import CorpusLayout
1971
+ from atif_embed.application.prune import prune_orphans
1972
+ from atif_embed.infrastructure.lance_store import LanceVectorStore
1973
+
1974
+ settings = _corpus_settings(None, corpus_root)
1975
+ embed_settings = _embed_settings(fmt)
1976
+ lance_uri = embed_settings.resolve_lance_uri(settings.corpus_root)
1977
+ pending = read_sink_pending(CorpusLayout(corpus_root=settings.corpus_root).sink_pending_path)
1978
+ if pending and not dry_run:
1979
+ emit_error(
1980
+ ClassifiedError(
1981
+ kind="lake_unavailable",
1982
+ exit_code=EXIT_CODES["lake_unavailable"],
1983
+ message=f"{len(pending)} session(s) of this corpus are pending a lake write, "
1984
+ "so their stored rows would look orphaned",
1985
+ hint="let the next materialize pass land them, then prune again",
1986
+ ),
1987
+ fmt,
1988
+ )
1989
+ raise SystemExit(EXIT_CODES["lake_unavailable"])
1990
+ report = prune_orphans(
1991
+ settings.corpus_root,
1992
+ lake=lake_steps_port(),
1993
+ store=LanceVectorStore(lance_uri, dim=embed_settings.embedding_dim),
1994
+ dry_run=dry_run,
1995
+ )
1996
+ if report is None:
1997
+ emit_error(
1998
+ ClassifiedError(
1999
+ kind="lake_unavailable",
2000
+ exit_code=EXIT_CODES["lake_unavailable"],
2001
+ message="--prune-orphans checks the store against the lake, and there is no "
2002
+ "usable lake for this corpus",
2003
+ hint="run `atif-sql lake rebuild`",
2004
+ ),
2005
+ fmt,
2006
+ )
2007
+ raise SystemExit(EXIT_CODES["lake_unavailable"])
2008
+ report["pending_sessions"] = len(pending)
2009
+ emit_json(report, fmt)
2010
+
2011
+
2012
+ # ---------------------------------------------------------------------------
2013
+ # search
2014
+ # ---------------------------------------------------------------------------
2015
+
2016
+
2017
+ @app.command
2018
+ def search(
2019
+ query_text: str,
2020
+ /,
2021
+ *,
2022
+ k: Annotated[int, cyclopts.Parameter(name=["-k", "--k"])] = 10,
2023
+ session_id: str | None = None,
2024
+ corpus_root: Path | None = None,
2025
+ all_corpora: bool = False,
2026
+ lake: Annotated[bool, cyclopts.Parameter(negative="--no-lake")] = True,
2027
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
2028
+ ) -> None:
2029
+ """Semantic top-k nearest-neighbor search over step embeddings.
2030
+
2031
+ Pipeline
2032
+ --------
2033
+ 1. Embed ``query_text`` with the selected provider: Cohere Embed v4
2034
+ ``search_query`` mode (float) by default, or EmbeddingGemma 2's
2035
+ ``SearchQuery`` prompt on this machine with ``ATIF_SQL_EMBED_PROVIDER=gemma``.
2036
+ 2. DuckDB cosine-kNN against the Lance-backed ``message_embeddings`` view
2037
+ (``register_vss`` guard-before-bind: a store written by a different
2038
+ provider raises instead of returning garbage scores).
2039
+ 3. Join back to ``steps`` via ``source_uuids[0]`` for a 200-char snippet.
2040
+
2041
+ Where the steps come from
2042
+ -------------------------
2043
+ Like ``query``: the lake's ``steps`` table when the lake holds this
2044
+ corpus, else (after a one-line warning, or with ``--no-lake``) the
2045
+ per-session artifacts. ``--all-corpora`` searches every corpus the lake
2046
+ holds, each through its own store, and adds a ``corpus`` column; it
2047
+ needs the lake (exit 78 without one).
2048
+
2049
+ Prereq
2050
+ ------
2051
+ The Lance store must exist. If it's empty or missing, the command exits
2052
+ with code 2 and a hint. Run ``atif-sql embed --all --no-dry-run`` to
2053
+ populate.
2054
+
2055
+ Flags
2056
+ -----
2057
+ --k N Top-k (default 10).
2058
+ --session-id ID Confine the kNN to one session.
2059
+ --corpus-root Override the materialized corpus root.
2060
+ --all-corpora Search every corpus the lake holds.
2061
+ --no-lake Read the per-session artifacts even when a lake exists.
2062
+
2063
+ Output columns
2064
+ --------------
2065
+ uuid, session_id, snippet, sim (cosine similarity ∈ [-1, 1]), plus
2066
+ corpus under ``--all-corpora``. Sorted by cosine distance ascending —
2067
+ highest sim first.
2068
+
2069
+ Exit codes: 0 success, 2 no_embeddings, 64 --all-corpora with --no-lake,
2070
+ 65 embedding_mismatch (the store was written by another provider), 70
2071
+ runtime (the query embedding failed), 77 root_refused (uid 0 without ``ATIF_SQL_ALLOW_ROOT=1``), 78
2072
+ extension_missing (a store exists but the lance DuckDB extension is not
2073
+ installed; run ``atif-sql embed --install-extension``), lake_unavailable
2074
+ (``--all-corpora`` without a usable lake) or terminal_state (the gemma
2075
+ provider is selected and the ``local`` extra is not installed).
2076
+ """
2077
+ _refuse_root("search", fmt)
2078
+
2079
+ import duckdb
2080
+
2081
+ from atif_cli.duck_errors import (
2082
+ REGISTRATION_ERRORS,
2083
+ classify_duckdb_error,
2084
+ classify_registration_error,
2085
+ )
2086
+ from atif_duck.infrastructure.registry import lance_extension_installed, register
2087
+ from atif_embed.application.embed import embed_query
2088
+ from atif_embed.domain.errors import DomainError
2089
+
2090
+ if all_corpora and not lake:
2091
+ err = ClassifiedError(
2092
+ kind="invalid_input",
2093
+ exit_code=EXIT_CODES["invalid_input"],
2094
+ message="--all-corpora reads the lake, and --no-lake turns it off",
2095
+ hint="drop one of the two flags",
2096
+ )
2097
+ emit_error(err, fmt)
2098
+ raise SystemExit(err.exit_code)
2099
+
2100
+ settings = _corpus_settings(None, corpus_root)
2101
+ embed_settings = _embed_settings(fmt)
2102
+ expected_model, expected_dim = embed_settings.expected_embedding_identity()
2103
+ lance_uri = embed_settings.resolve_lance_uri(settings.corpus_root)
2104
+
2105
+ con = duckdb.connect(":memory:")
2106
+ try:
2107
+ try:
2108
+ lake_reader = (
2109
+ _attach_query_lake(
2110
+ con, settings.corpus_root, all_corpora=all_corpora, fmt=fmt, command="search"
2111
+ )
2112
+ if lake
2113
+ else None
2114
+ )
2115
+ stores = _embedding_stores(lake_reader, embed_settings)
2116
+ register(
2117
+ con,
2118
+ settings.corpus_root,
2119
+ lance_uri=lance_uri,
2120
+ expected_model=expected_model,
2121
+ expected_dim=expected_dim,
2122
+ lake=lake_reader,
2123
+ lance_uris=stores,
2124
+ )
2125
+ except REGISTRATION_ERRORS as exc:
2126
+ err = classify_registration_error(exc)
2127
+ emit_error(err, fmt)
2128
+ raise SystemExit(err.exit_code) from exc
2129
+
2130
+ if any(uri.is_dir() for uri in stores or [lance_uri]) and not lance_extension_installed(
2131
+ con
2132
+ ):
2133
+ emit_error(
2134
+ ClassifiedError(
2135
+ kind="extension_missing",
2136
+ exit_code=EXIT_CODES["extension_missing"],
2137
+ message="the lance DuckDB extension is not installed, so the embeddings "
2138
+ f"store at {lance_uri} cannot be read",
2139
+ hint="run: atif-sql embed --install-extension (a one-time download)",
2140
+ ),
2141
+ fmt,
2142
+ )
2143
+ raise SystemExit(EXIT_CODES["extension_missing"])
2144
+
2145
+ row = con.execute("SELECT count(*) FROM message_embeddings").fetchone()
2146
+ if not row or int(row[0]) == 0:
2147
+ emit_error(
2148
+ ClassifiedError(
2149
+ kind="no_embeddings",
2150
+ exit_code=EXIT_CODES["no_embeddings"],
2151
+ message="No embeddings yet.",
2152
+ hint="Run: atif-sql embed --all --no-dry-run (or --limit N)",
2153
+ ),
2154
+ fmt,
2155
+ )
2156
+ raise SystemExit(EXIT_CODES["no_embeddings"])
2157
+
2158
+ try:
2159
+ qv = embed_query(query_text, settings=embed_settings)
2160
+ except DomainError as exc:
2161
+ kind = "terminal_state" if exc.terminal else "runtime_error"
2162
+ emit_error(
2163
+ ClassifiedError(kind=kind, exit_code=EXIT_CODES[kind], message=str(exc)), fmt
2164
+ )
2165
+ raise SystemExit(EXIT_CODES[kind]) from exc
2166
+ try:
2167
+ columns, rows = search_rows(
2168
+ con, qv, k=k, session_id=session_id, with_corpus=all_corpora
2169
+ )
2170
+ except duckdb.Error as exc:
2171
+ err = classify_duckdb_error(exc)
2172
+ emit_error(err, fmt)
2173
+ raise SystemExit(err.exit_code) from exc
2174
+ emit_rows(columns, rows, fmt)
2175
+ finally:
2176
+ con.close()
2177
+
2178
+
2179
+ #: ``search``'s kNN, one text per shape. Rank by cosine similarity
2180
+ #: descending: ORDER BY array_cosine_distance (== 1 - sim) ASC is what
2181
+ #: triggers the cosine HNSW index lookup. Using array_distance here (L2) would
2182
+ #: silently bypass the index AND give wrong ranks: the raw int8-cast-to-float
2183
+ #: document vectors have magnitudes in the thousands while the query vector is
2184
+ #: unit-normalized — only cosine is magnitude-invariant. ``{dim}`` is the only
2185
+ #: placeholder, filled with ``int(len(vector))``.
2186
+ _SEARCH_SQL = """
2187
+ WITH qv AS (SELECT CAST(? AS FLOAT[{dim}]) AS v)
2188
+ SELECT me.uuid AS uuid,
2189
+ s.session_id AS session_id,
2190
+ substr(s.message, 1, 200) AS snippet,
2191
+ array_cosine_similarity(me.embedding, (SELECT v FROM qv)) AS sim{corpus}
2192
+ FROM message_embeddings me
2193
+ JOIN steps s
2194
+ ON json_extract_string(s.source_uuids, '$[0]') = me.uuid{corpus_join}
2195
+ {session_filter}
2196
+ ORDER BY array_cosine_distance(me.embedding, (SELECT v FROM qv)) ASC
2197
+ LIMIT ?
2198
+ """
2199
+
2200
+
2201
+ def search_rows(
2202
+ con: Any,
2203
+ query_vector: list[float],
2204
+ *,
2205
+ k: int,
2206
+ session_id: str | None = None,
2207
+ with_corpus: bool = False,
2208
+ ) -> tuple[list[str], list[tuple[Any, ...]]]:
2209
+ """Run ``search``'s kNN on a registered connection; return ``(columns, rows)``.
2210
+
2211
+ Separate from the command so the same statement runs over a query vector
2212
+ the caller already has (the tests' fixed vectors; no Bedrock call).
2213
+ """
2214
+ params: list[object] = [query_vector]
2215
+ if session_id is not None:
2216
+ params.append(session_id)
2217
+ params.append(k)
2218
+ sql = _SEARCH_SQL.format(
2219
+ dim=len(query_vector),
2220
+ corpus=",\n c.corpus AS corpus"
2221
+ if with_corpus
2222
+ else "",
2223
+ corpus_join="\n JOIN sessions c ON c.session_id = s.session_id" if with_corpus else "",
2224
+ session_filter="WHERE s.session_id = ?" if session_id is not None else "",
2225
+ )
2226
+ cursor = con.execute(sql, params)
2227
+ columns = [d[0] for d in cursor.description or ()]
2228
+ return columns, cursor.fetchall()
2229
+
2230
+
2231
+ # ---------------------------------------------------------------------------
2232
+ # examples
2233
+ # ---------------------------------------------------------------------------
2234
+
2235
+ #: Header line stamped onto the examples output. Honest by construction:
2236
+ #: packages/atif-duck/tests/test_examples.py executes every one of these
2237
+ #: statements against the fixture corpus, so a release with a broken example
2238
+ #: cannot ship.
2239
+ _EXAMPLES_NOTE = "test-executed against this version"
2240
+
2241
+
2242
+ @app.command
2243
+ def examples(
2244
+ *,
2245
+ category: str | None = None,
2246
+ requires: str | None = None,
2247
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
2248
+ ) -> None:
2249
+ """List tested example queries for every view and macro.
2250
+
2251
+ Examples are DERIVED from the static atif-duck catalog (never hardcoded
2252
+ per object) and each one is executed by atif-duck's test suite against a
2253
+ fixture corpus — the listing is proof-carrying, not documentation that
2254
+ can rot. Answers from :mod:`atif_duck.domain` alone: no DuckDB import,
2255
+ no connection, lean-import fast like ``schema``.
2256
+
2257
+ Output
2258
+ ------
2259
+ TTY: a table grouped by ``requires`` (core / analytics / vss). Piped or
2260
+ ``--format json``: ``{"note": ..., "examples": [{name, sql, description,
2261
+ requires, category}, ...]}``.
2262
+
2263
+ Flags
2264
+ -----
2265
+ --category Filter: view | table-macro | scalar-macro.
2266
+ --requires Filter: core | analytics | vss.
2267
+
2268
+ Exit codes: 0 ok, 64 unknown --category/--requires value.
2269
+ """
2270
+ from atif_duck.domain.examples import (
2271
+ CATEGORY_VALUES,
2272
+ REQUIRES_VALUES,
2273
+ build_examples,
2274
+ )
2275
+
2276
+ for flag, value, allowed in (
2277
+ ("--category", category, CATEGORY_VALUES),
2278
+ ("--requires", requires, REQUIRES_VALUES),
2279
+ ):
2280
+ if value is not None and value not in allowed:
2281
+ emit_error(
2282
+ ClassifiedError(
2283
+ kind="invalid_input",
2284
+ exit_code=EXIT_CODES["invalid_input"],
2285
+ message=f"unknown {flag} value: {value!r}",
2286
+ hint=f"one of: {', '.join(allowed)}",
2287
+ ),
2288
+ fmt,
2289
+ )
2290
+ raise SystemExit(EXIT_CODES["invalid_input"])
2291
+
2292
+ selected = [
2293
+ example
2294
+ for example in build_examples()
2295
+ if (category is None or example.category == category)
2296
+ and (requires is None or example.requires == requires)
2297
+ ]
2298
+
2299
+ if resolve_format(fmt) is OutputFormat.TABLE:
2300
+ print(f"# {len(selected)} examples — {_EXAMPLES_NOTE}")
2301
+ for group in ("core", "analytics", "vss"):
2302
+ grouped = [e for e in selected if e.requires == group]
2303
+ if not grouped:
2304
+ continue
2305
+ print(f"\n[{group}]")
2306
+ width = max(len(e.name) for e in grouped)
2307
+ for example in grouped:
2308
+ print(f" {example.name:<{width}} {example.sql}")
2309
+ print(f" {'':<{width}} # {example.description}")
2310
+ return
2311
+ emit_json(
2312
+ {
2313
+ "note": _EXAMPLES_NOTE,
2314
+ "examples": [
2315
+ {
2316
+ "name": e.name,
2317
+ "sql": e.sql,
2318
+ "description": e.description,
2319
+ "requires": e.requires,
2320
+ "category": e.category,
2321
+ }
2322
+ for e in selected
2323
+ ],
2324
+ },
2325
+ fmt,
2326
+ )
2327
+
2328
+
2329
+ # ---------------------------------------------------------------------------
2330
+ # schema
2331
+ # ---------------------------------------------------------------------------
2332
+
2333
+
2334
+ @app.command
2335
+ def schema(
2336
+ *,
2337
+ fmt: Annotated[OutputFormat, cyclopts.Parameter(name="--format")] = OutputFormat.AUTO,
2338
+ ) -> None:
2339
+ """List every view (with columns) and every macro signature, with what each requires.
2340
+
2341
+ The canonical catalog for composing ``query`` calls. Answers from the
2342
+ static :mod:`atif_duck.domain.catalog` dicts (core :data:`VIEW_SCHEMA`
2343
+ and :data:`MACRO_SIGNATURES`, then :data:`ANALYTICS_VIEW_SCHEMA` and
2344
+ :data:`ANALYTICS_MACRO_SIGNATURES`) with no DuckDB import, no
2345
+ connection, no view registration; sub-50ms by construction (drift
2346
+ against the real DDL is caught by atif-duck's CI tests).
2347
+
2348
+ Each object carries ``requires``: ``core`` binds on any corpus,
2349
+ ``analytics`` once ``atif-sql analyze`` has written its parquet, ``vss``
2350
+ once ``atif-sql embed`` has built the store.
2351
+ """
2352
+ from atif_duck.domain.catalog import (
2353
+ ANALYTICS_MACRO_SIGNATURES,
2354
+ ANALYTICS_VIEW_SCHEMA,
2355
+ MACRO_SIGNATURES,
2356
+ VIEW_SCHEMA,
2357
+ )
2358
+ from atif_duck.domain.examples import object_requires
2359
+
2360
+ views = {**VIEW_SCHEMA, **ANALYTICS_VIEW_SCHEMA}
2361
+ macros = {**MACRO_SIGNATURES, **ANALYTICS_MACRO_SIGNATURES}
2362
+ examples_hint = "tested example queries: atif-sql examples (or: atif-sql query --examples)"
2363
+ if resolve_format(fmt) is OutputFormat.TABLE:
2364
+ for name, cols in views.items():
2365
+ print(f"\n{name} ({len(cols)} cols, requires: {object_requires(name)})")
2366
+ for col, col_type in cols:
2367
+ print(f" {col:<28} {col_type}")
2368
+ print(f"\nMacros ({len(macros)})")
2369
+ for macro, params in macros.items():
2370
+ call = f"{macro}({', '.join(params)})"
2371
+ print(f" {call:<44} requires: {object_requires(macro)}")
2372
+ print(f"\n{examples_hint}")
2373
+ return
2374
+ emit_json(
2375
+ {
2376
+ "views": {
2377
+ name: [{"column": c, "type": t} for c, t in cols] for name, cols in views.items()
2378
+ },
2379
+ "view_requires": {name: object_requires(name) for name in views},
2380
+ "macros": [
2381
+ {"name": n, "params": list(p), "requires": object_requires(n)}
2382
+ for n, p in macros.items()
2383
+ ],
2384
+ "examples_hint": examples_hint,
2385
+ },
2386
+ fmt,
2387
+ )
2388
+
2389
+
2390
+ #: Env var selecting the stderr log level. `logger.remove()` below drops the
2391
+ #: default handler, and with it the `LOGURU_LEVEL` that would have
2392
+ #: parameterized it, so this is the ONE knob reaching the workspace's INFO and
2393
+ #: DEBUG surface — measured 2026-08-28 over `packages/*/src`, 110 `logger.info`
2394
+ #: and 41 `logger.debug` sites that are otherwise unreachable through the CLI.
2395
+ #: Named to the pydantic-settings prefix every other setting uses.
2396
+ LOG_LEVEL_ENV: str = "ATIF_SQL_LOG_LEVEL"
2397
+ DEFAULT_LOG_LEVEL: str = "WARNING"
2398
+
2399
+
2400
+ def _stderr_log_level() -> str:
2401
+ """The level :func:`main` gives its stderr sink, read from the environment.
2402
+
2403
+ Upper-cased, and an unset or blank value reads as the default rather than
2404
+ as a request for level ``""``.
2405
+ """
2406
+ return os.environ.get(LOG_LEVEL_ENV, DEFAULT_LOG_LEVEL).strip().upper() or DEFAULT_LOG_LEVEL
2407
+
2408
+
2409
+ def main() -> None:
2410
+ """Console-script entry point.
2411
+
2412
+ Default stderr stays quiet for routine reads: loguru's default DEBUG sink
2413
+ is replaced with WARNING-and-up, so a piped read emits data on stdout and
2414
+ nothing on stderr unless something is actually wrong. Set
2415
+ :data:`LOG_LEVEL_ENV` to widen it.
2416
+
2417
+ An unknown level falls back to WARNING and says so, rather than replacing
2418
+ the caller's command with a loguru traceback — a bad log level is not a
2419
+ reason to refuse to run.
2420
+
2421
+ A usage error cyclopts refuses (an unknown option, a missing or uncoercible
2422
+ value) exits ``EXIT_CODES["invalid_input"]`` (64) after cyclopts prints its
2423
+ error panel, never cyclopts' own 2, which the table gives to ``empty_session``.
2424
+ """
2425
+ from loguru import logger
2426
+
2427
+ level = _stderr_log_level()
2428
+ logger.remove()
2429
+ try:
2430
+ logger.add(sys.stderr, level=level)
2431
+ except ValueError:
2432
+ logger.add(sys.stderr, level=DEFAULT_LOG_LEVEL)
2433
+ logger.warning(
2434
+ "{}={!r} is not a loguru level; using {}", LOG_LEVEL_ENV, level, DEFAULT_LOG_LEVEL
2435
+ )
2436
+ try:
2437
+ app(sys.argv[1:], exit_on_error=False)
2438
+ except cyclopts.CycloptsError:
2439
+ # cyclopts 5 exits 2 on a usage error (unknown option, missing or uncoercible
2440
+ # value), and EXIT_CODES gives 2 to empty_session and no_embeddings. A malformed
2441
+ # flag is invalid_input. cyclopts has already printed its error panel.
2442
+ raise SystemExit(EXIT_CODES["invalid_input"]) from None
2443
+
2444
+
2445
+ __all__ = [
2446
+ "EXIT_CODES",
2447
+ "ClassifiedError",
2448
+ "analyze",
2449
+ "app",
2450
+ "convert",
2451
+ "embed",
2452
+ "examples",
2453
+ "main",
2454
+ "materialize",
2455
+ "query",
2456
+ "schema",
2457
+ "search",
2458
+ "status",
2459
+ ]