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
@@ -0,0 +1,2194 @@
1
+ # SPDX-License-Identifier: Apache-2.0
2
+
3
+ """DuckDB view and macro registry over the materialized ATIF corpus.
4
+
5
+ Wires a DuckDB connection to a ``<corpus_root>/sessions/<id>/`` tree (per
6
+ docs/CONTRACT.md: ``trajectory.json``, ``edges.jsonl``, ``session_events.jsonl``,
7
+ each stored as ``<name>.zst`` or, in an older corpus, plain, plus
8
+ ``loss_report.json`` and ``meta.json``) and exposes it as a stable set of SQL views and analytical
9
+ macros. Every view reads ATIF-v1.7, never raw Claude Code JSONL, so a
10
+ column's meaning comes from the ATIF schema rather than from a transcript
11
+ field that happens to share its name.
12
+
13
+ Design notes
14
+ ------------
15
+ * Raw readers are ``CREATE TEMP TABLE`` over ``read_json`` with an explicit
16
+ ``columns={...}`` projection. In DuckDB 1.5+ the dict is a *strict
17
+ projection filter*: every field any downstream view touches must appear in
18
+ the dict or it silently disappears (the strict-filter lesson). Listing the
19
+ fields also skips JSON schema inference — the dominant cost on a large
20
+ corpus, otherwise paid on every ``DESCRIBE``/view bind.
21
+ * ``filename=true`` + a path regexp derive ``session_id_path`` from the
22
+ ``sessions/<id>/`` directory, which is the canonical session key (the
23
+ in-document ``session_id`` is advisory; the directory name is what
24
+ atif-corpus keyed the materialization on).
25
+ * ``trajectory.json`` is ONE document per file (``format='auto'``, compact
26
+ JSON), NOT newline-delimited; ``steps`` is projected as a ``JSON[]``
27
+ column so views unnest it lazily at query time. ``edges.jsonl`` is
28
+ ``format='newline_delimited'``. DuckDB decompresses a ``.zst`` file on its
29
+ own, so the readers take the stored file of each session
30
+ (:mod:`atif_duck.infrastructure.stored_artifacts`), and the path columns
31
+ strip ``.zst`` so a row names the same logical file before and after a
32
+ slim.
33
+ * Two sources per session, chosen per session. A session whose
34
+ ``meta.json`` names the current ``columnar_schema`` and whose parquet
35
+ artifacts are present (:mod:`atif_duck.domain.columnar`) is read through
36
+ lazy views over ``read_parquet``, so no JSON is parsed for it at query
37
+ time; every other session is read from its trajectory exactly as
38
+ before, unless the caller passes ``stage_columnar`` (the lake loader),
39
+ which stages parquet for those sessions first. When a corpus holds both kinds the two branches ``UNION ALL`` into
40
+ one relation per raw reader (``v_raw_trajectories``, ``v_raw_steps``,
41
+ ``v_raw_tool_calls``, ``v_raw_tool_results``), and the business views
42
+ never learn which branch a row came from. The parquet branch is lazy on
43
+ purpose: loading the tool-result text eagerly measured at 2.1 s per
44
+ process on a 300-session corpus, which is what the JSON path already
45
+ costs, whereas a lazy view answers the panel queries in about 20 ms.
46
+ :func:`register_raw` returns which sessions took which path and every
47
+ parquet the views will open, so the query sandbox can grant exactly those
48
+ files.
49
+ * All views use ``CREATE OR REPLACE`` so callers may safely re-register.
50
+ * No corpus path is ever part of a statement's text. The TEMP TABLE readers
51
+ take their glob or file list as a bound parameter (``read_json(?)``); the
52
+ parquet readers are relations built through the connection's
53
+ ``read_parquet`` API (the file list is a Python value) and registered as
54
+ views, because ``CREATE VIEW`` itself cannot be prepared. Every remaining interpolation is
55
+ a module constant, a catalog constant, or a projection expression, and the
56
+ AST test in ``tests/test_sql_text_boundaries.py`` fails on anything else.
57
+ * A session directory whose name fails the boundary in
58
+ :mod:`atif_duck.domain.session_id` registers nothing: the name is the one
59
+ piece of outside text that becomes a path, so it is checked before any path
60
+ is built, once per registration, and reported in
61
+ :attr:`RawSources.rejected_session_ids`.
62
+ * register-or-fail-loud: every registration function logs via
63
+ ``logger.exception`` and re-raises on any DDL failure.
64
+
65
+ Semantics a transcript-shaped reader gets wrong (documented per-view below)
66
+ --------------------------------------------------------------------------
67
+ * ``steps``, not ``messages``, is the primary per-turn surface: one row
68
+ per ATIF step (harbor bundles all events sharing an assistant
69
+ ``message.id`` into one step, so step count < raw message count by design).
70
+ * ``messages`` is a uuid-keyed view reconstructed from ``edges.jsonl``,
71
+ which carries the raw-record identity the trajectory itself cannot
72
+ provide (fidelity gap 7: UUID_NOT_PRESERVED).
73
+ * Token columns follow ATIF ``Metrics`` semantics: ``prompt_tokens`` is the
74
+ TOTAL input (non-cached + cache_read + cache_creation), ``cached_tokens``
75
+ is the cache-read subset, and ``cache_creation`` is recovered from
76
+ ``metrics.extra`` (fidelity gap 6: CACHE_SPLIT_PARTIAL).
77
+ """
78
+
79
+ from __future__ import annotations
80
+
81
+ from dataclasses import dataclass
82
+ from typing import TYPE_CHECKING
83
+
84
+ from loguru import logger
85
+
86
+ from atif_duck.domain.artifacts import EDGES_JSONL, TRAJECTORY_JSON
87
+ from atif_duck.domain.catalog import DEFAULT_PRICING, VIEW_SCHEMA
88
+ from atif_duck.domain.columnar import (
89
+ META_COLUMNAR_KEY,
90
+ SESSION_COLUMNS,
91
+ SESSION_EVENTS_JSONL,
92
+ SESSION_EVENTS_PARQUET,
93
+ SESSION_PARQUET,
94
+ STEPS_PARQUET,
95
+ TOOL_CALLS_PARQUET,
96
+ TOOL_RESULTS_PARQUET,
97
+ )
98
+ from atif_duck.domain.embedding_guard import EmbeddingProviderMismatch, ensure_store_matches
99
+ from atif_duck.domain.raw_readers import (
100
+ CORPUS_COLUMN,
101
+ EDGE_COLUMNS,
102
+ LOSS_REPORT_COLUMNS,
103
+ META_COLUMNS,
104
+ RAW_EDGES,
105
+ RAW_LOSS_REPORTS,
106
+ RAW_META,
107
+ RAW_SESSION_EVENTS,
108
+ RAW_STEPS,
109
+ RAW_TOOL_CALLS,
110
+ RAW_TOOL_RESULTS,
111
+ RAW_TRAJECTORIES,
112
+ TRAJECTORY_RELATION_COLUMNS,
113
+ )
114
+ from atif_duck.domain.session_id import session_id_rejection
115
+ from atif_duck.domain.sql_literal import SqlFragment, sql_literal
116
+ from atif_duck.infrastructure.columnar import ColumnarCoverage, session_has_columnar
117
+ from atif_duck.infrastructure.projections import (
118
+ CALL_COLUMNS,
119
+ EVENT_COLUMNS,
120
+ RESULT_COLUMNS,
121
+ WHOLE_STEP,
122
+ render,
123
+ step_columns,
124
+ step_key_columns,
125
+ )
126
+ from atif_duck.infrastructure.stored_artifacts import decoded_size, stored_artifact
127
+
128
+ if TYPE_CHECKING:
129
+ from collections.abc import Callable, Collection, Iterable, Mapping, Sequence
130
+ from pathlib import Path
131
+
132
+ import duckdb
133
+
134
+ from atif_duck.infrastructure.lake import LakeReader
135
+
136
+ # ---------------------------------------------------------------------------
137
+ # Raw-reader object names
138
+ # ---------------------------------------------------------------------------
139
+ #: The materialized raw readers. Named constants rather than inline literals
140
+ #: so a DDL interpolation site cannot typo a second table into existence.
141
+ #: Deliberately absent from :data:`atif_duck.domain.catalog.VIEW_NAMES`:
142
+ #: they describe corpus files rather than being queryable business surface.
143
+ #:
144
+ #: ``v_raw_trajectories`` carries the trajectory's top-level members minus
145
+ #: ``steps``; the step-level surface is ``v_raw_steps`` /
146
+ #: ``v_raw_tool_calls`` / ``v_raw_tool_results``, each already in its view's
147
+ #: column shape whichever source it came from. The JSON-path TEMP TABLE that
148
+ #: still holds ``steps JSON[]`` is ``v_raw_trajectories_json``.
149
+ _RAW_TRAJECTORIES_TABLE: str = RAW_TRAJECTORIES
150
+ _RAW_TRAJECTORIES_JSON_TABLE: str = "v_raw_trajectories_json"
151
+ _RAW_STEPS_TABLE: str = RAW_STEPS
152
+ _RAW_TOOL_CALLS_TABLE: str = RAW_TOOL_CALLS
153
+ _RAW_TOOL_RESULTS_TABLE: str = RAW_TOOL_RESULTS
154
+ _RAW_EDGES_TABLE: str = RAW_EDGES
155
+ _RAW_LOSS_REPORTS_TABLE: str = RAW_LOSS_REPORTS
156
+ _RAW_META_TABLE: str = RAW_META
157
+ #: ``session_events`` rows from either source, in the view's column shape. The
158
+ #: JSON path's TEMP TABLE is the parsed ``session_events.jsonl`` lines.
159
+ _RAW_SESSION_EVENTS_TABLE: str = RAW_SESSION_EVENTS
160
+ _RAW_SESSION_EVENTS_JSON_TABLE: str = "v_raw_session_events_json"
161
+ #: One row, one column: the corpus's ``sessions/`` directory, inserted as a
162
+ #: bound parameter. The columnar branch derives ``trajectory_path`` from it,
163
+ #: so the path is data rather than statement text, and it reads the same for a
164
+ #: session whose parquet the lake loader staged elsewhere.
165
+ _RAW_SESSIONS_DIR_TABLE: str = "v_raw_sessions_dir"
166
+
167
+ #: The columnar branch's parquet readers: one parameterized ``read_parquet``
168
+ #: relation per artifact kind, registered as a view under these names and
169
+ #: selected from by the union views above. They exist only while at least one
170
+ #: session is read from parquet.
171
+ _RAW_SESSIONS_PARQUET_TABLE: str = "v_raw_sessions_parquet"
172
+ _RAW_STEPS_PARQUET_TABLE: str = "v_raw_steps_parquet"
173
+ _RAW_TOOL_CALLS_PARQUET_TABLE: str = "v_raw_tool_calls_parquet"
174
+ _RAW_TOOL_RESULTS_PARQUET_TABLE: str = "v_raw_tool_results_parquet"
175
+ _RAW_SESSION_EVENTS_PARQUET_TABLE: str = "v_raw_session_events_parquet"
176
+
177
+
178
+ @dataclass(frozen=True, slots=True)
179
+ class RawSources:
180
+ """Which source each complete session was registered from.
181
+
182
+ Returned by :func:`register_raw` (and :func:`register`) so the
183
+ composition root can grant the sandbox exactly the files the views read
184
+ lazily, and so a test can prove the columnar path was taken rather than
185
+ silently falling back to JSON.
186
+ """
187
+
188
+ #: Sessions read through lazy views over their parquet artifacts.
189
+ columnar_session_ids: tuple[str, ...]
190
+ #: Sessions read from ``trajectory.json`` (no or stale columnar artifacts).
191
+ json_session_ids: tuple[str, ...]
192
+ #: Every parquet a registered raw reader opens at caller-query time.
193
+ lazy_read_paths: tuple[Path, ...]
194
+ #: Session directory names that failed the session id boundary and so
195
+ #: contribute to no view (sorted). Logged once each at registration.
196
+ rejected_session_ids: tuple[str, ...] = ()
197
+ #: True when the raw relations are views over the lake rather than over
198
+ #: the per-session artifacts (:func:`atif_duck.infrastructure.lake.register_lake_raw`).
199
+ from_lake: bool = False
200
+ #: The corpus a lake registration is scoped to (``None``: every corpus,
201
+ #: or not a lake registration).
202
+ lake_corpus: str | None = None
203
+ #: Sessions read from parquet the caller staged from their trajectory
204
+ #: (``stage_columnar``, the lake loader's path); also in
205
+ #: :attr:`columnar_session_ids`.
206
+ staged_session_ids: tuple[str, ...] = ()
207
+
208
+ @property
209
+ def coverage(self) -> ColumnarCoverage:
210
+ """The same counts ``atif-sql status`` reports."""
211
+ return ColumnarCoverage(
212
+ columnar_sessions=len(self.columnar_session_ids),
213
+ json_sessions=len(self.json_session_ids),
214
+ )
215
+
216
+
217
+ #: Ceiling for a ``read_json`` ``maximum_object_size``. Live trajectory.json
218
+ #: files reach 436 MB because harbor inlines subagent sidechains and tool
219
+ #: outputs; 1 GiB is about 2.3x headroom over the largest observed document.
220
+ _OBJECT_SIZE_CAP: int = 1_073_741_824
221
+
222
+ #: Floor for the same bound: DuckDB's own default.
223
+ _OBJECT_SIZE_FLOOR: int = 16_777_216
224
+
225
+ #: Headroom over the largest file when sizing the bound: a quarter of the file
226
+ #: plus one MiB, so a file that grows a little between the stat and the read
227
+ #: still parses.
228
+ _OBJECT_SIZE_HEADROOM_DIVISOR: int = 4
229
+ _OBJECT_SIZE_HEADROOM_BYTES: int = 1_048_576
230
+
231
+ #: The DuckDB extension that reads the Lance embeddings store.
232
+ LANCE_EXTENSION: str = "lance"
233
+
234
+
235
+ def _object_size_bound(paths: Iterable[Path]) -> int:
236
+ """``maximum_object_size`` for a ``read_json`` over ``paths``, sized from the files.
237
+
238
+ The bound is not free. DuckDB's eager JSON reader reserves about twice
239
+ this many bytes PER THREAD before it parses anything, so the 1 GiB
240
+ constant this replaced cost 2 GiB a thread: over 300 ``edges.jsonl``
241
+ files of under 8 MB each it exhausted a 6 GB ``memory_limit`` at four
242
+ threads and a 25 GB one at sixteen, while 16 MiB read the same rows in
243
+ 0.10 s (measured on the 300-session panel corpus). An earlier note here
244
+ claimed lowering it changed nothing; that measurement held the thread
245
+ count at one.
246
+
247
+ A newline-delimited file's objects are its lines and a one-document file
248
+ IS its object, so the largest file present bounds every object either
249
+ reader meets. That size plus headroom, floored at DuckDB's default and
250
+ capped at :data:`_OBJECT_SIZE_CAP`, is the bound. A compressed file is
251
+ measured by its decompressed size (read from its frame header), because
252
+ that's what the reader holds. A path that vanished between the listing
253
+ and the stat counts as zero.
254
+ """
255
+ largest = 0
256
+ for path in paths:
257
+ largest = max(largest, decoded_size(path))
258
+ with_headroom = largest + largest // _OBJECT_SIZE_HEADROOM_DIVISOR + _OBJECT_SIZE_HEADROOM_BYTES
259
+ return max(_OBJECT_SIZE_FLOOR, min(_OBJECT_SIZE_CAP, with_headroom))
260
+
261
+
262
+ # Explicit projection for ``v_raw_trajectories``.
263
+ #
264
+ # ATIF-v1.7 ``Trajectory`` is ``extra='forbid'``, so the top level is a
265
+ # closed set; we project exactly what the views touch. ``steps`` stays
266
+ # ``JSON[]`` (not a deep STRUCT) because Step.message is a union
267
+ # (str | ContentPart[]) and Step.extra / metrics.extra are free-form —
268
+ # a STRUCT projection would force one shape and silently null the other.
269
+ _TRAJECTORY_COLUMNS: dict[str, str] = {
270
+ "schema_version": "VARCHAR",
271
+ "session_id": "VARCHAR",
272
+ "trajectory_id": "VARCHAR",
273
+ "agent": "STRUCT(name VARCHAR, version VARCHAR, model_name VARCHAR, extra JSON)",
274
+ "steps": "JSON[]",
275
+ "final_metrics": (
276
+ "STRUCT(total_prompt_tokens BIGINT, total_completion_tokens BIGINT, "
277
+ "total_cached_tokens BIGINT, total_cost_usd DOUBLE, total_steps BIGINT, "
278
+ "extra JSON)"
279
+ ),
280
+ "extra": "JSON",
281
+ }
282
+
283
+ # The explicit ``read_json`` projections for ``v_raw_edges``,
284
+ # ``v_raw_loss_reports`` and ``v_raw_meta`` live in
285
+ # :mod:`atif_duck.domain.raw_readers`, beside the lake tables derived from them.
286
+ _EDGE_COLUMNS: dict[str, str] = EDGE_COLUMNS
287
+ _LOSS_REPORT_COLUMNS: dict[str, str] = LOSS_REPORT_COLUMNS
288
+ _META_COLUMNS: dict[str, str] = META_COLUMNS
289
+
290
+
291
+ def _typed_empty(columns: Sequence[tuple[str, str]]) -> SqlFragment:
292
+ """A zero-row SELECT carrying the declared column names and types."""
293
+ projected = ", ".join(f"CAST(NULL AS {sql_type}) AS {name}" for name, sql_type in columns)
294
+ return SqlFragment(f"SELECT {projected} WHERE false")
295
+
296
+
297
+ def _typed_parquet_select(reader: str, columns: Sequence[tuple[str, str]]) -> SqlFragment:
298
+ """SELECT the declared columns, cast to their declared types, from a parquet reader.
299
+
300
+ ``reader`` is one of the ``_RAW_*_PARQUET_TABLE`` names bound by
301
+ :func:`_bind_parquet_reader`; the paths themselves never appear here. The
302
+ casts are belt and braces: the producer wrote these exact types, and a
303
+ ``UNION ALL`` against the JSON branch needs both sides to agree on every
304
+ column, so the declared type is asserted here rather than inferred from
305
+ the first file.
306
+ """
307
+ projected = ", ".join(f"CAST({name} AS {sql_type}) AS {name}" for name, sql_type in columns)
308
+ return SqlFragment(f"SELECT {projected} FROM {reader}") # noqa: S608 # nosec B608 - reader is a module constant; columns are catalog constants
309
+
310
+
311
+ def _union_or_empty(
312
+ branches: Sequence[SqlFragment], columns: Sequence[tuple[str, str]]
313
+ ) -> SqlFragment:
314
+ """Join the present source branches, or an empty typed relation if none is."""
315
+ return SqlFragment("\nUNION ALL\n".join(branches)) if branches else _typed_empty(columns)
316
+
317
+
318
+ def _json_steps_select() -> SqlFragment:
319
+ """The JSON path's ``steps`` projection over ``v_raw_trajectories_json``."""
320
+ return SqlFragment(
321
+ f"SELECT t.session_id_path AS session_id,\n {render(step_columns(WHOLE_STEP))}\n" # nosec B608 - projection constants over a module-constant table
322
+ f"FROM {_RAW_TRAJECTORIES_JSON_TABLE} t, UNNEST(t.steps) AS s(step)"
323
+ )
324
+
325
+
326
+ def _json_tool_calls_select() -> SqlFragment:
327
+ """The JSON path's ``tool_calls`` projection.
328
+
329
+ The inner query narrows each step to (ids, ts, calls list) BEFORE the
330
+ second UNNEST: a step's JSON string is megabytes on harbor trajectories
331
+ that inline subagent sidechains, and unnesting it directly replicates
332
+ that string once per tool call.
333
+ """
334
+ return SqlFragment(
335
+ f"SELECT session_id, step_id, ts, {render(CALL_COLUMNS)}\n" # noqa: S608 # nosec B608 - constants only
336
+ "FROM (\n"
337
+ f" SELECT t.session_id_path AS session_id, {render(step_key_columns(WHOLE_STEP))},\n"
338
+ f" {WHOLE_STEP.json('tool_calls', '[*]')} AS calls\n"
339
+ f" FROM {_RAW_TRAJECTORIES_JSON_TABLE} t, UNNEST(t.steps) AS s(step)\n"
340
+ f" WHERE {WHOLE_STEP.json('tool_calls')} IS NOT NULL\n"
341
+ ") step_calls, UNNEST(calls) AS c(call)"
342
+ )
343
+
344
+
345
+ def _json_tool_results_select() -> SqlFragment:
346
+ """The JSON path's ``tool_results`` projection (same early narrowing)."""
347
+ return SqlFragment(
348
+ f"SELECT session_id, step_id, ts, {render(RESULT_COLUMNS)}\n" # noqa: S608 # nosec B608 - constants only
349
+ "FROM (\n"
350
+ f" SELECT t.session_id_path AS session_id, {render(step_key_columns(WHOLE_STEP))},\n"
351
+ f" {WHOLE_STEP.json('observation', '.results[*]')} AS results\n"
352
+ f" FROM {_RAW_TRAJECTORIES_JSON_TABLE} t, UNNEST(t.steps) AS s(step)\n"
353
+ f" WHERE {WHOLE_STEP.json('observation', '.results')} IS NOT NULL\n"
354
+ ") step_results, UNNEST(results) AS r(res)"
355
+ )
356
+
357
+
358
+ def _events_json_select() -> SqlFragment:
359
+ """The JSON path's ``session_events`` rows over ``v_raw_session_events_json``."""
360
+ return SqlFragment(
361
+ f"SELECT {_catalog_columns('session_events')} FROM {_RAW_SESSION_EVENTS_JSON_TABLE}" # noqa: S608 # nosec B608 - catalog columns over a module-constant table
362
+ )
363
+
364
+
365
+ def _register_events_json(con: duckdb.DuckDBPyConnection, paths: Sequence[Path]) -> None:
366
+ """Parse the JSON-path sessions' ``session_events.jsonl`` into a TEMP TABLE.
367
+
368
+ ``read_json_objects`` hands each line over whole as the ``ev`` JSON value,
369
+ so the typed columns come from :data:`EVENT_COLUMNS`, the same expressions
370
+ the columnar producer applies, and the two sources agree byte for byte.
371
+ The file list is the statement's one parameter.
372
+ """
373
+ bound = _object_size_bound(paths)
374
+ con.execute(
375
+ f"""
376
+ CREATE OR REPLACE TEMP TABLE {_RAW_SESSION_EVENTS_JSON_TABLE} AS
377
+ SELECT regexp_extract(filename, '/sessions/([^/]+)/session_events\\.jsonl(\\.zst)?$', 1)
378
+ AS session_id,
379
+ {render(EVENT_COLUMNS)}
380
+ FROM (
381
+ SELECT json AS ev, filename
382
+ FROM read_json_objects(
383
+ ?,
384
+ format='newline_delimited',
385
+ filename=true,
386
+ maximum_object_size={int(bound)}
387
+ )
388
+ );
389
+ """, # noqa: S608 # nosec B608 - file list is a bound parameter; table and projections are constants; the bound is an int
390
+ [[str(path) for path in paths]],
391
+ )
392
+ logger.debug(
393
+ "Registered {} over {} session_events.jsonl file(s)",
394
+ _RAW_SESSION_EVENTS_JSON_TABLE,
395
+ len(paths),
396
+ )
397
+
398
+
399
+ #: ``v_raw_trajectories`` columns, in order: the trajectory's top-level
400
+ #: members minus ``steps``, plus the two path-derived keys (the ``corpus``
401
+ #: column the union view appends is not part of either branch).
402
+ _TRAJECTORY_RELATION_COLUMNS: tuple[tuple[str, str], ...] = TRAJECTORY_RELATION_COLUMNS
403
+
404
+
405
+ def _json_trajectories_select() -> SqlFragment:
406
+ names = ", ".join(name for name, _ in _TRAJECTORY_RELATION_COLUMNS)
407
+ return SqlFragment(f"SELECT {names} FROM {_RAW_TRAJECTORIES_JSON_TABLE}") # noqa: S608 # nosec B608 - constants only
408
+
409
+
410
+ def _columnar_trajectories_select() -> SqlFragment:
411
+ """``session.parquet`` rows in ``v_raw_trajectories`` shape.
412
+
413
+ ``trajectory_path`` is the logical ATIF document path
414
+ (``<sessions_dir>/<id>/trajectory.json``), the value the JSON branch
415
+ reports too, so the ``sessions`` view's path column reads the same
416
+ whichever branch produced the row and whether the document is stored
417
+ compressed. The sessions directory comes from the one-row
418
+ :data:`_RAW_SESSIONS_DIR_TABLE`, filled by bound parameter, so the corpus
419
+ root never has to be spliced into the statement, and a session whose
420
+ parquet the lake loader staged outside the corpus still names its own
421
+ document.
422
+ """
423
+ casts = ", ".join(
424
+ f"CAST({name} AS {sql_type}) AS {name}"
425
+ for name, sql_type in SESSION_COLUMNS
426
+ if name != "session_id_path"
427
+ )
428
+ path_expr = (
429
+ f"(SELECT sessions_dir FROM {_RAW_SESSIONS_DIR_TABLE}) || '/' || session_id_path " # noqa: S608 # nosec B608 - module constant
430
+ f"|| '/{TRAJECTORY_JSON}'"
431
+ )
432
+ return SqlFragment(
433
+ f"SELECT {casts}, {path_expr} AS trajectory_path, session_id_path " # noqa: S608 # nosec B608 - catalog constants over a module-constant reader
434
+ f"FROM {_RAW_SESSIONS_PARQUET_TABLE}"
435
+ )
436
+
437
+
438
+ def _render_columns_clause(columns: dict[str, str]) -> SqlFragment:
439
+ """Render a ``columns={...}`` clause body for ``read_json``.
440
+
441
+ Keys are bare DuckDB identifiers (no quoting); values are SQL type
442
+ strings wrapped in single quotes. Both halves come from code-side
443
+ constants — never user input — so escaping is defensive only.
444
+ """
445
+ return SqlFragment(", ".join(f"{name}: {sql_literal(typ)}" for name, typ in columns.items()))
446
+
447
+
448
+ def _catalog_columns(view_name: str) -> SqlFragment:
449
+ """The catalog's column names for ``view_name``, comma-joined in catalog order."""
450
+ return SqlFragment(", ".join(name for name, _ in VIEW_SCHEMA[view_name]))
451
+
452
+
453
+ def _bind_parquet_reader(
454
+ con: duckdb.DuckDBPyConnection,
455
+ reader: str,
456
+ paths: Sequence[Path],
457
+ *,
458
+ with_filename: bool = False,
459
+ ) -> None:
460
+ """Register ``reader`` as a lazy ``read_parquet`` relation over ``paths``.
461
+
462
+ ``CREATE VIEW ... read_parquet(?)`` is refused ("this type of statement
463
+ can't be prepared"), so the file list goes through the connection's own
464
+ ``read_parquet`` relation API, which takes it as a Python value, and that
465
+ relation is registered as the view. It plans as the same ``PARQUET_SCAN``
466
+ a hand-written view would (projection and filter pushdown included) and
467
+ opens the files per query, not at registration. Measured on the 300
468
+ session panel corpus it matches the inlined-literal view on both wall
469
+ time and peak memory; a relation built from ``sql("... read_parquet(?)",
470
+ params=...)`` instead did NOT (about 3x the memory), so that shape is
471
+ deliberately not used here.
472
+ """
473
+ files = [str(path) for path in paths]
474
+ con.read_parquet(files, filename=with_filename).create_view(reader, replace=True)
475
+
476
+
477
+ # ---------------------------------------------------------------------------
478
+ # Raw readers
479
+ # ---------------------------------------------------------------------------
480
+
481
+
482
+ def _gate_session_dirs(
483
+ con: duckdb.DuckDBPyConnection,
484
+ sessions_dir: Path,
485
+ wanted: Collection[str] | None = None,
486
+ ) -> tuple[str, ...]:
487
+ """Apply the two per-directory gates; return the names the boundary rejected.
488
+
489
+ Walks ``sessions_dir`` once, right after ``v_raw_meta`` is read and
490
+ before any per-session path is built:
491
+
492
+ * A directory whose name fails the session id boundary
493
+ (:mod:`atif_duck.domain.session_id`) is REJECTED: its row is deleted
494
+ from ``v_raw_meta`` (by bound parameter, so the offending name is data
495
+ even here), which keeps it out of every other reader through the meta
496
+ gate and out of :func:`_split_sources`, and it is logged once. A corpus
497
+ an older version wrote may hold such a dir; it is reported in
498
+ :attr:`RawSources.rejected_session_ids` and never read.
499
+ * A well-named directory without a ``meta.json`` is INCOMPLETE (torn
500
+ writer, partial cleanup). The meta gate already excludes it; this makes
501
+ the exclusion OBSERVABLE, because a session silently missing from
502
+ ``sessions`` is much harder to diagnose than a logged skip.
503
+ """
504
+ if not sessions_dir.is_dir():
505
+ return ()
506
+ # The one interpolation is the module constant _RAW_META_TABLE.
507
+ rows = con.execute(f"SELECT session_id_path FROM {_RAW_META_TABLE}").fetchall() # noqa: S608 # nosec B608 - module constant
508
+ with_meta = {row[0] for row in rows}
509
+ rejected: list[str] = []
510
+ # A filtered registration is handed names rather than a listing; either
511
+ # way a name is checked before it becomes a path.
512
+ names = (
513
+ sorted(entry.name for entry in sessions_dir.iterdir() if entry.is_dir())
514
+ if wanted is None
515
+ else sorted(set(wanted))
516
+ )
517
+ for name in names:
518
+ rejection = session_id_rejection(name)
519
+ if rejection is not None:
520
+ logger.warning(
521
+ "Rejecting session dir {!r} ({}); nothing from it is registered", name, rejection
522
+ )
523
+ con.execute(f"DELETE FROM {_RAW_META_TABLE} WHERE session_id_path = ?", [name]) # noqa: S608 # nosec B608 - module constant; the name is bound
524
+ rejected.append(name)
525
+ continue
526
+ session_dir = sessions_dir / name
527
+ if name not in with_meta and session_dir.is_dir():
528
+ logger.warning(
529
+ "Skipping incomplete session dir {} (no meta.json — "
530
+ "crashed writer or partial cleanup); excluded from all views",
531
+ session_dir,
532
+ )
533
+ return tuple(sorted(rejected))
534
+
535
+
536
+ def _split_sources(
537
+ con: duckdb.DuckDBPyConnection, sessions_dir: Path, *, use_session_parquet: bool = True
538
+ ) -> tuple[list[str], dict[str, Path]]:
539
+ """Partition the meta-bearing sessions into columnar ids and ``{JSON id: trajectory file}``.
540
+
541
+ A session is columnar when :func:`session_has_columnar` says so (and
542
+ ``use_session_parquet`` is on); every other session with a stored
543
+ trajectory (``trajectory.json.zst`` or ``trajectory.json``) is JSON. A
544
+ meta-bearing dir with neither readable source is skipped with a warning
545
+ rather than failing the whole registration: ``read_json`` over an
546
+ explicit path list errors on one missing file, and one broken session
547
+ must not take the corpus offline.
548
+
549
+ Runs after :func:`_gate_session_dirs`, so every id read here has passed
550
+ the session id boundary and may become a path component.
551
+ """
552
+ # The one interpolation is the module constant _RAW_META_TABLE plus the
553
+ # contract's meta key.
554
+ rows = con.execute(
555
+ f"SELECT session_id_path, {META_COLUMNAR_KEY} FROM {_RAW_META_TABLE} ORDER BY 1" # noqa: S608 # nosec B608 - module constants
556
+ ).fetchall()
557
+ columnar_ids: list[str] = []
558
+ json_ids: dict[str, Path] = {}
559
+ for session_id, columnar_schema in rows:
560
+ session_dir = sessions_dir / str(session_id)
561
+ if use_session_parquet and session_has_columnar(session_dir, columnar_schema):
562
+ columnar_ids.append(str(session_id))
563
+ continue
564
+ trajectory = stored_artifact(session_dir, TRAJECTORY_JSON)
565
+ if trajectory is not None:
566
+ json_ids[str(session_id)] = trajectory
567
+ else:
568
+ logger.warning(
569
+ "Skipping session dir {} (meta.json present but no readable "
570
+ "trajectory or columnar artifacts); excluded from all views",
571
+ session_dir,
572
+ )
573
+ return columnar_ids, json_ids
574
+
575
+
576
+ def _typed_empty_table(
577
+ con: duckdb.DuckDBPyConnection, table: str, columns: Sequence[tuple[str, str]]
578
+ ) -> None:
579
+ """Create ``table`` as a zero-row TEMP TABLE of ``columns`` (a filtered pass found no file)."""
580
+ con.execute(f"CREATE OR REPLACE TEMP TABLE {table} AS {_typed_empty(columns)};") # nosec B608 - module constant over catalog constants
581
+
582
+
583
+ def _json_reader_columns(columns: dict[str, str], path_column: str) -> tuple[tuple[str, str], ...]:
584
+ """What ``SELECT *, filename AS <path>, <key>`` over a ``read_json(filename=true)`` yields."""
585
+ return (
586
+ *columns.items(),
587
+ ("filename", "VARCHAR"),
588
+ (path_column, "VARCHAR"),
589
+ ("session_id_path", "VARCHAR"),
590
+ )
591
+
592
+
593
+ def _existing(paths: Iterable[Path]) -> list[str]:
594
+ return [str(path) for path in paths if path.is_file()]
595
+
596
+
597
+ def register_raw(
598
+ con: duckdb.DuckDBPyConnection,
599
+ corpus_root: Path,
600
+ *,
601
+ session_ids: Collection[str] | None = None,
602
+ use_session_parquet: bool = True,
603
+ stage_columnar: Callable[[Sequence[str]], Mapping[str, Path]] | None = None,
604
+ ) -> RawSources:
605
+ """Create the raw readers over ``corpus_root``.
606
+
607
+ ``v_raw_meta`` / ``v_raw_edges`` / ``v_raw_loss_reports`` are TEMP TABLEs
608
+ over ``read_json`` of the CONTRACT corpus layout
609
+ ``<corpus_root>/sessions/<session_id>/...``, with ``session_id_path``
610
+ derived from the directory component via regexp over ``filename``. Each
611
+ glob or file list is ONE bound parameter, never statement text. (Meta and
612
+ loss reports go through a glob rather than an explicit file list because
613
+ DuckDB reads 300 small files about twice as fast through a glob; the cost
614
+ is that a corpus root containing a glob metacharacter such as ``[`` or a
615
+ backslash is not readable, which was already so. Edges are a list of each
616
+ session's stored file, since the file may carry either spelling.)
617
+
618
+ ``v_raw_trajectories`` / ``v_raw_steps`` / ``v_raw_tool_calls`` /
619
+ ``v_raw_tool_results`` are VIEWs that union up to two sources, chosen
620
+ per session by :func:`_split_sources`: the session's parquet artifacts
621
+ (read lazily through the ``_RAW_*_PARQUET_TABLE`` readers, each a
622
+ ``read_parquet`` relation over a Python file list, no JSON touched at
623
+ query time) or its ``trajectory.json`` (parsed once into the TEMP TABLE
624
+ ``v_raw_trajectories_json`` from a bound file list and unnested per
625
+ query, as every version before the columnar artifacts did). A corpus
626
+ with no columnar session registers exactly as before; a corpus with no
627
+ JSON session never opens a ``trajectory.json``. ``v_raw_trajectories``
628
+ ends with a ``corpus`` column: the corpus directory's name, the same
629
+ value the lake stores, so the ``sessions`` view names its corpus on
630
+ either path.
631
+
632
+ STORED NAMES: ``trajectory.json``, ``edges.jsonl`` and
633
+ ``session_events.jsonl`` may each be stored as ``<name>.zst``
634
+ (:mod:`atif_duck.domain.artifacts`); each session's stored file is
635
+ resolved once and DuckDB decompresses it. Their path columns
636
+ (``trajectory_path``, ``edges_path``) name the logical artifact either
637
+ way.
638
+
639
+ TORN-SET GUARD: ``meta.json`` is written last by the corpus writer, so
640
+ its presence marks a session dir as complete. Every other reader is
641
+ restricted to meta-bearing session dirs (the JSON path list is built
642
+ from ``v_raw_meta``; edges and loss use a semi-join), so a session dir
643
+ missing its meta (crashed writer) contributes nothing to any view
644
+ instead of a partial artifact set.
645
+
646
+ SESSION ID BOUNDARY: a session dir whose name fails
647
+ :mod:`atif_duck.domain.session_id` is removed from ``v_raw_meta`` by
648
+ :func:`_gate_session_dirs` before any per-session path is built, so the
649
+ meta gate excludes it everywhere. A corpus an older version wrote may
650
+ hold such a dir; it is logged once and reported, never read.
651
+
652
+ Parameters
653
+ ----------
654
+ con
655
+ Open DuckDB connection.
656
+ corpus_root
657
+ Materialized corpus root (the directory containing ``sessions/``).
658
+ session_ids
659
+ Register only these sessions (the lake writer loads a batch at a
660
+ time). The globs become explicit lists of the files that exist, and a
661
+ batch with no complete session registers every relation empty rather
662
+ than failing. ``None`` (the default) registers the whole corpus.
663
+ use_session_parquet
664
+ Read a session from its own parquet artifacts when they're current
665
+ (the default). ``False`` ignores them, so every session is read from
666
+ its trajectory: ``corpus slim`` verifies the lake that way before it
667
+ deletes them.
668
+ stage_columnar
669
+ Called once with the ids that would be read from their trajectory
670
+ through DuckDB's JSON reader; returns ``{id: directory}`` for each
671
+ one it wrote the five columnar files for (from the same trajectory).
672
+ Those sessions are read from the staged parquet instead. The lake
673
+ loader passes one: parsing a large trajectory in DuckDB needs far
674
+ more than the writer's memory cap, and the columnar producer types it
675
+ in bounded memory into the same rows.
676
+
677
+ Returns
678
+ -------
679
+ RawSources
680
+ Which sessions took which path, every parquet the views open at
681
+ query time (for the sandbox's file allowlist), and the rejected names.
682
+
683
+ Raises
684
+ ------
685
+ duckdb.Error
686
+ If any DDL fails (including an empty/absent corpus on an unfiltered
687
+ pass — ``read_json`` errors on a glob with zero matches, which is the
688
+ honest failure mode for "nothing materialized yet"). Logged via
689
+ ``logger.exception`` before re-raise.
690
+ """
691
+ sessions_dir = corpus_root / "sessions"
692
+ wanted = None if session_ids is None else sorted(set(session_ids))
693
+ # A filtered pass names its files. Every id is checked against the
694
+ # boundary before it becomes a path, and only files that exist are
695
+ # listed, because read_json errors on one missing file in a list.
696
+ wanted_ok = [sid for sid in wanted or () if session_id_rejection(sid) is None]
697
+ meta_source: str | list[str] = (
698
+ str(sessions_dir / "*" / "meta.json")
699
+ if wanted is None
700
+ else _existing(sessions_dir / sid / "meta.json" for sid in wanted_ok)
701
+ )
702
+
703
+ # meta.json is written LAST by atif-corpus — its presence marks a
704
+ # session dir as COMPLETE. The edges and loss readers are restricted to
705
+ # meta-bearing session dirs via this semi-join predicate, so a torn
706
+ # session dir (crashed writer, partial cleanup) is invisible to every
707
+ # view rather than surfacing a partial artifact set.
708
+ # The one interpolation is the module constant _RAW_META_TABLE.
709
+ meta_gate = f"session_id_path IN (SELECT session_id_path FROM {_RAW_META_TABLE})" # noqa: S608 # nosec B608 - module constant
710
+
711
+ try:
712
+ # meta FIRST: every other reader derives from it. The glob (or file
713
+ # list) is the statement's one parameter.
714
+ if meta_source:
715
+ con.execute(
716
+ f"""
717
+ CREATE OR REPLACE TEMP TABLE {_RAW_META_TABLE} AS
718
+ SELECT *,
719
+ filename AS meta_path,
720
+ regexp_extract(filename, '/sessions/([^/]+)/meta\\.json$', 1)
721
+ AS session_id_path
722
+ FROM read_json(
723
+ ?,
724
+ format='auto',
725
+ filename=true,
726
+ columns={{{_render_columns_clause(_META_COLUMNS)}}}
727
+ );
728
+ """, # noqa: S608 # nosec B608 - glob is a bound parameter; table and columns are constants
729
+ [meta_source],
730
+ )
731
+ else:
732
+ _typed_empty_table(
733
+ con, _RAW_META_TABLE, _json_reader_columns(_META_COLUMNS, "meta_path")
734
+ )
735
+ logger.debug("Registered {} from {}", _RAW_META_TABLE, "glob" if wanted is None else "list")
736
+ rejected = _gate_session_dirs(con, sessions_dir, wanted)
737
+
738
+ columnar_ids, json_trajectories = _split_sources(
739
+ con, sessions_dir, use_session_parquet=use_session_parquet
740
+ )
741
+ columnar_dirs = {sid: sessions_dir / sid for sid in columnar_ids}
742
+ staged: dict[str, Path] = {}
743
+ if stage_columnar is not None and json_trajectories:
744
+ staged = dict(stage_columnar(list(json_trajectories)))
745
+ for sid, directory in staged.items():
746
+ columnar_dirs[sid] = directory
747
+ del json_trajectories[sid]
748
+ columnar_ids = sorted(columnar_dirs)
749
+ json_ids = list(json_trajectories)
750
+ logger.info(
751
+ "register_raw: {} session(s) read from columnar artifacts ({} staged), "
752
+ "{} from their trajectory",
753
+ len(columnar_ids),
754
+ len(staged),
755
+ len(json_ids),
756
+ )
757
+ con.execute(
758
+ f"CREATE OR REPLACE TEMP TABLE {_RAW_SESSIONS_DIR_TABLE} AS " # nosec B608 - module constant; the directory is a bound parameter
759
+ "SELECT CAST(? AS VARCHAR) AS sessions_dir",
760
+ [str(sessions_dir)],
761
+ )
762
+
763
+ trajectory_branches: list[SqlFragment] = []
764
+ steps_branches: list[SqlFragment] = []
765
+ calls_branches: list[SqlFragment] = []
766
+ results_branches: list[SqlFragment] = []
767
+ events_branches: list[SqlFragment] = []
768
+ lazy_paths: list[Path] = []
769
+
770
+ parquet_readers: tuple[tuple[str, str, bool], ...] = (
771
+ (_RAW_SESSIONS_PARQUET_TABLE, SESSION_PARQUET, True),
772
+ (_RAW_STEPS_PARQUET_TABLE, STEPS_PARQUET, False),
773
+ (_RAW_TOOL_CALLS_PARQUET_TABLE, TOOL_CALLS_PARQUET, False),
774
+ (_RAW_TOOL_RESULTS_PARQUET_TABLE, TOOL_RESULTS_PARQUET, False),
775
+ (_RAW_SESSION_EVENTS_PARQUET_TABLE, SESSION_EVENTS_PARQUET, False),
776
+ )
777
+ if columnar_ids:
778
+ for reader, artifact, with_filename in parquet_readers:
779
+ paths = [columnar_dirs[sid] / artifact for sid in columnar_ids]
780
+ _bind_parquet_reader(con, reader, paths, with_filename=with_filename)
781
+ lazy_paths.extend(paths)
782
+ trajectory_branches.append(_columnar_trajectories_select())
783
+ steps_branches.append(
784
+ _typed_parquet_select(_RAW_STEPS_PARQUET_TABLE, VIEW_SCHEMA["steps"])
785
+ )
786
+ calls_branches.append(
787
+ _typed_parquet_select(_RAW_TOOL_CALLS_PARQUET_TABLE, VIEW_SCHEMA["tool_calls"])
788
+ )
789
+ results_branches.append(
790
+ _typed_parquet_select(_RAW_TOOL_RESULTS_PARQUET_TABLE, VIEW_SCHEMA["tool_results"])
791
+ )
792
+ events_branches.append(
793
+ _typed_parquet_select(
794
+ _RAW_SESSION_EVENTS_PARQUET_TABLE, VIEW_SCHEMA["session_events"]
795
+ )
796
+ )
797
+ else:
798
+ # A connection re-registered after every columnar session lost
799
+ # its artifacts must not keep the previous generation's readers.
800
+ for reader, _, _ in parquet_readers:
801
+ con.execute(f"DROP VIEW IF EXISTS {reader};")
802
+
803
+ if json_ids:
804
+ # One trajectory document per file -> format='auto' (NOT NDJSON).
805
+ # The explicit projection keeps `steps` as a lazy JSON[] column.
806
+ # An explicit path list rather than a glob: the glob would parse
807
+ # the columnar sessions' trajectory.json too, which is the cost
808
+ # this whole arrangement exists to avoid. The list is the
809
+ # statement's one parameter.
810
+ # A compressed document is decompressed by the reader (detected
811
+ # from the .zst suffix, per file); trajectory_path drops the
812
+ # suffix so it names the logical document either way.
813
+ trajectory_paths = [json_trajectories[sid] for sid in json_ids]
814
+ json_files = [str(path) for path in trajectory_paths]
815
+ trajectory_bound = _object_size_bound(trajectory_paths)
816
+ con.execute(
817
+ f"""
818
+ CREATE OR REPLACE TEMP TABLE {_RAW_TRAJECTORIES_JSON_TABLE} AS
819
+ SELECT *,
820
+ regexp_replace(filename, '\\.zst$', '') AS trajectory_path,
821
+ regexp_extract(filename, '/sessions/([^/]+)/trajectory\\.json(\\.zst)?$', 1)
822
+ AS session_id_path
823
+ FROM read_json(
824
+ ?,
825
+ format='auto',
826
+ filename=true,
827
+ columns={{{_render_columns_clause(_TRAJECTORY_COLUMNS)}}},
828
+ maximum_object_size={int(trajectory_bound)}
829
+ );
830
+ """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns are constants; the bound is an int
831
+ [json_files],
832
+ )
833
+ logger.debug(
834
+ "Registered {} over {} trajectory.json file(s)",
835
+ _RAW_TRAJECTORIES_JSON_TABLE,
836
+ len(json_ids),
837
+ )
838
+ trajectory_branches.append(_json_trajectories_select())
839
+ steps_branches.append(_json_steps_select())
840
+ calls_branches.append(_json_tool_calls_select())
841
+ results_branches.append(_json_tool_results_select())
842
+ else:
843
+ # A connection re-registered after every session gained its
844
+ # artifacts must not keep the previous generation's parsed JSON.
845
+ con.execute(f"DROP TABLE IF EXISTS {_RAW_TRAJECTORIES_JSON_TABLE};")
846
+
847
+ # A JSON-path session's events come from its session_events.jsonl
848
+ # (stored compressed or plain). A session materialized before that
849
+ # artifact existed has none and contributes no rows; an empty one (no
850
+ # kept records) is skipped here rather than handed to the reader.
851
+ event_files = [
852
+ path
853
+ for path in (
854
+ stored_artifact(sessions_dir / sid, SESSION_EVENTS_JSONL) for sid in json_ids
855
+ )
856
+ if path is not None and decoded_size(path) > 0
857
+ ]
858
+ if event_files:
859
+ _register_events_json(con, event_files)
860
+ events_branches.append(_events_json_select())
861
+ else:
862
+ con.execute(f"DROP TABLE IF EXISTS {_RAW_SESSION_EVENTS_JSON_TABLE};")
863
+
864
+ for table, branches, columns in (
865
+ (_RAW_TRAJECTORIES_TABLE, trajectory_branches, _TRAJECTORY_RELATION_COLUMNS),
866
+ (_RAW_STEPS_TABLE, steps_branches, VIEW_SCHEMA["steps"]),
867
+ (_RAW_TOOL_CALLS_TABLE, calls_branches, VIEW_SCHEMA["tool_calls"]),
868
+ (_RAW_TOOL_RESULTS_TABLE, results_branches, VIEW_SCHEMA["tool_results"]),
869
+ (_RAW_SESSION_EVENTS_TABLE, events_branches, VIEW_SCHEMA["session_events"]),
870
+ ):
871
+ # Table names are module constants; every branch is built above
872
+ # from catalog constants over module-constant readers.
873
+ body = _union_or_empty(branches, columns)
874
+ if table == _RAW_TRAJECTORIES_TABLE:
875
+ body = _with_corpus(body, corpus_root)
876
+ con.execute(f"CREATE OR REPLACE VIEW {table} AS\n{body};") # nosec B608 - constants only
877
+ logger.debug("Registered {} ({} source branch(es))", table, len(branches))
878
+
879
+ # edges.jsonl is one line per RAW record -> newline_delimited.
880
+ # The record's own `source_file` (the raw transcript path) is kept;
881
+ # the edges.jsonl path itself is aliased to `edges_path`, without the
882
+ # .zst suffix a compressed file carries. A list of each complete
883
+ # session's stored file rather than a glob: a glob per spelling would
884
+ # fail on the one with no match, and would read a session caught
885
+ # mid-slim twice.
886
+ meta_ids = [
887
+ str(row[0])
888
+ for row in con.execute(
889
+ f"SELECT session_id_path FROM {_RAW_META_TABLE} ORDER BY 1" # noqa: S608 # nosec B608 - module constant
890
+ ).fetchall()
891
+ ]
892
+ edge_paths = [
893
+ path
894
+ for path in (stored_artifact(sessions_dir / sid, EDGES_JSONL) for sid in meta_ids)
895
+ if path is not None
896
+ ]
897
+ edges_source: list[str] = [str(path) for path in edge_paths]
898
+ edges_bound = _object_size_bound(edge_paths)
899
+ if wanted is None:
900
+ loss_source: str | list[str] = str(sessions_dir / "*" / "loss_report.json")
901
+ else:
902
+ loss_source = _existing(sessions_dir / sid / "loss_report.json" for sid in wanted_ok)
903
+ if edges_source:
904
+ con.execute(
905
+ f"""
906
+ CREATE OR REPLACE TEMP TABLE {_RAW_EDGES_TABLE} AS
907
+ SELECT * FROM (
908
+ SELECT *,
909
+ regexp_replace(filename, '\\.zst$', '') AS edges_path,
910
+ regexp_extract(filename, '/sessions/([^/]+)/edges\\.jsonl(\\.zst)?$', 1)
911
+ AS session_id_path
912
+ FROM read_json(
913
+ ?,
914
+ format='newline_delimited',
915
+ filename=true,
916
+ columns={{{_render_columns_clause(_EDGE_COLUMNS)}}},
917
+ maximum_object_size={int(edges_bound)}
918
+ )
919
+ ) WHERE {meta_gate};
920
+ """, # noqa: S608 # nosec B608 - file list is a bound parameter; table/columns/gate are constants; the bound is an int
921
+ [edges_source],
922
+ )
923
+ else:
924
+ _typed_empty_table(
925
+ con, _RAW_EDGES_TABLE, _json_reader_columns(_EDGE_COLUMNS, "edges_path")
926
+ )
927
+ logger.debug("Registered {}", _RAW_EDGES_TABLE)
928
+
929
+ if loss_source:
930
+ con.execute(
931
+ f"""
932
+ CREATE OR REPLACE TEMP TABLE {_RAW_LOSS_REPORTS_TABLE} AS
933
+ SELECT * FROM (
934
+ SELECT *,
935
+ filename AS report_path,
936
+ regexp_extract(filename, '/sessions/([^/]+)/loss_report\\.json$', 1)
937
+ AS session_id_path
938
+ FROM read_json(
939
+ ?,
940
+ format='auto',
941
+ filename=true,
942
+ columns={{{_render_columns_clause(_LOSS_REPORT_COLUMNS)}}}
943
+ )
944
+ ) WHERE {meta_gate};
945
+ """, # noqa: S608 # nosec B608 - glob is a bound parameter; table/columns/gate are constants
946
+ [loss_source],
947
+ )
948
+ else:
949
+ _typed_empty_table(
950
+ con,
951
+ _RAW_LOSS_REPORTS_TABLE,
952
+ _json_reader_columns(_LOSS_REPORT_COLUMNS, "report_path"),
953
+ )
954
+ logger.debug("Registered {}", _RAW_LOSS_REPORTS_TABLE)
955
+ except Exception:
956
+ # register-or-fail-loud — any DuckDB error must surface to the caller.
957
+ logger.exception("Failed to register raw readers over {}", corpus_root)
958
+ raise
959
+ return RawSources(
960
+ columnar_session_ids=tuple(columnar_ids),
961
+ json_session_ids=tuple(json_ids),
962
+ lazy_read_paths=tuple(lazy_paths),
963
+ rejected_session_ids=rejected,
964
+ staged_session_ids=tuple(sorted(staged)),
965
+ )
966
+
967
+
968
+ def _with_corpus(body: SqlFragment, corpus_root: Path) -> SqlFragment:
969
+ """Append the ``corpus`` column (the corpus directory's name) to the trajectory rows.
970
+
971
+ ``CREATE VIEW`` cannot be prepared, so the name enters the statement
972
+ through :func:`sql_literal`, as ``ATTACH`` paths do.
973
+ """
974
+ return SqlFragment(
975
+ f"SELECT *, {sql_literal(corpus_root.name)} AS {CORPUS_COLUMN} FROM (\n{body}\n)" # noqa: S608 # nosec B608 - the constant-built union plus a sql_literal
976
+ )
977
+
978
+
979
+ # ---------------------------------------------------------------------------
980
+ # Views
981
+ # ---------------------------------------------------------------------------
982
+
983
+
984
+ def register_views(con: duckdb.DuckDBPyConnection) -> None:
985
+ """Create the business-level views on top of the raw readers.
986
+
987
+ Must be called after :func:`register_raw`. Creates, in dependency order:
988
+ ``steps``, ``sessions``, ``messages``, ``tool_calls``, ``tool_results``,
989
+ ``todo_events``, ``todo_state_current``, ``subagent_spawns``,
990
+ ``task_creations``, ``task_updates``, ``tasks_state_current``,
991
+ ``skill_invocations``, ``skill_usage``, ``subagent_steps``,
992
+ ``loss_reports``, ``session_events``.
993
+
994
+ Parameters
995
+ ----------
996
+ con
997
+ Open DuckDB connection with raw readers already registered.
998
+
999
+ Raises
1000
+ ------
1001
+ duckdb.Error
1002
+ If any view DDL fails. Logged via ``logger.exception`` before
1003
+ re-raise.
1004
+ """
1005
+ try:
1006
+ # One row per ATIF step, and the per-turn surface the contract names
1007
+ # ("messages-parity view name: steps"). The column expressions live
1008
+ # in :mod:`atif_duck.infrastructure.projections` (shared with the
1009
+ # columnar producer, so both sources agree byte for byte) and are
1010
+ # applied by register_raw when it builds ``v_raw_steps``; this view
1011
+ # pins the catalog's column order over that relation.
1012
+ con.execute(
1013
+ f"CREATE OR REPLACE VIEW steps AS SELECT {_catalog_columns('steps')} " # noqa: S608 # nosec B608 - catalog columns over a module-constant table
1014
+ f"FROM {_RAW_STEPS_TABLE};"
1015
+ )
1016
+ logger.debug("Registered view: steps")
1017
+
1018
+ # One row per materialized session. ``agent`` / ``agent_version`` come
1019
+ # from the trajectory's own ``agent`` struct, which harbor stamps per
1020
+ # adapter — ``claude-code`` with the Claude Code version, ``codex`` with
1021
+ # the codex-cli version. That is the ONE place the producing agent is
1022
+ # recorded inside an artifact, so every agent-aware query reads this
1023
+ # column rather than inferring the agent from a corpus path.
1024
+ #
1025
+ # cwd / git_branch come from the harbor adapter's ``agent.extra``; the
1026
+ # key differs per agent and both are read here because ``agent.extra``
1027
+ # is free-form: Claude Code writes LISTS in first-seen order (``cwds`` /
1028
+ # ``git_branches``) while Codex writes one ``cwd`` string and a ``git``
1029
+ # struct. A Claude Code session spanning two cwds reports the one it
1030
+ # started in; the later ones stay in ``agent.extra``.
1031
+ #
1032
+ # ``corpus`` names the corpus the session belongs to (the corpus
1033
+ # directory's name). Scoped to one corpus it is one value; over the
1034
+ # lake with every corpus attached it is how a query tells them apart.
1035
+ con.execute(
1036
+ """
1037
+ CREATE OR REPLACE VIEW sessions AS
1038
+ SELECT
1039
+ t.session_id_path AS session_id,
1040
+ t.agent.name AS agent,
1041
+ t.agent.version AS agent_version,
1042
+ coalesce(
1043
+ json_extract_string(t.agent.extra, '$.cwds[0]'),
1044
+ json_extract_string(t.agent.extra, '$.cwd')
1045
+ ) AS cwd,
1046
+ coalesce(
1047
+ json_extract_string(t.agent.extra, '$.git_branches[0]'),
1048
+ json_extract_string(t.agent.extra, '$.git.branch')
1049
+ ) AS git_branch,
1050
+ s.started_at,
1051
+ s.ended_at,
1052
+ s.agent_steps,
1053
+ s.step_count,
1054
+ t.agent.model_name AS model_name,
1055
+ t.final_metrics.total_cost_usd AS total_cost_usd,
1056
+ CAST(json_extract(t.final_metrics.extra, '$.reported_cost_usd') AS DOUBLE)
1057
+ AS reported_cost_usd,
1058
+ t.trajectory_path,
1059
+ t.corpus
1060
+ FROM v_raw_trajectories t
1061
+ LEFT JOIN (
1062
+ SELECT
1063
+ session_id,
1064
+ min(ts) AS started_at,
1065
+ max(ts) AS ended_at,
1066
+ count(*) FILTER (WHERE source = 'agent') AS agent_steps,
1067
+ count(*) AS step_count
1068
+ FROM steps
1069
+ GROUP BY session_id
1070
+ ) s ON s.session_id = t.session_id_path;
1071
+ """
1072
+ )
1073
+ logger.debug("Registered view: sessions")
1074
+
1075
+ # COMPAT view: uuid-keyed raw-record identity reconstructed from
1076
+ # edges.jsonl. The trajectory cannot provide this (fidelity gap 7:
1077
+ # harbor never lands the event uuid in Step.extra), so parent-chain
1078
+ # walks and per-record counts run against the census-derived edges.
1079
+ # ``type`` carries the raw record role (user/assistant/system/...);
1080
+ # a separate role column would be redundant with it.
1081
+ con.execute(
1082
+ """
1083
+ CREATE OR REPLACE VIEW messages AS
1084
+ SELECT
1085
+ uuid,
1086
+ parent_uuid,
1087
+ session_id_path AS session_id,
1088
+ ts,
1089
+ type,
1090
+ is_sidechain,
1091
+ coalesce(is_compact_summary, false) AS is_compact_summary,
1092
+ message_id,
1093
+ source_file
1094
+ FROM v_raw_edges;
1095
+ """
1096
+ )
1097
+ logger.debug("Registered view: messages")
1098
+
1099
+ # One row per ToolCall in any step. The columns are renamed off the
1100
+ # ATIF field names: function_name -> tool_name, tool_call_id ->
1101
+ # tool_use_id, arguments -> tool_input. ``tool_input`` is an already
1102
+ # parsed JSON value, NOT a JSON string, so a query must not
1103
+ # json_extract it twice. Expressions: projections.CALL_COLUMNS.
1104
+ con.execute(
1105
+ f"CREATE OR REPLACE VIEW tool_calls AS SELECT {_catalog_columns('tool_calls')} " # noqa: S608 # nosec B608 - catalog columns over a module-constant table
1106
+ f"FROM {_RAW_TOOL_CALLS_TABLE};"
1107
+ )
1108
+ logger.debug("Registered view: tool_calls")
1109
+
1110
+ # One row per ObservationResult. ``source_call_id`` is ATIF's join
1111
+ # key back to the tool_calls array, surfaced under the SAME name
1112
+ # ``tool_use_id`` that ``tool_calls`` exposes, so the two views join
1113
+ # with `USING (tool_use_id)` and neither side needs a rename.
1114
+ # Expressions: projections.RESULT_COLUMNS.
1115
+ con.execute(
1116
+ f"CREATE OR REPLACE VIEW tool_results AS SELECT {_catalog_columns('tool_results')} " # noqa: S608 # nosec B608 - catalog columns over a module-constant table
1117
+ f"FROM {_RAW_TOOL_RESULTS_TABLE};"
1118
+ )
1119
+ logger.debug("Registered view: tool_results")
1120
+
1121
+ # DuckDB's UNNEST requires a LIST, and the ``$.todos[*]`` wildcard
1122
+ # path yields JSON[] that UNNEST accepts natively — a bare
1123
+ # ``$.todos`` yields a scalar JSON value that it rejects.
1124
+ # ``step_id`` is the snapshot tie-break: steps are unique and ordered
1125
+ # per session, so no uuid column is needed to sequence the snapshots.
1126
+ con.execute(
1127
+ """
1128
+ CREATE OR REPLACE VIEW todo_events AS
1129
+ SELECT
1130
+ tc.session_id,
1131
+ tc.ts AS written_at,
1132
+ tc.step_id,
1133
+ json_extract_string(todo, '$.content') AS subject,
1134
+ json_extract_string(todo, '$.status') AS status,
1135
+ json_extract_string(todo, '$.activeForm') AS active_form,
1136
+ row_number() OVER (
1137
+ PARTITION BY tc.session_id
1138
+ ORDER BY tc.ts, tc.step_id
1139
+ ) AS snapshot_ix
1140
+ FROM tool_calls tc,
1141
+ UNNEST(json_extract(tc.tool_input, '$.todos[*]')) AS t(todo)
1142
+ WHERE tc.tool_name = 'TodoWrite';
1143
+ """
1144
+ )
1145
+ logger.debug("Registered view: todo_events")
1146
+
1147
+ # Latest snapshot wins per (session_id, subject) — TodoWrite always
1148
+ # rewrites the full list, so the newest row is the current state.
1149
+ con.execute(
1150
+ """
1151
+ CREATE OR REPLACE VIEW todo_state_current AS
1152
+ SELECT session_id, subject, status, active_form, written_at
1153
+ FROM (
1154
+ SELECT *,
1155
+ row_number() OVER (
1156
+ PARTITION BY session_id, subject
1157
+ ORDER BY snapshot_ix DESC
1158
+ ) AS rn
1159
+ FROM todo_events
1160
+ )
1161
+ WHERE rn = 1;
1162
+ """
1163
+ )
1164
+ logger.debug("Registered view: todo_state_current")
1165
+
1166
+ # Subagent launchers: ``Task`` (pre-v2.1.63) and ``Agent`` (v2.1.63+).
1167
+ # Input shape: {subagent_type, description, prompt, run_in_background?}.
1168
+ con.execute(
1169
+ """
1170
+ CREATE OR REPLACE VIEW subagent_spawns AS
1171
+ SELECT
1172
+ session_id,
1173
+ ts AS spawned_at,
1174
+ step_id,
1175
+ tool_use_id,
1176
+ tool_name AS spawn_tool,
1177
+ json_extract_string(tool_input, '$.subagent_type') AS subagent_type,
1178
+ json_extract_string(tool_input, '$.description') AS description,
1179
+ json_extract_string(tool_input, '$.prompt') AS prompt,
1180
+ json_extract_string(tool_input, '$.run_in_background') AS run_in_background
1181
+ FROM tool_calls
1182
+ WHERE tool_name IN ('Task', 'Agent');
1183
+ """
1184
+ )
1185
+ logger.debug("Registered view: subagent_spawns")
1186
+
1187
+ # Persistent task creation: ``TaskCreate`` (Claude Code v2.1.16+) and
1188
+ # the SDK-py mirror ``mcp__tasks__task_create``. Distinct from
1189
+ # subagent_spawns — no subagent_type / prompt fields.
1190
+ con.execute(
1191
+ """
1192
+ CREATE OR REPLACE VIEW task_creations AS
1193
+ SELECT
1194
+ session_id,
1195
+ ts AS created_at,
1196
+ step_id,
1197
+ tool_use_id,
1198
+ tool_name AS create_tool,
1199
+ json_extract_string(tool_input, '$.subject') AS subject,
1200
+ json_extract_string(tool_input, '$.description') AS description,
1201
+ json_extract_string(tool_input, '$.activeForm') AS active_form,
1202
+ json_extract(tool_input, '$.metadata') AS metadata
1203
+ FROM tool_calls
1204
+ WHERE tool_name IN ('TaskCreate', 'mcp__tasks__task_create');
1205
+ """
1206
+ )
1207
+ logger.debug("Registered view: task_creations")
1208
+
1209
+ # Task lifecycle updates. The native tool spells the key ``taskId``
1210
+ # (camel) and the mcp variant spells it ``id``, so both are extracted
1211
+ # and COALESCEd: reading only one silently drops half the updates.
1212
+ con.execute(
1213
+ """
1214
+ CREATE OR REPLACE VIEW task_updates AS
1215
+ SELECT
1216
+ session_id,
1217
+ ts AS updated_at,
1218
+ step_id,
1219
+ tool_use_id,
1220
+ tool_name AS update_tool,
1221
+ COALESCE(
1222
+ json_extract_string(tool_input, '$.taskId'),
1223
+ json_extract_string(tool_input, '$.id')
1224
+ ) AS task_id,
1225
+ json_extract_string(tool_input, '$.status') AS status,
1226
+ json_extract(tool_input, '$.addBlockedBy') AS add_blocked_by,
1227
+ json_extract_string(tool_input, '$.owner') AS owner
1228
+ FROM tool_calls
1229
+ WHERE tool_name IN ('TaskUpdate', 'mcp__tasks__task_update');
1230
+ """
1231
+ )
1232
+ logger.debug("Registered view: task_updates")
1233
+
1234
+ # Latest status per (session_id, task_id). The runtime assigns the
1235
+ # task id and returns it only in the tool RESULT, so we recover it
1236
+ # from ``tool_results`` via the source_call_id join, parsing the
1237
+ # "Task #N" / {taskId} shapes the result text carries, and falling
1238
+ # back to per-session creation order when neither shape matches.
1239
+ #
1240
+ # NULLIF is required because DuckDB's regexp_extract returns '' (not
1241
+ # NULL) on no-match, and '' satisfies COALESCE — without it every later
1242
+ # fallback branch is unreachable and non-"Task #N" results all collapse
1243
+ # onto the same empty task_id, colliding in the latest_status join.
1244
+ #
1245
+ # The creation-order fallback orders by ``created_at`` and then by
1246
+ # ``step_id, tool_use_id``: several TaskCreate calls in one step share a
1247
+ # timestamp, and a window with a tied ORDER BY numbers them in physical
1248
+ # scan order, which differs between the JSON reader (a hash join
1249
+ # reorders the unnested calls) and the parquet reader. The tiebreaker
1250
+ # makes the number a function of the data, not of the source format.
1251
+ con.execute(
1252
+ """
1253
+ CREATE OR REPLACE VIEW tasks_state_current AS
1254
+ WITH creates AS (
1255
+ SELECT
1256
+ tc.session_id,
1257
+ tc.created_at,
1258
+ tc.subject,
1259
+ tc.active_form,
1260
+ tc.tool_use_id,
1261
+ COALESCE(
1262
+ NULLIF(
1263
+ regexp_extract(
1264
+ CAST(tr.content AS VARCHAR), 'Task #(\\d+)', 1
1265
+ ),
1266
+ ''
1267
+ ),
1268
+ json_extract_string(tr.content, '$.taskId'),
1269
+ CAST(row_number() OVER (
1270
+ PARTITION BY tc.session_id
1271
+ ORDER BY tc.created_at, tc.step_id, tc.tool_use_id
1272
+ ) AS VARCHAR)
1273
+ ) AS task_id
1274
+ FROM task_creations tc
1275
+ LEFT JOIN tool_results tr USING (tool_use_id)
1276
+ ),
1277
+ latest_status AS (
1278
+ SELECT session_id, task_id, status, updated_at,
1279
+ row_number() OVER (
1280
+ PARTITION BY session_id, task_id
1281
+ ORDER BY updated_at DESC
1282
+ ) AS rn
1283
+ FROM task_updates
1284
+ WHERE task_id IS NOT NULL
1285
+ )
1286
+ SELECT
1287
+ c.session_id,
1288
+ c.task_id,
1289
+ c.subject,
1290
+ c.active_form,
1291
+ COALESCE(ls.status, 'pending') AS status,
1292
+ c.created_at,
1293
+ ls.updated_at AS last_updated_at
1294
+ FROM creates c
1295
+ LEFT JOIN latest_status ls
1296
+ ON ls.session_id = c.session_id
1297
+ AND ls.task_id = c.task_id
1298
+ AND ls.rn = 1;
1299
+ """
1300
+ )
1301
+ logger.debug("Registered view: tasks_state_current")
1302
+
1303
+ # Every Skill / slash-command invocation, unioned across both shapes:
1304
+ # * ``tool`` — the assistant invokes the built-in ``Skill`` tool with
1305
+ # ``arguments.skill = '<name>'``. Lives in ``tool_calls`` already.
1306
+ # * ``slash_command`` — the user types ``/<name>``, which Claude Code
1307
+ # serializes into the message text as
1308
+ # ``<command-name>/<name></command-name>`` (sometimes with
1309
+ # ``<command-args>``). ATIF flattens user content to one text
1310
+ # message per step, so ONE regex over ``steps.message`` covers both
1311
+ # raw serializations (list-typed content blocks and bare VARCHAR
1312
+ # content) without branching on the content type.
1313
+ # ``skill_id`` is the raw identifier, NOT a normalized name:
1314
+ # ``erpaval`` and ``personal-plugins:erpaval`` are distinct rows, so a
1315
+ # per-skill aggregate has to decide for itself whether to fold them.
1316
+ cmd_name_re = SqlFragment("<command-name>/([A-Za-z0-9_:.-]+)</command-name>")
1317
+ args_re = SqlFragment("<command-args>([^<]*)</command-args>")
1318
+ con.execute(
1319
+ f"""
1320
+ CREATE OR REPLACE VIEW skill_invocations AS
1321
+ SELECT
1322
+ tc.session_id,
1323
+ tc.ts,
1324
+ tc.step_id,
1325
+ 'tool' AS source,
1326
+ json_extract_string(tc.tool_input, '$.skill') AS skill_id,
1327
+ json_extract_string(tc.tool_input, '$.args') AS args,
1328
+ tc.tool_use_id
1329
+ FROM tool_calls tc
1330
+ WHERE tc.tool_name = 'Skill'
1331
+ AND json_extract_string(tc.tool_input, '$.skill') IS NOT NULL
1332
+ UNION ALL
1333
+ SELECT
1334
+ s.session_id,
1335
+ s.ts,
1336
+ s.step_id,
1337
+ 'slash_command' AS source,
1338
+ regexp_extract(s.message, '{cmd_name_re}', 1) AS skill_id,
1339
+ NULLIF(regexp_extract(s.message, '{args_re}', 1), '') AS args,
1340
+ NULL AS tool_use_id
1341
+ FROM steps s
1342
+ WHERE s.source = 'user'
1343
+ AND s.message LIKE '%<command-name>/%'
1344
+ AND regexp_extract(s.message, '{cmd_name_re}', 1) != '';
1345
+ """ # noqa: S608 # nosec B608 - the two interpolated regexes are local literals defined above
1346
+ )
1347
+ logger.debug("Registered view: skill_invocations")
1348
+
1349
+ # Labels derive from the skill_id string alone — there is no skills
1350
+ # catalog in the corpus to join against. The heuristic: a skill_id
1351
+ # WITHOUT a ':' has no plugin namespace and is labelled builtin
1352
+ # (`erpaval`, `review`, ...); a `plugin:skill` id splits into plugin +
1353
+ # skill_name and is not builtin. ``is_builtin`` therefore OVER-marks a
1354
+ # user-local non-plugin skill as builtin. That is tolerable only
1355
+ # because the rank/mix macros use the flag to damp slash-only noise;
1356
+ # a query that treats it as ground truth about provenance is wrong.
1357
+ con.execute(
1358
+ """
1359
+ CREATE OR REPLACE VIEW skill_usage AS
1360
+ SELECT
1361
+ si.session_id,
1362
+ si.ts,
1363
+ si.step_id,
1364
+ si.source,
1365
+ si.skill_id,
1366
+ si.args,
1367
+ si.tool_use_id,
1368
+ CASE WHEN strpos(si.skill_id, ':') > 0
1369
+ THEN substr(si.skill_id, strpos(si.skill_id, ':') + 1)
1370
+ ELSE si.skill_id
1371
+ END AS skill_name,
1372
+ CASE WHEN strpos(si.skill_id, ':') > 0
1373
+ THEN split_part(si.skill_id, ':', 1)
1374
+ ELSE NULL
1375
+ END AS plugin,
1376
+ strpos(si.skill_id, ':') = 0 AS is_builtin
1377
+ FROM skill_invocations si;
1378
+ """
1379
+ )
1380
+ logger.debug("Registered view: skill_usage")
1381
+
1382
+ # Harbor INLINES subagent transcripts into the flat step list
1383
+ # (fidelity gap 4), marking them only via ``extra.is_sidechain`` — so
1384
+ # the subagent surface is a filter over ``steps`` and there is no
1385
+ # separate subagent raw reader to read. A consequence worth knowing:
1386
+ # a session-level aggregate over ``steps`` already includes subagent
1387
+ # work unless it filters ``is_sidechain`` out.
1388
+ con.execute(
1389
+ """
1390
+ CREATE OR REPLACE VIEW subagent_steps AS
1391
+ SELECT *
1392
+ FROM steps
1393
+ WHERE is_sidechain;
1394
+ """
1395
+ )
1396
+ logger.debug("Registered view: subagent_steps")
1397
+
1398
+ # One row per subagent, declared by the converter in
1399
+ # ``trajectory.extra.subagents`` (from each ``agent-*.meta.json``
1400
+ # sidecar, else from the spawning call's own result). The converter
1401
+ # lists every agent id a sidechain step carries, so joining the step
1402
+ # activity onto the declared list loses no agent. ``description``
1403
+ # falls back to the spawning call's own input when the sidecar has
1404
+ # none.
1405
+ con.execute(
1406
+ """
1407
+ CREATE OR REPLACE VIEW subagents AS
1408
+ WITH declared AS (
1409
+ SELECT
1410
+ t.session_id_path AS session_id,
1411
+ json_extract_string(a, '$.agent_id') AS agent_id,
1412
+ json_extract_string(a, '$.agent_type') AS agent_type,
1413
+ json_extract_string(a, '$.description') AS description,
1414
+ json_extract_string(a, '$.parent_tool_call_id') AS parent_tool_call_id,
1415
+ json_extract_string(a, '$.link_source') AS link_source,
1416
+ json_extract(a, '$.spawn_depth')::BIGINT AS spawn_depth,
1417
+ json_extract_string(a, '$.parent_agent_id') AS parent_agent_id
1418
+ FROM v_raw_trajectories t,
1419
+ UNNEST(json_extract(t.extra, '$.subagents[*]')) AS s(a)
1420
+ WHERE json_extract(t.extra, '$.subagents') IS NOT NULL
1421
+ ),
1422
+ activity AS (
1423
+ SELECT session_id, agent_id,
1424
+ min(ts) AS first_ts,
1425
+ max(ts) AS last_ts,
1426
+ count(*) AS step_count
1427
+ FROM steps
1428
+ WHERE agent_id IS NOT NULL
1429
+ GROUP BY session_id, agent_id
1430
+ )
1431
+ SELECT
1432
+ d.session_id,
1433
+ d.agent_id,
1434
+ coalesce(d.agent_type, sp.subagent_type) AS agent_type,
1435
+ coalesce(d.description, sp.description) AS description,
1436
+ d.parent_tool_call_id,
1437
+ sp.step_id AS parent_step_id,
1438
+ d.link_source,
1439
+ d.spawn_depth,
1440
+ d.parent_agent_id,
1441
+ a.first_ts,
1442
+ a.last_ts,
1443
+ coalesce(a.step_count, 0) AS step_count
1444
+ FROM declared d
1445
+ LEFT JOIN activity a
1446
+ ON a.session_id = d.session_id AND a.agent_id = d.agent_id
1447
+ LEFT JOIN subagent_spawns sp
1448
+ ON sp.session_id = d.session_id AND sp.tool_use_id = d.parent_tool_call_id;
1449
+ """
1450
+ )
1451
+ logger.debug("Registered view: subagents")
1452
+
1453
+ # One row per attachment the converter moved to the blob store:
1454
+ # tool-result images from ``tool_results.images``, pasted user images
1455
+ # from ``steps.images``. ``blob_path`` is relative to the corpus root
1456
+ # (``blobs/sha256/<ab>/<sha256>.<ext>``), the layout atif-corpus writes.
1457
+ con.execute(
1458
+ """
1459
+ CREATE OR REPLACE VIEW images AS
1460
+ WITH listed AS (
1461
+ SELECT session_id, step_id, ts, 'tool_result' AS origin, tool_use_id,
1462
+ img
1463
+ FROM tool_results, UNNEST(json_extract(images, '$[*]')) AS i(img)
1464
+ WHERE images IS NOT NULL
1465
+ UNION ALL
1466
+ SELECT session_id, step_id, ts, 'user_message' AS origin,
1467
+ CAST(NULL AS VARCHAR) AS tool_use_id, img
1468
+ FROM steps, UNNEST(json_extract(images, '$[*]')) AS i(img)
1469
+ WHERE images IS NOT NULL
1470
+ )
1471
+ SELECT
1472
+ session_id,
1473
+ step_id,
1474
+ ts,
1475
+ origin,
1476
+ tool_use_id,
1477
+ json_extract_string(img, '$.sha256') AS sha256,
1478
+ json_extract_string(img, '$.media_type') AS media_type,
1479
+ json_extract(img, '$.bytes')::BIGINT AS size_bytes,
1480
+ json_extract(img, '$.width')::BIGINT AS width,
1481
+ json_extract(img, '$.height')::BIGINT AS height,
1482
+ 'blobs/sha256/' || substr(json_extract_string(img, '$.sha256'), 1, 2) || '/'
1483
+ || json_extract_string(img, '$.sha256') || '.'
1484
+ || json_extract_string(img, '$.extension') AS blob_path
1485
+ FROM listed;
1486
+ """
1487
+ )
1488
+ logger.debug("Registered view: images")
1489
+
1490
+ # Per-session conversion loss accounting from atif-converter's
1491
+ # LossReport: record counts by raw type plus the fidelity gaps this
1492
+ # session actually exhibits. Any count taken from the other views is
1493
+ # only as complete as this view says the conversion was, so a
1494
+ # discrepancy between a raw record count and a step count is explained
1495
+ # here rather than being a defect in the view that reports it.
1496
+ con.execute(
1497
+ """
1498
+ CREATE OR REPLACE VIEW loss_reports AS
1499
+ SELECT
1500
+ session_id_path AS session_id,
1501
+ records_total,
1502
+ records_converted,
1503
+ records_captured,
1504
+ records_dropped,
1505
+ gaps_observed,
1506
+ record_counts,
1507
+ subagent_files_found,
1508
+ subagent_files_convertible,
1509
+ workflow_subagent_files_found,
1510
+ report_path
1511
+ FROM v_raw_loss_reports;
1512
+ """
1513
+ )
1514
+ logger.debug("Registered view: loss_reports")
1515
+
1516
+ # One row per kept non-message record (hooks, injected context, API
1517
+ # errors, compaction boundaries, cost-state, mode changes; Codex
1518
+ # compactions and aborted turns), from session_events.jsonl or its
1519
+ # parquet. Order within a session by ``seq``: cost-state and mode rows
1520
+ # have no timestamp. ``parent_uuid`` is the raw record's parentUuid, so
1521
+ # a row joins the step that owns that uuid in ``steps.source_uuids``.
1522
+ con.execute(
1523
+ f"CREATE OR REPLACE VIEW session_events AS SELECT {_catalog_columns('session_events')} " # noqa: S608 # nosec B608 - catalog columns over a module-constant table
1524
+ f"FROM {_RAW_SESSION_EVENTS_TABLE};"
1525
+ )
1526
+ logger.debug("Registered view: session_events")
1527
+ except Exception:
1528
+ # register-or-fail-loud — any DuckDB error must surface to the caller.
1529
+ logger.exception("Failed to register derived views")
1530
+ raise
1531
+
1532
+
1533
+ # ---------------------------------------------------------------------------
1534
+ # VSS
1535
+ # ---------------------------------------------------------------------------
1536
+
1537
+
1538
+ def _lance_table_present(con: duckdb.DuckDBPyConnection) -> bool:
1539
+ """True iff the ATTACHed ``lance_store`` catalog exposes ``embeddings``.
1540
+
1541
+ Probed via ``duckdb_tables()`` rather than a speculative SELECT: a Lance
1542
+ directory that exists with metadata but no embeddings table (legitimate
1543
+ intermediate state after a ``lancedb.connect`` that never created a
1544
+ table) ATTACHes cleanly but then blows up at view-bind time with a
1545
+ catalog error. The right gate is "is the table actually there?".
1546
+ """
1547
+ row = con.execute(
1548
+ """
1549
+ SELECT count(*)
1550
+ FROM duckdb_tables()
1551
+ WHERE database_name = 'lance_store' AND table_name = 'embeddings'
1552
+ """
1553
+ ).fetchone()
1554
+ return row is not None and int(row[0]) > 0
1555
+
1556
+
1557
+ def lance_extension_installed(con: duckdb.DuckDBPyConnection) -> bool:
1558
+ """True when the lance extension is present in this connection's extension directory.
1559
+
1560
+ Read from ``duckdb_extensions()``, which lists the local directory and
1561
+ never reaches the network, so the answer costs a few milliseconds and no
1562
+ download.
1563
+ """
1564
+ row = con.execute(
1565
+ "SELECT installed FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION]
1566
+ ).fetchone()
1567
+ return row is not None and bool(row[0])
1568
+
1569
+
1570
+ def load_lance_extension(con: duckdb.DuckDBPyConnection) -> bool:
1571
+ """``LOAD`` the lance extension when it is installed; never install it.
1572
+
1573
+ Returns ``False`` when the extension is absent, and executes no
1574
+ ``INSTALL`` in that case whatever ``autoinstall_known_extensions`` says,
1575
+ because the installed check runs first. Registration at query time must
1576
+ never reach the network: the extension is 242 MB, fetched by name from
1577
+ the default repository, and ``query`` runs unattended.
1578
+ """
1579
+ if not lance_extension_installed(con):
1580
+ return False
1581
+ con.execute(f"LOAD {LANCE_EXTENSION};")
1582
+ return True
1583
+
1584
+
1585
+ def install_lance_extension(con: duckdb.DuckDBPyConnection) -> str:
1586
+ """``INSTALL`` then ``LOAD`` the lance extension; return its install path.
1587
+
1588
+ The one place atif-duck reaches the network, so it belongs to explicit
1589
+ commands only (``atif-sql embed --install-extension``, or a real embed
1590
+ run, which reaches Bedrock anyway). A query-time registration calls
1591
+ :func:`load_lance_extension` instead.
1592
+ """
1593
+ con.execute(f"INSTALL {LANCE_EXTENSION};")
1594
+ con.execute(f"LOAD {LANCE_EXTENSION};")
1595
+ row = con.execute(
1596
+ "SELECT install_path FROM duckdb_extensions() WHERE extension_name = ?", [LANCE_EXTENSION]
1597
+ ).fetchone()
1598
+ return str(row[0]) if row and row[0] is not None else ""
1599
+
1600
+
1601
+ def register_vss(
1602
+ con: duckdb.DuckDBPyConnection,
1603
+ *,
1604
+ lance_uri: Path,
1605
+ expected_model: str | None = None,
1606
+ expected_dim: int | None = None,
1607
+ dim: int = 1024,
1608
+ ) -> bool:
1609
+ """Bind ``message_embeddings`` over a LanceDB local dataset.
1610
+
1611
+ LanceDB stores embeddings + its IVF_HNSW_SQ index in one place (written
1612
+ by atif-embed's backfill); reads come back through DuckDB via the lance
1613
+ core extension (``LOAD lance; ATTACH (TYPE LANCE)``). The extension is
1614
+ only LOADed, never INSTALLed, here: :func:`load_lance_extension` checks
1615
+ the local extension directory first and a missing extension degrades to
1616
+ the empty fallback table with a warning, so registration never reaches
1617
+ the network. Installing is an explicit act
1618
+ (:func:`install_lance_extension`, behind ``atif-sql embed``). When no
1619
+ store directory exists the extension is not loaded at all, which is the
1620
+ state of every corpus that has not run ``embed``.
1621
+ The store probe runs through DuckDB itself. atif-duck declares no lancedb
1622
+ dependency: lancedb belongs to atif-embed, which writes the store, and the
1623
+ independence contract forbids an import edge between the two packages — so
1624
+ the lance extension is the only thing atif-duck may read the store with.
1625
+
1626
+ Parameters
1627
+ ----------
1628
+ con
1629
+ Open DuckDB connection.
1630
+ lance_uri
1631
+ Local LanceDB dataset directory.
1632
+ expected_model
1633
+ When supplied (by the composition root), the active embedder's
1634
+ ``model_id``. Read against the store's stamped ``model`` column via
1635
+ the fail-loud provider guard: a mismatch raises
1636
+ :class:`atif_duck.domain.embedding_guard.EmbeddingProviderMismatch`
1637
+ rather than letting a cross-provider query silently return garbage
1638
+ cosine scores (guard-before-bind).
1639
+ expected_dim
1640
+ The active embedder's dimension, checked alongside
1641
+ ``expected_model``. ``None`` trusts ``model_id`` alone.
1642
+ dim
1643
+ Fixed-length embedding dimension for the empty-fallback table. When a
1644
+ populated Lance store is present its own stamped ``dim`` drives the
1645
+ ``CAST(embedding AS FLOAT[dim])`` view instead, so a store written by
1646
+ a different-width provider binds correctly regardless of this
1647
+ argument.
1648
+
1649
+ Returns
1650
+ -------
1651
+ bool
1652
+ ``True`` when the Lance table is reachable through the
1653
+ ``message_embeddings`` view; ``False`` when no embeddings exist yet,
1654
+ or the store exists but the lance extension is not installed (the
1655
+ name is created as an empty TABLE with the right schema so a
1656
+ downstream ``CREATE MACRO semantic_search`` can still bind).
1657
+ """
1658
+ dim_i = int(dim)
1659
+
1660
+ import duckdb as _duckdb
1661
+
1662
+ attached = False
1663
+ if lance_uri.is_dir() and not load_lance_extension(con):
1664
+ logger.warning(
1665
+ "Lance store at {} cannot be read: the {} extension is not installed, so "
1666
+ "message_embeddings binds empty. Run `atif-sql embed --install-extension` "
1667
+ "(a one-time download) to enable vector search.",
1668
+ lance_uri,
1669
+ LANCE_EXTENSION,
1670
+ )
1671
+ elif lance_uri.is_dir():
1672
+ try:
1673
+ # ATTACH is one of the statement kinds DuckDB will not prepare, so
1674
+ # this is the one corpus path that still enters statement text; it
1675
+ # does so only through sql_literal, and the boundary test registers
1676
+ # a store under a quoted path to keep it that way.
1677
+ con.execute(
1678
+ f"ATTACH IF NOT EXISTS {sql_literal(str(lance_uri))} AS lance_store (TYPE LANCE);"
1679
+ )
1680
+ attached = True
1681
+ except _duckdb.Error:
1682
+ # register-or-fail-loud applies to DDL over a healthy store; a
1683
+ # directory the lance extension cannot ATTACH is equivalent to
1684
+ # "no store yet" and falls through to the empty-fallback table.
1685
+ logger.exception("Lance ATTACH failed for {}; treating as empty store", lance_uri)
1686
+
1687
+ if not attached or not _lance_table_present(con):
1688
+ logger.warning(
1689
+ "No Lance embeddings table at {}; creating empty message_embeddings "
1690
+ "table so semantic_search binds. Run `atif-sql embed --all --no-dry-run` "
1691
+ "to backfill.",
1692
+ lance_uri,
1693
+ )
1694
+ con.execute(
1695
+ f"""
1696
+ CREATE OR REPLACE TABLE message_embeddings (
1697
+ uuid VARCHAR PRIMARY KEY,
1698
+ model VARCHAR,
1699
+ dim INTEGER,
1700
+ embedding FLOAT[{dim_i}],
1701
+ embedded_at TIMESTAMPTZ
1702
+ );
1703
+ """ # nosec B608 - the only interpolation is the width, through int() coercion
1704
+ )
1705
+ return False
1706
+
1707
+ # Fail-loud provider/dimension guard: the store stamps ``(model, dim)``
1708
+ # on every row; read it back and refuse to bind a view over vectors
1709
+ # written by a different provider (even at matching width, cross-model
1710
+ # vectors live in incompatible spaces and produce numerically valid but
1711
+ # garbage cosine scores). The stored width also drives the CAST below.
1712
+ row = con.execute("SELECT model, dim FROM lance_store.main.embeddings LIMIT 1;").fetchone()
1713
+ if row is not None:
1714
+ stored_model, stored_dim = str(row[0]), int(row[1])
1715
+ if expected_model is not None:
1716
+ ensure_store_matches(
1717
+ stored_model=stored_model,
1718
+ stored_dim=stored_dim,
1719
+ expected_model=expected_model,
1720
+ expected_dim=expected_dim,
1721
+ )
1722
+ dim_i = stored_dim
1723
+
1724
+ # Project the embeddings table as a top-level view named
1725
+ # ``message_embeddings``, casting the embedding column to FLOAT[dim] so
1726
+ # the fixed-size ARRAY type is what downstream vector functions see.
1727
+ con.execute(
1728
+ f"""
1729
+ CREATE OR REPLACE VIEW message_embeddings AS
1730
+ SELECT
1731
+ uuid,
1732
+ model,
1733
+ dim,
1734
+ CAST(embedding AS FLOAT[{dim_i}]) AS embedding,
1735
+ embedded_at
1736
+ FROM lance_store.main.embeddings;
1737
+ """ # noqa: S608 # nosec B608 - the only interpolation is the width, through int() coercion
1738
+ )
1739
+ count_row = con.execute("SELECT count(*) FROM message_embeddings;").fetchone()
1740
+ count = int(count_row[0]) if count_row else 0
1741
+ logger.debug("Bound message_embeddings over Lance ({} rows, dim={})", count, dim_i)
1742
+ return True
1743
+
1744
+
1745
+ def _attached_lance_table_present(con: duckdb.DuckDBPyConnection, alias: str) -> bool:
1746
+ """True iff the Lance catalog attached as ``alias`` exposes ``embeddings``."""
1747
+ row = con.execute(
1748
+ "SELECT count(*) FROM duckdb_tables() WHERE database_name = ? AND table_name = 'embeddings'",
1749
+ [alias],
1750
+ ).fetchone()
1751
+ return row is not None and int(row[0]) > 0
1752
+
1753
+
1754
+ def register_vss_stores(
1755
+ con: duckdb.DuckDBPyConnection,
1756
+ *,
1757
+ lance_uris: Sequence[Path],
1758
+ expected_model: str | None = None,
1759
+ expected_dim: int | None = None,
1760
+ ) -> bool:
1761
+ """Bind ``message_embeddings`` over the union of several Lance stores (one per corpus).
1762
+
1763
+ ``query --all-corpora`` and ``search --all-corpora`` read every corpus
1764
+ the lake holds, and each corpus keeps its own store, so the view is the
1765
+ ``UNION ALL`` of every store that exists. Its shape is
1766
+ :func:`register_vss`'s, so ``semantic_search`` binds unchanged. Every
1767
+ store passes the same provider guard, and they must agree on the vector
1768
+ width. With one store or none this is :func:`register_vss`.
1769
+ """
1770
+ if not lance_uris:
1771
+ msg = "register_vss_stores needs at least one store path"
1772
+ raise ValueError(msg)
1773
+ present = [uri for uri in lance_uris if uri.is_dir()]
1774
+ if len(present) <= 1 or not lance_extension_installed(con):
1775
+ # One store, none, or no extension to read them with: register_vss
1776
+ # covers each (a missing extension binds the empty table and warns).
1777
+ return register_vss(
1778
+ con,
1779
+ lance_uri=present[0] if present else lance_uris[0],
1780
+ expected_model=expected_model,
1781
+ expected_dim=expected_dim,
1782
+ )
1783
+ con.execute(f"LOAD {LANCE_EXTENSION};")
1784
+ widths: set[int] = set()
1785
+ selects: list[str] = []
1786
+ for index, uri in enumerate(present):
1787
+ alias = f"lance_store_{int(index)}"
1788
+ con.execute(f"ATTACH IF NOT EXISTS {sql_literal(str(uri))} AS {alias} (TYPE LANCE);")
1789
+ if not _attached_lance_table_present(con, alias):
1790
+ logger.info("No Lance embeddings table at {}; leaving it out", uri)
1791
+ continue
1792
+ row = con.execute(f"SELECT model, dim FROM {alias}.main.embeddings LIMIT 1;").fetchone() # noqa: S608 # nosec B608 - the alias is an int-derived name
1793
+ if row is None:
1794
+ continue
1795
+ stored_model, stored_dim = str(row[0]), int(row[1])
1796
+ if expected_model is not None:
1797
+ ensure_store_matches(
1798
+ stored_model=stored_model,
1799
+ stored_dim=stored_dim,
1800
+ expected_model=expected_model,
1801
+ expected_dim=expected_dim,
1802
+ )
1803
+ widths.add(stored_dim)
1804
+ selects.append(alias)
1805
+ if len(widths) > 1:
1806
+ msg = f"the corpora's Lance stores disagree on the vector width: {sorted(widths)}"
1807
+ raise EmbeddingProviderMismatch(msg)
1808
+ if not selects:
1809
+ return register_vss(con, lance_uri=present[0], expected_model=None)
1810
+ width = int(widths.pop())
1811
+ union = " UNION ALL ".join(
1812
+ f"SELECT uuid, model, dim, CAST(embedding AS FLOAT[{width}]) AS embedding, embedded_at " # noqa: S608 # nosec B608 - int width and int-derived aliases
1813
+ f"FROM {alias}.main.embeddings"
1814
+ for alias in selects
1815
+ )
1816
+ con.execute(f"CREATE OR REPLACE VIEW message_embeddings AS {union};")
1817
+ logger.debug("Bound message_embeddings over {} Lance stores (dim={})", len(selects), width)
1818
+ return True
1819
+
1820
+
1821
+ # ---------------------------------------------------------------------------
1822
+ # Macros
1823
+ # ---------------------------------------------------------------------------
1824
+
1825
+
1826
+ def _pricing_values_clause(pricing: dict[str, tuple[float, float]]) -> SqlFragment:
1827
+ """Render a pricing dict as an inline SQL ``VALUES`` row list.
1828
+
1829
+ Parameters
1830
+ ----------
1831
+ pricing
1832
+ Mapping of ``model_name -> (input_rate, output_rate)`` per 1M tokens.
1833
+
1834
+ Returns
1835
+ -------
1836
+ str
1837
+ Comma-separated ``('model', in, out)`` rows. Emits a sentinel row
1838
+ that matches no real model if ``pricing`` is empty (DuckDB rejects
1839
+ empty ``VALUES`` lists).
1840
+ """
1841
+ if not pricing:
1842
+ return SqlFragment(f"({sql_literal('__no_pricing__')}, 0.0, 0.0)")
1843
+ # The model NAME is escaped and the two rates are coerced with `float()`
1844
+ # rather than interpolated as they arrive. Both are the declared type of the
1845
+ # `pricing` parameter, and atif-cli never passes one — but this is a library
1846
+ # meant to be embedded in-process, so a caller that reaches DuckDB with a
1847
+ # `str` rate past the type checker gets a ValueError here instead of an
1848
+ # injected VALUES row.
1849
+ rows = [
1850
+ f"({sql_literal(model)}, {float(in_rate)}, {float(out_rate)})"
1851
+ for model, (in_rate, out_rate) in sorted(pricing.items())
1852
+ ]
1853
+ return SqlFragment(", ".join(rows))
1854
+
1855
+
1856
+ def register_macros(
1857
+ con: duckdb.DuckDBPyConnection,
1858
+ pricing: dict[str, tuple[float, float]] | None = None,
1859
+ *,
1860
+ skip_vss: bool = False,
1861
+ ) -> None:
1862
+ """Create the SQL macros declared in the static catalog.
1863
+
1864
+ Every signature here must match its ``MACRO_SIGNATURES`` entry in
1865
+ :mod:`atif_duck.domain.catalog`: the catalog is what ``atif-sql schema``
1866
+ prints and what the derived examples are built from, and the drift test
1867
+ fails on any divergence.
1868
+
1869
+ Must be called after :func:`register_views` — DuckDB binds macro bodies
1870
+ at CREATE time, so every referenced view must already exist. The
1871
+ ``semantic_search`` macro additionally requires ``message_embeddings``
1872
+ (view or empty fallback table from :func:`register_vss`) to be bound
1873
+ first; ``skip_vss=True`` skips it for connections that never call
1874
+ :func:`register_vss`.
1875
+
1876
+ Parameters
1877
+ ----------
1878
+ con
1879
+ Open DuckDB connection with views already registered.
1880
+ pricing
1881
+ Optional pricing override; falls back to
1882
+ :data:`atif_duck.domain.catalog.DEFAULT_PRICING`.
1883
+ skip_vss
1884
+ When ``True``, skip the ``semantic_search`` macro registration.
1885
+
1886
+ Raises
1887
+ ------
1888
+ duckdb.Error
1889
+ If any macro DDL fails. Logged via ``logger.exception`` before
1890
+ re-raise.
1891
+ """
1892
+ pricing_rows = _pricing_values_clause(pricing if pricing is not None else DEFAULT_PRICING)
1893
+
1894
+ try:
1895
+ # ``ago('14 days')`` -> ``current_timestamp - INTERVAL 14 DAY``.
1896
+ # The CAST handles every interval-unit shape DuckDB recognizes.
1897
+ con.execute(
1898
+ """
1899
+ CREATE OR REPLACE MACRO ago(interval_text) AS (
1900
+ current_timestamp - CAST(interval_text AS INTERVAL)
1901
+ );
1902
+ """
1903
+ )
1904
+
1905
+ con.execute(
1906
+ """
1907
+ CREATE OR REPLACE MACRO model_used(sid) AS (
1908
+ SELECT any_value(model_name)
1909
+ FROM steps
1910
+ WHERE session_id = sid AND model_name IS NOT NULL
1911
+ );
1912
+ """
1913
+ )
1914
+
1915
+ # COORDINATE SPACE: ATIF's ``prompt_tokens`` is the TOTAL input
1916
+ # (non-cached + cache_read + cache_creation), so it is NOT the
1917
+ # billable uncached base. The charged base here is
1918
+ # (prompt_tokens - cached_tokens) = non-cached + cache_creation.
1919
+ # Billing cache_creation at the input rate and leaving cache reads
1920
+ # uncharged is an approximation: it over-charges a cache write and
1921
+ # under-charges a cache read relative to a provider's real rate card.
1922
+ # The prefix match strips dated model suffixes
1923
+ # (``claude-haiku-4-5-20251001`` -> ``claude-haiku-4-5``).
1924
+ #
1925
+ # LEFT JOIN plus ``unpriced_steps`` is the honest-accounting shape: an
1926
+ # inner join silently drops steps whose model has no pricing row, so a
1927
+ # session mixing priced and unpriced models would return a PARTIAL
1928
+ # number indistinguishable from a complete one. ``est_cost_usd``
1929
+ # covers the priced steps only and is meaningful ONLY when
1930
+ # ``unpriced_steps = 0``.
1931
+ #
1932
+ # Both counters filter on ``model_name IS NOT NULL``: USER steps carry
1933
+ # no model and cost nothing, so counting them as unpriced would put
1934
+ # every session with a user turn above zero and make a real pricing
1935
+ # gap indistinguishable from an ordinary conversation.
1936
+ con.execute(
1937
+ f"""
1938
+ CREATE OR REPLACE MACRO cost_estimate(sid) AS TABLE (
1939
+ SELECT sum(
1940
+ (coalesce(s.prompt_tokens, 0) - coalesce(s.cached_tokens, 0))
1941
+ * p.in_rate
1942
+ + coalesce(s.completion_tokens, 0) * p.out_rate
1943
+ ) / 1e6 AS est_cost_usd,
1944
+ count(*) FILTER (
1945
+ WHERE s.model_name IS NOT NULL AND p.model IS NOT NULL
1946
+ ) AS priced_steps,
1947
+ count(*) FILTER (
1948
+ WHERE s.model_name IS NOT NULL AND p.model IS NULL
1949
+ ) AS unpriced_steps
1950
+ FROM steps s
1951
+ LEFT JOIN (VALUES {pricing_rows}) p(model, in_rate, out_rate)
1952
+ ON regexp_replace(s.model_name, '-\\d{{8}}$', '') = p.model
1953
+ WHERE s.session_id = sid
1954
+ );
1955
+ """ # noqa: S608 # nosec B608 - pricing model names escaped by sql_literal; the rates are floats
1956
+ )
1957
+
1958
+ con.execute(
1959
+ """
1960
+ CREATE OR REPLACE MACRO tool_rank(last_n_days) AS TABLE (
1961
+ SELECT tool_name, count(*) AS n
1962
+ FROM tool_calls
1963
+ WHERE ts >= current_timestamp - (last_n_days * INTERVAL 1 DAY)
1964
+ AND tool_name IS NOT NULL
1965
+ GROUP BY 1
1966
+ ORDER BY n DESC
1967
+ );
1968
+ """
1969
+ )
1970
+
1971
+ con.execute(
1972
+ """
1973
+ CREATE OR REPLACE MACRO todo_velocity(sid) AS (
1974
+ SELECT count(*) FILTER (WHERE status = 'completed')::DOUBLE
1975
+ / NULLIF(count(DISTINCT subject), 0)
1976
+ FROM todo_state_current
1977
+ WHERE session_id = sid
1978
+ );
1979
+ """
1980
+ )
1981
+
1982
+ # SCOPE: fan-out counts spawn INTENT — the number of Task/Agent
1983
+ # launches the session issued — not the number of subagent transcripts
1984
+ # that resulted. ATIF inlines sidechains, so there is no per-subagent
1985
+ # file to census, and a launch that produced no steps still counts
1986
+ # here. Read this as an upper bound on subagents realized.
1987
+ con.execute(
1988
+ """
1989
+ CREATE OR REPLACE MACRO subagent_fanout(sid) AS (
1990
+ SELECT count(*)
1991
+ FROM subagent_spawns
1992
+ WHERE session_id = sid
1993
+ );
1994
+ """
1995
+ )
1996
+
1997
+ # Semantic top-k nearest neighbors over ``message_embeddings``.
1998
+ # ``query_vec`` must have the bound view's width.
1999
+ #
2000
+ # COSINE, not L2, throughout: stored document vectors are int8-cast
2001
+ # and un-normalized (observed norms 1206-1419 on the live store) while
2002
+ # ``embed_query`` returns unit-norm probes, so L2 ranks partly by
2003
+ # MAGNITUDE and puts lower-similarity rows above higher ones. ``sim``
2004
+ # and ``distance`` are the two halves of the SAME metric
2005
+ # (distance = 1 - sim); mixing cosine ``sim`` with L2 ``distance``
2006
+ # would make the two columns disagree about which hit is nearest.
2007
+ #
2008
+ # ``list_cosine_*`` rather than ``array_cosine_*``: the list variants
2009
+ # accept a fixed-size ARRAY and a variable-length LIST on either side,
2010
+ # so a caller-computed probe binds without knowing the store's width
2011
+ # at DDL time. Neither variant pushes down into the Lance index (both
2012
+ # plan a full __LANCE_TABLE_SCAN), so this costs no index lookup.
2013
+ if not skip_vss:
2014
+ con.execute(
2015
+ """
2016
+ CREATE OR REPLACE MACRO semantic_search(query_vec, k) AS TABLE (
2017
+ SELECT me.uuid,
2018
+ list_cosine_similarity(me.embedding, query_vec) AS sim,
2019
+ list_cosine_distance(me.embedding, query_vec) AS distance
2020
+ FROM message_embeddings me
2021
+ ORDER BY list_cosine_distance(me.embedding, query_vec)
2022
+ LIMIT k
2023
+ );
2024
+ """
2025
+ )
2026
+ else:
2027
+ logger.debug("Skipped semantic_search macro (skip_vss=True)")
2028
+
2029
+ # Skill / slash-command leaderboard over the last N days.
2030
+ con.execute(
2031
+ """
2032
+ CREATE OR REPLACE MACRO skill_rank(last_n_days) AS TABLE (
2033
+ SELECT skill_id,
2034
+ skill_name,
2035
+ plugin,
2036
+ is_builtin,
2037
+ count(*) AS n,
2038
+ count(DISTINCT session_id) AS sessions
2039
+ FROM skill_usage
2040
+ WHERE ts >= current_timestamp - (last_n_days * INTERVAL 1 DAY)
2041
+ GROUP BY 1, 2, 3, 4
2042
+ ORDER BY n DESC
2043
+ );
2044
+ """
2045
+ )
2046
+
2047
+ # How is each skill invoked? Built-ins excluded — they're almost
2048
+ # always slash-only and would drown everything else out.
2049
+ con.execute(
2050
+ """
2051
+ CREATE OR REPLACE MACRO skill_source_mix(last_n_days) AS TABLE (
2052
+ SELECT skill_id,
2053
+ skill_name,
2054
+ count(*) FILTER (WHERE source = 'tool') AS n_tool,
2055
+ count(*) FILTER (WHERE source = 'slash_command') AS n_slash,
2056
+ count(*) AS n_total
2057
+ FROM skill_usage
2058
+ WHERE ts >= current_timestamp - (last_n_days * INTERVAL 1 DAY)
2059
+ AND NOT is_builtin
2060
+ GROUP BY 1, 2
2061
+ ORDER BY n_total DESC
2062
+ );
2063
+ """
2064
+ )
2065
+
2066
+ _semantic_search_part = "" if skip_vss else "semantic_search, "
2067
+ logger.debug(
2068
+ "Registered macros: ago, model_used, cost_estimate, tool_rank, "
2069
+ f"todo_velocity, subagent_fanout, {_semantic_search_part}skill_rank, "
2070
+ "skill_source_mix"
2071
+ )
2072
+ except Exception:
2073
+ # register-or-fail-loud — any DuckDB error must surface to the caller.
2074
+ logger.exception("Failed to register macros")
2075
+ raise
2076
+
2077
+
2078
+ # ---------------------------------------------------------------------------
2079
+ # Entry point
2080
+ # ---------------------------------------------------------------------------
2081
+
2082
+
2083
+ def analytics_roots(corpus_root: Path, lake: LakeReader | None) -> tuple[Path, ...]:
2084
+ """The corpus roots whose ``analytics/`` parquets a registration binds.
2085
+
2086
+ ``corpus_root`` alone, unless ``lake`` spans every corpus: then every
2087
+ corpus the lake holds, in name order. atif-cli grants the query sandbox
2088
+ the same set.
2089
+ """
2090
+ if lake is not None and lake.corpus is None and lake.corpus_roots:
2091
+ return lake.corpus_roots
2092
+ return (corpus_root,)
2093
+
2094
+
2095
+ def register(
2096
+ con: duckdb.DuckDBPyConnection,
2097
+ corpus_root: Path,
2098
+ pricing: dict[str, tuple[float, float]] | None = None,
2099
+ *,
2100
+ skip_vss: bool = False,
2101
+ lance_uri: Path | None = None,
2102
+ expected_model: str | None = None,
2103
+ expected_dim: int | None = None,
2104
+ lake: LakeReader | None = None,
2105
+ lance_uris: Sequence[Path] | None = None,
2106
+ ) -> RawSources:
2107
+ """Register raw readers, views, VSS, and macros over ``corpus_root``, in order.
2108
+
2109
+ Every call re-scans the corpus, and re-parses every session that has no
2110
+ columnar artifacts into TEMP tables, so cost is O(JSON sessions) per
2111
+ connection: reuse one connection per process rather than registering per
2112
+ query. Returns :func:`register_raw`'s :class:`RawSources` so the caller
2113
+ knows which path each session took and which parquet files the views
2114
+ will open lazily.
2115
+
2116
+ Order matters: raw TEMP tables first (views bind against them at CREATE
2117
+ time), then views, then VSS (``semantic_search``'s body binds against
2118
+ ``message_embeddings`` at CREATE time), then macros (macro bodies bind
2119
+ against views at CREATE time), then the authorship surface
2120
+ (``step_author`` and the ``user_steps`` / ``human_turns`` /
2121
+ ``session_outcomes`` views over ``steps``), then the v2 analytics views +
2122
+ macros (which bind against the analytics parquets AND the base views —
2123
+ ``friction_rate`` counts ``human_turns``, ``conflicts_over_time`` joins
2124
+ ``messages``).
2125
+
2126
+ Parameters
2127
+ ----------
2128
+ con
2129
+ Open DuckDB connection.
2130
+ corpus_root
2131
+ Materialized corpus root (the directory containing ``sessions/``).
2132
+ pricing
2133
+ Optional pricing override for :func:`register_macros`.
2134
+ skip_vss
2135
+ When ``True``, neither :func:`register_vss` nor the
2136
+ ``semantic_search`` macro registers. The embed backfill needs this:
2137
+ it WRITES the store that ``message_embeddings`` reads, so a
2138
+ connection opened to run the backfill cannot bind that view first.
2139
+ Default ``False`` — an empty or absent store degrades to the empty
2140
+ fallback table, so registration is otherwise always safe.
2141
+ lance_uri
2142
+ Local LanceDB dataset directory; defaults to
2143
+ ``<corpus_root>/embeddings_lance`` (atif-embed's per-corpus default).
2144
+ expected_model, expected_dim
2145
+ Active embedder identity for :func:`register_vss`'s fail-loud guard.
2146
+ lake
2147
+ A lake already attached by
2148
+ :func:`atif_duck.infrastructure.lake.attach_lake_for_query`. When given,
2149
+ the raw relations are views over its tables (scoped as it says)
2150
+ instead of readers over ``corpus_root``'s per-session artifacts, and
2151
+ nothing is loaded eagerly. ``corpus_root`` still locates the
2152
+ default embeddings store, and the analytics parquets unless the lake
2153
+ spans every corpus: then each analytics view reads every lake
2154
+ corpus's parquets, so ``--all-corpora`` covers them as it covers the
2155
+ transcripts.
2156
+ lance_uris
2157
+ Several stores to bind ``message_embeddings`` over at once (one per
2158
+ corpus, for ``--all-corpora``); see :func:`register_vss_stores`.
2159
+ Overrides ``lance_uri``.
2160
+
2161
+ Raises
2162
+ ------
2163
+ duckdb.Error
2164
+ If any registration step fails (register-or-fail-loud).
2165
+ """
2166
+ from atif_duck.infrastructure.analytics import (
2167
+ register_analytics,
2168
+ register_analytics_macros,
2169
+ )
2170
+ from atif_duck.infrastructure.authorship import register_authorship
2171
+
2172
+ if lake is not None:
2173
+ from atif_duck.infrastructure.lake import register_lake_raw
2174
+
2175
+ sources = register_lake_raw(con, lake)
2176
+ else:
2177
+ sources = register_raw(con, corpus_root)
2178
+ register_views(con)
2179
+ if not skip_vss and lance_uris is not None:
2180
+ register_vss_stores(
2181
+ con, lance_uris=lance_uris, expected_model=expected_model, expected_dim=expected_dim
2182
+ )
2183
+ elif not skip_vss:
2184
+ register_vss(
2185
+ con,
2186
+ lance_uri=lance_uri if lance_uri is not None else corpus_root / "embeddings_lance",
2187
+ expected_model=expected_model,
2188
+ expected_dim=expected_dim,
2189
+ )
2190
+ register_macros(con, pricing=pricing, skip_vss=skip_vss)
2191
+ register_authorship(con)
2192
+ registered = register_analytics(con, *analytics_roots(corpus_root, lake))
2193
+ register_analytics_macros(con, registered)
2194
+ return sources