ai-plays-jackbox 0.4.1__py3-none-any.whl

ai_plays_jackbox/bot/standalone/drawful2.py

@@ -0,0 +1,159 @@
+import random
+
+from loguru import logger
+
+from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
+
+_DRAWING_PROMPT_TEMPLATE = """
+You are playing Drawful 2.
+
+Generate an image with the following prompt: {prompt}
+
+When generating your response, follow these rules:
+- Your personality is: {personality}
+- Make sure to implement your personality somehow into the drawing, but keep the prompt in mind
+- The image must be a simple sketch
+- The image must have a white background and use black for the lines
+- Avoid intricate details
+"""
+
+
+class Drawful2Bot(JackBoxBotBase):
+    _drawing_completed: bool = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def _player_operation_key(self) -> str:
+        return f"player:{self._player_id}"
+
+    def _is_player_operation_key(self, operation_key: str) -> bool:
+        return operation_key == self._player_operation_key
+
+    @property
+    def _room_operation_key(self) -> str:
+        return "room"
+
+    def _is_room_operation_key(self, operation_key: str) -> bool:
+        return operation_key == self._room_operation_key
+
+    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_text = self._html_to_text(prompt.get("html", "")) if prompt is not None else ""
+
+        match room_state:
+            case "Draw":
+                colors = data.get("colors", ["#fb405a", "#7a2259"])
+                selected_color = colors[0]
+                canvas_height = int(data.get("size", {}).get("height", 320))
+                canvas_width = int(data.get("size", {}).get("width", 320))
+                lines = self._generate_drawing(prompt_text, canvas_height, canvas_width)
+                object_key = data.get("objectKey", "")
+                if object_key != "":
+                    if not self._drawing_completed:
+                        self._object_update(
+                            object_key,
+                            {
+                                "lines": [{"color": selected_color, "thickness": 1, "points": l} for l in lines],
+                                "submit": True,
+                            },
+                        )
+                        # This prevents the bot from trying to draw multiple times
+                        self._drawing_completed = True
+
+            case "EnterSingleText":
+                # We need to reset this once we're entering options
+                self._drawing_completed = False
+                # Listen, the bot can't see the drawing
+                # so they're just going to say something
+                text_key = data.get("textKey", "")
+                self._text_update(text_key, self._generate_random_response())
+
+            case "MakeSingleChoice":
+                # Bot still can't see the drawing
+                # so just pick something
+                if data.get("type", "single") == "repeating":
+                    pass
+                choices = data.get("choices", [])
+                choices_as_ints = [i for i in range(0, len(choices))]
+                selected_choice = random.choice(choices_as_ints)
+                self._client_send({"action": "choose", "choice": selected_choice})
+
+    def _handle_room_operation(self, data: dict):
+        pass
+
+    def _generate_drawing(self, prompt: str, canvas_height: int, canvas_width: int) -> list[str]:
+        logger.info("Generating drawing...")
+        image_prompt = _DRAWING_PROMPT_TEMPLATE.format(prompt=prompt, personality=self._personality)
+        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, canvas_height, canvas_width)
+
+    def _generate_random_response(self) -> str:
+        possible_responses = [
+            "Abstract awkward silence",
+            "Abstract existential dread",
+            "Abstract late-stage capitalism",
+            "Abstract lost hope",
+            "Abstract the void",
+            "Baby Yoda, but weird",
+            "Barbie, but weird",
+            "Confused dentist",
+            "Confused gym teacher",
+            "Confused lawyer",
+            "Confused therapist",
+            "DJ in trouble",
+            "Definitely banana",
+            "Definitely blob",
+            "Definitely potato",
+            "Definitely spaghetti",
+            "Taylor Swift, but weird",
+            "Waluigi, but weird",
+            "banana with feelings",
+            "chicken riding a scooter",
+            "cloud with feelings",
+            "confused gym teacher",
+            "confused therapist",
+            "dentist in trouble",
+            "duck + existential dread",
+            "duck riding a scooter",
+            "excited janitor",
+            "ferret + awkward silence",
+            "giraffe + awkward silence",
+            "giraffe filing taxes",
+            "hamster + existential dread",
+            "hamster + lost hope",
+            "hamster riding a scooter",
+            "janitor in trouble",
+            "joyful octopus",
+            "lawyer in trouble",
+            "llama + awkward silence",
+            "llama + late-stage capitalism",
+            "lonely dentist",
+            "lonely hamster",
+            "lonely janitor",
+            "lonely pirate",
+            "mango with feelings",
+            "pirate in trouble",
+            "sad DJ",
+            "sad hamster",
+            "spaghetti with feelings",
+            "terrified duck",
+            "terrified ferret",
+            "terrified lawyer",
+        ]
+        chosen_response = random.choice(possible_responses)
+        return chosen_response

