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.
- atif_analytics/__init__.py +9 -0
- atif_analytics/application/__init__.py +3 -0
- atif_analytics/application/analyze.py +158 -0
- atif_analytics/application/prompts.py +750 -0
- atif_analytics/application/use_cases/__init__.py +3 -0
- atif_analytics/application/use_cases/_shared.py +264 -0
- atif_analytics/application/use_cases/classify.py +430 -0
- atif_analytics/application/use_cases/conflicts.py +448 -0
- atif_analytics/application/use_cases/friction.py +634 -0
- atif_analytics/application/use_cases/perceived.py +451 -0
- atif_analytics/domain/__init__.py +3 -0
- atif_analytics/domain/authorship.py +166 -0
- atif_analytics/domain/config.py +30 -0
- atif_analytics/domain/costs.py +50 -0
- atif_analytics/domain/friction.py +117 -0
- atif_analytics/domain/layout.py +94 -0
- atif_analytics/domain/models.py +385 -0
- atif_analytics/domain/ports.py +100 -0
- atif_analytics/domain/transcript.py +359 -0
- atif_analytics/infrastructure/__init__.py +3 -0
- atif_analytics/infrastructure/corpus_reader.py +546 -0
- atif_analytics/infrastructure/parquet_cache.py +186 -0
- atif_analytics/infrastructure/settings.py +114 -0
- atif_analytics/infrastructure/sqlite_state/__init__.py +3 -0
- atif_analytics/infrastructure/sqlite_state/checkpointer.py +326 -0
- atif_analytics/infrastructure/sqlite_state/retry_queue.py +296 -0
- atif_analytics/py.typed +0 -0
- atif_cli/__init__.py +3 -0
- atif_cli/__main__.py +8 -0
- atif_cli/app.py +2459 -0
- atif_cli/converter_adapter.py +151 -0
- atif_cli/corpus.py +409 -0
- atif_cli/cron.py +311 -0
- atif_cli/duck_errors.py +97 -0
- atif_cli/embed_lake.py +101 -0
- atif_cli/errors.py +89 -0
- atif_cli/lake.py +373 -0
- atif_cli/lake_sessions.py +270 -0
- atif_cli/output.py +297 -0
- atif_cli/py.typed +0 -0
- atif_converter/__init__.py +3 -0
- atif_converter/application/__init__.py +3 -0
- atif_converter/application/convert_and_audit.py +294 -0
- atif_converter/application/convert_codex.py +210 -0
- atif_converter/domain/__init__.py +9 -0
- atif_converter/domain/agents.py +40 -0
- atif_converter/domain/atif/__init__.py +55 -0
- atif_converter/domain/atif/agent.py +37 -0
- atif_converter/domain/atif/content.py +161 -0
- atif_converter/domain/atif/final_metrics.py +45 -0
- atif_converter/domain/atif/metrics.py +49 -0
- atif_converter/domain/atif/observation.py +21 -0
- atif_converter/domain/atif/observation_result.py +47 -0
- atif_converter/domain/atif/step.py +144 -0
- atif_converter/domain/atif/subagent_trajectory_ref.py +90 -0
- atif_converter/domain/atif/tool_call.py +36 -0
- atif_converter/domain/atif/trajectory.py +211 -0
- atif_converter/domain/atif/trajectory_validator.py +299 -0
- atif_converter/domain/blobs.py +510 -0
- atif_converter/domain/claude_code_conversion.py +1077 -0
- atif_converter/domain/codex_conversion.py +1036 -0
- atif_converter/domain/codex_edges.py +192 -0
- atif_converter/domain/codex_enrichment.py +419 -0
- atif_converter/domain/codex_fidelity.py +124 -0
- atif_converter/domain/codex_nested_calls.py +165 -0
- atif_converter/domain/edges.py +113 -0
- atif_converter/domain/enrichment.py +412 -0
- atif_converter/domain/errors.py +60 -0
- atif_converter/domain/fidelity.py +171 -0
- atif_converter/domain/model_prices.json +8538 -0
- atif_converter/domain/pricing.py +790 -0
- atif_converter/domain/result_signals.py +774 -0
- atif_converter/domain/schema_version.py +90 -0
- atif_converter/domain/session_events.py +408 -0
- atif_converter/infrastructure/__init__.py +3 -0
- atif_converter/infrastructure/census.py +89 -0
- atif_converter/infrastructure/claude_code_converter.py +127 -0
- atif_converter/infrastructure/codex_adapter.py +132 -0
- atif_converter/infrastructure/codex_census.py +100 -0
- atif_converter/infrastructure/codex_converter.py +82 -0
- atif_converter/infrastructure/harbor_adapter.py +187 -0
- atif_converter/infrastructure/raw_records.py +426 -0
- atif_converter/infrastructure/source_archive.py +154 -0
- atif_converter/py.typed +0 -0
- atif_corpus/__init__.py +13 -0
- atif_corpus/application/__init__.py +7 -0
- atif_corpus/application/materialize.py +1826 -0
- atif_corpus/domain/__init__.py +9 -0
- atif_corpus/domain/agents.py +33 -0
- atif_corpus/domain/generation.py +34 -0
- atif_corpus/domain/layout.py +168 -0
- atif_corpus/domain/ports.py +175 -0
- atif_corpus/domain/session_id.py +68 -0
- atif_corpus/domain/sessions.py +228 -0
- atif_corpus/domain/slug.py +72 -0
- atif_corpus/domain/source_layout.py +143 -0
- atif_corpus/domain/watermark.py +76 -0
- atif_corpus/infrastructure/__init__.py +7 -0
- atif_corpus/infrastructure/atomic.py +266 -0
- atif_corpus/infrastructure/compress_artifacts.py +161 -0
- atif_corpus/infrastructure/fake_converter.py +174 -0
- atif_corpus/infrastructure/scanner.py +316 -0
- atif_corpus/infrastructure/settings.py +130 -0
- atif_corpus/infrastructure/source_archive.py +202 -0
- atif_corpus/py.typed +0 -0
- atif_duck/__init__.py +43 -0
- atif_duck/domain/__init__.py +32 -0
- atif_duck/domain/artifacts.py +66 -0
- atif_duck/domain/authorship.py +165 -0
- atif_duck/domain/catalog.py +626 -0
- atif_duck/domain/columnar.py +149 -0
- atif_duck/domain/embedding_guard.py +88 -0
- atif_duck/domain/examples.py +265 -0
- atif_duck/domain/lake.py +427 -0
- atif_duck/domain/lake_sessions.py +97 -0
- atif_duck/domain/raw_readers.py +140 -0
- atif_duck/domain/session_id.py +69 -0
- atif_duck/domain/sql_literal.py +50 -0
- atif_duck/infrastructure/__init__.py +34 -0
- atif_duck/infrastructure/analytics.py +469 -0
- atif_duck/infrastructure/authorship.py +139 -0
- atif_duck/infrastructure/columnar.py +641 -0
- atif_duck/infrastructure/lake.py +1383 -0
- atif_duck/infrastructure/lake_sessions.py +194 -0
- atif_duck/infrastructure/lake_settings.py +63 -0
- atif_duck/infrastructure/lake_steps.py +225 -0
- atif_duck/infrastructure/projections.py +221 -0
- atif_duck/infrastructure/registry.py +2194 -0
- atif_duck/infrastructure/stored_artifacts.py +115 -0
- atif_duck/py.typed +0 -0
- atif_embed/__init__.py +12 -0
- atif_embed/application/__init__.py +3 -0
- atif_embed/application/embed.py +293 -0
- atif_embed/application/prune.py +80 -0
- atif_embed/domain/__init__.py +3 -0
- atif_embed/domain/discovery.py +70 -0
- atif_embed/domain/embedding_guard.py +83 -0
- atif_embed/domain/errors.py +94 -0
- atif_embed/domain/ports.py +223 -0
- atif_embed/domain/sql_literal.py +49 -0
- atif_embed/domain/text_stamp.py +71 -0
- atif_embed/infrastructure/__init__.py +3 -0
- atif_embed/infrastructure/cohere_bedrock.py +433 -0
- atif_embed/infrastructure/corpus_text_rows.py +273 -0
- atif_embed/infrastructure/gemma_local.py +355 -0
- atif_embed/infrastructure/lake_text_rows.py +203 -0
- atif_embed/infrastructure/lance_store.py +440 -0
- atif_embed/infrastructure/providers.py +30 -0
- atif_embed/infrastructure/settings.py +154 -0
- atif_embed/py.typed +0 -0
- atif_models/__init__.py +14 -0
- atif_models/domain/__init__.py +9 -0
- atif_models/domain/ports.py +143 -0
- atif_models/domain/registry.py +151 -0
- atif_models/domain/schema.py +109 -0
- atif_models/infrastructure/__init__.py +8 -0
- atif_models/infrastructure/openai_bedrock.py +314 -0
- atif_models/infrastructure/settings.py +85 -0
- atif_models/py.typed +0 -0
- atif_sql-0.2.0.dist-info/METADATA +272 -0
- atif_sql-0.2.0.dist-info/RECORD +164 -0
- atif_sql-0.2.0.dist-info/WHEEL +4 -0
- atif_sql-0.2.0.dist-info/entry_points.txt +2 -0
- 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
|