feat: claude opus

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

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 <question> 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"],

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:]

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 <question>"""
     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):

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__":