aisbf 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- aisbf/__init__.py +7 -2
- aisbf/config.py +32 -3
- aisbf/handlers.py +193 -0
- aisbf-0.1.2.data/data/share/aisbf/aisbf/__init__.py +77 -0
- aisbf-0.1.2.data/data/share/aisbf/aisbf/config.py +173 -0
- aisbf-0.1.2.data/data/share/aisbf/aisbf/handlers.py +383 -0
- aisbf-0.1.2.data/data/share/aisbf/aisbf/models.py +67 -0
- aisbf-0.1.2.data/data/share/aisbf/aisbf/providers.py +301 -0
- aisbf-0.1.2.data/data/share/aisbf/autoselect.json +17 -0
- aisbf-0.1.2.data/data/share/aisbf/autoselect.md +89 -0
- aisbf-0.1.2.data/data/share/aisbf/main.py +214 -0
- aisbf-0.1.2.data/data/share/aisbf/providers.json +156 -0
- aisbf-0.1.2.data/data/share/aisbf/requirements.txt +11 -0
- aisbf-0.1.2.data/data/share/aisbf/rotations.json +94 -0
- {aisbf-0.1.0.dist-info → aisbf-0.1.2.dist-info}/METADATA +14 -3
- aisbf-0.1.2.dist-info/RECORD +22 -0
- aisbf-0.1.2.dist-info/entry_points.txt +2 -0
- aisbf-0.1.0.dist-info/RECORD +0 -10
- {aisbf-0.1.0.dist-info → aisbf-0.1.2.dist-info}/WHEEL +0 -0
- {aisbf-0.1.0.dist-info → aisbf-0.1.2.dist-info}/licenses/LICENSE.txt +0 -0
- {aisbf-0.1.0.dist-info → aisbf-0.1.2.dist-info}/top_level.txt +0 -0
aisbf/__init__.py
CHANGED
|
@@ -23,7 +23,7 @@ Why did the programmer quit his job? Because he didn't get arrays!
|
|
|
23
23
|
A modular proxy server for managing multiple AI provider integrations.
|
|
24
24
|
"""
|
|
25
25
|
|
|
26
|
-
from .config import config, Config, ProviderConfig, RotationConfig, AppConfig
|
|
26
|
+
from .config import config, Config, ProviderConfig, RotationConfig, AppConfig, AutoselectConfig, AutoselectModelInfo
|
|
27
27
|
from .models import (
|
|
28
28
|
Message,
|
|
29
29
|
ChatCompletionRequest,
|
|
@@ -41,7 +41,7 @@ from .providers import (
|
|
|
41
41
|
get_provider_handler,
|
|
42
42
|
PROVIDER_HANDLERS
|
|
43
43
|
)
|
|
44
|
-
from .handlers import RequestHandler, RotationHandler
|
|
44
|
+
from .handlers import RequestHandler, RotationHandler, AutoselectHandler
|
|
45
45
|
|
|
46
46
|
__version__ = "0.1.0"
|
|
47
47
|
__all__ = [
|
|
@@ -51,6 +51,8 @@ __all__ = [
|
|
|
51
51
|
"ProviderConfig",
|
|
52
52
|
"RotationConfig",
|
|
53
53
|
"AppConfig",
|
|
54
|
+
"AutoselectConfig",
|
|
55
|
+
"AutoselectModelInfo",
|
|
54
56
|
# Models
|
|
55
57
|
"Message",
|
|
56
58
|
"ChatCompletionRequest",
|
|
@@ -58,6 +60,8 @@ __all__ = [
|
|
|
58
60
|
"Model",
|
|
59
61
|
"Provider",
|
|
60
62
|
"ErrorTracking",
|
|
63
|
+
"AutoselectModelInfo",
|
|
64
|
+
"AutoselectConfig",
|
|
61
65
|
# Providers
|
|
62
66
|
"BaseProviderHandler",
|
|
63
67
|
"GoogleProviderHandler",
|
|
@@ -69,4 +73,5 @@ __all__ = [
|
|
|
69
73
|
# Handlers
|
|
70
74
|
"RequestHandler",
|
|
71
75
|
"RotationHandler",
|
|
76
|
+
"AutoselectHandler",
|
|
72
77
|
]
|
aisbf/config.py
CHANGED
|
@@ -38,9 +38,20 @@ class ProviderConfig(BaseModel):
|
|
|
38
38
|
class RotationConfig(BaseModel):
|
|
39
39
|
providers: List[Dict]
|
|
40
40
|
|
|
41
|
+
class AutoselectModelInfo(BaseModel):
|
|
42
|
+
model_id: str
|
|
43
|
+
description: str
|
|
44
|
+
|
|
45
|
+
class AutoselectConfig(BaseModel):
|
|
46
|
+
model_name: str
|
|
47
|
+
description: str
|
|
48
|
+
fallback: str
|
|
49
|
+
available_models: List[AutoselectModelInfo]
|
|
50
|
+
|
|
41
51
|
class AppConfig(BaseModel):
|
|
42
52
|
providers: Dict[str, ProviderConfig]
|
|
43
53
|
rotations: Dict[str, RotationConfig]
|
|
54
|
+
autoselect: Dict[str, AutoselectConfig]
|
|
44
55
|
error_tracking: Dict[str, Dict]
|
|
45
56
|
|
|
46
57
|
class Config:
|
|
@@ -48,13 +59,14 @@ class Config:
|
|
|
48
59
|
self._ensure_config_directory()
|
|
49
60
|
self._load_providers()
|
|
50
61
|
self._load_rotations()
|
|
62
|
+
self._load_autoselect()
|
|
51
63
|
self._initialize_error_tracking()
|
|
52
64
|
|
|
53
65
|
def _get_config_source_dir(self):
|
|
54
66
|
"""Get the directory containing default config files"""
|
|
55
67
|
# Try installed location first
|
|
56
68
|
installed_dirs = [
|
|
57
|
-
Path('/usr/
|
|
69
|
+
Path('/usr/share/aisbf'),
|
|
58
70
|
Path.home() / '.local' / 'share' / 'aisbf',
|
|
59
71
|
]
|
|
60
72
|
|
|
@@ -64,7 +76,7 @@ class Config:
|
|
|
64
76
|
|
|
65
77
|
# Fallback to source tree config directory
|
|
66
78
|
# This is for development mode
|
|
67
|
-
source_dir = Path(__file__).parent.parent
|
|
79
|
+
source_dir = Path(__file__).parent.parent / 'config'
|
|
68
80
|
if source_dir.exists() and (source_dir / 'providers.json').exists():
|
|
69
81
|
return source_dir
|
|
70
82
|
|
|
@@ -90,7 +102,7 @@ class Config:
|
|
|
90
102
|
return
|
|
91
103
|
|
|
92
104
|
# Copy default config files if they don't exist
|
|
93
|
-
for config_file in ['providers.json', 'rotations.json']:
|
|
105
|
+
for config_file in ['providers.json', 'rotations.json', 'autoselect.json']:
|
|
94
106
|
src = source_dir / config_file
|
|
95
107
|
dst = config_dir / config_file
|
|
96
108
|
|
|
@@ -126,6 +138,20 @@ class Config:
|
|
|
126
138
|
data = json.load(f)
|
|
127
139
|
self.rotations = {k: RotationConfig(**v) for k, v in data['rotations'].items()}
|
|
128
140
|
|
|
141
|
+
def _load_autoselect(self):
|
|
142
|
+
autoselect_path = Path.home() / '.aisbf' / 'autoselect.json'
|
|
143
|
+
if not autoselect_path.exists():
|
|
144
|
+
# Fallback to source config if user config doesn't exist
|
|
145
|
+
try:
|
|
146
|
+
source_dir = self._get_config_source_dir()
|
|
147
|
+
autoselect_path = source_dir / 'autoselect.json'
|
|
148
|
+
except FileNotFoundError:
|
|
149
|
+
raise FileNotFoundError("Could not find autoselect.json configuration file")
|
|
150
|
+
|
|
151
|
+
with open(autoselect_path) as f:
|
|
152
|
+
data = json.load(f)
|
|
153
|
+
self.autoselect = {k: AutoselectConfig(**v) for k, v in data.items()}
|
|
154
|
+
|
|
129
155
|
def _initialize_error_tracking(self):
|
|
130
156
|
self.error_tracking = {}
|
|
131
157
|
for provider_id in self.providers:
|
|
@@ -141,4 +167,7 @@ class Config:
|
|
|
141
167
|
def get_rotation(self, rotation_id: str) -> RotationConfig:
|
|
142
168
|
return self.rotations.get(rotation_id)
|
|
143
169
|
|
|
170
|
+
def get_autoselect(self, autoselect_id: str) -> AutoselectConfig:
|
|
171
|
+
return self.autoselect.get(autoselect_id)
|
|
172
|
+
|
|
144
173
|
config = Config()
|
aisbf/handlers.py
CHANGED
|
@@ -23,7 +23,9 @@ Why did the programmer quit his job? Because he didn't get arrays!
|
|
|
23
23
|
Request handlers for AISBF.
|
|
24
24
|
"""
|
|
25
25
|
import asyncio
|
|
26
|
+
import re
|
|
26
27
|
from typing import Dict, List, Optional
|
|
28
|
+
from pathlib import Path
|
|
27
29
|
from fastapi import HTTPException, Request
|
|
28
30
|
from fastapi.responses import JSONResponse, StreamingResponse
|
|
29
31
|
from .models import ChatCompletionRequest, ChatCompletionResponse
|
|
@@ -188,3 +190,194 @@ class RotationHandler:
|
|
|
188
190
|
})
|
|
189
191
|
|
|
190
192
|
return all_models
|
|
193
|
+
|
|
194
|
+
class AutoselectHandler:
|
|
195
|
+
def __init__(self):
|
|
196
|
+
self.config = config
|
|
197
|
+
self._skill_file_content = None
|
|
198
|
+
|
|
199
|
+
def _get_skill_file_content(self) -> str:
|
|
200
|
+
"""Load the autoselect.md skill file content"""
|
|
201
|
+
if self._skill_file_content is None:
|
|
202
|
+
# Try installed locations first
|
|
203
|
+
installed_dirs = [
|
|
204
|
+
Path('/usr/share/aisbf'),
|
|
205
|
+
Path.home() / '.local' / 'share' / 'aisbf',
|
|
206
|
+
]
|
|
207
|
+
|
|
208
|
+
for installed_dir in installed_dirs:
|
|
209
|
+
skill_file = installed_dir / 'autoselect.md'
|
|
210
|
+
if skill_file.exists():
|
|
211
|
+
with open(skill_file) as f:
|
|
212
|
+
self._skill_file_content = f.read()
|
|
213
|
+
return self._skill_file_content
|
|
214
|
+
|
|
215
|
+
# Fallback to source tree config directory
|
|
216
|
+
source_dir = Path(__file__).parent.parent / 'config'
|
|
217
|
+
skill_file = source_dir / 'autoselect.md'
|
|
218
|
+
if skill_file.exists():
|
|
219
|
+
with open(skill_file) as f:
|
|
220
|
+
self._skill_file_content = f.read()
|
|
221
|
+
return self._skill_file_content
|
|
222
|
+
|
|
223
|
+
raise FileNotFoundError("Could not find autoselect.md skill file")
|
|
224
|
+
|
|
225
|
+
return self._skill_file_content
|
|
226
|
+
|
|
227
|
+
def _build_autoselect_prompt(self, user_prompt: str, autoselect_config) -> str:
|
|
228
|
+
"""Build the prompt for model selection"""
|
|
229
|
+
skill_content = self._get_skill_file_content()
|
|
230
|
+
|
|
231
|
+
# Build the available models list
|
|
232
|
+
models_list = ""
|
|
233
|
+
for model_info in autoselect_config.available_models:
|
|
234
|
+
models_list += f"<model><model_id>{model_info.model_id}</model_id><model_description>{model_info.description}</model_description></model>\n"
|
|
235
|
+
|
|
236
|
+
# Build the complete prompt
|
|
237
|
+
prompt = f"""{skill_content}
|
|
238
|
+
|
|
239
|
+
<aisbf_user_prompt>{user_prompt}</aisbf_user_prompt>
|
|
240
|
+
<aisbf_autoselect_list>
|
|
241
|
+
{models_list}
|
|
242
|
+
</aisbf_autoselect_list>
|
|
243
|
+
<aisbf_autoselect_fallback>{autoselect_config.fallback}</aisbf_autoselect_fallback>
|
|
244
|
+
"""
|
|
245
|
+
return prompt
|
|
246
|
+
|
|
247
|
+
def _extract_model_selection(self, response: str) -> Optional[str]:
|
|
248
|
+
"""Extract the model_id from the autoselection response"""
|
|
249
|
+
match = re.search(r'<aisbf_model_autoselection>(.*?)</aisbf_model_autoselection>', response, re.DOTALL)
|
|
250
|
+
if match:
|
|
251
|
+
return match.group(1).strip()
|
|
252
|
+
return None
|
|
253
|
+
|
|
254
|
+
async def _get_model_selection(self, prompt: str) -> str:
|
|
255
|
+
"""Send the autoselect prompt to a model and get the selection"""
|
|
256
|
+
# Use the first available provider/model for the selection
|
|
257
|
+
# This is a simple implementation - could be enhanced to use a specific selection model
|
|
258
|
+
rotation_handler = RotationHandler()
|
|
259
|
+
|
|
260
|
+
# Create a minimal request for model selection
|
|
261
|
+
selection_request = {
|
|
262
|
+
"messages": [{"role": "user", "content": prompt}],
|
|
263
|
+
"temperature": 0.1, # Low temperature for more deterministic selection
|
|
264
|
+
"max_tokens": 100, # We only need a short response
|
|
265
|
+
"stream": False
|
|
266
|
+
}
|
|
267
|
+
|
|
268
|
+
# Use the fallback rotation for the selection
|
|
269
|
+
try:
|
|
270
|
+
response = await rotation_handler.handle_rotation_request("general", selection_request)
|
|
271
|
+
content = response.get('choices', [{}])[0].get('message', {}).get('content', '')
|
|
272
|
+
model_id = self._extract_model_selection(content)
|
|
273
|
+
return model_id
|
|
274
|
+
except Exception as e:
|
|
275
|
+
# If selection fails, we'll handle it in the main handler
|
|
276
|
+
return None
|
|
277
|
+
|
|
278
|
+
async def handle_autoselect_request(self, autoselect_id: str, request_data: Dict) -> Dict:
|
|
279
|
+
"""Handle an autoselect request"""
|
|
280
|
+
autoselect_config = self.config.get_autoselect(autoselect_id)
|
|
281
|
+
if not autoselect_config:
|
|
282
|
+
raise HTTPException(status_code=400, detail=f"Autoselect {autoselect_id} not found")
|
|
283
|
+
|
|
284
|
+
# Extract the user prompt from the request
|
|
285
|
+
user_messages = request_data.get('messages', [])
|
|
286
|
+
if not user_messages:
|
|
287
|
+
raise HTTPException(status_code=400, detail="No messages provided")
|
|
288
|
+
|
|
289
|
+
# Build a string representation of the user prompt
|
|
290
|
+
user_prompt = ""
|
|
291
|
+
for msg in user_messages:
|
|
292
|
+
role = msg.get('role', 'user')
|
|
293
|
+
content = msg.get('content', '')
|
|
294
|
+
if isinstance(content, list):
|
|
295
|
+
# Handle complex content (e.g., with images)
|
|
296
|
+
content = str(content)
|
|
297
|
+
user_prompt += f"{role}: {content}\n"
|
|
298
|
+
|
|
299
|
+
# Build the autoselect prompt
|
|
300
|
+
autoselect_prompt = self._build_autoselect_prompt(user_prompt, autoselect_config)
|
|
301
|
+
|
|
302
|
+
# Get the model selection
|
|
303
|
+
selected_model_id = await self._get_model_selection(autoselect_prompt)
|
|
304
|
+
|
|
305
|
+
# Validate the selected model
|
|
306
|
+
if not selected_model_id:
|
|
307
|
+
# Fallback to the configured fallback model
|
|
308
|
+
selected_model_id = autoselect_config.fallback
|
|
309
|
+
else:
|
|
310
|
+
# Check if the selected model is in the available models list
|
|
311
|
+
available_ids = [m.model_id for m in autoselect_config.available_models]
|
|
312
|
+
if selected_model_id not in available_ids:
|
|
313
|
+
selected_model_id = autoselect_config.fallback
|
|
314
|
+
|
|
315
|
+
# Now proxy the actual request to the selected rotation
|
|
316
|
+
rotation_handler = RotationHandler()
|
|
317
|
+
return await rotation_handler.handle_rotation_request(selected_model_id, request_data)
|
|
318
|
+
|
|
319
|
+
async def handle_autoselect_streaming_request(self, autoselect_id: str, request_data: Dict):
|
|
320
|
+
"""Handle an autoselect streaming request"""
|
|
321
|
+
autoselect_config = self.config.get_autoselect(autoselect_id)
|
|
322
|
+
if not autoselect_config:
|
|
323
|
+
raise HTTPException(status_code=400, detail=f"Autoselect {autoselect_id} not found")
|
|
324
|
+
|
|
325
|
+
# Extract the user prompt from the request
|
|
326
|
+
user_messages = request_data.get('messages', [])
|
|
327
|
+
if not user_messages:
|
|
328
|
+
raise HTTPException(status_code=400, detail="No messages provided")
|
|
329
|
+
|
|
330
|
+
# Build a string representation of the user prompt
|
|
331
|
+
user_prompt = ""
|
|
332
|
+
for msg in user_messages:
|
|
333
|
+
role = msg.get('role', 'user')
|
|
334
|
+
content = msg.get('content', '')
|
|
335
|
+
if isinstance(content, list):
|
|
336
|
+
content = str(content)
|
|
337
|
+
user_prompt += f"{role}: {content}\n"
|
|
338
|
+
|
|
339
|
+
# Build the autoselect prompt
|
|
340
|
+
autoselect_prompt = self._build_autoselect_prompt(user_prompt, autoselect_config)
|
|
341
|
+
|
|
342
|
+
# Get the model selection
|
|
343
|
+
selected_model_id = await self._get_model_selection(autoselect_prompt)
|
|
344
|
+
|
|
345
|
+
# Validate the selected model
|
|
346
|
+
if not selected_model_id:
|
|
347
|
+
selected_model_id = autoselect_config.fallback
|
|
348
|
+
else:
|
|
349
|
+
available_ids = [m.model_id for m in autoselect_config.available_models]
|
|
350
|
+
if selected_model_id not in available_ids:
|
|
351
|
+
selected_model_id = autoselect_config.fallback
|
|
352
|
+
|
|
353
|
+
# Now proxy the actual streaming request to the selected rotation
|
|
354
|
+
rotation_handler = RotationHandler()
|
|
355
|
+
|
|
356
|
+
async def stream_generator():
|
|
357
|
+
try:
|
|
358
|
+
response = await rotation_handler.handle_rotation_request(
|
|
359
|
+
selected_model_id,
|
|
360
|
+
{**request_data, "stream": True}
|
|
361
|
+
)
|
|
362
|
+
for chunk in response:
|
|
363
|
+
yield f"data: {chunk}\n\n".encode('utf-8')
|
|
364
|
+
except Exception as e:
|
|
365
|
+
yield f"data: {str(e)}\n\n".encode('utf-8')
|
|
366
|
+
|
|
367
|
+
return StreamingResponse(stream_generator(), media_type="text/event-stream")
|
|
368
|
+
|
|
369
|
+
async def handle_autoselect_model_list(self, autoselect_id: str) -> List[Dict]:
|
|
370
|
+
"""List available models for an autoselect endpoint"""
|
|
371
|
+
autoselect_config = self.config.get_autoselect(autoselect_id)
|
|
372
|
+
if not autoselect_config:
|
|
373
|
+
raise HTTPException(status_code=400, detail=f"Autoselect {autoselect_id} not found")
|
|
374
|
+
|
|
375
|
+
# Return the available models that can be selected
|
|
376
|
+
return [
|
|
377
|
+
{
|
|
378
|
+
"id": model_info.model_id,
|
|
379
|
+
"name": model_info.model_id,
|
|
380
|
+
"description": model_info.description
|
|
381
|
+
}
|
|
382
|
+
for model_info in autoselect_config.available_models
|
|
383
|
+
]
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyleft (C) 2026 Stefy Lanza <stefy@nexlab.net>
|
|
3
|
+
|
|
4
|
+
AISBF - AI Service Broker Framework || AI Should Be Free
|
|
5
|
+
|
|
6
|
+
A modular proxy server for managing multiple AI provider integrations.
|
|
7
|
+
|
|
8
|
+
This program is free software: you can redistribute it and/or modify
|
|
9
|
+
it under the terms of the GNU General Public License as published by
|
|
10
|
+
the Free Software Foundation, either version 3 of the License, or
|
|
11
|
+
(at your option) any later version.
|
|
12
|
+
|
|
13
|
+
This program is distributed in the hope that it will be useful,
|
|
14
|
+
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
15
|
+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
16
|
+
GNU General Public License for more details.
|
|
17
|
+
|
|
18
|
+
You should have received a copy of the GNU General Public License
|
|
19
|
+
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
20
|
+
|
|
21
|
+
Why did the programmer quit his job? Because he didn't get arrays!
|
|
22
|
+
|
|
23
|
+
A modular proxy server for managing multiple AI provider integrations.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
from .config import config, Config, ProviderConfig, RotationConfig, AppConfig, AutoselectConfig, AutoselectModelInfo
|
|
27
|
+
from .models import (
|
|
28
|
+
Message,
|
|
29
|
+
ChatCompletionRequest,
|
|
30
|
+
ChatCompletionResponse,
|
|
31
|
+
Model,
|
|
32
|
+
Provider,
|
|
33
|
+
ErrorTracking
|
|
34
|
+
)
|
|
35
|
+
from .providers import (
|
|
36
|
+
BaseProviderHandler,
|
|
37
|
+
GoogleProviderHandler,
|
|
38
|
+
OpenAIProviderHandler,
|
|
39
|
+
AnthropicProviderHandler,
|
|
40
|
+
OllamaProviderHandler,
|
|
41
|
+
get_provider_handler,
|
|
42
|
+
PROVIDER_HANDLERS
|
|
43
|
+
)
|
|
44
|
+
from .handlers import RequestHandler, RotationHandler, AutoselectHandler
|
|
45
|
+
|
|
46
|
+
__version__ = "0.1.0"
|
|
47
|
+
__all__ = [
|
|
48
|
+
# Config
|
|
49
|
+
"config",
|
|
50
|
+
"Config",
|
|
51
|
+
"ProviderConfig",
|
|
52
|
+
"RotationConfig",
|
|
53
|
+
"AppConfig",
|
|
54
|
+
"AutoselectConfig",
|
|
55
|
+
"AutoselectModelInfo",
|
|
56
|
+
# Models
|
|
57
|
+
"Message",
|
|
58
|
+
"ChatCompletionRequest",
|
|
59
|
+
"ChatCompletionResponse",
|
|
60
|
+
"Model",
|
|
61
|
+
"Provider",
|
|
62
|
+
"ErrorTracking",
|
|
63
|
+
"AutoselectModelInfo",
|
|
64
|
+
"AutoselectConfig",
|
|
65
|
+
# Providers
|
|
66
|
+
"BaseProviderHandler",
|
|
67
|
+
"GoogleProviderHandler",
|
|
68
|
+
"OpenAIProviderHandler",
|
|
69
|
+
"AnthropicProviderHandler",
|
|
70
|
+
"OllamaProviderHandler",
|
|
71
|
+
"get_provider_handler",
|
|
72
|
+
"PROVIDER_HANDLERS",
|
|
73
|
+
# Handlers
|
|
74
|
+
"RequestHandler",
|
|
75
|
+
"RotationHandler",
|
|
76
|
+
"AutoselectHandler",
|
|
77
|
+
]
|
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyleft (C) 2026 Stefy Lanza <stefy@nexlab.net>
|
|
3
|
+
|
|
4
|
+
AISBF - AI Service Broker Framework || AI Should Be Free
|
|
5
|
+
|
|
6
|
+
Configuration management for AISBF.
|
|
7
|
+
|
|
8
|
+
This program is free software: you can redistribute it and/or modify
|
|
9
|
+
it under the terms of the GNU General Public License as published by
|
|
10
|
+
the Free Software Foundation, either version 3 of the License, or
|
|
11
|
+
(at your option) any later version.
|
|
12
|
+
|
|
13
|
+
This program is distributed in the hope that it will be useful,
|
|
14
|
+
but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
15
|
+
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
16
|
+
GNU General Public License for more details.
|
|
17
|
+
|
|
18
|
+
You should have received a copy of the GNU General Public License
|
|
19
|
+
along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
20
|
+
|
|
21
|
+
Why did the programmer quit his job? Because he didn't get arrays!
|
|
22
|
+
|
|
23
|
+
Configuration management for AISBF.
|
|
24
|
+
"""
|
|
25
|
+
from typing import Dict, List, Optional
|
|
26
|
+
from pydantic import BaseModel, Field
|
|
27
|
+
import json
|
|
28
|
+
import shutil
|
|
29
|
+
from pathlib import Path
|
|
30
|
+
|
|
31
|
+
class ProviderConfig(BaseModel):
|
|
32
|
+
id: str
|
|
33
|
+
name: str
|
|
34
|
+
endpoint: str
|
|
35
|
+
type: str
|
|
36
|
+
api_key_required: bool
|
|
37
|
+
|
|
38
|
+
class RotationConfig(BaseModel):
|
|
39
|
+
providers: List[Dict]
|
|
40
|
+
|
|
41
|
+
class AutoselectModelInfo(BaseModel):
|
|
42
|
+
model_id: str
|
|
43
|
+
description: str
|
|
44
|
+
|
|
45
|
+
class AutoselectConfig(BaseModel):
|
|
46
|
+
model_name: str
|
|
47
|
+
description: str
|
|
48
|
+
fallback: str
|
|
49
|
+
available_models: List[AutoselectModelInfo]
|
|
50
|
+
|
|
51
|
+
class AppConfig(BaseModel):
|
|
52
|
+
providers: Dict[str, ProviderConfig]
|
|
53
|
+
rotations: Dict[str, RotationConfig]
|
|
54
|
+
autoselect: Dict[str, AutoselectConfig]
|
|
55
|
+
error_tracking: Dict[str, Dict]
|
|
56
|
+
|
|
57
|
+
class Config:
|
|
58
|
+
def __init__(self):
|
|
59
|
+
self._ensure_config_directory()
|
|
60
|
+
self._load_providers()
|
|
61
|
+
self._load_rotations()
|
|
62
|
+
self._load_autoselect()
|
|
63
|
+
self._initialize_error_tracking()
|
|
64
|
+
|
|
65
|
+
def _get_config_source_dir(self):
|
|
66
|
+
"""Get the directory containing default config files"""
|
|
67
|
+
# Try installed location first
|
|
68
|
+
installed_dirs = [
|
|
69
|
+
Path('/usr/share/aisbf'),
|
|
70
|
+
Path.home() / '.local' / 'share' / 'aisbf',
|
|
71
|
+
]
|
|
72
|
+
|
|
73
|
+
for installed_dir in installed_dirs:
|
|
74
|
+
if installed_dir.exists() and (installed_dir / 'providers.json').exists():
|
|
75
|
+
return installed_dir
|
|
76
|
+
|
|
77
|
+
# Fallback to source tree config directory
|
|
78
|
+
# This is for development mode
|
|
79
|
+
source_dir = Path(__file__).parent.parent / 'config'
|
|
80
|
+
if source_dir.exists() and (source_dir / 'providers.json').exists():
|
|
81
|
+
return source_dir
|
|
82
|
+
|
|
83
|
+
# Last resort: try the old location in the package directory
|
|
84
|
+
package_dir = Path(__file__).parent
|
|
85
|
+
if (package_dir / 'providers.json').exists():
|
|
86
|
+
return package_dir
|
|
87
|
+
|
|
88
|
+
raise FileNotFoundError("Could not find configuration files")
|
|
89
|
+
|
|
90
|
+
def _ensure_config_directory(self):
|
|
91
|
+
"""Ensure ~/.aisbf/ directory exists and copy default config files if needed"""
|
|
92
|
+
config_dir = Path.home() / '.aisbf'
|
|
93
|
+
|
|
94
|
+
# Create config directory if it doesn't exist
|
|
95
|
+
config_dir.mkdir(exist_ok=True)
|
|
96
|
+
|
|
97
|
+
# Get the source directory for default config files
|
|
98
|
+
try:
|
|
99
|
+
source_dir = self._get_config_source_dir()
|
|
100
|
+
except FileNotFoundError:
|
|
101
|
+
print("Warning: Could not find default configuration files")
|
|
102
|
+
return
|
|
103
|
+
|
|
104
|
+
# Copy default config files if they don't exist
|
|
105
|
+
for config_file in ['providers.json', 'rotations.json', 'autoselect.json']:
|
|
106
|
+
src = source_dir / config_file
|
|
107
|
+
dst = config_dir / config_file
|
|
108
|
+
|
|
109
|
+
if not dst.exists() and src.exists():
|
|
110
|
+
shutil.copy2(src, dst)
|
|
111
|
+
print(f"Created default config file: {dst}")
|
|
112
|
+
|
|
113
|
+
def _load_providers(self):
|
|
114
|
+
providers_path = Path.home() / '.aisbf' / 'providers.json'
|
|
115
|
+
if not providers_path.exists():
|
|
116
|
+
# Fallback to source config if user config doesn't exist
|
|
117
|
+
try:
|
|
118
|
+
source_dir = self._get_config_source_dir()
|
|
119
|
+
providers_path = source_dir / 'providers.json'
|
|
120
|
+
except FileNotFoundError:
|
|
121
|
+
raise FileNotFoundError("Could not find providers.json configuration file")
|
|
122
|
+
|
|
123
|
+
with open(providers_path) as f:
|
|
124
|
+
data = json.load(f)
|
|
125
|
+
self.providers = {k: ProviderConfig(**v) for k, v in data['providers'].items()}
|
|
126
|
+
|
|
127
|
+
def _load_rotations(self):
|
|
128
|
+
rotations_path = Path.home() / '.aisbf' / 'rotations.json'
|
|
129
|
+
if not rotations_path.exists():
|
|
130
|
+
# Fallback to source config if user config doesn't exist
|
|
131
|
+
try:
|
|
132
|
+
source_dir = self._get_config_source_dir()
|
|
133
|
+
rotations_path = source_dir / 'rotations.json'
|
|
134
|
+
except FileNotFoundError:
|
|
135
|
+
raise FileNotFoundError("Could not find rotations.json configuration file")
|
|
136
|
+
|
|
137
|
+
with open(rotations_path) as f:
|
|
138
|
+
data = json.load(f)
|
|
139
|
+
self.rotations = {k: RotationConfig(**v) for k, v in data['rotations'].items()}
|
|
140
|
+
|
|
141
|
+
def _load_autoselect(self):
|
|
142
|
+
autoselect_path = Path.home() / '.aisbf' / 'autoselect.json'
|
|
143
|
+
if not autoselect_path.exists():
|
|
144
|
+
# Fallback to source config if user config doesn't exist
|
|
145
|
+
try:
|
|
146
|
+
source_dir = self._get_config_source_dir()
|
|
147
|
+
autoselect_path = source_dir / 'autoselect.json'
|
|
148
|
+
except FileNotFoundError:
|
|
149
|
+
raise FileNotFoundError("Could not find autoselect.json configuration file")
|
|
150
|
+
|
|
151
|
+
with open(autoselect_path) as f:
|
|
152
|
+
data = json.load(f)
|
|
153
|
+
self.autoselect = {k: AutoselectConfig(**v) for k, v in data.items()}
|
|
154
|
+
|
|
155
|
+
def _initialize_error_tracking(self):
|
|
156
|
+
self.error_tracking = {}
|
|
157
|
+
for provider_id in self.providers:
|
|
158
|
+
self.error_tracking[provider_id] = {
|
|
159
|
+
'failures': 0,
|
|
160
|
+
'last_failure': None,
|
|
161
|
+
'disabled_until': None
|
|
162
|
+
}
|
|
163
|
+
|
|
164
|
+
def get_provider(self, provider_id: str) -> ProviderConfig:
|
|
165
|
+
return self.providers.get(provider_id)
|
|
166
|
+
|
|
167
|
+
def get_rotation(self, rotation_id: str) -> RotationConfig:
|
|
168
|
+
return self.rotations.get(rotation_id)
|
|
169
|
+
|
|
170
|
+
def get_autoselect(self, autoselect_id: str) -> AutoselectConfig:
|
|
171
|
+
return self.autoselect.get(autoselect_id)
|
|
172
|
+
|
|
173
|
+
config = Config()
|