From 9b125913ae62991b6d20f00025484a63d5b372db Mon Sep 17 00:00:00 2001
From: Alex Yang <463909821@qq.com>
Date: Fri, 16 Aug 2024 14:58:17 +0800
Subject: [PATCH] wenxin character model supports prompt

---
 bot/baidu/baidu_wenxin.py | 7 ++++++-
 config.py                 | 1 +
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/bot/baidu/baidu_wenxin.py b/bot/baidu/baidu_wenxin.py
index ff78e75..dd660af 100644
--- a/bot/baidu/baidu_wenxin.py
+++ b/bot/baidu/baidu_wenxin.py
@@ -19,6 +19,11 @@ class BaiduWenxinBot(Bot):
     def __init__(self):
         super().__init__()
         wenxin_model = conf().get("baidu_wenxin_model")
+        self.prompt_enabled = conf().get("baidu_wenxin_prompt_enabled")
+        if self.prompt_enabled:
+            self.prompt = conf().get("character_desc", "")
+            if self.prompt == "":
+                logger.warn("[BAIDU] Although you enabled model prompt, character_desc is not specified.")
         if wenxin_model is not None:
             wenxin_model = conf().get("baidu_wenxin_model") or "eb-instant"
         else:
@@ -84,7 +89,7 @@ class BaiduWenxinBot(Bot):
             headers = {
                 'Content-Type': 'application/json'
             }
-            payload = {'messages': session.messages}
+            payload = {'messages': session.messages, 'system': self.prompt} if self.prompt_enabled else {'messages': session.messages}
             response = requests.request("POST", url, headers=headers, data=json.dumps(payload))
             response_text = json.loads(response.text)
             logger.info(f"[BAIDU] response text={response_text}")
diff --git a/config.py b/config.py
index b6ae49e..330afca 100644
--- a/config.py
+++ b/config.py
@@ -69,6 +69,7 @@ available_setting = {
     "baidu_wenxin_model": "eb-instant",  # 默认使用ERNIE-Bot-turbo模型
     "baidu_wenxin_api_key": "",  # Baidu api key
     "baidu_wenxin_secret_key": "",  # Baidu secret key
+    "baidu_wenxin_prompt_enabled": False,  # Enable prompt if you are using ernie character model
     # 讯飞星火API
     "xunfei_app_id": "",  # 讯飞应用ID
     "xunfei_api_key": "",  # 讯飞 API key