Skip to content

Commit

Permalink
refactor: disable type-var error code (#2868)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
ogabrielluiz authored Jul 22, 2024
1 parent 724f540 commit decc979
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 12 deletions.
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down
29 changes: 17 additions & 12 deletions src/backend/base/langflow/components/models/VertexAiModel.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -37,8 +39,8 @@ def build_model(self) -> LanguageModel:
)
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
Expand All @@ -52,15 +54,18 @@ def build_model(self) -> LanguageModel:
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,
),
)
1 change: 1 addition & 0 deletions src/backend/base/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ markers = ["async_test"]
namespace_packages = true
mypy_path = "langflow"
ignore_missing_imports = true
disable_error_code = ["type-var"]


[tool.ruff]
Expand Down

0 comments on commit decc979

Please sign in to comment.