"""
flask_mab
~~~~~~~~~
This module implements all the extension logic for
Multi-armed bandit experiments on Flask apps.
:copyright: (c) 2013 by `Mark Grey <http://www.deacondesperado.com>`_.
:license: BSD, see LICENSE for more details.
"""
from flask import current_app, g, request
import json
import flask_mab.storage
import types
from bunch import Bunch
from flask import _request_ctx_stack
from functools import wraps
try:
from flask import _app_ctx_stack as stack
except ImportError:
from flask import _request_ctx_stack as stack
__version__ = "1.1.1"
[docs]def choose_arm(bandit):
"""Route decorator for registering an impression conveinently
:param bandit: The bandit/experiment to register for
:type bandit: string
"""
def decorator(func):
#runs @ service init
if not hasattr(func, 'bandits'):
func.bandits = []
func.bandits.append(bandit)
@wraps(func)
def wrapper(*args, **kwargs):
#runs at endpoint hit
add_args = []
for bandit in func.bandits:
#Fetch from request first here?
arm_id, arm_value = suggest_arm_for(bandit)
add_args.append((bandit, arm_value))
kwargs.update(add_args)
return func(*args, **kwargs)
return wrapper
return decorator
[docs]def reward_endpt(bandit, reward_val=1):
"""Route decorator for rewards.
:param bandit: The bandit/experiment to register rewards
for using arm found in cookie.
:type bandit: string
:param reward: The amount of reward this endpoint should
give its winning arm
:type reward: float
"""
def decorator(func):
if not hasattr(func, 'rewards'):
func.rewards = []
func.rewards.append((bandit, reward_val))
@wraps(func)
def wrapper(*args, **kwargs):
for bandit, reward_amt in func.rewards:
if bandit in request.bandits.keys():
request.bandits_reward.add((bandit, request.bandits[bandit], reward_amt))
return func(*args, **kwargs)
return wrapper
return decorator
[docs]class BanditMiddleware(object):
"""The main flask extension.
Sets up all the necessary tracking for the bandit experiments
"""
def __init__(self, app=None):
"""Attach MAB logic to a Flask application
:param app: An optional Flask application
"""
if app is not None:
self.init_app(app)
def _register_storage(self, app):
storage_engine = getattr(
flask_mab.storage,
app.config.get('MAB_STORAGE_ENGINE', 'BanditStorage'))
storage_opts = app.config.get('MAB_STORAGE_OPTS', tuple())
storage_backend = storage_engine(*storage_opts)
app.extensions['mab'].bandit_storage = storage_backend
[docs] def init_app(self, app):
"""Attach Multi Armed Bandit to application and configure
:param app: A flask application instance
"""
app.config.setdefault('MAB_COOKIE_NAME', 'MAB')
app.config.setdefault('MAB_COOKIE_PATH', '/')
app.config.setdefault('MAB_COOKIE_TTL', None)
app.config.setdefault('MAB_DEBUG_HEADERS', True)
if not hasattr(app, 'extensions'):
app.extensions = {}
app.extensions['mab'] = Bunch()
self._register_storage(app)
if hasattr(app, 'teardown_appcontext'):
app.teardown_appcontext(self.teardown)
else:
app.teardown_request(self.teardown)
app.extensions['mab'].bandits = {}
app.extensions['mab'].reward_endpts = []
app.extensions['mab'].pull_endpts = []
app.extensions['mab'].debug_headers = app.config.get('MAB_DEBUG_HEADERS', True)
app.extensions['mab'].cookie_name = app.config.get('MAB_COOKIE_NAME', "MAB")
self._init_detection(app)
[docs] def teardown(self, *args, **kwargs):
"""Stub for old flask versions
"""
pass
def _init_detection(self, app):
"""
Attaches all request before/after handlers for bandits.
Nested functions are as follows
* detect_last_bandits: Loads any arms already assigned to this user
from the cookie.
* persist_bandits: Saves bandits down to storage engine at the end
of the request
* remember_bandit_arms: Sets the cookie for all requests that pulled an arm
* send_debug_header: Attaches a header for the MAB to the HTTP response for easier debugging
"""
@app.before_request
def detect_last_bandits():
bandits = request.cookies.get(app.extensions['mab'].cookie_name)
request.bandits_save = False
request.bandits_reward = set()
if bandits:
request.bandits = json.loads(bandits)
else:
request.bandits = {}
@app.after_request
def persist_bandits(response):
app.extensions['mab'].bandit_storage.save(app.extensions['mab'].bandits)
return response
@app.after_request
def remember_bandit_arms(response):
if request.bandits_save:
for bandit_id,arm in request.bandits.items():
#hook event for saving an impression here
app.extensions['mab'].bandits[bandit_id].pull_arm(arm)
for bandit_id, arm, reward_amt in request.bandits_reward:
try:
app.extensions['mab'].bandits[bandit_id].reward_arm(arm, reward_amt)
#hook event for saving a reward line here
except KeyError:
raise MABConfigException("Bandit %s not found" % bandit_id)
response.set_cookie(
app.extensions['mab'].cookie_name,
json.dumps(request.bandits))
return response
@app.after_request
def send_debug_header(response):
if app.extensions['mab'].debug_headers and request.bandits_save:
response.headers['X-MAB-Debug'] = "STORE; "+';'.join(
['%s:%s' % (key, val) for key, val in request.bandits.items()])
elif app.extensions['mab'].debug_headers:
response.headers['X-MAB-Debug'] = "SAVED; "+';'.join(['%s:%s' % (key, val) for key, val in request.bandits.items()])
return response
app.add_bandit = types.MethodType(add_bandit, app)
[docs]def add_bandit(app, name, bandit=None):
"""Attach a bandit for an experiment
:param name: The name of the experiment, will be used for lookups
:param bandit: The bandit to use for this experiment
:type bandit: Bandit
"""
saved_bandits = app.extensions['mab'].bandit_storage.load()
if name in saved_bandits.keys():
app.extensions['mab'].bandits[name] = saved_bandits[name]
else:
app.extensions['mab'].bandits[name] = bandit
[docs]def suggest_arm_for(key):
"""Get an experimental outcome by id. The primary way the implementor interfaces with their
experiments.
Suggests arms if not in cookie, using cookie val if present
:param key: The bandit/experiment to get a suggested arm for
:type key: string
:param also_pull: Should we register a pull/impression at the same time as suggesting
:raises KeyError: in case requested experiment does not exist
"""
app = current_app
try:
#Try to get the selected bandits from cookie
arm = app.extensions['mab'].bandits[key][request.bandits[key]]
return arm["id"], arm["value"]
except (AttributeError, TypeError, KeyError) as err:
#Assign an arm for a new client
try:
arm = app.extensions['mab'].bandits[key].suggest_arm()
request.bandits[key] = arm["id"]
request.bandits_save = True
return arm["id"], arm["value"]
except KeyError:
raise MABConfigException("Bandit %s not found" % key)
[docs]class MABConfigException(Exception):
"""Raised when internal state in MAB setup is invalid"""
pass
