From decc979e996288c37b5220ccf6d1cd9d0790ad19 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Mon, 22 Jul 2024 12:03:36 -0300
Subject: [PATCH] refactor: disable type-var error code (#2868)

* fix: disable error code "type-var" in mypy configuration

* refactor(VertexAiModel.py): restructure the return statement to improve readability and maintainability by using type casting for LanguageModel
---
 pyproject.toml                                |  1 +
 .../components/models/VertexAiModel.py        | 29 +++++++++++--------
 src/backend/base/pyproject.toml               |  1 +
 3 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 19d554a5d..5905f3607 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -192,6 +192,7 @@ line-length = 120
 [tool.mypy]
 plugins = ["pydantic.mypy"]
 follow_imports = "silent"
+disable_error_code = ["type-var"]
 
 [build-system]
 requires = ["poetry-core"]
diff --git a/src/backend/base/langflow/components/models/VertexAiModel.py b/src/backend/base/langflow/components/models/VertexAiModel.py
index 763342dfc..b52650b0b 100644
--- a/src/backend/base/langflow/components/models/VertexAiModel.py
+++ b/src/backend/base/langflow/components/models/VertexAiModel.py
@@ -1,3 +1,5 @@
+from typing import cast
+
 from langflow.base.models.model import LCModelComponent
 from langflow.field_typing import LanguageModel
 from langflow.inputs import MessageTextInput
@@ -37,8 +39,8 @@ class ChatVertexAIComponent(LCModelComponent):
             )
         location = self.location or None
         if self.credentials:
-            from google.oauth2 import service_account
             from google.cloud import aiplatform
+            from google.oauth2 import service_account
 
             credentials = service_account.Credentials.from_service_account_file(self.credentials)
             project = self.project or credentials.project_id
@@ -52,15 +54,18 @@ class ChatVertexAIComponent(LCModelComponent):
             project = self.project or None
             credentials = None
 
-        return ChatVertexAI(
-            credentials=credentials,
-            location=location,
-            project=project,
-            max_output_tokens=self.max_output_tokens,
-            max_retries=self.max_retries,
-            model_name=self.model_name,
-            temperature=self.temperature,
-            top_k=self.top_k,
-            top_p=self.top_p,
-            verbose=self.verbose,
+        return cast(
+            LanguageModel,
+            ChatVertexAI(
+                credentials=credentials,
+                location=location,
+                project=project,
+                max_output_tokens=self.max_output_tokens,
+                max_retries=self.max_retries,
+                model_name=self.model_name,
+                temperature=self.temperature,
+                top_k=self.top_k,
+                top_p=self.top_p,
+                verbose=self.verbose,
+            ),
         )
diff --git a/src/backend/base/pyproject.toml b/src/backend/base/pyproject.toml
index 9e51aac66..e2787617d 100644
--- a/src/backend/base/pyproject.toml
+++ b/src/backend/base/pyproject.toml
@@ -97,6 +97,7 @@ markers = ["async_test"]
 namespace_packages = true
 mypy_path = "langflow"
 ignore_missing_imports = true
+disable_error_code = ["type-var"]
 
 
 [tool.ruff]