OmniRoute/src/shared/utils/tiktokenCounter.ts

import type { Tiktoken } from "js-tiktoken";
import { createRequire } from "module";

const _require = createRequire(import.meta.url);

export type TokenizerEncoding = "cl100k_base" | "o200k_base";

export interface TokenizerContext {
  provider?: string | null;
  model?: string | null;
}

export function tokenizerContextFromBody(body: unknown): TokenizerContext {
  if (!body || typeof body !== "object" || Array.isArray(body)) return {};
  const record = body as Record<string, unknown>;
  return {
    provider: typeof record.provider === "string" ? record.provider : undefined,
    model: typeof record.model === "string" ? record.model : undefined,
  };
}

const encoders = new Map<TokenizerEncoding, Tiktoken>();

function normalize(value: unknown): string {
  return typeof value === "string" ? value.trim().toLowerCase() : "";
}

export function isCodexTokenizerContext(context?: TokenizerContext): boolean {
  const provider = normalize(context?.provider);
  const model = normalize(context?.model);
  return (
    provider === "codex" ||
    provider === "cx" ||
    model.startsWith("codex/") ||
    model.startsWith("cx/") ||
    model.includes("codex")
  );
}

export function resolveTokenizerEncoding(context?: TokenizerContext): TokenizerEncoding {
  return isCodexTokenizerContext(context) ? "o200k_base" : "cl100k_base";
}

function getEncoder(encoding: TokenizerEncoding): Tiktoken {
  const cached = encoders.get(encoding);
  if (cached) return cached;
  let tiktoken: { getEncoding: (name: string) => Tiktoken } | null = null;
  try {
    tiktoken ??= _require("js-tiktoken") as { getEncoding: (name: string) => Tiktoken };
    const created = tiktoken.getEncoding(encoding);
    encoders.set(encoding, created);
    return created;
  } catch {
    throw new Error(`js-tiktoken not available: cannot create encoder for ${encoding}`);
  }
}

/**
 * Exact token count for a string using the selected offline tokenizer.
 * Existing callers retain cl100k_base; Codex callers may pass provider/model context
 * to use o200k_base.
 * Defensive: never throws in a counting path — falls back to a char heuristic.
 */
export function countTextTokens(text: string, context?: TokenizerContext): number {
  if (!text || typeof text !== "string") return 0;
  try {
    return getEncoder(resolveTokenizerEncoding(context)).encode(text).length;
  } catch {
    return Math.ceil(text.length / 4);
  }
}