ai-plays-jackbox 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

ai_plays_jackbox/__init__.py

@@ -1 +0,0 @@
-from .run import run

ai_plays_jackbox/bot/__init__.py

@@ -1 +0,0 @@
-from .bot_base import JackBoxBotBase

ai_plays_jackbox/bot/bot_base.py

@@ -6,9 +6,10 @@ from typing import Optional
 from urllib import parse
 from uuid import uuid4
 
+import cv2
 import html2text
+import numpy as np
 from loguru import logger
-from ollama import Options, chat
 from pydantic import BaseModel, Field, field_validator
 from websocket import WebSocketApp
 
@@ -118,8 +119,8 @@ class JackBoxBotBase(ABC):
         elif server_message.opcode == "object" or server_message.opcode == "text":
             if server_message.opcode == "object":
                 operation = ObjectOperation(**server_message.result)
-            if server_message.opcode == "text":
-                operation = TextOperation(**server_message.result)
+            elif server_message.opcode == "text":
+                operation = TextOperation(**server_message.result)  # type: ignore
 
             if self._is_player_operation_key(operation.key):
                 self._handle_player_operation(operation.json_data)
@@ -141,16 +142,49 @@ class JackBoxBotBase(ABC):
     def _send_ws(self, opcode: str, params: dict) -> None:
         self._message_sequence += 1
         message = {"seq": self._message_sequence, "opcode": opcode, "params": params}
-        self._ws.send(json.dumps(message))
+        if self._ws is not None:
+            self._ws.send(json.dumps(message))
+        else:
+            raise Exception("Websocket connection has not been initialized")
 
     def _client_send(self, request: dict) -> None:
-        self._message_sequence += 1
         params = {"from": self._player_id, "to": 1, "body": request}
         self._send_ws("client/send", params)
 
+    def _object_update(self, key: str, val: dict) -> None:
+        params = {"key": key, "val": val}
+        self._send_ws("object/update", params)
+
+    def _text_update(self, key: str, val: str) -> None:
+        params = {"key": key, "val": val}
+        self._send_ws("text/update", params)
+
     def _html_to_text(self, html: str) -> str:
         return html2text.html2text(html)
 
+    def _image_bytes_to_polylines(self, image_bytes: bytes, canvas_height: int, canvas_width: int) -> list[str]:
+        # Let's edge trace the outputted image to contours
+        image_array = np.frombuffer(image_bytes, dtype=np.uint8)
+        image = cv2.imdecode(image_array, flags=1)
+        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        edges = cv2.Canny(gray_image, threshold1=100, threshold2=200)
+        contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Figure out scaling factor
+        height, width = gray_image.shape
+        scale_x = canvas_width / width
+        scale_y = canvas_height / height
+        scale_factor = min(scale_x, scale_y)
+
+        # generate the polylines from the contours
+        polylines = []
+        for contour in contours:
+            if len(contour) > 1:  # Only include contours with 2 or more points
+                polyline = [f"{int(point[0][0] * scale_factor)},{int(point[0][1] * scale_factor)}" for point in contour]  # type: ignore
+                polylines.append("|".join(polyline))
+
+        return polylines
+
     def __del__(self):
         self.disconnect()
 
@@ -169,9 +203,9 @@ class TextOperation(BaseModel):
     value: str = Field(alias="val")
     version: int
 
-    @field_validator("json_data")
+    @field_validator("json_data")  # type: ignore
     def set_json_data(cls, value, values: dict):
-        return json.loads(values.get("value"))
+        return json.loads(values.get("value", ""))
 
 
 class ObjectOperation(BaseModel):
@@ -181,6 +215,6 @@ class ObjectOperation(BaseModel):
     value: str = Field(default="")
     version: int
 
-    @field_validator("value")
+    @field_validator("value")  # type: ignore
     def set_value(cls, value, values: dict):
         return json.dumps(values.get("json_data"))

ai_plays_jackbox/bot/bot_factory.py

@@ -1,30 +1,28 @@
-from typing import Optional
-
 from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
+from ai_plays_jackbox.bot.jackbox5.mad_verse_city import MadVerseCityBot
+from ai_plays_jackbox.bot.jackbox5.patently_stupid import PatentlyStupidBot
+from ai_plays_jackbox.bot.jackbox6.joke_boat import JokeBoatBot
 from ai_plays_jackbox.bot.jackbox7.quiplash3 import Quiplash3Bot
