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

fix: changes based on PR comments

Browse source
This commit is contained in:
Andrej Milicevic 2025-11-11 13:52:34 +01:00
parent 322c134e7e
commit ed2d687135
3 changed files with 12 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
cognee/context_global_variables.py
View file

@ -16,8 +16,8 @@ vector_db_config = ContextVar("vector_db_config", default=None)
graph_db_config = ContextVar("graph_db_config", default=None) graph_db_config = ContextVar("graph_db_config", default=None)
session_user = ContextVar("session_user", default=None) session_user = ContextVar("session_user", default=None)
vector_dbs_with_multi_user_support = ["lancedb"] VECTOR_DBS_WITH_MULTI_USER_SUPPORT = ["lancedb"]
graph_dbs_with_multi_user_support = ["kuzu"] GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["kuzu"]
async def set_session_user_context_variable(user): async def set_session_user_context_variable(user):
@ -28,8 +28,8 @@ def multi_user_support_possible():
graph_db_config = get_graph_context_config() graph_db_config = get_graph_context_config()
vector_db_config = get_vectordb_context_config() vector_db_config = get_vectordb_context_config()
return ( return (
graph_db_config["graph_database_provider"] in graph_dbs_with_multi_user_support graph_db_config["graph_database_provider"] in GRAPH_DBS_WITH_MULTI_USER_SUPPORT
and vector_db_config["vector_db_provider"] in vector_dbs_with_multi_user_support and vector_db_config["vector_db_provider"] in VECTOR_DBS_WITH_MULTI_USER_SUPPORT
) )

4
cognee/infrastructure/databases/utils/constants.py
View file

@ -1,4 +0,0 @@
VECTOR_DBS_WITH_MULTI_USER_SUPPORT = ["lancedb", "falkor"]
GRAPH_DBS_WITH_MULTI_USER_SUPPORT = ["kuzu", "falkor"]
HYBRID_DBS = ["falkor"]

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

@ -11,12 +11,6 @@ from cognee.infrastructure.databases.graph.config import get_graph_config
from cognee.modules.data.methods import get_unique_dataset_id from cognee.modules.data.methods import get_unique_dataset_id
from cognee.modules.users.models import DatasetDatabase from cognee.modules.users.models import DatasetDatabase
from cognee.modules.users.models import User from cognee.modules.users.models import User
from .constants import (
GRAPH_DBS_WITH_MULTI_USER_SUPPORT,
VECTOR_DBS_WITH_MULTI_USER_SUPPORT,
HYBRID_DBS,
)
async def get_or_create_dataset_database( async def get_or_create_dataset_database(
dataset: Union[str, UUID], dataset: Union[str, UUID],
@ -42,15 +36,15 @@ async def get_or_create_dataset_database(
vector_config = get_vectordb_config() vector_config = get_vectordb_config()
graph_config = get_graph_config() graph_config = get_graph_config()
graph_db_name = f"{dataset_id}.pkl" if graph_config.graph_database_provider == "kuzu":
graph_db_name = f"{dataset_id}.pkl"
if graph_config.graph_database_provider in HYBRID_DBS:
vector_db_name = graph_db_name
else: else:
if vector_config.vector_db_provider == "lancedb": graph_db_name = dataset_id
vector_db_name = f"{dataset_id}.lance.db"
else: if vector_config.vector_db_provider == "lancedb":
vector_db_name = f"{dataset_id}.db" vector_db_name = f"{dataset_id}.lance.db"
else:
vector_db_name = dataset_id
async with db_engine.get_async_session() as session: async with db_engine.get_async_session() as session:
# Create dataset if it doesn't exist # Create dataset if it doesn't exist
@ -66,20 +60,6 @@ async def get_or_create_dataset_database(
if existing: if existing:
return existing return existing
# Check if we support multi-user for this provider. If not, use default
if graph_config.graph_database_provider not in GRAPH_DBS_WITH_MULTI_USER_SUPPORT:
raise EnvironmentError(
f"Multi-user is currently not supported for the graph database provider: {graph_config.graph_database_provider}. "
f"Supported providers are: {', '.join(GRAPH_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these"
f"providers, or disable BACKEND_ACCESS_CONTROL"
)
if vector_config.vector_db_provider not in VECTOR_DBS_WITH_MULTI_USER_SUPPORT:
raise EnvironmentError(
f"Multi-user is currently not supported for the vector database provider: {vector_config.vector_db_provider}. "
f"Supported providers are: {', '.join(VECTOR_DBS_WITH_MULTI_USER_SUPPORT)}. Either use one of these"
f"providers, or disable BACKEND_ACCESS_CONTROL"
)
# If there are no existing rows build a new row # If there are no existing rows build a new row
record = DatasetDatabase( record = DatasetDatabase(
owner_id=user.id, owner_id=user.id,