From 161fc6cdf0fa0774635bd6ed2df8661b6ff3cf97 Mon Sep 17 00:00:00 2001 From: jimmyzhuu Date: Wed, 29 Apr 2026 16:19:27 +0800 Subject: [PATCH] feat: add qianfan chat bot --- models/bot_factory.py | 4 + models/qianfan/__init__.py | 1 + models/qianfan/qianfan_bot.py | 157 ++++++++++++++++++++++++++++++ models/qianfan/qianfan_session.py | 57 +++++++++++ tests/test_qianfan_provider.py | 94 ++++++++++++++++++ 5 files changed, 313 insertions(+) create mode 100644 models/qianfan/__init__.py create mode 100644 models/qianfan/qianfan_bot.py create mode 100644 models/qianfan/qianfan_session.py diff --git a/models/bot_factory.py b/models/bot_factory.py index 632a9052..824aed04 100644 --- a/models/bot_factory.py +++ b/models/bot_factory.py @@ -21,6 +21,10 @@ def create_bot(bot_type): from models.deepseek.deepseek_bot import DeepSeekBot return DeepSeekBot() + elif bot_type == const.QIANFAN: + from models.qianfan.qianfan_bot import QianfanBot + return QianfanBot() + elif bot_type in (const.OPENAI, const.CHATGPT, const.CUSTOM): # OpenAI-compatible API from models.chatgpt.chat_gpt_bot import ChatGPTBot return ChatGPTBot() diff --git a/models/qianfan/__init__.py b/models/qianfan/__init__.py new file mode 100644 index 00000000..14668141 --- /dev/null +++ b/models/qianfan/__init__.py @@ -0,0 +1 @@ +# encoding:utf-8 diff --git a/models/qianfan/qianfan_bot.py b/models/qianfan/qianfan_bot.py new file mode 100644 index 00000000..8e7fe687 --- /dev/null +++ b/models/qianfan/qianfan_bot.py @@ -0,0 +1,157 @@ +# encoding:utf-8 + +import time + +import requests +from bridge.context import ContextType +from bridge.reply import Reply, ReplyType +from common import const +from common.log import logger +from config import conf, load_config +from models.bot import Bot +from models.openai_compatible_bot import OpenAICompatibleBot +from models.session_manager import SessionManager +from .qianfan_session import QianfanSession + +DEFAULT_API_BASE = "https://qianfan.baidubce.com/v2" +DEFAULT_MODEL = const.ERNIE_45_TURBO_128K + + +class QianfanBot(Bot, OpenAICompatibleBot): + def __init__(self): + super().__init__() + model = self._resolve_model() + self.sessions = SessionManager(QianfanSession, model=model) + self.args = { + "model": model, + "temperature": conf().get("temperature", 0.7), + "top_p": conf().get("top_p", 1.0), + "frequency_penalty": conf().get("frequency_penalty", 0.0), + "presence_penalty": conf().get("presence_penalty", 0.0), + } + + def _resolve_model(self): + model = conf().get("model") or DEFAULT_MODEL + if model == const.QIANFAN: + return DEFAULT_MODEL + return model + + @property + def api_key(self): + return conf().get("qianfan_api_key") + + @property + def api_base(self): + url = conf().get("qianfan_api_base") or DEFAULT_API_BASE + url = url.rstrip("/") + suffix = "/chat/completions" + if url.endswith(suffix): + url = url[:-len(suffix)] + return url.rstrip("/") + + def get_api_config(self): + return { + "api_key": self.api_key, + "api_base": self.api_base, + "model": self._resolve_model(), + "default_temperature": conf().get("temperature", 0.7), + "default_top_p": conf().get("top_p", 1.0), + "default_frequency_penalty": conf().get("frequency_penalty", 0.0), + "default_presence_penalty": conf().get("presence_penalty", 0.0), + } + + def _build_headers(self): + return { + "Content-Type": "application/json", + "Authorization": "Bearer {}".format(self.api_key), + } + + def reply(self, query, context=None): + if context.type == ContextType.TEXT: + logger.info("[QIANFAN] 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("[QIANFAN] session query={}".format(session.messages)) + + reply_content = self.reply_text(session, args=self.args.copy()) + logger.debug( + "[QIANFAN] 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("[QIANFAN] reply {} used 0 tokens.".format(reply_content)) + return reply + else: + reply = Reply(ReplyType.ERROR, "Bot不支持处理{}类型的消息".format(context.type)) + return reply + + def reply_text(self, session, args=None, retry_count=0): + try: + body = dict(args) if args else dict(self.args) + body["messages"] = session.messages + response = requests.post( + "{}/chat/completions".format(self.api_base), + headers=self._build_headers(), + json=body, + timeout=conf().get("request_timeout", 180), + ) + if response.status_code == 200: + data = response.json() + return { + "total_tokens": data["usage"]["total_tokens"], + "completion_tokens": data["usage"]["completion_tokens"], + "content": data["choices"][0]["message"]["content"], + } + + error = response.json().get("error", {}) + logger.error( + "[QIANFAN] chat failed, status_code={}, msg={}".format( + response.status_code, error.get("message") + ) + ) + result = {"completion_tokens": 0, "content": "提问太快啦，请休息一下再问我吧"} + need_retry = False + if response.status_code >= 500: + need_retry = retry_count < 2 + elif response.status_code == 401: + result["content"] = "授权失败，请检查 Qianfan API Key 是否正确" + elif response.status_code == 429: + result["content"] = "请求过于频繁，请稍后再试" + need_retry = retry_count < 2 + + if need_retry: + time.sleep(3) + return self.reply_text(session, args, retry_count + 1) + return result + except Exception as e: + logger.exception(e) + if retry_count < 2: + return self.reply_text(session, args, retry_count + 1) + return {"completion_tokens": 0, "content": "我现在有点累了，等会再来吧"} diff --git a/models/qianfan/qianfan_session.py b/models/qianfan/qianfan_session.py new file mode 100644 index 00000000..b5e19153 --- /dev/null +++ b/models/qianfan/qianfan_session.py @@ -0,0 +1,57 @@ +from models.session_manager import Session +from common.log import logger + + +class QianfanSession(Session): + def __init__(self, session_id, system_prompt=None, model="ernie-4.5-turbo-128k"): + 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) + + +def num_tokens_from_messages(messages, model): + tokens = 0 + for msg in messages: + content = msg.get("content", "") + if isinstance(content, str): + tokens += len(content) + elif isinstance(content, list): + for block in content: + if isinstance(block, dict): + tokens += len(block.get("text", "")) + return tokens diff --git a/tests/test_qianfan_provider.py b/tests/test_qianfan_provider.py index fa58dc38..b4e5f522 100644 --- a/tests/test_qianfan_provider.py +++ b/tests/test_qianfan_provider.py @@ -99,5 +99,99 @@ class TestQianfanConstantsAndRouting(unittest.TestCase): ) +class TestQianfanBot(unittest.TestCase): + def _fake_conf(self, values=None): + data = { + "model": "ernie-4.5-turbo-128k", + "qianfan_api_key": "test-qianfan-key", + "qianfan_api_base": "https://qianfan.baidubce.com/v2", + "temperature": 0.7, + "top_p": 1.0, + "frequency_penalty": 0.0, + "presence_penalty": 0.0, + "request_timeout": 180, + "clear_memory_commands": ["#清除记忆"], + "conversation_max_tokens": 1000, + "expires_in_seconds": 3600, + } + if values: + data.update(values) + fake_conf = MagicMock() + fake_conf.get.side_effect = lambda key, default=None: data.get(key, default) + return fake_conf + + def test_bot_factory_returns_qianfan_bot(self): + from common import const + from models.bot_factory import create_bot + + fake_conf = self._fake_conf() + with patch("models.qianfan.qianfan_bot.conf", return_value=fake_conf): + with patch("models.qianfan.qianfan_bot.SessionManager"): + bot = create_bot(const.QIANFAN) + + from models.qianfan.qianfan_bot import QianfanBot + self.assertIsInstance(bot, QianfanBot) + + def test_default_model_uses_ernie_when_model_is_provider_alias(self): + fake_conf = self._fake_conf({"model": "qianfan"}) + with patch("models.qianfan.qianfan_bot.conf", return_value=fake_conf): + with patch("models.qianfan.qianfan_bot.SessionManager"): + from models.qianfan.qianfan_bot import QianfanBot + + bot = QianfanBot() + + self.assertEqual(bot.args["model"], "ernie-4.5-turbo-128k") + + def test_reply_text_posts_openai_compatible_payload(self): + fake_conf = self._fake_conf() + fake_response = MagicMock() + fake_response.status_code = 200 + fake_response.json.return_value = { + "choices": [{"message": {"content": "你好，我是文心。"}}], + "usage": {"total_tokens": 12, "completion_tokens": 6}, + } + session = MagicMock() + session.messages = [{"role": "user", "content": "你好"}] + + with patch("models.qianfan.qianfan_bot.conf", return_value=fake_conf): + with patch("models.qianfan.qianfan_bot.SessionManager"): + from models.qianfan.qianfan_bot import QianfanBot + + bot = QianfanBot() + with patch("models.qianfan.qianfan_bot.requests.post", return_value=fake_response) as post: + result = bot.reply_text(session) + + self.assertEqual(result["content"], "你好，我是文心。") + self.assertEqual(result["total_tokens"], 12) + self.assertEqual(result["completion_tokens"], 6) + post.assert_called_once() + url = post.call_args.args[0] + kwargs = post.call_args.kwargs + self.assertEqual(url, "https://qianfan.baidubce.com/v2/chat/completions") + self.assertEqual(kwargs["headers"]["Authorization"], "Bearer test-qianfan-key") + self.assertEqual(kwargs["json"]["model"], "ernie-4.5-turbo-128k") + self.assertEqual(kwargs["json"]["messages"], [{"role": "user", "content": "你好"}]) + + def test_reply_text_returns_auth_error_for_401(self): + fake_conf = self._fake_conf() + fake_response = MagicMock() + fake_response.status_code = 401 + fake_response.json.return_value = {"error": {"message": "invalid api key"}} + fake_response.text = '{"error":{"message":"invalid api key"}}' + session = MagicMock() + session.messages = [{"role": "user", "content": "你好"}] + + with patch("models.qianfan.qianfan_bot.conf", return_value=fake_conf): + with patch("models.qianfan.qianfan_bot.SessionManager"): + from models.qianfan.qianfan_bot import QianfanBot + + bot = QianfanBot() + with patch("models.qianfan.qianfan_bot.requests.post", return_value=fake_response): + result = bot.reply_text(session) + + self.assertEqual(result["completion_tokens"], 0) + self.assertEqual(result["content"], "授权失败，请检查 Qianfan API Key 是否正确") + + if __name__ == "__main__": unittest.main()