From bbcd8baf3a0b0b6ddd6cac94e12977c301ab0cd5 Mon Sep 17 00:00:00 2001 From: Andrej Milicevic Date: Tue, 28 Oct 2025 17:56:32 +0100 Subject: [PATCH] feature: add multi-user for Falkor db --- .../infrastructure/databases/graph/config.py | 4 +++ .../databases/graph/get_graph_engine.py | 1 + .../utils/get_or_create_dataset_database.py | 35 +++++++++++++++---- .../databases/vector/create_vector_engine.py | 4 +-- .../modules/users/models/DatasetDatabase.py | 12 ++++--- 5 files changed, 42 insertions(+), 14 deletions(-) diff --git a/cognee/infrastructure/databases/graph/config.py b/cognee/infrastructure/databases/graph/config.py index b7907313c..23687b359 100644 --- a/cognee/infrastructure/databases/graph/config.py +++ b/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, } diff --git a/cognee/infrastructure/databases/graph/get_graph_engine.py b/cognee/infrastructure/databases/graph/get_graph_engine.py index 217f63070..70c27aab3 100644 --- a/cognee/infrastructure/databases/graph/get_graph_engine.py +++ b/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. diff --git a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py index 2b9b00569..0af94fd3a 100644 --- a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +++ b/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 - # TODO: not we support multi user for that db. If not, set to default, which is lance and/or kuzu. + # Check if we support multi-user for this provider. If not, use default + 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, - graph_database_provider=graph_config.graph_database_provider, - vector_database_url=vector_config.vector_db_url, - graph_database_url=graph_config.graph_database_url, + vector_database_provider=vector_provider, + graph_database_provider=graph_provider, + vector_database_url=vector_url, + graph_database_url=graph_url, + vector_database_key=vector_key, + graph_database_key=graph_key, ) try: diff --git a/cognee/infrastructure/databases/vector/create_vector_engine.py b/cognee/infrastructure/databases/vector/create_vector_engine.py index 7e3fb367f..35bbc110a 100644 --- a/cognee/infrastructure/databases/vector/create_vector_engine.py +++ b/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": diff --git a/cognee/modules/users/models/DatasetDatabase.py b/cognee/modules/users/models/DatasetDatabase.py index 3d3899f4c..25d610ab9 100644 --- a/cognee/modules/users/models/DatasetDatabase.py +++ b/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) - graph_database_provider = Column(String, unique=True, nullable=False) + vector_database_provider = Column(String, unique=False, nullable=False) + graph_database_provider = Column(String, unique=False, nullable=False) - vector_database_url = Column(String, unique=True, nullable=True) - graph_database_url = Column(String, unique=True, nullable=True) + vector_database_url = Column(String, unique=False, 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))