diff options
Diffstat (limited to 'supergenerator.py')
| -rw-r--r-- | supergenerator.py | 383 |
1 files changed, 118 insertions, 265 deletions
diff --git a/supergenerator.py b/supergenerator.py index 1377247..4df694b 100644 --- a/supergenerator.py +++ b/supergenerator.py @@ -2,233 +2,24 @@ import random import re import shelve import time +import ai +import json + from configparser import DEFAULTSECT from pprint import pprint -from openai import AsyncOpenAI - from statistics import * from colors import * +from helpers import * from config import * +from vars import * -last_command_time = {} -last_error_time = {} -last_block_time = {} -alo_command = {} -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 "provider error" in text: - return { - "type": "provider", - "message": "врат ошибка на стороне провайдера модели. смени провайдера или модель.", - } - elif "rate limit" in text: - return { - "type": "quota", - "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, - markdown=DEFAULT_MARKDOWN, - prompts={}, -): - 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, - "markdown": markdown, - "prompts": prompts, - } - - -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={}, -): - 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, - } - - -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": {}, - } - 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 - +from aiogram.exceptions import TelegramBadRequest, TelegramForbiddenError async def generate( - prompt: str, bot=None, user_id=None, chat_id=None, message_id=None, images=None + prompt: str, bot=None, user_id=None, chat_id=None, message_id=None, images=None, msg=None, process=None, message=None, inline_id=None ) -> str: if user_id: if user_id not in user_contexts: @@ -265,9 +56,11 @@ async def generate( f"2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит" return text - if stream_cf == "native" and user_id == chat_id and settings.get("nostream") == False: + if stream_cf == "native" and user_id == chat_id and settings.get("nostream", False) == False: + stream = True + elif stream_cf == "edit" and chat_id and message_id and settings.get("nostream", False) == False: stream = True - elif stream_cf == "edit" and chat_id and message_id and settings.get("nostream") == False: + elif stream_cf == "edit" or stream_cf == "native" and inline_id: stream = True else: stream = False @@ -281,10 +74,20 @@ async def generate( else: for image in images: try: - image_text = await image2text(image) + if process and msg and message: + process[image["process"]] = "make" + await msg.edit_text(make_graph(message, process)) + image_text = await image2text(image["data"]) prompt = f"<фото>{image_text}</фото>\n{prompt}" except Exception as e: + if process and msg and message: + process[image["process"]] = "error" + await msg.edit_text(make_graph(message, process)) print(f"{RED} -- err photo - {e} : {config.get('model', '')}{RESET}") + else: + if process and msg and message: + process[image["process"]] = "done" + await msg.edit_text(make_graph(message, process)) content += [{"type": "text", "text": prompt}] @@ -322,15 +125,19 @@ async def generate( add_one_to("stats/gens") try: - client = providers[model[0]] + provider = PROVIDERS[model[0]] except Exception: text = "❌произошла ошибка, возможные проблемы:\n" \ "1. у вас выбрана удалённая модель или провайдер\n" \ "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит" return text - + + if process and msg and message: + process["text"] = "make" + await msg.edit_text(make_graph(message, process)) + text = await openai_generate_text( - oai=client, + api=provider, bot=bot, model=model[1], settings=settings, @@ -340,6 +147,7 @@ async def generate( chat_id=chat_id, draft_id=chat_id, message_id=message_id, + inline_id=inline_id ) if user_id: @@ -352,8 +160,8 @@ async def generate( async def image2text(image) -> str: if IMAGE_MODEL == "": return "" - model = IMAGE_MODEL.split("*", 2) - client = providers[model[0]] + model = IMAGE_MODEL.split("*", 1) + api = PROVIDERS[model[0]] text = "" messages = [ { @@ -364,24 +172,26 @@ async def image2text(image) -> str: ], } ] - response = await client.responses.create( - model=model[1], input=messages, stream=False - ) + async for response in ai.responses( + url=api["url"], api_key=api["key"], model=model[1], input=messages, stream=False + ): + pass text = response.output_text return text async def openai_generate_text( - oai: AsyncOpenAI, + api, bot, model, settings, messages, stream, stream_cf, - chat_id, - draft_id, - message_id, + chat_id=None, + draft_id=None, + message_id=None, + inline_id=None ) -> str: try: extra_body = { @@ -389,56 +199,99 @@ async def openai_generate_text( } if settings.get("noexb"): extra_body = None - - completion = await oai.chat.completions.create( - model=model, - messages=messages, - stream=stream, - max_tokens=MAX_TOKENS, - extra_body=extra_body - ) + + alo = { + "url": api["url"], + "api_key": api["key"], + "model": model, + "messages": messages, + "stream": stream, + "max_tokens": MAX_TOKENS, + "extra": extra_body + } + if stream: message = "" + t_counter = 0 counter = 0 + counter_2 = 1 if not draft_id: draft_id = chat_id - async for event in completion: - content = event.choices[0].delta.content + async for chunk in ai.completions(**alo): + # chunk = json.loads(chunk) + try: + content = chunk["choices"][0]["delta"]["content"] + except: + continue if content: message += content - counter += 1 - if stream_cf == "native": - counter_limit = settings.get("n_tok") - else: - counter_limit = settings.get("e_tok") + counter += len(content) + t_counter += 1 - if counter % counter_limit == 0: + if counter > MAX_SYMBOLS: + break + + if int(counter / STREAMING_SPLIT_SYMBOLS) == counter_2: + counter_2 += 1 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}") + + if counter < 3950: + display_text += "▋" + + try: + if stream_cf == "native" and chat_id and draft_id: + await bot.send_message_draft( + chat_id=chat_id, + draft_id=draft_id, + text=display_text, + ) + elif stream_cf == "edit" and chat_id and message_id: + await bot.edit_message_text( + chat_id=chat_id, + message_id=message_id, + text=display_text, + ) + elif stream_cf == "edit" or stream_cf == "native" and inline_id: + await bot.edit_message_text( + inline_message_id=inline_id, + text=display_text, + ) + except TelegramBadRequest as e: + if "can't parse entities" in str(e): + if stream_cf == "native" and chat_id and draft_id: + await bot.send_message_draft( + chat_id=chat_id, + draft_id=draft_id, + text=display_text, + parse_mode=None, + ) + elif stream_cf == "edit" and chat_id and message_id: + await bot.edit_message_text( + chat_id=chat_id, + message_id=message_id, + text=display_text, + parse_mode=None, + ) + elif stream_cf == "edit" or stream_cf == "native" and inline_id: + await bot.edit_message_text( + inline_message_id=inline_id, + text=display_text, + parse_mode=None, + ) + else: + print(f"{YELLOW} -- {e}{RESET}") + except Exception as e: + print(f"{YELLOW} -- {e}{RESET}") text = message else: - text = completion.choices[0].message.content + async for completion in ai.completions(**alo): + pass + text = completion["choices"][0]["message"]["content"] except Exception as e: error = parse_response(str(e).lower()) print(f"{RED} -- {error['type']} - {e}{RESET}") |