feature: add multi-user for Falkor db

cognee/infrastructure/databases/graph/config.py

@@ -26,6 +26,7 @@ class GraphConfig(BaseSettings):
     - graph_database_username
     - graph_database_password
     - graph_database_port
+    - graph_database_key
     - graph_file_path
     - graph_model
     - graph_topology
@@ -41,6 +42,7 @@ class GraphConfig(BaseSettings):
     graph_database_username: str = ""
     graph_database_password: str = ""
     graph_database_port: int = 123
+    graph_database_key: str = ""
     graph_file_path: str = ""
     graph_filename: str = ""
     graph_model: object = KnowledgeGraph
@@ -90,6 +92,7 @@ class GraphConfig(BaseSettings):
             "graph_database_username": self.graph_database_username,
             "graph_database_password": self.graph_database_password,
             "graph_database_port": self.graph_database_port,
+            "graph_database_key": self.graph_database_key,
             "graph_file_path": self.graph_file_path,
             "graph_model": self.graph_model,
             "graph_topology": self.graph_topology,
@@ -116,6 +119,7 @@ class GraphConfig(BaseSettings):
             "graph_database_username": self.graph_database_username,
             "graph_database_password": self.graph_database_password,
             "graph_database_port": self.graph_database_port,
+            "graph_database_key": self.graph_database_key,
             "graph_file_path": self.graph_file_path,
         }
 

cognee/infrastructure/databases/graph/get_graph_engine.py

@@ -33,6 +33,7 @@ def create_graph_engine(
     graph_database_username="",
     graph_database_password="",
     graph_database_port="",
+    graph_database_key="",
 ):
     """
     Create a graph engine based on the specified provider type.

cognee/infrastructure/databases/utils/get_or_create_dataset_database.py

@@ -14,11 +14,14 @@ from cognee.modules.users.models import User
 
 
 # TODO: Find a better place to define these
-default_vector_db_name = "lance.db"
 default_vector_db_provider = "lancedb"
 default_graph_db_provider = "kuzu"
 default_vector_db_url = None
 default_graph_db_url = None
+default_vector_db_key = None
+default_graph_db_key = None
+vector_dbs_with_multi_user_support = ["lancedb", "falkor"]
+graph_dbs_with_multi_user_support = ["kuzu", "falkor"]
 
 async def get_or_create_dataset_database(
     dataset: Union[str, UUID],
@@ -61,8 +64,24 @@ async def get_or_create_dataset_database(
         if existing:
             return existing
 
-        # TODO: Set the vector and graph database stuff (name, provider, etc.) based on the whether or
+        # Check if we support multi-user for this provider. If not, use default
-        # TODO: not we support multi user for that db. If not, set to default, which is lance and/or kuzu.
+        if graph_config.graph_database_provider in graph_dbs_with_multi_user_support:
+            graph_provider = graph_config.graph_database_provider
+            graph_url = graph_config.graph_database_url
+            graph_key = graph_config.graph_database_key
+        else:
+            graph_provider = default_graph_db_provider
+            graph_url = default_graph_db_url
+            graph_key = default_graph_db_key
+
+        if vector_config.vector_db_provider in vector_dbs_with_multi_user_support:
+            vector_provider = vector_config.vector_db_provider
+            vector_url = vector_config.vector_db_url
+            vector_key = vector_config.vector_db_key
+        else:
+            vector_provider = default_vector_db_provider
+            vector_url = default_vector_db_url
+            vector_key = default_vector_db_key
 
         # If there are no existing rows build a new row
         record = DatasetDatabase(
@@ -70,10 +89,12 @@ async def get_or_create_dataset_database(
             dataset_id=dataset_id,
             vector_database_name=vector_db_name,
             graph_database_name=graph_db_name,
-            vector_database_provider=vector_config.vector_db_provider,
+            vector_database_provider=vector_provider,
-            graph_database_provider=graph_config.graph_database_provider,
+            graph_database_provider=graph_provider,
-            vector_database_url=vector_config.vector_db_url,
+            vector_database_url=vector_url,
-            graph_database_url=graph_config.graph_database_url,
+            graph_database_url=graph_url,
+            vector_database_key=vector_key,
+            graph_database_key=graph_key,
         )
 
         try:

cognee/infrastructure/databases/vector/create_vector_engine.py

@@ -1,6 +1,6 @@
 from .supported_databases import supported_databases
 from .embeddings import get_embedding_engine
-from cognee.infrastructure.databases.graph.config import get_graph_config
+from cognee.infrastructure.databases.graph.config import get_graph_context_config
 
 from functools import lru_cache
 
@@ -46,7 +46,7 @@ def create_vector_engine(
             url=vector_db_url,
             api_key=vector_db_key,
             embedding_engine=embedding_engine,
-            graph_name=get_graph_config().graph_database_name
+            graph_name=get_graph_context_config()["graph_database_name"],
         )
 
     if vector_db_provider == "pgvector":

cognee/modules/users/models/DatasetDatabase.py

@@ -12,15 +12,17 @@ class DatasetDatabase(Base):
         UUID, ForeignKey("datasets.id", ondelete="CASCADE"), primary_key=True, index=True
     )
 
-    # TODO: Why is this unique? Isn't it fact that two or more datasets can have the same vector and graph store?
     vector_database_name = Column(String, unique=True, nullable=False)
     graph_database_name = Column(String, unique=True, nullable=False)
 
-    vector_database_provider = Column(String, unique=True, nullable=False)
+    vector_database_provider = Column(String, unique=False, nullable=False)
-    graph_database_provider = Column(String, unique=True, nullable=False)
+    graph_database_provider = Column(String, unique=False, nullable=False)
 
-    vector_database_url = Column(String, unique=True, nullable=True)
+    vector_database_url = Column(String, unique=False, nullable=True)
-    graph_database_url = Column(String, unique=True, nullable=True)
+    graph_database_url = Column(String, unique=False, nullable=True)
+
+    vector_database_key = Column(String, unique=False, nullable=True)
+    graph_database_key = Column(String, unique=False, nullable=True)
 
     created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(timezone.utc))
     updated_at = Column(DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc))