dlightrag-memory 2.0.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.
@@ -0,0 +1,1311 @@
1
+ # Copyright 2025-2026 Hanlian Lu. SPDX-License-Identifier: Apache-2.0
2
+ """PostgreSQL adapter for the Memory facade (P3/P4 substrate).
3
+
4
+ Owns its own schema, namespace, and migration registry: the package never
5
+ shares DlightRAG's answer-run migrations or tables. The three recall legs are
6
+ implemented here behind the neutral ports:
7
+
8
+ - exact: ``normalized_body`` btree equality (Python-side NFKC normalization)
9
+ - sparse: pg_textsearch BM25 with the corpus-tuned k1/b and both textsearch
10
+ configs (``simple`` + ``public.jiebacfg``), merged by best score into one
11
+ ranking so a record never double-counts in fusion
12
+ - dense: optional ``halfvec`` column + HNSW index when a TextEmbedder is bound
13
+
14
+ Dense is opt-in: with the NullEmbedder the adapter runs exact + sparse only.
15
+ A changed embedder fingerprint leaves old rows out of the dense leg (exact and
16
+ sparse still reach them); automatic re-embedding is a P4 decision.
17
+ """
18
+
19
+ from __future__ import annotations
20
+
21
+ import json
22
+ import uuid
23
+ from collections.abc import Awaitable, Callable
24
+ from dataclasses import replace
25
+ from datetime import UTC, datetime
26
+ from typing import Any, Protocol
27
+
28
+ from dlightrag_memory._storage.pg_bm25 import (
29
+ build_bm25_sql,
30
+ ensure_bm25_indexes,
31
+ extension_bootstrap_sql,
32
+ text_configs_available,
33
+ )
34
+ from dlightrag_memory.errors import MemoryWriteRejectedError
35
+ from dlightrag_memory.models import (
36
+ MemoryOperation,
37
+ MemoryOperationReceipt,
38
+ MemoryProvenance,
39
+ MemoryRecord,
40
+ )
41
+ from dlightrag_memory.normalize import normalized_body
42
+ from dlightrag_memory.ports import (
43
+ NullEmbedder,
44
+ SearchCandidate,
45
+ TextEmbedder,
46
+ Vector,
47
+ )
48
+ from dlightrag_memory.recall import recall_recency
49
+ from dlightrag_memory.store import (
50
+ OperationGuard,
51
+ operation_change_id,
52
+ operation_fingerprint,
53
+ operation_receipt,
54
+ operation_record_id,
55
+ )
56
+
57
+
58
+ class PGConnection(Protocol):
59
+ """The duck-typed asyncpg surface the PostgreSQL adapter consumes."""
60
+
61
+ async def fetch(self, query: str, *args: Any) -> list[Any]: ...
62
+
63
+ async def fetchrow(self, query: str, *args: Any) -> Any: ...
64
+
65
+ async def fetchval(self, query: str, *args: Any) -> Any: ...
66
+
67
+ async def execute(self, query: str, *args: Any) -> Any: ...
68
+
69
+ def transaction(self) -> Any: ...
70
+
71
+
72
+ class PGPool(Protocol):
73
+ """A pool whose ``acquire()`` yields one connection context.
74
+
75
+ Duck-typed: both a bound asyncpg pool and a lazy pool holder satisfy it.
76
+ """
77
+
78
+ def acquire(self) -> Any: ...
79
+
80
+
81
+ _RECORDS_TABLE = """
82
+ CREATE TABLE IF NOT EXISTS dlightrag_memory_records (
83
+ owner_id TEXT NOT NULL,
84
+ memory_id UUID NOT NULL,
85
+ kind TEXT NOT NULL,
86
+ body TEXT NOT NULL,
87
+ normalized_body TEXT NOT NULL,
88
+ origin_kind TEXT NOT NULL,
89
+ origin_id TEXT NOT NULL,
90
+ run_id TEXT,
91
+ session_id TEXT,
92
+ status TEXT NOT NULL,
93
+ supersedes_id UUID,
94
+ embedding_fingerprint TEXT,
95
+ created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
96
+ updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
97
+ PRIMARY KEY (owner_id, memory_id),
98
+ CONSTRAINT dlightrag_memory_records_kind_check
99
+ CHECK (kind IN ('preference', 'fact')),
100
+ CONSTRAINT dlightrag_memory_records_origin_check
101
+ CHECK (origin_kind IN ('answer_run', 'management', 'mcp', 'undo')),
102
+ CONSTRAINT dlightrag_memory_records_status_check
103
+ CHECK (status IN ('active', 'superseded', 'forgotten')),
104
+ CONSTRAINT dlightrag_memory_records_body_check
105
+ CHECK (char_length(body) BETWEEN 1 AND 500)
106
+ )
107
+ """
108
+
109
+ _OPERATIONS_TABLE = """
110
+ CREATE TABLE IF NOT EXISTS dlightrag_memory_operations (
111
+ owner_id TEXT NOT NULL,
112
+ change_id UUID NOT NULL,
113
+ idempotency_key TEXT NOT NULL,
114
+ request_fingerprint TEXT NOT NULL,
115
+ operation TEXT NOT NULL,
116
+ outcome TEXT NOT NULL,
117
+ mutation_scope TEXT,
118
+ receipt JSONB NOT NULL,
119
+ before_records JSONB NOT NULL DEFAULT '[]'::jsonb,
120
+ undone_by UUID,
121
+ created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
122
+ PRIMARY KEY (owner_id, change_id),
123
+ UNIQUE (owner_id, idempotency_key),
124
+ CONSTRAINT dlightrag_memory_operations_action_check
125
+ CHECK (operation IN ('remember', 'forget', 'undo')),
126
+ CONSTRAINT dlightrag_memory_operations_outcome_check
127
+ CHECK (outcome IN ('changed', 'unchanged', 'conflict'))
128
+ )
129
+ """
130
+
131
+ _RECORD_INDEXES = (
132
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_records_recall "
133
+ "ON dlightrag_memory_records (owner_id, status, updated_at DESC)",
134
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_records_list "
135
+ "ON dlightrag_memory_records (owner_id, status, updated_at DESC, memory_id DESC)",
136
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_records_exact "
137
+ "ON dlightrag_memory_records (owner_id, normalized_body)",
138
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_records_purge "
139
+ "ON dlightrag_memory_records (status, updated_at) "
140
+ "WHERE status = 'superseded'",
141
+ )
142
+
143
+ _RECORD_COLUMNS = """
144
+ owner_id, memory_id, kind, body, normalized_body, origin_kind, origin_id, run_id,
145
+ session_id, status, supersedes_id, embedding_fingerprint, created_at, updated_at
146
+ """
147
+
148
+ _OPERATION_INDEXES = (
149
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_operations_scope "
150
+ "ON dlightrag_memory_operations (owner_id, mutation_scope, outcome)",
151
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_operations_retention "
152
+ "ON dlightrag_memory_operations (created_at)",
153
+ )
154
+
155
+
156
+ def _vector_text(vector: Vector) -> str:
157
+ return "[" + ",".join(f"{float(value):g}" for value in vector) + "]"
158
+
159
+
160
+ class PostgresMemoryStore:
161
+ """Owner-isolated Memory Records in one package-owned PostgreSQL schema."""
162
+
163
+ def __init__(
164
+ self,
165
+ pool: PGPool | None = None,
166
+ *,
167
+ dsn: str | None = None,
168
+ pool_factory: Callable[[], Awaitable[Any]] | None = None,
169
+ embedder: TextEmbedder = NullEmbedder(),
170
+ ) -> None:
171
+ if pool is None and not dsn and pool_factory is None:
172
+ raise ValueError("PostgresMemoryStore needs a pool, a dsn, or a pool factory")
173
+ self._dsn = dsn
174
+ self._pool = pool
175
+ self._pool_factory = pool_factory
176
+ self._owned_pool: Any = None
177
+ self._embedder = embedder
178
+ self._dense = not isinstance(embedder, NullEmbedder)
179
+ self._bm25_indexes: tuple[str, ...] = ()
180
+ self._operation_pool = pool # test hook: backdate rows directly
181
+ self._initialized = False
182
+
183
+ async def aclose(self) -> None:
184
+ """Close the adapter-owned pool; safe to call more than once."""
185
+ owned = self._owned_pool
186
+ self._owned_pool = None
187
+ if owned is not None:
188
+ await owned.close()
189
+
190
+ async def _acquire_context(self) -> Any:
191
+ """Return one pool acquire context, creating an owned pool when needed."""
192
+ if self._pool is None:
193
+ if self._pool_factory is not None:
194
+ self._pool = await self._pool_factory()
195
+ else:
196
+ import asyncpg
197
+
198
+ if self._owned_pool is None:
199
+ self._owned_pool = await asyncpg.create_pool(self._dsn)
200
+ self._pool = self._owned_pool
201
+ return self._pool.acquire()
202
+
203
+ async def initialize(self) -> None:
204
+ async def operation(conn: PGConnection) -> None:
205
+ await conn.execute(_RECORDS_TABLE)
206
+ await conn.execute(_OPERATIONS_TABLE)
207
+ for statement in (*_RECORD_INDEXES, *_OPERATION_INDEXES):
208
+ await conn.execute(statement)
209
+ for statement in extension_bootstrap_sql():
210
+ await conn.execute(statement)
211
+ available = await text_configs_available(conn)
212
+ self._bm25_indexes = await ensure_bm25_indexes(conn, available=available)
213
+ if self._dense:
214
+ dim = int(self._embedder.dim)
215
+ if dim < 1:
216
+ raise ValueError("embedder dim must be positive for the dense leg")
217
+ await conn.execute(_embedding_column_sql(dim))
218
+ await conn.execute(_embedding_index_sql())
219
+
220
+ acquire = await self._acquire_context()
221
+ async with acquire as conn:
222
+ await operation(conn)
223
+ await self.verify()
224
+
225
+ async def verify(self) -> None:
226
+ """Validate the writer's schema and load search-index facts, no DDL."""
227
+
228
+ async def operation(conn: PGConnection) -> None:
229
+ rows = await conn.fetch(
230
+ "SELECT table_name FROM information_schema.tables "
231
+ "WHERE table_name = ANY($1::text[])",
232
+ ["dlightrag_memory_records", "dlightrag_memory_operations"],
233
+ )
234
+ found = {str(row["table_name"]) for row in rows}
235
+ missing = {"dlightrag_memory_records", "dlightrag_memory_operations"} - found
236
+ if missing:
237
+ raise RuntimeError(
238
+ f"Memory schema is missing {', '.join(sorted(missing))}; "
239
+ "initialize it on the writer first"
240
+ )
241
+ columns = await conn.fetch(
242
+ "SELECT table_name, column_name FROM information_schema.columns "
243
+ "WHERE table_name = ANY($1::text[])",
244
+ ["dlightrag_memory_records", "dlightrag_memory_operations"],
245
+ )
246
+ by_table: dict[str, set[str]] = {}
247
+ for row in columns:
248
+ by_table.setdefault(str(row["table_name"]), set()).add(str(row["column_name"]))
249
+ if "confidence" in by_table.get("dlightrag_memory_records", set()):
250
+ raise RuntimeError("Memory schema still contains the removed confidence column")
251
+ required = {
252
+ "dlightrag_memory_records": {"origin_kind", "origin_id", "normalized_body"},
253
+ "dlightrag_memory_operations": {
254
+ "change_id",
255
+ "request_fingerprint",
256
+ "receipt",
257
+ "before_records",
258
+ "undone_by",
259
+ },
260
+ }
261
+ for table, names in required.items():
262
+ missing_columns = names - by_table.get(table, set())
263
+ if missing_columns:
264
+ raise RuntimeError(f"{table} is missing {', '.join(sorted(missing_columns))}")
265
+ self._bm25_indexes = await ensure_bm25_indexes(conn, verify_only=True)
266
+
267
+ acquire = await self._acquire_context()
268
+ async with acquire as conn:
269
+ await operation(conn)
270
+ self._initialized = True
271
+
272
+ def _embedder_fingerprint(self) -> str:
273
+ """One canonical embedding-space identity for TEXT persistence."""
274
+ return self._embedder.embedding_fingerprint
275
+
276
+ async def apply_operation(
277
+ self,
278
+ operation: MemoryOperation,
279
+ *,
280
+ guard: OperationGuard | None = None,
281
+ ) -> MemoryOperationReceipt:
282
+ """Atomically settle one operation, its receipt, journal, and mutation cap."""
283
+ change_id = operation_change_id(operation)
284
+ fingerprint = operation_fingerprint(operation)
285
+ embedding = (
286
+ await self._embedding(operation.body.strip())
287
+ if operation.action == "remember" and self._dense
288
+ else None
289
+ )
290
+
291
+ async def settle(conn: PGConnection) -> MemoryOperationReceipt:
292
+ async with conn.transaction():
293
+ await conn.fetchval(_LOCK_OWNER, operation.owner_id)
294
+ if guard is not None:
295
+ await guard(conn)
296
+ replay = await conn.fetchrow(
297
+ _SELECT_OPERATION, operation.owner_id, _uuid(change_id, label="change_id")
298
+ )
299
+ if replay is not None:
300
+ if str(replay["request_fingerprint"]) != fingerprint:
301
+ raise MemoryWriteRejectedError(
302
+ "Memory idempotency key was reused with different input."
303
+ )
304
+ return _receipt_row(replay)
305
+
306
+ if operation.action == "remember":
307
+ receipt, before = await self._settle_remember(
308
+ conn,
309
+ operation=operation,
310
+ change_id=change_id,
311
+ embedding=embedding,
312
+ )
313
+ elif operation.action == "forget":
314
+ receipt, before = await self._settle_forget(
315
+ conn, operation=operation, change_id=change_id
316
+ )
317
+ else:
318
+ receipt, before = await self._settle_undo(
319
+ conn, operation=operation, change_id=change_id
320
+ )
321
+
322
+ if receipt.changed and operation.mutation_scope is not None:
323
+ used = int(
324
+ await conn.fetchval(
325
+ _COUNT_SCOPE_MUTATIONS,
326
+ operation.owner_id,
327
+ operation.mutation_scope,
328
+ )
329
+ or 0
330
+ )
331
+ if used >= int(operation.mutation_limit or 0):
332
+ raise MemoryWriteRejectedError(
333
+ "This Answer Run reached its Memory mutation limit."
334
+ )
335
+
336
+ await conn.execute(
337
+ _INSERT_OPERATION,
338
+ operation.owner_id,
339
+ _uuid(change_id, label="change_id"),
340
+ operation.idempotency_key,
341
+ fingerprint,
342
+ receipt.action,
343
+ receipt.outcome,
344
+ operation.mutation_scope,
345
+ json.dumps(_receipt_json(receipt), ensure_ascii=False),
346
+ json.dumps([_record_json(record) for record in before], ensure_ascii=False),
347
+ receipt.created_at,
348
+ )
349
+ if receipt.action == "undo" and receipt.changed and receipt.target_change_id:
350
+ await conn.execute(
351
+ _MARK_OPERATION_UNDONE,
352
+ operation.owner_id,
353
+ _uuid(receipt.target_change_id, label="target_change_id"),
354
+ _uuid(change_id, label="change_id"),
355
+ )
356
+ return receipt
357
+
358
+ return await self._write(settle)
359
+
360
+ async def _settle_remember(
361
+ self,
362
+ conn: PGConnection,
363
+ *,
364
+ operation: MemoryOperation,
365
+ change_id: str,
366
+ embedding: Vector | None,
367
+ ) -> tuple[MemoryOperationReceipt, tuple[MemoryRecord, ...]]:
368
+ now = datetime.now(UTC)
369
+ body = operation.body.strip()
370
+ duplicate_row = await conn.fetchrow(
371
+ _SELECT_ACTIVE_NORMALIZED_FOR_UPDATE,
372
+ operation.owner_id,
373
+ normalized_body(body),
374
+ )
375
+ if duplicate_row is not None:
376
+ duplicate = _row(duplicate_row)
377
+ outcome = (
378
+ "conflict"
379
+ if operation.supersedes_id and duplicate.memory_id != operation.supersedes_id
380
+ else "unchanged"
381
+ )
382
+ return (
383
+ operation_receipt(
384
+ operation,
385
+ change_id,
386
+ outcome,
387
+ memory_ids=(duplicate.memory_id,),
388
+ kind=duplicate.kind,
389
+ body=duplicate.body,
390
+ now=now,
391
+ ),
392
+ (),
393
+ )
394
+
395
+ before: tuple[MemoryRecord, ...] = ()
396
+ if operation.supersedes_id:
397
+ old_row = await conn.fetchrow(
398
+ _SELECT_ONE_FOR_UPDATE,
399
+ operation.owner_id,
400
+ _uuid(operation.supersedes_id, label="supersedes_id"),
401
+ )
402
+ if old_row is None or str(old_row["status"]) != "active":
403
+ return (
404
+ operation_receipt(operation, change_id, "conflict", body=body, now=now),
405
+ (),
406
+ )
407
+ old = _row(old_row)
408
+ before = (old,)
409
+ await conn.execute(
410
+ _MARK_SUPERSEDED,
411
+ operation.owner_id,
412
+ _uuid(operation.supersedes_id, label="supersedes_id"),
413
+ )
414
+
415
+ memory_id = operation_record_id(operation.owner_id, change_id)
416
+ record = MemoryRecord(
417
+ owner_id=operation.owner_id,
418
+ memory_id=memory_id,
419
+ kind=operation.kind or "fact",
420
+ body=body,
421
+ provenance=operation.provenance,
422
+ status="active",
423
+ supersedes_id=operation.supersedes_id,
424
+ created_at=now,
425
+ updated_at=now,
426
+ )
427
+ await _insert_record(self, conn, record=record, embedding=embedding)
428
+ return (
429
+ operation_receipt(
430
+ operation,
431
+ change_id,
432
+ "changed",
433
+ memory_ids=(memory_id,),
434
+ kind=record.kind,
435
+ body=record.body,
436
+ now=now,
437
+ ),
438
+ before,
439
+ )
440
+
441
+ async def _settle_forget(
442
+ self,
443
+ conn: PGConnection,
444
+ *,
445
+ operation: MemoryOperation,
446
+ change_id: str,
447
+ ) -> tuple[MemoryOperationReceipt, tuple[MemoryRecord, ...]]:
448
+ now = datetime.now(UTC)
449
+ if operation.memory_id:
450
+ rows = await conn.fetch(
451
+ _SELECT_ACTIVE_ID_FOR_UPDATE,
452
+ operation.owner_id,
453
+ _uuid(operation.memory_id, label="memory_id"),
454
+ )
455
+ else:
456
+ rows = await conn.fetch(
457
+ _SELECT_ACTIVE_NORMALIZED_ALL_FOR_UPDATE,
458
+ operation.owner_id,
459
+ normalized_body(operation.body),
460
+ )
461
+ matches = tuple(_row(row) for row in rows)
462
+ if not matches:
463
+ return (operation_receipt(operation, change_id, "unchanged", now=now), ())
464
+ ids = [_uuid(record.memory_id, label="memory_id") for record in matches]
465
+ await conn.execute(_MARK_FORGOTTEN_IDS, operation.owner_id, ids)
466
+ first = matches[0]
467
+ return (
468
+ operation_receipt(
469
+ operation,
470
+ change_id,
471
+ "changed",
472
+ memory_ids=tuple(record.memory_id for record in matches),
473
+ kind=first.kind,
474
+ body=first.body,
475
+ now=now,
476
+ ),
477
+ matches,
478
+ )
479
+
480
+ async def _settle_undo(
481
+ self,
482
+ conn: PGConnection,
483
+ *,
484
+ operation: MemoryOperation,
485
+ change_id: str,
486
+ ) -> tuple[MemoryOperationReceipt, tuple[MemoryRecord, ...]]:
487
+ now = datetime.now(UTC)
488
+ target_id = operation.target_change_id or ""
489
+ target = await conn.fetchrow(
490
+ _SELECT_OPERATION_FOR_UPDATE,
491
+ operation.owner_id,
492
+ _uuid(target_id, label="target_change_id"),
493
+ )
494
+ if target is None or target["undone_by"] is not None:
495
+ return (
496
+ operation_receipt(
497
+ operation,
498
+ change_id,
499
+ "conflict",
500
+ target_change_id=target_id,
501
+ now=now,
502
+ ),
503
+ (),
504
+ )
505
+ target_receipt = _receipt_row(target)
506
+ before = _records_json(target["before_records"])
507
+ if not target_receipt.changed or target_receipt.action == "undo":
508
+ return (
509
+ operation_receipt(
510
+ operation,
511
+ change_id,
512
+ "conflict",
513
+ target_change_id=target_id,
514
+ now=now,
515
+ ),
516
+ (),
517
+ )
518
+
519
+ if target_receipt.action == "remember":
520
+ current_id = target_receipt.memory_id or ""
521
+ current_row = await conn.fetchrow(
522
+ _SELECT_ONE_FOR_UPDATE,
523
+ operation.owner_id,
524
+ _uuid(current_id, label="memory_id"),
525
+ )
526
+ if current_row is None or str(current_row["status"]) != "active":
527
+ return (
528
+ operation_receipt(
529
+ operation,
530
+ change_id,
531
+ "conflict",
532
+ target_change_id=target_id,
533
+ now=now,
534
+ ),
535
+ (),
536
+ )
537
+ current = _row(current_row)
538
+ if target_receipt.supersedes_id and before:
539
+ await conn.execute(
540
+ _MARK_SUPERSEDED,
541
+ operation.owner_id,
542
+ _uuid(current_id, label="memory_id"),
543
+ )
544
+ restored_id = operation_record_id(operation.owner_id, change_id)
545
+ restored = replace(
546
+ before[0],
547
+ memory_id=restored_id,
548
+ provenance=operation.provenance,
549
+ status="active",
550
+ supersedes_id=current_id,
551
+ created_at=now,
552
+ updated_at=now,
553
+ )
554
+ await _insert_record(self, conn, record=restored, embedding=None)
555
+ return (
556
+ operation_receipt(
557
+ operation,
558
+ change_id,
559
+ "changed",
560
+ memory_ids=(restored_id,),
561
+ kind=restored.kind,
562
+ body=restored.body,
563
+ supersedes_id=current_id,
564
+ target_change_id=target_id,
565
+ now=now,
566
+ ),
567
+ (current,),
568
+ )
569
+ await conn.execute(
570
+ _MARK_FORGOTTEN_IDS,
571
+ operation.owner_id,
572
+ [_uuid(current_id, label="memory_id")],
573
+ )
574
+ return (
575
+ operation_receipt(
576
+ operation,
577
+ change_id,
578
+ "changed",
579
+ memory_ids=(current_id,),
580
+ kind=current.kind,
581
+ body=current.body,
582
+ target_change_id=target_id,
583
+ now=now,
584
+ ),
585
+ (current,),
586
+ )
587
+
588
+ # Forget target: preflight the full batch set-wise before mutating
589
+ # anything. The persisted journal must be well formed (nonempty, every
590
+ # before record belongs to this owner, memory ids are unique and, in
591
+ # order, exactly the target receipt's memory_ids), every target id
592
+ # must still exist for this owner as a forgotten row, and no active
593
+ # record outside the batch may share a normalized body (siblings
594
+ # inside the batch compensate the exact prior forget and are never
595
+ # conflicts).
596
+ if (
597
+ not before
598
+ or any(old.owner_id != operation.owner_id for old in before)
599
+ or len({old.memory_id for old in before}) != len(before)
600
+ or tuple(old.memory_id for old in before) != target_receipt.memory_ids
601
+ ):
602
+ return (
603
+ operation_receipt(
604
+ operation,
605
+ change_id,
606
+ "conflict",
607
+ target_change_id=target_id,
608
+ now=now,
609
+ ),
610
+ (),
611
+ )
612
+ target_ids = [_uuid(old.memory_id, label="memory_id") for old in before]
613
+ current_rows = await conn.fetch(_SELECT_IDS_FOR_UPDATE, operation.owner_id, target_ids)
614
+ current_by_id = {str(row["memory_id"]): row for row in current_rows}
615
+ if any(
616
+ current_by_id.get(old.memory_id) is None
617
+ or str(current_by_id[old.memory_id]["status"]) != "forgotten"
618
+ for old in before
619
+ ):
620
+ return (
621
+ operation_receipt(
622
+ operation,
623
+ change_id,
624
+ "conflict",
625
+ target_change_id=target_id,
626
+ now=now,
627
+ ),
628
+ (),
629
+ )
630
+ bodies = sorted({normalized_body(old.body) for old in before})
631
+ if await conn.fetchval(
632
+ _SELECT_ACTIVE_NORMALIZED_CONFLICT_EXISTS, operation.owner_id, bodies
633
+ ):
634
+ return (
635
+ operation_receipt(
636
+ operation,
637
+ change_id,
638
+ "conflict",
639
+ target_change_id=target_id,
640
+ now=now,
641
+ ),
642
+ (),
643
+ )
644
+
645
+ restored_records: list[MemoryRecord] = []
646
+ for index, old in enumerate(before):
647
+ restored_records.append(
648
+ replace(
649
+ old,
650
+ memory_id=operation_record_id(operation.owner_id, change_id, index=index),
651
+ provenance=operation.provenance,
652
+ status="active",
653
+ supersedes_id=old.memory_id,
654
+ created_at=now,
655
+ updated_at=now,
656
+ )
657
+ )
658
+ if restored_records:
659
+ restored_rows = await conn.fetch(
660
+ _INSERT_RESTORED_BATCH,
661
+ operation.owner_id,
662
+ operation.provenance.origin_kind,
663
+ operation.provenance.origin_id,
664
+ self._embedder_fingerprint() if self._dense else None,
665
+ now,
666
+ _restore_batch_json(restored_records),
667
+ )
668
+ restored_ids = {str(row["memory_id"]) for row in restored_rows}
669
+ if len(restored_rows) != len(restored_records) or restored_ids != {
670
+ record.memory_id for record in restored_records
671
+ }:
672
+ raise ValueError("memory id already exists with different content")
673
+ first = restored_records[0] if restored_records else None
674
+ return (
675
+ operation_receipt(
676
+ operation,
677
+ change_id,
678
+ "changed",
679
+ memory_ids=tuple(record.memory_id for record in restored_records),
680
+ kind=None if first is None else first.kind,
681
+ body="" if first is None else first.body,
682
+ target_change_id=target_id,
683
+ now=now,
684
+ ),
685
+ before,
686
+ )
687
+
688
+ async def clear_owner(
689
+ self,
690
+ *,
691
+ owner_id: str,
692
+ guard: OperationGuard | None = None,
693
+ ) -> int:
694
+ async def operation(conn: PGConnection) -> int:
695
+ async with conn.transaction():
696
+ await conn.fetchval(_LOCK_OWNER, owner_id)
697
+ if guard is not None:
698
+ await guard(conn)
699
+ await conn.execute(_CLEAR_OPERATIONS, owner_id)
700
+ records = await conn.execute(_CLEAR_RECORDS, owner_id)
701
+ return _command_count(records)
702
+
703
+ return await self._write(operation)
704
+
705
+ async def count_active(self, *, owner_id: str) -> int:
706
+ async def operation(conn: PGConnection) -> int:
707
+ return int(await conn.fetchval(_COUNT_ACTIVE, owner_id) or 0)
708
+
709
+ return await self._read(operation)
710
+
711
+ async def insert(self, record: MemoryRecord) -> None:
712
+ embedding = await self._embedding(record.body) if self._dense else None
713
+
714
+ async def operation(conn: PGConnection) -> None:
715
+ async with conn.transaction():
716
+ if embedding is None:
717
+ await conn.execute(_INSERT, *_insert_params(self, record=record))
718
+ else:
719
+ await conn.execute(
720
+ _INSERT_WITH_EMBEDDING,
721
+ *_insert_params(self, record=record),
722
+ _vector_text(embedding),
723
+ )
724
+ current = await conn.fetchrow(
725
+ _SELECT_ONE, record.owner_id, _uuid(record.memory_id, label="memory_id")
726
+ )
727
+ if current is None or _row(current) != record:
728
+ raise ValueError("memory id already exists with different content")
729
+
730
+ await self._write(operation)
731
+
732
+ async def supersede(self, *, owner_id: str, old_id: str, new: MemoryRecord) -> None:
733
+ if new.owner_id != owner_id:
734
+ raise ValueError("supersede cannot change owner")
735
+ embedding = await self._embedding(new.body) if self._dense else None
736
+
737
+ async def operation(conn: PGConnection) -> None:
738
+ async with conn.transaction():
739
+ tag = await conn.execute(
740
+ _MARK_SUPERSEDED, owner_id, _uuid(old_id, label="memory_id")
741
+ )
742
+ if str(tag).endswith(" 0"):
743
+ raise KeyError(old_id)
744
+ if embedding is None:
745
+ await conn.execute(_INSERT, *_insert_params(self, record=new))
746
+ else:
747
+ await conn.execute(
748
+ _INSERT_WITH_EMBEDDING,
749
+ *_insert_params(self, record=new),
750
+ _vector_text(embedding),
751
+ )
752
+
753
+ await self._write(operation)
754
+
755
+ async def forget(self, *, owner_id: str, memory_id: str) -> bool:
756
+ async def operation(conn: PGConnection) -> bool:
757
+ async with conn.transaction():
758
+ tag = await conn.execute(_DELETE, owner_id, _uuid(memory_id, label="memory_id"))
759
+ return not str(tag).endswith(" 0")
760
+
761
+ return await self._write(operation)
762
+
763
+ async def forget_matching(self, *, owner_id: str, body: str) -> int:
764
+ async def operation(conn: PGConnection) -> int:
765
+ async with conn.transaction():
766
+ result = await conn.execute(_DELETE_BODY, owner_id, body.strip())
767
+ return int(str(result).rsplit(" ", 1)[-1])
768
+
769
+ return await self._write(operation)
770
+
771
+ async def forget_all(self, *, owner_id: str) -> int:
772
+ async def operation(conn: PGConnection) -> int:
773
+ async with conn.transaction():
774
+ result = await conn.execute(_DELETE_ALL, owner_id)
775
+ return int(str(result).rsplit(" ", 1)[-1])
776
+
777
+ return await self._write(operation)
778
+
779
+ async def get(self, *, owner_id: str, memory_id: str) -> MemoryRecord | None:
780
+ async def operation(conn: PGConnection) -> MemoryRecord | None:
781
+ row = await conn.fetchrow(_SELECT_ONE, owner_id, _uuid(memory_id, label="memory_id"))
782
+ return None if row is None else _row(row)
783
+
784
+ return await self._read(operation)
785
+
786
+ async def list_active_page(
787
+ self,
788
+ *,
789
+ owner_id: str,
790
+ after: tuple[datetime, str] | None = None,
791
+ limit: int = 50,
792
+ ) -> tuple[tuple[MemoryRecord, ...], tuple[datetime, str] | None]:
793
+ """Keyset page over (updated_at DESC, memory_id) for owner browse."""
794
+ cap = max(1, min(int(limit), 100))
795
+
796
+ async def operation(
797
+ conn: PGConnection,
798
+ ) -> tuple[tuple[MemoryRecord, ...], tuple[datetime, str] | None]:
799
+ if after is None:
800
+ rows = await conn.fetch(_SELECT_ACTIVE_PAGE, owner_id, cap + 1)
801
+ else:
802
+ rows = await conn.fetch(
803
+ _SELECT_ACTIVE_PAGE_AFTER, owner_id, after[0], after[1], cap + 1
804
+ )
805
+ page = tuple(_row(row) for row in rows[:cap])
806
+ if len(rows) <= cap:
807
+ return page, None
808
+ last = _row(rows[cap - 1])
809
+ return page, (recall_recency(last), last.memory_id)
810
+
811
+ return await self._read(operation)
812
+
813
+ async def purge_superseded(self, *, older_than: datetime) -> int:
814
+ async def operation(conn: PGConnection) -> int:
815
+ async with conn.transaction():
816
+ operations = await conn.execute(_PURGE_OPERATIONS, older_than)
817
+ records = await conn.execute(_PURGE_RECORDS, older_than)
818
+ return _command_count(operations) + _command_count(records)
819
+
820
+ return await self._write(operation)
821
+
822
+ async def search_candidates(
823
+ self, *, owner_id: str, query: str, limit: int
824
+ ) -> tuple[SearchCandidate, ...]:
825
+ """Return leg-tagged candidates in per-leg rank order, no cross-leg merge.
826
+
827
+ The sparse leg merges both BM25 configs by best score into ONE ranking
828
+ (a record matching both configs must not double-count in RRF); the
829
+ façade fuses exact/sparse/dense with RRF.
830
+ """
831
+ cap = max(1, min(int(limit), 100))
832
+ key = normalized_body(query)
833
+
834
+ async def operation(conn: PGConnection) -> tuple[SearchCandidate, ...]:
835
+ candidates: list[SearchCandidate] = []
836
+ exact_rows = await conn.fetch(_SEARCH_EXACT, owner_id, key, cap)
837
+ candidates.extend(
838
+ SearchCandidate(record=_row(row), leg="exact", score=2.0) for row in exact_rows
839
+ )
840
+ sparse_by_id: dict[str, SearchCandidate] = {}
841
+ for bm25_index in self._bm25_indexes:
842
+ rows = await conn.fetch(
843
+ build_bm25_sql(index_name=bm25_index, limit=cap), query, owner_id
844
+ )
845
+ for row in rows:
846
+ record = _row(row)
847
+ score = float(row["score"])
848
+ existing = sparse_by_id.get(record.memory_id)
849
+ if existing is None or score > existing.score:
850
+ sparse_by_id[record.memory_id] = SearchCandidate(
851
+ record=record, leg="sparse", score=score
852
+ )
853
+ candidates.extend(
854
+ sorted(
855
+ sparse_by_id.values(),
856
+ key=lambda candidate: candidate.score,
857
+ reverse=True,
858
+ )[:cap]
859
+ )
860
+ if self._dense:
861
+ vector = await self._query_embedding(query)
862
+ dense_rows = await conn.fetch(
863
+ _SEARCH_DENSE,
864
+ owner_id,
865
+ self._embedder_fingerprint(),
866
+ _vector_text(vector),
867
+ cap,
868
+ )
869
+ candidates.extend(
870
+ SearchCandidate(record=_row(row), leg="dense", score=float(row["score"]))
871
+ for row in dense_rows
872
+ )
873
+ return tuple(candidates)
874
+
875
+ return await self._read(operation)
876
+
877
+ async def _embedding(self, text: str) -> Vector:
878
+ (vector,) = await self._embedder.embed_documents((text,))
879
+ return vector
880
+
881
+ async def _query_embedding(self, text: str) -> Vector:
882
+ """Embed one query with the port's query context (asymmetric-aware)."""
883
+ return await self._embedder.embed_query(text)
884
+
885
+ async def _write(self, operation: Any) -> Any:
886
+ async with await self._acquire_context() as conn:
887
+ return await operation(conn)
888
+
889
+ async def _read(self, operation: Any) -> Any:
890
+ async with await self._acquire_context() as conn:
891
+ return await operation(conn)
892
+
893
+
894
+ async def _insert_record(
895
+ store: PostgresMemoryStore,
896
+ conn: PGConnection,
897
+ *,
898
+ record: MemoryRecord,
899
+ embedding: Vector | None,
900
+ ) -> None:
901
+ if embedding is None:
902
+ await conn.execute(_INSERT, *_insert_params(store, record=record))
903
+ else:
904
+ await conn.execute(
905
+ _INSERT_WITH_EMBEDDING,
906
+ *_insert_params(store, record=record),
907
+ _vector_text(embedding),
908
+ )
909
+
910
+
911
+ def _receipt_json(receipt: MemoryOperationReceipt) -> dict[str, Any]:
912
+ return {
913
+ "action": receipt.action,
914
+ "body": receipt.body,
915
+ "change_id": receipt.change_id,
916
+ "created_at": None if receipt.created_at is None else receipt.created_at.isoformat(),
917
+ "kind": receipt.kind,
918
+ "memory_ids": list(receipt.memory_ids),
919
+ "mutation_scope": receipt.mutation_scope,
920
+ "outcome": receipt.outcome,
921
+ "provenance": _provenance_json(receipt.provenance),
922
+ "supersedes_id": receipt.supersedes_id,
923
+ "target_change_id": receipt.target_change_id,
924
+ }
925
+
926
+
927
+ def _receipt_row(row: Any) -> MemoryOperationReceipt:
928
+ value = _json_object(row["receipt"])
929
+ return MemoryOperationReceipt(
930
+ change_id=str(value["change_id"]),
931
+ action=str(value["action"]), # type: ignore[arg-type]
932
+ outcome=str(value["outcome"]), # type: ignore[arg-type]
933
+ memory_ids=tuple(str(item) for item in value.get("memory_ids", [])),
934
+ provenance=_provenance_from_json(value["provenance"]),
935
+ kind=str(value["kind"]) if value.get("kind") is not None else None, # type: ignore[arg-type]
936
+ body=str(value.get("body") or ""),
937
+ supersedes_id=(
938
+ str(value["supersedes_id"]) if value.get("supersedes_id") is not None else None
939
+ ),
940
+ target_change_id=(
941
+ str(value["target_change_id"]) if value.get("target_change_id") is not None else None
942
+ ),
943
+ mutation_scope=(
944
+ str(value["mutation_scope"]) if value.get("mutation_scope") is not None else None
945
+ ),
946
+ created_at=_datetime_value(value.get("created_at")),
947
+ )
948
+
949
+
950
+ def _record_json(record: MemoryRecord) -> dict[str, Any]:
951
+ return {
952
+ "body": record.body,
953
+ "created_at": None if record.created_at is None else record.created_at.isoformat(),
954
+ "kind": record.kind,
955
+ "memory_id": record.memory_id,
956
+ "owner_id": record.owner_id,
957
+ "provenance": _provenance_json(record.provenance),
958
+ "status": record.status,
959
+ "supersedes_id": record.supersedes_id,
960
+ "updated_at": None if record.updated_at is None else record.updated_at.isoformat(),
961
+ }
962
+
963
+
964
+ def _records_json(value: Any) -> tuple[MemoryRecord, ...]:
965
+ rows = _json_array(value)
966
+ return tuple(
967
+ MemoryRecord(
968
+ owner_id=str(row["owner_id"]),
969
+ memory_id=str(row["memory_id"]),
970
+ kind=str(row["kind"]), # type: ignore[arg-type]
971
+ body=str(row["body"]),
972
+ provenance=_provenance_from_json(row["provenance"]),
973
+ status=str(row["status"]), # type: ignore[arg-type]
974
+ supersedes_id=(
975
+ str(row["supersedes_id"]) if row.get("supersedes_id") is not None else None
976
+ ),
977
+ created_at=_datetime_value(row.get("created_at")),
978
+ updated_at=_datetime_value(row.get("updated_at")),
979
+ )
980
+ for row in rows
981
+ )
982
+
983
+
984
+ def _provenance_json(provenance: MemoryProvenance) -> dict[str, Any]:
985
+ return {
986
+ "origin_kind": provenance.origin_kind,
987
+ "origin_id": provenance.origin_id,
988
+ "run_id": provenance.run_id,
989
+ "session_id": provenance.session_id,
990
+ }
991
+
992
+
993
+ def _provenance_from_json(value: Any) -> MemoryProvenance:
994
+ row = _json_object(value)
995
+ return MemoryProvenance(
996
+ origin_kind=str(row["origin_kind"]), # type: ignore[arg-type]
997
+ origin_id=str(row["origin_id"]),
998
+ run_id=str(row["run_id"]) if row.get("run_id") is not None else None,
999
+ session_id=str(row["session_id"]) if row.get("session_id") is not None else None,
1000
+ )
1001
+
1002
+
1003
+ def _json_object(value: Any) -> dict[str, Any]:
1004
+ parsed = json.loads(value) if isinstance(value, str) else value
1005
+ if not isinstance(parsed, dict):
1006
+ raise ValueError("Memory operation receipt is not an object")
1007
+ return parsed
1008
+
1009
+
1010
+ def _json_array(value: Any) -> list[dict[str, Any]]:
1011
+ parsed = json.loads(value) if isinstance(value, str) else value
1012
+ if not isinstance(parsed, list) or not all(isinstance(item, dict) for item in parsed):
1013
+ raise ValueError("Memory operation before-records is not an array")
1014
+ return parsed
1015
+
1016
+
1017
+ def _datetime_value(value: Any) -> datetime | None:
1018
+ if value is None or isinstance(value, datetime):
1019
+ return value
1020
+ return datetime.fromisoformat(str(value))
1021
+
1022
+
1023
+ def _command_count(tag: Any) -> int:
1024
+ return int(str(tag).rsplit(" ", 1)[-1])
1025
+
1026
+
1027
+ def _insert_params(store: PostgresMemoryStore, *, record: MemoryRecord) -> tuple[Any, ...]:
1028
+ return (
1029
+ record.owner_id,
1030
+ _uuid(record.memory_id, label="memory_id"),
1031
+ record.kind,
1032
+ record.body,
1033
+ normalized_body(record.body),
1034
+ record.provenance.origin_kind,
1035
+ record.provenance.origin_id,
1036
+ record.provenance.run_id,
1037
+ record.provenance.session_id,
1038
+ record.status,
1039
+ _uuid(record.supersedes_id, label="supersedes_id") if record.supersedes_id else None,
1040
+ store._embedder_fingerprint() if store._dense else None, # noqa: SLF001
1041
+ record.created_at,
1042
+ record.updated_at,
1043
+ )
1044
+
1045
+
1046
+ def _restore_batch_json(records: list[MemoryRecord]) -> str:
1047
+ """Encode one undo restoration batch as a single JSONB recordset parameter."""
1048
+ return json.dumps(
1049
+ [
1050
+ {
1051
+ "body": record.body,
1052
+ "kind": record.kind,
1053
+ "memory_id": record.memory_id,
1054
+ "normalized_body": normalized_body(record.body),
1055
+ "run_id": record.provenance.run_id,
1056
+ "session_id": record.provenance.session_id,
1057
+ "supersedes_id": record.supersedes_id,
1058
+ }
1059
+ for record in records
1060
+ ],
1061
+ ensure_ascii=False,
1062
+ )
1063
+
1064
+
1065
+ def _embedding_column_sql(dim: int) -> str: # noqa: S608 - dim is a validated int
1066
+ return f"ALTER TABLE dlightrag_memory_records ADD COLUMN IF NOT EXISTS embedding halfvec({dim})"
1067
+
1068
+
1069
+ def _embedding_index_sql() -> str:
1070
+ return (
1071
+ "CREATE INDEX IF NOT EXISTS idx_dlightrag_memory_records_dense "
1072
+ "ON dlightrag_memory_records USING hnsw (embedding halfvec_cosine_ops)"
1073
+ )
1074
+
1075
+
1076
+ _INSERT = """
1077
+ INSERT INTO dlightrag_memory_records (
1078
+ owner_id, memory_id, kind, body, normalized_body, origin_kind, origin_id, run_id,
1079
+ session_id, status, supersedes_id, embedding_fingerprint, created_at, updated_at
1080
+ )
1081
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14)
1082
+ ON CONFLICT (owner_id, memory_id) DO NOTHING
1083
+ """
1084
+
1085
+ _INSERT_WITH_EMBEDDING = """
1086
+ INSERT INTO dlightrag_memory_records (
1087
+ owner_id, memory_id, kind, body, normalized_body, origin_kind, origin_id, run_id,
1088
+ session_id, status, supersedes_id, embedding_fingerprint, embedding, created_at, updated_at
1089
+ )
1090
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $15::halfvec, $13, $14)
1091
+ ON CONFLICT (owner_id, memory_id) DO NOTHING
1092
+ """
1093
+
1094
+ _LOCK_OWNER = "SELECT pg_advisory_xact_lock(hashtext($1))"
1095
+
1096
+ _SELECT_OPERATION = """
1097
+ SELECT request_fingerprint, receipt, before_records, undone_by
1098
+ FROM dlightrag_memory_operations
1099
+ WHERE owner_id = $1 AND change_id = $2
1100
+ """
1101
+
1102
+ _SELECT_OPERATION_FOR_UPDATE = _SELECT_OPERATION + " FOR UPDATE"
1103
+
1104
+ _INSERT_OPERATION = """
1105
+ INSERT INTO dlightrag_memory_operations (
1106
+ owner_id, change_id, idempotency_key, request_fingerprint, operation, outcome,
1107
+ mutation_scope, receipt, before_records, created_at
1108
+ )
1109
+ VALUES ($1, $2, $3, $4, $5, $6, $7, $8::jsonb, $9::jsonb, $10)
1110
+ """
1111
+
1112
+ _COUNT_SCOPE_MUTATIONS = """
1113
+ SELECT COUNT(*)
1114
+ FROM dlightrag_memory_operations
1115
+ WHERE owner_id = $1 AND mutation_scope = $2 AND outcome = 'changed'
1116
+ """
1117
+
1118
+ _MARK_OPERATION_UNDONE = """
1119
+ UPDATE dlightrag_memory_operations
1120
+ SET undone_by = $3
1121
+ WHERE owner_id = $1 AND change_id = $2 AND undone_by IS NULL
1122
+ """
1123
+
1124
+ _SELECT_ACTIVE_NORMALIZED_FOR_UPDATE = f"""
1125
+ SELECT {_RECORD_COLUMNS}
1126
+ FROM dlightrag_memory_records
1127
+ WHERE owner_id = $1 AND status = 'active' AND normalized_body = $2
1128
+ ORDER BY updated_at DESC
1129
+ LIMIT 1
1130
+ FOR UPDATE
1131
+ """ # noqa: S608
1132
+
1133
+ _SELECT_ACTIVE_NORMALIZED_ALL_FOR_UPDATE = f"""
1134
+ SELECT {_RECORD_COLUMNS}
1135
+ FROM dlightrag_memory_records
1136
+ WHERE owner_id = $1 AND status = 'active' AND normalized_body = $2
1137
+ ORDER BY updated_at DESC
1138
+ FOR UPDATE
1139
+ """ # noqa: S608
1140
+
1141
+ _SELECT_ACTIVE_ID_FOR_UPDATE = f"""
1142
+ SELECT {_RECORD_COLUMNS}
1143
+ FROM dlightrag_memory_records
1144
+ WHERE owner_id = $1 AND memory_id = $2 AND status = 'active'
1145
+ FOR UPDATE
1146
+ """ # noqa: S608
1147
+
1148
+ _SELECT_ONE_FOR_UPDATE = f"""
1149
+ SELECT {_RECORD_COLUMNS}
1150
+ FROM dlightrag_memory_records
1151
+ WHERE owner_id = $1 AND memory_id = $2
1152
+ FOR UPDATE
1153
+ """ # noqa: S608
1154
+
1155
+ _SELECT_IDS_FOR_UPDATE = f"""
1156
+ SELECT {_RECORD_COLUMNS}
1157
+ FROM dlightrag_memory_records
1158
+ WHERE owner_id = $1 AND memory_id = ANY($2::uuid[])
1159
+ FOR UPDATE
1160
+ """ # noqa: S608
1161
+
1162
+ _SELECT_ACTIVE_NORMALIZED_CONFLICT_EXISTS = """
1163
+ SELECT EXISTS (
1164
+ SELECT 1
1165
+ FROM dlightrag_memory_records
1166
+ WHERE owner_id = $1 AND status = 'active' AND normalized_body = ANY($2::text[])
1167
+ )
1168
+ """
1169
+
1170
+ _INSERT_RESTORED_BATCH = """
1171
+ INSERT INTO dlightrag_memory_records (
1172
+ owner_id, memory_id, kind, body, normalized_body, origin_kind, origin_id, run_id,
1173
+ session_id, status, supersedes_id, embedding_fingerprint, created_at, updated_at
1174
+ )
1175
+ SELECT
1176
+ $1,
1177
+ (record->>'memory_id')::uuid,
1178
+ record->>'kind',
1179
+ record->>'body',
1180
+ record->>'normalized_body',
1181
+ $2,
1182
+ $3,
1183
+ record->>'run_id',
1184
+ record->>'session_id',
1185
+ 'active',
1186
+ NULLIF(record->>'supersedes_id', '')::uuid,
1187
+ $4,
1188
+ $5,
1189
+ $5
1190
+ FROM jsonb_array_elements($6::jsonb) AS record
1191
+ ON CONFLICT (owner_id, memory_id) DO NOTHING
1192
+ RETURNING memory_id
1193
+ """
1194
+
1195
+ _MARK_FORGOTTEN_IDS = """
1196
+ UPDATE dlightrag_memory_records
1197
+ SET status = 'forgotten', updated_at = NOW()
1198
+ WHERE owner_id = $1 AND memory_id = ANY($2::uuid[]) AND status = 'active'
1199
+ """
1200
+
1201
+ _COUNT_ACTIVE = """
1202
+ SELECT COUNT(*) FROM dlightrag_memory_records
1203
+ WHERE owner_id = $1 AND status = 'active'
1204
+ """
1205
+
1206
+ _CLEAR_OPERATIONS = "DELETE FROM dlightrag_memory_operations WHERE owner_id = $1"
1207
+ _CLEAR_RECORDS = "DELETE FROM dlightrag_memory_records WHERE owner_id = $1"
1208
+
1209
+ _MARK_SUPERSEDED = """
1210
+ UPDATE dlightrag_memory_records
1211
+ SET status = 'superseded', updated_at = NOW()
1212
+ WHERE owner_id = $1 AND memory_id = $2 AND status = 'active'
1213
+ """
1214
+
1215
+ _DELETE = """
1216
+ UPDATE dlightrag_memory_records
1217
+ SET status = 'forgotten', updated_at = NOW()
1218
+ WHERE owner_id = $1 AND memory_id = $2 AND status != 'forgotten'
1219
+ """
1220
+
1221
+ _DELETE_BODY = """
1222
+ UPDATE dlightrag_memory_records
1223
+ SET status = 'forgotten', updated_at = NOW()
1224
+ WHERE owner_id = $1 AND body = $2 AND status != 'forgotten'
1225
+ """
1226
+
1227
+ _DELETE_ALL = """
1228
+ UPDATE dlightrag_memory_records
1229
+ SET status = 'forgotten', updated_at = NOW()
1230
+ WHERE owner_id = $1 AND status != 'forgotten'
1231
+ """
1232
+
1233
+ _SELECT_ONE = f"""
1234
+ SELECT {_RECORD_COLUMNS}
1235
+ FROM dlightrag_memory_records
1236
+ WHERE owner_id = $1 AND memory_id = $2
1237
+ """ # noqa: S608 - interpolates only the trusted _RECORD_COLUMNS constant
1238
+
1239
+ _SELECT_ACTIVE_PAGE = f"""
1240
+ SELECT {_RECORD_COLUMNS}
1241
+ FROM dlightrag_memory_records
1242
+ WHERE owner_id = $1 AND status = 'active'
1243
+ ORDER BY updated_at DESC, memory_id DESC
1244
+ LIMIT $2
1245
+ """ # noqa: S608 - interpolates only the trusted _RECORD_COLUMNS constant
1246
+
1247
+ _SELECT_ACTIVE_PAGE_AFTER = f"""
1248
+ SELECT {_RECORD_COLUMNS}
1249
+ FROM dlightrag_memory_records
1250
+ WHERE owner_id = $1 AND status = 'active'
1251
+ AND (updated_at, memory_id) < ($2, $3)
1252
+ ORDER BY updated_at DESC, memory_id DESC
1253
+ LIMIT $4
1254
+ """ # noqa: S608 - interpolates only the trusted _RECORD_COLUMNS constant
1255
+
1256
+ _PURGE_RECORDS = """
1257
+ DELETE FROM dlightrag_memory_records
1258
+ WHERE status != 'active' AND updated_at < $1
1259
+ """
1260
+
1261
+ _PURGE_OPERATIONS = """
1262
+ DELETE FROM dlightrag_memory_operations
1263
+ WHERE created_at < $1
1264
+ """
1265
+
1266
+ _SEARCH_EXACT = f"""
1267
+ SELECT {_RECORD_COLUMNS}
1268
+ FROM dlightrag_memory_records
1269
+ WHERE owner_id = $1 AND status = 'active' AND normalized_body = $2
1270
+ ORDER BY updated_at DESC
1271
+ LIMIT $3
1272
+ """ # noqa: S608 - interpolates only the trusted _RECORD_COLUMNS constant
1273
+
1274
+ _SEARCH_DENSE = f"""
1275
+ SELECT {_RECORD_COLUMNS}, 1 - (embedding <=> $3::halfvec) AS score
1276
+ FROM dlightrag_memory_records
1277
+ WHERE owner_id = $1 AND status = 'active'
1278
+ AND embedding_fingerprint = $2
1279
+ AND embedding IS NOT NULL
1280
+ ORDER BY embedding <=> $3::halfvec
1281
+ LIMIT $4
1282
+ """ # noqa: S608 - interpolates only the trusted _RECORD_COLUMNS constant
1283
+
1284
+
1285
+ def _uuid(value: str, *, label: str) -> uuid.UUID:
1286
+ try:
1287
+ return uuid.UUID(value)
1288
+ except ValueError as exc:
1289
+ raise ValueError(f"{label} must be a canonical UUID") from exc
1290
+
1291
+
1292
+ def _row(row: Any) -> MemoryRecord:
1293
+ return MemoryRecord(
1294
+ owner_id=str(row["owner_id"]),
1295
+ memory_id=str(row["memory_id"]),
1296
+ kind=str(row["kind"]), # type: ignore[arg-type]
1297
+ body=str(row["body"]),
1298
+ provenance=MemoryProvenance(
1299
+ origin_kind=str(row["origin_kind"]), # type: ignore[arg-type]
1300
+ origin_id=str(row["origin_id"]),
1301
+ run_id=str(row["run_id"]) if row["run_id"] is not None else None,
1302
+ session_id=str(row["session_id"]) if row["session_id"] is not None else None,
1303
+ ),
1304
+ status=str(row["status"]), # type: ignore[arg-type]
1305
+ supersedes_id=str(row["supersedes_id"]) if row["supersedes_id"] is not None else None,
1306
+ created_at=row["created_at"],
1307
+ updated_at=row["updated_at"],
1308
+ )
1309
+
1310
+
1311
+ __all__ = ["PostgresMemoryStore"]