mem0-auto-capture-hook/hook/handler.ts

import fs from "fs";
import os from "os";
import path from "path";

// ---------------------------------------------------------------------------
// Config — openclaw does NOT inject cfg into hook events.
// Read from env vars; fall back to ~/.openclaw/mem0.json for convenience.
// ---------------------------------------------------------------------------
function loadPluginCfg() {
  const cfgPath = path.join(os.homedir(), ".openclaw", "mem0.json");
  let fileCfg: Record<string, unknown> = {};
  try {
    if (fs.existsSync(cfgPath)) {
      fileCfg = JSON.parse(fs.readFileSync(cfgPath, "utf8"));
    }
  } catch {
    // ignore malformed file
  }
  return {
    baseUrl:
      process.env.MEM0_BASE_URL ||
      (fileCfg.baseUrl as string) ||
      "http://192.168.0.200:8420",
    userId:
      process.env.MEM0_USER_ID || (fileCfg.userId as string) || undefined,
    recallLimit: Number(
      process.env.MEM0_RECALL_LIMIT || fileCfg.recallLimit || 5
    ),
    captureTrigger: (
      process.env.MEM0_CAPTURE_TRIGGER ||
      (fileCfg.captureTrigger as string) ||
      "always"
    ) as "always" | "phrase" | "explicit",
    triggerPhrase:
      process.env.MEM0_TRIGGER_PHRASE ||
      (fileCfg.triggerPhrase as string) ||
      "please remember",
    autoCapture:
      (process.env.MEM0_AUTO_CAPTURE ||
        String(fileCfg.autoCapture ?? "true")) !== "false",
    autoRecall:
      (process.env.MEM0_AUTO_RECALL ||
        String(fileCfg.autoRecall ?? "true")) !== "false",
    recentKeep: Number(
      process.env.MEM0_RECENT_KEEP || fileCfg.recentKeep || 5
    ),
    // Knowledge-base settings
    knowledgeUserId:
      process.env.MEM0_KNOWLEDGE_USER_ID ||
      (fileCfg.knowledgeUserId as string) ||
      "knowledge_base",
    rerankThreshold: Number(
      process.env.MEM0_RERANK_THRESHOLD || fileCfg.rerankThreshold || 0.002
    ),
    knowledgeLimit: Number(
      process.env.MEM0_KNOWLEDGE_LIMIT || fileCfg.knowledgeLimit || 5
    ),
  };
}

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------
type HookEvent = {
  type: string;
  action: string;
  sessionKey: string;
  timestamp: Date;
  messages: string[] | any[];
  context: any;
};

interface MemoryResult {
  id: string;
  memory: string;
  score: number;
  rerank_score: number;
  metadata?: Record<string, any>;
}

interface SearchResponse {
  results: MemoryResult[];
}

// ---------------------------------------------------------------------------
// In-memory state — capped to avoid unbounded growth
// ---------------------------------------------------------------------------
const MAX_SESSIONS = 500;
const MAX_TRANSCRIPTS = 1000;

class LRUMap<K, V> extends Map<K, V> {
  private readonly maxSize: number;
  constructor(maxSize: number) {
    super();
    this.maxSize = maxSize;
  }
  set(key: K, value: V): this {
    if (this.size >= this.maxSize && !this.has(key)) {
      this.delete(this.keys().next().value!);
    }
    return super.set(key, value);
  }
}

const recentBySession = new LRUMap<string, string[]>(MAX_SESSIONS);
const transcriptByMessageId = new LRUMap<string, string>(MAX_TRANSCRIPTS);

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------
const LOCAL_STT_URL = "http://192.168.0.200:5005/transcribe";

function getAgentIdFromSessionKey(sessionKey?: string): string | undefined {
  if (!sessionKey) return undefined;
  const parts = sessionKey.split(":");
  if (parts.length >= 2 && parts[0] === "agent") return parts[1];
  console.warn("[mem0] unexpected sessionKey format:", sessionKey);
  return undefined;
}

