uzbekgpt/helpers.py

import re
import time
import shelve
import requests

from aiogram.methods import AnswerGuestQuery, SendPoll
from aiogram.types import InlineQueryResultArticle, InputTextMessageContent, InputRichMessageContent, InputRichMessage, Message
from aiogram.enums import ChatType

from vars import *
from config import *


def is_blocked(user_id):
    current_time = time.time()
    if user_id and user_id in last_command_time:
        time_diff = current_time - last_command_time[user_id][0]
        if user_id in last_error_time:
            time_error_diff = current_time - last_error_time[user_id]
        else:
            last_error_time[user_id] = current_time
            time_error_diff = 1984
        
        if time_diff < 3:
            alo = "🚫 *файл не вошёл.* хватит так быстро слать свои сообщения.\n" \
                f"подожди {round(3 - time_diff, 3)} секунд, брат😡😡"
        elif last_command_time[user_id][1] and time_diff < 5:
            alo = "🚫 *файл не вошёл.* я ещё не сгенерировал прошлое сообщение.\n" \
                f"подожди {round(10 - time_diff, 3)} секунд, чтобы я закрыл глаза " \
                "на это da."
        elif len(last_command_time[user_id][2]) > 1 and time_diff > 5:
            last_command_time[user_id] = [current_time, False, []]
            return "", False
        elif len(last_command_time[user_id][2]) > 1:
            alo = "🚫 *файл не вошёл.* хватит так быстро слать свои сообщения."
        else:
            return "", False

        
        
        if time_error_diff < 2:
            return "", True
        else:
            last_error_time[user_id] = current_time
            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().split(" ", 1)[0] == call for call in BOT_CALLS
                    )
                    for call in BOT_CALLS:
                        if call.startswith("@"):
                            if call.lower() in user_text.lower():
                                mentions_bot = True
                else:
                    mentions_bot = False
            else:
                mentions_bot = any(
                    user_text.lower().startswith(call.lower()) for call in BOT_CALLS
                )
                for call in BOT_CALLS:
                    if call.startswith("@"):
                        if call.lower() in user_text.lower():
                            mentions_bot = True
            if not (is_reply_to_bot or mentions_bot):
                return False
    return True

async def g_answer(
    message, 
    text, 
    title="da", 
    id="1", 
    parse_mode="Markdown",
    reply_markup=None,
    user_id=None
):
    input_message_content = InputTextMessageContent(
        message_text=text,
        parse_mode=parse_mode
    )
    if user_id:
        config = get_user_config(user_id)
        if config["markdown"] == "new":
            input_message_content=InputRichMessageContent(
                rich_message=InputRichMessage(markdown=text)
            )
    return await message.answer_guest_query(
        result=InlineQueryResultArticle(
            id=id,
            reply_markup=reply_markup,
            title=title,
            input_message_content=input_message_content
        )
    )


async def answer_poll(
    data,
    message,
    question,
    options,
    is_anonymous=True,
    poll_type="regular",
    allows_multiple_answers=False,
    correct_option_id=None,
    user_id=None
):
    return await data["bot"](
        SendPoll(
            chat_id=user_id,
            question=question,
            options=options,
            is_anonymous=is_anonymous,
            type=poll_type,
            allows_multiple_answers=allows_multiple_answers,
            correct_option_id=correct_option_id,
        )
    )


def send_to_vosk(audio):
    files = {
        "audio": (
            "voice.ogg",
            audio,
            "audio/ogg"
        )
    }

    r = requests.post(
        VOSK_API,
        files=files,
        timeout=15
    )

    return r.json()


async def mans(message, text):
    if message.guest_query_id:
        return await message.answer_guest_query(
            result=InlineQueryResultArticle(
                id=1,
                reply_markup=None,
                title="da",
                input_message_content=InputTextMessageContent(
                    message_text=text,
                    parse_mode="Markdown"
                )
            )
        )
    else:
        return await message.reply(text)
    
def graph_emoji(status):
    match status:
        case "done":
            return "✅"
        case "error":
            return "🚫"
        case "make":
            return "💭"
        case "wait":
            return "⏳"
        case _:
            return "🥺"


