Fix llama-index test after upgrading deps (#1192)

databricks · Feb 23, 2024 · 33dcf03 · 33dcf03
1 parent 521fb51
commit 33dcf03
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/lilac/sources/github_source.py b/lilac/sources/github_source.py
@@ -74,7 +74,7 @@ def scrub_github_token(self, github_token: str) -> str:
   @override
   def setup(self) -> None:
     try:
-      from llama_index import download_loader
+      from llama_index.core.readers import download_loader
 
     except ImportError:
       raise ImportError(

diff --git a/lilac/sources/llama_index_docs_source.py b/lilac/sources/llama_index_docs_source.py
@@ -2,7 +2,7 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Iterable, Iterator, Optional
 
 if TYPE_CHECKING:
-  from llama_index import Document
+  from llama_index.core.schema import Document
 
 from typing_extensions import override
 

diff --git a/lilac/sources/llama_index_docs_source_test.py b/lilac/sources/llama_index_docs_source_test.py
@@ -1,6 +1,6 @@
 """Tests for the Llama Index Docs source."""
 
-from llama_index import Document
+from llama_index.core.schema import Document
 
 from ..schema import schema
 from ..source import SourceSchema