fix: pass tools from getGenerativeModel to chatSession.

PiperOrigin-RevId: 613294961

Author: yyyu-google

src/models/generative_models.ts

@@ -304,6 +304,7 @@ export class GenerativeModelPreview {
       location: this.location,
       googleAuth: this.googleAuth,
       publisher_model_endpoint: this.publisherModelEndpoint,
+      tools: this.tools,
     };
 
     if (request) {

system_test/end_to_end_sample_test.ts

@@ -434,10 +434,10 @@ describe('generateContent', () => {
     });
     const response = result.response;
     const groundingMetadata = response.candidates[0].groundingMetadata;
-    expect(groundingMetadata).toBeDefined();
+    expect(!!groundingMetadata).toBeTruthy();
     if (groundingMetadata) {
-      expect(groundingMetadata.groundingAttributions).toBeTruthy();
-      expect(groundingMetadata.webSearchQueries).toBeTruthy();
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
     }
   });
   it('should return grounding metadata when passed GoogleSearchRetriever in generateContent', async () => {
@@ -450,10 +450,10 @@ describe('generateContent', () => {
     });
     const response = result.response;
     const groundingMetadata = response.candidates[0].groundingMetadata;
-    expect(groundingMetadata).toBeDefined();
+    expect(!!groundingMetadata).toBeTruthy();
     if (groundingMetadata) {
-      expect(groundingMetadata.groundingAttributions).toBeTruthy();
-      expect(groundingMetadata.webSearchQueries).toBeTruthy();
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
     }
   });
   it('in preview should return grounding metadata when passed GoogleSearchRetriever in getGenerativeModel', async () => {
@@ -466,10 +466,10 @@ describe('generateContent', () => {
     });
     const response = result.response;
     const groundingMetadata = response.candidates[0].groundingMetadata;
-    expect(groundingMetadata).toBeDefined();
+    expect(!!groundingMetadata).toBeTruthy();
     if (groundingMetadata) {
-      expect(groundingMetadata.groundingAttributions).toBeTruthy();
-      expect(groundingMetadata.webSearchQueries).toBeTruthy();
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
     }
   });
   it('in preview should return grounding metadata when passed GoogleSearchRetriever in generateContent', async () => {
@@ -482,10 +482,10 @@ describe('generateContent', () => {
     });
     const response = result.response;
     const groundingMetadata = response.candidates[0].groundingMetadata;
-    expect(groundingMetadata).toBeDefined();
+    expect(!!groundingMetadata).toBeTruthy();
     if (groundingMetadata) {
-      expect(groundingMetadata.groundingAttributions).toBeTruthy();
-      expect(groundingMetadata.webSearchQueries).toBeTruthy();
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
     }
   });
 });
@@ -514,6 +514,68 @@ describe('sendMessage', () => {
     );
     expect(chat.history.length).toBe(2);
   });
+  it('should return grounding metadata when passed GoogleSearchRetriever in getGenerativeModel', async () => {
+    const generativeTextModel = vertex_ai.getGenerativeModel({
+      model: 'gemini-pro',
+      tools: TOOLS_WITH_GOOGLE_SEARCH_RETRIEVAL,
+    });
+    const chat = generativeTextModel.startChat();
+    const result = await chat.sendMessage('Why is the sky blue?');
+    const response = result.response;
+    const groundingMetadata = response.candidates[0].groundingMetadata;
+    expect(!!groundingMetadata).toBeTruthy();
+    if (groundingMetadata) {
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
+    }
+  });
+  it('should return grounding metadata when passed GoogleSearchRetriever in startChat', async () => {
+    const generativeTextModel = vertex_ai.getGenerativeModel({
+      model: 'gemini-pro',
+    });
+    const chat = generativeTextModel.startChat({
+      tools: TOOLS_WITH_GOOGLE_SEARCH_RETRIEVAL,
+    });
+    const result = await chat.sendMessage('Why is the sky blue?');
+    const response = result.response;
+    const groundingMetadata = response.candidates[0].groundingMetadata;
+    expect(!!groundingMetadata).toBeTruthy();
+    if (groundingMetadata) {
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
+    }
+  });
+  it('in preview should return grounding metadata when passed GoogleSearchRetriever in getGenerativeModel', async () => {
+    const generativeTextModel = vertex_ai.preview.getGenerativeModel({
+      model: 'gemini-pro',
+      tools: TOOLS_WITH_GOOGLE_SEARCH_RETRIEVAL,
+    });
+    const chat = generativeTextModel.startChat();
+    const result = await chat.sendMessage('Why is the sky blue?');
+    const response = result.response;
+    const groundingMetadata = response.candidates[0].groundingMetadata;
+    expect(!!groundingMetadata).toBeTruthy();
+    if (groundingMetadata) {
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
+    }
+  });
+  it('in preview should return grounding metadata when passed GoogleSearchRetriever in startChat', async () => {
+    const generativeTextModel = vertex_ai.preview.getGenerativeModel({
+      model: 'gemini-pro',
+    });
+    const chat = generativeTextModel.startChat({
+      tools: TOOLS_WITH_GOOGLE_SEARCH_RETRIEVAL,
+    });
+    const result = await chat.sendMessage('Why is the sky blue?');
+    const response = result.response;
+    const groundingMetadata = response.candidates[0].groundingMetadata;
+    expect(!!groundingMetadata).toBeTruthy();
+    if (groundingMetadata) {
+      expect(!!groundingMetadata.groundingAttributions).toBeTruthy();
+      expect(!!groundingMetadata.webSearchQueries).toBeTruthy();
+    }
+  });
 });
 
 describe('sendMessageStream', () => {