def make_graph(message, process):
    text = ""
    replied = message.reply_to_message or None
    user_id = message.from_user.id
    config = get_user_config(user_id)
    im_model = IMAGE_MODEL.split("*", 1)
    model = config["model"].split("*", 1)
    photo_a = process.get("photo_a", None)
    photo = process.get("photo", None)
    stt = process.get("stt", None)
    t = process.get("text", None)
    
    if photo_a:
        text += f"{graph_emoji(photo_a)} анализ фото (в ответе)\n"
    if photo:
        text += f"{graph_emoji(photo)} анализ фото\n"
    if stt:
        text += f"{graph_emoji(stt)} stt (голос -> текст)\n"
        
    text += f"{graph_emoji(t)} запрос `{model[1]}` (`{model[0]}`)\n"
    
    if replied and replied.document:
        text += f"↳ {replied.document.file_size} файл (в ответе)\n"
    if message.document:
        text += f"↳ {message.document.file_size} файл\n"
        
    if replied and replied.photo:
        text += "↳ фото (в ответе)\n"
    if message.photo:
        text += "↳ фото\n"
        
    if replied and replied.location:
        text += "↳ геолокация (в ответе)\n"
    if message.location:
        text += "↳ геолокация\n"
        
    if replied and replied.poll:
        text += "↳ опрос (в ответе)\n"
    if message.poll:
        text += "↳ опрос\n"  
      
    if replied and replied.voice:
        text += "↳ голосовое сообщение (в ответе)\n"
    if message.voice:
        text += "↳ голосовое сообщение\n"
      
    if replied and replied.text:
        text += f"↳ {len(replied.text)} текст в ответе\n"

    return text


def clean_tail_repeats(text: str, max_repeat: int = 10):
    if not text:
        return text, 0

    original_length = len(text)
    cleaned = text

    pattern = rf"(.)\1{{{max_repeat},}}$"

    def limit_repeats(match):
        char = match.group(1)
        return char * max_repeat

    cleaned = re.sub(pattern, limit_repeats, cleaned)

    pattern_space = rf"((\S)\s+)\2{{{max_repeat},}}$"
    cleaned = re.sub(
        pattern_space,
        lambda m: (m.group(2) + " ") * min(max_repeat, len(m.group(1))),
        cleaned,
    )

    words = cleaned.split()
    if len(words) >= 2:
        for i in range(1, min(len(words), 10)):
            if all(words[-j] == words[-j - 1] for j in range(i)):
                unique_words = words[: -(i + 1)]
                cleaned = " ".join(unique_words + [words[-1]])
                break

    removed_count = original_length - len(cleaned)

    cleaned = re.sub(r" +", " ", cleaned)
    cleaned = cleaned.strip()

    return cleaned, removed_count


def get_key(provider, model_full):
    prov = PROVIDERS.get(provider, {})
    key = prov.get("key", "")
    if isinstance(key, list):
        parts = model_full.split("*")
        if len(parts) >= 3:
            try:
                idx = int(parts[2]) - 1
            except ValueError:
                return key[0]
            if 0 <= idx < len(key):
                return key[idx]
        return key[0]
    return key


_expanded_auto_models = None


def expand_auto_models():
    global _expanded_auto_models
    if _expanded_auto_models is not None:
        return _expanded_auto_models
    result = []
    for entry in AUTO_MODELS:
        parts = entry.split("*", 2)
        provider = parts[0]
        model = parts[1]
        suffix = parts[2] if len(parts) >= 3 else None
        prov_key = PROVIDERS.get(provider, {}).get("key", "")
        if isinstance(prov_key, list) and not suffix:
            for i in range(len(prov_key)):
                result.append(f"{provider}*{model}*{i + 1}")
        else:
            result.append(entry)
    _expanded_auto_models = result
    return result


