add example on how to use prediction api

Author: atulikumwenayo

src/datacustomcode/einstein_predictions/impl/default.py

@@ -31,5 +31,5 @@ def predict(self, request: PredictionRequest) -> PredictionResponse:
             version="v1",
             prediction_type=request.prediction_type,
             status_code=200,
-            data={"results": [{"prediction": {"predictedValue": "1"}}]},
+            data={"results": [{"prediction": {"predictedValue": 1.0}}]},
         )

src/datacustomcode/einstein_predictions/types.py

@@ -150,7 +150,7 @@ def set_prediction_type(
         self._prediction_type = prediction_type
         return self
 
-    def set_model(self, model_api_name: str) -> "PredictionRequestBuilder":
+    def set_model_api_name(self, model_api_name: str) -> "PredictionRequestBuilder":
         self._model_api_name = model_api_name
         return self
 

src/datacustomcode/templates/function/payload/entrypoint.py

@@ -2,6 +2,11 @@
 from typing import List
 from uuid import uuid4
 
+from datacustomcode.einstein_predictions.types import (
+    PredictionColumBuilder,
+    PredictionRequestBuilder,
+    PredictionType,
+)
 from datacustomcode.function import Runtime
 from datacustomcode.llm_gateway.types.generate_text_request_builder import (
     GenerateTextRequestBuilder,
@@ -38,6 +43,28 @@ def chunk_text(text: str, chunk_size: int = 1000) -> List[str]:
     return chunks
 
 
+def make_einstein_prediction(runtime: Runtime) -> None:
+    column = (
+        PredictionColumBuilder()
+        .set_column_name("col1")
+        .set_string_values(["str1", "str2"])
+        .build()
+    )
+    prediction_request = (
+        PredictionRequestBuilder()
+        .set_prediction_type(PredictionType.REGRESSION)
+        .set_model_api_name("regressionModel")
+        .set_prediction_columns([column])
+        .build()
+    )
+
+    prediction_response = runtime.einstein_predictions.predict(prediction_request)
+    print(
+        f"Einstein prediction results - success: {prediction_response.is_success} \
+           response data: {prediction_response.data}"
+    )
+
+
 def function(request: dict, runtime: Runtime) -> dict:
     logger.info("Inside Function")
     logger.info(request)
@@ -46,6 +73,8 @@ def function(request: dict, runtime: Runtime) -> dict:
     output_chunks = []
     current_seq_no = 1  # Start sequence number from 1
 
+    make_einstein_prediction(runtime)
+
     builder = GenerateTextRequestBuilder()
     llm_request = builder.set_prompt("Hello").set_model("modelName").build()
     llm_response = runtime.llm_gateway.generate_text(llm_request)

tests/test_einstein_predictions.py

@@ -142,7 +142,7 @@ def test_builder_creates_valid_request(self):
         request = (
             PredictionRequestBuilder()
             .set_prediction_type(PredictionType.CLUSTERING)
-            .set_model("cluster_model")
+            .set_model_api_name("cluster_model")
             .set_prediction_columns(
                 [PredictionColumn(column_name="test_col", double_values=[1.0])]
             )