143 lines
4.1 KiB
Python
143 lines
4.1 KiB
Python
import yaml
|
|
from pathlib import Path
|
|
from typing import Type, TypeVar, Any, Optional, Final
|
|
from dataclasses import dataclass, asdict, field
|
|
|
|
ConfigInst = TypeVar('ConfigInst', bound='Config')
|
|
AIConfigInst = TypeVar('AIConfigInst', bound='AIConfig')
|
|
OpenAIConfigInst = TypeVar('OpenAIConfigInst', bound='OpenAIConfig')
|
|
|
|
|
|
supported_ais: list[str] = ['openai']
|
|
default_ai_ID: str = 'default'
|
|
default_config_path = '.config.yaml'
|
|
|
|
|
|
class ConfigError(Exception):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class AIConfig:
|
|
"""
|
|
The base class of all AI configurations.
|
|
"""
|
|
ID: str
|
|
name: str
|
|
|
|
|
|
@dataclass
|
|
class OpenAIConfig(AIConfig):
|
|
"""
|
|
The OpenAI section of the configuration file.
|
|
"""
|
|
# the name must not be changed
|
|
name: Final[str] = 'openai'
|
|
# all members have default values, so we can easily create
|
|
# a default configuration
|
|
ID: str = 'default'
|
|
api_key: str = '0123456789'
|
|
system: str = 'You are an assistant'
|
|
model: str = 'gpt-3.5-turbo-16k'
|
|
temperature: float = 1.0
|
|
max_tokens: int = 4000
|
|
top_p: float = 1.0
|
|
frequency_penalty: float = 0.0
|
|
presence_penalty: float = 0.0
|
|
|
|
@classmethod
|
|
def from_dict(cls: Type[OpenAIConfigInst], source: dict[str, Any]) -> OpenAIConfigInst:
|
|
"""
|
|
Create OpenAIConfig from a dict.
|
|
"""
|
|
res = cls(
|
|
system=str(source['system']),
|
|
api_key=str(source['api_key']),
|
|
model=str(source['model']),
|
|
max_tokens=int(source['max_tokens']),
|
|
temperature=float(source['temperature']),
|
|
top_p=float(source['top_p']),
|
|
frequency_penalty=float(source['frequency_penalty']),
|
|
presence_penalty=float(source['presence_penalty'])
|
|
)
|
|
# overwrite default ID if provided
|
|
if 'ID' in source:
|
|
res.ID = source['ID']
|
|
return res
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|
|
|
|
|
|
def ai_config_instance(name: str, conf_dict: Optional[dict[str, Any]] = None) -> AIConfig:
|
|
"""
|
|
Creates an AIConfig instance of the given name.
|
|
"""
|
|
if name.lower() == 'openai':
|
|
if conf_dict is None:
|
|
return OpenAIConfig()
|
|
else:
|
|
return OpenAIConfig.from_dict(conf_dict)
|
|
else:
|
|
raise ConfigError(f"AI '{name}' is not supported")
|
|
|
|
|
|
def create_default_ai_configs() -> dict[str, AIConfig]:
|
|
"""
|
|
Create a dict containing default configurations for all supported AIs.
|
|
"""
|
|
return {ai_config_instance(name).ID: ai_config_instance(name) for name in supported_ais}
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
"""
|
|
The configuration file structure.
|
|
"""
|
|
# all members have default values, so we can easily create
|
|
# a default configuration
|
|
db: str = './db/'
|
|
ais: dict[str, AIConfig] = field(default_factory=create_default_ai_configs)
|
|
|
|
@classmethod
|
|
def from_dict(cls: Type[ConfigInst], source: dict[str, Any]) -> ConfigInst:
|
|
"""
|
|
Create Config from a dict.
|
|
"""
|
|
# create the correct AI type instances
|
|
ais: dict[str, AIConfig] = {}
|
|
for ID, conf in source['ais'].items():
|
|
# add the AI ID to the config (for easy internal access)
|
|
conf['ID'] = ID
|
|
ai_conf = ai_config_instance(conf['name'], conf)
|
|
ais[ID] = ai_conf
|
|
return cls(
|
|
db=str(source['db']),
|
|
ais=ais
|
|
)
|
|
|
|
@classmethod
|
|
def create_default(self, file_path: Path) -> None:
|
|
"""
|
|
Creates a default Config in the given file.
|
|
"""
|
|
conf = Config()
|
|
conf.to_file(file_path)
|
|
|
|
@classmethod
|
|
def from_file(cls: Type[ConfigInst], path: str) -> ConfigInst:
|
|
with open(path, 'r') as f:
|
|
source = yaml.load(f, Loader=yaml.FullLoader)
|
|
return cls.from_dict(source)
|
|
|
|
def to_file(self, file_path: Path) -> None:
|
|
# remove the AI name from the config (for a cleaner format)
|
|
data = self.as_dict()
|
|
for conf in data['ais'].values():
|
|
del (conf['ID'])
|
|
with open(file_path, 'w') as f:
|
|
yaml.dump(data, f, sort_keys=False)
|
|
|
|
def as_dict(self) -> dict[str, Any]:
|
|
return asdict(self)
|