134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
"""This module contains the configuration for the graph database."""
|
|
|
|
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
|
|
|
|
|
|
class GraphConfig(BaseSettings):
|
|
"""
|
|
Represents the configuration for a graph system, including parameters for graph file
|
|
storage and database connections.
|
|
|
|
Public methods:
|
|
- to_dict
|
|
- to_hashable_dict
|
|
|
|
Instance variables:
|
|
- graph_filename
|
|
- graph_database_provider
|
|
- graph_database_url
|
|
- graph_database_username
|
|
- graph_database_password
|
|
- graph_database_port
|
|
- graph_file_path
|
|
- graph_model
|
|
- graph_topology
|
|
- model_config
|
|
"""
|
|
|
|
# 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 = ""
|
|
graph_filename: str = ""
|
|
graph_model: object = KnowledgeGraph
|
|
graph_topology: object = KnowledgeGraph
|
|
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:
|
|
"""
|
|
Return the configuration as a dictionary.
|
|
|
|
This dictionary contains all the configurations related to the graph, which includes
|
|
details for file storage and database connectivity.
|
|
|
|
Returns:
|
|
--------
|
|
|
|
- dict: A dictionary representation of the configuration settings.
|
|
"""
|
|
return {
|
|
"graph_filename": self.graph_filename,
|
|
"graph_database_provider": self.graph_database_provider,
|
|
"graph_database_url": self.graph_database_url,
|
|
"graph_database_username": self.graph_database_username,
|
|
"graph_database_password": self.graph_database_password,
|
|
"graph_database_port": self.graph_database_port,
|
|
"graph_file_path": self.graph_file_path,
|
|
"graph_model": self.graph_model,
|
|
"graph_topology": self.graph_topology,
|
|
"model_config": self.model_config,
|
|
}
|
|
|
|
def to_hashable_dict(self) -> dict:
|
|
"""
|
|
Return a hashable dictionary with essential database configuration parameters.
|
|
|
|
This dictionary excludes certain non-hashable objects and focuses on unique identifiers
|
|
for database configurations.
|
|
|
|
Returns:
|
|
--------
|
|
|
|
- dict: A dictionary representation of the essential database configuration
|
|
settings.
|
|
"""
|
|
return {
|
|
"graph_database_provider": self.graph_database_provider,
|
|
"graph_database_url": self.graph_database_url,
|
|
"graph_database_username": self.graph_database_username,
|
|
"graph_database_password": self.graph_database_password,
|
|
"graph_database_port": self.graph_database_port,
|
|
"graph_file_path": self.graph_file_path,
|
|
}
|
|
|
|
|
|
@lru_cache
|
|
def get_graph_config():
|
|
"""
|
|
Retrieve the graph configuration. This function utilizes caching to return a singleton
|
|
instance of the GraphConfig class for efficiency.
|
|
|
|
It creates and returns a GraphConfig object, which contains various settings related to
|
|
graph configuration.
|
|
|
|
Returns:
|
|
--------
|
|
|
|
- GraphConfig: A GraphConfig instance containing the graph configuration settings.
|
|
"""
|
|
return GraphConfig()
|
|
|
|
|
|
def get_graph_context_config():
|
|
"""This function will get the appropriate graph db config based on async context.
|
|
This allows the use of multiple graph databases for different threads, async tasks and parallelization
|
|
"""
|
|
from cognee.context_global_variables import graph_db_config
|
|
|
|
if graph_db_config.get():
|
|
return graph_db_config.get()
|
|
return get_graph_config().to_hashable_dict()
|