use production llm gateway

atulikumwenayo · atulikumwenayo · commit afa12009a223 · 2026-04-24T16:54:01.000-04:00
diff --git a/src/datacustomcode/einstein_platform_client.py b/src/datacustomcode/einstein_platform_client.py
@@ -0,0 +1,75 @@
+# 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 (
+    Any,
+    Dict,
+    Optional,
+)
+
+from abc import ABC
+from loguru import logger
+
+from datacustomcode.token_provider import (
+    CredentialsTokenProvider,
+    SFCLITokenProvider,
+    TokenProvider,
+)
+
+
+class EinsteinPlatformClient(ABC):
+    EINSTEIN_PLATFORM_URL = "https://api.salesforce.com/einstein/platform/v1"
+    EINSTEIN_WARNING_MESSAGE = (
+        "If your code uses Einstein APIs, make sure you have "
+        'configured the SDK to use "client_credentials" auth type. '
+        "Refer to https://developer.salesforce.com/docs/ai/agentforce/"
+        "guide/agent-api-get-started.html#create-a-salesforce-app "
+        "to create your external client app."
+    )
+
+    def __init__(
+        self,
+        credentials_profile: Optional[str] = None,
+        sf_cli_org: Optional[str] = None,
+    ):
+        if sf_cli_org:
+            self._token_provider: TokenProvider = SFCLITokenProvider(sf_cli_org)
+            logger.debug(f"Using SF CLI token provider for org: {sf_cli_org}")
+        else:
+            profile = credentials_profile or "default"
+            self._token_provider = CredentialsTokenProvider(profile)
+            logger.debug(f"Using credentials token provider with profile: {profile}")
+        self.token_response = None
+
+    def get_headers(self):
+        if self.token_response is None:
+            self.token_response = self._token_provider.get_token()
+
+        return {
+            "Authorization": f"Bearer {self.token_response.access_token}",
+            "Content-Type": "application/json",
+            "x-sfdc-app-context": "EinsteinGPT",
+            "x-client-feature-id": "ai-platform-models-connected-app",
+        }
+
+    def parse_response(self, response):
+        response_data: Dict[str, Any] = {}
+        if response.content:
+            try:
+                response_data = response.json()
+            except ValueError:
+                logger.warning("Failed to parse response as JSON")
+                response_data = {"raw_response": response.text}
+        return response_data
diff --git a/src/datacustomcode/einstein_predictions/__init__.py b/src/datacustomcode/einstein_predictions/__init__.py
@@ -17,6 +17,6 @@
 from datacustomcode.einstein_predictions.impl.default import DefaultEinsteinPredictions
 
 __all__ = [
-    "EinsteinPredictions",
     "DefaultEinsteinPredictions",
+    "EinsteinPredictions",
 ]
diff --git a/src/datacustomcode/einstein_predictions/impl/default.py b/src/datacustomcode/einstein_predictions/impl/default.py
@@ -24,23 +24,17 @@
 from loguru import logger
 import requests
 
+from datacustomcode.einstein_platform_client import EinsteinPlatformClient
 from datacustomcode.einstein_predictions.base import EinsteinPredictions
 from datacustomcode.einstein_predictions.types import (
     PredictionRequest,
     PredictionResponse,
     PredictionType,
 )
-from datacustomcode.token_provider import (
-    CredentialsTokenProvider,
-    SFCLITokenProvider,
-    TokenProvider,
-)
 
 
-class DefaultEinsteinPredictions(EinsteinPredictions):
+class DefaultEinsteinPredictions(EinsteinPlatformClient, EinsteinPredictions):
     CONFIG_NAME = "DefaultEinsteinPredictions"
-    EINSTEIN_PLATFORM_URL = "https://api.salesforce.com/einstein/platform/v1"
-
     ENDPOINT_MAP: ClassVar[dict[PredictionType, str]] = {
         PredictionType.REGRESSION: "regression",
         PredictionType.CLUSTERING: "clustering",
@@ -55,21 +49,12 @@ def __init__(
         sf_cli_org: Optional[str] = None,
         **kwargs,
     ):
-        super().__init__(**kwargs)
-
-        if sf_cli_org:
-            self._token_provider: TokenProvider = SFCLITokenProvider(sf_cli_org)
-            logger.debug(f"Using SF CLI token provider for org: {sf_cli_org}")
-        else:
-            profile = credentials_profile or "default"
-            self._token_provider = CredentialsTokenProvider(profile)
-            logger.debug(f"Using credentials token provider with profile: {profile}")
+        EinsteinPlatformClient.__init__(
+            self, credentials_profile=credentials_profile, sf_cli_org=sf_cli_org
+        )
+        EinsteinPredictions.__init__(self, **kwargs)
 
     def predict(self, request: PredictionRequest) -> PredictionResponse:
-        """Make a prediction request to the Einstein Predictions API"""
-        token_response = self._token_provider.get_token()
-        access_token = token_response.access_token
-
         endpoint = self.ENDPOINT_MAP.get(request.prediction_type)
         if not endpoint:
             raise RuntimeError(
@@ -102,42 +87,24 @@ def predict(self, request: PredictionRequest) -> PredictionResponse:
         if request.settings:
             payload["settings"] = request.settings
 
-        headers = {
-            "Authorization": f"Bearer {access_token}",
-            "Content-Type": "application/json",
-            "x-sfdc-app-context": "EinsteinGPT",
-            "x-client-feature-id": "ai-platform-models-connected-app",
-        }
-
         logger.debug(f"Making Einstein prediction request to: {api_url}")
         try:
-            response = requests.post(api_url, json=payload, headers=headers, timeout=60)
+            response = requests.post(
+                api_url, json=payload, headers=self.get_headers(), timeout=180
+            )
             if not response.ok and not response.text:
                 error_msg = (
                     f"Einstein Prediction request failed: {api_url} - "
                     f"{response.status_code} {response.reason}. "
-                    "If your code uses Einstein APIs, make sure you have "
-                    'configured the SDK to use "client_credentials" auth type. '
-                    "Refer to https://developer.salesforce.com/docs/ai/agentforce/"
-                    "guide/agent-api-get-started.html#create-a-salesforce-app "
-                    "to create your external client app."
+                    f"{self.EINSTEIN_WARNING_MESSAGE}"
                 )
                 logger.error(error_msg)
         except requests.exceptions.RequestException as e:
-            logger.error(f"Prediction API request failed: {api_url} {e}")
-            raise RuntimeError(f"Prediction API request failed: {e}") from e
-
-        response_data: Dict[str, Any] = {}
-        if response.content:
-            try:
-                response_data = response.json()
-            except ValueError:
-                logger.warning("Failed to parse response as JSON")
-                response_data = {"raw_response": response.text}
+            logger.error(f"Einstein Prediction request failed: {api_url} {e}")
+            raise RuntimeError(f"Einstein Prediction request failed: {e}") from e
 
         return PredictionResponse(
-            version="v1",
             prediction_type=request.prediction_type,
             status_code=response.status_code,
-            data=response_data,
+            data=self.parse_response(response),
         )
diff --git a/src/datacustomcode/llm_gateway/default.py b/src/datacustomcode/llm_gateway/default.py
@@ -13,23 +13,69 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import (
+    Any,
+    Dict,
+    Optional,
+)
+
+from loguru import logger
+import requests
+
+from datacustomcode.einstein_platform_client import EinsteinPlatformClient
+
+
 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.types.generate_text_response_builder import (
-    GenerateTextResponseBuilder,
-)
 
 
-class DefaultLLMGateway(LLMGateway):
+class DefaultLLMGateway(EinsteinPlatformClient, LLMGateway):
     CONFIG_NAME = "DefaultLLMGateway"
 
+    def __init__(
+        self,
+        credentials_profile: Optional[str] = None,
+        sf_cli_org: Optional[str] = None,
+        **kwargs,
+    ):
+        EinsteinPlatformClient.__init__(
+            self, credentials_profile=credentials_profile, sf_cli_org=sf_cli_org
+        )
+        LLMGateway.__init__(self, **kwargs)
+
     def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
+        api_url = (
+            f"{self.EINSTEIN_PLATFORM_URL}/models/"
+            f"{request.model_name}/generations"
+        )
 
-        response_data = {
-            "version": "v1",
-            "status_code": 200,
-            "data": {"generation": {"generatedText": "Hello World"}},
+        payload: Dict[str, Any] = {
+            "prompt": request.prompt
         }
 
-        return GenerateTextResponseBuilder.build(response_data)
+        if request.localization:
+            payload["localization"] = request.localization
+        if request.tags:
+            payload["tags"] = request.tags
+
+        logger.debug(f"Making Generate text request: {api_url}")
+        try:
+            response = requests.post(
+                api_url, json=payload, headers=self.get_headers(), timeout=180
+            )
+            if not response.ok and not response.text:
+                error_msg = (
+                    f"Generate text request failed: {api_url} - "
+                    f"{response.status_code} {response.reason}. "
+                    f"{self.EINSTEIN_WARNING_MESSAGE}"
+                )
+                logger.error(error_msg)
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Generate text request failed: {api_url} {e}")
+            raise RuntimeError(f"Generate text request failed: {e}") from e
+
+        return GenerateTextResponse(
+            status_code=response.status_code,
+            data=self.parse_response(response)
+        )
diff --git a/tests/test_llm_gateway.py b/tests/test_llm_gateway.py
@@ -3,8 +3,6 @@
 from pydantic import ValidationError
 import pytest
 
-from datacustomcode.llm_gateway.base import LLMGateway
-from datacustomcode.llm_gateway.default import DefaultLLMGateway
 from datacustomcode.llm_gateway.types.generate_text_request import GenerateTextRequest
 from datacustomcode.llm_gateway.types.generate_text_request_builder import (
     GenerateTextRequestBuilder,
@@ -210,28 +208,3 @@ def test_builder_with_minimal_dict(self):
         response = GenerateTextResponseBuilder.build(response_dict)
         assert response.status_code == 200
         assert response.version == "v1"  # Default value
-
-
-class TestDefaultLLMGateway:
-    """Test DefaultLLMGateway implementation."""
-
-    def test_default_gateway_is_llm_gateway(self):
-        """Test DefaultLLMGateway inherits from LLMGateway."""
-        gateway = DefaultLLMGateway()
-        assert isinstance(gateway, LLMGateway)
-
-    def test_generate_text_returns_response(self):
-        """Test generate_text returns GenerateTextResponse."""
-        gateway = DefaultLLMGateway()
-        request = GenerateTextRequest(model_name="gpt-4", prompt="Hello")
-        response = gateway.generate_text(request)
-        assert isinstance(response, GenerateTextResponse)
-
-    def test_generate_text_success_response(self):
-        """Test generate_text returns successful response."""
-        gateway = DefaultLLMGateway()
-        request = GenerateTextRequest(model_name="gpt-4", prompt="Hello")
-        response = gateway.generate_text(request)
-        assert response.is_success is True
-        assert response.status_code == 200
-        assert len(response.text) > 0

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,6 @@`
`17`	`17`	`from datacustomcode.einstein_predictions.impl.default import DefaultEinsteinPredictions`
`18`	`18`
`19`	`19`	`__all__ = [`
`20`		`- "EinsteinPredictions",`
`21`	`20`	`"DefaultEinsteinPredictions",`
	`21`	`+ "EinsteinPredictions",`
`22`	`22`	`]`