From 1f2aece91a752959f65298da3ebbb65bb239c793 Mon Sep 17 00:00:00 2001
From: atulikumwenayo <atulikumwenayo@salesforce.com>
Date: Mon, 4 May 2026 16:40:47 -0400
Subject: [PATCH] make llm gateway impl configurable

---
 src/datacustomcode/function/runtime.py     | 17 ++++++++++++-----
 src/datacustomcode/llm_gateway/__init__.py |  8 ++++++++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/datacustomcode/function/runtime.py b/src/datacustomcode/function/runtime.py
index ffda532..df9cb1e 100644
--- a/src/datacustomcode/function/runtime.py
+++ b/src/datacustomcode/function/runtime.py
@@ -21,7 +21,8 @@
 from datacustomcode.einstein_predictions_config import einstein_predictions_config
 from datacustomcode.file.path.default import DefaultFindFilePath
 from datacustomcode.function.base import BaseRuntime
-from datacustomcode.llm_gateway.default import DefaultLLMGateway
+from datacustomcode.llm_gateway.base import LLMGateway
+from datacustomcode.llm_gateway_config import llm_gateway_config
 
 
 class Runtime(BaseRuntime):
@@ -46,7 +47,7 @@ def __new__(cls):
                 raise RuntimeError(
                     "Runtime can only be instantiated once by the SDK.\n\n"
                     "Do not instantiate it yourself. Accept it as a parameter:\n\n"
-                    "  from datacustomcode.runtime.function.RunTime import Function\n"
+                    "  from datacustomcode.function.runtime import Runtime\n"
                     "  \n"
                     "  def function(request: dict, runtime: Runtime) -> dict:\n"
                     "      response = {...}\n"
@@ -65,13 +66,19 @@ def __init__(self) -> None:
         super().__init__()
 
         # Initialize resources
-        self._llm_gateway = DefaultLLMGateway()
+        self._llm_gateway: Optional[LLMGateway] = None
         self._file = DefaultFindFilePath()
         self._einstein_predictions: Optional[EinsteinPredictions] = None
 
     @property
-    def llm_gateway(self) -> DefaultLLMGateway:
-        """Access LLM operations."""
+    def llm_gateway(self) -> LLMGateway:
+        if self._llm_gateway is None:
+            if llm_gateway_config.llm_gateway_config is None:
+                raise RuntimeError(
+                    "LLM Gateway is not configured. "
+                    "Add 'llm_gateway_config' section to config.yaml"
+                )
+            self._llm_gateway = llm_gateway_config.llm_gateway_config.to_object()
         return self._llm_gateway
 
     @property
diff --git a/src/datacustomcode/llm_gateway/__init__.py b/src/datacustomcode/llm_gateway/__init__.py
index 93988ff..ea8b4af 100644
--- a/src/datacustomcode/llm_gateway/__init__.py
+++ b/src/datacustomcode/llm_gateway/__init__.py
@@ -12,3 +12,11 @@
 # 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.default import DefaultLLMGateway
+
+__all__ = [
+    "DefaultLLMGateway",
+    "LLMGateway",
+]