cognee 0.2.2.dev1__py3-none-any.whl → 0.2.3.dev1__py3-none-any.whl

distributed/pyproject.toml

@@ -0,0 +1,186 @@
+[project]
+name = "cognee"
+version = "0.2.2.dev0"
+description = "Cognee - is a library for enriching LLM context with a semantic layer for better understanding and reasoning."
+authors = [
+    { name = "Vasilije Markovic" },
+    { name = "Boris Arzentar" },
+]
+requires-python = ">=3.10,<=3.13"
+readme = "README.md"
+license = "Apache-2.0"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Topic :: Software Development :: Libraries",
+    "Operating System :: MacOS :: MacOS X",
+    "Operating System :: POSIX :: Linux",
+    "Operating System :: Microsoft :: Windows",
+]
+dependencies = [
+    "openai>=1.80.1,<2",
+    "python-dotenv>=1.0.1,<2.0.0",
+    "pydantic>=2.11.7,<3.0.0",
+    "pydantic-settings>=2.10.1,<3",
+    "typing_extensions>=4.12.2,<5.0.0",
+    "nltk>=3.9.1,<4.0.0",
+    "numpy>=1.26.4, <=4.0.0",
+    "pandas>=2.2.2,<3.0.0",
+    # Note: New s3fs and boto3 versions don't work well together
+    # Always use comaptible fixed versions of these two dependencies
+    "s3fs[boto3]==2025.3.2",
+    "sqlalchemy>=2.0.39,<3.0.0",
+    "aiosqlite>=0.20.0,<1.0.0",
+    "tiktoken>=0.8.0,<1.0.0",
+    "litellm>=1.57.4, <1.71.0",
+    "instructor>=1.9.1,<2.0.0",
+    "langfuse>=2.32.0,<3",
+    "filetype>=1.2.0,<2.0.0",
+    "aiohttp>=3.11.14,<4.0.0",
+    "aiofiles>=23.2.1,<24.0.0",
+    "rdflib>=7.1.4,<7.2.0",
+    "pypdf>=4.1.0,<6.0.0",
+    "jinja2>=3.1.3,<4",
+    "matplotlib>=3.8.3,<4",
+    "networkx>=3.4.2,<4",
+    "lancedb>=0.24.0,<1.0.0",
+    "alembic>=1.13.3,<2",
+    "pre-commit>=4.0.1,<5",
+    "scikit-learn>=1.6.1,<2",
+    "limits>=4.4.1,<5",
+    "fastapi>=0.115.7,<1.0.0",
+    "python-multipart>=0.0.20,<1.0.0",
+    "fastapi-users[sqlalchemy]>=14.0.1,<15.0.0",
+    "dlt[sqlalchemy]>=1.9.0,<2",
+    "sentry-sdk[fastapi]>=2.9.0,<3",
+    "structlog>=25.2.0,<26",
+    "pympler>=1.1,<2.0.0",
+    "onnxruntime>=1.0.0,<2.0.0",
+    "pylance>=0.22.0,<1.0.0",
+    "kuzu (==0.11.0)"
+]
+
+[project.optional-dependencies]
+api = [
+    "uvicorn>=0.34.0,<1.0.0",
+    "gunicorn>=20.1.0,<24",
+    "websockets>=15.0.1,<16.0.0"
+]
+distributed = [
+    "modal>=1.0.5,<2.0.0",
+]
+
+neo4j = ["neo4j>=5.28.0,<6"]
+postgres = [
+    "psycopg2>=2.9.10,<3",
+    "pgvector>=0.3.5,<0.4",
+    "asyncpg>=0.30.0,<1.0.0",
+]
+postgres-binary = [
+    "psycopg2-binary>=2.9.10,<3.0.0",
+    "pgvector>=0.3.5,<0.4",
+    "asyncpg>=0.30.0,<1.0.0",
+]
+notebook = ["notebook>=7.1.0,<8"]
+langchain = [
+    "langsmith>=0.2.3,<1.0.0",
+    "langchain_text_splitters>=0.3.2,<1.0.0",
+]
+llama-index = ["llama-index-core>=0.12.11,<0.13"]
+gemini = ["google-generativeai>=0.8.4,<0.9"]
+huggingface = ["transformers>=4.46.3,<5"]
+ollama = ["transformers>=4.46.3,<5"]
+mistral = ["mistral-common>=1.5.2,<2"]
+anthropic = ["anthropic>=0.26.1,<0.27"]
+deepeval = ["deepeval>=2.0.1,<3"]
+posthog = ["posthog>=3.5.0,<4"]
+falkordb = ["falkordb>=1.0.9,<2.0.0"]
+groq = ["groq>=0.8.0,<1.0.0"]
+chromadb = [
+    "chromadb>=0.3.0,<0.7",
+    "pypika==0.48.8",
+]
+docs = ["unstructured[csv, doc, docx, epub, md, odt, org, ppt, pptx, rst, rtf, tsv, xlsx]>=0.18.1,<19"]
+codegraph = [
+    "fastembed<=0.6.0 ; python_version < '3.13'",
+    "transformers>=4.46.3,<5",
+    "tree-sitter>=0.24.0,<0.25",
+    "tree-sitter-python>=0.23.6,<0.24",
+]
+evals = [
+    "plotly>=6.0.0,<7",
+    "gdown>=5.2.0,<6",
+]
+gui = [
+    "pyside6>=6.8.3,<7",
+    "qasync>=0.27.1,<0.28",
+]
+graphiti = ["graphiti-core>=0.7.0,<0.8"]
+# Note: New s3fs and boto3 versions don't work well together
+# Always use comaptible fixed versions of these two dependencies
+aws = ["s3fs[boto3]==2025.3.2"]
+dev = [
+    "pytest>=7.4.0,<8",
+    "pytest-cov>=6.1.1,<7.0.0",
+    "pytest-asyncio>=0.21.1,<0.22",
+    "coverage>=7.3.2,<8",
+    "mypy>=1.7.1,<2",
+    "notebook>=7.1.0,<8",
+    "deptry>=0.20.0,<0.21",
+    "pylint>=3.0.3,<4",
+    "ruff>=0.9.2,<1.0.0",
+    "tweepy>=4.14.0,<5.0.0",
+    "gitpython>=3.1.43,<4",
+    "mkdocs-material>=9.5.42,<10",
+    "mkdocs-minify-plugin>=0.8.0,<0.9",
+    "mkdocstrings[python]>=0.26.2,<0.27",
+]
+debug = ["debugpy>=1.8.9,<2.0.0"]
+
+[project.urls]
+Homepage = "https://www.cognee.ai"
+Repository = "https://github.com/topoteretes/cognee"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+exclude = [
+  "/bin",
+  "/dist",
+  "/.data",
+  "/.github",
+  "/alembic",
+  "/deployment",
+  "/cognee-mcp",
+  "/cognee-frontend",
+  "/examples",
+  "/helm",
+  "/licenses",
+  "/logs",
+  "/notebooks",
+  "/profiling",
+  "/tests",
+  "/tools",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["cognee", "distributed"]
+
+[tool.ruff]
+line-length = 100
+exclude = [
+    "migrations/",  # Ignore migrations directory
+    "notebooks/",       # Ignore notebook files
+    "build/",           # Ignore build directory
+    "cognee/pipelines.py",
+    "cognee/modules/users/models/Group.py",
+    "cognee/modules/users/models/ACL.py",
+    "cognee/modules/pipelines/models/Task.py",
+    "cognee/modules/data/models/Dataset.py"
+]
+
+[tool.ruff.lint]
+ignore = ["F401"]

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)

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)

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)

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)

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}")

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

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)

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/infrastructure/databases/exceptions/EmbeddingException.py

@@ -1,20 +0,0 @@
-from cognee.exceptions import CogneeApiError
-from fastapi import status
-
-
-class EmbeddingException(CogneeApiError):
-    """
-    Custom exception for handling embedding-related errors.
-
-    This exception class is designed to indicate issues specifically related to embeddings
-    within the application. It extends the base exception class CogneeApiError and allows
-    for customization of the error message, name, and status code.
-    """
-
-    def __init__(
-        self,
-        message: str = "Embedding Exception.",
-        name: str = "EmbeddingException",
-        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
-    ):
-        super().__init__(message, name, status_code)