diff --git a/.gitignore b/.gitignore index 8a9f4f6..361fc8a 100644 --- a/.gitignore +++ b/.gitignore @@ -162,3 +162,4 @@ cython_debug/ *.png OUT_FOLDER/ cache/ +nohup.out diff --git a/README.md b/README.md index 1c2ce7e..7bf7e21 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,13 @@ for yihong0618's channel: https://t.me/hyi0618 ![image](https://github.com/yihong0618/tg_bot_collections/assets/15976103/29848d22-5289-4953-8ab0-4e84c16f79e3) +## Bot -> Gemini player + +1. visit https://makersuite.google.com/app/apikey get the key +2. export GOOGLE_GEMINI_KEY=${the_key} +3. use `gemini: ${message}` to ask + + ## HOW TO 1. pip install -r requirements.txt diff --git a/requirements.txt b/requirements.txt index f146ba1..f81b1d9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,5 @@ pyTelegramBotAPI==4.14.0 cairosvg github_poster prettymapp +google-generativeai==2.25.2 diff --git a/tg.py b/tg.py index cb1c089..0ebdd5d 100644 --- a/tg.py +++ b/tg.py @@ -5,6 +5,7 @@ import shutil import subprocess import traceback from tempfile import SpooledTemporaryFile +from os import environ import numpy as np import PIL @@ -16,10 +17,45 @@ from prettymapp.plotting import Plot as PrettyPlot from prettymapp.settings import STYLES from telebot import TeleBot # type: ignore from telebot.types import BotCommand, Message # type: ignore +import google.generativeai as genai PIL.Image.MAX_IMAGE_PIXELS = 933120000 MAX_IN_MEMORY = 10 * 1024 * 1024 # 10MiB +GOOGLE_GEMINI_KEY = environ.get("GOOGLE_GEMINI_KEY") + + +genai.configure(api_key=GOOGLE_GEMINI_KEY) +generation_config = { + "temperature": 0.9, + "top_p": 1, + "top_k": 1, + "max_output_tokens": 2048, +} + +safety_settings = [ + {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, + {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"}, + { + "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "threshold": "BLOCK_MEDIUM_AND_ABOVE", + }, + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "threshold": "BLOCK_MEDIUM_AND_ABOVE", + }, +] + + +def make_new_gemini_convo(): + model = genai.GenerativeModel( + model_name="gemini-pro", + generation_config=generation_config, + safety_settings=safety_settings, + ) + convo = model.start_chat() + return convo + class Plot(PrettyPlot): # memory leak fix for Plot. thanks @higuoxing https://github.com/higuoxing @@ -93,6 +129,7 @@ def main(): parser.add_argument("tg_token", help="tg token") options = parser.parse_args() print("Arg parse done.") + gemini_player_dict = {} # Init bot bot = TeleBot(options.tg_token) @@ -102,6 +139,7 @@ def main(): "github", "github poster: /github [-]" ), BotCommand("map", "pretty map: /map
"), + BotCommand("gemini", "Gemini : /gemini "), ] ) print("Bot init done.") @@ -190,6 +228,36 @@ def main(): bot.delete_message(reply_message.chat.id, reply_message.message_id) gc.collect() + @bot.message_handler(commands=["gemini"]) + @bot.message_handler(regexp="^gemini:") + def gemini_handler(message: Message): + reply_message = bot.reply_to( + message, + "Generating google gemini answer please wait:", + ) + m = message.text.strip().split(maxsplit=1)[1].strip() + player = None + # restart will lose all TODO + if str(message.from_user.id) not in gemini_player_dict: + player = make_new_gemini_convo() + gemini_player_dict[str(message.from_user.id)] = player + else: + player = gemini_player_dict[str(message.from_user.id)] + if len(player.history) > 10: + bot.reply_to(message, "Your hisotry length > 5 will only keep last 5") + player.history = player.history[2:] + try: + player.send_message(m) + try: + bot.reply_to(message, "Gemini answer:\n" + player.last.text, parse_mode='MarkdownV2') + except: + bot.reply_to(message, "Gemini answer:\n" + player.last.text) + + + except Exception as e: + traceback.print_exc() + bot.reply_to(message, "Something wrong please check") + # Start bot print("Starting tg collections bot.") bot.infinity_polling()