CREATE FUNCTION replace_value(data JSON, key text, value anyelement)
RETURNS JSON AS $$
    import json
    d = json.loads(data)
    d[key] = value
    return json.dumps(d)
$$ LANGUAGE plpython3u;

ERROR:  PL/Python functions cannot accept type anyelement

CREATE FUNCTION replace_value(data JSON, key text, value text)...
CREATE FUNCTION replace_value(data JSON, key text, value float)...
CREATE FUNCTION replace_value(data JSON, key text, value bigint)...