metona-ai-desktop/electron/harness/memory/manager.ts

/**
 * Memory Manager — 记忆管理器
 *
 * 基于 SQLite（better-sqlite3）的三层记忆系统。
 * 表结构由 DatabaseService 统一创建，此处不再重复。
 *
 * v0.2.0 增强：
 * - TF-IDF 语义检索替代 LIKE 关键词搜索
 * - 时间衰减策略：老旧记忆权重降低
 * - IDF 缓存：避免每次搜索重新计算
 *
 * @see docs/生产级通用 AI Agent 智能体桌面应用：完整设计与构建指南.html — 第六章
 * @see standard/开发规范.md — 使用 better-sqlite3（禁止自写数据库层）
 */

import { nanoid } from 'nanoid';
import type Database from 'better-sqlite3';
import log from 'electron-log';

export type MemoryType = 'episodic' | 'semantic' | 'working';
export type MemorySource = 'user_input' | 'tool_result' | 'agent_thought' | 'imported';

export interface MemoryItem {
  id: string;
  type: MemoryType;
  content: string;
  summary?: string;
  source: MemorySource;
  importance: number;
  sessionId?: string;
  createdAt: number;
  expiresAt?: number;
}

export interface SearchResult extends MemoryItem {
  score: number;
}

interface MemorySearchOptions {
  topK?: number;
  sessionId?: string;
  type?: MemoryType;
  minImportance?: number;
}

/** 分词：将文本拆分为词项（支持中英文） */
function tokenize(text: string): string[] {
  // 转小写，按非字母数字字符分割
  const lower = text.toLowerCase();
  // 英文词
  const words = lower.match(/[a-z][a-z0-9_-]{1,}/g) ?? [];
  // 中文 bigram（相邻两字组合），覆盖 CJK 统一表意、扩展 A、平假名/片假名、谚文
  const cjkChars = lower.match(/[\u4e00-\u9fff\u3400-\u4dbf\u3040-\u30ff\uac00-\ud7af]/g) ?? [];
  const bigrams: string[] = [];
  for (let i = 0; i < cjkChars.length - 1; i++) {
    bigrams.push(cjkChars[i] + cjkChars[i + 1]);
  }
  // 单字 CJK 补 unigram（避免单字文档无 token）
  if (cjkChars.length === 1) {
    bigrams.push(cjkChars[0]);
  }
  return [...words, ...bigrams];
}

/** 计算词频（TF） */
function computeTF(tokens: string[]): Map<string, number> {
  const tf = new Map<string, number>();
  for (const token of tokens) {
    tf.set(token, (tf.get(token) ?? 0) + 1);
  }
  // 归一化
  const total = tokens.length || 1;
  for (const [key, val] of tf) {
    tf.set(key, val / total);
  }
  return tf;
}

/** 计算余弦相似度的点积部分 */
function dotProduct(tf1: Map<string, number>, tf2: Map<string, number>, idf: Map<string, number>): number {
  let sum = 0;
  for (const [term, freq1] of tf1) {
    const freq2 = tf2.get(term);
    if (freq2 !== undefined) {
      const idfVal = idf.get(term) ?? 1;
      sum += freq1 * freq2 * idfVal * idfVal;
    }
  }
  return sum;
}

/** 计算向量模长 */
function vectorNorm(tf: Map<string, number>, idf: Map<string, number>): number {
  let sum = 0;
  for (const [term, freq] of tf) {
    const idfVal = idf.get(term) ?? 1;
    sum += (freq * idfVal) ** 2;
  }
  return Math.sqrt(sum);
}

/** 时间衰减权重：30 天半衰期（age=30 时 weight=0.5） */
function timeDecayWeight(createdAt: number, now: number = Date.now()): number {
  const ageDays = Math.max(0, (now - createdAt) / (24 * 60 * 60 * 1000));
  const halfLifeDays = 30;
  return Math.pow(0.5, ageDays / halfLifeDays);
}

/**
 * 记忆管理器
 */
