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()

markov/migrations/0004_alter_markovstate_context.py

@@ -0,0 +1,19 @@
+# Generated by Django 3.2.18 on 2023-02-19 19:00
+
+from django.db import migrations, models
+import django.db.models.deletion
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('markov', '0003_auto_20161112_2348'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='markovstate',
+            name='context',
+            field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='states', to='markov.markovcontext'),
+        ),
+    ]

markov/models.py

@@ -41,7 +41,7 @@ class MarkovState(models.Model):
     v = models.CharField(max_length=128)
 
     count = models.IntegerField(default=0)
-    context = models.ForeignKey(MarkovContext, on_delete=models.CASCADE)
+    context = models.ForeignKey(MarkovContext, on_delete=models.CASCADE, related_name='states')
 
     class Meta:
         """Options for the model itself."""