count tokens endpoint for bedrock

Author: narengogi

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/services/transformToProviderRequest.ts

@@ -70,7 +70,7 @@ const getValue = (configParam: string, params: Params, paramConfig: any) => {
 export const transformUsingProviderConfig = (
   providerConfig: ProviderConfig,
   params: Params,
-  providerOptions: Options
+  providerOptions?: Options
 ) => {
   const transformedRequest: { [key: string]: any } = {};