MagicCode/packages/core/src/context.ts

// Taille du contexte : estimation locale, lecture d'une taille en tokens, question posée au LLM.
import type { LLMProvider, Message, ToolSpec } from "./types";

const CHARS_PER_TOKEN = 4;
const MESSAGE_OVERHEAD_TOKENS = 4; // balises de rôle autour de chaque message

/** Estimation d'un texte (≈ 4 caractères par token). */
export function estimateTextTokens(text: string): number {
  return Math.ceil(text.length / CHARS_PER_TOKEN);
}

/**
 * Estimation brute de la taille d'une requête (prompt système, outils, messages) quand le serveur ne renvoie pas la
 * consommation réelle. Seul le contenu compte (texte, arguments et résultats d'outils), pas la structure JSON interne.
 * Elle est ensuite calibrée sur les valeurs réelles du serveur (voir AgentSession.tokenScale).
 */
export function estimateTokens(system: string, messages: Message[], tools: ToolSpec[] = []): number {
  let chars = system.length;
  for (const m of messages) {
    for (const b of m.content) {
      if (b.type === "text") chars += b.text.length;
      else if (b.type === "tool_use") chars += b.name.length + JSON.stringify(b.input).length;
      else chars += b.content.length;
    }
  }
  if (tools.length > 0) chars += JSON.stringify(tools).length;
  return Math.ceil(chars / CHARS_PER_TOKEN) + messages.length * MESSAGE_OVERHEAD_TOKENS;
}

/** Question posée au LLM après le choix du modèle (la consigne finale facilite la lecture de la réponse). */
export const CONTEXT_WINDOW_QUESTION =
  "Taille maximale du contexte pour ce modèle, en tokens ? Réponds uniquement par un nombre entier, sans autre texte.";

const MIN_CONTEXT = 1_000;
const MAX_CONTEXT = 100_000_000;

const MULTIPLIERS: Record<string, number> = {
  k: 1_000,
  mille: 1_000,
  m: 1_000_000,
  million: 1_000_000,
  millions: 1_000_000,
  milliard: 1_000_000_000,
  milliards: 1_000_000_000,
};

// 128 000 · 128,000 · 128.000 · 1.5 · 1,5 · 200000 — suivi éventuel d'un multiplicateur (k, M, mille, million…) ou de « token »
const NUMBER = /(\d{1,3}(?:[ .,  ]\d{3})+(?!\d)|\d+(?:[.,]\d+)?)\s*(milliards?|millions?|mille|[kKmM](?![A-Za-zÀ-ÿ])|(?=tokens?\b))?/gi;

function toNumber(raw: string): number {
  if (/^\d{1,3}([ .,  ]\d{3})+$/.test(raw)) return Number(raw.replace(/[ .,  ]/g, ""));
  return Number(raw.replace(",", "."));
}

/**
 * Extrait une taille de contexte (en tokens) de la réponse d'un LLM : « 128000 », « 128 000 tokens », « 200k », « 1M »,
 * « 1,5 million »… Une valeur avec unité l'emporte ; sinon la plus grande valeur plausible (1 000 à 100 000 000).
 * Retourne null si rien d'exploitable.
 */
export function parseTokenCount(text: string): number | null {
  let best: { value: number; unit: boolean } | null = null;
  for (const m of text.matchAll(NUMBER)) {
    const base = toNumber(m[1] ?? "");
    const suffix = (m[2] ?? "").toLowerCase();
    const value = Math.round(base * (MULTIPLIERS[suffix] ?? 1));
    if (!Number.isFinite(value) || value < MIN_CONTEXT || value > MAX_CONTEXT) continue;
    const unit = suffix !== "" || /^\s*tokens?\b/i.test(text.slice((m.index ?? 0) + m[0].length));
    if (!best || (unit && !best.unit) || (unit === best.unit && value > best.value)) best = { value, unit };
  }
  return best ? best.value : null;
}

export interface ContextWindowAnswer {
  /** Taille maximale en tokens, null si la réponse n'a pas pu être lue. */
  tokens: number | null;
  /** Réponse brute du LLM (utile pour informer l'utilisateur en cas d'échec). */
  answer: string;
}

/** Demande au LLM la taille maximale de son contexte. Lève une erreur si l'appel échoue (réseau, HTTP, annulation). */
export async function queryContextWindow(
  provider: LLMProvider,
  signal?: AbortSignal,
): Promise<ContextWindowAnswer> {
  let answer = "";
  for await (const ev of provider.chat(
    {
      system: "Tu réponds de façon brève et exacte.",
      messages: [{ role: "user", content: [{ type: "text", text: CONTEXT_WINDOW_QUESTION }] }],
      tools: [],
    },
    signal,
  )) {
    if (ev.type === "text_delta") answer += ev.text;
  }
  return { tokens: parseTokenCount(answer), answer: answer.trim() };
}