Fix types to match Vertex API and convert functionCall() to functionCalls() (#284)

Author: hsubox76

packages/vertexai/package.json

@@ -36,12 +36,10 @@
     "build": "rollup -c",
     "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",
-    "test:browser": "karma start --single-run",
-    "test:node": "TS_NODE_COMPILER_OPTIONS='{\"module\":\"commonjs\"}' nyc --reporter lcovonly -- mocha src/**/*.test.* --config ../../config/mocharc.node.js"
+    "testsetup": "yarn ts-node ./test-utils/convert-mocks.ts",
+    "test": "run-p --npm-path npm lint test:browser",
+    "test:ci": "yarn testsetup && node ../../scripts/run_tests_in_ci.js -s test",
+    "test:browser": "yarn testsetup && karma start --single-run"
   },
   "peerDependencies": {
     "@firebase/app": "0.x",

packages/vertexai/rollup.config.js

@@ -52,7 +52,10 @@ const browserBuilds = [
     output: [{ file: pkg.esm5, format: 'es', sourcemap: true }],
     plugins: [
       ...es5BuildPlugins,
-      replace(generateBuildTargetReplaceConfig('esm', 5))
+      replace({
+        ...generateBuildTargetReplaceConfig('esm', 5),
+        __PACKAGE_VERSION__: pkg.version
+      })
     ],
     external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   },
@@ -65,7 +68,10 @@ const browserBuilds = [
     },
     plugins: [
       ...es2017BuildPlugins,
-      replace(generateBuildTargetReplaceConfig('esm', 2017)),
+      replace({
+        ...generateBuildTargetReplaceConfig('esm', 2017),
+        __PACKAGE_VERSION__: pkg.version
+      }),
       emitModulePackageFile()
     ],
     external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
@@ -79,7 +85,10 @@ const browserBuilds = [
     },
     plugins: [
       ...es2017BuildPlugins,
-      replace(generateBuildTargetReplaceConfig('cjs', 2017))
+      replace({
+        ...generateBuildTargetReplaceConfig('cjs', 2017),
+        __PACKAGE_VERSION__: pkg.version
+      })
     ],
     external: id => deps.some(dep => id === dep || id.startsWith(`${dep}/`))
   }

packages/vertexai/src/methods/count-tokens.ts

