Function calling

packages/main/src/models/generative-model.ts

@@ -35,6 +35,7 @@ import {
   RequestOptions,
   SafetySetting,
   StartChatParams,
+  Tool,
 } from "../../types";
 import { ChatSession } from "../methods/chat-session";
 import { countTokens } from "../methods/count-tokens";
@@ -53,6 +54,7 @@ export class GenerativeModel {
   generationConfig: GenerationConfig;
   safetySettings: SafetySetting[];
   requestOptions: RequestOptions;
+  tools?: Tool[];
 
   constructor(
     public apiKey: string,
@@ -68,6 +70,7 @@ export class GenerativeModel {
     }
     this.generationConfig = modelParams.generationConfig || {};
     this.safetySettings = modelParams.safetySettings || [];
+    this.tools = modelParams.tools;
     this.requestOptions = requestOptions || {};
   }
 
@@ -85,6 +88,7 @@ export class GenerativeModel {
       {
         generationConfig: this.generationConfig,
         safetySettings: this.safetySettings,
+        tools: this.tools,
         ...formattedParams,
       },
       this.requestOptions,
@@ -107,6 +111,7 @@ export class GenerativeModel {
       {
         generationConfig: this.generationConfig,
         safetySettings: this.safetySettings,
+        tools: this.tools,
         ...formattedParams,
       },
       this.requestOptions,
@@ -121,7 +126,10 @@ export class GenerativeModel {
     return new ChatSession(
       this.apiKey,
       this.model,
-      startChatParams,
+      {
+        tools: this.tools,
+        ...startChatParams,
+      },
       this.requestOptions,
     );
   }

packages/main/test-integration/node/count-tokens.test.ts

@@ -0,0 +1,49 @@
+/**
+ * @license
+ * Copyright 2023 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 { GoogleGenerativeAI, HarmBlockThreshold, HarmCategory } from "../..";
+
+use(chaiAsPromised);
+
+/**
+ * Integration tests against live backend.
+ */
+
+describe("countTokens", function () {
+  this.timeout(60e3);
+  this.slow(10e3);
+  it("counts tokens right", async () => {
+    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
+    const model = genAI.getGenerativeModel({
+      model: "gemini-pro",
+      safetySettings: [
+        {
+          category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+          threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+        },
+      ],
+    });
+    const response1 = await model.countTokens("count me");
+    const response2 = await model.countTokens({
+      contents: [{ role: "user", parts: [{ text: "count me" }] }],
+    });
+    expect(response1.totalTokens).to.equal(3);
+    expect(response2.totalTokens).to.equal(3);
+  });
+});

packages/main/test-integration/node/embed-content.test.ts

@@ -0,0 +1,43 @@
+/**
+ * @license
+ * Copyright 2023 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 { GoogleGenerativeAI } from "../..";
+
+use(chaiAsPromised);
+
+/**
+ * Integration tests against live backend.
+ */
+
+describe("embedContent", function () {
+  this.timeout(60e3);
+  this.slow(10e3);
+  it("embeds a single Content object", async () => {
+    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
+    const model = genAI.getGenerativeModel({
+      model: "embedding-001",
+    });
+    const response1 = await model.embedContent("embed me");
+    const response2 = await model.embedContent({
+      content: { role: "user", parts: [{ text: "embed me" }] },
+    });
+    expect(response1.embedding).to.not.be.empty;
+    expect(response1).to.eql(response2);
+  });
+});

packages/main/test-integration/node/generate-content-multimodal.test.ts

@@ -0,0 +1,68 @@
+/**
+ * @license
+ * Copyright 2023 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 * as fs from "fs";
+import { join } from "path";
+
+import { expect, use } from "chai";
+import * as chaiAsPromised from "chai-as-promised";
+import { GoogleGenerativeAI, HarmBlockThreshold, HarmCategory } from "../..";
+
+use(chaiAsPromised);
+
+/**
+ * Integration tests against live backend.
+ */
+
+describe("generateContent", function () {
+  this.timeout(60e3);
+  this.slow(10e3);
+  it("non-streaming, image buffer provided", async () => {
+    const imageBuffer = fs.readFileSync(
+      join(__dirname, "../../test-utils/cat.png"),
+    );
+    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
+    const base64Image = imageBuffer.toString("base64");
+    const model = genAI.getGenerativeModel({
+      model: "gemini-pro-vision",
+      safetySettings: [
+        {
+          category: HarmCategory.HARM_CATEGORY_HARASSMENT,
+          threshold: HarmBlockThreshold.BLOCK_ONLY_HIGH,
+        },
+      ],
+    });
+    const result = await model.generateContent({
+      contents: [
+        {
+          role: "user",
+          parts: [
+            { text: "Is it a cat?" },
+            {
+              inlineData: {
+                mimeType: "image/png",
+                data: base64Image,
+              },
+            },
+          ],
+        },
+      ],
+    });
+    const response = result.response;
+    expect(response.text()).to.not.be.empty;
+  });
+});

packages/main/test-integration/node/generate-content-tools.test.ts

@@ -0,0 +1,183 @@
+/**
+ * @license
+ * Copyright 2023 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 { FunctionDeclarationSchemaType, GoogleGenerativeAI } from "../..";
+
+use(chaiAsPromised);
+
+/**
+ * Integration tests against live backend.
+ */
+
+describe("generateContent", function () {
+  this.timeout(60e3);
+  this.slow(10e3);
+  // This test can be flaky
+  // eslint-disable-next-line no-restricted-properties
+  it("non-streaming, tools usage", async () => {
+    const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY || "");
+    const model = genAI.getGenerativeModel(
+      {
+        model: "gemini-pro",
+        tools: [
+          {
+            functionDeclarations: [
+              {
+                name: "find_movies",
+                description:
+                  "find movie titles currently playing in theaters based on any description, genre, title words, etc.",
+                parameters: {
+                  type: FunctionDeclarationSchemaType.OBJECT,
+                  properties: {
+                    location: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description:
+                        "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                    },
+                    description: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description:
+                        "Any kind of description including category or genre, title words, attributes, etc.",
+                    },
+                  },
+                  required: ["description"],
+                },
+              },
+              {
+                name: "find_theaters",
+                description:
+                  "find theaters based on location and optionally movie title which are is currently playing in theaters",
+                parameters: {
+                  type: FunctionDeclarationSchemaType.OBJECT,
+                  properties: {
+                    location: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description:
+                        "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                    },
+                    movie: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description: "Any movie title",
+                    },
+                  },
+                  required: ["location"],
+                },
+              },
+              {
+                name: "get_showtimes",
+                description:
+                  "Find the start times for movies playing in a specific theater",
+                parameters: {
+                  type: FunctionDeclarationSchemaType.OBJECT,
+                  properties: {
+                    location: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description:
+                        "The city and state, e.g. San Francisco, CA or a zip code e.g. 95616",
+                    },
+                    movie: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description: "Any movie title",
+                    },
+                    theater: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description: "Name of the theater",
+                    },
+                    date: {
+                      type: FunctionDeclarationSchemaType.STRING,
+                      description: "Date for requested showtime",
+                    },
+                  },
+                  required: ["location", "movie", "theater", "date"],
+                },
+              },
+            ],
+          },
+        ],
+      },
+      { apiVersion: "v1beta" },
+    );
+
+    const src1 = {
+      role: "user",
+      parts: [
+        {
+          text: "Which theaters in Mountain View show Barbie movie?",
+        },
+      ],
+    };
+    const exp1 = {
+      role: "model",
+      parts: [
+        {
+          functionCall: {
+            name: "find_theaters",
+            args: {
+              location: "Mountain View, CA",
+              movie: "Barbie",
+            },
+          },
+        },
+      ],
+    };
+
+    const src2 = {
+      role: "function",
+      parts: [
+        {
+          functionResponse: {
+            name: "find_theaters",
+            response: {
+              name: "find_theaters",
+              content: {
+                movie: "Barbie",
+                theaters: [
+                  {
+                    name: "AMC Mountain View 16",
+                    address: "2000 W El Camino Real, Mountain View, CA 94040",
+                  },
+                  {
+                    name: "Regal Edwards 14",
+                    address: "245 Castro St, Mountain View, CA 94040",
+                  },
+                ],
+              },
+            },
+          },
+        },
+      ],
+    };
+
+    const result1 = await model.generateContent({
+      contents: [src1],
+    });
+    const response1 = result1.response;
+    expect(response1.candidates.length).to.equal(1);
+    expect(response1.candidates[0].content.role).to.equal("model");
+    expect(response1.candidates[0].content.parts.length).to.equal(1);
+    expect(response1.candidates[0].content).to.deep.equal(exp1);
+
+    const result3 = await model.generateContent({
+      contents: [src1, exp1, src2],
+    });
+    const response3 = result3.response;
+    expect(response3.text()).include("AMC Mountain View 16");
+    expect(response3.text()).include("Regal Edwards 14");
+  });
+});