path: root/supergenerator.py

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

from pprint import pprint

from openai import AsyncOpenAI

from statistics import *
from colors import *
from config import *

last_command_time = {}
last_error_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": "сервер петух и оборвал соединение. подожди или смени провайдера в настройках узбекгпт.",
        }
    elif "invalid token" 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_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_mode=DEFAULT_STREAM, 
    call_mode=DEFAULT_CALL, 
    prompt_name=DEFAULT_PROMPT,
    notify_mode=DEFAULT_NOTIFY,
):
    with shelve.open("users") as db:
        db[str(user_id)] = {
            "model": model_name,
            "stream": stream_mode,
            "call": call_mode,
            "prompt": prompt_name,
            "notify": notify_mode,
        }


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
):
    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,
            }

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, "prompt": DEFAULT_PROMPT, "notify": DEFAULT_NOTIFY},
        )


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

        config = get_user_config(user_id)
        if "*" in config.get("model", ""):
            model = config.get("model").split("*", 2)
        else:
            print(f"{RED} -- err model - failed to get prov&model in \"{config.get('model','')}\"{RESET}")
            return "произошла ошибка при парсинге вашей выбранной модели. либо хост бота гандон либо тебе надо сменить модель в настройках!!"

        settings = MODELS.get(model[0]).get(model[1])
        
        stream_cf = config["stream"]
        system_prompt_name = config["prompt"]
        if system_prompt_name in SYSTEM_PROMPTS:
            system_prompt = SYSTEM_PROMPTS[system_prompt_name] + ADD_TO_PROMPT
        else:
            text = "❌произошла ошибка, возможные проблемы:\n" \
                f"1. у вас выбран удалённый сист промпт (`{system_prompt_name}`)\n" \
                f"2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
            return text

        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

        content = []
        
        if images:
            if settings.get("vision"):
                for image in images:
                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}}]
            else:
                for image in images:
                    try:
                        image_text = await image2text(image)
                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
                    except Exception as e:
                        print(f"{RED} -- err photo - {e} : {config.get('model', '')}{RESET}")

        content += [{"type": "text", "text": prompt}]
        
        user_contexts[user_id].append({"role": "user", "content": content})
        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)
        settings = MODELS.get(model[0]).get(model[1])
        stream = False
        stream_cf = "none"
        content = []
        if images:
            if settings.get("vision"):
                for image in images:
                    content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image}"}}]
            else:
                for image in images:
                    try:
                        image_text = await image2text(image)
                        prompt = f"<фото>{image_text}</фото>\n{prompt}"
                    except Exception as e:
                        print(f"{RED} -- err photo - {e} : {model}{RESET}")

        content += [{"type": "text", "text": prompt}]
        messages = [{"role": "system", "content": SYSTEM_PROMPTS[DEFAULT_PROMPT]+ADD_TO_PROMPT}] + [
            {"role": "user", "content": content}
        ]

    text = ""

    add_one_to("stats/gens")

    try:
        client = providers[model[0]]
    except Exception:
        text = "❌произошла ошибка, возможные проблемы:\n" \
            "1. у вас выбрана удалённая модель или провайдер\n" \
            "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит"
        return text
        
    text = await openai_generate_text(
        oai=client,
        bot=bot,
        model=model[1],
        settings=settings,
        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": [{"type": "text", "text": 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} обрезано>*"
        if text == "":
            text = "(пустой ответ)"
        return text