xunfei_spark_bot.py 8.8KB

hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
hace 1 año
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268
  1. # encoding:utf-8
  2. import requests, json
  3. from bot.bot import Bot
  4. from bot.session_manager import SessionManager
  5. from bot.baidu.baidu_wenxin_session import BaiduWenxinSession
  6. from bridge.context import ContextType, Context
  7. from bridge.reply import Reply, ReplyType
  8. from common.log import logger
  9. from config import conf
  10. from common import const
  11. import time
  12. import _thread as thread
  13. import datetime
  14. from datetime import datetime
  15. from wsgiref.handlers import format_date_time
  16. from urllib.parse import urlencode
  17. import base64
  18. import ssl
  19. import hashlib
  20. import hmac
  21. import json
  22. from time import mktime
  23. from urllib.parse import urlparse
  24. import websocket
  25. import queue
  26. import threading
  27. import random
  28. # 消息队列 map
  29. queue_map = dict()
  30. # 响应队列 map
  31. reply_map = dict()
  32. class XunFeiBot(Bot):
  33. def __init__(self):
  34. super().__init__()
  35. self.app_id = conf().get("xunfei_app_id")
  36. self.api_key = conf().get("xunfei_api_key")
  37. self.api_secret = conf().get("xunfei_api_secret")
  38. # 默认使用v2.0版本: "generalv2"
  39. # v1.5版本为 "general"
  40. # v3.0版本为: "generalv3"
  41. self.domain = "generalv3"
  42. # 默认使用v2.0版本: "ws://spark-api.xf-yun.com/v2.1/chat"
  43. # v1.5版本为: "ws://spark-api.xf-yun.com/v1.1/chat"
  44. # v3.0版本为: "ws://spark-api.xf-yun.com/v3.1/chat"
  45. # v3.5版本为: "wss://spark-api.xf-yun.com/v3.5/chat"
  46. self.spark_url = "wss://spark-api.xf-yun.com/v3.5/chat"
  47. self.host = urlparse(self.spark_url).netloc
  48. self.path = urlparse(self.spark_url).path
  49. # 和wenxin使用相同的session机制
  50. self.sessions = SessionManager(BaiduWenxinSession, model=const.XUNFEI)
  51. def reply(self, query, context: Context = None) -> Reply:
  52. if context.type == ContextType.TEXT:
  53. logger.info("[XunFei] query={}".format(query))
  54. session_id = context["session_id"]
  55. request_id = self.gen_request_id(session_id)
  56. reply_map[request_id] = ""
  57. session = self.sessions.session_query(query, session_id)
  58. threading.Thread(target=self.create_web_socket,
  59. args=(session.messages, request_id)).start()
  60. depth = 0
  61. time.sleep(0.1)
  62. t1 = time.time()
  63. usage = {}
  64. while depth <= 300:
  65. try:
  66. data_queue = queue_map.get(request_id)
  67. if not data_queue:
  68. depth += 1
  69. time.sleep(0.1)
  70. continue
  71. data_item = data_queue.get(block=True, timeout=0.1)
  72. if data_item.is_end:
  73. # 请求结束
  74. del queue_map[request_id]
  75. if data_item.reply:
  76. reply_map[request_id] += data_item.reply
  77. usage = data_item.usage
  78. break
  79. reply_map[request_id] += data_item.reply
  80. depth += 1
  81. except Exception as e:
  82. depth += 1
  83. continue
  84. t2 = time.time()
  85. logger.info(
  86. f"[XunFei-API] response={reply_map[request_id]}, time={t2 - t1}s, usage={usage}"
  87. )
  88. self.sessions.session_reply(reply_map[request_id], session_id,
  89. usage.get("total_tokens"))
  90. reply = Reply(ReplyType.TEXT, reply_map[request_id])
  91. del reply_map[request_id]
  92. return reply
  93. else:
  94. reply = Reply(ReplyType.ERROR,
  95. "Bot不支持处理{}类型的消息".format(context.type))
  96. return reply
  97. def create_web_socket(self, prompt, session_id, temperature=0.5):
  98. logger.info(f"[XunFei] start connect, prompt={prompt}")
  99. websocket.enableTrace(False)
  100. wsUrl = self.create_url()
  101. ws = websocket.WebSocketApp(wsUrl,
  102. on_message=on_message,
  103. on_error=on_error,
  104. on_close=on_close,
  105. on_open=on_open)
  106. data_queue = queue.Queue(1000)
  107. queue_map[session_id] = data_queue
  108. ws.appid = self.app_id
  109. ws.question = prompt
  110. ws.domain = self.domain
  111. ws.session_id = session_id
  112. ws.temperature = temperature
  113. ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})
  114. def gen_request_id(self, session_id: str):
  115. return session_id + "_" + str(int(time.time())) + "" + str(
  116. random.randint(0, 100))
  117. # 生成url
  118. def create_url(self):
  119. # 生成RFC1123格式的时间戳
  120. now = datetime.now()
  121. date = format_date_time(mktime(now.timetuple()))
  122. # 拼接字符串
  123. signature_origin = "host: " + self.host + "\n"
  124. signature_origin += "date: " + date + "\n"
  125. signature_origin += "GET " + self.path + " HTTP/1.1"
  126. # 进行hmac-sha256进行加密
  127. signature_sha = hmac.new(self.api_secret.encode('utf-8'),
  128. signature_origin.encode('utf-8'),
  129. digestmod=hashlib.sha256).digest()
  130. signature_sha_base64 = base64.b64encode(signature_sha).decode(
  131. encoding='utf-8')
  132. authorization_origin = f'api_key="{self.api_key}", algorithm="hmac-sha256", headers="host date request-line", ' \
  133. f'signature="{signature_sha_base64}"'
  134. authorization = base64.b64encode(
  135. authorization_origin.encode('utf-8')).decode(encoding='utf-8')
  136. # 将请求的鉴权参数组合为字典
  137. v = {"authorization": authorization, "date": date, "host": self.host}
  138. # 拼接鉴权参数,生成url
  139. url = self.spark_url + '?' + urlencode(v)
  140. # 此处打印出建立连接时候的url,参考本demo的时候可取消上方打印的注释,比对相同参数时生成的url与自己代码生成的url是否一致
  141. return url
  142. def gen_params(self, appid, domain, question):
  143. """
  144. 通过appid和用户的提问来生成请参数
  145. """
  146. data = {
  147. "header": {
  148. "app_id": appid,
  149. "uid": "1234"
  150. },
  151. "parameter": {
  152. "chat": {
  153. "domain": domain,
  154. "random_threshold": 0.5,
  155. "max_tokens": 2048,
  156. "auditing": "default"
  157. }
  158. },
  159. "payload": {
  160. "message": {
  161. "text": question
  162. }
  163. }
  164. }
  165. return data
  166. class ReplyItem:
  167. def __init__(self, reply, usage=None, is_end=False):
  168. self.is_end = is_end
  169. self.reply = reply
  170. self.usage = usage
  171. # 收到websocket错误的处理
  172. def on_error(ws, error):
  173. logger.error(f"[XunFei] error: {str(error)}")
  174. # 收到websocket关闭的处理
  175. def on_close(ws, one, two):
  176. data_queue = queue_map.get(ws.session_id)
  177. data_queue.put("END")
  178. # 收到websocket连接建立的处理
  179. def on_open(ws):
  180. logger.info(f"[XunFei] Start websocket, session_id={ws.session_id}")
  181. thread.start_new_thread(run, (ws, ))
  182. def run(ws, *args):
  183. data = json.dumps(
  184. gen_params(appid=ws.appid,
  185. domain=ws.domain,
  186. question=ws.question,
  187. temperature=ws.temperature))
  188. ws.send(data)
  189. # Websocket 操作
  190. # 收到websocket消息的处理
  191. def on_message(ws, message):
  192. data = json.loads(message)
  193. code = data['header']['code']
  194. if code != 0:
  195. logger.error(f'请求错误: {code}, {data}')
  196. ws.close()
  197. else:
  198. choices = data["payload"]["choices"]
  199. status = choices["status"]
  200. content = choices["text"][0]["content"]
  201. data_queue = queue_map.get(ws.session_id)
  202. if not data_queue:
  203. logger.error(
  204. f"[XunFei] can't find data queue, session_id={ws.session_id}")
  205. return
  206. reply_item = ReplyItem(content)
  207. if status == 2:
  208. usage = data["payload"].get("usage")
  209. reply_item = ReplyItem(content, usage)
  210. reply_item.is_end = True
  211. ws.close()
  212. data_queue.put(reply_item)
  213. def gen_params(appid, domain, question, temperature=0.5):
  214. """
  215. 通过appid和用户的提问来生成请参数
  216. """
  217. data = {
  218. "header": {
  219. "app_id": appid,
  220. "uid": "1234"
  221. },
  222. "parameter": {
  223. "chat": {
  224. "domain": domain,
  225. "temperature": temperature,
  226. "random_threshold": 0.5,
  227. "max_tokens": 2048,
  228. "auditing": "default"
  229. }
  230. },
  231. "payload": {
  232. "message": {
  233. "text": question
  234. }
  235. }
  236. }
  237. return data