Skip to content

Flesh out Slack frontend, with help from Cursor (CodeRabbit) #4

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
cmsj wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Flesh out Slack frontend, with help from Cursor (CodeRabbit) #4

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
174 changes: 167 additions & 7 deletions src/dreambot/frontend/slack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
"""Slack frontend for Dreambot."""
import asyncio
import base64
import io
import os
Comment on lines +3 to +5
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Remove unused import io.

The import io is unused in this file. Please remove it to adhere to best practices and static analysis recommendations:

 import base64
-import io
 import os
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
import base64
import io
import os
import base64
import os
🧰 Tools
🪛 Ruff (0.8.2)

4-4: io imported but unused

Remove unused import: io

(F401)

import traceback

from typing import Any
Expand Down Expand Up @@ -49,9 +52,23 @@ async def on_message(body: dict[str, Any]): # type: ignore
await self.on_message(body)
except Exception as exc:
self.logger.error("Error handling Slack message: %s", exc)

self.is_booted = True # FIXME: This really should happen in response to some kind of "connected" event, but Slack doesn't seem to have one

# Add an event handler for when the app is ready
@self.slack.event("app_mention") # type: ignore
async def on_app_mention(body: dict[str, Any]): # type: ignore
# This is just a dummy handler to ensure the app is ready
pass

# Start the handler
await self.handler.start_async()

# Verify the connection by making a test API call
auth_test = await self.slack.client.auth_test() # type: ignore
self.logger.info("Connected to Slack workspace: %s", auth_test["team"])

