91 lines
5.3 KiB
Python
91 lines
5.3 KiB
Python
"""Clean up learned chains with speaker nicks (from the bridge) or self (because the bridge broke the regex)."""
|
|
from django.core.management import BaseCommand
|
|
|
|
from ircbot.models import IrcChannel
|
|
from markov.models import MarkovContext, MarkovState
|
|
|
|
|
|
class Command(BaseCommand):
|
|
"""Find markov chains that erroneously have speaker/self nicks and remove them."""
|
|
|
|
def handle(self, *args, **kwargs):
|
|
"""Scan the DB, looking for bad chains, and repair them."""
|
|
candidate_channels = IrcChannel.objects.exclude(discord_bridge='')
|
|
markov_contexts = MarkovContext.objects.filter(markovtarget__name__in=list(candidate_channels))
|
|
for context in markov_contexts:
|
|
self.stdout.write(self.style.NOTICE(f"scanning context {context}..."))
|
|
# get starting states that look like they came over the bridge
|
|
bridge_states = context.states.filter(k1=MarkovState._start1, k2=MarkovState._start2,
|
|
v__regex=r'<[A-Za-z0-9_]+>', context=context)
|
|
self._chain_remover(context, bridge_states)
|
|
# get states that look like mentions
|
|
for target in context.markovtarget_set.all():
|
|
if target.channel.server.additional_addressed_nicks:
|
|
all_nicks = '|'.join(target.channel.server.additional_addressed_nicks.split('\n') +
|
|
[target.channel.server.nickname])
|
|
else:
|
|
all_nicks = target.channel.server.nickname
|
|
mention_regex = r'^(({nicks})[:,]|@({nicks}))$'.format(nicks=all_nicks)
|
|
mention_states = context.states.filter(k1=MarkovState._start1, k2=MarkovState._start2,
|
|
v__regex=mention_regex, context=context)
|
|
self._chain_remover(context, mention_states)
|
|
|
|
def _chain_remover(self, context, start_states):
|
|
"""Remove a given k from markov states, deleting the found states after rebuilding subsequent states.
|
|
|
|
As in, if trying to remove A,B -> X, then B,X -> C and X,C -> D must be rebuilt (A,B -> C / B,C -> D)
|
|
then the three states with X deleted.
|
|
"""
|
|
for start_state in start_states:
|
|
self.stdout.write(self.style.NOTICE(f" diving into {start_state}..."))
|
|
# find the states that build off of the start
|
|
second_states = context.states.filter(k1=start_state.k2, k2=start_state.v)
|
|
|
|
for second_state in second_states:
|
|
self.stdout.write(self.style.NOTICE(f" diving into {second_state}..."))
|
|
# find the third states
|
|
leaf_states = context.states.filter(k1=second_state.k2, k2=second_state.v)
|
|
|
|
for leaf_state in leaf_states:
|
|
self.stdout.write(self.style.NOTICE(f" upserting state based on {leaf_state}"))
|
|
# get/update state without the nick from the bridge
|
|
try:
|
|
updated_leaf = MarkovState.objects.get(k1=second_state.k1, k2=leaf_state.k2, v=leaf_state.v,
|
|
context=context)
|
|
updated_leaf.count += leaf_state.count
|
|
updated_leaf.save()
|
|
self.stdout.write(self.style.SUCCESS(f" updated count for {updated_leaf}"))
|
|
except MarkovState.DoesNotExist:
|
|
new_leaf = MarkovState.objects.create(k1=second_state.k1, k2=leaf_state.k2, v=leaf_state.v,
|
|
context=context)
|
|
new_leaf.count = leaf_state.count
|
|
new_leaf.save()
|
|
self.stdout.write(self.style.SUCCESS(f" created {new_leaf}"))
|
|
|
|
# remove the migrated leaf state
|
|
self.stdout.write(self.style.SUCCESS(f" deleting {leaf_state}"))
|
|
leaf_state.delete()
|
|
|
|
# take care of the new middle state
|
|
self.stdout.write(self.style.NOTICE(f" upserting state based on {second_state}"))
|
|
try:
|
|
updated_second = MarkovState.objects.get(k1=start_state.k1, k2=start_state.k2, v=second_state.v,
|
|
context=context)
|
|
updated_second.count += second_state.count
|
|
updated_second.save()
|
|
self.stdout.write(self.style.SUCCESS(f" updated count for {updated_second}"))
|
|
except MarkovState.DoesNotExist:
|
|
new_second = MarkovState.objects.create(k1=start_state.k1, k2=start_state.k2, v=second_state.v,
|
|
context=context)
|
|
new_second.count = second_state.count
|
|
new_second.save()
|
|
self.stdout.write(self.style.SUCCESS(f" created {new_second}"))
|
|
|
|
# remove the migrated second state
|
|
self.stdout.write(self.style.SUCCESS(f" deleting {second_state}"))
|
|
second_state.delete()
|
|
|
|
# remove the dead end original start
|
|
self.stdout.write(self.style.SUCCESS(f" deleting {start_state}"))
|
|
start_state.delete()
|