Search grounding

.changeset/shy-panthers-thank.md

@@ -0,0 +1,5 @@
+---
+"@google/generative-ai": minor
+---
+
+Add support for search grounding

common/api-review/generative-ai-server.api.md

@@ -471,8 +471,10 @@ export interface TextPart {
     text: string;
 }
 
+// Warning: (ae-forgotten-export) The symbol "GoogleSearchRetrievalTool" needs to be exported by the entry point index.d.ts
+//
 // @public
-export type Tool = FunctionDeclarationsTool | CodeExecutionTool;
+export type Tool = FunctionDeclarationsTool | CodeExecutionTool | GoogleSearchRetrievalTool;
 
 // @public
 export interface ToolConfig {

common/api-review/generative-ai.api.md

@@ -155,6 +155,14 @@ export interface CountTokensResponse {
     totalTokens: number;
 }
 
+// @public
+export enum DynamicRetrievalMode {
+    // (undocumented)
+    MODE_DYNAMIC = "MODE_DYNAMIC",
+    // (undocumented)
+    MODE_UNSPECIFIED = "MODE_UNSPECIFIED"
+}
+
 // @public
 export interface EmbedContentRequest {
     // (undocumented)
@@ -378,6 +386,8 @@ export interface GenerateContentCandidate {
     finishMessage?: string;
     // (undocumented)
     finishReason?: FinishReason;
+    // Warning: (ae-forgotten-export) The symbol "GroundingMetadata" needs to be exported by the entry point index.d.ts
+    groundingMetadata?: GroundingMetadata;
     // (undocumented)
     index: number;
     logprobsResult?: LogprobsResult;
@@ -734,8 +744,10 @@ export interface TextPart {
     text: string;
 }
 
+// Warning: (ae-forgotten-export) The symbol "GoogleSearchRetrievalTool" needs to be exported by the entry point index.d.ts
+//
 // @public
-export type Tool = FunctionDeclarationsTool | CodeExecutionTool;
+export type Tool = FunctionDeclarationsTool | CodeExecutionTool | GoogleSearchRetrievalTool;
 
 // @public
 export interface ToolConfig {

docs/reference/main/generative-ai.dynamicretrievalmode.md

@@ -0,0 +1,21 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [DynamicRetrievalMode](./generative-ai.dynamicretrievalmode.md)
+
+## DynamicRetrievalMode enum
+
+The mode of the predictor to be used in dynamic retrieval.
+
+**Signature:**
+
+```typescript
+export declare enum DynamicRetrievalMode 
+```
+
+## Enumeration Members
+
+|  Member | Value | Description |
+|  --- | --- | --- |
+|  MODE\_DYNAMIC | <code>&quot;MODE_DYNAMIC&quot;</code> |  |
+|  MODE\_UNSPECIFIED | <code>&quot;MODE_UNSPECIFIED&quot;</code> |  |
+

docs/reference/main/generative-ai.generatecontentcandidate.groundingmetadata.md

@@ -0,0 +1,13 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [@google/generative-ai](./generative-ai.md) &gt; [GenerateContentCandidate](./generative-ai.generatecontentcandidate.md) &gt; [groundingMetadata](./generative-ai.generatecontentcandidate.groundingmetadata.md)
+
+## GenerateContentCandidate.groundingMetadata property
+
+Search gounding metadata.
+
+**Signature:**
+
+```typescript
+groundingMetadata?: GroundingMetadata;
+```

docs/reference/main/generative-ai.generatecontentcandidate.md

@@ -21,6 +21,7 @@ export interface GenerateContentCandidate
 |  [content](./generative-ai.generatecontentcandidate.content.md) |  | [Content](./generative-ai.content.md) |  |
 |  [finishMessage?](./generative-ai.generatecontentcandidate.finishmessage.md) |  | string | _(Optional)_ |
 |  [finishReason?](./generative-ai.generatecontentcandidate.finishreason.md) |  | [FinishReason](./generative-ai.finishreason.md) | _(Optional)_ |
+|  [groundingMetadata?](./generative-ai.generatecontentcandidate.groundingmetadata.md) |  | GroundingMetadata | _(Optional)_ Search gounding metadata. |
 |  [index](./generative-ai.generatecontentcandidate.index.md) |  | number |  |
 |  [logprobsResult?](./generative-ai.generatecontentcandidate.logprobsresult.md) |  | [LogprobsResult](./generative-ai.logprobsresult.md) | _(Optional)_ Log-likelihood scores for the response tokens and top tokens. |
 |  [safetyRatings?](./generative-ai.generatecontentcandidate.safetyratings.md) |  | [SafetyRating](./generative-ai.safetyrating.md)<!-- -->\[\] | _(Optional)_ |

docs/reference/main/generative-ai.md

@@ -21,6 +21,7 @@
 |  Enumeration | Description |
 |  --- | --- |
 |  [BlockReason](./generative-ai.blockreason.md) | Reason that a prompt was blocked. |
+|  [DynamicRetrievalMode](./generative-ai.dynamicretrievalmode.md) | The mode of the predictor to be used in dynamic retrieval. |
 |  [ExecutableCodeLanguage](./generative-ai.executablecodelanguage.md) |  |
 |  [FinishReason](./generative-ai.finishreason.md) | Reason that a candidate finished. |
 |  [FunctionCallingMode](./generative-ai.functioncallingmode.md) |  |

docs/reference/main/generative-ai.tool.md

@@ -9,7 +9,7 @@ Defines a tool that model can call to access external knowledge.
 **Signature:**
 
 ```typescript
-export declare type Tool = FunctionDeclarationsTool | CodeExecutionTool;
+export declare type Tool = FunctionDeclarationsTool | CodeExecutionTool | GoogleSearchRetrievalTool;
 ```
 **References:** [FunctionDeclarationsTool](./generative-ai.functiondeclarationstool.md)<!-- -->, [CodeExecutionTool](./generative-ai.codeexecutiontool.md)
 

docs/reference/server/generative-ai.tool.md

@@ -9,7 +9,7 @@ Defines a tool that model can call to access external knowledge.
 **Signature:**
 
 ```typescript
-export declare type Tool = FunctionDeclarationsTool | CodeExecutionTool;
+export declare type Tool = FunctionDeclarationsTool | CodeExecutionTool | GoogleSearchRetrievalTool;
 ```
 **References:** [FunctionDeclarationsTool](./generative-ai.functiondeclarationstool.md)<!-- -->, [CodeExecutionTool](./generative-ai.codeexecutiontool.md)
 

src/methods/generate-content.test.ts

@@ -15,7 +15,7 @@
  * limitations under the License.
  */
 
-import { expect, use } from "chai";
+import { assert, expect, use } from "chai";
 import { match, restore, stub } from "sinon";
 import * as sinonChai from "sinon-chai";
 import * as chaiAsPromised from "chai-as-promised";
@@ -203,4 +203,70 @@ describe("generateContent()", () => {
     ).to.be.rejectedWith(/400.*invalid argument/);
     expect(makeRequestStub).to.be.called;
   });
+  it("search-grounding", async () => {
+    const mockResponse = getMockResponse("unary-success-search-grounding.json");
+    const makeRequestStub = stub(request, "makeModelRequest").resolves(
+      mockResponse as Response,
+    );
+    const expectedGourdingMetadata = {
+      searchEntryPoint: {
+        renderedContent: "test_rendered_content",
+      },
+      groundingChunks: [
+        {
+          web: {
+            uri: "test_uri_1",
+            title: "test_title_1",
+          },
+        },
+        {
+          web: {
+            uri: "test_uri_2",
+            title: "test_title_2",
+          },
+        },
+      ],
+      groundingSupports: [
+        {
+          segment: {
+            endIndex: 76,
+            text: "The current stock price for Alphabet Inc (Google) Class C (GOOG) is $166.94.",
+          },
+          groundingChunkIndices: [0, 1],
+          confidenceScores: [0.9863024, 0.9863024],
+        },
+        {
+          segment: {
+            startIndex: 77,
+            endIndex: 135,
+            text: "This represents a decrease of -0.88% in the past 24 hours.",
+          },
+          groundingChunkIndices: [0, 1],
+          confidenceScores: [0.98734087, 0.98734087],
+        },
+        {
+          segment: {
+            startIndex: 136,
+            endIndex: 239,
+            text: "Please note that stock prices can fluctuate frequently and this price is valid as of October 2nd, 2024.",
+          },
+          groundingChunkIndices: [0],
+          confidenceScores: [0.6393517],
+        },
+      ],
+      webSearchQueries: ["what is the current google stock price"],
+    };
+    const result = await generateContent("key", "model", fakeRequestParams);
+    expect(makeRequestStub).to.be.calledWith(
+      "model",
+      request.Task.GENERATE_CONTENT,
+      "key",
+      false,
+      match.any,
+    );
+    assert.equal(
+      JSON.stringify(result.response.candidates[0].groundingMetadata),
+      JSON.stringify(expectedGourdingMetadata),
+    );
+  });
 });

src/requests/stream-reader.test.ts

@@ -132,6 +132,86 @@ describe("processStream", () => {
       args: { city: "San Jose" },
     });
   });
