uzbekgpt/helpers.py

import re
import time
import shelve
import requests
from datetime import datetime

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)
        if len(parts) < 2:
            continue
        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": "лимит токенов закончился у провайдера врат. подожди или смени модель.",
        }
    elif "response payload is not completed" in text or "not enough data to satisfy" in text or "server disconnected" in text or "connection reset" in text or "serverdisconnectederror" in text:
        return {
            "type": "incomplete",
            "message": "сервер петух и оборвал ответ на середине. подожди или смени модель.",
        }
    elif "timeout" in text or "timed out" in text:
        return {
            "type": "timeout",
            "message": "сервер думал слишком долго и умер. подожди или смени модель.",
        }
    elif "not found" in text or " 404" in text or "404 " in text or "no such model" in text or "does not exist" in text:
        return {
            "type": "notfound",
            "message": "модель не найдена у провайдера. смени модель в настройках.",
        }
    elif "forbidden" in text or "unauthorized" in text or " 403" in text or "403 " in text or " 401" in text or "401 " in text or "invalid api key" in text or "incorrect api key" in text:
        return {
            "type": "key",
            "message": "ключ или доступ к провайдеру сдох. смени модель или напиши хостеру.",
        }
    elif " 429" in text or "429 " in text or "exhausted" in text or "quota" in text:
        return {
            "type": "quota",
            "message": "лимит запросов у провайдера закончился. подожди или смени модель.",
        }
    elif " 500" in text or "500 " in text or " 502" in text or "502 " in text or " 503" in text or "503 " in text or "bad gateway" in text:
        return {
            "type": "int",
            "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


def get_group_config(group_id):
    with shelve.open("groups") as db:
        config = db.get(str(group_id), {})
        defaults = {
            "post_prompt": SYSTEM_PROMPTS[DEFAULT_PROMPT]["text"]+ADD_TO_PROMPT+OLD_MD,
            "post_reply": "da",
            "polls": "da",
            "lurk": "da",
            "lurk_chance": 10,
            "lurk_reply_chance": 2,
        }
        updated = False
        for key, value in defaults.items():
            if key not in config:
                config[key] = value
                updated = True
        if updated:
            db[str(group_id)] = config
    return config


def set_group_setting(
    group_id,
    post_prompt=None,
    post_reply=None,
    polls=None,
    lurk=None,
    lurk_chance=None,
    lurk_reply_chance=None,
):
    with shelve.open("groups") as db:
        group_id = str(group_id)
        if group_id not in db:
            db[group_id] = {}

        settings = db[group_id]

        if post_prompt is not None: settings["post_prompt"] = post_prompt
        if post_reply is not None: settings["post_reply"] = post_reply
        if polls is not None: settings["polls"] = polls
        if lurk is not None: settings["lurk"] = lurk
        if lurk_chance is not None: settings["lurk_chance"] = lurk_chance
        if lurk_reply_chance is not None: settings["lurk_reply_chance"] = lurk_reply_chance

        db[group_id] = settings


def apply_time_tags(text):
    now = datetime.now()
    text = text.replace("{{date}}", now.strftime("%d.%m.%Y"))
    text = text.replace("{{time}}", now.strftime("%H:%M"))
    text = text.replace("{{now}}", now.strftime("%d.%m.%Y %H:%M"))
    return text


def web_search(query):
    r = requests.get(
        "https://lite.duckduckgo.com/lite/",
        params={"q": query},
        headers={"User-Agent": "Mozilla/5.0"},
        timeout=10
    )
    r.raise_for_status()
    links = re.findall(r"<a[^>]*class=['\"]result-link['\"][^>]*>(.*?)</a>", r.text, re.DOTALL)
    snippets = re.findall(r"<td[^>]*class=['\"]result-snippet['\"][^>]*>(.*?)</td>", r.text, re.DOTALL)
    results = []
    for i in range(min(5, len(links))):
        title = re.sub(r"<[^>]+>", "", links[i]).strip()
        snippet = re.sub(r"<[^>]+>", "", snippets[i]).strip()[:300] if i < len(snippets) else ""
        results.append(f"{title}\n{snippet}")
    if not results:
        return "ничего не нашлось"
    return "\n\n".join(results)


def wiki_search(query):
    r = requests.get(
        "https://ru.wikipedia.org/w/rest.php/v1/search/page",
        params={"q": query, "limit": 3},
        headers={"User-Agent": "uzbekgpt/1.0"},
        timeout=10
    )
    r.raise_for_status()
    results = []
    for p in r.json().get("pages", []):
        title = p.get("title", "")
        excerpt = re.sub(r"<[^>]+>", "", p.get("excerpt", ""))
        results.append(f"{title}\n{excerpt[:400]}")
    if not results:
        return "ничего не нашлось в википедии"
    return "\n\n".join(results)


def extract_tool(text):
    m = re.search(r"<search>(.*?)</search>", text, re.DOTALL)
    if m:
        return "search", m.group(1).strip()
    m = re.search(r"<wiki>(.*?)</wiki>", text, re.DOTALL)
    if m:
        return "wiki", m.group(1).strip()
    return None


def add_tokens(user_id, input_t, output_t):
    with shelve.open("tokens") as db:
        key = str(user_id)
        data = db.get(key, {"input": 0, "output": 0, "requests": 0})
        data["input"] += input_t
        data["output"] += output_t
        data["requests"] += 1
        db[key] = data


def get_tokens(user_id):
    with shelve.open("tokens") as db:
        return db.get(str(user_id), {"input": 0, "output": 0, "requests": 0})


async def alert_admin(bot, provider_name, error_text):
    current = time.time()
    if provider_name in provider_alert_time and current - provider_alert_time[provider_name] < 600:
        return
    provider_alert_time[provider_name] = current
    try:
        await bot.send_message(ADMIN_ID, f"🚨 провайдер `{provider_name}` кажется упал:\n{str(error_text)[:300]}")
    except Exception as e:
        print(f" -- alert admin - {e}")