@@ -29,10 +29,16 @@ export async function countTokens(
   params: CountTokensRequest,
   requestOptions?: RequestOptions
 ): Promise<CountTokensResponse> {
-  const url = new RequestUrl(model, Task.COUNT_TOKENS, apiSettings, false, {});
+  const url = new RequestUrl(
+    model,
+    Task.COUNT_TOKENS,
+    apiSettings,
+    false,
+    requestOptions
+  );
   const response = await makeRequest(
     url,
-    JSON.stringify({ ...params, model }),
+    JSON.stringify(params),
     requestOptions
   );
   return response.json();

packages/vertexai/src/methods/embed-content.ts

@@ -104,7 +104,7 @@ describe('generateContent()', () => {
     );
     expect(result.response.text()).to.include('Quantum mechanics is');
     expect(
-      result.response.candidates?.[0].citationMetadata?.citationSources.length
+      result.response.candidates?.[0].citationMetadata?.citations.length
     ).to.equal(1);
     expect(makeRequestStub).to.be.calledWith(
       match.instanceOf(request.RequestUrl),

packages/vertexai/src/methods/generate-content.test.ts

@@ -20,12 +20,8 @@ import {
   generateContentStream
 } from '../methods/generate-content';
 import {
-  BatchEmbedContentsRequest,
-  BatchEmbedContentsResponse,
   CountTokensRequest,
   CountTokensResponse,
-  EmbedContentRequest,
-  EmbedContentResponse,
   GenerateContentRequest,
   GenerateContentResult,
   GenerateContentStreamResult,
@@ -39,11 +35,7 @@ import {
 } from '../types';
 import { ChatSession } from '../methods/chat-session';
 import { countTokens } from '../methods/count-tokens';
-import { batchEmbedContents, embedContent } from '../methods/embed-content';
-import {
-  formatEmbedContentInput,
-  formatGenerateContentInput
-} from '../requests/request-helpers';
+import { formatGenerateContentInput } from '../requests/request-helpers';
 import { Vertex } from '../public-types';
 import { ERROR_FACTORY, VertexError } from '../errors';
 import { ApiSettings } from '../types/internal';
@@ -163,28 +155,4 @@ export class GenerativeModel {
     const formattedParams = formatGenerateContentInput(request);
     return countTokens(this._apiSettings, this.model, formattedParams);
   }
-
-  /**
-   * Embeds the provided content.
-   */
-  async embedContent(
-    request: EmbedContentRequest | string | Array<string | Part>
-  ): Promise<EmbedContentResponse> {
-    const formattedParams = formatEmbedContentInput(request);
-    return embedContent(this._apiSettings, this.model, formattedParams);
-  }
-
-  /**
-   * Embeds an array of {@link EmbedContentRequest}s.
-   */
-  async batchEmbedContents(
-    batchEmbedContentRequest: BatchEmbedContentsRequest
-  ): Promise<BatchEmbedContentsResponse> {
-    return batchEmbedContents(
-      this._apiSettings,
-      this.model,
-      batchEmbedContentRequest,
-      this.requestOptions
-    );
-  }
 }

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

@@ -15,12 +15,7 @@
  * limitations under the License.
  */
 
-import {
-  Content,
-  EmbedContentRequest,
-  GenerateContentRequest,
-  Part
-} from '../types';
+import { Content, GenerateContentRequest, Part } from '../types';
 import { ERROR_FACTORY, VertexError } from '../errors';
 
 export function formatNewContent(
@@ -96,13 +91,3 @@ export function formatGenerateContentInput(
     return { contents: [content] };
   }
 }
-
-export function formatEmbedContentInput(
-  params: EmbedContentRequest | string | Array<string | Part>
-): EmbedContentRequest {
-  if (typeof params === 'string' || Array.isArray(params)) {
-    const content = formatNewContent(params);
-    return { content };
-  }
-  return params;
-}

packages/vertexai/src/requests/request-helpers.ts

@@ -90,6 +90,16 @@ describe('request methods', () => {
         '/v100omega/projects/my-project/locations/us-central1/models/model-name'
       );
     });
+    it('custom baseUrl', async () => {
+      const url = new RequestUrl(
+        'models/model-name',
+        Task.GENERATE_CONTENT,
+        fakeApiSettings,
+        false,
+        { baseUrl: 'https://my.special.endpoint' }
+      );
+      expect(url.toString()).to.include('https://my.special.endpoint');
+    });
     it('non-stream - tunedModels/', async () => {
       const url = new RequestUrl(
         'tunedModels/model-name',

packages/vertexai/src/requests/request.test.ts

@@ -18,24 +18,19 @@
 import { RequestOptions } from '../types';
 import { ERROR_FACTORY, VertexError } from '../errors';
 import { ApiSettings } from '../types/internal';
+import { version } from '../../package.json';
 
-const BASE_URL = 'https://staging-firebaseml.sandbox.googleapis.com';
+const DEFAULT_BASE_URL = 'https://firebaseml.googleapis.com';
 
 export const DEFAULT_API_VERSION = 'v2beta';
 
-/**
- * We can't `require` package.json if this runs on web. We will use rollup to
- * swap in the version number here at build time.
- */
-const PACKAGE_VERSION = '__PACKAGE_VERSION__';
+const PACKAGE_VERSION = version;
 const PACKAGE_LOG_HEADER = 'firebase-vertexai-js';
 
 export enum Task {
   GENERATE_CONTENT = 'generateContent',
   STREAM_GENERATE_CONTENT = 'streamGenerateContent',
-  COUNT_TOKENS = 'countTokens',
-  EMBED_CONTENT = 'embedContent',
-  BATCH_EMBED_CONTENTS = 'batchEmbedContents'
+  COUNT_TOKENS = 'countTokens'
 }
 
 export class RequestUrl {
@@ -48,7 +43,8 @@ export class RequestUrl {
   ) {}
   toString(): string {
     const apiVersion = this.requestOptions?.apiVersion || DEFAULT_API_VERSION;
-    let url = `${BASE_URL}/${apiVersion}`;
+    const baseUrl = this.requestOptions?.baseUrl || DEFAULT_BASE_URL;
+    let url = `${baseUrl}/${apiVersion}`;
     url += `/projects/${this.apiSettings.project}`;
     url += `/locations/${this.apiSettings.location}`;
     url += `/${this.model}`;
@@ -58,6 +54,17 @@ export class RequestUrl {
     }
     return url;
   }
+
+  /**
+   * If the model needs to be passed to the backend, it needs to
+   * include project and location path.
+   */
+  get fullModelString(): string {
+    let modelString = `projects/${this.apiSettings.project}`;
+    modelString += `/locations/${this.apiSettings.location}`;
+    modelString += `/${this.model}`;
+    return modelString;
+  }
 }
 
 /**

packages/vertexai/src/requests/request.ts

@@ -39,6 +39,7 @@ const fakeResponseText: GenerateContentResponse = {
     }
   ]
 };
+
 const fakeResponseFunctionCall: GenerateContentResponse = {
   candidates: [
     {
@@ -61,6 +62,38 @@ const fakeResponseFunctionCall: GenerateContentResponse = {
   ]
 };
 
+const fakeResponseFunctionCalls: GenerateContentResponse = {
+  candidates: [
+    {
+      index: 0,
+      content: {
+        role: 'model',
+        parts: [
+          {
+            functionCall: {
+              name: 'find_theaters',
+              args: {
+                location: 'Mountain View, CA',
+                movie: 'Barbie'
+              }
+            }
+          },
+          {
+            functionCall: {
+              name: 'find_times',
+              args: {
+                location: 'Mountain View, CA',
+                movie: 'Barbie',
+                time: '20:00'
+              }
+            }
+          }
+        ]
+      }
+    }
+  ]
+};
+
 const badFakeResponse: GenerateContentResponse = {
   promptFeedback: {
     blockReason: BlockReason.SAFETY,
@@ -79,9 +112,16 @@ describe('response-helpers methods', () => {
     });
     it('good response functionCall', async () => {
       const enhancedResponse = addHelpers(fakeResponseFunctionCall);
-      expect(enhancedResponse.functionCall()).to.deep.equal(
+      expect(enhancedResponse.functionCalls()).to.deep.equal([
         fakeResponseFunctionCall.candidates?.[0].content.parts[0].functionCall
-      );
+      ]);
+    });
+    it('good response functionCalls', async () => {
+      const enhancedResponse = addHelpers(fakeResponseFunctionCalls);
+      expect(enhancedResponse.functionCalls()).to.deep.equal([
+        fakeResponseFunctionCalls.candidates?.[0].content.parts[0].functionCall,
+        fakeResponseFunctionCalls.candidates?.[0].content.parts[1].functionCall
+      ]);
     });
     it('bad response safety', async () => {
       const enhancedResponse = addHelpers(badFakeResponse);