feat: yi large and md

Signed-off-by: yihong0618 <zouzou0208@gmail.com>

.gitignore

@@ -167,3 +167,4 @@ nohup.out
 *.mp3
 *.mp4
 *.pdf
+.pdm-python

handlers/llama.py

@@ -93,12 +93,12 @@ def llama_pro_handler(message: Message, bot: TeleBot) -> None:
 
     player_message = []
     # restart will lose all TODO
-    if str(message.from_user.id) not in llama_player_dict:
-        llama_player_dict[str(message.from_user.id)] = (
+    if str(message.from_user.id) not in llama_pro_player_dict:
+        llama_pro_player_dict[str(message.from_user.id)] = (
             player_message  # for the imuutable list
         )
     else:
-        player_message = llama_player_dict[str(message.from_user.id)]
+        player_message = llama_pro_player_dict[str(message.from_user.id)]
     if m.strip() == "clear":
         bot.reply_to(
             message,

handlers/qwen.py

@@ -93,12 +93,12 @@ def qwen_pro_handler(message: Message, bot: TeleBot) -> None:
 
     player_message = []
     # restart will lose all TODO
-    if str(message.from_user.id) not in qwen_player_dict:
-        qwen_player_dict[str(message.from_user.id)] = (
+    if str(message.from_user.id) not in qwen_pro_player_dict:
+        qwen_pro_player_dict[str(message.from_user.id)] = (
             player_message  # for the imuutable list
         )
     else:
-        player_message = qwen_player_dict[str(message.from_user.id)]
+        player_message = qwen_pro_player_dict[str(message.from_user.id)]
     if m.strip() == "clear":
         bot.reply_to(
             message,

handlers/useful.py

@@ -0,0 +1,19 @@
+# useful md for myself and you.
+
+from telebot import TeleBot
+from telebot.types import Message
+
+
+from . import *
+
+
+def md_handler(message: Message, bot: TeleBot):
+    """pretty md: /md <address>"""
+    who = "Markdown"
+    reply_id = bot_reply_first(message, who, bot)
+    bot_reply_markdown(reply_id, who, message.text.strip(), bot)
+
+
+def register(bot: TeleBot) -> None:
+    bot.register_message_handler(md_handler, commands=["md"], pass_bot=True)
+    bot.register_message_handler(md_handler, regexp="^md:", pass_bot=True)

handlers/yi.py

@@ -1,15 +1,18 @@
 from os import environ
+import time
 
 from openai import OpenAI
 import requests
 from telebot import TeleBot
 from telebot.types import Message
+from telegramify_markdown import convert
 
 from . import *
 
 YI_BASE_URL = environ.get("YI_BASE_URL")
 YI_API_KEY = environ.get("YI_API_KEY")
 YI_MODEL = "yi-34b-chat-200k"
+YI_PRO_MODEL = "yi-large"
 
 client = OpenAI(
     # defaults to os.environ.get("OPENAI_API_KEY")
@@ -86,6 +89,76 @@ def yi_handler(message: Message, bot: TeleBot) -> None:
     bot_reply_markdown(reply_id, who, yi_reply_text, bot)
 
 
+def yi_pro_handler(message: Message, bot: TeleBot) -> None:
+    """yi_pro : /yi_pro <question>"""
+    m = message.text.strip()
+
+    player_message = []
+    # restart will lose all TODO
+    if str(message.from_user.id) not in yi_player_dict:
+        yi_player_dict[str(message.from_user.id)] = (
+            player_message  # for the imuutable list
+        )
+    else:
+        player_message = yi_player_dict[str(message.from_user.id)]
+    if m.strip() == "clear":
+        bot.reply_to(
+            message,
+            "just clear your yi messages history",
+        )
+        player_message.clear()
+        return
+    if m[:4].lower() == "new ":
+        m = m[4:].strip()
+        player_message.clear()
+    m = enrich_text_with_urls(m)
+
+    who = "yi Pro"
+    reply_id = bot_reply_first(message, who, bot)
+
+    player_message.append({"role": "user", "content": m})
+    # keep the last 5, every has two ask and answer.
+    if len(player_message) > 10:
+        player_message = player_message[2:]
+
+    try:
+        r = client.chat.completions.create(
+            messages=player_message,
+            max_tokens=8192,
+            model=YI_PRO_MODEL,
+            stream=True,
+        )
+        s = ""
+        start = time.time()
+        for chunk in r:
+            if chunk.choices[0].delta.content is None:
+                break
+            s += chunk.choices[0].delta.content
+            # 0.7 is enough for yi3 here its very fast
+            if time.time() - start > 0.7:
+                start = time.time()
+                bot_reply_markdown(reply_id, who, s, bot, split_text=False)
+
+        if not bot_reply_markdown(reply_id, who, s, bot):
+            # maybe not complete
+            # maybe the same message
+            player_message.clear()
+            return
+
+        player_message.append(
+            {
+                "role": "assistant",
+                "content": convert(s),
+            }
+        )
+
+    except Exception as e:
+        print(e)
+        bot_reply_markdown(reply_id, who, "answer wrong", bot)
+        player_message.clear()
+        return
+
+
 def yi_photo_handler(message: Message, bot: TeleBot) -> None:
     s = message.caption
     prompt = s.strip()
@@ -139,6 +212,8 @@ if YI_API_KEY and YI_BASE_URL:
     def register(bot: TeleBot) -> None:
         bot.register_message_handler(yi_handler, commands=["yi"], pass_bot=True)
         bot.register_message_handler(yi_handler, regexp="^yi:", pass_bot=True)
+        bot.register_message_handler(yi_handler, commands=["yi_pro"], pass_bot=True)
+        bot.register_message_handler(yi_handler, regexp="^yi_pro:", pass_bot=True)
         bot.register_message_handler(
             yi_photo_handler,
             content_types=["photo"],