Coverage for apps/streamchat/streamchat_job.py: 92%
124 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamChat job to generate a video chat response.
3"""
4import asyncio
5import json
6import sys
7import time
8import aiofiles
9import unicodedata
11from PIL import Image
13from typing import List, override
14from typing import Any
15from typing import Dict
16from typing import Optional
18from chat_prompts import IMG_PROMPT
19from chat_prompts import IMG_NEG_PROMPT
20from chat_prompts import CHAT_PROMPT
21from chat_prompts import VIDEO_PROMPT
22from chat_prompts import VIDEO_NEG_PROMPT
24# Local relative imports
25sys.path.append("..") # noqa: E402
26sys.path.append("../..") # noqa: E402
28from streamwise_job import StreamWiseJob
29from streamwise_job import JobStatus
31from lmm_service_manager import LMMServiceManager
33from character import Character
35from file_utils import save_base64_as_binary
37from console_utils import bytes_to_human
40class StreamChatJob(StreamWiseJob):
41 """A job to generate a video chat response."""
43 def __init__(
44 self,
45 job_id: str,
46 service_manager: LMMServiceManager,
47 config: Dict[str, Any] = {},
48 ) -> None:
49 super().__init__(
50 "streamchat",
51 job_id,
52 service_manager,
53 config)
54 self.character: Optional[Character] = None
55 self.image: Optional[Image.Image] = None
56 self.messages = [{
57 "role": "system",
58 "content": CHAT_PROMPT
59 }]
61 def get_config_gender(self) -> str:
62 return self.get_config_str(
63 "gender_prompt",
64 "female")
66 @override
67 async def generate(
68 self,
69 job_config: Dict[str, Any],
70 ) -> None:
71 await self.gen_chat_base()
73 async def gen_chat_base(
74 self,
75 ) -> None:
76 """Generate a chat video."""
77 async with self.job_status_handler():
78 self.character = Character(
79 name="Assistant",
80 gender=self.get_config_gender(),
81 speech_speed=self.get_config_float("speech_speed", 1.1),
82 )
84 # Generate the base main image
85 img_prompt = IMG_PROMPT
86 if self.character.gender:
87 img_prompt += "The gender of the character is a " + self.character.gender + "."
88 style_prompt = self.get_config_str("style_prompt", "")
89 if style_prompt:
90 img_prompt += "The style of the image is: " + style_prompt + "."
91 scene_prompt = self.get_config_str("scene_prompt", "")
92 if scene_prompt:
93 img_prompt += "The scene is: " + scene_prompt + "."
94 custom_prompt = self.get_config_str("custom_prompt", "")
95 if custom_prompt:
96 img_prompt += "Additional details: " + custom_prompt + "."
98 width, height = self.width, self.height
99 image = await self.gen.gen_image(
100 img_prompt,
101 neg_prompt=IMG_NEG_PROMPT,
102 width=width,
103 height=height,
104 task_id="main_image",
105 deadline=self.get_submission_time(),
106 )
107 if image is None:
108 raise Exception("Image generation failed.")
109 self.image = image
110 width, height = image.size
111 image_path = f"{self.job_path}/main_image.png"
112 image.save(image_path)
113 self.logger.info(f"Image with {width}x{height} pixels saved to '{image_path}'.")
115 def _get_msg_id(self) -> int:
116 """Get the current message ID."""
117 # Exclude system message and count pairs (Q/A)
118 return (len(self.messages) - 1) // 2
120 async def transcribe_audio(
121 self,
122 audio_path: str,
123 ) -> str:
124 msg_id = self._get_msg_id()
125 audio_transcript, _ = await self.gen.gen_audio_transcript(
126 audio_path,
127 task_id=f"chat{msg_id:03d}",
128 )
129 return audio_transcript
131 async def gen_chat(
132 self,
133 user_message: str
134 ) -> Dict[str, str]:
135 """
136 Generate a chat response given a user message.
137 """
138 async with self.job_status_handler():
139 msg_id = self._get_msg_id()
141 # Generate text response
142 response_text = await self.gen_chat_text(
143 user_message=user_message,
144 msg_id=msg_id,
145 )
147 await self.save_status(JobStatus.RUNNING)
149 # Generate audio response
150 audio_task = asyncio.create_task(
151 self.gen_chat_audio(
152 response_text=response_text,
153 msg_id=msg_id,
154 )
155 )
156 await self.save_status(JobStatus.RUNNING)
158 # Generate video response
159 # TODO await task somewhere?
160 # video_task =
161 asyncio.create_task(
162 self.gen_chat_video(
163 audio_task=audio_task,
164 msg_id=msg_id,
165 )
166 )
168 """
169 video_binary = await video_task
170 if video_binary is None:
171 raise Exception("Video generation failed.")
172 self.logger.info(f"Generated video with {bytes_to_human(len(video_binary))}.")
173 """
175 return {
176 "id": str(msg_id),
177 "reply": response_text,
178 }
180 async def gen_chat_text(
181 self,
182 user_message: str,
183 msg_id: int,
184 ) -> str:
185 message_path = f"{self.job_path}/chat{msg_id:03d}_message.txt"
186 async with aiofiles.open(message_path, "w") as file:
187 await file.write(user_message)
189 self.messages.append({
190 "role": "user",
191 "content": user_message
192 })
194 prompt_path = f"{self.job_path}/chat{msg_id:03d}_prompt.jsonl"
195 async with aiofiles.open(prompt_path, "w") as file:
196 for msg in self.messages:
197 json_line = json.dumps(
198 msg,
199 ensure_ascii=False,
200 separators=(",", ":")
201 )
202 await file.write(json_line + "\n")
204 # Generate the text response using the LLM
205 response_text = await self.gen.gen_text(
206 self.messages,
207 task_id=f"chat{msg_id:03d}",
208 )
210 self.logger.info(f"[{msg_id}] Generated response: {response_text}")
211 response_path = f"{self.job_path}/chat{msg_id:03d}_response.txt"
212 async with aiofiles.open(response_path, "w") as file:
213 await file.write(response_text)
215 self.messages.append({
216 "role": "assistant",
217 "content": response_text
218 })
220 return response_text
222 async def gen_chat_audio(
223 self,
224 response_text: str,
225 msg_id: int,
226 ) -> str:
227 lang_code = "a" # American English TODO
228 response_text_clean = response_text.replace("\n", " ").strip()
229 response_text_clean = remove_emojis(response_text_clean)
230 assert self.character is not None, "Character must be set before generating audio"
231 audio_base64 = await self.gen.gen_audio(
232 text=response_text_clean,
233 voice=self.character.voice,
234 speed=self.character.speech_speed,
235 lang_code=lang_code,
236 task_id=f"{msg_id:03d}",
237 deadline=time.time(), # Now
238 )
239 if audio_base64 is None:
240 raise Exception("Audio generation failed.")
242 self.logger.info(f"[{msg_id}] Generated audio with {bytes_to_human(len(audio_base64))}.")
244 audio_path = f"{self.job_path}/chat{msg_id:03d}.wav"
245 await save_base64_as_binary(
246 audio_path,
247 audio_base64)
248 return audio_base64
250 async def gen_chat_video(
251 self,
252 audio_task: asyncio.Task,
253 msg_id: int,
254 ) -> None:
255 """
256 Generate chat video from base image and audio.
257 """
258 if self.image is None:
259 raise Exception("Base image not found for video generation.")
261 audio_base64 = await audio_task
262 if audio_base64 is None:
263 raise Exception("Audio generation failed.")
265 video_binary = await self.gen.gen_video_audio_from_img(
266 img=self.image,
267 audio_base64=audio_base64,
268 prompt=VIDEO_PROMPT,
269 neg_prompt=VIDEO_NEG_PROMPT,
270 width=self.width,
271 height=self.height,
272 steps=self.get_num_steps(),
273 task_id=f"{msg_id:03d}",
274 deadline=time.time(), # Now
275 )
276 if video_binary is None:
277 raise Exception("Video generation failed.")
279 self.logger.info(f"[{msg_id}] Generated video with {bytes_to_human(len(video_binary))}.")
281 video_path = f"{self.job_path}/chat{msg_id:03d}.mp4"
282 async with aiofiles.open(video_path, "wb") as file:
283 await file.write(video_binary)
285 async def get_chat_history(self) -> List[Dict[str, str]]:
286 """Get the chat history."""
287 history = []
288 for msg in self.messages:
289 history.append({
290 "role": msg["role"],
291 "content": msg["content"]
292 })
293 return history
296def remove_emojis(
297 text: str
298) -> str:
299 """
300 Remove emojis from the given text.
301 Avoid TTS speaking them.
302 """
303 return "".join(
304 ch for ch in text
305 if not unicodedata.category(ch).startswith("So")
306 )