Merge pull request #9 from tc-mateus/fix/request-pipeline

fix: change prediction to pipeline service

Author: butzist

dhl_sdk/_constants.py

@@ -15,7 +15,7 @@
 DATASETS_URL = "api/db/v2/datasets"
 MODELS_URL = "api/db/v2/pipelineJobs"
 TEMPLATES_URL = "api/db/v2/pipelineJobTemplates"
-PREDICT_URL = "api/pipeline/v1/predictors"
+PREDICT_URL = "api/pipeline/v1/pipeline"
 
 
 PROCESS_UNIT_MAP = {

dhl_sdk/_input_processing.py

@@ -20,7 +20,9 @@
     _validate_spectra_format,
 )
 from dhl_sdk._utils import (
-    PredictionRequest,
+    Metadata,
+    PipelineStage,
+    PredictionPipelineRequest,
     PredictionRequestConfig,
     Predictions,
     PredictionResponse,
@@ -64,6 +66,8 @@ def variables(self) -> list[Variable]:
 
 
 class Model(Protocol):
+    id: str
+
     @property
     def dataset(self) -> Dataset:
         ...
@@ -313,9 +317,21 @@ def format(self) -> list[dict]:
             else:
                 instances[0].append(None)
 
-        json_data = PredictionRequest(
-            instances=instances, config=self.prediction_config
-        ).model_dump(by_alias=True, exclude_none=True, exclude=["sampleId", "steps"])
+        json_data = PredictionPipelineRequest(
+            instances=instances,
+            metadata=Metadata(
+                variables=[{"id": var.id} for var in input_variables],
+            ),
+            stages=[PipelineStage(config=self.prediction_config, id=self.model.id)],
+        ).model_dump(
+            by_alias=True,
+            exclude_none=True,
+            include={
+                "instances": {"__all__": {"__all__": {"timestamps", "values"}}},
+                "metadata": True,
+                "stages": True,
+            },
+        )
 
         return [json_data]
 
@@ -413,9 +429,23 @@ def format(self) -> list[dict]:
             else:
                 instances[0].append(None)
 
-        json_data = PredictionRequest(
-            instances=instances, config=self.prediction_config
-        ).model_dump(by_alias=True, exclude_none=True, exclude="sampleId")
+        json_data = PredictionPipelineRequest(
+            instances=instances,
+            metadata=Metadata(
+                variables=[{"id": var.id} for var in input_variables],
+            ),
+            stages=[PipelineStage(config=self.prediction_config, id=self.model.id)],
+        ).model_dump(
+            by_alias=True,
+            exclude_none=True,
+            include={
+                "instances": {
+                    "__all__": {"__all__": {"timestamps", "values", "steps"}}
+                },
+                "metadata": True,
+                "stages": True,
+            },
+        )
 
         return [json_data]
 

dhl_sdk/_spectra_utils.py

@@ -6,7 +6,13 @@
 
 import numpy as np
 
-from dhl_sdk._utils import Instance, PredictionRequest
+from dhl_sdk._utils import (
+    Instance,
+    Metadata,
+    PipelineStage,
+    PredictionPipelineRequest,
+    SpectraPredictionConfig,
+)
 from dhl_sdk.exceptions import InvalidSpectraException
 
 # Type Aliases
@@ -23,7 +29,7 @@ class Dataset(Protocol):
     def variables(self) -> list:
         ...
 
-    def get_spectrum_index(self) -> int:
+    def get_spectra_index(self) -> int:
         ...
 
 
@@ -106,7 +112,7 @@ def _convert_to_request(
     # get number of vars in model from config
     variables = model.dataset.variables
     n_vars = len(variables)
-    spectrum_index = model.dataset.get_spectrum_index()
+    spectrum_index = model.dataset.get_spectra_index()
 
     request_data = []
     # handle pagination
@@ -123,7 +129,22 @@ def _convert_to_request(
                         )
                         break
 
-        json_data = PredictionRequest(instances=[instance]).model_dump(by_alias=True)
+        json_data = PredictionPipelineRequest(
+            instances=[instance],
+            metadata=Metadata(
+                variables=[{"id": var.id} for var in model.dataset.variables],
+            ),
+            stages=[PipelineStage(config=SpectraPredictionConfig(), id=model.id)],
+        ).model_dump(
+            by_alias=True,
+            exclude_none=True,
+            include={
+                "instances": True,
+                "metadata": True,
+                "stages": True,
+            },
+        )
+
         request_data.append(json_data)
 
     return request_data

dhl_sdk/_utils.py

@@ -4,7 +4,7 @@
 import urllib.parse as urlparse
 from datetime import datetime
 from functools import reduce
-from typing import Optional, Union
+from typing import Literal, Optional, Union
 
 import numpy as np
 from pydantic import BaseModel, Field, model_validator
@@ -91,12 +91,50 @@ def new(
         )
 
 
+class SpectraPredictionConfig(BaseModel):
+    """Pydantic class representing Spectra Prediction Config"""
+
+    prediction_mode: Literal["classic", "onlySpectra"] = Field(
+        default="classic", alias="predictionMode"
+    )
+
+
+class OnlyId(BaseModel):
+    """Pydantic class representing a sctuc with only the id"""
+
+    id: str
+
+
 class PredictionRequest(BaseModel):
     """Pydantic class representing the expected Predict Request"""
 
     instances: list[list[Optional[Instance]]]
     metadata: Optional[dict] = None
-    config: Optional[PredictionRequestConfig] = None
+    config: Optional[Union[PredictionRequestConfig, SpectraPredictionConfig]] = None
+
+
+class Metadata(BaseModel):
+    """Pydantic class representing Metadata for Predict Request"""
+
+    experiments: list[Optional[OnlyId]] = [None]
+    variables: list[OnlyId]
+
+
+class PipelineStage(BaseModel):
+    """Pydantic class representing the Prediction Pipeline Stage"""
+
+    config: Union[PredictionRequestConfig, SpectraPredictionConfig]
+    id: str
+    merge_strategy: str = Field(default="merge", alias="mergeStrategy")
+    type: str = Field(default="predict")
+
+
+class PredictionPipelineRequest(BaseModel):
+    """Pydantic class representing the expected Predict Request"""
+
+    instances: list[list[Optional[Instance]]]
+    metadata: Metadata
+    stages: list[PipelineStage] = None
 
 
 class PredictionResponse(BaseModel):

dhl_sdk/entities.py

@@ -119,13 +119,20 @@ def __init__(self, **data):
         super().__init__(**data)
         self._client = data["client"]
 
-    def get_spectrum_index(self) -> int:
-        """Get the index of the spectrum variable"""
+    def get_spectra_index(self) -> int:
+        """Get the index of the spectra variable"""
         for index, variable in enumerate(self.variables):
             if variable.variant == "spectrum":
                 return index
         raise ValueError("No spectrum variable found in dataset")
 
+    def get_spectra_code(self) -> str:
+        """Get variable code of spectra variable"""
+        for variable in self.variables:
+            if variable.variant == "spectrum":
+                return variable.code
+        raise ValueError("No spectrum variable found in dataset")
+
     @staticmethod
     def requests(client: Client) -> CRUDClient["SpectraDataset"]:
         # pylint: disable=missing-function-docstring
@@ -163,12 +170,10 @@ def get_predictions(self, preprocessor: Preprocessor) -> dict:
                 "The provided inputs failed the validation step"
             )
 
-        predict_url = f"{PREDICT_URL}/{self.id}/predict"
-
         predictions = []
         for prediction_data in json_data:
             try:
-                response = self._client.post(predict_url, prediction_data)
+                response = self._client.post(PREDICT_URL, prediction_data)
                 response.raise_for_status()
 
                 # in case of an error in the response (not HTTP)
@@ -295,7 +300,13 @@ def predict(
             spectra=spectra, inputs=inputs, model=self
         )
 
-        return super().get_predictions(spectra_processing_strategy)
+        predictions = super().get_predictions(spectra_processing_strategy)
+
+        spectra_code = self.dataset.get_spectra_code()
+        if spectra_code in predictions:
+            predictions.pop(spectra_code)
+
+        return predictions
 
     @property
     def inputs(self) -> list[str]:
@@ -316,7 +327,7 @@ def spectra_size(self) -> int:
 
     def _get_spectra_size(self) -> int:
         """Get the size of the spectra from variable information in the API"""
-        spectrum = self.dataset.variables[self.dataset.get_spectrum_index()]
+        spectrum = self.dataset.variables[self.dataset.get_spectra_index()]
         return spectrum.size
 
     @staticmethod

examples.ipynb

@@ -807,7 +807,7 @@
     "\n",
     "prediction_config = PredictionConfig(model_confidence=50)\n",
     "\n",
-    "result = model_hist.predict(timestamps, steps, inputs, timestamps_unit=\"s\",  config = prediction_config)"
+    "result = model_hist.predict(timestamps, steps, inputs, timestamps_unit=\"s\", config = prediction_config)"
    ]
   },
   {

tests/test_utils.py

@@ -33,9 +33,11 @@ def setUp(self):
             "spectrum": {"xAxis": {"dimension": 4}},
         }
         self.model_no_inputs = Mock()
+        self.model_no_inputs.id = "model-id-1"
         self.model_no_inputs.inputs = []
         self.model_no_inputs.dataset.variables = [Variable(**spectrum_var)]
         self.model_with_inputs = Mock()
+        self.model_with_inputs.id = "model-id-2"
         self.model_with_inputs.dataset.variables = [
             Variable(**spectrum_var),
             Variable(id="id-123", code="var1", variant="numeric", name="variable 1"),
@@ -213,7 +215,7 @@ def test_validation_with_input(self):
     def test_convert_to_request(self):
         model = self.model_with_inputs
         model.spectra_size = 4
-        model.dataset.get_spectrum_index.return_value = 0
+        model.dataset.get_spectra_index.return_value = 0
 
         spectra = [[1.0, 2.0, 3.0, 3.0], [4.0, 5.0, 6.0, 6.0], [7.0, 8.0, 9.0, 9.0]]
         inputs = {"var1": [0, 1, 0], "var2": [1, 1, 1]}
@@ -231,7 +233,7 @@ def test_convert_to_request(self):
     def test_convert_request_noinput(self):
         model = self.model_no_inputs
         model.spectra_size = 4
-        model.dataset.get_spectrum_index.return_value = 0
+        model.dataset.get_spectra_index.return_value = 0
 
         spectra = [[1.0, 2.0, 3.0, 3.0], [4.0, 5.0, 6.0, 6.0], [7.0, 8.0, 9.0, 9.0]]
         processor = SpectraPreprocessor(spectra=spectra, model=model, inputs=None)
@@ -283,6 +285,7 @@ def setUp(self):
             "group": {"code": "X"},
         }
         self.model = Mock()
+        self.model.id = "model-id-1"
         self.model.dataset.variables = [
             Variable(**var1),
             Variable(**var2),