@lota-sdk/core 0.4.8 → 0.4.10

package/src/workers/memory-consolidation.worker.ts

@@ -1,18 +1,67 @@
 import type { SandboxedJob } from 'bullmq'
+import { Schema, Effect } from 'effect'
 import { BoundQuery, eq, inside } from 'surrealdb'
 
 import { MEMORY } from '../config/constants'
 import { serverLogger } from '../config/logger'
 import { ensureRecordId, recordIdToString } from '../db/record-id'
 import type { RecordIdInput } from '../db/record-id'
-import { databaseService } from '../db/service'
+import type { SurrealDBService } from '../db/service'
 import { TABLES } from '../db/tables'
+import { effectTryPromise } from '../effect/helpers'
+import { DatabaseServiceTag } from '../effect/services'
 import type { MemoryConsolidationJob } from '../queues/memory-consolidation.queue'
+import { nowDate, unsafeDateFrom } from '../utils/date-time'
+import { getErrorMessage } from '../utils/errors'
 import { initializeSandboxedWorkerRuntime } from './bootstrap'
 import { toSandboxedWorkerError } from './utils/sandbox-error'
 import { createTracedWorkerProcessor } from './worker-utils'
 
-await initializeSandboxedWorkerRuntime()
+class MemoryConsolidationError extends Schema.TaggedErrorClass<MemoryConsolidationError>()('MemoryConsolidationError', {
+  stage: Schema.Literals([
+    'read-memories',
+    'read-neighbors',
+    'archive-memory',
+    'relate-memories',
+    'write-history',
+    'load-stale',
+    'update-stale',
+    'load-middle-nodes',
+    'read-existing-relation',
+    'write-relation',
+    'update-middle-node',
+    'decay-standard',
+    'decay-ephemeral',
+    'cleanup-relations',
+    'load-scope-ids',
+  ]),
+  message: Schema.String,
+  cause: Schema.Defect,
+}) {}
+
+function toMemoryConsolidationError(
+  stage: MemoryConsolidationError['stage'],
+  cause: unknown,
+): MemoryConsolidationError {
+  return new MemoryConsolidationError({ stage, message: getErrorMessage(cause), cause })
+}
+
+function tryMemoryConsolidationEffect<A, R = never>(
+  stage: MemoryConsolidationError['stage'],
+  thunk: () => PromiseLike<A> | Effect.Effect<A, unknown, R>,
+): Effect.Effect<A, MemoryConsolidationError, R> {
+  return effectTryPromise(thunk, (cause) => toMemoryConsolidationError(stage, cause))
+}
+
+const runtime = await initializeSandboxedWorkerRuntime()
+
+const memoryConsolidationDatabaseService: SurrealDBService = await runtime.runPromise(
+  Effect.service(DatabaseServiceTag),
+)
+
+function db(): SurrealDBService {
+  return memoryConsolidationDatabaseService
+}
 
 const MEMORY_TABLE = TABLES.MEMORY
 const MEMORY_RELATION_TABLE = TABLES.MEMORY_RELATION
@@ -39,272 +88,311 @@ function isContentSubsumed(shorter: string, longer: string): boolean {
   return aWords.size > 0 && overlap / aWords.size >= 0.8
 }
 
