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

cognee/api/client.py

@@ -74,7 +74,9 @@ if CORS_ALLOWED_ORIGINS:
         origin.strip() for origin in CORS_ALLOWED_ORIGINS.split(",") if origin.strip()
     ]
 else:
-    allowed_origins = []  # Block all except explicitly set origins
+    allowed_origins = [
+        "http://localhost:3000",
+    ]  # Block all except explicitly set origins
 
 app.add_middleware(
     CORSMiddleware,

cognee/api/v1/datasets/routers/get_datasets_router.py

@@ -290,7 +290,7 @@ def get_datasets_router() -> APIRouter:
         if dataset is None:
             raise DatasetNotFoundError(message=f"Dataset ({str(dataset_id)}) not found.")
 
-        graph_data = await get_formatted_graph_data(dataset)
+        graph_data = await get_formatted_graph_data(dataset.id, user.id)
 
         return graph_data
 
@@ -353,7 +353,7 @@ def get_datasets_router() -> APIRouter:
 
     @router.get("/status", response_model=dict[str, PipelineRunStatus])
     async def get_dataset_status(
-        datasets: Annotated[List[UUID], Query(alias="dataset")] = None,
+        datasets: Annotated[List[UUID], Query(alias="dataset")] = [],
         user: User = Depends(get_authenticated_user),
     ):
         """

cognee/infrastructure/databases/graph/kuzu/adapter.py

@@ -72,11 +72,36 @@ class KuzuAdapter(GraphDBInterface):
 
                 run_sync(file_storage.ensure_directory_exists())
 
-                self.db = Database(
-                    self.db_path,
-                    buffer_pool_size=256 * 1024 * 1024,  # 256MB buffer pool
-                    max_db_size=1024 * 1024 * 1024,
-                )
+                try:
+                    self.db = Database(
+                        self.db_path,
+                        buffer_pool_size=2048 * 1024 * 1024,  # 2048MB buffer pool
+                        max_db_size=4096 * 1024 * 1024,
+                    )
+                except RuntimeError:
+                    from .kuzu_migrate import read_kuzu_storage_version
+                    import kuzu
+
+                    kuzu_db_version = read_kuzu_storage_version(self.db_path)
+                    if (
+                        kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2"
+                    ) and kuzu_db_version != kuzu.__version__:
+                        # Try to migrate kuzu database to latest version
+                        from .kuzu_migrate import kuzu_migration
+
+                        kuzu_migration(
+                            new_db=self.db_path + "_new",
+                            old_db=self.db_path,
+                            new_version=kuzu.__version__,
+                            old_version=kuzu_db_version,
+                            overwrite=True,
+                        )
+
+                    self.db = Database(
+                        self.db_path,
+                        buffer_pool_size=2048 * 1024 * 1024,  # 2048MB buffer pool
+                        max_db_size=4096 * 1024 * 1024,
+                    )
 
             self.db.init_database()
             self.connection = Connection(self.db)
@@ -1438,11 +1463,8 @@ class KuzuAdapter(GraphDBInterface):
         It raises exceptions for failures occurring during deletion processes.
         """
         try:
-            # Use DETACH DELETE to remove both nodes and their relationships in one operation
-            await self.query("MATCH (n:Node) DETACH DELETE n")
-            logger.info("Cleared all data from graph while preserving structure")
-
             if self.connection:
+                self.connection.close()
                 self.connection = None
             if self.db:
                 self.db.close()

cognee/infrastructure/databases/graph/kuzu/kuzu_migrate.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+Kuzu Database Migration Script
+
+This script migrates Kuzu databases between different versions by:
+1. Setting up isolated Python environments for each Kuzu version
+2. Exporting data from the source database using the old version
+3. Importing data into the target database using the new version
+4. If overwrite is enabled target database will replace source database and source database will have the prefix _old
+5. If delete-old is enabled target database will be renamed to source database and source database will be deleted
+
+The script automatically handles:
+- Environment setup (creates virtual environments as needed)
+- Export/import validation
+- Error handling and reporting
+
+Usage Examples:
+    # Basic migration from 0.9.0 to 0.11.0
+    python kuzu_migrate.py --old-version 0.9.0 --new-version 0.11.0 --old-db /path/to/old/database --new-db /path/to/new/database
+
+Requirements:
+- Python 3.7+
+- Internet connection (to download Kuzu packages)
+- Sufficient disk space for virtual environments and temporary exports
+
+Notes:
+- Can only be used to migrate to newer Kuzu versions, from 0.11.0 onwards
+"""
+
+import tempfile
+import sys
+import struct
+import shutil
+import subprocess
+import argparse
+import os
+
+
+kuzu_version_mapping = {
+    34: "0.7.0",
+    35: "0.7.1",
+    36: "0.8.2",
+    37: "0.9.0",
+    38: "0.10.1",
+    39: "0.11.0",
+}
+
+
+def read_kuzu_storage_version(kuzu_db_path: str) -> int:
+    """
+    Reads the Kùzu storage version code from the first catalog.bin file bytes.
+
+    :param kuzu_db_path: Path to the Kuzu database file/directory.
+    :return: Storage version code as an integer.
+    """
+    if os.path.isdir(kuzu_db_path):
+        kuzu_version_file_path = os.path.join(kuzu_db_path, "catalog.kz")
+        if not os.path.isfile(kuzu_version_file_path):
+            raise FileExistsError("Kuzu catalog.kz file does not exist")
+    else:
+        kuzu_version_file_path = kuzu_db_path
+
+    with open(kuzu_version_file_path, "rb") as f:
+        # Skip the 3-byte magic "KUZ" and one byte of padding
+        f.seek(4)
+        # Read the next 8 bytes as a little-endian unsigned 64-bit integer
+        data = f.read(8)
+        if len(data) < 8:
+            raise ValueError(
+                f"File '{kuzu_version_file_path}' does not contain a storage version code."
+            )
+        version_code = struct.unpack("<Q", data)[0]
+
+    if kuzu_version_mapping.get(version_code):
+        return kuzu_version_mapping[version_code]
+    else:
+        ValueError("Could not map version_code to proper Kuzu version.")
+
+
+def ensure_env(version: str, export_dir) -> str:
+    """
+    Create (if needed) a venv at .kuzu_envs/{version} and install kuzu=={version}.
+    Returns the path to the venv's python executable.
+    """
+    # Use temp directory to create venv
+    kuzu_envs_dir = os.path.join(export_dir, ".kuzu_envs")
+
+    # venv base under the script directory
+    base = os.path.join(kuzu_envs_dir, version)
+    py_bin = os.path.join(base, "bin", "python")
+    # If environment already exists clean it
+    if os.path.isfile(py_bin):
+        shutil.rmtree(base)
+
+    print(f"→ Setting up venv for Kùzu {version}...", file=sys.stderr)
+    # Create venv
+    # NOTE: Running python in debug mode can cause issues with creating a virtual environment from that python instance
+    subprocess.run([sys.executable, "-m", "venv", base], check=True)
+    # Install the specific Kùzu version
+    subprocess.run([py_bin, "-m", "pip", "install", "--upgrade", "pip"], check=True)
+    subprocess.run([py_bin, "-m", "pip", "install", f"kuzu=={version}"], check=True)
+    return py_bin
+
+
+def run_migration_step(python_exe: str, db_path: str, cypher: str):
+    """
+    Uses the given python_exe to execute a short snippet that
+    connects to the Kùzu database and runs a Cypher command.
+    """
+    snippet = f"""
+import kuzu
+db = kuzu.Database(r"{db_path}")
+conn = kuzu.Connection(db)
+conn.execute(r\"\"\"{cypher}\"\"\")
+"""
+    proc = subprocess.run([python_exe, "-c", snippet], capture_output=True, text=True)
+    if proc.returncode != 0:
+        print(f"[ERROR] {cypher} failed:\n{proc.stderr}", file=sys.stderr)
+        sys.exit(proc.returncode)
+
+
+def kuzu_migration(new_db, old_db, new_version, old_version=None, overwrite=None, delete_old=None):
+    """
+    Main migration function that handles the complete migration process.
+    """
+    print(f"🔄 Migrating Kuzu database from {old_version} to {new_version}", file=sys.stderr)
+    print(f"📂 Source: {old_db}", file=sys.stderr)
+    print("", file=sys.stderr)
+
+    # If version of old kuzu db is not provided try to determine it based on file info
+    if not old_version:
+        old_version = read_kuzu_storage_version(old_db)
+
+    # Check if old database exists
+    if not os.path.exists(old_db):
+        print(f"Source database '{old_db}' does not exist.", file=sys.stderr)
+        sys.exit(1)
+
+    # Prepare target - ensure parent directory exists but remove target if it exists
+    parent_dir = os.path.dirname(new_db)
+    if parent_dir:
+        os.makedirs(parent_dir, exist_ok=True)
+
+    if os.path.exists(new_db):
+        raise FileExistsError(
+            "File already exists at new database location, remove file or change new database file path to continue"
+        )
+
+    # Use temp directory for all processing, it will be cleaned up after with statement
+    with tempfile.TemporaryDirectory() as export_dir:
+        # Set up environments
+        print(f"Setting up Kuzu {old_version} environment...", file=sys.stderr)
+        old_py = ensure_env(old_version, export_dir)
+        print(f"Setting up Kuzu {new_version} environment...", file=sys.stderr)
+        new_py = ensure_env(new_version, export_dir)
+
+        export_file = os.path.join(export_dir, "kuzu_export")
+        print(f"Exporting old DB → {export_dir}", file=sys.stderr)
+        run_migration_step(old_py, old_db, f"EXPORT DATABASE '{export_file}'")
+        print("Export complete.", file=sys.stderr)
+
+        # Check if export files were created and have content
+        schema_file = os.path.join(export_file, "schema.cypher")
+        if not os.path.exists(schema_file) or os.path.getsize(schema_file) == 0:
+            raise ValueError(f"Schema file not found: {schema_file}")
+
+        print(f"Importing into new DB at {new_db}", file=sys.stderr)
+        run_migration_step(new_py, new_db, f"IMPORT DATABASE '{export_file}'")
+        print("Import complete.", file=sys.stderr)
+
+    # Rename new kuzu database to old kuzu database name if enabled
+    if overwrite or delete_old:
+        # Remove kuzu lock from migrated DB
+        lock_file = new_db + ".lock"
+        if os.path.exists(lock_file):
+            os.remove(lock_file)
+        rename_databases(old_db, old_version, new_db, delete_old)
+
+    print("✅ Kuzu graph database migration finished successfully!")
+
+
+def rename_databases(old_db: str, old_version: str, new_db: str, delete_old: bool):
+    """
+    When overwrite is enabled, back up the original old_db (file with .lock and .wal or directory)
+    by renaming it to *_old, and replace it with the newly imported new_db files.
+
+    When delete_old is enabled replace the old database with the new one and delete old database
+    """
+    base_dir = os.path.dirname(old_db)
+    name = os.path.basename(old_db.rstrip(os.sep))
+    # Add _old_ and version info to backup graph database
+    backup_database_name = f"{name}_old_" + old_version.replace(".", "_")
+    backup_base = os.path.join(base_dir, backup_database_name)
+
+    if os.path.isfile(old_db):
+        # File-based database: handle main file and accompanying lock/WAL
+        for ext in ["", ".wal"]:
+            src = old_db + ext
+            dst = backup_base + ext
+            if os.path.exists(src):
+                if delete_old:
+                    os.remove(src)
+                else:
+                    os.rename(src, dst)
+                    print(f"Renamed '{src}' to '{dst}'", file=sys.stderr)
+    elif os.path.isdir(old_db):
+        # Directory-based Kuzu database
+        backup_dir = backup_base
+        if delete_old:
+            shutil.rmtree(old_db)
+        else:
+            os.rename(old_db, backup_dir)
+            print(f"Renamed directory '{old_db}' to '{backup_dir}'", file=sys.stderr)
+    else:
+        print(f"Original database path '{old_db}' not found for renaming.", file=sys.stderr)
+        sys.exit(1)
+
+    # Now move new files into place
+    for ext in ["", ".wal"]:
+        src_new = new_db + ext
+        dst_new = os.path.join(base_dir, name + ext)
+        if os.path.exists(src_new):
+            os.rename(src_new, dst_new)
+            print(f"Renamed '{src_new}' to '{dst_new}'", file=sys.stderr)
+
+
+def main():
+    p = argparse.ArgumentParser(
+        description="Migrate Kùzu DB via PyPI versions",
+        epilog="""
+Examples:
+  %(prog)s --old-version 0.9.0 --new-version 0.11.0 \\
+    --old-db /path/to/old/db --new-db /path/to/new/db --overwrite
+
+Note: This script will create temporary virtual environments in .kuzu_envs/ directory
+to isolate different Kuzu versions.
+        """,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument(
+        "--old-version",
+        required=False,
+        default=None,
+        help="Source Kuzu version (e.g., 0.9.0). If not provided automatic kuzu version detection will be attempted.",
+    )
+    p.add_argument("--new-version", required=True, help="Target Kuzu version (e.g., 0.11.0)")
+    p.add_argument("--old-db", required=True, help="Path to source database directory")
+    p.add_argument(
+        "--new-db",
+        required=True,
+        help="Path to target database directory, it can't be the same path as the old database. Use the overwrite flag if you want to replace the old database with the new one.",
+    )
+    p.add_argument(
+        "--overwrite",
+        required=False,
+        action="store_true",
+        default=False,
+        help="Rename new-db to the old-db name and location, keeps old-db as backup if delete-old is not True",
+    )
+    p.add_argument(
+        "--delete-old",
+        required=False,
+        action="store_true",
+        default=False,
+        help="When overwrite and delete-old is True old-db will not be stored as backup",
+    )
+
+    args = p.parse_args()
+
+    kuzu_migration(
+        new_db=args.new_db,
+        old_db=args.old_db,
+        new_version=args.new_version,
+        old_version=args.old_version,
+        overwrite=args.overwrite,
+        delete_old=args.delete_old,
+    )
+
+
+if __name__ == "__main__":
+    main()

cognee/infrastructure/databases/graph/neo4j_driver/adapter.py

@@ -33,7 +33,7 @@ from .neo4j_metrics_utils import (
 from .deadlock_retry import deadlock_retry
 
 
-logger = get_logger("Neo4jAdapter", level=ERROR)
+logger = get_logger("Neo4jAdapter")
 
 BASE_LABEL = "__Node__"
 
@@ -870,34 +870,52 @@ class Neo4jAdapter(GraphDBInterface):
 
             A tuple containing two lists: nodes and edges with their properties.
         """
-        query = "MATCH (n) RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties"
+        import time
 
-        result = await self.query(query)
+        start_time = time.time()
 
-        nodes = [
-            (
-                record["properties"]["id"],
-                record["properties"],
+        try:
+            # Retrieve nodes
+            query = "MATCH (n) RETURN ID(n) AS id, labels(n) AS labels, properties(n) AS properties"
+            result = await self.query(query)
+
+            nodes = []
+            for record in result:
+                nodes.append(
+                    (
+                        record["properties"]["id"],
+                        record["properties"],
+                    )
+                )
+
+            # Retrieve edges
+            query = """
+            MATCH (n)-[r]->(m)
+            RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
+            """
+            result = await self.query(query)
+
+            edges = []
+            for record in result:
+                edges.append(
+                    (
+                        record["properties"]["source_node_id"],
+                        record["properties"]["target_node_id"],
+                        record["type"],
+                        record["properties"],
+                    )
+                )
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"Retrieved {len(nodes)} nodes and {len(edges)} edges in {retrieval_time:.2f} seconds"
             )
-            for record in result
-        ]
 
-        query = """
-        MATCH (n)-[r]->(m)
-        RETURN ID(n) AS source, ID(m) AS target, TYPE(r) AS type, properties(r) AS properties
-        """
-        result = await self.query(query)
-        edges = [
-            (
-                record["properties"]["source_node_id"],
-                record["properties"]["target_node_id"],
-                record["type"],
-                record["properties"],
-            )
-            for record in result
-        ]
+            return (nodes, edges)
 
-        return (nodes, edges)
+        except Exception as e:
+            logger.error(f"Error during graph data retrieval: {str(e)}")
+            raise
 
     async def get_nodeset_subgraph(
         self, node_type: Type[Any], node_name: List[str]
@@ -918,50 +936,71 @@ class Neo4jAdapter(GraphDBInterface):
             - Tuple[List[Tuple[int, dict]], List[Tuple[int, int, str, dict]]}: A tuple
               containing nodes and edges in the requested subgraph.
         """
-        label = node_type.__name__
+        import time
 
-        query = f"""
-        UNWIND $names AS wantedName
-        MATCH (n:`{label}`)
-        WHERE n.name = wantedName
-        WITH collect(DISTINCT n) AS primary
-        UNWIND primary AS p
-        OPTIONAL MATCH (p)--(nbr)
-        WITH primary, collect(DISTINCT nbr) AS nbrs
-        WITH primary + nbrs AS nodelist
-        UNWIND nodelist AS node
-        WITH collect(DISTINCT node) AS nodes
-        MATCH (a)-[r]-(b)
-        WHERE a IN nodes AND b IN nodes
-        WITH nodes, collect(DISTINCT r) AS rels
-        RETURN
-          [n IN nodes |
-             {{ id: n.id,
-                properties: properties(n) }}] AS rawNodes,
-          [r IN rels  |
-             {{ type: type(r),
-                properties: properties(r) }}] AS rawRels
-        """
+        start_time = time.time()
 
-        result = await self.query(query, {"names": node_name})
-        if not result:
-            return [], []
+        try:
+            label = node_type.__name__
 
-        raw_nodes = result[0]["rawNodes"]
-        raw_rels = result[0]["rawRels"]
+            query = f"""
+            UNWIND $names AS wantedName
+            MATCH (n:`{label}`)
+            WHERE n.name = wantedName
+            WITH collect(DISTINCT n) AS primary
+            UNWIND primary AS p
+            OPTIONAL MATCH (p)--(nbr)
+            WITH primary, collect(DISTINCT nbr) AS nbrs
+            WITH primary + nbrs AS nodelist
+            UNWIND nodelist AS node
+            WITH collect(DISTINCT node) AS nodes
+            MATCH (a)-[r]-(b)
+            WHERE a IN nodes AND b IN nodes
+            WITH nodes, collect(DISTINCT r) AS rels
+            RETURN
+              [n IN nodes |
+                 {{ id: n.id,
+                    properties: properties(n) }}] AS rawNodes,
+              [r IN rels  |
+                 {{ type: type(r),
+                    properties: properties(r) }}] AS rawRels
+            """
 
-        nodes = [(n["properties"]["id"], n["properties"]) for n in raw_nodes]
-        edges = [
-            (
-                r["properties"]["source_node_id"],
-                r["properties"]["target_node_id"],
-                r["type"],
-                r["properties"],
+            result = await self.query(query, {"names": node_name})
+
+            if not result:
+                return [], []
+
+            raw_nodes = result[0]["rawNodes"]
+            raw_rels = result[0]["rawRels"]
+
+            # Process nodes
+            nodes = []
+            for n in raw_nodes:
+                nodes.append((n["properties"]["id"], n["properties"]))
+
+            # Process edges
+            edges = []
+            for r in raw_rels:
+                edges.append(
+                    (
+                        r["properties"]["source_node_id"],
+                        r["properties"]["target_node_id"],
+                        r["type"],
+                        r["properties"],
+                    )
+                )
+
+            retrieval_time = time.time() - start_time
+            logger.info(
+                f"Retrieved {len(nodes)} nodes and {len(edges)} edges for {node_type.__name__} in {retrieval_time:.2f} seconds"
             )
-            for r in raw_rels
-        ]
 
-        return nodes, edges
+            return nodes, edges
+
+        except Exception as e:
+            logger.error(f"Error during nodeset subgraph retrieval: {str(e)}")
+            raise
 
     async def get_filtered_graph_data(self, attribute_filters):
         """
@@ -1011,8 +1050,8 @@ class Neo4jAdapter(GraphDBInterface):
 
         edges = [
             (
-                record["source"],
-                record["target"],
+                record["properties"]["source_node_id"],
+                record["properties"]["target_node_id"],
                 record["type"],
                 record["properties"],
             )

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -49,9 +49,16 @@ class SQLAlchemyAdapter:
 
                 run_sync(self.pull_from_s3())
 
-        self.engine = create_async_engine(
-            connection_string, poolclass=NullPool if "sqlite" in connection_string else None
-        )
+        if "sqlite" in connection_string:
+            self.engine = create_async_engine(
+                connection_string,
+                poolclass=NullPool,
+            )
+        else:
+            self.engine = create_async_engine(
+                connection_string, pool_size=12, max_overflow=12, poolclass=None
+            )
+
         self.sessionmaker = async_sessionmaker(bind=self.engine, expire_on_commit=False)
 
     async def push_to_s3(self) -> None:

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -19,7 +19,7 @@ def create_vector_engine(
     for each provider, raising an EnvironmentError if any are missing, or ImportError if the
     ChromaDB package is not installed.
 
-    Supported providers include: Weaviate, Qdrant, pgvector, FalkorDB, ChromaDB, and
+    Supported providers include: Qdrant, pgvector, FalkorDB, ChromaDB, and
     LanceDB.
 
     Parameters:
@@ -30,7 +30,7 @@ def create_vector_engine(
           providers.
         - vector_db_key (str): The API key or access token for the vector database instance.
         - vector_db_provider (str): The name of the vector database provider to use (e.g.,
-          'weaviate', 'qdrant').
+          'qdrant', 'pgvector').
 
     Returns:
     --------
@@ -48,15 +48,7 @@ def create_vector_engine(
             embedding_engine=embedding_engine,
         )
 
-    if vector_db_provider == "weaviate":
-        from .weaviate_db import WeaviateAdapter
-
-        if not (vector_db_url and vector_db_key):
-            raise EnvironmentError("Missing requred Weaviate credentials!")
-
-        return WeaviateAdapter(vector_db_url, vector_db_key, embedding_engine=embedding_engine)
-
-    elif vector_db_provider == "qdrant":
+    if vector_db_provider == "qdrant":
         if not (vector_db_url and vector_db_key):
             raise EnvironmentError("Missing requred Qdrant credentials!")
 

cognee/modules/data/models/Data.py

@@ -18,12 +18,12 @@ class Data(Base):
     mime_type = Column(String)
     raw_data_location = Column(String)
     owner_id = Column(UUID, index=True)
-    tenant_id = Column(UUID, index=True, default=None)
+    tenant_id = Column(UUID, index=True, nullable=True)
     content_hash = Column(String)
     external_metadata = Column(JSON)
     node_set = Column(JSON, nullable=True)  # Store NodeSet as JSON list of strings
     token_count = Column(Integer)
-    data_size = Column(Integer)  # File size in bytes
+    data_size = Column(Integer, nullable=True)  # File size in bytes
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))
 

cognee/modules/data/processing/document_types/UnstructuredDocument.py

@@ -18,11 +18,8 @@ class UnstructuredDocument(Document):
             except ModuleNotFoundError:
                 raise UnstructuredLibraryImportError
 
-            if self.raw_data_location.startswith("s3://"):
-                async with open_data_file(self.raw_data_location, mode="rb") as f:
-                    elements = partition(file=f, content_type=self.mime_type)
-            else:
-                elements = partition(self.raw_data_location, content_type=self.mime_type)
+            async with open_data_file(self.raw_data_location, mode="rb") as f:
+                elements = partition(file=f, content_type=self.mime_type)
 
             in_memory_file = StringIO("\n\n".join([str(el) for el in elements]))
             in_memory_file.seek(0)