cognee 0.3.5__py3-none-any.whl → 0.3.7__py3-none-any.whl

{cognee-0.3.5.dist-info → cognee-0.3.7.dist-info}/entry_points.txt

@@ -1,2 +1,3 @@
 [console_scripts]
+cognee = cognee.cli._cognee:main
 cognee-cli = cognee.cli._cognee:main

distributed/Dockerfile

@@ -29,6 +29,3 @@ RUN poetry install --extras neo4j --extras postgres --extras aws --extras distri
 
 COPY cognee/ /app/cognee
 COPY distributed/ /app/distributed
-RUN chmod +x /app/distributed/entrypoint.sh
-
-ENTRYPOINT ["/app/distributed/entrypoint.sh"]

distributed/entrypoint.py

@@ -10,6 +10,7 @@ from distributed.app import app
 from distributed.queues import add_nodes_and_edges_queue, add_data_points_queue
 from distributed.workers.graph_saving_worker import graph_saving_worker
 from distributed.workers.data_point_saving_worker import data_point_saving_worker
+from distributed.signal import QueueSignal
 
 logger = get_logger()
 
@@ -23,13 +24,14 @@ async def main():
     await add_nodes_and_edges_queue.clear.aio()
     await add_data_points_queue.clear.aio()
 
-    number_of_graph_saving_workers = 1  # Total number of graph_saving_worker to spawn
-    number_of_data_point_saving_workers = 5  # Total number of graph_saving_worker to spawn
+    number_of_graph_saving_workers = 1  # Total number of graph_saving_worker to spawn (MAX 1)
+    number_of_data_point_saving_workers = (
+        10  # Total number of graph_saving_worker to spawn (MAX 10)
+    )
 
-    results = []
     consumer_futures = []
 
-    # await prune.prune_data()  # We don't want to delete files on s3
+    await prune.prune_data()  # This prunes the data from the file storage
     # Delete DBs and saved files from metastore
     await prune.prune_system(metadata=True)
 
@@ -45,16 +47,28 @@ async def main():
         worker_future = data_point_saving_worker.spawn()
         consumer_futures.append(worker_future)
 
+    """ Example: Setting and adding S3 path as input
     s3_bucket_path = os.getenv("S3_BUCKET_PATH")
     s3_data_path = "s3://" + s3_bucket_path
 
     await cognee.add(s3_data_path, dataset_name="s3-files")
+    """
+    await cognee.add(
+        [
+            "Audi is a German car manufacturer",
+            "The Netherlands is next to Germany",
+            "Berlin is the capital of Germany",
+            "The Rhine is a major European river",
+            "BMW produces luxury vehicles",
+        ],
+        dataset_name="s3-files",
+    )
 
     await cognee.cognify(datasets=["s3-files"])
 
-    # Push empty tuple into the queue to signal the end of data.
-    await add_nodes_and_edges_queue.put.aio(())
-    await add_data_points_queue.put.aio(())
+    # Put Processing end signal into the queues to stop the consumers
+    await add_nodes_and_edges_queue.put.aio(QueueSignal.STOP)
+    await add_data_points_queue.put.aio(QueueSignal.STOP)
 
     for consumer_future in consumer_futures:
         try:
@@ -64,8 +78,6 @@ async def main():
         except Exception as e:
             logger.error(e)
 
-    print(results)
-
 
 if __name__ == "__main__":
     asyncio.run(main())

distributed/signal.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+
+class QueueSignal(str, Enum):
+    STOP = "STOP"

distributed/workers/data_point_saving_worker.py

@@ -1,16 +1,17 @@
+import os
 import modal
 import asyncio
 from sqlalchemy.exc import OperationalError, DBAPIError
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 from distributed.app import app
+from distributed.signal import QueueSignal
 from distributed.modal_image import image
 from distributed.queues import add_data_points_queue
 
 from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.vector import get_vector_engine
 
-
 logger = get_logger("data_point_saving_worker")
 
 
@@ -39,54 +40,83 @@ def is_deadlock_error(error):
     return False
 
 
+secret_name = os.environ.get("MODAL_SECRET_NAME", "distributed_cognee")
+
+
 @app.function(
     retries=3,
     image=image,
     timeout=86400,
-    max_containers=5,
-    secrets=[modal.Secret.from_name("distributed_cognee")],
+    max_containers=10,
+    secrets=[modal.Secret.from_name(secret_name)],
 )
 async def data_point_saving_worker():
     print("Started processing of data points; starting vector engine queue.")
     vector_engine = get_vector_engine()