export class MemoryManager {
  /** IDF 缓存：词项 -> 文档频率 */
  private idfCache = new Map<string, number>();
  /** 缓存的记忆总数 */
  private cachedDocCount = 0;
  /** 缓存最后更新时间 */
  private cacheUpdatedAt = 0;
  /** 缓存有效期（5 分钟） */
  private readonly CACHE_TTL = 5 * 60 * 1000;

  constructor(private getDB: () => Database.Database) {}

  /**
   * 初始化（表结构由 DatabaseService 创建）
   */
  initialize(): void {
    log.info('MemoryManager initialized (v0.2.0: TF-IDF enabled)');
  }

  /**
   * 更新 IDF 缓存
   *
   * v0.3.0 增强：
   * - 原子替换缓存（先构建新数据再替换，避免中间不一致状态）
   * - 错误处理（数据库查询失败时保留旧缓存，不更新时间戳）
   */
  private updateIdfCache(): void {
    const now = Date.now();
    if (now - this.cacheUpdatedAt < this.CACHE_TTL && this.cachedDocCount > 0) {
      return; // 缓存未过期
    }

    const db = this.getDB();

    try {
      // v0.3.0: 先构建新缓存数据，再原子替换
      const newIdfCache = new Map<string, number>();

      // 获取所有记忆内容（episodic + semantic + working）
      const episodicRows = db.prepare('SELECT content, summary FROM episodic_memories').all() as Array<{ content: string; summary: string | null }>;
      const semanticRows = db.prepare('SELECT value FROM semantic_memories').all() as Array<{ value: string }>;
      const workingRows = db.prepare('SELECT value FROM working_memories').all() as Array<{ value: string }>;

      const allDocs = [
        ...episodicRows.map((r) => r.content + ' ' + (r.summary ?? '')),
        ...semanticRows.map((r) => r.value),
        ...workingRows.map((r) => r.value),
      ];

      const newDocCount = allDocs.length;
      const docFreq = new Map<string, number>();

      for (const doc of allDocs) {
        const tokens = new Set(tokenize(doc));
        for (const token of tokens) {
          docFreq.set(token, (docFreq.get(token) ?? 0) + 1);
        }
      }

      // IDF = log((N+1)/(df+1)) + 1（Sklearn 风格平滑），确保非负
      for (const [term, df] of docFreq) {
        newIdfCache.set(term, Math.log((newDocCount + 1) / (df + 1)) + 1);
      }

      // v0.3.0: 原子替换 — 只有新数据完全准备好后才替换旧缓存
      this.idfCache = newIdfCache;
      this.cachedDocCount = newDocCount;
      this.cacheUpdatedAt = now;
    } catch (error) {
      // v0.3.0: 数据库查询失败时保留旧缓存，不更新 cacheUpdatedAt
      // 这样下次 search() 会再次尝试更新
      log.error('MemoryManager: Failed to update IDF cache, keeping stale cache:', error);
    }
  }

