feat: set default database name based on provider (#1026)
<!-- .github/pull_request_template.md --> ## Description Default graph name will consider provider when being created ## DCO Affirmation I affirm that all code in every commit of this pull request conforms to the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
commit
e44840c601
2 changed files with 25 additions and 12 deletions
|
|
@ -3,6 +3,8 @@
|
||||||
import os
|
import os
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||||
|
import pydantic
|
||||||
|
from pydantic import Field
|
||||||
from cognee.shared.data_models import KnowledgeGraph
|
from cognee.shared.data_models import KnowledgeGraph
|
||||||
from cognee.root_dir import get_absolute_path
|
from cognee.root_dir import get_absolute_path
|
||||||
|
|
||||||
|
|
@ -29,18 +31,33 @@ class GraphConfig(BaseSettings):
|
||||||
- model_config
|
- model_config
|
||||||
"""
|
"""
|
||||||
|
|
||||||
graph_filename: str = "cognee_graph"
|
# Using Field we are able to dynamically load current GRAPH_DATABASE_PROVIDER value in the model validator part
|
||||||
graph_database_provider: str = "kuzu"
|
# and determine default graph db file and path based on this parameter if no values are provided
|
||||||
|
graph_database_provider: str = Field("kuzu", env="GRAPH_DATABASE_PROVIDER")
|
||||||
|
|
||||||
graph_database_url: str = ""
|
graph_database_url: str = ""
|
||||||
graph_database_username: str = ""
|
graph_database_username: str = ""
|
||||||
graph_database_password: str = ""
|
graph_database_password: str = ""
|
||||||
graph_database_port: int = 123
|
graph_database_port: int = 123
|
||||||
graph_file_path: str = os.path.join(
|
graph_file_path: str = ""
|
||||||
os.path.join(get_absolute_path(".cognee_system"), "databases"), graph_filename
|
graph_filename: str = ""
|
||||||
)
|
|
||||||
graph_model: object = KnowledgeGraph
|
graph_model: object = KnowledgeGraph
|
||||||
graph_topology: 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
|
||||||
|
# If no specific graph_filename or path are provided
|
||||||
|
@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 = os.path.join(get_absolute_path(".cognee_system"), "databases")
|
||||||
|
values.graph_file_path = os.path.join(base, values.graph_filename)
|
||||||
|
return values
|
||||||
|
|
||||||
def to_dict(self) -> dict:
|
def to_dict(self) -> dict:
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
|
|
@ -46,11 +46,7 @@ class KuzuAdapter(GraphDBInterface):
|
||||||
def _initialize_connection(self) -> None:
|
def _initialize_connection(self) -> None:
|
||||||
"""Initialize the Kuzu database connection and schema."""
|
"""Initialize the Kuzu database connection and schema."""
|
||||||
try:
|
try:
|
||||||
try:
|
os.makedirs(self.db_path, exist_ok=True)
|
||||||
os.makedirs(self.db_path, exist_ok=True)
|
|
||||||
except FileExistsError:
|
|
||||||
os.remove(self.db_path)
|
|
||||||
os.makedirs(self.db_path, exist_ok=True)
|
|
||||||
|
|
||||||
self.db = Database(self.db_path)
|
self.db = Database(self.db_path)
|
||||||
self.db.init_database()
|
self.db.init_database()
|
||||||
|
|
@ -79,7 +75,7 @@ class KuzuAdapter(GraphDBInterface):
|
||||||
logger.debug("Kuzu database initialized successfully")
|
logger.debug("Kuzu database initialized successfully")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to initialize Kuzu database: {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]:
|
async def query(self, query: str, params: Optional[dict] = None) -> List[Tuple]:
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue