Added test

Author: gsiddh

packages/vertexai/src/methods/chrome-adapter.test.ts

@@ -157,7 +157,7 @@ describe('ChromeAdapter', () => {
             Promise.resolve({
               available: 'after-download'
             }),
-          create: () => {}
+          create: () => { }
         }
       } as AI;
       const downloadPromise = new Promise<AILanguageModel>(() => {
@@ -182,7 +182,7 @@ describe('ChromeAdapter', () => {
             Promise.resolve({
               available: 'after-download'
             }),
-          create: () => {}
+          create: () => { }
         }
       } as AI;
       let resolveDownload;
@@ -298,4 +298,48 @@ describe('ChromeAdapter', () => {
       });
     });
   });
+  describe('countTokens', () => {
+    it('Extracts initial prompts and counts tokens', async () => {
+      const aiProvider = {
+        languageModel: {
+          create: () => Promise.resolve({})
+        }
+      } as AI;
+      const expectedCount = 10;
+      const countPromptTokensStub = stub().resolves(expectedCount);
+      const factoryStub = stub(aiProvider.languageModel, 'create').resolves({
+        countPromptTokens: countPromptTokensStub
+      } as unknown as AILanguageModel);
+      const text = ['first', 'second', 'third'];
+      const onDeviceParams = {
+        initialPrompts: [{ role: 'user', content: text[0] }]
+      } as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      const response = await adapter.countTokens({
+        contents: [
+          { role: 'model', parts: [{ text: text[1] }] },
+          { role: 'user', parts: [{ text: text[2] }] }
+        ]
+      });
+      expect(factoryStub).to.have.been.calledOnceWith({
+        initialPrompts: [
+          { role: 'user', content: text[0] },
+          // Asserts tail is passed as initial prompts, and
+          // role is normalized from model to assistant.
+          { role: 'assistant', content: text[1] }
+        ]
+      });
+      expect(countPromptTokensStub).to.have.been.calledOnceWith({
+        role: 'user',
+        content: text[2]
+      });
+      expect(await response.json()).to.deep.equal({
+        totalTokens: expectedCount
+      });
+    });
+  });
 });

packages/vertexai/src/methods/chrome-adapter.ts

@@ -105,13 +105,15 @@ export class ChromeAdapter {
   }
   async countTokens(request: CountTokensRequest): Promise<Response> {
     const options = this.onDeviceParams || {};
-    const prompts = ChromeAdapter.toInitialPrompts(request.contents);
+    options.initialPrompts ??= [];
+    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
+    const currentPrompt = extractedInitialPrompts.pop()!;
+    options.initialPrompts.push(...extractedInitialPrompts);
     const session = await this.session(options);
-    const tokenCount = await session.countPromptTokens(prompts);
+    const tokenCount = await session.countPromptTokens(currentPrompt);
     return {
       json: async () => ({
         totalTokens: tokenCount,
-        totalBillableCharacters: 0,
       })
     } as Response;
   }