  /**
   * TF-IDF 相似度搜索
   */
  private tfidfSearch(query: string, options: MemorySearchOptions): SearchResult[] {
    const db = this.getDB();
    this.updateIdfCache();

    const queryTokens = tokenize(query);
    if (queryTokens.length === 0) return [];

    const queryTF = computeTF(queryTokens);
    const queryNorm = vectorNorm(queryTF, this.idfCache);
    if (queryNorm === 0) return [];

    const { topK = 5, type, minImportance = 0 } = options;
    const results: SearchResult[] = [];
    const now = Date.now();

    // 搜索 episodic 记忆
    if (!type || type === 'episodic') {
      const rows = db.prepare(`
        SELECT * FROM episodic_memories WHERE importance >= ?
        ORDER BY importance DESC, created_at DESC LIMIT ?
      `).all(minImportance, topK * 3) as Array<{
        id: string; session_id: string | null; content: string; summary: string | null;
        source: string; importance: number; created_at: number; expires_at: number | null;
      }>;

      for (const row of rows) {
        const docText = row.content + ' ' + (row.summary ?? '');
        const docTokens = tokenize(docText);
        const docTF = computeTF(docTokens);
        const docNorm = vectorNorm(docTF, this.idfCache);

        if (docNorm === 0) continue;

        const dotProd = dotProduct(queryTF, docTF, this.idfCache);
        const cosineSim = dotProd / (queryNorm * docNorm);

        // 时间衰减
        const decayWeight = timeDecayWeight(row.created_at, now);
        // 最终分数 = 余弦相似度 * 时间衰减 * 重要度权重
        const finalScore = cosineSim * decayWeight * (0.5 + row.importance * 0.5);

        if (finalScore > 0) {
          results.push({
            id: row.id, type: 'episodic', content: row.content,
            summary: row.summary ?? undefined, source: row.source as MemorySource,
            importance: row.importance, sessionId: row.session_id ?? undefined,
            createdAt: row.created_at, expiresAt: row.expires_at ?? undefined,
            score: finalScore,
          });
        }
      }
    }

    // 搜索 semantic 记忆
    if (!type || type === 'semantic') {
      const rows = db.prepare(`
        SELECT * FROM semantic_memories WHERE confidence >= ?
        ORDER BY confidence DESC, access_count DESC LIMIT ?
      `).all(minImportance, Math.ceil(topK * 1.5)) as Array<{
        id: string; key: string; value: string; category: string | null;
        confidence: number; source_session: string | null; created_at: number;
      }>;

      for (const row of rows) {
        const docText = row.key + ' ' + row.value;
        const docTokens = tokenize(docText);
        const docTF = computeTF(docTokens);
        const docNorm = vectorNorm(docTF, this.idfCache);

        if (docNorm === 0) continue;

        const dotProd = dotProduct(queryTF, docTF, this.idfCache);
        const cosineSim = dotProd / (queryNorm * docNorm);

        const decayWeight = timeDecayWeight(row.created_at, now);
        const finalScore = cosineSim * decayWeight * (0.5 + row.confidence * 0.5);

        if (finalScore > 0) {
          results.push({
            id: row.id, type: 'semantic', content: row.value,
            source: 'imported', importance: row.confidence,
            sessionId: row.source_session ?? undefined,
            createdAt: row.created_at, score: finalScore,
          });
        }
      }
    }

    // 搜索 working 记忆
    if (!type || type === 'working') {
      const rows = db.prepare(`
        SELECT * FROM working_memories ORDER BY updated_at DESC LIMIT ?
      `).all(topK * 3) as Array<{
        id: string; session_id: string; task_id: string;
        key: string; value: string; updated_at: number;
      }>;

      for (const row of rows) {
        const docText = row.key + ' ' + row.value;
        const docTokens = tokenize(docText);
        const docTF = computeTF(docTokens);
        const docNorm = vectorNorm(docTF, this.idfCache);

        if (docNorm === 0) continue;

        const dotProd = dotProduct(queryTF, docTF, this.idfCache);
        const cosineSim = dotProd / (queryNorm * docNorm);

        const decayWeight = timeDecayWeight(row.updated_at, now);
        const finalScore = cosineSim * decayWeight * 0.5;

        if (finalScore > 0) {
          results.push({
            id: row.id, type: 'working', content: row.value,
            source: 'agent_thought', importance: 0.5,
            sessionId: row.session_id, createdAt: row.updated_at,
            score: finalScore,
          });
        }
      }
    }

    return results.sort((a, b) => b.score - a.score).slice(0, topK);
  }

