From 55e8dcf48a09b5e0301ef075e031a460eb764982 Mon Sep 17 00:00:00 2001
From: musistudio <m@musiiot.top>
Date: Mon, 6 Oct 2025 08:35:26 +0800
Subject: [PATCH] add count_tokens api

---
 src/server.ts       | 7 +++++++
 src/utils/router.ts | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/server.ts b/src/server.ts
index 7483913..6ca3e4d 100644
--- a/src/server.ts
+++ b/src/server.ts
@@ -5,10 +5,17 @@ import { join } from "path";
 import fastifyStatic from "@fastify/static";
 import { readdirSync, statSync, readFileSync, writeFileSync, existsSync } from "fs";
 import { homedir } from "os";
+import {calculateTokenCount} from "./utils/router";
 
 export const createServer = (config: any): Server => {
   const server = new Server(config);
 
+  server.app.post("/v1/messages/count_tokens", async (req, reply) => {
+    const {messages, tools, system} = req.body;
+    const tokenCount = calculateTokenCount(messages, system, tools);
+    return { "input_tokens": tokenCount }
+  });
+
   // Add endpoint to read config.json with access control
   server.app.get("/api/config", async (req, reply) => {
     return await readConfigFile();
diff --git a/src/utils/router.ts b/src/utils/router.ts
index 2d54503..c9a3b54 100644
--- a/src/utils/router.ts
+++ b/src/utils/router.ts
@@ -9,7 +9,7 @@ import { readFile } from 'fs/promises'
 
 const enc = get_encoding("cl100k_base");
 
-const calculateTokenCount = (
+export const calculateTokenCount = (
   messages: MessageParam[],
   system: any,
   tools: Tool[]