Port code from GoogleAI repo (#283)

Author: hsubox76

.gitignore

@@ -98,3 +98,6 @@ toc/
 .terraform.lock.hcl
 *.tfstate
 *.tfstate.*
+
+# generated test case text data
+mocks-lookup.ts

packages/vertexai/README.md

@@ -1,4 +1,5 @@
 # @firebase/vertexai
 
+This is the Firebase Vertex AI component of the Firebase JS SDK.
 
-  
+**This package is not intended for direct usage, and should only be used via the officially supported [firebase](https://www.npmjs.com/package/firebase) package.**

packages/vertexai/package.json

@@ -34,8 +34,9 @@
     "lint": "eslint -c .eslintrc.js '**/*.ts' --ignore-path '../../.gitignore'",
     "lint:fix": "eslint --fix -c .eslintrc.js '**/*.ts' --ignore-path '../../.gitignore'",
     "build": "rollup -c",
-    "build:deps": "lerna run --scope @firebase/template --include-dependencies build",
+    "build:deps": "lerna run --scope @firebase/vertexai --include-dependencies build",
     "dev": "rollup -c -w",
+    "pretest": "yarn ts-node ./test-utils/convert-mocks.ts",
     "test": "run-p --npm-path npm lint test:all",
     "test:ci": "node ../../scripts/run_tests_in_ci.js -s test:all",
     "test:all": "run-p --npm-path npm test:browser test:node",
@@ -55,7 +56,7 @@
   },
   "license": "Apache-2.0",
   "devDependencies": {
-    "@firebase/app": "0.9.28",
+    "@firebase/app": "0.9.29",
     "@rollup/plugin-json": "4.1.0",
     "rollup": "2.79.1",
     "rollup-plugin-replace": "2.2.0",

packages/vertexai/src/api.test.ts

@@ -0,0 +1,65 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { ModelParams } from './types';
+import { getGenerativeModel } from './api';
+import { expect } from 'chai';
+import { Vertex } from './public-types';
+import { GenerativeModel } from './models/generative-model';
+import { VertexError } from './errors';
+
+const fakeVertex: Vertex = {
+  app: {
+    name: 'DEFAULT',
+    automaticDataCollectionEnabled: true,
+    options: {
+      apiKey: 'key',
+      projectId: 'my-project'
+    }
+  },
+  region: 'us-central1'
+};
+
+describe('Top level API', () => {
+  it('getGenerativeModel throws if no model is provided', () => {
+    expect(() => getGenerativeModel(fakeVertex, {} as ModelParams)).to.throw(
+      VertexError.NO_MODEL
+    );
+  });
+  it('getGenerativeModel throws if no apiKey is provided', () => {
+    const fakeVertexNoApiKey = {
+      ...fakeVertex,
+      app: { options: { projectId: 'my-project' } }
+    } as Vertex;
+    expect(() =>
+      getGenerativeModel(fakeVertexNoApiKey, { model: 'my-model' })
+    ).to.throw(VertexError.NO_API_KEY);
+  });
+  it('getGenerativeModel throws if no projectId is provided', () => {
+    const fakeVertexNoProject = {
+      ...fakeVertex,
+      app: { options: { apiKey: 'my-key' } }
+    } as Vertex;
+    expect(() =>
+      getGenerativeModel(fakeVertexNoProject, { model: 'my-model' })
+    ).to.throw(VertexError.NO_PROJECT_ID);
+  });
+  it('getGenerativeModel gets a GenerativeModel', () => {
+    const genModel = getGenerativeModel(fakeVertex, { model: 'my-model' });
+    expect(genModel).to.be.an.instanceOf(GenerativeModel);
+    expect(genModel.model).to.equal('publishers/google/models/my-model');
+  });
+});

packages/vertexai/src/api.ts

@@ -21,6 +21,13 @@ import { getModularInstance } from '@firebase/util';
 import { VERTEX_TYPE } from './constants';
 import { VertexService } from './factory';
 import { Vertex } from './public-types';
+import { ERROR_FACTORY, VertexError } from './errors';
+import { ModelParams, RequestOptions } from './types';
+import { GenerativeModel } from './models/generative-model';
+
+export { ChatSession } from './methods/chat-session';
+
+export { GenerativeModel };
 
 declare module '@firebase/component' {
   interface NameServiceMapping {
@@ -47,24 +54,13 @@ export function getVertex(app: FirebaseApp = getApp()): Vertex {
   return vertexProvider.initialize();
 }
 
-export function getGenerativeModel(vertex: Vertex): GenerativeModel {
-  return new GenerativeModel(vertex, {});
-}
-
-// Just a stub
-class GenerativeModel {
-  private _apiKey?: string;
-  constructor(vertex: Vertex, modelParams: {}) {
-    if (!vertex.app.options.apiKey) {
-      // throw error
-    } else {
-      this._apiKey = vertex.app.options.apiKey;
-      //TODO: remove when we use this
-      console.log(this._apiKey);
-    }
-    //TODO: do something with modelParams
-    console.log(modelParams);
+export function getGenerativeModel(
+  vertex: Vertex,
+  modelParams: ModelParams,
+  requestOptions?: RequestOptions
+): GenerativeModel {
+  if (!modelParams.model) {
+    throw ERROR_FACTORY.create(VertexError.NO_MODEL);
   }
+  return new GenerativeModel(vertex, modelParams, requestOptions);
 }
-
-//TODO: add all top-level exportable methods and classes

packages/vertexai/src/constants.ts

@@ -16,3 +16,5 @@
  */
 
 export const VERTEX_TYPE = 'vertex';
+
+export const DEFAULT_REGION = 'us-central1';

packages/vertexai/src/errors.ts

@@ -0,0 +1,63 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { ErrorFactory, ErrorMap } from '@firebase/util';
+import { GenerateContentResponse } from './types';
+
+export const enum VertexError {
+  FETCH_ERROR = 'fetch-error',
+  INVALID_CONTENT = 'invalid-content',
+  NO_API_KEY = 'no-api-key',
+  NO_MODEL = 'no-model',
+  NO_PROJECT_ID = 'no-project-id',
+  PARSE_FAILED = 'parse-failed',
+  RESPONSE_ERROR = 'response-error'
+}
+
+const ERRORS: ErrorMap<VertexError> = {
+  [VertexError.FETCH_ERROR]: `Error fetching from {$url}: {$message}`,
+  [VertexError.INVALID_CONTENT]: `Content formatting error: {$message}`,
+  [VertexError.NO_API_KEY]:
+    `The "apiKey" field is empty in the local Firebase config. Firebase VertexAI requires this field to` +
+    `contain a valid API key.`,
+  [VertexError.NO_PROJECT_ID]:
+    `The "projectId" field is empty in the local Firebase config. Firebase VertexAI requires this field to` +
+    `contain a valid project ID.`,
+  [VertexError.NO_MODEL]:
+    `Must provide a model name. ` +
+    `Example: genai.getGenerativeModel({ model: 'my-model-name' })`,
+  [VertexError.PARSE_FAILED]: `Parsing failed: {$message}`,
+  [VertexError.RESPONSE_ERROR]:
+    `Response error: {$message}. Response body stored in ` +
+    `error.customData.response`
+};
+
+interface ErrorParams {
+  [VertexError.FETCH_ERROR]: { url: string; message: string };
+  [VertexError.INVALID_CONTENT]: { message: string };
+  [VertexError.PARSE_FAILED]: { message: string };
+  [VertexError.RESPONSE_ERROR]: {
+    message: string;
+    response: GenerateContentResponse;
+  };
+}
+
+export const ERROR_FACTORY = new ErrorFactory<VertexError, ErrorParams>(
+  'vertex',
+  'Vertex',
+  ERRORS
+);

packages/vertexai/src/factory.ts

@@ -22,6 +22,7 @@ import {
   FirebaseAppCheckInternal
 } from '@firebase/app-check-interop-types';
 import { Provider } from '@firebase/component';
+import { DEFAULT_REGION } from './constants';
 
 export function factory(
   app: FirebaseApp,
@@ -33,6 +34,7 @@ export function factory(
 
 export class VertexService implements Vertex, _FirebaseService {
   appCheck: FirebaseAppCheckInternal | null;
+  region: string;
 
   constructor(
     public app: FirebaseApp,
@@ -41,10 +43,7 @@ export class VertexService implements Vertex, _FirebaseService {
   ) {
     const appCheck = appCheckProvider?.getImmediate({ optional: true });
     this.appCheck = appCheck || null;
-    // This needs to go into the url
-    if (this.options?.region) {
-      console.log(this.options?.region);
-    }
+    this.region = this.options?.region || DEFAULT_REGION;
   }
 
   _delete(): Promise<void> {