mirror of
https://github.com/alterware/aw-bot.git
synced 2025-11-18 00:52:09 +00:00
feat: basic image upload to AI model (still no context window)
This commit is contained in:
@@ -17,6 +17,7 @@ COPY aw.py .
|
|||||||
COPY LICENSE .
|
COPY LICENSE .
|
||||||
|
|
||||||
ENV BOT_TOKEN=""
|
ENV BOT_TOKEN=""
|
||||||
|
ENV GOOGLE_API_KEY=""
|
||||||
|
|
||||||
# Where the database will be stored
|
# Where the database will be stored
|
||||||
ENV BOT_DATA_DIR=""
|
ENV BOT_DATA_DIR=""
|
||||||
|
|||||||
@@ -1,20 +1,46 @@
|
|||||||
import os
|
import os
|
||||||
|
import requests
|
||||||
|
from google.genai import types
|
||||||
from google import genai
|
from google import genai
|
||||||
|
|
||||||
API_KEY = os.getenv("GOOGLE_API_KEY")
|
API_KEY = os.getenv("GOOGLE_API_KEY")
|
||||||
|
|
||||||
|
|
||||||
async def forward_to_google_api(prompt):
|
async def forward_to_google_api(prompt, image_object=None):
|
||||||
"""
|
"""
|
||||||
Forwards the message content and optional image URL to a Google API.
|
Forwards the message content and optional image object to a Google API.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
prompt (discord.Message): The message object to forward.
|
prompt (discord.Message): The message object to forward.
|
||||||
|
image_object (tuple, optional): A tuple containing the image URL and its MIME type (e.g., ("url", "image/jpeg")).
|
||||||
"""
|
"""
|
||||||
|
if not API_KEY:
|
||||||
|
await prompt.reply(
|
||||||
|
"Google API key is not set. Please contact the administrator.",
|
||||||
|
mention_author=True,
|
||||||
|
)
|
||||||
|
return
|
||||||
|
|
||||||
client = genai.Client(api_key=API_KEY)
|
client = genai.Client(api_key=API_KEY)
|
||||||
|
|
||||||
|
input = [prompt.content]
|
||||||
|
if image_object:
|
||||||
|
try:
|
||||||
|
image_url, mime_type = image_object
|
||||||
|
image = requests.get(image_url)
|
||||||
|
image.raise_for_status()
|
||||||
|
input.append(types.Part.from_bytes(data=image.content, mime_type=mime_type))
|
||||||
|
except requests.RequestException:
|
||||||
|
await prompt.reply(f"Failed to fetch the image", mention_author=True)
|
||||||
|
return
|
||||||
|
|
||||||
response = client.models.generate_content(
|
response = client.models.generate_content(
|
||||||
model="gemini-2.0-flash", contents=prompt.content
|
model="gemini-2.0-flash",
|
||||||
|
contents=input,
|
||||||
|
config=types.GenerateContentConfig(
|
||||||
|
max_output_tokens=400,
|
||||||
|
system_instruction="You are a Discord chat bot named 'AlterWare' who helps users. You should limit your answers to be less than 2000 characters.",
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
await prompt.reply(
|
await prompt.reply(
|
||||||
|
|||||||
@@ -317,7 +317,19 @@ async def handle_message(message, bot):
|
|||||||
staff_role = message.guild.get_role(STAFF_ROLE_ID)
|
staff_role = message.guild.get_role(STAFF_ROLE_ID)
|
||||||
member = message.guild.get_member(message.author.id)
|
member = message.guild.get_member(message.author.id)
|
||||||
if staff_role in member.roles:
|
if staff_role in member.roles:
|
||||||
await forward_to_google_api(message)
|
image_object = None
|
||||||
|
|
||||||
|
for attachment in message.attachments:
|
||||||
|
if attachment.filename.lower().endswith(
|
||||||
|
".jpg"
|
||||||
|
) or attachment.filename.lower().endswith(".jpeg"):
|
||||||
|
image_object = (attachment.url, "image/jpeg")
|
||||||
|
break
|
||||||
|
elif attachment.filename.lower().endswith(".png"):
|
||||||
|
image_object = (attachment.url, "image/png")
|
||||||
|
break
|
||||||
|
|
||||||
|
await forward_to_google_api(message, image_object)
|
||||||
return
|
return
|
||||||
|
|
||||||
# Too many mentions
|
# Too many mentions
|
||||||
|
|||||||
Reference in New Issue
Block a user