# Set the booted flag
self.is_booted = True
self.logger.info("Slack connection established")
except SlackApiError as exc:
self.logger.error("Slack API error: %s", exc)
except Exception as exc:
Expand All @@ -78,7 +95,112 @@ async def callback_receive_workload(self, queue_name: str, message: dict[str, An
Returns:
bool: True if the message should be ack'd in NATS, False otherwise.
"""
# FIXME: Implement
self.logger.info("Received message for queue %s", queue_name)

# Get the channel to send the message to
channel_id = message["channel"]

# Prepare the message content
reply_content = ""
reply_blocks = []

if "reply-image" in message:
# For image replies, we need to upload the image to Slack
image_bytes = base64.standard_b64decode(message["reply-image"])
filename = self.clean_filename(message["prompt"], suffix=".png", output_dir=self.options["output_dir"])

# Save the image to disk
with open(os.path.join(self.options["output_dir"], filename), "wb") as image_file:
image_file.write(image_bytes)

# Upload the image to Slack
try:
with open(os.path.join(self.options["output_dir"], filename), "rb") as image_file:
upload_result = await self.slack.client.files_upload_v2(
file=image_file,
filename=filename,
title="Dream result",
initial_comment="I dreamed this:"
)

# Get the file URL from the upload result
file_url = upload_result["file"]["url_private"]

# Create a message with the image
reply_blocks = [
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": "I dreamed this:"
}
},
{
"type": "image",
"image_url": file_url,
"alt_text": "Dream result"
}
]

self.logger.info("OUTPUT: %s %s", self.log_slug(message), "Image uploaded")
except Exception as exc:
self.logger.error("Failed to upload image: %s", exc)
reply_content = "Failed to upload dream image."
self.logger.error("OUTPUT: %s %s", self.log_slug(message), reply_content)
Comment on lines +107 to +149
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Apr 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Image handling implementation looks good, but consider adding safeguards.

The image handling code thoroughly processes base64-encoded images, but consider adding:

  1. A file format check before decoding and saving
  2. A file size limit check before uploading to Slack
  3. Error handling for the base64 decoding operation

Additionally, the code calls self.clean_filename() which isn't defined in the shown file - ensure this method exists and properly sanitizes filenames.

🏁 Script executed:

#!/bin/bash
# Search for the clean_filename method implementation
rg -A 5 "def clean_filename" --type py

Length of output: 585

Enhancement: Add Additional Safeguards for Image Handling

The image handling logic is generally solid; however, please strengthen it by:

  • Adding a file format validation: Verify that the incoming file content (even if base64 encoded) conforms to the expected PNG format before decoding and saving.
  • Including a file size check: Implement a limit to ensure that excessively large images aren’t processed or uploaded to Slack.
  • Enhancing error handling for base64 decoding: Wrap the base64 decoding step in its own try/except block to catch decoding errors and provide clear error messaging.
  • Verifying filename sanitization: The method clean_filename() is defined in src/dreambot/shared/worker.py and appears to handle basic sanitization. Please review it to confirm it adequately tackles edge cases (e.g., extremely long filenames or disallowed characters).

elif "reply-text" in message:
reply_content = message["reply-text"]
self.logger.info("OUTPUT: %s %s", self.log_slug(message), message["reply-text"])
elif "reply-none" in message:
self.logger.info("SILENCE FOR %s %s", self.log_slug(message), message["reply-none"])
return True
elif "error" in message:
reply_content = f"Dream sequence collapsed: {message['error']}"
self.logger.error("OUTPUT: %s %s", self.log_slug(message), reply_content)
elif "usage" in message:
reply_content = f"{message['usage']}"
self.logger.info("OUTPUT: %s %s", self.log_slug(message), message["usage"])
else:
reply_content = "Dream sequence collapsed, unknown reason."
self.logger.error("Unknown workload message: %s", message)

# Send the message to the channel
try:
if reply_blocks:
await self.slack.client.chat_postMessage(
channel=channel_id,
blocks=reply_blocks
)
else:
await self.slack.client.chat_postMessage(
channel=channel_id,
text=reply_content
)

# Add a reaction to the original message if we have the message ID
if "origin_message" in message:
try:
await self.slack.client.reactions_add(
channel=channel_id,
timestamp=message["origin_message"],
name="thumbsup"
)
except Exception as exc:
self.logger.error("Failed to add reaction: %s", exc)
except Exception as exc:
self.logger.error("Failed to send reply: %s", exc)
traceback.print_exc()

# Try to add a thumbs down reaction to the original message
if "origin_message" in message:
try:
await self.slack.client.reactions_add(
channel=channel_id,
timestamp=message["origin_message"],
name="thumbsdown"
)
except Exception as exc:
self.logger.error("Failed to add reaction: %s", exc)

return True

async def on_message(self, msg: dict[str, Any]):
Expand Down Expand Up @@ -109,27 +231,65 @@ async def on_message(self, msg: dict[str, Any]):
"frontend": "slack",
"channel": msg["event"]["channel"],
"user": msg["event"]["user"],
"origin_message": msg["event"]["ts"], # Slack uses timestamps as message IDs
"trigger": trigger,
"prompt": prompt,
}

# Get user information
if reply["user"] not in self.user_name_cache:
user_info = await self.slack.client.users_info(user=reply["user"]) # type: ignore
self.logger.error("********* FOUND USER INFO: %s", user_info)
self.user_name_cache[reply["user"]] = user_info["user"]["real_name"]
reply["user_name"] = self.user_name_cache[reply["user"]]

# FIXME: Get channel_name and server_name here
# Get channel information
channel_id = reply["channel"]
if channel_id not in self.channel_name_cache:
try:
channel_info = await self.slack.client.conversations_info(channel=channel_id) # type: ignore
self.channel_name_cache[channel_id] = channel_info["channel"]["name"]
except Exception as exc:
self.logger.error("Failed to get channel info: %s", exc)
self.channel_name_cache[channel_id] = "DM"

reply["channel_name"] = self.channel_name_cache[channel_id]

# Get workspace (server) information
try:
workspace_info = await self.slack.client.team_info() # type: ignore
reply["server_name"] = workspace_info["team"]["name"]
reply["server_id"] = workspace_info["team"]["id"]
except Exception as exc:
self.logger.error("Failed to get workspace info: %s", exc)
reply["server_name"] = "Slack"
reply["server_id"] = "unknown"

# Check if the message has an image
if "files" in msg["event"] and msg["event"]["files"]:
for file in msg["event"]["files"]:
if file["filetype"] in ["png", "jpg", "jpeg", "gif"]:
reply["image_url"] = file["url_private"]
break

self.logger.info("INPUT: %s %s", self.log_slug(reply), text)

# Publish the trigger
try:
await self.callback_send_workload(reply)
# FIXME: Add thumbs-up reaction
# Add a thumbs-up reaction to the message
await self.slack.client.reactions_add(
channel=channel_id,
timestamp=msg["event"]["ts"],
name="thumbsup"
)
except Exception:
traceback.print_exc()
# FIXME: Add thumbs-down reaction
# Add a thumbs-down reaction to the message
await self.slack.client.reactions_add(
channel=channel_id,
timestamp=msg["event"]["ts"],
name="thumbsdown"
)

def log_slug(self, resp: dict[str, str]) -> str:
"""Return a string to identify a message in logs."""
Expand Down