swarms/apps/BingBot/cogs/edgegpt.py

import os
import discord
import json
from typing import Optional
from EdgeGPT.ImageGen import ImageGenAsync, ImageGen
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
from dotenv import load_dotenv

load_dotenv()

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

    auth_cookie = os.environ.get("AUTH_COOKIE")
    auth_cookie_SRCHHPGUSR = os.environ.get("AUTH_COOKIE_SRCHHPGUSR")
    # auth_cookie_SRCHHPGUSR = os.environ.get("AUTH_COOKIE_SRCHHPGUSR")
    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))