Prototype LocalModel

Author: erikeldridge

packages/vertexai/package.json

@@ -58,6 +58,7 @@
   "devDependencies": {
     "@firebase/app": "0.11.4",
     "@rollup/plugin-json": "6.1.0",
+    "@types/dom-chromium-ai": "0.0.6",
     "rollup": "2.79.2",
     "rollup-plugin-replace": "2.2.0",
     "rollup-plugin-typescript2": "0.36.0",

packages/vertexai/src/api.test.ts

@@ -16,7 +16,13 @@
  */
 import { ImagenModelParams, ModelParams, VertexAIErrorCode } from './types';
 import { VertexAIError } from './errors';
-import { ImagenModel, getGenerativeModel, getImagenModel } from './api';
+import {
+  ImagenModel,
+  getGenerativeModel,
+  getHybridModel,
+  getImagenModel,
+  getLocalModel
+} from './api';
 import { expect } from 'chai';
 import { VertexAI } from './public-types';
 import { GenerativeModel } from './models/generative-model';
@@ -167,4 +173,38 @@ describe('Top level API', () => {
     expect(genModel).to.be.an.instanceOf(ImagenModel);
     expect(genModel.model).to.equal('publishers/google/models/my-model');
   });
+  it('getLocalModel', async () => {
+    const languageModel = {
+      prompt: (s: string) => Promise.resolve(s)
+    } as AILanguageModel;
+    const aiProvider = {
+      languageModel: {
+        create: () => Promise.resolve(languageModel)
+      }
+    } as AI;
+    const model = getLocalModel(aiProvider);
+    const expectedText = 'hello';
+    const response = await model.generateContent(expectedText);
+    expect(response.response.text()).to.equal(expectedText);
+  });
+  it('getHybridModel', async () => {
+    const languageModel = {
+      prompt: (s: string) => Promise.resolve(s)
+    } as AILanguageModel;
+    const aiProvider = {
+      languageModel: {
+        create: () => Promise.resolve(languageModel)
+      }
+    } as AI;
+    const model = getHybridModel(
+      getGenerativeModel(fakeVertexAI, {
+        model: 'my-model'
+      }),
+      getLocalModel(aiProvider)
+    );
+    const chat = model.startChat();
+    const expectedText = 'hello';
+    const response = await chat.sendMessage(expectedText);
+    expect(response.response.text()).to.equal(expectedText);
+  });
 });

packages/vertexai/src/api.ts

@@ -28,7 +28,13 @@ import {
   VertexAIErrorCode
 } from './types';
 import { VertexAIError } from './errors';
-import { VertexAIModel, GenerativeModel, ImagenModel } from './models';
+import {
+  VertexAIModel,
+  GenerativeModel,
+  ImagenModel,
+  HybridModel,
+  LocalModel
+} from './models';
 
 export { ChatSession } from './methods/chat-session';
 export * from './requests/schema-builder';
@@ -109,3 +115,14 @@ export function getImagenModel(
   }
   return new ImagenModel(vertexAI, modelParams, requestOptions);
 }
+
+export function getLocalModel(aiProvider: AI = window.ai): LocalModel {
+  return new LocalModel(aiProvider);
+}
+
+export function getHybridModel(
+  remoteModel: GenerativeModel,
+  localModel: LocalModel
+): HybridModel {
+  return new HybridModel(remoteModel, localModel);
+}

packages/vertexai/src/models/generative-model.ts

@@ -148,3 +148,74 @@ export class GenerativeModel extends VertexAIModel {
     return countTokens(this._apiSettings, this.model, formattedParams);
   }
 }
+
+interface ChatMethods {
+  sendMessage(
+    request: string | Array<string | Part>
+  ): Promise<GenerateContentResult>;
+}
+class HybridChat implements ChatMethods {
+  constructor(
+    private remoteModel: GenerativeModel,
+    private localModel: LocalModel
+  ) {}
+  async sendMessage(
+    request: string | Array<string | Part>
+  ): Promise<GenerateContentResult> {
+    if (await this.localModel.isSupported(request)) {
+      return this.localModel.generateContent(request);
+    }
+    return this.remoteModel.generateContent(request);
+  }
+}
+interface GenModelMethods {
+  generateContent(
+    request: GenerateContentRequest | string | Array<string | Part>
+  ): Promise<GenerateContentResult>;
+}
+/**
+ * Normalizes Chrome API, if available, to Vertex API
+ */
+export class LocalModel implements GenModelMethods {
+  constructor(private aiProvider?: AI) {}
+  async generateContent(
+    request: GenerateContentRequest | string | Array<string | Part>
+  ): Promise<GenerateContentResult> {
+    const session = await this.session();
+    if (typeof request !== 'string') {
+      throw new Error('unsupported request format');
+    }
+    const result = await session.prompt(request);
+    return {
+      response: {
+        text: () => result,
+        functionCalls: () => undefined
+      }
+    } as GenerateContentResult;
+  }
+  async isSupported(
+    request: string | Array<string | Part> | GenerateContentRequest
+  ): Promise<boolean> {
+    return typeof request === 'string';
+  }
+  private async session(): Promise<AILanguageModel> {
+    return this.aiProvider!.languageModel.create();
+  }
+}
+export class HybridModel implements GenModelMethods {
+  constructor(
+    private remoteModel: GenerativeModel,
+    private localModel: LocalModel
+  ) {}
+  async generateContent(
+    request: GenerateContentRequest | string | Array<string | Part>
+  ): Promise<GenerateContentResult> {
+    if (await this.localModel.isSupported(request)) {
+      return this.localModel.generateContent(request);
+    }
+    return this.remoteModel.generateContent(request);
+  }
+  startChat(): HybridChat {
+    return new HybridChat(this.remoteModel, this.localModel);
+  }
+}

yarn.lock

@@ -2949,6 +2949,11 @@
   resolved "https://registry.npmjs.org/@types/deep-eql/-/deep-eql-4.0.2.tgz#334311971d3a07121e7eb91b684a605e7eea9cbd"
   integrity sha512-c9h9dVVMigMPc4bwTvC5dxqtqJZwQPePsWjPlpSOnojbor6pGqdk541lfA7AqFQr5pB1BRdq0juY9db81BwyFw==
 
+"@types/[email protected]":
+  version "0.0.6"
+  resolved "https://registry.npmjs.org/@types/dom-chromium-ai/-/dom-chromium-ai-0.0.6.tgz#0c9e5712d8db3d26586cd9f175001b509cd2e514"
+  integrity sha512-/jUGe9a3BLzsjjg18Olk/Ul64PZ0P4aw8uNxrXeXVTni5PSxyCfyhHb4UohsXNVByOnwYGzlqUcb3vYKVsG4mg==
+
 "@types/eslint-scope@^3.7.7":
   version "3.7.7"
   resolved "https://registry.npmjs.org/@types/eslint-scope/-/eslint-scope-3.7.7.tgz#3108bd5f18b0cdb277c867b3dd449c9ed7079ac5"