diff --git a/build/example_config.toml b/build/config_example.toml similarity index 100% rename from build/example_config.toml rename to build/config_example.toml diff --git a/build/entrypoint.sh b/build/entrypoint.sh index d9f359c3a..6d930b828 100755 --- a/build/entrypoint.sh +++ b/build/entrypoint.sh @@ -1,11 +1,5 @@ #!/bin/bash -# If a custom config is specified, use it -if [[ -f /podman_compose/build/config.toml ]]; then - printf "Using custom config.toml\n" - cp /podman_compose/build/config.toml /kai/kai/config.toml -fi - until PGPASSWORD="${KAI__INCIDENT_STORE__ARGS__PASSWORD}" pg_isready -q -h "${KAI__INCIDENT_STORE__ARGS__HOST}" -U "${KAI__INCIDENT_STORE__ARGS__USER}" -d "${KAI__INCIDENT_STORE__ARGS__DATABASE}"; do sleep 1 done @@ -34,7 +28,15 @@ if [[ ${MODE} != "importer" ]]; then sleep 5 fi fi - PYTHONPATH="/kai/kai" python /kai/kai/server.py + + # If a custom config is specified, use it + if [[ -f /podman_compose/build/config.toml ]]; then + printf "Using custom config.toml\n" + PYTHONPATH="/kai/kai" python /kai/kai/server.py --config-file /podman_compose/build/config.toml + else + PYTHONPATH="/kai/kai" python /kai/kai/server.py + fi + else cd /kai || exit python ./kai/hub_importer.py --loglevel "${KAI__LOG_LEVEL}" --config_filepath ./kai/config.toml "${KAI__HUB_URL}" "${KAI__IMPORTER_ARGS}" diff --git a/kai/config.toml b/kai/config_default.toml similarity index 100% rename from kai/config.toml rename to kai/config_default.toml diff --git a/kai/config_local.toml b/kai/config_local.toml new file mode 100644 index 000000000..ccd2e9d17 --- /dev/null +++ b/kai/config_local.toml @@ -0,0 +1,6 @@ +[incident_store.args] +provider = "postgresql" +host = "127.0.0.1" +database = "kai" +user = "kai" +password = "dog8code" diff --git a/kai/models/kai_config.py b/kai/models/kai_config.py index 28ada28bf..3557968c1 100644 --- a/kai/models/kai_config.py +++ b/kai/models/kai_config.py @@ -5,7 +5,14 @@ import yaml from pydantic import BaseModel, Field, model_validator -from pydantic_settings import BaseSettings, SettingsConfigDict +from pydantic.fields import FieldInfo +from pydantic_settings import ( + BaseSettings, + PydanticBaseSettingsSource, + SettingsConfigDict, +) + +from kai.constants import PATH_KAI """ https://docs.pydantic.dev/2.0/migration/#required-optional-and-nullable-fields @@ -154,8 +161,50 @@ class KaiConfigModels(BaseModel): # Main config -# TODO: Evaluate the usage of pydantic-settings to simplify command line -# argument management. +class TomlConfigSettingsSource(PydanticBaseSettingsSource): + """ + Helper class to load a TOML file and convert it to a dictionary for + pydantic-settings. + """ + + def __init__(self, settings_cls: type[BaseSettings], str_path: str): + self.settings_cls = settings_cls + self.config = settings_cls.model_config + + self.str_path = str_path + + if not os.path.exists(str_path): + self.file_content_toml = {} + else: + with open(str_path, "r") as f: + self.file_content_toml = tomllib.loads(f.read()) + + def get_field_value( + self, field: FieldInfo, field_name: str + ) -> tuple[Any, str, bool]: + return self.file_content_toml.get(field_name), field_name, False + + def prepare_field_value( + self, field_name: str, field: FieldInfo, value: Any, value_is_complex: bool + ) -> Any: + return value + + def __call__(self) -> dict[str, Any]: + d: dict[str, Any] = {} + + for field_name, field in self.settings_cls.model_fields.items(): + field_value, field_key, value_is_complex = self.get_field_value( + field, field_name + ) + field_value = self.prepare_field_value( + field_name, field, field_value, value_is_complex + ) + if field_value is not None: + d[field_key] = field_value + + return d + + class KaiConfig(BaseSettings): model_config = SettingsConfigDict(env_prefix="KAI__", env_nested_delimiter="__") @@ -179,6 +228,38 @@ class KaiConfig(BaseSettings): default_factory=lambda: [SolutionConsumerKind.DIFF_ONLY] ) + @classmethod + def settings_customise_sources( + cls, + settings_cls: type[BaseSettings], + init_settings: PydanticBaseSettingsSource, + env_settings: PydanticBaseSettingsSource, + dotenv_settings: PydanticBaseSettingsSource, + file_secret_settings: PydanticBaseSettingsSource, + ) -> tuple[PydanticBaseSettingsSource, ...]: + """ + Config is loaded with the following priority (higher overrides lower): + + - Command line args (not implemented) + - Config file that is declared on the command line / via init arguments. + - Environment vars + - Local config file (config_local.toml) + - Global config file (config_default.toml) + - Default field values + """ + return ( + init_settings, + env_settings, + dotenv_settings, + file_secret_settings, + TomlConfigSettingsSource( + settings_cls, os.path.join(PATH_KAI, "config_local.toml") + ), + TomlConfigSettingsSource( + settings_cls, os.path.join(PATH_KAI, "config_default.toml") + ), + ) + @staticmethod def model_validate_filepath(filepath: str): """ diff --git a/kai/server.py b/kai/server.py index 5f3b3f741..0c88e13e2 100644 --- a/kai/server.py +++ b/kai/server.py @@ -2,15 +2,15 @@ """This module is intended to facilitate using Konveyor with LLMs.""" +import argparse import logging -import os import pprint -from functools import lru_cache +from functools import cache +from typing import Optional from aiohttp import web from gunicorn.app.wsgiapp import WSGIApplication -from kai.constants import PATH_KAI from kai.kai_logging import initLoggingFromConfig from kai.models.kai_config import KaiConfig from kai.routes import kai_routes @@ -26,12 +26,22 @@ # the same manner as `git stash apply` -@lru_cache +@cache def get_config(): - if not os.path.exists(os.path.join(PATH_KAI, "config.toml")): - raise FileNotFoundError("Config file not found.") - - return KaiConfig.model_validate_filepath(os.path.join(PATH_KAI, "config.toml")) + parser = argparse.ArgumentParser() + parser.add_argument( + "--config-file", + help="Path to an optional config file.", + type=Optional[str], + default=None, + required=False, + ) + args = parser.parse_args() + + if args.config_file: + return KaiConfig.model_validate_filepath(args.config_file) + + return KaiConfig() def app() -> web.Application: