make llm gateway implementation configurable

Author: atulikumwenayo

src/datacustomcode/llm_gateway/base.py

@@ -14,9 +14,11 @@
 # limitations under the License.
 from __future__ import annotations
 
-from abc import abstractmethod
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 
+from datacustomcode.mixin import UserExtendableNamedConfigMixin
+
 if TYPE_CHECKING:
     from datacustomcode.llm_gateway.types.generate_text_request import (
         GenerateTextRequest,
@@ -26,8 +28,10 @@
     )
 
 
-class LLMGateway:
-    def __init__(self) -> None:
+class LLMGateway(ABC, UserExtendableNamedConfigMixin):
+    CONFIG_NAME: str
+
+    def __init__(self, **kwargs):
         pass
 
     @abstractmethod

src/datacustomcode/llm_gateway/default.py

@@ -22,6 +22,8 @@
 
 
 class DefaultLLMGateway(LLMGateway):
+    CONFIG_NAME = "DefaultLLMGateway"
+
     def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
 
         response_data = {

src/datacustomcode/llm_gateway_config.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import (
+    ClassVar,
+    Generic,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from datacustomcode.common_config import (
+    BaseConfig,
+    BaseObjectConfig,
+    default_config_file,
+)
+from datacustomcode.llm_gateway.base import LLMGateway
+
+_E = TypeVar("_E", bound=LLMGateway)
+
+
+class LLMGatewayObjectConfig(BaseObjectConfig, Generic[_E]):
+    type_base: ClassVar[Type[LLMGateway]] = LLMGateway  # type: ignore[type-abstract]
+
+    def to_object(self) -> _E:
+        type_ = self.type_base.subclass_from_config_name(self.type_config_name)
+        return cast(_E, type_(**self.options))
+
+
+class LLMGatewayConfig(BaseConfig):
+    llm_gateway_config: Union[
+        LLMGatewayObjectConfig[LLMGateway], None
+    ] = None
+
+    def update(self, other: "LLMGatewayConfig") -> "LLMGatewayConfig":
+        def merge(
+            config_a: Union[LLMGatewayObjectConfig, None],
+            config_b: Union[LLMGatewayObjectConfig, None],
+        ) -> Union[LLMGatewayObjectConfig, None]:
+            if config_a is not None and config_a.force:
+                return config_a
+            if config_b:
+                return config_b
+            return config_a
+
+        self.llm_gateway_config = merge(
+            self.llm_gateway_config, other.llm_gateway_config
+        )
+        return self
+
+
+# Global LLM Gateway config instance
+llm_gateway_config = LLMGatewayConfig()
+llm_gateway_config.load(default_config_file())

tests/test_llm_gateway_config_update.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import tempfile
+
+import yaml
+
+from datacustomcode.llm_gateway.default import DefaultLLMGateway
+from datacustomcode.llm_gateway_config import (
+    LLMGatewayConfig,
+    LLMGatewayObjectConfig,
+)
+
+
+class TestLLMGatewayConfigUpdate:
+    def test_update_replaces_config_without_force(self):
+        config1 = LLMGatewayConfig(
+            llm_gateway_config=LLMGatewayObjectConfig(
+                type_config_name="OldImplementation", options={"old": True}
+            )
+        )
+
+        config2 = LLMGatewayConfig(
+            llm_gateway_config=LLMGatewayObjectConfig(
+                type_config_name="NewImplementation", options={"new": True}
+            )
+        )
+
+        config1.update(config2)
+
+        assert config1.llm_gateway_config.type_config_name == "NewImplementation"
+        assert config1.llm_gateway_config.options == {"new": True}
+
+    def test_update_respects_force_flag(self):
+        config1 = LLMGatewayConfig(
+            llm_gateway_config=LLMGatewayObjectConfig(
+                type_config_name="ForcedImplementation",
+                options={"forced": True},
+                force=True,
+            )
+        )
+
+        config2 = LLMGatewayConfig(
+            llm_gateway_config=LLMGatewayObjectConfig(
+                type_config_name="NewImplementation", options={"new": True}
+            )
+        )
+
+        config1.update(config2)
+
+        assert (
+            config1.llm_gateway_config.type_config_name == "ForcedImplementation"
+        )
+        assert config1.llm_gateway_config.options == {"forced": True}
+        assert config1.llm_gateway_config.force is True
+
+
+class TestLLMGatewayConfigLoad:
+    def test_load_from_yaml_file(self):
+        config_data = {
+            "llm_gateway_config": {"type_config_name": "DefaultLLMGateway"}
+        }
+
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
+            yaml.dump(config_data, f)
+            temp_file = f.name
+
+        try:
+            config = LLMGatewayConfig()
+            config.load(temp_file)
+
+            assert config.llm_gateway_config is not None
+            assert (
+                config.llm_gateway_config.type_config_name == "DefaultLLMGateway"
+            )
+            llm_gateway = config.llm_gateway_config.to_object()
+            assert llm_gateway is not None
+            assert isinstance(llm_gateway, DefaultLLMGateway)
+        finally:
+            os.unlink(temp_file)

tests/test_runtime_llm_gateway.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datacustomcode.llm_gateway.base import LLMGateway
+from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
+from datacustomcode.llm_gateway.types.generate_text_response import GenerateTextResponse
+from datacustomcode.llm_gateway_config import LLMGatewayObjectConfig
+
+
+class TestCustomLLMGatewayImplementation:
+    def test_custom_implementation_is_discoverable(self):
+        class CustomLLMGateway(LLMGateway):
+            CONFIG_NAME = "CustomLLMGateway"
+
+            def __init__(self, custom_param: str = "default", **kwargs):
+                super().__init__(**kwargs)
+                self.custom_param = custom_param
+
+            def generate_text(
+                self, request: GenerateTextRequest
+            ) -> GenerateTextResponse:
+                return GenerateTextResponse(
+                    version="v1",
+                    status_code=200,
+                    data={"generation": {"generatedText": "Custom response"}},
+                )
+
+        available_names = LLMGateway.available_config_names()
+        assert "CustomLLMGateway" in available_names
+
+        cls = LLMGateway.subclass_from_config_name("CustomLLMGateway")
+        assert cls == CustomLLMGateway
+
+        # Verify we can create via config
+        llm_config = LLMGatewayObjectConfig(
+            type_config_name="CustomLLMGateway",
+            options={"custom_param": "my_value"},
+        )
+        instance = llm_config.to_object()
+        assert isinstance(instance, CustomLLMGateway)
+        assert instance.custom_param == "my_value"
+
+        request = GenerateTextRequest(model_name="test-model", prompt="Hello")
+        response = instance.generate_text(request)
+        assert response.is_success is True
+        assert response.data["generation"]["generatedText"] == "Custom response"