uzbekgpt/main.py

import asyncio
import base64
import io
import sys
import time

import gconfig

# from aiogram.methods import SendMessageDraft
from datetime import datetime
from aiogram import Bot, Dispatcher, F
from aiogram.client.default import DefaultBotProperties
from aiogram.types import Update
from aiogram.methods import AnswerGuestQuery, SendRichMessage
from aiogram.types import InlineQueryResultArticle, InputTextMessageContent
from aiogram import BaseMiddleware


from aiogram.enums import ChatType
from aiogram.exceptions import TelegramForbiddenError
from aiogram.types import (
    ChosenInlineResult,
    InlineKeyboardButton,
    InlineKeyboardMarkup,
    InlineQuery,
    InlineQueryResultArticle,
    InputTextMessageContent,
    InputRichMessageContent,
    InputRichMessage,
    Message,
)
from callbacks import *
from statistics import *
from colors import *
from helpers import *
from commands import *
from dotenv import load_dotenv
from random import uniform, shuffle

try:
    from config import *
except:
    gconfig.copy_config()
finally:
    from config import *

from supergenerator import *

load_dotenv()

print(f"{BOLD}{GREEN} -- uzbekgpt - привет как дела. {RESET}")


if not API_TOKEN:
    print(f"{RED}{BOLD} !! suka API_TOKEN поставь в .env либо pashol naxxuy {RESET}")
    sys.exit(1)

try:
    bot = Bot(token=API_TOKEN, default=DefaultBotProperties(parse_mode="Markdown"))
    dp = Dispatcher()
    dp.include_router(router)
except Exception as e:
    print(f"{RED}{BOLD} !! вирус у вас тут ошибка: {e}{RESET}")
    sys.exit(1)


def parse_tools(text):
    ...
    return text


@dp.update.outer_middleware()
async def guest_middleware(handler, event, data):
    # print("e: ", event)
    if event.guest_message:
        message = event.guest_message
    elif event.message:
        if event.message.text and any(event.message.text.startswith(cmd) for cmd in commands):
            return await handler(event, data)
        message = event.message
    else:
        return await handler(event, data)

    chat_type = message.sender_chat.type if message.sender_chat else ""
    user_id = message.sender_chat.id if message.sender_chat else message.from_user.id
    is_channel = (
        ChatType.CHANNEL == chat_type
        and message.from_user
        and message.from_user.id == 777000
    )
    current_time = time.time()
    config = get_user_config(user_id)
    stream = config.get("stream")
    user_text = message.text or message.caption or ""
    
    if not is_replied_bot(message, is_channel, user_text, config) and event.message:
        return

    if user_id:
        await asyncio.sleep(uniform(0,0.3))
            
        text, blocked = is_blocked(user_id)
        if not text and blocked:
            return
        if blocked:
            if event.guest_message:
                await g_answer(data, message, text, title="ошибка врат")
            elif event.message:
                error_message = await message.reply(text)
                await asyncio.sleep(7)
                try:
                    await error_message.delete()
                except Exception as e:
                    print(f"{YELLOW} -- del err - {e}{RESET}")
        
            return

        ensure_user(user_id)
        
        if not user_id in last_command_time:
            last_command_time[user_id] = [current_time, True, [message.message_id]]
        elif len(last_command_time[user_id][2]) == 0:
            last_command_time[user_id] = [current_time, True, [message.message_id]]
        else:
            last_command_time[user_id][0] = current_time
            last_command_time[user_id][1] = True
            last_command_time[user_id][2] += [message.message_id]
            
        await asyncio.sleep(uniform(0,0.5))

        if event.message:
            try:
                await bot.send_chat_action(chat_id=message.chat.id, action="typing")
            except TelegramForbiddenError:
                return
        
        msg = None
        replied = message.reply_to_message or None
        if replied:
            replied_text = message.reply_to_message.text or message.reply_to_message.caption or None
        else:
            replied_text = None
            
        prompt = ""
        images = []
        
        if message.document:
            if message.document.file_size > 4096:
                if event.guest_message:
                    await g_answer(data, message, "слишком много данных не хочу отвечать☝️☝️")
                elif event.message:
                    await message.reply("слишком много данных не хочу отвечать☝️☝️")
                return
        
            file_bytes = io.BytesIO()
        
            try:
                await bot.download(message.document, destination=file_bytes)
        
                file_bytes.seek(0)
                file_content = file_bytes.read().decode("utf-8", errors="ignore")
                file_bytes.close()
        
                prompt = f"<файл>{file_content}</файл>"
        
            except UnicodeDecodeError:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл брат❌❌")
                elif event.message:
                    await message.reply("файл не вошёл брат❌❌")
                return
        
            except Exception:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл ало ошибка❌❌")
                elif event.message:
                    await message.reply("файл не вошёл ало ошибка❌❌")
                return
        
        if replied and replied.document:
            if replied.document.file_size > 4096:
                if event.guest_message:
                    await g_answer(data, message, "слишком много данных не хочу отвечать☝️☝️")
                elif event.message:
                    await message.reply("слишком много данных не хочу отвечать☝️☝️")
                return
        
            file_bytes = io.BytesIO()
        
            try:
                await bot.download(replied.document, destination=file_bytes)
                file_bytes.seek(0)
                file_content = file_bytes.read().decode("utf-8", errors="ignore")
                prompt = f"<файл>{file_content}</файл>{user_text}"
            except UnicodeDecodeError:
                if event.guest_message:
                    await g_answer(data, message, "файл не вошёл❌")
                elif event.message:
                    await message.reply("файл не вошёл❌")
                return
            except Exception as e:
                print(f"{RED} -- {e}{RESET}")
                return
            
        if replied and replied.location:
            loc = replied.location
            prompt += (
                f"<ответ на локацию>\n"
                f"{loc.latitude}, {loc.longitude}\n"
                f"</ответ на локацию>\n"
            )
        
        if message.location:
            loc = message.location
            prompt += (
                f"<локация>\n"
                f"{loc.latitude}, {loc.longitude}\n"
                f"</локация>\n"
            )
            
        if replied and replied.poll:
            poll = replied.poll
            if poll.allows_multiple_answers:
                prompt += "<множественный "
            else:
                prompt += "<одиночный "
            prompt += f"опрос в ответе: {poll.question}>\n"
            for option in poll.options:
                prompt += f"- {option.text}\n"
            prompt += "</опрос в ответе>"

        if message.poll:
            poll = message.poll
            if poll.allows_multiple_answers:
                prompt += "<множественный "
            else:
                prompt += "<одиночный "
            prompt += f"опрос: {poll.question}>\n"
            for option in poll.options:
                prompt += f"- {option.text}\n"
            prompt += "</опрос>"  
    
        if replied and replied.photo and IMAGE_MODEL != "":
            try:
                photo = replied.photo[-1]
                file_bytes = await bot.download(photo)
                image_bytes = file_bytes.read()
                b64_image = base64.b64encode(image_bytes).decode("utf-8")
                images += [b64_image]
            except Exception as e:
                print(f"{RED} -- photo - {e}{RESET}")
        
        if message.photo and IMAGE_MODEL != "":
            try:
                if event.message:
                    if not msg:
                        msg = await message.reply("👀")
                photo = message.photo[-1]
                file_bytes = await bot.download(photo)
                image_bytes = file_bytes.read()
                b64_image = base64.b64encode(image_bytes).decode("utf-8")
                images += [b64_image]
            except Exception as e:
                print(f"{RED} -- photo - {e}{RESET}")
        
        prompt = f"{prompt}\n{user_text}"
        if replied_text:
            prompt = f"<ответ на>{replied_text}</ответ> {prompt}"
        prompt = prompt[:MAX_PROMPT]

        if config["markdown"] == "" and message.from_user.id != 777000:
            last_command_time[user_id][1] = False
            last_command_time[user_id][2] = []
            alloh_text = """👆*врат важный вопрос*👆если ты видишь сообщение сверху то нажми "новый стиль",   
а если у тебя тг старый и ты не видишь, то нажми "старый стиль". потому что тг обновил
фукция маркdown.✅✅✅"""
            alloh_one_text = """👇*перед моим ответом ответь на очень важный вопрос*👇
ты хочешь использовать новый стиль выделения текста телеграм 12.8+ (6.9.3+ десктоп) или 
хочешь старый для старых версий телеграм и телеграм х? эту настройку можно будет изменить.✅ """
            reply_markup_url = InlineKeyboardMarkup(inline_keyboard=[
                    [InlineKeyboardButton(text="новый стиль", url=f"https://t.me/{BOT_USERNAME}?start=new_md")],
                    [InlineKeyboardButton(text="старый стиль", url=f"https://t.me/{BOT_USERNAME}?start=old_md")],
                ]
            )
            reply_markup = InlineKeyboardMarkup(inline_keyboard=[
                    [InlineKeyboardButton(text="новый стиль", callback_data=f"set_markdown%{user_id}%new")],
                    [InlineKeyboardButton(text="старый стиль", callback_data=f"set_markdown%{user_id}%old")],
                ]
            )
            if event.guest_message:
                await g_answer(data, message, alloh_one_text, reply_markup=reply_markup_url)
            elif event.message:
                await message.reply_rich(InputRichMessage(markdown="# ало\n```\nкак дела\n```"))
                await message.answer(alloh_text, reply_markup=reply_markup)
            return
        
        if event.message:
            try:
                await bot.send_chat_action(chat_id=message.chat.id, action="typing")
            except TelegramForbiddenError:
                return
    
        if stream == "edit":
            if event.message:
                if msg:
                    message_id = msg.message_id
                else:
                    msg = await message.reply("💬")
                    message_id = msg.message_id
            elif event.guest_message:
                stream = "none"
        else:
            message_id = None
            
        if len(last_command_time[user_id][2]) == 1:
            if is_channel and event.message:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=message.chat.id,
                    user_id=None,
                    message_id=message_id,
                    images=images,
                )
            elif event.message:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=message.chat.id,
                    user_id=user_id,
                    message_id=message_id,
                    images=images,
                )
            else:
                result = await generate(
                    bot=bot,
                    prompt=prompt,
                    chat_id=None,
                    user_id=user_id,
                    message_id=None,
                    images=images,
                )
        
            result = parse_tools(result)

            if event.guest_message:
                add_one_to("stats/guest")
        else:
            print(f"{YELLOW} -- err in blocking 2 {RESET}")
            return
    else:
        return

    last_command_time[user_id][1] = False
    last_command_time[user_id][2] = []

    if ":" in result and ("poll" in result or "quiz" in result):
        pw = ["[poll]", "poll", "quiz"]
        rl = result.split("\n")
        for i, line in enumerate(rl):
            if any(line.startswith(w) for w in pw):
                rl = rl[i:]
                break
        poll_type = rl[0].split(":")[0].replace("[", "").replace("]", "")
        poll_name = rl[0].split(":")[1]
        if len(rl) > 1:
            options = rl[1].split(" =| ")
        else:
            options = ["варианты не указаны"]
        correct_option = options[0]
        correct_option_id = None
        is_anon = True
        allows_multiple_answers = False
        
        if event.guest_message:
            mid = user_id        
        elif message.chat.type == "private":
            mid = user_id
        elif message.chat.type in ("group", "supergroup"):
            mid = message.chat.id
        else:
            return

        if poll_type == "quiz":
            shuffle(options)
            correct_option_id = options.index(correct_option)
        elif poll_type == "poll":
            poll_type = "regular"
            allows_multiple_answers = False
        elif poll_type == "multipoll":
            poll_type = "regular"
            allows_multiple_answers = True
        else:
            poll_type = "regular"

        

        try:
            result = await answer_poll(
                data=data,
                message=message,
                question=poll_name,
                options=options,
                is_anonymous=is_anon,
                poll_type=poll_type,
                allows_multiple_answers=allows_multiple_answers,
                correct_option_id=correct_option_id,
                user_id=mid
            )
        except Exception as e:
            if event.guest_message:
                await g_answer(data, message, "гостевой режим не работает с опросами. либо добавь меня " \
                    "в группу, либо напиши мне в личку чтобы я выдал тебе опрос. ✅"
                )
            else:
                print(f"{RED} -- err poll {e} {RESET}")
            return

        if result and event.guest_message:
            await g_answer(data, message, "отправил опрос в личку, так как гостевой режим не работает с " \
                "опросами✅✅"
            )

    
        return

    kwargs = (
        {"rich_message": InputRichMessage(markdown=result)}
        if config["markdown"] == "new"
        else {"text": result}
    )
    
    try:
        if event.guest_message:
            await g_answer(data, message, result, user_id=user_id)
        elif event.message:
            if msg:
                if stream == "native":
                    await msg.delete()
                    if config["markdown"] == "new":
                        await message.reply_rich(**kwargs)
                    else:
                        await message.reply(**kwargs)
                else:
                    await msg.edit_text(**kwargs)
            else:
                if config["markdown"] == "new":
                    await message.reply_rich(**kwargs)
                else:
                    await message.reply(**kwargs)
    except Exception as e:
        kwargs["parse_mode"] = None
        try:
            if event.guest_message:
                await g_answer(data, message, result)
            elif event.message:
                if msg:
                    if stream == "native":
                        await msg.delete()
                        if config["markdown"] == "new":
                            await message.reply_rich(**kwargs)
                        else:
                            await message.reply(**kwargs)
                    else:
                        await msg.edit_text(**kwargs)
                else:
                    if config["markdown"] == "new":
                        await message.reply_rich(**kwargs)
                    else:
                        await message.reply(**kwargs)
        except Exception as e:
            print(f"{RED} -- send guest - {e}{RESET}")
        else:
            print(f"{YELLOW} -- markdown - {e}{RESET}")
    return