ai_plays_jackbox/cli/main.py

@@ -0,0 +1,117 @@
+import argparse
+
+from ai_plays_jackbox.bot.bot_personality import JackBoxBotVariant
+from ai_plays_jackbox.constants import (
+    DEFAULT_NUM_OF_BOTS,
+    DEFAULT_TEMPERATURE,
+    DEFAULT_TOP_P,
+)
+from ai_plays_jackbox.llm.chat_model_factory import CHAT_MODEL_PROVIDERS
+from ai_plays_jackbox.run import run
+
+
+def _validate_room_code(string_to_check: str) -> str:
+    if not string_to_check.isalpha() or len(string_to_check) != 4:
+        raise argparse.ArgumentTypeError("Must be 4 letters")
+    return string_to_check
+
+
+def _validate_num_of_bots(string_to_check: str) -> int:
+    try:
+        number_value = int(string_to_check)
+    except ValueError:
+        raise argparse.ArgumentTypeError("Must number 1-10")
+    if number_value < 1 or number_value > 10:
+        raise argparse.ArgumentTypeError("Must number 1-10")
+    return number_value
+
+
+def _validate_temperature(string_to_check: str) -> float:
+    try:
+        number_value = float(string_to_check)
+    except ValueError:
+        raise argparse.ArgumentTypeError("Must number 0.1-2.0")
+    if number_value <= 0 or number_value > 2.0:
+        raise argparse.ArgumentTypeError("Must number 0.1-2.0")
+    return number_value
+
+
+def _validate_top_p(string_to_check: str) -> float:
+    try:
+        number_value = float(string_to_check)
+    except ValueError:
+        raise argparse.ArgumentTypeError("Must number 0.0-1.0")
+    if number_value < 0 or number_value > 1.0:
+        raise argparse.ArgumentTypeError("Must number 0.0-1.0")
+    return number_value
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--room-code",
+        required=True,
+        help="The JackBox room code",
+        type=_validate_room_code,
+        metavar="WXYZ",
+    )
+    parser.add_argument(
+        "--chat-model-provider",
+        required=True,
+        help="Choose which chat model platform to use",
+        choices=list(CHAT_MODEL_PROVIDERS.keys()),
+        type=str,
+    )
+    parser.add_argument(
+        "--chat-model-name",
+        required=False,
+        help="Choose which chat model to use (Will default to default for provider)",
+        type=str,
+    )
+    parser.add_argument(
+        "--num-of-bots",
+        required=False,
+        default=DEFAULT_NUM_OF_BOTS,
+        help="How many bots to have play",
+        type=_validate_num_of_bots,
+        metavar=str(DEFAULT_NUM_OF_BOTS),
+    )
+    parser.add_argument(
+        "--bots-in-play",
+        required=False,
+        nargs="*",
+        help="Which bots are in play?",
+        choices=[variant.name for variant in JackBoxBotVariant],
+        type=str,
+    )
+    parser.add_argument(
+        "--temperature",
+        required=False,
+        default=DEFAULT_TEMPERATURE,
+        help="Temperature for Gen AI",
+        type=_validate_temperature,
+        metavar=str(DEFAULT_TEMPERATURE),
+    )
+    parser.add_argument(
+        "--top-p",
+        required=False,
+        default=DEFAULT_TOP_P,
+        help="Top P for Gen AI",
+        type=_validate_top_p,
+        metavar=str(DEFAULT_TOP_P),
+    )
+    args = parser.parse_args()
+
+    run(
+        args.room_code.upper(),
+        args.chat_model_provider,
+        chat_model_name=args.chat_model_name,
+        num_of_bots=args.num_of_bots,
+        bots_in_play=args.bots_in_play,
+        chat_model_temperature=args.temperature,
+        chat_model_top_p=args.top_p,
+    )
+
+
+if __name__ == "__main__":
+    main()

ai_plays_jackbox/constants.py

@@ -0,0 +1,4 @@
+ECAST_HOST = "ecast.jackboxgames.com"
+DEFAULT_TEMPERATURE = 0.5
+DEFAULT_TOP_P = 0.9
+DEFAULT_NUM_OF_BOTS = 4

ai_plays_jackbox/llm/__init__.py

@@ -0,0 +1 @@
+

ai_plays_jackbox/llm/chat_model.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from typing import Optional
+
+
+class ChatModel(ABC):
+    _model: str
+    _chat_model_temperature: float
+    _chat_model_top_p: float
+
+    def __init__(self, model: str, chat_model_temperature: float = 0.5, chat_model_top_p: float = 0.9):
+        self._model = model
+        self._chat_model_temperature = chat_model_temperature
+        self._chat_model_top_p = chat_model_top_p
+
+    @classmethod
+    @abstractmethod
+    def get_default_model(self) -> str:
+        pass
+
+    @abstractmethod
+    def generate_text(
+        self,
+        prompt: str,
+        instructions: str,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> str:
+        pass
+
+    @abstractmethod
+    def generate_sketch(
+        self,
+        prompt: str,
+        instructions: str,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> bytes:
+        pass

ai_plays_jackbox/llm/chat_model_factory.py

@@ -0,0 +1,35 @@
+from typing import Optional
+
+from ai_plays_jackbox.llm.chat_model import ChatModel
+from ai_plays_jackbox.llm.gemini_model import GeminiModel
+from ai_plays_jackbox.llm.ollama_model import OllamaModel
+from ai_plays_jackbox.llm.openai_model import OpenAIModel
+
+CHAT_MODEL_PROVIDERS: dict[str, type[ChatModel]] = {
+    "openai": OpenAIModel,
+    "gemini": GeminiModel,
+    "ollama": OllamaModel,
+}
+
+
+class ChatModelFactory:
+    @staticmethod
+    def get_chat_model(
+        chat_model_provider: str,
+        chat_model_name: Optional[str] = None,
+        chat_model_temperature: float = 0.5,
+        chat_model_top_p: float = 0.9,
+    ) -> ChatModel:
+        chat_model_provider = chat_model_provider.lower()
+        if chat_model_provider not in CHAT_MODEL_PROVIDERS.keys():
+            raise ValueError(f"Unknown chat model provider: {chat_model_provider}")
+
+        return CHAT_MODEL_PROVIDERS[chat_model_provider](
+            (
+                chat_model_name
+                if chat_model_name is not None
+                else CHAT_MODEL_PROVIDERS[chat_model_provider].get_default_model()
+            ),
+            chat_model_temperature=chat_model_temperature,
+            chat_model_top_p=chat_model_top_p,
+        )

ai_plays_jackbox/llm/gemini_model.py

@@ -0,0 +1,86 @@
+import os
+from typing import Optional
+
+from google import genai
+from google.genai.types import GenerateContentConfig
+from loguru import logger
+
+from ai_plays_jackbox.llm.chat_model import ChatModel
+
+
+class GeminiModel(ChatModel):
+    _gemini_vertex_ai_client: genai.Client
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._gemini_vertex_ai_client = genai.Client(
+            vertexai=bool(os.environ.get("GOOGLE_GENAI_USE_VERTEXAI")),
+            api_key=os.environ.get("GOOGLE_GEMINI_DEVELOPER_API_KEY"),
+            project=os.environ.get("GOOGLE_CLOUD_PROJECT"),
+            location=os.environ.get("GOOGLE_CLOUD_LOCATION"),
+        )
+
+        # Check connection and if model exists, this will hard fail if connection can't be made
+        # Or if the model is not found
+        _ = self._gemini_vertex_ai_client.models.get(model=self._model)
+
+    @classmethod
+    def get_default_model(cls):
+        return "gemini-2.0-flash-001"
+
+    def generate_text(
+        self,
+        prompt: str,
+        instructions: str,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> str:
+        if temperature is None:
+            temperature = self._chat_model_temperature
+        if top_p is None:
+            top_p = self._chat_model_top_p
+
+        chat_response = self._gemini_vertex_ai_client.models.generate_content(
+            model=self._model,
+            contents=prompt,
+            config=GenerateContentConfig(
+                system_instruction=[instructions],
+                max_output_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+            ),
+        )
+
+        text = str(chat_response.text).strip().replace("\n", "")
+        logger.info(f"Generated text: {text}")
+        return text
+
+    def generate_sketch(
+        self,
+        prompt: str,
+        instructions: str,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> bytes:
+        image_gen_response = self._gemini_vertex_ai_client.models.generate_content(
+            model="gemini-2.0-flash-preview-image-generation",
+            contents=prompt,
+            config=GenerateContentConfig(
+                system_instruction=[instructions],
+                temperature=temperature,
+                top_p=top_p,
+                response_modalities=["IMAGE"],
+            ),
+        )
+
+        if (
+            image_gen_response.candidates
+            and image_gen_response.candidates[0].content
+            and image_gen_response.candidates[0].content.parts
+        ):
+            for part in image_gen_response.candidates[0].content.parts:
+                if part.inline_data is not None and part.inline_data.data is not None:
+                    return part.inline_data.data
+
+        return b""

ai_plays_jackbox/llm/ollama_model.py

@@ -0,0 +1,53 @@
+from typing import Optional
+
+from loguru import logger
+from ollama import Options, chat, show
+
+from ai_plays_jackbox.llm.chat_model import ChatModel
+
+
+class OllamaModel(ChatModel):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # Check connection and if model exists, this will hard fail if connection can't be made
+        # Or if the model is not found
+        _ = show(self._model)
+
+    @classmethod
+    def get_default_model(cls):
+        return "gemma3:12b"
+
+    def generate_text(
+        self,
+        prompt: str,
+        instructions: str,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> str:
+        if temperature is None:
+            temperature = self._chat_model_temperature
+        if top_p is None:
+            top_p = self._chat_model_top_p
+
+        instructions_formatted = {"role": "system", "content": instructions}
+        chat_response = chat(
+            model=self._model,
+            messages=[instructions_formatted, {"role": "user", "content": prompt}],
+            stream=False,
+            options=Options(num_predict=max_tokens, temperature=temperature, top_p=top_p),
+        )
+        text = str(chat_response.message.content).strip().replace("\n", " ")
+        logger.info(f"Generated text: {text}")
+        return text
+
+    def generate_sketch(
+        self,
+        prompt: str,
+        instructions: str,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> bytes:
+        raise Exception("Ollama model not supported yet for sketches")

ai_plays_jackbox/llm/openai_model.py

@@ -0,0 +1,86 @@
+import base64
+import os
+from typing import Optional
+
+from loguru import logger
+from openai import OpenAI
+from openai.types.chat import (
+    ChatCompletionDeveloperMessageParam,
+    ChatCompletionUserMessageParam,
+)
+from openai.types.responses import Response
+
+from ai_plays_jackbox.llm.chat_model import ChatModel
+
+
+class OpenAIModel(ChatModel):
+    _open_ai_client: OpenAI
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._open_ai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+
+        # Check connection and if model exists, this will hard fail if connection can't be made
+        # Or if the model is not found
+        _ = self._open_ai_client.models.retrieve(self._model)
+
+    @classmethod
+    def get_default_model(cls):
+        return "gpt-4o-mini"
+
+    def generate_text(
+        self,
+        prompt: str,
+        instructions: str,
+        max_tokens: Optional[int] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> str:
+        if temperature is None:
+            temperature = self._chat_model_temperature
+        if top_p is None:
+            top_p = self._chat_model_top_p
+
+        chat_response = self._open_ai_client.chat.completions.create(
+            model=self._model,
+            messages=[
+                ChatCompletionDeveloperMessageParam(content=instructions, role="developer"),
+                ChatCompletionUserMessageParam(content=prompt, role="user"),
+            ],
+            stream=False,
+            max_completion_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+        )
+        text = str(chat_response.choices[0].message.content).strip().replace("\n", "")
+        logger.info(f"Generated text: {text}")
+        return text
+
+    def generate_sketch(
+        self,
+        prompt: str,
+        instructions: str,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+    ) -> bytes:
+        image_gen_response: Response = self._open_ai_client.responses.create(
+            model=self._model,
+            instructions=instructions,
+            input=prompt,
+            temperature=temperature,
+            top_p=top_p,
+            tools=[
+                {
+                    "type": "image_generation",
+                    "quality": "low",
+                    "size": "1024x1024",
+                }
+            ],
+        )
+        # Save the image to a file
+        image_data = [output.result for output in image_gen_response.output if output.type == "image_generation_call"]
+        image_base64 = ""
+        if image_data:
+            image_base64 = str(image_data[0])
+
+        return base64.b64decode(image_base64)

ai_plays_jackbox/room/room.py

@@ -0,0 +1,87 @@
+import random
+import threading
+from time import sleep
+from typing import Optional
+
+import requests
+from loguru import logger
+
+from ai_plays_jackbox.bot.bot_base import JackBoxBotBase
+from ai_plays_jackbox.bot.bot_factory import JackBoxBotFactory
+from ai_plays_jackbox.bot.bot_personality import JackBoxBotVariant
+from ai_plays_jackbox.constants import ECAST_HOST
+from ai_plays_jackbox.llm.chat_model import ChatModel
+
+
+class JackBoxRoom:
+    _bots: list[JackBoxBotBase]
+
+    def __init__(self):
+        self._bots = []
+        self._lock = threading.Lock()
+
+    def play(
+        self,
+        room_code: str,
+        chat_model: ChatModel,
+        num_of_bots: int = 4,
+        bots_in_play: Optional[list] = None,
+    ):
+        room_type = self._get_room_type(room_code)
+        if not room_type:
+            logger.error(f"Unable to find room {room_code}")
+            return
+        logger.info(f"We're playing {room_type}!")
+        bot_factory = JackBoxBotFactory()
+        if bots_in_play is None or len(bots_in_play) == 0:
+            bots_to_make = random.sample(list(JackBoxBotVariant), num_of_bots)
+        else:
+            bots_in_play_variants = [variant for variant in JackBoxBotVariant if variant.name in bots_in_play]
+            bots_to_make = random.choices(bots_in_play_variants, k=num_of_bots)
+
+        for b in bots_to_make:
+            bot = bot_factory.get_bot(
+                room_type,
+                chat_model,
+                name=b.value.name,
+                personality=b.value.personality,
+            )
+            self._bots.append(bot)
+            with self._lock:
+                bot.connect(room_code)
+            sleep(0.5)
+
+        try:
+            while True:
+                sleep(1)
+                if self.is_finished():
+                    print("All bots disconnected, ending...")
+                    break
+        except KeyboardInterrupt:
+            self.end()
+
+    def is_finished(self) -> bool:
+        for b in self._bots:
+            if not b.is_disconnected():
+                return False
+        return True
+
+    def end(self):
+        for b in self._bots:
+            with self._lock:
+                b.disconnect()
+
+    def _get_room_type(self, room_code: str) -> str:
+        try:
+            response = requests.request(
+                "GET",
+                f"https://{ECAST_HOST}/api/v2/rooms/{room_code}",
+                headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; rv:68.0) Gecko/20100101 Firefox/68.0"},
+            )
+            response.raise_for_status()
+            response_data = response.json()
+        except requests.HTTPError:
+            if response.status_code != 404:
+                logger.error(f"Received {response.status_code} when trying to get room type.")
+            return ""
+        return response_data["body"]["appTag"]