Merge pull request Portkey-AI#461 from narengogi/fix/vertex-ai-gemini…

…-role-alternation-fix vertex-ai/gemini Combine messages with the same role in succession to handle multi turn messages
shivam-pareek · Jul 19, 2024 · d68542c · d68542c
2 parents 491d265 + ae11fb9
commit d68542c
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 13 deletions.
diff --git a/src/providers/google-vertex-ai/chatComplete.ts b/src/providers/google-vertex-ai/chatComplete.ts
@@ -16,6 +16,7 @@ import {
 import {
   GoogleMessage,
   GoogleMessageRole,
+  GoogleToolConfig,
   SYSTEM_INSTRUCTION_DISABLED_MODELS,
   transformOpenAIRoleToGoogleRole,
   transformToolChoiceForGemini,
@@ -135,16 +136,15 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
 
           // @NOTE: This takes care of the "Please ensure that multiturn requests alternate between user and model."
           // error that occurs when we have multiple user messages in a row.
-          const shouldAppendEmptyModeChat =
-            lastRole === 'user' &&
-            role === 'user' &&
-            !params.model?.includes('vision');
+          const shouldCombineMessages =
+            lastRole === role && !params.model?.includes('vision');
 
-          if (shouldAppendEmptyModeChat) {
-            messages.push({ role: 'model', parts: [{ text: '' }] });
+          if (shouldCombineMessages) {
+            messages[messages.length - 1].parts.push(...parts);
+          } else {
+            messages.push({ role, parts });
           }
 
-          messages.push({ role, parts });
           lastRole = role;
         });
 
@@ -258,12 +258,16 @@ export const VertexGoogleChatCompleteConfig: ProviderConfig = {
         ) {
           allowedFunctionNames.push(params.tool_choice.function.name);
         }
-        return {
-          functionCallingConfig: {
+        const toolConfig: GoogleToolConfig = {
+          function_calling_config: {
             mode: transformToolChoiceForGemini(params.tool_choice),
-            allowedFunctionNames,
           },
         };
+        if (allowedFunctionNames.length > 0) {
+          toolConfig.function_calling_config.allowed_function_names =
+            allowedFunctionNames;
+        }
+        return toolConfig;
       }
     },
   },

diff --git a/src/providers/google/chatComplete.ts b/src/providers/google/chatComplete.ts
@@ -73,6 +73,13 @@ export interface GoogleMessage {
   parts: GoogleMessagePart[];
 }
 
+export interface GoogleToolConfig {
+  function_calling_config: {
+    mode: GoogleToolChoiceType | undefined;
+    allowed_function_names?: string[];
+  };
+}
+
 export const transformOpenAIRoleToGoogleRole = (
   role: OpenAIMessageRole
 ): GoogleMessageRole => {
@@ -285,12 +292,16 @@ export const GoogleChatCompleteConfig: ProviderConfig = {
         ) {
           allowedFunctionNames.push(params.tool_choice.function.name);
         }
-        return {
-          functionCallingConfig: {
+        const toolConfig: GoogleToolConfig = {
+          function_calling_config: {
             mode: transformToolChoiceForGemini(params.tool_choice),
-            allowedFunctionNames,
           },
         };
+        if (allowedFunctionNames.length > 0) {
+          toolConfig.function_calling_config.allowed_function_names =
+            allowedFunctionNames;
+        }
+        return toolConfig;
       }
     },
   },