From a3581689f2e935d965ce98682bdd13152c88bccb Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 22 Oct 2024 11:53:33 +0200 Subject: [PATCH] refactor: Add setting of database configs through dictionary Added the ability to set database configurations through dictionary for ease of use. Updated test_pgvector to use this way of setting configuration Refactor #COG-170 --- cognee/api/v1/config/config.py | 24 ++++++++++++++++++++++++ cognee/tests/test_pgvector.py | 21 +++++++++++++-------- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py index d5acfd08e..6b32752cf 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -130,6 +130,30 @@ class config(): vector_db_config = get_vectordb_config() vector_db_config.vector_engine_provider = vector_engine_provider + @staticmethod + def set_relational_db_config(config_dict: dict): + """ + Updates the relational db config with values from config_dict. + """ + relational_db_config = get_relational_config() + for key, value in config_dict.items(): + if hasattr(relational_db_config, key): + object.__setattr__(relational_db_config, key, value) + else: + raise AttributeError(f"'{key}' is not a valid attribute of the config.") + + @staticmethod + def set_vector_db_config(config_dict: dict): + """ + Updates the vector db config with values from config_dict. + """ + vector_db_config = get_vectordb_config() + for key, value in config_dict.items(): + if hasattr(vector_db_config, key): + object.__setattr__(vector_db_config, key, value) + else: + raise AttributeError(f"'{key}' is not a valid attribute of the config.") + @staticmethod def set_vector_db_key(db_key: str): vector_db_config = get_vectordb_config() diff --git a/cognee/tests/test_pgvector.py b/cognee/tests/test_pgvector.py index 8d839ab55..9105bd8e5 100644 --- a/cognee/tests/test_pgvector.py +++ b/cognee/tests/test_pgvector.py @@ -9,14 +9,19 @@ from cognee.api.v1.search import SearchType logging.basicConfig(level=logging.DEBUG) async def main(): - #TODO: Should this be set on pipeline side or in the test? - cognee.config.set_vector_engine_provider("pgvector") - cognee.config.set_db_provider("postgres") - cognee.config.set_db_name("cognee_db") - cognee.config.set_db_host("127.0.0.1") - cognee.config.set_db_port("5432") - cognee.config.set_db_username("cognee") - cognee.config.set_db_password("cognee") + cognee.config.set_vector_db_config({ "vector_db_url": "", + "vector_db_key": "", + "vector_engine_provider": "pgvector" + } + ) + cognee.config.set_relational_db_config({"db_path": "", + "db_name": "cognee_db", + "db_host": "127.0.0.1", + "db_port": "5432", + "db_username": "cognee", + "db_password": "cognee", + "db_provider": "postgres"} + ) data_directory_path = str(pathlib.Path(os.path.join(pathlib.Path(__file__).parent, ".data_storage/test_pgvector")).resolve()) cognee.config.data_root_directory(data_directory_path)