+from ai_plays_jackbox.bot.standalone.drawful2 import Drawful2Bot
 from ai_plays_jackbox.llm.chat_model import ChatModel
-from ai_plays_jackbox.llm.ollama_model import OllamaModel
+
+BOT_TYPES: dict[str, type[JackBoxBotBase]] = {
+    "quiplash3": Quiplash3Bot,
+    "patentlystupid": PatentlyStupidBot,
+    "drawful2international": Drawful2Bot,
+    "rapbattle": MadVerseCityBot,
+    "jokeboat": JokeBoatBot,
+}
 
 
 class JackBoxBotFactory:
     @staticmethod
     def get_bot(
         room_type: str,
+        chat_model: ChatModel,
         name: str = "FunnyBot",
         personality: str = "You are the funniest bot ever.",
-        chat_model: Optional[ChatModel] = None,
     ) -> JackBoxBotBase:
-        if chat_model is None:
-            chat_model = OllamaModel()
-        if room_type == "quiplash3":
-            return Quiplash3Bot(name=name, personality=personality, chat_model=chat_model)
-        # elif room_type == "ridictionary":
-        #     return DictionariumBot(name=name, personality=personality)
-        # elif room_type == "patentlystupid":
-        #     return PatentlyStupidBot(name=name, personality=personality, model=model)
-        # elif room_type == "fourbage":
-        #     return Fibbage4Bot(name=name, personality=personality, model=model)
-        # elif room_type == "rapbattle":
-        #     return MadVerseCityBot(name=name, personality=personality, model=model)
-        else:
+        if room_type not in BOT_TYPES.keys():
             raise ValueError(f"Unknown room type: {room_type}")
+        return BOT_TYPES[room_type](name=name, personality=personality, chat_model=chat_model)

ai_plays_jackbox/bot/bot_personality.py

