@shadowforge0/aquifer-memory 1.0.2 → 1.2.1

package/core/aquifer.js

@@ -9,6 +9,7 @@ const entity = require('./entity');
 const { hybridRank } = require('./hybrid-rank');
 const { summarize } = require('../pipeline/summarize');
 const { extractEntities } = require('../pipeline/extract-entities');
+const { createEmbedder } = require('../pipeline/embed');
 
 // ---------------------------------------------------------------------------
 // Schema name validation
@@ -54,48 +55,104 @@ function buildRerankDocument(row, maxChars) {
   return text;
 }
 
+// ---------------------------------------------------------------------------
+// resolveEmbedFn — v1.2.0 embed autodetect (explicit > object > env > null)
+// ---------------------------------------------------------------------------
+
+function resolveEmbedFn(embedConfig, env) {
+  if (embedConfig && typeof embedConfig.fn === 'function') {
+    return embedConfig.fn;
+  }
+  if (embedConfig && embedConfig.provider) {
+    const embedder = createEmbedder(embedConfig);
+    return (texts) => embedder.embedBatch(texts);
+  }
+  const provider = env.EMBED_PROVIDER;
+  if (!provider) return null;
+
+  const opts = { provider };
+  if (provider === 'ollama') {
+    opts.ollamaUrl = env.OLLAMA_URL || env.AQUIFER_EMBED_BASE_URL || 'http://localhost:11434';
+    opts.model = env.AQUIFER_EMBED_MODEL || 'bge-m3';
+  } else if (provider === 'openai') {
+    opts.openaiApiKey = env.OPENAI_API_KEY;
+    if (!opts.openaiApiKey) {
+      throw new Error('EMBED_PROVIDER=openai requires OPENAI_API_KEY');
+    }
+    opts.openaiModel = env.AQUIFER_EMBED_MODEL || 'text-embedding-3-small';
+    if (env.AQUIFER_EMBED_DIM) opts.openaiDimensions = Number(env.AQUIFER_EMBED_DIM);
+  } else {
+    throw new Error(`EMBED_PROVIDER=${provider} not supported by autodetect (use 'ollama' or 'openai', or pass config.embed.fn explicitly)`);
+  }
+  const embedder = createEmbedder(opts);
+  return (texts) => embedder.embedBatch(texts);
+}
+
 // ---------------------------------------------------------------------------
 // createAquifer
 // ---------------------------------------------------------------------------
 
