fix: VertexAI KeyError 'location'

langflow-ai · Jul 26, 2024 · 2811507 · 2811507
1 parent 53084bc
commit 2811507
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py b/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py
@@ -17,7 +17,7 @@ class VertexAIEmbeddingsComponent(LCModelComponent):
             value="",
             file_types=["json"],
         ),
-        MessageTextInput(name="location", display_name="Location", advanced=True),
+        MessageTextInput(name="location", display_name="Location", value="us-central1", advanced=True),
         MessageTextInput(name="project", display_name="Project", info="The project ID.", advanced=True),
         IntInput(name="max_output_tokens", display_name="Max Output Tokens", advanced=True),
         IntInput(name="max_retries", display_name="Max Retries", value=1, advanced=True),
@@ -53,7 +53,7 @@ def build_embeddings(self) -> Embeddings:
         return VertexAIEmbeddings(
             credentials=gcloud_credentials,
             location=self.location,
-            max_output_tokens=self.max_output_tokens,
+            max_output_tokens=self.max_output_tokens or None,
             max_retries=self.max_retries,
             model_name=self.model_name,
             n=self.n,
@@ -62,6 +62,6 @@ def build_embeddings(self) -> Embeddings:
             stop=self.stop_sequences or None,
             streaming=self.streaming,
             temperature=self.temperature,
-            top_k=self.top_k,
+            top_k=self.top_k or None,
             top_p=self.top_p,
         )
diff --git a/src/backend/base/langflow/components/models/VertexAiModel.py b/src/backend/base/langflow/components/models/VertexAiModel.py
@@ -21,7 +21,7 @@ class ChatVertexAIComponent(LCModelComponent):
         ),
         MessageTextInput(name="model_name", display_name="Model Name", value="gemini-1.5-pro"),
         StrInput(name="project", display_name="Project", info="The project ID.", advanced=True),
-        StrInput(name="location", display_name="Location", advanced=True),
+        StrInput(name="location", display_name="Location", value="us-central1", advanced=True),
         IntInput(name="max_output_tokens", display_name="Max Output Tokens", advanced=True),
         IntInput(name="max_retries", display_name="Max Retries", value=1, advanced=True),
         FloatInput(name="temperature", value=0.0, display_name="Temperature"),
@@ -60,11 +60,11 @@ def build_model(self) -> LanguageModel:
                 credentials=credentials,
                 location=location,
                 project=project,
-                max_output_tokens=self.max_output_tokens,
+                max_output_tokens=self.max_output_tokens or None,
                 max_retries=self.max_retries,
                 model_name=self.model_name,
                 temperature=self.temperature,
-                top_k=self.top_k,
+                top_k=self.top_k or None,
                 top_p=self.top_p,
                 verbose=self.verbose,
             ),