Context & Memory 实现

// context/types.ts

export interface ContextConfig {
  maxTotalTokens: number;      // 总窗口上限
  systemPromptBudget: number;  // system prompt 预算
  profileBudget: number;       // 用户画像预算
  taskBudget: number;          // 当前任务预算
  historyBudget: number;       // 历史摘要预算
  recentBudget: number;        // 最近消息预算
  outputReserve: number;       // 输出保留空间
}

export const DEFAULT_CONTEXT_CONFIG: ContextConfig = {
  maxTotalTokens: 16000,
  systemPromptBudget: 2000,
  profileBudget: 500,
  taskBudget: 4000,
  historyBudget: 2000,
  recentBudget: 3000,
  outputReserve: 4000,
};

export interface ContextWindow {
  systemPrompt: string;
  messages: Message[];
  tokenCount: number;
  layers: {
    system: number;
    profile: number;
    task: number;
    history: number;
    recent: number;
  };
}

// context/manager.ts

import { ContextConfig, ContextWindow, DEFAULT_CONTEXT_CONFIG } from './types';
import { Compressor } from './compressor';
import { TokenBudget } from './budget';

export class ContextManager {
  private config: ContextConfig;
  private compressor: Compressor;
  private budget: TokenBudget;

  private systemPrompt: string = '';
  private profileBlock: string = '';
  private taskBlock: string = '';
  private historySummary: string = '';
  private recentMessages: Message[] = [];
  private allMessages: Message[] = [];

  constructor(config: Partial<ContextConfig> = {}) {
    this.config = { ...DEFAULT_CONTEXT_CONFIG, ...config };
    this.compressor = new Compressor();
    this.budget = new TokenBudget(this.config);
  }

  setSystemPrompt(prompt: string): void {
    this.systemPrompt = this.compressor.truncate(prompt, this.config.systemPromptBudget);
  }

  setProfile(profile: string): void {
    this.profileBlock = this.compressor.truncate(profile, this.config.profileBudget);
  }

  setTaskContext(context: string): void {
    this.taskBlock = this.compressor.truncate(context, this.config.taskBudget);
  }

  addMessage(message: Message): void {
    this.allMessages.push(message);
    this.recentMessages.push(message);

    // 滑动窗口：只保留最近 N 条
    while (this.countTokens(this.recentMessages) > this.config.recentBudget) {
      const evicted = this.recentMessages.shift()!;
      // 被踢出的消息进入历史摘要队列
      this.appendToHistory(evicted);
    }
  }

  addToolResult(toolCallId: string, result: string): void {
    // 工具输出可能很长，先压缩再加入
    const compressed = this.compressor.compressToolOutput(result, 1000);
    this.addMessage({
      role: 'tool',
      toolCallId,
      content: compressed,
    });
  }

  build(): ContextWindow {
    const fullSystem = this.buildSystemBlock();
    const messages = this.buildMessages();
    const tokenCount = this.countTokens(messages) + this.estimateTokens(fullSystem);

    return {
      systemPrompt: fullSystem,
      messages,
      tokenCount,
      layers: {
        system: this.estimateTokens(this.systemPrompt),
        profile: this.estimateTokens(this.profileBlock),
        task: this.estimateTokens(this.taskBlock),
        history: this.estimateTokens(this.historySummary),
        recent: this.countTokens(this.recentMessages),
      },
    };
  }

  needsCompaction(): boolean {
    const current = this.build().tokenCount;
    const threshold = this.config.maxTotalTokens - this.config.outputReserve;
    return current > threshold * 0.9; // 90% 时预警
  }

  async autoCompact(queryEngine: QueryEngine): Promise<void> {
    if (!this.needsCompaction()) return;

    // Level 1: 压缩工具输出
    this.compressRecentToolOutputs();

    if (this.needsCompaction()) {
      // Level 2: 压缩历史为摘要
      this.historySummary = await this.compressor.summarizeHistory(
        this.historySummary,
        this.recentMessages.slice(0, -3), // 保留最近 3 条不压缩
        queryEngine,
      );
      this.recentMessages = this.recentMessages.slice(-3);
    }

    if (this.needsCompaction()) {
      // Level 3: 激进压缩——任务上下文也压缩
      this.taskBlock = await this.compressor.summarize(this.taskBlock, queryEngine, 1000);
    }
  }

