diff --git a/commands.py b/commands.py index 67c6ff4..98e80dd 100644 --- a/commands.py +++ b/commands.py @@ -52,6 +52,7 @@ commands = [ "/reset", "/tben", "/group", + "/cost", ] @router.message(Command("stats")) @@ -560,6 +561,21 @@ async def clear_handler(message: Message): print(f"{RED} -- {e}{RESET}") +@router.message(Command("cost")) +async def cost_handler(message: Message): + if message.sender_chat: + user_id = message.sender_chat.id + else: + user_id = message.from_user.id + t = get_tokens(user_id) + await message.reply( + f"💸 *твои токены:*\n" + f"вход: `{t['input']}`\n" + f"выход: `{t['output']}`\n" + f"запросов: `{t['requests']}`" + ) + + @router.message(Command("markpidor")) async def alo(message: Message): await message.reply_rich(InputRichMessage(markdown="# ало")) diff --git a/helpers.py b/helpers.py index 2f7dcfa..8ec5882 100644 --- a/helpers.py +++ b/helpers.py @@ -2,6 +2,7 @@ import re import time import shelve import requests +from datetime import datetime from aiogram.methods import AnswerGuestQuery, SendPoll from aiogram.types import InlineQueryResultArticle, InputTextMessageContent, InputRichMessageContent, InputRichMessage, Message @@ -571,3 +572,85 @@ def set_group_setting( db[group_id] = settings + +def apply_time_tags(text): + now = datetime.now() + text = text.replace("{{date}}", now.strftime("%d.%m.%Y")) + text = text.replace("{{time}}", now.strftime("%H:%M")) + text = text.replace("{{now}}", now.strftime("%d.%m.%Y %H:%M")) + return text + + +def web_search(query): + r = requests.get( + "https://lite.duckduckgo.com/lite/", + params={"q": query}, + headers={"User-Agent": "Mozilla/5.0"}, + timeout=10 + ) + r.raise_for_status() + links = re.findall(r"]*class=['\"]result-link['\"][^>]*>(.*?)", r.text, re.DOTALL) + snippets = re.findall(r"]*class=['\"]result-snippet['\"][^>]*>(.*?)", r.text, re.DOTALL) + results = [] + for i in range(min(5, len(links))): + title = re.sub(r"<[^>]+>", "", links[i]).strip() + snippet = re.sub(r"<[^>]+>", "", snippets[i]).strip()[:300] if i < len(snippets) else "" + results.append(f"{title}\n{snippet}") + if not results: + return "ничего не нашлось" + return "\n\n".join(results) + + +def wiki_search(query): + r = requests.get( + "https://ru.wikipedia.org/w/rest.php/v1/search/page", + params={"q": query, "limit": 3}, + headers={"User-Agent": "uzbekgpt/1.0"}, + timeout=10 + ) + r.raise_for_status() + results = [] + for p in r.json().get("pages", []): + title = p.get("title", "") + excerpt = re.sub(r"<[^>]+>", "", p.get("excerpt", "")) + results.append(f"{title}\n{excerpt[:400]}") + if not results: + return "ничего не нашлось в википедии" + return "\n\n".join(results) + + +def extract_tool(text): + m = re.search(r"(.*?)", text, re.DOTALL) + if m: + return "search", m.group(1).strip() + m = re.search(r"(.*?)", text, re.DOTALL) + if m: + return "wiki", m.group(1).strip() + return None + + +def add_tokens(user_id, input_t, output_t): + with shelve.open("tokens") as db: + key = str(user_id) + data = db.get(key, {"input": 0, "output": 0, "requests": 0}) + data["input"] += input_t + data["output"] += output_t + data["requests"] += 1 + db[key] = data + + +def get_tokens(user_id): + with shelve.open("tokens") as db: + return db.get(str(user_id), {"input": 0, "output": 0, "requests": 0}) + + +async def alert_admin(bot, provider_name, error_text): + current = time.time() + if provider_name in provider_alert_time and current - provider_alert_time[provider_name] < 600: + return + provider_alert_time[provider_name] = current + try: + await bot.send_message(ADMIN_ID, f"🚨 провайдер `{provider_name}` кажется упал:\n{str(error_text)[:300]}") + except Exception as e: + print(f" -- alert admin - {e}") + diff --git a/prompts.py b/prompts.py index c95ea1f..4bc9a46 100644 --- a/prompts.py +++ b/prompts.py @@ -72,3 +72,14 @@ OLD_MD = """ NEW_MD = """ иногда используй Markdown для разметки (можно почти весь, даже заголовки можно) """ + +TOOLS_PROMPT = """ +инструменты: +если не знаешь ответ или нужны свежие данные из интернета используй поиск: +ответь сообщением только с тегом и запросом внутри, без другого текста: +запрос +или для статьи из википедии: +запрос +после получения результатов ответишь нормально пользователю. +не используй поиск если знаешь ответ или вопрос простой. +""" diff --git a/supergenerator.py b/supergenerator.py index 16121b7..1e1a229 100644 --- a/supergenerator.py +++ b/supergenerator.py @@ -1,4 +1,3 @@ -import json import ai from statistics import * @@ -6,6 +5,7 @@ from colors import * from helpers import * from config import * from vars import * +from prompts import TOOLS_PROMPT from aiogram.exceptions import TelegramBadRequest @@ -13,6 +13,8 @@ from aiogram.exceptions import TelegramBadRequest async def generate( prompt: str, bot=None, user_id=None, chat_id=None, message_id=None, images=None, msg=None, process=None, message=None, inline_id=None, system_prompt=None ) -> str: + prompt = apply_time_tags(prompt) + if user_id: if user_id not in user_contexts: user_contexts[user_id] = [] @@ -48,6 +50,8 @@ async def generate( f"2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит" return text + system_prompt += TOOLS_PROMPT + content = [] if images and not is_auto: @@ -98,11 +102,15 @@ async def generate( sp = system_prompt else: sp = SYSTEM_PROMPTS[DEFAULT_PROMPT]["text"]+ADD_TO_PROMPT+OLD_MD + sp += TOOLS_PROMPT messages = [{"role": "system", "content": sp}] + [ {"role": "user", "content": content} ] text = "" + total_in = 0 + total_out = 0 + tools_used = [] if str(user_id) in temp_benned: return temp_benned[str(user_id)] @@ -149,6 +157,7 @@ async def generate( auto_content += [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image['data']}"}}] auto_messages = [{"role": "system", "content": messages[0]["content"]}, {"role": "user", "content": auto_content}] + messages[2:] + sent_messages = auto_messages text = await openai_generate_text( api=provider, @@ -164,6 +173,10 @@ async def generate( message_id=message_id, inline_id=inline_id ) + total_in += text.get("input") or 0 + total_out += text.get("output") or 0 + if text["type"] == "error" and text.get("etype") in ["connection", "timeout", "int", "incomplete"]: + await alert_admin(bot, provider_name, text["text"]) last_type = text["type"] if text == "": text = {"type": "error", "text": "все модели в автовыборе вернули ошибку"} @@ -177,6 +190,7 @@ async def generate( "2. конфиг бота настроен неправильно и надо связаться с тем кто его хостит" settings = MODELS.get(provider_name, {}).get(model_name, {}) key = get_key(provider_name, model_str) + sent_messages = messages if stream_cf == "native" and user_id == chat_id and settings.get("nostream", False) == False: stream = True @@ -205,8 +219,61 @@ async def generate( message_id=message_id, inline_id=inline_id ) + total_in += text.get("input") or 0 + total_out += text.get("output") or 0 + if text["type"] == "error" and text.get("etype") in ["connection", "timeout", "int", "incomplete"]: + await alert_admin(bot, provider_name, text["text"]) + + for _ in range(2): + if text.get("type") != "da": + break + tool = extract_tool(text.get("text", "")) + if not tool: + break + kind, query = tool + try: + if kind == "search": + tool_result = web_search(query) + else: + tool_result = wiki_search(query) + tools_used.append(kind) + except Exception as e: + print(f"{RED} -- tool {kind} - {e}{RESET}") + tool_result = f"поиск не сработал ({e})" + tool_exchange = [ + {"role": "assistant", "content": [{"type": "text", "text": text["text"]}]}, + {"role": "user", "content": [{"type": "text", "text": f"<результаты поиска>\n{tool_result[:3000]}\n\nответь на вопрос используя эти данные. тег поиска больше не пиши."}]}, + ] + sent_messages += tool_exchange + if user_id: + user_contexts[user_id] = (user_contexts[user_id] + tool_exchange)[-MAX_CONTEXT:] + text = await openai_generate_text( + api=provider, + api_key=key, + bot=bot, + model=model_name, + settings=settings, + messages=sent_messages, + stream=False, + stream_cf=stream_cf, + chat_id=chat_id, + draft_id=chat_id, + message_id=message_id, + inline_id=inline_id + ) + total_in += text.get("input") or 0 + total_out += text.get("output") or 0 + + if tools_used and text.get("type") == "da": + tool_names = [] + if "search" in tools_used: + tool_names.append("поиск (duckduckgo)") + if "wiki" in tools_used: + tool_names.append("википедия") + text["text"] = "🔍 использован: " + ", ".join(tool_names) + "\n\n" + text["text"] if user_id: + add_tokens(user_id, total_in, total_out) user_contexts[user_id].append({"role": "assistant", "content": [{"type": "text", "text": text["text"]}]}) user_contexts[user_id] = user_contexts[user_id][-MAX_CONTEXT:] @@ -259,6 +326,8 @@ async def openai_generate_text( inline_id=None, api_key=None ) -> str: + usage_in = 0 + usage_out = 0 try: extra_body = { "think": "false" @@ -288,13 +357,19 @@ async def openai_generate_text( try: content = chunk["choices"][0]["delta"]["content"] except: + if isinstance(chunk, dict) and chunk.get("usage"): + u = chunk["usage"] + usage_in = u.get("prompt_tokens", 0) or 0 + usage_out = u.get("completion_tokens", 0) or 0 + continue err_text = extract_error(chunk) if err_text: error = parse_response(err_text.lower()) print(f"{RED} -- {error['type']} - {err_text}{RESET}") return { "type": "error", - "text": error["message"] + "text": error["message"], + "etype": error["type"] } if isinstance(chunk, dict) and "choices" not in chunk: stream_error = str(chunk)[:300] @@ -371,7 +446,8 @@ async def openai_generate_text( print(f"{RED} -- {error['type']} - {stream_error}{RESET}") return { "type": "error", - "text": error["message"] + "text": error["message"], + "etype": error["type"] } else: async for completion in ai.completions(**alo): @@ -385,35 +461,27 @@ async def openai_generate_text( print(f"{RED} -- {error['type']} - {err_text}{RESET}") return { "type": "error", - "text": error["message"] + "text": error["message"], + "etype": error["type"] } error = parse_response(str(completion).lower()) print(f"{RED} -- {error['type']} - {completion}{RESET}") return { "type": "error", - "text": error["message"] + "text": error["message"], + "etype": error["type"] } + if isinstance(completion, dict) and completion.get("usage"): + u = completion["usage"] + usage_in = u.get("prompt_tokens", 0) or 0 + usage_out = u.get("completion_tokens", 0) or 0 except Exception as e: error = parse_response(str(e).lower()) print(f"{RED} -- {error['type']} - {e}{RESET}") - err_str = str(e) - try: - status_code = int(err_str.split(" ", 1)[0]) - except ValueError: - status_code = 0 - if 400 <= status_code < 500: - debug_alo = {k: v for k, v in alo.items() if k != "api_key"} - return { - "type": "error", - "text": error["message"] + "\n\n" \ - f"модель: `{model}`\n" \ - f"статус: `{status_code}`\n" \ - f"raw: `{err_str[:500]}`\n" \ - f"запрос:\n```{json.dumps(debug_alo, ensure_ascii=False)[:1200]}\n```" - } return { "type": "error", - "text": error["message"] + "text": error["message"], + "etype": error["type"] } else: text, s = clean_tail_repeats(text) @@ -422,9 +490,12 @@ async def openai_generate_text( if text == "": return { "type": "error", - "text": "модель вернула пустой ответ" + "text": "модель вернула пустой ответ", + "etype": "empty" } return { "type": "da", - "text": text + "text": text, + "input": usage_in, + "output": usage_out } diff --git a/tokens b/tokens new file mode 100644 index 0000000..e60a583 Binary files /dev/null and b/tokens differ diff --git a/vars.py b/vars.py index 120d5a0..be0b41c 100644 --- a/vars.py +++ b/vars.py @@ -12,3 +12,5 @@ user_contexts = {} temp_benned = {} # последние сообщения групп для подглядывания group_last_messages = {} +# чтобы не спамить админу алертами при падении провайдера +provider_alert_time = {}