添加Ollama支持

bot/bot_factory.py

@@ -67,6 +67,10 @@ def create_bot(bot_type):
     elif bot_type == const.MiniMax:
         from bot.minimax.minimax_bot import MinimaxBot
         return MinimaxBot()
+
+    elif bot_type == const.OLLAMA:
+        from bot.ollama.ollama_bot import OllamaBot
+        return OllamaBot()
 
 
     raise RuntimeError

bot/ollama/ollama_bot.py

@@ -0,0 +1,125 @@
+# encoding:utf-8
+
+from bot.bot import Bot
+from bot.ollama.ollama_session import OllamaSession
+from bot.session_manager import SessionManager
+from bridge.context import ContextType
+from bridge.reply import Reply, ReplyType
+from common.log import logger
+from common.token_bucket import TokenBucket
+from config import conf, load_config
+from bot.ollama.ollama_session import num_tokens_from_messages
+
+import ollama
+
+class OllamaBot(Bot):
+    def __init__(self):
+        super().__init__()
+        # set the default api_key
+        if conf().get("rate_limit_chatgpt"):
+            self.tb4chatgpt = TokenBucket(conf().get("rate_limit_chatgpt", 20))
+        conf_model = conf().get("model") or "gpt-3.5-turbo"
+        self.sessions = SessionManager(OllamaSession, model=conf_model)
+
+        self.client = ollama.Client(
+            host=conf().get("ollama_host", "http://localhost:11434")   
+        )
+
+        self.args = {
+            "model": conf_model,  # 对话模型的名称
+            "temperature": conf().get("temperature", 0.9),  # 值在[0,1]之间，越大表示回复越具有不确定性
+            # "max_tokens":4096,  # 回复最大的字符数
+            "top_p": conf().get("top_p", 1),
+            "frequency_penalty": conf().get("frequency_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容
+            "presence_penalty": conf().get("presence_penalty", 0.0),  # [-2,2]之间，该值越大则更倾向于产生不同的内容
+            "request_timeout": conf().get("request_timeout", None),  # 请求超时时间，openai接口默认设置为600，对于难问题一般需要较长时间
+            "timeout": conf().get("request_timeout", None),  # 重试超时时间，在这个时间内，将会自动重试
+        }
+
+
+    def reply(self, query, context=None):
+        # acquire reply content
+        if context.type == ContextType.TEXT:
+            logger.info("[CHATGPT] query={}".format(query))
+
+            session_id = context["session_id"]
+            reply = None
+            clear_memory_commands = conf().get("clear_memory_commands", ["#清除记忆"])
+            if query in clear_memory_commands:
+                self.sessions.clear_session(session_id)
+                reply = Reply(ReplyType.INFO, "记忆已清除")
+            elif query == "#清除所有":
+                self.sessions.clear_all_session()
+                reply = Reply(ReplyType.INFO, "所有人记忆已清除")
+            elif query == "#更新配置":
+                load_config()
+                reply = Reply(ReplyType.INFO, "配置已更新")
+            if reply:
+                return reply
+            session = self.sessions.session_query(query, session_id)
+            logger.debug("[OLLAMA] session query={}".format(session.messages))
+
+            model = context.get("gpt_model")
+            new_args = None
+            if model:
+                new_args = self.args.copy()
+                new_args["model"] = model
+
+            reply_content = self.reply_text(session, args=new_args)
+            logger.debug(
+                "[OLLAMA] new_query={}, session_id={}, reply_cont={}, completion_tokens={}".format(
+                    session.messages,
+                    session_id,
+                    reply_content["content"],
+                    reply_content["completion_tokens"],
+                )
+            )
+            if reply_content["completion_tokens"] == 0 and len(reply_content["content"]) > 0:
+                reply = Reply(ReplyType.ERROR, reply_content["content"])
+            elif reply_content["completion_tokens"] > 0:
+                self.sessions.session_reply(reply_content["content"], session_id, reply_content["total_tokens"])
+                reply = Reply(ReplyType.TEXT, reply_content["content"])
+            else:
+                reply = Reply(ReplyType.ERROR, reply_content["content"])
+                logger.debug("[OLLAMA] reply {} used 0 tokens.".format(reply_content))
+            return reply
+
+        elif context.type == ContextType.IMAGE_CREATE:
+            ok, retstring = self.create_img(query, 0)
+            reply = None
+            if ok:
+                reply = Reply(ReplyType.IMAGE_URL, retstring)
+            else:
+                reply = Reply(ReplyType.ERROR, retstring)
+            return reply
+        else:
+            reply = Reply(ReplyType.ERROR, "Bot不支持处理{}类型的消息".format(context.type))
+            return reply
+
+    def reply_text(self, session: OllamaSession, args=None) -> dict:
+        """
+        call Ollama's ChatCompletion to get the answer
+        :param session: a conversation session
+        :param session_id: session id
+        :param retry_count: retry count
+        :return: {}
+        """
+        try:
+            # if api_key == None, the default openai.api_key will be used
+            if args is None:
+                args = self.args
+            response = self.client.chat(messages=session.messages, model=args["model"], options=args)
+            # 计算 total_tokens
+            total_tokens = num_tokens_from_messages(session.messages, args["model"])
+            # 计算 completion_tokens
+            completion_tokens = num_tokens_from_messages([{"role": "assistant", "content": response["message"]["content"]}])
+            return {
+                "total_tokens": total_tokens + completion_tokens,
+                "completion_tokens": completion_tokens,
+                "content": response["message"]["content"],
+            }
+        except Exception as e:
+            result = {"completion_tokens": 0, "content": "我现在有点累了，等会再来吧"}
+            logger.exception("[CHATGPT] Exception: {}".format(e))
+            self.sessions.clear_session(session.session_id)
+            return result

bot/ollama/ollama_session.py

@@ -0,0 +1,87 @@
+from bot.session_manager import Session
+from common.log import logger
+
+"""
+    e.g.  [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Who won the world series in 2020?"},
+        {"role": "assistant", "content": "The Los Angeles Dodgers won the World Series in 2020."},
+        {"role": "user", "content": "Where was it played?"}
+    ]
+"""
+
+
+class OllamaSession(Session):
+    def __init__(self, session_id, system_prompt=None, model="gpt-3.5-turbo"):
+        super().__init__(session_id, system_prompt)
+        self.model = model
+        self.reset()
+
+    def discard_exceeding(self, max_tokens, cur_tokens=None):
+        precise = True
+        try:
+            cur_tokens = self.calc_tokens()
+        except Exception as e:
+            precise = False
+            if cur_tokens is None:
+                raise e
+            logger.debug("Exception when counting tokens precisely for query: {}".format(e))
+        while cur_tokens > max_tokens:
+            if len(self.messages) > 2:
+                self.messages.pop(1)
+            elif len(self.messages) == 2 and self.messages[1]["role"] == "assistant":
+                self.messages.pop(1)
+                if precise:
+                    cur_tokens = self.calc_tokens()
+                else:
+                    cur_tokens = cur_tokens - max_tokens
+                break
+            elif len(self.messages) == 2 and self.messages[1]["role"] == "user":
+                logger.warn("user message exceed max_tokens. total_tokens={}".format(cur_tokens))
+                break
+            else:
+                logger.debug("max_tokens={}, total_tokens={}, len(messages)={}".format(max_tokens, cur_tokens, len(self.messages)))
+                break
+            if precise:
+                cur_tokens = self.calc_tokens()
+            else:
+                cur_tokens = cur_tokens - max_tokens
+        return cur_tokens
+
+    def calc_tokens(self):
+        return num_tokens_from_messages(self.messages, self.model)
+
+
+# refer to https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+def num_tokens_from_messages(messages, model="gpt-3.5-turbo"):
+    """Returns the number of tokens used by a list of messages."""
+    import tiktoken
+
+    model = "gpt-3.5-turbo"
+
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+    except KeyError:
+        logger.debug("Warning: model not found. Using cl100k_base encoding.")
+        encoding = tiktoken.get_encoding("cl100k_base")
+
+    tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
+    tokens_per_name = -1  # if there's a name, the role is omitted
+
+    num_tokens = 0
+    for message in messages:
+        num_tokens += tokens_per_message
+        for key, value in message.items():
+            num_tokens += len(encoding.encode(value))
+            if key == "name":
+                num_tokens += tokens_per_name
+    num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
+    return num_tokens
+
+
+def num_tokens_by_character(messages):
+    """Returns the number of tokens used by a list of messages."""
+    tokens = 0
+    for msg in messages:
+        tokens += len(msg["content"])
+    return tokens

common/const.py

@@ -9,6 +9,7 @@
 CLAUDEAPI= "claudeAPI"  # 通过Claude api调用模型
 QWEN = "qwen"  # 旧版通义模型
 QWEN_DASHSCOPE = "dashscope"  # 通义新版sdk和api key
+OLLAMA = "ollama"
 
 
 GEMINI = "gemini"  # gemini-1.0-pro

config-template.json

@@ -1,5 +1,7 @@
 {
   "channel_type": "wx",
+  "bot_type": "ollama",
+  "ollama_host": "http://localhost:11434",
   "model": "",
   "open_ai_api_key": "YOUR API KEY",
   "claude_api_key": "YOUR API KEY",

config.py

@@ -179,6 +179,8 @@
     "Minimax_api_key": "",
     "Minimax_group_id": "",
     "Minimax_base_url": "",
+
+    "ollama_host": "http://localhost:11434",  # ollama host
 }
 
 

plugins/godcmd/godcmd.py

@@ -313,7 +313,7 @@ def on_handle_context(self, e_context: EventContext):
                     except Exception as e:
                         ok, result = False, "你没有设置私有GPT模型"
                 elif cmd == "reset":
-                    if bottype in [const.OPEN_AI, const.CHATGPT, const.CHATGPTONAZURE, const.LINKAI, const.BAIDU, const.XUNFEI, const.QWEN, const.GEMINI, const.ZHIPU_AI, const.CLAUDEAPI]:
+                    if bottype in [const.OPEN_AI, const.CHATGPT, const.CHATGPTONAZURE, const.LINKAI, const.BAIDU, const.XUNFEI, const.QWEN, const.GEMINI, const.ZHIPU_AI, const.CLAUDEAPI, const.OLLAMA]:
                         bot.sessions.clear_session(session_id)
                         if Bridge().chat_bots.get(bottype):
                             Bridge().chat_bots.get(bottype).sessions.clear_session(session_id)
@@ -339,7 +339,7 @@ def on_handle_context(self, e_context: EventContext):
                             ok, result = True, "配置已重载"
                         elif cmd == "resetall":
                             if bottype in [const.OPEN_AI, const.CHATGPT, const.CHATGPTONAZURE, const.LINKAI,
-                                           const.BAIDU, const.XUNFEI, const.QWEN, const.GEMINI, const.ZHIPU_AI, const.MOONSHOT]:
+                                           const.BAIDU, const.XUNFEI, const.QWEN, const.GEMINI, const.ZHIPU_AI, const.MOONSHOT, const.OLLAMA]:
                                 channel.cancel_all_session()
                                 bot.sessions.clear_all_session()
                                 ok, result = True, "重置所有会话成功"