openaitelegram.py

import openai
import json
import argparse
import asyncio
import os
import logging
import urllib.request
from PIL import Image
from io import BytesIO
from aiogram import Bot, types, executor
from aiogram.dispatcher import Dispatcher, middlewares
from aiogram.utils.executor import start_polling
from aiogram.types import InputFile, ReplyKeyboardMarkup, ReplyKeyboardRemove, KeyboardButton
from aiogram.dispatcher.middlewares import LifetimeControllerMiddleware
from misc.middleware.throttling import rate_limit
from misc.text.prompt_frameworks import DAN_prompt
from api_key import key, bot_token

# key = openai.api_key = os.environ.get("OPENAI_KEY")
# bot_token = os.environ.get("BOT_TOKEN")

bot = Bot(token=bot_token)
dp = Dispatcher(bot)

dp.middleware.setup(LifetimeControllerMiddleware())


available_models = {
    'davinci03': "text-davinci-003",
    'davinci02': "text-davinci-002",
    'curie': "text-curie-001",
    'babbage': "text-babbage-001",
    'ada': "text-ada-001"
}


def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="text-davinci-003",
                        help="GPT-3 model to use. Available options: davinci03, davinci02, curie, babbage, ada")
    return parser.parse_args()


@dp.message_handler(commands=['start'])
async def start_message(message: types.Message):
    await message.reply(
        "Welcome to the GPT-3 bot! Use the command /gpt, /gpt4 or /dan for text and, /dalle for image, followed by your prompt to generate a response. Use the command /help to see the available models.")


@dp.message_handler(commands=['gpt'])
@rate_limit(10, key="gpt")
async def handle_gpt(message):
    args = get_arguments()
    models = args.model
    prompt = message.text.replace("/gpt ", "")
    if "--model" in prompt:
        models = prompt.split("--model")[-1].strip()

    response = openai.Completion.create(
        engine=models, prompt=prompt, temperature=1, max_tokens=3600)

    response_str = json.dumps(response)
    json_response = json.loads(response_str)
    model_response = json_response['choices'][0]['text']

    await message.reply(model_response)

@dp.message_handler(commands=['gpt4'])
@rate_limit(10, key="gpt4")
async def handle_chat(message):
    user_prompt = message.text.replace("/gpt4 ", "")
    api_message = [
        {"role": "system", "content": "You are an assistant"},
        {"role": "user", "content": user_prompt},
        {"role": "assistant", "content": "Proceed to answer as a helpful assistant"}
    ]

    chat = openai.ChatCompletion.create(
        model="gpt-4",
        messages=api_message
    )

    response_str = json.dumps(chat)
    json_response = json.loads(response_str)
    gpt4_response = json_response['choices'][0]['message']['content']

    await message.reply(gpt4_response)

@dp.message_handler(commands=['dan'])
@rate_limit(10, key="dan")
async def handle_chat(message):
    user_prompt = message.text.replace("/dan ", "")
    api_message = [
        {"role": "system", "content": f"This is you, DAN: {DAN_prompt}"},
        {"role": "user", "content": user_prompt},
        {"role": "assistant", "content": "I understand. Proceed to answer as DAN"}
    ]

    chat = openai.ChatCompletion.create(
        model="gpt-4",
        messages=api_message
    )

    response_str = json.dumps(chat)
    json_response = json.loads(response_str)
    dan_response = json_response['choices'][0]['message']['content']

    await message.reply(dan_response)


@dp.message_handler(commands=['dalle'])
@rate_limit(10, key="dalle")
async def handle_dalle(message):
    dalle_prompt = message.text.replace("/dalle ", "")
    chat_id = message.chat.id

    photo_gen = openai.Image.create(
        prompt=dalle_prompt,
        n=1,
        size="256x256"
    )

    image_url = photo_gen['data'][0]['url']
    pngUrl = urllib.request.urlopen(image_url)

    await bot.send_photo(chat_id, photo=pngUrl, allow_sending_without_reply=True)


@dp.message_handler(commands=['model'])
async def handle_model(message: types.Message):
    # Choose model from dict
    model = message.text.split()[1]
    models = {'davinci03': "text-davinci-003", 'davinci02': "text-davinci-002",
              'curie': "text-curie-001", 'babbage': "text-babbage-001", 'ada': "text-ada-001"}
    if model not in models:
        await message.reply("Invalid model selected. Please use /help to see available models.")
    else:
        global current_model
        current_model = models[model]
        await message.reply(f"Model changed to {current_model}.")


@dp.message_handler(commands=['help'])
async def handle_help(message: types.Message):
    # Set help message
    available_models = "davinci03, davinci02, curie, babbage, ada"
    await message.reply("Gpt-4 use command /dan or /gpt4. Use the command /gpt or /dan for text. /dalle, followed by prompt to generate an image.")
    await message.reply(f"Use /model option followed by the desired model when running the script. Available models: {available_models}")


@dp.message_handler(commands=['count_users'])
async def count_users(message: types.Message):
    # Get the chat ID from the message
    chat_id = message.chat.id

    # Get the chat members count
    chat = await bot.get_chat(chat_id)
    count = await chat.get_members_count()

    # Send the count back to the user
    await message.reply(f"There are {count} users in this chat.")

executor.start_polling(dp)