From 97e9070c574bb1872642eacdbefa93834bd79f59 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 20 Dec 2024 10:57:29 +0000 Subject: [PATCH] feat: Add persistent site exclusion settings (#652) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add SITE_EXCLUSIONS section to config - Implement get_excluded_sites and set_excluded_sites methods - Add comprehensive tests for site exclusion functionality This implements persistent storage of site exclusions in settings, allowing users to maintain their exclusion preferences across sessions. Link to Devin run: https://app.devin.ai/sessions/121045305ac0458bbdf2566092dbc1b2 Co-Authored-By: Erkin Alp Güney --- sample.config.toml | 6 ++++-- src/config.py | 25 ++++++++++++++++-------- tests/test_config.py | 45 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 66 insertions(+), 10 deletions(-) create mode 100644 tests/test_config.py diff --git a/sample.config.toml b/sample.config.toml index 62cde107..3df9eae8 100644 --- a/sample.config.toml +++ b/sample.config.toml @@ -25,10 +25,12 @@ OLLAMA = "http://127.0.0.1:11434" LM_STUDIO = "http://localhost:1234/v1" OPENAI = "https://api.openai.com/v1" - [LOGGING] LOG_REST_API = "true" LOG_PROMPTS = "false" [TIMEOUT] -INFERENCE = 60 \ No newline at end of file +INFERENCE = 60 + +[SITE_EXCLUSIONS] +EXCLUDED_SITES = [] diff --git a/src/config.py b/src/config.py index a3303118..bd70c72e 100644 --- a/src/config.py +++ b/src/config.py @@ -22,23 +22,23 @@ def _load_config(self): # check if all the keys are present in the config file with open("sample.config.toml", "r") as f: sample_config = toml.load(f) - + with open("config.toml", "r+") as f: config = toml.load(f) - + # Update the config with any missing keys and their keys of keys for key, value in sample_config.items(): config.setdefault(key, value) if isinstance(value, dict): for sub_key, sub_value in value.items(): config[key].setdefault(sub_key, sub_value) - + f.seek(0) toml.dump(config, f) f.truncate() - + self.config = config - + def get_config(self): return self.config @@ -59,7 +59,7 @@ def get_google_search_api_endpoint(self): def get_ollama_api_endpoint(self): return self.config["API_ENDPOINTS"]["OLLAMA"] - + def get_lmstudio_api_endpoint(self): return self.config["API_ENDPOINTS"]["LM_STUDIO"] @@ -107,10 +107,19 @@ def get_logging_rest_api(self): def get_logging_prompts(self): return self.config["LOGGING"]["LOG_PROMPTS"] == "true" - + def get_timeout_inference(self): return self.config["TIMEOUT"]["INFERENCE"] + def get_excluded_sites(self): + return self.config.get("SITE_EXCLUSIONS", {}).get("EXCLUDED_SITES", []) + + def set_excluded_sites(self, sites): + if "SITE_EXCLUSIONS" not in self.config: + self.config["SITE_EXCLUSIONS"] = {} + self.config["SITE_EXCLUSIONS"]["EXCLUDED_SITES"] = sites + self.save_config() + def set_bing_api_key(self, key): self.config["API_KEYS"]["BING"] = key self.save_config() @@ -134,7 +143,7 @@ def set_google_search_api_endpoint(self, endpoint): def set_ollama_api_endpoint(self, endpoint): self.config["API_ENDPOINTS"]["OLLAMA"] = endpoint self.save_config() - + def set_lmstudio_api_endpoint(self, endpoint): self.config["API_ENDPOINTS"]["LM_STUDIO"] = endpoint self.save_config() diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 00000000..a6fb52f3 --- /dev/null +++ b/tests/test_config.py @@ -0,0 +1,45 @@ +import os +import pytest +from src.config import Config + +@pytest.fixture +def config(): + # Create a temporary config for testing + if os.path.exists("config.toml"): + os.rename("config.toml", "config.toml.bak") + + yield Config() + + # Restore original config + if os.path.exists("config.toml.bak"): + os.rename("config.toml.bak", "config.toml") + else: + os.remove("config.toml") + +def test_excluded_sites_empty_by_default(config): + """Test that excluded sites list is empty by default.""" + assert config.get_excluded_sites() == [] + +def test_set_excluded_sites(config): + """Test setting and getting excluded sites.""" + test_sites = ["example.com", "test.org"] + config.set_excluded_sites(test_sites) + assert config.get_excluded_sites() == test_sites + +def test_excluded_sites_persistence(config): + """Test that excluded sites persist after saving.""" + test_sites = ["example.com", "test.org"] + config.set_excluded_sites(test_sites) + + # Create new config instance to test persistence + new_config = Config() + assert new_config.get_excluded_sites() == test_sites + +def test_update_excluded_sites(config): + """Test updating excluded sites list.""" + initial_sites = ["example.com"] + config.set_excluded_sites(initial_sites) + + updated_sites = ["example.com", "test.org"] + config.set_excluded_sites(updated_sites) + assert config.get_excluded_sites() == updated_sites