import os

os.environ.setdefault("USE_TF", "0")

import laya


class EndpointHandler:
    def __init__(self, path=""):
        self.agent = laya.load(path)

    def __call__(self, data):
        payload = data.get("inputs", data)

        state = payload.get("state")
        questions = payload.get("questions")

        if state is None or questions is None:
            raise ValueError(
                "Request must contain inputs.state and inputs.questions"
            )

        return self.agent.predict(state, questions)