gsavin commited on
Commit
3c0992c
·
1 Parent(s): 1af8d03

chore: move llm init to generate prompt step

Browse files
src/agent/image_agent.py CHANGED
@@ -66,9 +66,6 @@ class ChangeScene(BaseModel):
66
  )
67
  scene_description: Optional[str] = None
68
 
69
-
70
- image_prompt_generator_llm = create_light_llm(0.1).with_structured_output(ChangeScene)
71
-
72
  async def generate_image_prompt(user_hash: str, scene_description: str, last_choice = "No choice yet") -> ChangeScene:
73
  """
74
  Generates a detailed image prompt string based on a scene description.
@@ -86,6 +83,8 @@ async def generate_image_prompt(user_hash: str, scene_description: str, last_cho
86
  last_choice=last_choice,
87
  scene_description=scene_description
88
  )
 
 
89
 
90
  response = await with_retries(lambda: image_prompt_generator_llm.ainvoke(
91
  [
 
66
  )
67
  scene_description: Optional[str] = None
68
 
 
 
 
69
  async def generate_image_prompt(user_hash: str, scene_description: str, last_choice = "No choice yet") -> ChangeScene:
70
  """
71
  Generates a detailed image prompt string based on a scene description.
 
83
  last_choice=last_choice,
84
  scene_description=scene_description
85
  )
86
+
87
+ image_prompt_generator_llm = create_light_llm(0.1).with_structured_output(ChangeScene)
88
 
89
  response = await with_retries(lambda: image_prompt_generator_llm.ainvoke(
90
  [
src/agent/music_agent.py CHANGED
@@ -38,9 +38,6 @@ class MusicPrompt(BaseModel):
38
  prompt: str
39
 
40
 
41
- llm = create_light_llm(0.1).with_structured_output(MusicPrompt)
42
-
43
-
44
  async def generate_music_prompt(user_hash: str, scene_description: str, last_choice = "No choice yet") -> str:
45
  logger.info(f"Generating music prompt for the current scene: {scene_description}")
46
 
@@ -55,6 +52,8 @@ async def generate_music_prompt(user_hash: str, scene_description: str, last_cho
55
  scene_description=scene_description
56
  )
57
 
 
 
58
  response = await with_retries(lambda: llm.ainvoke(
59
  [SystemMessage(content=system_prompt), HumanMessage(content=scene)]
60
  ))
 
38
  prompt: str
39
 
40
 
 
 
 
41
  async def generate_music_prompt(user_hash: str, scene_description: str, last_choice = "No choice yet") -> str:
42
  logger.info(f"Generating music prompt for the current scene: {scene_description}")
43
 
 
52
  scene_description=scene_description
53
  )
54
 
55
+ llm = create_light_llm(0.1).with_structured_output(MusicPrompt)
56
+
57
  response = await with_retries(lambda: llm.ainvoke(
58
  [SystemMessage(content=system_prompt), HumanMessage(content=scene)]
59
  ))