ai-plays-jackbox 0.0.1__py3-none-any.whl → 0.1.0__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.py

@@ -1,6 +1,12 @@
 import argparse
 
-from ai_plays_jackbox import run
+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:
@@ -49,42 +55,49 @@ def cli():
         metavar="WXYZ",
     )
     parser.add_argument(
-        "--chat-model-name",
+        "--chat-model-provider",
         required=True,
         help="Choose which chat model platform to use",
-        choices=("ollama", "openai", "gemini"),
+        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=4,
-        help="How many bots to have play (Defaults to 4)",
+        default=DEFAULT_NUM_OF_BOTS,
+        help="How many bots to have play",
         type=_validate_num_of_bots,
-        metavar="4",
+        metavar=str(DEFAULT_NUM_OF_BOTS),
     )
     parser.add_argument(
         "--temperature",
         required=False,
-        default=0.5,
-        help="Temperature for Gen AI (Defaults to 0.5)",
+        default=DEFAULT_TEMPERATURE,
+        help="Temperature for Gen AI",
         type=_validate_temperature,
-        metavar="0.5",
+        metavar=str(DEFAULT_TEMPERATURE),
     )
     parser.add_argument(
         "--top-p",
         required=False,
-        default=0.9,
-        help="Top P for Gen AI (Defaults to 0.9)",
+        default=DEFAULT_TOP_P,
+        help="Top P for Gen AI",
         type=_validate_top_p,
-        metavar="0.9",
+        metavar=str(DEFAULT_TOP_P),
     )
     args = parser.parse_args()
 
     run(
         args.room_code.upper(),
-        num_of_bots=args.num_of_bots,
+        args.chat_model_provider,
         chat_model_name=args.chat_model_name,
+        num_of_bots=args.num_of_bots,
         chat_model_temperature=args.temperature,
         chat_model_top_p=args.top_p,
     )

ai_plays_jackbox/constants.py

@@ -1 +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/chat_model.py

@@ -3,13 +3,20 @@ from typing import Optional
 
 
 class ChatModel(ABC):
-    _chat_model_temperature: float = 0.5
-    _chat_model_top_p: float = 0.9
+    _model: str
+    _chat_model_temperature: float
+    _chat_model_top_p: float
 
-    def __init__(self, chat_model_temperature: float = 0.5, chat_model_top_p: float = 0.9):
+    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,
@@ -20,3 +27,13 @@ class ChatModel(ABC):
         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

@@ -1,30 +1,35 @@
+from typing import Optional
+
 from ai_plays_jackbox.llm.chat_model import ChatModel
-from ai_plays_jackbox.llm.gemini_vertex_ai import GeminiVertextAIModel
+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_name: str,
+        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_name = chat_model_name.lower()
-        if chat_model_name == "ollama":
-            return OllamaModel(
-                model="gemma3:12b", chat_model_temperature=chat_model_temperature, chat_model_top_p=chat_model_top_p
-            )
-        if chat_model_name == "openai":
-            return OpenAIModel(
-                model="gpt-4o-mini", chat_model_temperature=chat_model_temperature, chat_model_top_p=chat_model_top_p
-            )
-        if chat_model_name == "gemini":
-            return GeminiVertextAIModel(
-                model="gemini-2.0-flash-001",
-                chat_model_temperature=chat_model_temperature,
-                chat_model_top_p=chat_model_top_p,
-            )
-        else:
-            raise ValueError(f"Unknown chat model type: {chat_model_name}")
+        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(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

@@ -1,20 +1,23 @@
 from typing import Optional
 
 from loguru import logger
-from ollama import Options, chat, list
+from ollama import Options, chat, show
 
 from ai_plays_jackbox.llm.chat_model import ChatModel
 
 
 class OllamaModel(ChatModel):
-    _model: str
 
-    def __init__(self, model: str = "gemma3:12b", *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._model = model
 
-        # Check connection, this will hard fail if connection can't be made
-        _ = list()
+        # 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,
@@ -36,6 +39,15 @@ class OllamaModel(ChatModel):
             stream=False,
             options=Options(num_predict=max_tokens, temperature=temperature, top_p=top_p),
         )
-        text = chat_response.message.content.strip().replace("\n", " ")
+        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

@@ -1,23 +1,32 @@
+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):
-    _model: str
     _open_ai_client: OpenAI
 
-    def __init__(self, model: str = "gpt-4o-mini", *args, **kwargs):
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._open_ai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
-        self._model = model
 
-        # Check connection, this will hard fail if connection can't be made
-        _ = self._open_ai_client.models.list()
+        # 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,
@@ -32,15 +41,46 @@ class OpenAIModel(ChatModel):
         if top_p is None:
             top_p = self._chat_model_top_p
 
-        instructions_formatted = {"role": "developer", "content": instructions}
         chat_response = self._open_ai_client.chat.completions.create(
             model=self._model,
-            messages=[instructions_formatted, {"role": "user", "content": prompt}],
+            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 = chat_response.choices[0].message.content.strip().replace("\n", "")
+        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.py

@@ -11,7 +11,6 @@ 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
-from ai_plays_jackbox.llm.ollama_model import OllamaModel
 
 
 class JackBoxRoom:
@@ -24,12 +23,10 @@ class JackBoxRoom:
     def play(
         self,
         room_code: str,
+        chat_model: ChatModel,
         num_of_bots: int = 4,
         bots_in_play: Optional[list] = None,
-        chat_model: Optional[ChatModel] = None,
     ):
-        if chat_model is None:
-            chat_model = OllamaModel()
         room_type = self._get_room_type(room_code)
         if not room_type:
             logger.error(f"Unable to find room {room_code}")
@@ -45,9 +42,9 @@ class JackBoxRoom:
         for b in bots_to_make:
             bot = bot_factory.get_bot(
                 room_type,
+                chat_model,
                 name=b.value.name,
                 personality=b.value.personality,
-                chat_model=chat_model,
             )
             self._bots.append(bot)
             with self._lock:

ai_plays_jackbox/run.py

@@ -6,21 +6,18 @@ from ai_plays_jackbox.room import JackBoxRoom
 
 def run(
     room_code: str,
+    chat_model_provider: str,
+    chat_model_name: Optional[str] = None,
     num_of_bots: int = 4,
     bots_in_play: Optional[list] = None,
-    chat_model_name: str = "ollama",
     chat_model_temperature: float = 0.5,
     chat_model_top_p: float = 0.9,
 ):
-    """Will run a set of bots through a game of JackBox given a room code.
-
-    Args:
-        room_code (str): The room code.
-        num_of_bots (int, optional): The number of bots to participate. Defaults to 4.
-        chat_model (str, optional): The chat model to use to generate responses. Defaults to "ollama".
-    """
     chat_model = ChatModelFactory.get_chat_model(
-        chat_model_name, chat_model_temperature=chat_model_temperature, chat_model_top_p=chat_model_top_p
+        chat_model_provider,
+        chat_model_name=chat_model_name,
+        chat_model_temperature=chat_model_temperature,
+        chat_model_top_p=chat_model_top_p,
     )
     room = JackBoxRoom()
-    room.play(room_code, num_of_bots=num_of_bots, bots_in_play=bots_in_play, chat_model=chat_model)
+    room.play(room_code, chat_model, num_of_bots=num_of_bots, bots_in_play=bots_in_play)