Cog 337 llama index support

.github/workflows/test_cognee_llama_index_notebook.yml

@@ -0,0 +1,63 @@
+name: test | llama index notebook
+
+on:
+  workflow_dispatch:
+  pull_request:
+    branches:
+      - main
+    types: [labeled, synchronize]
+
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  RUNTIME__LOG_LEVEL: ERROR
+
+jobs:
+  get_docs_changes:
+    name: docs changes
+    uses: ./.github/workflows/get_docs_changes.yml
+
+  run_notebook_test:
+    name: test
+    needs: get_docs_changes
+    if: needs.get_docs_changes.outputs.changes_outside_docs == 'true' && github.event.label.name == 'run-checks'
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Check out
+        uses: actions/checkout@master
+
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11.x'
+
+      - name: Install Poetry
+        uses: snok/[email protected]
+        with:
+          virtualenvs-create: true
+          virtualenvs-in-project: true
+          installer-parallel: true
+
+      - name: Install dependencies
+        run: |
+          poetry install --no-interaction --all-extras --no-root
+          poetry add jupyter --no-interaction
+
+      - name: Execute Jupyter Notebook
+        env:
+          ENV: 'dev'
+          LLM_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          GRAPHISTRY_USERNAME: ${{ secrets.GRAPHISTRY_USERNAME }}
+          GRAPHISTRY_PASSWORD: ${{ secrets.GRAPHISTRY_PASSWORD }}
+        run: |
+          poetry run jupyter nbconvert \
+          --to notebook \
+          --execute notebooks/cognee_llama_index.ipynb \
+          --output executed_notebook.ipynb \
+          --ExecutePreprocessor.timeout=1200

cognee/api/v1/add/add_v2.py

@@ -21,4 +21,4 @@ async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_nam
     pipeline = run_tasks(tasks, data, "add_pipeline")
 
     async for result in pipeline:
-        print(result)
+        print(result)

cognee/modules/users/permissions/methods/check_permission_on_documents.py

@@ -33,4 +33,4 @@ async def check_permission_on_documents(user: User, permission_type: str, docume
         has_permissions = all(document_id in resource_ids for document_id in document_ids)
 
         if not has_permissions:
-            raise PermissionDeniedException(f"User {user.username} does not have {permission_type} permission on documents")
+            raise PermissionDeniedException(f"User {user.email} does not have {permission_type} permission on documents")

cognee/tasks/ingestion/__init__.py

@@ -1,2 +1,4 @@
 from .ingest_data import ingest_data
 from .save_data_to_storage import save_data_to_storage
+from .save_data_item_to_storage import save_data_item_to_storage
+from .save_data_item_with_metadata_to_storage import save_data_item_with_metadata_to_storage

cognee/tasks/ingestion/ingest_data.py

@@ -3,7 +3,7 @@
 
 from cognee.shared.utils import send_telemetry
 from cognee.modules.users.models import User
-from cognee.infrastructure.databases.relational import get_relational_config, get_relational_engine
+from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.data.methods import create_dataset
 from cognee.modules.users.permissions.methods import give_permission_on_document
 from .get_dlt_destination import get_dlt_destination

cognee/tasks/ingestion/ingest_data_with_metadata.py

@@ -0,0 +1,92 @@
+import dlt
+import cognee.modules.ingestion as ingestion
+from typing import Any
+from cognee.shared.utils import send_telemetry
+from cognee.modules.users.models import User
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.data.methods import create_dataset
+from cognee.modules.users.permissions.methods import give_permission_on_document
+from .get_dlt_destination import get_dlt_destination
+from .save_data_item_with_metadata_to_storage import save_data_item_with_metadata_to_storage
+
+async def ingest_data_with_metadata(data: Any, dataset_name: str, user: User):
+    destination = get_dlt_destination()
+
+    pipeline = dlt.pipeline(
+        pipeline_name = "file_load_from_filesystem",
+        destination = destination,
+    )
+
+    @dlt.resource(standalone = True, merge_key = "id")
+    async def data_resources(data: Any, user: User):
+        if not isinstance(data, list):
+            # Convert data to a list as we work with lists further down.
+            data = [data]
+
+        # Process data
+        for data_item in data:
+
+            file_path = save_data_item_with_metadata_to_storage(data_item, dataset_name)
+
+            # Ingest data and add metadata
+            with open(file_path.replace("file://", ""), mode = "rb") as file:
+                classified_data = ingestion.classify(file)
+
+                data_id = ingestion.identify(classified_data)
+
+                file_metadata = classified_data.get_metadata()
+
+                from sqlalchemy import select
+                from cognee.modules.data.models import Data
+
+                db_engine = get_relational_engine()
+
+                async with db_engine.get_async_session() as session:
+                    dataset = await create_dataset(dataset_name, user.id, session)
+
+                    data_point = (await session.execute(
+                        select(Data).filter(Data.id == data_id)
+                    )).scalar_one_or_none()
+
+                    if data_point is not None:
+                        data_point.name = file_metadata["name"]
+                        data_point.raw_data_location = file_metadata["file_path"]
+                        data_point.extension = file_metadata["extension"]
+                        data_point.mime_type = file_metadata["mime_type"]
+
+                        await session.merge(data_point)
+                        await session.commit()
+                    else:
+                        data_point = Data(
+                            id = data_id,
+                            name = file_metadata["name"],
+                            raw_data_location = file_metadata["file_path"],
+                            extension = file_metadata["extension"],
+                            mime_type = file_metadata["mime_type"],
+                        )
+
+                        dataset.data.append(data_point)
+                        await session.commit()
+
+                yield {
+                    "id": data_id,
+                    "name": file_metadata["name"],
+                    "file_path": file_metadata["file_path"],
+                    "extension": file_metadata["extension"],
+                    "mime_type": file_metadata["mime_type"],
+                }
+
+                await give_permission_on_document(user, data_id, "read")
+                await give_permission_on_document(user, data_id, "write")
+
+
+    send_telemetry("cognee.add EXECUTION STARTED", user_id = user.id)
+    run_info = pipeline.run(
+        data_resources(data, user),
+        table_name = "file_metadata",
+        dataset_name = dataset_name,
+        write_disposition = "merge",
+    )
+    send_telemetry("cognee.add EXECUTION COMPLETED", user_id = user.id)
+
+    return run_info

cognee/tasks/ingestion/save_data_item_to_storage.py

@@ -0,0 +1,20 @@
+from typing import Union, BinaryIO
+from cognee.modules.ingestion import save_data_to_file
+
+def save_data_item_to_storage(data_item: Union[BinaryIO, str], dataset_name: str) -> str:
+
+    # data is a file object coming from upload.
+    if hasattr(data_item, "file"):
+        file_path = save_data_to_file(data_item.file, dataset_name, filename=data_item.filename)
+
+    elif isinstance(data_item, str):
+        # data is a file path
+        if data_item.startswith("file://") or data_item.startswith("/"):
+            file_path = data_item.replace("file://", "")
+        # data is text
+        else:
+            file_path = save_data_to_file(data_item, dataset_name)
+    else:
+        raise ValueError(f"Data type not supported: {type(data_item)}")
+
+    return file_path

cognee/tasks/ingestion/save_data_item_with_metadata_to_storage.py

@@ -0,0 +1,28 @@
+from typing import Union, BinaryIO, Any
+from cognee.modules.ingestion import save_data_to_file
+
+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)
+
+    # data is a file object coming from upload.
+    elif hasattr(data_item, "file"):
+        file_path = save_data_to_file(data_item.file, dataset_name, filename=data_item.filename)
+
+    elif isinstance(data_item, str):
+        # data is a file path
+        if data_item.startswith("file://") or data_item.startswith("/"):
+            file_path = data_item.replace("file://", "")
+        # data is text
+        else:
+            file_path = save_data_to_file(data_item, dataset_name)
+    else:
+        raise ValueError(f"Data type not supported: {type(data_item)}")
+
+    return file_path

