Skip to content

Commit 97e9070

Browse files
feat: Add persistent site exclusion settings (stitionai#652)
- 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 <[email protected]>
1 parent 3b98ed3 commit 97e9070

File tree

3 files changed

+66
-10
lines changed

3 files changed

+66
-10
lines changed

Diff for: sample.config.toml

+4-2
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,12 @@ OLLAMA = "http://127.0.0.1:11434"
2525
LM_STUDIO = "http://localhost:1234/v1"
2626
OPENAI = "https://api.openai.com/v1"
2727

28-
2928
[LOGGING]
3029
LOG_REST_API = "true"
3130
LOG_PROMPTS = "false"
3231

3332
[TIMEOUT]
34-
INFERENCE = 60
33+
INFERENCE = 60
34+
35+
[SITE_EXCLUSIONS]
36+
EXCLUDED_SITES = []

Diff for: src/config.py

+17-8
Original file line numberDiff line numberDiff line change
@@ -22,23 +22,23 @@ def _load_config(self):
2222
# check if all the keys are present in the config file
2323
with open("sample.config.toml", "r") as f:
2424
sample_config = toml.load(f)
25-
25+
2626
with open("config.toml", "r+") as f:
2727
config = toml.load(f)
28-
28+
2929
# Update the config with any missing keys and their keys of keys
3030
for key, value in sample_config.items():
3131
config.setdefault(key, value)
3232
if isinstance(value, dict):
3333
for sub_key, sub_value in value.items():
3434
config[key].setdefault(sub_key, sub_value)
35-
35+
3636
f.seek(0)
3737
toml.dump(config, f)
3838
f.truncate()
39-
39+
4040
self.config = config
41-
41+
4242
def get_config(self):
4343
return self.config
4444

@@ -59,7 +59,7 @@ def get_google_search_api_endpoint(self):
5959

6060
def get_ollama_api_endpoint(self):
6161
return self.config["API_ENDPOINTS"]["OLLAMA"]
62-
62+
6363
def get_lmstudio_api_endpoint(self):
6464
return self.config["API_ENDPOINTS"]["LM_STUDIO"]
6565

@@ -107,10 +107,19 @@ def get_logging_rest_api(self):
107107

108108
def get_logging_prompts(self):
109109
return self.config["LOGGING"]["LOG_PROMPTS"] == "true"
110-
110+
111111
def get_timeout_inference(self):
112112
return self.config["TIMEOUT"]["INFERENCE"]
113113

114+
def get_excluded_sites(self):
115+
return self.config.get("SITE_EXCLUSIONS", {}).get("EXCLUDED_SITES", [])
116+
117+
def set_excluded_sites(self, sites):
118+
if "SITE_EXCLUSIONS" not in self.config:
119+
self.config["SITE_EXCLUSIONS"] = {}
120+
self.config["SITE_EXCLUSIONS"]["EXCLUDED_SITES"] = sites
121+
self.save_config()
122+
114123
def set_bing_api_key(self, key):
115124
self.config["API_KEYS"]["BING"] = key
116125
self.save_config()
@@ -134,7 +143,7 @@ def set_google_search_api_endpoint(self, endpoint):
134143
def set_ollama_api_endpoint(self, endpoint):
135144
self.config["API_ENDPOINTS"]["OLLAMA"] = endpoint
136145
self.save_config()
137-
146+
138147
def set_lmstudio_api_endpoint(self, endpoint):
139148
self.config["API_ENDPOINTS"]["LM_STUDIO"] = endpoint
140149
self.save_config()

Diff for: tests/test_config.py

+45
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import os
2+
import pytest
3+
from src.config import Config
4+
5+
@pytest.fixture
6+
def config():
7+
# Create a temporary config for testing
8+
if os.path.exists("config.toml"):
9+
os.rename("config.toml", "config.toml.bak")
10+
11+
yield Config()
12+
13+
# Restore original config
14+
if os.path.exists("config.toml.bak"):
15+
os.rename("config.toml.bak", "config.toml")
16+
else:
17+
os.remove("config.toml")
18+
19+
def test_excluded_sites_empty_by_default(config):
20+
"""Test that excluded sites list is empty by default."""
21+
assert config.get_excluded_sites() == []
22+
23+
def test_set_excluded_sites(config):
24+
"""Test setting and getting excluded sites."""
25+
test_sites = ["example.com", "test.org"]
26+
config.set_excluded_sites(test_sites)
27+
assert config.get_excluded_sites() == test_sites
28+
29+
def test_excluded_sites_persistence(config):
30+
"""Test that excluded sites persist after saving."""
31+
test_sites = ["example.com", "test.org"]
32+
config.set_excluded_sites(test_sites)
33+
34+
# Create new config instance to test persistence
35+
new_config = Config()
36+
assert new_config.get_excluded_sites() == test_sites
37+
38+
def test_update_excluded_sites(config):
39+
"""Test updating excluded sites list."""
40+
initial_sites = ["example.com"]
41+
config.set_excluded_sites(initial_sites)
42+
43+
updated_sites = ["example.com", "test.org"]
44+
config.set_excluded_sites(updated_sites)
45+
assert config.get_excluded_sites() == updated_sites

0 commit comments

Comments
 (0)