From 458f9f4850a48bb6320767c0abdc0bec48a73a48 Mon Sep 17 00:00:00 2001 From: Rita Agarwala Date: Mon, 11 May 2026 20:37:38 +0530 Subject: [PATCH 1/3] Adding example for using prediction API --- .../chunking_with_prediction/config.json | 3 + .../chunking_with_prediction/entrypoint.py | 252 ++++++++++++++++++ .../chunking_with_prediction/tests/test.json | 40 +++ 3 files changed, 295 insertions(+) create mode 100644 src/datacustomcode/templates/function/example/chunking_with_prediction/config.json create mode 100644 src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py create mode 100644 src/datacustomcode/templates/function/example/chunking_with_prediction/tests/test.json diff --git a/src/datacustomcode/templates/function/example/chunking_with_prediction/config.json b/src/datacustomcode/templates/function/example/chunking_with_prediction/config.json new file mode 100644 index 0000000..2f911e9 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_prediction/config.json @@ -0,0 +1,3 @@ +{ + "entryPoint": "entrypoint.py" +} diff --git a/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py new file mode 100644 index 0000000..e474622 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, Salesforce, Inc. +# SPDX-License-Identifier: Apache-2 + +""" +Housing Sale Price Prediction with Einstein Regression + +This example uses Einstein regression model to predict housing sale prices +based on property features like Year_Built__c. + +Model: YH_Regression_Python_Predicted_SalePrice_CM_12l_ATC937af934 +Type: Regression +Input: Year_Built__c (numeric) +Output: Predicted_SalePrice +""" + +import logging +from typing import Any, Dict, Optional + +from datacustomcode.function import Runtime +from datacustomcode.function.feature_types.chunking import ( + ChunkType, + SearchIndexChunkingV1Output, + SearchIndexChunkingV1Request, + SearchIndexChunkingV1Response, +) + +from datacustomcode.einstein_predictions.types import ( + PredictionColumBuilder, + PredictionRequestBuilder, + PredictionType, +) + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +# Configuration +PREDICTION_MODEL_NAME = "YH_Regression_Python_Predicted_SalePrice_CM_12l_ATC937af934" + + +def predict_sale_price( + features: Dict[str, Any], + runtime: Runtime, +) -> Optional[float]: + """Predict housing sale price using Einstein regression model. + + Args: + features: Extracted housing features (numeric and string) + runtime: Runtime with prediction client + + Returns: + Predicted sale price or None if prediction fails + """ + try: + # Build prediction columns - handle both numeric and string values + prediction_columns = [] + + for column_name, value in features.items(): + if isinstance(value, str): + # String values (e.g., Garage_Qual__c) + column = ( + PredictionColumBuilder() + .set_column_name(column_name) + .set_string_values([value]) + .build() + ) + elif isinstance(value, (int, float)): + # Numeric values + column = ( + PredictionColumBuilder() + .set_column_name(column_name) + .set_double_values([float(value)]) + .build() + ) + else: + # Skip unsupported types + logger.warning(f"Skipping field {column_name} with unsupported type {type(value)}") + continue + + prediction_columns.append(column) + + # Build regression prediction request + prediction_request = ( + PredictionRequestBuilder() + .set_prediction_type(PredictionType.REGRESSION) + .set_model_api_name(PREDICTION_MODEL_NAME) + .set_prediction_columns(prediction_columns) + .build() + ) + + prediction_response = runtime.einstein_predictions.predict(prediction_request) + + if not prediction_response.is_success: + logger.error(f"Prediction failed: {prediction_response.data}") + return None + + # Parse regression response + results = prediction_response.data.get("results", []) + if not results: + logger.warning("No results in prediction response") + return None + + first_result = results[0] + prediction_type = first_result.get("type") + + if prediction_type != "RegressionPredictionSuccess": + logger.error(f"Unexpected prediction type: {prediction_type}") + logger.error(f"Full result: {first_result}") + return None + + prediction_data = first_result.get("prediction", {}) + predicted_value = prediction_data.get("value") + + if predicted_value is None: + logger.warning("No predicted value in response") + return None + + logger.info(f"Predicted sale price: ${predicted_value:,.2f}") + + # Log top contributors (which features influenced the price most) + top_contributors = prediction_data.get("topContributors", []) + if top_contributors: + logger.info(f"Top price contributors: {top_contributors}") + + return float(predicted_value) + + except Exception as e: + logger.error(f"Prediction failed with error: {e}", exc_info=True) + return None + + +def enrich_property_with_price( + source_dmo_fields: Dict[str, Any], + runtime: Runtime, +) -> Dict[str, str]: + """Enrich property data with predicted sale price. + + Args: + source_dmo_fields: Property features from source DMO + runtime: Runtime for predictions + + Returns: + Citations dictionary with predicted price + """ + citations = {} + + # Copy original fields to citations + if source_dmo_fields: + for key, value in source_dmo_fields.items(): + citations[key] = str(value) + + # Get price prediction - pass source_dmo_fields directly as features + predicted_price = predict_sale_price(source_dmo_fields, runtime) + + if predicted_price is not None: + citations["predicted_sale_price"] = f"${predicted_price:,.2f}" + citations["predicted_sale_price_raw"] = str(predicted_price) + citations["prediction_status"] = "success" + else: + citations["predicted_sale_price"] = "N/A" + citations["prediction_status"] = "failed" + + return citations + + +def function( + request: SearchIndexChunkingV1Request, runtime: Runtime +) -> SearchIndexChunkingV1Response: + """Housing price prediction using Einstein regression. + + Predicts sale prices for properties based on Year_Built__c feature + and adds predictions to citations for real estate data enrichment. + + Input format: + { + "input": [ + { + "text": "Beautiful 3BR house built in 1990", + "metadata": { + "source_dmo_fields": { + "Year_Built__c": 1990, + "address": "123 Main St", + "city": "San Francisco" + } + } + } + ] + } + + Output format: + { + "output": [ + { + "text": "Beautiful 3BR house built in 1990", + "seq_no": 1, + "citations": { + "Year_Built__c": "1990", + "address": "123 Main St", + "predicted_sale_price": "$350,000.00", + "predicted_sale_price_raw": "350000.0", + "prediction_status": "success" + } + } + ] + } + + Args: + request: Input properties to enrich + runtime: Runtime with prediction API access + + Returns: + Properties enriched with predicted sale prices + """ + logger.info( + f"Processing {len(request.input)} properties for price prediction" + ) + + enriched_properties = [] + seq_no = 1 + + for doc_idx, doc in enumerate(request.input): + text = doc.text + metadata = doc.metadata + + logger.info(f"Property {doc_idx + 1}: {text[:100]}...") + + # Get source_dmo_fields + source_dmo_fields = {} + if metadata and metadata.source_dmo_fields: + source_dmo_fields = dict(metadata.source_dmo_fields) + + # Enrich with price prediction - pass source_dmo_fields directly + citations = enrich_property_with_price(source_dmo_fields, runtime) + + # Create output + property_output = SearchIndexChunkingV1Output( + chunk_type=ChunkType.TEXT, + text=text.strip(), + seq_no=seq_no, + citations=citations, + ) + enriched_properties.append(property_output) + + logger.info( + f"Property {seq_no}: Predicted price = " + f"{citations.get('predicted_sale_price', 'N/A')}" + ) + seq_no += 1 + + logger.info(f"Total properties enriched: {len(enriched_properties)}") + + return SearchIndexChunkingV1Response(output=enriched_properties) \ No newline at end of file diff --git a/src/datacustomcode/templates/function/example/chunking_with_prediction/tests/test.json b/src/datacustomcode/templates/function/example/chunking_with_prediction/tests/test.json new file mode 100644 index 0000000..9968b39 --- /dev/null +++ b/src/datacustomcode/templates/function/example/chunking_with_prediction/tests/test.json @@ -0,0 +1,40 @@ +{ + "input": [ + { + "text": "Luxury 5-bedroom house with 5000 sq ft living area, built in 2023", + "metadata": { + "type": "text", + "source_dmo_fields": { + "First_Flr_SF__c": 2600, + "Full_Bath__c": 4, + "Garage_Cars__c": 4, + "Garage_Qual__c": "good", + "Gr_Liv_Area__c": 5000, + "Lot_Area__c": 3000, + "Overall_Cond__c": 10, + "Second_Flr_SF__c": 2400, + "Total_Bsmt_SF__c": 0, + "Year_Built__c": 2023 + } + } + }, + { + "text": "Spacious 4-bedroom family home with 3500 sq ft living space, built in 2020", + "metadata": { + "type": "text", + "source_dmo_fields": { + "First_Flr_SF__c": 2000, + "Full_Bath__c": 3, + "Garage_Cars__c": 3, + "Garage_Qual__c": "excellent", + "Gr_Liv_Area__c": 3500, + "Lot_Area__c": 8000, + "Overall_Cond__c": 9, + "Second_Flr_SF__c": 1500, + "Total_Bsmt_SF__c": 1000, + "Year_Built__c": 2020 + } + } + } + ] +} \ No newline at end of file From 0fb162195fe8872a908cf9f53ebca3cf0ca0239d Mon Sep 17 00:00:00 2001 From: Rita Agarwala Date: Mon, 11 May 2026 21:13:48 +0530 Subject: [PATCH 2/3] Fixing lint error --- .../chunking_with_prediction/entrypoint.py | 42 ++++++++----------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py index e474622..a8a5354 100644 --- a/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py +++ b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py @@ -15,8 +15,17 @@ """ import logging -from typing import Any, Dict, Optional +from typing import ( + Any, + Dict, + Optional, +) +from datacustomcode.einstein_predictions.types import ( + PredictionColumBuilder, + PredictionRequestBuilder, + PredictionType, +) from datacustomcode.function import Runtime from datacustomcode.function.feature_types.chunking import ( ChunkType, @@ -25,12 +34,6 @@ SearchIndexChunkingV1Response, ) -from datacustomcode.einstein_predictions.types import ( - PredictionColumBuilder, - PredictionRequestBuilder, - PredictionType, -) - logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) @@ -74,7 +77,9 @@ def predict_sale_price( ) else: # Skip unsupported types - logger.warning(f"Skipping field {column_name} with unsupported type {type(value)}") + logger.warning( + f"Skipping field {column_name} with unsupported type {type(value)}" + ) continue prediction_columns.append(column) @@ -95,6 +100,10 @@ def predict_sale_price( return None # Parse regression response + if prediction_response.data is None: + logger.warning("Prediction response data is None") + return None + results = prediction_response.data.get("results", []) if not results: logger.warning("No results in prediction response") @@ -179,8 +188,6 @@ def function( "metadata": { "source_dmo_fields": { "Year_Built__c": 1990, - "address": "123 Main St", - "city": "San Francisco" } } } @@ -195,7 +202,6 @@ def function( "seq_no": 1, "citations": { "Year_Built__c": "1990", - "address": "123 Main St", "predicted_sale_price": "$350,000.00", "predicted_sale_price_raw": "350000.0", "prediction_status": "success" @@ -211,9 +217,6 @@ def function( Returns: Properties enriched with predicted sale prices """ - logger.info( - f"Processing {len(request.input)} properties for price prediction" - ) enriched_properties = [] seq_no = 1 @@ -222,9 +225,6 @@ def function( text = doc.text metadata = doc.metadata - logger.info(f"Property {doc_idx + 1}: {text[:100]}...") - - # Get source_dmo_fields source_dmo_fields = {} if metadata and metadata.source_dmo_fields: source_dmo_fields = dict(metadata.source_dmo_fields) @@ -241,12 +241,6 @@ def function( ) enriched_properties.append(property_output) - logger.info( - f"Property {seq_no}: Predicted price = " - f"{citations.get('predicted_sale_price', 'N/A')}" - ) seq_no += 1 - logger.info(f"Total properties enriched: {len(enriched_properties)}") - - return SearchIndexChunkingV1Response(output=enriched_properties) \ No newline at end of file + return SearchIndexChunkingV1Response(output=enriched_properties) From 5c7116d200c2c69d4762df34c8a86d87251831d3 Mon Sep 17 00:00:00 2001 From: Rita Agarwala Date: Tue, 12 May 2026 15:28:49 +0530 Subject: [PATCH 3/3] Fixing key to be retrieved --- .../example/chunking_with_prediction/entrypoint.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py index a8a5354..df9b780 100644 --- a/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py +++ b/src/datacustomcode/templates/function/example/chunking_with_prediction/entrypoint.py @@ -93,8 +93,13 @@ def predict_sale_price( .build() ) + logger.info(f"Model: {PREDICTION_MODEL_NAME}") + prediction_response = runtime.einstein_predictions.predict(prediction_request) + logger.info(f"Response status code: {prediction_response.status_code}") + logger.info(f"Response data: {prediction_response.data}") + if not prediction_response.is_success: logger.error(f"Prediction failed: {prediction_response.data}") return None @@ -118,7 +123,7 @@ def predict_sale_price( return None prediction_data = first_result.get("prediction", {}) - predicted_value = prediction_data.get("value") + predicted_value = prediction_data.get("predictedValue") if predicted_value is None: logger.warning("No predicted value in response")