Revert "feat: refactor API key management and update client usage in audio and image generation"

This reverts commit 64ad372380e1c14b443720d14b377e55352f60f9.

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/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,33 +8,40 @@ 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:
         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}")
@@ -44,8 +49,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,
     )
         
 
@@ -78,8 +86,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]
     
 
@@ -117,4 +125,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

@@ -366,4 +366,5 @@ with gr.Blocks(
         outputs=[audio_out],
     )
 
+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