diff --git a/handlers/claude.py b/handlers/claude.py index 160f8e9..1aaf141 100644 --- a/handlers/claude.py +++ b/handlers/claude.py @@ -1,11 +1,12 @@ from os import environ from pathlib import Path +import time from anthropic import Anthropic, APITimeoutError from telebot import TeleBot from telebot.types import Message -import telegramify_markdown +from telegramify_markdown import convert from telegramify_markdown.customize import markdown_symbol markdown_symbol.head_level_1 = "📌" # If you want, Customizing the head level 1 symbol @@ -13,12 +14,14 @@ markdown_symbol.link = "🔗" # If you want, Customizing the link symbol ANTHROPIC_API_KEY = environ.get("ANTHROPIC_API_KEY") ANTHROPIC_BASE_URL = environ.get("ANTHROPIC_BASE_URL") -ANTHROPIC_MODEL = "claude-3-haiku-20240307" # change model here you can use claude-3-opus-20240229 but for now its slow +ANTHROPIC_MODEL = "claude-3-haiku-20240307" +ANTHROPIC_PRO_MODEL = "claude-3-opus-20240229" client = Anthropic(base_url=ANTHROPIC_BASE_URL, api_key=ANTHROPIC_API_KEY, timeout=20) # Global history cache claude_player_dict = {} +claude_pro_player_dict = {} def claude_handler(message: Message, bot: TeleBot) -> None: @@ -32,6 +35,13 @@ def claude_handler(message: Message, bot: TeleBot) -> None: ) else: player_message = claude_player_dict[str(message.from_user.id)] + if m.strip() == "clear": + bot.reply_to( + message, + "just clear you claude messages history", + ) + player_message.clear() + return player_message.append({"role": "user", "content": m}) # keep the last 5, every has two ask and answer. @@ -45,7 +55,7 @@ def claude_handler(message: Message, bot: TeleBot) -> None: # tricky player_message.pop() r = client.messages.create( - max_tokens=1024, messages=player_message, model=ANTHROPIC_MODEL + max_tokens=1024, messages=player_message, model=ANTHROPIC_PRO_MODEL ) if not r.content: claude_reply_text = "Claude did not answer." @@ -72,16 +82,95 @@ def claude_handler(message: Message, bot: TeleBot) -> None: try: bot.reply_to( message, - "Claude answer:\n" + telegramify_markdown.convert(claude_reply_text), + "Claude answer:\n" + convert(claude_reply_text), parse_mode="MarkdownV2", ) - return except: print("wrong markdown format") bot.reply_to( message, "claude answer:\n\n" + claude_reply_text, ) + + +def claude_pro_handler(message: Message, bot: TeleBot) -> None: + """claude_pro : /claude_pro TODO refactor""" + m = message.text.strip() + player_message = [] + if str(message.from_user.id) not in claude_pro_player_dict: + claude_pro_player_dict[str(message.from_user.id)] = ( + player_message # for the imuutable list + ) + else: + player_message = claude_pro_player_dict[str(message.from_user.id)] + if m.strip() == "clear": + bot.reply_to( + message, + "just clear you claude opus messages history", + ) + player_message.clear() + return + + 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:] + + claude_reply_text = "" + try: + if len(player_message) > 2: + if player_message[-1]["role"] == player_message[-2]["role"]: + # tricky + player_message.pop() + r = client.messages.create( + max_tokens=1024, + messages=player_message, + model=ANTHROPIC_PRO_MODEL, + stream=True, + ) + s = "" + start = time.time() + is_send = True + reply_id = None + for e in r: + if e.type == "content_block_delta": + s += e.delta.text + if time.time() - start > 1.1: + start = time.time() + if is_send: + reply_id = bot.reply_to( + message, + convert(s), + parse_mode="MarkdownV2", + ) + is_send = False + else: + try: + # maybe the same message + bot.edit_message_text( + message_id=reply_id.message_id, + chat_id=reply_id.chat.id, + text=convert(s), + parse_mode="MarkdownV2", + ) + except: + pass + + player_message.append( + { + "role": "assistant", + "content": convert(s), + } + ) + + except APITimeoutError: + bot.reply_to( + message, + "claude answer:\n" + "claude answer timeout", + parse_mode="MarkdownV2", + ) + # pop my user + player_message.pop() return @@ -140,6 +229,12 @@ def claude_photo_handler(message: Message, bot: TeleBot) -> None: def register(bot: TeleBot) -> None: bot.register_message_handler(claude_handler, commands=["claude"], pass_bot=True) bot.register_message_handler(claude_handler, regexp="^claude:", pass_bot=True) + bot.register_message_handler( + claude_pro_handler, commands=["claude_pro"], pass_bot=True + ) + bot.register_message_handler( + claude_pro_handler, regexp="^claude_pro:", pass_bot=True + ) bot.register_message_handler( claude_photo_handler, content_types=["photo"], diff --git a/handlers/gemini.py b/handlers/gemini.py index 3d22862..8c6549c 100644 --- a/handlers/gemini.py +++ b/handlers/gemini.py @@ -52,6 +52,13 @@ def gemini_handler(message: Message, bot: TeleBot) -> None: gemini_player_dict[str(message.from_user.id)] = player else: player = gemini_player_dict[str(message.from_user.id)] + if m.strip() == "clear": + bot.reply_to( + message, + "just clear you gemini messages history", + ) + player.history.clear() + return # keep the last 5, every has two ask and answer. if len(player.history) > 10: player.history = player.history[2:] diff --git a/handlers/yi.py b/handlers/yi.py index 1e6ed39..04a312b 100644 --- a/handlers/yi.py +++ b/handlers/yi.py @@ -22,7 +22,6 @@ client = OpenAI( api_key=YI_API_KEY, base_url=YI_BASE_URL, ) -print(client.base_url) # Global history cache yi_player_dict = {} @@ -31,6 +30,7 @@ yi_player_dict = {} def yi_handler(message: Message, bot: TeleBot) -> None: """yi : /yi """ m = message.text.strip() + player_message = [] # restart will lose all TODO if str(message.from_user.id) not in yi_player_dict: @@ -39,6 +39,13 @@ def yi_handler(message: Message, bot: TeleBot) -> None: ) else: player_message = yi_player_dict[str(message.from_user.id)] + if m.strip() == "clear": + bot.reply_to( + message, + "just clear you yi messages history", + ) + player_message.clear() + return player_message.append({"role": "user", "content": m}) # keep the last 5, every has two ask and answer. @@ -52,6 +59,7 @@ def yi_handler(message: Message, bot: TeleBot) -> None: # tricky player_message.pop() r = client.chat.completions.create(messages=player_message, model=YI_MODEL) + content = r.choices[0].message.content.encode("utf8").decode() if not content: yi_reply_text = "yi did not answer." @@ -61,7 +69,7 @@ def yi_handler(message: Message, bot: TeleBot) -> None: player_message.append( { "role": "assistant", - "content": content, + "content": yi_reply_text, } ) @@ -82,14 +90,12 @@ def yi_handler(message: Message, bot: TeleBot) -> None: "yi answer:\n" + telegramify_markdown.convert(yi_reply_text), parse_mode="MarkdownV2", ) - return except: print("wrong markdown format") bot.reply_to( message, "yi answer:\n\n" + yi_reply_text, ) - return def _image_to_data_uri(file_path): diff --git a/tg.py b/tg.py index a2fecc9..528ba88 100644 --- a/tg.py +++ b/tg.py @@ -1,5 +1,4 @@ import argparse - from telebot import TeleBot from handlers import list_available_commands, load_handlers @@ -32,7 +31,7 @@ def main(): # Start bot print("Starting tg collections bot.") - bot.infinity_polling() + bot.infinity_polling(timeout=10, long_polling_timeout=5) if __name__ == "__main__":