  /**
   * 存储记忆
   *
   * v0.3.0 修复：
   * - switch 添加 default 分支，未知 type 抛错而非静默失败
   * - working 类型使用 item.id（若提供）或生成唯一 key，避免同 session 多次存储互相覆盖
   * - semantic 类型使用 item.summary 作为 key（若提供），支持更新已有记忆
   */
  store(item: Omit<MemoryItem, 'id' | 'createdAt'> & { id?: string }): string {
    const db = this.getDB();
    const id = item.id ?? `mem_${nanoid(12)}`;
    const now = Date.now();
    const importance = item.importance ?? this.calculateImportance(item);

    switch (item.type) {
      case 'episodic':
        db.prepare(`
          INSERT INTO episodic_memories (id, session_id, content, summary, source, importance, created_at)
          VALUES (?, ?, ?, ?, ?, ?, ?)
        `).run(id, item.sessionId ?? null, item.content, item.summary ?? null, item.source, importance, now);
        break;
      case 'semantic':
        // v0.3.0 修复：使用 summary 作为 key（若提供），支持更新已有语义记忆
        db.prepare(`
          INSERT OR REPLACE INTO semantic_memories (id, key, value, category, confidence, source_session, created_at, updated_at, access_count)
          VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0)
        `).run(id, item.summary ?? id, item.content, 'general', importance, item.sessionId ?? null, now, now);
        break;
      case 'working':
        // v0.3.0 修复：使用 summary 作为 key（若提供），避免硬编码 'default' 导致覆盖
        db.prepare(`
          INSERT OR REPLACE INTO working_memories (id, session_id, task_id, key, value, updated_at)
          VALUES (?, ?, ?, ?, ?, ?)
        `).run(id, item.sessionId ?? 'default', 'default', item.summary ?? id, item.content, now);
        break;
      default:
        // v0.3.0 修复：未知 type 抛错而非静默失败
        throw new Error(`Unknown memory type: ${(item as { type: string }).type}`);
    }

    // 使 IDF 缓存失效
    this.cacheUpdatedAt = 0;

    log.debug(`Memory stored: ${id} (${item.type})`);
    return id;
  }

  /**
   * 检索记忆（v0.2.0: TF-IDF 语义检索 + 时间衰减）
   *
   * v0.2.0 变更：
   * - 使用 TF-IDF 余弦相似度替代 LIKE 关键词搜索
   * - 支持中英文分词（英文按词，中文按 bigram）
   * - 时间衰减：30 天半衰期，老旧记忆权重降低
   * - IDF 缓存：5 分钟有效期，避免重复计算
   */
  search(query: string, options: MemorySearchOptions = {}): SearchResult[] {
    const db = this.getDB();
    const { topK = 5, type, minImportance = 0 } = options;
    // v0.3.0 修复：拦截空 query 和纯空格 query
    if (!query || !query.trim()) return [];

    // v0.2.0: 优先使用 TF-IDF 语义搜索
    const tfidfResults = this.tfidfSearch(query, options);
    if (tfidfResults.length > 0) {
      return tfidfResults;
    }

    // 回退：如果 TF-IDF 没有结果（如 IDF 缓存为空），使用 LIKE 关键词搜索
    // 转义 LIKE 通配符，避免用户输入的 % 和 _ 影响匹配
    // v0.3.0 修复: 反斜杠也需转义，否则含 \ 的搜索（如 Windows 路径）会导致 SQLite LIKE 报错
    const escapedQuery = query.replace(/[%_\\]/g, '\\$&');
    const pattern = `%${escapedQuery}%`;
    const results: SearchResult[] = [];

    // 搜索情节记忆
    if (!type || type === 'episodic') {
      const rows = db.prepare(`
        SELECT * FROM episodic_memories
        WHERE (content LIKE ? ESCAPE '\\' OR summary LIKE ? ESCAPE '\\') AND importance >= ?
        ORDER BY importance DESC, created_at DESC LIMIT ?
      `).all(pattern, pattern, minImportance, topK) as Array<{
        id: string; session_id: string | null; content: string; summary: string | null;
        source: string; importance: number; created_at: number; expires_at: number | null;
      }>;
      for (const row of rows) {
        results.push({
          id: row.id, type: 'episodic', content: row.content,
          summary: row.summary ?? undefined, source: row.source as MemorySource,
          importance: row.importance, sessionId: row.session_id ?? undefined,
          createdAt: row.created_at, expiresAt: row.expires_at ?? undefined,
          score: row.importance * timeDecayWeight(row.created_at),
        });
      }
    }

    // 搜索语义记忆
    if (!type || type === 'semantic') {
      const rows = db.prepare(`
        SELECT * FROM semantic_memories
        WHERE (key LIKE ? ESCAPE '\\' OR value LIKE ? ESCAPE '\\') AND confidence >= ?
        ORDER BY confidence DESC, access_count DESC LIMIT ?
      `).all(pattern, pattern, minImportance, Math.ceil(topK / 2)) as Array<{
        id: string; key: string; value: string; category: string | null;
        confidence: number; source_session: string | null; created_at: number;
      }>;
      for (const row of rows) {
        results.push({
          id: row.id, type: 'semantic', content: row.value,
          source: 'imported', importance: row.confidence,
          sessionId: row.source_session ?? undefined,
          createdAt: row.created_at, score: row.confidence * timeDecayWeight(row.created_at),
        });
      }
    }

    // 搜索工作记忆
    // v0.3.0 修复：LIKE 回退路径也需添加 !type 分支（与 tfidfSearch 保持一致）
    if (!type || type === 'working') {
      const rows = db.prepare(`
        SELECT * FROM working_memories
        WHERE (key LIKE ? ESCAPE '\\' OR value LIKE ? ESCAPE '\\')
        ORDER BY updated_at DESC LIMIT ?
      `).all(pattern, pattern, topK) as Array<{
        id: string; session_id: string; task_id: string;
        key: string; value: string; updated_at: number;
      }>;
      for (const row of rows) {
        results.push({
          id: row.id, type: 'working', content: row.value,
          source: 'agent_thought', importance: 0.5,
          sessionId: row.session_id, createdAt: row.updated_at,
          score: 0.3 * timeDecayWeight(row.updated_at),
        });
      }
    }

    return results.sort((a, b) => b.score - a.score).slice(0, topK);
  }

