Files
bot-on-anything/model/baidu/yiyan_model.py
2023-03-18 12:25:32 +08:00

108 lines
3.3 KiB
Python

# encoding:utf-8
from model.model import Model
from config import model_conf
from common import const
from common import log
import requests
import time
import json
sessions = {}
class YiyanModel(Model):
def __init__(self):
self.acs_token = model_conf(const.BAIDU).get('acs_token')
self.cookie = model_conf(const.BAIDU).get('cookie')
self.base_url = 'https://yiyan.baidu.com/eb'
def reply(self, query, context=None):
user_id = context.get('session_id') or context.get('from_user_id')
# 1.create session
chat_session_id = sessions.get(user_id)
if not chat_session_id:
self.new_session(context)
sessions[user_id] = context['chat_session_id']
else:
context['chat_session_id'] = chat_session_id
# 2.create chat
context['query'] = query
self.new_chat(context)
# 3.query
context['reply'] = ''
self.query(context, 0, 0)
return context['reply']
def new_session(self, context):
data = {
"sessionName": "test session",
"timestamp": int(time.time() * 1000),
"deviceType": "pc"
}
res = requests.post(url=self.base_url+'/session/new', headers=self._create_header(), json=data)
# print(res.headers)
context['chat_session_id'] = res.json()['data']['sessionId']
def new_chat(self, context):
headers = self._create_header()
headers['Acs-Token'] = self.acs_token
data = {
"sessionId": context.get('chat_session_id'),
"text": context['query'],
"parentChatId": 0,
"type": 10,
"timestamp": int(time.time() * 1000),
"deviceType": "pc",
"code": 0,
"msg": ""
}
res = requests.post(url=self.base_url+'/chat/new', headers=headers, json=data)
print(res.text)
context['chat_id'] = res.json()['data']['botChat']['id']
context['parent_chat_id'] = res.json()['data']['botChat']['parent']
def query(self, context, sentence_id, count):
headers = self._create_header()
headers['Acs-Token'] = self.acs_token
data = {
"chatId": context['chat_id'],
"parentChatId": context['parent_chat_id'],
"sentenceId": sentence_id,
"stop": 0,
"timestamp": 1679068791405,
"deviceType": "pc"
}
res = requests.post(url=self.base_url + '/chat/query', headers=headers, json=data).json()
if res['data']['text'] != '':
context['reply'] += res['data']['text']
if res['data']['is_end'] == 1:
return
if count > 10:
return
if not res['data']['text']:
time.sleep(1)
return self.query(context, sentence_id, count+1)
def _create_header(self):
headers = {
'Host': 'yiyan.baidu.com',
'Origin': 'https://yiyan.baidu.com',
'Referer': 'https://yiyan.baidu.com',
'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_13_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/105.0.0.0 Safari/537.36',
'Content-Type': 'application/json',
'Cookie': self.cookie
}
return headers