Updated how config priority is done

Signed-off-by: JonahSussman <sussmanjonah@gmail.com>

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}"

kai/config_local.toml

@@ -0,0 +1,6 @@
+[incident_store.args]
+provider = "postgresql"
+host = "127.0.0.1"
+database = "kai"
+user = "kai"
+password = "dog8code"

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):
         """

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: