Merge branch 'main' into py_linkerai_single

This commit is contained in:
hrz
2025-06-05 09:26:00 +08:00
72 changed files with 3629 additions and 714 deletions
@@ -2,9 +2,6 @@ import http.client
import json
import asyncio
from typing import Optional, Tuple, List
import opuslib_next
import wave
import io
import os
import uuid
import hmac
@@ -16,6 +13,7 @@ import time
from datetime import datetime
from config.logger import setup_logging
from core.providers.asr.base import ASRProviderBase
from core.providers.asr.dto.dto import InterfaceType
TAG = __name__
logger = setup_logging()
@@ -92,6 +90,7 @@ class AccessToken:
class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool):
super().__init__()
self.interface_type = InterfaceType.NON_STREAM
"""阿里云ASR初始化"""
# 新增空值判断逻辑
self.access_key_id = config.get("access_key_id")
@@ -165,21 +164,6 @@ class ASRProvider(ASRProviderBase):
request += "&enable_voice_detection=false"
return request
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1) # 单声道
wf.setsampwidth(2) # 16-bit
wf.setframerate(self.sample_rate)
wf.writeframes(b"".join(pcm_data))
logger.bind(tag=TAG).debug(f"音频文件已保存至: {file_path}")
return file_path
async def _send_request(self, pcm_data: bytes) -> Optional[str]:
"""发送请求到阿里云ASR服务"""
try:
@@ -229,7 +213,7 @@ class ASRProvider(ASRProviderBase):
return None
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""将语音数据转换为文本"""
if self._is_token_expired():
@@ -239,7 +223,7 @@ class ASRProvider(ASRProviderBase):
file_path = None
try:
# 解码Opus为PCM
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -1,18 +1,10 @@
import base64
import hashlib
import hmac
import json
import time
from datetime import datetime, timezone
import os
import uuid
from typing import Optional, Tuple, List
import wave
import opuslib_next
from aip import AipSpeech
from core.providers.asr.base import ASRProviderBase
from config.logger import setup_logging
from core.providers.asr.dto.dto import InterfaceType
TAG = __name__
logger = setup_logging()
@@ -21,6 +13,7 @@ logger = setup_logging()
class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool = True):
super().__init__()
self.interface_type = InterfaceType.NON_STREAM
self.app_id = config.get("app_id")
self.api_key = config.get("api_key")
self.secret_key = config.get("secret_key")
@@ -36,22 +29,8 @@ class ASRProvider(ASRProviderBase):
# 确保输出目录存在
os.makedirs(self.output_dir, exist_ok=True)
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""将语音数据转换为文本"""
if not opus_data:
@@ -66,7 +45,7 @@ class ASRProvider(ASRProviderBase):
return None, file_path
# 将Opus音频数据解码为PCM
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
+95 -10
View File
@@ -1,7 +1,19 @@
from abc import ABC, abstractmethod
from typing import Optional, Tuple, List
import os
import wave
import copy
import uuid
import queue
import asyncio
import traceback
import threading
import opuslib_next
from abc import ABC, abstractmethod
from config.logger import setup_logging
from typing import Optional, Tuple, List
from core.handle.receiveAudioHandle import startToChat
from core.handle.reportHandle import enqueue_asr_report
from core.utils.util import remove_punctuation_and_length
from core.handle.receiveAudioHandle import handleAudioMessage
TAG = __name__
logger = setup_logging()
@@ -9,24 +21,97 @@ logger = setup_logging()
class ASRProviderBase(ABC):
def __init__(self):
self.audio_format = "opus"
pass
# 打开音频通道
# 这里默认是非流式的处理方式
# 流式处理方式请在子类中重写
async def open_audio_channels(self, conn):
# tts 消化线程
conn.asr_priority_thread = threading.Thread(
target=self.asr_text_priority_thread, args=(conn,), daemon=True
)
conn.asr_priority_thread.start()
# 有序处理ASR音频
def asr_text_priority_thread(self, conn):
while not conn.stop_event.is_set():
try:
message = conn.asr_audio_queue.get(timeout=1)
future = asyncio.run_coroutine_threadsafe(
handleAudioMessage(conn, message),
conn.loop,
)
future.result()
except queue.Empty:
continue
except Exception as e:
logger.bind(tag=TAG).error(
f"处理ASR文本失败: {str(e)}, 类型: {type(e).__name__}, 堆栈: {traceback.format_exc()}"
)
continue
# 接收音频
# 这里默认是非流式的处理方式
# 流式处理方式请在子类中重写
async def receive_audio(self, conn, audio, audio_have_voice):
if conn.client_listen_mode == "auto" or conn.client_listen_mode == "realtime":
have_voice = audio_have_voice
else:
have_voice = conn.client_have_voice
# 如果本次没有声音,本段也没声音,就把声音丢弃了
conn.asr_audio.append(audio)
if have_voice == False and conn.client_have_voice == False:
conn.asr_audio = conn.asr_audio[-10:]
return
# 如果本段有声音,且已经停止了
if conn.client_voice_stop:
asr_audio_task = copy.deepcopy(conn.asr_audio)
conn.asr_audio.clear()
# 音频太短了,无法识别
conn.reset_vad_states()
if len(asr_audio_task) > 15:
await self.handle_voice_stop(conn, asr_audio_task)
# 处理语音停止
async def handle_voice_stop(self, conn, asr_audio_task):
raw_text, _ = await self.speech_to_text(
asr_audio_task, conn.session_id, conn.audio_format
) # 确保ASR模块返回原始文本
conn.logger.bind(tag=TAG).info(f"识别文本: {raw_text}")
text_len, _ = remove_punctuation_and_length(raw_text)
self.stop_ws_connection()
if text_len > 0:
# 使用自定义模块进行上报
await startToChat(conn, raw_text)
enqueue_asr_report(conn, raw_text, asr_audio_task)
def stop_ws_connection(self):
pass
@abstractmethod
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
pass
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
@abstractmethod
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""将语音数据转换为文本"""
pass
def set_audio_format(self, format: str) -> None:
"""设置音频格式"""
self.audio_format = format
@staticmethod
def decode_opus(opus_data: List[bytes]) -> bytes:
"""将Opus音频数据解码为PCM数据"""
@@ -1,17 +1,14 @@
import time
import io
import wave
import os
from typing import Optional, Tuple, List
import uuid
import websockets
import json
import gzip
import opuslib_next
from core.providers.asr.base import ASRProviderBase
import websockets
from config.logger import setup_logging
from typing import Optional, Tuple, List
from core.providers.asr.base import ASRProviderBase
from core.providers.asr.dto.dto import InterfaceType
TAG = __name__
logger = setup_logging()
@@ -86,6 +83,7 @@ def parse_response(res):
class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool):
super().__init__()
self.interface_type = InterfaceType.NON_STREAM
self.appid = config.get("appid")
self.cluster = config.get("cluster")
self.access_token = config.get("access_token")
@@ -102,20 +100,6 @@ class ASRProvider(ASRProviderBase):
# 确保输出目录存在
os.makedirs(self.output_dir, exist_ok=True)
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
@staticmethod
def _generate_header(
message_type=CLIENT_FULL_REQUEST, message_type_specific_flags=NO_SEQUENCE
@@ -244,14 +228,14 @@ class ASRProvider(ASRProviderBase):
yield data[offset:data_len], True
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""将语音数据转换为文本"""
file_path = None
try:
# 合并所有opus数据包
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -0,0 +1,344 @@
import json
import gzip
import uuid
import asyncio
import websockets
import opuslib_next
from core.providers.asr.base import ASRProviderBase
from config.logger import setup_logging
from core.providers.asr.dto.dto import InterfaceType
TAG = __name__
logger = setup_logging()
class ASRProvider(ASRProviderBase):
def __init__(self, config, delete_audio_file):
super().__init__()
self.interface_type = InterfaceType.STREAM
self.config = config
self.text = ""
self.max_retries = 3
self.retry_delay = 2
self.decoder = opuslib_next.Decoder(16000, 1)
self.asr_ws = None
self.forward_task = None
self.is_processing = False # 添加处理状态标志
# 配置参数
self.appid = str(config.get("appid"))
self.cluster = config.get("cluster")
self.access_token = config.get("access_token")
self.boosting_table_name = config.get("boosting_table_name", "")
self.correct_table_name = config.get("correct_table_name", "")
self.output_dir = config.get("output_dir", "tmp/")
self.delete_audio_file = delete_audio_file
# 火山引擎ASR配置
self.ws_url = "wss://openspeech.bytedance.com/api/v3/sauc/bigmodel"
self.uid = config.get("uid", "streaming_asr_service")
self.workflow = config.get(
"workflow", "audio_in,resample,partition,vad,fe,decode,itn,nlu_punctuate"
)
self.result_type = config.get("result_type", "single")
self.format = config.get("format", "pcm")
self.codec = config.get("codec", "pcm")
self.rate = config.get("sample_rate", 16000)
self.language = config.get("language", "zh-CN")
self.bits = config.get("bits", 16)
self.channel = config.get("channel", 1)
self.auth_method = config.get("auth_method", "token")
self.secret = config.get("secret", "access_secret")
async def open_audio_channels(self, conn):
await super().open_audio_channels(conn)
async def receive_audio(self, conn, audio, audio_have_voice):
conn.asr_audio.append(audio)
conn.asr_audio = conn.asr_audio[-10:]
# 如果本次有声音,且之前没有建立连接
if audio_have_voice and self.asr_ws is None and not self.is_processing:
try:
self.is_processing = True
# 建立新的WebSocket连接
headers = self.token_auth() if self.auth_method == "token" else None
logger.bind(tag=TAG).info(f"正在连接ASR服务,headers: {headers}")
self.asr_ws = await websockets.connect(
self.ws_url,
additional_headers=headers,
max_size=1000000000,
ping_interval=None,
ping_timeout=None,
close_timeout=10,
)
# 发送初始化请求
request_params = self.construct_request(str(uuid.uuid4()))
try:
payload_bytes = str.encode(json.dumps(request_params))
payload_bytes = gzip.compress(payload_bytes)
full_client_request = self.generate_header()
full_client_request.extend((len(payload_bytes)).to_bytes(4, "big"))
full_client_request.extend(payload_bytes)
logger.bind(tag=TAG).info(f"发送初始化请求: {request_params}")
await self.asr_ws.send(full_client_request)
# 等待初始化响应
init_res = await self.asr_ws.recv()
result = self.parse_response(init_res)
logger.bind(tag=TAG).info(f"收到初始化响应: {result}")
# 检查初始化响应
if "code" in result and result["code"] != 1000:
error_msg = f"ASR服务初始化失败: {result.get('payload_msg', {}).get('message', '未知错误')}"
if "payload_msg" in result:
error_msg += f"\n详细错误信息: {json.dumps(result['payload_msg'], ensure_ascii=False)}"
logger.bind(tag=TAG).error(error_msg)
raise Exception(error_msg)
except Exception as e:
logger.bind(tag=TAG).error(f"发送初始化请求失败: {str(e)}")
if hasattr(e, "__cause__") and e.__cause__:
logger.bind(tag=TAG).error(f"错误原因: {str(e.__cause__)}")
raise e
# 启动接收ASR结果的异步任务
self.forward_task = asyncio.create_task(self._forward_asr_results(conn))
# 发送缓存的音频数据
if conn.asr_audio and len(conn.asr_audio) > 0:
for cached_audio in conn.asr_audio[-10:]:
try:
pcm_frame = self.decoder.decode(cached_audio, 960)
payload = gzip.compress(pcm_frame)
audio_request = bytearray(
self.generate_audio_default_header()
)
audio_request.extend(len(payload).to_bytes(4, "big"))
audio_request.extend(payload)
await self.asr_ws.send(audio_request)
except Exception as e:
logger.bind(tag=TAG).info(
f"发送缓存音频数据时发生错误: {e}"
)
except Exception as e:
logger.bind(tag=TAG).error(f"建立ASR连接失败: {str(e)}")
if hasattr(e, "__cause__") and e.__cause__:
logger.bind(tag=TAG).error(f"错误原因: {str(e.__cause__)}")
if self.asr_ws:
await self.asr_ws.close()
self.asr_ws = None
self.is_processing = False
return
# 发送当前音频数据
if self.asr_ws and self.is_processing:
try:
pcm_frame = self.decoder.decode(audio, 960)
payload = gzip.compress(pcm_frame)
audio_request = bytearray(self.generate_audio_default_header())
audio_request.extend(len(payload).to_bytes(4, "big"))
audio_request.extend(payload)
await self.asr_ws.send(audio_request)
except Exception as e:
logger.bind(tag=TAG).info(f"发送音频数据时发生错误: {e}")
async def _forward_asr_results(self, conn):
try:
while self.asr_ws and not conn.stop_event.is_set():
try:
response = await self.asr_ws.recv()
result = self.parse_response(response)
logger.bind(tag=TAG).debug(f"收到ASR结果: {result}")
if "payload_msg" in result:
payload = result["payload_msg"]
if "result" in payload:
utterances = payload["result"].get("utterances", [])
# 检查duration和空文本的情况
if (
payload.get("audio_info", {}).get("duration", 0) > 2000
and not utterances
and not payload["result"].get("text")
):
logger.bind(tag=TAG).error(f"识别文本:空")
self.text = ""
conn.reset_vad_states()
await self.handle_voice_stop(conn, None)
break
for utterance in utterances:
if utterance.get("definite", False):
self.text = utterance["text"]
logger.bind(tag=TAG).info(
f"识别到文本: {self.text}"
)
conn.reset_vad_states()
await self.handle_voice_stop(conn, None)
break
elif "error" in payload:
error_msg = payload.get("error", "未知错误")
logger.bind(tag=TAG).error(f"ASR服务返回错误: {error_msg}")
break
except websockets.ConnectionClosed:
logger.bind(tag=TAG).info("ASR服务连接已关闭")
self.is_processing = False
break
except Exception as e:
logger.bind(tag=TAG).error(f"处理ASR结果时发生错误: {str(e)}")
if hasattr(e, "__cause__") and e.__cause__:
logger.bind(tag=TAG).error(f"错误原因: {str(e.__cause__)}")
self.is_processing = False
break
except Exception as e:
logger.bind(tag=TAG).error(f"ASR结果转发任务发生错误: {str(e)}")
if hasattr(e, "__cause__") and e.__cause__:
logger.bind(tag=TAG).error(f"错误原因: {str(e.__cause__)}")
finally:
if self.asr_ws:
await self.asr_ws.close()
self.asr_ws = None
self.is_processing = False
def stop_ws_connection(self):
if self.asr_ws:
asyncio.create_task(self.asr_ws.close())
self.asr_ws = None
self.is_processing = False
def construct_request(self, reqid):
req = {
"app": {
"appid": self.appid,
"cluster": self.cluster,
"token": self.access_token,
},
"user": {"uid": self.uid},
"request": {
"reqid": reqid,
"workflow": self.workflow,
"show_utterances": True,
"result_type": self.result_type,
"sequence": 1,
"boosting_table_name": self.boosting_table_name,
"correct_table_name": self.correct_table_name,
"end_window_size": 200,
},
"audio": {
"format": self.format,
"codec": self.codec,
"rate": self.rate,
"language": self.language,
"bits": self.bits,
"channel": self.channel,
"sample_rate": self.rate,
},
}
logger.bind(tag=TAG).debug(
f"构造请求参数: {json.dumps(req, ensure_ascii=False)}"
)
return req
def token_auth(self):
return {
"X-Api-App-Key": self.appid,
"X-Api-Access-Key": self.access_token,
"X-Api-Resource-Id": "volc.bigasr.sauc.duration",
"X-Api-Connect-Id": str(uuid.uuid4()),
"Host": "openspeech.bytedance.com",
}
def generate_header(
self,
version=0x01,
message_type=0x01,
message_type_specific_flags=0x00,
serial_method=0x01,
compression_type=0x01,
reserved_data=0x00,
extension_header: bytes = b"",
):
header = bytearray()
header_size = int(len(extension_header) / 4) + 1
header.append((version << 4) | header_size)
header.append((message_type << 4) | message_type_specific_flags)
header.append((serial_method << 4) | compression_type)
header.append(reserved_data)
header.extend(extension_header)
return header
def generate_audio_default_header(self):
return self.generate_header(
version=0x01,
message_type=0x02,
message_type_specific_flags=0x00,
serial_method=0x01,
compression_type=0x01,
)
def generate_last_audio_default_header(self):
return self.generate_header(
version=0x01,
message_type=0x02,
message_type_specific_flags=0x02,
serial_method=0x01,
compression_type=0x01,
)
def parse_response(self, res: bytes) -> dict:
try:
# 检查响应长度
if len(res) < 4:
logger.bind(tag=TAG).error(f"响应数据长度不足: {len(res)}")
return {"error": "响应数据长度不足"}
# 获取消息头
header = res[:4]
message_type = header[1] >> 4
# 如果是错误响应
if message_type == 0x0F: # SERVER_ERROR_RESPONSE
code = int.from_bytes(header[4:8], "big", signed=False)
error_msg = res[8:].decode("utf-8")
return {"code": code, "error": error_msg}
# 获取JSON数据(跳过12字节头部)
try:
json_data = res[12:].decode("utf-8")
result = json.loads(json_data)
logger.bind(tag=TAG).debug(f"成功解析JSON响应: {result}")
return {"payload_msg": result}
except (UnicodeDecodeError, json.JSONDecodeError) as e:
logger.bind(tag=TAG).error(f"JSON解析失败: {str(e)}")
logger.bind(tag=TAG).error(f"原始数据: {res}")
raise
except Exception as e:
logger.bind(tag=TAG).error(f"解析响应失败: {str(e)}")
logger.bind(tag=TAG).error(f"原始响应数据: {res.hex()}")
raise
async def speech_to_text(self, opus_data, session_id, audio_format):
result = self.text
self.text = "" # 清空text
return result, None
async def close(self):
"""资源清理方法"""
if self.asr_ws:
await self.asr_ws.close()
self.asr_ws = None
if self.forward_task:
self.forward_task.cancel()
try:
await self.forward_task
except asyncio.CancelledError:
pass
self.forward_task = None
self.is_processing = False
@@ -0,0 +1,9 @@
from enum import Enum
from typing import Union, Optional
class InterfaceType(Enum):
# 接口类型
STREAM = "STREAM" # 流式接口
NON_STREAM = "NON_STREAM" # 非流式接口
LOCAL = "LOCAL" # 本地服务
@@ -1,15 +1,14 @@
import time
import wave
import os
import sys
import io
from config.logger import setup_logging
from typing import Optional, Tuple, List
import uuid
from core.providers.asr.base import ASRProviderBase
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import shutil
from core.providers.asr.dto.dto import InterfaceType
TAG = __name__
logger = setup_logging()
@@ -38,6 +37,7 @@ class CaptureOutput:
class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool):
super().__init__()
self.interface_type = InterfaceType.LOCAL
self.model_dir = config.get("model_dir")
self.output_dir = config.get("output_dir") # 修正配置键名
self.delete_audio_file = delete_audio_file
@@ -53,22 +53,8 @@ class ASRProvider(ASRProviderBase):
# device="cuda:0", # 启用GPU加速
)
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""语音转文本主处理逻辑"""
file_path = None
@@ -77,7 +63,7 @@ class ASRProvider(ASRProviderBase):
while retry_count < MAX_RETRIES:
try:
# 合并所有opus数据包
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -1,11 +1,8 @@
from typing import Optional, Tuple, List
import opuslib_next
from core.providers.asr.base import ASRProviderBase
import os
from core.providers.asr.dto.dto import InterfaceType
import ssl
import json
import uuid
import wave
import websockets
from config.logger import setup_logging
import asyncio
@@ -23,6 +20,7 @@ class ASRProvider(ASRProviderBase):
:param delete_audio_file: Boolean to indicate whether to delete audio files after processing.
"""
super().__init__()
self.interface_type = InterfaceType.NON_STREAM
self.host = config.get("host", "localhost")
self.port = config.get("port", 10095)
self.api_key = config.get("api_key", "none")
@@ -43,20 +41,6 @@ class ASRProvider(ASRProviderBase):
self.ssl_context.check_hostname = False
self.ssl_context.verify_mode = ssl.CERT_NONE
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
async def _receive_responses(self, ws) -> None:
"""
Asynchronous generator to receive messages from the WebSocket.
@@ -116,7 +100,7 @@ class ASRProvider(ASRProviderBase):
logger.bind(tag=TAG).debug(f"Sent end message: {end_message}")
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""
Convert speech data to text using FunASR.
@@ -125,7 +109,7 @@ class ASRProvider(ASRProviderBase):
:return: Tuple containing recognized text and optional timestamp.
"""
file_path = None
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -5,8 +5,7 @@ import sys
import io
from config.logger import setup_logging
from typing import Optional, Tuple, List
import uuid
import opuslib_next
from core.providers.asr.dto.dto import InterfaceType
from core.providers.asr.base import ASRProviderBase
import numpy as np
@@ -38,6 +37,7 @@ class CaptureOutput:
class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool):
super().__init__()
self.interface_type = InterfaceType.LOCAL
self.model_dir = config.get("model_dir")
self.output_dir = config.get("output_dir")
self.delete_audio_file = delete_audio_file
@@ -84,20 +84,6 @@ class ASRProvider(ASRProviderBase):
use_itn=True,
)
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
def read_wave(self, wave_filename: str) -> Tuple[np.ndarray, int]:
"""
Args:
@@ -123,14 +109,14 @@ class ASRProvider(ASRProviderBase):
return samples_float32, f.getframerate()
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""语音转文本主处理逻辑"""
file_path = None
try:
# 保存音频文件
start_time = time.time()
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -5,11 +5,8 @@ import json
import time
from datetime import datetime, timezone
import os
import uuid
from typing import Optional, Tuple, List
import wave
import opuslib_next
from core.providers.asr.dto.dto import InterfaceType
import requests
from core.providers.asr.base import ASRProviderBase
from config.logger import setup_logging
@@ -25,6 +22,7 @@ class ASRProvider(ASRProviderBase):
def __init__(self, config: dict, delete_audio_file: bool = True):
super().__init__()
self.interface_type = InterfaceType.NON_STREAM
self.secret_id = config.get("secret_id")
self.secret_key = config.get("secret_key")
self.output_dir = config.get("output_dir")
@@ -33,22 +31,8 @@ class ASRProvider(ASRProviderBase):
# 确保输出目录存在
os.makedirs(self.output_dir, exist_ok=True)
def save_audio_to_file(self, pcm_data: List[bytes], session_id: str) -> str:
"""PCM数据保存为WAV文件"""
module_name = __name__.split(".")[-1]
file_name = f"asr_{module_name}_{session_id}_{uuid.uuid4()}.wav"
file_path = os.path.join(self.output_dir, file_name)
with wave.open(file_path, "wb") as wf:
wf.setnchannels(1)
wf.setsampwidth(2) # 2 bytes = 16-bit
wf.setframerate(16000)
wf.writeframes(b"".join(pcm_data))
return file_path
async def speech_to_text(
self, opus_data: List[bytes], session_id: str
self, opus_data: List[bytes], session_id: str, audio_format="opus"
) -> Tuple[Optional[str], Optional[str]]:
"""将语音数据转换为文本"""
if not opus_data:
@@ -63,7 +47,7 @@ class ASRProvider(ASRProviderBase):
return None, file_path
# 将Opus音频数据解码为PCM
if self.audio_format == "pcm":
if audio_format == "pcm":
pcm_data = opus_data
else:
pcm_data = self.decode_opus(opus_data)
@@ -122,6 +122,8 @@ class IntentProvider(IntentProviderBase):
async def detect_intent(self, conn, dialogue_history: List[Dict], text: str) -> str:
if not self.llm:
raise ValueError("LLM provider not set")
if conn.func_handler is None:
return '{"function_call": {"name": "continue_chat"}}'
# 记录整体开始时间
total_start_time = time.time()
@@ -148,9 +150,8 @@ class IntentProvider(IntentProviderBase):
self.clean_cache()
if self.promot == "":
if hasattr(conn, "func_handler"):
functions = conn.func_handler.get_functions()
self.promot = self.get_intent_system_prompt(functions)
functions = conn.func_handler.get_functions()
self.promot = self.get_intent_system_prompt(functions)
music_config = initialize_music_handler(conn)
music_file_names = music_config["music_file_names"]
@@ -52,6 +52,7 @@ class TTSProviderBase(ABC):
)
self.first_sentence_punctuations = (
"",
"",
"~",
"",
",",
@@ -178,6 +179,9 @@ class TTSProviderBase(ABC):
while not self.conn.stop_event.is_set():
try:
message = self.tts_text_queue.get(timeout=1)
if self.conn.client_abort:
logger.bind(tag=TAG).info("收到打断信息,终止TTS文本处理线程")
continue
if message.sentence_type == SentenceType.FIRST:
# 初始化参数
self.tts_stop_request = False
@@ -216,6 +220,7 @@ class TTSProviderBase(ABC):
logger.bind(tag=TAG).error(
f"处理TTS文本失败: {str(e)}, 类型: {type(e).__name__}, 堆栈: {traceback.format_exc()}"
)
continue
def _audio_play_priority_thread(self):
while not self.conn.stop_event.is_set():
@@ -3,13 +3,13 @@ import uuid
import json
import queue
import asyncio
import threading
import traceback
import websockets
from config.logger import setup_logging
from core.utils import opus_encoder_utils
from core.utils.util import check_model_key
from core.providers.tts.base import TTSProviderBase
from core.handle.abortHandle import handleAbortMessage
from core.providers.tts.dto.dto import SentenceType, ContentType, InterfaceType
TAG = __name__
@@ -138,6 +138,7 @@ class Response:
class TTSProvider(TTSProviderBase):
def __init__(self, config, delete_audio_file):
super().__init__(config, delete_audio_file)
self.ws = None
self.interface_type = InterfaceType.DUAL_STREAM
self.appId = config.get("appid")
self.access_token = config.get("access_token")
@@ -152,73 +153,106 @@ class TTSProvider(TTSProviderBase):
self.authorization = config.get("authorization")
self.header = {"Authorization": f"{self.authorization}{self.access_token}"}
self.enable_two_way = True
self.start_connection_flag = False
self.tts_text = ""
# 合成文字语音后,播放的音频文件列表
self.before_stop_play_files = []
self.opus_encoder = opus_encoder_utils.OpusEncoderUtils(
sample_rate=16000, channels=1, frame_size_ms=60
)
check_model_key("TTS", self.access_token)
# 添加会话状态控制
self._session_lock = asyncio.Lock() # 会话操作的并发锁
self._current_session_id = None # 当前会话ID
self._session_started = False # 会话是否已开始
self._session_finished = False # 会话是否已结束
###################################################################################
# 火山双流式TTS重写父类的方法--开始
###################################################################################
async def open_audio_channels(self, conn):
await super().open_audio_channels(conn)
ws_header = {
"X-Api-App-Key": self.appId,
"X-Api-Access-Key": self.access_token,
"X-Api-Resource-Id": self.resource_id,
"X-Api-Connect-Id": uuid.uuid4(),
}
self.ws = await websockets.connect(
self.ws_url, additional_headers=ws_header, max_size=1000000000
)
tts_priority = threading.Thread(
target=self._start_monitor_tts_response_thread, daemon=True
)
tts_priority.start()
try:
await super().open_audio_channels(conn)
except Exception as e:
logger.bind(tag=TAG).error(f"Failed to open audio channels: {str(e)}")
self.ws = None
raise
async def _ensure_connection(self):
"""建立新的WebSocket连接"""
try:
logger.bind(tag=TAG).info("开始建立新连接...")
ws_header = {
"X-Api-App-Key": self.appId,
"X-Api-Access-Key": self.access_token,
"X-Api-Resource-Id": self.resource_id,
"X-Api-Connect-Id": uuid.uuid4(),
}
self.ws = await websockets.connect(
self.ws_url, additional_headers=ws_header, max_size=1000000000
)
logger.bind(tag=TAG).info("WebSocket连接建立成功")
return self.ws
except Exception as e:
logger.bind(tag=TAG).error(f"建立连接失败: {str(e)}")
self.ws = None
raise
def tts_text_priority_thread(self):
"""火山引擎双流式TTS的文本处理线程"""
logger.bind(tag=TAG).info("TTS文本处理线程启动")
while not self.conn.stop_event.is_set():
try:
logger.bind(tag=TAG).debug("等待TTS文本队列消息...")
message = self.tts_text_queue.get(timeout=1)
logger.bind(tag=TAG).debug(
f"TTS任务|{message.sentence_type.name} {message.content_type.name}"
f"收到TTS任务|{message.sentence_type.name} {message.content_type.name} | 会话ID: {self.conn.sentence_id}"
)
if self.conn.client_abort:
logger.bind(tag=TAG).info("收到打断信息,终止TTS文本处理线程")
continue
if message.sentence_type == SentenceType.FIRST:
# 初始化参数
future = asyncio.run_coroutine_threadsafe(
self.start_session(self.conn.sentence_id), loop=self.conn.loop
)
future.result()
self.tts_audio_first_sentence = True
self.before_stop_play_files.clear()
elif ContentType.TEXT == message.content_type:
if message.content_detail:
try:
logger.bind(tag=TAG).info("开始启动TTS会话...")
future = asyncio.run_coroutine_threadsafe(
self.text_to_speak(message.content_detail, None),
self.start_session(self.conn.sentence_id),
loop=self.conn.loop,
)
future.result()
self.tts_audio_first_sentence = True
self.before_stop_play_files.clear()
logger.bind(tag=TAG).info("TTS会话启动成功")
except Exception as e:
logger.bind(tag=TAG).error(f"启动TTS会话失败: {str(e)}")
continue
elif ContentType.TEXT == message.content_type:
if message.content_detail:
try:
logger.bind(tag=TAG).debug(
f"开始发送TTS文本: {message.content_detail}"
)
future = asyncio.run_coroutine_threadsafe(
self.text_to_speak(message.content_detail, None),
loop=self.conn.loop,
)
future.result()
logger.bind(tag=TAG).debug("TTS文本发送成功")
except Exception as e:
logger.bind(tag=TAG).error(f"发送TTS文本失败: {str(e)}")
continue
elif ContentType.FILE == message.content_type:
logger.bind(tag=TAG).info(
f"添加音频文件到待播放列表: {message.content_file}"
)
self.before_stop_play_files.append(
(message.content_file, message.content_detail)
)
if message.sentence_type == SentenceType.LAST:
future = asyncio.run_coroutine_threadsafe(
self.finish_session(self.conn.sentence_id), loop=self.conn.loop
)
future.result()
try:
logger.bind(tag=TAG).info("开始结束TTS会话...")
future = asyncio.run_coroutine_threadsafe(
self.finish_session(self.conn.sentence_id),
loop=self.conn.loop,
)
future.result()
except Exception as e:
logger.bind(tag=TAG).error(f"结束TTS会话失败: {str(e)}")
continue
except queue.Empty:
continue
@@ -226,96 +260,220 @@ class TTSProvider(TTSProviderBase):
logger.bind(tag=TAG).error(
f"处理TTS文本失败: {str(e)}, 类型: {type(e).__name__}, 堆栈: {traceback.format_exc()}"
)
continue
async def text_to_speak(self, text, _):
# 发送文本
await self.send_text(self.speaker, text, self.conn.sentence_id)
return
"""发送文本到TTS服务"""
try:
# 建立新连接
if self.ws is None:
await handleAbortMessage(self.conn)
logger.bind(tag=TAG).error(f"WebSocket连接不存在,终止发送文本")
return
# 发送文本
await self.send_text(self.speaker, text, self.conn.sentence_id)
return
except Exception as e:
logger.bind(tag=TAG).error(f"发送TTS文本失败: {str(e)}")
if self.ws:
try:
await self.ws.close()
except:
pass
self.ws = None
raise
###################################################################################
# 火山双流式TTS重写父类的方法--结束
###################################################################################
def _start_monitor_tts_response_thread(self):
# 初始化链接
asyncio.run_coroutine_threadsafe(
self._start_monitor_tts_response(), loop=self.conn.loop
)
async def start_session(self, session_id):
logger.bind(tag=TAG).info(f"开始会话~~{session_id}")
try:
# 建立新连接
await self._ensure_connection()
# 启动监听任务
self._monitor_task = asyncio.create_task(self._start_monitor_tts_response())
header = Header(
message_type=FULL_CLIENT_REQUEST,
message_type_specific_flags=MsgTypeFlagWithEvent,
serial_method=JSON,
).as_bytes()
optional = Optional(
event=EVENT_StartSession, sessionId=session_id
).as_bytes()
payload = self.get_payload_bytes(
event=EVENT_StartSession, speaker=self.speaker
)
await self.send_event(header, optional, payload)
logger.bind(tag=TAG).info("会话启动请求已发送")
except Exception as e:
logger.bind(tag=TAG).error(f"启动会话失败: {str(e)}")
# 确保清理资源
if hasattr(self, "_monitor_task"):
try:
self._monitor_task.cancel()
await self._monitor_task
except:
pass
self._monitor_task = None
if self.ws:
try:
await self.ws.close()
except:
pass
self.ws = None
raise
async def finish_session(self, session_id):
logger.bind(tag=TAG).info(f"关闭会话~~{session_id}")
try:
if self.ws:
header = Header(
message_type=FULL_CLIENT_REQUEST,
message_type_specific_flags=MsgTypeFlagWithEvent,
serial_method=JSON,
).as_bytes()
optional = Optional(
event=EVENT_FinishSession, sessionId=session_id
).as_bytes()
payload = str.encode("{}")
await self.send_event(header, optional, payload)
logger.bind(tag=TAG).info("会话结束请求已发送")
# 等待监听任务完成
if hasattr(self, "_monitor_task"):
try:
await self._monitor_task
except Exception as e:
logger.bind(tag=TAG).error(
f"等待监听任务完成时发生错误: {str(e)}"
)
finally:
self._monitor_task = None
# 关闭连接
await self.close()
except Exception as e:
logger.bind(tag=TAG).error(f"关闭会话失败: {str(e)}")
# 确保清理资源
if hasattr(self, "_monitor_task"):
try:
self._monitor_task.cancel()
await self._monitor_task
except:
pass
self._monitor_task = None
if self.ws:
try:
await self.ws.close()
except:
pass
self.ws = None
raise
async def close(self):
"""资源清理方法"""
if self.ws:
try:
await self.ws.close()
except:
pass
self.ws = None
async def _start_monitor_tts_response(self):
"""监听TTS响应"""
opus_datas_cache = []
# 添加标志来区分是否是第一句话
is_first_sentence = True
while not self.conn.stop_event.is_set():
try:
# 确保 `recv()` 运行在同一个 event loop
msg = await self.ws.recv()
res = self.parser_response(msg)
self.print_response(res, "send_text res:")
try:
while not self.conn.stop_event.is_set():
try:
# 确保 `recv()` 运行在同一个 event loop
msg = await self.ws.recv()
res = self.parser_response(msg)
self.print_response(res, "send_text res:")
if res.optional.event == EVENT_TTSSentenceStart:
json_data = json.loads(res.payload.decode("utf-8"))
self.tts_text = json_data.get("text", "")
logger.bind(tag=TAG).debug(f"句子语音生成开始: {self.tts_text}")
self.tts_audio_queue.put((SentenceType.FIRST, [], self.tts_text))
opus_datas_cache = []
elif (
res.optional.event == EVENT_TTSResponse
and res.header.message_type == AUDIO_ONLY_RESPONSE
):
logger.bind(tag=TAG).debug(f"推送数据到队列里面~~")
opus_datas = self.wav_to_opus_data_audio_raw(res.payload)
logger.bind(tag=TAG).debug(
f"推送数据到队列里面帧数~~{len(opus_datas)}"
)
if is_first_sentence:
# 第一句话直接发送
# 检查客户端是否中止
if self.conn.client_abort:
logger.bind(tag=TAG).info("收到打断信息,终止监听TTS响应")
break
if res.optional.event == EVENT_TTSSentenceStart:
json_data = json.loads(res.payload.decode("utf-8"))
self.tts_text = json_data.get("text", "")
logger.bind(tag=TAG).debug(f"句子语音生成开始: {self.tts_text}")
self.tts_audio_queue.put(
(SentenceType.MIDDLE, opus_datas, self.tts_text)
(SentenceType.FIRST, [], self.tts_text)
)
else:
# 后续句子缓存
opus_datas_cache = opus_datas_cache + opus_datas
elif res.optional.event == EVENT_TTSSentenceEnd:
logger.bind(tag=TAG).info(f"句子语音生成成功:{self.tts_text}")
if not is_first_sentence:
# 只有非第一句话才发送缓存的数据
self.tts_audio_queue.put(
(SentenceType.MIDDLE, opus_datas_cache, self.tts_text)
opus_datas_cache = []
elif (
res.optional.event == EVENT_TTSResponse
and res.header.message_type == AUDIO_ONLY_RESPONSE
):
logger.bind(tag=TAG).debug(f"推送数据到队列里面~~")
opus_datas = self.wav_to_opus_data_audio_raw(res.payload)
logger.bind(tag=TAG).debug(
f"推送数据到队列里面帧数~~{len(opus_datas)}"
)
# 第一句话结束后,将标志设置为False
is_first_sentence = False
elif res.optional.event == EVENT_SessionFinished:
logger.bind(tag=TAG).debug(f"会话结束~~")
for tts_file, text in self.before_stop_play_files:
if tts_file and os.path.exists(tts_file):
audio_datas = self._process_audio_file(tts_file)
if is_first_sentence:
# 第一句话直接发送
self.tts_audio_queue.put(
(SentenceType.MIDDLE, audio_datas, text)
(SentenceType.MIDDLE, opus_datas, self.tts_text)
)
self.before_stop_play_files.clear()
self.tts_audio_queue.put((SentenceType.LAST, [], None))
opus_datas_cache = []
is_first_sentence = True
continue
except websockets.ConnectionClosed:
break # 连接关闭时退出监听
except Exception as e:
logger.bind(tag=TAG).error(f"Error in _start_monitor_tts_response: {e}")
traceback.print_exc()
continue
else:
# 后续句子缓存
opus_datas_cache = opus_datas_cache + opus_datas
elif res.optional.event == EVENT_TTSSentenceEnd:
logger.bind(tag=TAG).info(f"句子语音生成成功:{self.tts_text}")
if not is_first_sentence:
# 只有非第一句话才发送缓存的数据
self.tts_audio_queue.put(
(SentenceType.MIDDLE, opus_datas_cache, self.tts_text)
)
# 第一句话结束后,将标志设置为False
is_first_sentence = False
elif res.optional.event == EVENT_SessionFinished:
logger.bind(tag=TAG).debug(f"会话结束~~")
for tts_file, text in self.before_stop_play_files:
if tts_file and os.path.exists(tts_file):
audio_datas = self._process_audio_file(tts_file)
self.tts_audio_queue.put(
(SentenceType.MIDDLE, audio_datas, text)
)
self.before_stop_play_files.clear()
self.tts_audio_queue.put((SentenceType.LAST, [], None))
break
except websockets.ConnectionClosed:
logger.bind(tag=TAG).warning("WebSocket连接已关闭")
break
except Exception as e:
logger.bind(tag=TAG).error(
f"Error in _start_monitor_tts_response: {e}"
)
traceback.print_exc()
break
finally:
# 确保清理资源
if self.ws:
try:
await self.ws.close()
except:
pass
self.ws = None
async def send_event(
self, header: bytes, optional: bytes | None = None, payload: bytes = None
):
full_client_request = bytearray(header)
if optional is not None:
full_client_request.extend(optional)
if payload is not None:
payload_size = len(payload).to_bytes(4, "big", signed=True)
full_client_request.extend(payload_size)
full_client_request.extend(payload)
await self.ws.send(full_client_request)
try:
full_client_request = bytearray(header)
if optional is not None:
full_client_request.extend(optional)
if payload is not None:
payload_size = len(payload).to_bytes(4, "big", signed=True)
full_client_request.extend(payload_size)
full_client_request.extend(payload)
await self.ws.send(full_client_request)
except websockets.ConnectionClosed:
logger.bind(tag=TAG).error(f"ConnectionClosed")
raise
async def send_text(self, speaker: str, text: str, session_id):
header = Header(
@@ -438,92 +596,6 @@ class TTSProvider(TTSProviderBase):
)
)
async def finish_connection(self):
header = Header(
message_type=FULL_CLIENT_REQUEST,
message_type_specific_flags=MsgTypeFlagWithEvent,
serial_method=JSON,
).as_bytes()
optional = Optional(event=EVENT_FinishConnection).as_bytes()
payload = str.encode("{}")
await self.send_event(header, optional, payload)
return
async def start_session(self, session_id):
logger.bind(tag=TAG).debug(f"开始会话~~{session_id}")
async with self._session_lock:
# 如果已有会话未结束,先关闭它
if self._session_started and not self._session_finished:
logger.bind(tag=TAG).warning(
f"发现未关闭的会话 {self._current_session_id},正在关闭..."
)
await self.finish_session(self._current_session_id)
# 重置会话状态
self._current_session_id = session_id
self._session_started = True
self._session_finished = False
header = Header(
message_type=FULL_CLIENT_REQUEST,
message_type_specific_flags=MsgTypeFlagWithEvent,
serial_method=JSON,
).as_bytes()
optional = Optional(
event=EVENT_StartSession, sessionId=session_id
).as_bytes()
payload = self.get_payload_bytes(
event=EVENT_StartSession, speaker=self.speaker
)
await self.send_event(header, optional, payload)
async def finish_session(self, session_id):
logger.bind(tag=TAG).debug(f"关闭会话~~{session_id}")
async with self._session_lock:
# 检查会话状态
if not self._session_started:
logger.bind(tag=TAG).warning(f"尝试关闭未开始的会话 {session_id}")
return
if self._session_finished:
logger.bind(tag=TAG).warning(f"会话 {session_id} 已经关闭")
return
if self._current_session_id != session_id:
logger.bind(tag=TAG).warning(
f"尝试关闭错误的会话 {session_id},当前会话为 {self._current_session_id}"
)
return
header = Header(
message_type=FULL_CLIENT_REQUEST,
message_type_specific_flags=MsgTypeFlagWithEvent,
serial_method=JSON,
).as_bytes()
optional = Optional(
event=EVENT_FinishSession, sessionId=session_id
).as_bytes()
payload = str.encode("{}")
await self.send_event(header, optional, payload)
# 更新会话状态
self._session_finished = True
self._session_started = False
self._current_session_id = None
async def reset(self):
# 关闭之前的对话
if self.start_connection_flag:
await self.finish_connection()
self.start_connection_flag = False
await self.start_connection()
self.start_connection_flag = True
async def close(self):
"""资源清理方法"""
await self.finish_connection()
await self.ws.close()
def wav_to_opus_data_audio_raw(self, raw_data_var, is_end=False):
opus_datas = self.opus_encoder.encode_pcm_to_opus(raw_data_var, is_end)
return opus_datas
@@ -0,0 +1,12 @@
from abc import ABC, abstractmethod
from config.logger import setup_logging
TAG = __name__
logger = setup_logging()
class VLLMProviderBase(ABC):
@abstractmethod
def response(self, question, base64_image):
"""VLLM response generator"""
pass
@@ -0,0 +1,65 @@
import openai
import json
from config.logger import setup_logging
from core.utils.util import check_model_key
from core.providers.vllm.base import VLLMProviderBase
TAG = __name__
logger = setup_logging()
class VLLMProvider(VLLMProviderBase):
def __init__(self, config):
self.model_name = config.get("model_name")
self.api_key = config.get("api_key")
if "base_url" in config:
self.base_url = config.get("base_url")
else:
self.base_url = config.get("url")
param_defaults = {
"max_tokens": (500, int),
"temperature": (0.7, lambda x: round(float(x), 1)),
"top_p": (1.0, lambda x: round(float(x), 1)),
}
for param, (default, converter) in param_defaults.items():
value = config.get(param)
try:
setattr(
self,
param,
converter(value) if value not in (None, "") else default,
)
except (ValueError, TypeError):
setattr(self, param, default)
check_model_key("VLLM", self.api_key)
self.client = openai.OpenAI(api_key=self.api_key, base_url=self.base_url)
def response(self, question, base64_image):
try:
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": question},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64_image}"
},
},
],
}
]
response = self.client.chat.completions.create(
model=self.model_name, messages=messages, stream=False
)
return response.choices[0].message.content
except Exception as e:
logger.bind(tag=TAG).error(f"Error in response generation: {e}")
raise