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