293 lines
8.6 KiB
Python
293 lines
8.6 KiB
Python
|
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
|