feat: support gemini close #402 #412 #403 #338 (#420)

* feat: support gemini close #402 #412 #403 #338 Signed-off-by: yihong0618 <zouzou0208@gmail.com> * fix: versions things Signed-off-by: yihong0618 <zouzou0208@gmail.com> * fix: lint Signed-off-by: yihong0618 <zouzou0208@gmail.com> --------- Signed-off-by: yihong0618 <zouzou0208@gmail.com>
yihong0618 · Jan 7, 2024 · e4b30b4 · e4b30b4
1 parent 05a2027
commit e4b30b4
Show file tree

Hide file tree

Showing 8 changed files with 866 additions and 733 deletions.
diff --git a/README.md b/README.md
@@ -16,6 +16,7 @@ Play ChatGPT and other LLM with Xiaomi AI Speaker
 - ChatGPT
 - New Bing
 - [ChatGLM](http://open.bigmodel.cn/)
+- [Gemini](https://makersuite.google.com/app/apikey)
 - [Bard](https://github.com/dsdanielpark/Bard-API)
 
 ## Windows 获取小米音响DID
@@ -73,6 +74,8 @@ xiaogpt --hardware LX06  --mute_xiaoai --stream
 # 如果你想使用 gpt3 ai
 export OPENAI_API_KEY=${your_api_key}
 xiaogpt --hardware LX06  --mute_xiaoai --use_gpt3
+# 如果你想使用 google 的 gemini
+xiaogpt --hardware LX06  --mute_xiaoai --use_gemini --gemini_key ${gemini_key}
 # 如果你想用 edge-tts
 xiaogpt --hardware LX06 --cookie ${cookie} --use_chatgpt_api --tts edge
 # 如果你想使用 LangChain + SerpApi 实现上网检索或其他本地服务（目前仅支持 stream 模式）
@@ -101,6 +104,8 @@ python3 xiaogpt.py --hardware LX06  --mute_xiaoai --use_gpt3
 python3 xiaogpt.py --hardware LX06  --mute_xiaoai --use_glm --glm_key ${glm_key}
 # 如果你想使用 google 的 bard
 python3 xiaogpt.py --hardware LX06  --mute_xiaoai --use_bard --bard_token ${bard_token}
+# 如果你想使用 google 的 gemini
+python3 xiaogpt.py --hardware LX06  --mute_xiaoai --use_gemini --gemini_key ${gemini_key}
 # 如果你想使用 LangChain+SerpApi 实现上网检索或其他本地服务（目前仅支持 stream 模式）
 export OPENAI_API_KEY=${your_api_key}
 export SERPAPI_API_KEY=${your_serpapi_key}
@@ -153,6 +158,7 @@ Bard-API [参考](https://github.com/dsdanielpark/Bard-API)
 | openai_key            | openai的apikey                                                          |                                                                                                           |
 | serpapi_api_key       | serpapi的key 参考 [SerpAPI](https://serpapi.com/)                       |                                                                                                           |
 | glm_key               | chatglm 的 apikey                                                       |                                                                                                           |
+| gemini_key            | gemini 的 apikey [参与](https://makersuite.google.com/app/apikey)                                                      |                                                                                                           |
 | bard_token            | bard 的 token 参考 [Bard-API](https://github.com/dsdanielpark/Bard-API) |                                                                                                           |
 | cookie                | 小爱账户cookie （如果用上面密码登录可以不填）                             |                                                                                                           |
 | mi_did                | 设备did                                                                 |                                                                                                           |

diff --git a/pdm.lock b/pdm.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -10,7 +10,7 @@ classifiers = [
     "Operating System :: OS Independent",
     "Programming Language :: Python :: 3",
 ]
-requires-python = ">=3.8.1"
+requires-python = ">=3.9"
 dependencies = [
     "miservice_fork",
     "openai>=1",
@@ -23,6 +23,7 @@ dependencies = [
     "langchain>=0.0.343",
     "beautifulsoup4>=4.12.0",
     "google-search-results>=2.4.2",
+    "google-generativeai",
     "numexpr>=2.8.6",
 ]
 license = {text = "MIT"}

diff --git a/xiao_config.json.example b/xiao_config.json.example
@@ -4,6 +4,7 @@
   "password": "",
   "openai_key": "",
   "glm_key": "",
+  "gemini_key": "",
   "bard_token": "",
   "serpapi_api_key": "",
   "cookie": "",

diff --git a/xiaogpt/bot/__init__.py b/xiaogpt/bot/__init__.py
@@ -6,6 +6,7 @@
 from xiaogpt.bot.newbing_bot import NewBingBot
 from xiaogpt.bot.glm_bot import GLMBot
 from xiaogpt.bot.bard_bot import BardBot
+from xiaogpt.bot.gemini_bot import GeminiBot
 from xiaogpt.bot.langchain_bot import LangChainBot
 from xiaogpt.config import Config
 
@@ -15,6 +16,7 @@
     "chatgptapi": ChatGPTBot,
     "glm": GLMBot,
     "bard": BardBot,
+    "gemini": GeminiBot,
     "langchain": LangChainBot,
 }
 
@@ -32,6 +34,7 @@ def get_bot(config: Config) -> BaseBot:
     "NewBingBot",
     "GLMBot",
     "BardBot",
+    "GeminiBot",
     "get_bot",
     "LangChainBot",
 ]
diff --git a/xiaogpt/bot/gemini_bot.py b/xiaogpt/bot/gemini_bot.py
@@ -0,0 +1,64 @@
+"""Google Gemini bot"""
+from __future__ import annotations
+
+from typing import Any
+
+from rich import print
+import google.generativeai as genai
+from google.generativeai.types.generation_types import StopCandidateException
+
+from xiaogpt.bot.base_bot import BaseBot, ChatHistoryMixin
+
+generation_config = {
+    "temperature": 0.7,
+    "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",
+    },
+]
+
+
+class GeminiBot(ChatHistoryMixin, BaseBot):
+    name = "Gemini"
+
+    def __init__(self, gemini_key: str) -> None:
+        genai.configure(api_key=gemini_key)
+        self.history = []
+        model = genai.GenerativeModel(
+            model_name="gemini-pro",
+            generation_config=generation_config,
+            safety_settings=safety_settings,
+        )
+        self.convo = model.start_chat()
+
+    @classmethod
+    def from_config(cls, config):
+        return cls(gemini_key=config.gemini_key)
+
+    async def ask(self, query, **options):
+        self.convo.send_message(query)
+        message = self.convo.last.text.strip()
+        print(message)
+        if len(self.convo.history) > 10:
+            self.convo.history = self.convo.history[2:]
+        return message
+
+    async def ask_stream(self, query: str, **options: Any):
+        if len(self.convo.history) > 10:
+            self.convo.history = self.convo.history[2:]
+        response = self.convo.send_message(query, stream=True)
+        for chunk in response:
+            print(chunk.text)
+            yield chunk.text
diff --git a/xiaogpt/cli.py b/xiaogpt/cli.py
@@ -32,6 +32,11 @@ def main():
         dest="glm_key",
         help="chatglm api key",
     )
+    parser.add_argument(
+        "--gemini_key",
+        dest="gemini_key",
+        help="gemini api key",
+    )
     parser.add_argument(
         "--bard_token",
         dest="bard_token",
@@ -132,6 +137,13 @@ def main():
         const="bard",
         help="if use bard",
     )
+    bot_group.add_argument(
+        "--use_gemini",
+        dest="bot",
+        action="store_const",
+        const="gemini",
+        help="if use gemini",
+    )
     parser.add_argument(
         "--bing_cookie_path",
         dest="bing_cookie_path",
@@ -141,7 +153,7 @@ def main():
         "--bot",
         dest="bot",
         help="bot type",
-        choices=["gpt3", "chatgptapi", "newbing", "glm", "bard", "langchain"],
+        choices=["gpt3", "chatgptapi", "newbing", "glm", "bard", "gemini", "langchain"],
     )
     parser.add_argument(
         "--config",

diff --git a/xiaogpt/config.py b/xiaogpt/config.py
@@ -62,6 +62,7 @@ class Config:
     password: str = os.getenv("MI_PASS", "")
     openai_key: str = os.getenv("OPENAI_API_KEY", "")
     glm_key: str = os.getenv("CHATGLM_KEY", "")
+    gemini_key: str = os.getenv("GEMINI_KEY", "")  # keep the old rule
     bard_token: str = os.getenv("BARD_TOKEN", "")
     serpapi_api_key: str = os.getenv("SERPAPI_API_KEY", "")
     proxy: str | None = None
@@ -147,6 +148,8 @@ def read_from_file(cls, config_path: str) -> dict:
                     key, value = "bot", "newbing"
                 elif key == "use_glm":
                     key, value = "bot", "glm"
+                elif key == "use_gemini":
+                    key, value = "bot", "gemini"
                 elif key == "use_bard":
                     key, value = "bot", "bard"
                 elif key == "use_langchain":