1 files changed, 135 insertions, 0 deletions
diff --git a/velascobot2/parrot.py b/velascobot2/parrot.py
new file mode 100644
index 0000000..c7ba217
--- /dev/null
+++ b/velascobot2/parrot.py
@@ -0,0 +1,135 @@
+import random
+import logging
+from velascobot2.database import Database
+from velascobot2.utils import GraphicType, Response, ChatData
+
+logger = logging.getLogger(__name__)
+
+
+def _rewrite(text: str) -> list[str]:
+    """
+    Splits strings into lists of words delimited by whitespaces.
+    Extra whitespaces are inserted around newline characters so they are
+    included as their own Markov chain element, to avoid polluting with
+    very similar words that would only differ in having a whitespace
+    attached or not
+    """
+    words = text.replace("\n", " \n ").split(" ")
+    return [word.strip(" ") for word in words if len(word.strip(" ")) > 0]
+
+
+def _normalize_key(*words: str) -> str:
+    """Returns a case insensitive n-gram key by joining them with whitespaces"""
+    return " ".join(word.casefold() for word in words)
+
+
+def _triplets(wordlist: list[str]):
+    """Generates triplets of words from the given data string. So if our string
+    were "What a lovely day", we'd generate (What, a, lovely) and then
+    (a, lovely, day).
+    """
+    if len(wordlist) < 3:
+        return
+
+    for i in range(len(wordlist) - 2):
+        yield (wordlist[i], wordlist[i + 1], wordlist[i + 2])
+
+
+class Parrot:
+    def __init__(self, bot_name: str):
+        self.db = Database(bot_name)
+
+    def parse_text(self, chat_id: str, text: str) -> None:
+        logger.debug(f"{chat_id} - Processing text: {text}")
+        words = [w for w in _rewrite(text) if w]
+        words.append("")  # end of message
+        words = ["", ""] + words  # start of message
+
+        for key, word, count in self.parse_words(words):
+            self.db.add_words(chat_id, key, word, count)
+
+        self.db.incr_processed_messages(chat_id)
+
+    def parse_words(self, words: list[str]):
+        word_counts: dict[tuple[str, str], int] = {}
+
+        for w1, w2, w3 in _triplets(words):
+            key = _normalize_key(w1, w2)
+            word_counts[(key, w3)] = word_counts.get((key, w3), 0) + 1
+
+        for (key, word), count in word_counts.items():
+            yield (key, word, count)
+
+    def parse_graphic(
+        self, chat_id: str, graph_type: GraphicType, graph_id: str
+    ):
+        logger.debug(f"{chat_id} - Parsing graphic {graph_type}: {graph_id}")
+        assert graph_type in GraphicType
+        self.db.add_graphic(chat_id, graph_type, graph_id)
+
+    def give_text(self, chat_id: str) -> str:
+        w1, w2 = "", ""
+        w3 = "invalid"
+        words = []
+        while w3:
+            key = _normalize_key(w1, w2)
+            weighted_next_words = self.db.get_next_words(chat_id, key)
+            w3 = self.choose_word(weighted_next_words)
+            words.append(w3)
+            w1, w2 = w2, w3
+        return " ".join(words)
+
+    def give_graphic(self, chat_id: str) -> tuple[GraphicType, str]:
+        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:
+        words, weights = zip(*weighted_words)
+        return random.choices(words, weights=weights)[0]
+
+    def choose_graphic(
+        self, weighted_graphics: list[tuple[str, str, int]]
+    ) -> tuple[GraphicType, str]:
+        graphics = [
+            (graph_type, graph_id)
+            for graph_type, graph_id, _ in weighted_graphics
+        ]
+        weights = [count for _, _, count in weighted_graphics]
+        chosen_type, chosen_id = random.choices(graphics, weights=weights)[0]
+        return GraphicType(chosen_type), chosen_id
+
+    def choose_response_type(self, chat_id: str) -> Response:
+        graphics_count = self.db.count_graphics(chat_id)
+        words_count = self.db.count_words(chat_id, _normalize_key("", ""))
+        total = graphics_count + words_count
+
+        if random.random() < graphics_count / total:
+            return Response.GRAPHIC
+        else:
+            return Response.TEXT
+
+    def get_all_chats(self) -> list[ChatData]:
+        return self.db.get_all_chats()
+
+    def get_chat_data(self, chat_id: str) -> ChatData | None:
+        return self.db.get_chat_data(chat_id)
+
+    def ensure_chat_exists(
+        self, chat_id: str, chat_type: str, chat_name: str
+    ) -> bool:
+        return self.db.ensure_chat_exists(chat_id, chat_type, chat_name)
+
+    def update_chat_name(self, chat_id: str, chat_name: str):
+        self.db.update_chat_name(chat_id, chat_name)
+
+    def update_chat_answer_chance(self, chat_id: str, answer_chance: float):
+        self.db.update_chat_answer_chance(chat_id, answer_chance)
+
+    def update_chat_message_period(self, chat_id: str, period: int):
+        self.db.update_chat_message_period(chat_id, period)
+
+    def update_chat_set_restricted(self, chat_id: str, restricted: bool):
+        self.db.update_chat_set_restricted(chat_id, restricted)
+
+    def update_chat_set_silenced(self, chat_id: str, silenced: bool):
+        self.db.update_chat_set_silenced(chat_id, silenced)