function pushRecent(
  sessionKey: string,
  text: string,
  keep: number
): string[] {
  const list = recentBySession.get(sessionKey) || [];
  list.push(text);
  while (list.length > keep) list.shift();
  recentBySession.set(sessionKey, list);
  return list;
}

function isMediaPlaceholder(text: string): boolean {
  return /<media:[^>]+>/i.test(text);
}

function extractMessageText(context: any): string {
  const transcript =
    typeof context?.transcript === "string" ? context.transcript.trim() : "";
  if (transcript) return transcript;
  const candidates = [context?.bodyForAgent, context?.content, context?.body];
  for (const candidate of candidates) {
    if (typeof candidate === "string") {
      const trimmed = candidate.trim();
      if (trimmed && !isMediaPlaceholder(trimmed)) return trimmed;
    }
  }
  return "";
}

function getAudioPath(context: any): string | undefined {
  if (typeof context?.mediaPath === "string") return context.mediaPath;
  if (
    Array.isArray(context?.mediaPaths) &&
    typeof context.mediaPaths[0] === "string"
  ) {
    return context.mediaPaths[0];
  }
  if (typeof context?.media?.path === "string") return context.media.path;
  return undefined;
}

async function transcribeAudio(localPath: string): Promise<string> {
  const buffer = fs.readFileSync(localPath);
  const blob = new Blob([buffer]);
  const form = new FormData();
  form.append("file", blob, "audio.ogg");
  const res = await fetch(LOCAL_STT_URL, { method: "POST", body: form });
  if (!res.ok) throw new Error(`STT failed: ${res.status}`);
  const data = await res.json();
  const text =
    typeof data?.text === "string" ? data.text.trim() : "";
  if (!text) throw new Error("STT returned empty transcript");
  return text;
}

// ---------------------------------------------------------------------------
// mem0 search helpers — each returns null on failure so callers can degrade
// gracefully when one endpoint is down.
// ---------------------------------------------------------------------------
async function mem0SearchMemories(
  baseUrl: string,
  userId: string,
  query: string,
  limit: number
): Promise<MemoryResult[] | null> {
  try {
    const res = await fetch(`${baseUrl}/memories/search`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ query, userId }),
    });
    if (!res.ok) {
      console.error(`[mem0-recall] /memories/search returned ${res.status}`);
      return null;
    }
    const data: SearchResponse = await res.json();
    return Array.isArray(data?.results) ? data.results.slice(0, limit) : [];
  } catch (err) {
    console.error("[mem0-recall] /memories/search failed:", err);
    return null;
  }
}

async function mem0SearchKnowledge(
  baseUrl: string,
  knowledgeUserId: string,
  query: string,
  limit: number
): Promise<MemoryResult[] | null> {
  try {
    const res = await fetch(`${baseUrl}/knowledge/search`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({ query, userId: knowledgeUserId }),
    });
    if (!res.ok) {
      console.error(`[mem0-recall] /knowledge/search returned ${res.status}`);
      return null;
    }
    const data: SearchResponse = await res.json();
    return Array.isArray(data?.results) ? data.results.slice(0, limit) : [];
  } catch (err) {
    console.error("[mem0-recall] /knowledge/search failed:", err);
    return null;
  }
}

// ---------------------------------------------------------------------------
// Result filtering & formatting
// ---------------------------------------------------------------------------
function filterByRerank(
  results: MemoryResult[],
  threshold: number
): MemoryResult[] {
  return results.filter(
    (r) => typeof r.rerank_score === "number" && r.rerank_score >= threshold
  );
}

