宜博 / lemonboss

/**
 * MemoryService — Agent 自动记忆读写
 *
 * Phase 1 (MVP):
 *   - 对话结束 → 自动生成 session 摘要 → short-term/
 *   - 对话结束 → 关键事实追加 → memory.md
 *   - 新对话开始 → 加载最近 N 个 session 摘要注入 prompt
 *
 * Phase 2:
 *   - 海马体关键词匹配 → 检索相关小脑方案
 *   - 小脑方案固化 → cerebellum/ 存储成功方案
 *   - 海马体索引自动更新
 */

const fs = require('fs');
const path = require('path');

const DATA_DIR = path.resolve(__dirname, '../../data/agents');
const MAX_SHORT_TERM = 10;            // 最多保留最近 N 个 session 摘要
const MAX_MEMORY_MD_SIZE = 3000;      // memory.md 最大字符数
const INJECT_SESSION_COUNT = 5;       // prompt 注入最近 N 个 session
const HIPPOCAMPUS_MIN_SCORE = 0.3;    // 海马体匹配最低分
const MAX_CEREBELLUM_INJECT = 2;      // 最多注入 N 个小脑方案

// ════════════════════════════════════════
//  1. 对话结束 → 自动保存 Session 摘要
// ════════════════════════════════════════

/**
 * 从 messages.log 和 meta 中提取信息，生成 session 摘要，写入 short-term/
 *
 * @param {string} agentId
 * @param {string} convId
 * @param {object} meta - { status, stats, createdAt, updatedAt, title }
 * @param {string} lastOutput - Agent 最后的输出文本
 */
function saveSessionSummary(agentId, convId, meta, lastOutput) {
  try {
    const agentDir = path.join(DATA_DIR, agentId);
    const shortTermDir = path.join(agentDir, 'memory', 'short-term');

    if (!fs.existsSync(shortTermDir)) {
      fs.mkdirSync(shortTermDir, { recursive: true });
    }

    // 提取任务描述 (从 title 或 messages.log 的第一条)
    const task = extractTask(meta.title || '', agentId, convId);

    // 提取结果摘要 (Agent 最后输出的前 500 字)
    const resultSummary = extractResultSummary(lastOutput);

    // 提取提到的文件路径
    const files = extractMentionedFiles(lastOutput);

    // 计算持续时间
    const createdAt = new Date(meta.createdAt);
    const updatedAt = new Date(meta.updatedAt || Date.now());
    const durationSec = Math.round((updatedAt.getTime() - createdAt.getTime()) / 1000);
    const durationStr = durationSec > 60
      ? `${Math.round(durationSec / 60)}分钟`
      : `${durationSec}秒`;

    // 构造摘要
    const timestamp = new Date().toISOString().replace('T', ' ').substring(0, 19);
    const summary = [
      `# Session 摘要`,
      ``,
      `- **时间**: ${timestamp}`,
      `- **对话ID**: ${convId}`,
      `- **状态**: ${meta.status || 'completed'}`,
      `- **工具调用**: ${meta.stats?.toolCallCount || 0} 次`,
      `- **耗时**: ${durationStr}`,
      ``,
      `## 任务`,
      task || '(无明确任务描述)',
      ``,
      `## 结果`,
      resultSummary || '(无输出)',
    ];

    if (files.length > 0) {
      summary.push('', '## 涉及文件', ...files.map(f => `- ${f}`));
    }

    summary.push('');

    // 写入文件
    const filename = `session-${Date.now()}.md`;
    const filepath = path.join(shortTermDir, filename);
    fs.writeFileSync(filepath, summary.join('\n'), 'utf-8');

    // 清理过旧的 session 文件（保留最近 MAX_SHORT_TERM 个）
    cleanupShortTerm(shortTermDir);

    // 追加关键事实到 memory.md
    appendToMemoryMd(agentId, task, resultSummary, meta);

    console.log(`[Memory] ✅ Session 摘要已保存: ${agentId}/${filename}`);
    return filepath;
  } catch (err) {
    console.error(`[Memory] ⚠️ Session 摘要保存失败:`, err.message);
    return null;
  }
}

