fix: gemini function calling

ax-llm · Nov 3, 2024 · df0237d · df0237d
1 parent 01b4e2c
commit df0237d
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 36 deletions.
diff --git a/src/ax/ai/google-gemini/api.ts b/src/ax/ai/google-gemini/api.ts
@@ -208,38 +208,40 @@ export class AxAIGoogleGemini extends AxBaseAI<
           }
 
           case 'assistant': {
-            if ('content' in msg && typeof msg.content === 'string') {
-              const parts: Extract<
-                AxAIGoogleGeminiChatRequest['contents'][0],
-                { role: 'model' }
-              >['parts'] = [{ text: msg.content }];
+            let parts: Extract<
+              AxAIGoogleGeminiChatRequest['contents'][0],
+              { role: 'model' }
+            >['parts'] = [];
+
+            if (msg.functionCalls) {
+              parts = msg.functionCalls.map((f) => {
+                const args =
+                  typeof f.function.params === 'string'
+                    ? JSON.parse(f.function.params)
+                    : f.function.params;
+                return {
+                  functionCall: {
+                    name: f.function.name,
+                    args: args
+                  }
+                };
+              });
+
+              if (!parts) {
+                throw new Error('Function call is empty');
+              }
+
               return {
                 role: 'model' as const,
                 parts
               };
             }
 
-            let parts: Extract<
-              AxAIGoogleGeminiChatRequest['contents'][0],
-              { role: 'model' }
-            >['parts'] = [];
-
-            if ('functionCalls' in msg) {
-              parts =
-                msg.functionCalls?.map((f) => {
-                  const args =
-                    typeof f.function.params === 'string'
-                      ? JSON.parse(f.function.params)
-                      : f.function.params;
-                  return {
-                    functionCall: {
-                      name: f.function.name,
-                      args: args
-                    }
-                  };
-                }) ?? [];
+            if (!msg.content) {
+              throw new Error('Assistant content is empty');
             }
 
+            parts = [{ text: msg.content }];
             return {
               role: 'model' as const,
               parts
@@ -276,11 +278,11 @@ export class AxAIGoogleGemini extends AxBaseAI<
     let tools: AxAIGoogleGeminiChatRequest['tools'] | undefined = [];
 
     if (req.functions && req.functions.length > 0) {
-      tools.push({ functionDeclarations: req.functions });
+      tools.push({ function_declarations: req.functions });
     }
 
     if (this.options?.codeExecution) {
-      tools.push({ codeExecution: {} });
+      tools.push({ code_execution: {} });
     }
 
     if (tools.length === 0) {

diff --git a/src/ax/ai/google-gemini/types.ts b/src/ax/ai/google-gemini/types.ts
@@ -79,8 +79,8 @@ export type AxAIGoogleGeminiToolFunctionDeclaration = {
 };
 
 export type AxAIGoogleGeminiTool = {
-  functionDeclarations?: AxAIGoogleGeminiToolFunctionDeclaration[];
-  codeExecution?: object;
+  function_declarations?: AxAIGoogleGeminiToolFunctionDeclaration[];
+  code_execution?: object;
 };
 
 export type AxAIGoogleGeminiToolConfig = {

diff --git a/src/examples/food-search.ts b/src/examples/food-search.ts
@@ -140,10 +140,14 @@ const functions: AxFunction[] = [
   }
 ];
 
+// const ai = new AxAI({
+//   name: 'openai',
+//   apiKey: process.env.OPENAI_APIKEY as string
+// });
+
 const ai = new AxAI({
-  name: 'openai',
-  apiKey: process.env.OPENAI_APIKEY as string
-  //   config: { model: AxAIOpenAIModel.GPT4OMini }
+  name: 'google-gemini',
+  apiKey: process.env.GOOGLE_APIKEY as string
 });
 
 // ai.setOptions({ debug: true });
@@ -158,11 +162,6 @@ const ai = new AxAI({
 //   apiKey: process.env.COHERE_APIKEY as string
 // });
 
-// const ai = new AxAI({
-//   name: 'google-gemini',
-//   apiKey: process.env.GOOGLE_APIKEY as string
-// });
-
 // const ai = new AxAI({
 //   name: 'anthropic',
 //   apiKey: process.env.ANTHROPIC_APIKEY as string