from chatterbot.storage import StorageAdapter
class Query(object):
def __init__(self, query={}):
self.query = query
def value(self):
return self.query.copy()
def raw(self, data):
query = self.query.copy()
query.update(data)
return Query(query)
def statement_text_equals(self, statement_text):
query = self.query.copy()
query['text'] = statement_text
return Query(query)
def statement_text_not_in(self, statements):
query = self.query.copy()
if 'text' not in query:
query['text'] = {}
if '$nin' not in query['text']:
query['text']['$nin'] = []
query['text']['$nin'].extend(statements)
return Query(query)
def statement_response_list_contains(self, statement_text):
query = self.query.copy()
if 'in_response_to' not in query:
query['in_response_to'] = {}
if '$elemMatch' not in query['in_response_to']:
query['in_response_to']['$elemMatch'] = {}
query['in_response_to']['$elemMatch']['text'] = statement_text
return Query(query)
def statement_response_list_equals(self, response_list):
query = self.query.copy()
query['in_response_to'] = response_list
return Query(query)
class MongoDatabaseAdapter(StorageAdapter):
"""
The MongoDatabaseAdapter is an interface that allows
ChatterBot to store statements in a MongoDB database.
:keyword database: The name of the database you wish to connect to.
:type database: str
.. code-block:: python
database='chatterbot-database'
:keyword database_uri: The URI of a remote instance of MongoDB.
:type database_uri: str
.. code-block:: python
database_uri='mongodb://example.com:8100/'
"""
def __init__(self, **kwargs):
super(MongoDatabaseAdapter, self).__init__(**kwargs)
from pymongo import MongoClient
from pymongo.errors import OperationFailure
self.database_name = self.kwargs.get(
'database', 'chatterbot-database'
)
self.database_uri = self.kwargs.get(
'database_uri', 'mongodb://localhost:27017/'
)
# Use the default host and port
self.client = MongoClient(self.database_uri)
# Increase the sort buffer to 42M if possible
try:
self.client.admin.command({'setParameter': 1, 'internalQueryExecMaxBlockingSortBytes': 44040192})
except OperationFailure:
pass
# Specify the name of the database
self.database = self.client[self.database_name]
# The mongo collection of statement documents
self.statements = self.database['statements']
# The mongo collection of conversation documents
self.conversations = self.database['conversations']
# Set a requirement for the text attribute to be unique
self.statements.create_index('text', unique=True)
self.base_query = Query()
def get_statement_model(self):
"""
Return the class for the statement model.
"""
from chatterbot.conversation import Statement
# Create a storage-aware statement
statement = Statement
statement.storage = self
return statement
def get_response_model(self):
"""
Return the class for the response model.
"""
from chatterbot.conversation import Response
# Create a storage-aware response
response = Response
response.storage = self
return response
def count(self):
return self.statements.count()
def find(self, statement_text):
Statement = self.get_model('statement')
query = self.base_query.statement_text_equals(statement_text)
values = self.statements.find_one(query.value())
if not values:
return None
del values['text']
# Build the objects for the response list
values['in_response_to'] = self.deserialize_responses(
values.get('in_response_to', [])
)
return Statement(statement_text, **values)
def deserialize_responses(self, response_list):
"""
Takes the list of response items and returns
the list converted to Response objects.
"""
Statement = self.get_model('statement')
Response = self.get_model('response')
proxy_statement = Statement('')
for response in response_list:
text = response['text']
del response['text']
proxy_statement.add_response(
Response(text, **response)
)
return proxy_statement.in_response_to
def mongo_to_object(self, statement_data):
"""
Return Statement object when given data
returned from Mongo DB.
"""
Statement = self.get_model('statement')
statement_text = statement_data['text']
del statement_data['text']
statement_data['in_response_to'] = self.deserialize_responses(
statement_data.get('in_response_to', [])
)
return Statement(statement_text, **statement_data)
def filter(self, **kwargs):
"""
Returns a list of statements in the database
that match the parameters specified.
"""
import pymongo
query = self.base_query
order_by = kwargs.pop('order_by', None)
# Convert Response objects to data
if 'in_response_to' in kwargs:
serialized_responses = []
for response in kwargs['in_response_to']:
serialized_responses.append({'text': response})
query = query.statement_response_list_equals(serialized_responses)
del kwargs['in_response_to']
if 'in_response_to__contains' in kwargs:
query = query.statement_response_list_contains(
kwargs['in_response_to__contains']
)
del kwargs['in_response_to__contains']
query = query.raw(kwargs)
matches = self.statements.find(query.value())
if order_by:
direction = pymongo.ASCENDING
# Sort so that newer datetimes appear first
if order_by == 'created_at':
direction = pymongo.DESCENDING
matches = matches.sort(order_by, direction)
results = []
for match in list(matches):
results.append(self.mongo_to_object(match))
return results
def update(self, statement):
from pymongo import UpdateOne
from pymongo.errors import BulkWriteError
data = statement.serialize()
operations = []
update_operation = UpdateOne(
{'text': statement.text},
{'$set': data},
upsert=True
)
operations.append(update_operation)
# Make sure that an entry for each response is saved
for response_dict in data.get('in_response_to', []):
response_text = response_dict.get('text')
# $setOnInsert does nothing if the document is not created
update_operation = UpdateOne(
{'text': response_text},
{'$set': response_dict},
upsert=True
)
operations.append(update_operation)
try:
self.statements.bulk_write(operations, ordered=False)
except BulkWriteError as bwe:
# Log the details of a bulk write error
self.logger.error(str(bwe.details))
return statement
def create_conversation(self):
"""
Create a new conversation.
"""
conversation_id = self.conversations.insert_one({}).inserted_id
return conversation_id
def get_latest_response(self, conversation_id):
"""
Returns the latest response in a conversation if it exists.
Returns None if a matching conversation cannot be found.
"""
from pymongo import DESCENDING
statements = list(self.statements.find({
'conversations.id': conversation_id
}).sort('conversations.created_at', DESCENDING))
if not statements:
return None
return self.mongo_to_object(statements[-2])
def add_to_conversation(self, conversation_id, statement, response):
"""
Add the statement and response to the conversation.
"""
from datetime import datetime, timedelta
self.statements.update_one(
{
'text': statement.text
},
{
'$push': {
'conversations': {
'id': conversation_id,
'created_at': datetime.utcnow()
}
}
}
)
self.statements.update_one(
{
'text': response.text
},
{
'$push': {
'conversations': {
'id': conversation_id,
# Force the response to be at least one millisecond after the input statement
'created_at': datetime.utcnow() + timedelta(milliseconds=1)
}
}
}
)
def get_random(self):
"""
Returns a random statement from the database
"""
from random import randint
count = self.count()
if count < 1:
raise self.EmptyDatabaseException()
random_integer = randint(0, count - 1)
statements = self.statements.find().limit(1).skip(random_integer)
return self.mongo_to_object(list(statements)[0])
def remove(self, statement_text):
"""
Removes the statement that matches the input text.
Removes any responses from statements if the response text matches the
input text.
"""
for statement in self.filter(in_response_to__contains=statement_text):
statement.remove_response(statement_text)
self.update(statement)
self.statements.delete_one({'text': statement_text})
def get_response_statements(self):
"""
Return only statements that are in response to another statement.
A statement must exist which lists the closest matching statement in the
in_response_to field. Otherwise, the logic adapter may find a closest
matching statement that does not have a known response.
"""
response_query = self.statements.aggregate([{'$group': {'_id': '$in_response_to.text'}}])
responses = []
for r in response_query:
try:
responses.extend(r['_id'])
except TypeError:
pass
_statement_query = {
'text': {
'$in': responses
}
}
_statement_query.update(self.base_query.value())
statement_query = self.statements.find(_statement_query)
statement_objects = []
for statement in list(statement_query):
statement_objects.append(self.mongo_to_object(statement))
return statement_objects
def drop(self):
"""
Remove the database.
"""
self.client.drop_database(self.database_name)