feat: rpt client updated to 1.6

packages/gen/gen_ai_hub/proxy/native/sap/models.py

@@ -1,23 +1,53 @@
-from typing import Optional, Literal, Union, Any
-from pydantic import BaseModel, RootModel, model_validator
+from typing import Annotated, Optional, Literal, Union, Any
+from pydantic import BaseModel, Field, RootModel, model_validator
 
 
 class TargetColumn(BaseModel):
     """Represents a target column in data.
 
     :param name: Name of the target column.
     :type name: str
-    :param prediction_placeholder: Placeholder string denoting where predictions will be inserted.
-        Defaults to ``"[PREDICT]"``.
-    :type prediction_placeholder: str
+    :param prediction_placeholder: The prediction placeholder in any column for which to predict a value. The model will predict a value for all table cells containing this value.
+    :type prediction_placeholder: Optional[Union[str, int, float]]
     :param task_type: Task type of the target column.
         One of ``"classification"`` or ``"regression"``. Defaults to ``None``.
     :type task_type: Optional[Literal["classification", "regression"]]
+    :param top_k: How many predictions to output for this classification column.If not provided, only a single prediction is returned. Only relevant for classification.
+    :type top_k: Optional[int]
     """
 
     name: str
-    prediction_placeholder: str = "[PREDICT]"
+    prediction_placeholder: Optional[Union[str, int, float]]
     task_type: Optional[Literal["classification", "regression"]] = None
+    top_k: Optional[int] = None
+
+
+class ExplanationConfig(BaseModel):
+    """
+    Configuration for explainability outputs.
+
+    :param top_column_scores: For how many columns to output column scores (optional, default is 0). 0 by default (no explainability). Max value is 20.
+    :type top_column_scores: int
+    :param top_relevant_context_rows: For how many context rows to return indices per query row (optional, default is 0). 0 by default (no explainability). Max value is 20.
+    :type top_relevant_context_rows: int
+    """
+
+    top_column_scores: int = 0
+    top_relevant_context_rows: int = 0
+
+
+class ExplanationResult(BaseModel):
+    """
+    Explanation data for predictions.
+
+    :param top_column_scores: Column scores per query row extracted from the model (higher means more weight was put on this column).
+    :type top_column_scores: Optional[list[dict[str, Union[int, float]]]]
+    :param top_relevant_context_rows: 2D array where each subarray contains indices of most relevant context rows for that query row. The first dimension indexes query rows, the second dimension indexes all rows as a sequential integer index.
+    :type top_relevant_context_rows: Optional[list[list[int]]]
+    """
+
+    top_column_scores: Optional[list[dict[str, Union[int, float]]]]
+    top_relevant_context_rows: Optional[list[list[int]]]
 
 
 class PredictionConfig(BaseModel):
@@ -26,19 +56,32 @@ class PredictionConfig(BaseModel):
 
     :param target_columns: List of target columns to predict.
     :type target_columns: list[TargetColumn]
+    :param explanations: Optional configuration for explainability outputs (column scores and relevant context rows).
+    :type explanations: Optional[ExplanationConfig]
+    :param context_mode: Context mode for predictions. Set it to \"default\" for the best balance between accuracy and latency/cost. Set it to \"deep\" for higher accuracy with >8k context rows at increased latency and cost (only for \"sap-rpt-1.6-large\").
+    :type context_mode: Literal['default', 'deep']
     """
 
     target_columns: list[TargetColumn]
+    explanations: Optional[ExplanationConfig] = None
+    context_mode: Literal['default', 'deep'] = 'default'
 
 
 class DataType(BaseModel):
     """Schema definition for a column.
 
     :param dtype: The data type of the column.
-    :type dtype: Literal["string", "numeric", "date"]
+    :type dtype: Literal['string', 'numeric', 'date',
+        'time', 'boolean', 'largestring', 'uuid', 'integer', 'int16',
+        'int32', 'int64', 'uint8', 'decimal',
+        'double', 'datetime', 'timestamp']
     """
 
-    dtype: Literal["string", "numeric", "date"]
+    dtype: Literal['string', 'numeric', 'date',
+                   'time', 'boolean', 'largestring',
+                   'uuid', 'integer', 'int16',
+                   'int32', 'int64', 'uint8', 'decimal',
+                   'double', 'datetime', 'timestamp']
 
 
 class RPTRequest(BaseModel):
@@ -109,12 +152,15 @@ class ResponseMetadata(BaseModel):
     :type num_predictions: int
     :param num_query_rows: Number of query rows for which a prediction was made.
     :type num_query_rows: int
+    :param context_mode: The context mode used for this prediction request.
+    :type context_mode: Literal['default', 'deep']
     """
 
     num_rows: int
     num_columns: int
     num_predictions: int
     num_query_rows: int
+    context_mode: Literal['default', 'deep'] = 'default'
 
 
 class ResponseStatus(BaseModel):
@@ -135,12 +181,15 @@ class PredictionItem(BaseModel):
 
     :param prediction: The predicted value.
     :type prediction: Union[str, float]
