From 77fefcb4b51a39bcabec518af6ca5672e5700ae9 Mon Sep 17 00:00:00 2001 From: yihong0618 Date: Mon, 13 May 2024 14:03:04 +0800 Subject: [PATCH] feat: yi large and md Signed-off-by: yihong0618 --- .gitignore | 1 + handlers/llama.py | 6 ++-- handlers/qwen.py | 6 ++-- handlers/useful.py | 19 ++++++++++++ handlers/yi.py | 75 ++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 101 insertions(+), 6 deletions(-) create mode 100644 handlers/useful.py diff --git a/.gitignore b/.gitignore index 4a571f8..6c05fbe 100644 --- a/.gitignore +++ b/.gitignore @@ -167,3 +167,4 @@ nohup.out *.mp3 *.mp4 *.pdf +.pdm-python diff --git a/handlers/llama.py b/handlers/llama.py index 8693098..8e64fa0 100644 --- a/handlers/llama.py +++ b/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, diff --git a/handlers/qwen.py b/handlers/qwen.py index 0605951..40c0bd8 100644 --- a/handlers/qwen.py +++ b/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, diff --git a/handlers/useful.py b/handlers/useful.py new file mode 100644 index 0000000..25e50fd --- /dev/null +++ b/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
""" + 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) diff --git a/handlers/yi.py b/handlers/yi.py index 495db24..ebdc00a 100644 --- a/handlers/yi.py +++ b/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 """ + 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"],