Add unit and integration tests

Author: dlarocque

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

@@ -113,6 +113,46 @@ describe('Generate Content', () => {
         ).to.be.closeTo(4, TOKEN_COUNT_DELTA);
       });
 
+      it('generateContent: google search grounding', async () => {
+        const model = getGenerativeModel(testConfig.ai, {
+          model: testConfig.model,
+          generationConfig: commonGenerationConfig,
+          safetySettings: commonSafetySettings,
+          tools: [{ googleSearch: {} }]
+        });
+
+        const result = await model.generateContent(
+          'What is the speed of light in a vaccuum in meters per second?'
+        );
+        const response = result.response;
+        const trimmedText = response.text().trim();
+        const groundingMetadata = response.candidates?.[0].groundingMetadata;
+        expect(trimmedText).to.contain('299,792,458');
+        expect(groundingMetadata).to.exist;
+        expect(groundingMetadata!.searchEntryPoint?.renderedContent).to.contain(
+          'div'
+        );
+        expect(
+          groundingMetadata!.groundingChunks
+        ).to.have.length.greaterThanOrEqual(1);
+        groundingMetadata!.groundingChunks!.forEach(groundingChunk => {
+          expect(groundingChunk.web).to.exist;
+          expect(groundingChunk.web!.uri).to.exist;
+        });
+        expect(
+          groundingMetadata?.groundingSupports
+        ).to.have.length.greaterThanOrEqual(1);
+        groundingMetadata!.groundingSupports!.forEach(groundingSupport => {
+          expect(
+            groundingSupport.groundingChunkIndices
+          ).to.have.length.greaterThanOrEqual(1);
+          expect(groundingSupport.segment).to.exist;
+          expect(groundingSupport.segment?.endIndex).to.exist;
+          expect(groundingSupport.segment?.text).to.exist;
+          // Since partIndex and startIndex are commonly 0, they may be omitted from responses.
+        });
+      });
+
       it('generateContentStream: text input, text output', async () => {
         const model = getGenerativeModel(testConfig.ai, {
           model: testConfig.model,

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

@@ -193,6 +193,53 @@ describe('generateContent()', () => {
       match.any
     );
   });
+  it('google search grounding', async () => {
+    const mockResponse = getMockResponse(
+      'vertexAI',
+      'unary-success-google-search-grounding.json'
+    );
+    const makeRequestStub = stub(request, 'makeRequest').resolves(
+      mockResponse as Response
+    );
+    const result = await generateContent(
+      fakeApiSettings,
+      'model',
+      fakeRequestParams
+    );
+    expect(result.response.text()).to.include('The temperature is 67°F (19°C)');
+    const groundingMetadata = result.response.candidates?.[0].groundingMetadata;
+    expect(groundingMetadata).to.not.be.undefined;
+    expect(groundingMetadata!.searchEntryPoint?.renderedContent).to.contain(
+      'div'
+    );
+    expect(groundingMetadata!.groundingChunks?.length).to.equal(2);
+    expect(groundingMetadata!.groundingChunks?.[0].web?.uri).to.contain(
+      'https://vertexaisearch.cloud.google.com'
+    );
+    expect(groundingMetadata!.groundingChunks?.[0].web?.title).to.equal(
+      'accuweather.com'
+    );
+    expect(groundingMetadata!.groundingSupports?.length).to.equal(3);
+    expect(
+      groundingMetadata!.groundingSupports?.[0].groundingChunkIndices
+    ).to.deep.equal([0]);
+    expect(groundingMetadata!.groundingSupports?.[0].segment).to.deep.equal({
+      endIndex: 56,
+      text: 'The current weather in London, United Kingdom is cloudy.'
+    });
+    expect(groundingMetadata!.groundingSupports?.[0].segment?.partIndex).to.be
+      .undefined;
+    expect(groundingMetadata!.groundingSupports?.[0].segment?.startIndex).to.be
+      .undefined;
+
+    expect(makeRequestStub).to.be.calledWith(
+      'model',
+      Task.GENERATE_CONTENT,
+      fakeApiSettings,
+      false,
+      match.any
+    );
+  });
   it('blocked prompt', async () => {
     const mockResponse = getMockResponse(
       'vertexAI',