handler.py

from typing import Dict, List, Any
import numpy as np
from concrete.ml.deployment import FHEModelServer


def from_json(python_object):
    if "__class__" in python_object:
        return bytes(python_object["__value__"])


def to_json(python_object):
    if isinstance(python_object, bytes):
        return {"__class__": "bytes", "__value__": list(python_object)}
    raise TypeError(repr(python_object) + " is not JSON serializable")


class EndpointHandler:
    def __init__(self, path=""):

        # For server
        self.fhemodel_server = FHEModelServer(path + "/compiled_model")

        # Simulate a database of keys
        self.key_database = {}

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str`)
              date (:obj: `str`)
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """

        # Get method
        method = data.pop("method", data)

        if method is "save_key":

            # Get keys
            evaluation_keys = from_json(data.pop("evaluation_keys", data))

            uid = np.random.randint()

            while uid is in self.key_database.keys():
                uid = np.random.randint()

            self.key_database[uid] = evaluation_keys

            return uid

        assert method == "inference":

        uid = data.pop("uid", data)

        # Get inputs
        encrypted_inputs = from_json(data.pop("encrypted_inputs", data))

        # Find key in the database
        evaluation_keys = self.key_database[uid]

        # Run CML prediction
        encrypted_prediction = self.fhemodel_server.run(encrypted_inputs, evaluation_keys)

        return to_json(encrypted_prediction)