import json
import boto3
import os

sagemaker = boto3.client('sagemaker-runtime')
s3 = boto3.client('s3')

CONFIGFILE_S3_KEY = 'mlops-staging/config/ml-pipeline-config.json'
config_file= "/tmp/ml-pipeline-config.json"
bucket = os.environ['StagingBucket']

s3.download_file(bucket, CONFIGFILE_S3_KEY, config_file)
f = open(config_file, 'r')
config = json.load(f)
ep_name = config['InferenceEPName']

def handler(event, context):
    input = bytes(event['features'], 'utf-8')
    response = sagemaker.invoke_endpoint(
                        EndpointName= ep_name,
                        Body= input,
                        ContentType='text/csv',
                        Accept='text/csv')
              
    prediction = response['Body'].read().decode("utf-8")

    return {"prediction": prediction}