test(ai): update unit test for firebase ai

Author: russellwheatley

packages/ai/__tests__/api.test.ts

@@ -15,20 +15,19 @@
  * limitations under the License.
  */
 import { describe, expect, it } from '@jest/globals';
-import { getApp, type ReactNativeFirebase } from '../../app/lib';
+import { type ReactNativeFirebase } from '../../app/lib';
 
-import { ModelParams, VertexAIErrorCode } from '../lib/types';
-import { VertexAIError } from '../lib/errors';
-import { getGenerativeModel, getVertexAI } from '../lib/index';
+import { ModelParams, AIErrorCode } from '../lib/types';
+import { AIError } from '../lib/errors';
+import { getGenerativeModel } from '../lib/index';
 
-import { VertexAI } from '../lib/public-types';
+import { AI } from '../lib/public-types';
 import { GenerativeModel } from '../lib/models/generative-model';
 
-import '../../auth/lib';
-import '../../app-check/lib';
-import { getAuth } from '../../auth/lib';
+import { AI_TYPE } from '../lib/constants';
+import { VertexAIBackend } from '../lib/backend';
 
-const fakeVertexAI: VertexAI = {
+const fakeAI: AI = {
   app: {
     name: 'DEFAULT',
     options: {
@@ -37,66 +36,76 @@ const fakeVertexAI: VertexAI = {
       projectId: 'my-project',
     },
   } as ReactNativeFirebase.FirebaseApp,
+  backend: new VertexAIBackend('us-central1'),
   location: 'us-central1',
 };
 
 describe('Top level API', () => {
-  it('should allow auth and app check instances to be passed in', () => {
-    const app = getApp();
-    const auth = getAuth();
-    const appCheck = app.appCheck();
-
-    getVertexAI(app, { appCheck, auth });
-  });
-
   it('getGenerativeModel throws if no model is provided', () => {
     try {
-      getGenerativeModel(fakeVertexAI, {} as ModelParams);
+      getGenerativeModel(fakeAI, {} as ModelParams);
     } catch (e) {
-      expect((e as VertexAIError).code).toContain(VertexAIErrorCode.NO_MODEL);
-      expect((e as VertexAIError).message).toContain(
+      expect((e as AIError).code).toContain(AIErrorCode.NO_MODEL);
+      expect((e as AIError).message).toContain(
         `VertexAI: Must provide a model name. Example: ` +
-          `getGenerativeModel({ model: 'my-model-name' }) (vertexAI/${VertexAIErrorCode.NO_MODEL})`,
+          `getGenerativeModel({ model: 'my-model-name' }) (vertexAI/${AIErrorCode.NO_MODEL})`,
       );
     }
   });
 
   it('getGenerativeModel throws if no apiKey is provided', () => {
     const fakeVertexNoApiKey = {
-      ...fakeVertexAI,
-      app: { options: { projectId: 'my-project' } },
-    } as VertexAI;
+      ...fakeAI,
+      app: { options: { projectId: 'my-project', appId: 'my-appid' } },
+    } as AI;
     try {
       getGenerativeModel(fakeVertexNoApiKey, { model: 'my-model' });
     } catch (e) {
-      expect((e as VertexAIError).code).toContain(VertexAIErrorCode.NO_API_KEY);
-      expect((e as VertexAIError).message).toBe(
+      expect((e as AIError).code).toContain(AIErrorCode.NO_API_KEY);
+      expect((e as AIError).message).toBe(
         `VertexAI: The "apiKey" field is empty in the local ` +
           `Firebase config. Firebase VertexAI requires this field to` +
-          ` contain a valid API key. (vertexAI/${VertexAIErrorCode.NO_API_KEY})`,
+          ` contain a valid API key. (vertexAI/${AIErrorCode.NO_API_KEY})`,
       );
     }
   });
 
   it('getGenerativeModel throws if no projectId is provided', () => {
     const fakeVertexNoProject = {
-      ...fakeVertexAI,
+      ...fakeAI,
       app: { options: { apiKey: 'my-key' } },
-    } as VertexAI;
+    } as AI;
     try {
       getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
     } catch (e) {
-      expect((e as VertexAIError).code).toContain(VertexAIErrorCode.NO_PROJECT_ID);
-      expect((e as VertexAIError).message).toBe(
+      expect((e as AIError).code).toContain(AIErrorCode.NO_PROJECT_ID);
+      expect((e as AIError).message).toBe(
         `VertexAI: The "projectId" field is empty in the local` +
           ` Firebase config. Firebase VertexAI requires this field ` +
-          `to contain a valid project ID. (vertexAI/${VertexAIErrorCode.NO_PROJECT_ID})`,
+          `to contain a valid project ID. (vertexAI/${AIErrorCode.NO_PROJECT_ID})`,
+      );
+    }
+  });
+
+  it('getGenerativeModel throws if no appId is provided', () => {
+    const fakeVertexNoProject = {
+      ...fakeAI,
+      app: { options: { apiKey: 'my-key', projectId: 'my-projectid' } },
+    } as AI;
+    try {
+      getGenerativeModel(fakeVertexNoProject, { model: 'my-model' });
+    } catch (e) {
+      expect((e as AIError).code).toContain(AIErrorCode.NO_APP_ID);
+      expect((e as AIError).message).toBe(
+        `AI: The "appId" field is empty in the local` +
+          ` Firebase config. Firebase AI requires this field ` +
+          `to contain a valid app ID. (${AI_TYPE}/${AIErrorCode.NO_APP_ID})`,
       );
     }
   });
 
   it('getGenerativeModel gets a GenerativeModel', () => {
-    const genModel = getGenerativeModel(fakeVertexAI, { model: 'my-model' });
+    const genModel = getGenerativeModel(fakeAI, { model: 'my-model' });
     expect(genModel).toBeInstanceOf(GenerativeModel);
     expect(genModel.model).toBe('publishers/google/models/my-model');
   });

packages/ai/__tests__/backend.test.ts

@@ -16,8 +16,9 @@
  */
 import { describe, it, expect } from '@jest/globals';
 import { GoogleAIBackend, VertexAIBackend } from '../lib/backend';
-import { BackendType } from 'lib/public-types';
-import { DEFAULT_LOCATION } from 'lib/constants';
+import { BackendType } from '../lib/public-types';
+import { DEFAULT_LOCATION } from '../lib/constants';
+
 describe('Backend', () => {
   describe('GoogleAIBackend', () => {
     it('sets backendType to GOOGLE_AI', () => {

packages/ai/__tests__/chat-session.test.ts

@@ -21,11 +21,14 @@ import { GenerateContentStreamResult } from '../lib/types';
 import { ChatSession } from '../lib/methods/chat-session';
 import { ApiSettings } from '../lib/types/internal';
 import { RequestOptions } from '../lib/types/requests';
+import { VertexAIBackend } from '../lib/backend';
 
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1',
+  backend: new VertexAIBackend(),
 };
 
 const requestOptions: RequestOptions = {

packages/ai/__tests__/count-tokens.test.ts

@@ -21,11 +21,14 @@ import { countTokens } from '../lib/methods/count-tokens';
 import { CountTokensRequest } from '../lib/types';
 import { ApiSettings } from '../lib/types/internal';
 import { Task } from '../lib/requests/request';
+import { GoogleAIBackend } from '../lib/backend';
 
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
   location: 'us-central1',
+  appId: '',
+  backend: new GoogleAIBackend(),
 };
 
 const fakeRequestParams: CountTokensRequest = {

packages/ai/__tests__/generate-content.test.ts

@@ -14,7 +14,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { describe, expect, it, afterEach, jest } from '@jest/globals';
+import { describe, expect, it, afterEach, jest, beforeEach } from '@jest/globals';
 import { getMockResponse } from './test-utils/mock-response';
 import * as request from '../lib/requests/request';
 import { generateContent } from '../lib/methods/generate-content';
@@ -27,11 +27,25 @@ import {
 } from '../lib/types';
 import { ApiSettings } from '../lib/types/internal';
 import { Task } from '../lib/requests/request';
+import { GoogleAIBackend, VertexAIBackend } from '../lib/backend';
+import { SpiedFunction } from 'jest-mock';
+import { AIError } from '../lib/errors';
+import { mapGenerateContentRequest } from '../lib/googleai-mappers';
 
 const fakeApiSettings: ApiSettings = {
   apiKey: 'key',
   project: 'my-project',
+  appId: 'my-appid',
   location: 'us-central1',
+  backend: new VertexAIBackend(),
+};
+
+const fakeGoogleAIApiSettings: ApiSettings = {
+  apiKey: 'key',
+  project: 'my-project',
+  appId: 'my-appid',
+  location: 'us-central1',
+  backend: new GoogleAIBackend(),
 };
 
 const fakeRequestParams: GenerateContentRequest = {
@@ -48,6 +62,19 @@ const fakeRequestParams: GenerateContentRequest = {
   ],
 };
 
+const fakeGoogleAIRequestParams: GenerateContentRequest = {
+  contents: [{ parts: [{ text: 'hello' }], role: 'user' }],
+  generationConfig: {
+    topK: 16,
+  },
+  safetySettings: [
+    {
+      category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+      threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+    },
+  ],
+};
+
 describe('generateContent()', () => {
   afterEach(() => {
     jest.restoreAllMocks();
@@ -88,6 +115,28 @@ describe('generateContent()', () => {
     );
   });
 
+  it('long response with token details', async () => {
+    const mockResponse = getMockResponse('unary-success-basic-response-long-usage-metadata.json');
+    const makeRequestStub = jest
+      .spyOn(request, 'makeRequest')
+      .mockResolvedValue(mockResponse as Response);
+    const result = await generateContent(fakeApiSettings, 'model', fakeRequestParams);
+    expect(result.response.usageMetadata?.totalTokenCount).toEqual(1913);
+    expect(result.response.usageMetadata?.candidatesTokenCount).toEqual(76);
+    expect(result.response.usageMetadata?.promptTokensDetails?.[0]?.modality).toEqual('IMAGE');
+    expect(result.response.usageMetadata?.promptTokensDetails?.[0]?.tokenCount).toEqual(1806);
+    expect(result.response.usageMetadata?.candidatesTokensDetails?.[0]?.modality).toEqual('TEXT');
+    expect(result.response.usageMetadata?.candidatesTokensDetails?.[0]?.tokenCount).toEqual(76);
+    expect(makeRequestStub).toHaveBeenCalledWith(
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
+      expect.anything(),
+      undefined,
+    );
+  });
+
   it('citations', async () => {
     const mockResponse = getMockResponse('unary-success-citations.json');
     const makeRequestStub = jest
@@ -201,4 +250,54 @@ describe('generateContent()', () => {
     );
     expect(mockFetch).toHaveBeenCalled();
   });
+
+  describe('googleAI', () => {
+    let makeRequestStub: SpiedFunction<typeof request.makeRequest>;
+
+    beforeEach(() => {
+      makeRequestStub = jest.spyOn(request, 'makeRequest');
+    });
+
+    afterEach(() => {
+      jest.restoreAllMocks();
+    });
+
+    it('throws error when method is defined', async () => {
+      const mockResponse = getMockResponse('unary-success-basic-reply-short.txt');
+      makeRequestStub.mockResolvedValue(mockResponse as Response);
+
+      const requestParamsWithMethod: GenerateContentRequest = {
+        contents: [{ parts: [{ text: 'hello' }], role: 'user' }],
+        safetySettings: [
+          {
+            category: HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
+            threshold: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+            method: HarmBlockMethod.SEVERITY, // Unsupported in Google AI.
+          },
+        ],
+      };
+
+      // Expect generateContent to throw a AIError that method is not supported.
+      await expect(
+        generateContent(fakeGoogleAIApiSettings, 'model', requestParamsWithMethod),
+      ).rejects.toThrow(AIError);
+      expect(makeRequestStub).not.toHaveBeenCalled();
+    });
+
+    it('maps request to GoogleAI format', async () => {
+      const mockResponse = getMockResponse('unary-success-basic-reply-short.txt');
+      makeRequestStub.mockResolvedValue(mockResponse as Response);
+
+      await generateContent(fakeGoogleAIApiSettings, 'model', fakeGoogleAIRequestParams);
+
+      expect(makeRequestStub).toHaveBeenCalledWith(
+        'model',
+        Task.GENERATE_CONTENT,
+        fakeGoogleAIApiSettings,
+        false,
+        JSON.stringify(mapGenerateContentRequest(fakeGoogleAIRequestParams)),
+        undefined,
+      );
+    });
+  });
 });