gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feature: add multi-user for Falkor db

Browse source
This commit is contained in:
Andrej Milicevic 2025-10-28 17:56:32 +01:00
parent 813ee94836
commit bbcd8baf3a
5 changed files with 42 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
cognee/infrastructure/databases/graph/config.py
View file

@ -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,
}

1
cognee/infrastructure/databases/graph/get_graph_engine.py
View file

@ -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.

35
cognee/infrastructure/databases/utils/get_or_create_dataset_database.py
View file

@ -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:

4
cognee/infrastructure/databases/vector/create_vector_engine.py
View file

@ -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":

12
cognee/modules/users/models/DatasetDatabase.py
View file

@ -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))