-    :param confidence: Confidence score for classification tasks. Defaults to ``None``.
+    :param confidence: Confidence score for classification tasks, in the range ``[0.0, 1.0]``. Defaults to ``None``.
     :type confidence: Optional[float]
+    :param confidence_interval: Lower and upper bounds of the confidence interval. Null for classification predictions.
+    :type confidence_interval: Optional[tuple[float, float]]
     """
 
     prediction: Union[str, float]
-    confidence: Optional[float] = None
+    confidence: Optional[float] = Field(default=None, ge=0.0, le=1.0)
+    confidence_interval: Optional[Annotated[list[float], Field(min_length=2, max_length=2)]] = None
 
 
 class Prediction(RootModel[dict[str, Union[list[PredictionItem], Any]]]):
@@ -166,13 +215,16 @@ class RPTResponse(BaseModel):
     :type status: ResponseStatus
     :param predictions: Prediction data returned by the service.
     :type predictions: list[Prediction]
+    :param explanations: Explanation data containing context row and column scores.
+    :type explanations: Optional[ExplanationResult]
     :param metadata: Metadata about the request/response.
     :type metadata: ResponseMetadata
     """
 
     id: str
     status: ResponseStatus
     predictions: list[Prediction]
+    explanations: Optional[ExplanationResult] = None
     metadata: ResponseMetadata
 
 

packages/gen/integration_tests/constants.py

@@ -23,4 +23,4 @@
 PERPLEXITY_TEST_MODEL = "sonar"
 PERPLEXITY_SONAR_DEEP_RESEARCH_TEST_MODEL = "sonar-deep-research"
 COHERE_COMMAND_A_TEST_MODEL = "cohere--command-a-reasoning"
-SAP_RPT_1_SMALL_TEST_MODEL = "sap-rpt-1-small"
+SAP_RPT_1_6_SMALL_TEST_MODEL = "sap-rpt-1.6"

packages/gen/integration_tests/native_clients/test_sap_rpt.py

@@ -1,6 +1,6 @@
 import unittest
 
-from integration_tests.constants import SAP_RPT_1_SMALL_TEST_MODEL
+from integration_tests.constants import SAP_RPT_1_6_SMALL_TEST_MODEL
 from integration_tests.setup_aicore import TestCaseStandardSetupMixin
 from gen_ai_hub.proxy.native.sap.client import RPTClient
 from gen_ai_hub.proxy.native.sap.models import RPTRequest, RPTResponse, PredictionConfig, TargetColumn
@@ -138,15 +138,15 @@ def setUp(self) -> None:
         self.client = RPTClient(proxy_client=self.proxy_client)
 
     def test_client_find_url_by_model_name(self):
