socaichat/langgraph/src/agent/utils/truncate-messages.ts

import { BaseMessage, SystemMessage, HumanMessage } from "@langchain/core/messages";

const MAX_RECENT_MESSAGES = 20;

/**
 * Estimate token count using a simple character-based heuristic.
 * ~4 chars per token for English, ~2 chars per token for Chinese.
 * We use ~3 as a rough average for mixed content.
 */
function estimateTokens(message: BaseMessage): number {
  const content = typeof message.content === "string"
    ? message.content
    : JSON.stringify(message.content);
  return Math.ceil(content.length / 3);
}

/**
 * Truncate messages to fit within a token budget.
 * Strategy:
 * - Always keep system messages (if any)
 * - Keep the most recent N messages
 * - If the history exceeds maxTokens, drop older messages and inject a summary note
 */
export function truncateMessages(
  messages: BaseMessage[],
  maxTokens = 60000,
): BaseMessage[] {
  if (messages.length === 0) return messages;

  // Separate system messages from the rest
  const systemMessages = messages.filter((m) => m instanceof SystemMessage);
  const nonSystemMessages = messages.filter((m) => !(m instanceof SystemMessage));

  // Calculate total tokens
  const totalTokens = messages.reduce((sum, m) => sum + estimateTokens(m), 0);

  // If within budget, return as-is
  if (totalTokens <= maxTokens) {
    return messages;
  }

  // Keep only recent messages
  const recentMessages = nonSystemMessages.slice(-MAX_RECENT_MESSAGES);

  // Check if recent messages alone fit
  const recentTokens =
    systemMessages.reduce((sum, m) => sum + estimateTokens(m), 0) +
    recentMessages.reduce((sum, m) => sum + estimateTokens(m), 0);

  if (recentTokens <= maxTokens) {
    const droppedCount = nonSystemMessages.length - recentMessages.length;
    if (droppedCount > 0) {
      const summaryNote = new HumanMessage(
        `[系统提示：为保持对话质量，已省略前 ${droppedCount} 条较早的消息。如需回溯之前的内容，请告诉我。]`,
      );
      return [...systemMessages, summaryNote, ...recentMessages];
    }
    return [...systemMessages, ...recentMessages];
  }

  // Still too large — progressively trim from the front of recent messages
  const result: BaseMessage[] = [...systemMessages];
  let budget = maxTokens - systemMessages.reduce((sum, m) => sum + estimateTokens(m), 0);

  const trimmed: BaseMessage[] = [];
  for (let i = recentMessages.length - 1; i >= 0; i--) {
    const tokens = estimateTokens(recentMessages[i]);
    if (budget - tokens >= 0) {
      trimmed.unshift(recentMessages[i]);
      budget -= tokens;
    } else {
      break;
    }
  }

  const droppedCount = nonSystemMessages.length - trimmed.length;
  if (droppedCount > 0) {
    result.push(
      new HumanMessage(
        `[系统提示：为保持对话质量，已省略前 ${droppedCount} 条较早的消息。如需回溯之前的内容，请告诉我。]`,
      ),
    );
  }
  result.push(...trimmed);

  return result;
}