// ════════════════════════════════════════
//  2. 新对话开始 → 加载近期记忆
// ════════════════════════════════════════

/**
 * 加载最近 N 个 session 摘要，用于注入 system prompt
 *
 * @param {string} agentId
 * @returns {string} 格式化的近期记忆文本
 */
function loadRecentMemory(agentId) {
  try {
    const shortTermDir = path.join(DATA_DIR, agentId, 'memory', 'short-term');
    if (!fs.existsSync(shortTermDir)) return '';

    // 按文件名倒序（时间戳降序）取最近 N 个
    const files = fs.readdirSync(shortTermDir)
      .filter(f => f.startsWith('session-') && f.endsWith('.md'))
      .sort()
      .reverse()
      .slice(0, INJECT_SESSION_COUNT);

    if (files.length === 0) return '';

    const sessions = files.map(f => {
      const content = fs.readFileSync(path.join(shortTermDir, f), 'utf-8');
      return compactSession(content);
    });

    return [
      `# 近期记忆（最近 ${sessions.length} 次对话）`,
      '',
      ...sessions,
    ].join('\n');
  } catch (err) {
    console.error(`[Memory] ⚠️ 加载近期记忆失败:`, err.message);
    return '';
  }
}

// ════════════════════════════════════════
//  3. memory.md 自动追加
// ════════════════════════════════════════

/**
 * 追加关键事实到 memory.md
 * 只在有实质操作（toolCallCount > 0）时追加
 */
function appendToMemoryMd(agentId, task, result, meta) {
  if (!meta.stats?.toolCallCount || meta.stats.toolCallCount === 0) return;

  const memoryPath = path.join(DATA_DIR, agentId, 'memory', 'memory.md');
  let content = '';
  try {
    content = fs.readFileSync(memoryPath, 'utf-8');
  } catch { /* 文件不存在 */ }

  const timestamp = new Date().toISOString().substring(0, 10);
  const taskOneLine = (task || '').split('\n')[0].substring(0, 80);
  const resultOneLine = (result || '').split('\n')[0].substring(0, 100);

  // 构造新条目
  const entry = `- [${timestamp}] ${taskOneLine} → ${resultOneLine}`;

  // 追加到"关键事实"段
  if (content.includes('## 关键事实')) {
    content = content.replace(
      /## 关键事实\n([\s\S]*?)(?=\n## |$)/,
      (match, existing) => {
        const trimmed = existing.trim();
        if (trimmed === '（尚无记录）') {
          return `## 关键事实\n${entry}\n`;
        }
        return `## 关键事实\n${trimmed}\n${entry}\n`;
      }
    );
  } else {
    // 没有"关键事实"段，追加到末尾
    content += `\n\n## 关键事实\n${entry}\n`;
  }

  // 大小控制：超过限制时截断最早的条目
  if (content.length > MAX_MEMORY_MD_SIZE) {
    content = truncateMemory(content);
  }

  fs.writeFileSync(memoryPath, content, 'utf-8');
}

// ════════════════════════════════════════
//  4. Phase 2: 海马体检索 (Hippocampus)
// ════════════════════════════════════════

/**
 * 根据任务文本检索海马体，返回匹配的小脑方案
 *
 * @param {string} agentId
 * @param {string} taskText - 用户的任务文本
 * @returns {string} 格式化的小脑方案文本（用于注入 prompt）
 */
