diff options
Diffstat (limited to 'velascobot2/parrot.py')
| -rw-r--r-- | velascobot2/parrot.py | 6 |
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] |