cognee 0.2.1.dev7__py3-none-any.whl → 0.2.2.dev1__py3-none-any.whl

cognee/tasks/graph/extract_graph_from_code.py

@@ -1,7 +1,8 @@
 import asyncio
 from typing import Type, List
 from pydantic import BaseModel
-from cognee.modules.data.extraction.knowledge_graph import extract_content_graph
+
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
 from cognee.tasks.storage import add_data_points
 
@@ -17,7 +18,7 @@ async def extract_graph_from_code(
         - Graph nodes are stored using the `add_data_points` function for later retrieval or analysis.
     """
     chunk_graphs = await asyncio.gather(
-        *[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
+        *[LLMGateway.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
     )
 
     for chunk_index, chunk in enumerate(data_chunks):

cognee/tasks/graph/extract_graph_from_data.py

@@ -3,15 +3,15 @@ from typing import Type, List
 from pydantic import BaseModel
 
 from cognee.infrastructure.databases.graph import get_graph_engine
+from cognee.tasks.storage.add_data_points import add_data_points
 from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
-from cognee.modules.data.extraction.knowledge_graph import extract_content_graph
 from cognee.modules.graph.utils import (
     expand_with_nodes_and_edges,
     retrieve_existing_edges,
 )
 from cognee.shared.data_models import KnowledgeGraph
-from cognee.tasks.storage.add_data_points import add_data_points
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 
 async def integrate_chunk_graphs(
@@ -40,6 +40,7 @@ async def integrate_chunk_graphs(
 
     if len(graph_nodes) > 0:
         await add_data_points(graph_nodes)
+
     if len(graph_edges) > 0:
         await graph_engine.add_edges(graph_edges)
 
@@ -55,7 +56,7 @@ async def extract_graph_from_data(
     Extracts and integrates a knowledge graph from the text content of document chunks using a specified graph model.
     """
     chunk_graphs = await asyncio.gather(
-        *[extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
+        *[LLMGateway.extract_content_graph(chunk.text, graph_model) for chunk in data_chunks]
     )
 
     # Note: Filter edges with missing source or target nodes

cognee/tasks/graph/infer_data_ontology.py

@@ -15,19 +15,19 @@ from pydantic import BaseModel
 
 from cognee.modules.graph.exceptions import EntityNotFoundError
 from cognee.modules.ingestion.exceptions import IngestionError
-from cognee.infrastructure.llm.prompts import read_query_prompt
-from cognee.infrastructure.llm.get_llm_client import get_llm_client
+
 from cognee.infrastructure.data.chunking.config import get_chunk_config
 from cognee.infrastructure.data.chunking.get_chunking_engine import get_chunk_engine
 from cognee.infrastructure.databases.graph.get_graph_engine import get_graph_engine
 from cognee.infrastructure.files.utils.extract_text_from_file import extract_text_from_file
 from cognee.infrastructure.files.utils.guess_file_type import guess_file_type, FileTypeException
-from cognee.modules.data.extraction.knowledge_graph.add_model_class_to_graph import (
+from cognee.modules.data.methods.add_model_class_to_graph import (
     add_model_class_to_graph,
 )
 from cognee.tasks.graph.models import NodeModel, GraphOntology
 from cognee.shared.data_models import KnowledgeGraph
 from cognee.modules.engine.utils import generate_node_id, generate_node_name
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 
 logger = get_logger("task:infer_data_ontology")
 
@@ -52,11 +52,10 @@ async def extract_ontology(content: str, response_model: Type[BaseModel]):
 
         The structured ontology extracted from the content.
     """
-    llm_client = get_llm_client()
 
-    system_prompt = read_query_prompt("extract_ontology.txt")
+    system_prompt = LLMGateway.read_query_prompt("extract_ontology.txt")
 
-    ontology = await llm_client.acreate_structured_output(content, system_prompt, response_model)
+    ontology = await LLMGateway.acreate_structured_output(content, system_prompt, response_model)
 
     return ontology
 

cognee/tasks/ingestion/data_item_to_text_file.py

@@ -0,0 +1,79 @@
+import os
+from urllib.parse import urlparse
+from typing import List, Tuple
+from pathlib import Path
+import tempfile
+
+from cognee.infrastructure.loaders.LoaderInterface import LoaderInterface
+from cognee.modules.ingestion.exceptions import IngestionError
+from cognee.infrastructure.loaders import get_loader_engine
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.files.utils.open_data_file import open_data_file
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+logger = get_logger(__name__)
+
+
+class SaveDataSettings(BaseSettings):
+    accept_local_file_path: bool = True
+
+    model_config = SettingsConfigDict(env_file=".env", extra="allow")
+
+
+settings = SaveDataSettings()
+
+
+async def pull_from_s3(file_path, destination_file) -> None:
+    async with open_data_file(file_path) as file:
+        while True:
+            chunk = file.read(8192)
+            if not chunk:
+                break
+            destination_file.write(chunk)
+
+
+async def data_item_to_text_file(
+    data_item_path: str, preferred_loaders: List[str]
+) -> Tuple[str, LoaderInterface]:
+    if isinstance(data_item_path, str):
+        parsed_url = urlparse(data_item_path)
+
+        # data is s3 file path
+        if parsed_url.scheme == "s3":
+            # TODO: Rework this to work with file streams and not saving data to temp storage
+            # Note: proper suffix information is needed for OpenAI to handle mp3 files
+            path_info = Path(parsed_url.path)
+            with tempfile.NamedTemporaryFile(mode="wb", suffix=path_info.suffix) as temp_file:
+                await pull_from_s3(data_item_path, temp_file)
+                temp_file.flush()  # Data needs to be saved to local storage
+                loader = get_loader_engine()
+                return await loader.load_file(temp_file.name, preferred_loaders), loader.get_loader(
+                    temp_file.name, preferred_loaders
+                )
+
+        # data is local file path
+        elif parsed_url.scheme == "file":
+            if settings.accept_local_file_path:
+                loader = get_loader_engine()
+                return await loader.load_file(data_item_path, preferred_loaders), loader.get_loader(
+                    data_item_path, preferred_loaders
+                )
+            else:
+                raise IngestionError(message="Local files are not accepted.")
+
+        # data is an absolute file path
+        elif data_item_path.startswith("/") or (
+            os.name == "nt" and len(data_item_path) > 1 and data_item_path[1] == ":"
+        ):
+            # Handle both Unix absolute paths (/path) and Windows absolute paths (C:\path)
+            if settings.accept_local_file_path:
+                loader = get_loader_engine()
+                return await loader.load_file(data_item_path, preferred_loaders), loader.get_loader(
+                    data_item_path, preferred_loaders
+                )
+            else:
+                raise IngestionError(message="Local files are not accepted.")
+
+    # data is not a supported type
+    raise IngestionError(message=f"Data type not supported: {type(data_item_path)}")

cognee/tasks/ingestion/ingest_data.py

@@ -1,16 +1,16 @@
 import json
 import inspect
-from os import path
 from uuid import UUID
 from typing import Union, BinaryIO, Any, List, Optional
 
 import cognee.modules.ingestion as ingestion
-from cognee.infrastructure.files.utils.open_data_file import open_data_file
 from cognee.infrastructure.databases.relational import get_relational_engine
 from cognee.modules.data.models import Data
 from cognee.modules.users.models import User
 from cognee.modules.users.methods import get_default_user
 from cognee.modules.users.permissions.methods import get_specific_user_permission_datasets
+from cognee.infrastructure.files.utils.open_data_file import open_data_file
+from cognee.infrastructure.files.utils.get_data_file_path import get_data_file_path
 from cognee.modules.data.methods import (
     get_authorized_existing_datasets,
     get_dataset_data,
@@ -18,6 +18,7 @@ from cognee.modules.data.methods import (
 )
 
 from .save_data_item_to_storage import save_data_item_to_storage
+from .data_item_to_text_file import data_item_to_text_file
 
 
 async def ingest_data(
@@ -26,6 +27,7 @@ async def ingest_data(
     user: User,
     node_set: Optional[List[str]] = None,
     dataset_id: UUID = None,
+    preferred_loaders: List[str] = None,
 ):
     if not user:
         user = await get_default_user()
@@ -42,6 +44,7 @@ async def ingest_data(
         user: User,
         node_set: Optional[List[str]] = None,
         dataset_id: UUID = None,
+        preferred_loaders: List[str] = None,
     ):
         new_datapoints = []
         existing_data_points = []
@@ -74,71 +77,96 @@ async def ingest_data(
         dataset_data_map = {str(data.id): True for data in dataset_data}
 
         for data_item in data:
-            file_path = await save_data_item_to_storage(data_item)
+            # Get file path of data item or create a file it doesn't exist
+            original_file_path = await save_data_item_to_storage(data_item)
+
+            # Transform file path to be OS usable
+            actual_file_path = get_data_file_path(original_file_path)
 
-            # Ingest data and add metadata
-            async with open_data_file(file_path) as file:
+            # Store all input data as text files in Cognee data storage
+            cognee_storage_file_path, loader_engine = await data_item_to_text_file(
+                actual_file_path, preferred_loaders
+            )
+
+            # Find metadata from original file
+            async with open_data_file(original_file_path) as file:
                 classified_data = ingestion.classify(file)
 
-                # data_id is the hash of file contents + owner id to avoid duplicate data
+                # data_id is the hash of original file contents + owner id to avoid duplicate data
                 data_id = ingestion.identify(classified_data, user)
+                original_file_metadata = classified_data.get_metadata()
 
-                file_metadata = classified_data.get_metadata()
-
-                from sqlalchemy import select
-
-                db_engine = get_relational_engine()
-
-                # Check to see if data should be updated
-                async with db_engine.get_async_session() as session:
-                    data_point = (
-                        await session.execute(select(Data).filter(Data.id == data_id))
-                    ).scalar_one_or_none()
-
-                ext_metadata = get_external_metadata_dict(data_item)
-
-                if node_set:
-                    ext_metadata["node_set"] = node_set
-
-                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"]
-                    data_point.owner_id = user.id
-                    data_point.content_hash = file_metadata["content_hash"]
-                    data_point.file_size = file_metadata["file_size"]
-                    data_point.external_metadata = ext_metadata
-                    data_point.node_set = json.dumps(node_set) if node_set else None
-                    data_point.tenant_id = user.tenant_id if user.tenant_id else None
-
-                    # Check if data is already in dataset
-                    if str(data_point.id) in dataset_data_map:
-                        existing_data_points.append(data_point)
-                    else:
-                        dataset_new_data_points.append(data_point)
-                        dataset_data_map[str(data_point.id)] = True
+            # Find metadata from Cognee data storage text file
+            async with open_data_file(cognee_storage_file_path) as file:
+                classified_data = ingestion.classify(file)
+                storage_file_metadata = classified_data.get_metadata()
+
+            from sqlalchemy import select
+
+            db_engine = get_relational_engine()
+
+            # Check to see if data should be updated
+            async with db_engine.get_async_session() as session:
+                data_point = (
+                    await session.execute(select(Data).filter(Data.id == data_id))
+                ).scalar_one_or_none()
+
+            # TODO: Maybe allow getting of external metadata through ingestion loader?
+            ext_metadata = get_external_metadata_dict(data_item)
+
+            if node_set:
+                ext_metadata["node_set"] = node_set
+
+            if data_point is not None:
+                data_point.name = original_file_metadata["name"]
+                data_point.raw_data_location = cognee_storage_file_path
+                data_point.original_data_location = original_file_metadata["file_path"]
+                data_point.extension = storage_file_metadata["extension"]
+                data_point.mime_type = storage_file_metadata["mime_type"]
+                data_point.original_extension = original_file_metadata["extension"]
+                data_point.original_mime_type = original_file_metadata["mime_type"]
+                data_point.loader_engine = loader_engine.loader_name
+                data_point.owner_id = user.id
+                data_point.content_hash = original_file_metadata["content_hash"]
+                data_point.raw_content_hash = storage_file_metadata["content_hash"]
+                data_point.file_size = original_file_metadata["file_size"]
+                data_point.external_metadata = ext_metadata
+                data_point.node_set = json.dumps(node_set) if node_set else None
+                data_point.tenant_id = user.tenant_id if user.tenant_id else None
+
+                # Check if data is already in dataset
+                if str(data_point.id) in dataset_data_map:
+                    existing_data_points.append(data_point)
                 else:
-                    if str(data_id) in dataset_data_map:
-                        continue
-
-                    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"],
-                        owner_id=user.id,
-                        content_hash=file_metadata["content_hash"],
-                        external_metadata=ext_metadata,
-                        node_set=json.dumps(node_set) if node_set else None,
-                        data_size=file_metadata["file_size"],
-                        tenant_id=user.tenant_id if user.tenant_id else None,
-                        token_count=-1,
-                    )
-
-                    new_datapoints.append(data_point)
+                    dataset_new_data_points.append(data_point)
                     dataset_data_map[str(data_point.id)] = True
+            else:
+                if str(data_id) in dataset_data_map:
+                    continue
+
+                data_point = Data(
+                    id=data_id,
+                    name=original_file_metadata["name"],
+                    raw_data_location=cognee_storage_file_path,
+                    original_data_location=original_file_metadata["file_path"],
+                    extension=storage_file_metadata["extension"],
+                    mime_type=storage_file_metadata["mime_type"],
+                    original_extension=original_file_metadata["extension"],
+                    original_mime_type=original_file_metadata["mime_type"],
+                    loader_engine=loader_engine.loader_name,
+                    owner_id=user.id,
+                    content_hash=original_file_metadata["content_hash"],
+                    raw_content_hash=storage_file_metadata["content_hash"],
+                    external_metadata=ext_metadata,
+                    node_set=json.dumps(node_set) if node_set else None,
+                    data_size=original_file_metadata["file_size"],
+                    tenant_id=user.tenant_id if user.tenant_id else None,
+                    pipeline_status={},
+                    token_count=-1,
+                )
+
+                new_datapoints.append(data_point)
+                dataset_data_map[str(data_point.id)] = True
 
         async with db_engine.get_async_session() as session:
             if dataset not in session:
@@ -160,4 +188,6 @@ async def ingest_data(
 
         return existing_data_points + dataset_new_data_points + new_datapoints
 
-    return await store_data_to_dataset(data, dataset_name, user, node_set, dataset_id)
+    return await store_data_to_dataset(
+        data, dataset_name, user, node_set, dataset_id, preferred_loaders
+    )

cognee/tasks/ingestion/resolve_data_directories.py

@@ -40,6 +40,9 @@ async def resolve_data_directories(
                     if include_subdirectories:
                         base_path = item if item.endswith("/") else item + "/"
                         s3_keys = fs.glob(base_path + "**")
+                        # If path is not directory attempt to add item directly
+                        if not s3_keys:
+                            s3_keys = fs.ls(item)
                     else:
                         s3_keys = fs.ls(item)
                     # Filter out keys that represent directories using fs.isdir

cognee/tasks/repo_processor/get_repo_file_dependencies.py

@@ -103,6 +103,9 @@ async def get_repo_file_dependencies(
           extraction of dependencies (default is False). (default False)
     """
 
+    if isinstance(repo_path, list) and len(repo_path) == 1:
+        repo_path = repo_path[0]
+
     if not os.path.exists(repo_path):
         raise FileNotFoundError(f"Repository path {repo_path} does not exist.")
 

cognee/tasks/storage/index_data_points.py

@@ -38,7 +38,7 @@ async def index_data_points(data_points: list[DataPoint]):
         index_name = index_name_and_field[:first_occurence]
         field_name = index_name_and_field[first_occurence + 1 :]
         try:
-            # In case the ammount if indexable points is too large we need to send them in batches
+            # In case the amount of indexable points is too large we need to send them in batches
             batch_size = 100
             for i in range(0, len(indexable_points), batch_size):
                 batch = indexable_points[i : i + batch_size]

cognee/tasks/storage/index_graph_edges.py

@@ -1,3 +1,4 @@
+from cognee.modules.engine.utils.generate_edge_id import generate_edge_id
 from cognee.shared.logging_utils import get_logger, ERROR
 from collections import Counter
 
@@ -49,7 +50,9 @@ async def index_graph_edges(batch_size: int = 1024):
     )
 
     for text, count in edge_types.items():
-        edge = EdgeType(relationship_name=text, number_of_edges=count)
+        edge = EdgeType(
+            id=generate_edge_id(edge_id=text), relationship_name=text, number_of_edges=count
+        )
         data_point_type = type(edge)
 
         for field_name in edge.metadata["index_fields"]:

cognee/tasks/summarization/summarize_code.py

@@ -3,8 +3,7 @@ from typing import AsyncGenerator, Union
 from uuid import uuid5
 
 from cognee.infrastructure.engine import DataPoint
-from cognee.modules.data.extraction.extract_summary import extract_code_summary
-
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from .models import CodeSummary
 
 
@@ -17,7 +16,7 @@ async def summarize_code(
     code_data_points = [file for file in code_graph_nodes if hasattr(file, "source_code")]
 
     file_summaries = await asyncio.gather(
-        *[extract_code_summary(file.source_code) for file in code_data_points]
+        *[LLMGateway.extract_code_summary(file.source_code) for file in code_data_points]
     )
 
     file_summaries_map = {

cognee/tasks/summarization/summarize_text.py

@@ -2,8 +2,9 @@ import asyncio
 from typing import Type
 from uuid import uuid5
 from pydantic import BaseModel
-from cognee.modules.data.extraction.extract_summary import extract_summary
+
 from cognee.modules.chunking.models.DocumentChunk import DocumentChunk
+from cognee.infrastructure.llm.LLMGateway import LLMGateway
 from cognee.modules.cognify.config import get_cognify_config
 from .models import TextSummary
 
@@ -42,7 +43,7 @@ async def summarize_text(
         summarization_model = cognee_config.summarization_model
 
     chunk_summaries = await asyncio.gather(
-        *[extract_summary(chunk.text, summarization_model) for chunk in data_chunks]
+        *[LLMGateway.extract_summary(chunk.text, summarization_model) for chunk in data_chunks]
     )
 
     summaries = [

cognee/tests/test_cognee_server_start.py

@@ -6,6 +6,7 @@ import signal
 import requests
 from pathlib import Path
 import sys
+import uuid
 
 
 class TestCogneeServerStart(unittest.TestCase):
@@ -47,7 +48,7 @@ class TestCogneeServerStart(unittest.TestCase):
         """Test that the server is running and can accept connections."""
         # Test health endpoint
         health_response = requests.get("http://localhost:8000/health", timeout=15)
-        self.assertEqual(health_response.status_code, 200)
+        self.assertIn(health_response.status_code, [200, 503])
 
         # Test root endpoint
         root_response = requests.get("http://localhost:8000/", timeout=15)
@@ -74,7 +75,8 @@ class TestCogneeServerStart(unittest.TestCase):
         file_path = Path(os.path.join(Path(__file__).parent, "test_data/example.png"))
         headers = {"Authorization": auth_var}
 
-        form_data = {"datasetName": "test"}
+        dataset_name = f"test_{uuid.uuid4().hex[:8]}"
+        form_data = {"datasetName": dataset_name}
 
         file = {
             "data": (
@@ -83,8 +85,11 @@ class TestCogneeServerStart(unittest.TestCase):
             )
         }
 
+        payload = {"datasets": [dataset_name]}
+
         add_response = requests.post(url, headers=headers, data=form_data, files=file, timeout=50)
-        add_response.raise_for_status()  # raise if HTTP 4xx/5xx
+        if add_response.status_code not in [200, 201, 409]:
+            add_response.raise_for_status()
 
         # Cognify request
         url = "http://127.0.0.1:8000/api/v1/cognify"
@@ -93,10 +98,9 @@ class TestCogneeServerStart(unittest.TestCase):
             "Content-Type": "application/json",
         }
 
-        payload = {"datasets": ["test"]}
-
         cognify_response = requests.post(url, headers=headers, json=payload, timeout=150)
-        cognify_response.raise_for_status()  # raises on HTTP 4xx/5xx
+        if cognify_response.status_code not in [200, 201, 409]:
+            cognify_response.raise_for_status()
 
         # TODO: Add test to verify cognify pipeline is complete before testing search
 
@@ -111,7 +115,8 @@ class TestCogneeServerStart(unittest.TestCase):
         payload = {"searchType": "GRAPH_COMPLETION", "query": "What's in the document?"}
 
         search_response = requests.post(url, headers=headers, json=payload, timeout=50)
-        search_response.raise_for_status()  # raises on HTTP 4xx/5xx
+        if search_response.status_code not in [200, 201, 409]:
+            search_response.raise_for_status()
 
 
 if __name__ == "__main__":

cognee/tests/test_deduplication.py

@@ -26,8 +26,8 @@ async def test_deduplication():
     explanation_file_path2 = os.path.join(
         pathlib.Path(__file__).parent, "test_data/Natural_language_processing_copy.txt"
     )
-    await cognee.add([explanation_file_path], dataset_name)
-    await cognee.add([explanation_file_path2], dataset_name2)
+    await cognee.add([explanation_file_path], dataset_name, incremental_loading=False)
+    await cognee.add([explanation_file_path2], dataset_name2, incremental_loading=False)
 
     result = await relational_engine.get_all_data_from_table("data")
     assert len(result) == 1, "More than one data entity was found."

cognee/tests/test_deletion.py

@@ -12,7 +12,21 @@ async def main():
     await cognee.prune.prune_data()
     await cognee.prune.prune_system(metadata=True)
 
-    text_1 = """
+    pdf_document = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/artificial-intelligence.pdf"
+    )
+
+    txt_document = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Natural_language_processing_copy.txt"
+    )
+
+    audio_document = os.path.join(pathlib.Path(__file__).parent, "test_data/text_to_speech.mp3")
+
+    image_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.png")
+
+    unstructured_document = os.path.join(pathlib.Path(__file__).parent, "test_data/example.pptx")
+
+    text_document_as_literal = """
     1. Audi
     Audi is known for its modern designs and advanced technology. Founded in the early 1900s, the brand has earned a reputation for precision engineering and innovation. With features like the Quattro all-wheel-drive system, Audi offers a range of vehicles from stylish sedans to high-performance sports cars.
 
@@ -31,27 +45,54 @@ async def main():
     Each of these car manufacturer contributes to Germany's reputation as a leader in the global automotive industry, showcasing a blend of innovation, performance, and design excellence.
     """
 
-    text_2 = """
-    1. Apple
-    Apple is renowned for its innovative consumer electronics and software. Its product lineup includes the iPhone, iPad, Mac computers, and wearables like the Apple Watch. Known for its emphasis on sleek design and user-friendly interfaces, Apple has built a loyal customer base and created a seamless ecosystem that integrates hardware, software, and services.
+    ################### HARD DELETE
 
-    2. Google
-    Founded in 1998, Google started as a search engine and quickly became the go-to resource for finding information online. Over the years, the company has diversified its offerings to include digital advertising, cloud computing, mobile operating systems (Android), and various web services like Gmail and Google Maps. Google's innovations have played a major role in shaping the internet landscape.
+    # Add documents and get dataset information
+    add_result = await cognee.add(
+        [
+            pdf_document,
+            txt_document,
+            text_document_as_literal,
+            unstructured_document,
+            audio_document,
+            image_document,
+        ]
+    )
+    dataset_id = add_result.dataset_id
 
-    3. Microsoft
-    Microsoft Corporation has been a dominant force in software for decades. Its Windows operating system and Microsoft Office suite are staples in both business and personal computing. In recent years, Microsoft has expanded into cloud computing with Azure, gaming with the Xbox platform, and even hardware through products like the Surface line. This evolution has helped the company maintain its relevance in a rapidly changing tech world.
+    await cognee.cognify()
 
-    4. Amazon
-    What began as an online bookstore has grown into one of the largest e-commerce platforms globally. Amazon is known for its vast online marketplace, but its influence extends far beyond retail. With Amazon Web Services (AWS), the company has become a leader in cloud computing, offering robust solutions that power websites, applications, and businesses around the world. Amazon's constant drive for innovation continues to reshape both retail and technology sectors.
+    from cognee.infrastructure.databases.graph import get_graph_engine
 
-    5. Meta
-    Meta, originally known as Facebook, revolutionized social media by connecting billions of people worldwide. Beyond its core social networking service, Meta is investing in the next generation of digital experiences through virtual and augmented reality technologies, with projects like Oculus. The company's efforts signal a commitment to evolving digital interaction and building the metaverse—a shared virtual space where users can connect and collaborate.
+    graph_engine = await get_graph_engine()
+    nodes, edges = await graph_engine.get_graph_data()
+    assert len(nodes) > 10 and len(edges) > 10, "Graph database is not loaded."
 
-    Each of these companies has significantly impacted the technology landscape, driving innovation and transforming everyday life through their groundbreaking products and services.
-    """
+    # Get the data IDs from the dataset
+    dataset_data = await get_dataset_data(dataset_id)
+    assert len(dataset_data) > 0, "Dataset should contain data"
+
+    # Delete each document using its ID
+    for data_item in dataset_data:
+        await cognee.delete(data_item.id, dataset_id, mode="hard")
+
+    nodes, edges = await graph_engine.get_graph_data()
+
+    assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with hard delete."
+
+    ################### SOFT DELETE
 
     # Add documents and get dataset information
-    add_result = await cognee.add([text_1, text_2])
+    add_result = await cognee.add(
+        [
+            pdf_document,
+            txt_document,
+            text_document_as_literal,
+            unstructured_document,
+            audio_document,
+            image_document,
+        ]
+    )
     dataset_id = add_result.dataset_id
 
     await cognee.cognify()
@@ -68,11 +109,11 @@ async def main():
 
     # Delete each document using its ID
     for data_item in dataset_data:
-        await cognee.delete(data_item.id, dataset_id, mode="hard")
+        await cognee.delete(data_item.id, dataset_id, mode="soft")
 
     nodes, edges = await graph_engine.get_graph_data()
 
-    assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted."
+    assert len(nodes) == 0 and len(edges) == 0, "Document is not deleted with soft delete."
 
 
 if __name__ == "__main__":