114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
# encoding = utf8
|
|
import json
|
|
import logging
|
|
import time
|
|
import uuid
|
|
|
|
import requests
|
|
# import websocket
|
|
from websocket import create_connection
|
|
|
|
from nlp.nlp_base import NLPBase
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PPTWebsocket:
|
|
def __init__(self, token):
|
|
self.__token = token
|
|
self._response = ''
|
|
self._requesting = False
|
|
headers = {
|
|
"Authorization": f"Bearer {token}",
|
|
}
|
|
header_list = [f"{key}: {value}" for key, value in headers.items()]
|
|
sid = uuid.uuid4()
|
|
url = f'ws://crm.yuekexinxi.com/llm/chat?agent_id=4&app_key=fjV0TjBd1Te9zMWVmp4UuUycHJ3YLEv2&session_id={sid}'
|
|
# websocket.enableTrace(True)
|
|
self._ws = create_connection(
|
|
url,
|
|
header=header_list
|
|
)
|
|
result = self._ws.recv()
|
|
print("接收结果：", result)
|
|
|
|
def __request(self, question):
|
|
print("获取连接状态：", self._ws.getstatus())
|
|
self._requesting = True
|
|
self._ws.send(question)
|
|
|
|
result = self._ws.recv()
|
|
self._requesting = False
|
|
print("接收结果：", result)
|
|
return json.loads(result)
|
|
|
|
def request(self, question, handle, callback):
|
|
t = time.time()
|
|
self._requesting = True
|
|
logger.info(f'-------ppt ask:{question}')
|
|
self._response = self.__request(question)
|
|
content = self._response['content']['text']
|
|
callback(content)
|
|
|
|
self._requesting = False
|
|
logger.info(f'-------dou_bao nlp time:{time.time() - t:.4f}s')
|
|
|
|
def close(self):
|
|
pass
|
|
# if self._response is not None and self._requesting:
|
|
# self._response.close()
|
|
|
|
def aclose(self):
|
|
pass
|
|
# if self._response is not None and self._requesting:
|
|
# self._response.close()
|
|
logger.info('PPTWebsocket close')
|
|
|
|
|
|
class PPT(NLPBase):
|
|
def __init__(self, context, split, callback=None):
|
|
super().__init__(context, split, callback)
|
|
logger.info("PPT init")
|
|
self.__token = '' # 'c9635f9e-0f9e-4ca1-ac90-8af25a541b74'
|
|
self._ppt = None# PPTHttp(self.__token)
|
|
self._request_token()
|
|
|
|
def _request_token(self):
|
|
url = 'http://crm.yuekexinxi.com/llm/platform/token'
|
|
data = {
|
|
"app_key": "fjV0TjBd1Te9zMWVmp4UuUycHJ3YLEv2",
|
|
"secret_key": "0Fhy1a9JFaSAiSuu7DdtrZG1Mcoz1Gpz"
|
|
}
|
|
response = requests.post(url, json=data, stream=True)
|
|
if response.status_code != 200:
|
|
logger.error(f'request token failed: {response.text}')
|
|
return
|
|
|
|
try:
|
|
content = json.loads(response.text)
|
|
logger.info(content)
|
|
code = content.get('code')
|
|
if 200 != code:
|
|
logger.warning(f"content code {code} != 200")
|
|
return
|
|
self.__token = content.get('data')
|
|
self._ppt = PPTWebsocket(self.__token)
|
|
|
|
except Exception as e:
|
|
logger.error(f'response error: {e}')
|
|
|
|
def _request(self, question):
|
|
self._ppt.request(question, self._split_handle, self._on_callback)
|
|
|
|
def _on_close(self):
|
|
if self._ppt is not None:
|
|
self._ppt.close()
|
|
logger.info('AsyncArk close')
|
|
|
|
def on_clear_cache(self, *args, **kwargs):
|
|
super().on_clear_cache(*args, **kwargs)
|
|
if self._ppt is not None:
|
|
self._ppt.aclose()
|
|
logger.info('DouBao clear_cache')
|
|
|