From 0ddc3af8bf69163ce1e5304e67ef7e0004e4533a Mon Sep 17 00:00:00 2001 From: F4ria Date: Mon, 25 Dec 2023 15:32:56 +0800 Subject: [PATCH] feat: extracting gemini's reply text from the exception if available. --- handlers/gemini.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/handlers/gemini.py b/handlers/gemini.py index 2244a18..27cdf25 100644 --- a/handlers/gemini.py +++ b/handlers/gemini.py @@ -3,6 +3,7 @@ from pathlib import Path import re import google.generativeai as genai +from google.generativeai.types.generation_types import StopCandidateException from telebot import TeleBot from telebot.types import Message @@ -161,8 +162,18 @@ def gemini_handler(message: Message, bot: TeleBot) -> None: # keep the last 5, every has two ask and answer. if len(player.history) > 10: player.history = player.history[2:] - player.send_message(m) - gemini_reply_text = player.last.text.strip() + + try: + player.send_message(m) + gemini_reply_text = player.last.text.strip() + except StopCandidateException as e: + match = re.search(r'content\s*{\s*parts\s*{\s*text:\s*"([^"]+)"', str(e)) + if match: + gemini_reply_text = match.group(1) + gemini_reply_text = re.sub(r"\\n", "\n", gemini_reply_text) + else: + raise e + try: bot.reply_to( message,