ChatMastermind/chatmastermind/configuration.py

import yaml
from pathlib import Path
from typing import Type, TypeVar, Any, Optional, ClassVar
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_config_path = '.config.yaml'


class ConfigError(Exception):
    pass


def str_presenter(dumper: yaml.Dumper, data: str) -> yaml.ScalarNode:
    """
    Changes the YAML dump style to multiline syntax for multiline strings.
    """
    if len(data.splitlines()) > 1:
        return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
    return dumper.represent_scalar('tag:yaml.org,2002:str', data)


yaml.add_representer(str, str_presenter)


@dataclass
class AIConfig:
    """
    The base class of all AI configurations.
    """
    # the name of the AI the config class represents
    # -> it's a class variable and thus not part of the
    #    dataclass constructor
    name: ClassVar[str]
    # a user-defined ID for an AI configuration entry
    ID: str

    # the name must not be changed
    def __setattr__(self, name: str, value: Any) -> None:
        if name == 'name':
            raise AttributeError("'{name}' is not allowed to be changed")
        else:
            super().__setattr__(name, value)


@dataclass
class OpenAIConfig(AIConfig):
    """
    The OpenAI section of the configuration file.
    """
    name: ClassVar[str] = 'openai'

    # all members have default values, so we can easily create
    # a default configuration
    ID: str = 'myopenai'
    api_key: str = '0123456789'
    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
    system: str = 'You are an assistant'

    @classmethod
    def from_dict(cls: Type[OpenAIConfigInst], source: dict[str, Any]) -> OpenAIConfigInst:
        """
        Create OpenAIConfig from a dict.
        """
        res = cls(
            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']),
            system=str(source['system'])
        )
        # overwrite default ID if provided
        if 'ID' in source:
            res.ID = source['ID']
        return res


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"Unknown AI '{name}'")


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 (with the same format as the config file).
        """
        # 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]:
        res = asdict(self)
        # add the AI name manually (as first element)
        # (not done by 'asdict' because it's a class variable)
        for ID, conf in res['ais'].items():
            res['ais'][ID] = {**{'name': self.ais[ID].name}, **conf}
        return res