Ollama Support

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

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

bridge/bridge.py

@@ -13,11 +13,15 @@
 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"),
         }
+        # 判断是否使用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"]:

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"

config.py

@@ -6,17 +6,21 @@
 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使用的代理
+    # 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版本

requirements.txt

@@ -8,3 +8,6 @@ Pillow
 pre-commit
 web.py
 linkai>=0.0.5.0
+ollama
+tiktoken
+google-generativeai