Xory/huginn
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

i guess this counts as img & pdf recog

Browse source
This commit is contained in:
xory 2025-08-15 13:31:56 +00:00
parent 65fbb69733
commit 2750b5fa78
1 changed files with 42 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

36
main.py
View file

@ -5,6 +5,7 @@ from discord import app_commands
from discord.ext import commands from discord.ext import commands
from tools import searxng, open_url, run_command from tools import searxng, open_url, run_command
from traceback import print_exc from traceback import print_exc
from typing import Optional
import asyncio import asyncio
import os import os
import io import io
@ -38,11 +39,22 @@ async def generation(interaction: discord.Interaction) -> None:
@bot.tree.command(name="ask", description="ai thing yes 👍") @bot.tree.command(name="ask", description="ai thing yes 👍")
@app_commands.allowed_installs(guilds=False, users=True) @app_commands.allowed_installs(guilds=False, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True) @app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def ask(interaction: discord.Interaction, prompt: str) -> None: async def ask(interaction: discord.Interaction, prompt: str, attachment: Optional[discord.Attachment]) -> None:
await interaction.response.defer() await interaction.response.defer()
attachment_text: str | None = None
attachment_data: bytes | None = None
content_type: str | None = None
if attachment:
attachment_data = await attachment.read()
content_type = attachment.content_type
if content_type and ('text/plain' in content_type or 'text/markdown' in content_type):
attachment_text = attachment_data.decode("utf-8")
response: types.GenerateContentResponse | None = None response: types.GenerateContentResponse | None = None
for _ in range(5): for _ in range(5):
try: try:
if not attachment:
response = await client.aio.models.generate_content( response = await client.aio.models.generate_content(
model="gemini-2.5-flash", model="gemini-2.5-flash",
contents=[ contents=[
@ -50,16 +62,38 @@ async def ask(interaction: discord.Interaction, prompt: str) -> None:
], ],
config=config config=config
) )
elif attachment and attachment_text:
response = await client.aio.models.generate_content(
model="gemini-2.5-flash",
contents=[
prompt,
types.Part.from_text(text=attachment_text)
],
config=config
)
elif content_type and attachment_data and ("image/" in content_type or "application/pdf" in content_type):
response = await client.aio.models.generate_content(
model="gemini-2.5-flash",
contents=[
prompt,
types.Part.from_bytes(
data=attachment_data, mime_type=content_type)
],
config=config
)
break break
except: except:
print_exc() print_exc()
await asyncio.sleep(10) await asyncio.sleep(10)
continue continue
if not response: if not response:
print("e")
await interaction.edit_original_response(content="`[E] API Error. Please try again later.`") await interaction.edit_original_response(content="`[E] API Error. Please try again later.`")
return return
if not response.text: if not response.text:
await interaction.edit_original_response(content="`[E] Model returned no response`") await interaction.edit_original_response(content="`[E] Model returned no response`")
generation: str = response.text or "" generation: str = response.text or ""
if len(generation) > 2000: if len(generation) > 2000:
buffer = io.BytesIO(generation.encode("utf-8")) buffer = io.BytesIO(generation.encode("utf-8"))