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