count-tokens.ts

Author: russellwheatley

packages/ai/lib/methods/count-tokens.ts

@@ -18,19 +18,28 @@
 import { CountTokensRequest, CountTokensResponse, RequestOptions } from '../types';
 import { Task, makeRequest } from '../requests/request';
 import { ApiSettings } from '../types/internal';
+import { BackendType } from '../public-types';
+import * as GoogleAIMapper from '../googleai-mappers';
 
 export async function countTokens(
   apiSettings: ApiSettings,
   model: string,
   params: CountTokensRequest,
   requestOptions?: RequestOptions,
 ): Promise<CountTokensResponse> {
+  let body: string = '';
+  if (apiSettings.backend.backendType === BackendType.GOOGLE_AI) {
+    const mappedParams = GoogleAIMapper.mapCountTokensRequest(params, model);
+    body = JSON.stringify(mappedParams);
+  } else {
+    body = JSON.stringify(params);
+  }
   const response = await makeRequest(
     model,
     Task.COUNT_TOKENS,
     apiSettings,
     false,
-    JSON.stringify(params),
+    body,
     requestOptions,
   );
   return response.json();