path: root/supergenerator.py

import random
import re
import shelve
import time
from configparser import DEFAULTSECT

from openai import AsyncOpenAI

from colors import *
from config import *

last_command_time = {}
last_block_time = {}
user_contexts = {}


def get_clients():
    provs = {}
    for provider, settings in PROVIDERS.items():
        if settings.get("module") == "openai":
            provs[provider] = AsyncOpenAI(
                base_url=settings.get("url"), api_key=settings.get("key")
            )
    return provs


providers = get_clients()


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 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": "сервер петух и оборвал соединение. подожди или смени провайдера в настройках узбекгпт.",
        }
    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_user_ids():
    ids = []
    with shelve.open("users") as db:
        for user_id in db.keys():
            try:
                ids.append(int(user_id))
            except ValueError:
                pass
    return ids


def set_user_config(
    user_id, model_name, stream_mode=DEFAULT_STREAM, call_mode=DEFAULT_CALL
):
    with shelve.open("users") as db:
        db[str(user_id)] = {
            "model": model_name,
            "stream": stream_mode,
            "call": call_mode,
        }


def get_user_config(user_id):
    with shelve.open("users") as db:
        return db.get(
            str(user_id),
            {"model": DEFAULT_MODEL, "stream": DEFAULT_STREAM, "call": DEFAULT_CALL},
        )


async def generate(
    prompt: str, bot=None, user_id=None, chat_id=None, message_id=None
) -> str:
    if user_id:
        if user_id not in user_contexts:
            user_contexts[user_id] = []

        config = get_user_config(user_id)
        model = config.get("model").split("*", 2)
        stream_cf = config["stream"]

        if stream_cf == "native" and user_id == chat_id:
            stream = True
        elif stream_cf == "edit" and chat_id and message_id:
            stream = True
        else:
            stream = False

        user_contexts[user_id].append({"role": "user", "content": prompt})
        user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]

        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + user_contexts[
            user_id
        ]
    else:
        model = DEFAULT_MODEL.split("*", 2)
        stream = False
        stream_cf = "none"
        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + [
            {"role": "user", "content": prompt}
        ]

    text = ""

    client = providers[model[0]]
    text = await openai_generate_text(
        oai=client,
        bot=bot,
        model=model[1],
        settings=MODELS.get(model[0]).get(model[1]),
        messages=messages,
        stream=stream,
        stream_cf=stream_cf,
        chat_id=chat_id,
        draft_id=chat_id,
        message_id=message_id,
    )

    if user_id:
        user_contexts[user_id].append({"role": "assistant", "content": text})
        user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:]

    return text


async def image2text(image) -> str:
    if IMAGE_MODEL == "":
        return ""
    model = IMAGE_MODEL.split("*", 2)
    client = providers[model[0]]
    text = ""
    messages = [
        {
            "role": "user",
            "content": [
                {"type": "input_text", "text": IMAGE_PROMPT},
                {"type": "input_image", "image_url": f"data:image/jpeg;base64,{image}"},
            ],
        }
    ]
    response = await client.responses.create(
        model=model[1], input=messages, stream=False
    )
    text = response.output_text
    return text


async def openai_generate_text(
    oai: AsyncOpenAI,
    bot,
    model,
    settings,
    messages,
    stream,
    stream_cf,
    chat_id,
    draft_id,
    message_id,
) -> str:
    try:
        completion = await oai.chat.completions.create(
            model=model,
            messages=messages,
            stream=stream,
            max_tokens=MAX_TOKENS,
            extra_body={
                "think": False
            }
        )
        if stream:
            message = ""
            counter = 0
            if not draft_id:
                draft_id = chat_id
            async for event in completion:
                content = event.choices[0].delta.content
                if content:
                    message += content
                    counter += 1
                    if stream_cf == "native":
                        counter_limit = settings.get("n_tok")
                    else:
                        counter_limit = settings.get("e_tok")

                    if counter % counter_limit == 0:
                        display_text, s = clean_tail_repeats(message)
                        if s > 0:
                            display_text += f"*... <{s} обрезано>*"
                        if display_text.count("```") % 2 != 0:
                            display_text += "\n```"
                        try:
                            if stream_cf == "native":
                                await bot.send_message_draft(
                                    chat_id=chat_id,
                                    draft_id=draft_id,
                                    text=display_text,
                                    parse_mode="Markdown",
                                )
                            elif stream_cf == "edit":
                                await bot.edit_message_text(
                                    chat_id=chat_id,
                                    message_id=message_id,
                                    text=display_text,
                                    parse_mode="Markdown",
                                )
                        except Exception as e:
                            print(f"{YELLOW} -- {e}{RESET}")

            text = message
        else:
            text = completion.choices[0].message.content
    except Exception as e:
        error = parse_response(str(e))
        print(f"{RED} -- {error['type']} - {e}{RESET}")
        return error["message"]
    else:
        text, s = clean_tail_repeats(text)
        if s > 0:
            text += f"*... <{s} обрезано>*"
        return text