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

refactor: Add setting of database configs through dictionary

Browse source 
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
This commit is contained in:
Igor Ilic 2024-10-22 11:53:33 +02:00
parent 4a73505e23
commit a3581689f2
2 changed files with 37 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
cognee/api/v1/config/config.py
View file

@ -130,6 +130,30 @@ class config():
vector_db_config = get_vectordb_config() vector_db_config = get_vectordb_config()
vector_db_config.vector_engine_provider = vector_engine_provider 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 @staticmethod
def set_vector_db_key(db_key: str): def set_vector_db_key(db_key: str):
vector_db_config = get_vectordb_config() vector_db_config = get_vectordb_config()

21
cognee/tests/test_pgvector.py
View file

@ -9,14 +9,19 @@ from cognee.api.v1.search import SearchType
logging.basicConfig(level=logging.DEBUG) logging.basicConfig(level=logging.DEBUG)
async def main(): async def main():
#TODO: Should this be set on pipeline side or in the test? cognee.config.set_vector_db_config({ "vector_db_url": "",
cognee.config.set_vector_engine_provider("pgvector") "vector_db_key": "",
cognee.config.set_db_provider("postgres") "vector_engine_provider": "pgvector"
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_relational_db_config({"db_path": "",
cognee.config.set_db_username("cognee") "db_name": "cognee_db",
cognee.config.set_db_password("cognee") "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()) 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) cognee.config.data_root_directory(data_directory_path)