From 87d09ad6a79fbd894d020bebcf112cde54ea71e8 Mon Sep 17 00:00:00 2001 From: vasilije Date: Sat, 17 Jan 2026 21:21:50 +0000 Subject: [PATCH 1/2] fix: Add generic set() method to config class for CLI The CLI config set command was failing because config.set() didn't exist. Added a generic set() method that maps config keys to their specific setter methods. Also added comprehensive test suite covering all configuration options. Co-Authored-By: Claude Sonnet 4.5 --- cognee/api/v1/config/config.py | 34 +++ .../unit/api/v1/config/test_config_set.py | 193 ++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 cognee/tests/unit/api/v1/config/test_config_set.py diff --git a/cognee/api/v1/config/config.py b/cognee/api/v1/config/config.py index 464753438..551840616 100644 --- a/cognee/api/v1/config/config.py +++ b/cognee/api/v1/config/config.py @@ -176,3 +176,37 @@ class config: def set_vector_db_url(db_url: str): vector_db_config = get_vectordb_config() vector_db_config.vector_db_url = db_url + + @staticmethod + def set(key: str, value): + """ + Generic setter that maps configuration keys to their specific setter methods. + This enables CLI commands like 'cognee config set llm_api_key '. + """ + # Map configuration keys to their setter methods + setter_mapping = { + "llm_provider": "set_llm_provider", + "llm_model": "set_llm_model", + "llm_api_key": "set_llm_api_key", + "llm_endpoint": "set_llm_endpoint", + "graph_database_provider": "set_graph_database_provider", + "vector_db_provider": "set_vector_db_provider", + "vector_db_url": "set_vector_db_url", + "vector_db_key": "set_vector_db_key", + "chunk_size": "set_chunk_size", + "chunk_overlap": "set_chunk_overlap", + "chunk_strategy": "set_chunk_strategy", + "chunk_engine": "set_chunk_engine", + "classification_model": "set_classification_model", + "summarization_model": "set_summarization_model", + "graph_model": "set_graph_model", + "system_root_directory": "system_root_directory", + "data_root_directory": "data_root_directory", + } + + if key not in setter_mapping: + raise InvalidConfigAttributeError(attribute=key) + + method_name = setter_mapping[key] + method = getattr(config, method_name) + method(value) diff --git a/cognee/tests/unit/api/v1/config/test_config_set.py b/cognee/tests/unit/api/v1/config/test_config_set.py new file mode 100644 index 000000000..40918f9fe --- /dev/null +++ b/cognee/tests/unit/api/v1/config/test_config_set.py @@ -0,0 +1,193 @@ +""" +Tests for the config.set() method to ensure CLI config commands work correctly. +""" + +import pytest +from unittest.mock import patch, MagicMock +from cognee.api.v1.config.config import config +from cognee.api.v1.exceptions.exceptions import InvalidConfigAttributeError + + +class TestConfigSet: + """Test the config.set() method for various configuration keys.""" + + def test_set_llm_api_key(self): + """Test setting LLM API key""" + test_key = "sk-test-key-123" + + with patch("cognee.api.v1.config.config.get_llm_config") as mock_get_llm_config: + mock_llm_config = MagicMock() + mock_get_llm_config.return_value = mock_llm_config + + config.set("llm_api_key", test_key) + + assert mock_llm_config.llm_api_key == test_key + + def test_set_llm_provider(self): + """Test setting LLM provider""" + test_provider = "anthropic" + + with patch("cognee.api.v1.config.config.get_llm_config") as mock_get_llm_config: + mock_llm_config = MagicMock() + mock_get_llm_config.return_value = mock_llm_config + + config.set("llm_provider", test_provider) + + assert mock_llm_config.llm_provider == test_provider + + def test_set_llm_model(self): + """Test setting LLM model""" + test_model = "gpt-4o" + + with patch("cognee.api.v1.config.config.get_llm_config") as mock_get_llm_config: + mock_llm_config = MagicMock() + mock_get_llm_config.return_value = mock_llm_config + + config.set("llm_model", test_model) + + assert mock_llm_config.llm_model == test_model + + def test_set_llm_endpoint(self): + """Test setting LLM endpoint""" + test_endpoint = "https://api.example.com" + + with patch("cognee.api.v1.config.config.get_llm_config") as mock_get_llm_config: + mock_llm_config = MagicMock() + mock_get_llm_config.return_value = mock_llm_config + + config.set("llm_endpoint", test_endpoint) + + assert mock_llm_config.llm_endpoint == test_endpoint + + def test_set_graph_database_provider(self): + """Test setting graph database provider""" + test_provider = "neo4j" + + with patch("cognee.api.v1.config.config.get_graph_config") as mock_get_graph_config: + mock_graph_config = MagicMock() + mock_get_graph_config.return_value = mock_graph_config + + config.set("graph_database_provider", test_provider) + + assert mock_graph_config.graph_database_provider == test_provider + + def test_set_vector_db_provider(self): + """Test setting vector database provider""" + test_provider = "chromadb" + + with patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config: + mock_vector_config = MagicMock() + mock_get_vectordb_config.return_value = mock_vector_config + + config.set("vector_db_provider", test_provider) + + assert mock_vector_config.vector_db_provider == test_provider + + def test_set_vector_db_url(self): + """Test setting vector database URL""" + test_url = "http://localhost:8000" + + with patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config: + mock_vector_config = MagicMock() + mock_get_vectordb_config.return_value = mock_vector_config + + config.set("vector_db_url", test_url) + + assert mock_vector_config.vector_db_url == test_url + + def test_set_vector_db_key(self): + """Test setting vector database key""" + test_key = "test-key-123" + + with patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config: + mock_vector_config = MagicMock() + mock_get_vectordb_config.return_value = mock_vector_config + + config.set("vector_db_key", test_key) + + assert mock_vector_config.vector_db_key == test_key + + def test_set_chunk_size(self): + """Test setting chunk size""" + test_size = 2000 + + with patch("cognee.api.v1.config.config.get_chunk_config") as mock_get_chunk_config: + mock_chunk_config = MagicMock() + mock_get_chunk_config.return_value = mock_chunk_config + + config.set("chunk_size", test_size) + + assert mock_chunk_config.chunk_size == test_size + + def test_set_chunk_overlap(self): + """Test setting chunk overlap""" + test_overlap = 20 + + with patch("cognee.api.v1.config.config.get_chunk_config") as mock_get_chunk_config: + mock_chunk_config = MagicMock() + mock_get_chunk_config.return_value = mock_chunk_config + + config.set("chunk_overlap", test_overlap) + + assert mock_chunk_config.chunk_overlap == test_overlap + + def test_set_invalid_key(self): + """Test that setting an invalid key raises InvalidConfigAttributeError""" + with pytest.raises(InvalidConfigAttributeError): + config.set("invalid_key", "some_value") + + def test_set_multiple_keys(self): + """Test setting multiple configuration keys in sequence""" + with patch("cognee.api.v1.config.config.get_llm_config") as mock_get_llm_config: + mock_llm_config = MagicMock() + mock_get_llm_config.return_value = mock_llm_config + + # Set multiple keys + config.set("llm_api_key", "test-key") + config.set("llm_provider", "openai") + config.set("llm_model", "gpt-4o") + + # Verify all were set + assert mock_llm_config.llm_api_key == "test-key" + assert mock_llm_config.llm_provider == "openai" + assert mock_llm_config.llm_model == "gpt-4o" + + def test_set_system_root_directory(self): + """Test setting system root directory""" + test_dir = "/tmp/test" + + with patch("cognee.api.v1.config.config.get_base_config") as mock_get_base_config, \ + patch("cognee.api.v1.config.config.get_relational_config") as mock_get_relational_config, \ + patch("cognee.api.v1.config.config.get_graph_config") as mock_get_graph_config, \ + patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config: + + mock_base_config = MagicMock() + mock_base_config.system_root_directory = "" + mock_get_base_config.return_value = mock_base_config + + mock_relational_config = MagicMock() + mock_get_relational_config.return_value = mock_relational_config + + mock_graph_config = MagicMock() + mock_graph_config.graph_filename = "cognee.db" + mock_get_graph_config.return_value = mock_graph_config + + mock_vector_config = MagicMock() + mock_vector_config.vector_db_provider = "lancedb" + mock_get_vectordb_config.return_value = mock_vector_config + + config.set("system_root_directory", test_dir) + + assert mock_base_config.system_root_directory == test_dir + + def test_set_data_root_directory(self): + """Test setting data root directory""" + test_dir = "/tmp/data" + + with patch("cognee.api.v1.config.config.get_base_config") as mock_get_base_config: + mock_base_config = MagicMock() + mock_get_base_config.return_value = mock_base_config + + config.set("data_root_directory", test_dir) + + assert mock_base_config.data_root_directory == test_dir From 2a83ae3de131ce5eb19a031a428a5f0da0abc63f Mon Sep 17 00:00:00 2001 From: vasilije Date: Sun, 18 Jan 2026 12:47:39 +0000 Subject: [PATCH 2/2] added fixes to cli --- cognee/tests/unit/api/v1/config/test_config_set.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/cognee/tests/unit/api/v1/config/test_config_set.py b/cognee/tests/unit/api/v1/config/test_config_set.py index 40918f9fe..b2e04f810 100644 --- a/cognee/tests/unit/api/v1/config/test_config_set.py +++ b/cognee/tests/unit/api/v1/config/test_config_set.py @@ -156,11 +156,14 @@ class TestConfigSet: """Test setting system root directory""" test_dir = "/tmp/test" - with patch("cognee.api.v1.config.config.get_base_config") as mock_get_base_config, \ - patch("cognee.api.v1.config.config.get_relational_config") as mock_get_relational_config, \ - patch("cognee.api.v1.config.config.get_graph_config") as mock_get_graph_config, \ - patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config: - + with ( + patch("cognee.api.v1.config.config.get_base_config") as mock_get_base_config, + patch( + "cognee.api.v1.config.config.get_relational_config" + ) as mock_get_relational_config, + patch("cognee.api.v1.config.config.get_graph_config") as mock_get_graph_config, + patch("cognee.api.v1.config.config.get_vectordb_config") as mock_get_vectordb_config, + ): mock_base_config = MagicMock() mock_base_config.system_root_directory = "" mock_get_base_config.return_value = mock_base_config