diff --git a/main.py b/main.py index ea2c2dd..1663b4f 100644 --- a/main.py +++ b/main.py @@ -4,6 +4,8 @@ from dotenv import load_dotenv from discord import app_commands from discord.ext import commands from tools import searxng, open_url, run_command +from traceback import print_exc +import asyncio import os import io import discord @@ -38,11 +40,24 @@ async def generation(interaction: discord.Interaction) -> None: @app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True) async def ask(interaction: discord.Interaction, prompt: str) -> None: await interaction.response.defer() - response = await client.aio.models.generate_content( - model="gemini-2.5-flash", - contents=prompt, - config=config - ) + response: types.GenerateContentResponse | None = None + for _ in range(5): + try: + response = await client.aio.models.generate_content( + model="gemini-2.5-flash", + contents=[ + prompt + ], + config=config + ) + break + except: + print_exc() + await asyncio.sleep(10) + continue + if not response: + await interaction.edit_original_response(content="`[E] API Error. Please try again later.`") + return if not response.text: await interaction.edit_original_response(content="`[E] Model returned no response`") generation: str = response.text or "" diff --git a/tools.py b/tools.py index 8b66676..f12e2a1 100644 --- a/tools.py +++ b/tools.py @@ -33,13 +33,6 @@ SUPPORTED_TEXT_MIMETYPES = [ "application/x-yaml", ] -SUPPORTED_IMAGE_DOCUMENT_MIMETYPES = [ - "application/pdf", - "image/png", - "image/apng", - "image/jpeg" -] - async def searxng(query: str) -> list: """ @@ -89,7 +82,7 @@ async def searxng(query: str) -> list: return results -async def open_url(url: str) -> dict | types.Part: +async def open_url(url: str) -> dict: """ Opens a URL and returns its full content (if it's HTML, it will be converted to clean Markdown). Use this when a `search` result's content is insufficient or when a user provides a direct URL to analyze. @@ -116,19 +109,13 @@ async def open_url(url: str) -> dict | types.Part: content_type = response.content_type.split(";")[0].strip() content_length = response.content_length or 0 - if content_type not in SUPPORTED_TEXT_MIMETYPES + SUPPORTED_IMAGE_DOCUMENT_MIMETYPES: + if content_type not in SUPPORTED_TEXT_MIMETYPES: return { "content_type": content_type, "content_length": content_length, "content": None, } - if content_type in SUPPORTED_IMAGE_DOCUMENT_MIMETYPES: - return types.Part.from_bytes( - data=await response.read(), - mime_type=content_type - ) - if "text/html" in content_type: content = markdownify(await response.text()) if len(content) > 262144: