feat: gemini bot

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

.gitignore

@@ -162,3 +162,4 @@ cython_debug/
 *.png
 OUT_FOLDER/
 cache/
+nohup.out

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

requirements.txt

@@ -2,4 +2,5 @@ pyTelegramBotAPI==4.14.0
 cairosvg
 github_poster
 prettymapp
+google-generativeai==2.25.2
 

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 <github_user_name> [<start>-<end>]"
             ),
             BotCommand("map", "pretty map: /map <address>"),
+            BotCommand("gemini", "Gemini : /gemini <question>"),
         ]
     )
     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()