function formatKnowledgeCitation(meta: Record<string, any> = {}): string {
  const parts: string[] = [];
  if (meta.source_file) parts.push(meta.source_file);
  if (meta.chapter != null) parts.push(`ch.${meta.chapter}`);
  if (meta.page_start != null && meta.page_end != null) {
    parts.push(`pp.${meta.page_start}-${meta.page_end}`);
  } else if (meta.page_start != null) {
    parts.push(`p.${meta.page_start}`);
  }
  return parts.length > 0 ? `(from: ${parts.join(", ")})` : "";
}

function buildInjectionBlock(
  personalResults: MemoryResult[] | null,
  knowledgeResults: MemoryResult[] | null,
  threshold: number
): string {
  const sections: string[] = [];

  // Personal memories
  const personal = personalResults
    ? filterByRerank(personalResults, threshold)
    : null;

  if (personal === null) {
    // endpoint was down — omit the section entirely, already logged
  } else if (personal.length > 0) {
    const lines = personal
      .map((r) => `- ${r.memory}`)
      .join("\n");
    sections.push(`[MEMORY - Personal]\n${lines}`);
  }

  // Knowledge base — sort by rerank_score descending
  const knowledge = knowledgeResults
    ? filterByRerank(knowledgeResults, threshold).sort(
        (a, b) => b.rerank_score - a.rerank_score
      )
    : null;

  if (knowledge === null) {
    // endpoint was down — omit the section entirely, already logged
  } else if (knowledge.length > 0) {
    const lines = knowledge
      .map((r) => {
        const citation = formatKnowledgeCitation(r.metadata || {});
        return citation ? `- ${citation} ${r.memory}` : `- ${r.memory}`;
      })
      .join("\n");
    sections.push(`[MEMORY - Knowledge Base]\n${lines}`);
  }

  return sections.join("\n\n");
}

// Build a quick hash to deduplicate captures
function simpleHash(text: string): string {
  let h = 0;
  for (let i = 0; i < text.length; i++) {
    h = (Math.imul(31, h) + text.charCodeAt(i)) | 0;
  }
  return h.toString(36);
}

const recentlyCaptured = new LRUMap<string, number>(200);
const CAPTURE_DEDUP_MS = 60_000;