def parse_response(text: str):
    if "connection error." in text:
        return {
            "type": "connection",
            "message": "ошибка подключения к провайдеру. подожди или смени провайдера в настройках узбекгпт.",
        }
    elif "model quota exceeded" in text:
        return {
            "type": "quota",
            "message": "лимит токенов закончился. подожди или смени модель.",
        }
    elif "tier capacity exceeded." in text:
        return {"type": "tier", "message": "попробуй ещё раз!"}
    elif "internal server" in text:
        return {
            "type": "int",
            "message": "произошла ошибка на стороне провайдера. подожди или смени провайдера в настройках узбекгпт.",
        }
    elif "(incomplete chunked read)" in text:
        return {
            "type": "incomplete",
            "message": "сервер петух и оборвал соединение. подожди или смени провайдера в настройках узбекгпт.",
        }
    elif "provider error" in text:
        return {
            "type": "provider",
            "message": "врат ошибка на стороне провайдера модели. смени провайдера или модель.",
        }
    elif "rate limit" in text:
        return {
            "type": "quota",
            "message": "лимит токенов закончился у провайдера врат. подожди или смени модель.",
        }
    else:
        return {
            "type": "unknown",
            "message": f"произошла какая-то ошибка при генерации... {text}",
        }


def galockinator(text):
    for _ in range(random.randint(1, 3)):
        if random.random() < 0.3:
            text += "☝️"
        else:
            text += "✅"

    return text


def remove_think_tags(text):
    result = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    result = re.sub(r"<thought>.*?</thought>", "", text, flags=re.DOTALL)
    return result


def get_all_users(only_notify=False):
    with shelve.open("users") as db:
        return {k: db[k] for k in db.keys() if not only_notify or db[k].get("notify") == "da"}


def get_all_user_ids(only_notify=False):
    ids = []
    with shelve.open("users") as db:
        for user_id in db.keys():
            if only_notify and db[user_id].get("notify") != "da":
                continue
            try:
                ids.append(int(user_id))
            except ValueError:
                pass
    return ids


def set_user_config(
    user_id, 
    model_name, 
    stream=DEFAULT_STREAM, 
    call=DEFAULT_CALL, 
    prompt=DEFAULT_PROMPT,
    notify=DEFAULT_NOTIFY,
    markdown=DEFAULT_MARKDOWN,
    prompts={},
):
    with shelve.open("users") as db:
        db[str(user_id)] = {
            "model": model_name,
            "stream": stream,
            "call": call,
            "prompt": prompt,
            "notify": notify,
            "markdown": markdown,
            "prompts": prompts,
        }



def set_user_setting(
    user_id, 
    model=None, 
    stream=None, 
    call=None, 
    prompt=None,
    notify=None,
    markdown=None,
    prompts=None,
):
    with shelve.open("users") as db:
        user_id = str(user_id)
        if user_id not in db:
            db[user_id] = {}
        
        settings = db[user_id]
        
        if model is not None: settings["model"] = model
        if stream is not None: settings["stream"] = stream
        if call is not None: settings["call"] = call
        if prompt is not None: settings["prompt"] = prompt
        if notify is not None: settings["notify"] = notify
        if markdown is not None: settings["markdown"] = markdown
        if prompts is not None: settings["prompts"] = prompts

        db[user_id] = settings


def delete_user_from_db(user_id):
    with shelve.open("users") as db:
        if str(user_id) in db:
            del db[str(user_id)]


def reset_user_setting(user_id, key, value):
    with shelve.open("users") as db:
        user = db[str(user_id)]
        user[key] = value
        db[str(user_id)] = user


def ensure_user(
    user_id,
    model_name=DEFAULT_MODEL,
    stream_mode=DEFAULT_STREAM,
    call_mode=DEFAULT_CALL,
    prompt_name=DEFAULT_PROMPT,
    notify_mode=DEFAULT_NOTIFY,
    markdown=DEFAULT_MARKDOWN,
    prompts={},
):
    with shelve.open("users") as db:
        key = str(user_id)

        if key not in db:
            db[key] = {
                "model": model_name,
                "stream": stream_mode,
                "call": call_mode,
                "prompt": prompt_name,
                "notify": notify_mode,
                "markdown": markdown,
                "prompts": prompts,
            }


def get_user_config(user_id):
    with shelve.open("users") as db:
        config = db.get(str(user_id), {})
        defaults = {
            "model": DEFAULT_MODEL,
            "stream": DEFAULT_STREAM,
            "call": DEFAULT_CALL,
            "prompt": DEFAULT_PROMPT,
            "notify": DEFAULT_NOTIFY,
            "markdown": DEFAULT_MARKDOWN,
            "prompts": {},
        }
        updated = False
        for key, value in defaults.items():
            if key not in config:
                config[key] = value
                updated = True
        if updated:
            db[str(user_id)] = config
    return config