diff options
Diffstat (limited to 'supergenerator.py')
| -rw-r--r-- | supergenerator.py | 293 |
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 |