  /**
   * 获取工作记忆
   */
  getWorkingMemory(sessionId: string, taskId: string = 'default'): Map<string, string> {
    const db = this.getDB();
    const rows = db.prepare(`
      SELECT key, value FROM working_memories WHERE session_id = ? AND task_id = ?
    `).all(sessionId, taskId) as Array<{ key: string; value: string }>;
    return new Map(rows.map((r) => [r.key, r.value]));
  }

  /**
   * 更新工作记忆
   */
  setWorkingMemory(sessionId: string, taskId: string, key: string, value: string): void {
    const db = this.getDB();
    db.prepare(`
      INSERT OR REPLACE INTO working_memories (id, session_id, task_id, key, value, updated_at)
      VALUES (?, ?, ?, ?, ?, ?)
    `).run(`wm_${nanoid(8)}`, sessionId, taskId, key, value, Date.now());
  }

  /**
   * 清除工作记忆
   */
  clearWorkingMemory(sessionId: string, taskId?: string): void {
    const db = this.getDB();
    if (taskId) {
      db.prepare('DELETE FROM working_memories WHERE session_id = ? AND task_id = ?').run(sessionId, taskId);
    } else {
      db.prepare('DELETE FROM working_memories WHERE session_id = ?').run(sessionId);
    }
  }

  /**
   * 清理过期记忆
   */
  cleanupExpired(): number {
    const db = this.getDB();
    const result = db.prepare('DELETE FROM episodic_memories WHERE expires_at IS NOT NULL AND expires_at < ?').run(Date.now());
    return result.changes;
  }

  private calculateImportance(item: Omit<MemoryItem, 'id' | 'createdAt'>): number {
    let score = 0.5;
    if (item.source === 'user_input') score += 0.2;
    if (item.source === 'tool_result') score += 0.1;
    if (item.content.length > 200) score += 0.1;
    return Math.min(1, Math.max(0, score));
  }
}