sqbyl-runtime 0.1.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.
- sqbyl_runtime/__init__.py +40 -0
- sqbyl_runtime/context.py +319 -0
- sqbyl_runtime/cost.py +265 -0
- sqbyl_runtime/db/__init__.py +49 -0
- sqbyl_runtime/db/connection.py +172 -0
- sqbyl_runtime/db/dialects.py +329 -0
- sqbyl_runtime/db/errors.py +35 -0
- sqbyl_runtime/db/guard.py +99 -0
- sqbyl_runtime/export.py +230 -0
- sqbyl_runtime/fingerprint.py +154 -0
- sqbyl_runtime/llm/__init__.py +43 -0
- sqbyl_runtime/llm/_errors.py +30 -0
- sqbyl_runtime/llm/anthropic_client.py +141 -0
- sqbyl_runtime/llm/base.py +177 -0
- sqbyl_runtime/llm/factory.py +40 -0
- sqbyl_runtime/llm/mock.py +74 -0
- sqbyl_runtime/llm/openai_client.py +204 -0
- sqbyl_runtime/llm/replay.py +73 -0
- sqbyl_runtime/models/__init__.py +50 -0
- sqbyl_runtime/models/assets.py +36 -0
- sqbyl_runtime/models/base.py +23 -0
- sqbyl_runtime/models/judges.py +19 -0
- sqbyl_runtime/models/release.py +109 -0
- sqbyl_runtime/models/selection.py +26 -0
- sqbyl_runtime/models/semantics.py +100 -0
- sqbyl_runtime/pipeline.py +288 -0
- sqbyl_runtime/py.typed +0 -0
- sqbyl_runtime/runtime.py +202 -0
- sqbyl_runtime/schema.py +41 -0
- sqbyl_runtime/selection.py +348 -0
- sqbyl_runtime/state/__init__.py +26 -0
- sqbyl_runtime/state/layout.py +46 -0
- sqbyl_runtime/state/traces.py +127 -0
- sqbyl_runtime/state/usage.py +164 -0
- sqbyl_runtime-0.1.0.dist-info/METADATA +57 -0
- sqbyl_runtime-0.1.0.dist-info/RECORD +37 -0
- sqbyl_runtime-0.1.0.dist-info/WHEEL +4 -0
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
"""sqbyl-runtime — the minimal, shippable sqbyl runtime.
|
|
2
|
+
|
|
3
|
+
Contains only what a production app needs to embed a released agent: release
|
|
4
|
+
``load()``, ``ask()``, the ``LLMClient`` seam, and structured logging. None of
|
|
5
|
+
the dev toolkit (eval, synth, Coach, judges, console) lives here or is importable
|
|
6
|
+
from here — that one-way dependency arrow is enforced by import-linter in CI.
|
|
7
|
+
|
|
8
|
+
from sqbyl_runtime import load
|
|
9
|
+
agent = load("revenue-analytics.v3.json", db=DATABASE_URL, model="claude-opus-4-8")
|
|
10
|
+
agent.ask("How many orders shipped last month?") # → AgentResult
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from sqbyl_runtime.export import (
|
|
14
|
+
McpServer,
|
|
15
|
+
answer_dict,
|
|
16
|
+
as_callable,
|
|
17
|
+
langchain_tool,
|
|
18
|
+
serve_mcp_stdio,
|
|
19
|
+
)
|
|
20
|
+
from sqbyl_runtime.runtime import (
|
|
21
|
+
Agent,
|
|
22
|
+
ModelMismatchWarning,
|
|
23
|
+
SchemaMismatchWarning,
|
|
24
|
+
load,
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
__version__ = "0.0.0"
|
|
28
|
+
|
|
29
|
+
__all__ = [
|
|
30
|
+
"Agent",
|
|
31
|
+
"McpServer",
|
|
32
|
+
"ModelMismatchWarning",
|
|
33
|
+
"SchemaMismatchWarning",
|
|
34
|
+
"answer_dict",
|
|
35
|
+
"as_callable",
|
|
36
|
+
"langchain_tool",
|
|
37
|
+
"load",
|
|
38
|
+
"serve_mcp_stdio",
|
|
39
|
+
"__version__",
|
|
40
|
+
]
|
sqbyl_runtime/context.py
ADDED
|
@@ -0,0 +1,319 @@
|
|
|
1
|
+
"""The context compiler (spec §5 steps 1-2, plan 2.1).
|
|
2
|
+
|
|
3
|
+
Turn a project's knowledge (semantics + examples + trusted assets + instructions)
|
|
4
|
+
plus a question into the prompt the agent runtime sends to Claude. Two halves:
|
|
5
|
+
|
|
6
|
+
* a **stable system block** — instructions, annotated schema/semantics, trusted
|
|
7
|
+
assets, and few-shot examples. It does not vary with the question, so it is the
|
|
8
|
+
prompt-cache unit (``cache_system=True`` downstream).
|
|
9
|
+
* a **question turn** — the one varying part, sent as the user message.
|
|
10
|
+
|
|
11
|
+
This lives in ``sqbyl-runtime`` because ``ask()`` compiles context at inference
|
|
12
|
+
time (the same compiler runs in dev and in a shipped release). The small-project
|
|
13
|
+
path is "include everything"; large-schema shortlisting (Phase 9) runs as a
|
|
14
|
+
separate :mod:`~sqbyl_runtime.selection` step *before* this renderer and hands the
|
|
15
|
+
narrowed table set in — so the renderer itself stays a pure function of its inputs
|
|
16
|
+
(deterministic, snapshot-testable): no clocks, no IDs, no LLM calls, stable ordering.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from __future__ import annotations
|
|
20
|
+
|
|
21
|
+
from collections.abc import Sequence
|
|
22
|
+
|
|
23
|
+
from pydantic import BaseModel, Field
|
|
24
|
+
|
|
25
|
+
from sqbyl_runtime.llm.base import LLMClient, Usage
|
|
26
|
+
from sqbyl_runtime.models import (
|
|
27
|
+
Column,
|
|
28
|
+
Dialect,
|
|
29
|
+
Example,
|
|
30
|
+
Profile,
|
|
31
|
+
ReleaseArtifact,
|
|
32
|
+
ScalarBound,
|
|
33
|
+
SelectionConfig,
|
|
34
|
+
TableSemantics,
|
|
35
|
+
TrustedAsset,
|
|
36
|
+
)
|
|
37
|
+
from sqbyl_runtime.selection import LLMCallHook, ValueMatch, select_context
|
|
38
|
+
|
|
39
|
+
# Past this many tables "include everything" stops being viable and Phase 9's
|
|
40
|
+
# LLM/lexical shortlisting is needed; until then we still include all but flag it.
|
|
41
|
+
_INCLUDE_ALL_TABLE_LIMIT = 30
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
class CompiledContext(BaseModel):
|
|
45
|
+
"""The compiled prompt plus what went into it (for citation + caching).
|
|
46
|
+
|
|
47
|
+
``usage`` carries any tokens the *selection* step spent (LLM shortlisting on a
|
|
48
|
+
large schema); the pipeline folds it into the run's total so it's metered and
|
|
49
|
+
budgeted like generation (invariant 5). It is ``Usage()`` for the include-all /
|
|
50
|
+
lexical paths, which spend nothing.
|
|
51
|
+
"""
|
|
52
|
+
|
|
53
|
+
system: str
|
|
54
|
+
user: str
|
|
55
|
+
selected_tables: list[str] = Field(default_factory=list)
|
|
56
|
+
offered_assets: list[str] = Field(default_factory=list)
|
|
57
|
+
notes: list[str] = Field(default_factory=list)
|
|
58
|
+
usage: Usage = Field(default_factory=Usage)
|
|
59
|
+
# The selection strategy that actually ran (rewritten to ``include_all`` on a fallback)
|
|
60
|
+
# and whether a narrowing strategy degraded to include-all — carried out so the pipeline
|
|
61
|
+
# can put the drop rationale on the run's trace/result (invariant 7, transparency).
|
|
62
|
+
selection_strategy: str = "include_all"
|
|
63
|
+
selection_fell_back: bool = False
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class ProjectKnowledge(BaseModel):
|
|
67
|
+
"""Everything the agent reasons over, decoupled from where it came from.
|
|
68
|
+
|
|
69
|
+
Both a dev project (loaded from files) and a shipped ``ReleaseArtifact`` produce
|
|
70
|
+
one of these, so the runtime pipeline is identical in dev and in production.
|
|
71
|
+
"""
|
|
72
|
+
|
|
73
|
+
dialect: Dialect
|
|
74
|
+
semantics: list[TableSemantics] = Field(default_factory=list)
|
|
75
|
+
instructions: str = ""
|
|
76
|
+
examples: list[Example] = Field(default_factory=list)
|
|
77
|
+
trusted_assets: list[TrustedAsset] = Field(default_factory=list)
|
|
78
|
+
selection: SelectionConfig = Field(default_factory=SelectionConfig)
|
|
79
|
+
|
|
80
|
+
@classmethod
|
|
81
|
+
def from_release(cls, release: ReleaseArtifact) -> ProjectKnowledge:
|
|
82
|
+
return cls(
|
|
83
|
+
dialect=release.dialect,
|
|
84
|
+
semantics=release.semantics,
|
|
85
|
+
instructions=release.instructions,
|
|
86
|
+
examples=release.examples,
|
|
87
|
+
trusted_assets=release.trusted_assets,
|
|
88
|
+
selection=release.selection,
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
def compile(
|
|
92
|
+
self,
|
|
93
|
+
question: str,
|
|
94
|
+
*,
|
|
95
|
+
llm: LLMClient | None = None,
|
|
96
|
+
model: str | None = None,
|
|
97
|
+
on_llm_call: LLMCallHook | None = None,
|
|
98
|
+
) -> CompiledContext:
|
|
99
|
+
"""Select the relevant subset (LLM/lexical for large schemas) then render it.
|
|
100
|
+
|
|
101
|
+
``llm`` and ``model`` power the ``llm``/``llm_lexical`` selection strategies;
|
|
102
|
+
without them (or on the include-all/lexical strategies) selection spends nothing
|
|
103
|
+
and this is a pure function of the project's files. ``on_llm_call`` lets a caller
|
|
104
|
+
(the pipeline) trace the shortlisting call as its own GenAI span.
|
|
105
|
+
"""
|
|
106
|
+
return compile_context(
|
|
107
|
+
question,
|
|
108
|
+
dialect=self.dialect,
|
|
109
|
+
semantics=self.semantics,
|
|
110
|
+
instructions=self.instructions,
|
|
111
|
+
examples=self.examples,
|
|
112
|
+
trusted_assets=self.trusted_assets,
|
|
113
|
+
selection=self.selection,
|
|
114
|
+
llm=llm,
|
|
115
|
+
model=model,
|
|
116
|
+
on_llm_call=on_llm_call,
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def compile_context(
|
|
121
|
+
question: str,
|
|
122
|
+
*,
|
|
123
|
+
dialect: Dialect,
|
|
124
|
+
semantics: Sequence[TableSemantics],
|
|
125
|
+
instructions: str = "",
|
|
126
|
+
examples: Sequence[Example] = (),
|
|
127
|
+
trusted_assets: Sequence[TrustedAsset] = (),
|
|
128
|
+
selection: SelectionConfig | None = None,
|
|
129
|
+
llm: LLMClient | None = None,
|
|
130
|
+
model: str | None = None,
|
|
131
|
+
on_llm_call: LLMCallHook | None = None,
|
|
132
|
+
) -> CompiledContext:
|
|
133
|
+
"""Compile the question + project knowledge into a system/user prompt pair.
|
|
134
|
+
|
|
135
|
+
Runs the :mod:`~sqbyl_runtime.selection` step first (include-all for small
|
|
136
|
+
projects; lexical/LLM shortlisting past ``max_tables``), then renders only the
|
|
137
|
+
surviving tables and the examples that reference them, plus any value-match hints.
|
|
138
|
+
"""
|
|
139
|
+
selection = selection or SelectionConfig()
|
|
140
|
+
picked = select_context(
|
|
141
|
+
question,
|
|
142
|
+
semantics=semantics,
|
|
143
|
+
config=selection,
|
|
144
|
+
llm=llm,
|
|
145
|
+
model=model,
|
|
146
|
+
on_llm_call=on_llm_call,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
chosen_names = set(picked.tables)
|
|
150
|
+
selected = [t for t in semantics if t.table in chosen_names]
|
|
151
|
+
notes = list(picked.notes)
|
|
152
|
+
if len(semantics) > _INCLUDE_ALL_TABLE_LIMIT and selection.strategy == "include_all":
|
|
153
|
+
notes.append(
|
|
154
|
+
f"{len(semantics)} tables exceeds the include-everything limit "
|
|
155
|
+
f"({_INCLUDE_ALL_TABLE_LIMIT}); set selection.strategy to lexical/llm to narrow"
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
# When selection narrowed the schema, keep only examples that reference a surviving
|
|
159
|
+
# table (an example about a dropped table is noise for this question). Include-all
|
|
160
|
+
# keeps every example.
|
|
161
|
+
shown_examples = _relevant_examples(examples, selected, narrowed=len(selected) < len(semantics))
|
|
162
|
+
|
|
163
|
+
system = _render_system(
|
|
164
|
+
dialect=dialect,
|
|
165
|
+
instructions=instructions,
|
|
166
|
+
semantics=selected,
|
|
167
|
+
trusted_assets=trusted_assets,
|
|
168
|
+
examples=shown_examples,
|
|
169
|
+
value_matches=picked.value_matches,
|
|
170
|
+
)
|
|
171
|
+
return CompiledContext(
|
|
172
|
+
system=system,
|
|
173
|
+
user=_render_question(question),
|
|
174
|
+
selected_tables=[t.table for t in selected],
|
|
175
|
+
offered_assets=[a.name for a in trusted_assets],
|
|
176
|
+
notes=notes,
|
|
177
|
+
usage=picked.usage,
|
|
178
|
+
selection_strategy=picked.strategy,
|
|
179
|
+
selection_fell_back=picked.fell_back,
|
|
180
|
+
)
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
def _relevant_examples(
|
|
184
|
+
examples: Sequence[Example], selected: Sequence[TableSemantics], *, narrowed: bool
|
|
185
|
+
) -> list[Example]:
|
|
186
|
+
"""Keep examples whose SQL names a selected table; on include-all keep them all.
|
|
187
|
+
|
|
188
|
+
Best-effort substring match on table names — an example that mentions none of the
|
|
189
|
+
surviving tables is dropped only when we actually narrowed, so a project that never
|
|
190
|
+
narrows behaves exactly as before.
|
|
191
|
+
"""
|
|
192
|
+
if not narrowed:
|
|
193
|
+
return list(examples)
|
|
194
|
+
names = [t.table for t in selected]
|
|
195
|
+
kept = [ex for ex in examples if any(_names_table(ex.sql, name) for name in names)]
|
|
196
|
+
# Never strip every example to zero on a narrow — keep the originals if none match,
|
|
197
|
+
# since a few-shot example is cheap and dropping all of them helps nothing.
|
|
198
|
+
return kept or list(examples)
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
def _names_table(sql: str, table: str) -> bool:
|
|
202
|
+
"""Whether ``sql`` references ``table`` as a whole word (case-insensitive)."""
|
|
203
|
+
import re
|
|
204
|
+
|
|
205
|
+
return re.search(rf"\b{re.escape(table)}\b", sql, re.IGNORECASE) is not None
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
# --- rendering (deterministic) --------------------------------------------------
|
|
209
|
+
|
|
210
|
+
|
|
211
|
+
def _render_system(
|
|
212
|
+
*,
|
|
213
|
+
dialect: Dialect,
|
|
214
|
+
instructions: str,
|
|
215
|
+
semantics: Sequence[TableSemantics],
|
|
216
|
+
trusted_assets: Sequence[TrustedAsset],
|
|
217
|
+
examples: Sequence[Example],
|
|
218
|
+
value_matches: Sequence[ValueMatch] = (),
|
|
219
|
+
) -> str:
|
|
220
|
+
blocks: list[str] = [
|
|
221
|
+
f"You are a careful {dialect.value} SQL analyst. Answer questions by writing a "
|
|
222
|
+
"single read-only SELECT, grounded in the semantic layer below. Prefer measures, "
|
|
223
|
+
"filters, and trusted assets over ad-hoc arithmetic.",
|
|
224
|
+
]
|
|
225
|
+
if instructions.strip():
|
|
226
|
+
# Rendered verbatim: instructions.md is author-owned markdown (it brings its
|
|
227
|
+
# own headings), so we don't wrap it in another.
|
|
228
|
+
blocks.append(instructions.strip())
|
|
229
|
+
blocks.append(_render_tables(semantics))
|
|
230
|
+
if value_matches:
|
|
231
|
+
blocks.append(_render_value_matches(value_matches))
|
|
232
|
+
if trusted_assets:
|
|
233
|
+
blocks.append(_render_trusted_assets(trusted_assets))
|
|
234
|
+
if examples:
|
|
235
|
+
blocks.append(_render_examples(examples))
|
|
236
|
+
return "\n\n".join(blocks).strip() + "\n"
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _render_value_matches(value_matches: Sequence[ValueMatch]) -> str:
|
|
240
|
+
"""Grounding hints mapping question literals to canonical declared values (§5.1)."""
|
|
241
|
+
lines = ["# Value hints (question terms mapped to declared column values)"]
|
|
242
|
+
for match in value_matches:
|
|
243
|
+
lines.append(f'- "{match.term}" → {match.table}.{match.column} = {match.value!r}')
|
|
244
|
+
return "\n".join(lines)
|
|
245
|
+
|
|
246
|
+
|
|
247
|
+
def _render_tables(semantics: Sequence[TableSemantics]) -> str:
|
|
248
|
+
lines = ["# Schema"]
|
|
249
|
+
for table in semantics:
|
|
250
|
+
header = f"## {table.table}"
|
|
251
|
+
if table.description:
|
|
252
|
+
header += f" — {table.description}"
|
|
253
|
+
lines.append(header)
|
|
254
|
+
if table.synonyms:
|
|
255
|
+
lines.append(f"synonyms: {', '.join(table.synonyms)}")
|
|
256
|
+
lines.append("columns:")
|
|
257
|
+
for col in table.columns:
|
|
258
|
+
lines.append(_render_column(col))
|
|
259
|
+
for join in table.joins:
|
|
260
|
+
conf = "" if join.confidence is None else f" [confidence {join.confidence:.2f}]"
|
|
261
|
+
lines.append(f"join: {join.type} -> {join.to} ON {join.on}{conf}")
|
|
262
|
+
for measure in table.measures:
|
|
263
|
+
desc = f" — {measure.description}" if measure.description else ""
|
|
264
|
+
lines.append(f"measure {measure.name}: {measure.sql}{desc}")
|
|
265
|
+
for filt in table.filters:
|
|
266
|
+
desc = f" — {filt.description}" if filt.description else ""
|
|
267
|
+
lines.append(f"filter {filt.name}: {filt.sql}{desc}")
|
|
268
|
+
return "\n".join(lines)
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
def _render_column(col: Column) -> str:
|
|
272
|
+
parts = [f"- {col.name} ({col.type})"]
|
|
273
|
+
if col.description:
|
|
274
|
+
parts.append(f" — {col.description}")
|
|
275
|
+
if col.synonyms:
|
|
276
|
+
parts.append(f" [synonyms: {', '.join(col.synonyms)}]")
|
|
277
|
+
hint = _profile_hint(col.profile, col.sample_values)
|
|
278
|
+
if hint:
|
|
279
|
+
parts.append(f" {hint}")
|
|
280
|
+
return "".join(parts)
|
|
281
|
+
|
|
282
|
+
|
|
283
|
+
def _profile_hint(profile: Profile | bool | None, sample_values: list[ScalarBound] | None) -> str:
|
|
284
|
+
"""A compact grounding hint a human would eyeball: range or representative values."""
|
|
285
|
+
if sample_values:
|
|
286
|
+
shown = ", ".join(str(v) for v in sample_values)
|
|
287
|
+
return f"[values: {shown}]"
|
|
288
|
+
if isinstance(profile, Profile) and profile.min is not None and profile.max is not None:
|
|
289
|
+
return f"[range: {profile.min}..{profile.max}]"
|
|
290
|
+
return ""
|
|
291
|
+
|
|
292
|
+
|
|
293
|
+
def _render_trusted_assets(assets: Sequence[TrustedAsset]) -> str:
|
|
294
|
+
lines = ["# Trusted assets (prefer these over ad-hoc SQL; cite the one you use)"]
|
|
295
|
+
for asset in assets:
|
|
296
|
+
params = ", ".join(f"{p.name} {p.type}" for p in asset.params)
|
|
297
|
+
header = f"## {asset.name}({params})"
|
|
298
|
+
if asset.description:
|
|
299
|
+
header += f" — {asset.description}"
|
|
300
|
+
lines.append(header)
|
|
301
|
+
lines.append(asset.sql.strip())
|
|
302
|
+
return "\n".join(lines)
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def _render_examples(examples: Sequence[Example]) -> str:
|
|
306
|
+
lines = ["# Examples"]
|
|
307
|
+
for ex in examples:
|
|
308
|
+
lines.append(f"Q: {ex.question}")
|
|
309
|
+
lines.append("SQL:")
|
|
310
|
+
lines.append(ex.sql.strip())
|
|
311
|
+
return "\n".join(lines)
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
def _render_question(question: str) -> str:
|
|
315
|
+
return (
|
|
316
|
+
f"Question: {question.strip()}\n\n"
|
|
317
|
+
"Write a single read-only SELECT that answers it. Think briefly about which "
|
|
318
|
+
"tables, measures, and trusted assets apply, then give the SQL."
|
|
319
|
+
)
|
sqbyl_runtime/cost.py
ADDED
|
@@ -0,0 +1,265 @@
|
|
|
1
|
+
"""Token pricing + cost estimation + the live spend meter (invariant 5).
|
|
2
|
+
|
|
3
|
+
This module is the single place that turns a :class:`Usage` (or a planned call
|
|
4
|
+
count) into dollars, and the single place that *tracks* dollars as they're spent:
|
|
5
|
+
|
|
6
|
+
* :func:`price_usage` / :func:`estimate_cost` — the pricing primitives.
|
|
7
|
+
* :class:`CostEstimate` — the structured up-front estimate a paid command prints
|
|
8
|
+
(and ``--dry-run`` returns without spending a cent).
|
|
9
|
+
* :class:`SpendMeter` — the live tally a command runs work against: it meters every
|
|
10
|
+
call to ``.sqbyl/usage.db`` and enforces a ``--budget`` cap (hard-stop in ``--auto``,
|
|
11
|
+
pause-and-ask in guided mode via the CLI).
|
|
12
|
+
|
|
13
|
+
Rates are list prices in USD per **million** tokens. They are approximate and easy
|
|
14
|
+
to update; treat them as estimates, not invoices.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import threading
|
|
20
|
+
from dataclasses import dataclass
|
|
21
|
+
from typing import TYPE_CHECKING
|
|
22
|
+
|
|
23
|
+
from pydantic import BaseModel
|
|
24
|
+
|
|
25
|
+
from sqbyl_runtime.llm.base import Usage
|
|
26
|
+
|
|
27
|
+
if TYPE_CHECKING:
|
|
28
|
+
from sqbyl_runtime.state.usage import UsageStore
|
|
29
|
+
|
|
30
|
+
# Costs are floats; a hair of tolerance keeps a cap from tripping on rounding noise
|
|
31
|
+
# (e.g. an estimate of exactly the budget must be allowed to run).
|
|
32
|
+
_EPSILON = 1e-9
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@dataclass(frozen=True)
|
|
36
|
+
class ModelRate:
|
|
37
|
+
"""USD per million tokens for one model."""
|
|
38
|
+
|
|
39
|
+
input: float
|
|
40
|
+
output: float
|
|
41
|
+
cache_write: float
|
|
42
|
+
cache_read: float
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
# Approximate list prices (USD / 1M tokens). Update as pricing changes.
|
|
46
|
+
# OpenAI caches automatically with no separate cache-*write* charge, so cache_write == input
|
|
47
|
+
# for those rows; cache_read is the discounted cached-input rate.
|
|
48
|
+
MODEL_RATES: dict[str, ModelRate] = {
|
|
49
|
+
# Anthropic
|
|
50
|
+
"claude-opus-4-8": ModelRate(input=15.0, output=75.0, cache_write=18.75, cache_read=1.5),
|
|
51
|
+
"claude-sonnet-4-6": ModelRate(input=3.0, output=15.0, cache_write=3.75, cache_read=0.3),
|
|
52
|
+
"claude-haiku-4-5-20251001": ModelRate(input=1.0, output=5.0, cache_write=1.25, cache_read=0.1),
|
|
53
|
+
# OpenAI (gpt-5 family)
|
|
54
|
+
"gpt-5": ModelRate(input=1.25, output=10.0, cache_write=1.25, cache_read=0.125),
|
|
55
|
+
"gpt-5-mini": ModelRate(input=0.25, output=2.0, cache_write=0.25, cache_read=0.025),
|
|
56
|
+
"gpt-5-nano": ModelRate(input=0.05, output=0.4, cache_write=0.05, cache_read=0.005),
|
|
57
|
+
}
|
|
58
|
+
# Fallback when a model id isn't in the table — priced as the flagship so estimates never
|
|
59
|
+
# silently read as $0. Note this same rate prices *actual* metered spend too, so a real call
|
|
60
|
+
# on an unrecognized (often cheaper) model is billed in usage.db as an upper bound, not an
|
|
61
|
+
# exact invoice — add the model to MODEL_RATES to reconcile it precisely.
|
|
62
|
+
_DEFAULT_RATE = MODEL_RATES["claude-opus-4-8"]
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def rate_for(model: str) -> ModelRate:
|
|
66
|
+
return MODEL_RATES.get(model, _DEFAULT_RATE)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def price_usage(usage: Usage, model: str) -> float:
|
|
70
|
+
"""Dollar cost of one metered call's :class:`Usage`."""
|
|
71
|
+
rate = rate_for(model)
|
|
72
|
+
return (
|
|
73
|
+
usage.input_tokens * rate.input
|
|
74
|
+
+ usage.output_tokens * rate.output
|
|
75
|
+
+ usage.cache_creation_input_tokens * rate.cache_write
|
|
76
|
+
+ usage.cache_read_input_tokens * rate.cache_read
|
|
77
|
+
) / 1_000_000.0
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
def estimate_cost(
|
|
81
|
+
*,
|
|
82
|
+
model: str,
|
|
83
|
+
calls: int,
|
|
84
|
+
avg_input_tokens: int,
|
|
85
|
+
avg_output_tokens: int,
|
|
86
|
+
) -> float:
|
|
87
|
+
"""A rough up-front estimate for a planned batch of ``calls`` (no cache assumed)."""
|
|
88
|
+
rate = rate_for(model)
|
|
89
|
+
per_call = (avg_input_tokens * rate.input + avg_output_tokens * rate.output) / 1_000_000.0
|
|
90
|
+
return per_call * calls
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
# ── the structured up-front estimate (spec §5.5, §9) ────────────────────────────────────
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
class EstimateItem(BaseModel):
|
|
97
|
+
"""One planned step in a command's cost estimate — the rows of the SAM-style plan."""
|
|
98
|
+
|
|
99
|
+
label: str
|
|
100
|
+
model: str
|
|
101
|
+
calls: int
|
|
102
|
+
avg_input_tokens: int
|
|
103
|
+
avg_output_tokens: int
|
|
104
|
+
|
|
105
|
+
@property
|
|
106
|
+
def cost_usd(self) -> float:
|
|
107
|
+
return estimate_cost(
|
|
108
|
+
model=self.model,
|
|
109
|
+
calls=self.calls,
|
|
110
|
+
avg_input_tokens=self.avg_input_tokens,
|
|
111
|
+
avg_output_tokens=self.avg_output_tokens,
|
|
112
|
+
)
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
class CostEstimate(BaseModel):
|
|
116
|
+
"""A planned batch of paid work, itemized. Printed before spending and returned by
|
|
117
|
+
``--dry-run`` / ``sqbyl cost`` — the "here's the plan and the estimate" of spec §5.5.
|
|
118
|
+
|
|
119
|
+
It is an *estimate*, not an invoice: no prompt-cache savings are assumed, so the real
|
|
120
|
+
metered spend should come in at or under it.
|
|
121
|
+
"""
|
|
122
|
+
|
|
123
|
+
items: list[EstimateItem] = []
|
|
124
|
+
|
|
125
|
+
@property
|
|
126
|
+
def total_usd(self) -> float:
|
|
127
|
+
return sum(item.cost_usd for item in self.items)
|
|
128
|
+
|
|
129
|
+
@property
|
|
130
|
+
def calls(self) -> int:
|
|
131
|
+
return sum(item.calls for item in self.items)
|
|
132
|
+
|
|
133
|
+
def render(self, *, indent: str = " ") -> str:
|
|
134
|
+
"""The itemized plan as a right-aligned ``label … ~$cost`` table."""
|
|
135
|
+
if not self.items:
|
|
136
|
+
return f"{indent}(no paid work planned)"
|
|
137
|
+
width = max(len(item.label) for item in self.items)
|
|
138
|
+
lines = [
|
|
139
|
+
f"{indent}{item.label.ljust(width)} ~${item.cost_usd:.4f} "
|
|
140
|
+
f"({item.calls}× {item.model})"
|
|
141
|
+
for item in self.items
|
|
142
|
+
]
|
|
143
|
+
lines.append(f"{indent}{'─' * (width + 20)}")
|
|
144
|
+
lines.append(f"{indent}{'estimated total'.ljust(width)} ~${self.total_usd:.4f}")
|
|
145
|
+
return "\n".join(lines)
|
|
146
|
+
|
|
147
|
+
|
|
148
|
+
# ── the live spend meter + budget cap (invariant 5) ─────────────────────────────────────
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
class BudgetError(RuntimeError):
|
|
152
|
+
"""A planned paid step would push spend past the hard cap.
|
|
153
|
+
|
|
154
|
+
Raised by :meth:`SpendMeter.guard` in ``--auto``/hard mode so a headless run stops
|
|
155
|
+
at the cap rather than silently overspending (spec §9). In guided mode the CLI checks
|
|
156
|
+
:meth:`SpendMeter.would_exceed` first and pauses to ask instead.
|
|
157
|
+
"""
|
|
158
|
+
|
|
159
|
+
def __init__(self, *, spent: float, budget: float, attempted: float) -> None:
|
|
160
|
+
self.spent = spent
|
|
161
|
+
self.budget = budget
|
|
162
|
+
self.attempted = attempted
|
|
163
|
+
super().__init__(
|
|
164
|
+
f"budget ${budget:.4f} would be exceeded: ${spent:.4f} spent + "
|
|
165
|
+
f"~${attempted:.4f} planned"
|
|
166
|
+
)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
class SpendMeter:
|
|
170
|
+
"""A running tally of paid spend, capped by an optional ``--budget`` (spec §9).
|
|
171
|
+
|
|
172
|
+
One meter spans a whole command. Every call is priced, added to the tally, and — when a
|
|
173
|
+
store is attached — durably recorded to ``.sqbyl/usage.db``, so the ledger total always
|
|
174
|
+
reconciles with the meter exactly. The cap is advisory via :meth:`would_exceed` (for a
|
|
175
|
+
guided pause) and enforced via :meth:`guard` (a hard stop that raises :class:`BudgetError`).
|
|
176
|
+
|
|
177
|
+
**Concurrency.** The lock makes each *field access* atomic — :meth:`record` and
|
|
178
|
+
:attr:`spent` are safe to call from many threads and the ledger stays exact. It does
|
|
179
|
+
**not** make the cap a concurrency-safe admission gate: :meth:`guard`/:meth:`would_exceed`
|
|
180
|
+
and the later :meth:`record` are separate steps, so N threads can each pass the check
|
|
181
|
+
before any records and collectively overshoot by up to N−1 calls. Serial callers (the CLI
|
|
182
|
+
commands) are fine; a parallel fan-out must bound dispatch with the orchestrator's own
|
|
183
|
+
pre-dispatch budget gate (Phase 6.1) rather than trusting :meth:`guard` alone.
|
|
184
|
+
"""
|
|
185
|
+
|
|
186
|
+
def __init__(
|
|
187
|
+
self,
|
|
188
|
+
*,
|
|
189
|
+
budget: float | None = None,
|
|
190
|
+
store: UsageStore | None = None,
|
|
191
|
+
command: str | None = None,
|
|
192
|
+
content_hash: str | None = None,
|
|
193
|
+
) -> None:
|
|
194
|
+
self._budget = budget
|
|
195
|
+
self._store = store
|
|
196
|
+
self._command = command
|
|
197
|
+
self._content_hash = content_hash
|
|
198
|
+
self._spent = 0.0
|
|
199
|
+
self._lock = threading.Lock()
|
|
200
|
+
|
|
201
|
+
@property
|
|
202
|
+
def budget(self) -> float | None:
|
|
203
|
+
return self._budget
|
|
204
|
+
|
|
205
|
+
@property
|
|
206
|
+
def spent(self) -> float:
|
|
207
|
+
with self._lock:
|
|
208
|
+
return self._spent
|
|
209
|
+
|
|
210
|
+
@property
|
|
211
|
+
def remaining(self) -> float | None:
|
|
212
|
+
"""Dollars left under the cap, or ``None`` when uncapped."""
|
|
213
|
+
if self._budget is None:
|
|
214
|
+
return None
|
|
215
|
+
with self._lock:
|
|
216
|
+
return self._budget - self._spent
|
|
217
|
+
|
|
218
|
+
def would_exceed(self, next_cost: float) -> bool:
|
|
219
|
+
"""Would spending ``next_cost`` more push past the cap? (Never True when uncapped.)"""
|
|
220
|
+
if self._budget is None:
|
|
221
|
+
return False
|
|
222
|
+
with self._lock:
|
|
223
|
+
return self._spent + next_cost > self._budget + _EPSILON
|
|
224
|
+
|
|
225
|
+
def guard(self, next_cost: float) -> None:
|
|
226
|
+
"""Hard-stop precondition: raise :class:`BudgetError` if the next step won't fit."""
|
|
227
|
+
if self.would_exceed(next_cost):
|
|
228
|
+
with self._lock:
|
|
229
|
+
spent = self._spent
|
|
230
|
+
assert self._budget is not None
|
|
231
|
+
raise BudgetError(spent=spent, budget=self._budget, attempted=next_cost)
|
|
232
|
+
|
|
233
|
+
def record(
|
|
234
|
+
self,
|
|
235
|
+
usage: Usage,
|
|
236
|
+
*,
|
|
237
|
+
model: str,
|
|
238
|
+
role: str | None = None,
|
|
239
|
+
run_id: str | None = None,
|
|
240
|
+
) -> float:
|
|
241
|
+
"""Price one completed call's ``usage``, add it to the tally, and persist it.
|
|
242
|
+
|
|
243
|
+
Returns the call's dollar cost. This is metering *after* a call — the cap is
|
|
244
|
+
enforced *before* dispatch via :meth:`guard`/:meth:`would_exceed`, so a recorded
|
|
245
|
+
call can legitimately carry the tally over budget by one call's worth.
|
|
246
|
+
"""
|
|
247
|
+
cost = price_usage(usage, model)
|
|
248
|
+
with self._lock:
|
|
249
|
+
self._spent += cost
|
|
250
|
+
store = self._store
|
|
251
|
+
if store is not None:
|
|
252
|
+
from sqbyl_runtime.state.usage import UsageRecord
|
|
253
|
+
|
|
254
|
+
store.record(
|
|
255
|
+
UsageRecord.from_usage(
|
|
256
|
+
usage,
|
|
257
|
+
model=model,
|
|
258
|
+
command=self._command,
|
|
259
|
+
role=role,
|
|
260
|
+
cost_usd=cost,
|
|
261
|
+
run_id=run_id,
|
|
262
|
+
content_hash=self._content_hash,
|
|
263
|
+
)
|
|
264
|
+
)
|
|
265
|
+
return cost
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
"""Read-only database access — the dialect seam, the SQL guard, and the connection.
|
|
2
|
+
|
|
3
|
+
This is the only door to a SQL database in sqbyl. It lives in ``sqbyl-runtime``
|
|
4
|
+
because the shipped agent executes its generated SQL through exactly this layer
|
|
5
|
+
(spec §5 step 5), so read-only enforcement travels with the runtime.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from sqbyl_runtime.db.connection import Database, QueryResult, resolve_url
|
|
11
|
+
from sqbyl_runtime.db.dialects import (
|
|
12
|
+
BigQueryAdapter,
|
|
13
|
+
DialectAdapter,
|
|
14
|
+
DuckDBAdapter,
|
|
15
|
+
MySQLAdapter,
|
|
16
|
+
PostgresAdapter,
|
|
17
|
+
PrivilegeReport,
|
|
18
|
+
SnowflakeAdapter,
|
|
19
|
+
SQLiteAdapter,
|
|
20
|
+
adapter_for,
|
|
21
|
+
)
|
|
22
|
+
from sqbyl_runtime.db.errors import (
|
|
23
|
+
StaticValidationError,
|
|
24
|
+
UnparseableSqlError,
|
|
25
|
+
WritablePrivilegeWarning,
|
|
26
|
+
WriteAttemptError,
|
|
27
|
+
)
|
|
28
|
+
from sqbyl_runtime.db.guard import assert_read_only, is_read_only
|
|
29
|
+
|
|
30
|
+
__all__ = [
|
|
31
|
+
"BigQueryAdapter",
|
|
32
|
+
"Database",
|
|
33
|
+
"DialectAdapter",
|
|
34
|
+
"DuckDBAdapter",
|
|
35
|
+
"MySQLAdapter",
|
|
36
|
+
"PostgresAdapter",
|
|
37
|
+
"PrivilegeReport",
|
|
38
|
+
"QueryResult",
|
|
39
|
+
"SQLiteAdapter",
|
|
40
|
+
"SnowflakeAdapter",
|
|
41
|
+
"StaticValidationError",
|
|
42
|
+
"UnparseableSqlError",
|
|
43
|
+
"WritablePrivilegeWarning",
|
|
44
|
+
"WriteAttemptError",
|
|
45
|
+
"adapter_for",
|
|
46
|
+
"assert_read_only",
|
|
47
|
+
"is_read_only",
|
|
48
|
+
"resolve_url",
|
|
49
|
+
]
|