first commit

1 files changed, 293 insertions, 0 deletions

diff --git a/supergenerator.py b/supergenerator.py
new file mode 100644
index 0000000..ef1ddeb
--- /dev/null
+++ b/supergenerator.py
@@ -0,0 +1,293 @@
+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