From 2df129f91bc7773c43d241e744d73e0fcb943624 Mon Sep 17 00:00:00 2001
From: vijaykammili <121184261+vijaykammili@users.noreply.github.com>
Date: Thu, 19 Dec 2024 04:40:14 -0700
Subject: [PATCH] Vertex Gemini Custom Model Changes (#3737)

---
 .../ChatGoogleVertexAI/ChatGoogleVertexAI.ts  | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/packages/components/nodes/chatmodels/ChatGoogleVertexAI/ChatGoogleVertexAI.ts b/packages/components/nodes/chatmodels/ChatGoogleVertexAI/ChatGoogleVertexAI.ts
index 607f3ba50..5b813984b 100644
--- a/packages/components/nodes/chatmodels/ChatGoogleVertexAI/ChatGoogleVertexAI.ts
+++ b/packages/components/nodes/chatmodels/ChatGoogleVertexAI/ChatGoogleVertexAI.ts
@@ -23,6 +23,10 @@ class ChatVertexAI extends LcChatVertexAI implements IVisionChatModal {
 
     constructor(id: string, fields?: ChatVertexAIInput) {
         // @ts-ignore
+        if (fields?.model) {
+            fields.modelName = fields.model
+            delete fields.model
+        }
         super(fields ?? {})
         this.id = id
         this.configuredModel = fields?.modelName || ''
@@ -61,7 +65,7 @@ class GoogleVertexAI_ChatModels implements INode {
     constructor() {
         this.label = 'ChatGoogleVertexAI'
         this.name = 'chatGoogleVertexAI'
-        this.version = 5.0
+        this.version = 5.1
         this.type = 'ChatGoogleVertexAI'
         this.icon = 'GoogleVertex.svg'
         this.category = 'Chat Models'
@@ -89,6 +93,14 @@ class GoogleVertexAI_ChatModels implements INode {
                 type: 'asyncOptions',
                 loadMethod: 'listModels'
             },
+            {
+                label: 'Custom Model Name',
+                name: 'customModelName',
+                type: 'string',
+                placeholder: 'gemini-1.5-pro-exp-0801',
+                description: 'Custom model name to use. If provided, it will override the model selected',
+                additionalParams: true
+            },
             {
                 label: 'Temperature',
                 name: 'temperature',
@@ -163,7 +175,6 @@ class GoogleVertexAI_ChatModels implements INode {
                 throw new Error(
                     'Error: More than one component has been inputted. Please use only one of the following: Google Application Credential File Path or Google Credential JSON Object'
                 )
-
             if (googleApplicationCredentialFilePath && !googleApplicationCredential)
                 authOptions.keyFile = googleApplicationCredentialFilePath
             else if (!googleApplicationCredentialFilePath && googleApplicationCredential)
@@ -174,6 +185,7 @@ class GoogleVertexAI_ChatModels implements INode {
 
         const temperature = nodeData.inputs?.temperature as string
         const modelName = nodeData.inputs?.modelName as string
+        const customModelName = nodeData.inputs?.customModelName as string
         const maxOutputTokens = nodeData.inputs?.maxOutputTokens as string
         const topP = nodeData.inputs?.topP as string
         const cache = nodeData.inputs?.cache as BaseCache
@@ -190,11 +202,10 @@ class GoogleVertexAI_ChatModels implements INode {
 
         const obj: ChatVertexAIInput = {
             temperature: parseFloat(temperature),
-            model: modelName,
+            modelName: customModelName || modelName,
             streaming: streaming ?? true
         }
         if (Object.keys(authOptions).length !== 0) obj.authOptions = authOptions
-
         if (maxOutputTokens) obj.maxOutputTokens = parseInt(maxOutputTokens, 10)
         if (topP) obj.topP = parseFloat(topP)
         if (cache) obj.cache = cache