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]