@@ -27,7 +27,6 @@ class JackBoxBotVariant(Enum):
     SORRYBOT = JackBoxBotPersonality(
         name="SorryBot",
         personality="You are embarrassed by your answers and feel the need to apologize profusely to the rest of the group for them.",
-        voice="Mouse",
     )
     HOSTAGEBOT = JackBoxBotPersonality(
         name="HostageBot",

ai_plays_jackbox/bot/jackbox5/bot_base.py

@@ -0,0 +1,26 @@
+from abc import ABC
+from typing import Optional
+
+from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
+
+
+class JackBox5BotBase(JackBoxBotBase, ABC):
+    _actual_player_operation_key: Optional[str] = None
+
+    @property
+    def _player_operation_key(self):
+        return f"bc:customer:"
+
+    def _is_player_operation_key(self, operation_key: str) -> bool:
+        if self._actual_player_operation_key is None and self._player_operation_key in operation_key:
+            self._actual_player_operation_key = operation_key
+            return True
+        else:
+            return self._actual_player_operation_key == operation_key
+
+    @property
+    def _room_operation_key(self):
+        return "bc:room"
+
+    def _is_room_operation_key(self, operation_key: str) -> bool:
+        return operation_key == self._room_operation_key

ai_plays_jackbox/bot/jackbox5/mad_verse_city.py

@@ -0,0 +1,121 @@
+import random
+from typing import Optional
+
+import demoji
+from loguru import logger
+
+from ai_plays_jackbox.bot.jackbox5.bot_base import JackBox5BotBase
+
+_WORD_PROMPT_TEMPLATE = """
+You are playing Mad Verse City. You are being asked to come up with a word.
+
+{prompt}
+
+When generating your response, follow these rules:
+- You response must be {max_length} characters or less. It must be a singular word
+- Your personality is: {personality}
+- Do not include quotes in your response or any newlines, just the response itself.
+- Some suggestions for the word are {suggestions}
+"""
+
+_RHYME_PROMPT_TEMPLATE = """
+You are playing Mad Verse City.
+
+{prompt}
+
+When generating your response, follow these rules:
+- Your response must rhyme with {rhyme_word}. It cannot be the same word as that.
+- You response must be {max_length} characters or less.
+- Your personality is: {personality}
+- Do not include quotes in your response or any newlines, just the response itself.
+"""
+
+
+class MadVerseCityBot(JackBox5BotBase):
+    _actual_player_operation_key: Optional[str] = None
+    _current_word: str
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _handle_welcome(self, data: dict):
+        pass
+
+    def _handle_player_operation(self, data: dict):
+        if not data:
+            return
+        room_state = data.get("state", None)
+        if not room_state:
+            return
+
+        prompt: dict[str, str] = data.get("prompt", {})
+        prompt_html = prompt.get("html", "")
+        clean_prompt = self._html_to_text(prompt_html)
+
+        max_length = data.get("maxLength", 40)
+        suggestions = data.get("suggestions", [])
+        if not suggestions:
+            suggestions = []
+        choices: list[dict] = data.get("choices", [])
+
+        match room_state:
+            case "EnterSingleText":
+                if "Give me a" in clean_prompt:
+                    word = self._generate_word(clean_prompt, max_length - 10, suggestions)
+                    self._current_word = demoji.replace(word, "")
+                    self._client_send({"action": "write", "entry": word})
+
+                if "Now, write a line to rhyme with" in clean_prompt:
+                    rhyme = self._generate_rhyme(clean_prompt, max_length - 10, self._current_word)
+                    rhyme = demoji.replace(rhyme, "")
+                    self._client_send({"action": "write", "entry": rhyme})
+
+            case "MakeSingleChoice":
+                if not choices:
+                    pass
+                if "Who won this battle" in data.get("prompt", {}).get("text", "") and data.get("chosen", None) != 0:
+                    choice_indexes = [i for i in range(0, len(choices))]
+                    choice = random.choice(choice_indexes)
+                    self._client_send({"action": "choose", "choice": choice})
+
+                if (
+                    "Press this button to skip the tutorial" in data.get("prompt", {}).get("text", "")
+                    and data.get("chosen", None) != 0
+                ):
+                    self._client_send({"action": "choose", "choice": 0})
+
+    def _handle_room_operation(self, data: dict):
+        pass
+
+    def _generate_word(self, prompt: str, max_length: int, suggestions: list[str]) -> str:
+        logger.info("Generating word...")
+        formatted_prompt = _WORD_PROMPT_TEMPLATE.format(
+            personality=self._personality,
+            prompt=prompt,
+            max_length=max_length,
+            suggestions=", ".join(suggestions),
+        )
+
+        word = self._chat_model.generate_text(
+            formatted_prompt,
+            "",
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return word
+
+    def _generate_rhyme(self, prompt: str, max_length: int, rhyme_word: str) -> str:
+        logger.info("Generating rhyme...")
+        formatted_prompt = _RHYME_PROMPT_TEMPLATE.format(
+            personality=self._personality,
+            prompt=prompt,
+            max_length=max_length,
+            rhyme_word=rhyme_word,
+        )
+        rhyme = self._chat_model.generate_text(
+            formatted_prompt,
+            "",
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return rhyme

ai_plays_jackbox/bot/jackbox5/patently_stupid.py

@@ -0,0 +1,167 @@
+import random
+from typing import Optional
+
+from loguru import logger
+
+from ai_plays_jackbox.bot.jackbox5.bot_base import JackBox5BotBase
+
+_ISSUE_PROMPT_PROMPT_TEMPLATE = """
+You are playing Patently Stupid. You need to fill in the given prompt.
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- You response must be 45 letters or less.
+- Do not include quotes in your response.
+
+Fill in the blank:
+
+{prompt}
+"""
+
+_SOLUTION_TITLE_PROMPT_TEMPLATE = """
+You are playing Patently Stupid. The issue you are trying to solve is "{issue}"
+
+I need you to generate a title for an invention that would solve the solution.
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- Your response must be 3 words or less.
+- Do not include quotes in your response.
+- Respond with only the title, nothing else. No newlines, etc.
+"""
+
+_SOLUTION_TAGLINE_PROMPT_TEMPLATE = """
+You are playing Patently Stupid. The issue you are trying to solve is "{issue}" and the invention that is going to solve it is called "{title}"
+
+I need you to generate a tagline for the invention.
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- Your response must be 30 characters or less.
+- Do not include quotes in your response.
+- Respond with only the tagline, nothing else. No newlines, etc.
+"""
+
+_SOLUTION_IMAGE_PROMPT_TEMPLATE = """
+You are playing Patently Stupid. The issue you are trying to solve is "{issue}" and the invention that is going to solve it is called "{title}"
+
+I need an drawing of this new invention.
+
+When generating your response, follow these rules:
+- The image must be a simple sketch
+- The image must have a white background and use black for the lines
+- Avoid intricate details
+"""
+
+
+class PatentlyStupidBot(JackBox5BotBase):
+    _issue_to_solve: Optional[str] = None
+    _solution_title: Optional[str] = None
+    _solution_tagline: Optional[str] = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _handle_welcome(self, data: dict):
+        pass
+
+    def _handle_player_operation(self, data: dict):
+        if not data:
+            return
+        room_state = data.get("state", None)
+        if not room_state:
+            return
+        entry = data.get("entry")
+        prompt: dict[str, str] = data.get("prompt", {})
+        prompt_html = prompt.get("html", "")
+        clean_prompt = self._html_to_text(prompt_html)
+        choices: list[dict] = data.get("choices", [])
+
+        if room_state == "EnterSingleText" and not bool(entry):
+            if "Write a title" in clean_prompt:
+                self._client_send({"action": "write", "entry": self._solution_title})
+            if "Write a tagline" in clean_prompt:
+                self._client_send({"action": "write", "entry": self._solution_tagline})
+            if "Fill in the Blank" in clean_prompt:
+                issue_fill_in = self._generate_issue(clean_prompt)
+                self._client_send({"action": "write", "entry": issue_fill_in})
+
+        elif room_state == "MakeSingleChoice":
+            if "Present your idea!" in clean_prompt:
+                done_text_html = data.get("doneText", {}).get("html", "")
+                if done_text_html == "":
+                    self._client_send({"action": "choose", "choice": 1})
+            elif "Invest in the best!" in clean_prompt:
+                filtered_choices = [c for c in choices if not c.get("disabled", True)]
+                if filtered_choices:
+                    choice = random.choice(filtered_choices)
+                    self._client_send({"action": "choose", "choice": choice.get("index", 0)})
+            elif "Press to skip this presentation." in clean_prompt:
+                pass
+            else:
+                choice_indexes = [i for i in range(0, len(choices))]
+                choice = random.choice(choice_indexes)  # type: ignore
+                self._client_send({"action": "choose", "choice": choice})
+
+        elif room_state == "Draw":
+            logger.info(data)
+            self._issue_to_solve = self._html_to_text(data["popup"]["html"])
+            self._solution_title = self._generate_title()
+            self._solution_tagline = self._generate_tagline()
+
+            lines = self._generate_drawing()
+            self._client_send(
+                {
+                    "action": "submit",
+                    "lines": [{"color": "#000000", "thickness": 6, "points": l} for l in lines],
+                },
+            )
+
+    def _handle_room_operation(self, data: dict):
+        pass
+
+    def _generate_issue(self, prompt: str) -> str:
+        formatted_prompt = _ISSUE_PROMPT_PROMPT_TEMPLATE.format(personality=self._personality, prompt=prompt)
+        issue = self._chat_model.generate_text(
+            formatted_prompt,
+            "",
+            max_tokens=10,
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return issue
+
+    def _generate_title(self) -> str:
+        prompt = _SOLUTION_TITLE_PROMPT_TEMPLATE.format(personality=self._personality, issue=self._issue_to_solve)
+        title = self._chat_model.generate_text(
+            prompt,
+            "",
+            max_tokens=2,
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return title
+
+    def _generate_tagline(self) -> str:
+        prompt = _SOLUTION_TAGLINE_PROMPT_TEMPLATE.format(
+            personality=self._personality, issue=self._issue_to_solve, title=self._solution_title
+        )
+        tagline = self._chat_model.generate_text(
+            prompt,
+            "",
+            max_tokens=12,
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return tagline
+
+    def _generate_drawing(self) -> list[str]:
+        logger.info("Generating drawing...")
+        image_prompt = _SOLUTION_IMAGE_PROMPT_TEMPLATE.format(issue=self._issue_to_solve, title=self._solution_title)
+        image_bytes = self._chat_model.generate_sketch(
+            image_prompt,
+            "",
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return self._image_bytes_to_polylines(image_bytes, 475, 475)

ai_plays_jackbox/bot/jackbox6/bot_base.py

@@ -1,4 +1,4 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 
 from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
 

ai_plays_jackbox/bot/jackbox6/joke_boat.py

@@ -0,0 +1,105 @@
+import random
+
+from loguru import logger
+
+from ai_plays_jackbox.bot.jackbox6.bot_base import JackBox6BotBase
+
+_TOPIC_PROMPT_TEMPLATE = """
+You are playing Joke Boat.
+
+You are being asked to come up with a topic that is {placeholder}.
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- You response must be {max_length} characters or less
+- Your response should be a single word.
+- Do not include quotes in your response or any newlines, just the response itself
+"""
+
+_PUNCHLINE_INSTRUCTIONS_TEMPLATE = """
+You are playing Joke Boat. You need to fill in the given prompt with a punchline.
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- You response must be {max_length} characters or less
+- Do not include quotes in your response or any newlines, just the response itself
+"""
+
+
+class JokeBoatBot(JackBox6BotBase):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _handle_welcome(self, data: dict):
+        pass
+
+    def _handle_player_operation(self, data: dict):
+        if not data:
+            return
+        room_state = data.get("state", None)
+        if not room_state:
+            return
+
+        prompt = data.get("prompt", {})
+        prompt_html = prompt.get("html", "")
+        clean_prompt = self._html_to_text(prompt_html)
+
+        match room_state:
+            case "MakeSingleChoice":
+                choices: list[dict] = data.get("choices", [])
+                choice_type = data.get("choiceType", "")
+                choice_indexes = [i for i in range(0, len(choices))]
+                selected_choice = random.choice(choice_indexes)
+
+                if choice_type == "ChooseAuthorReady":
+                    selected_choice = 1
+                if choice_type == "Skip":
+                    selected_choice = 0
+
+                if data.get("chosen", None) is None:
+                    self._client_send({"action": "choose", "choice": selected_choice})
+
+            case "EnterSingleText":
+                if "Write as many topics as you can." in clean_prompt:
+                    placeholder = data.get("placeholder", "")
+                    max_length = data.get("maxLength", 42)
+                    topic = self._generate_topic(placeholder, max_length)
+                    self._client_send({"action": "write", "entry": topic})
+                if "Write your punchline" in clean_prompt or "Write the punchline to this joke" in clean_prompt:
+                    max_length = data.get("maxLength", 80)
+                    punchline = self._generate_punchline(clean_prompt, max_length)
+                    self._client_send({"action": "write", "entry": punchline})
+
+    def _handle_room_operation(self, data: dict):
+        pass
+
+    def _generate_topic(self, placeholder: str, max_length: int) -> str:
+        logger.info("Generating topic...")
+        formatted_prompt = _TOPIC_PROMPT_TEMPLATE.format(
+            personality=self._personality,
+            placeholder=placeholder,
+            max_length=max_length,
+        )
+
+        topic = self._chat_model.generate_text(
+            formatted_prompt,
+            "",
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return topic[:max_length]
+
+    def _generate_punchline(self, prompt: str, max_length: int) -> str:
+        logger.info("Generating punchline...")
+        formatted_instructions = _PUNCHLINE_INSTRUCTIONS_TEMPLATE.format(
+            personality=self._personality,
+            prompt=prompt,
+            max_length=max_length,
+        )
+        punchline = self._chat_model.generate_text(
+            prompt,
+            formatted_instructions,
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
+        )
+        return punchline[:max_length]

ai_plays_jackbox/bot/jackbox7/bot_base.py

@@ -1,4 +1,4 @@
-from abc import ABC, abstractmethod
+from abc import ABC
 
 from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
 

ai_plays_jackbox/bot/jackbox7/quiplash3.py

@@ -75,7 +75,11 @@ class Quiplash3Bot(JackBox7BotBase):
             max_tokens = 32
             instructions = _FINAL_QUIP_PROMPT_INSTRUCTIONS_TEMPLATE.format(personality=self._personality)
         quip = self._chat_model.generate_text(
-            prompt, instructions=instructions, max_tokens=max_tokens, temperature=0.7, top_p=0.7
+            prompt,
+            instructions,
+            max_tokens=max_tokens,
+            temperature=self._chat_model._chat_model_temperature,
+            top_p=self._chat_model._chat_model_top_p,
         )
         return quip
 
@@ -84,7 +88,7 @@ class Quiplash3Bot(JackBox7BotBase):
         instructions = _QUIP_CHOICE_PROMPT_INSTRUCTIONS_TEMPLATE.format(prompt=prompt)
         response = self._chat_model.generate_text(
             f"Vote for your favorite response. Your options are: {choices_str}",
-            instructions=instructions,
+            instructions,
             max_tokens=1,
         )
         try:
@@ -100,5 +104,5 @@ class Quiplash3Bot(JackBox7BotBase):
             return choosen_prompt - 1
 
     def _choose_random_favorite(self, choices: list[dict]) -> int:
-        choices = [i for i in range(0, len(choices))]
-        return random.choice(choices)
+        choices_as_ints = [i for i in range(0, len(choices))]
+        return random.choice(choices_as_ints)