feat: add util for LLM retries

src/agent/image_agent.py

@@ -2,9 +2,10 @@ from pydantic import BaseModel, Field
 from typing import Literal, Optional
 from agent.llm import create_light_llm
 from langchain_core.messages import SystemMessage, HumanMessage
-from agent.redis_state import get_user_state, set_user_state
+from agent.redis_state import get_user_state
 from agent.prompts import GAME_STATE_PROMPT
 import logging
+from agent.utils import with_retries
 
 logger = logging.getLogger(__name__)
 
@@ -86,11 +87,11 @@ async def generate_image_prompt(user_hash: str, scene_description: str, last_cho
         scene_description=scene_description
     )
 
-    response = await image_prompt_generator_llm.ainvoke(
+    response = await with_retries(lambda: image_prompt_generator_llm.ainvoke(
         [
             SystemMessage(content=IMAGE_GENERATION_SYSTEM_PROMPT),
             HumanMessage(content=scene),
         ]
-    )
+    ))
     logger.info(f"Image prompt generated")
     return response

src/agent/music_agent.py

@@ -3,7 +3,8 @@ from agent.llm import create_light_llm
 from agent.prompts import GAME_STATE_PROMPT
 from langchain_core.messages import SystemMessage, HumanMessage
 import logging
-from agent.redis_state import get_user_state, set_user_state
+from agent.redis_state import get_user_state
+from agent.utils import with_retries
 
 logger = logging.getLogger(__name__)
 
@@ -54,8 +55,8 @@ async def generate_music_prompt(user_hash: str, scene_description: str, last_cho
         scene_description=scene_description
     )
     
-    response = await llm.ainvoke(
+    response = await with_retries(lambda: llm.ainvoke(
         [SystemMessage(content=system_prompt), HumanMessage(content=scene)]
-    )
-    logger.info(f"Music prompt generated")
+    ))
+    logger.info("Music prompt generated")
     return response.prompt

src/agent/tools.py

@@ -18,6 +18,7 @@ from agent.models import (
 )
 from agent.prompts import ENDING_CHECK_PROMPT, SCENE_PROMPT, STORY_FRAME_PROMPT
 from agent.redis_state import get_user_state, set_user_state
+from agent.utils import with_retries
 from images.image_generator import modify_image, generate_image
 from agent.image_agent import ChangeScene
 
@@ -43,7 +44,7 @@ async def generate_story_frame(
         character=character,
         genre=genre,
     )
-    resp: StoryFrameLLM = await llm.ainvoke(prompt)
+    resp: StoryFrameLLM = await with_retries(lambda: llm.ainvoke(prompt))
     story_frame = StoryFrame(
         lore=resp.lore,
         goal=resp.goal,
@@ -77,10 +78,10 @@ async def generate_scene(
         history="; ".join(f"{c.scene_id}:{c.choice_text}" for c in state.user_choices),
         last_choice=last_choice,
     )
-    resp: SceneLLM = await llm.ainvoke(prompt)
+    resp: SceneLLM = await with_retries(lambda: llm.ainvoke(prompt))
     if len(resp.choices) < 2:
-        resp = await llm.ainvoke(
-            prompt + "\nThe scene must contain exactly two choices."
+        resp = await with_retries(
+            lambda: llm.ainvoke(prompt + "\nThe scene must contain exactly two choices.")
         )
     scene_id = str(uuid.uuid4())
     choices = [
@@ -163,7 +164,7 @@ async def check_ending(
         history=history,
         endings=",".join(f"{e.id}:{e.condition}" for e in state.story_frame.endings),
     )
-    resp: EndingCheckResult = await llm.ainvoke(prompt)
+    resp: EndingCheckResult = await with_retries(lambda: llm.ainvoke(prompt))
     if resp.ending_reached and resp.ending:
         state.ending = resp.ending
         await set_user_state(user_hash, state)

src/agent/utils.py

@@ -0,0 +1,31 @@
+"""Agent-related utilities."""
+
+import asyncio
+import logging
+from typing import Awaitable, Callable, TypeVar
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+async def with_retries(
+    awaitable_factory: Callable[[], Awaitable[T]],
+    retries: int = 3,
+    timeout: int = 15,
+) -> T:
+    """Execute an awaitable with retries and timeout.
+
+    :param awaitable_factory: A function that returns an awaitable.
+    :param retries: Maximum number of retries.
+    :param timeout: Timeout in seconds for each attempt.
+    :return: The result of the awaitable.
+    """
+    last_exception = None
+    for attempt in range(retries):
+        try:
+            return await asyncio.wait_for(awaitable_factory(), timeout=timeout)
+        except Exception as e:
+            logger.warning(f"Attempt {attempt + 1}/{retries} failed with error: {e}")
+            last_exception = e
+    raise last_exception from last_exception 

src/images/image_generator.py

@@ -8,6 +8,7 @@ import asyncio
 import gradio as gr
 from config import settings
 from services.google import GoogleClientFactory
+from agent.utils import with_retries
 
 logger = logging.getLogger(__name__)
 
@@ -49,16 +50,15 @@ async def generate_image(prompt: str) -> tuple[str, str] | None:
 
     try:
         async with GoogleClientFactory.image() as client:
-            response = await asyncio.wait_for(
-                client.models.generate_content(
+            response = await with_retries(
+                lambda: 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
@@ -116,8 +116,8 @@ async def modify_image(image_path: str, modification_prompt: str) -> str | None:
             input_image = Image.open(image_path)
 
             # Make the API call with both text and image
-            response = await asyncio.wait_for(
-                client.models.generate_content(
+            response = await with_retries(
+                lambda: client.models.generate_content(
                     model="gemini-2.0-flash-preview-image-generation",
                     contents=[modification_prompt, input_image],
                     config=types.GenerateContentConfig(
@@ -125,7 +125,6 @@ async def modify_image(image_path: str, modification_prompt: str) -> str | None:
                         safety_settings=safety_settings,
                     ),
                 ),
-                settings.request_timeout,
             )
 
         # Process the response parts