function retrieveRelevantPlans(agentId, taskText) {
  try {
    const agentDir = path.join(DATA_DIR, agentId);
    const index = loadHippocampusIndex(agentDir);
    if (index.length === 0) return '';

    // 从任务文本提取关键词
    const taskKeywords = extractKeywords(taskText);
    if (taskKeywords.length === 0) return '';

    // 计算每个索引条目的匹配分数
    const scored = index.map(entry => {
      const entryKeywords = (entry.keywords || []).map(k => k.toLowerCase());
      const hits = taskKeywords.filter(kw =>
        entryKeywords.some(ek => ek.includes(kw) || kw.includes(ek))
      );
      const score = entryKeywords.length > 0
        ? hits.length / entryKeywords.length  // 命中了 entry 的多少比例关键词
        : 0;
      return { ...entry, score, hits };
    });

    // 过滤 + 排序
    const matches = scored
      .filter(e => e.score >= HIPPOCAMPUS_MIN_SCORE)
      .sort((a, b) => {
        // 先按分数，再按置信度，最后按使用次数
        if (b.score !== a.score) return b.score - a.score;
        if ((b.confidence || 0) !== (a.confidence || 0)) return (b.confidence || 0) - (a.confidence || 0);
        return (b.useCount || 0) - (a.useCount || 0);
      })
      .slice(0, MAX_CEREBELLUM_INJECT);

    if (matches.length === 0) return '';

    // 加载对应的小脑方案文件
    const cerebellumDir = path.join(agentDir, 'memory', 'cerebellum');
    const plans = [];

    for (const match of matches) {
      const planPath = path.join(cerebellumDir, match.planFile);
      if (fs.existsSync(planPath)) {
        const planContent = fs.readFileSync(planPath, 'utf-8').trim();
        plans.push(`### 经验方案: ${match.scenario} (置信度 ${Math.round((match.confidence || 0) * 100)}%, 使用 ${match.useCount || 0} 次)\n\n${planContent}`);

        // 更新原始索引中的使用次数 + 最后使用时间
        const origEntry = index.find(e => e.planFile === match.planFile);
        if (origEntry) {
          origEntry.useCount = (origEntry.useCount || 0) + 1;
          origEntry.lastUsed = new Date().toISOString().substring(0, 10);
        }
      }
    }

    if (plans.length > 0) {
      // 保存更新后的使用统计
      saveHippocampusIndex(agentDir, index);
      console.log(`[Memory] 🐴 海马体命中 ${plans.length} 个方案 (agent: ${agentId})`);
    }

    if (plans.length === 0) return '';

    return [
      `# 经验参考（海马体检索）`,
      '',
      '> 以下是你过去处理类似任务的成功经验。可以参考，但不必完全照搬。',
      '',
      ...plans,
    ].join('\n');
  } catch (err) {
    console.error(`[Memory] ⚠️ 海马体检索失败:`, err.message);
    return '';
  }
}

// ════════════════════════════════════════
//  5. Phase 2: 小脑固化 (Cerebellum)
// ════════════════════════════════════════

/**
 * 保存一个成功方案到小脑，并更新海马体索引
 *
 * @param {string} agentId
 * @param {object} plan - { scenario, keywords, steps, notes }
 * @returns {string|null} 保存成功返回文件路径
 */
