From fbf344042d37db192926498215c3556e635be5c6 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Tue, 8 Jul 2025 11:43:44 +0200 Subject: [PATCH 1/2] LCORE-293: support for customization options --- src/configuration.py | 9 +++++++++ src/models/config.py | 20 ++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/src/configuration.py b/src/configuration.py index d6cf2a7a..137d304a 100644 --- a/src/configuration.py +++ b/src/configuration.py @@ -6,6 +6,7 @@ import yaml from models.config import ( Configuration, + Customization, LLamaStackConfiguration, UserDataCollection, ServiceConfiguration, @@ -90,5 +91,13 @@ def authentication_configuration(self) -> Optional[AuthenticationConfiguration]: ), "logic error: configuration is not loaded" return self._configuration.authentication + @property + def customization(self) -> Optional[Customization]: + """Return customization configuration.""" + assert ( + self._configuration is not None + ), "logic error: configuration is not loaded" + return self._configuration.customization + configuration: AppConfig = AppConfig() diff --git a/src/models/config.py b/src/models/config.py index b67a090d..3bef0ccf 100644 --- a/src/models/config.py +++ b/src/models/config.py @@ -7,6 +7,8 @@ import constants +from utils import checks + class TLSConfiguration(BaseModel): """TLS configuration.""" @@ -123,6 +125,23 @@ def check_authentication_model(self) -> Self: return self +class Customization(BaseModel): + """Service customization.""" + + system_prompt_path: Optional[FilePath] = None + system_prompt: Optional[str] = None + + @model_validator(mode="after") + def check_authentication_model(self) -> Self: + """Load system prompt from file.""" + if self.system_prompt_path is not None: + checks.file_check(self.system_prompt_path, "system prompt") + self.system_prompt = checks.get_attribute_from_file( + dict(self), "system_prompt_path" + ) + return self + + class Configuration(BaseModel): """Global service configuration.""" @@ -134,6 +153,7 @@ class Configuration(BaseModel): authentication: Optional[AuthenticationConfiguration] = ( AuthenticationConfiguration() ) + customization: Optional[Customization] = None def dump(self, filename: str = "configuration.json") -> None: """Dump actual configuration into JSON file.""" From 5eeb059302344e2ec80f87622f0fcfb8a51432b7 Mon Sep 17 00:00:00 2001 From: Pavel Tisnovsky Date: Tue, 8 Jul 2025 11:45:46 +0200 Subject: [PATCH 2/2] Updated unit tests accordingly --- tests/unit/__init__.py | 4 +++ tests/unit/app/endpoints/test_config.py | 1 + tests/unit/app/endpoints/test_info.py | 1 + tests/unit/app/endpoints/test_models.py | 3 ++ tests/unit/app/endpoints/test_query.py | 1 + tests/unit/models/test_config.py | 9 +++++ tests/unit/test_configuration.py | 47 +++++++++++++++++++++++++ tests/unit/utils/test_common.py | 6 ++++ 8 files changed, 72 insertions(+) diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py index 268ef062..d16b811f 100644 --- a/tests/unit/__init__.py +++ b/tests/unit/__init__.py @@ -27,6 +27,10 @@ "k8s_ca_cert_path": None, "k8s_cluster_api": None, }, + "customization": { + "system_prompt_path": None, + "system_prompt": None, + }, } # NOTE(lucasagomes): Configuration must be initialized before importing diff --git a/tests/unit/app/endpoints/test_config.py b/tests/unit/app/endpoints/test_config.py index bb9ab122..4a2cd3c1 100644 --- a/tests/unit/app/endpoints/test_config.py +++ b/tests/unit/app/endpoints/test_config.py @@ -43,6 +43,7 @@ def test_config_endpoint_handler_configuration_loaded(mocker): "authentication": { "module": "noop", }, + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) diff --git a/tests/unit/app/endpoints/test_info.py b/tests/unit/app/endpoints/test_info.py index 9b36cbfa..7dc1d46b 100644 --- a/tests/unit/app/endpoints/test_info.py +++ b/tests/unit/app/endpoints/test_info.py @@ -22,6 +22,7 @@ def test_info_endpoint(mocker): "user_data_collection": { "feedback_disabled": True, }, + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) diff --git a/tests/unit/app/endpoints/test_models.py b/tests/unit/app/endpoints/test_models.py index c9943ea2..b7451579 100644 --- a/tests/unit/app/endpoints/test_models.py +++ b/tests/unit/app/endpoints/test_models.py @@ -45,6 +45,7 @@ def test_models_endpoint_handler_improper_llama_stack_configuration(mocker): "transcripts_disabled": True, }, "mcp_servers": [], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -82,6 +83,7 @@ def test_models_endpoint_handler_configuration_loaded(mocker): "user_data_collection": { "feedback_disabled": True, }, + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -114,6 +116,7 @@ def test_models_endpoint_handler_unable_to_retrieve_models_list(mocker): "user_data_collection": { "feedback_disabled": True, }, + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) diff --git a/tests/unit/app/endpoints/test_query.py b/tests/unit/app/endpoints/test_query.py index 6375353c..38afa72b 100644 --- a/tests/unit/app/endpoints/test_query.py +++ b/tests/unit/app/endpoints/test_query.py @@ -41,6 +41,7 @@ def setup_configuration(): "transcripts_disabled": True, }, "mcp_servers": [], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) diff --git a/tests/unit/models/test_config.py b/tests/unit/models/test_config.py index 7636888b..254ae783 100644 --- a/tests/unit/models/test_config.py +++ b/tests/unit/models/test_config.py @@ -250,6 +250,7 @@ def test_configuration_empty_mcp_servers() -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=[], + customization=None, ) assert cfg is not None assert cfg.mcp_servers == [] @@ -270,6 +271,7 @@ def test_configuration_single_mcp_server() -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=[mcp_server], + customization=None, ) assert cfg is not None assert len(cfg.mcp_servers) == 1 @@ -296,6 +298,7 @@ def test_configuration_multiple_mcp_servers() -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=mcp_servers, + customization=None, ) assert cfg is not None assert len(cfg.mcp_servers) == 3 @@ -317,6 +320,7 @@ def test_dump_configuration(tmp_path) -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=[], + customization=None, ) assert cfg is not None dump_file = tmp_path / "test.json" @@ -370,6 +374,7 @@ def test_dump_configuration(tmp_path) -> None: "k8s_ca_cert_path": None, "k8s_cluster_api": None, }, + "customization": None, } @@ -388,6 +393,7 @@ def test_dump_configuration_with_one_mcp_server(tmp_path) -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=mcp_servers, + customization=None, ) dump_file = tmp_path / "test.json" cfg.dump(dump_file) @@ -442,6 +448,7 @@ def test_dump_configuration_with_one_mcp_server(tmp_path) -> None: "k8s_ca_cert_path": None, "k8s_cluster_api": None, }, + "customization": None, } @@ -462,6 +469,7 @@ def test_dump_configuration_with_more_mcp_servers(tmp_path) -> None: feedback_disabled=True, feedback_storage=None ), mcp_servers=mcp_servers, + customization=None, ) dump_file = tmp_path / "test.json" cfg.dump(dump_file) @@ -532,6 +540,7 @@ def test_dump_configuration_with_more_mcp_servers(tmp_path) -> None: "k8s_ca_cert_path": None, "k8s_cluster_api": None, }, + "customization": None, } diff --git a/tests/unit/test_configuration.py b/tests/unit/test_configuration.py index 61eba233..8ba212aa 100644 --- a/tests/unit/test_configuration.py +++ b/tests/unit/test_configuration.py @@ -49,6 +49,7 @@ def test_init_from_dict() -> None: "feedback_disabled": True, }, "mcp_servers": [], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -110,6 +111,7 @@ def test_init_from_dict_with_mcp_servers() -> None: "url": "https://api.example.com", }, ], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -216,6 +218,7 @@ def test_mcp_servers_property_empty() -> None: "feedback_disabled": True, }, "mcp_servers": [], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -251,6 +254,7 @@ def test_mcp_servers_property_with_servers() -> None: "url": "http://localhost:8080", }, ], + "customization": None, } cfg = AppConfig() cfg.init_from_dict(config_dict) @@ -306,3 +310,46 @@ def test_mcp_servers_not_loaded(): AssertionError, match="logic error: configuration is not loaded" ): cfg.mcp_servers + + +def test_load_configuration_with_customization(tmpdir) -> None: + """Test loading configuration from YAML file with customization.""" + system_prompt_filename = tmpdir / "system_prompt.txt" + with open(system_prompt_filename, "w") as fout: + fout.write("this is system prompt") + + cfg_filename = tmpdir / "config.yaml" + with open(cfg_filename, "w") as fout: + fout.write( + f""" +name: test service +service: + host: localhost + port: 8080 + auth_enabled: false + workers: 1 + color_log: true + access_log: true +llama_stack: + use_as_library_client: false + url: http://localhost:8321 + api_key: test-key +user_data_collection: + feedback_disabled: true +mcp_servers: + - name: filesystem-server + url: http://localhost:3000 + - name: git-server + provider_id: custom-git-provider + url: https://git.example.com/mcp +customization: + system_prompt_path: {system_prompt_filename} + """ + ) + + cfg = AppConfig() + cfg.load_configuration(cfg_filename) + + assert cfg.customization is not None + assert cfg.customization.system_prompt is not None + assert cfg.customization.system_prompt == "this is system prompt" diff --git a/tests/unit/utils/test_common.py b/tests/unit/utils/test_common.py index df969dc9..40c02e56 100644 --- a/tests/unit/utils/test_common.py +++ b/tests/unit/utils/test_common.py @@ -42,6 +42,7 @@ async def test_register_mcp_servers_empty_list(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=[], + customization=None, ) # Call the function await register_mcp_servers_async(mock_logger, config) @@ -80,6 +81,7 @@ async def test_register_mcp_servers_single_server_not_registered(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=[mcp_server], + customization=None, ) # Call the function @@ -122,6 +124,7 @@ async def test_register_mcp_servers_single_server_already_registered(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=[mcp_server], + customization=None, ) # Call the function @@ -167,6 +170,7 @@ async def test_register_mcp_servers_multiple_servers_mixed_registration(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=mcp_servers, + customization=None, ) # Call the function @@ -219,6 +223,7 @@ async def test_register_mcp_servers_with_custom_provider(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=[mcp_server], + customization=None, ) # Call the function @@ -267,6 +272,7 @@ async def test_register_mcp_servers_async_with_library_client(mocker): ), user_data_collection=UserDataCollection(feedback_disabled=True), mcp_servers=[mcp_server], + customization=None, ) # Call the async function