Fix llama_index mypy errors

zenml-io · Nov 23, 2023 · d0af0c4 · d0af0c4
1 parent 1d7c33f
commit d0af0c4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/zenml/integrations/llama_index/materializers/document_materializer.py b/src/zenml/integrations/llama_index/materializers/document_materializer.py
@@ -29,8 +29,12 @@
     Document = Any
     LCDocument = Any
 else:
-    from langchain.docstore.document import Document as LCDocument
-    from llama_index.readers.schema.base import Document
+    from langchain.docstore.document import (  # type: ignore[import]
+        Document as LCDocument,
+    )
+    from llama_index.readers.schema.base import (  # type: ignore[import]
+        Document,
+    )
 
 
 class LlamaIndexDocumentMaterializer(LangchainDocumentMaterializer):

diff --git a/src/zenml/integrations/llama_index/materializers/gpt_index_materializer.py b/src/zenml/integrations/llama_index/materializers/gpt_index_materializer.py
@@ -39,8 +39,12 @@
     GPTFaissIndex = Any
     T = TypeVar("T", bound=Any)
 else:
-    from llama_index.indices.base import BaseGPTIndex
-    from llama_index.indices.vector_store import GPTFaissIndex
+    from llama_index.indices.base import (  # type: ignore[import]
+        BaseGPTIndex,
+    )
+    from llama_index.indices.vector_store import (  # type: ignore[import]
+        GPTFaissIndex,
+    )
 
     T = TypeVar("T", bound=BaseGPTIndex[Any])