memory-braid 0.2.0 → 0.3.3

package/src/entities.ts

@@ -0,0 +1,354 @@
+import os from "node:os";
+import path from "node:path";
+import { normalizeWhitespace } from "./chunking.js";
+import type { MemoryBraidConfig } from "./config.js";
+import { MemoryBraidLogger } from "./logger.js";
+
+type NerPipeline = (text: string, options?: Record<string, unknown>) => Promise<unknown>;
+
+type NerRecord = {
+  word?: unknown;
+  entity_group?: unknown;
+  entity?: unknown;
+  score?: unknown;
+};
+
+export type ExtractedEntity = {
+  text: string;
+  type: "person" | "organization" | "location" | "misc";
+  score: number;
+  canonicalUri: string;
+};
+
+function summarizeEntityTypes(entities: ExtractedEntity[]): Record<string, number> {
+  const summary: Record<string, number> = {};
+  for (const entity of entities) {
+    summary[entity.type] = (summary[entity.type] ?? 0) + 1;
+  }
+  return summary;
+}
+
+function resolveStateDir(explicitStateDir?: string): string {
+  const resolved =
+    explicitStateDir?.trim() ||
+    process.env.OPENCLAW_STATE_DIR?.trim() ||
+    path.join(os.homedir(), ".openclaw");
+  return path.resolve(resolved);
+}
+
+export function resolveEntityModelCacheDir(stateDir?: string): string {
+  return path.join(resolveStateDir(stateDir), "memory-braid", "models", "entity-extraction");
+}
+
+function slugify(value: string): string {
+  const ascii = value
+    .normalize("NFKD")
+    .replace(/[\u0300-\u036f]/g, "");
+  const slug = ascii
+    .toLowerCase()
+    .replace(/[^a-z0-9]+/g, "-")
+    .replace(/^-+|-+$/g, "");
+  return slug || "unknown";
+}
+
+export function buildCanonicalEntityUri(
+  type: ExtractedEntity["type"],
+  text: string,
+): string {
+  return `entity://${type}/${slugify(text)}`;
+}
+
+function normalizeEntityType(raw: unknown): ExtractedEntity["type"] {
+  const label = typeof raw === "string" ? raw.toUpperCase() : "";
+  if (label.includes("PER")) {
+    return "person";
+  }
+  if (label.includes("ORG")) {
+    return "organization";
+  }
+  if (label.includes("LOC") || label.includes("GPE")) {
+    return "location";
+  }
+  return "misc";
+}
+
+function normalizeEntityText(raw: unknown): string {
+  if (typeof raw !== "string") {
+    return "";
+  }
+  return normalizeWhitespace(raw.replace(/^##/, "").replace(/^▁/, ""));
+}
+
+type EntityExtractionOptions = {
+  stateDir?: string;
+};
+
+export class EntityExtractionManager {
+  private readonly cfg: MemoryBraidConfig["entityExtraction"];
+  private readonly log: MemoryBraidLogger;
+  private stateDir?: string;
+  private pipelinePromise: Promise<NerPipeline | null> | null = null;
+
+  constructor(
+    cfg: MemoryBraidConfig["entityExtraction"],
+    log: MemoryBraidLogger,
+    options?: EntityExtractionOptions,
+  ) {
+    this.cfg = cfg;
+    this.log = log;
+    this.stateDir = options?.stateDir;
+  }
+
+  setStateDir(stateDir?: string): void {
+    const next = stateDir?.trim();
+    if (!next || next === this.stateDir) {
+      return;
+    }
+    this.stateDir = next;
+    this.pipelinePromise = null;
+  }
+
+  getStatus(): {
+    enabled: boolean;
+    provider: MemoryBraidConfig["entityExtraction"]["provider"];
+    model: string;
+    minScore: number;
+    maxEntitiesPerMemory: number;
+    cacheDir: string;
+  } {
+    return {
+      enabled: this.cfg.enabled,
+      provider: this.cfg.provider,
+      model: this.cfg.model,
+      minScore: this.cfg.minScore,
+      maxEntitiesPerMemory: this.cfg.maxEntitiesPerMemory,
+      cacheDir: resolveEntityModelCacheDir(this.stateDir),
+    };
+  }
+
+  async warmup(params?: {
+    runId?: string;
+    reason?: string;
+    forceReload?: boolean;
+    text?: string;
+  }): Promise<{
+    ok: boolean;
+    cacheDir: string;
+    model: string;
+    entities: number;
+    durMs: number;
+    error?: string;
+  }> {
+    const startedAt = Date.now();
+    if (!this.cfg.enabled) {
+      return {
+        ok: false,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        model: this.cfg.model,
+        entities: 0,
+        durMs: Date.now() - startedAt,
+        error: "entity_extraction_disabled",
+      };
+    }
+
+    const pipeline = await this.ensurePipeline(params?.forceReload);
+    if (!pipeline) {
+      return {
+        ok: false,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        model: this.cfg.model,
+        entities: 0,
+        durMs: Date.now() - startedAt,
+        error: "model_load_failed",
+      };
+    }
+
+    try {
+      const entities = await this.extractWithPipeline({
+        pipeline,
+        text: params?.text ?? this.cfg.startup.warmupText,
+      });
+      this.log.info("memory_braid.entity.warmup", {
+        runId: params?.runId,
+        reason: params?.reason ?? "manual",
+        provider: this.cfg.provider,
+        model: this.cfg.model,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        entities: entities.length,
+        entityTypes: summarizeEntityTypes(entities),
+        sampleEntityUris: entities.slice(0, 5).map((entry) => entry.canonicalUri),
+        durMs: Date.now() - startedAt,
+      });
+      return {
+        ok: true,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        model: this.cfg.model,
+        entities: entities.length,
+        durMs: Date.now() - startedAt,
+      };
+    } catch (err) {
+      const message = err instanceof Error ? err.message : String(err);
+      this.log.warn("memory_braid.entity.warmup", {
+        runId: params?.runId,
+        reason: params?.reason ?? "manual",
+        provider: this.cfg.provider,
+        model: this.cfg.model,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        error: message,
+      });
+      return {
+        ok: false,
+        cacheDir: resolveEntityModelCacheDir(this.stateDir),
+        model: this.cfg.model,
+        entities: 0,
+        durMs: Date.now() - startedAt,
+        error: message,
+      };
+    }
+  }
+
+  async extract(params: { text: string; runId?: string }): Promise<ExtractedEntity[]> {
+    if (!this.cfg.enabled) {
+      return [];
+    }
+
+    const text = normalizeWhitespace(params.text);
+    if (!text) {
+      return [];
+    }
+
+    const pipeline = await this.ensurePipeline();
+    if (!pipeline) {
+      return [];
+    }
+
+    try {
+      const entities = await this.extractWithPipeline({ pipeline, text });
+      this.log.debug("memory_braid.entity.extract", {
+        runId: params.runId,
+        provider: this.cfg.provider,
+        model: this.cfg.model,
+        entities: entities.length,
+        entityTypes: summarizeEntityTypes(entities),
+        sampleEntityUris: entities.slice(0, 5).map((entry) => entry.canonicalUri),
+      });
+      return entities;
+    } catch (err) {
+      this.log.warn("memory_braid.entity.extract", {
+        runId: params.runId,
+        provider: this.cfg.provider,
+        model: this.cfg.model,
+        error: err instanceof Error ? err.message : String(err),
+      });
+      return [];
+    }
+  }
+
+  private async ensurePipeline(forceReload = false): Promise<NerPipeline | null> {
+    if (!this.cfg.enabled) {
+      return null;
+    }
+
+    if (forceReload) {
+      this.pipelinePromise = null;
+    }
+
+    if (this.pipelinePromise) {
+      return this.pipelinePromise;
+    }
+
+    this.pipelinePromise = this.loadPipeline();
+    return this.pipelinePromise;
+  }
+
+  private async loadPipeline(): Promise<NerPipeline | null> {
+    const cacheDir = resolveEntityModelCacheDir(this.stateDir);
+    this.log.info("memory_braid.entity.model_load", {
+      provider: this.cfg.provider,
+      model: this.cfg.model,
+      cacheDir,
+    });
+
+    try {
+      const mod = (await import("@xenova/transformers")) as {
+        env?: Record<string, unknown>;
+        pipeline?: (
+          task: string,
+          model: string,
+          options?: Record<string, unknown>,
+        ) => Promise<unknown>;
+      };
+
+      if (!mod.pipeline) {
+        throw new Error("@xenova/transformers pipeline export not found");
+      }
+
+      if (mod.env) {
+        mod.env.cacheDir = cacheDir;
+        mod.env.allowRemoteModels = true;
+        mod.env.allowLocalModels = true;
+        mod.env.useFS = true;
+      }
+
+      const classifier = await mod.pipeline("token-classification", this.cfg.model, {
+        quantized: true,
+      });
+
+      if (typeof classifier !== "function") {
+        throw new Error("token-classification pipeline is not callable");
+      }
+
+      return classifier as NerPipeline;
+    } catch (err) {
+      this.log.error("memory_braid.entity.model_load", {
+        provider: this.cfg.provider,
+        model: this.cfg.model,
+        cacheDir,
+        error: err instanceof Error ? err.message : String(err),
+      });
+      return null;
+    }
+  }
+
+  private async extractWithPipeline(params: {
+    pipeline: NerPipeline;
+    text: string;
+  }): Promise<ExtractedEntity[]> {
+    const raw = await params.pipeline(params.text, {
+      aggregation_strategy: "simple",
+    });
+    const rows = Array.isArray(raw) ? raw : [];
+
+    const deduped = new Map<string, ExtractedEntity>();
+    for (const row of rows) {
+      if (!row || typeof row !== "object") {
+        continue;
+      }
+      const record = row as NerRecord;
+      const entityText = normalizeEntityText(record.word);
+      if (!entityText) {
+        continue;
+      }
+      const score = typeof record.score === "number" ? Math.max(0, Math.min(1, record.score)) : 0;
+      if (score < this.cfg.minScore) {
+        continue;
+      }
+
+      const type = normalizeEntityType(record.entity_group ?? record.entity);
+      const canonicalUri = buildCanonicalEntityUri(type, entityText);
+      const current = deduped.get(canonicalUri);
+      if (!current || score > current.score) {
+        deduped.set(canonicalUri, {
+          text: entityText,
+          type,
+          score,
+          canonicalUri,
+        });
+      }
+    }
+
+    return Array.from(deduped.values())
+      .sort((a, b) => b.score - a.score)
+      .slice(0, this.cfg.maxEntitiesPerMemory);
+  }
+}

package/src/extract.ts

@@ -3,6 +3,8 @@ import type { MemoryBraidConfig } from "./config.js";
 import { MemoryBraidLogger } from "./logger.js";
 import type { ExtractedCandidate } from "./types.js";
 
+type MlProvider = "openai" | "anthropic" | "gemini";
+
 const HEURISTIC_PATTERNS = [
   /remember|remember that|keep in mind|note that/i,
   /i prefer|prefer to|don't like|do not like|hate|love/i,
@@ -145,14 +147,11 @@ function parseJsonObjectArray(raw: string): Array<Record<string, unknown>> {
 }
 
 async function callMlEnrichment(params: {
-  provider: "openai" | "anthropic" | "gemini";
+  provider: MlProvider;
   model: string;
   timeoutMs: number;
   candidates: ExtractedCandidate[];
 }): Promise<Array<Record<string, unknown>>> {
-  const controller = new AbortController();
-  const timer = setTimeout(() => controller.abort(), params.timeoutMs);
-
   const prompt = [
     "Classify the memory candidates.",
     "Return ONLY JSON array.",
@@ -160,6 +159,52 @@ async function callMlEnrichment(params: {
     "Category one of: preference, decision, fact, task, other.",
     JSON.stringify(params.candidates.map((candidate, index) => ({ index, text: candidate.text }))),
   ].join("\n");
+  return callMlJson({
+    provider: params.provider,
+    model: params.model,
+    timeoutMs: params.timeoutMs,
+    prompt,
+  });
+}
+
+async function callMlExtraction(params: {
+  provider: MlProvider;
+  model: string;
+  timeoutMs: number;
+  maxItems: number;
+  messages: Array<{ role: string; text: string }>;
+}): Promise<Array<Record<string, unknown>>> {
+  const recent = params.messages.slice(-30).map((item) => ({
+    role: item.role,
+    text: item.text,
+  }));
+
+  const prompt = [
+    "Extract durable user memories from this conversation.",
+    "Return ONLY JSON array.",
+    "Each item: {text:string, category:string, score:number}.",
+    "Category one of: preference, decision, fact, task, other.",
+    "Keep each text concise and atomic.",
+    `Maximum items: ${params.maxItems}.`,
+    JSON.stringify(recent),
+  ].join("\n");
+
+  return callMlJson({
+    provider: params.provider,
+    model: params.model,
+    timeoutMs: params.timeoutMs,
+    prompt,
+  });
+}
+
+async function callMlJson(params: {
+  provider: MlProvider;
+  model: string;
+  timeoutMs: number;
+  prompt: string;
+}): Promise<Array<Record<string, unknown>>> {
+  const controller = new AbortController();
+  const timer = setTimeout(() => controller.abort(), params.timeoutMs);
 
   try {
     if (params.provider === "openai") {
@@ -183,7 +228,7 @@ async function callMlEnrichment(params: {
             },
             {
               role: "user",
-              content: prompt,
+              content: params.prompt,
             },
           ],
         }),
@@ -212,7 +257,7 @@ async function callMlEnrichment(params: {
           model: params.model,
           max_tokens: 1000,
           temperature: 0,
-          messages: [{ role: "user", content: prompt }],
+          messages: [{ role: "user", content: params.prompt }],
         }),
         signal: controller.signal,
       });
@@ -236,7 +281,7 @@ async function callMlEnrichment(params: {
         },
         body: JSON.stringify({
           generationConfig: { temperature: 0 },
-          contents: [{ role: "user", parts: [{ text: prompt }] }],
+          contents: [{ role: "user", parts: [{ text: params.prompt }] }],
         }),
         signal: controller.signal,
       },
@@ -251,6 +296,19 @@ async function callMlEnrichment(params: {
   }
 }
 
+function normalizeCategory(value: unknown, fallback: ExtractedCandidate["category"] = "other"): ExtractedCandidate["category"] {
+  if (
+    value === "preference" ||
+    value === "decision" ||
+    value === "fact" ||
+    value === "task" ||
+    value === "other"
+  ) {
+    return value;
+  }
+  return fallback;
+}
+
 function applyMlResult(
   candidates: ExtractedCandidate[],
   result: Array<Record<string, unknown>>,
@@ -282,14 +340,7 @@ function applyMlResult(
     if (!keep) {
       continue;
     }
-    const category =
-      ml.category === "preference" ||
-      ml.category === "decision" ||
-      ml.category === "fact" ||
-      ml.category === "task" ||
-      ml.category === "other"
-        ? (ml.category as ExtractedCandidate["category"])
-        : candidate.category;
+    const category = normalizeCategory(ml.category, candidate.category);
     const score = typeof ml.score === "number" ? Math.max(0, Math.min(1, ml.score)) : candidate.score;
     out.push({
       ...candidate,
@@ -301,6 +352,39 @@ function applyMlResult(
   return out;
 }
 
+function applyMlExtractionResult(
+  result: Array<Record<string, unknown>>,
+  maxItems: number,
+): ExtractedCandidate[] {
+  const out: ExtractedCandidate[] = [];
+  const seen = new Set<string>();
+
+  for (const item of result) {
+    const rawText = typeof item.text === "string" ? item.text : "";
+    const text = normalizeWhitespace(rawText);
+    if (!text || text.length < 20 || text.length > 3000) {
+      continue;
+    }
+    const key = sha256(normalizeForHash(text));
+    if (seen.has(key)) {
+      continue;
+    }
+    seen.add(key);
+
+    out.push({
+      text,
+      category: normalizeCategory(item.category),
+      score: typeof item.score === "number" ? Math.max(0, Math.min(1, item.score)) : 0.5,
+      source: "ml",
+    });
+    if (out.length >= maxItems) {
+      break;
+    }
+  }
+
+  return out;
+}
+
 export async function extractCandidates(params: {
   messages: unknown[];
   cfg: MemoryBraidConfig;
@@ -308,43 +392,86 @@ export async function extractCandidates(params: {
   runId?: string;
 }): Promise<ExtractedCandidate[]> {
   const normalized = normalizeMessages(params.messages);
-  const heuristic = pickHeuristicCandidates(normalized, params.cfg.capture.ml.maxItemsPerRun);
+  const heuristic = pickHeuristicCandidates(normalized, params.cfg.capture.maxItemsPerRun);
 
   params.log.debug("memory_braid.capture.extract", {
     runId: params.runId,
+    mode: params.cfg.capture.mode,
+    maxItemsPerRun: params.cfg.capture.maxItemsPerRun,
     totalMessages: normalized.length,
     heuristicCandidates: heuristic.length,
   });
 
-  if (
-    params.cfg.capture.extraction.mode !== "heuristic_plus_ml" ||
-    !params.cfg.capture.ml.provider ||
-    !params.cfg.capture.ml.model
-  ) {
+  if (params.cfg.capture.mode === "local") {
+    params.log.debug("memory_braid.capture.mode", {
+      runId: params.runId,
+      mode: params.cfg.capture.mode,
+      decision: "heuristic_only",
+      candidates: heuristic.length,
+    });
+    return heuristic;
+  }
+
+  if (!params.cfg.capture.ml.provider || !params.cfg.capture.ml.model) {
+    params.log.warn("memory_braid.capture.ml", {
+      runId: params.runId,
+      reason: "missing_provider_or_model",
+      mode: params.cfg.capture.mode,
+      hasProvider: Boolean(params.cfg.capture.ml.provider),
+      hasModel: Boolean(params.cfg.capture.ml.model),
+      fallback: "heuristic",
+      candidates: heuristic.length,
+    });
     return heuristic;
   }
 
   try {
-    const ml = await callMlEnrichment({
+    if (params.cfg.capture.mode === "hybrid") {
+      const ml = await callMlEnrichment({
+        provider: params.cfg.capture.ml.provider,
+        model: params.cfg.capture.ml.model,
+        timeoutMs: params.cfg.capture.ml.timeoutMs,
+        candidates: heuristic,
+      });
+      const enriched = applyMlResult(heuristic, ml);
+      params.log.debug("memory_braid.capture.ml", {
+        runId: params.runId,
+        mode: params.cfg.capture.mode,
+        provider: params.cfg.capture.ml.provider,
+        model: params.cfg.capture.ml.model,
+        requested: heuristic.length,
+        returned: ml.length,
+        enriched: enriched.length,
+        fallbackUsed: ml.length === 0,
+      });
+      return enriched;
+    }
+
+    const mlExtractedRaw = await callMlExtraction({
       provider: params.cfg.capture.ml.provider,
       model: params.cfg.capture.ml.model,
       timeoutMs: params.cfg.capture.ml.timeoutMs,
-      candidates: heuristic,
+      maxItems: params.cfg.capture.maxItemsPerRun,
+      messages: normalized,
     });
-    const enriched = applyMlResult(heuristic, ml);
+    const mlExtracted = applyMlExtractionResult(mlExtractedRaw, params.cfg.capture.maxItemsPerRun);
     params.log.debug("memory_braid.capture.ml", {
       runId: params.runId,
+      mode: params.cfg.capture.mode,
       provider: params.cfg.capture.ml.provider,
       model: params.cfg.capture.ml.model,
-      requested: heuristic.length,
-      returned: ml.length,
-      enriched: enriched.length,
+      returned: mlExtractedRaw.length,
+      extracted: mlExtracted.length,
+      fallbackUsed: mlExtracted.length === 0,
     });
-    return enriched;
+    return mlExtracted.length > 0 ? mlExtracted : heuristic;
   } catch (err) {
     params.log.warn("memory_braid.capture.ml", {
       runId: params.runId,
+      mode: params.cfg.capture.mode,
       error: err instanceof Error ? err.message : String(err),
+      fallback: "heuristic",
+      candidates: heuristic.length,
     });
     return heuristic;
   }