From c4700b1027270e8fd0ace205cfdcc41f8060137f Mon Sep 17 00:00:00 2001
From: Igor Ilic <igorilic03@gmail.com>
Date: Fri, 15 Nov 2024 17:16:36 +0100
Subject: [PATCH] fix: Add install of llama-index to notebook

Removed additional references to llama-index from core cognee lib.
Added llama-index-core install from notebook

Fix COG-337
---
 .../save_data_item_with_metadata_to_storage.py      | 11 +++++++----
 notebooks/cognee_llama_index.ipynb                  | 13 +++++++++++--
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
index acb128088..ec29edb89 100644
--- a/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
+++ b/cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py
@@ -1,9 +1,12 @@
-from llama_index.core import Document
-from typing import Union, BinaryIO
+from typing import Union, BinaryIO, Any
 from cognee.modules.ingestion import save_data_to_file
-from .transform_data import get_data_from_llama_index
 
-def save_data_item_with_metadata_to_storage(data_item: Union[BinaryIO, Document, str], dataset_name: str) -> str:
+def save_data_item_with_metadata_to_storage(data_item: Union[BinaryIO, str, Any], dataset_name: str) -> str:
+    # Dynamic import is used because the llama_index module is optional. 
+    # For the same reason Any is accepted as a data item
+    from llama_index.core import Document
+    from .transform_data import get_data_from_llama_index
+
     # Check if data is of type Document or any of it's subclasses
     if isinstance(data_item, Document):
         file_path = get_data_from_llama_index(data_item, dataset_name)
diff --git a/notebooks/cognee_llama_index.ipynb b/notebooks/cognee_llama_index.ipynb
index 04c93486d..15e17163d 100644
--- a/notebooks/cognee_llama_index.ipynb
+++ b/notebooks/cognee_llama_index.ipynb
@@ -7,6 +7,15 @@
     "## Cognee GraphRAG with LlamaIndex Documents"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%pip install llama-index-core"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -43,7 +52,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -62,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [