class MenuApplicationWorker(ApplicationWorker):
flush_sessions = False
@inlineCallbacks
def startWorker(self):
"""
Setup session manager.
"""
self.redis_config = self.config.get('redis_config', {})
self.redis_server = redis.Redis(**self.redis_config)
self.session_manager = SessionManager(
r_server=self.redis_server,
prefix="%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=getattr(self, 'MAX_SESSION_LENGTH', None)
)
if self.flush_sessions:
self.redis_server.flushdb()
yield super(MenuApplicationWorker, self).startWorker()
@inlineCallbacks
def stopWorker(self):
"""
Stop session manager.
"""
yield self.session_manager.stop()
yield super(MenuApplicationWorker, self).stopWorker()
def consume_user_message(self, message):
"""
Dynamically contruct menu based on user input.
If no session is found self.initial_menu is constructed.
"""
user_id = message.user()
session = self.session_manager.load_session(user_id)
try:
if not session:
session = self.session_manager.create_session(user_id)
menu_class = self.initial_menu
else:
if 'class' in session:
menu_class = load_class_by_string(session['class'])
else:
menu_class = self.initial_menu
try:
menu = menu_class(message, session, int(message['content']))
except ValueError:
menu = menu_class(message, session, '')
self.session_manager.save_session(user_id, session)
self.reply_to(message, menu.generate_menu(message['content']))
except Exception, e:
self.reply_to(message, str(e))
def setup_application(self):
if "yaml_template" in self.config:
with open(self.config["yaml_template"], "rb") as f:
self.yaml_template = f.read()
else:
self.yaml_template = resource_string(__name__,
"toy_decision_tree.yaml")
self.r_server = redis.Redis(**self.config.get('redis', {}))
self.session_manager = SessionManager(
self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
class MenuApplicationWorker(ApplicationWorker):
flush_sessions = False
@inlineCallbacks
def startWorker(self):
"""
Setup session manager.
"""
self.redis_config = self.config.get('redis_config', {})
self.redis_server = redis.Redis(**self.redis_config)
self.session_manager = SessionManager(
r_server=self.redis_server,
prefix="%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=getattr(self, 'MAX_SESSION_LENGTH', None))
if self.flush_sessions:
self.redis_server.flushdb()
yield super(MenuApplicationWorker, self).startWorker()
@inlineCallbacks
def stopWorker(self):
"""
Stop session manager.
"""
yield self.session_manager.stop()
yield super(MenuApplicationWorker, self).stopWorker()
def consume_user_message(self, message):
"""
Dynamically contruct menu based on user input.
If no session is found self.initial_menu is constructed.
"""
user_id = message.user()
session = self.session_manager.load_session(user_id)
try:
if not session:
session = self.session_manager.create_session(user_id)
menu_class = self.initial_menu
else:
if 'class' in session:
menu_class = load_class_by_string(session['class'])
else:
menu_class = self.initial_menu
try:
menu = menu_class(message, session, int(message['content']))
except ValueError:
menu = menu_class(message, session, '')
self.session_manager.save_session(user_id, session)
self.reply_to(message, menu.generate_menu(message['content']))
except Exception, e:
self.reply_to(message, str(e))
def startWorker(self):
"""
Setup session manager.
"""
self.redis_config = self.config.get('redis_config', {})
self.redis_server = redis.Redis(**self.redis_config)
self.session_manager = SessionManager(
r_server=self.redis_server,
prefix="%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=getattr(self, 'MAX_SESSION_LENGTH', None))
if self.flush_sessions:
self.redis_server.flushdb()
yield super(MenuApplicationWorker, self).startWorker()
class DynamicMenuApplicationWorker(ApplicationWorker):
@inlineCallbacks
def startWorker(self):
"""
Setup session manager.
"""
self.redis_config = self.config.get("redis_config", {})
self.redis_server = redis.Redis(**self.redis_config)
self.session_manager = SessionManager(
r_server=self.redis_server,
prefix="%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=getattr(self, "MAX_SESSION_LENGTH", None),
)
self.redis_server.flushdb()
yield super(DynamicMenuApplicationWorker, self).startWorker()
@inlineCallbacks
def stopWorker(self):
"""
Stop session manager.
"""
yield self.session_manager.stop()
yield super(DynamicMenuApplicationWorker, self).stopWorker()
def consume_user_message(self, message):
"""
Dynamically contruct menu based on user input.
If no session is found self.initial_menu is used as frist menu.
"""
user_id = message.user()
session = self.session_manager.load_session(user_id)
try:
if not session:
session = self.session_manager.create_session(user_id)
menu_class = self.initial_menu
else:
if "class" in session:
menu_class = class_from_str(session["class"])
else:
menu_class = self.initial_menu
try:
menu = menu_class(message, session, int(message["content"]))
except ValueError:
menu = menu_class(message, session, "")
self.session_manager.save_session(user_id, session)
self.reply_to(message, menu.generate_menu(message["content"]))
except Exception, e:
self.reply_to(message, str(e))
def startWorker(self):
self.session_manager = SessionManager(
get_deploy_int(self._amqp_client.vhost),
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
yield super(WikipediaWorker, self).startWorker()
def setup_application(self):
if "yaml_template" in self.config:
with open(self.config["yaml_template"], "rb") as f:
self.yaml_template = f.read()
else:
self.yaml_template = resource_string(__name__,
"toy_decision_tree.yaml")
self.r_server = redis.Redis(**self.config.get('redis', {}))
self.session_manager = SessionManager(self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
def startWorker(self):
self.sms_transport = self.config.get("sms_transport", None)
self.override_sms_address = self.config.get("override_sms_address", None)
db = get_deploy_int(self._amqp_client.vhost)
self.r_server = redis.Redis("localhost", db=db)
self.session_manager = SessionManager(
self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH,
)
yield super(WikipediaWorker, self).startWorker()
def startWorker(self):
"""
Setup session manager.
"""
self.redis_config = self.config.get("redis_config", {})
self.redis_server = redis.Redis(**self.redis_config)
self.session_manager = SessionManager(
r_server=self.redis_server,
prefix="%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=getattr(self, "MAX_SESSION_LENGTH", None),
)
self.redis_server.flushdb()
yield super(DynamicMenuApplicationWorker, self).startWorker()
class SessionManagerTestCase(TestCase):
def setUp(self):
self.fake_redis = FakeRedis()
self.sm = SessionManager(self.fake_redis, prefix="test")
def tearDown(self):
self.sm.stop()
self.fake_redis.teardown()
def test_active_sessions(self):
def get_sessions():
return sorted(self.sm.active_sessions())
def ids():
return [x[0] for x in get_sessions()]
self.assertEqual(ids(), [])
self.sm.create_session("u1")
self.assertEqual(ids(), ["u1"])
# 10 seconds later
self.sm.create_session("u2", created_at=time.time() + 10)
self.assertEqual(ids(), ["u1", "u2"])
s1, s2 = get_sessions()
self.assertTrue(s1[1]['created_at'] < s2[1]['created_at'])
def test_schedule_session_expiry(self):
self.sm.max_session_length = 60.0
self.sm.create_session("u1")
def test_create_and_retrieve_session(self):
session = self.sm.create_session("u1")
self.assertEqual(sorted(session.keys()), ['created_at'])
self.assertTrue(time.time() - float(session['created_at']) < 10.0)
loaded = self.sm.load_session("u1")
self.assertEqual(loaded, session)
def test_save_session(self):
test_session = {"foo": 5, "bar": "baz"}
self.sm.create_session("u1")
self.sm.save_session("u1", test_session)
session = self.sm.load_session("u1")
self.assertTrue(session.pop('created_at') is not None)
# Redis saves & returns all session values as strings
self.assertEqual(session, dict([map(str, kvs) for kvs
in test_session.items()]))
def test_lazy_clearing(self):
self.sm.save_session('user_id', {})
self.assertEqual(list(self.sm.active_sessions()), [])
def setUp(self):
self.fake_redis = FakeRedis()
self.sm = SessionManager(self.fake_redis, prefix="test")
class DecisionTreeWorker(ApplicationWorker):
"""Demo application that serves a series of questions.
Configuration options:
:type worker_name: str
:param worker_name:
Worker name. Used as part of the redis prefix for session data.
:type yaml_template: str
:param yaml_template:
Name of file containing the YAML template for the decision tree.
Optional. If left out, a demo decision tree is read from
vumi.demos/toy_decision_tree.yaml.
"""
MAX_SESSION_LENGTH = 3 * 60
def validate_config(self):
if "worker_name" not in self.config:
raise ConfigError("DecisionTreeWorker requires a worker_name"
" in its configuration.")
def setup_application(self):
if "yaml_template" in self.config:
with open(self.config["yaml_template"], "rb") as f:
self.yaml_template = f.read()
else:
self.yaml_template = resource_string(__name__,
"toy_decision_tree.yaml")
self.r_server = redis.Redis(**self.config.get('redis', {}))
self.session_manager = SessionManager(
self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
def teardown_application(self):
self.session_manager.stop()
def consume_user_message(self, msg):
user_id = msg.user()
response = ''
continue_session = False
if not self.yaml_template:
log.err("yaml_template is missing")
return
decision_tree = self.get_decision_tree(user_id)
if not decision_tree.is_started():
decision_tree.start()
elif not decision_tree.is_completed():
decision_tree.answer(msg.payload['content'])
if not decision_tree.is_completed():
response += decision_tree.question()
continue_session = True
self.save_decision_tree(user_id, decision_tree)
else:
response += decision_tree.finish() or ''
self.post_result(decision_tree)
self.delete_decision_tree(user_id)
self.reply_to(msg, response, continue_session)
def post_result(self, tree):
log.msg(tree.dump_json_data())
def get_initial_data(self, tree):
return tree.get_initial_data()
def get_decision_tree(self, user_id):
data = self.session_manager.load_session(user_id)
if not data:
data = self.session_manager.create_session(user_id)
if 'decision_tree' in data:
return yaml.safe_load(data['decision_tree'])
else:
return self.setup_new_decision_tree()
def save_decision_tree(self, user_id, tree):
data = {}
data['decision_tree'] = yaml.safe_dump(tree)
self.session_manager.save_session(user_id, data)
def delete_decision_tree(self, user_id):
self.session_manager.clear_session(user_id)
def setup_new_decision_tree(self):
decision_tree = TraversedDecisionTree()
decision_tree.load_yaml_template(self.yaml_template)
json_string = self.get_initial_data(decision_tree)
decision_tree.load_json_data(json_string)
return decision_tree
class DecisionTreeWorker(ApplicationWorker):
"""Demo application that serves a series of questions.
Configuration options:
:type worker_name: str
:param worker_name:
Worker name. Used as part of the redis prefix for session data.
:type yaml_template: str
:param yaml_template:
Name of file containing the YAML template for the decision tree.
Optional. If left out, a demo decision tree is read from
vumi.demos/toy_decision_tree.yaml.
"""
MAX_SESSION_LENGTH = 3 * 60
def validate_config(self):
if "worker_name" not in self.config:
raise ConfigError("DecisionTreeWorker requires a worker_name"
" in its configuration.")
def setup_application(self):
if "yaml_template" in self.config:
with open(self.config["yaml_template"], "rb") as f:
self.yaml_template = f.read()
else:
self.yaml_template = resource_string(__name__,
"toy_decision_tree.yaml")
self.r_server = redis.Redis(**self.config.get('redis', {}))
self.session_manager = SessionManager(self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
def teardown_application(self):
self.session_manager.stop()
def consume_user_message(self, msg):
user_id = msg.user()
response = ''
continue_session = False
if not self.yaml_template:
log.err("yaml_template is missing")
return
decision_tree = self.get_decision_tree(user_id)
if not decision_tree.is_started():
decision_tree.start()
elif not decision_tree.is_completed():
decision_tree.answer(msg.payload['content'])
if not decision_tree.is_completed():
response += decision_tree.question()
continue_session = True
self.save_decision_tree(user_id, decision_tree)
else:
response += decision_tree.finish() or ''
self.post_result(decision_tree)
self.delete_decision_tree(user_id)
self.reply_to(msg, response, continue_session)
def post_result(self, tree):
log.msg(tree.dump_json_data())
def get_initial_data(self, tree):
return tree.get_initial_data()
def get_decision_tree(self, user_id):
data = self.session_manager.load_session(user_id)
if not data:
data = self.session_manager.create_session(user_id)
if 'decision_tree' in data:
return yaml.safe_load(data['decision_tree'])
else:
return self.setup_new_decision_tree()
def save_decision_tree(self, user_id, tree):
data = {}
data['decision_tree'] = yaml.safe_dump(tree)
self.session_manager.save_session(user_id, data)
def delete_decision_tree(self, user_id):
self.session_manager.clear_session(user_id)
def setup_new_decision_tree(self):
decision_tree = TraversedDecisionTree()
decision_tree.load_yaml_template(self.yaml_template)
json_string = self.get_initial_data(decision_tree)
decision_tree.load_json_data(json_string)
return decision_tree
class SessionManagerTestCase(TestCase):
def setUp(self):
self.sm = SessionManager(db=0, prefix="test")
self.sm.r_server = FakeRedis() # stub out redis
def tearDown(self):
self.sm.stop()
self.sm.r_server.teardown() # teardown fake redis
def test_active_sessions(self):
def get_sessions():
return sorted(self.sm.active_sessions())
def ids():
return [x[0] for x in get_sessions()]
self.assertEqual(ids(), [])
self.sm.create_session("u1")
self.assertEqual(ids(), ["u1"])
self.sm.create_session("u2")
self.assertEqual(ids(), ["u1", "u2"])
s1, s2 = get_sessions()
self.assertTrue(s1[1]['created_at'] < s2[1]['created_at'])
def test_schedule_session_expiry(self):
self.sm.max_session_length = 60.0
self.sm.create_session("u1")
def test_create_and_retreive_session(self):
session = self.sm.create_session("u1")
self.assertEqual(sorted(session.keys()), ['created_at'])
self.assertTrue(time.time() - session['created_at'] < 10.0)
loaded = self.sm.load_session("u1")
self.assertEqual(loaded, session)
def test_save_session(self):
test_session = {"foo": 5, "bar": "baz"}
self.sm.create_session("u1")
self.sm.save_session("u1", test_session)
session = self.sm.load_session("u1")
self.assertTrue(session.pop('created_at') is not None)
self.assertEqual(session, test_session)
class WikipediaWorker(ApplicationWorker):
MAX_SESSION_LENGTH = 3 * 60
@inlineCallbacks
def startWorker(self):
self.sms_transport = self.config.get("sms_transport", None)
self.override_sms_address = self.config.get("override_sms_address", None)
db = get_deploy_int(self._amqp_client.vhost)
self.r_server = redis.Redis("localhost", db=db)
self.session_manager = SessionManager(
self.r_server,
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH,
)
yield super(WikipediaWorker, self).startWorker()
@inlineCallbacks
def stopWorker(self):
yield self.session_manager.stop()
yield super(WikipediaWorker, self).stopWorker()
@inlineCallbacks
def consume_user_message(self, msg):
user_id = msg.user()
session = self.session_manager.load_session(user_id)
if (not session) or (msg["content"] is None):
session = self.session_manager.create_session(user_id)
session["state"] = "new"
pfunc = getattr(self, "process_message_%s" % (session["state"],))
session = yield pfunc(msg, session)
if session["state"] is not None:
self.session_manager.save_session(user_id, session)
else:
self.session_manager.clear_session(user_id)
def process_message_new(self, msg, session):
self.reply_to(msg, "What would you like to search Wikipedia for?", True)
session["state"] = "searching"
return session
@inlineCallbacks
def process_message_searching(self, msg, session):
query = msg["content"].strip()
results = yield WikipediaAPI().search(query)
if results:
session["results"] = json.dumps(results)
self.reply_to(msg, format_options(results), True)
session["state"] = "sections"
else:
self.reply_to(msg, "Sorry, no Wikipedia results for %s" % query, False)
session["state"] = None
returnValue(session)
def select_option(self, options, msg):
response = msg["content"].strip()
if response.isdigit():
try:
return options[int(response) - 1]
except (KeyError, IndexError):
pass
self.reply_to(msg, "Sorry, invalid selection. Please restart and try again", False)
return None
@inlineCallbacks
def process_message_sections(self, msg, session):
selection = self.select_option(json.loads(session["results"]), msg)
if not selection:
session["state"] = None
returnValue(session)
session["page"] = selection
results = yield WikipediaAPI().get_sections(selection)
results = [selection] + results
session["results"] = json.dumps(results)
self.reply_to(msg, format_options(results), True)
session["state"] = "content"
returnValue(session)
@inlineCallbacks
def process_message_content(self, msg, session):
selection = self.select_option(json.loads(session["results"]), msg)
if not selection:
session["state"] = None
returnValue(session)
content = yield WikipediaAPI().get_content(session["page"], int(msg["content"].strip()) - 1)
ussd_cont = "%s...\n(Full content sent by SMS.)" % (content[:100],)
self.reply_to(msg, ussd_cont, False)
if self.sms_transport:
bmsg = msg.reply(content[:250])
bmsg["transport_name"] = self.sms_transport
if self.override_sms_address:
bmsg["to_addr"] = self.override_sms_address
self.transport_publisher.publish_message(bmsg, routing_key="%s.outbound" % (self.sms_transport,))
session["state"] = None
returnValue(session)
def setUp(self):
self.sm = SessionManager(db=0, prefix="test")
self.sm.r_server = FakeRedis() # stub out redis
def setUp(self):
self.fake_redis = FakeRedis()
self.add_cleanup(self.fake_redis.teardown)
self.sm = SessionManager(self.fake_redis, prefix="test")
self.add_cleanup(self.sm.stop)
def setUp(self):
self.fake_redis = FakeRedis()
self.add_cleanup(self.fake_redis.teardown)
self.sm = SessionManager(self.fake_redis, prefix="test")
self.add_cleanup(self.sm.stop)
class TestSessionManager(VumiTestCase):
def setUp(self):
self.fake_redis = FakeRedis()
self.add_cleanup(self.fake_redis.teardown)
self.sm = SessionManager(self.fake_redis, prefix="test")
self.add_cleanup(self.sm.stop)
def test_active_sessions(self):
def get_sessions():
return sorted(self.sm.active_sessions())
def ids():
return [x[0] for x in get_sessions()]
self.assertEqual(ids(), [])
self.sm.create_session("u1")
self.assertEqual(ids(), ["u1"])
# 10 seconds later
self.sm.create_session("u2", created_at=time.time() + 10)
self.assertEqual(ids(), ["u1", "u2"])
s1, s2 = get_sessions()
self.assertTrue(s1[1]['created_at'] < s2[1]['created_at'])
def test_schedule_session_expiry(self):
self.sm.max_session_length = 60.0
self.sm.create_session("u1")
def test_create_and_retrieve_session(self):
session = self.sm.create_session("u1")
self.assertEqual(sorted(session.keys()), ['created_at'])
self.assertTrue(time.time() - float(session['created_at']) < 10.0)
loaded = self.sm.load_session("u1")
self.assertEqual(loaded, session)
def test_save_session(self):
test_session = {"foo": 5, "bar": "baz"}
self.sm.create_session("u1")
self.sm.save_session("u1", test_session)
session = self.sm.load_session("u1")
self.assertTrue(session.pop('created_at') is not None)
# Redis saves & returns all session values as strings
self.assertEqual(session,
dict([map(str, kvs) for kvs in test_session.items()]))
def test_lazy_clearing(self):
self.sm.save_session('user_id', {})
self.assertEqual(list(self.sm.active_sessions()), [])
class WikipediaWorker(ApplicationWorker):
MAX_SESSION_LENGTH = 3 * 60
@inlineCallbacks
def startWorker(self):
self.session_manager = SessionManager(
get_deploy_int(self._amqp_client.vhost),
"%(worker_name)s:%(transport_name)s" % self.config,
max_session_length=self.MAX_SESSION_LENGTH)
yield super(WikipediaWorker, self).startWorker()
@inlineCallbacks
def stopWorker(self):
yield self.session_manager.stop()
yield super(WikipediaWorker, self).stopWorker()
def consume_user_message(self, msg):
user_id = msg.user()
session = self.session_manager.load_session(user_id)
if session and msg['content'] is not None:
self.resume_wikipedia_session(msg, session)
else:
session = self.session_manager.create_session(user_id)
self.new_wikipedia_session(msg, session)
def new_wikipedia_session(self, msg, session):
self.reply_to(msg, "What would you like to search Wikipedia for?",
True)
def resume_wikipedia_session(self, msg, session):
response = msg['content'].strip()
if response.isdigit():
self.handle_selection(msg, session, int(response))
else:
self.handle_search(msg, session, response)
@inlineCallbacks
def handle_search(self, msg, session, query):
results = yield OpenSearch().search(query)
if results:
session['results'] = json.dumps(results)
self.reply_to(msg, pretty_print_results(results), True)
self.session_manager.save_session(msg.user(), session)
else:
self.reply_to(msg, 'Sorry, no Wikipedia results for %s' % query,
False)
self.session_manager.clear_session(msg.user())
def handle_selection(self, msg, session, number):
try:
results = json.loads(session['results'])
interest = results[number - 1]
self.reply_to(msg,
'%s: %s...\nFull text will be delivered to you via SMS' %
(interest['text'], interest['description'][:100]), False)
except (KeyError, IndexError):
self.reply_to(msg,
'Sorry, invalid selection. Please restart and try again',
False)
self.session_manager.clear_session(msg.user())
