cdryzun/tg_bot_collections
1
0
Fork 0
mirror of https://github.com/cdryzun/tg_bot_collections.git synced 2025-04-29 00:27:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: claude opus

Browse Source 
Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
yihong0618 2024-03-18 15:09:36 +08:00
parent 2952b81521
commit 68ab644e4a
4 changed files with 118 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

105
handlers/claude.py
View File

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

7
handlers/gemini.py
View File

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

14
handlers/yi.py
View File

@ -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):

3
tg.py
View File

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