+
+  it("streaming response - searchGrounding", async () => {
+    const fakeResponse = getMockResponseStreaming(
+      "streaming-success-search-grounding.txt",
+    );
+    const result = processStream(fakeResponse as Response);
+    const aggregatedResponse = await result.response;
+    const expectedGroundingMetadata = {
+      searchEntryPoint: {
+        renderedContent: "test_rendered_content",
+      },
+      groundingChunks: [
+        {
+          web: {
+            uri: "test_uri_1",
+            title: "test_title_1",
+          },
+        },
+        {
+          web: {
+            uri: "test_uri_2",
+            title: "test_title_2",
+          },
+        },
+      ],
+      groundingSupports: [
+        {
+          segment: {
+            endIndex: 41,
+            text: "The current stock price for Alphabet Inc.",
+          },
+          groundingChunkIndices: [0],
+          confidenceScores: [0.68925554],
+        },
+        {
+          segment: {
+            startIndex: 42,
+            endIndex: 82,
+            text: "(Google) Class C (GOOG) is \\$166.79 USD.",
+          },
+          groundingChunkIndices: [1, 0],
+          confidenceScores: [0.92251855, 0.92251855],
+        },
+        {
+          segment: {
+            startIndex: 83,
+            endIndex: 147,
+            text: "This price reflects a decrease of -0.97% over the last 24 hours.",
+          },
+          groundingChunkIndices: [0],
+          confidenceScores: [0.9831334],
+        },
+        {
+          segment: {
+            startIndex: 150,
+            endIndex: 199,
+            text: "Please note that stock prices can change rapidly.",
+          },
+          groundingChunkIndices: [1],
+          confidenceScores: [0.6181941],
+        },
+        {
+          segment: {
+            startIndex: 201,
+            endIndex: 267,
+            text: "This information is current as of October 2, 2024, at 5:58 PM UTC.",
+          },
+          groundingChunkIndices: [1],
+          confidenceScores: [0.6107691],
+        },
+      ],
+      webSearchQueries: ["what is the current google stock price"],
+    };
+    expect(aggregatedResponse.text()).to.contains("$166.79 USD");
+    console.log(aggregatedResponse.candidates[0]);
+    expect(aggregatedResponse.candidates[0].groundingMetadata).to.be.deep.equal(
+      expectedGroundingMetadata,
+    );
+  });
+
   it("candidate had finishReason", async () => {
     const fakeResponse = getMockResponseStreaming(
       "streaming-failure-finish-reason-safety.txt",

src/requests/stream-reader.ts

@@ -153,6 +153,8 @@ export function aggregateResponses(
         // Keep overwriting, the last one will be final
         aggregatedResponse.candidates[i].citationMetadata =
           candidate.citationMetadata;
+        aggregatedResponse.candidates[i].groundingMetadata =
+          candidate.groundingMetadata;
         aggregatedResponse.candidates[i].finishReason = candidate.finishReason;
         aggregatedResponse.candidates[i].finishMessage =
           candidate.finishMessage;

test-integration/node/generate-content-search-grounding.test.ts

@@ -0,0 +1,67 @@
+/**
+ * @license
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { expect, use } from "chai";
+import * as chaiAsPromised from "chai-as-promised";
+import {
+  DynamicRetrievalMode,
+  GoogleGenerativeAI,
+  HarmBlockThreshold,
+  HarmCategory,
+} from "../..";
+
+use(chaiAsPromised);
+
+/**
+ * Integration tests against live backend.
+ */
+
+describe("generateContent", function () {
+  this.timeout(60e3);
+  this.slow(10e3);
+  it("non-streaming, get Google stock price from google search", async () => {
+    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
+    const model = genAI.getGenerativeModel(
+      {
+        model: "gemini-1.5-pro",
+        safetySettings: [
+          {
+            category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+            threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+          },
+        ],
+        generationConfig: {},
+        tools: [
+          {
+            googleSearchRetrieval: {
+              dynamicRetrievalConfig: {
+                mode: DynamicRetrievalMode.MODE_DYNAMIC,
+                dynamicThreshold: 0.7,
+              },
+            },
+          },
+        ],
+      },
+      { apiVersion: "v1beta" },
+    );
+    const result = await model.generateContent(
+      "What is the current Google stock price?",
+    );
+    const response = result.response;
+    expect(response.text()).to.not.be.empty;
+  });
+});

test-utils/mock-responses/streaming-success-search-grounding.txt

@@ -0,0 +1,14 @@
+data: {"candidates": [{"content": {"parts": [{"text": "The"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 1, "totalTokenCount": 9}}
+
+data: {"candidates": [{"content": {"parts": [{"text": " current stock price for Alphabet Inc. (Google) Class C (GOOG)"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 17, "totalTokenCount": 25}}
+
+data: {"candidates": [{"content": {"parts": [{"text": " is \\$166.79 USD. This price reflects a decrease of"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 33, "totalTokenCount": 41}}
+
+data: {"candidates": [{"content": {"parts": [{"text": " -0.97% over the last 24 hours. \n\nPlease note that stock prices can change rapidly.  This information is current as of"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 65, "totalTokenCount": 73}}
+
+data: {"candidates": [{"content": {"parts": [{"text": " October 2, 2024, at 5:58 PM UTC.  You can find the most up-to-date information on"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 97, "totalTokenCount": 105}}
+
+data: {"candidates": [{"content": {"parts": [{"text": " financial websites like Google Finance or TradingView.\n"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 106, "totalTokenCount": 114}}
+
+data: {"candidates": [{"content": {}, "index": 0, "groundingMetadata": {"searchEntryPoint": {"renderedContent": "test_rendered_content"}, "groundingChunks": [{"web": {"uri": "test_uri_1", "title": "test_title_1"}}, {"web": {"uri": "test_uri_2", "title": "test_title_2"}}], "groundingSupports": [{"segment": {"endIndex": 41, "text": "The current stock price for Alphabet Inc."}, "groundingChunkIndices": [0], "confidenceScores": [0.68925554]}, {"segment": {"startIndex": 42, "endIndex": 82, "text": "(Google) Class C (GOOG) is \\$166.79 USD."}, "groundingChunkIndices": [1, 0], "confidenceScores": [0.92251855, 0.92251855]}, {"segment": {"startIndex": 83, "endIndex": 147, "text": "This price reflects a decrease of -0.97% over the last 24 hours."}, "groundingChunkIndices": [0], "confidenceScores": [0.9831334]}, {"segment": {"startIndex": 150, "endIndex": 199, "text": "Please note that stock prices can change rapidly."}, "groundingChunkIndices": [1], "confidenceScores": [0.6181941]}, {"segment": {"startIndex": 201, "endIndex": 267, "text": "This information is current as of October 2, 2024, at 5:58 PM UTC."}, "groundingChunkIndices": [1], "confidenceScores": [0.6107691]}], "webSearchQueries": ["what is the current google stock price"]}}], "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 106, "totalTokenCount": 114}}
+