телифон связь поиск в интернате

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="# ало"))

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"<a[^>]*class=['\"]result-link['\"][^>]*>(.*?)</a>", r.text, re.DOTALL)
+    snippets = re.findall(r"<td[^>]*class=['\"]result-snippet['\"][^>]*>(.*?)</td>", 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"<search>(.*?)</search>", text, re.DOTALL)
+    if m:
+        return "search", m.group(1).strip()
+    m = re.search(r"<wiki>(.*?)</wiki>", 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}")
+

prompts.py

@@ -72,3 +72,14 @@ OLD_MD = """
 NEW_MD = """
 иногда используй Markdown для разметки (можно почти весь, даже заголовки можно)
 """
+
+TOOLS_PROMPT = """
+инструменты:
+если не знаешь ответ или нужны свежие данные из интернета используй поиск:
+ответь сообщением только с тегом и запросом внутри, без другого текста:
+<search>запрос</search>
+или для статьи из википедии:
+<wiki>запрос</wiki>
+после получения результатов ответишь нормально пользователю.
+не используй поиск если знаешь ответ или вопрос простой.
+"""

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
         }

vars.py

@@ -12,3 +12,5 @@ user_contexts = {}
 temp_benned = {}
 # последние сообщения групп для подглядывания
 group_last_messages = {}
+# чтобы не спамить админу алертами при падении провайдера
+provider_alert_time = {}