TweaksHEADmain

1 files changed, 5 insertions, 1 deletions

diff --git a/velascobot2/parrot.py b/velascobot2/parrot.py
index be99e20..3e1db94 100644
--- a/velascobot2/parrot.py
+++ b/velascobot2/parrot.py
@@ -73,6 +73,8 @@ class Parrot:
             key = _normalize_key(w1, w2)
             weighted_next_words = self.db.get_next_words(chat_id, key)
             w3 = self.choose_word(weighted_next_words)
+            if w3 is None:
+                raise Exception(f"Got a nonexistent word. Generated text so far: {' '.join(words)}. Dropping text")
             words.append(w3)
             w1, w2 = w2, w3
         return " ".join(words)
@@ -81,7 +83,9 @@ class Parrot:
         weighted_graphics = self.db.get_graphics(chat_id)
         return self.choose_graphic(weighted_graphics)
 
-    def choose_word(self, weighted_words: list[tuple[str, int]]) -> str:
+    def choose_word(self, weighted_words: list[tuple[str, int]]) -> str | None:
+        if len(weighted_words) == 0:
+            return None
         words, weights = zip(*weighted_words)
         return random.choices(words, weights=weights)[0]