diff --git a/services/dma-api/src/dma_api/repository.py b/services/dma-api/src/dma_api/repository.py index 963c676..f4cfa82 100644 --- a/services/dma-api/src/dma_api/repository.py +++ b/services/dma-api/src/dma_api/repository.py @@ -141,77 +141,89 @@ def initialize(self) -> None: def create_or_get(self, record: MemoryRecord, idempotency_key: str) -> tuple[MemoryRecord, bool]: """Create a record, or return the result of an exact idempotent replay.""" - with self._transaction() as connection: - connection.execute("BEGIN IMMEDIATE") - existing = connection.execute( - """ - SELECT memory_id FROM idempotency_keys - WHERE tenant_id = ? AND operation = 'remember' AND idempotency_key = ? - """, - (record.tenant_id, idempotency_key), - ).fetchone() - if existing is not None: - return self._get_by_id(connection, existing["memory_id"]), False - - if record.type is MemoryType.SEMANTIC: - duplicate = self._find_normalized_semantic(connection, record) - if duplicate is not None: - updated = MemoryRecord( - id=duplicate.id, - tenant_id=duplicate.tenant_id, - agent_id=duplicate.agent_id, - content=record.content, - type=record.type, - version=duplicate.version + 1, - created_at=duplicate.created_at, - updated_at=record.updated_at, - expires_at=record.expires_at, - metadata=record.metadata, + max_retries = 3 + for attempt in range(max_retries): + try: + with self._connect() as connection: + connection.execute("BEGIN IMMEDIATE") + existing = connection.execute( + """ + SELECT memory_id FROM idempotency_keys + WHERE tenant_id = ? AND operation = 'remember' AND idempotency_key = ? + """, + (record.tenant_id, idempotency_key), + ).fetchone() + if existing is not None: + return self._get_by_id(connection, existing["memory_id"]), False + + if record.type is MemoryType.SEMANTIC: + duplicate = self._find_normalized_semantic(connection, record) + if duplicate is not None: + updated = MemoryRecord( + id=duplicate.id, + tenant_id=duplicate.tenant_id, + agent_id=duplicate.agent_id, + content=record.content, + type=record.type, + version=duplicate.version + 1, + created_at=duplicate.created_at, + updated_at=record.updated_at, + expires_at=record.expires_at, + metadata=record.metadata, + ) + connection.execute( + """UPDATE memories SET content = ?, version = ?, updated_at = ?, expires_at = ?, metadata_json = ? WHERE id = ?""", + (updated.content, updated.version, updated.updated_at.isoformat(), updated.expires_at.isoformat() if updated.expires_at else None, json.dumps(updated.metadata, separators=(",", ":"), sort_keys=True), updated.id), + ) + connection.execute("DELETE FROM memory_search WHERE memory_id = ?", (updated.id,)) + connection.execute("INSERT INTO memory_search (content, memory_id, tenant_id, agent_id, type) VALUES (?, ?, ?, ?, ?)", (updated.content, updated.id, updated.tenant_id, updated.agent_id, updated.type.value)) + connection.execute("INSERT INTO idempotency_keys (tenant_id, operation, idempotency_key, memory_id) VALUES (?, 'remember', ?, ?)", (record.tenant_id, idempotency_key, updated.id)) + return updated, False + + connection.execute( + """ + INSERT INTO memories ( + id, tenant_id, agent_id, content, type, version, + created_at, updated_at, expires_at, metadata_json + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + record.id, + record.tenant_id, + record.agent_id, + record.content, + record.type.value, + record.version, + record.created_at.isoformat(), + record.updated_at.isoformat(), + record.expires_at.isoformat() if record.expires_at else None, + json.dumps(record.metadata, separators=(",", ":"), sort_keys=True), + ), ) connection.execute( - """UPDATE memories SET content = ?, version = ?, updated_at = ?, expires_at = ?, metadata_json = ? WHERE id = ?""", - (updated.content, updated.version, self._utc_isoformat(updated.updated_at), self._utc_isoformat(updated.expires_at) if updated.expires_at else None, json.dumps(updated.metadata, separators=(",", ":"), sort_keys=True), updated.id), + """ + INSERT INTO memory_search (content, memory_id, tenant_id, agent_id, type) + VALUES (?, ?, ?, ?, ?) + """, + (record.content, record.id, record.tenant_id, record.agent_id, record.type.value), ) - connection.execute("DELETE FROM memory_search WHERE memory_id = ?", (updated.id,)) - connection.execute("INSERT INTO memory_search (content, memory_id, tenant_id, agent_id, type) VALUES (?, ?, ?, ?, ?)", (updated.content, updated.id, updated.tenant_id, updated.agent_id, updated.type.value)) - connection.execute("INSERT INTO idempotency_keys (tenant_id, operation, idempotency_key, memory_id) VALUES (?, 'remember', ?, ?)", (record.tenant_id, idempotency_key, updated.id)) - return updated, False - - connection.execute( - """ - INSERT INTO memories ( - id, tenant_id, agent_id, content, type, version, - created_at, updated_at, expires_at, metadata_json - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, - ( - record.id, - record.tenant_id, - record.agent_id, - record.content, - record.type.value, - record.version, - self._utc_isoformat(record.created_at), - self._utc_isoformat(record.updated_at), - self._utc_isoformat(record.expires_at) if record.expires_at else None, - json.dumps(record.metadata, separators=(",", ":"), sort_keys=True), - ), - ) - connection.execute( - """ - INSERT INTO memory_search (content, memory_id, tenant_id, agent_id, type) - VALUES (?, ?, ?, ?, ?) - """, - (record.content, record.id, record.tenant_id, record.agent_id, record.type.value), - ) - connection.execute( - """ - INSERT INTO idempotency_keys (tenant_id, operation, idempotency_key, memory_id) - VALUES (?, 'remember', ?, ?) - """, - (record.tenant_id, idempotency_key, record.id), - ) - return record, True + connection.execute( + """ + INSERT INTO idempotency_keys (tenant_id, operation, idempotency_key, memory_id) + VALUES (?, 'remember', ?, ?) + """, + (record.tenant_id, idempotency_key, record.id), + ) + return record, True + except sqlite3.OperationalError as error: + if "database is locked" in str(error): + if attempt == max_retries - 1: + raise + import time + time.sleep(0.1 * (2 ** attempt)) + continue + raise + raise RuntimeError("Unreachable code reached") def _find_normalized_semantic(self, connection: sqlite3.Connection, record: MemoryRecord) -> MemoryRecord | None: rows = connection.execute("SELECT * FROM memories WHERE tenant_id = ? AND agent_id = ? AND type = 'semantic'", (record.tenant_id, record.agent_id)).fetchall()