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

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 → room/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

@@ -1,26 +1,23 @@
 from typing import Optional
 
 from ai_plays_jackbox.llm.chat_model_factory import ChatModelFactory
-from ai_plays_jackbox.room import JackBoxRoom
+from ai_plays_jackbox.room.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)

ai_plays_jackbox/scripts/lint.py

@@ -0,0 +1,18 @@
+import subprocess
+
+
+def run():
+    commands = [
+        ["autoflake", "--in-place", "--recursive", "--remove-all-unused-imports", "--verbose", "ai_plays_jackbox"],
+        ["isort", "--profile", "black", "--project=ai_plays_jackbox", "ai_plays_jackbox"],
+        ["black", "-l", "120", "ai_plays_jackbox"],
+        ["mypy", "ai_plays_jackbox"],
+    ]
+
+    for cmd in commands:
+        print(f"\n>>> Running: {' '.join(cmd)}")
+        subprocess.run(["poetry", "run"] + cmd, check=True)
+
+
+if __name__ == "__main__":
+    run()

ai_plays_jackbox/ui/main.py

@@ -0,0 +1,12 @@
+from nicegui import app, ui
+
+from ai_plays_jackbox.ui.startup import startup
+
+
+def main(reload: bool = False):
+    app.on_startup(startup)
+    ui.run(favicon="🤖", reload=reload)
+
+
+if __name__ in {"__main__", "__mp_main__"}:
+    main(True)

ai_plays_jackbox/ui/startup.py

