management command to remove nicks from chains due to bridge

markov/management/__init__.py

@@ -0,0 +1 @@
+"""Management operations for the markov plugin and models."""

markov/management/commands/__init__.py

@@ -0,0 +1 @@
+"""Management commands for the markov plugin and models."""

markov/management/commands/fix_chains_with_nicks_via_bridge.py

@@ -0,0 +1,69 @@
+"""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
+            start_states = context.states.filter(k1=MarkovState._start1, k2=MarkovState._start2,
+                                                 v__regex=r'<.*>')
+            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)
+                            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)
+                        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()