forked from killua/TakwayDisplayPlatform
285 lines
10 KiB
Python
285 lines
10 KiB
Python
from utils.xf_asr_utils import xf_asr_websocket_factory, make_first_frame, make_continue_frame, make_last_frame, parse_xfasr_recv
|
||
from .model import Assistant
|
||
from .abstract import *
|
||
from .public import *
|
||
from utils.vits_utils import TextToSpeech
|
||
from config import Config
|
||
import aiohttp
|
||
import asyncio
|
||
import struct
|
||
import base64
|
||
import json
|
||
|
||
# ----------- 初始化vits ----------- #
|
||
vits = TextToSpeech()
|
||
# ---------------------------------- #
|
||
|
||
|
||
|
||
#------ 具体 ASR, LLM, TTS 类 ------ #
|
||
FIRST_FRAME = 1
|
||
CONTINUE_FRAME =2
|
||
LAST_FRAME =3
|
||
|
||
class XF_ASR(ASR):
|
||
def __init__(self):
|
||
self.websocket = None
|
||
self.current_message = ""
|
||
self.audio = ""
|
||
self.status = FIRST_FRAME
|
||
self.segment_duration_threshold = 25 #超时时间为25秒
|
||
self.segment_start_time = None
|
||
|
||
async def stream_recognize(self, chunk):
|
||
if self.websocket is None: #如果websocket未建立，则建立一个新的连接
|
||
self.websocket = await xf_asr_websocket_factory()
|
||
if self.segment_start_time is None: #如果是第一段，则记录开始时间
|
||
self.segment_start_time = asyncio.get_event_loop().time()
|
||
if chunk['meta_info']['is_end']:
|
||
self.status = LAST_FRAME
|
||
audio_data = chunk['audio']
|
||
self.audio += audio_data
|
||
|
||
if self.status == FIRST_FRAME: #发送第一帧
|
||
await self.websocket.send(make_first_frame(audio_data))
|
||
self.status = CONTINUE_FRAME
|
||
elif self.status == CONTINUE_FRAME: #发送中间帧
|
||
await self.websocket.send(make_continue_frame(audio_data))
|
||
elif self.status == LAST_FRAME: #发送最后一帧
|
||
await self.websocket.send(make_last_frame(audio_data))
|
||
self.current_message += parse_xfasr_recv(json.loads(await self.websocket.recv()))
|
||
if self.current_message == "":
|
||
raise AsrResultNoneError()
|
||
await self.websocket.close()
|
||
print("语音识别结束，用户消息：", self.current_message)
|
||
return [{"text":self.current_message, "audio":self.audio}]
|
||
|
||
current_time = asyncio.get_event_loop().time()
|
||
if current_time - self.segment_start_time > self.segment_duration_threshold: #超时，发送最后一帧并重置状态
|
||
await self.websocket.send(make_last_frame())
|
||
self.current_message += parse_xfasr_recv(await self.websocket.recv())
|
||
await self.websocket.close()
|
||
self.websocket = await xf_asr_websocket_factory()
|
||
self.status = FIRST_FRAME
|
||
self.segment_start_time = current_time
|
||
return []
|
||
|
||
class MINIMAX_LLM(LLM):
|
||
def __init__(self):
|
||
self.token = 0
|
||
|
||
async def chat(self, assistant, prompt, db):
|
||
llm_info = json.loads(assistant.llm_info)
|
||
messages = json.loads(assistant.messages)
|
||
messages.append({"role":"user","content":prompt})
|
||
assistant.messages = json.dumps(messages)
|
||
payload = json.dumps({
|
||
"model": llm_info['model'],
|
||
"stream": True,
|
||
"messages": messages,
|
||
"max_tokens": llm_info['max_tokens'],
|
||
"temperature": llm_info['temperature'],
|
||
"top_p": llm_info['top_p'],
|
||
})
|
||
headers = {
|
||
'Authorization': f"Bearer {Config.MINIMAX_LLM.API_KEY}",
|
||
'Content-Type': 'application/json'
|
||
}
|
||
async with aiohttp.ClientSession() as client:
|
||
async with client.post(Config.MINIMAX_LLM.URL, headers=headers, data=payload) as response:
|
||
async for chunk in response.content.iter_any():
|
||
try:
|
||
chunk_msg = self.__parseChunk(chunk)
|
||
msg_frame = {"is_end":False,"code":200,"msg":chunk_msg}
|
||
yield msg_frame
|
||
except LLMResponseEnd:
|
||
msg_frame = {"is_end":True,"msg":""}
|
||
if self.token > llm_info['max_tokens'] * 0.8:
|
||
msg_frame['msg'] = 'max_token reached'
|
||
as_query = db.query(Assistant).filter(Assistant.id == assistant.id).first()
|
||
as_query.messages = json.dumps([{'role':'system','content':assistant.system_prompt},{'role':'user','content':prompt}])
|
||
db.commit()
|
||
assistant.messages = as_query.messages
|
||
yield msg_frame
|
||
|
||
|
||
def __parseChunk(self, llm_chunk):
|
||
result = ""
|
||
data=json.loads(llm_chunk.decode('utf-8')[6:])
|
||
if data["object"] == "chat.completion": #如果是结束帧
|
||
self.token = data['usage']['total_tokens']
|
||
raise LLMResponseEnd()
|
||
elif data['object'] == 'chat.completion.chunk':
|
||
for choice in data['choices']:
|
||
result += choice['delta']['content']
|
||
else:
|
||
raise UnkownLLMFrame()
|
||
return result
|
||
|
||
class VITS_TTS(TTS):
|
||
def __init__(self):
|
||
pass
|
||
|
||
def synthetize(self, assistant, text):
|
||
tts_info = json.loads(assistant.tts_info)
|
||
return vits.synthesize(text, tts_info)
|
||
# --------------------------------- #
|
||
|
||
|
||
|
||
# ------ ASR, LLM, TTS 工厂类 ------ #
|
||
class ASRFactory:
|
||
def create_asr(self,asr_type:str) -> ASR:
|
||
if asr_type == 'XF':
|
||
return XF_ASR()
|
||
|
||
class LLMFactory:
|
||
def create_llm(self,llm_type:str) -> LLM:
|
||
if llm_type == 'MINIMAX':
|
||
return MINIMAX_LLM()
|
||
|
||
class TTSFactory:
|
||
def create_tts(self,tts_type:str) -> TTS:
|
||
if tts_type == 'VITS':
|
||
return VITS_TTS()
|
||
# --------------------------------- #
|
||
|
||
|
||
|
||
# ----------- 具体服务类 ----------- #
|
||
# 从单说话人的asr_results中取出第一个结果作为prompt
|
||
class BasicPromptService(PromptService):
|
||
def prompt_process(self, prompt:str, **kwargs):
|
||
if 'asr_result' in kwargs:
|
||
return kwargs['asr_result'][0]['text'], kwargs
|
||
else:
|
||
raise NoAsrResultsError()
|
||
|
||
class UserAudioRecordService(UserAudioService):
|
||
def user_audio_process(self, audio:str, **kwargs):
|
||
audio_decoded = base64.b64decode(audio)
|
||
kwargs['recorder'].user_audio += audio_decoded
|
||
return audio,kwargs
|
||
|
||
class TTSAudioRecordService(TTSAudioService):
|
||
def tts_audio_process(self, audio:bytes, **kwargs):
|
||
kwargs['recorder'].tts_audio += audio
|
||
return audio,kwargs
|
||
# --------------------------------- #
|
||
|
||
|
||
|
||
# ------------ 服务链类 ----------- #
|
||
class UserAudioServiceChain:
|
||
def __init__(self):
|
||
self.services = []
|
||
def add_service(self, service:UserAudioService):
|
||
self.services.append(service)
|
||
def user_audio_process(self, audio, **kwargs):
|
||
for service in self.services:
|
||
audio, kwargs = service.user_audio_process(audio, **kwargs)
|
||
return audio
|
||
|
||
class PromptServiceChain:
|
||
def __init__(self):
|
||
self.services = []
|
||
def add_service(self, service:PromptService):
|
||
self.services.append(service)
|
||
def prompt_process(self, asr_results):
|
||
kwargs = {"asr_result":asr_results}
|
||
prompt = ""
|
||
for service in self.services:
|
||
prompt, kwargs = service.prompt_process(prompt, **kwargs)
|
||
return prompt
|
||
|
||
class LLMMsgServiceChain:
|
||
def __init__(self, ):
|
||
self.services = []
|
||
self.spliter = SentenceSegmentation()
|
||
def add_service(self, service:LLMMsgService):
|
||
self.services.append(service)
|
||
def llm_msg_process(self, llm_msg):
|
||
kwargs = {}
|
||
llm_chunks = self.spliter.split(llm_msg)
|
||
for service in self.services:
|
||
llm_chunks , kwargs = service.llm_msg_process(llm_chunks, **kwargs)
|
||
return llm_chunks
|
||
|
||
class TTSAudioServiceChain:
|
||
def __init__(self):
|
||
self.services = []
|
||
def add_service(self, service:TTSAudioService):
|
||
self.services.append(service)
|
||
def tts_audio_process(self, audio, **kwargs):
|
||
for service in self.services:
|
||
audio, kwargs = service.tts_audio_process(audio, **kwargs)
|
||
return audio
|
||
# --------------------------------- #
|
||
|
||
class Agent():
|
||
def __init__(self, asr_type:str, llm_type:str, tts_type:str):
|
||
asrFactory = ASRFactory()
|
||
llmFactory = LLMFactory()
|
||
ttsFactory = TTSFactory()
|
||
self.recorder = None
|
||
|
||
self.user_audio_service_chain = UserAudioServiceChain() #创建用户音频处理服务链
|
||
self.user_audio_service_chain.add_service(UserAudioRecordService()) #添加用户音频记录服务
|
||
|
||
self.asr = asrFactory.create_asr(asr_type)
|
||
|
||
self.prompt_service_chain = PromptServiceChain()
|
||
self.prompt_service_chain.add_service(BasicPromptService())
|
||
|
||
self.llm = llmFactory.create_llm(llm_type)
|
||
|
||
self.llm_chunk_service_chain = LLMMsgServiceChain()
|
||
|
||
self.tts = ttsFactory.create_tts(tts_type)
|
||
|
||
self.tts_audio_service_chain = TTSAudioServiceChain()
|
||
self.tts_audio_service_chain.add_service(TTSAudioRecordService())
|
||
|
||
def init_recorder(self,user_id):
|
||
self.recorder = Recorder(user_id)
|
||
|
||
# 对用户输入的音频进行预处理
|
||
def user_audio_process(self, audio, db):
|
||
return self.user_audio_service_chain.user_audio_process(audio, recorder=self.recorder)
|
||
|
||
# 进行流式语音识别
|
||
async def stream_recognize(self, chunk, db):
|
||
return await self.asr.stream_recognize(chunk)
|
||
|
||
# 进行Prompt加工
|
||
def prompt_process(self, asr_results, db):
|
||
return self.prompt_service_chain.prompt_process(asr_results)
|
||
|
||
# 进行大模型调用
|
||
async def chat(self, assistant ,prompt, db):
|
||
return self.llm.chat(assistant, prompt, db)
|
||
|
||
# 对大模型的返回进行处理
|
||
def llm_msg_process(self, llm_chunk, db):
|
||
return self.llm_chunk_service_chain.llm_msg_process(llm_chunk)
|
||
|
||
# 进行TTS合成
|
||
def synthetize(self, assistant, text, db):
|
||
return self.tts.synthetize(assistant, text)
|
||
|
||
# 对合成后的音频进行处理
|
||
def tts_audio_process(self, audio, db):
|
||
return self.tts_audio_service_chain.tts_audio_process(audio, recorder=self.recorder)
|
||
|
||
# 编码
|
||
def encode(self, text, audio):
|
||
text_resp = {"type":"text","code":200,"msg":text}
|
||
text_resp_bytes = json.dumps(text_resp,ensure_ascii=False).encode('utf-8')
|
||
header = struct.pack('!II',len(text_resp_bytes),len(audio))
|
||
final_resp = header + text_resp_bytes + audio
|
||
return final_resp
|
||
|
||
# 保存音频
|
||
def save(self):
|
||
self.recorder.write()
|
||
|