cognee/tasks/ingestion/save_data_to_storage.py

@@ -1,5 +1,5 @@
 from typing import Union, BinaryIO
-from cognee.modules.ingestion import save_data_to_file
+from cognee.tasks.ingestion.save_data_item_to_storage import save_data_item_to_storage
 
 def save_data_to_storage(data: Union[BinaryIO, str], dataset_name) -> list[str]:
     if not isinstance(data, list):
@@ -9,19 +9,7 @@ def save_data_to_storage(data: Union[BinaryIO, str], dataset_name) -> list[str]:
     file_paths = []
 
     for data_item in data:
-        # data is a file object coming from upload.
-        if hasattr(data_item, "file"):
-            file_path = save_data_to_file(data_item.file, dataset_name, filename = data_item.filename)
-            file_paths.append(file_path)
-
-        if isinstance(data_item, str):
-            # data is a file path
-            if data_item.startswith("file://") or data_item.startswith("/"):
-                file_paths.append(data_item.replace("file://", ""))
-
-            # data is text
-            else:
-                file_path = save_data_to_file(data_item, dataset_name)
-                file_paths.append(file_path)
+        file_path = save_data_item_to_storage(data_item, dataset_name)
+        file_paths.append(file_path)
 
     return file_paths

cognee/tasks/ingestion/transform_data.py

@@ -0,0 +1,18 @@
+from llama_index.core import Document
+from llama_index.core.schema import ImageDocument
+from cognee.modules.ingestion import save_data_to_file
+from typing import Union
+
+def get_data_from_llama_index(data_point: Union[Document, ImageDocument], dataset_name: str) -> str:
+    # Specific type checking is used to ensure it's not a child class from Document
+    if type(data_point) == Document:
+        file_path = data_point.metadata.get("file_path")
+        if file_path is None:
+            file_path = save_data_to_file(data_point.text, dataset_name)
+            return file_path
+        return file_path
+    elif type(data_point) == ImageDocument:
+        if data_point.image_path is None:
+            file_path = save_data_to_file(data_point.text, dataset_name)
+            return file_path
+        return data_point.image_path

examples/python/simple_example.py

@@ -27,8 +27,7 @@ async def main():
 
     # Query cognee for insights on the added text
     search_results = await cognee.search(
-        SearchType.INSIGHTS,
-        {'query': 'Tell me about NLP'}
+        SearchType.INSIGHTS, query='Tell me about NLP'
     )
 
     # Display search results