function saveCerebellumPlan(agentId, plan) {
  try {
    const agentDir = path.join(DATA_DIR, agentId);
    const cerebellumDir = path.join(agentDir, 'memory', 'cerebellum');

    if (!fs.existsSync(cerebellumDir)) {
      fs.mkdirSync(cerebellumDir, { recursive: true });
    }

    // 生成文件名: 基于 scenario 的 slug
    const slug = (plan.scenario || 'plan')
      .toLowerCase()
      .replace(/[^a-z0-9\u4e00-\u9fff]+/g, '-')
      .replace(/^-|-$/g, '')
      .substring(0, 40);
    const filename = `${slug}.md`;
    const filepath = path.join(cerebellumDir, filename);

    // 构造方案内容
    const content = [
      `# 方案: ${plan.scenario}`,
      '',
      `> 创建时间: ${new Date().toISOString().substring(0, 10)}`,
      '',
    ];

    if (plan.conditions) {
      content.push('## 适用条件', ...plan.conditions.map(c => `- ${c}`), '');
    }

    if (plan.steps) {
      content.push('## 执行步骤', ...plan.steps.map((s, i) => `${i + 1}. ${s}`), '');
    }

    if (plan.notes) {
      content.push('## 注意事项', ...plan.notes.map(n => `- ${n}`), '');
    }

    if (plan.rawContent) {
      content.push('## 详细内容', plan.rawContent, '');
    }

    fs.writeFileSync(filepath, content.join('\n'), 'utf-8');

    // 更新海马体索引
    const index = loadHippocampusIndex(agentDir);

    // 检查是否已有同名条目
    const existing = index.findIndex(e => e.planFile === filename);
    const entry = {
      scenario: plan.scenario,
      keywords: plan.keywords || [],
      planFile: filename,
      confidence: plan.confidence || 0.7,
      useCount: 0,
      lastUsed: null,
      createdAt: new Date().toISOString().substring(0, 10),
    };

    if (existing >= 0) {
      // 更新已有条目（保留使用次数）
      entry.useCount = index[existing].useCount || 0;
      entry.lastUsed = index[existing].lastUsed;
      index[existing] = entry;
    } else {
      index.push(entry);
    }

    saveHippocampusIndex(agentDir, index);

    // 同步更新 hippocampus.md (人类可读版)
    syncHippocampusMd(agentDir, index);

    console.log(`[Memory] 🧩 小脑方案已固化: ${agentId}/cerebellum/${filename}`);
    return filepath;
  } catch (err) {
    console.error(`[Memory] ⚠️ 小脑方案保存失败:`, err.message);
    return null;
  }
}

// ════════════════════════════════════════
//  海马体索引 JSON 读写
// ════════════════════════════════════════

/**
 * 加载海马体索引 (JSON 格式，比 md 表格更可靠)
 * 文件: memory/hippocampus-index.json
 */
function loadHippocampusIndex(agentDir) {
  const indexPath = path.join(agentDir, 'memory', 'hippocampus-index.json');
  try {
    return JSON.parse(fs.readFileSync(indexPath, 'utf-8'));
  } catch {
    return [];
  }
}

/** 保存海马体索引 */
function saveHippocampusIndex(agentDir, index) {
  const indexPath = path.join(agentDir, 'memory', 'hippocampus-index.json');
  const dir = path.dirname(indexPath);
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
  fs.writeFileSync(indexPath, JSON.stringify(index, null, 2), 'utf-8');
}

/**
 * 同步 hippocampus.md — 人类可读的 markdown 表格
 * (JSON 是程序真正使用的，md 只是给人看)
 */
function syncHippocampusMd(agentDir, index) {
  const mdPath = path.join(agentDir, 'memory', 'hippocampus.md');
  const rows = index.map((e, i) =>
    `| ${i + 1} | ${e.scenario} | ${(e.keywords || []).join(', ')} | ${e.planFile} | ${e.confidence} | ${e.useCount || 0} | ${e.lastUsed || '-'} |`
  );

  const content = [
    '# 海马体 — 经验记忆索引',
    '',
    '> 记录过往处理过的情景模式。遇到相似任务时，优先调用小脑中的固化方案。',
    '> ⚠️ 此文件由系统自动生成，勿手动编辑。真实数据在 hippocampus-index.json。',
    '',
    '| # | 情景模式 | 匹配关键词 | 方案文件 | 置信度 | 使用次数 | 上次使用 |',
    '|---|---------|-----------|---------|--------|---------|----------|',
    ...(rows.length > 0 ? rows : ['| | （暂无经验记录） | | | | | |']),
    '',
  ].join('\n');

  fs.writeFileSync(mdPath, content, 'utf-8');
}

// ════════════════════════════════════════
//  关键词提取
// ════════════════════════════════════════

/** 中文停用词表（常见无意义词） */
const STOP_WORDS = new Set([
  '的', '了', '在', '是', '我', '有', '和', '就', '不', '人', '都', '一', '一个',
  '上', '也', '很', '到', '说', '要', '去', '你', '会', '着', '没有', '看', '好',
  '自己', '这', '他', '她', '它', '那', '些', '所以', '因为', '但是', '而且', '或者',
  '如果', '虽然', '所有', '可以', '应该', '需要', '能够', '已经', '正在', '请',
  '把', '被', '让', '给', '从', '对', '将', '于', '之', '等', '以', '及', '与',
  'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been', 'being',
  'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could',
  'should', 'may', 'might', 'can', 'shall', 'to', 'of', 'in', 'for',
  'on', 'with', 'at', 'by', 'from', 'as', 'into', 'about', 'and', 'or',
  'but', 'not', 'this', 'that', 'it', 'its', 'my', 'your', 'his', 'her',
]);

/**
 * 从文本中提取关键词
 * 策略: 中文按 2-4 字切片 + 英文按空格分词 + 去停用词
 */
function extractKeywords(text) {
  if (!text || text.length < 2) return [];

  const keywords = new Set();

  // 1. 英文词 (≥2字母)
  const englishWords = text.match(/[a-zA-Z]{2,}/g) || [];
  englishWords.forEach(w => {
    const lower = w.toLowerCase();
    if (!STOP_WORDS.has(lower) && lower.length >= 3) {
      keywords.add(lower);
    }
  });

  // 2. 中文：提取连续中文字符序列，再分割成 2-4 字的 ngram
  const chineseSegments = text.match(/[\u4e00-\u9fff]{2,}/g) || [];
  for (const seg of chineseSegments) {
    // 完整段 (如果 ≤6 字)
    if (seg.length <= 6 && !STOP_WORDS.has(seg)) {
      keywords.add(seg);
    }
    // 2-gram
    for (let i = 0; i <= seg.length - 2; i++) {
      const bi = seg.substring(i, i + 2);
      if (!STOP_WORDS.has(bi)) keywords.add(bi);
    }
    // 3-gram (如果段够长)
    if (seg.length >= 3) {
      for (let i = 0; i <= seg.length - 3; i++) {
        const tri = seg.substring(i, i + 3);
        if (!STOP_WORDS.has(tri)) keywords.add(tri);
      }
    }
  }

  return [...keywords].slice(0, 30);  // 最多 30 个关键词
}

// ════════════════════════════════════════
//  Phase 1 内部工具函数
// ════════════════════════════════════════

/** 从 title 提取任务描述 */
function extractTask(title, agentId, convId) {
  // 去掉 [TASK from xxx] 前缀
  let task = title.replace(/\[TASK from \w+\]\s*/g, '').trim();
  // 去掉任务编号行
  task = task.replace(/任务编号:\s*T-\w+\s*/g, '').trim();
  task = task.replace(/优先级:\s*P\d+\s*/g, '').trim();

  // 如果 title 太短，尝试从 messages.log 读取
  if (task.length < 10) {
    try {
      const logPath = path.join(DATA_DIR, agentId, 'conversations', convId, 'messages.log');
      const log = fs.readFileSync(logPath, 'utf-8');
      const userMsg = log.match(/\[USER\]\n([\s\S]*?)(?:\n\[AGENT\]|$)/);
      if (userMsg) task = userMsg[1].trim().substring(0, 300);
    } catch { /* ignore */ }
  }

  return task.substring(0, 300);
}

/** 从 Agent 输出提取结果摘要 */
function extractResultSummary(output) {
  if (!output) return '';
  // 优先取 [RESULT: DONE]...[/RESULT] 块
  const resultBlock = output.match(/\[RESULT:\s*DONE\]([\s\S]*?)\[\/RESULT\]/);
  if (resultBlock) return resultBlock[1].trim().substring(0, 500);

  // 否则取最后 500 字
  const trimmed = output.trim();
  if (trimmed.length > 500) {
    return '...' + trimmed.substring(trimmed.length - 500);
  }
  return trimmed;
}

