Merge pull request #4 from DeltaZN/feature/unique-session-ids

feat: implement Google API key management and refactor client usage i…

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,59 +1,71 @@
 import asyncio
-from google import genai
 from google.genai import types
-from config import settings
 import wave
 import queue
 import logging
 import io
 import time
+from config import settings
+from services.google import GoogleClientFactory
 
 logger = logging.getLogger(__name__)
 
-client = genai.Client(api_key=settings.gemini_api_key.get_secret_value(), http_options={'api_version': 'v1alpha'})
+
+
 
 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))
-
-        # 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)
+        logger.info(
+            f"Music generation already started for user hash {user_hash}, skipping new generation"
         )
-        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()
-        }
-        
+        return
+    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 asyncio.wait_for(
+                session.set_weighted_prompts(
+                    prompts=[types.WeightedPrompt(text=music_tone, weight=1.0)]
+                ),
+                settings.request_timeout,
+            )
+            await asyncio.wait_for(
+                session.set_music_generation_config(
+                    config=types.LiveMusicGenerationConfig(bpm=90, temperature=1.0)
+                ),
+                settings.request_timeout,
+            )
+            await asyncio.wait_for(session.play(), settings.request_timeout)
+            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}")
-    session = sessions.get(user_hash, {}).get('session')
+    session = sessions.get(user_hash, {}).get("session")
     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)]
+        ),
+        settings.request_timeout,
     )
-        
+
 
 SAMPLE_RATE = 48000
 NUM_CHANNELS = 2  # Stereo
 SAMPLE_WIDTH = 2  # 16-bit audio -> 2 bytes per sample
 
+
 async def receive_audio(session, user_hash):
     """Process incoming audio from the music generation."""
     while True:
@@ -61,7 +73,7 @@ async def receive_audio(session, user_hash):
             async for message in session.receive():
                 if message.server_content and message.server_content.audio_chunks:
                     audio_data = message.server_content.audio_chunks[0].data
-                    queue = sessions[user_hash]['queue']
+                    queue = sessions[user_hash]["queue"]
                     # audio_data is already bytes (raw PCM)
                     await asyncio.to_thread(queue.put, audio_data)
                 await asyncio.sleep(10**-12)
@@ -69,42 +81,47 @@ async def receive_audio(session, user_hash):
             logger.error(f"Error in receive_audio: {e}")
             break
 
+
 sessions = {}
 
+
 async def start_music_generation(user_hash: str, music_tone: str):
     """Start the music generation in a separate thread."""
     await generate_music(user_hash, music_tone, receive_audio)
-    
+
+
 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()
+        session = sessions[user_hash]["session"]
+        await asyncio.wait_for(session.stop(), settings.request_timeout)
+        await asyncio.wait_for(session.close(), settings.request_timeout)
         del sessions[user_hash]
-    
+
 
 def update_audio(user_hash):
     """Continuously stream audio from the queue as WAV bytes."""
     if user_hash == "":
         return
-    
+
     logger.info(f"Starting audio update loop for user hash: {user_hash}")
     while True:
         if user_hash not in sessions:
             time.sleep(0.5)
             continue
-        queue = sessions[user_hash]['queue']
-        pcm_data = queue.get() # This is raw PCM audio bytes
-        
+        queue = sessions[user_hash]["queue"]
+        pcm_data = queue.get()  # This is raw PCM audio bytes
+
         if not isinstance(pcm_data, bytes):
-            logger.warning(f"Expected bytes from audio_queue, got {type(pcm_data)}. Skipping.")
+            logger.warning(
+                f"Expected bytes from audio_queue, got {type(pcm_data)}. Skipping."
+            )
             continue
 
         # Lyria provides stereo, 16-bit PCM at 48kHz.
         # Ensure the number of bytes is consistent with stereo 16-bit audio.
         # Each frame = NUM_CHANNELS * SAMPLE_WIDTH bytes.
