MagicCode/packages/core/src/context.ts

// Taille du contexte : comptage local des tokens (js-tiktoken), lecture d'une taille en tokens, question posée au LLM.
import { Tiktoken } from "js-tiktoken/lite";
import o200k_base from "js-tiktoken/ranks/o200k_base";
import type { ContentBlock, LLMProvider, Message, ToolSpec } from "./types";

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

// Encodeur o200k_base (js-tiktoken, version « lite » : seul ce vocabulaire est embarqué), créé à la première utilisation
let encoder: Tiktoken | undefined;

/** Nombre de tokens d'un texte (js-tiktoken, o200k_base). Un texte qui ressemble à un jeton spécial est compté comme du texte. */
export function countTextTokens(text: string): number {
  if (text === "") return 0;
  encoder ??= new Tiktoken(o200k_base);
  return encoder.encode(text, [], []).length;
}

// Un bloc de la conversation ne change plus une fois ajouté : son compte est mémorisé (le contexte est recompté avant
// CHAQUE saisie de l'utilisateur, sans re-tokeniser toute la conversation à chaque fois).
const blockCounts = new WeakMap<object, number>();

function countBlock(block: ContentBlock): number {
  let n = blockCounts.get(block);
  if (n === undefined) {
    n =
      block.type === "text"
        ? countTextTokens(block.text)
        : block.type === "tool_use"
          ? countTextTokens(block.name + JSON.stringify(block.input))
          : countTextTokens(block.content);
    blockCounts.set(block, n);
  }
  return n;
}

/** countTextTokens avec mémoire de la dernière valeur (prompt système et outils changent rarement). */
function lastValueCounter(): (text: string) => number {
  let lastText: string | undefined;
  let lastCount = 0;
  return (text) => {
    if (text !== lastText) {
      lastCount = countTextTokens(text);
      lastText = text;
    }
    return lastCount;
  };
}
const countSystem = lastValueCounter();
const countTools = lastValueCounter();

export interface ContextParts {
  system: string;
  messages: Message[];
  tools?: ToolSpec[];
  /** Textes ajoutés au contexte mais pas encore envoyés (commandes « ! », fichiers cités). */
  pending?: string[];
}

/**
 * Nombre de tokens du contexte : prompt système + définition des outils + messages (+ contexte en attente).
 * Calcul local et déterministe, SOMME des comptes de chaque élément : ajouter un message ne fait donc qu'augmenter le
 * total (il ne diminue que si la conversation est réduite : /clear, /compact, annulation d'un tour, autre session).
 */
export function countContextTokens(parts: ContextParts): number {
  let total = countSystem(parts.system);
  if (parts.tools && parts.tools.length > 0) total += countTools(JSON.stringify(parts.tools));
  for (const m of parts.messages) {
    total += MESSAGE_OVERHEAD_TOKENS;
    for (const b of m.content) total += countBlock(b);
  }
  for (const p of parts.pending ?? []) total += countTextTokens(p);
  return total;
}

/** 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() };
}