Only init vertex if the model is enabled

slack-samples · Sep 25, 2024 · ea8744c · ea8744c
1 parent e2171c1
commit ea8744c
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ai/providers/vertexai.py b/ai/providers/vertexai.py
@@ -61,7 +61,11 @@ class VertexAPI(BaseAPIProvider):
 
     def __init__(self):
         self.enabled = bool(os.environ.get("VERTEX_AI_PROJECT_ID", ""))
-        vertexai.init(project=os.environ.get("VERTEX_AI_PROJECT_ID"), location=os.environ.get("VERTEX_AI_LOCATION"))
+        if self.enabled:
+            vertexai.init(
+                project=os.environ.get("VERTEX_AI_PROJECT_ID"),
+                location=os.environ.get("VERTEX_AI_LOCATION"),
+            )
 
     def set_model(self, model_name: str):
         if model_name not in self.MODELS.keys():