/** 提取输出中提到的文件路径 */
function extractMentionedFiles(output) {
  if (!output) return [];
  // 匹配常见文件路径模式
  const patterns = [
    /(?:写入|创建|保存|输出)[到至：:]\s*[`"']?([^\s`"']+\.\w+)/g,
    /(?:data\/\w+\/[^\s]+)/g,
    /(?:workspace\/[^\s]+)/g,
  ];

  const files = new Set();
  for (const pat of patterns) {
    let m;
    while ((m = pat.exec(output)) !== null) {
      const f = m[1] || m[0];
      if (f.length > 5 && f.length < 200) files.add(f);
    }
  }
  return [...files].slice(0, 10);  // 最多 10 个
}

/** 压缩 session 内容为紧凑格式（用于 prompt 注入） */
function compactSession(content) {
  // 提取关键字段
  const time = content.match(/时间\**[：:]\s*(.+)/)?.[1] || '';
  const status = content.match(/状态\**[：:]\s*(.+)/)?.[1] || '';
  const tools = content.match(/工具调用\**[：:]\s*(.+)/)?.[1] || '';
  const duration = content.match(/耗时\**[：:]\s*(.+)/)?.[1] || '';

  // 提取任务和结果段
  const taskMatch = content.match(/## 任务\n([\s\S]*?)(?=\n## |$)/);
  const resultMatch = content.match(/## 结果\n([\s\S]*?)(?=\n## |$)/);
  const filesMatch = content.match(/## 涉及文件\n([\s\S]*?)(?=\n## |$)/);

  const task = taskMatch?.[1]?.trim().substring(0, 100) || '';
  const result = resultMatch?.[1]?.trim().substring(0, 200) || '';
  const filesList = filesMatch?.[1]?.trim() || '';

  const parts = [`### ${time} [${status}] (${tools}, ${duration})`];
  if (task) parts.push(`- 任务: ${task}`);
  if (result) parts.push(`- 结果: ${result}`);
  if (filesList) parts.push(`- 文件: ${filesList.replace(/\n/g, ', ').replace(/- /g, '')}`);
  parts.push('');

  return parts.join('\n');
}

/** 清理过旧的 short-term 文件 */
function cleanupShortTerm(dir) {
  try {
    const files = fs.readdirSync(dir)
      .filter(f => f.startsWith('session-') && f.endsWith('.md'))
      .sort();

    // 保留最近 MAX_SHORT_TERM 个
    while (files.length > MAX_SHORT_TERM) {
      const oldest = files.shift();
      fs.unlinkSync(path.join(dir, oldest));
    }
  } catch { /* ignore */ }
}

/** memory.md 超大时截断最早条目 */
function truncateMemory(content) {
  const lines = content.split('\n');
  // 找到 ## 关键事实 段，删除最早的条目
  const factIdx = lines.findIndex(l => l.startsWith('## 关键事实'));
  if (factIdx === -1) return content;

  // 找到条目行（以 - [ 开头）
  const entries = [];
  for (let i = factIdx + 1; i < lines.length; i++) {
    if (lines[i].startsWith('## ')) break;
    if (lines[i].startsWith('- [')) entries.push(i);
  }

  // 删除最早的 1/3 条目
  const removeCount = Math.max(1, Math.floor(entries.length / 3));
  const removeIndices = new Set(entries.slice(0, removeCount));
  const newLines = lines.filter((_, i) => !removeIndices.has(i));

  return newLines.join('\n');
}

module.exports = {
  // Phase 1
  saveSessionSummary,
  loadRecentMemory,
  appendToMemoryMd,
  // Phase 2
  retrieveRelevantPlans,
  saveCerebellumPlan,
  extractKeywords,
  // 内部 (测试用)
  loadHippocampusIndex,
  saveHippocampusIndex,
};