semora-store-pg 0.1.0__tar.gz
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.
- semora_store_pg-0.1.0/.gitignore +28 -0
- semora_store_pg-0.1.0/PKG-INFO +17 -0
- semora_store_pg-0.1.0/pyproject.toml +31 -0
- semora_store_pg-0.1.0/src/semora_store_pg/__init__.py +17 -0
- semora_store_pg-0.1.0/src/semora_store_pg/postgres.py +369 -0
- semora_store_pg-0.1.0/src/semora_store_pg/py.typed +0 -0
- semora_store_pg-0.1.0/src/semora_store_pg/transcript.py +195 -0
|
@@ -0,0 +1,28 @@
|
|
|
1
|
+
.venv/
|
|
2
|
+
.pytest_cache/
|
|
3
|
+
.mypy_cache/
|
|
4
|
+
.ruff_cache/
|
|
5
|
+
__pycache__/
|
|
6
|
+
*.py[cod]
|
|
7
|
+
*.egg-info/
|
|
8
|
+
build/
|
|
9
|
+
dist/
|
|
10
|
+
.coverage
|
|
11
|
+
htmlcov/
|
|
12
|
+
.env
|
|
13
|
+
.env.*
|
|
14
|
+
!.env.example
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
# Local tool/editor state — machine-specific, never pushed.
|
|
18
|
+
.claude/
|
|
19
|
+
.codecanvas/
|
|
20
|
+
.vscode/
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
# Superpowers design/spec scratch — working notes, not project documentation.
|
|
24
|
+
docs/superpowers/
|
|
25
|
+
|
|
26
|
+
# 로컬 자격증명 — 절대 커밋 금지.
|
|
27
|
+
a.txt
|
|
28
|
+
*.token
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: semora-store-pg
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: PostgreSQL adapters for Semora storage-semantic contracts.
|
|
5
|
+
Project-URL: Homepage, https://github.com/donggyun112/semora
|
|
6
|
+
Project-URL: Source, https://github.com/donggyun112/semora
|
|
7
|
+
Project-URL: Changelog, https://github.com/donggyun112/semora/blob/main/CHANGELOG.md
|
|
8
|
+
Author: donggyun112
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
Classifier: Development Status :: 4 - Beta
|
|
11
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
14
|
+
Classifier: Typing :: Typed
|
|
15
|
+
Requires-Python: >=3.12
|
|
16
|
+
Requires-Dist: psycopg[binary,pool]<4,>=3.3
|
|
17
|
+
Requires-Dist: semora-store==0.1.0
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling>=1.27"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "semora-store-pg"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "PostgreSQL adapters for Semora storage-semantic contracts."
|
|
9
|
+
requires-python = ">=3.12"
|
|
10
|
+
license = "MIT"
|
|
11
|
+
authors = [{ name = "donggyun112" }]
|
|
12
|
+
classifiers = [
|
|
13
|
+
"Development Status :: 4 - Beta",
|
|
14
|
+
"License :: OSI Approved :: MIT License",
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Programming Language :: Python :: 3.12",
|
|
17
|
+
"Typing :: Typed",
|
|
18
|
+
]
|
|
19
|
+
urls = { Homepage = "https://github.com/donggyun112/semora", Source = "https://github.com/donggyun112/semora", Changelog = "https://github.com/donggyun112/semora/blob/main/CHANGELOG.md" }
|
|
20
|
+
dependencies = [
|
|
21
|
+
"semora-store==0.1.0",
|
|
22
|
+
"psycopg[binary,pool]>=3.3,<4",
|
|
23
|
+
]
|
|
24
|
+
# Not `semora`. This implements the ledger protocol and nothing above it, which is what keeps a
|
|
25
|
+
# database driver out of the default install.
|
|
26
|
+
|
|
27
|
+
[tool.uv.sources]
|
|
28
|
+
semora-store = { workspace = true }
|
|
29
|
+
|
|
30
|
+
[tool.hatch.build.targets.wheel]
|
|
31
|
+
packages = ["src/semora_store_pg"]
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""Expose PostgreSQL-backed step ledger and transcript implementations."""
|
|
2
|
+
|
|
3
|
+
from .postgres import (
|
|
4
|
+
SCHEMA,
|
|
5
|
+
PostgresSteps,
|
|
6
|
+
)
|
|
7
|
+
from .transcript import (
|
|
8
|
+
TRANSCRIPT_SCHEMA,
|
|
9
|
+
PostgresTranscript,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
__all__ = [
|
|
13
|
+
"SCHEMA",
|
|
14
|
+
"TRANSCRIPT_SCHEMA",
|
|
15
|
+
"PostgresSteps",
|
|
16
|
+
"PostgresTranscript",
|
|
17
|
+
]
|
|
@@ -0,0 +1,369 @@
|
|
|
1
|
+
"""Implement the durable step ledger with PostgreSQL."""
|
|
2
|
+
|
|
3
|
+
from typing import Any, Self
|
|
4
|
+
|
|
5
|
+
from psycopg import AsyncConnection
|
|
6
|
+
from psycopg.rows import dict_row
|
|
7
|
+
from psycopg.types.json import Jsonb
|
|
8
|
+
from psycopg_pool import AsyncConnectionPool
|
|
9
|
+
from semora_store import (
|
|
10
|
+
EffectConflict,
|
|
11
|
+
ExecutionContext,
|
|
12
|
+
ExecutionTransition,
|
|
13
|
+
Fenced,
|
|
14
|
+
InputRecord,
|
|
15
|
+
Step,
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
__all__ = ["SCHEMA", "PostgresSteps"]
|
|
19
|
+
|
|
20
|
+
SCHEMA = """
|
|
21
|
+
create table if not exists ledger_step (
|
|
22
|
+
run_id text not null,
|
|
23
|
+
key text not null,
|
|
24
|
+
status text not null check (status in ('running', 'done')),
|
|
25
|
+
value jsonb,
|
|
26
|
+
attempt integer not null default 1,
|
|
27
|
+
started_at timestamptz not null default now(),
|
|
28
|
+
finished_at timestamptz,
|
|
29
|
+
primary key (run_id, key)
|
|
30
|
+
);
|
|
31
|
+
|
|
32
|
+
create table if not exists ledger_run_lease (
|
|
33
|
+
run_id text primary key,
|
|
34
|
+
owner text not null,
|
|
35
|
+
token bigint not null default 1,
|
|
36
|
+
expires_at timestamptz not null
|
|
37
|
+
);
|
|
38
|
+
|
|
39
|
+
create table if not exists ledger_input (
|
|
40
|
+
sequence bigserial primary key,
|
|
41
|
+
run_id text not null,
|
|
42
|
+
input_id text not null,
|
|
43
|
+
status text not null,
|
|
44
|
+
value jsonb not null,
|
|
45
|
+
submitted_at timestamptz not null default now(),
|
|
46
|
+
admitted_at timestamptz,
|
|
47
|
+
unique (run_id, input_id)
|
|
48
|
+
);
|
|
49
|
+
|
|
50
|
+
alter table ledger_input drop constraint if exists ledger_input_status_check;
|
|
51
|
+
alter table ledger_input add constraint ledger_input_status_check
|
|
52
|
+
check (status in ('pending', 'claimed', 'admitted', 'discarded'));
|
|
53
|
+
|
|
54
|
+
drop index if exists ledger_input_pending;
|
|
55
|
+
create index if not exists ledger_input_pending
|
|
56
|
+
on ledger_input (run_id, sequence) where status in ('pending', 'claimed');
|
|
57
|
+
|
|
58
|
+
create index if not exists ledger_step_running
|
|
59
|
+
on ledger_step (started_at) where status = 'running';
|
|
60
|
+
"""
|
|
61
|
+
"""The three tables. The partial indexes serve recovery and operator stuck-work queries."""
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class PostgresSteps:
|
|
65
|
+
"""Persist steps, leases, and queued inputs in PostgreSQL."""
|
|
66
|
+
|
|
67
|
+
def __init__(self, pool: AsyncConnectionPool[AsyncConnection[Any]]) -> None:
|
|
68
|
+
"""Initialize the ledger with a caller-owned connection pool."""
|
|
69
|
+
self._pool = pool
|
|
70
|
+
|
|
71
|
+
def for_execution(self, context: ExecutionContext) -> Self:
|
|
72
|
+
"""Return this adapter because its default schema is scope-neutral."""
|
|
73
|
+
del context
|
|
74
|
+
return self
|
|
75
|
+
|
|
76
|
+
async def read(self, run_id: str, key: str) -> Step:
|
|
77
|
+
"""Return the persisted state of a step."""
|
|
78
|
+
async with (
|
|
79
|
+
self._pool.connection() as connection,
|
|
80
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
81
|
+
):
|
|
82
|
+
await cursor.execute(
|
|
83
|
+
"select status, value from ledger_step where run_id = %s and key = %s",
|
|
84
|
+
(run_id, key),
|
|
85
|
+
)
|
|
86
|
+
row = await cursor.fetchone()
|
|
87
|
+
if row is None:
|
|
88
|
+
return Step("absent")
|
|
89
|
+
if row["status"] == "running":
|
|
90
|
+
return Step("running")
|
|
91
|
+
return Step("done", row["value"])
|
|
92
|
+
|
|
93
|
+
async def start(self, run_id: str, key: str, token: int = 0) -> bool:
|
|
94
|
+
"""Atomically record new running intent after validating the fencing token."""
|
|
95
|
+
async with self._pool.connection() as connection:
|
|
96
|
+
await self._fence(connection, run_id, token)
|
|
97
|
+
async with connection.cursor() as cursor:
|
|
98
|
+
await cursor.execute(
|
|
99
|
+
"""
|
|
100
|
+
insert into ledger_step (run_id, key, status)
|
|
101
|
+
values (%s, %s, 'running')
|
|
102
|
+
on conflict (run_id, key) do nothing
|
|
103
|
+
returning 1
|
|
104
|
+
""",
|
|
105
|
+
(run_id, key),
|
|
106
|
+
)
|
|
107
|
+
return await cursor.fetchone() is not None
|
|
108
|
+
|
|
109
|
+
async def finish_effect(self, run_id: str, key: str, value: Any, token: int = 0) -> None:
|
|
110
|
+
"""Complete a running effect without replacing a committed result."""
|
|
111
|
+
async with self._pool.connection() as connection:
|
|
112
|
+
await self._fence(connection, run_id, token)
|
|
113
|
+
await self._complete_effect(connection, run_id, key, value, "running")
|
|
114
|
+
|
|
115
|
+
async def _complete_effect(
|
|
116
|
+
self,
|
|
117
|
+
connection: AsyncConnection[Any],
|
|
118
|
+
run_id: str,
|
|
119
|
+
key: str,
|
|
120
|
+
value: Any,
|
|
121
|
+
expected: str,
|
|
122
|
+
) -> None:
|
|
123
|
+
"""Apply one conditional effect completion inside the caller's transaction."""
|
|
124
|
+
async with connection.cursor(row_factory=dict_row) as cursor:
|
|
125
|
+
if expected == "absent":
|
|
126
|
+
await cursor.execute(
|
|
127
|
+
"""
|
|
128
|
+
insert into ledger_step (run_id, key, status, value, finished_at)
|
|
129
|
+
values (%s, %s, 'done', %s, now())
|
|
130
|
+
on conflict (run_id, key) do nothing
|
|
131
|
+
returning 1
|
|
132
|
+
""",
|
|
133
|
+
(run_id, key, Jsonb(value)),
|
|
134
|
+
)
|
|
135
|
+
else:
|
|
136
|
+
await cursor.execute(
|
|
137
|
+
"""
|
|
138
|
+
update ledger_step
|
|
139
|
+
set status = 'done', value = %s, finished_at = now()
|
|
140
|
+
where run_id = %s and key = %s and status = 'running'
|
|
141
|
+
returning 1
|
|
142
|
+
""",
|
|
143
|
+
(Jsonb(value), run_id, key),
|
|
144
|
+
)
|
|
145
|
+
if await cursor.fetchone() is not None:
|
|
146
|
+
return
|
|
147
|
+
await cursor.execute(
|
|
148
|
+
"select status, value from ledger_step where run_id = %s and key = %s",
|
|
149
|
+
(run_id, key),
|
|
150
|
+
)
|
|
151
|
+
row = await cursor.fetchone()
|
|
152
|
+
if row is not None and row["status"] == "done":
|
|
153
|
+
if row["value"] == value:
|
|
154
|
+
return
|
|
155
|
+
raise EffectConflict(run_id, key, "a different result is already committed")
|
|
156
|
+
status = "absent" if row is None else str(row["status"])
|
|
157
|
+
raise EffectConflict(run_id, key, f"expected {expected!r}, found {status!r}")
|
|
158
|
+
|
|
159
|
+
async def write_control(self, run_id: str, key: str, value: Any, token: int = 0) -> None:
|
|
160
|
+
"""Upsert mutable framework control state."""
|
|
161
|
+
async with self._pool.connection() as connection:
|
|
162
|
+
await self._fence(connection, run_id, token)
|
|
163
|
+
async with connection.cursor() as cursor:
|
|
164
|
+
await cursor.execute(
|
|
165
|
+
"""
|
|
166
|
+
insert into ledger_step (run_id, key, status, value, finished_at)
|
|
167
|
+
values (%s, %s, 'done', %s, now())
|
|
168
|
+
on conflict (run_id, key) do update
|
|
169
|
+
set status = 'done', value = excluded.value, finished_at = now()
|
|
170
|
+
""",
|
|
171
|
+
(run_id, key, Jsonb(value)),
|
|
172
|
+
)
|
|
173
|
+
|
|
174
|
+
async def finish(self, run_id: str, key: str, value: Any, token: int = 0) -> None:
|
|
175
|
+
"""Compatibility alias for ``write_control``."""
|
|
176
|
+
await self.write_control(run_id, key, value, token)
|
|
177
|
+
|
|
178
|
+
async def _fence(self, connection: AsyncConnection[Any], run_id: str, token: int) -> None:
|
|
179
|
+
"""Reject stale lease holders within the caller's transaction."""
|
|
180
|
+
if not token:
|
|
181
|
+
return
|
|
182
|
+
async with connection.cursor(row_factory=dict_row) as cursor:
|
|
183
|
+
await cursor.execute("select token from ledger_run_lease where run_id = %s", (run_id,))
|
|
184
|
+
row = await cursor.fetchone()
|
|
185
|
+
issued = row["token"] if row else 0
|
|
186
|
+
if token < issued:
|
|
187
|
+
raise Fenced(run_id, token, issued)
|
|
188
|
+
|
|
189
|
+
async def forget(self, run_id: str, key: str, token: int = 0) -> None:
|
|
190
|
+
"""Remove an unfinished step after validating the fencing token."""
|
|
191
|
+
async with self._pool.connection() as connection:
|
|
192
|
+
await self._fence(connection, run_id, token)
|
|
193
|
+
async with connection.cursor() as cursor:
|
|
194
|
+
await cursor.execute(
|
|
195
|
+
"delete from ledger_step where run_id = %s and key = %s and status = 'running'",
|
|
196
|
+
(run_id, key),
|
|
197
|
+
)
|
|
198
|
+
|
|
199
|
+
async def acquire(self, run_id: str, owner: str, ttl_seconds: float = 60.0) -> int:
|
|
200
|
+
"""Acquire or renew a run lease and return its fencing token, or zero on contention."""
|
|
201
|
+
async with (
|
|
202
|
+
self._pool.connection() as connection,
|
|
203
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
204
|
+
):
|
|
205
|
+
await cursor.execute(
|
|
206
|
+
"""
|
|
207
|
+
insert into ledger_run_lease (run_id, owner, token, expires_at)
|
|
208
|
+
values (%s, %s, 1, now() + make_interval(secs => %s))
|
|
209
|
+
on conflict (run_id) do update
|
|
210
|
+
set owner = excluded.owner,
|
|
211
|
+
expires_at = excluded.expires_at,
|
|
212
|
+
token = ledger_run_lease.token
|
|
213
|
+
+ (ledger_run_lease.owner <> excluded.owner)::int
|
|
214
|
+
where ledger_run_lease.owner = excluded.owner
|
|
215
|
+
or ledger_run_lease.expires_at < now()
|
|
216
|
+
returning token
|
|
217
|
+
""",
|
|
218
|
+
(run_id, owner, ttl_seconds),
|
|
219
|
+
)
|
|
220
|
+
row = await cursor.fetchone()
|
|
221
|
+
return int(row["token"]) if row else 0
|
|
222
|
+
|
|
223
|
+
async def release(self, run_id: str, owner: str) -> None:
|
|
224
|
+
"""Expire an owned lease without resetting its fencing token.
|
|
225
|
+
|
|
226
|
+
Args:
|
|
227
|
+
run_id: Durable run identifier.
|
|
228
|
+
owner: Current lease owner.
|
|
229
|
+
"""
|
|
230
|
+
async with self._pool.connection() as connection, connection.cursor() as cursor:
|
|
231
|
+
await cursor.execute(
|
|
232
|
+
"""
|
|
233
|
+
update ledger_run_lease
|
|
234
|
+
set owner = '', expires_at = to_timestamp(0)
|
|
235
|
+
where run_id = %s and owner = %s
|
|
236
|
+
""",
|
|
237
|
+
(run_id, owner),
|
|
238
|
+
)
|
|
239
|
+
|
|
240
|
+
async def enqueue_input(self, run_id: str, input_id: str, value: dict[str, Any]) -> bool:
|
|
241
|
+
"""Append an input idempotently and report whether it was inserted."""
|
|
242
|
+
async with self._pool.connection() as connection, connection.cursor() as cursor:
|
|
243
|
+
await cursor.execute(
|
|
244
|
+
"""
|
|
245
|
+
insert into ledger_input (run_id, input_id, status, value)
|
|
246
|
+
values (%s, %s, 'pending', %s)
|
|
247
|
+
on conflict (run_id, input_id) do nothing
|
|
248
|
+
""",
|
|
249
|
+
(run_id, input_id, Jsonb(value)),
|
|
250
|
+
)
|
|
251
|
+
inserted = cursor.rowcount > 0
|
|
252
|
+
return inserted
|
|
253
|
+
|
|
254
|
+
async def list_inputs(self, run_id: str) -> list[InputRecord]:
|
|
255
|
+
"""Return a run's inputs in submission order."""
|
|
256
|
+
async with (
|
|
257
|
+
self._pool.connection() as connection,
|
|
258
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
259
|
+
):
|
|
260
|
+
await cursor.execute(
|
|
261
|
+
"""
|
|
262
|
+
select input_id, status, value, sequence
|
|
263
|
+
from ledger_input
|
|
264
|
+
where run_id = %s
|
|
265
|
+
order by sequence
|
|
266
|
+
""",
|
|
267
|
+
(run_id,),
|
|
268
|
+
)
|
|
269
|
+
rows = await cursor.fetchall()
|
|
270
|
+
return [
|
|
271
|
+
InputRecord(row["input_id"], row["status"], row["value"], int(row["sequence"]))
|
|
272
|
+
for row in rows
|
|
273
|
+
]
|
|
274
|
+
|
|
275
|
+
async def claim_input(self, run_id: str, input_id: str, token: int = 0) -> None:
|
|
276
|
+
"""Mark an input as claimed unless it is already terminal."""
|
|
277
|
+
async with self._pool.connection() as connection:
|
|
278
|
+
await self._fence(connection, run_id, token)
|
|
279
|
+
async with connection.cursor() as cursor:
|
|
280
|
+
await cursor.execute(
|
|
281
|
+
"""
|
|
282
|
+
update ledger_input
|
|
283
|
+
set status = 'claimed'
|
|
284
|
+
where run_id = %s and input_id = %s
|
|
285
|
+
and status not in ('admitted', 'discarded')
|
|
286
|
+
""",
|
|
287
|
+
(run_id, input_id),
|
|
288
|
+
)
|
|
289
|
+
|
|
290
|
+
async def admit_inputs(self, run_id: str, input_ids: list[str], token: int = 0) -> None:
|
|
291
|
+
"""Mark the selected inputs as admitted."""
|
|
292
|
+
if not input_ids:
|
|
293
|
+
return
|
|
294
|
+
async with self._pool.connection() as connection:
|
|
295
|
+
await self._fence(connection, run_id, token)
|
|
296
|
+
async with connection.cursor() as cursor:
|
|
297
|
+
await cursor.execute(
|
|
298
|
+
"""
|
|
299
|
+
update ledger_input
|
|
300
|
+
set status = 'admitted', admitted_at = now()
|
|
301
|
+
where run_id = %s and input_id = any(%s) and status <> 'discarded'
|
|
302
|
+
""",
|
|
303
|
+
(run_id, input_ids),
|
|
304
|
+
)
|
|
305
|
+
|
|
306
|
+
async def discard_inputs(self, run_id: str, input_ids: list[str], token: int = 0) -> None:
|
|
307
|
+
"""Make screened-out inputs terminal without deleting their idempotency keys."""
|
|
308
|
+
if not input_ids:
|
|
309
|
+
return
|
|
310
|
+
async with self._pool.connection() as connection:
|
|
311
|
+
await self._fence(connection, run_id, token)
|
|
312
|
+
async with connection.cursor() as cursor:
|
|
313
|
+
await cursor.execute(
|
|
314
|
+
"""
|
|
315
|
+
update ledger_input
|
|
316
|
+
set status = 'discarded'
|
|
317
|
+
where run_id = %s and input_id = any(%s)
|
|
318
|
+
""",
|
|
319
|
+
(run_id, input_ids),
|
|
320
|
+
)
|
|
321
|
+
|
|
322
|
+
async def commit_transition(
|
|
323
|
+
self,
|
|
324
|
+
run_id: str,
|
|
325
|
+
transition: ExecutionTransition,
|
|
326
|
+
token: int = 0,
|
|
327
|
+
) -> set[str]:
|
|
328
|
+
"""Atomically apply typed effect, control, and input changes."""
|
|
329
|
+
inserted: set[str] = set()
|
|
330
|
+
async with self._pool.connection() as connection:
|
|
331
|
+
await self._fence(connection, run_id, token)
|
|
332
|
+
seen: set[str] = set()
|
|
333
|
+
for effect in transition.effects:
|
|
334
|
+
if effect.key in seen:
|
|
335
|
+
raise ValueError(f"effect {effect.key!r} appears twice in one transition")
|
|
336
|
+
seen.add(effect.key)
|
|
337
|
+
overlap = seen.intersection(transition.controls)
|
|
338
|
+
if overlap:
|
|
339
|
+
raise ValueError(
|
|
340
|
+
f"transition classifies keys as both effect and control: {sorted(overlap)}"
|
|
341
|
+
)
|
|
342
|
+
for effect in transition.effects:
|
|
343
|
+
await self._complete_effect(
|
|
344
|
+
connection, run_id, effect.key, effect.value, effect.expected
|
|
345
|
+
)
|
|
346
|
+
async with connection.cursor() as cursor:
|
|
347
|
+
for key, value in transition.controls.items():
|
|
348
|
+
await cursor.execute(
|
|
349
|
+
"""
|
|
350
|
+
insert into ledger_step (run_id, key, status, value, finished_at)
|
|
351
|
+
values (%s, %s, 'done', %s, now())
|
|
352
|
+
on conflict (run_id, key) do update
|
|
353
|
+
set status = 'done', value = excluded.value, finished_at = now()
|
|
354
|
+
""",
|
|
355
|
+
(run_id, key, Jsonb(value)),
|
|
356
|
+
)
|
|
357
|
+
for input_id, value in transition.inputs:
|
|
358
|
+
await cursor.execute(
|
|
359
|
+
"""
|
|
360
|
+
insert into ledger_input (run_id, input_id, status, value)
|
|
361
|
+
values (%s, %s, 'pending', %s)
|
|
362
|
+
on conflict (run_id, input_id) do nothing
|
|
363
|
+
returning input_id
|
|
364
|
+
""",
|
|
365
|
+
(run_id, input_id, Jsonb(value)),
|
|
366
|
+
)
|
|
367
|
+
if await cursor.fetchone() is not None:
|
|
368
|
+
inserted.add(input_id)
|
|
369
|
+
return inserted
|
|
File without changes
|
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
"""PostgreSQL implementation of the append-only transcript store.
|
|
2
|
+
|
|
3
|
+
Entries are persisted verbatim as JSONB. Generated columns provide indexed transcript metadata,
|
|
4
|
+
while run and model-usage tables store operational accounting data.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from typing import Any, Self
|
|
8
|
+
|
|
9
|
+
from psycopg import AsyncConnection
|
|
10
|
+
from psycopg.rows import dict_row
|
|
11
|
+
from psycopg.types.json import Jsonb
|
|
12
|
+
from psycopg_pool import AsyncConnectionPool
|
|
13
|
+
from semora_store import MODEL_USAGE_FIELDS, RUN_FIELDS, ExecutionContext, check_fields
|
|
14
|
+
|
|
15
|
+
__all__ = ["TRANSCRIPT_SCHEMA", "PostgresTranscript"]
|
|
16
|
+
|
|
17
|
+
TRANSCRIPT_SCHEMA = """
|
|
18
|
+
create table if not exists ledger_transcript (
|
|
19
|
+
entry jsonb not null,
|
|
20
|
+
seq bigserial primary key,
|
|
21
|
+
ts timestamptz not null,
|
|
22
|
+
conversation_id text generated always as (entry ->> 'conversation_id') stored,
|
|
23
|
+
uuid text generated always as (entry ->> 'uuid') stored,
|
|
24
|
+
parent_uuid text generated always as (entry ->> 'parent_uuid') stored,
|
|
25
|
+
type text generated always as (entry ->> 'type') stored,
|
|
26
|
+
schema_version text generated always as (entry ->> 'schema_version') stored,
|
|
27
|
+
run_id text generated always as (entry -> 'metadata' ->> 'run_id') stored,
|
|
28
|
+
unique (conversation_id, uuid)
|
|
29
|
+
);
|
|
30
|
+
|
|
31
|
+
create index if not exists ledger_transcript_replay
|
|
32
|
+
on ledger_transcript (conversation_id, seq);
|
|
33
|
+
|
|
34
|
+
create index if not exists ledger_transcript_run
|
|
35
|
+
on ledger_transcript (run_id) where run_id is not null;
|
|
36
|
+
|
|
37
|
+
create table if not exists ledger_run (
|
|
38
|
+
run_id text primary key,
|
|
39
|
+
conversation_id text,
|
|
40
|
+
stop_reason text,
|
|
41
|
+
tool_calls integer,
|
|
42
|
+
interrupted_mid_turn boolean,
|
|
43
|
+
started_at timestamptz,
|
|
44
|
+
ended_at timestamptz
|
|
45
|
+
);
|
|
46
|
+
|
|
47
|
+
create index if not exists ledger_run_conversation
|
|
48
|
+
on ledger_run (conversation_id, started_at desc);
|
|
49
|
+
|
|
50
|
+
create table if not exists ledger_run_model (
|
|
51
|
+
run_id text not null,
|
|
52
|
+
model text not null,
|
|
53
|
+
prompt_tokens bigint,
|
|
54
|
+
completion_tokens bigint,
|
|
55
|
+
total_tokens bigint,
|
|
56
|
+
cached_tokens bigint,
|
|
57
|
+
cache_write_tokens bigint,
|
|
58
|
+
cost_usd numeric(12, 6),
|
|
59
|
+
primary key (run_id, model)
|
|
60
|
+
);
|
|
61
|
+
|
|
62
|
+
create index if not exists ledger_run_model_by_model
|
|
63
|
+
on ledger_run_model (model);
|
|
64
|
+
"""
|
|
65
|
+
"""Schema for transcript entries, run metadata, and per-model usage records."""
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
class PostgresTranscript:
|
|
69
|
+
"""Persist conversation entries and per-run cost in PostgreSQL."""
|
|
70
|
+
|
|
71
|
+
def __init__(self, pool: AsyncConnectionPool[AsyncConnection[Any]]) -> None:
|
|
72
|
+
"""Initialize the store with a caller-owned connection pool."""
|
|
73
|
+
self._pool = pool
|
|
74
|
+
|
|
75
|
+
def for_execution(self, context: ExecutionContext) -> Self:
|
|
76
|
+
"""Return this adapter because its default schema is scope-neutral."""
|
|
77
|
+
del context
|
|
78
|
+
return self
|
|
79
|
+
|
|
80
|
+
async def append(self, entry: dict[str, Any]) -> bool:
|
|
81
|
+
"""Append an entry verbatim and report whether it was new."""
|
|
82
|
+
async with self._pool.connection() as connection, connection.cursor() as cursor:
|
|
83
|
+
await cursor.execute(
|
|
84
|
+
"""
|
|
85
|
+
insert into ledger_transcript (entry, ts)
|
|
86
|
+
values (%s, coalesce((%s)::timestamptz, now()))
|
|
87
|
+
on conflict (conversation_id, uuid) do nothing
|
|
88
|
+
returning seq
|
|
89
|
+
""",
|
|
90
|
+
(Jsonb(entry), entry.get("timestamp")),
|
|
91
|
+
)
|
|
92
|
+
inserted = await cursor.fetchone()
|
|
93
|
+
return inserted is not None
|
|
94
|
+
|
|
95
|
+
async def read(self, conversation_id: str, *, limit: int | None = None) -> list[dict[str, Any]]:
|
|
96
|
+
"""Return entries in arrival order, optionally limited to the newest tail."""
|
|
97
|
+
async with (
|
|
98
|
+
self._pool.connection() as connection,
|
|
99
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
100
|
+
):
|
|
101
|
+
if limit is None:
|
|
102
|
+
await cursor.execute(
|
|
103
|
+
"""
|
|
104
|
+
select entry from ledger_transcript
|
|
105
|
+
where conversation_id = %s
|
|
106
|
+
order by seq
|
|
107
|
+
""",
|
|
108
|
+
(conversation_id,),
|
|
109
|
+
)
|
|
110
|
+
else:
|
|
111
|
+
await cursor.execute(
|
|
112
|
+
"""
|
|
113
|
+
select entry from (
|
|
114
|
+
select entry, seq from ledger_transcript
|
|
115
|
+
where conversation_id = %s
|
|
116
|
+
order by seq desc
|
|
117
|
+
limit %s
|
|
118
|
+
) tail
|
|
119
|
+
order by seq
|
|
120
|
+
""",
|
|
121
|
+
(conversation_id, max(limit, 0)),
|
|
122
|
+
)
|
|
123
|
+
rows = await cursor.fetchall()
|
|
124
|
+
return [row["entry"] for row in rows]
|
|
125
|
+
|
|
126
|
+
async def record_run(self, run_id: str, fields: dict[str, Any]) -> None:
|
|
127
|
+
"""Merge validated fields into a run record, creating it when absent."""
|
|
128
|
+
await self._upsert("ledger_run", ("run_id",), (run_id,), fields, RUN_FIELDS)
|
|
129
|
+
|
|
130
|
+
async def record_model_usage(self, run_id: str, model: str, counts: dict[str, Any]) -> None:
|
|
131
|
+
"""Merge validated token counts for one model of one run."""
|
|
132
|
+
await self._upsert(
|
|
133
|
+
"ledger_run_model", ("run_id", "model"), (run_id, model), counts, MODEL_USAGE_FIELDS
|
|
134
|
+
)
|
|
135
|
+
|
|
136
|
+
async def read_run(self, run_id: str) -> dict[str, Any] | None:
|
|
137
|
+
"""Return a run record with unwritten null fields omitted."""
|
|
138
|
+
async with (
|
|
139
|
+
self._pool.connection() as connection,
|
|
140
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
141
|
+
):
|
|
142
|
+
await cursor.execute(
|
|
143
|
+
f"select {', '.join(sorted(RUN_FIELDS))} from ledger_run where run_id = %s",
|
|
144
|
+
(run_id,),
|
|
145
|
+
)
|
|
146
|
+
row = await cursor.fetchone()
|
|
147
|
+
return None if row is None else {k: v for k, v in row.items() if v is not None}
|
|
148
|
+
|
|
149
|
+
async def read_model_usage(self, run_id: str) -> dict[str, dict[str, Any]]:
|
|
150
|
+
"""This run's token counts keyed by the model that spent them."""
|
|
151
|
+
async with (
|
|
152
|
+
self._pool.connection() as connection,
|
|
153
|
+
connection.cursor(row_factory=dict_row) as cursor,
|
|
154
|
+
):
|
|
155
|
+
await cursor.execute(
|
|
156
|
+
f"""
|
|
157
|
+
select model, {", ".join(sorted(MODEL_USAGE_FIELDS))}
|
|
158
|
+
from ledger_run_model where run_id = %s
|
|
159
|
+
""",
|
|
160
|
+
(run_id,),
|
|
161
|
+
)
|
|
162
|
+
rows = await cursor.fetchall()
|
|
163
|
+
return {
|
|
164
|
+
row["model"]: {k: v for k, v in row.items() if k != "model" and v is not None}
|
|
165
|
+
for row in rows
|
|
166
|
+
}
|
|
167
|
+
|
|
168
|
+
async def _upsert(
|
|
169
|
+
self,
|
|
170
|
+
table: str,
|
|
171
|
+
keys: tuple[str, ...],
|
|
172
|
+
key_values: tuple[Any, ...],
|
|
173
|
+
fields: dict[str, Any],
|
|
174
|
+
allowed: frozenset[str],
|
|
175
|
+
) -> None:
|
|
176
|
+
"""Merge allow-listed fields into a keyed record."""
|
|
177
|
+
check_fields(table, fields, allowed)
|
|
178
|
+
columns = sorted(fields)
|
|
179
|
+
named = ", ".join((*keys, *columns))
|
|
180
|
+
placeholders = ", ".join(["%s"] * (len(keys) + len(columns)))
|
|
181
|
+
# No column to merge means the row only has to exist — an insert that keeps what is there.
|
|
182
|
+
merge = (
|
|
183
|
+
", ".join(f"{column} = excluded.{column}" for column in columns)
|
|
184
|
+
if columns
|
|
185
|
+
else f"{keys[0]} = {table}.{keys[0]}"
|
|
186
|
+
)
|
|
187
|
+
async with self._pool.connection() as connection, connection.cursor() as cursor:
|
|
188
|
+
await cursor.execute(
|
|
189
|
+
f"""
|
|
190
|
+
insert into {table} ({named})
|
|
191
|
+
values ({placeholders})
|
|
192
|
+
on conflict ({", ".join(keys)}) do update set {merge}
|
|
193
|
+
""",
|
|
194
|
+
(*key_values, *(fields[column] for column in columns)),
|
|
195
|
+
)
|