feat: implement Google API key management and refactor client usage in audio and image generation

src/agent/llm.py

@@ -4,31 +4,17 @@ import logging
 from langchain_google_genai import ChatGoogleGenerativeAI
 
 from config import settings
+from services.google import ApiKeyPool
 
 logger = logging.getLogger(__name__)
 
-_API_KEYS: list[str] = []
-_current_key_idx = 0
+_pool = ApiKeyPool()
 MODEL_NAME = "gemini-2.5-flash-preview-05-20"
 
 
 def _get_api_key() -> str:
-    """Return an API key using round-robin selection."""
-    global _API_KEYS, _current_key_idx
-
-    if not _API_KEYS:
-        keys_str = settings.gemini_api_key.get_secret_value()
-        if keys_str:
-            _API_KEYS = [k.strip() for k in keys_str.split(",") if k.strip()]
-        if not _API_KEYS:
-            msg = "Google API keys are not configured or invalid"
-            logger.error(msg)
-            raise ValueError(msg)
-
-    key = _API_KEYS[_current_key_idx]
-    _current_key_idx = (_current_key_idx + 1) % len(_API_KEYS)
-    logger.debug("Using Google API key index %s", _current_key_idx)
-    return key
+    """Return an API key using round-robin selection in a thread-safe way."""
+    return _pool.get_key_sync()
 
 
 def create_llm(

src/agent/redis_state.py

@@ -5,9 +5,12 @@ from __future__ import annotations
 import json
 import msgpack
 import redis.asyncio as redis
+import logging
 
 from agent.models import UserState
 
+logger = logging.getLogger(__name__)
+
 
 class UserRepository:
     """Repository for storing UserState objects in Redis."""
@@ -18,6 +21,7 @@ class UserRepository:
     async def get(self, user_id: str) -> UserState:
         """Return user state for the given id, creating it if absent."""
         key = f"llmgamehub:{user_id}"
+        logger.debug("Fetching state for %s", user_id)
         data = await self.redis.hget(key, "data")
         if data is None:
             return UserState()
@@ -27,12 +31,14 @@ class UserRepository:
     async def set(self, user_id: str, state: UserState) -> None:
         """Persist updated user state."""
         key = f"llmgamehub:{user_id}"
+        logger.debug("Saving state for %s", user_id)
         packed = msgpack.packb(json.loads(state.json()))
         await self.redis.hset(key, mapping={"data": packed})
 
     async def reset(self, user_id: str) -> None:
         """Remove stored state for a user."""
         key = f"llmgamehub:{user_id}"
+        logger.debug("Resetting state for %s", user_id)
         await self.redis.delete(key)
 
 
@@ -40,12 +46,15 @@ _repo = UserRepository()
 
 
 async def get_user_state(user_hash: str) -> UserState:
+    logger.debug("get_user_state for %s", user_hash)
     return await _repo.get(user_hash)
 
 
 async def set_user_state(user_hash: str, state: UserState) -> None:
+    logger.debug("set_user_state for %s", user_hash)
     await _repo.set(user_hash, state)
 
 
 async def reset_user_state(user_hash: str) -> None:
+    logger.debug("reset_user_state for %s", user_hash)
     await _repo.reset(user_hash)

src/audio/audio_generator.py

@@ -1,7 +1,5 @@
 import asyncio
-from google import genai
 from google.genai import types
-from config import settings
 import wave
 import queue
 import logging
@@ -10,34 +8,41 @@ import time
 
 logger = logging.getLogger(__name__)
 
-client = genai.Client(api_key=settings.gemini_api_key.get_secret_value(), http_options={'api_version': 'v1alpha'})
+from services.google import GoogleClientFactory
 
 async def generate_music(user_hash: str, music_tone: str, receive_audio):
     if user_hash in sessions:
         logger.info(f"Music generation already started for user hash {user_hash}, skipping new generation")
         return
-    async with (
-        client.aio.live.music.connect(model='models/lyria-realtime-exp') as session,
-        asyncio.TaskGroup() as tg,
-    ):
-        # Set up task to receive server messages.
-        tg.create_task(receive_audio(session, user_hash))
+    async with GoogleClientFactory.audio() as client:
+        async with (
+            client.live.music.connect(model='models/lyria-realtime-exp') as session,
+            asyncio.TaskGroup() as tg,
+        ):
+            # Set up task to receive server messages.
+            tg.create_task(receive_audio(session, user_hash))
 
-        # Send initial prompts and config
-        await session.set_weighted_prompts(
-          prompts=[
-            types.WeightedPrompt(text=music_tone, weight=1.0),
-          ]
-        )
-        await session.set_music_generation_config(
-          config=types.LiveMusicGenerationConfig(bpm=90, temperature=1.0)
-        )
-        await session.play()
-        logger.info(f"Started music generation for user hash {user_hash}, music tone: {music_tone}")
-        sessions[user_hash] = {
-            'session': session,
-            'queue': queue.Queue()
-        }
+            # Send initial prompts and config
+            await asyncio.wait_for(
+                session.set_weighted_prompts(
+                    prompts=[types.WeightedPrompt(text=music_tone, weight=1.0)]
+                ),
+                40,
+            )
+            await asyncio.wait_for(
+                session.set_music_generation_config(
+                    config=types.LiveMusicGenerationConfig(bpm=90, temperature=1.0)
+            ),
+            40,
+            )
+            await asyncio.wait_for(session.play(), 40)
+            logger.info(
+                f"Started music generation for user hash {user_hash}, music tone: {music_tone}"
+            )
+            sessions[user_hash] = {
+                'session': session,
+                'queue': queue.Queue()
+            }
         
 async def change_music_tone(user_hash: str, new_tone):
     logger.info(f"Changing music tone to {new_tone}")
@@ -45,8 +50,11 @@ async def change_music_tone(user_hash: str, new_tone):
     if not session:
         logger.error(f"No session found for user hash {user_hash}")
         return
-    await session.set_weighted_prompts(
-        prompts=[types.WeightedPrompt(text=new_tone, weight=1.0)]
+    await asyncio.wait_for(
+        session.set_weighted_prompts(
+            prompts=[types.WeightedPrompt(text=new_tone, weight=1.0)]
+        ),
+        40,
     )
         
 
@@ -79,8 +87,8 @@ async def cleanup_music_session(user_hash: str):
     if user_hash in sessions:
         logger.info(f"Cleaning up music session for user hash {user_hash}")
         session = sessions[user_hash]['session']
-        await session.stop()
-        await session.close()
+        await asyncio.wait_for(session.stop(), 40)
+        await asyncio.wait_for(session.close(), 40)
         del sessions[user_hash]
     
 
@@ -122,4 +130,4 @@ def update_audio(user_hash):
             wf.setframerate(SAMPLE_RATE)
             wf.writeframes(pcm_data)
         wav_bytes = wav_buffer.getvalue()
-        yield wav_bytes
+        yield wav_bytes

src/images/image_generator.py

@@ -1,17 +1,15 @@
-from google import genai
 from google.genai import types
 import os
 from PIL import Image
 from io import BytesIO
 from datetime import datetime
-from config import settings
 import logging
 import asyncio
 import gradio as gr
 
-logger = logging.getLogger(__name__)
+from services.google import GoogleClientFactory
 
-client = genai.Client(api_key=settings.gemini_api_key.get_secret_value()).aio
+logger = logging.getLogger(__name__)
 
 safety_settings = [
     types.SafetySetting(
@@ -50,14 +48,18 @@ async def generate_image(prompt: str) -> tuple[str, str] | None:
     logger.info(f"Generating image with prompt: {prompt}")
 
     try:
-        response = await client.models.generate_content(
-            model="gemini-2.0-flash-preview-image-generation",
-            contents=prompt,
-            config=types.GenerateContentConfig(
-                response_modalities=["TEXT", "IMAGE"],
-                safety_settings=safety_settings,
-            ),
-        )
+        async with GoogleClientFactory.image() as client:
+            response = await asyncio.wait_for(
+                client.models.generate_content(
+                    model="gemini-2.0-flash-preview-image-generation",
+                    contents=prompt,
+                    config=types.GenerateContentConfig(
+                        response_modalities=["TEXT", "IMAGE"],
+                        safety_settings=safety_settings,
+                    ),
+                ),
+                40,
+            )
 
         # Process the response parts
         image_saved = False
@@ -108,23 +110,23 @@ async def modify_image(image_path: str, modification_prompt: str) -> str | None:
         logger.error(f"Error: Image file not found at {image_path}")
         return None
 
-    key = settings.gemini_api_key.get_secret_value()
-
-    client = genai.Client(api_key=key).aio
-
     try:
-        # Load the input image
-        input_image = Image.open(image_path)
-
-        # Make the API call with both text and image
-        response = await client.models.generate_content(
-            model="gemini-2.0-flash-preview-image-generation",
-            contents=[modification_prompt, input_image],
-            config=types.GenerateContentConfig(
-                response_modalities=["TEXT", "IMAGE"],
-                safety_settings=safety_settings,
-            ),
-        )
+        async with GoogleClientFactory.image() as client:
+            # Load the input image
+            input_image = Image.open(image_path)
+
+            # Make the API call with both text and image
+            response = await asyncio.wait_for(
+                client.models.generate_content(
+                    model="gemini-2.0-flash-preview-image-generation",
+                    contents=[modification_prompt, input_image],
+                    config=types.GenerateContentConfig(
+                        response_modalities=["TEXT", "IMAGE"],
+                        safety_settings=safety_settings,
+                    ),
+                ),
+                40,
+            )
 
         # Process the response parts
         image_saved = False

src/main.py

@@ -345,13 +345,6 @@ with gr.Blocks(
         ],
     )
 
-    game_choices.change(
-        fn=update_scene,
-        inputs=[local_storage, game_choices],
-        outputs=[game_text, game_image, game_choices, custom_choice],
-        concurrency_limit=CONCURRENCY_LIMIT,
-    )
-
     custom_choice.submit(
         fn=update_scene,
         inputs=[local_storage, custom_choice],
@@ -367,9 +360,10 @@ with gr.Blocks(
     )
     local_storage.change(
         fn=update_audio,
-        inputs=[local_storage],
+        inputs=[],
         outputs=[audio_out],
         concurrency_limit=CONCURRENCY_LIMIT,
     )
 
+demo.queue()
 demo.launch(ssr_mode=False)

src/services/google.py

@@ -0,0 +1,76 @@
+import asyncio
+import logging
+from contextlib import asynccontextmanager
+from google import genai
+import threading
+
+from config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class ApiKeyPool:
+    """Manage Google API keys with round-robin selection."""
+
+    def __init__(self) -> None:
+        self._keys: list[str] | None = None
+        self._index = 0
+        self._lock = asyncio.Lock()
+        self._sync_lock = threading.Lock()
+
+    def _load_keys(self) -> None:
+        keys_raw = (
+            getattr(settings, "gemini_api_keys", None) or settings.gemini_api_key
+        )
+        keys_str = keys_raw.get_secret_value()
+        keys = [k.strip() for k in keys_str.split(',') if k.strip()] if keys_str else []
+        if not keys:
+            msg = "Google API keys are not configured or invalid"
+            logger.error(msg)
+            raise ValueError(msg)
+        self._keys = keys
+
+    async def get_key(self) -> str:
+        async with self._lock:
+            if self._keys is None:
+                self._load_keys()
+            key = self._keys[self._index]
+            self._index = (self._index + 1) % len(self._keys)
+            logger.debug("Using Google API key index %s", self._index)
+            return key
+
+    def get_key_sync(self) -> str:
+        """Synchronous helper for environments without an event loop."""
+        with self._sync_lock:
+            if self._keys is None:
+                self._load_keys()
+            key = self._keys[self._index]
+            self._index = (self._index + 1) % len(self._keys)
+            logger.debug("Using Google API key index %s", self._index)
+            return key
+
+
+class GoogleClientFactory:
+    """Factory for thread-safe creation of Google GenAI clients."""
+
+    _pool = ApiKeyPool()
+
+    @classmethod
+    @asynccontextmanager
+    async def image(cls):
+        key = await cls._pool.get_key()
+        client = genai.Client(api_key=key)
+        try:
+            yield client.aio
+        finally:
+            pass
+
+    @classmethod
+    @asynccontextmanager
+    async def audio(cls):
+        key = await cls._pool.get_key()
+        client = genai.Client(api_key=key, http_options={"api_version": "v1alpha"})
+        try:
+            yield client.aio
+        finally:
+            pass