From 402ada6a198a3d961508ed388f9d8dffc3a73a9f Mon Sep 17 00:00:00 2001
From: Stefan Nica <stefan@zenml.io>
Date: Thu, 23 Nov 2023 10:58:39 +0100
Subject: [PATCH] Fix llama_index mypy errors

---
 .../llama_index/materializers/document_materializer.py    | 4 +++-
 .../llama_index/materializers/gpt_index_materializer.py   | 8 ++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/zenml/integrations/llama_index/materializers/document_materializer.py b/src/zenml/integrations/llama_index/materializers/document_materializer.py
index 9829232a50..2df8e22a1a 100644
--- a/src/zenml/integrations/llama_index/materializers/document_materializer.py
+++ b/src/zenml/integrations/llama_index/materializers/document_materializer.py
@@ -30,7 +30,9 @@
     LCDocument = Any
 else:
     from langchain.docstore.document import Document as LCDocument
-    from llama_index.readers.schema.base import Document
+    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
index 51d7070926..c04c368609 100644
--- 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])