cognee 0.2.1.dev1__tar.gz → 0.2.1.dev2__tar.gz

{cognee-0.2.1.dev1 → cognee-0.2.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognee
-Version: 0.2.1.dev1
+Version: 0.2.1.dev2
 Summary: Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning.
 Project-URL: Homepage, https://www.cognee.ai
 Project-URL: Repository, https://github.com/topoteretes/cognee

cognee-0.2.1.dev2/distributed/Dockerfile

@@ -0,0 +1,34 @@
+FROM python:3.11-slim
+
+# Set environment variables
+ENV PIP_NO_CACHE_DIR=true
+ENV PATH="${PATH}:/root/.poetry/bin"
+ENV PYTHONPATH=/app
+ENV RUN_MODE=modal
+ENV SKIP_MIGRATIONS=true
+ENV COGNEE_DISTRIBUTED=true
+
+# System dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    libpq-dev \
+    git \
+    curl \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+COPY pyproject.toml poetry.lock README.md /app/
+
+RUN pip install poetry
+
+RUN poetry config virtualenvs.create false
+
+RUN poetry install --extras neo4j --extras postgres --extras aws --extras distributed --no-root
+
+COPY cognee/ /app/cognee
+COPY distributed/ /app/distributed
+RUN chmod +x /app/distributed/entrypoint.sh
+
+ENTRYPOINT ["/app/distributed/entrypoint.sh"]

cognee-0.2.1.dev2/distributed/app.py

@@ -0,0 +1,4 @@
+from modal import App
+
+
+app = App("cognee_modal_distributed")

cognee-0.2.1.dev2/distributed/entrypoint.py

@@ -0,0 +1,71 @@
+import os
+import asyncio
+
+import cognee
+from cognee.api.v1.prune import prune
+from cognee.shared.logging_utils import get_logger
+from cognee.modules.engine.operations.setup import setup
+
+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
+
+logger = get_logger()
+
+
+os.environ["COGNEE_DISTRIBUTED"] = "True"
+
+
+@app.local_entrypoint()
+async def main():
+    # Clear queues
+    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
+
+    results = []
+    consumer_futures = []
+
+    # await prune.prune_data()  # We don't want to delete files on s3
+    # Delete DBs and saved files from metastore
+    await prune.prune_system(metadata=True)
+
+    await setup()
+
+    # Start graph_saving_worker functions
+    for _ in range(number_of_graph_saving_workers):
+        worker_future = graph_saving_worker.spawn()
+        consumer_futures.append(worker_future)
+
+    # Start data_point_saving_worker functions
+    for _ in range(number_of_data_point_saving_workers):
+        worker_future = data_point_saving_worker.spawn()
+        consumer_futures.append(worker_future)
+
+    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.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(())
+
+    for consumer_future in consumer_futures:
+        try:
+            print("Finished but waiting for saving workers to finish.")
+            consumer_final = consumer_future.get()
+            print(f"All workers are done: {consumer_final}")
+        except Exception as e:
+            logger.error(e)
+
+    print(results)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

cognee-0.2.1.dev2/distributed/entrypoint.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+
+echo "$@"
+
+exec "$@" # Runs the command passed to the entrypoint script.

cognee-0.2.1.dev2/distributed/modal_image.py

@@ -0,0 +1,11 @@
+import pathlib
+from os import path
+from modal import Image
+from logging import getLogger
+
+logger = getLogger("modal_image_creation")
+
+image = Image.from_dockerfile(
+    path=pathlib.Path(path.join(path.dirname(__file__), "Dockerfile")).resolve(),
+    force_build=False,
+).add_local_python_source("cognee", "entrypoint")

cognee-0.2.1.dev2/distributed/queues.py

@@ -0,0 +1,5 @@
+from modal import Queue
+
+
+add_nodes_and_edges_queue = Queue.from_name("add_nodes_and_edges_queue", create_if_missing=True)
+add_data_points_queue = Queue.from_name("add_data_points_queue", create_if_missing=True)

cognee-0.2.1.dev2/distributed/tasks/queued_add_data_points.py

@@ -0,0 +1,13 @@
+async def queued_add_data_points(collection_name, data_points_batch):
+    from grpclib import GRPCError
+    from ..queues import add_data_points_queue
+
+    try:
+        await add_data_points_queue.put.aio((collection_name, data_points_batch))
+    except GRPCError:
+        first_half, second_half = (
+            data_points_batch[: len(data_points_batch) // 2],
+            data_points_batch[len(data_points_batch) // 2 :],
+        )
+        await queued_add_data_points(collection_name, first_half)
+        await queued_add_data_points(collection_name, second_half)

cognee-0.2.1.dev2/distributed/tasks/queued_add_edges.py

@@ -0,0 +1,13 @@
+async def queued_add_edges(edge_batch):
+    from grpclib import GRPCError
+    from ..queues import add_nodes_and_edges_queue
+
+    try:
+        await add_nodes_and_edges_queue.put.aio(([], edge_batch))
+    except GRPCError:
+        first_half, second_half = (
+            edge_batch[: len(edge_batch) // 2],
+            edge_batch[len(edge_batch) // 2 :],
+        )
+        await queued_add_edges(first_half)
+        await queued_add_edges(second_half)

cognee-0.2.1.dev2/distributed/tasks/queued_add_nodes.py

@@ -0,0 +1,13 @@
+async def queued_add_nodes(node_batch):
+    from grpclib import GRPCError
+    from ..queues import add_nodes_and_edges_queue
+
+    try:
+        await add_nodes_and_edges_queue.put.aio((node_batch, []))
+    except GRPCError:
+        first_half, second_half = (
+            node_batch[: len(node_batch) // 2],
+            node_batch[len(node_batch) // 2 :],
+        )
+        await queued_add_nodes(first_half)
+        await queued_add_nodes(second_half)

cognee-0.2.1.dev2/distributed/test.py

@@ -0,0 +1,28 @@
+from modal import App
+
+app = App("cognee_distributed_test")
+
+
+@app.function()
+def sum_distributed(numbers: list):
+    result = sum(numbers)
+
+    return result
+
+
+@app.local_entrypoint()
+def main():
+    sum = 0
+    numbers = range(100)
+    batch_size = 10
+
+    local_sum = sum_distributed.local(numbers=numbers)
+
+    print(f"Local sum: {local_sum}")
+
+    batches = [list(numbers[i : i + batch_size]) for i in range(0, len(numbers), batch_size)]
+
+    for result in sum_distributed.map(batches):
+        sum += result
+
+    print(f"Distributed sum: {sum}")

cognee-0.2.1.dev2/distributed/utils.py

@@ -0,0 +1,19 @@
+import os
+from functools import wraps
+
+
+def override_distributed(new_func):
+    def decorator(func):
+        @wraps(func)
+        async def wrapper(self, *args, distributed=None, **kwargs):
+            default_distributed_value = os.getenv("COGNEE_DISTRIBUTED", "False").lower() == "true"
+            distributed = default_distributed_value if distributed is None else distributed
+
+            if distributed:
+                return await new_func(*args, **kwargs)
+            else:
+                return await func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

cognee-0.2.1.dev2/distributed/workers/data_point_saving_worker.py

@@ -0,0 +1,93 @@
+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.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")
+
+
+class VectorDatabaseDeadlockError(Exception):
+    message = "A deadlock occurred while trying to add data points to the vector database."
+
+
+def is_deadlock_error(error):
+    # SQLAlchemy-wrapped asyncpg
+    try:
+        import asyncpg
+
+        if isinstance(error.orig, asyncpg.exceptions.DeadlockDetectedError):
+            return True
+    except ImportError:
+        pass
+
+    # PostgreSQL: SQLSTATE 40P01 = deadlock_detected
+    if hasattr(error.orig, "pgcode") and error.orig.pgcode == "40P01":
+        return True
+
+    # SQLite: It doesn't support real deadlocks but may simulate them as "database is locked"
+    if "database is locked" in str(error.orig).lower():
+        return True
+
+    return False
+
+
+@app.function(
+    retries=3,
+    image=image,
+    timeout=86400,
+    max_containers=5,
+    secrets=[modal.Secret.from_name("distributed_cognee")],
+)
+async def data_point_saving_worker():
+    print("Started processing of data points; starting vector engine queue.")
+    vector_engine = get_vector_engine()
+
+    while 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
+                        )
+                    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("Finished adding data points.")
+
+        else:
+            print("No jobs, go to sleep.")
+            await asyncio.sleep(5)

cognee-0.2.1.dev2/distributed/workers/graph_saving_worker.py

@@ -0,0 +1,104 @@
+import modal
+import asyncio
+from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
+
+from distributed.app import app
+from distributed.modal_image import image
+from distributed.queues import add_nodes_and_edges_queue
+
+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")
+
+
+class GraphDatabaseDeadlockError(Exception):
+    message = "A deadlock occurred while trying to add data points to the vector database."
+
+
+def is_deadlock_error(error):
+    graph_config = get_graph_config()
+
+    if graph_config.graph_database_provider == "neo4j":
+        # Neo4j
+        from neo4j.exceptions import TransientError
+
+        if isinstance(error, TransientError) and (
+            error.code == "Neo.TransientError.Transaction.DeadlockDetected"
+        ):
+            return True
+
+    # Kuzu
+    if "deadlock" in str(error).lower() or "cannot acquire lock" in str(error).lower():
+        return True
+
+    return False
+
+
+@app.function(
+    retries=3,
+    image=image,
+    timeout=86400,
+    max_containers=5,
+    secrets=[modal.Secret.from_name("distributed_cognee")],
+)
+async def graph_saving_worker():
+    print("Started processing of nodes and edges; starting graph engine queue.")
+    graph_engine = await get_graph_engine()
+
+    while 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)
+            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)

{cognee-0.2.1.dev1 → cognee-0.2.1.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cognee"
-version = "0.2.1-dev1"
+version = "0.2.1-dev2"
 description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
 authors = [
     { name = "Vasilije Markovic" },
@@ -157,7 +157,6 @@ exclude = [
   "/.data",
   "/.github",
   "/alembic",
-  "/distributed",
   "/deployment",
   "/cognee-mcp",
   "/cognee-frontend",