update:优化了websocket对数据库连接的获取与释放

This commit is contained in:
killua committed 2024-06-20 10:34:26 +08:00
1 parent ce46c0f35b
commit 5a47440e0f
8 files changed
+91 -54

No files matched your search

+1 -1
View File
@@ -10,7 +10,7 @@ class ASR(ABC):
class LLM(ABC):
@abstractmethod
def chat(self, assistant, prompt, db):
def chat(self, assistant, prompt):
pass
class TTS(ABC):
+24 -26
View File
@@ -56,20 +56,18 @@ class XF_ASR(ASR):
await self.websocket.send(make_continue_frame(audio_data))
elif self.status == LAST_FRAME: #发送最后一帧
await self.websocket.send(make_last_frame(audio_data))
logger.debug("发送完毕")
self.current_message += parse_xfasr_recv(json.loads(await self.websocket.recv()))
if self.current_message == "":
raise AsrResultNoneError()
if "进入沉默模式" in self.current_message:
if self.current_message in ["啊"]:
raise SideNoiseError()
if "闭嘴" in self.current_message:
self.is_slience = True
asyncio.create_task(self.websocket.close())
raise EnterSlienceMode()
if "退出沉默模式" in self.current_message:
self.is_slience = False
self.current_message = "已退出沉默模式"
if self.is_slience:
asyncio.create_task(self.websocket.close())
raise SlienceMode()
asyncio.create_task(self.websocket.close())
logger.debug(f"ASR结果: {self.current_message}")
return [{"text":self.current_message, "audio":self.audio}]
current_time = asyncio.get_event_loop().time()
@@ -86,7 +84,7 @@ class MINIMAX_LLM(LLM):
def __init__(self):
self.token = 0
async def chat(self, assistant, prompt, db):
async def chat(self, assistant, prompt):
llm_info = json.loads(assistant.llm_info)
messages = json.loads(assistant.messages)
messages.append({'role':'user','content':prompt})
@@ -111,12 +109,10 @@ class MINIMAX_LLM(LLM):
yield msg_frame
except LLMResponseEnd:
msg_frame = {"is_end":True,"code":200,"msg":""}
assistant.token = self.token
if self.token > llm_info['max_tokens'] * 0.8: #如果token超过80%,则重置session
msg_frame['code'] = '201'
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
assistant.messages = json.dumps([{'role':'system','content':assistant.system_prompt},{'role':'user','content':prompt}])
yield msg_frame
@@ -147,7 +143,7 @@ class VOLCENGINE_LLM(LLM):
self.token = 0
self.client = Ark(api_key=Config.VOLCENGINE_LLM.API_KEY)
async def chat(self, assistant, prompt, db):
async def chat(self, assistant, prompt):
llm_info = json.loads(assistant.llm_info)
model = self.__get_model(llm_info)
messages = json.loads(assistant.messages)
@@ -167,18 +163,20 @@ class VOLCENGINE_LLM(LLM):
msg_frame = {"is_end":False,"code":200,"msg":chunk_msg}
yield msg_frame
except LLMResponseEnd:
msg_frame = {"is_end":True,"code":20-0,"msg":""}
msg_frame = {"is_end":True,"code":200,"msg":""}
assistant.token = self.token
if self.token > llm_info['max_tokens'] * 0.8: #如果token超过80%,则重置session
msg_frame['code'] = '201'
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
assistant.messages = json.dumps([{'role':'system','content':assistant.system_prompt},{'role':'user','content':prompt}])
yield msg_frame
def __get_model(self, llm_info):
if llm_info['model'] == 'doubao-4k-lite':
return Config.VOLCENGINE_LLM.DOUBAO_LITE_4k
elif llm_info['model'] == 'doubao-32k-lite':
return Config.VOLCENGINE_LLM.DOUBAO_LITE_32k
elif llm_info['model'] == 'doubao-32k-pro':
return Config.VOLCENGINE_LLM.DOUBAO_PRO_32k
else:
raise UnknownVolcEngineModelError()
@@ -322,31 +320,31 @@ class Agent():
self.recorder = Recorder(user_id)
# 对用户输入的音频进行预处理
def user_audio_process(self, audio, db):
def user_audio_process(self, audio):
return self.user_audio_service_chain.user_audio_process(audio, recorder=self.recorder)
# 进行流式语音识别
async def stream_recognize(self, chunk, db):
async def stream_recognize(self, chunk):
return await self.asr.stream_recognize(chunk)
# 进行Prompt加工
def prompt_process(self, asr_results, db):
def prompt_process(self, asr_results):
return self.prompt_service_chain.prompt_process(asr_results)
# 进行大模型调用
async def chat(self, assistant ,prompt, db):
return self.llm.chat(assistant, prompt, db)
async def chat(self, assistant ,prompt):
return self.llm.chat(assistant, prompt)
# 对大模型的返回进行处理
def llm_msg_process(self, llm_chunk, db):
def llm_msg_process(self, llm_chunk):
return self.llm_msg_service_chain.llm_msg_process(llm_chunk)
# 进行TTS合成
def synthetize(self, assistant, text, db):
def synthetize(self, assistant, text):
return self.tts.synthetize(assistant, text)
# 对合成后的音频进行处理
def tts_audio_process(self, audio, db):
def tts_audio_process(self, audio):
return self.tts_audio_service_chain.tts_audio_process(audio, recorder=self.recorder)
# 编码
+5
View File
@@ -46,6 +46,11 @@ class SessionNotFoundError(Exception):
super().__init__(message)
self.message = message
class LlmResultNoneError(Exception):
def __init__(self, message="LLM Result is None!"):
super().__init__(message)
self.message = message
# 大模型返回结束(非异常)
class LLMResponseEnd(Exception):
def __init__(self, message="LLM Response End!"):
+11 -2
View File
@@ -1,8 +1,10 @@
from sqlalchemy import create_engine, Column, Integer, String, CHAR
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager
from config import Config
engine = create_engine(Config.SQLITE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
@@ -30,10 +32,17 @@ class Assistant(Base):
Base.metadata.create_all(bind=engine)
def get_db():
@contextmanager
def get_db_context():
db = SessionLocal()
try:
yield db
finally:
db.close()
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()