  getStats(): string {
    const window = this.build();
    const used = window.tokenCount;
    const max = this.config.maxTotalTokens;
    return `Context: ${used}/${max} tokens (${Math.round(used / max * 100)}%) | ` +
      `system:${window.layers.system} profile:${window.layers.profile} ` +
      `task:${window.layers.task} history:${window.layers.history} recent:${window.layers.recent}`;
  }

  private buildSystemBlock(): string {
    const parts: string[] = [this.systemPrompt];
    if (this.profileBlock) parts.push(`\n## 用户背景\n${this.profileBlock}`);
    if (this.historySummary) parts.push(`\n## 已完成诊断摘要\n${this.historySummary}`);
    return parts.join('\n');
  }

  private buildMessages(): Message[] {
    const messages: Message[] = [];

    // 任务上下文作为第一条 user 消息注入
    if (this.taskBlock) {
      messages.push({ role: 'user', content: this.taskBlock });
      messages.push({ role: 'assistant', content: '好的，我已了解当前任务上下文。请继续。' });
    }

    // 最近消息
    messages.push(...this.recentMessages);

    return messages;
  }

  private appendToHistory(message: Message): void {
    // 简单追加，等 autoCompact 时统一压缩
    if (message.role === 'assistant' && message.content) {
      this.historySummary += `\n- ${message.content.slice(0, 100)}`;
    }
  }

  private compressRecentToolOutputs(): void {
    this.recentMessages = this.recentMessages.map(msg => {
      if (msg.role === 'tool' && msg.content && msg.content.length > 500) {
        return { ...msg, content: this.compressor.compressToolOutput(msg.content, 500) };
      }
      return msg;
    });
  }

  private countTokens(messages: Message[]): number {
    return messages.reduce((sum, m) => sum + this.estimateTokens(m.content ?? ''), 0);
  }

  private estimateTokens(text: string): number {
    // 粗估：中文 ~1.5 token/字，英文 ~0.25 token/字
    const cjk = (text.match(/[一-鿿]/g) || []).length;
    const ascii = text.length - cjk;
    return Math.ceil(cjk * 1.5 + ascii * 0.25);
  }
}

// context/compressor.ts

export class Compressor {
  truncate(text: string, maxTokens: number): string {
    const estimated = this.estimateTokens(text);
    if (estimated <= maxTokens) return text;

    // 按比例截断
    const ratio = maxTokens / estimated;
    const maxChars = Math.floor(text.length * ratio * 0.9); // 留 10% 余量
    return text.slice(0, maxChars) + '\n\n[... 已截断]';
  }

  compressToolOutput(output: string, maxTokens: number): string {
    const estimated = this.estimateTokens(output);
    if (estimated <= maxTokens) return output;

    // 策略：保留 JSON 的 key 结构 + 值的前 N 字符
    try {
      const parsed = JSON.parse(output);
      return JSON.stringify(this.compressObject(parsed, maxTokens), null, 0);
    } catch {
      // 非 JSON，直接截断
      return this.truncate(output, maxTokens);
    }
  }

  async summarizeHistory(
    existingSummary: string,
    messages: Message[],
    queryEngine: QueryEngine,
  ): Promise<string> {
    const content = messages
      .filter(m => m.role === 'assistant' && m.content)
      .map(m => m.content!.slice(0, 200))
      .join('\n');

    if (!content) return existingSummary;

    const response = await queryEngine.query({
      task: 'knowledge_summary',
      messages: [{
        role: 'user',
        content: `请将以下对话历史压缩为简洁摘要（保留关键结论和数据，删除过程细节）：\n\n已有摘要：\n${existingSummary}\n\n新增内容：\n${content}`,
      }],
      systemPrompt: '你是一个对话压缩器。输出尽量简洁，用 bullet points，保留数字和关键结论。',
      maxTokens: 500,
    });

    return response.content ?? existingSummary;
  }

