import test from "node:test";
import assert from "node:assert/strict";
import { splitStreamContent } from "../src/stream.js";
import { sampleVocabulary } from "../src/sampling.js";
test("thinking delimiter split over token boundaries never leaks into final content", () => {
  const state = { thinking: true, pending: "" };
  let reasoning = "",
    content = "";
  for (const text of ["Consider 1. <", "/th", "ink", ">The answer."]) {
    const item = splitStreamContent(state, text);
    reasoning += item.reasoning_content || "";
    content += item.content || "";
  }
  assert.equal(reasoning, "Consider 1. ");
  assert.equal(content, "The answer.");
});
test("Top K disabled samples beyond the first 200 vocabulary entries", () => {
  const logits = new Float32Array(300).fill(0);
  let outside = false;
  for (let seed = 0; seed < 30; seed++) {
    const id = sampleVocabulary(
      logits,
      { temperature: 1, topP: 1, repeatPenalty: 1, seed },
      [],
      7,
    );
    outside ||= id >= 200;
  }
  assert.ok(outside);
});
test("greedy sampling respects repetition penalty and deterministic tie breaks", () => {
  assert.equal(
    sampleVocabulary(
      [1, 2, 3],
      { temperature: 0, topP: 1, repeatPenalty: 2, seed: 1 },
      [2],
      1,
    ),
    1,
  );
});
import { verifySampledDraft, drawDistribution } from "../src/sampling.js";
test("speculative rejection preserves target probabilities despite a different draft", () => {
  const p = new Map([
      [0, 0.2],
      [1, 0.7],
      [2, 0.1],
    ]),
    q = new Map([
      [0, 0.6],
      [1, 0.1],
      [2, 0.3],
    ]),
    counts = [0, 0, 0];
  // Exhaustively integrate three independent uniform grids, avoiding a flaky RNG test.
  for (let d = 0; d < 100; d++)
    for (let a = 0; a < 100; a++)
      for (let r = 0; r < 100; r++) {
        const token = drawDistribution(q, (d + 0.5) / 100);
        counts[
          verifySampledDraft(p, q, token, (a + 0.5) / 100, (r + 0.5) / 100)
            .token
        ]++;
      }
  for (let i = 0; i < 3; i++)
    assert.ok(Math.abs(counts[i] / 1e6 - p.get(i)) < 0.005);
});
import { Tokenizer } from "@huggingface/tokenizers";
import { TokenStream } from "../src/token-stream.js";
import { readFileSync } from "node:fs";
test("streaming tokenizer preserves split Unicode, whitespace, and tool special tokens", () => {
  const tokenizer = new Tokenizer(
    JSON.parse(readFileSync("public/tokenizer/tokenizer.json")),
    JSON.parse(readFileSync("public/tokenizer/tokenizer_config.json")),
  );
  for (const input of [
    "Hello, 世界 🌌 café\n\n  foo\tbar",
    '<function name="create_html"><param name="html"><![CDATA[<h1>你好🚀</h1>]]></param></function>',
  ]) {
    const ids = tokenizer.encode(input, { add_special_tokens: false }).ids,
      stream = new TokenStream(tokenizer);
    const text = ids.map((id) => stream.push(id)).join("") + stream.flush();
    assert.equal(text, tokenizer.decode(ids, { skip_special_tokens: false }));
    assert.equal(text, input);
  }
});