ai-terminal-assistant/src/context/compaction.ts

import { generateText, type ModelMessage, type LanguageModel } from 'ai';
import { TokenCounter } from './token-counter.js';
import {
  SUMMARY_MARKER,
  type CompressionConfig,
  DEFAULT_COMPRESSION_CONFIG,
} from './types.js';

/**
 * 摘要生成系统提示词
 */
const COMPACTION_SYSTEM_PROMPT = `你是一个专门生成对话摘要的助手。你的任务是将对话历史压缩成一个简洁但信息完整的摘要。

摘要应该包含：
1. 已完成的工作和关键结果
2. 当前正在进行的任务
3. 涉及的重要文件和代码
4. 用户的关键需求和约束
5. 下一步需要做的事情

要求：
- 保留关键技术细节（文件路径、函数名、配置等）
- 使用简洁的列表格式
- 不要遗漏重要信息
- 使用中文回复`;

/**
 * 摘要生成用户提示词
 */
const COMPACTION_USER_PROMPT = `请总结上面的对话。这个摘要将是对话继续时唯一可用的上下文，所以要保留关键信息，包括：完成了什么、正在进行的工作、涉及的文件、下一步计划、以及用户的关键需求或约束。要简洁但足够详细，以便工作可以无缝继续。`;

/**
 * 检查消息是否为摘要消息
 */
export function isSummaryMessage(message: ModelMessage): boolean {
  if (typeof message.content === 'string') {
    return message.content.includes(SUMMARY_MARKER);
  }
  if (Array.isArray(message.content)) {
    return message.content.some(
      (part) =>
        typeof part === 'object' &&
        'text' in part &&
        typeof part.text === 'string' &&
        part.text.includes(SUMMARY_MARKER)
    );
  }
  return false;
}

/**
 * 创建摘要消息
 */
function createSummaryMessage(summary: string): ModelMessage {
  return {
    role: 'assistant',
    content: `${SUMMARY_MARKER}\n## 对话摘要\n\n${summary}\n${SUMMARY_MARKER}`,
  };
}

/**
 * Compaction 策略：使用 AI 生成对话摘要
 *
 * 逻辑：
 * 1. 将历史消息（排除最近保护的部分）发送给 AI
 * 2. AI 生成摘要
 * 3. 用摘要消息替换旧消息
 * 4. 保留最近的消息不变
 *
 * @param messages 消息数组
 * @param model 语言模型
 * @param config 压缩配置
 * @returns 压缩后的消息数组和释放的 tokens
 */
export async function compact(
  messages: ModelMessage[],
  model: LanguageModel,
  config: CompressionConfig = DEFAULT_COMPRESSION_CONFIG
): Promise<{ messages: ModelMessage[]; freedTokens: number }> {
  const { pruneProtect } = config;

  // 计算需要保护的消息数量
  let protectedTokens = 0;
  let protectedCount = 0;

  for (let i = messages.length - 1; i >= 0; i--) {
    const tokens = TokenCounter.estimateMessage(messages[i]);
    if (protectedTokens + tokens > pruneProtect) {
      break;
    }
    protectedTokens += tokens;
    protectedCount++;
  }

  // 确保至少保护最后 2 条消息（除非 pruneProtect 为 0，表示强制压缩模式）
  if (pruneProtect > 0) {
    protectedCount = Math.max(protectedCount, 2);
  } else {
    // 强制压缩模式：至少保护 1 条消息
    protectedCount = Math.max(protectedCount, 1);
  }

  // 分割消息：需要压缩的部分 vs 保护的部分
  const toCompact = messages.slice(0, messages.length - protectedCount);
  const toKeep = messages.slice(messages.length - protectedCount);

  // 如果没有需要压缩的消息，直接返回
  if (toCompact.length === 0) {
    return { messages, freedTokens: 0 };
  }

  // 检查是否已有摘要消息
  const existingSummaryIndex = toCompact.findIndex(isSummaryMessage);
  const messagesForSummary =
    existingSummaryIndex >= 0 ? toCompact.slice(existingSummaryIndex) : toCompact;

  // 计算压缩前的 tokens
  const beforeTokens = TokenCounter.estimateMessages(toCompact);

  try {
    // 调用 AI 生成摘要
    const result = await generateText({
      model,
      system: COMPACTION_SYSTEM_PROMPT,
      messages: [
        ...messagesForSummary,
        {
          role: 'user',
          content: COMPACTION_USER_PROMPT,
        },
      ],
      maxOutputTokens: 2000,
    });

    const summaryMessage = createSummaryMessage(result.text);
    const afterTokens = TokenCounter.estimateMessage(summaryMessage);

    // 返回：摘要 + 保护的消息
    return {
      messages: [summaryMessage, ...toKeep],
      freedTokens: beforeTokens - afterTokens,
    };
  } catch (error) {
    console.error('生成摘要失败:', error);
    // 失败时返回原消息
    return { messages, freedTokens: 0 };
  }
}

/**
 * 简单压缩：不使用 AI，直接截断旧消息
 * 用于没有模型可用或快速压缩的场景
 */
export function simpleCompact(
  messages: ModelMessage[],
  config: CompressionConfig = DEFAULT_COMPRESSION_CONFIG
): { messages: ModelMessage[]; freedTokens: number } {
  const { pruneProtect } = config;

  // 计算需要保留的消息
  let keptTokens = 0;
  let keepFromIndex = messages.length;

  for (let i = messages.length - 1; i >= 0; i--) {
    const tokens = TokenCounter.estimateMessage(messages[i]);
    if (keptTokens + tokens > pruneProtect) {
      break;
    }
    keptTokens += tokens;
    keepFromIndex = i;
  }

  // 确保至少保留最后 N 条消息（强制模式下保留 1 条，否则保留 2 条）
  const minKeep = pruneProtect > 0 ? 2 : 1;
  keepFromIndex = Math.min(keepFromIndex, messages.length - minKeep);

  const removed = messages.slice(0, keepFromIndex);
  const kept = messages.slice(keepFromIndex);

  if (removed.length === 0) {
    return { messages, freedTokens: 0 };
  }

  // 创建简单摘要
  const simpleSummary: ModelMessage = {
    role: 'assistant',
    content: `${SUMMARY_MARKER}\n[对话历史已压缩，共移除 ${removed.length} 条消息]\n${SUMMARY_MARKER}`,
  };

  const freedTokens = TokenCounter.estimateMessages(removed);

  return {
    messages: [simpleSummary, ...kept],
    freedTokens,
  };
}