edgegpt.py

import discord
import json
from typing import Optional
from EdgeGPT.ImageGen import ImageGenAsync
from EdgeGPT.EdgeGPT import Chatbot
from discord import app_commands
from core.classes import Cog_Extension
from src import log
from src.imageCreate import create_image, get_using_create, set_using_create
from src.response import send_message, get_using_send, set_using_send

logger = log.setup_logger(__name__)

users_chatbot = {}
users_image_generator = {}
user_conversation_style = {}

async def init_chatbot(user_id):
    with open("./cookies.json", encoding="utf-8") as file:
        cookie_json = json.load(file)
    for cookie in cookie_json:
        if cookie.get("name") == "_U":
            auth_cookie = cookie.get("value")
            break
    users_chatbot[user_id] = UserChatbot(cookies=cookie_json)
    users_image_generator[user_id] = ImageGenAsync(auth_cookie, quiet=True)
    user_conversation_style[user_id] = "balanced"

class UserChatbot:
    def __init__(self, cookies):
        self.chatbot = Chatbot(cookies=cookies)

    async def send_message(self, interaction, message, conversation_style):
        await send_message(self.chatbot, interaction, message, conversation_style)

    async def create_image(self, interaction, prompt: str, image_generator):
        await create_image(interaction, prompt, image_generator)

    async def reset(self):
        await self.chatbot.reset()

class EdgeGPT(Cog_Extension):
    # Chat with Bing
    @app_commands.command(name="bing", description="Have a chat with Bing")
    async def bing(self, interaction: discord.Interaction, *, message: str):
        try:
            using = await get_using_send(interaction.user.id)
        except:
            await set_using_send(interaction.user.id, False)
            using = await get_using_send(interaction.user.id)
        if not using:
            await interaction.response.defer(ephemeral=False, thinking=True) 
            username = str(interaction.user)
            usermessage = message
            channel = str(interaction.channel)
            user_id = interaction.user.id
            if user_id not in users_chatbot:
                await init_chatbot(interaction.user.id)
            conversation_style = user_conversation_style[user_id]
            logger.info(f"\x1b[31m{username}\x1b[0m : '{usermessage}' ({channel}) [Style: {conversation_style}]")
            await users_chatbot[user_id].send_message(interaction, usermessage, conversation_style)
        else:
            await interaction.response.defer(ephemeral=True, thinking=True)
            await interaction.followup.send("> **Please wait for your last conversation to finish.**")

    # Reset Bing conversation
    @app_commands.command(name="reset", description="Reset Bing conversation")
    async def reset(self, interaction: discord.Interaction):
        await interaction.response.defer(ephemeral=True, thinking=True)
        user_id = interaction.user.id
        try:
            await users_chatbot[user_id].reset()
            await interaction.followup.send("> **Info: Reset finish.**")
            logger.warning("\x1b[31mBing has been successfully reset\x1b[0m")
        except:
            await interaction.followup.send(f"> **You don't have any conversation yet.**")
            logger.exception("Bing reset failed.")

    # Switch conversation style
    @app_commands.command(name="switch_style", description="Switch conversation style")
    @app_commands.choices(style=[app_commands.Choice(name="Creative", value="creative"), app_commands.Choice(name="Balanced", value="balanced"), app_commands.Choice(name="Precise", value="precise")])
    async def switch_style(self, interaction: discord.Interaction, style: app_commands.Choice[str]):
        await interaction.response.defer(ephemeral=True, thinking=True)
        user_id = interaction.user.id
        if user_id not in users_chatbot:
            await init_chatbot(user_id)
        user_conversation_style[user_id] = style.value
        await interaction.followup.send(f"> **Info: successfull switch conversation style to {style.value}.**")
        logger.warning(f"\x1b[31mConversation style has been successfully switch to {style.value}\x1b[0m")
    
    # Set and delete personal Bing Cookies
    @app_commands.command(name="bing_cookies", description="Set or delete Bing Cookies")
    @app_commands.choices(choice=[app_commands.Choice(name="set", value="set"), app_commands.Choice(name="delete", value="delete")])
    async def cookies_setting(self, interaction: discord.Interaction, choice: app_commands.Choice[str], cookies_file: Optional[discord.Attachment]=None):
        await interaction.response.defer(ephemeral=True, thinking=True)
        user_id = interaction.user.id
        if choice.value == "set":
            try:
                content = json.loads(await cookies_file.read())
                for cookie in content:
                    if cookie.get("name") == "_U":
                        auth_cookie = cookie.get("value")
                        break
                users_image_generator[user_id] = ImageGenAsync(auth_cookie, quiet=True)
                users_chatbot[user_id] = UserChatbot(cookies=content)
                user_conversation_style[user_id] = "balanced" 
                await interaction.followup.send("> **Upload successful!**")
                logger.warning(f"\x1b[31m{interaction.user} set Bing Cookies successful\x1b[0m")
            except:
                await interaction.followup.send("> **Please upload your Bing Cookies.**")
        else:
            try:
                del users_chatbot[user_id]
                del users_image_generator[user_id]
                del user_conversation_style[user_id]
                await interaction.followup.send("> **Delete finish.**")
                logger.warning(f"\x1b[31m{interaction.user} delete Cookies\x1b[0m")
            except:
                await interaction.followup.send("> **You don't have any Bing Cookies.**")

    # Create images
    @app_commands.command(name="create_image", description="generate image by Bing image creator")
    async def create_image(self, interaction: discord.Interaction, *, prompt: str):
        user_id = interaction.user.id
        if interaction.user.id not in users_chatbot:
            await init_chatbot(user_id)
        try:
            using = await get_using_create(user_id)
        except:
            await set_using_create(user_id, False)
            using = await get_using_create(user_id)
        if not using: 
            logger.info(f"\x1b[31m{interaction.user}\x1b[0m : '{prompt}' ({interaction.channel}) [BingImageCreator]")
            await users_chatbot[user_id].create_image(interaction, prompt, users_image_generator[user_id])
        else:
            await interaction.response.defer(ephemeral=True, thinking=True)
            await interaction.followup.send("> **Please wait for your last image to create finish.**")

async def setup(bot):
    await bot.add_cog(EdgeGPT(bot))