From c335bf681a45ffd523a6607abd07f9f945fde789 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 22 Feb 2024 14:57:03 -0300
Subject: [PATCH] Fix import error in ChatVertexAIComponent

---
 src/backend/langflow/components/models/VertexAi.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/backend/langflow/components/models/VertexAi.py b/src/backend/langflow/components/models/VertexAi.py
index 2555aee0f..65f64578f 100644
--- a/src/backend/langflow/components/models/VertexAi.py
+++ b/src/backend/langflow/components/models/VertexAi.py
@@ -1,7 +1,7 @@
 from typing import List, Optional
 
-from langchain_google_vertexai import ChatVertexAI
 from langchain_core.messages.base import BaseMessage
+
 from langflow import CustomComponent
 from langflow.field_typing import Text
 
@@ -62,7 +62,7 @@ class ChatVertexAIComponent(CustomComponent):
 
     def build(
         self,
-        inputs:str,
+        inputs: str,
         credentials: Optional[str],
         project: str,
         examples: Optional[List[BaseMessage]] = [],
@@ -74,7 +74,13 @@ class ChatVertexAIComponent(CustomComponent):
         top_p: float = 0.95,
         verbose: bool = False,
     ) -> Text:
-        output =  ChatVertexAI(
+        try:
+            from langchain_google_vertexai import ChatVertexAI
+        except ImportError:
+            raise ImportError(
+                "To use the ChatVertexAI model, you need to install the langchain-google-vertexai package."
+            )
+        output = ChatVertexAI(
             credentials=credentials,
             examples=examples,
             location=location,