@router.inline_query()
async def inline_handler(inline_query: InlineQuery):
    user_id = inline_query.from_user.id
    ensure_user(user_id)
    
    button = InlineKeyboardButton(text="ало жди 📞", callback_data="pasholnaxxuy")

    if inline_query.query == "clear":
        result = [
            InlineQueryResultArticle(
                id="1",
                title="очистить контекст",
                description="нажми сюда чтоб узбэкгпт прочистил свои мозги✅",
                input_message_content=InputTextMessageContent(
                    message_text="я промыл себе мозги и контекст очищен✅"
                ),
            )
        ]
    else:
        result = [
            InlineQueryResultArticle(
                id="1",
                title="генерация думат",
                description=f'думать "{inline_query.query}" ',
                input_message_content=InputTextMessageContent(
                    message_text="я думаю врат... как дела✅"
                ),
                reply_markup=InlineKeyboardMarkup(inline_keyboard=[[button]]),
            )
        ]

    await inline_query.answer(results=result, cache_time=0)


@router.chosen_inline_result()
async def chosen_inline_result_handler(chosen_result: ChosenInlineResult):
    inline_message_id = chosen_result.inline_message_id
    user_id = chosen_result.from_user.id

    add_one_to("stats/inline")
    
    if chosen_result.query == "clear":
        user_contexts[user_id] = []
    else:
        current_time = time.time()
        alo, blocked = is_blocked(user_id)
        if not alo and blocked:
            return
        if alo and blocked:
            await bot.edit_message_text(
                text=alo,
                inline_message_id=inline_message_id
            )
            return
        
        if not user_id in last_command_time:
            last_command_time[user_id] = [current_time, True, [inline_message_id]]
        elif len(last_command_time[user_id][2]) == 0:
            last_command_time[user_id] = [current_time, True, [inline_message_id]]
        else:
            last_command_time[user_id][0] = current_time
            last_command_time[user_id][1] = True
            last_command_time[user_id][2] += [inline_message_id]
        
        if inline_message_id:
            try:
                result = await generate(
                    prompt=chosen_result.query[:MAX_PROMPT],
                    user_id=chosen_result.from_user.id,
                )
                last_command_time[user_id][1] = False
                last_command_time[user_id][2] = []
            except Exception as e:
                error = parse_response(str(e))
                await bot.edit_message_text(
                    text=error["message"],
                    inline_message_id=inline_message_id,
                )
                print(f"{RED} -- {error['type']} -- {error['message']}")
                return
            if result:
                try:
                    await bot.edit_message_text(
                        text=result, inline_message_id=inline_message_id
                    )
                except Exception as e:
                    print(f"{RED} -- {e}{RESET}")
                    if "MESSAGE_TOO_LONG" not in str(e):
                        await bot.edit_message_text(
                            text=result,
                            inline_message_id=inline_message_id,
                            parse_mode="None",
                        )
                    else:
                        await bot.edit_message_text(
                            text=result[:800],
                            inline_message_id=inline_message_id,
                            parse_mode="None",
                        )


async def main():
    ensure_path("stats/gens")
    ensure_path("stats/inline")
    ensure_path("stats/guest")
    ensure_path("stats/ben")
    ensure_path("stats/reporn")
    await bot.delete_webhook(drop_pending_updates=True)
    await dp.start_polling(
        bot,
        allowed_updates = [
            "message",
            "callback_query",
            "guest_message",
            "inline_query",
            "chosen_inline_result",
        ],
    )


if __name__ == "__main__":
    asyncio.run(main())