Integration tests

Author: dlarocque

packages/vertexai/integration/chat.test.ts

@@ -1,6 +1,4 @@
-import { initializeApp } from "@firebase/app";
-import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, Modality, SafetySetting, getGenerativeModel, getVertexAI } from "../src";
-import { expect } from "chai";
+import { Content, GenerationConfig, HarmBlockMethod, HarmBlockThreshold, HarmCategory, SafetySetting } from "../src";
 
 // TODO (dlarocque): Use seperate Firebase config specifically for Vertex AI
 // TODO (dlarocque): Load this from environment variables, so we can set the config as a 
@@ -16,17 +14,17 @@ export const config = {
   measurementId: "G-1VL38N8YFE"
 };
 
-initializeApp(config);
-const MODEL_NAME = 'gemini-1.5-pro'; 
+export const MODEL_NAME = 'gemini-1.5-pro'; 
 
-let generationConfig: GenerationConfig = {
+/// TODO (dlarocque): Fix the naming on these.
+export const generationConfig: GenerationConfig = {
   temperature: 0,
   topP: 0,
   topK: 1,
   responseMimeType: 'text/plain'
 }
 
-let safetySettings: SafetySetting[] = [
+export const safetySettings: SafetySetting[] = [
   {
     category: HarmCategory.HARM_CATEGORY_HARASSMENT,
     threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
@@ -47,35 +45,11 @@ let safetySettings: SafetySetting[] = [
   },
 ];
 
-let systemInstruction: Content = {
+export const systemInstruction: Content = {
   role: 'system',
   parts: [
     {
       text: 'You are a friendly and helpful assistant.'
     }
   ]
-};
-
-describe('VertexAIService', () => {
-  it('CountTokens text', async () => {
-    const vertexAI = getVertexAI(); 
-    const model = getGenerativeModel(
-      vertexAI, 
-      {
-        model: MODEL_NAME, 
-        generationConfig,
-        systemInstruction,
-        safetySettings
-      }
-    );
-
-    let response = await model.countTokens('Why is the sky blue?');
-
-    expect(response.totalTokens).to.equal(6);
-    expect(response.totalBillableCharacters).to.equal(16);
-    expect(response.promptTokensDetails).to.not.be.null;
-    expect(response.promptTokensDetails!.length).to.equal(1);
-    expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
-    expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
-  });
-});
+};

packages/vertexai/integration/index.test.ts renamed to packages/vertexai/integration/constants.ts

@@ -0,0 +1,31 @@
+import { expect } from "chai";
+import { Modality, getGenerativeModel, getVertexAI } from "../src";
+import { MODEL_NAME, generationConfig, systemInstruction, safetySettings, config } from "./constants";
+import { initializeApp } from "@firebase/app";
+
+describe('Count Tokens', () => {
+
+  before(() => initializeApp(config))
+
+  it('CountTokens text', async () => {
+    const vertexAI = getVertexAI(); 
+    const model = getGenerativeModel(
+      vertexAI, 
+      {
+        model: MODEL_NAME, 
+        generationConfig,
+        systemInstruction,
+        safetySettings
+      }
+    );
+
+    let response = await model.countTokens('Why is the sky blue?');
+
+    expect(response.totalTokens).to.equal(6);
+    expect(response.totalBillableCharacters).to.equal(16);
+    expect(response.promptTokensDetails).to.not.be.null;
+    expect(response.promptTokensDetails!.length).to.equal(1);
+    expect(response.promptTokensDetails![0].modality).to.equal(Modality.TEXT);
+    expect(response.promptTokensDetails![0].tokenCount).to.equal(6);
+  });
+});

packages/vertexai/integration/count-tokens.test.ts

@@ -0,0 +1,37 @@
+import { expect } from "chai";
+import { getGenerativeModel, getVertexAI } from "../src";
+import { MODEL_NAME, generationConfig, systemInstruction, safetySettings, config } from "./constants";
+import { initializeApp } from "@firebase/app";
+
+// Token counts are only expected to differ by at most this number of tokens.
+// Set to 1 for whitespace that is not always present.
+const TOKEN_COUNT_DELTA = 1;
+
+describe('Generate Content', () => {
+
+  before(() => initializeApp(config))
+
+  it('generateContent', async () => {
+    const vertexAI = getVertexAI(); 
+    const model = getGenerativeModel(
+      vertexAI, 
+      {
+        model: MODEL_NAME, 
+        generationConfig,
+        systemInstruction,
+        safetySettings
+      }
+    );
+
+    const result = await model.generateContent("Where is Google headquarters located? Answer with the city name only.");
+    const response = result.response;
+    
+    const trimmedText = response.text().trim();
+    expect(trimmedText).to.equal('Mountain View');
+
+    console.log(JSON.stringify(response));
+
+    expect(response.usageMetadata).to.not.be.null;
+    expect(response.usageMetadata!.promptTokenCount).to.be.closeTo(21, TOKEN_COUNT_DELTA);
+  });
+});