-async function deduplicateScope(scopeId: string): Promise<number> {
-  const memoryRows = await databaseService.query<{
-    id: RecordIdInput
-    content: string
-    importance: number
-    embedding: number[]
-    updatedAt?: string | Date | number | null
-    createdAt: string | Date | number
-  }>(
-    new BoundQuery(
-      `SELECT id, content, importance, embedding, updatedAt, createdAt
-     FROM ${MEMORY_TABLE}
-     WHERE scopeId = $scopeId AND archivedAt IS NONE
-     ORDER BY createdAt DESC
-     LIMIT $limit`,
-      { scopeId, limit: MAX_MEMORIES_PER_SCOPE },
-    ),
-  )
-
-  const memories = memoryRows.map((row) => ({ ...row, id: toMemoryId(row.id) }))
-  if (memories.length < 2) return 0
-
-  const archived = new Set<string>()
-  let mergeCount = 0
-
-  for (const memory of memories) {
-    if (archived.has(memory.id)) continue
-
-    const candidateLimit = 20
-    const neighborStatements = await databaseService.queryAll<{
-      id: RecordIdInput
-      content: string
-      importance: number
-      similarity: number
-      updatedAt?: string | Date | number | null
-      createdAt: string | Date | number
-    }>(
-      new BoundQuery(
-        `LET $candidateRows = (
-         SELECT
-           id,
-           vector::distance::knn() AS distance
-         FROM ${MEMORY_TABLE}
-         WHERE scopeId = $scopeId
-           AND archivedAt IS NONE
-           AND id != $memoryId
-           AND embedding <|${candidateLimit}|> $embedding
-         ORDER BY distance ASC
-         LIMIT ${candidateLimit}
-       );
-
-       SELECT
-         id,
-         content,
-         importance,
-         updatedAt,
-         createdAt,
-         vector::similarity::cosine(embedding, $embedding) AS similarity
-       FROM ${MEMORY_TABLE}
-       WHERE id IN $candidateRows.id
-       ORDER BY similarity DESC
-       LIMIT ${candidateLimit}`,
-        { scopeId, memoryId: ensureRecordId(memory.id, TABLES.MEMORY), embedding: memory.embedding },
+function deduplicateScopeEffect(scopeId: string) {
+  return Effect.gen(function* () {
+    const memoryRows = yield* tryMemoryConsolidationEffect('read-memories', () =>
+      db().query<{
+        id: RecordIdInput
+        content: string
+        importance: number
+        embedding: number[]
+        updatedAt?: string | Date | number | null
+        createdAt: string | Date | number
+      }>(
+        new BoundQuery(
+          `SELECT id, content, importance, embedding, updatedAt, createdAt
+           FROM ${MEMORY_TABLE}
+           WHERE scopeId = $scopeId AND archivedAt IS NONE
+           ORDER BY createdAt DESC
+           LIMIT $limit`,
+          { scopeId, limit: MAX_MEMORIES_PER_SCOPE },
+        ),
       ),
     )
 
-    const neighbors = (neighborStatements.at(-1) ?? []).map((row) => ({ ...row, id: toMemoryId(row.id) }))
+    const memories = memoryRows.map((row) => ({ ...row, id: toMemoryId(row.id) }))
+    if (memories.length < 2) return 0
+
+    const archived = new Set<string>()
+    let mergeCount = 0
+
+    for (const memory of memories) {
+      if (archived.has(memory.id)) continue
+
+      const candidateLimit = 20
+      const neighborStatements = yield* tryMemoryConsolidationEffect('read-neighbors', () =>
+        db().queryAll<{
+          id: RecordIdInput
+          content: string
+          importance: number
+          similarity: number
+          updatedAt?: string | Date | number | null
+          createdAt: string | Date | number
+        }>(
+          new BoundQuery(
+            `LET $candidateRows = (
+               SELECT
+                 id,
+                 vector::distance::knn() AS distance
+               FROM ${MEMORY_TABLE}
+               WHERE scopeId = $scopeId
+                 AND archivedAt IS NONE
+                 AND id != $memoryId
+                 AND embedding <|${candidateLimit}|> $embedding
+               ORDER BY distance ASC
+               LIMIT ${candidateLimit}
+             );
+
+             SELECT
+               id,
+               content,
+               importance,
+               updatedAt,
+               createdAt,
+               vector::similarity::cosine(embedding, $embedding) AS similarity
+             FROM ${MEMORY_TABLE}
+             WHERE id IN $candidateRows.id
+             ORDER BY similarity DESC
+             LIMIT ${candidateLimit}`,
+            { scopeId, memoryId: ensureRecordId(memory.id, TABLES.MEMORY), embedding: memory.embedding },
+          ),
+        ),
+      )
 
-    for (const neighbor of neighbors) {
-      if (neighbor.similarity < SOFT_SIMILARITY_THRESHOLD) break
-      if (archived.has(neighbor.id)) continue
+      const neighbors = (neighborStatements.at(-1) ?? []).map((row) => ({ ...row, id: toMemoryId(row.id) }))
 
-      const isHardMatch = neighbor.similarity >= HARD_SIMILARITY_THRESHOLD
-      if (!isHardMatch) {
-        const [shorter, longer] =
-          memory.content.length <= neighbor.content.length
-            ? [memory.content, neighbor.content]
-            : [neighbor.content, memory.content]
-        if (!isContentSubsumed(shorter, longer)) continue
-      }
+      for (const neighbor of neighbors) {
+        if (neighbor.similarity < SOFT_SIMILARITY_THRESHOLD) break
+        if (archived.has(neighbor.id)) continue
 
-      const memoryCreatedAt = new Date(memory.createdAt).getTime()
-      const neighborCreatedAt = new Date(neighbor.createdAt).getTime()
-      const keepMemory =
-        memory.importance > neighbor.importance ||
-        (memory.importance === neighbor.importance && memoryCreatedAt >= neighborCreatedAt)
-      const winner = keepMemory ? memory : neighbor
-      const loser = keepMemory ? neighbor : memory
-
-      await databaseService.relate(
-        ensureRecordId(winner.id, TABLES.MEMORY),
-        MEMORY_RELATION_TABLE,
-        ensureRecordId(loser.id, TABLES.MEMORY),
-        { relationType: RELATION_SUPERSEDES, confidence: 1.0 },
-      )
-
-      await databaseService.query(
-        new BoundQuery(
-          `UPDATE ${MEMORY_TABLE}
-         SET archivedAt = time::now(), validUntil = time::now()
-         WHERE id = $loserId AND archivedAt IS NONE`,
-          { loserId: ensureRecordId(loser.id, TABLES.MEMORY) },
-        ),
-      )
+        const isHardMatch = neighbor.similarity >= HARD_SIMILARITY_THRESHOLD
+        if (!isHardMatch) {
+          const [shorter, longer] =
+            memory.content.length <= neighbor.content.length
+              ? [memory.content, neighbor.content]
+              : [neighbor.content, memory.content]
+          if (!isContentSubsumed(shorter, longer)) continue
+        }
 
-      await databaseService.insert<Record<string, unknown>>(MEMORY_HISTORY_TABLE, {
-        memoryId: ensureRecordId(loser.id, TABLES.MEMORY),
-        prevValue: loser.content,
-        event: 'DELETE',
-      })
+        const memoryCreatedAt = unsafeDateFrom(memory.createdAt).getTime()
+        const neighborCreatedAt = unsafeDateFrom(neighbor.createdAt).getTime()
+        const keepMemory =
+          memory.importance > neighbor.importance ||
+          (memory.importance === neighbor.importance && memoryCreatedAt >= neighborCreatedAt)
+        const winner = keepMemory ? memory : neighbor
+        const loser = keepMemory ? neighbor : memory
+
+        yield* tryMemoryConsolidationEffect('archive-memory', () =>
+          db().query(
+            new BoundQuery(
+              `UPDATE ${MEMORY_TABLE}
+               SET archivedAt = time::now(), validUntil = time::now()
+               WHERE id = $loserId AND archivedAt IS NONE`,
+              { loserId: ensureRecordId(loser.id, TABLES.MEMORY) },
+            ),
+          ),
+        )
+        yield* tryMemoryConsolidationEffect('relate-memories', () =>
+          db().relate(
+            ensureRecordId(winner.id, TABLES.MEMORY),
+            MEMORY_RELATION_TABLE,
+            ensureRecordId(loser.id, TABLES.MEMORY),
+            { relationType: RELATION_SUPERSEDES, confidence: 1.0 },
+          ),
+        )
+        yield* tryMemoryConsolidationEffect('write-history', () =>
+          db().insert<Record<string, unknown>>(MEMORY_HISTORY_TABLE, {
+            memoryId: ensureRecordId(loser.id, TABLES.MEMORY),
+            prevValue: loser.content,
+            event: 'DELETE',
+          }),
+        )
 
-      archived.add(loser.id)
-      mergeCount++
+        archived.add(loser.id)
+        mergeCount++
 
-      if (!keepMemory) break
+        if (!keepMemory) continue
+      }
     }
-  }
 
-  return mergeCount
+    return mergeCount
+  })
 }
 
-async function pruneStaleMemories(): Promise<number> {
-  const stale = await databaseService.query<{ id: RecordIdInput }>(
-    new BoundQuery(
-      `SELECT id FROM ${MEMORY_TABLE}
-     WHERE accessCount = 0
-       AND createdAt < time::now() - 90d
-       AND archivedAt IS NONE
-       AND importance < 0.5
-     LIMIT 200`,
-    ),
-  )
+function pruneStaleMemoriesEffect() {
+  return Effect.gen(function* () {
+    const stale = yield* tryMemoryConsolidationEffect('load-stale', () =>
+      db().query<{ id: RecordIdInput }>(
+        new BoundQuery(
+          `SELECT id FROM ${MEMORY_TABLE}
+           WHERE accessCount = 0
+             AND createdAt < time::now() - 90d
+             AND archivedAt IS NONE
+             AND importance < 0.5
+             LIMIT 200`,
+        ),
+      ),
+    )
 
-  if (stale.length === 0) return 0
+    if (stale.length === 0) return 0
 
-  const staleIds = stale.map((row) => ensureRecordId(row.id, TABLES.MEMORY))
-  await databaseService.updateWhere(MEMORY_TABLE, inside('id', staleIds), { archivedAt: new Date() })
+    const staleIds = stale.map((row) => ensureRecordId(row.id, TABLES.MEMORY))
+    yield* tryMemoryConsolidationEffect('update-stale', () =>
+      db().updateWhere(MEMORY_TABLE, inside('id', staleIds), { archivedAt: nowDate() }),
+    )
 
-  return stale.length
+    return stale.length
+  })
 }
 
-async function collapseSupersedeCh(): Promise<number> {
-  const middleNodes = await databaseService.query<{
-    middleId: RecordIdInput
-    predecessors: RecordIdInput[]
-    successors: RecordIdInput[]
-  }>(
-    new BoundQuery(
-      `SELECT
-       id AS middleId,
-       <-${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']<-${MEMORY_TABLE}.id AS predecessors,
-       ->${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']->${MEMORY_TABLE}.id AS successors
-     FROM ${MEMORY_TABLE}
-     WHERE archivedAt IS NONE
-       AND count(->${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']) > 0
-       AND count(<-${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']) > 0
-     LIMIT ${MEMORY.MAX_KNN_LIMIT}`,
-    ),
-  )
-
-  let collapsed = 0
-
-  for (const node of middleNodes) {
-    const predecessors = node.predecessors
-    const successors = node.successors
-
-    for (const predId of predecessors) {
-      for (const succId of successors) {
-        const predRef = ensureRecordId(predId, TABLES.MEMORY)
-        const succRef = ensureRecordId(succId, TABLES.MEMORY)
-        const existing = await databaseService.query<{ id: RecordIdInput }>(
-          new BoundQuery(
-            `SELECT id FROM ${MEMORY_RELATION_TABLE}
-           WHERE in = $predId AND out = $succId AND relationType = '${RELATION_SUPERSEDES}'
-           LIMIT 1`,
-            { predId: predRef, succId: succRef },
-          ),
-        )
+function collapseSupersededChainEffect() {
+  return Effect.gen(function* () {
+    const middleNodes = yield* tryMemoryConsolidationEffect('load-middle-nodes', () =>
+      db().query<{ middleId: RecordIdInput; predecessors: RecordIdInput[]; successors: RecordIdInput[] }>(
+        new BoundQuery(
+          `SELECT
+             id AS middleId,
+             <-${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']<-${MEMORY_TABLE}.id AS predecessors,
+             ->${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']->${MEMORY_TABLE}.id AS successors
+           FROM ${MEMORY_TABLE}
+           WHERE archivedAt IS NONE
+             AND count(->${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']) > 0
+             AND count(<-${MEMORY_RELATION_TABLE}[WHERE relationType = '${RELATION_SUPERSEDES}']) > 0
+             LIMIT ${MEMORY.MAX_KNN_LIMIT}`,
+        ),
+      ),
+    )
 
-        if (existing.length === 0) {
-          await databaseService.relate(predRef, MEMORY_RELATION_TABLE, succRef, {
-            relationType: RELATION_SUPERSEDES,
-            confidence: 1.0,
-          })
+    let collapsed = 0
+
+    for (const node of middleNodes) {
+      for (const predId of node.predecessors) {
+        for (const succId of node.successors) {
+          const predRef = ensureRecordId(predId, TABLES.MEMORY)
+          const succRef = ensureRecordId(succId, TABLES.MEMORY)
+          const existing = yield* tryMemoryConsolidationEffect('read-existing-relation', () =>
+            db().query<{ id: RecordIdInput }>(
+              new BoundQuery(
+                `SELECT id FROM ${MEMORY_RELATION_TABLE}
+                 WHERE in = $predId AND out = $succId AND relationType = '${RELATION_SUPERSEDES}'
+                 LIMIT 1`,
+                { predId: predRef, succId: succRef },
+              ),
+            ),
+          )
+
+          if (existing.length === 0) {
+            yield* tryMemoryConsolidationEffect('write-relation', () =>
+              db().relate(predRef, MEMORY_RELATION_TABLE, succRef, {
+                relationType: RELATION_SUPERSEDES,
+                confidence: 1.0,
+              }),
+            )
+          }
         }
       }
-    }
 
-    await databaseService.updateWhere(MEMORY_TABLE, eq('id', ensureRecordId(node.middleId, TABLES.MEMORY)), {
-      archivedAt: new Date(),
-    })
+      yield* tryMemoryConsolidationEffect('update-middle-node', () =>
+        db().updateWhere(MEMORY_TABLE, eq('id', ensureRecordId(node.middleId, TABLES.MEMORY)), {
+          archivedAt: nowDate(),
+        }),
+      )
 
-    collapsed++
-  }
+      collapsed++
+    }
 
-  return collapsed
+    return collapsed
+  })
 }
 
-async function decayImportance(): Promise<number> {
-  const standardResult = await databaseService.query<{ count: number }>(
-    new BoundQuery(
-      `UPDATE ${MEMORY_TABLE}
-     SET importance = math::max([0.1, importance * 0.95])
-     WHERE lastAccessedAt IS NOT NONE
-       AND lastAccessedAt < time::now() - 30d
-       AND archivedAt IS NONE
-       AND importance > 0.1
-       AND (durability = 'standard' OR durability IS NONE)
-     RETURN count() AS count`,
-    ),
-  )
-
-  const ephemeralResult = await databaseService.query<{ count: number }>(
-    new BoundQuery(
-      `UPDATE ${MEMORY_TABLE}
-     SET importance = math::max([0.1, importance * 0.85])
-     WHERE lastAccessedAt IS NOT NONE
-       AND lastAccessedAt < time::now() - 30d
-       AND archivedAt IS NONE
-       AND importance > 0.1
-       AND durability = 'ephemeral'
-     RETURN count() AS count`,
-    ),
-  )
+function decayImportanceEffect() {
+  return Effect.gen(function* () {
+    const [standardResult, ephemeralResult] = yield* Effect.all([
+      tryMemoryConsolidationEffect('decay-standard', () =>
+        db().query<{ count: number }>(
+          new BoundQuery(
+            `UPDATE ${MEMORY_TABLE}
+             SET importance = math::max([0.1, importance * 0.95])
+             WHERE lastAccessedAt IS NOT NONE
+               AND lastAccessedAt < time::now() - 30d
+               AND archivedAt IS NONE
+               AND importance > 0.1
+               AND (durability = 'standard' OR durability IS NONE)
+               RETURN count() AS count`,
+          ),
+        ),
+      ),
+      tryMemoryConsolidationEffect('decay-ephemeral', () =>
+        db().query<{ count: number }>(
+          new BoundQuery(
+            `UPDATE ${MEMORY_TABLE}
+             SET importance = math::max([0.1, importance * 0.85])
+             WHERE lastAccessedAt IS NOT NONE
+               AND lastAccessedAt < time::now() - 30d
+               AND archivedAt IS NONE
+               AND importance > 0.1
+               AND durability = 'ephemeral'
+               RETURN count() AS count`,
+          ),
+        ),
+      ),
+    ])
 
-  return (standardResult[0]?.count ?? 0) + (ephemeralResult[0]?.count ?? 0)
+    return (standardResult[0]?.count ?? 0) + (ephemeralResult[0]?.count ?? 0)
+  })
 }
 
-async function cleanupOrphanedRelations(): Promise<number> {
-  const result = await databaseService.query<{ count: number }>(
-    new BoundQuery(
-      `DELETE ${MEMORY_RELATION_TABLE}
-     WHERE in.archivedAt IS NOT NONE OR out.archivedAt IS NOT NONE
-     RETURN count() AS count`,
-    ),
-  )
-  return result[0]?.count ?? 0
+function cleanupOrphanedRelationsEffect() {
+  return Effect.gen(function* () {
+    const result = yield* tryMemoryConsolidationEffect('cleanup-relations', () =>
+      db().query<{ count: number }>(
+        new BoundQuery(
+          `DELETE ${MEMORY_RELATION_TABLE}
+           WHERE in.archivedAt IS NOT NONE OR out.archivedAt IS NOT NONE
+           RETURN count() AS count`,
+        ),
+      ),
+    )
+    return result[0]?.count ?? 0
+  })
 }
 
-const handler = async (job: SandboxedJob<MemoryConsolidationJob>) => {
-  const targetScope = job.data.scopeId
-
-  try {
-    let totalMerged = 0
-    if (targetScope) {
-      totalMerged = await deduplicateScope(targetScope)
-    } else {
-      const scopeIds = await databaseService.query<string>(
-        new BoundQuery(`SELECT VALUE scopeId FROM ${MEMORY_TABLE} WHERE archivedAt IS NONE GROUP BY scopeId`),
-      )
-
-      const results = await Promise.all(scopeIds.map(deduplicateScope))
-      totalMerged = results.reduce((a, b) => a + b, 0)
-    }
-
-    const pruned = await pruneStaleMemories()
-    const collapsed = await collapseSupersedeCh()
-    const decayed = await decayImportance()
-    const orphaned = await cleanupOrphanedRelations()
+const handler = (job: SandboxedJob<MemoryConsolidationJob>) =>
+  Effect.runPromise(
+    Effect.gen(function* () {
+      const targetScope = job.data.scopeId
+
+      let totalMerged = 0
+      if (targetScope) {
+        totalMerged = yield* deduplicateScopeEffect(targetScope)
+      } else {
+        const scopeIds = yield* tryMemoryConsolidationEffect('load-scope-ids', () =>
+          db().query<string>(
+            new BoundQuery(`SELECT VALUE scopeId FROM ${MEMORY_TABLE} WHERE archivedAt IS NONE GROUP BY scopeId`),
+          ),
+        )
 
-    serverLogger.info`Memory consolidation complete (merged: ${totalMerged}, pruned: ${pruned}, collapsed: ${collapsed}, decayed: ${decayed}, orphaned relations: ${orphaned})`
-  } catch (error) {
-    const serialized = toSandboxedWorkerError(error, 'Memory consolidation job failed')
-    serverLogger.error`${serialized.message}`
-    throw serialized
-  }
-}
+        const results = yield* Effect.forEach(scopeIds, deduplicateScopeEffect, { concurrency: 2 })
+        totalMerged = results.reduce((a, b) => a + b, 0)
+      }
+      const [pruned, collapsed, decayed, orphaned] = yield* Effect.all([
+        pruneStaleMemoriesEffect(),
+        collapseSupersededChainEffect(),
+        decayImportanceEffect(),
+        cleanupOrphanedRelationsEffect(),
+      ])
+
+      yield* Effect.sync(() => {
+        serverLogger.info`Memory consolidation complete (merged: ${totalMerged}, pruned: ${pruned}, collapsed: ${collapsed}, decayed: ${decayed}, orphaned relations: ${orphaned})`
+      })
+    }).pipe(
+      Effect.catch((error) => {
+        const serialized = toSandboxedWorkerError(error, 'Memory consolidation job failed')
+        return Effect.sync(() => {
+          serverLogger.error`${serialized.message}`
+        }).pipe(Effect.andThen(Effect.fail(serialized)))
+      }),
+    ),
+  )
 
 export default createTracedWorkerProcessor('memory-consolidation', handler)