add a new unified route for count_tokens endpoint

src/handlers/messagesCountTokensHandler.ts

@@ -0,0 +1,57 @@
+import { RouterError } from '../errors/RouterError';
+import {
+  constructConfigFromRequestHeaders,
+  tryTargetsRecursively,
+} from './handlerUtils';
+import { Context } from 'hono';
+
+/**
+ * Handles the '/messages' API request by selecting the appropriate provider(s) and making the request to them.
+ *
+ * @param {Context} c - The Cloudflare Worker context.
+ * @returns {Promise<Response>} - The response from the provider.
+ * @throws Will throw an error if no provider options can be determined or if the request to the provider(s) fails.
+ * @throws Will throw an 500 error if the handler fails due to some reasons
+ */
+export async function messagesCountTokensHandler(
+  c: Context
+): Promise<Response> {
+  try {
+    let request = await c.req.json();
+    let requestHeaders = Object.fromEntries(c.req.raw.headers);
+    const camelCaseConfig = constructConfigFromRequestHeaders(requestHeaders);
+    const tryTargetsResponse = await tryTargetsRecursively(
+      c,
+      camelCaseConfig ?? {},
+      request,
+      requestHeaders,
+      'messagesCountTokens',
+      'POST',
+      'config'
+    );
+
+    return tryTargetsResponse;
+  } catch (err: any) {
+    console.log('messagesCountTokens error', err.message);
+    let statusCode = 500;
+    let errorMessage = 'Something went wrong';
+
+    if (err instanceof RouterError) {
+      statusCode = 400;
+      errorMessage = err.message;
+    }
+
+    return new Response(
+      JSON.stringify({
+        status: 'failure',
+        message: errorMessage,
+      }),
+      {
+        status: statusCode,
+        headers: {
+          'content-type': 'application/json',
+        },
+      }
+    );
+  }
+}

src/index.ts

@@ -36,6 +36,7 @@ import { messagesHandler } from './handlers/messagesHandler';
 // Config
 import conf from '../conf.json';
 import modelResponsesHandler from './handlers/modelResponsesHandler';
+import { messagesCountTokensHandler } from './handlers/messagesCountTokensHandler';
 
 // Create a new Hono server instance
 const app = new Hono();
@@ -126,6 +127,12 @@ app.onError((err, c) => {
  */
 app.post('/v1/messages', requestValidator, messagesHandler);
 
+app.post(
+  '/v1/messages/count_tokens',
+  requestValidator,
+  messagesCountTokensHandler
+);
+
 /**
  * POST route for '/v1/chat/completions'.
  * Handles requests by passing them to the chatCompletionsHandler.

src/providers/anthropic/api.ts

@@ -31,6 +31,8 @@ const AnthropicAPIConfig: ProviderAPIConfig = {
         return '/messages';
       case 'messages':
         return '/messages';
+      case 'messagesCountTokens':
+        return '/messages/count_tokens';
       default:
         return '';
     }

src/providers/anthropic/index.ts

@@ -19,6 +19,7 @@ const AnthropicConfig: ProviderConfigs = {
   complete: AnthropicCompleteConfig,
   chatComplete: AnthropicChatCompleteConfig,
   messages: AnthropicMessagesConfig,
+  messagesCountTokens: AnthropicMessagesConfig,
   api: AnthropicAPIConfig,
   responseTransforms: {
     'stream-complete': AnthropicCompleteStreamChunkTransform,

src/providers/google-vertex-ai/api.ts

@@ -177,6 +177,8 @@ export const GoogleApiConfig: ProviderAPIConfig = {
           mappedFn === 'stream-messages'
         ) {
           return `${projectRoute}/publishers/${provider}/models/${model}:streamRawPredict`;
+        } else if (mappedFn === 'messagesCountTokens') {
+          return `${projectRoute}/publishers/${provider}/models/count-tokens:rawPredict`;
         }
       }
 

src/providers/google-vertex-ai/index.ts

@@ -51,6 +51,7 @@ import {
   VertexAnthropicMessagesConfig,
   VertexAnthropicMessagesResponseTransform,
 } from './messages';
+import { VertexAnthropicMessagesCountTokensConfig } from './messagesCountTokens';
 
 const VertexConfig: ProviderConfigs = {
   api: VertexApiConfig,
@@ -117,6 +118,7 @@ const VertexConfig: ProviderConfigs = {
           createBatch: GoogleBatchCreateConfig,
           createFinetune: baseConfig.createFinetune,
           messages: VertexAnthropicMessagesConfig,
+          messagesCountTokens: VertexAnthropicMessagesCountTokensConfig,
           responseTransforms: {
             'stream-chatComplete':
               VertexAnthropicChatCompleteStreamChunkTransform,

src/providers/google-vertex-ai/messagesCountTokens.ts

@@ -0,0 +1,14 @@
+import { MessageCreateParamsBase } from '../../types/MessagesRequest';
+import { getMessagesConfig } from '../anthropic-base/messages';
+
+export const VertexAnthropicMessagesCountTokensConfig = {
+  ...getMessagesConfig({}),
+  model: {
+    param: 'model',
+    required: true,
+    transform: (params: MessageCreateParamsBase) => {
+      let model = params.model ?? '';
+      return model.replace('anthropic.', '');
+    },
+  },
+};

src/providers/types.ts

@@ -105,7 +105,8 @@ export type endpointStrings =
   | 'getModelResponse'
   | 'deleteModelResponse'
   | 'listResponseInputItems'
-  | 'messages';
+  | 'messages'
+  | 'messagesCountTokens';
 
 /**
  * A collection of API configurations for multiple AI providers.