@@ -0,0 +1,248 @@
+from multiprocessing import Process, Queue
+
+import psutil
+from loguru import logger
+from nicegui import app, ui
+
+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
+
+LOG_QUEUE: Queue = Queue()
+LOG_DISPLAY = None
+SELECT_ALL_BOT_VARIANTS = None
+BOT_VARIANT_CHECKBOX_STATES: dict = {}
+
+
+def _format_log(record):
+    thread_name = record["thread"].name
+    color = "red"
+    colored_name = f"<{color}>{thread_name:<12}</{color}>"
+
+    return (
+        f"<green>{record['time']:YYYY-MM-DD HH:mm:ss}</green> | "
+        f"<cyan>{record['level']:<8}</cyan> | "
+        f"{colored_name} | "
+        f"{record['message']}\n"
+    )
+
+
+def _build_log_display():
+    global LOG_DISPLAY
+    with ui.row().classes("w-full"):
+        LOG_DISPLAY = ui.log(max_lines=100).classes("h-64 overflow-auto bg-black text-white")
+    ui.timer(interval=0.5, callback=_poll_log_queue)
+
+
+def _poll_log_queue():
+    global LOG_DISPLAY
+    try:
+        while not LOG_QUEUE.empty():
+            log_msg = LOG_QUEUE.get_nowait()
+            LOG_DISPLAY.push(log_msg)
+    except Exception as e:
+        LOG_DISPLAY.push(f"[ERROR] Failed to read log: {e}")
+
+
+def _start(
+    room_code: str,
+    chat_model_provider: str,
+    chat_model_name: str,
+    num_of_bots: int,
+    bots_in_play: list[str],
+    temperature: float,
+    top_p: float,
+    log_queue: Queue,
+):
+    logger.add(lambda msg: log_queue.put(msg), format=_format_log)
+
+    try:
+        run(
+            room_code.strip().upper(),
+            chat_model_provider,
+            chat_model_name=chat_model_name,
+            num_of_bots=num_of_bots,
+            bots_in_play=bots_in_play,
+            chat_model_temperature=temperature,
+            chat_model_top_p=top_p,
+        )
+    except Exception as e:
+        logger.exception("Bot startup failed")
+
+
+def _is_game_process_alive():
+    game_pid = app.storage.general.get("game_pid", None)
+    is_game_alive = game_pid is not None and psutil.pid_exists(game_pid) and psutil.Process(game_pid).is_running()
+    if not is_game_alive:
+        app.storage.general["game_pid"] = None
+    return is_game_alive
+
+
+def _handle_start_click(
+    room_code: str,
+    chat_model_provider: str,
+    chat_model_name: str,
+    num_of_bots: int,
+    temperature: float,
+    top_p: float,
+):
+    global BOT_VARIANT_CHECKBOX_STATES
+
+    if not _is_game_process_alive():
+        logger.info("Starting...")
+        game_thread = Process(
+            target=_start,
+            args=(
+                room_code,
+                chat_model_provider,
+                chat_model_name,
+                num_of_bots,
+                [k for k, v in BOT_VARIANT_CHECKBOX_STATES.items() if v.value],
+                temperature,
+                top_p,
+                LOG_QUEUE,
+            ),
+            daemon=True,
+        )
+        game_thread.start()
+        app.storage.general["game_pid"] = game_thread.pid
+
+
+def _select_all_bot_variants_changed():
+    for checkbox in BOT_VARIANT_CHECKBOX_STATES.values():
+        checkbox.value = SELECT_ALL_BOT_VARIANTS.value
+
+
+def _sync_select_all_bot_variants():
+    all_checked = all(cb.value for cb in BOT_VARIANT_CHECKBOX_STATES.values())
+    SELECT_ALL_BOT_VARIANTS.value = all_checked
+
+
+def _setup_bot_variant_display():
+    global SELECT_ALL_BOT_VARIANTS
+    with ui.list().props("bordered separator").classes("w-full"):
+        with ui.item_label("Bot Personalities").props("header").classes("text-bold"):
+            SELECT_ALL_BOT_VARIANTS = ui.checkbox(text="Select All", value=True)
+            SELECT_ALL_BOT_VARIANTS.on("update:model-value", lambda e: _select_all_bot_variants_changed())
+        ui.separator()
+        with ui.element("div").classes("overflow-y-auto h-64"):
+            for variant in list(JackBoxBotVariant):
+                with ui.item():
+                    with ui.item_section().props("avatar"):
+                        cb = ui.checkbox(value=True)
+                        cb.on("update:model-value", lambda e: _sync_select_all_bot_variants())
+                        BOT_VARIANT_CHECKBOX_STATES[variant.name] = cb
+                    with ui.item_section():
+                        ui.item_label(variant.value.name)
+                        ui.item_label(variant.value.personality).props("caption")
+
+
+def startup():
+    ui.page_title("AI Plays JackBox")
+    ui.label("🤖 AI Plays JackBox").classes("text-2xl font-bold")
+
+    _build_log_display()
+
+    with ui.grid(columns=16).classes("w-full gap-0"):
+        with ui.column().classes("col-span-1"):
+            pass
+        with ui.column().classes("col-span-7"):
+            with ui.row():
+                ui.label("Number of Bots")
+                num_of_bots_label = ui.label(str(DEFAULT_NUM_OF_BOTS))
+                num_of_bots = ui.slider(
+                    min=1,
+                    max=10,
+                    value=DEFAULT_NUM_OF_BOTS,
+                    step=1,
+                    on_change=lambda e: num_of_bots_label.set_text(f"{e.value}"),
+                )
+                chat_model_provider = ui.select(
+                    list(CHAT_MODEL_PROVIDERS.keys()),
+                    label="Chat Model Provider",
+                    value=list(CHAT_MODEL_PROVIDERS.keys())[0],
+                    on_change=lambda e: chat_model_name.set_value(CHAT_MODEL_PROVIDERS[e.value].get_default_model()),
+                ).classes("w-1/3")
+
+                chat_model_name = ui.input(
+                    label="Chat Model Name",
+                    value=CHAT_MODEL_PROVIDERS[chat_model_provider.value].get_default_model(),
+                ).classes("w-1/3")
+
+                room_code = (
+                    ui.input(
+                        label="Room Code",
+                        placeholder="ABCD",
+                        validation={
+                            "must be letters only": lambda value: value.isalpha(),
+                            "must be 4 letters": lambda value: len(value) == 4,
+                        },
+                    )
+                    .props("uppercase")
+                    .classes("w-1/4")
+                )
+                start_button = (
+                    ui.button(
+                        "Start Bots",
+                        on_click=lambda _: _handle_start_click(
+                            room_code.value,
+                            chat_model_provider.value,
+                            chat_model_name.value,
+                            num_of_bots.value,
+                            temperature.value,
+                            top_p.value,
+                        ),
+                    )
+                    .bind_enabled_from(room_code, "error", lambda error: room_code.value and not error)
+                    .classes("w-full")
+                )
+                ui.timer(
+                    interval=0.5,
+                    callback=lambda: start_button.props(
+                        f"color={'blue' if _is_game_process_alive() else 'green'}"
+                    ).set_text("Running..." if _is_game_process_alive() else "Start Bots"),
+                )
+
+                ui.label("Advanced Options").classes("w-full text-xl font-bold")
+
+                ui.label("Temperature").classes("w-1/4").tooltip(
+                    """
+                    What sampling temperature to use, between 0 and 2. Higher values like 0.8 will
+                    make the output more random, while lower values like 0.2 will make it more
+                    focused and deterministic. We generally recommend altering this or `top_p` but
+                    not both."""
+                )
+                temperature_label = ui.label(str(DEFAULT_TEMPERATURE)).classes("w-1/6")
+                temperature = ui.slider(
+                    min=0.0,
+                    max=2.0,
+                    value=DEFAULT_TEMPERATURE,
+                    step=0.1,
+                    on_change=lambda e: temperature_label.set_text(f"{e.value}"),
+                ).classes("w-1/2")
+
+                ui.label("Top P").classes("w-1/4").tooltip(
+                    """
+                    An alternative to sampling with temperature, called nucleus sampling, where the
+                    model considers the results of the tokens with top_p probability mass. So 0.1
+                    means only the tokens comprising the top 10% probability mass are considered."""
+                )
+                top_p_label = ui.label(str(DEFAULT_TOP_P)).classes("w-1/6")
+                top_p = ui.slider(
+                    min=0.0,
+                    max=1.0,
+                    value=DEFAULT_TOP_P,
+                    step=0.1,
+                    on_change=lambda e: top_p_label.set_text(f"{e.value}"),
+                ).classes("w-1/2")
+
+        with ui.column().classes("col-span-1"):
+            pass
+
+        with ui.column().classes("col-span-6"):
+            _setup_bot_variant_display()