feat: use more precise prompts for music and image generation

src/agent/image_agent.py

@@ -2,6 +2,8 @@ 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.state import get_user_state, set_user_state
+from agent.prompts import GAME_STATE_PROMPT
 import logging
 
 logger = logging.getLogger(__name__)
@@ -66,12 +68,24 @@ class ChangeScene(BaseModel):
 
 image_prompt_generator_llm = create_light_llm(0.1).with_structured_output(ChangeScene)
 
-async def generate_image_prompt(scene_description: str, request_id: str) -> ChangeScene:
+async def generate_image_prompt(user_hash: str, scene_description: str) -> ChangeScene:
     """
     Generates a detailed image prompt string based on a scene description.
     This prompt is intended for use with an AI image generation model.
     """
     logger.info(f"Generating image prompt for the current scene: {request_id}")
+
+    state = get_user_state(user_hash)
+    scene = GAME_STATE_PROMPT.format(
+        lore=state.story_frame.lore,
+        goal=state.story_frame.goal,
+        milestones=",".join(m.id for m in state.story_frame.milestones),
+        endings=",".join(e.id for e in state.story_frame.endings),
+        history="; ".join(f"{c.scene_id}:{c.choice_text}" for c in state.user_choices),
+        last_choice=last_choice,
+        scene_description=scene_description
+    )
+
     response = await image_prompt_generator_llm.ainvoke(
         [
             SystemMessage(content=IMAGE_GENERATION_SYSTEM_PROMPT),

src/agent/llm_graph.py

@@ -60,7 +60,7 @@ async def node_init_game(state: GraphState) -> GraphState:
     first_scene = await generate_scene.ainvoke(
         {"user_hash": state.user_hash, "last_choice": "start"}
     )
-    change_scene = await generate_image_prompt(first_scene["description"], state.user_hash)
+    change_scene = await generate_image_prompt(state.user_hash, first_scene["description"])
     logger.info(f"Change scene: {change_scene}")
     await generate_scene_image.ainvoke(
         {
@@ -94,7 +94,7 @@ async def node_player_step(state: GraphState) -> GraphState:
                 "last_choice": state.choice_text,
             }
         )
-        change_scene = await generate_image_prompt(next_scene["description"], state.user_hash)
+        change_scene = await generate_image_prompt(state.user_hash, next_scene["description"])
         current_image = None
         if scene_id and scene_id in user_state.scenes:
             current_image = user_state.scenes[scene_id].image
@@ -107,8 +107,9 @@ async def node_player_step(state: GraphState) -> GraphState:
                 "change_scene": change_scene,
             }
         )
-        music_task = change_music_tone(state.user_hash, next_scene["music"])
-        await asyncio.gather(image_task, music_task)
+        music_task = generate_music_prompt(next_scene["description"])
+        _, music_prompt = await asyncio.gather(image_task, music_task)
+        asyncio.create_task(change_music_tone(state.user_hash, music_prompt))
         state.scene = next_scene
     return state
 

src/agent/music_agent.py

@@ -1,7 +1,9 @@
 from pydantic import BaseModel
 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.state import get_user_state, set_user_state
 
 logger = logging.getLogger(__name__)
 
@@ -38,10 +40,22 @@ class MusicPrompt(BaseModel):
 llm = create_light_llm(0.1).with_structured_output(MusicPrompt)
 
 
-async def generate_music_prompt(scene_description: str, request_id: str) -> str:
-    logger.info(f"Generating music prompt for the current scene: {request_id}")
+async def generate_music_prompt(user_hash: str, scene_description: str) -> str:
+    logger.info(f"Generating music prompt for the current scene: {scene_description}")
+
+    state = get_user_state(user_hash)
+    scene = GAME_STATE_PROMPT.format(
+        lore=state.story_frame.lore,
+        goal=state.story_frame.goal,
+        milestones=",".join(m.id for m in state.story_frame.milestones),
+        endings=",".join(e.id for e in state.story_frame.endings),
+        history="; ".join(f"{c.scene_id}:{c.choice_text}" for c in state.user_choices),
+        last_choice=last_choice,
+        scene_description=scene_description
+    )
+    
     response = await llm.ainvoke(
-        [SystemMessage(content=system_prompt), HumanMessage(content=scene_description)]
+        [SystemMessage(content=system_prompt), HumanMessage(content=scene)]
     )
-    logger.info(f"Music prompt generated: {request_id}")
+    logger.info(f"Music prompt generated")
     return response.prompt

src/agent/prompts.py

@@ -13,6 +13,22 @@ Return ONLY a JSON object with:
 Translate the lore, goal, milestones and endings to the language which is used in the game and setting description.
 """
 
+GAME_STATE_PROMPT = """
+---Game Settings START---
+Lore: {lore}
+Goal: {goal}
+Milestones: {milestones}
+Endings: {endings}
+---Game Settings END---
+
+---User's actions START---
+History: {history}
+Last choice: {last_choice}
+---User's actions END---
+
+Game response to user's action: {scene_description}
+"""
+
 SCENE_PROMPT = """You are an AI agent for a visual novel game. 
 Your role is to process incoming data and generate the next scene description and choices.
 Translate the scene description and choices into a language which is used in the Game Settings.

src/agent/runner.py

@@ -56,7 +56,7 @@ async def process_step(
         ending_desc = ending_info.get("description") or ending_info.get(
             "condition", ""
         )
-        change_scene = await generate_image_prompt(ending_desc, user_hash)
+        change_scene = await generate_image_prompt(user_hash, ending_desc)
         if change_scene.change_scene == "no_change":
             change_scene.change_scene = "change_completely"
             if not change_scene.scene_description: