fix mypy

langflow-ai · Jul 2, 2024 · 5407133 · 5407133
1 parent fc86f39
commit 5407133
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 5 deletions.
diff --git a/src/backend/base/langflow/base/models/model.py b/src/backend/base/langflow/base/models/model.py
@@ -167,7 +167,7 @@ def get_chat_result(
             raise e
 
     @abstractmethod
-    def build_model(self) -> LanguageModel:
+    def build_model(self) -> LanguageModel:  # type: ignore
         """
         Implement this method to build the model.
         """
diff --git a/src/backend/base/langflow/components/models/CohereModel.py b/src/backend/base/langflow/components/models/CohereModel.py
@@ -33,7 +33,7 @@ class CohereComponent(LCModelComponent):
         ),
     ]
 
-    def build_model(self) -> LanguageModel | BaseChatModel:
+    def build_model(self) -> LanguageModel:
         cohere_api_key = self.cohere_api_key
         temperature = self.temperature
 

diff --git a/src/backend/base/langflow/components/vectorstores/AstraDB.py b/src/backend/base/langflow/components/vectorstores/AstraDB.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from loguru import logger
 
 from langchain_core.vectorstores import VectorStore
@@ -23,7 +25,7 @@ class AstraVectorStoreComponent(LCVectorStoreComponent):
     documentation: str = "https://python.langchain.com/docs/integrations/vectorstores/astradb"
     icon: str = "AstraDB"
 
-    _cached_vectorstore: VectorStore = None
+    _cached_vectorstore: Optional[VectorStore] = None
 
     inputs = [
         StrInput(

diff --git a/src/backend/base/langflow/components/vectorstores/Cassandra.py b/src/backend/base/langflow/components/vectorstores/Cassandra.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from langchain_community.vectorstores import Cassandra
 
@@ -24,7 +24,7 @@ class CassandraVectorStoreComponent(LCVectorStoreComponent):
     documentation = "https://python.langchain.com/docs/modules/data_connection/vectorstores/integrations/cassandra"
     icon = "Cassandra"
 
-    _cached_vectorstore: Cassandra = None
+    _cached_vectorstore: Optional[Cassandra] = None
 
     inputs = [
         MessageTextInput(