Implement getNumTokens on GoogleBaseLLM and ChatGoogleBase

Author: hans00

libs/langchain-google-common/src/chat_models.ts

@@ -306,6 +306,11 @@ export abstract class ChatGoogleBase<AuthOptions>
     return copyAIModelParams(this, options);
   }
 
+  async getNumTokens(messages: BaseMessage[]): Promise<number> {
+    const parameters = this.invocationParams();
+    return this.connection.requestCountTokens(messages, parameters);
+  }
+
   async _generate(
     messages: BaseMessage[],
     options: this["ParsedCallOptions"],

libs/langchain-google-common/src/connection.ts

@@ -96,9 +96,9 @@ export abstract class GoogleConnection<
     data: unknown | undefined,
     _options: CallOptions,
     requestHeaders: Record<string, string> = {},
-    callMethod?: string
+    callMethod: string | undefined = undefined
   ): Promise<GoogleAbstractedClientOps> {
-    const url = await this.buildUrl();
+    const url = await this.buildUrl(callMethod);
     const method = this.buildMethod();
     const infoHeaders = (await this._clientInfoHeaders()) ?? {};
     const additionalHeaders = (await this.additionalHeaders()) ?? {};
@@ -127,17 +127,25 @@ export abstract class GoogleConnection<
   async _request(
     data: unknown | undefined,
     options: CallOptions,
-    requestHeaders: Record<string, string> = {},
-    callMethod?: string
+    requestHeaders: Record<string, string> = {}
   ): Promise<ResponseType> {
-    const opts = await this._buildOpts(data, options, requestHeaders, callMethod);
+    const opts = await this._buildOpts(data, options, requestHeaders);
     const callResponse = await this.caller.callWithOptions(
       { signal: options?.signal },
       async () => this.client.request(opts)
     );
     const response: unknown = callResponse; // Done for typecast safety, I guess
     return <ResponseType>response;
   }
+
+  async _requestCountTokens(data: unknown | undefined): Promise<number> {
+    const opts = await this._buildOpts(data, {} as CallOptions, {}, 'countTokens');
+    const { totalTokens } = await this.caller.callWithOptions(
+      {},
+      async () => this.client.request(opts)
+    ) as { totalTokens: number };
+    return totalTokens;
+  }
 }
 
 export abstract class GoogleHostConnection<
@@ -334,11 +342,11 @@ export abstract class GoogleAIConnection<
     return url;
   }
 
-  async buildUrlVertex(): Promise<string> {
+  async buildUrlVertex(callMethod?: string): Promise<string> {
     if (this.isApiKey) {
-      return this.buildUrlVertexExpress();
+      return this.buildUrlVertexExpress(callMethod);
     } else {
-      return this.buildUrlVertexLocation();
+      return this.buildUrlVertexLocation(callMethod);
     }
   }
 
@@ -356,6 +364,14 @@ export abstract class GoogleAIConnection<
     parameters: GoogleAIModelRequestParams
   ): Promise<unknown>;
 
+  async requestCountTokens(
+    input: InputType,
+    parameters: GoogleAIModelRequestParams
+  ): Promise<number> {
+    const data = await this.formatData(input, parameters);
+    return await this._requestCountTokens(data);
+  }
+
   async request(
     input: InputType,
     parameters: GoogleAIModelRequestParams,
@@ -428,12 +444,6 @@ export abstract class AbstractGoogleLLMConnection<
   ): Promise<unknown> {
     return this.api.formatData(input, parameters);
   }
-
-  async getNumTokens(input: MessageType) {
-    const data = this.formatData(input)
-    const { totalTokens } = await this._request(data, {}, {}, 'countTokens');
-    return totalTokens
-  }
 }
 
 export interface GoogleCustomEventInfo {

libs/langchain-google-common/src/llms.ts

@@ -177,6 +177,11 @@ export abstract class GoogleBaseLLM<AuthOptions>
     return prompt;
   }
 
+  async getNumTokens(input: string): Promise<number> {
+    const parameters = copyAIModelParams(this, undefined);
+    return this.connection.requestCountTokens(input, parameters);
+  }
+
   /**
    * For some given input string and options, return a string output.
    *