MatrixGPT/matrix_gpt/config.py

import copy
from pathlib import Path
from types import NoneType
from typing import Union

import bison

OPENAI_DEFAULT_SYSTEM_PROMPT = ""
OPENAI_DEFAULT_INJECTED_SYSTEM_PROMPT = ""

config_scheme = bison.Scheme(
    bison.Option('store_path', default='bot-store/', field_type=str),
    bison.DictOption('auth', scheme=bison.Scheme(
        bison.Option('username', field_type=str, required=True),
        bison.Option('password', field_type=str, required=True),
        bison.Option('homeserver', field_type=str, required=True),
        bison.Option('device_id', field_type=str, required=True),
    )),
    bison.ListOption('allowed_to_chat', default=['all']),
    bison.ListOption('allowed_to_thread', default=['all']),
    bison.ListOption('allowed_to_invite', default=['all']),
    bison.ListOption('autojoin_rooms', default=[]),
    bison.ListOption('whitelist_rooms', default=[]),
    bison.ListOption('blacklist_rooms', default=[]),
    bison.Option('reply_in_thread', default=True, field_type=bool),
    bison.Option('set_avatar', default=True, field_type=bool),
    bison.Option('response_timeout', default=120, field_type=int),
    bison.ListOption('command', member_scheme=bison.Scheme(
        bison.Option('trigger', field_type=str, required=True),
        bison.Option('model', field_type=str, required=True),
        bison.ListOption('allowed_to_chat', default=['all']),
        bison.ListOption('allowed_to_thread', default=['all']),
        bison.Option('max_tokens', field_type=int, default=0, required=False),
    )),
    bison.DictOption('openai', scheme=bison.Scheme(
        bison.Option('api_key', field_type=str, required=True),
        bison.Option('api_base', field_type=[str, NoneType], default=None, required=False),
        bison.Option('api_retries', field_type=int, default=2),
        bison.Option('temperature', field_type=float, default=0.5),
        bison.Option('system_prompt', field_type=[str, NoneType], default=OPENAI_DEFAULT_SYSTEM_PROMPT),
        bison.Option('injected_system_prompt', field_type=[str, NoneType], default=OPENAI_DEFAULT_INJECTED_SYSTEM_PROMPT),
    )),
    bison.DictOption('logging', scheme=bison.Scheme(
        bison.Option('log_level', field_type=str, default='info'),
        bison.Option('log_full_response', field_type=bool, default=True),
    )),
)


class ConfigManager:
    def __init__(self):
        self._config = bison.Bison(scheme=config_scheme)
        self._command_prefixes = {}
        self._loaded = False

    def load(self, path: Path):
        if self._loaded:
            raise Exception('Already loaded')
        self._config.config_name = 'config'
        self._config.config_format = bison.bison.YAML
        self._config.add_config_paths(str(path.parent))
        self._config.parse()
        self._command_prefixes = self._generate_command_prefixes()
        self._loaded = True

    def validate(self):
        self._config.validate()

    @property
    def config(self):
        return copy.deepcopy(self._config)

    def _generate_command_prefixes(self):
        command_prefixes = {}
        for item in self._config.config['command']:
            command_prefixes[item['trigger']] = item
        return command_prefixes

    @property
    def command_prefixes(self):
        return self._command_prefixes

    def get(self, key, default=None):
        return copy.copy(self._config.get(key, default))

    def __setitem__(self, key, item):
        raise Exception

    def __getitem__(self, key):
        return self._config.config[key]

    def __repr__(self):
        return repr(self._config.config)

    def __len__(self):
        return len(self._config.config)

    def __delitem__(self, key):
        raise Exception


global_config = ConfigManager()