uzbekgpt/main.py

import asyncio
import base64
import io
import sys

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
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,
    Message,
)
from callbacks import *
from statistics import *
from colors import *
from commands import *
from dotenv import load_dotenv

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


def is_blocked(user_id):
    current_time = time.time()
    if user_id:
        if user_id in last_command_time:
            time_diff = current_time - last_command_time[user_id]
            if time_diff < 3:
                alo = "🚫 *файл не вошёл.* хватит так быстро слать свои сообщения.\n" \
                    f"подожди {round(3 - time_diff, 3)} секунд, брат😡😡"
                return alo, True
    return "", False


def is_replied_bot(message, is_channel, user_text, config=None):
    if message.chat.type in [ChatType.GROUP, ChatType.SUPERGROUP, ChatType.CHANNEL]:
        if not is_channel:
            is_reply_to_bot = (
                message.reply_to_message
                and message.reply_to_message.from_user
                and message.reply_to_message.from_user.is_bot
                and message.reply_to_message.from_user.username == BOT_USERNAME
            )
            if config:
                if config.get("call") == "da":
                    mentions_bot = any(
                        user_text.lower().startswith(call.lower()) for call in BOT_CALLS
                    )
                else:
                    mentions_bot = False
            else:
                mentions_bot = any(
                    user_text.lower().startswith(call.lower()) for call in BOT_CALLS
                )
            if not (is_reply_to_bot or mentions_bot):
                return False
    return True

@dp.update.outer_middleware()
async def guest_middleware(handler, event, data):
    # print("e: ", event)
    if event.guest_message:
        message = event.guest_message
        user_id = message.sender_chat.id if message.sender_chat else message.from_user.id
        current_time = time.time()
        
        if user_id:
            text, blocked = is_blocked(user_id)
            if blocked:
                await data["bot"](
                    AnswerGuestQuery(
                        guest_query_id=message.guest_query_id,
                        result=InlineQueryResultArticle(
                            id="1",
                            title="err",
                            input_message_content=InputTextMessageContent(
                                message_text=text
                            ),
                        ),
                    )
                )
        
                return

            ensure_user(user_id)
            
            last_command_time[user_id] = current_time
            user_text = message.text or message.caption or ""
            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 = ""
            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")
                    image_text = await image2text(image=b64_image)
                except Exception as e:
                    image_text = "произошла ошибка при анализе фото"
                    print(f"{RED} -- photo - {e}{RESET}")
                prompt = f"{prompt}<фото в ответе>{image_text}</фото в ответе>\n"
            if message.photo and IMAGE_MODEL != "":
                try:
                    photo = message.photo[-1]
                    file_bytes = await bot.download(photo)
                    image_bytes = file_bytes.read()
                    b64_image = base64.b64encode(image_bytes).decode("utf-8")
                    image_text = await image2text(image=b64_image)
                except Exception as e:
                    image_text = "произошла ошибка при анализе фото"
                    print(f"{RED} -- photo - {e}{RESET}")
                prompt = f"{prompt}<фото>{image_text}</фото>\n"
        
            prompt = f"{prompt}\n{user_text}"
            if replied_text:
                prompt = f"<ответ на>{replied_text}</ответ> {prompt}"
            prompt = prompt[:MAX_PROMPT]
                    
            result = await generate(
                bot=bot,
                prompt=prompt,
                chat_id=None,
                user_id=user_id,
                message_id=None,
            )
            add_one_to("stats/guest")
        try:
            await data["bot"](
                AnswerGuestQuery(
                    guest_query_id=message.guest_query_id,
                    result=InlineQueryResultArticle(
                        id="1",
                        title="ok",
                        input_message_content=InputTextMessageContent(
                            message_text=result
                        ),
                    ),
                )
            )
        except Exception as e:
            try:
                await data["bot"](
                    AnswerGuestQuery(
                        guest_query_id=message.guest_query_id,
                        result=InlineQueryResultArticle(
                            id="1",
                            title="ok",
                            input_message_content=InputTextMessageContent(
                                message_text=result,
                                parse_mode="None"
                            ),
                        ),
                    )
                )
            except Exception as e:
                print(f"{RED} -- send guest - {e}{RESET}")
            else:
                print(f"{YELLOW} -- markdown - {e}{RESET}")

        return

    return await handler(event, data)
    
    
