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,74 @@
     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"])
+@app.get("/authenticated-route")
+async def authenticated_route(user: User = Depends(current_active_user)):
+    return {"message": f"Hello {user.email}!"}
+
+@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 +342,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/cognify/cognify_v2.py

@@ -1,8 +1,14 @@
 import asyncio
+import hashlib
 import logging
+import uuid
 from typing import Union
 
+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 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 +31,21 @@
 
 update_status_lock = asyncio.Lock()
 
-async def cognify(datasets: Union[str, list[str]] = None, root_node_id: str = None):
+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, user_id:str="default_user"):
+    session: AsyncSession = async_session_maker()
+    user = await fastapi_users.get_user_manager.get(user_id)
+    user_permissions = await get_user_permissions(user, session)
+    hash_object = hashlib.sha256(user.encode())
+    hashed_user_id = hash_object.hexdigest()
+    required_permission = "write"
+    if required_permission not in user_permissions:
+        raise PermissionDeniedException("Not enough permissions")
+
     relational_config = get_relationaldb_config()
     db_engine = relational_config.database_engine
     create_task_status_table()
@@ -60,7 +80,7 @@ async def run_cognify_pipeline(dataset_name: str, files: list[dict]):
                 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(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

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)

cognee/infrastructure/data/models/Dataset.py

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

cognee/infrastructure/data/models/DatasetData.py

@@ -1,9 +1,9 @@
 from uuid import uuid4
 from datetime import datetime, timezone
 from sqlalchemy import Column, DateTime, UUID, ForeignKey
-from cognee.infrastructure.databases.relational import ModelBase
+from cognee.infrastructure.databases.relational import Base
 
-class DatasetData(ModelBase):
+class DatasetData(Base):
     __tablename__ = "dataset_data"
 
     id = Column(UUID, primary_key = True, default = uuid4())

cognee/infrastructure/databases/relational/ModelBase.py

@@ -1,3 +1,7 @@
-from sqlalchemy.orm import declarative_base
 
-ModelBase = declarative_base()
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    pass
+

cognee/infrastructure/databases/relational/__init__.py

@@ -1,4 +1,4 @@
-from .ModelBase import ModelBase
+from .ModelBase import Base
 from .DatabaseEngine import DatabaseEngine
 from .sqlite.SqliteEngine import SqliteEngine
 from .duckdb.DuckDBAdapter import DuckDBAdapter

cognee/infrastructure/databases/relational/config.py

@@ -6,19 +6,22 @@
 
 class RelationalConfig(BaseSettings):
     db_path: str =  os.path.join(get_absolute_path(".cognee_system"), "databases")
-    db_name: str =  "cognee.db"
+    db_name: str =  "cognee_db"
     db_host: str =  "localhost"
     db_port: str =  "5432"
     db_user: str = "cognee"
     db_password: str =  "cognee"
-    database_engine: object = create_relational_engine(db_path, db_name)
+    db_provider: str = "postgresql+asyncpg"
+    # database_engine: object = create_relational_engine(db_path, db_name, db_provider)
     db_file_path: str = os.path.join(db_path, db_name)
 
+
     model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
 
     def create_engine(self):
         self.db_file_path = os.path.join(self.db_path, self.db_name)
-        self.database_engine = create_relational_engine(self.db_path, self.db_name)
+        self.database_engine = create_relational_engine(self.db_path, self.db_name, self.db_provider, self.db_host, self.db_port, self.db_user, self.db_password)
+        return self.database_engine
 
     def to_dict(self) -> dict:
         return {
@@ -29,6 +32,7 @@ def to_dict(self) -> dict:
             "db_user": self.db_user,
             "db_password": self.db_password,
             "db_engine": self.database_engine,
+            "db_provider": self.db_provider,
         }
 
 @lru_cache

cognee/infrastructure/databases/relational/create_relational_engine.py

@@ -1,10 +1,40 @@
+from enum import Enum
+
+from cognee.infrastructure.databases.relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
 from cognee.infrastructure.files.storage import LocalStorage
 from cognee.infrastructure.databases.relational import DuckDBAdapter
 
-def create_relational_engine(db_path: str, db_name: str):
+
+class DBProvider(Enum):
+    DUCKDB = "duckdb"
+    POSTGRES = "postgresql+asyncpg"
+
+def create_relational_engine(db_path: str, db_name: str, db_provider:str, db_host:str, db_port:str, db_user:str, db_password:str):
     LocalStorage.ensure_directory_exists(db_path)
 
-    return DuckDBAdapter(
-        db_name = db_name,
-        db_path = db_path,
-    )
+    provider = DBProvider(db_provider)
+
+    if provider == DBProvider.DUCKDB:
+        # return DuckDBAdapter(
+        #     db_name = db_name,
+        #     db_path = db_path,
+        # )
+        return SQLAlchemyAdapter(
+            db_name = db_name,
+            db_path = db_path,
+            db_type = db_provider,
+            db_host=db_host,
+            db_port=db_port,
+            db_user=db_user,
+            db_password=db_password
+        )
+    elif provider == DBProvider.POSTGRES:
+        return SQLAlchemyAdapter(
+            db_name = db_name,
+            db_path = db_path,
+            db_type = db_provider,
+            db_host= db_host,
+            db_port= db_port,
+            db_user= db_user,
+            db_password= db_password
+        )