diff --git a/README.md b/README.md index abb8332..636f125 100644 --- a/README.md +++ b/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 diff --git a/handlers/chatgpt.py b/handlers/chatgpt.py index 3234100..b9fb99c 100644 --- a/handlers/chatgpt.py +++ b/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: - bot.register_message_handler(chatgpt_handler, commands=["gpt"], pass_bot=True) - bot.register_message_handler(chatgpt_handler, regexp="^gpt:", pass_bot=True) - bot.register_message_handler( - chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True - ) - bot.register_message_handler(chatgpt_pro_handler, regexp="^gpt_pro:", 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, - ) +if CHATGPT_API_KEY: + + def register(bot: TeleBot) -> None: + bot.register_message_handler(chatgpt_handler, commands=["gpt"], pass_bot=True) + bot.register_message_handler(chatgpt_handler, regexp="^gpt:", pass_bot=True) + bot.register_message_handler( + chatgpt_pro_handler, commands=["gpt_pro"], pass_bot=True + ) + bot.register_message_handler( + chatgpt_pro_handler, regexp="^gpt_pro:", 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, + ) diff --git a/handlers/claude.py b/handlers/claude.py index 3ee5fc1..ebe8582 100644 --- a/handlers/claude.py +++ b/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: - 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"], - func=lambda m: m.caption and m.caption.startswith(("claude:", "/claude")), - pass_bot=True, - ) +if ANTHROPIC_API_KEY: + + 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"], + func=lambda m: m.caption and m.caption.startswith(("claude:", "/claude")), + pass_bot=True, + ) diff --git a/handlers/llama.py b/handlers/llama.py new file mode 100644 index 0000000..f9c3dde --- /dev/null +++ b/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 """ + 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 """ + 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 + ) diff --git a/handlers/sd.py b/handlers/sd.py index cbd2a5b..266eb25 100644 --- a/handlers/sd.py +++ b/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(): - api_host = "https://api.stability.ai" - url = f"{api_host}/v1/user/balance" +def get_user_balance(): + api_host = "https://api.stability.ai" + 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: - print("Non-200 response: " + str(response.text)) + if response.status_code != 200: + print("Non-200 response: " + str(response.text)) - # Do something with the payload... - payload = response.json() - return payload["credits"] + # Do something with the payload... + payload = response.json() + 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: - with open("sd3.jpeg", "wb") as file: - file.write(response.content) - return True +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: + 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
""" + 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())) - return False + bot.reply_to(message, "prompt error") + except Exception as e: + print(e) + bot.reply_to(message, "sd3 error") - def sd_handler(message: Message, bot: TeleBot): - """pretty sd3: /sd3
""" - 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") + +if SD_API_KEY: def register(bot: TeleBot) -> None: bot.register_message_handler(sd_handler, commands=["sd3"], pass_bot=True) diff --git a/handlers/yi.py b/handlers/yi.py index bd3dbef..495db24 100644 --- a/handlers/yi.py +++ b/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: - 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_photo_handler, - content_types=["photo"], - func=lambda m: m.caption and m.caption.startswith(("yi:", "/yi")), - pass_bot=True, - ) +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_photo_handler, + content_types=["photo"], + func=lambda m: m.caption and m.caption.startswith(("yi:", "/yi")), + pass_bot=True, + ) diff --git a/requirements.txt b/requirements.txt index e9e4923..9d4d497 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,5 @@ telegramify-markdown openai requests urlextract -requests \ No newline at end of file +requests +groq \ No newline at end of file