-        # If len(pcm_data) is not a multiple of (NUM_CHANNELS * SAMPLE_WIDTH), 
+        # If len(pcm_data) is not a multiple of (NUM_CHANNELS * SAMPLE_WIDTH),
         # it might indicate an incomplete chunk or an issue.
         bytes_per_frame = NUM_CHANNELS * SAMPLE_WIDTH
         if len(pcm_data) % bytes_per_frame != 0:
@@ -113,13 +130,13 @@ def update_audio(user_hash):
                 f"bytes_per_frame ({bytes_per_frame}). This might cause issues with WAV formatting."
             )
             # Depending on strictness, you might want to skip this chunk:
-            # continue 
+            # continue
 
         wav_buffer = io.BytesIO()
-        with wave.open(wav_buffer, 'wb') as wf:
+        with wave.open(wav_buffer, "wb") as wf:
             wf.setnchannels(NUM_CHANNELS)
-            wf.setsampwidth(SAMPLE_WIDTH) # Corresponds to 16-bit audio
+            wf.setsampwidth(SAMPLE_WIDTH)  # Corresponds to 16-bit audio
             wf.setframerate(SAMPLE_RATE)
             wf.writeframes(pcm_data)
         wav_bytes = wav_buffer.getvalue()
-        yield wav_bytes
+        yield wav_bytes

src/config.py

@@ -29,6 +29,6 @@ class AppSettings(BaseAppSettings):
     top_p: float = 0.95
     temperature: float = 0.5
     pregenerate_next_scene: bool = True
-    
+    request_timeout: int = 20
 
 settings = AppSettings()

src/images/image_generator.py

@@ -1,18 +1,16 @@
-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
+from config import settings
+from services.google import GoogleClientFactory
 
 logger = logging.getLogger(__name__)
 
-client = genai.Client(api_key=settings.gemini_api_key.get_secret_value()).aio
-
 safety_settings = [
     types.SafetySetting(
         category="HARM_CATEGORY_HARASSMENT",
@@ -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,
+                    ),
+                ),
+                settings.request_timeout,
+            )
 
         # 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,
+                    ),
+                ),
+                settings.request_timeout,
+            )
 
         # Process the response parts
         image_saved = False

src/main.py

@@ -136,7 +136,7 @@ with gr.Blocks(
     with gr.Column(visible=False, elem_id="loading-indicator") as loading_indicator:
         gr.HTML("<div class='loading-text'>🚀 Starting your adventure...</div>")
 
-    local_storage = gr.BrowserState("", "user_hash")
+    ls_user_hash = gr.BrowserState("", "user_hash")
 
     # Constructor Interface (visible by default)
     with gr.Column(
@@ -313,7 +313,7 @@ with gr.Blocks(
     start_btn.click(
         fn=start_game_with_music,
         inputs=[
-            local_storage,
+            ls_user_hash,
             setting_description,
             char_name,
             char_age,
@@ -330,13 +330,14 @@ with gr.Blocks(
             game_image,
             game_choices,
             custom_choice,
+            ls_user_hash,
         ],
         concurrency_limit=CONCURRENCY_LIMIT,
     )
 
     back_btn.click(
         fn=return_to_constructor,
-        inputs=[local_storage],
+        inputs=[ls_user_hash],
         outputs=[
             loading_indicator,
             constructor_interface,
@@ -345,16 +346,9 @@ 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],
+        inputs=[ls_user_hash, custom_choice],
         outputs=[game_text, game_image, game_choices, custom_choice],
         concurrency_limit=CONCURRENCY_LIMIT,
     )
@@ -363,13 +357,15 @@ with gr.Blocks(
     demo.load(
         fn=generate_user_hash,
         inputs=[],
-        outputs=[local_storage],
+        outputs=[ls_user_hash],
     )
-    local_storage.change(
+    ls_user_hash.change(
         fn=update_audio,
-        inputs=[local_storage],
+        inputs=[ls_user_hash],
         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