Merge pull request #1293 from narengogi/feat/unified-count-tokens-endpoint

add a new unified route for count_tokens endpoint

Author: VisargD

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/count_tokens' 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/bedrock/api.ts

@@ -280,6 +280,9 @@ const BedrockAPIConfig: BedrockAPIConfigInterface = {
       case 'cancelFinetune': {
         return `/model-customization-jobs/${jobId}/stop`;
       }
+      case 'messagesCountTokens': {
+        return `/model/${uriEncodedModel}/count-tokens`;
+      }
       default:
         return '';
     }

src/providers/bedrock/countTokens.ts

@@ -0,0 +1,45 @@
+import { ProviderConfig } from '../types';
+import { BedrockMessagesParams } from './types';
+import { transformUsingProviderConfig } from '../../services/transformToProviderRequest';
+import { BedrockConverseMessagesConfig } from './messages';
+import { Params } from '../../types/requestBody';
+import { BEDROCK } from '../../globals';
+import { BedrockErrorResponseTransform } from './chatComplete';
+import { generateInvalidProviderResponseError } from '../utils';
+
+// https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CountTokens.html#API_runtime_CountTokens_RequestSyntax
+export const BedrockConverseMessageCountTokensConfig: ProviderConfig = {
+  messages: {
+    param: 'input',
+    required: true,
+    transform: (params: BedrockMessagesParams) => {
+      return {
+        converse: transformUsingProviderConfig(
+          BedrockConverseMessagesConfig,
+          params as Params
+        ),
+      };
+    },
+  },
+};
+
+// https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_CountTokens.html#API_runtime_CountTokens_ResponseSyntax
+export const BedrockConverseMessageCountTokensResponseTransform = (
+  response: any,
+  responseStatus: number
+) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    return (
+      BedrockErrorResponseTransform(response) ||
+      generateInvalidProviderResponseError(response, BEDROCK)
+    );
+  }
+
+  if ('inputTokens' in response) {
+    return {
+      input_tokens: response.inputTokens,
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, BEDROCK);
+};

src/providers/bedrock/index.ts

@@ -80,6 +80,10 @@ import {
   BedrockConverseMessagesStreamChunkTransform,
   BedrockMessagesResponseTransform,
 } from './messages';
+import {
+  BedrockConverseMessageCountTokensConfig,
+  BedrockConverseMessageCountTokensResponseTransform,
+} from './countTokens';
 
 const BedrockConfig: ProviderConfigs = {
   api: BedrockAPIConfig,
@@ -110,8 +114,6 @@ const BedrockConfig: ProviderConfigs = {
             responseTransforms: {
               'stream-complete': BedrockAnthropicCompleteStreamChunkTransform,
               complete: BedrockAnthropicCompleteResponseTransform,
-              messages: BedrockMessagesResponseTransform,
-              'stream-messages': BedrockConverseMessagesStreamChunkTransform,
             },
           };
           break;
@@ -201,24 +203,40 @@ const BedrockConfig: ProviderConfigs = {
             },
           };
       }
-      if (!config.chatComplete) {
-        config.chatComplete = BedrockConverseChatCompleteConfig;
-      }
-      if (!config.messages) {
-        config.messages = BedrockConverseMessagesConfig;
-      }
-      if (!config.responseTransforms?.['stream-chatComplete']) {
-        config.responseTransforms = {
-          ...(config.responseTransforms ?? {}),
-          'stream-chatComplete': BedrockChatCompleteStreamChunkTransform,
-        };
-      }
-      if (!config.responseTransforms?.chatComplete) {
-        config.responseTransforms = {
-          ...(config.responseTransforms ?? {}),
+
+      // defaults
+      config = {
+        ...config,
+        ...(!config.chatComplete && {
+          chatComplete: BedrockConverseChatCompleteConfig,
+        }),
+        ...(!config.messages && {
+          messages: BedrockConverseMessagesConfig,
+        }),
+        ...(!config.messagesCountTokens && {
+          messagesCountTokens: BedrockConverseMessageCountTokensConfig,
+        }),
+      };
+
+      config.responseTransforms = {
+        ...(config.responseTransforms ?? {}),
+        ...(!config.responseTransforms?.chatComplete && {
           chatComplete: BedrockChatCompleteResponseTransform,
-        };
-      }
+        }),
+        ...(!config.responseTransforms?.['stream-chatComplete'] && {
+          'stream-chatComplete': BedrockChatCompleteStreamChunkTransform,
+        }),
+        ...(!config.responseTransforms?.messages && {
+          messages: BedrockMessagesResponseTransform,
+        }),
+        ...(!config.responseTransforms?.['stream-messages'] && {
+          'stream-messages': BedrockConverseMessagesStreamChunkTransform,
+        }),
+        ...(!config.responseTransforms?.messagesCountTokens && {
+          messagesCountTokens:
+            BedrockConverseMessageCountTokensResponseTransform,
+        }),
+      };
     }
 
     const commonResponseTransforms = {

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

@@ -178,6 +178,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';
 import {
   GetMistralAIChatCompleteResponseTransform,
   GetMistralAIChatCompleteStreamChunkTransform,
@@ -122,6 +123,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.', '');
+    },
+  },
+};