Refactor API response handling to include similarity attributes in output

app.py

@@ -180,11 +180,19 @@ def run_pred_pipeline(input: PredictionInput):
             payload = dict(run_id=task_run_id)
             response = requests.get(api_url, headers=headers, data=json.dumps(payload))
             output_json = json.loads(response.json()['notebook_output']['result'])
-            output_dict = process_api_response(output_json)
+            temp_predictions_dict, sample_sim_attr = process_api_response(output_json)
+            data_out = {
+                        "status" : "success",
+                        "data" : {
+                            "id": input.dict()['id'],
+                            "predictions": temp_predictions_dict,
+                            "similarity": sample_sim_attr
+                        }
+                    }
             #nb_output = output_json['prediction']
             break;
-        print()
-    return output_dict
+        
+    return data_out
 
 
 @app.get("/get_prediction_from_databricks")

helper.py

@@ -4,7 +4,8 @@ import json
 
 def process_api_response(json_response):
     # Use json.loads() to parse the JSON string
-    input_data = json.loads(json_response)
+    print("Type of json_response : ", type(json_response))
+    input_data = json_response
 
     # Extract predictions and similarity attributes
     predictions = input_data.get("predictions", {})
@@ -19,16 +20,18 @@ def process_api_response(json_response):
     sample_sim_attr = sim_attr_df.to_dict()
 
     # Construct final output dictionary
-    data_out = {
-        "status": "success",
-        "data": {
-            "id": input_data.get("id", "default_id"),  # Assuming an 'id' key exists in the input
-            "predictions": temp_predictions_dict,
-            "similarity": sample_sim_attr
-        }
-    }
+    return temp_predictions_dict, sample_sim_attr   
+
+    # data_out = {
+    #     "status": "success",
+    #     "data": {
+    #         "id": input_data.get("id", "default_id"),  # Assuming an 'id' key exists in the input
+    #         "predictions": temp_predictions_dict,
+    #         "similarity": sample_sim_attr
+    #     }
+    # }
 
-    return data_out
+    # return data_out
 
 def get_sample_similarity_attr():
     sample_sim = {