Applying previous commit.
This commit is contained in:
parent
f3ec18a6a5
commit
da5a4eddaa
@ -18,6 +18,7 @@ import shutil # For backing up files
|
||||
from typing import Optional, List, Dict, Any, Tuple # For type hinting
|
||||
import asyncio
|
||||
import aiofiles
|
||||
import re
|
||||
|
||||
# Google Generative AI Imports (using Vertex AI backend)
|
||||
from google import genai
|
||||
@ -446,6 +447,55 @@ class AIModerationCog(commands.Cog):
|
||||
print(f"Error processing video: {e}")
|
||||
return None, None
|
||||
|
||||
async def process_url_attachment(self, url: str) -> tuple[str, bytes, str, str]:
|
||||
"""Fetch an attachment from a direct link."""
|
||||
import aiohttp
|
||||
|
||||
try:
|
||||
cleaned_url = url.strip("<>")
|
||||
filename = cleaned_url.split("/")[-1].split("?")[0]
|
||||
_, ext = os.path.splitext(filename.lower())
|
||||
if ext in self.image_extensions:
|
||||
attachment_type = "image"
|
||||
elif ext in self.gif_extensions:
|
||||
attachment_type = "gif"
|
||||
elif ext in self.video_extensions:
|
||||
attachment_type = "video"
|
||||
else:
|
||||
return None, None, None, None
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(cleaned_url) as resp:
|
||||
if resp.status != 200:
|
||||
print(
|
||||
f"Failed to fetch URL attachment {cleaned_url}: {resp.status}"
|
||||
)
|
||||
return None, None, None, None
|
||||
data = await resp.read()
|
||||
mime_type = resp.headers.get(
|
||||
"Content-Type", f"image/{ext.lstrip('.')}"
|
||||
)
|
||||
return mime_type, data, attachment_type, filename
|
||||
except Exception as e:
|
||||
print(f"Error processing URL attachment {url}: {e}")
|
||||
return None, None, None, None
|
||||
|
||||
def extract_direct_attachment_urls(self, text: str) -> List[str]:
|
||||
"""Return a list of direct image/video URLs found in the text."""
|
||||
|
||||
urls = re.findall(r"https?://\S+", text or "")
|
||||
allowed_exts = (
|
||||
self.image_extensions + self.gif_extensions + self.video_extensions
|
||||
)
|
||||
results = []
|
||||
for u in urls:
|
||||
cleaned = u.strip("<>")
|
||||
path = cleaned.split("?")[0]
|
||||
_, ext = os.path.splitext(path.lower())
|
||||
if ext in allowed_exts:
|
||||
results.append(cleaned)
|
||||
return results
|
||||
|
||||
# --- AI Moderation Command Group ---
|
||||
aimod_group = app_commands.Group(
|
||||
name="aimod", description="AI Moderation commands."
|
||||
@ -2077,8 +2127,13 @@ CRITICAL: Do NOT output anything other than the required JSON response.
|
||||
if message.author.bot:
|
||||
print(f"Ignoring message {message.id} from bot.")
|
||||
return
|
||||
# Ignore messages without content or attachments
|
||||
if not message.content and not message.attachments:
|
||||
link_urls = (
|
||||
self.extract_direct_attachment_urls(message.content)
|
||||
if message.content
|
||||
else []
|
||||
)
|
||||
# Ignore messages without content, attachments, or direct attachment links
|
||||
if not message.content and not message.attachments and not link_urls:
|
||||
print(f"Ignoring message {message.id} with no content or attachments.")
|
||||
return
|
||||
# Ignore DMs
|
||||
@ -2125,6 +2180,27 @@ CRITICAL: Do NOT output anything other than the required JSON response.
|
||||
f"Processed {len(image_data_list)} attachments for message {message.id}"
|
||||
)
|
||||
|
||||
# Check for direct link attachments in the message content
|
||||
if link_urls:
|
||||
processed_links = 0
|
||||
for url in link_urls:
|
||||
mime_type, image_bytes, attachment_type, filename = (
|
||||
await self.process_url_attachment(url)
|
||||
)
|
||||
if mime_type and image_bytes and attachment_type:
|
||||
image_data_list.append(
|
||||
(mime_type, image_bytes, attachment_type, filename)
|
||||
)
|
||||
processed_links += 1
|
||||
print(
|
||||
f"Processed linked attachment: {filename} as {attachment_type}"
|
||||
)
|
||||
|
||||
if processed_links > 0:
|
||||
print(
|
||||
f"Processed {processed_links} linked attachments for message {message.id}"
|
||||
)
|
||||
|
||||
# Only proceed with AI analysis if there's text to analyze or attachments
|
||||
if not message_content and not image_data_list:
|
||||
print(
|
||||
|
Loading…
x
Reference in New Issue
Block a user