-        url = self.client._get_url(model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        url = self.client._get_url(model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsNotNone(url)
 
     def test_client_find_url_by_model_name_and_version(self):
-        url = self.client._get_url(model_name=SAP_RPT_1_SMALL_TEST_MODEL, model_version="latest")
+        url = self.client._get_url(model_name=SAP_RPT_1_6_SMALL_TEST_MODEL, model_version="latest")
         self.assertIsNotNone(url)
 
     def test_client_find_url_by_config_name(self):
-        url = self.client._get_url(config_name="sap-rpt-1-small-latest")
+        url = self.client._get_url(config_name="sap-rpt-1.6-latest")
         self.assertIsNotNone(url)
 
     def test_client_find_url_with_invalid_model_name(self):
@@ -155,7 +155,7 @@ def test_client_find_url_with_invalid_model_name(self):
 
     def test_predict_by_row(self):
         body = RPTRequest(**request_by_row_dict)
-        response = self.client.predict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
         self.assertEqual(response.predictions[0]["ID"], "35")
@@ -165,7 +165,7 @@ def test_predict_by_row(self):
 
     def test_predict_by_columns(self):
         body = RPTRequest(**request_by_columns_dict)
-        response = self.client.predict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
         self.assertEqual(response.metadata.num_columns, 5)
@@ -175,7 +175,7 @@ def test_predict_by_columns(self):
 
     def test_predict_with_api_url(self):
         body = RPTRequest(**request_by_columns_dict)
-        deployment_url = self.client._get_url(model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        deployment_url = self.client._get_url(model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         response = self.client.predict(body=body, deployment_url=deployment_url)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
@@ -187,11 +187,11 @@ def test_regression_prediction(self):
         body = RPTRequest(
             prediction_config=PredictionConfig(
                 target_columns=[
-                    TargetColumn(name="DISCOUNT_RATE", task_type="regression")
+                    TargetColumn(name="DISCOUNT_RATE", task_type="regression", prediction_placeholder="[PREDICT]")
                 ]),
             rows=rows_regression
         )
-        response = self.client.predict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
         self.assertEqual(response.metadata.num_predictions, 2)
@@ -201,13 +201,81 @@ def test_timeout_error(self):
         body = RPTRequest(
             prediction_config=PredictionConfig(
                 target_columns=[
-                    TargetColumn(name="DISCOUNT_RATE", task_type="regression")
+                    TargetColumn(name="DISCOUNT_RATE", task_type="regression", prediction_placeholder="[PREDICT]")
                 ]),
             rows=rows_regression
         )
 
         with self.assertRaises(Exception):
-            self.client.predict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL, timeout=0.001)
+            self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL, timeout=0.001)
+
+
+
+    def test_predict_classification_with_top_k(self):
+        body = RPTRequest(**{
+            **request_by_columns_dict,
+            'prediction_config': {
+                'target_columns': [
+                    {
+                        'name': 'COSTCENTER',
+                        'prediction_placeholder': '[PREDICT]',
+                        'task_type': 'classification',
+                        'top_k': 2
+                    }
+                ]
+            }
+        })
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertEqual(response.status.code, 0)
+        costcenter_predictions = response.predictions[0]['COSTCENTER']
+        self.assertGreaterEqual(len(costcenter_predictions), 1)
+
+    def test_predict_response_includes_context_mode(self):
+        body = RPTRequest(**request_by_columns_dict)
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertIn(response.metadata.context_mode, ['default', 'deep'])
+
+    def test_regression_prediction_includes_confidence_interval(self):
+        body = RPTRequest(
+            prediction_config=PredictionConfig(
+                target_columns=[
+                    TargetColumn(name='DISCOUNT_RATE', task_type='regression', prediction_placeholder='[PREDICT]')
+                ]),
+            rows=rows_regression
+        )
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertEqual(response.status.code, 0)
+        discount_predictions = response.predictions[0]['DISCOUNT_RATE']
+        self.assertIsNotNone(discount_predictions[0].confidence_interval)
+        self.assertEqual(len(discount_predictions[0].confidence_interval), 2)
+        self.assertIsNone(discount_predictions[0].confidence)
+
+    def test_predict_with_explanations(self):
+        body = RPTRequest(**{
+            **request_by_columns_dict,
+            'prediction_config': {
+                'target_columns': [
+                    {
+                        'name': 'COSTCENTER',
+                        'prediction_placeholder': '[PREDICT]',
+                        'task_type': 'classification'
+                    }
+                ],
+                'explanations': {
+                    'top_column_scores': 3,
+                    'top_relevant_context_rows': 2
+                }
+            }
+        })
+        response = self.client.predict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertEqual(response.status.code, 0)
+        self.assertIsNotNone(response.explanations)
+        self.assertIsNotNone(response.explanations.top_column_scores)
+        self.assertIsNotNone(response.explanations.top_relevant_context_rows)
 
 
 class AsyncRPTClientTests(TestCaseStandardSetupMixin, unittest.IsolatedAsyncioTestCase):
@@ -217,7 +285,7 @@ def setUp(self) -> None:
 
     async def test_apredict_by_row(self):
         body = RPTRequest(**request_by_row_dict)
-        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
         self.assertEqual(response.predictions[0]["ID"], "35")
@@ -227,10 +295,33 @@ async def test_apredict_by_row(self):
 
     async def test_apredict_by_columns(self):
         body = RPTRequest(**request_by_columns_dict)
-        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_SMALL_TEST_MODEL)
+        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
         self.assertIsInstance(response, RPTResponse)
         self.assertEqual(response.status.code, 0)
         self.assertEqual(response.metadata.num_columns, 5)
         self.assertEqual(response.metadata.num_predictions, 1)
         self.assertIn("COSTCENTER", response.predictions[0].model_dump())
         self.assertNotIn("ID", response.predictions[0].model_dump())
+
+    async def test_apredict_response_includes_context_mode(self):
+        body = RPTRequest(**request_by_columns_dict)
+        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertIn(response.metadata.context_mode, ['default', 'deep'])
+
+    async def test_apredict_regression_includes_confidence_interval(self):
+        body = RPTRequest(
+            prediction_config=PredictionConfig(
+                target_columns=[
+                    TargetColumn(name='DISCOUNT_RATE', task_type='regression', prediction_placeholder='[PREDICT]')
+                ]),
+            rows=rows_regression
+        )
+        response = await self.client.apredict(body=body, model_name=SAP_RPT_1_6_SMALL_TEST_MODEL)
+        self.assertIsInstance(response, RPTResponse)
+        self.assertEqual(response.status.code, 0)
+        discount_predictions = response.predictions[0]['DISCOUNT_RATE']
+        self.assertIsNotNone(discount_predictions[0].confidence_interval)
+        self.assertEqual(len(discount_predictions[0].confidence_interval), 2)
+        self.assertIsNone(discount_predictions[0].confidence)
+

packages/gen/integration_tests/setup_aicore.py

@@ -175,7 +175,7 @@ def get_standard_models():
         ("sonar", "latest"),
         ("sonar-deep-research", "latest"),
         ("cohere--command-a-reasoning", "latest"),
-        ("sap-rpt-1-small","latest"),
+        ("sap-rpt-1.6","latest"),
     ]