-function createAquifer(config) {
-  if (!config || !config.db) {
-    throw new Error('config.db (pg.Pool or connection string) is required');
+function createAquifer(config = {}) {
+  // v1.2.0: db falls back to DATABASE_URL / AQUIFER_DB_URL env so hosts can
+  // call createAquifer() with zero args for install-and-go.
+  const dbInput = config.db !== undefined
+    ? config.db
+    : (process.env.DATABASE_URL || process.env.AQUIFER_DB_URL || null);
+
+  if (!dbInput) {
+    throw new Error(
+      'Aquifer requires a database: pass config.db (pg.Pool or connection string), '
+      + 'or set DATABASE_URL / AQUIFER_DB_URL in the environment.'
+    );
   }
 
-  const schema = config.schema || 'aquifer';
+  const schema = config.schema || process.env.AQUIFER_SCHEMA || 'aquifer';
   validateSchema(schema);
 
   if (config.tenantId === '') throw new Error('config.tenantId must not be empty');
-  const tenantId = config.tenantId || 'default';
+  const tenantId = config.tenantId || process.env.AQUIFER_TENANT_ID || 'default';
 
   // Pool management
   let pool;
   let ownsPool = false;
-  if (typeof config.db === 'string') {
-    pool = new Pool({ connectionString: config.db });
+  if (typeof dbInput === 'string') {
+    pool = new Pool({ connectionString: dbInput });
     ownsPool = true;
   } else {
-    pool = config.db;
+    pool = dbInput;
     ownsPool = !!config.ownsPool;  // allow factory to claim ownership
   }
 
   // Embed config (lazy — only required for recall/enrich)
-  const embedFn = config.embed && typeof config.embed.fn === 'function' ? config.embed.fn : null;
-  let embedDim = config.embed ? (config.embed.dim || null) : null;
-
+  // v1.2.0 fallback chain:
+  //   1. config.embed.fn (explicit function)
+  //   2. config.embed.provider (build via createEmbedder)
+  //   3. EMBED_PROVIDER env + provider-specific key (zero-arg install-and-go)
+  //   4. null — defer to requireEmbed() at call time
+  const embedFn = resolveEmbedFn(config.embed, process.env);
   function requireEmbed(op) {
-    if (!embedFn) throw new Error(`Aquifer.${op}() requires config.embed.fn (async (texts) => number[][])`);
+    if (!embedFn) throw new Error(`Aquifer.${op}() requires config.embed.fn or EMBED_PROVIDER env (async (texts) => number[][])`);
   }
 
   // LLM config (optional — only needed for enrich with built-in summarize)
-  const llmFn = config.llm && typeof config.llm.fn === 'function' ? config.llm.fn : null;
+  // v1.2.0: falls back to AQUIFER_LLM_PROVIDER env + provider-specific key.
+  const { resolveLlmFn } = require('../consumers/shared/llm-autodetect');
+  const llmFn = resolveLlmFn(config.llm, process.env);
 
   // Summarize config
   const summarizePromptFn = config.summarize && config.summarize.prompt ? config.summarize.prompt : null;
 
+  // Enrich stale-claim window: a 'processing' session older than this is
+  // reclaimable by a concurrent enrich() caller (covers crashed workers).
+  const staleEnrichMinutes = Number.isFinite(config.staleEnrichMinutes)
+    ? Math.max(1, Math.floor(config.staleEnrichMinutes))
+    : 10;
+
   // Entity config
   let entitiesEnabled = config.entities && config.entities.enabled === true;
+
+  // Facts config (opt-in consolidation lifecycle)
+  let factsEnabled = config.facts && config.facts.enabled === true;
   const mergeCall = config.entities && config.entities.mergeCall !== undefined ? config.entities.mergeCall : true;
   const entityPromptFn = config.entities && config.entities.prompt ? config.entities.prompt : null;
   const entityScope = (config.entities && config.entities.scope) || 'default';
@@ -187,21 +244,38 @@ function createAquifer(config) {
     // --- lifecycle ---
 
     async migrate() {
-      // 1. Run base DDL
-      const baseSql = loadSql('001-base.sql', schema);
-      await pool.query(baseSql);
+      // Advisory lock prevents concurrent migrations across processes.
+      // Lock key is derived from schema name to allow parallel migration
+      // of different schemas in the same database.
+      const lockKey = Buffer.from(`aquifer:${schema}`).reduce((h, b) => (h * 31 + b) & 0x7fffffff, 0);
+      await pool.query('SELECT pg_advisory_lock($1)', [lockKey]);
+      try {
+        // 1. Run base DDL
+        const baseSql = loadSql('001-base.sql', schema);
+        await pool.query(baseSql);
+
+        // 2. If entities enabled, run entity DDL
+        if (entitiesEnabled) {
+          const entitySql = loadSql('002-entities.sql', schema);
+          await pool.query(entitySql);
+        }
 
-      // 2. If entities enabled, run entity DDL
-      if (entitiesEnabled) {
-        const entitySql = loadSql('002-entities.sql', schema);
-        await pool.query(entitySql);
-      }
+        // 3. Trust + feedback (always, not gated by entities)
+        const trustSql = loadSql('003-trust-feedback.sql', schema);
+        await pool.query(trustSql);
 
-      // 3. Trust + feedback (always, not gated by entities)
-      const trustSql = loadSql('003-trust-feedback.sql', schema);
-      await pool.query(trustSql);
+        // 4. Facts / consolidation (opt-in)
+        if (factsEnabled) {
+          const factsSql = loadSql('004-facts.sql', schema);
+          await pool.query(factsSql);
+        }
 
-      migrated = true;
+        migrated = true;
+      } finally {
+        await pool.query('SELECT pg_advisory_unlock($1)', [lockKey]).catch((err) => {
+          console.warn(`[aquifer] failed to release migration advisory lock for schema "${schema}": ${err.message}`);
+        });
+      }
     },
 
     async close() {
@@ -216,7 +290,7 @@ function createAquifer(config) {
       sources.set(name, {
         type: opts.type || 'custom',
         search: opts.search || null,
-        weight: opts.weight !== null && opts.weight !== undefined ? opts.weight : 1.0,
+        weight: opts.weight !== undefined && opts.weight !== undefined ? opts.weight : 1.0,
       });
     },
 
@@ -229,6 +303,32 @@ function createAquifer(config) {
       }
     },
 
+    async enableFacts() {
+      factsEnabled = true;
+      // Run the facts DDL (idempotent — all CREATE/ALTER use IF NOT EXISTS).
+      // Safe to call repeatedly; also safe to call before migrate() (will no-op
+      // until base schema exists, which enrich/commit will materialize).
+      await ensureMigrated();
+      const factsSql = loadSql('004-facts.sql', schema);
+      await pool.query(factsSql);
+    },
+
+    async consolidate(sessionId, opts = {}) {
+      if (!factsEnabled) throw new Error('aquifer.consolidate() requires enableFacts() first');
+      await ensureMigrated();
+      const { applyConsolidation } = require('../pipeline/consolidation');
+      const agentId = opts.agentId || 'agent';
+      return applyConsolidation(pool, {
+        actions: opts.actions || [],
+        agentId,
+        sessionId,
+        schema,
+        tenantId,
+        normalizeSubject: opts.normalizeSubject || null,
+        recapOverview: opts.recapOverview || '',
+      });
+    },
+
     // --- write path ---
 
     async commit(sessionId, messages, opts = {}) {
@@ -293,17 +393,19 @@ function createAquifer(config) {
       const postProcess = opts.postProcess || null;  // async (ctx) => void
       const optModel = 'model' in opts ? opts.model : undefined; // undefined = no override
 
-      // 1. Optimistic lock: claim session for processing
-      //    Also reclaim stale 'processing' sessions (stuck > 10 min = likely killed process)
-      const STALE_MINUTES = 10;
+      // 1. Optimistic lock: claim session for processing.
+      //    Also reclaim stale 'processing' sessions (likely killed worker).
+      //    Stale window is config.staleEnrichMinutes (default 10).
       const claimResult = await pool.query(
         `UPDATE ${qi(schema)}.sessions
         SET processing_status = 'processing', processing_started_at = NOW()
         WHERE session_id = $1 AND agent_id = $2 AND tenant_id = $3
           AND (processing_status IN ('pending', 'failed')
-               OR (processing_status = 'processing' AND (processing_started_at IS NULL OR processing_started_at < NOW() - INTERVAL '${STALE_MINUTES} minutes')))
+               OR (processing_status = 'processing'
+                   AND (processing_started_at IS NULL
+                        OR processing_started_at < NOW() - make_interval(mins => $4))))
         RETURNING *`,
-        [sessionId, agentId, tenantId]
+        [sessionId, agentId, tenantId, staleEnrichMinutes]
       );
       const session = claimResult.rows[0];
       if (!session) {
@@ -324,34 +426,46 @@ function createAquifer(config) {
       // 2. Extract user turns
       const turns = storage.extractUserTurns(normalized);
 
+      // Collected across pre-tx and tx phases; any non-empty warnings demote
+      // the final status from 'succeeded' to 'partial' (see step 8 below).
+      const warnings = [];
+
       // 3. Summarize (custom or built-in)
       let summaryResult = null;
       let entityRaw = null;
       let extra = null;
 
       if (!skipSummary && normalized.length > 0) {
-        if (customSummaryFn) {
-          // Custom pipeline: caller handles LLM call and parsing
-          summaryResult = await customSummaryFn(normalized);
-          if (summaryResult.entityRaw) entityRaw = summaryResult.entityRaw;
-          if (summaryResult.extra) extra = summaryResult.extra;
-        } else {
-          // Built-in pipeline
-          const doMergeEntities = entitiesEnabled && mergeCall && !skipEntities;
-          summaryResult = await summarize(normalized, {
-            llmFn,
-            promptFn: summarizePromptFn,
-            mergeEntities: doMergeEntities,
-          });
-          if (summaryResult.entityRaw) {
-            entityRaw = summaryResult.entityRaw;
+        // Pre-transaction failures (customSummaryFn / summarize throws) would
+        // otherwise bubble out and leave the session stuck in 'processing'
+        // until stale reclaim. Capture as a warning so status ends 'partial',
+        // keeping parity with how embed/entity-extract failures are treated.
+        try {
+          if (customSummaryFn) {
+            // Custom pipeline: caller handles LLM call and parsing
+            summaryResult = await customSummaryFn(normalized);
+            if (summaryResult && summaryResult.entityRaw) entityRaw = summaryResult.entityRaw;
+            if (summaryResult && summaryResult.extra) extra = summaryResult.extra;
+          } else {
+            // Built-in pipeline
+            const doMergeEntities = entitiesEnabled && mergeCall && !skipEntities;
+            summaryResult = await summarize(normalized, {
+              llmFn,
+              promptFn: summarizePromptFn,
+              mergeEntities: doMergeEntities,
+            });
+            if (summaryResult.entityRaw) {
+              entityRaw = summaryResult.entityRaw;
+            }
           }
+        } catch (e) {
+          warnings.push(`summary step failed: ${e.message}`);
+          summaryResult = null;
         }
       }
 
       // 4. Pre-compute all LLM/embed results BEFORE opening transaction
       //    (avoids holding pool connection during slow LLM/embed calls)
-      const warnings = [];
       let summaryEmbedding = null;
       let turnVectors = null;
       let parsedEntities = [];
@@ -485,7 +599,11 @@ function createAquifer(config) {
         await client.query('ROLLBACK').catch(() => {});
         try {
           await storage.markStatus(pool, session.id, 'failed', err.message, { schema });
-        } catch (_) { /* swallow */ }
+        } catch (markErr) {
+          // Secondary failure: session is stuck in 'processing' until stale reclaim.
+          // Surface so operators notice and don't silently rely on the timeout.
+          console.warn(`[aquifer] enrich failed for session ${sessionId} AND markStatus('failed') also failed: ${markErr.message}`);
+        }
         throw err;
       } finally {
         client.release();
@@ -683,7 +801,7 @@ function createAquifer(config) {
               entityScoreBySession.set(row.session_id, parseInt(row.entity_count) / maxCount);
             }
           }
-        } catch (_) { /* entity search failure non-fatal */ }
+        } catch { /* entity search failure non-fatal */ }
       }
 
       // 3. Run search paths in parallel (conditioned on mode)
@@ -738,7 +856,7 @@ function createAquifer(config) {
       for (const r of [...filteredFts, ...filteredEmb, ...filteredTurn]) {
         const sid = r.session_id || String(r.id);
         const ss = typeof r.structured_summary === 'string'
-          ? (() => { try { return JSON.parse(r.structured_summary); } catch (_) { return null; } })()
+          ? (() => { try { return JSON.parse(r.structured_summary); } catch { return null; } })()
           : r.structured_summary;
         if (ss && Array.isArray(ss.open_loops) && ss.open_loops.length > 0) {
           openLoopSet.add(sid);
@@ -751,7 +869,7 @@ function createAquifer(config) {
       const externalPromises = [];
       for (const [name, sourceConfig] of sources) {
         if (typeof sourceConfig.search === 'function') {
-          const w = sourceConfig.weight !== null && sourceConfig.weight !== undefined ? sourceConfig.weight : 1.0;
+          const w = sourceConfig.weight !== undefined && sourceConfig.weight !== undefined ? sourceConfig.weight : 1.0;
           externalPromises.push(
             Promise.race([
               sourceConfig.search(query, opts),
@@ -822,7 +940,7 @@ function createAquifer(config) {
       if (sessionRowIds.length > 0) {
         try {
           await storage.recordAccess(pool, sessionRowIds, { schema });
-        } catch (_) { /* access recording non-fatal */ }
+        } catch { /* access recording non-fatal */ }
       }
 
       // 8. Format results
@@ -833,7 +951,6 @@ function createAquifer(config) {
         startedAt: r.started_at,
         summaryText: r.summary_text || null,
         structuredSummary: r.structured_summary || null,
-        summarySnippet: r.summary_snippet || null,
         matchedTurnText: r.matched_turn_text || null,
         matchedTurnIndex: r.matched_turn_index || null,
         score: r._rerankScore ?? r._score,
@@ -904,35 +1021,31 @@ function createAquifer(config) {
       return { id: result.rows[0].id, sessionId, agentId, status: 'skipped' };
     },
 
-    async getSessionFull(sessionId) {
-      const result = await pool.query(
-        `SELECT * FROM ${qi(schema)}.sessions
-        WHERE session_id = $1 AND tenant_id = $2
-        LIMIT 1`,
-        [sessionId, tenantId]
-      );
-      const session = result.rows[0];
-      if (!session) return null;
-
-      const sumResult = await pool.query(
-        `SELECT * FROM ${qi(schema)}.session_summaries
-        WHERE session_row_id = $1
-        LIMIT 1`,
-        [session.id]
-      );
-
-      return {
-        session,
-        summary: sumResult.rows[0] || null,
-      };
-    },
-
     // --- public config accessor ---
 
     getConfig() {
       return { schema, tenantId };
     },
 
+    // v1.2.0: expose the internal pool so host persona layers can reuse it
+    // for host-owned tables (e.g. daily_entries). Read-only — callers should
+    // not call pool.end() on it; use aquifer.close() for that.
+    getPool() {
+      return pool;
+    },
+
+    // v1.2.0: expose resolved LLM function. May be null if no llm.fn was
+    // supplied and AQUIFER_LLM_PROVIDER env is unset. Persona layers that
+    // implement custom summaryFn can reuse this instead of wiring their own.
+    getLlmFn() {
+      return llmFn;
+    },
+
+    // v1.2.0: expose resolved embed function (may be null same as LLM).
+    getEmbedFn() {
+      return embedFn;
+    },
+
     // --- admin query helpers ---
 
     async getStats() {
@@ -965,7 +1078,7 @@ function createAquifer(config) {
           [tenantId]
         );
         entityCount = entResult.rows[0]?.count || 0;
-      } catch (_) { /* entities table may not exist */ }
+      } catch { /* entities table may not exist */ }
 
       return {
         sessions: Object.fromEntries(sessions.rows.map(r => [r.processing_status, r.count])),
@@ -1024,7 +1137,11 @@ function createAquifer(config) {
       const maxChars = opts.maxChars || 4000;
       const format = opts.format || 'structured';
 
-      const where = [`s.tenant_id = $1`, `s.processing_status = 'succeeded'`];
+      // 'partial' sessions have a summary but recorded warnings during enrich;
+      // they are user-visible content, not in-progress — bootstrap must include
+      // them alongside 'succeeded'. 'pending' / 'processing' have no summary
+      // yet and are correctly excluded.
+      const where = [`s.tenant_id = $1`, `s.processing_status IN ('succeeded', 'partial')`];
       const params = [tenantId];
 
       if (agentId) {
@@ -1037,7 +1154,10 @@ function createAquifer(config) {
       }
 
       params.push(lookbackDays);
-      where.push(`s.started_at > now() - ($${params.length} || ' days')::interval`);
+      // upsertSession sets ended_at on every commit but started_at / last_message_at
+      // only when the caller supplies them — fall back through both so sessions
+      // committed without explicit timestamps remain reachable.
+      where.push(`COALESCE(s.last_message_at, s.ended_at, s.started_at) > now() - ($${params.length} || ' days')::interval`);
 
       params.push(limit);
 
@@ -1047,7 +1167,7 @@ function createAquifer(config) {
          FROM ${qi(schema)}.sessions s
          JOIN ${qi(schema)}.session_summaries ss ON ss.session_row_id = s.id
          WHERE ${where.join(' AND ')}
-         ORDER BY s.started_at DESC
+         ORDER BY COALESCE(s.last_message_at, s.ended_at, s.started_at) DESC
          LIMIT $${params.length}`,
         params
       );
@@ -1126,8 +1246,6 @@ function formatBootstrapText(data, maxChars) {
   }
 
   let truncated = false;
-  const parts = [];
-
   // Build session lines (newest first, truncate from oldest if over budget)
   const sessionLines = [];
   for (const s of data.sessions) {

package/core/entity.js

@@ -222,27 +222,40 @@ async function upsertEntityRelations(pool, {
 }) {
   if (!pairs || pairs.length === 0) return { upserted: 0 };
   const ts = occurredAt || new Date().toISOString();
-  let upserted = 0;
 
+  // Filter and normalize pairs
+  const validPairs = [];
   for (const { srcEntityId, dstEntityId } of pairs) {
     if (!srcEntityId || !dstEntityId || srcEntityId === dstEntityId) continue;
+    validPairs.push({
+      lo: Math.min(srcEntityId, dstEntityId),
+      hi: Math.max(srcEntityId, dstEntityId),
+    });
+  }
 
-    const lo = Math.min(srcEntityId, dstEntityId);
-    const hi = Math.max(srcEntityId, dstEntityId);
-
-    await pool.query(
-      `INSERT INTO ${qi(schema)}.entity_relations
-        (src_entity_id, dst_entity_id, co_occurrence_count, first_seen_at, last_seen_at)
-      VALUES ($1, $2, 1, $3, $3)
-      ON CONFLICT (src_entity_id, dst_entity_id) DO UPDATE SET
-        co_occurrence_count = ${qi(schema)}.entity_relations.co_occurrence_count + 1,
-        last_seen_at        = GREATEST(${qi(schema)}.entity_relations.last_seen_at, EXCLUDED.last_seen_at)`,
-      [lo, hi, ts]
-    );
-    upserted++;
+  if (validPairs.length === 0) return { upserted: 0 };
+
+  // Batch insert: multi-row VALUES
+  const valueClauses = [];
+  const params = [];
+
+  for (const { lo, hi } of validPairs) {
+    const off = params.length;
+    params.push(lo, hi, ts);
+    valueClauses.push(`($${off+1}, $${off+2}, 1, $${off+3}, $${off+3})`);
   }
 
-  return { upserted };
+  await pool.query(
+    `INSERT INTO ${qi(schema)}.entity_relations
+      (src_entity_id, dst_entity_id, co_occurrence_count, first_seen_at, last_seen_at)
+    VALUES ${valueClauses.join(',\n')}
+    ON CONFLICT (src_entity_id, dst_entity_id) DO UPDATE SET
+      co_occurrence_count = ${qi(schema)}.entity_relations.co_occurrence_count + 1,
+      last_seen_at        = GREATEST(${qi(schema)}.entity_relations.last_seen_at, EXCLUDED.last_seen_at)`,
+    params
+  );
+
+  return { upserted: validPairs.length };
 }
 
 // ---------------------------------------------------------------------------
@@ -373,8 +386,7 @@ async function resolveEntities(pool, {
     if (!normQ || seen.has(normQ)) continue;
     seen.set(normQ, true);
 
-    const escaped = _escapeIlike(normQ);
-    const result = await pool.query(
+  const result = await pool.query(
       `SELECT id, name, normalized_name
       FROM ${qi(schema)}.entities
       WHERE status = 'active'