feat: added functionCallingConfig config option for gemini models (#…

…1227)
firebase · Nov 11, 2024 · c77d14e · c77d14e
1 parent ea03813
commit c77d14e
Show file tree

Hide file tree

Showing 2 changed files with 95 additions and 5 deletions.
diff --git a/js/plugins/googleai/src/gemini.ts b/js/plugins/googleai/src/gemini.ts
@@ -16,6 +16,7 @@
 
 import {
   FileDataPart,
+  FunctionCallingMode,
   FunctionCallPart,
   FunctionDeclaration,
   FunctionResponsePart,
@@ -30,29 +31,30 @@ import {
   SchemaType,
   StartChatParams,
   Tool,
+  ToolConfig,
 } from '@google/generative-ai';
 import {
-  GENKIT_CLIENT_HEADER,
   Genkit,
+  GENKIT_CLIENT_HEADER,
   GenkitError,
   JSONSchema,
   z,
 } from 'genkit';
 import {
   CandidateData,
   GenerationCommonConfigSchema,
+  getBasicUsageStats,
   MediaPart,
   MessageData,
   ModelAction,
   ModelInfo,
   ModelMiddleware,
+  modelRef,
   ModelReference,
   Part,
   ToolDefinitionSchema,
   ToolRequestPart,
   ToolResponsePart,
-  getBasicUsageStats,
-  modelRef,
 } from 'genkit/model';
 import {
   downloadRequestMedia,
@@ -79,6 +81,12 @@ const SafetySettingsSchema = z.object({
 export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
   safetySettings: z.array(SafetySettingsSchema).optional(),
   codeExecution: z.union([z.boolean(), z.object({}).strict()]).optional(),
+  functionCallingConfig: z
+    .object({
+      mode: z.enum(['MODE_UNSPECIFIED', 'AUTO', 'ANY', 'NONE']).optional(),
+      allowedFunctionNames: z.array(z.string()).optional(),
+    })
+    .optional(),
 });
 
 export const gemini10Pro = modelRef({
@@ -531,7 +539,7 @@ export function defineGoogleAIModel(
       if (apiVersion) {
         options.baseUrl = baseUrl;
       }
-      const requestConfig = {
+      const requestConfig: z.infer<typeof GeminiConfigSchema> = {
         ...defaultConfig,
         ...request.config,
       };
@@ -565,7 +573,6 @@ export function defineGoogleAIModel(
       }
 
       const tools: Tool[] = [];
-
       if (request.tools?.length) {
         tools.push({
           functionDeclarations: request.tools.map(toGeminiTool),
@@ -581,6 +588,19 @@ export function defineGoogleAIModel(
         });
       }
 
+      let toolConfig: ToolConfig | undefined;
+      if (requestConfig.functionCallingConfig) {
+        toolConfig = {
+          functionCallingConfig: {
+            allowedFunctionNames:
+              requestConfig.functionCallingConfig.allowedFunctionNames,
+            mode: toGeminiFunctionMode(
+              requestConfig.functionCallingConfig.mode
+            ),
+          },
+        };
+      }
+
       //  cannot use tools with json mode
       const jsonMode =
         (request.output?.format === 'json' ||
@@ -605,6 +625,7 @@ export function defineGoogleAIModel(
         systemInstruction,
         generationConfig,
         tools,
+        toolConfig,
         history: messages
           .slice(0, -1)
           .map((message) => toGeminiMessage(message, model)),
@@ -668,3 +689,27 @@ export function defineGoogleAIModel(
     }
   );
 }
+
+function toGeminiFunctionMode(
+  genkitMode: string | undefined
+): FunctionCallingMode | undefined {
+  if (genkitMode === undefined) {
+    return undefined;
+  }
+  switch (genkitMode) {
+    case 'MODE_UNSPECIFIED': {
+      return FunctionCallingMode.MODE_UNSPECIFIED;
+    }
+    case 'ANY': {
+      return FunctionCallingMode.ANY;
+    }
+    case 'AUTO': {
+      return FunctionCallingMode.AUTO;
+    }
+    case 'NONE': {
+      return FunctionCallingMode.NONE;
+    }
+    default:
+      throw new Error(`unsupported function calling mode: ${genkitMode}`);
+  }
+}
diff --git a/js/plugins/vertexai/src/gemini.ts b/js/plugins/vertexai/src/gemini.ts
@@ -16,6 +16,7 @@
 
 import {
   Content,
+  FunctionCallingMode,
   FunctionDeclaration,
   FunctionDeclarationSchemaType,
   Part as GeminiPart,
@@ -25,6 +26,7 @@ import {
   HarmBlockThreshold,
   HarmCategory,
   StartChatParams,
+  ToolConfig,
   VertexAI,
 } from '@google-cloud/vertexai';
 import { GENKIT_CLIENT_HEADER, Genkit, JSONSchema, z } from 'genkit';
@@ -71,6 +73,12 @@ export const GeminiConfigSchema = GenerationCommonConfigSchema.extend({
   location: z.string().optional(),
   vertexRetrieval: VertexRetrievalSchema.optional(),
   googleSearchRetrieval: GoogleSearchRetrievalSchema.optional(),
+  functionCallingConfig: z
+    .object({
+      mode: z.enum(['MODE_UNSPECIFIED', 'AUTO', 'ANY', 'NONE']).optional(),
+      allowedFunctionNames: z.array(z.string()).optional(),
+    })
+    .optional(),
 });
 
 export const gemini10Pro = modelRef({
@@ -485,6 +493,18 @@ export function defineGeminiModel(
         ? [{ functionDeclarations: request.tools?.map(toGeminiTool) }]
         : [];
 
+      let toolConfig: ToolConfig | undefined;
+      if (request?.config?.functionCallingConfig) {
+        toolConfig = {
+          functionCallingConfig: {
+            allowedFunctionNames:
+              request.config.functionCallingConfig.allowedFunctionNames,
+            mode: toGeminiFunctionMode(
+              request.config.functionCallingConfig.mode
+            ),
+          },
+        };
+      }
       // Cannot use tools and function calling at the same time
       const jsonMode =
         (request.output?.format === 'json' || !!request.output?.schema) &&
@@ -493,6 +513,7 @@ export function defineGeminiModel(
       const chatRequest: StartChatParams = {
         systemInstruction,
         tools,
+        toolConfig,
         history: messages
           .slice(0, -1)
           .map((message) => toGeminiMessage(message, model)),
@@ -589,3 +610,27 @@ export function defineGeminiModel(
     }
   );
 }
+
+function toGeminiFunctionMode(
+  genkitMode: string | undefined
+): FunctionCallingMode | undefined {
+  if (genkitMode === undefined) {
+    return undefined;
+  }
+  switch (genkitMode) {
+    case 'MODE_UNSPECIFIED': {
+      return FunctionCallingMode.MODE_UNSPECIFIED;
+    }
+    case 'ANY': {
+      return FunctionCallingMode.ANY;
+    }
+    case 'AUTO': {
+      return FunctionCallingMode.AUTO;
+    }
+    case 'NONE': {
+      return FunctionCallingMode.NONE;
+    }
+    default:
+      throw new Error(`unsupported function calling mode: ${genkitMode}`);
+  }
+}