added boilerplate for user management

cognee/api/client.py

@@ -13,6 +13,8 @@
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 
+from cognee.infrastructure.databases.relational.user_authentication.routers import permission_router
+
 # Set up logging
 logging.basicConfig(
     level=logging.INFO,  # Set the logging level (e.g., DEBUG, INFO, WARNING, ERROR, CRITICAL)
@@ -26,8 +28,13 @@
         traces_sample_rate = 1.0,
         profiles_sample_rate = 1.0,
     )
-
-app = FastAPI(debug = os.getenv("ENV") != "prod")
+from contextlib import asynccontextmanager
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Not needed if you setup a migration system like Alembic
+    await create_db_and_tables()
+    yield
+app = FastAPI(debug = os.getenv("ENV") != "prod", lifespan=lifespan)
 
 origins = [
     "http://frontend:3000",
@@ -43,6 +50,71 @@
     allow_headers=["*"],
 )
 
+from contextlib import asynccontextmanager
+
+from fastapi import Depends, FastAPI
+
+from cognee.infrastructure.databases.relational.user_authentication.authentication_db import User, create_db_and_tables
+from cognee.infrastructure.databases.relational.user_authentication.schemas import UserCreate, UserRead, UserUpdate
+from cognee.infrastructure.databases.relational.user_authentication.users import auth_backend, current_active_user, fastapi_users
+
+app.include_router(
+    fastapi_users.get_auth_router(auth_backend), prefix="/auth/jwt", tags=["auth"]
+)
+app.include_router(
+    fastapi_users.get_register_router(UserRead, UserCreate),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_reset_password_router(),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_verify_router(UserRead),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_users_router(UserRead, UserUpdate),
+    prefix="/users",
+    tags=["users"],
+)
+
+app.include_router(permission_router, prefix="/manage", tags=["management"])
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Not needed if you setup a migration system like Alembic
+    await create_db_and_tables()
+    yield
+app.include_router(
+    fastapi_users.get_auth_router(auth_backend), prefix="/auth/jwt", tags=["auth"]
+)
+app.include_router(
+    fastapi_users.get_register_router(UserRead, UserCreate),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_reset_password_router(),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_verify_router(UserRead),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_users_router(UserRead, UserUpdate),
+    prefix="/users",
+    tags=["users"],
+)
+
+
+
 @app.get("/")
 async def root():
     """
@@ -267,8 +339,8 @@ def start_api_server(host: str = "0.0.0.0", port: int = 8000):
         relational_config.create_engine()
 
         from cognee.modules.data.deletion import prune_system, prune_data
-        asyncio.run(prune_data())
-        asyncio.run(prune_system(metadata = True))
+        # asyncio.run(prune_data())
+        # asyncio.run(prune_system(metadata = True))
 
         uvicorn.run(app, host = host, port = port)
     except Exception as e:

cognee/api/v1/add/add.py

@@ -3,7 +3,11 @@
 import asyncio
 import dlt
 import duckdb
+from fastapi_users import fastapi_users
+
 import cognee.modules.ingestion as ingestion
+from cognee.infrastructure.databases.relational.user_authentication.users import give_permission_document, \
+    get_async_session_context, current_active_user, create_default_user
 from cognee.infrastructure.files.storage import LocalStorage
 from cognee.modules.ingestion import get_matched_datasets, save_data_to_file
 from cognee.shared.utils import send_telemetry
@@ -48,7 +52,7 @@ async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_nam
 
     return []
 
-async def add_files(file_paths: List[str], dataset_name: str):
+async def add_files(file_paths: List[str], dataset_name: str,  user_id: str = "default_user"):
     base_config = get_base_config()
     data_directory_path = base_config.data_root_directory
 
@@ -82,12 +86,17 @@ async def add_files(file_paths: List[str], dataset_name: str):
     )
 
     @dlt.resource(standalone = True, merge_key = "id")
-    def data_resources(file_paths: str):
+    def data_resources(file_paths: str, user_id: str = user_id):
         for file_path in file_paths:
             with open(file_path.replace("file://", ""), mode = "rb") as file:
                 classified_data = ingestion.classify(file)
 
                 data_id = ingestion.identify(classified_data)
+                async with get_async_session_context() as session:
+                    if user_id is None:
+                        current_active_user = create_default_user()
+
+                    give_permission_document(current_active_user, data_id, "write", session= session)
 
                 file_metadata = classified_data.get_metadata()
 

cognee/api/v1/authenticate_user/__init__.py

@@ -0,0 +1 @@
+from .authenticate_user import authenticate_user

cognee/api/v1/authenticate_user/authenticate_user.py

@@ -0,0 +1,9 @@
+from cognee.infrastructure.databases.relational.user_authentication.users import authenticate_user_method
+
+
+async def authenticate_user():
+    """
+    This function is used to authenticate a user.
+    """
+    output = await authenticate_user_method()
+    return output

cognee/api/v1/cognify/cognify_v2.py

@@ -1,8 +1,18 @@
 import asyncio
+import hashlib
 import logging
+import uuid
 from typing import Union
 
+from fastapi_users import fastapi_users
+from sqlalchemy.ext.asyncio import AsyncSession
+
 from cognee.infrastructure.databases.graph import get_graph_config
+from cognee.infrastructure.databases.relational.user_authentication.authentication_db import async_session_maker
+from cognee.infrastructure.databases.relational.user_authentication.users import has_permission_document, \
+    get_user_permissions, get_async_session_context, fast_api_users_init
+# from cognee.infrastructure.databases.relational.user_authentication.authentication_db import async_session_maker
+# from cognee.infrastructure.databases.relational.user_authentication.users import get_user_permissions, fastapi_users
 from cognee.modules.cognify.config import get_cognify_config
 from cognee.infrastructure.databases.relational.config import get_relationaldb_config
 from cognee.modules.data.processing.document_types.AudioDocument import AudioDocument
@@ -25,7 +35,13 @@
 
 update_status_lock = asyncio.Lock()
 
+class PermissionDeniedException(Exception):
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(self.message)
+
 async def cognify(datasets: Union[str, list[str]] = None, root_node_id: str = None):
+
     relational_config = get_relationaldb_config()
     db_engine = relational_config.database_engine
     create_task_status_table()
@@ -35,68 +51,79 @@ async def cognify(datasets: Union[str, list[str]] = None, root_node_id: str = No
 
 
     async def run_cognify_pipeline(dataset_name: str, files: list[dict]):
-        async with update_status_lock:
-            task_status = get_task_status([dataset_name])
-
-            if dataset_name in task_status and task_status[dataset_name] == "DATASET_PROCESSING_STARTED":
-                logger.info(f"Dataset {dataset_name} is being processed.")
-                return
-
-            update_task_status(dataset_name, "DATASET_PROCESSING_STARTED")
-        try:
-            cognee_config = get_cognify_config()
-            graph_config = get_graph_config()
-            root_node_id = None
-
-            if graph_config.infer_graph_topology and graph_config.graph_topology_task:
-                from cognee.modules.topology.topology import TopologyEngine
-                topology_engine = TopologyEngine(infer=graph_config.infer_graph_topology)
-                root_node_id = await topology_engine.add_graph_topology(files = files)
-            elif graph_config.infer_graph_topology and not graph_config.infer_graph_topology:
-                from cognee.modules.topology.topology import TopologyEngine
-                topology_engine = TopologyEngine(infer=graph_config.infer_graph_topology)
-                await topology_engine.add_graph_topology(graph_config.topology_file_path)
-            elif not graph_config.graph_topology_task:
-                root_node_id = "ROOT"
-
-            tasks = [
-                Task(process_documents, parent_node_id = root_node_id, task_config = { "batch_size": 10 }), # Classify documents and save them as a nodes in graph db, extract text chunks based on the document type
-                Task(establish_graph_topology, topology_model = KnowledgeGraph), # Set the graph topology for the document chunk data
-                Task(expand_knowledge_graph, graph_model = KnowledgeGraph), # Generate knowledge graphs from the document chunks and attach it to chunk nodes
-                Task(filter_affected_chunks, collection_name = "chunks"), # Find all affected chunks, so we don't process unchanged chunks
-                Task(
-                    save_data_chunks,
-                    collection_name = "chunks",
-                ), # Save the document chunks in vector db and as nodes in graph db (connected to the document node and between each other)
-                run_tasks_parallel([
-                    Task(
-                        summarize_text_chunks,
-                        summarization_model = cognee_config.summarization_model,
-                        collection_name = "chunk_summaries",
-                    ), # Summarize the document chunks
-                    Task(
-                        classify_text_chunks,
-                        classification_model = cognee_config.classification_model,
-                    ),
-                ]),
-                Task(remove_obsolete_chunks), # Remove the obsolete document chunks.
-            ]
-
-            pipeline = run_tasks(tasks, [
-                PdfDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "pdf" else
-                AudioDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "audio" else
-                ImageDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "image" else
-                TextDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"])
-                for file in files
-            ])
-
-            async for result in pipeline:
-                print(result)
-
-            update_task_status(dataset_name, "DATASET_PROCESSING_FINISHED")
-        except Exception as error:
-            update_task_status(dataset_name, "DATASET_PROCESSING_ERROR")
-            raise error
+
+        for file in files:
+            file["id"] = str(uuid.uuid4())
+            file["name"] = file["name"].replace(" ", "_")
+
+            async with get_async_session_context() as session:
+                active_user = await fast_api_users_init()
+
+                out = await has_permission_document(active_user.current_user(active=True), file["id"], "write", session)
+
+
+                async with update_status_lock:
+                    task_status = get_task_status([dataset_name])
+
+                    if dataset_name in task_status and task_status[dataset_name] == "DATASET_PROCESSING_STARTED":
+                        logger.info(f"Dataset {dataset_name} is being processed.")
+                        return
+
+                    update_task_status(dataset_name, "DATASET_PROCESSING_STARTED")
+                try:
+                    cognee_config = get_cognify_config()
+                    graph_config = get_graph_config()
+                    root_node_id = None
+
+                    if graph_config.infer_graph_topology and graph_config.graph_topology_task:
+                        from cognee.modules.topology.topology import TopologyEngine
+                        topology_engine = TopologyEngine(infer=graph_config.infer_graph_topology)
+                        root_node_id = await topology_engine.add_graph_topology(files = files)
+                    elif graph_config.infer_graph_topology and not graph_config.infer_graph_topology:
+                        from cognee.modules.topology.topology import TopologyEngine
+                        topology_engine = TopologyEngine(infer=graph_config.infer_graph_topology)
+                        await topology_engine.add_graph_topology(graph_config.topology_file_path)
+                    elif not graph_config.graph_topology_task:
+                        root_node_id = "ROOT"
+
+                    tasks = [
+                        Task(process_documents, parent_node_id = root_node_id, task_config = { "batch_size": 10 }, user_id = hashed_user_id, user_permissions=user_permissions), # Classify documents and save them as a nodes in graph db, extract text chunks based on the document type
+                        Task(establish_graph_topology, topology_model = KnowledgeGraph), # Set the graph topology for the document chunk data
+                        Task(expand_knowledge_graph, graph_model = KnowledgeGraph), # Generate knowledge graphs from the document chunks and attach it to chunk nodes
+                        Task(filter_affected_chunks, collection_name = "chunks"), # Find all affected chunks, so we don't process unchanged chunks
+                        Task(
+                            save_data_chunks,
+                            collection_name = "chunks",
+                        ), # Save the document chunks in vector db and as nodes in graph db (connected to the document node and between each other)
+                        run_tasks_parallel([
+                            Task(
+                                summarize_text_chunks,
+                                summarization_model = cognee_config.summarization_model,
+                                collection_name = "chunk_summaries",
+                            ), # Summarize the document chunks
+                            Task(
+                                classify_text_chunks,
+                                classification_model = cognee_config.classification_model,
+                            ),
+                        ]),
+                        Task(remove_obsolete_chunks), # Remove the obsolete document chunks.
+                    ]
+
+                    pipeline = run_tasks(tasks, [
+                        PdfDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "pdf" else
+                        AudioDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "audio" else
+                        ImageDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"]) if file["extension"] == "image" else
+                        TextDocument(title=f"{file['name']}.{file['extension']}", file_path=file["file_path"])
+                        for file in files
+                    ])
+
+                    async for result in pipeline:
+                        print(result)
+
+                    update_task_status(dataset_name, "DATASET_PROCESSING_FINISHED")
+                except Exception as error:
+                    update_task_status(dataset_name, "DATASET_PROCESSING_ERROR")
+                    raise error
 
 
     existing_datasets = db_engine.get_datasets()

cognee/api/v1/create_user/__init__.py

@@ -0,0 +1 @@
+from .create_user import create_user

cognee/api/v1/create_user/create_user.py

@@ -0,0 +1,22 @@
+from cognee.infrastructure.databases.relational.user_authentication.users import create_user_method
+
+
+
+async def create_user(email: str, password: str, is_superuser: bool = False):
+    output = await create_user_method(email=email, password=password, is_superuser=is_superuser)
+    return output
+
+
+if __name__ == "__main__":
+    import asyncio
+    # Define an example user
+    example_email = "[email protected]"
+    example_password = "securepassword123"
+    example_is_superuser = False
+
+    # Create an event loop and run the create_user function
+    loop = asyncio.get_event_loop()
+    result = loop.run_until_complete(create_user(example_email, example_password, example_is_superuser))
+
+    # Print the result
+    print(result)

cognee/api/v1/reset_user_password/__init__.py

@@ -0,0 +1 @@
+from .create_user import create_user

cognee/api/v1/verify_user_token/__init__.py

@@ -0,0 +1 @@
+from .verify_user_token import verify_user_token

cognee/api/v1/verify_user_token/verify_user_token.py

@@ -0,0 +1,8 @@
+from cognee.infrastructure.databases.relational.user_authentication.users import user_check_token
+
+
+
+async def verify_user_token(token: str):
+
+    output = await user_check_token(token=token)
+    return output

cognee/infrastructure/data/models/Data.py

@@ -2,10 +2,10 @@
 from datetime import datetime, timezone
 from sqlalchemy.orm import relationship, MappedColumn, Mapped
 from sqlalchemy import Column, String, DateTime, UUID, Text, JSON
-from cognee.infrastructure.databases.relational import ModelBase
+from cognee.infrastructure.databases.relational import Base
 from .DatasetData import DatasetData
 
-class Data(ModelBase):
+class Data(Base):
     __tablename__ = "data"
 
     id = Column(UUID, primary_key = True)