Improve image generation

seratch · May 22, 2024 · aed46a7 · aed46a7
1 parent 8988c5e
commit aed46a7
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 33 deletions.
diff --git a/app/bolt_listeners.py b/app/bolt_listeners.py
@@ -7,6 +7,7 @@
 from slack_bolt import App, Ack, BoltContext, BoltResponse
 from slack_bolt.request.payload_utils import is_event
 from slack_sdk.web import WebClient
+from slack_sdk.errors import SlackApiError
 
 from app.env import (
     OPENAI_TIMEOUT_SECONDS,
@@ -64,6 +65,7 @@
     build_image_generation_wip_modal,
     build_image_generation_result_modal,
     build_image_generation_text_modal,
+    build_image_generation_result_blocks,
 )
 
 
@@ -772,7 +774,7 @@ def display_image_generation_result(
             "OPENAI_IMAGE_GENERATION_MODEL", OPENAI_IMAGE_GENERATION_MODEL
         )
         text = "\n".join(map(lambda s: f">{s}", prompt.split("\n")))
-        view = build_image_generation_result_modal(
+        blocks = build_image_generation_result_blocks(
             text=text,
             spent_seconds=str(round(spent_seconds, 2)),
             image_url=image_url,
@@ -781,22 +783,48 @@ def display_image_generation_result(
             quality=quality,
             style=style,
         )
-        client.views_update(view_id=payload["id"], view=view)
+        client.chat_postMessage(
+            channel=context.actor_user_id,
+            text=f"Here is the generated image URL: {image_url}",
+            blocks=blocks,
+        )
+        client.views_update(
+            view_id=payload["id"],
+            view=build_image_generation_result_modal(blocks),
+        )
 
     except (APITimeoutError, TimeoutError):
+        client.chat_postMessage(
+            channel=context.actor_user_id,
+            text=TIMEOUT_ERROR_MESSAGE,
+        )
         client.views_update(
             view_id=payload["id"],
             view=build_image_generation_text_modal(TIMEOUT_ERROR_MESSAGE),
         )
-    except Exception as e:
-        logger.exception(f"Failed to share a generated image: {e}")
+    except SlackApiError as e:
+        logger.exception(f"Failed to call Slack APIs for image generation: {e}")
         client.views_update(
             view_id=payload["id"],
             view=build_image_generation_text_modal(
                 f"{text}\n\n:warning: My apologies! "
-                f"An error occurred while generating an image: {e}"
+                f"An error occurred while calling Slack APIs: {e}"
             ),
         )
+    except Exception as e:
+        logger.exception(f"Failed to share a generated image: {e}")
+        error = (
+            f"{text}\n\n:warning: My apologies! "
+            f"An error occurred while generating an image: {e}"
+        )
+        client.chat_postMessage(
+            channel=context.actor_user_id,
+            text=error,
+        )
+        client.views_update(
+            view_id=payload["id"],
+            view=build_image_generation_text_modal(error),
+        )
 
 
 #

diff --git a/app/slack_ui.py b/app/slack_ui.py
@@ -684,10 +684,24 @@ def build_image_generation_input_modal(prompt: str) -> dict:
 
 
 def build_image_generation_wip_modal() -> dict:
-    return build_image_generation_text_modal("Working on this now ... :hourglass:")
+    return build_image_generation_text_modal(
+        "Working on this now ... :hourglass:\n\n"
+        "Once the image is ready, this app will send it to you in a DM. "
+        "If you don't want to wait here, you can close this modal at any time."
+    )
+
+
+def build_image_generation_result_modal(blocks: list) -> dict:
+    return {
+        "type": "modal",
+        "callback_id": "image-generation",
+        "title": {"type": "plain_text", "text": "Image Generation"},
+        "close": {"type": "plain_text", "text": "Close"},
+        "blocks": blocks,
+    }
 
 
-def build_image_generation_result_modal(
+def build_image_generation_result_blocks(
     *,
     text: str,
     size: str,
@@ -696,34 +710,28 @@ def build_image_generation_result_modal(
     spent_seconds: str,
     image_url: str,
     model: str,
-) -> dict:
-    return {
-        "type": "modal",
-        "callback_id": "image-generation",
-        "title": {"type": "plain_text", "text": "Image Generation"},
-        "close": {"type": "plain_text", "text": "Close"},
-        "blocks": [
-            {
-                "type": "section",
-                "text": {
-                    "type": "mrkdwn",
-                    "text": f"{text}\n>\n>(model: {model}, size: {size}, quality: {quality}, style: {style})\n\n"
-                    f"Here is the image content URL generated by the above prompt "
-                    f"(time spent: {spent_seconds} seconds):"
-                    f" <{image_url}|Click>",
-                },
+) -> list[dict]:
+    return [
+        {
+            "type": "section",
+            "text": {
+                "type": "mrkdwn",
+                "text": f"{text}\n>\n>(model: {model}, size: {size}, quality: {quality}, style: {style})\n\n"
+                f"Here is the image content URL generated by the above prompt "
+                f"(time spent: {spent_seconds} seconds):"
+                f" <{image_url}|Click>",
             },
-            {
-                "type": "image",
-                "title": {
-                    "type": "plain_text",
-                    "text": f"This image was generated using ChatGPT's {model} model",
-                },
-                "image_url": image_url,
-                "alt_text": f"Generated by {model}",
+        },
+        {
+            "type": "image",
+            "title": {
+                "type": "plain_text",
+                "text": f"This image was generated using ChatGPT's {model} model",
             },
-        ],
-    }
+            "image_url": image_url,
+            "alt_text": f"Generated by {model}",
+        },
+    ]
 
 
 def build_image_generation_text_modal(section_text: str) -> dict: