mirror of
https://github.com/cdryzun/tg_bot_collections.git
synced 2025-04-29 00:27:09 +08:00
feat: support gemini pro
Signed-off-by: yihong0618 <zouzou0208@gmail.com>
This commit is contained in:
parent
e65ee0d6e9
commit
c87d7f99ea
@ -1,13 +1,20 @@
|
||||
from os import environ
|
||||
import re
|
||||
import time
|
||||
|
||||
import google.generativeai as genai
|
||||
from google.generativeai.types.generation_types import StopCandidateException
|
||||
from telebot import TeleBot
|
||||
from telebot.types import Message
|
||||
|
||||
from telegramify_markdown import convert
|
||||
from telegramify_markdown.customize import markdown_symbol
|
||||
|
||||
from . import *
|
||||
|
||||
markdown_symbol.head_level_1 = "📌" # If you want, Customizing the head level 1 symbol
|
||||
markdown_symbol.link = "🔗" # If you want, Customizing the link symbol
|
||||
|
||||
GOOGLE_GEMINI_KEY = environ.get("GOOGLE_GEMINI_KEY")
|
||||
|
||||
genai.configure(api_key=GOOGLE_GEMINI_KEY)
|
||||
@ -27,11 +34,16 @@ safety_settings = [
|
||||
|
||||
# Global history cache
|
||||
gemini_player_dict = {}
|
||||
gemini_pro_player_dict = {}
|
||||
|
||||
|
||||
def make_new_gemini_convo():
|
||||
def make_new_gemini_convo(is_pro=False):
|
||||
model_name = "models/gemini-1.0-pro-latest"
|
||||
if is_pro:
|
||||
model_name = "models/gemini-1.5-pro-latest"
|
||||
|
||||
model = genai.GenerativeModel(
|
||||
model_name="gemini-pro",
|
||||
model_name=model_name,
|
||||
generation_config=generation_config,
|
||||
safety_settings=safety_settings,
|
||||
)
|
||||
@ -87,6 +99,75 @@ def gemini_handler(message: Message, bot: TeleBot) -> None:
|
||||
bot_reply_markdown(reply_id, "Gemini", gemini_reply_text, bot)
|
||||
|
||||
|
||||
def gemini_pro_handler(message: Message, bot: TeleBot) -> None:
|
||||
"""Gemini : /gemini_pro <question>"""
|
||||
m = message.text.strip()
|
||||
player = None
|
||||
# restart will lose all TODO
|
||||
if str(message.from_user.id) not in gemini_pro_player_dict:
|
||||
player = make_new_gemini_convo(is_pro=True)
|
||||
gemini_pro_player_dict[str(message.from_user.id)] = player
|
||||
else:
|
||||
player = gemini_pro_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
|
||||
|
||||
# show something, make it more responsible
|
||||
reply_id = bot_reply_first(message, "Geminipro", bot)
|
||||
|
||||
# keep the last 5, every has two ask and answer.
|
||||
if len(player.history) > 10:
|
||||
player.history = player.history[2:]
|
||||
|
||||
try:
|
||||
r = player.send_message(m, stream=True)
|
||||
s = ""
|
||||
start = time.time()
|
||||
for e in r:
|
||||
s += e.text
|
||||
print(s)
|
||||
if time.time() - start > 1.7:
|
||||
start = time.time()
|
||||
try:
|
||||
# maybe the same message
|
||||
if not reply_id:
|
||||
continue
|
||||
bot.edit_message_text(
|
||||
message_id=reply_id.message_id,
|
||||
chat_id=reply_id.chat.id,
|
||||
text=convert(s),
|
||||
parse_mode="MarkdownV2",
|
||||
)
|
||||
except Exception as e:
|
||||
print(str(e))
|
||||
try:
|
||||
# maybe not complete
|
||||
# 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 Exception as e:
|
||||
player.history.clear()
|
||||
print(str(e))
|
||||
return
|
||||
except:
|
||||
bot.reply_to(
|
||||
message,
|
||||
"claude answer:\n" + "geminipro answer timeout",
|
||||
parse_mode="MarkdownV2",
|
||||
)
|
||||
player.history.clear()
|
||||
return
|
||||
|
||||
|
||||
def gemini_photo_handler(message: Message, bot: TeleBot) -> None:
|
||||
s = message.caption
|
||||
reply_message = bot.reply_to(
|
||||
@ -117,6 +198,12 @@ def gemini_photo_handler(message: Message, bot: TeleBot) -> None:
|
||||
def register(bot: TeleBot) -> None:
|
||||
bot.register_message_handler(gemini_handler, commands=["gemini"], pass_bot=True)
|
||||
bot.register_message_handler(gemini_handler, regexp="^gemini:", pass_bot=True)
|
||||
bot.register_message_handler(
|
||||
gemini_pro_handler, commands=["gemini_pro"], pass_bot=True
|
||||
)
|
||||
bot.register_message_handler(
|
||||
gemini_pro_handler, regexp="^gemini_pro:", pass_bot=True
|
||||
)
|
||||
bot.register_message_handler(
|
||||
gemini_photo_handler,
|
||||
content_types=["photo"],
|
||||
|
Loading…
x
Reference in New Issue
Block a user