model selection & minor sysprompt edits

This commit is contained in:
Xory 2025-10-23 21:34:01 +03:00
parent 118218cc53
commit ebd8d9ad81
3 changed files with 26 additions and 9 deletions

22
main.py
View file

@ -5,7 +5,7 @@ from discord import app_commands
from discord.ext import commands
from tools import searxng, run_command, open_url
from traceback import print_exc
from typing import Optional
from typing import Optional, Literal
import asyncio
import os
import io
@ -66,7 +66,7 @@ async def generation(interaction: discord.Interaction) -> None:
@bot.tree.command(name="ask", description="ai thing yes 👍")
@app_commands.allowed_installs(guilds=False, users=True)
@app_commands.allowed_contexts(guilds=True, dms=True, private_channels=True)
async def ask(interaction: discord.Interaction, prompt: str, attachment: Optional[discord.Attachment], ephemeral: Optional[bool]) -> None:
async def ask(interaction: discord.Interaction, prompt: str, attachment: Optional[discord.Attachment], ephemeral: Optional[bool], model: Optional[Literal["pro", "flash", "flash-lite"]]) -> None:
if not ephemeral:
await interaction.response.defer()
else:
@ -82,11 +82,23 @@ async def ask(interaction: discord.Interaction, prompt: str, attachment: Optiona
attachment_text = attachment_data.decode("utf-8")
response: types.GenerateContentResponse | None = None
model_name: str = ""
match model:
case "lite":
model_name = "gemini-flash-lite-latest"
case "pro":
model_name = "gemini-2.5-pro"
case _:
model_name = "gemini-flash-latest"
print(f"[d] using {model_name}") # TODO: delete this
for _ in range(5):
try:
if not attachment:
response = await client.aio.models.generate_content(
model="gemini-2.5-flash",
model=model_name,
contents=[
prompt
],
@ -94,7 +106,7 @@ async def ask(interaction: discord.Interaction, prompt: str, attachment: Optiona
)
elif attachment and attachment_text:
response = await client.aio.models.generate_content(
model="gemini-2.5-flash",
model=model_name,
contents=[
prompt,
types.Part.from_text(text=attachment_text)
@ -103,7 +115,7 @@ async def ask(interaction: discord.Interaction, prompt: str, attachment: Optiona
)
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",
model=model_name,
contents=[
prompt,
types.Part.from_bytes(