+    # Defines how many data packets do we glue together from the modal queue before embedding call and ingestion
+    BATCH_SIZE = 25
+    stop_seen = False
 
     while True:
+        if stop_seen:
+            print("Finished processing all data points; stopping vector engine queue consumer.")
+            return True
+
         if await add_data_points_queue.len.aio() != 0:
             try:
-                add_data_points_request = await add_data_points_queue.get.aio(block=False)
-            except modal.exception.DeserializationError as error:
-                logger.error(f"Deserialization error: {str(error)}")
-                continue
-
-            if len(add_data_points_request) == 0:
-                print("Finished processing all data points; stopping vector engine queue.")
-                return True
-
-            if len(add_data_points_request) == 2:
-                (collection_name, data_points) = add_data_points_request
-
-                print(f"Adding {len(data_points)} data points to '{collection_name}' collection.")
-
-                @retry(
-                    retry=retry_if_exception_type(VectorDatabaseDeadlockError),
-                    stop=stop_after_attempt(3),
-                    wait=wait_exponential(multiplier=2, min=1, max=6),
-                )
-                async def add_data_points():
-                    try:
-                        await vector_engine.create_data_points(
-                            collection_name, data_points, distributed=False
+                print("Remaining elements in queue:")
+                print(await add_data_points_queue.len.aio())
+
+                # collect batched requests
+                batched_points = {}
+                for _ in range(min(BATCH_SIZE, await add_data_points_queue.len.aio())):
+                    add_data_points_request = await add_data_points_queue.get.aio(block=False)
+
+                    if not add_data_points_request:
+                        continue
+
+                    if add_data_points_request == QueueSignal.STOP:
+                        await add_data_points_queue.put.aio(QueueSignal.STOP)
+                        stop_seen = True
+                        break
+
+                    if len(add_data_points_request) == 2:
+                        collection_name, data_points = add_data_points_request
+                        if collection_name not in batched_points:
+                            batched_points[collection_name] = []
+                        batched_points[collection_name].extend(data_points)
+                    else:
+                        print("NoneType or invalid request detected.")
+
+                if batched_points:
+                    for collection_name, data_points in batched_points.items():
+                        print(
+                            f"Adding {len(data_points)} data points to '{collection_name}' collection."
                         )
-                    except DBAPIError as error:
-                        if is_deadlock_error(error):
-                            raise VectorDatabaseDeadlockError()
-                    except OperationalError as error:
-                        if is_deadlock_error(error):
-                            raise VectorDatabaseDeadlockError()
 
-                await add_data_points()
+                        @retry(
+                            retry=retry_if_exception_type(VectorDatabaseDeadlockError),
+                            stop=stop_after_attempt(3),
+                            wait=wait_exponential(multiplier=2, min=1, max=6),
+                        )
+                        async def add_data_points():
+                            try:
+                                await vector_engine.create_data_points(
+                                    collection_name, data_points, distributed=False
+                                )
+                            except DBAPIError as error:
+                                if is_deadlock_error(error):
+                                    raise VectorDatabaseDeadlockError()
+                            except OperationalError as error:
+                                if is_deadlock_error(error):
+                                    raise VectorDatabaseDeadlockError()
+
+                        await add_data_points()
+                        print(f"Finished adding data points to '{collection_name}'.")
 
-                print("Finished adding data points.")
+            except modal.exception.DeserializationError as error:
+                logger.error(f"Deserialization error: {str(error)}")
+                continue
 
         else:
             print("No jobs, go to sleep.")

distributed/workers/graph_saving_worker.py

@@ -1,8 +1,10 @@
+import os
 import modal
 import asyncio
 from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 from distributed.app import app
+from distributed.signal import QueueSignal
 from distributed.modal_image import image
 from distributed.queues import add_nodes_and_edges_queue
 
@@ -10,7 +12,6 @@ from cognee.shared.logging_utils import get_logger
 from cognee.infrastructure.databases.graph import get_graph_engine
 from cognee.infrastructure.databases.graph.config import get_graph_config
 
-
 logger = get_logger("graph_saving_worker")
 
 
@@ -37,68 +38,91 @@ def is_deadlock_error(error):
     return False
 
 
+secret_name = os.environ.get("MODAL_SECRET_NAME", "distributed_cognee")
+
+
 @app.function(
     retries=3,
     image=image,
     timeout=86400,
-    max_containers=5,
-    secrets=[modal.Secret.from_name("distributed_cognee")],
+    max_containers=1,
+    secrets=[modal.Secret.from_name(secret_name)],
 )
 async def graph_saving_worker():
     print("Started processing of nodes and edges; starting graph engine queue.")
     graph_engine = await get_graph_engine()
+    # Defines how many data packets do we glue together from the queue before ingesting them into the graph database
+    BATCH_SIZE = 25
+    stop_seen = False
 
     while True:
+        if stop_seen:
+            print("Finished processing all data points; stopping graph engine queue consumer.")
+            return True
+
         if await add_nodes_and_edges_queue.len.aio() != 0:
             try:
-                nodes_and_edges = await add_nodes_and_edges_queue.get.aio(block=False)
+                print("Remaining elements in queue:")
+                print(await add_nodes_and_edges_queue.len.aio())
+
+                all_nodes, all_edges = [], []
+                for _ in range(min(BATCH_SIZE, await add_nodes_and_edges_queue.len.aio())):
+                    nodes_and_edges = await add_nodes_and_edges_queue.get.aio(block=False)
+
+                    if not nodes_and_edges:
+                        continue
+
+                    if nodes_and_edges == QueueSignal.STOP:
+                        await add_nodes_and_edges_queue.put.aio(QueueSignal.STOP)
+                        stop_seen = True
+                        break
+
+                    if len(nodes_and_edges) == 2:
+                        nodes, edges = nodes_and_edges
+                        all_nodes.extend(nodes)
+                        all_edges.extend(edges)
+                    else:
+                        print("None Type detected.")
+
+                if all_nodes or all_edges:
+                    print(f"Adding {len(all_nodes)} nodes and {len(all_edges)} edges.")
+
+                    @retry(
+                        retry=retry_if_exception_type(GraphDatabaseDeadlockError),
+                        stop=stop_after_attempt(3),
+                        wait=wait_exponential(multiplier=2, min=1, max=6),
+                    )
+                    async def save_graph_nodes(new_nodes):
+                        try:
+                            await graph_engine.add_nodes(new_nodes, distributed=False)
+                        except Exception as error:
+                            if is_deadlock_error(error):
+                                raise GraphDatabaseDeadlockError()
+
+                    @retry(
+                        retry=retry_if_exception_type(GraphDatabaseDeadlockError),
+                        stop=stop_after_attempt(3),
+                        wait=wait_exponential(multiplier=2, min=1, max=6),
+                    )
+                    async def save_graph_edges(new_edges):
+                        try:
+                            await graph_engine.add_edges(new_edges, distributed=False)
+                        except Exception as error:
+                            if is_deadlock_error(error):
+                                raise GraphDatabaseDeadlockError()
+
+                    if all_nodes:
+                        await save_graph_nodes(all_nodes)
+
+                    if all_edges:
+                        await save_graph_edges(all_edges)
+
+                    print("Finished adding nodes and edges.")
+
             except modal.exception.DeserializationError as error:
                 logger.error(f"Deserialization error: {str(error)}")
                 continue
 
-            if len(nodes_and_edges) == 0:
-                print("Finished processing all nodes and edges; stopping graph engine queue.")
-                return True
-
-            if len(nodes_and_edges) == 2:
-                print(
-                    f"Adding {len(nodes_and_edges[0])} nodes and {len(nodes_and_edges[1])} edges."
-                )
-                nodes = nodes_and_edges[0]
-                edges = nodes_and_edges[1]
-
-                @retry(
-                    retry=retry_if_exception_type(GraphDatabaseDeadlockError),
-                    stop=stop_after_attempt(3),
-                    wait=wait_exponential(multiplier=2, min=1, max=6),
-                )
-                async def save_graph_nodes(new_nodes):
-                    try:
-                        await graph_engine.add_nodes(new_nodes, distributed=False)
-                    except Exception as error:
-                        if is_deadlock_error(error):
-                            raise GraphDatabaseDeadlockError()
-
-                @retry(
-                    retry=retry_if_exception_type(GraphDatabaseDeadlockError),
-                    stop=stop_after_attempt(3),
-                    wait=wait_exponential(multiplier=2, min=1, max=6),
-                )
-                async def save_graph_edges(new_edges):
-                    try:
-                        await graph_engine.add_edges(new_edges, distributed=False)
-                    except Exception as error:
-                        if is_deadlock_error(error):
-                            raise GraphDatabaseDeadlockError()
-
-                if nodes:
-                    await save_graph_nodes(nodes)
-
-                if edges:
-                    await save_graph_edges(edges)
-
-                print("Finished adding nodes and edges.")
-
         else:
             print("No jobs, go to sleep.")
             await asyncio.sleep(5)