feat(js): ai.countTokens support

js/ai/src/generate.ts

@@ -49,18 +49,22 @@ import {
 import { GenerateResponse } from './generate/response.js';
 import { Message } from './message.js';
 import {
+  GenerateRequestSchema,
   GenerateResponseChunkData,
   GenerateResponseData,
+  GenerationUsageSchema,
   ResolvedModel,
   resolveModel,
   type GenerateActionOptions,
   type GenerateRequest,
   type GenerationCommonConfigSchema,
+  type GenerationUsage,
   type MessageData,
   type MiddlewareRef,
   type ModelArgument,
   type ModelMiddlewareArgument,
   type Part,
+  type TokenCounterAction,
   type ToolRequestPart,
   type ToolResponsePart,
 } from './model.js';
@@ -796,3 +800,106 @@ export function tagAsPreamble(msgs?: MessageData[]): MessageData[] | undefined {
     },
   }));
 }
+
+/**
+ * Counts the tokens for a given generate request.
+ */
+export async function countTokens<
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+>(
+  registry: Registry,
+  options:
+    | GenerateOptions<O, CustomOptions>
+    | PromiseLike<GenerateOptions<O, CustomOptions>>
+): Promise<GenerationUsage> {
+  const resolvedOptions: GenerateOptions<O, CustomOptions> = {
+    ...(await Promise.resolve(options)),
+  };
+
+  const childRegistry = Registry.withParent(registry);
+
+  maybeRegisterDynamicTools(childRegistry, resolvedOptions);
+  maybeRegisterDynamicResources(childRegistry, resolvedOptions);
+
+  const resolvedModel = await resolveModel(
+    childRegistry,
+    resolvedOptions.model
+  );
+
+  const tools = await toolsToActionRefs(childRegistry, resolvedOptions.tools);
+  const resources = await resourcesToActionRefs(
+    childRegistry,
+    resolvedOptions.resources
+  );
+
+  const request = await toGenerateRequest(childRegistry, {
+    ...resolvedOptions,
+    tools,
+    resources,
+  });
+
+  request.config = {
+    ...(resolvedModel?.version ? { version: resolvedModel.version } : {}),
+    ...stripUndefinedOptions(resolvedModel?.config),
+    ...stripUndefinedOptions(request.config),
+  };
+  if (Object.keys(request.config || {}).length === 0) {
+    delete request.config;
+  }
+
+  const middlewareRefs = await normalizeMiddleware(
+    childRegistry,
+    resolvedOptions.use
+  );
+  const resolvedMiddleware = await resolveMiddleware(
+    childRegistry,
+    middlewareRefs
+  );
+  maybeRegisterDynamicMiddlewareTools(childRegistry, resolvedMiddleware);
+
+  let interceptedRequest = request;
+  if (resolvedMiddleware && resolvedMiddleware.length > 0) {
+    const dispatchModel = async (
+      index: number,
+      req: GenerateRequest
+    ): Promise<GenerateResponseData> => {
+      if (index === resolvedMiddleware.length) {
+        interceptedRequest = req;
+        // Return a dummy response to safely unwind the middleware chain without
+        // executing the actual model generation, since we only want to intercept the mutated request.
+        return {
+          message: { role: 'model', content: [] },
+          finishReason: 'stop',
+        };
+      }
+      const currentMiddleware = resolvedMiddleware[index];
+      if (currentMiddleware.model) {
+        return currentMiddleware.model(req, {}, async (modifiedReq) =>
+          dispatchModel(index + 1, modifiedReq || req)
+        );
+      } else {
+        return dispatchModel(index + 1, req);
+      }
+    };
+    await dispatchModel(0, request);
+  }
+
+  const counterActionName = `/model-token-counter/${resolvedModel.modelAction.__action.name}`;
+  const counterAction =
+    resolvedModel.modelAction.__tokenCounterAction ||
+    (await childRegistry.lookupAction<
+      typeof GenerateRequestSchema,
+      typeof GenerationUsageSchema,
+      TokenCounterAction<CustomOptions>
+    >(counterActionName));
+
+  if (!counterAction) {
+    throw new GenkitError({
+      status: 'NOT_FOUND',
+      message: `Model '${resolvedModel.modelAction.__action.name}' does not support token counting (model-token-counter action not found).`,
+    });
+  }
+
+  return await counterAction(interceptedRequest);
+}

