update: 封装了豆包模型

app/concrete.py

@@ -1,4 +1,5 @@
 from utils.xf_asr_utils import xf_asr_websocket_factory, make_first_frame, make_continue_frame, make_last_frame, parse_xfasr_recv
+from volcenginesdkarkruntime import Ark
 from .model import Assistant
 from .abstract import *
 from .public import *
@@ -98,9 +99,9 @@ class MINIMAX_LLM(LLM):
                         msg_frame = {"is_end":False,"code":200,"msg":chunk_msg}
                         yield msg_frame
                     except LLMResponseEnd:
-                        msg_frame = {"is_end":True,"msg":""}
+                        msg_frame = {"is_end":True,"code":200,"msg":""}
                         if self.token > llm_info['max_tokens'] * 0.8:  #如果token超过80%，则重置session
-                            msg_frame['msg'] = 'max_token reached'
+                            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()
@@ -130,6 +131,51 @@ class MINIMAX_LLM(LLM):
             logger.error(llm_chunk)
             raise AbnormalLLMFrame(f"error llm_chunk:{llm_chunk}")
 
+class VOLCENGINE_LLM(LLM):
+    def __init__(self):
+        self.token = 0
+        self.client = Ark(api_key=Config.VOLCENGINE_LLM.API_KEY)
+    
+    async def chat(self, assistant, prompt, db):
+        llm_info = json.loads(assistant.llm_info)
+        model = self.__get_model(llm_info)
+        messages = json.loads(assistant.messages)
+        messages.append({'role':'user','content':prompt})
+        stream = self.client.chat.completions.create(
+            model = model,
+            messages=messages,
+            stream=True,
+            stream_options={'include_usage': True}
+        )
+        for chunk in stream:
+            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,"code":20-0,"msg":""}
+                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
+                yield msg_frame
+
+    def __get_model(self, llm_info):
+        if llm_info['model'] == 'doubao-4k-lite':
+            return Config.VOLCENGINE_LLM.DOUBAO_LITE_4k
+        else:
+            raise UnknownVolcEngineModelError()
+    
+    def __parseChunk(self, llm_chunk):
+        if llm_chunk.usage:
+            self.token = llm_chunk.usage.total_tokens
+            raise LLMResponseEnd()
+        if not llm_chunk.choices:
+            raise AbnormalLLMFrame(f"error volcengine llm_chunk:{llm_chunk}")
+        return llm_chunk.choices[0].delta.content
+
 class VITS_TTS(TTS):
     def __init__(self):
         pass
@@ -151,6 +197,8 @@ class LLMFactory:
     def create_llm(self,llm_type:str) -> LLM:
         if llm_type == 'MINIMAX':
             return MINIMAX_LLM()
+        if llm_type == 'VOLCENGINE':
+            return VOLCENGINE_LLM()
 
 class TTSFactory:
     def create_tts(self,tts_type:str) -> TTS:

app/exception.py

@@ -10,6 +10,12 @@ class NoAsrResultsError(Exception):
         super().__init__(message)
         self.message = message
 
+# 未知的火山引擎模型
+class UnknownVolcEngineModelError(Exception):
+    def __init__(self, message="Unknown Volc Engine Model!"):
+        super().__init__(message)
+        self.message = message
+
 # 未知LLM返回帧
 class UnkownLLMFrame(Exception):
     def __init__(self, message="Unkown LLM Frame!"):