Add support for Vertex AI gemini and normal gemini count token

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

@@ -179,6 +179,10 @@ export const GoogleApiConfig: ProviderAPIConfig = {
         'imageGenerate',
         `${projectRoute}/publishers/${provider}/models/${model}:predict`,
       ],
+      [
+        'messagesCountTokens',
+        `${projectRoute}/publishers/${provider}/models/${model}:countTokens`,
+      ],
     ]);
 
     switch (provider) {

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

@@ -0,0 +1,55 @@
+import { GOOGLE_VERTEX_AI } from '../../globals';
+import { ErrorResponse, ProviderConfig } from '../types';
+import {
+  generateErrorResponse,
+  generateInvalidProviderResponseError,
+} from '../utils';
+import { VertexGoogleChatCompleteConfig } from './chatComplete';
+import { GoogleErrorResponse } from './types';
+
+// VertexGeminiCountTokensConfig reuses the same transforms as
+// VertexGoogleChatCompleteConfig for contents, systemInstruction, tools, and
+// tool_choice so that all token-contributing parameters are forwarded to the
+// Vertex AI Gemini countTokens endpoint.
+export const VertexGeminiCountTokensConfig: ProviderConfig = {
+  model: {
+    param: 'model',
+    required: true,
+  },
+  messages: VertexGoogleChatCompleteConfig.messages,
+  tools: VertexGoogleChatCompleteConfig.tools,
+  tool_choice: VertexGoogleChatCompleteConfig.tool_choice,
+};
+
+interface VertexGeminiCountTokensResponse {
+  totalTokens: number;
+  cachedContentTokenCount?: number;
+}
+
+// VertexGeminiCountTokensResponseTransform maps Vertex AI Gemini's
+// { totalTokens } to the gateway's unified { input_tokens } format.
+export const VertexGeminiCountTokensResponseTransform: (
+  response: VertexGeminiCountTokensResponse | GoogleErrorResponse,
+  responseStatus: number
+) => { input_tokens: number } | ErrorResponse = (response, responseStatus) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    const errorResponse = response as GoogleErrorResponse;
+    return generateErrorResponse(
+      {
+        message: errorResponse.error?.message ?? '',
+        type: errorResponse.error?.status ?? null,
+        param: null,
+        code: String(errorResponse.error?.code ?? ''),
+      },
+      GOOGLE_VERTEX_AI
+    );
+  }
+
+  if ('totalTokens' in response) {
+    return {
+      input_tokens: response.totalTokens,
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, GOOGLE_VERTEX_AI);
+};

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

@@ -53,6 +53,10 @@ import {
   VertexAnthropicMessagesResponseTransform,
 } from './messages';
 import { VertexAnthropicMessagesCountTokensConfig } from './messagesCountTokens';
+import {
+  VertexGeminiCountTokensConfig,
+  VertexGeminiCountTokensResponseTransform,
+} from './geminiCountTokens';
 import {
   GetMistralAIChatCompleteResponseTransform,
   GetMistralAIChatCompleteStreamChunkTransform,
@@ -112,13 +116,15 @@ const VertexConfig: ProviderConfigs = {
           api: GoogleApiConfig,
           embed: GoogleEmbedConfig,
           imageGenerate: GoogleImageGenConfig,
+          messagesCountTokens: VertexGeminiCountTokensConfig,
           createBatch: GoogleBatchCreateConfig,
           createFinetune: baseConfig.createFinetune,
           responseTransforms: {
             'stream-chatComplete': GoogleChatCompleteStreamChunkTransform,
             chatComplete: GoogleChatCompleteResponseTransform,
             embed: GoogleEmbedResponseTransform,
             imageGenerate: GoogleImageGenResponseTransform,
+            messagesCountTokens: VertexGeminiCountTokensResponseTransform,
             ...responseTransforms,
           },
           requestTransforms: {

src/providers/google/api.ts

@@ -23,6 +23,9 @@ export const GoogleApiConfig: ProviderAPIConfig = {
       case 'embed': {
         return `/${routeVersion}/models/${model}:embedContent?key=${apiKey}`;
       }
+      case 'messagesCountTokens': {
+        return `/${routeVersion}/models/${model}:countTokens?key=${apiKey}`;
+      }
       default:
         return '';
     }

src/providers/google/countTokens.ts

@@ -0,0 +1,49 @@
+import { GOOGLE } from '../../globals';
+import { ErrorResponse, ProviderConfig } from '../types';
+import { generateInvalidProviderResponseError } from '../utils';
+import {
+  GoogleChatCompleteConfig,
+  GoogleErrorResponse,
+  GoogleErrorResponseTransform,
+} from './chatComplete';
+
+// GoogleCountTokensConfig reuses the same transforms as GoogleChatCompleteConfig
+// for contents, systemInstruction, tools, and tool_choice so that all token-contributing
+// parameters are forwarded to the Gemini countTokens endpoint.
+// The model param is required but does not need a default for counting.
+export const GoogleCountTokensConfig: ProviderConfig = {
+  model: {
+    param: 'model',
+    required: true,
+  },
+  messages: GoogleChatCompleteConfig.messages,
+  tools: GoogleChatCompleteConfig.tools,
+  tool_choice: GoogleChatCompleteConfig.tool_choice,
+};
+
+interface GoogleCountTokensResponse {
+  totalTokens: number;
+  cachedContentTokenCount?: number;
+}
+
+// GoogleCountTokensResponseTransform maps Gemini's { totalTokens } to the
+// gateway's unified { input_tokens } format.
+export const GoogleCountTokensResponseTransform: (
+  response: GoogleCountTokensResponse | GoogleErrorResponse,
+  responseStatus: number
+) => { input_tokens: number } | ErrorResponse = (response, responseStatus) => {
+  if (responseStatus !== 200) {
+    const errorResponse = GoogleErrorResponseTransform(
+      response as GoogleErrorResponse
+    );
+    if (errorResponse) return errorResponse;
+  }
+
+  if ('totalTokens' in response) {
+    return {
+      input_tokens: response.totalTokens,
+    };
+  }
+
+  return generateInvalidProviderResponseError(response, GOOGLE);
+};

src/providers/google/index.ts

@@ -5,16 +5,22 @@ import {
   GoogleChatCompleteResponseTransform,
   GoogleChatCompleteStreamChunkTransform,
 } from './chatComplete';
+import {
+  GoogleCountTokensConfig,
+  GoogleCountTokensResponseTransform,
+} from './countTokens';
 import { GoogleEmbedConfig, GoogleEmbedResponseTransform } from './embed';
 
 const GoogleConfig: ProviderConfigs = {
   api: GoogleApiConfig,
   chatComplete: GoogleChatCompleteConfig,
   embed: GoogleEmbedConfig,
+  messagesCountTokens: GoogleCountTokensConfig,
   responseTransforms: {
     chatComplete: GoogleChatCompleteResponseTransform,
     'stream-chatComplete': GoogleChatCompleteStreamChunkTransform,
     embed: GoogleEmbedResponseTransform,
+    messagesCountTokens: GoogleCountTokensResponseTransform,
   },
 };