  async summarize(text: string, queryEngine: QueryEngine, maxTokens: number): Promise<string> {
    const response = await queryEngine.query({
      task: 'knowledge_summary',
      messages: [{ role: 'user', content: `请将以下内容压缩到 ${maxTokens} tokens 以内，保留关键信息：\n\n${text}` }],
      maxTokens: Math.min(maxTokens, 1000),
    });
    return response.content ?? this.truncate(text, maxTokens);
  }

  private compressObject(obj: any, maxTokens: number): any {
    if (typeof obj === 'string') {
      return obj.length > 100 ? obj.slice(0, 100) + '...' : obj;
    }
    if (Array.isArray(obj)) {
      // 数组只保留前 3 项
      return obj.slice(0, 3).map(item => this.compressObject(item, maxTokens));
    }
    if (typeof obj === 'object' && obj !== null) {
      const result: any = {};
      for (const [key, value] of Object.entries(obj)) {
        result[key] = this.compressObject(value, maxTokens);
      }
      return result;
    }
    return obj;
  }

  private estimateTokens(text: string): number {
    const cjk = (text.match(/[一-鿿]/g) || []).length;
    const ascii = text.length - cjk;
    return Math.ceil(cjk * 1.5 + ascii * 0.25);
  }
}

// memory/types.ts

export type MemoryType = 'user_profile' | 'diagnosis_summary' | 'weak_point' | 'preference';

export interface MemoryEntry {
  id: string;
  type: MemoryType;
  key: string;                   // 语义标识符
  value: unknown;                // 结构化数据
  confidence: number;            // 0-1，多次验证后提升
  createdAt: string;
  updatedAt: string;
  expiresAt?: string;            // 短期记忆有过期
  accessCount: number;           // 被召回次数
}

export interface UserProfile {
  targetRole?: string;           // 目标岗位
  techStack?: string[];          // 技术栈
  experience?: string;           // 工作经验
  weakDimensions?: string[];     // 历次诊断中暴露的弱项
  strongDimensions?: string[];   // 强项
  lastDiagnosisScore?: number;   // 上次诊断总分
  totalDiagnoses?: number;       // 累计诊断次数
  scoreHistory?: Array<{ date: string; score: number }>;
}

// memory/store.ts

import Database from 'better-sqlite3';
import { MemoryEntry, MemoryType, UserProfile } from './types';

export class MemoryStore {
  private db: Database.Database;

  constructor(dbPath: string) {
    this.db = new Database(dbPath);
    this.db.pragma('journal_mode = WAL');
    this.initSchema();
  }

