flynn/src/context/weighting.ts

import type { Message } from '../models/types.js';
import { getMessageText } from '../models/media.js';

export interface WeightedMessage {
  index: number;
  message: Message;
  score: number;
}

const TOOL_OUTCOME_PATTERN = /(tool|command|exit code|stack trace|traceback|error|failed|succeeded|output|result)/i;
const CORRECTION_PATTERN = /(actually|correction|instead|sorry|to clarify|i meant|wrong)/i;
const PREFERENCE_PATTERN = /(prefer|preference|always|never|please|timezone|call me|i like|i dislike|do not|don't)/i;

export function scoreMessageImportance(message: Message): number {
  const text = getMessageText(message).trim();
  if (text.length === 0) {
    return 0;
  }

  let score = 0.1;

  if (TOOL_OUTCOME_PATTERN.test(text)) {
    score += 0.45;
  }
  if (CORRECTION_PATTERN.test(text)) {
    score += 0.35;
  }
  if (PREFERENCE_PATTERN.test(text)) {
    score += 0.4;
  }
  if (text.includes('```')) {
    score += 0.2;
  }

  if (message.role === 'user') {
    score += 0.05;
  }

  if (text.length > 240) {
    score += 0.1;
  }

  return Math.max(0, Math.min(1, score));
}

export function rankMessagesByImportance(messages: Message[]): WeightedMessage[] {
  return messages.map((message, index) => {
    const base = scoreMessageImportance(message);
    const recencyBoost = messages.length > 1 ? (index / (messages.length - 1)) * 0.08 : 0;
    return {
      index,
      message,
      score: Math.max(0, Math.min(1, base + recencyBoost)),
    };
  });
}

export function selectImportantMessages(messages: Message[], opts: {
  threshold: number;
  maxMessages: number;
}): WeightedMessage[] {
  if (messages.length === 0 || opts.maxMessages <= 0) {
    return [];
  }

  const ranked = rankMessagesByImportance(messages)
    .filter(item => item.score >= opts.threshold)
    .sort((a, b) => b.score - a.score || b.index - a.index)
    .slice(0, opts.maxMessages)
    .sort((a, b) => a.index - b.index);

  return ranked;
}