estimateContextTokens function

ContextUsageEstimate estimateContextTokens(List<Message> messages)

Estimate context tokens for messages using provider usage when available.

Ported from pi's estimateContextTokens: the last assistant message with valid usage (not errored/aborted, non-zero tokens) anchors the estimate; everything after it is estimated heuristically. Without any usage, the whole list is estimated.

Implementation

ContextUsageEstimate estimateContextTokens(List<Message> messages) {
  ({Usage usage, int index})? usageInfo;
  for (var i = messages.length - 1; i >= 0; i--) {
    final usage = _assistantUsage(messages[i]);
    if (usage != null) {
      usageInfo = (usage: usage, index: i);
      break;
    }
  }

  if (usageInfo == null) {
    var estimated = 0;
    for (final message in messages) {
      estimated += estimateTokens(message);
    }
    return ContextUsageEstimate(
      tokens: estimated,
      usageTokens: 0,
      trailingTokens: estimated,
      lastUsageIndex: null,
    );
  }

  final usageTokens = calculateContextTokens(usageInfo.usage);
  var trailingTokens = 0;
  for (var i = usageInfo.index + 1; i < messages.length; i++) {
    trailingTokens += estimateTokens(messages[i]);
  }

  return ContextUsageEstimate(
    tokens: usageTokens + trailingTokens,
    usageTokens: usageTokens,
    trailingTokens: trailingTokens,
    lastUsageIndex: usageInfo.index,
  );
}