feat: Add tools for user avatar data and role color
Introduces `get_user_avatar_data` to retrieve a user's avatar as base64 encoded image data, allowing the AI to "see" the profile picture directly. This includes special handling in `process_requested_tools` to attach the image as a `types.Part` in the prompt. Also adds `get_user_highest_role_color` to fetch the color and details of a user's highest-positioned role.
This commit is contained in:
parent
d3cf350434
commit
cde506052f
GPG Key ID:
13E498CE010AC6FD
31
gurt/api.py
31
gurt/api.py
@ -653,7 +653,36 @@ async def process_requested_tools(cog: 'GurtCog', function_call: types.FunctionC
|
||||
|
||||
# --- Image URL Detection & Modification ---
|
||||
# Check specific tools and keys known to contain image URLs
|
||||
if function_name == "get_user_avatar_url" and isinstance(modified_result_content, dict):
|
||||
|
||||
# Special handling for get_user_avatar_data to directly use its base64 output
|
||||
if function_name == "get_user_avatar_data" and isinstance(modified_result_content, dict):
|
||||
base64_image_data = modified_result_content.get("base64_data")
|
||||
image_mime_type = modified_result_content.get("content_type")
|
||||
|
||||
if base64_image_data and image_mime_type:
|
||||
try:
|
||||
image_bytes = base64.b64decode(base64_image_data)
|
||||
# Validate MIME type (optional, but good practice)
|
||||
supported_image_mimes = ["image/png", "image/jpeg", "image/webp", "image/heic", "image/heif"]
|
||||
clean_mime_type = image_mime_type.split(';')[0].lower()
|
||||
|
||||
if clean_mime_type in supported_image_mimes:
|
||||
image_part = types.Part(data=image_bytes, mime_type=clean_mime_type)
|
||||
parts_to_return.append(image_part) # Corrected: Add to parts_to_return for this tool's response
|
||||
print(f"Added image part directly from get_user_avatar_data (MIME: {clean_mime_type}, {len(image_bytes)} bytes).")
|
||||
# Replace base64_data in the textual response to avoid sending it twice
|
||||
modified_result_content["base64_data"] = "[Image Content Attached In Prompt]"
|
||||
modified_result_content["content_type"] = f"[MIME type: {clean_mime_type} - Content Attached In Prompt]"
|
||||
else:
|
||||
print(f"Warning: MIME type '{clean_mime_type}' from get_user_avatar_data not in supported list. Not attaching image part.")
|
||||
modified_result_content["base64_data"] = "[Image Data Not Attached - Unsupported MIME Type]"
|
||||
except Exception as e:
|
||||
print(f"Error processing base64 data from get_user_avatar_data: {e}")
|
||||
modified_result_content["base64_data"] = f"[Error Processing Image Data: {e}]"
|
||||
# Prevent generic URL download logic from re-processing this avatar
|
||||
original_image_url = None # Explicitly nullify to skip URL download
|
||||
|
||||
elif function_name == "get_user_avatar_url" and isinstance(modified_result_content, dict):
|
||||
avatar_url_value = modified_result_content.get("avatar_url")
|
||||
if avatar_url_value and isinstance(avatar_url_value, str):
|
||||
original_image_url = avatar_url_value # Store original
|
||||
|
||||
@ -1610,6 +1610,35 @@ def create_tools_list():
|
||||
)
|
||||
# --- End User Profile Tool Declarations ---
|
||||
|
||||
# --- Get User Avatar Data ---
|
||||
tool_declarations.append(
|
||||
FunctionDeclaration(
|
||||
name="get_user_avatar_data",
|
||||
description="Gets the user's avatar URL, content type, and base64 encoded image data. This allows the AI to 'see' the profile picture.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "string", "description": "The User ID of the target user."}
|
||||
},
|
||||
"required": ["user_id"]
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
# --- Get User Highest Role Color ---
|
||||
tool_declarations.append(
|
||||
FunctionDeclaration(
|
||||
name="get_user_highest_role_color",
|
||||
description="Gets the color (hex and RGB) of the user's highest positioned role that has a color applied. Also returns role name and ID.",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"user_id": {"type": "string", "description": "The User ID of the target user."}
|
||||
},
|
||||
"required": ["user_id"]
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
return tool_declarations
|
||||
|
||||
|
||||
@ -12,6 +12,7 @@ import re
|
||||
import traceback # Added for error logging
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Any, Optional, Tuple, Union # Added Union
|
||||
import base64 # Added for avatar data encoding
|
||||
|
||||
# Third-party imports for tools
|
||||
from tavily import TavilyClient
|
||||
@ -2776,6 +2777,89 @@ async def get_user_profile_info(cog: commands.Cog, user_id: str) -> Dict[str, An
|
||||
|
||||
# --- End User Profile Tools ---
|
||||
|
||||
async def get_user_avatar_data(cog: commands.Cog, user_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Gets the user's avatar URL, content type, and base64 encoded image data.
|
||||
This allows the AI to "see" the profile picture.
|
||||
"""
|
||||
print(f"Executing get_user_avatar_data for user ID: {user_id}.")
|
||||
user_obj, error_resp = await _get_user_or_member(cog, user_id)
|
||||
if error_resp: return error_resp
|
||||
if not user_obj: return {"error": f"Failed to retrieve user object for ID {user_id}."}
|
||||
|
||||
avatar_asset = user_obj.display_avatar
|
||||
avatar_url = str(avatar_asset.url)
|
||||
|
||||
if not cog.session:
|
||||
return {"error": "aiohttp session not initialized in cog."}
|
||||
|
||||
try:
|
||||
async with cog.session.get(avatar_url) as response:
|
||||
if response.status == 200:
|
||||
image_bytes = await response.read()
|
||||
content_type = response.headers.get("Content-Type", "application/octet-stream")
|
||||
base64_data = base64.b64encode(image_bytes).decode('utf-8')
|
||||
return {
|
||||
"status": "success",
|
||||
"user_id": user_id,
|
||||
"avatar_url": avatar_url,
|
||||
"content_type": content_type,
|
||||
"base64_data": base64_data,
|
||||
"timestamp": datetime.datetime.now().isoformat()
|
||||
}
|
||||
else:
|
||||
error_message = f"Failed to fetch avatar image from {avatar_url}. Status: {response.status}"
|
||||
print(error_message)
|
||||
return {"error": error_message, "user_id": user_id, "avatar_url": avatar_url}
|
||||
except aiohttp.ClientError as e:
|
||||
error_message = f"Network error fetching avatar from {avatar_url}: {str(e)}"
|
||||
print(error_message)
|
||||
return {"error": error_message, "user_id": user_id, "avatar_url": avatar_url}
|
||||
except Exception as e:
|
||||
error_message = f"Unexpected error fetching avatar data for {user_id}: {str(e)}"
|
||||
print(error_message); traceback.print_exc()
|
||||
return {"error": error_message, "user_id": user_id}
|
||||
|
||||
async def get_user_highest_role_color(cog: commands.Cog, user_id: str) -> Dict[str, Any]:
|
||||
"""
|
||||
Gets the color of the user's highest positioned role that has a color applied.
|
||||
Returns the role name and hex color string.
|
||||
"""
|
||||
print(f"Executing get_user_highest_role_color for user ID: {user_id}.")
|
||||
user_obj, error_resp = await _get_user_or_member(cog, user_id)
|
||||
if error_resp: return error_resp
|
||||
if not user_obj: return {"error": f"Failed to retrieve user object for ID {user_id}."}
|
||||
|
||||
if not isinstance(user_obj, discord.Member):
|
||||
return {"error": f"User {user_id} is not a member of a server in the current context. Cannot get role color.", "user_id": user_id}
|
||||
|
||||
# Roles are already sorted by position, member.roles[0] is @everyone, member.roles[-1] is highest.
|
||||
# We need to iterate from highest to lowest that has a color.
|
||||
highest_colored_role = None
|
||||
for role in reversed(user_obj.roles): # Iterate from highest position downwards
|
||||
if role.color != discord.Color.default(): # Check if the role has a non-default color
|
||||
highest_colored_role = role
|
||||
break # Found the highest role with a color
|
||||
|
||||
if highest_colored_role:
|
||||
return {
|
||||
"status": "success",
|
||||
"user_id": user_id,
|
||||
"role_name": highest_colored_role.name,
|
||||
"role_id": str(highest_colored_role.id),
|
||||
"color_hex": str(highest_colored_role.color), # Returns like #RRGGBB
|
||||
"color_rgb": highest_colored_role.color.to_rgb(), # Returns (r, g, b) tuple
|
||||
"guild_id": str(user_obj.guild.id),
|
||||
"timestamp": datetime.datetime.now().isoformat()
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"status": "no_colored_role",
|
||||
"user_id": user_id,
|
||||
"message": "User has no roles with a custom color.",
|
||||
"guild_id": str(user_obj.guild.id),
|
||||
"timestamp": datetime.datetime.now().isoformat()
|
||||
}
|
||||
|
||||
# --- Tool Mapping ---
|
||||
# This dictionary maps tool names (used in the AI prompt) to their implementation functions.
|
||||
@ -2859,4 +2943,8 @@ TOOL_MAPPING = {
|
||||
"get_user_roles": get_user_roles,
|
||||
"get_user_profile_info": get_user_profile_info,
|
||||
# --- End User Profile Tools ---
|
||||
|
||||
# --- New Profile Picture and Role Color Tools ---
|
||||
"get_user_avatar_data": get_user_avatar_data,
|
||||
"get_user_highest_role_color": get_user_highest_role_color,
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user