From 7d9a0ab3f53c2bdce8f3c0c9b334b4a555937678 Mon Sep 17 00:00:00 2001 From: Xu Jianhui Date: Tue, 30 Apr 2024 11:12:50 +0800 Subject: [PATCH 1/2] Ollama Support --- bot/bot_factory.py | 3 ++ bot/ollama/ollama_bot.py | 72 ++++++++++++++++++++++++++++++++++++++++ bridge/bridge.py | 4 +++ common/const.py | 5 +-- config.py | 5 ++- requirements.txt | 1 + 6 files changed, 87 insertions(+), 3 deletions(-) create mode 100644 bot/ollama/ollama_bot.py diff --git a/bot/bot_factory.py b/bot/bot_factory.py index 3b31af02b..0883e8bfa 100644 --- a/bot/bot_factory.py +++ b/bot/bot_factory.py @@ -65,5 +65,8 @@ def create_bot(bot_type): from bot.moonshot.moonshot_bot import MoonshotBot return MoonshotBot() + elif bot_type == const.OLLAMA_SERVER: + from bot.ollama.ollama_bot import OllamaBot + return OllamaBot() raise RuntimeError diff --git a/bot/ollama/ollama_bot.py b/bot/ollama/ollama_bot.py new file mode 100644 index 000000000..2ba0e01ac --- /dev/null +++ b/bot/ollama/ollama_bot.py @@ -0,0 +1,72 @@ + +from bot.bot import Bot +from bot.session_manager import SessionManager +from bridge.context import ContextType, Context +from bridge.reply import Reply, ReplyType +from common.log import logger +from config import conf +from bot.baidu.baidu_wenxin_session import BaiduWenxinSession +import ollama + + +# Ollama对话模型API (可用) +class OllamaBot(Bot): + + def __init__(self): + super().__init__() + self.model = conf().get("model") or "gemma:7b" + # 复用文心的token计算方式 + self.sessions = SessionManager(BaiduWenxinSession, model=conf().get("model") or "gemma:7b") + + def reply(self, query, context: Context = None) -> Reply: + try: + if context.type != ContextType.TEXT: + logger.warn(f"[Ollama-{self.model}] Unsupported message type, type={context.type}") + return Reply(ReplyType.TEXT, None) + logger.info(f"[Ollama-{self.model}] query={query}") + session_id = context["session_id"] + session = self.sessions.session_query(query, session_id) + # 这里直接调用本地的Ollama服务 + response = ollama.chat( + model=self.model, + messages=self.filter_messages(session.messages)) + reply_text = response['message']['content'] + self.sessions.session_reply(reply_text, session_id) + logger.info(f"[Ollama-{self.model}] reply={reply_text}") + return Reply(ReplyType.TEXT, reply_text) + except Exception as e: + logger.error("f[Ollama-{self.model}] fetch reply error, may contain unsafe content") + logger.error(e) + return Reply(ReplyType.ERROR, f"Ollama failed{e}") + + def _convert_to_gemini_messages(self, messages: list): + res = [] + for msg in messages: + if msg.get("role") == "user": + role = "user" + elif msg.get("role") == "assistant": + role = "model" + else: + continue + res.append({ + "role": role, + "parts": [{"text": msg.get("content")}] + }) + return res + + @staticmethod + def filter_messages(messages: list): + res = [] + turn = "user" + if not messages: + return res + for i in range(len(messages) - 1, -1, -1): + message = messages[i] + if message.get("role") != turn: + continue + res.insert(0, message) + if turn == "user": + turn = "assistant" + elif turn == "assistant": + turn = "user" + return res diff --git a/bridge/bridge.py b/bridge/bridge.py index 52a0df2f9..b92b51ca3 100644 --- a/bridge/bridge.py +++ b/bridge/bridge.py @@ -18,6 +18,10 @@ def __init__(self): "text_to_voice": conf().get("text_to_voice", "google"), "translate": conf().get("translate", "baidu"), } + # 判断是否使用ollama服务 + if conf().get("use_ollama_server", False): + self.btype["chat"] = const.OLLAMA_SERVER + # 这边取配置的模型 model_type = conf().get("model") or const.GPT35 if model_type in ["text-davinci-003"]: diff --git a/common/const.py b/common/const.py index bb3587560..5095fccd5 100644 --- a/common/const.py +++ b/common/const.py @@ -17,7 +17,8 @@ GEMINI = "gemini" ZHIPU_AI = "glm-4" MOONSHOT = "moonshot" - +# Ollama +OLLAMA_SERVER = "ollama" # model CLAUDE3 = "claude-3-opus-20240229" @@ -33,7 +34,7 @@ MODEL_LIST = ["gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-4", "wenxin", "wenxin-4", "xunfei", "claude", "claude-3-opus-20240229", "gpt-4-turbo", "gpt-4-turbo-preview", "gpt-4-1106-preview", GPT4_TURBO_PREVIEW, GPT4_TURBO_01_25, QWEN, GEMINI, ZHIPU_AI, MOONSHOT, - QWEN_TURBO, QWEN_PLUS, QWEN_MAX] + QWEN_TURBO, QWEN_PLUS, QWEN_MAX, OLLAMA_SERVER] # channel FEISHU = "feishu" diff --git a/config.py b/config.py index f6d25da85..b9a29a382 100644 --- a/config.py +++ b/config.py @@ -15,8 +15,11 @@ # openai apibase，当use_azure_chatgpt为true时，需要设置对应的api base "open_ai_api_base": "https://api.openai.com/v1", "proxy": "", # openai使用的代理 + # use_ollama_server， 当 use_ollama_server 为true时，优先使用本地的Ollama服务。 + # 配置的model需要是已经安装的 ollama list 中的模型。例如 gemma:7b , wizardlm2:7b-fp16 + "use_ollama_server": True, # chatgpt模型， 当use_azure_chatgpt为true时，其名称为Azure上model deployment名称 - "model": "gpt-3.5-turbo", # 还支持 gpt-4, gpt-4-turbo, wenxin, xunfei, qwen + "model": "wizardlm2:7b-fp16", # 还支持 gpt-4, gpt-4-turbo, wenxin, xunfei, qwen. wizardlm2:7b-fp16(For Ollama server) "use_azure_chatgpt": False, # 是否使用azure的chatgpt "azure_deployment_id": "", # azure 模型部署名称 "azure_api_version": "", # azure api版本 diff --git a/requirements.txt b/requirements.txt index dc2477b79..c7435fb63 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ Pillow pre-commit web.py linkai>=0.0.5.0 +ollama \ No newline at end of file From eb7beddb05f800ba4c14ff0eeaed347c4f567551 Mon Sep 17 00:00:00 2001 From: Xu Jianhui Date: Tue, 7 May 2024 14:09:04 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E6=94=AF=E6=8C=81Gemini?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bridge/bridge.py | 2 +- config-template.json | 37 ------------------------------------- config.py | 3 ++- requirements.txt | 4 +++- 4 files changed, 6 insertions(+), 40 deletions(-) delete mode 100644 config-template.json diff --git a/bridge/bridge.py b/bridge/bridge.py index b92b51ca3..a30f886b7 100644 --- a/bridge/bridge.py +++ b/bridge/bridge.py @@ -13,7 +13,7 @@ class Bridge(object): def __init__(self): self.btype = { - "chat": const.CHATGPT, + "chat": conf().get("bot_type", const.CHATGPT), "voice_to_text": conf().get("voice_to_text", "openai"), "text_to_voice": conf().get("text_to_voice", "google"), "translate": conf().get("translate", "baidu"), diff --git a/config-template.json b/config-template.json deleted file mode 100644 index f3b253dbe..000000000 --- a/config-template.json +++ /dev/null @@ -1,37 +0,0 @@ -{ - "channel_type": "wx", - "model": "", - "open_ai_api_key": "YOUR API KEY", - "claude_api_key": "YOUR API KEY", - "text_to_image": "dall-e-2", - "voice_to_text": "openai", - "text_to_voice": "openai", - "proxy": "", - "hot_reload": false, - "single_chat_prefix": [ - "bot", - "@bot" - ], - "single_chat_reply_prefix": "[bot] ", - "group_chat_prefix": [ - "@bot" - ], - "group_name_white_list": [ - "ChatGPT测试群", - "ChatGPT测试群2" - ], - "image_create_prefix": [ - "画" - ], - "speech_recognition": true, - "group_speech_recognition": false, - "voice_reply_voice": false, - "conversation_max_tokens": 2500, - "expires_in_seconds": 3600, - "character_desc": "你是ChatGPT, 一个由OpenAI训练的大型语言模型, 你旨在回答并解决人们的任何问题，并且可以使用多种语言与人交流。", - "temperature": 0.7, - "subscribe_msg": "感谢您的关注！\n这里是AI智能助手，可以自由对话。\n支持语音对话。\n支持图片输入。\n支持图片输出，画字开头的消息将按要求创作图片。\n支持tool、角色扮演和文字冒险等丰富的插件。\n输入{trigger_prefix}#help 查看详细指令。", - "use_linkai": false, - "linkai_api_key": "", - "linkai_app_code": "" -} diff --git a/config.py b/config.py index b9a29a382..542262eee 100644 --- a/config.py +++ b/config.py @@ -6,12 +6,13 @@ import pickle from common.log import logger - +# Google 配置 pip install -q -U google-generativeai # 将所有可用的配置项写在字典里, 请使用小写字母 # 此处的配置值无实际意义，程序不会读取此处的配置，仅用于提示格式，请将配置加入到config.json中 available_setting = { # openai api配置 "open_ai_api_key": "", # openai api key + "bot_type": "", # 机器人类型，支持chatgpt, openai, baidu, xunfei, qwen, gemini, zhipu_ai, moonshot, ollama # openai apibase，当use_azure_chatgpt为true时，需要设置对应的api base "open_ai_api_base": "https://api.openai.com/v1", "proxy": "", # openai使用的代理 diff --git a/requirements.txt b/requirements.txt index c7435fb63..9868b6a38 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,6 @@ Pillow pre-commit web.py linkai>=0.0.5.0 -ollama \ No newline at end of file +ollama +tiktoken +google-generativeai \ No newline at end of file