2025-05-13 13:52:28 +07:00

81 lines
2.3 KiB
Python

import os
import openai
import json
import nextcord
from nextcord.ext import commands
Enabled = True
# Get these from .env file too
SquogAILink = os.environ.get("AI_URL") ## AI_URL = https://
SquogAIToken = os.environ.get("AI_TOKEN") ## AI_TOKEN = ...
SquogAI = openai.OpenAI(api_key=SquogAIToken, base_url=SquogAILink)
SquogDataFile = open("./userinfo/userinfo.json", "r")
SquogData = json.loads(SquogDataFile.read())
SquogDataFile.close()
SquogUserPreset = {"id": 999, "messages": [{"role": "system", "content": "You must act like a curious creature and never say you are a PC program."}]}
async def getResponse(SquogMessageList):
SquogResponse = SquogAI.chat.completions.create(messages=SquogMessageList, model="gpt-4o", temperature=0.9)
return json.loads(SquogResponse.model_dump_json())
async def find(listIterate: list, param: str, value: any):
found = None
for g in listIterate:
if g[param] == value:
found = g
return found
async def flush():
LocalSquogDataFile = open("./userinfo/userinfo.json", "w")
LocalSquogDataFile.write(json.dumps(SquogData))
LocalSquogDataFile.flush()
LocalSquogDataFile.close()
class AI(commands.Cog):
def __init__(self, bot: commands.Bot):
self.client = bot
@commands.Cog.listener()
async def on_message(self, message: nextcord.Message):
if not Enabled:
return
SquogMention = False
for Mention in message.mentions:
if Mention.id == self.client.user.id:
SquogMention = True
if not SquogMention:
return
message.channel.typing()
SquogUser = await find(SquogData, "id", message.author.id)
if not SquogUser:
SquogUser = SquogUserPreset.copy()
SquogUser["id"] = message.author.id
SquogData.append(SquogUser)
SquogUser["messages"].append({"role": "user", "content": message.content})
try:
SquogResponse = await getResponse(SquogUser["messages"])
SquogUser["messages"].append(SquogResponse["choices"][0]["message"])
await message.reply(SquogResponse["choices"][0]["message"]["content"])
await flush()
except:
await message.reply("Failed to load a reply.")
def setup(bot):
bot.add_cog(AI(bot))