feat: llama3 bot

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

README.md

@@ -49,6 +49,12 @@ Note, if you are using third party service, you need to `export ANTHROPIC_BASE_U
 
 Note, if you are using third party service, you need to `export OPENAI_API_BASE=${the_url}` to change the url.
 
+## Bot -> llama3
+
+1. visit https://console.groq.com/docs/quickstart get the key
+2. export GROQ_API_KEY=${the_key}
+3. use `llama_pro: ${message}` to ask
+
 
 ## HOW TO Install and Run
 

handlers/chatgpt.py

@@ -205,16 +205,20 @@ def chatgpt_photo_handler(message: Message, bot: TeleBot) -> None:
         bot_reply_markdown(reply_id, who, "answer wrong", bot)
 
 
-def register(bot: TeleBot) -> None:
+if CHATGPT_API_KEY:
-    bot.register_message_handler(chatgpt_handler, commands=["gpt"], pass_bot=True)
+
-    bot.register_message_handler(chatgpt_handler, regexp="^gpt:", pass_bot=True)
+    def register(bot: TeleBot) -> None:
-    bot.register_message_handler(
+        bot.register_message_handler(chatgpt_handler, commands=["gpt"], pass_bot=True)
-        chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True
+        bot.register_message_handler(chatgpt_handler, regexp="^gpt:", pass_bot=True)
-    )
+        bot.register_message_handler(
-    bot.register_message_handler(chatgpt_pro_handler, regexp="^gpt_pro:", pass_bot=True)
+            chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True
-    bot.register_message_handler(
+        )
-        chatgpt_photo_handler,
+        bot.register_message_handler(
-        content_types=["photo"],
+            chatgpt_pro_handler, regexp="^gpt_pro:", pass_bot=True
-        func=lambda m: m.caption and m.caption.startswith(("gpt:", "/gpt")),
+        )
-        pass_bot=True,
+        bot.register_message_handler(
-    )
+            chatgpt_photo_handler,
+            content_types=["photo"],
+            func=lambda m: m.caption and m.caption.startswith(("gpt:", "/gpt")),
+            pass_bot=True,
+        )

handlers/claude.py

@@ -19,7 +19,13 @@ ANTHROPIC_BASE_URL = environ.get("ANTHROPIC_BASE_URL")
 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)
+if environ.get("ANTHROPIC_BASE_URL"):
+    client = Anthropic(
+        base_url=ANTHROPIC_BASE_URL, api_key=ANTHROPIC_API_KEY, timeout=20
+    )
+else:
+    client = Anthropic(api_key=ANTHROPIC_API_KEY, timeout=20)
+
 
 # Global history cache
 claude_player_dict = {}
@@ -217,18 +223,20 @@ def claude_photo_handler(message: Message, bot: TeleBot) -> None:
         bot_reply_markdown(reply_id, who, "answer wrong", bot)
 
 
-def register(bot: TeleBot) -> None:
+if ANTHROPIC_API_KEY:
-    bot.register_message_handler(claude_handler, commands=["claude"], pass_bot=True)
+
-    bot.register_message_handler(claude_handler, regexp="^claude:", pass_bot=True)
+    def register(bot: TeleBot) -> None:
-    bot.register_message_handler(
+        bot.register_message_handler(claude_handler, commands=["claude"], pass_bot=True)
-        claude_pro_handler, commands=["claude_pro"], pass_bot=True
+        bot.register_message_handler(claude_handler, regexp="^claude:", pass_bot=True)
-    )
+        bot.register_message_handler(
-    bot.register_message_handler(
+            claude_pro_handler, commands=["claude_pro"], pass_bot=True
-        claude_pro_handler, regexp="^claude_pro:", pass_bot=True
+        )
-    )
+        bot.register_message_handler(
-    bot.register_message_handler(
+            claude_pro_handler, regexp="^claude_pro:", pass_bot=True
-        claude_photo_handler,
+        )
-        content_types=["photo"],
+        bot.register_message_handler(
-        func=lambda m: m.caption and m.caption.startswith(("claude:", "/claude")),
+            claude_photo_handler,
-        pass_bot=True,
+            content_types=["photo"],
-    )
+            func=lambda m: m.caption and m.caption.startswith(("claude:", "/claude")),
+            pass_bot=True,
+        )

handlers/llama.py

@@ -0,0 +1,170 @@
+from os import environ
+import time
+
+from openai import OpenAI
+from telebot import TeleBot
+from telebot.types import Message
+
+from . import *
+
+from groq import Groq
+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
+markdown_symbol.link = "🔗"  # If you want, Customizing the link symbol
+
+LLAMA_API_KEY = environ.get("GROQ_API_KEY")
+LLAMA_MODEL = "llama3-8b-8192"
+LLAMA_PRO_MODEL = "llama3-70b-8192"
+
+if LLAMA_API_KEY:
+    client = Groq(api_key=LLAMA_API_KEY)
+
+# Global history cache
+llama_player_dict = {}
+llama_pro_player_dict = {}
+
+
+def llama_handler(message: Message, bot: TeleBot) -> None:
+    """llama : /llama <question>"""
+    m = message.text.strip()
+
+    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)] = (
+            player_message  # for the imuutable list
+        )
+    else:
+        player_message = llama_player_dict[str(message.from_user.id)]
+    if m.strip() == "clear":
+        bot.reply_to(
+            message,
+            "just clear your llama 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 = "llama"
+    # show something, make it more responsible
+    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:]
+
+    llama_reply_text = ""
+    try:
+        r = client.chat.completions.create(
+            messages=player_message, max_tokens=8192, model=LLAMA_MODEL
+        )
+        content = r.choices[0].message.content.encode("utf8").decode()
+        if not content:
+            llama_reply_text = f"{who} did not answer."
+            player_message.pop()
+        else:
+            llama_reply_text = content
+            player_message.append(
+                {
+                    "role": "assistant",
+                    "content": llama_reply_text,
+                }
+            )
+
+    except Exception as e:
+        print(e)
+        bot_reply_markdown(reply_id, who, "answer wrong", bot)
+        # pop my user
+        player_message.pop()
+        return
+
+    # reply back as Markdown and fallback to plain text if failed.
+    bot_reply_markdown(reply_id, who, llama_reply_text, bot)
+
+
+def llama_pro_handler(message: Message, bot: TeleBot) -> None:
+    """llama_pro : /llama_pro <question>"""
+    m = message.text.strip()
+
+    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)] = (
+            player_message  # for the imuutable list
+        )
+    else:
+        player_message = llama_player_dict[str(message.from_user.id)]
+    if m.strip() == "clear":
+        bot.reply_to(
+            message,
+            "just clear your llama 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 = "llama 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=LLAMA_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
+            if time.time() - start > 2.0:
+                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
+
+
+if LLAMA_API_KEY:
+
+    def register(bot: TeleBot) -> None:
+        bot.register_message_handler(llama_handler, commands=["llama"], pass_bot=True)
+        bot.register_message_handler(llama_handler, regexp="^llama:", pass_bot=True)
+        bot.register_message_handler(
+            llama_pro_handler, commands=["llama_pro"], pass_bot=True
+        )
+        bot.register_message_handler(
+            llama_pro_handler, regexp="^llama_pro:", pass_bot=True
+        )

handlers/sd.py

@@ -8,62 +8,66 @@ from . import *
 
 SD_API_KEY = environ.get("SD3_KEY")
 
-if SD_API_KEY is not None:
 
-    def get_user_balance():
+def get_user_balance():
-        api_host = "https://api.stability.ai"
+    api_host = "https://api.stability.ai"
-        url = f"{api_host}/v1/user/balance"
+    url = f"{api_host}/v1/user/balance"
 
-        response = requests.get(url, headers={"Authorization": f"Bearer {SD_API_KEY}"})
+    response = requests.get(url, headers={"Authorization": f"Bearer {SD_API_KEY}"})
 
-        if response.status_code != 200:
+    if response.status_code != 200:
-            print("Non-200 response: " + str(response.text))
+        print("Non-200 response: " + str(response.text))
 
-        # Do something with the payload...
+    # Do something with the payload...
-        payload = response.json()
+    payload = response.json()
-        return payload["credits"]
+    return payload["credits"]
 
-    def generate_sd3_image(prompt):
-        response = requests.post(
-            f"https://api.stability.ai/v2beta/stable-image/generate/sd3",
-            headers={"authorization": f"Bearer {SD_API_KEY}", "accept": "image/*"},
-            files={"none": ""},
-            data={
-                "prompt": prompt,
-                "model": "sd3-turbo",
-                "output_format": "jpeg",
-            },
-        )
 
-        if response.status_code == 200:
+def generate_sd3_image(prompt):
-            with open("sd3.jpeg", "wb") as file:
+    response = requests.post(
-                file.write(response.content)
+        f"https://api.stability.ai/v2beta/stable-image/generate/sd3",
-            return True
+        headers={"authorization": f"Bearer {SD_API_KEY}", "accept": "image/*"},
+        files={"none": ""},
+        data={
+            "prompt": prompt,
+            "model": "sd3-turbo",
+            "output_format": "jpeg",
+        },
+    )
+
+    if response.status_code == 200:
+        with open("sd3.jpeg", "wb") as file:
+            file.write(response.content)
+        return True
+    else:
+        print(str(response.json()))
+        return False
+
+
+def sd_handler(message: Message, bot: TeleBot):
+    """pretty sd3: /sd3 <address>"""
+    credits = get_user_balance()
+    bot.reply_to(
+        message,
+        f"Generating pretty sd3-turbo image may take some time please left credits {credits} every try will cost 4 criedits wait:",
+    )
+    m = message.text.strip()
+    prompt = m.strip()
+    try:
+        r = generate_sd3_image(prompt)
+        if r:
+            with open(f"sd3.jpeg", "rb") as photo:
+                bot.send_photo(
+                    message.chat.id, photo, reply_to_message_id=message.message_id
+                )
         else:
-            print(str(response.json()))
+            bot.reply_to(message, "prompt error")
-            return False
+    except Exception as e:
+        print(e)
+        bot.reply_to(message, "sd3 error")
 
-    def sd_handler(message: Message, bot: TeleBot):
+
-        """pretty sd3: /sd3 <address>"""
+if SD_API_KEY:
-        credits = get_user_balance()
-        bot.reply_to(
-            message,
-            f"Generating pretty sd3-turbo image may take some time please left credits {credits} every try will cost 4 criedits wait:",
-        )
-        m = message.text.strip()
-        prompt = m.strip()
-        try:
-            r = generate_sd3_image(prompt)
-            if r:
-                with open(f"sd3.jpeg", "rb") as photo:
-                    bot.send_photo(
-                        message.chat.id, photo, reply_to_message_id=message.message_id
-                    )
-            else:
-                bot.reply_to(message, "prompt error")
-        except Exception as e:
-            print(e)
-            bot.reply_to(message, "sd3 error")
 
     def register(bot: TeleBot) -> None:
         bot.register_message_handler(sd_handler, commands=["sd3"], pass_bot=True)

handlers/yi.py

@@ -134,12 +134,14 @@ def yi_photo_handler(message: Message, bot: TeleBot) -> None:
         bot_reply_markdown(reply_id, who, "answer wrong", bot)
 
 
-def register(bot: TeleBot) -> None:
+if YI_API_KEY and YI_BASE_URL:
-    bot.register_message_handler(yi_handler, commands=["yi"], pass_bot=True)
+
-    bot.register_message_handler(yi_handler, regexp="^yi:", pass_bot=True)
+    def register(bot: TeleBot) -> None:
-    bot.register_message_handler(
+        bot.register_message_handler(yi_handler, commands=["yi"], pass_bot=True)
-        yi_photo_handler,
+        bot.register_message_handler(yi_handler, regexp="^yi:", pass_bot=True)
-        content_types=["photo"],
+        bot.register_message_handler(
-        func=lambda m: m.caption and m.caption.startswith(("yi:", "/yi")),
+            yi_photo_handler,
-        pass_bot=True,
+            content_types=["photo"],
-    )
+            func=lambda m: m.caption and m.caption.startswith(("yi:", "/yi")),
+            pass_bot=True,
+        )

requirements.txt

@@ -9,3 +9,4 @@ openai
 requests
 urlextract
 requests
+groq