@router.message(F.content_type.in_({"text", "photo", "document"}))
async def text_handler(message: Message):
    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
    )
    user_text = message.text or message.caption or ""
    msg = None
    config = get_user_config(user_id)
    stream = config.get("stream")
    current_time = time.time()
    
    if not is_replied_bot(message, is_channel, user_text, config):
        return

    if user_id:
        text, blocked = is_blocked(user_id)
        if blocked:
            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

        last_command_time[user_id] = current_time
        ensure_user(user_id)
        
    replied = message.reply_to_message
    prompt = ""

    try:
        await bot.send_chat_action(chat_id=message.chat.id, action="typing")
    except TelegramForbiddenError:
        return

    if message.document:
        if message.document.file_size > 4096:
            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:
            await message.reply("файл не вошёл брат❌❌")
            return

        except Exception:
            await message.reply("файл не вошёл ало❌")
            return

    if replied and replied.document:
        if replied.document.file_size > 4096:
            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:
            await message.reply("файл не вошёл❌")
            return
        except Exception as e:
            print(f"{RED} -- {e}{RESET}")
            return

    elif replied and replied.photo and IMAGE_MODEL != "":
        try:
            msg = await message.reply("👀")
        except Exception:
            return
        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")
            image_text = await image2text(image=b64_image)
        except Exception as e:
            image_text = "произошла ошибка при анализе фото"
            print(f"{RED} -- photo - {e}{RESET}")
        if replied.caption:
            prompt = f"<ответ на фото>{image_text}</ответ на фото>\n<подпись в ответе>{replied.caption}</подпись в ответе>\n"
        else:
            prompt = f"<ответ на фото>{image_text}</ответ на фото>\n<подпись в ответе>{replied.caption}\n"
        await msg.edit_text("💬")

    elif replied and replied.text:
        replied_text = message.reply_to_message.text
        prompt = f"<ответил на сообщение> {replied_text} </ответил на сообщение>\n"

    if message.photo and IMAGE_MODEL != "":
        try:
            msg = await message.reply("👀")
        except Exception:
            return
        try:
            photo = message.photo[-1]
            file_bytes = await bot.download(photo)
            image_bytes = file_bytes.read()
            b64_image = base64.b64encode(image_bytes).decode("utf-8")
            image_text = await image2text(image=b64_image)
        except Exception as e:
            image_text = "произошла ошибка при анализе фото"
            print(f"{RED} -- photo - {e}{RESET}")
        prompt = f"<фото>{image_text}</фото>\n"
        await msg.edit_text("💬")

    prompt = f"{prompt}\n{user_text}"
    prompt = prompt[:MAX_PROMPT]

    try:
        await bot.send_chat_action(chat_id=message.chat.id, action="typing")
    except TelegramForbiddenError:
        return

    if stream == "edit":
        if not msg:
            try:
                msg = await message.reply("💬")
            except Exception:
                return
        message_id = msg.message_id
    else:
        message_id = None

    if is_channel:
        result = await generate(
            bot=bot,
            prompt=prompt,
            chat_id=message.chat.id,
            user_id=None,
            message_id=message_id,
        )
    else:
        result = await generate(
            bot=bot,
            prompt=prompt,
            chat_id=message.chat.id,
            user_id=user_id,
            message_id=message_id,
        )

    result = parse_tools(result)
    
    try:
        if result:
            if msg:
                if stream == "native":
                    await msg.delete()
                    await message.reply(result)
                else:
                    await msg.edit_text(result)
            else:
                await message.reply(result)
    except Exception as e:
        try:
            if msg:
                if stream == "native":
                    await msg.delete()
                    await message.reply(result, parse_mode="None")
                else:
                    await msg.edit_text(result, parse_mode="None")
            else:
                await message.reply(result, parse_mode="None")
        except Exception as e:
            print(f"{RED} -- {e}{RESET}")
            return
        print(f"{YELLOW} -- {e}{RESET}")


@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 alo and blocked:
            await bot.edit_message_text(
                text=alo,
                inline_message_id=inline_message_id
            )
            return
        last_command_time[user_id] = current_time
        if inline_message_id:
            try:
                result = await generate(
                    prompt=chosen_result.query[:MAX_PROMPT],
                    user_id=chosen_result.from_user.id,
                )
            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:
                if result.startswith("<!block>") or result.endswith("<!block>"):
                    result = result.replace("<!block>", "")
                    last_block_time[user_id] = current_time
                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())