// ---------------------------------------------------------------------------
// Main handler
// ---------------------------------------------------------------------------
export default async function handler(event: HookEvent) {
  console.log(
    "[mem0-FIRE]",
    JSON.stringify(
      {
        type: event.type,
        action: event.action,
        sessionKey: event.sessionKey,
        contextKeys: Object.keys(event.context || {}),
        content: event.context?.content?.slice(0, 80),
        messagesLen: event.messages?.length,
      },
      null,
      2
    )
  );

  if (event.type !== "message") {
    console.log("[mem0-FIRE] bailed: type is not message, got:", event.type);
    return;
  }

  const pluginCfg = loadPluginCfg();
  const userId =
    getAgentIdFromSessionKey(event.sessionKey) ||
    pluginCfg.userId ||
    "default";

  // ── Audio transcription (runs once per messageId) ────────────────────────
  const messageId =
    event.context?.messageId || event.context?.metadata?.messageId;
  const audioPath = getAudioPath(event.context);
  const hasTranscript =
    typeof event.context?.transcript === "string" &&
    event.context.transcript.trim().length > 0;

  if (!hasTranscript && audioPath && fs.existsSync(audioPath)) {
    if (messageId && transcriptByMessageId.has(messageId)) {
      event.context.transcript = transcriptByMessageId.get(messageId);
    } else {
      try {
        const text = await transcribeAudio(audioPath);
        event.context.transcript = text;
        if (messageId) transcriptByMessageId.set(messageId, text);
      } catch (err) {
        console.error("[mem0-stt] failed:", err);
        if (
          isMediaPlaceholder(
            event.context?.bodyForAgent ||
              event.context?.content ||
              ""
          )
        ) {
          return;
        }
      }
    }
  }

  // Patch bodyForAgent with transcript for audio messages
  const transcriptNow =
    typeof event.context?.transcript === "string"
      ? event.context.transcript.trim()
      : "";
  if (transcriptNow && isMediaPlaceholder(event.context?.bodyForAgent || "")) {
    event.context.bodyForAgent = transcriptNow;
  }

  // ── Extract text (shared by both branches) ───────────────────────────────
  function resolveText(): string {
    let text = extractMessageText(event.context);
    if (!text && Array.isArray(event.messages) && event.messages.length > 0) {
      const m = event.messages[0];
      if (typeof m === "string") text = m.trim();
      else if (typeof m?.text === "string") text = m.text.trim();
      else if (typeof m?.body === "string") text = m.body.trim();
    }
    return text;
  }

  // ── Auto-recall: queries both endpoints in parallel ───────────────────────
  const runAutoRecall = async (text: string) => {
    if (!pluginCfg.autoRecall) return;

    const {
      baseUrl,
      knowledgeUserId,
      recallLimit,
      knowledgeLimit,
      rerankThreshold,
    } = pluginCfg;

    console.log("[mem0-auto-recall] query:", text.slice(0, 120));

    // Fire both searches in parallel — neither blocks the other
    const [personalResults, knowledgeResults] = await Promise.all([
      mem0SearchMemories(baseUrl, userId, text, recallLimit),
      mem0SearchKnowledge(baseUrl, knowledgeUserId, text, knowledgeLimit),
    ]);

    console.log("[mem0-auto-recall]", {
      userId,
      personalCount: personalResults?.length ?? "error",
      knowledgeCount: knowledgeResults?.length ?? "error",
      threshold: rerankThreshold,
    });

    const injectionBlock = buildInjectionBlock(
      personalResults,
      knowledgeResults,
      rerankThreshold
    );

    if (!injectionBlock) return; // nothing passed the threshold from either endpoint

    event.context.bodyForAgent = `${text}\n\n${injectionBlock}`;
    console.log("[mem0-injected-prompt]\n" + event.context.bodyForAgent);
  };

  // ── received: capture ────────────────────────────────────────────────────
  if (event.action === "received") {
    if (!pluginCfg.autoCapture) return;
    const text = resolveText();
    if (!text) return;
    if (isMediaPlaceholder(text) && !transcriptNow) return;

    const recent = pushRecent(
      event.sessionKey || "global",
      text,
      pluginCfg.recentKeep
    );

    const { captureTrigger, triggerPhrase } = pluginCfg;
    let shouldCapture = false;
    if (captureTrigger === "always") {
      shouldCapture = true;
    } else if (captureTrigger === "phrase") {
      shouldCapture = new RegExp(triggerPhrase, "i").test(text);
    } else {
      shouldCapture = /please\s+remember/i.test(text);
    }
    if (!shouldCapture) return;

    const captureText = recent.join("\n");
    const hash = simpleHash(captureText);
    const lastCapture = recentlyCaptured.get(hash);
    if (lastCapture && Date.now() - lastCapture < CAPTURE_DEDUP_MS) {
      console.log("[mem0-auto-capture] skipped duplicate");
      return;
    }
    recentlyCaptured.set(hash, Date.now());

    try {
      console.log("[mem0-auto-capture]", {
        userId,
        captureTrigger,
        text: captureText.slice(0, 160),
      });
      await fetch(`${pluginCfg.baseUrl}/memories`, {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ text: captureText, userId }),
      });
    } catch (err) {
      console.error("[mem0-auto-capture] write failed:", err);
    }
    return;
  }

  // ── preprocessed: recall ─────────────────────────────────────────────────
  if (event.action === "preprocessed") {
    const text = resolveText();
    if (!text) return;
    if (!transcriptNow && isMediaPlaceholder(text)) return;
    await runAutoRecall(text);
    return;
  }

  // ── transcribed: recall ──────────────────────────────────────────────────
  if (event.action === "transcribed") {
    const text = extractMessageText(event.context);
    if (!text) return;
    await runAutoRecall(text);
  }
}