Test content generation

Author: erikeldridge

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

@@ -173,13 +173,18 @@ describe('ChromeAdapter', () => {
       const createStub = stub(aiProvider.languageModel, 'create').resolves(
         {} as AILanguageModel
       );
-      const adapter = new ChromeAdapter(aiProvider, 'prefer_on_device');
+      const onDeviceParams = {} as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
       expect(
         await adapter.isAvailable({
           contents: [{ role: 'user', parts: [{ text: 'hi' }] }]
         })
       ).to.be.false;
-      expect(createStub).to.have.been.calledOnce;
+      expect(createStub).to.have.been.calledOnceWith(onDeviceParams);
     });
     it('avoids redundant downloads', async () => {
       const aiProvider = {
@@ -250,4 +255,67 @@ describe('ChromeAdapter', () => {
       ).to.be.false;
     });
   });
+  describe('generateContentOnDevice', () => {
+    it('Extracts and concats initial prompts', async () => {
+      const aiProvider = {
+        languageModel: {
+          create: () => Promise.resolve({})
+        }
+      } as AI;
+      const factoryStub = stub(aiProvider.languageModel, 'create').resolves({
+        prompt: s => Promise.resolve(s)
+      } 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.generateContentOnDevice({
+        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(response.text()).to.equal(text[2]);
+    });
+    it('Extracts system prompt', async () => {
+      const aiProvider = {
+        languageModel: {
+          create: () => Promise.resolve({})
+        }
+      } as AI;
+      const factoryStub = stub(aiProvider.languageModel, 'create').resolves({
+        prompt: s => Promise.resolve(s)
+      } as AILanguageModel);
+      const onDeviceParams = {
+        systemPrompt: 'be yourself'
+      } as AILanguageModelCreateOptionsWithSystemPrompt;
+      const adapter = new ChromeAdapter(
+        aiProvider,
+        'prefer_on_device',
+        onDeviceParams
+      );
+      const text = 'hi';
+      const response = await adapter.generateContentOnDevice({
+        contents: [{ role: 'user', parts: [{ text }] }]
+      });
+      expect(factoryStub).to.have.been.calledOnceWith({
+        initialPrompts: [],
+        systemPrompt: onDeviceParams.systemPrompt
+      });
+      expect(response.text()).to.equal(text);
+    });
+  });
 });

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

@@ -21,7 +21,6 @@ import {
   EnhancedGenerateContentResponse,
   GenerateContentRequest,
   InferenceMode,
-  Part,
   Role,
   TextPart
 } from '../types';
@@ -78,16 +77,13 @@ export class ChromeAdapter {
   async generateContentOnDevice(
     request: GenerateContentRequest
   ): Promise<EnhancedGenerateContentResponse> {
-    const initialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
+    const createOptions = this.onDeviceParams || {};
+    createOptions.initialPrompts ??= [];
+    const extractedInitialPrompts = ChromeAdapter.toInitialPrompts(request.contents);
     // Assumes validation asserted there is at least one initial prompt.
-    const prompt = initialPrompts.pop()!;
-    const systemPrompt = ChromeAdapter.toSystemPrompt(
-      request.systemInstruction
-    );
-    const session = await this.session({
-      initialPrompts,
-      systemPrompt
-    });
+    const prompt = extractedInitialPrompts.pop()!;
+    createOptions.initialPrompts.push(...extractedInitialPrompts);
+    const session = await this.session(createOptions);
     const result = await session.prompt(prompt.content);
     return {
       text: () => result,
@@ -152,33 +148,6 @@ export class ChromeAdapter {
         this.isDownloading = false;
       });
   }
-  private static toSystemPrompt(
-    prompt: string | Content | Part | undefined
-  ): string | undefined {
-    if (!prompt) {
-      return undefined;
-    }
-
-    if (typeof prompt === 'string') {
-      return prompt;
-    }
-
-    const systemContent = prompt as Content;
-    if (
-      systemContent.parts &&
-      systemContent.parts[0] &&
-      systemContent.parts[0].text
-    ) {
-      return systemContent.parts[0].text;
-    }
-
-    const systemPart = prompt as Part;
-    if (systemPart.text) {
-      return systemPart.text;
-    }
-
-    return undefined;
-  }
   private static toOnDeviceRole(role: Role): AILanguageModelPromptRole {
     return role === 'model' ? 'assistant' : 'user';
   }
@@ -192,9 +161,9 @@ export class ChromeAdapter {
     }));
   }
   private async session(
-    opts: AILanguageModelCreateOptionsWithSystemPrompt
+    options: AILanguageModelCreateOptionsWithSystemPrompt
   ): Promise<AILanguageModel> {
-    const newSession = await this.aiProvider!.languageModel.create(opts);
+    const newSession = await this.aiProvider!.languageModel.create(options);
     if (this.oldSession) {
       this.oldSession.destroy();
     }