js/ai/src/genkit-ai.ts

@@ -34,14 +34,19 @@ import {
   type EmbeddingBatch,
 } from './embedder.js';
 import {
+  countTokens,
   generate,
   generateStream,
   type GenerateOptions,
   type GenerateResponse,
   type GenerateStreamOptions,
   type GenerateStreamResponse,
 } from './generate.js';
-import { GenerationCommonConfigSchema, type Part } from './model-types.js';
+import {
+  GenerationCommonConfigSchema,
+  type GenerationUsage,
+  type Part,
+} from './model-types.js';
 
 /**
  * `GenkitAI` encapsulates Genkit's AI APIs.
@@ -263,6 +268,89 @@ export class GenkitAI {
     return generateStream(this.registry, options);
   }
 
+  /**
+   * Make a countTokens call to the default model with a simple text prompt.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-flash-latest'), // default model
+   * })
+   *
+   * const usage = await ai.countTokens('hi');
+   * ```
+   */
+  countTokens(strPrompt: string): Promise<GenerationUsage>;
+
+  /**
+   * Make a countTokens call to the default model with a multipart request.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   *   model: googleAI.model('gemini-flash-latest'), // default model
+   * })
+   *
+   * const usage = await ai.countTokens([
+   *   { media: {url: 'http://....'} },
+   *   { text: 'describe this image' }
+   * ]);
+   * ```
+   */
+  countTokens(parts: Part[]): Promise<GenerationUsage>;
+
+  /**
+   * Count tokens calculates the token usage of a generative model based on the provided prompt and configuration.
+   *
+   * See {@link GenerateOptions} for detailed information about available options.
+   *
+   * ```ts
+   * const ai = genkit({
+   *   plugins: [googleAI()],
+   * })
+   *
+   * const usage = await ai.countTokens({
+   *   system: 'talk like a pirate',
+   *   prompt: [
+   *     { media: { url: 'http://....' } },
+   *     { text: 'describe this image' }
+   *   ],
+   *   messages: conversationHistory,
+   *   tools: [ userInfoLookup ],
+   *   model: googleAI.model('gemini-flash-latest'),
+   * });
+   * ```
+   */
+  countTokens<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    opts:
+      | GenerateOptions<O, CustomOptions>
+      | PromiseLike<GenerateOptions<O, CustomOptions>>
+  ): Promise<GenerationUsage>;
+
+  async countTokens<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,
+  >(
+    options:
+      | string
+      | Part[]
+      | GenerateOptions<O, CustomOptions>
+      | PromiseLike<GenerateOptions<O, CustomOptions>>
+  ): Promise<GenerationUsage> {
+    if (typeof options === 'string' || Array.isArray(options)) {
+      options = { prompt: options };
+    }
+    return countTokens(
+      this.registry,
+      options as
+        | GenerateOptions<O, CustomOptions>
+        | PromiseLike<GenerateOptions<O, CustomOptions>>
+    );
+  }
+
   /**
    * Checks the status of of a given operation. Returns a new operation which will contain the updated status.
    *

js/ai/src/index.ts

@@ -44,6 +44,7 @@ export {
   GenerateResponseChunk,
   GenerationBlockedError,
   GenerationResponseError,
+  countTokens,
   generate,
   generateOperation,
   generateStream,
@@ -75,8 +76,10 @@ export {
   ModelResponseSchema,
   PartSchema,
   RoleSchema,
+  isModelAction,
   modelActionMetadata,
   modelRef,
+  registerModelAction,
   type GenerateRequest,
   type GenerateRequestData,
   type GenerateResponseChunkData,
@@ -91,6 +94,8 @@ export {
   type ModelResponseData,
   type Part,
   type Role,
+  type TokenCounterAction,
+  type TokenCounterMiddleware,
   type ToolRequestPart,
   type ToolResponsePart,
 } from './model.js';

js/ai/src/model-types.ts

@@ -266,7 +266,9 @@ export const GenerationUsageSchema = z.object({
   outputVideos: z.number().optional(),
   inputAudioFiles: z.number().optional(),
   outputAudioFiles: z.number().optional(),
+  /** @deprecated use `raw` instead */
   custom: z.record(z.number()).optional(),
+  raw: z.unknown().optional(),
   thoughtsTokens: z.number().optional(),
   cachedContentTokens: z.number().optional(),
 });