uzbekgpt/helpers.py

import re
import os
import time
import shelve
import requests
import html as _html
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={},
    calling={},
):
    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,
            "calling": calling,
        }



def set_user_setting(
    user_id, 
    model=None, 
    stream=None, 
    call=None, 
    prompt=None,
    notify=None,
    markdown=None,
    prompts=None,
    calling=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
        if calling is not None: settings["calling"] = calling

        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={},
    calling=[],
):
    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,
                "calling": calling,
            }


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": {},
            "calling": [],
        }
        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):
    headers = {"User-Agent": "Mozilla/5.0"}
    links = []
    snippets = []

    try:
        r = requests.get("https://lite.duckduckgo.com/lite/", params={"q": query}, headers=headers, 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)
        if not links:
            print(f" -- search lite get пусто: {r.status_code} len={len(r.text)}")
    except Exception as e:
        print(f" -- search lite get - {e}")

    if not links:
        try:
            r = requests.post("https://lite.duckduckgo.com/lite/", data={"q": query}, headers=headers, 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)
            if not links:
                print(f" -- search lite post пусто: {r.status_code} len={len(r.text)}")
        except Exception as e:
            print(f" -- search lite post - {e}")

    if not links:
        try:
            r = requests.post("https://html.duckduckgo.com/html/", data={"q": query}, headers=headers, timeout=10)
            r.raise_for_status()
            links = re.findall(r"<a[^>]*class=['\"]result__a['\"][^>]*>(.*?)</a>", r.text, re.DOTALL)
            snippets = re.findall(r"<a[^>]*class=['\"]result__snippet['\"][^>]*>(.*?)</a>", r.text, re.DOTALL)
            if not links:
                print(f" -- search html post пусто: {r.status_code} len={len(r.text)}")
        except Exception as e:
            print(f" -- search html post - {e}")

    results = []
    for i in range(min(5, len(links))):
        title = _html.unescape(re.sub(r"<[^>]+>", "", links[i])).strip()
        snippet = _html.unescape(re.sub(r"<[^>]+>", "", snippets[i])).strip()[:300] if i < len(snippets) else ""
        results.append(f"{title}\n{snippet}")
    if not results:
        try:
            return wiki_search(query)
        except Exception:
            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 = _html.unescape(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}")