фикс я хуй знает много чего я добавил

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}")