@xalia/agent 0.6.8 → 0.6.10

package/src/agent/tokenCounter.test.ts

@@ -0,0 +1,243 @@
+import { describe, it, expect, afterEach, beforeEach } from "vitest";
+import {
+  TokenCounter,
+  createTokenCounter,
+  countTokensQuick,
+  ModelDetector,
+} from "./tokenCounter";
+import type { MessageParam } from "./llm";
+import { readFileSync } from "fs";
+import { join } from "path";
+
+describe("ModelDetector", () => {
+  describe("detectProvider", () => {
+    it.each([
+      ["gpt-4o", "openai"],
+      ["gpt-4", "openai"],
+      ["openai/gpt-4o", "openai"],
+      ["o1-preview", "openai"],
+      ["claude-3-7-sonnet-20250219", "anthropic"],
+      ["anthropic/claude-sonnet-4.5", "anthropic"],
+      ["gemini-2.0-flash-exp", "google"],
+      ["google/gemini-pro", "google"],
+      ["unknown-model", "unknown"],
+      ["", "unknown"],
+    ])("should detect %s as %s", (model, expected) => {
+      expect(ModelDetector.detectProvider(model)).toBe(expected);
+    });
+  });
+
+  describe("normalizeModel", () => {
+    it.each([
+      ["openai/gpt-4o", "openai", "gpt-4o"],
+      ["anthropic/claude-3-opus", "anthropic", "claude-3-opus"],
+      ["google/gemini-pro", "google", "gemini-pro"],
+      ["gpt-4o-mini", "openai", "gpt-4o-mini"], // Preserves variants
+      ["o1-preview", "openai", "o1-preview"],
+    ] as const)("should normalize %s to %s", (input, provider, expected) => {
+      expect(ModelDetector.normalizeModel(input, provider)).toBe(expected);
+    });
+  });
+});
+
+describe("TokenCounter", () => {
+  let counter: TokenCounter | undefined;
+
+  afterEach(() => {
+    if (counter) {
+      counter.free();
+    }
+  });
+
+  describe("Basic counting", () => {
+    it("should count tokens for OpenAI models", () => {
+      counter = new TokenCounter("gpt-4o");
+      expect(counter.getProvider()).toBe("openai");
+      expect(counter.countTokens("Hello, world!")).toBeGreaterThan(0);
+      expect(counter.countTokens("")).toBe(0);
+    });
+
+    it("should count tokens for Anthropic models with +7 overhead", () => {
+      counter = new TokenCounter("anthropic/claude-sonnet-4.5");
+      expect(counter.getProvider()).toBe("anthropic");
+      // With +7 overhead, "Hello, world!" should be ~11 tokens
+      expect(counter.countTokens("Hello, world!")).toBeLessThan(15);
+    });
+
+    it("should handle provider prefixes", () => {
+      counter = new TokenCounter("openai/gpt-4o");
+      expect(counter.getProvider()).toBe("openai");
+      expect(counter.getModel()).toBe("openai/gpt-4o");
+    });
+
+    it("should fall back to approximation for unknown models", () => {
+      counter = new TokenCounter("unknown-model");
+      expect(counter.getProvider()).toBe("unknown");
+      const text = "Hello, world!";
+      const expected = Math.ceil(text.length / 3.5);
+      expect(counter.countTokens(text)).toBe(expected);
+    });
+  });
+
+  describe("Message counting", () => {
+    beforeEach(() => {
+      counter = new TokenCounter("gpt-4o");
+    });
+
+    it("should count simple message with overhead", () => {
+      if (!counter) throw new Error("Counter not initialized");
+      const message: MessageParam = {
+        role: "user",
+        content: "Hello, how are you?",
+      };
+      expect(counter.countMessageTokens(message)).toBeGreaterThan(4);
+    });
+
+    it("should count message with tool calls", () => {
+      if (!counter) throw new Error("Counter not initialized");
+      const message: MessageParam = {
+        role: "assistant",
+        content: "Let me help you.",
+        tool_calls: [
+          {
+            id: "call_123",
+            type: "function",
+            function: {
+              name: "get_file",
+              arguments: JSON.stringify({ name: "test.txt" }),
+            },
+          },
+        ],
+      };
+      expect(counter.countMessageTokens(message)).toBeGreaterThan(10);
+    });
+
+    it("should count multimodal message with images", () => {
+      if (!counter) throw new Error("Counter not initialized");
+      const message: MessageParam = {
+        role: "user",
+        content: [
+          { type: "text", text: "What's in this image?" },
+          {
+            type: "image_url",
+            image_url: { url: "data:image/png;base64,..." },
+          },
+        ],
+      };
+      // OpenAI images: ~170 tokens + text + overhead
+      expect(counter.countMessageTokens(message)).toBeGreaterThan(170);
+    });
+
+    it("should count conversation with overhead", () => {
+      if (!counter) throw new Error("Counter not initialized");
+      const messages: MessageParam[] = [
+        { role: "system", content: "You are a helpful assistant." },
+        { role: "user", content: "Hello!" },
+        { role: "assistant", content: "Hi there! How can I help?" },
+      ];
+      expect(counter.countMessagesTokens(messages)).toBeGreaterThan(20);
+      expect(counter.countMessagesTokens([])).toBe(0);
+    });
+  });
+
+  describe("Factory functions", () => {
+    it("should create counter with factory", () => {
+      const testCounter = createTokenCounter("gpt-4o");
+      expect(testCounter.getProvider()).toBe("openai");
+      testCounter.free();
+    });
+
+    it("should count tokens quickly", () => {
+      const count = countTokensQuick("Hello, world!", "gpt-4o");
+      expect(count).toBeGreaterThan(0);
+    });
+  });
+
+  describe("Edge cases", () => {
+    beforeEach(() => {
+      counter = new TokenCounter("gpt-4o");
+    });
+
+    it("should handle long text, special chars, and code", () => {
+      if (!counter) throw new Error("Counter not initialized");
+      expect(counter.countTokens("Hello ".repeat(1000))).toBeGreaterThan(1000);
+      expect(counter.countTokens("Hello! 你好 🌍")).toBeGreaterThan(0);
+      expect(
+        counter.countTokens('function hello() { return "world"; }')
+      ).toBeGreaterThan(5);
+    });
+  });
+});
+
+describe("Dev mode accurate counting", () => {
+  it("should throw error without dev mode", async () => {
+    const counter = new TokenCounter("anthropic/claude-sonnet-4.5");
+    await expect(counter.countTokensAccurate("test")).rejects.toThrow(
+      "Accurate token counting requires enableDevMode"
+    );
+    counter.free();
+  });
+
+  it("should compare p50k_base approximation vs Anthropic API", async () => {
+    const counter = new TokenCounter("claude-3-7-sonnet-20250219", {
+      enableDevMode: true,
+    });
+
+    // Test cases covering different scales
+    const readmePath = join(__dirname, "..", "..", "README.md");
+    const harryPotterPath = join(__dirname, "test_data", "harrypotter.txt");
+
+    const testCases = [
+      { text: "Hello, world!", label: "short" },
+      { text: "The quick brown fox jumps over the lazy dog.", label: "medium" },
+      { text: readFileSync(readmePath, "utf-8"), label: "README.md" },
+      {
+        text: readFileSync(harryPotterPath, "utf-8"),
+        label: "harrypotter.txt",
+      },
+    ];
+
+    if (process.env.ANTHROPIC_API_KEY) {
+      console.log(
+        "\n📊 Token Count Comparison (p50k_base+7 vs Anthropic API)\n"
+      );
+
+      for (const { text, label } of testCases) {
+        const accurate = await counter.countTokensAccurate(text);
+        const approximate = counter.countTokens(text);
+        const diff = accurate - approximate;
+        const accuracy = ((approximate / accurate) * 100).toFixed(1);
+
+        const size =
+          text.length > 1000
+            ? `${String(text.length)} chars`
+            : `${String(text.length)} chars`;
+
+        console.log(`${label.padEnd(20)} [${size}]`);
+        console.log(
+          `  Accurate:    ${accurate.toLocaleString().padStart(8)} tokens`
+        );
+        console.log(
+          `  Approximate: ${approximate.toLocaleString().padStart(8)} tokens`
+        );
+        const diffStr = (diff > 0 ? "+" : "") + String(diff);
+        console.log(`  Difference:  ${diffStr.padStart(8)} (${accuracy}%)\n`);
+
+        expect(accurate).toBeGreaterThan(0);
+        expect(approximate).toBeGreaterThan(0);
+      }
+    } else {
+      await expect(counter.countTokensAccurate("test")).rejects.toThrow();
+    }
+
+    counter.free();
+  });
+
+  it("should return regular count for non-Anthropic models", async () => {
+    const counter = new TokenCounter("gpt-4o", { enableDevMode: true });
+    const regular = counter.countTokens("Hello, world!");
+    const accurate = await counter.countTokensAccurate("Hello, world!");
+    expect(accurate).toBe(regular);
+    counter.free();
+  });
+});