diff --git a/cognee/infrastructure/databases/graph/config.py b/cognee/infrastructure/databases/graph/config.py index caa6608d3..34c144a3d 100644 --- a/cognee/infrastructure/databases/graph/config.py +++ b/cognee/infrastructure/databases/graph/config.py @@ -1,8 +1,10 @@ """This module contains the configuration for the graph database.""" -import os +from pathlib import Path from functools import lru_cache from pydantic_settings import BaseSettings, SettingsConfigDict +import pydantic +from pydantic import Field from cognee.shared.data_models import KnowledgeGraph from cognee.root_dir import get_absolute_path @@ -29,18 +31,32 @@ class GraphConfig(BaseSettings): - model_config """ - graph_filename: str = "cognee_graph.pkl" - graph_database_provider: str = "NETWORKX" + # Using Field we are able to dynamically load current GRAPH_DATABASE_PROVIDER value in the model validator part + # and determine default graph db file and path based on this parameter if no values are provided + graph_database_provider: str = Field("NETWORKX", env="GRAPH_DATABASE_PROVIDER") + graph_database_url: str = "" graph_database_username: str = "" graph_database_password: str = "" graph_database_port: int = 123 - graph_file_path: str = os.path.join( - os.path.join(get_absolute_path(".cognee_system"), "databases"), graph_filename - ) + graph_file_path: str = "" + graph_filename: str = "" graph_model: object = KnowledgeGraph graph_topology: object = KnowledgeGraph - model_config = SettingsConfigDict(env_file=".env", extra="allow") + model_config = SettingsConfigDict(env_file=".env", extra="allow", populate_by_name=True) + + # Model validator updates graph_filename and path dynamically after class creation based on current database provider + @pydantic.model_validator(mode="after") + def fill_derived(cls, values): + provider = values.graph_database_provider.lower() + # Set filename based on graph database provider if no filename is provided + if not values.graph_filename: + values.graph_filename = f"cognee_graph_{provider}" + # Set file path based on graph database provider if no file path is provided + if not values.graph_file_path: + base = Path(get_absolute_path(".cognee_system")) / "databases" + values.graph_file_path = base / values.graph_filename + return values def to_dict(self) -> dict: """ diff --git a/cognee/infrastructure/databases/graph/kuzu/adapter.py b/cognee/infrastructure/databases/graph/kuzu/adapter.py index 5b105c94a..527a3d682 100644 --- a/cognee/infrastructure/databases/graph/kuzu/adapter.py +++ b/cognee/infrastructure/databases/graph/kuzu/adapter.py @@ -46,11 +46,7 @@ class KuzuAdapter(GraphDBInterface): def _initialize_connection(self) -> None: """Initialize the Kuzu database connection and schema.""" try: - try: - os.makedirs(self.db_path, exist_ok=True) - except FileExistsError: - os.remove(self.db_path) - os.makedirs(self.db_path, exist_ok=True) + os.makedirs(self.db_path, exist_ok=True) self.db = Database(self.db_path) self.db.init_database() @@ -79,7 +75,7 @@ class KuzuAdapter(GraphDBInterface): logger.debug("Kuzu database initialized successfully") except Exception as e: logger.error(f"Failed to initialize Kuzu database: {e}") - raise + raise e async def query(self, query: str, params: Optional[dict] = None) -> List[Tuple]: """