Commit a971160

benny-dou <60535774+benny-dou@users.noreply.github.com>
2026-04-28 03:03:14
chore(ai): improve logging in image generation
1 parent 2445c41
Changed files (1)
src
ai
src/ai/images/openai_img.py
@@ -90,18 +90,20 @@ async def openai_image_generation(
                 func = openai.images.generate
             async for chunk in await func(**params):
                 resp = trim_none(chunk.model_dump())
-                if b64_json := resp.pop("b64_json", None):
+                b64_json = resp.pop("b64_json", None)
+                logger.trace(resp)
+                rtype = resp.pop("type", None)
+                is_completed = rtype in {"image_generation.completed", "image_edit.completed"}
+                if b64_json:
                     image_bytes = base64.b64decode(b64_json)
                     save_path = Path(DOWNLOAD_DIR) / f"{rand_string(10)}.{params.get('output_format', 'jpeg')}"
                     async with await anyio.open_file(save_path, "wb") as f:
                         await f.write(image_bytes)
-                    rtype = resp.pop("type", None)
-                    caption = f"{EMOJI_IMG_BOT}**{model_name}**\n" if rtype == "image_generation.completed" else f"⌛️**{model_name}** 中间结果...\n"
+                    caption = f"{EMOJI_IMG_BOT}**{model_name}**\n" if is_completed else f"⌛️**{model_name}** 中间结果...\n"
                     status_msg = await status_msg.edit_media(InputMediaPhoto(str(save_path), caption=caption + prettify(resp)))
                     save_path.unlink(missing_ok=True)
-                elif resp.pop("type", None) == "image_generation.completed":
+                elif is_completed:
                     await status_msg.edit_caption(f"{EMOJI_IMG_BOT}**{model_name}**\n{prettify(resp)}")
-                logger.trace(resp)
             if status_msg.photo:
                 return {"success": True}
         except Exception as e: