From bb553d087ae689f213297a79f29b1ff9cee17738 Mon Sep 17 00:00:00 2001
From: Junyan Xu <junyanxu5513@gmail.com>
Date: Mon, 7 Oct 2024 10:39:52 -0700
Subject: [PATCH] Add samples for log prob and search grounding (#273)

* Add samples for log prob and search grounding

* Fix the comment.

* Update search grounding and log probs example. Resolve comment from Christina

* Remove s

* Add import comments to the sample
---
 samples/README.md           |  2 ++
 samples/log_prob.js         | 45 ++++++++++++++++++++++++++++
 samples/search_grounding.js | 58 +++++++++++++++++++++++++++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 samples/log_prob.js
 create mode 100644 samples/search_grounding.js

diff --git a/samples/README.md b/samples/README.md
index 4d61ba4..0c37e0e 100644
--- a/samples/README.md
+++ b/samples/README.md
@@ -46,4 +46,6 @@ to comment out cases you do not want to run.
 | [function_calling.js](./function_calling.js)             | Using function calling |
 | [safety_settings.js](./safety_settings.js)               | Setting and using safety controls |
 | [system_instruction.js](./system_instruction.js)         | Setting system instructions |
+| [search_grounding.js](./search_grounding.js)             | Generate with google search grounding |
+| [log_prob.js](./log_prob.js)                             | Generate with log probability for each token |
 | [text_generation.js](./text_generation.js)               | Generating text |
diff --git a/samples/log_prob.js b/samples/log_prob.js
new file mode 100644
index 0000000..0e2573e
--- /dev/null
+++ b/samples/log_prob.js
@@ -0,0 +1,45 @@
+/**
+ * @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 {GoogleGenerativeAI} from "@google/generative-ai";
+
+async function enableLogProb() {
+    // [START log_prob]
+    // Make sure to include these imports:
+    // import {GoogleGenerativeAI} from "@google/generative-ai";
+    const genAI = new GoogleGenerativeAI(process.env.API_KEY);
+    const model = genAI.getGenerativeModel(
+        {
+            // Only 002 models + flash 1.5 8b models are enabled with log probs 
+            // option.
+            model: "gemini-1.5-flash-002",
+            generationConfig: {
+                responseLogprobs: true
+            },
+        },
+    );
+    const prompt = "Hello!";
+    const result = await model.generateContent(prompt);
+    console.log(result.response.candidates[0].logprobsResult);
+    // [END log_prob]
+}
+async function runAll() {
+    // Comment out or delete any sample cases you don't want to run.
+    await enableLogProb();
+}
+
+runAll();
diff --git a/samples/search_grounding.js b/samples/search_grounding.js
new file mode 100644
index 0000000..363df16
--- /dev/null
+++ b/samples/search_grounding.js
@@ -0,0 +1,58 @@
+/**
+ * @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 {
+  DynamicRetrievalMode,
+  GoogleGenerativeAI,
+} from "@google/generative-ai";
+
+async function searchGrounding() {
+  // [START search_grounding]
+  // Make sure to include these imports:
+  // import {
+  //  DynamicRetrievalMode,
+  //  GoogleGenerativeAI,
+  // } from "@google/generative-ai";
+  const genAI = new GoogleGenerativeAI(process.env.API_KEY);
+  const model = genAI.getGenerativeModel(
+    {
+      model: "gemini-1.5-flash",
+      tools: [
+        {
+          googleSearchRetrieval: {
+            dynamicRetrievalConfig: {
+              mode: DynamicRetrievalMode.MODE_DYNAMIC,
+              dynamicThreshold: 0.7,
+            },
+          },
+        },
+      ],
+    },
+    { apiVersion: "v1beta" },
+  );
+
+  const prompt = "What is the price of Google stock today?";
+  const result = await model.generateContent(prompt);
+  console.log(result.response.candidates[0].groundingMetadata);
+  // [END search_grounding]
+}
+async function runAll() {
+  // Comment out or delete any sample cases you don't want to run.
+  await searchGrounding();
+}
+
+runAll();