  private initSchema(): void {
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS memories (
        id TEXT PRIMARY KEY,
        type TEXT NOT NULL,
        key TEXT NOT NULL,
        value TEXT NOT NULL,
        confidence REAL NOT NULL DEFAULT 0.5,
        created_at TEXT NOT NULL DEFAULT (datetime('now')),
        updated_at TEXT NOT NULL DEFAULT (datetime('now')),
        expires_at TEXT,
        access_count INTEGER NOT NULL DEFAULT 0
      );

      CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(type);
      CREATE INDEX IF NOT EXISTS idx_memories_key ON memories(key);
      CREATE INDEX IF NOT EXISTS idx_memories_type_key ON memories(type, key);
    `);
  }

  create(type: MemoryType, key: string, value: unknown): string {
    const id = `${type}:${key}:${Date.now()}`;
    this.db.prepare(`
      INSERT INTO memories (id, type, key, value, confidence)
      VALUES (?, ?, ?, ?, 0.5)
    `).run(id, type, key, JSON.stringify(value));
    return id;
  }

  retrieve(opts: {
    type?: MemoryType;
    key?: string;
    limit?: number;
    minConfidence?: number;
  } = {}): MemoryEntry[] {
    let sql = 'SELECT * FROM memories WHERE 1=1';
    const params: any[] = [];

    if (opts.type) {
      sql += ' AND type = ?';
      params.push(opts.type);
    }
    if (opts.key) {
      sql += ' AND key LIKE ?';
      params.push(`%${opts.key}%`);
    }
    if (opts.minConfidence) {
      sql += ' AND confidence >= ?';
      params.push(opts.minConfidence);
    }
    sql += ' AND (expires_at IS NULL OR expires_at > datetime("now"))';
    sql += ' ORDER BY confidence DESC, updated_at DESC';

    if (opts.limit) {
      sql += ' LIMIT ?';
      params.push(opts.limit);
    }

    const rows = this.db.prepare(sql).all(...params);

    // 记录访问次数
    for (const row of rows) {
      this.db.prepare('UPDATE memories SET access_count = access_count + 1 WHERE id = ?')
        .run((row as any).id);
    }

    return rows.map(r => this.rowToEntry(r));
  }

  update(id: string, value: unknown, boostConfidence?: boolean): void {
    const updates: string[] = ['value = ?', 'updated_at = datetime("now")'];
    const params: any[] = [JSON.stringify(value)];

    if (boostConfidence) {
      updates.push('confidence = MIN(1.0, confidence + 0.1)');
    }

    this.db.prepare(`UPDATE memories SET ${updates.join(', ')} WHERE id = ?`)
      .run(...params, id);
  }

  delete(id: string): void {
    this.db.prepare('DELETE FROM memories WHERE id = ?').run(id);
  }

  deleteAll(type?: MemoryType): void {
    if (type) {
      this.db.prepare('DELETE FROM memories WHERE type = ?').run(type);
    } else {
      this.db.prepare('DELETE FROM memories').run();
    }
  }

  // 用户画像快捷方法
  getProfile(): UserProfile {
    const entries = this.retrieve({ type: 'user_profile' });
    const profile: UserProfile = {};
    for (const entry of entries) {
      Object.assign(profile, entry.value);
    }
    return profile;
  }

  updateProfile(patch: Partial<UserProfile>): void {
    const existing = this.retrieve({ type: 'user_profile', key: 'main' });
    if (existing.length > 0) {
      const merged = { ...(existing[0].value as UserProfile), ...patch };
      this.update(existing[0].id, merged, true);
    } else {
      this.create('user_profile', 'main', patch);
    }
  }

  // 清理过期记忆
  evictExpired(): number {
    const result = this.db.prepare(
      'DELETE FROM memories WHERE expires_at IS NOT NULL AND expires_at <= datetime("now")'
    ).run();
    return result.changes;
  }

  // 统计
  getStats(): { total: number; byType: Record<string, number> } {
    const total = (this.db.prepare('SELECT COUNT(*) as c FROM memories').get() as any).c;
    const byType = this.db.prepare(
      'SELECT type, COUNT(*) as c FROM memories GROUP BY type'
    ).all().reduce((acc: any, r: any) => ({ ...acc, [r.type]: r.c }), {});
    return { total, byType };
  }

  private rowToEntry(row: any): MemoryEntry {
    return {
      id: row.id,
      type: row.type,
      key: row.key,
      value: JSON.parse(row.value),
      confidence: row.confidence,
      createdAt: row.created_at,
      updatedAt: row.updated_at,
      expiresAt: row.expires_at,
      accessCount: row.access_count,
    };
  }
}

// memory/retriever.ts

export class MemoryRetriever {
  private store: MemoryStore;

  constructor(store: MemoryStore) {
    this.store = store;
  }

  retrieveForDiagnosis(question: string, dimension?: string): RetrievedMemories {
    // 1. 始终召回用户画像
    const profile = this.store.getProfile();

    // 2. 按维度召回历史弱点
    const weakPoints = this.store.retrieve({
      type: 'weak_point',
      key: dimension,
      limit: 3,
      minConfidence: 0.6,
    });

    // 3. 召回相关的历史诊断摘要
    const history = this.store.retrieve({
      type: 'diagnosis_summary',
      limit: 2,
    });

    return { profile, weakPoints, history };
  }

  formatForContext(memories: RetrievedMemories): string {
    const parts: string[] = [];

    if (memories.profile.targetRole) {
      parts.push(`目标岗位: ${memories.profile.targetRole}`);
    }
    if (memories.profile.techStack?.length) {
      parts.push(`技术栈: ${memories.profile.techStack.join(', ')}`);
    }
    if (memories.profile.weakDimensions?.length) {
      parts.push(`历史弱项: ${memories.profile.weakDimensions.join(', ')}`);
    }
    if (memories.weakPoints.length) {
      parts.push(`近期同维度弱点:\n${memories.weakPoints.map(w => `- ${(w.value as any).description}`).join('\n')}`);
    }

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

interface RetrievedMemories {
  profile: UserProfile;
  weakPoints: MemoryEntry[];
  history: MemoryEntry[];
}

// memory/triggers.ts

export class MemoryTriggers {
  private store: MemoryStore;

  constructor(store: MemoryStore) {
    this.store = store;
  }

  afterDiagnosis(report: DiagnosisReport): void {
    // 触发 1: 更新用户诊断历史
    const profile = this.store.getProfile();
    this.store.updateProfile({
      lastDiagnosisScore: report.summary.overallScore,
      totalDiagnoses: (profile.totalDiagnoses ?? 0) + 1,
      scoreHistory: [
        ...(profile.scoreHistory ?? []),
        { date: new Date().toISOString().slice(0, 10), score: report.summary.overallScore },
      ].slice(-20), // 只保留最近 20 次
    });

    // 触发 2: 更新弱项维度
    if (report.summary.topWeaknesses.length > 0) {
      this.updateWeakDimensions(report.summary.topWeaknesses);
    }

    // 触发 3: 诊断摘要（保留关键数据点，不保留全文）
    this.store.create('diagnosis_summary', `diag-${Date.now()}`, {
      date: new Date().toISOString().slice(0, 10),
      totalQuestions: report.summary.totalQuestions,
      overallScore: report.summary.overallScore,
      contentAvg: report.summary.contentAvg,
      speechAvg: report.summary.speechAvg,
      topWeaknesses: report.summary.topWeaknesses.slice(0, 3),
    });
  }

  afterSingleQuestion(question: string, diagnosis: ContentDiagnosis, dimension: string): void {
    // 只在低分题触发弱点记录
    if (diagnosis.overallScore < 50) {
      const existing = this.store.retrieve({ type: 'weak_point', key: dimension });
      const alreadyKnown = existing.some(e =>
        (e.value as any).description?.includes(diagnosis.keyMissing[0])
      );

      if (!alreadyKnown) {
        this.store.create('weak_point', dimension, {
          question: question.slice(0, 80),
          score: diagnosis.overallScore,
          description: diagnosis.keyMissing[0] ?? '回答深度不足',
          dimension,
        });
      } else {
        // 已知弱点再次出现——提升 confidence
        const match = existing.find(e =>
          (e.value as any).description?.includes(diagnosis.keyMissing[0])
        );
        if (match) this.store.update(match.id, match.value, true);
      }
    }
  }

  onUserInfo(info: Partial<UserProfile>): void {
    this.store.updateProfile(info);
  }

  private updateWeakDimensions(weaknesses: string[]): void {
    const profile = this.store.getProfile();
    const current = new Set(profile.weakDimensions ?? []);
    weaknesses.forEach(w => current.add(w));
    // 保留最近 5 个弱项
    this.store.updateProfile({
      weakDimensions: [...current].slice(-5),
    });
  }
}

// agent/loop.ts 中相关片段

async function agentLoop(input: string, session: Session): Promise<void> {
  const contextManager = session.contextManager;
  const memoryStore = session.memoryStore;
  const retriever = new MemoryRetriever(memoryStore);

  // 1. 召回记忆注入 Context
  const memories = retriever.retrieveForDiagnosis(input);
  contextManager.setProfile(retriever.formatForContext(memories));

  // 2. 正常 Agent Loop...
  contextManager.addMessage({ role: 'user', content: input });
  const window = contextManager.build();

  const response = await queryEngine.query({
    messages: window.messages,
    systemPrompt: window.systemPrompt,
    // ...
  });

  // 3. 检查是否需要压缩
  await contextManager.autoCompact(queryEngine);

  // 4. 诊断完成后触发记忆写入（通过 post-tool hook）
  // 见 Hook 层实现
}