slipstream/discordbot
slipstream/discordbot
1
1
Fork 0
Code Issues Pull Requests 3 Actions Packages Projects Releases Wiki Activity

feat: Integrate Google GenAI for content moderation

Browse Source 
Integrates Google Generative AI (Vertex AI backend) into the `Rule34Cog` to enable content moderation.

- Adds `google.genai` imports and initializes the GenAI client.
- Defines `STANDARD_SAFETY_SETTINGS` with all thresholds set to `BLOCK_NONE` as requested.
- Includes a helper function `_get_response_text` for safely extracting text from GenAI responses.
- Sets up the GenAI client using `PROJECT_ID` and `LOCATION` from `gurt.config`.
This commit is contained in:
Slipstream 2025-05-29 22:41:14 -06:00
parent 4d84ee7f6f
commit 5a618f80a0
Signed by: slipstream
GPG Key ID: 13E498CE010AC6FD
1 changed files with 193 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

212
cogs/rule34_cog.py
View File

@ -3,12 +3,77 @@ from discord.ext import commands, tasks
from discord import app_commands
import typing # Need this for Optional
import logging
import re # For _get_response_text, if copied
# Google Generative AI Imports (using Vertex AI backend)
from google import genai
from google.genai import types
from google.api_core import exceptions as google_exceptions
# Import project configuration for Vertex AI
from gurt.config import PROJECT_ID, LOCATION
from .gelbooru_watcher_base_cog import GelbooruWatcherBaseCog
# Setup logger for this cog
log = logging.getLogger(__name__)
# Define standard safety settings using google.generativeai types
# Set all thresholds to OFF as requested
STANDARD_SAFETY_SETTINGS = [
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold="BLOCK_NONE"),
types.SafetySetting(category=types.HarmCategory.HARM_CATEGORY_HARASSMENT, threshold="BLOCK_NONE"),
]
# --- Helper Function to Safely Extract Text (copied from teto_cog.py) ---
def _get_response_text(response: typing.Optional[types.GenerateContentResponse]) -> typing.Optional[str]:
"""
Safely extracts the text content from the first text part of a GenerateContentResponse.
Handles potential errors and lack of text parts gracefully.
"""
if not response:
# log.debug("[_get_response_text] Received None response object.")
return None
if hasattr(response, 'text') and response.text:
# log.debug("[_get_response_text] Found text directly in response.text attribute.")
return response.text
if not response.candidates:
# log.debug(f"[_get_response_text] Response object has no candidates. Response: {response}")
return None
try:
candidate = response.candidates[0]
if not hasattr(candidate, 'content') or not candidate.content:
# log.debug(f"[_get_response_text] Candidate 0 has no 'content'. Candidate: {candidate}")
return None
if not hasattr(candidate.content, 'parts') or not candidate.content.parts:
# log.debug(f"[_get_response_text] Candidate 0 content has no 'parts' or parts list is empty. types.Content: {candidate.content}")
return None
for i, part in enumerate(candidate.content.parts):
if hasattr(part, 'text') and part.text is not None:
if isinstance(part.text, str) and part.text.strip():
# log.debug(f"[_get_response_text] Found non-empty text in part {i}.")
return part.text
else:
# log.debug(f"[_get_response_text] types.Part {i} has 'text' attribute, but it's empty or not a string: {part.text!r}")
pass # Continue searching
# log.debug(f"[_get_response_text] No usable text part found in candidate 0 after iterating through all parts.")
return None
except (AttributeError, IndexError, TypeError) as e:
# log.warning(f"[_get_response_text] Error accessing response structure: {type(e).__name__}: {e}")
# log.debug(f"Problematic response object: {response}")
return None
except Exception as e:
# log.error(f"[_get_response_text] Unexpected error extracting text: {e}")
# log.debug(f"Response object during error: {response}")
return None
class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
# Define the command group specific to this cog
r34watch = app_commands.Group(name="r34watch", description="Manage Rule34 tag watchers for new posts.")
@ -18,27 +83,105 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
bot=bot,
cog_name="Rule34",
api_base_url="https://api.rule34.xxx/index.php",
default_tags="kasane_teto breast_milk", # Example default
default_tags="kasane_teto breast_milk", # Example default, will be overridden if tags are required
is_nsfw_site=True,
command_group_name="r34watch", # For potential use in base class messages
main_command_name="rule34" # For potential use in base class messages
)
# Initialize Google GenAI Client for Vertex AI
try:
if PROJECT_ID and LOCATION:
self.genai_client = genai.Client(
vertexai=True,
project=PROJECT_ID,
location=LOCATION,
)
log.info(f"Rule34Cog: Google GenAI Client initialized for Vertex AI project '{PROJECT_ID}' in location '{LOCATION}'.")
else:
self.genai_client = None
log.warning("Rule34Cog: PROJECT_ID or LOCATION not found in config. Google GenAI Client not initialized.")
except Exception as e:
self.genai_client = None
log.error(f"Rule34Cog: Error initializing Google GenAI Client for Vertex AI: {e}")
self.tag_transformer_model = "gemini-2.0-flash-lite-001" # Hardcoded as per request
# The __init__ in base class handles session creation and task starting.
async def _transform_tags_ai(self, user_tags: str) -> typing.Optional[str]:
"""Transforms user-provided tags into rule34-style tags using AI."""
if not self.genai_client:
log.warning("Rule34Cog: GenAI client not initialized, cannot transform tags.")
return None
if not user_tags:
return "" # Return empty if no tags provided to transform
system_prompt_text = (
"You are an AI assistant specialized in transforming user-provided text into tags suitable for rule34.xxx. "
"Your task is to convert natural language input into a space-separated list of tags, where multi-word concepts are joined by underscores. "
"For example, if the user provides 'hatsune miku blue hair', you should output 'hatsune_miku blue_hair'. "
"Only output the transformed tags. Do not add any other text, explanations, or greetings."
)
prompt_parts = [
types.Part(text=system_prompt_text),
types.Part(text=f"User input: \"{user_tags}\""),
types.Part(text="Transformed tags:"),
]
contents_for_api = [types.Content(role="user", parts=prompt_parts)]
generation_config = types.GenerateContentConfig(
temperature=0.2, # Low temperature for more deterministic output
max_output_tokens=256,
safety_settings=STANDARD_SAFETY_SETTINGS,
)
try:
log.debug(f"Rule34Cog: Sending to Vertex AI for tag transformation. Model: {self.tag_transformer_model}, Input: '{user_tags}'")
response = await self.genai_client.aio.models.generate_content(
model=f"publishers/google/models/{self.tag_transformer_model}",
contents=contents_for_api,
config=generation_config,
)
transformed_tags = _get_response_text(response)
if transformed_tags:
log.info(f"Rule34Cog: Tags transformed: '{user_tags}' -> '{transformed_tags.strip()}'")
return transformed_tags.strip()
else:
log.warning(f"Rule34Cog: AI tag transformation returned empty for input: '{user_tags}'. Response: {response}")
return None
except google_exceptions.GoogleAPICallError as e:
log.error(f"Rule34Cog: Vertex AI API call failed for tag transformation: {e}")
return None
except Exception as e:
log.error(f"Rule34Cog: Unexpected error during AI tag transformation: {e}")
return None
# --- Prefix Command ---
@commands.command(name="rule34")
async def rule34_prefix(self, ctx: commands.Context, *, tags: typing.Optional[str] = None):
async def rule34_prefix(self, ctx: commands.Context, *, tags: str):
"""Search for images on Rule34 with the provided tags."""
actual_tags = tags or self.default_tags
# The base class _prefix_command_logic needs to be adjusted or called differently.
# For now, let's replicate the simple call structure.
# We need to ensure the loading message is handled correctly.
loading_msg = await ctx.reply(f"Fetching data from {self.cog_name}, please wait...")
response = await self._fetch_posts_logic("prefix_internal", actual_tags, hidden=False)
if not tags: # Should not happen due to 'tags: str' but as a safeguard
await ctx.reply("Please provide tags to search for.")
return
loading_msg = await ctx.reply(f"Transforming tags and fetching data from {self.cog_name}, please wait...")
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None: # AI transformation failed
await loading_msg.edit(content=f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`")
return
if not transformed_tags: # AI returned empty
await loading_msg.edit(content=f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.")
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Prefix): Using final tags: '{final_tags}' from original: '{tags}'")
response = await self._fetch_posts_logic("prefix_internal", final_tags, hidden=False)
if isinstance(response, tuple):
content, all_results = response
view = self.GelbooruButtons(self, actual_tags, all_results, hidden=False)
view = self.GelbooruButtons(self, final_tags, all_results, hidden=False) # Pass final_tags to buttons
await loading_msg.edit(content=content, view=view)
elif isinstance(response, str): # Error
await loading_msg.edit(content=response, view=None)
@ -47,24 +190,51 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
# --- Slash Command ---
@app_commands.command(name="rule34", description="Get random image from Rule34 with specified tags")
@app_commands.describe(
tags="The tags to search for (e.g., 'kasane_teto rating:safe')",
tags="The tags to search for (e.g., 'hatsune miku rating:safe')", # Updated example
hidden="Set to True to make the response visible only to you (default: False)"
)
async def rule34_slash(self, interaction: discord.Interaction, tags: typing.Optional[str] = None, hidden: bool = False):
async def rule34_slash(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
"""Slash command version of rule34."""
actual_tags = tags or self.default_tags
await self._slash_command_logic(interaction, actual_tags, hidden)
await interaction.response.defer(thinking=True, ephemeral=hidden)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None:
await interaction.followup.send(f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`", ephemeral=True)
return
if not transformed_tags:
await interaction.followup.send(f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.", ephemeral=True)
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Slash): Using final tags: '{final_tags}' from original: '{tags}'")
# _slash_command_logic handles its own deferral/response
await self._slash_command_logic(interaction, final_tags, hidden, already_deferred=True)
# --- New Browse Command ---
@app_commands.command(name="rule34browse", description="Browse Rule34 results with navigation buttons")
@app_commands.describe(
tags="The tags to search for (e.g., 'kasane_teto rating:safe')",
tags="The tags to search for (e.g., 'hatsune miku rating:safe')", # Updated example
hidden="Set to True to make the response visible only to you (default: False)"
)
async def rule34_browse_slash(self, interaction: discord.Interaction, tags: typing.Optional[str] = None, hidden: bool = False):
async def rule34_browse_slash(self, interaction: discord.Interaction, tags: str, hidden: bool = False):
"""Browse Rule34 results with navigation buttons."""
actual_tags = tags or self.default_tags
await self._browse_slash_command_logic(interaction, actual_tags, hidden)
await interaction.response.defer(thinking=True, ephemeral=hidden)
transformed_tags = await self._transform_tags_ai(tags)
if transformed_tags is None:
await interaction.followup.send(f"Sorry, I couldn't transform the tags using AI. Please try again or use rule34-formatted tags directly. Original tags: `{tags}`", ephemeral=True)
return
if not transformed_tags:
await interaction.followup.send(f"Sorry, the AI couldn't understand the tags provided: `{tags}`. Please try rephrasing.", ephemeral=True)
return
final_tags = f"{transformed_tags} -ai_generated"
log.info(f"Rule34Cog (Browse): Using final tags: '{final_tags}' from original: '{tags}'")
# _browse_slash_command_logic handles its own deferral/response
await self._browse_slash_command_logic(interaction, final_tags, hidden, already_deferred=True)
# --- r34watch slash command group ---
# All subcommands will call the corresponding _watch_..._logic methods from the base class.
@ -79,7 +249,9 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
@app_commands.checks.has_permissions(manage_guild=True)
async def r34watch_add(self, interaction: discord.Interaction, tags: str, channel: typing.Union[discord.TextChannel, discord.ForumChannel], thread_target: typing.Optional[str] = None, post_title: typing.Optional[str] = None):
await interaction.response.defer(ephemeral=True) # Defer here before calling base logic
await self._watch_add_logic(interaction, tags, channel, thread_target, post_title)
final_tags = f"{tags} -ai_generated"
log.info(f"Rule34Cog (Watch Add): Using final tags for watch: '{final_tags}' from original: '{tags}'")
await self._watch_add_logic(interaction, final_tags, channel, thread_target, post_title)
@r34watch.command(name="request", description="Request a new Rule34 tag watch (requires moderator approval).")
@app_commands.describe(
@ -89,7 +261,9 @@ class Rule34Cog(GelbooruWatcherBaseCog): # Removed name="Rule34"
)
async def r34watch_request(self, interaction: discord.Interaction, tags: str, forum_channel: discord.ForumChannel, post_title: typing.Optional[str] = None):
await interaction.response.defer(ephemeral=True)
await self._watch_request_logic(interaction, tags, forum_channel, post_title)
final_tags = f"{tags} -ai_generated"
log.info(f"Rule34Cog (Watch Request): Using final tags for watch request: '{final_tags}' from original: '{tags}'")
await self._watch_request_logic(interaction, final_tags, forum_channel, post_title)
@r34watch.command(name="pending_list", description="Lists all pending Rule34 watch requests.")
@app_commands.checks.has_permissions(manage_guild=True)