Merge remote-tracking branch 'origin/predict' into run_with_json

Author: ritaagarwala-sf

src/datacustomcode/einstein_platform_client.py

@@ -20,6 +20,7 @@
 )
 
 from loguru import logger
+import requests
 
 from datacustomcode.token_provider import (
     CredentialsTokenProvider,
@@ -49,7 +50,7 @@ def __init__(
         self.token_response = None
         super().__init__(**kwargs)
 
-    def get_headers(self):
+    def _get_headers(self):
         if self.token_response is None:
             self.token_response = self._token_provider.get_token()
 
@@ -60,6 +61,23 @@ def get_headers(self):
             "x-client-feature-id": "ai-platform-models-connected-app",
         }
 
+    def make_post_request(self, url, payload):
+        try:
+            response = requests.post(
+                url, json=payload, headers=self._get_headers(), timeout=180
+            )
+            if not response.ok:
+                error_msg = (
+                    f"Request to {url} failed. "
+                    f"Reason: {response.status_code} {response.reason} - "
+                    f"Response body: {response.text}"
+                )
+                logger.error(error_msg)
+            return response
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Request to {url} failed: {e}")
+            raise RuntimeError(f"Request to {url} failed {e}") from e
+
     def parse_response(self, response):
         response_data: Dict[str, Any] = {}
         if response.content:

src/datacustomcode/einstein_predictions/impl/default.py

@@ -20,9 +20,6 @@
     List,
 )
 
-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 (
@@ -74,21 +71,7 @@ def predict(self, request: PredictionRequest) -> PredictionResponse:
         if request.settings:
             payload["settings"] = request.settings
 
-        logger.debug(f"Making Einstein prediction request to: {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"Einstein Prediction request failed: {api_url} - "
-                    f"{response.status_code} {response.reason}"
-                )
-                logger.error(error_msg)
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Einstein Prediction request failed: {api_url} {e}")
-            raise RuntimeError(f"Einstein Prediction request failed: {e}") from e
-
+        response = self.make_post_request(api_url, payload)
         return PredictionResponse(
             prediction_type=request.prediction_type,
             status_code=response.status_code,

src/datacustomcode/llm_gateway/default.py

@@ -15,9 +15,6 @@
 
 from typing import Any, Dict
 
-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
@@ -42,21 +39,7 @@ def generate_text(self, request: GenerateTextRequest) -> GenerateTextResponse:
         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}"
-                )
-                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
-
+        response = self.make_post_request(api_url, payload)
         response_dict = {
             "status_code": response.status_code,
             "data": self.parse_response(response),