407 lines
16 KiB
Python
407 lines
16 KiB
Python
from fastapi import FastAPI, HTTPException, Depends, Security, File, UploadFile
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from fastapi.security import APIKeyHeader
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
from kafka import KafkaProducer
|
|
from redis import Redis
|
|
import os
|
|
import json
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from dotenv import load_dotenv
|
|
import tempfile
|
|
import hashlib
|
|
from pydantic import BaseModel, Field
|
|
|
|
# 在文件顶部添加这个函数
|
|
def get_audio_hash(text):
|
|
return hashlib.md5(text.encode()).hexdigest()
|
|
|
|
# 加载 .env 文件
|
|
load_dotenv()
|
|
|
|
app = FastAPI()
|
|
v1_chat_app = FastAPI()
|
|
app.mount("/v1_chat", v1_chat_app)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
# 配置
|
|
KAFKA_BROKER = os.getenv('KAFKA_BROKER')
|
|
REDIS_HOST = os.getenv('REDIS_HOST')
|
|
REDIS_PORT = int(os.getenv('REDIS_PORT'))
|
|
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD')
|
|
REDIS_TTS_DB = int(os.getenv('REDIS_TTS_DB'))
|
|
REDIS_ASR_DB = int(os.getenv('REDIS_ASR_DB'))
|
|
REDIS_CHAT_DB = int(os.getenv('REDIS_CHAT_DB'))
|
|
REDIS_API_DB = int(os.getenv('REDIS_API_DB'))
|
|
REDIS_API_USAGE_DB = int(os.getenv('REDIS_API_USAGE_DB'))
|
|
REDIS_TASK_DB = int(os.getenv('REDIS_TASK_DB'))
|
|
|
|
|
|
# Redis 配置
|
|
REDIS_GIRL_DB = int(os.getenv('REDIS_GIRL_DB'))
|
|
REDIS_WOMAN_DB = int(os.getenv('REDIS_WOMAN_DB'))
|
|
REDIS_MAN_DB = int(os.getenv('REDIS_MAN_DB'))
|
|
REDIS_LEIJUN_DB = int(os.getenv('REDIS_LEIJUN_DB'))
|
|
REDIS_DUFU_DB = int(os.getenv('REDIS_DUFU_DB'))
|
|
REDIS_HEJIONG_DB = int(os.getenv('REDIS_HEJIONG_DB'))
|
|
REDIS_MAHUATENG_DB = int(os.getenv('REDIS_MAHUATENG_DB'))
|
|
REDIS_LIDAN_DB = int(os.getenv('REDIS_LIDAN_DB'))
|
|
REDIS_YUHUA_DB = int(os.getenv('REDIS_YUHUA_DB'))
|
|
REDIS_LIUZHENYUN_DB = int(os.getenv('REDIS_LIUZHENYUN_DB'))
|
|
REDIS_DABING_DB = int(os.getenv('REDIS_DABING_DB'))
|
|
REDIS_LUOXIANG_DB = int(os.getenv('REDIS_LUOXIANG_DB'))
|
|
REDIS_XUZHIYUAN_DB = int(os.getenv('REDIS_XUZHIYUAN_DB'))
|
|
|
|
KAFKA_TTS_TOPIC = os.getenv('KAFKA_TTS_TOPIC')
|
|
KAFKA_ASR_TOPIC = os.getenv('KAFKA_ASR_TOPIC')
|
|
KAFKA_CHAT_TOPIC = os.getenv('KAFKA_CHAT_TOPIC')
|
|
|
|
OUTPUT_PATH= os.getenv('OUTPUT_PATH')
|
|
|
|
# 初始化 Kafka Producer
|
|
producer = KafkaProducer(
|
|
bootstrap_servers=[KAFKA_BROKER],
|
|
value_serializer=lambda v: json.dumps(v).encode('utf-8')
|
|
)
|
|
|
|
# 初始化 Redis
|
|
redis_tts_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_TTS_DB)
|
|
redis_asr_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_ASR_DB)
|
|
redis_chat_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_CHAT_DB)
|
|
redis_api_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_API_DB)
|
|
redis_api_usage_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_API_USAGE_DB)
|
|
redis_task_client = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_TASK_DB)
|
|
|
|
redis_tts_girl = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_GIRL_DB)
|
|
redis_tts_woman = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_WOMAN_DB)
|
|
redis_tts_man = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_MAN_DB)
|
|
redis_tts_leijun = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_LEIJUN_DB)
|
|
redis_tts_dufu = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DUFU_DB)
|
|
redis_tts_hejiong = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_HEJIONG_DB)
|
|
redis_tts_mahuateng = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_MAHUATENG_DB)
|
|
redis_tts_lidan = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_LIDAN_DB)
|
|
redis_tts_yuhua = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_YUHUA_DB)
|
|
redis_tts_liuzhenyun = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_LIUZHENYUN_DB)
|
|
redis_tts_dabing = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_DABING_DB)
|
|
redis_tts_luoxiang = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_LUOXIANG_DB)
|
|
redis_tts_xuzhiyuan = Redis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, db=REDIS_XUZHIYUAN_DB)
|
|
|
|
# 创建一个音色到对应 Redis 客户端的映射
|
|
voice_to_redis = {
|
|
"default": redis_tts_girl,
|
|
"girl": redis_tts_girl,
|
|
"woman": redis_tts_woman,
|
|
"man": redis_tts_man,
|
|
"leijun": redis_tts_leijun,
|
|
"dufu": redis_tts_dufu,
|
|
"hejiong": redis_tts_hejiong,
|
|
"mahuateng": redis_tts_mahuateng,
|
|
"lidan": redis_tts_lidan,
|
|
"yuhua": redis_tts_yuhua,
|
|
"liuzhenyun": redis_tts_liuzhenyun,
|
|
"dabing": redis_tts_dabing,
|
|
"luoxiang": redis_tts_luoxiang,
|
|
"xuzhiyuan": redis_tts_xuzhiyuan
|
|
}
|
|
|
|
|
|
# 定义API密钥头部
|
|
api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
|
|
|
|
def get_audio_hash(text):
|
|
return hashlib.md5(text.encode()).hexdigest()
|
|
|
|
# 验证API密钥的函数
|
|
async def get_api_key(api_key: str = Security(api_key_header)):
|
|
if api_key and api_key.startswith("Bearer "):
|
|
key = api_key.split(" ")[1]
|
|
if key.startswith("obs-"):
|
|
return key
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="无效的API密钥",
|
|
headers={"WWW-Authenticate": "Bearer"},
|
|
)
|
|
|
|
async def verify_api_key(api_key: str = Depends(get_api_key)):
|
|
redis_key = f"api_key:{api_key}"
|
|
|
|
api_key_info = redis_api_client.hgetall(redis_key)
|
|
|
|
if not api_key_info:
|
|
raise HTTPException(status_code=401, detail="无效的API密钥")
|
|
|
|
api_key_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in api_key_info.items()}
|
|
|
|
if api_key_info.get('is_active') != '1':
|
|
raise HTTPException(status_code=401, detail="API密钥已停用")
|
|
|
|
expires_at = datetime.fromisoformat(api_key_info.get('expires_at'))
|
|
if datetime.now(timezone.utc) > expires_at:
|
|
raise HTTPException(status_code=401, detail="API密钥已过期")
|
|
|
|
usage_info = redis_api_usage_client.hgetall(redis_key)
|
|
usage_info = {k.decode('utf-8'): v.decode('utf-8') for k, v in usage_info.items()}
|
|
|
|
return {
|
|
"api_key": api_key,
|
|
**api_key_info,
|
|
**usage_info
|
|
}
|
|
|
|
async def update_token_usage(api_key: str, new_tokens_used: int, model_name: str):
|
|
redis_key = f"api_key:{api_key}"
|
|
current_time = datetime.now(timezone.utc).isoformat()
|
|
|
|
pipe = redis_api_usage_client.pipeline()
|
|
|
|
pipe.hincrby(redis_key, "tokens_used", new_tokens_used)
|
|
pipe.hset(redis_key, "last_used_at", current_time)
|
|
|
|
model_tokens_field = f"{model_name}_tokens_used"
|
|
model_last_used_field = f"{model_name}_last_used_at"
|
|
|
|
pipe.hincrby(redis_key, model_tokens_field, new_tokens_used)
|
|
pipe.hset(redis_key, model_last_used_field, current_time)
|
|
|
|
pipe.execute()
|
|
|
|
async def process_request(api_key_info: dict, model_name: str, tokens_required: int, task_data: dict, kafka_topic: str):
|
|
api_key = api_key_info['api_key']
|
|
usage_key = f"api_key:{api_key}"
|
|
total_tokens = int(redis_api_usage_client.hget(usage_key, "total_tokens") or 0)
|
|
tokens_used = int(redis_api_usage_client.hget(usage_key, "tokens_used") or 0)
|
|
|
|
if tokens_used + tokens_required > total_tokens:
|
|
raise HTTPException(status_code=403, detail="Token 余额不足")
|
|
|
|
# 更新 token 使用量
|
|
await update_token_usage(api_key, tokens_required, model_name)
|
|
|
|
# 发送任务到Kafka
|
|
producer.send(kafka_topic, task_data)
|
|
|
|
# 获取更新后的 token 使用情况
|
|
updated_api_key_info = await verify_api_key(api_key)
|
|
new_tokens_used = int(updated_api_key_info.get("tokens_used", 0))
|
|
model_tokens_used = int(updated_api_key_info.get(f"{model_name}_tokens_used", 0))
|
|
|
|
return {
|
|
"message": f"{model_name.upper()}请求已排队等待处理",
|
|
"tokens_used": tokens_required,
|
|
"total_tokens_used": new_tokens_used,
|
|
f"{model_name}_tokens_used": model_tokens_used,
|
|
"tokens_remaining": total_tokens - new_tokens_used
|
|
}
|
|
|
|
class TTSRequest(BaseModel):
|
|
text: str
|
|
voice: str = Field(..., description="选择的音色")
|
|
|
|
class ChatRequest(BaseModel):
|
|
session_id: str
|
|
query: str
|
|
model: str = "qwen2.5:3b"
|
|
|
|
|
|
@v1_chat_app.post("/tts")
|
|
async def tts_request(request: TTSRequest, api_key_info: dict = Depends(verify_api_key)):
|
|
task_id = str(uuid.uuid4())
|
|
text_hash = get_audio_hash(request.text)
|
|
|
|
# 验证音色选择
|
|
valid_voices = ["default", "girl", "woman", "man", "leijun", "dufu", "hejiong", "mahuateng", "lidan", "yuhua", "liuzhenyun", "dabing", "luoxiang", "xuzhiyuan"]
|
|
if request.voice not in valid_voices:
|
|
raise HTTPException(status_code=400, detail="无效的音色选择")
|
|
|
|
# 如果声音是 'default'，则将其视为 'girl'
|
|
voice = 'girl' if request.voice == 'default' else request.voice
|
|
|
|
# 使用对应音色的 Redis 客户端
|
|
redis_tts = voice_to_redis[request.voice]
|
|
|
|
# 检查是否已存在相同内容的音频文件
|
|
existing_audio_info = redis_tts.get(f"tts:{text_hash}")
|
|
if existing_audio_info:
|
|
existing_audio_path = json.loads(existing_audio_info)['path']
|
|
if os.path.exists(existing_audio_path):
|
|
return {
|
|
"message": "TTS请求已完成",
|
|
"task_id": task_id,
|
|
"status": "completed",
|
|
"audio_path": existing_audio_path
|
|
}
|
|
|
|
# 如果不存在，创建新的任务
|
|
task_data = {
|
|
"task_id": task_id,
|
|
"text": request.text,
|
|
"text_hash": text_hash,
|
|
"voice": request.voice,
|
|
"status": "queued",
|
|
"created_at": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
|
|
# 存储任务信息到Redis
|
|
redis_task_client.set(f"task_status:tts:{task_id}", "queued")
|
|
|
|
result = await process_request(api_key_info, "tts", 1, task_data, KAFKA_TTS_TOPIC)
|
|
result["task_id"] = task_id
|
|
return result
|
|
@v1_chat_app.post("/asr")
|
|
async def asr_request(audio: UploadFile = File(...), api_key_info: dict = Depends(verify_api_key)):
|
|
task_id = str(uuid.uuid4())
|
|
|
|
UPLOAD_DIR = "/obscura/task/audio_upload"
|
|
os.makedirs(UPLOAD_DIR, exist_ok=True)
|
|
file_path = os.path.join(UPLOAD_DIR, f"{task_id}.wav")
|
|
|
|
with open(file_path, "wb") as temp_audio:
|
|
content = await audio.read()
|
|
temp_audio.write(content)
|
|
|
|
task_data = {
|
|
'file_path': file_path,
|
|
'task_id': task_id,
|
|
'status': 'queued'
|
|
}
|
|
|
|
# 存储任务状态，使用一致的键名格式
|
|
redis_task_client.set(f"task_status:asr:{task_id}", "queued")
|
|
|
|
result = await process_request(api_key_info, "asr", 1, task_data, KAFKA_ASR_TOPIC)
|
|
result["task_id"] = task_id
|
|
return result
|
|
|
|
@v1_chat_app.post("/chat")
|
|
async def chat_request(request: ChatRequest, api_key_info: dict = Depends(verify_api_key)):
|
|
task_id = str(uuid.uuid4())
|
|
task_data = {
|
|
"task_id": task_id,
|
|
"session_id": request.session_id,
|
|
"query": request.query,
|
|
"model": request.model,
|
|
"status": "queued",
|
|
"created_at": datetime.now(timezone.utc).isoformat(),
|
|
}
|
|
|
|
# 设置任务状态为 "queued"
|
|
redis_task_client.set(f"chat:{task_id}:status", "queued")
|
|
|
|
result = await process_request(api_key_info, "chat", 1, task_data, KAFKA_CHAT_TOPIC)
|
|
result["task_id"] = task_id
|
|
return result
|
|
|
|
|
|
@v1_chat_app.get("/chat_result/{task_id}")
|
|
async def get_chat_result(task_id: str, api_key_info: dict = Depends(verify_api_key)):
|
|
# 从Redis任务数据库获取任务状态
|
|
task_status = redis_task_client.get(f"chat:{task_id}:status")
|
|
if task_status:
|
|
status = task_status.decode('utf-8')
|
|
if status == "completed":
|
|
# 从Redis任务数据库获取聊天结果
|
|
chat_result = redis_task_client.get(f"chat:{task_id}:result")
|
|
if chat_result:
|
|
result = json.loads(chat_result)
|
|
return {
|
|
"status": "completed",
|
|
"result": result
|
|
}
|
|
return {"status": status}
|
|
|
|
return {"status": "not_found"}
|
|
|
|
@v1_chat_app.get("/tts_result/{task_id}")
|
|
async def get_tts_result(task_id: str, api_key_info: dict = Depends(verify_api_key)):
|
|
task_status = redis_task_client.get(f"task_status:tts:{task_id}")
|
|
if task_status:
|
|
status = task_status.decode('utf-8')
|
|
if status == "completed":
|
|
task_info = redis_task_client.get(f"task_info:tts:{task_id}")
|
|
if task_info:
|
|
task_data = json.loads(task_info)
|
|
text_hash = task_data['text_hash']
|
|
voice = task_data['voice']
|
|
# 'default' 和 'girl' 都使用 girl 的 Redis
|
|
redis_tts = voice_to_redis['girl'] if voice in ['default', 'girl'] else voice_to_redis[voice]
|
|
|
|
audio_info = redis_tts.get(f"tts:{text_hash}")
|
|
if audio_info:
|
|
audio_path = json.loads(audio_info)['path']
|
|
return {
|
|
"status": "completed",
|
|
"audio_path": audio_path
|
|
}
|
|
return {"status": status}
|
|
|
|
return {"status": "not_found"}
|
|
|
|
@v1_chat_app.get("/asr_result/{task_id}")
|
|
async def get_asr_result(task_id: str, api_key_info: dict = Depends(verify_api_key)):
|
|
# 从Redis任务数据库获取任务状态，使用一致的键名格式
|
|
task_status = redis_task_client.get(f"task_status:asr:{task_id}")
|
|
if task_status:
|
|
status = task_status.decode('utf-8')
|
|
if status == "completed":
|
|
# 从Redis ASR结果数据库获取转录结果
|
|
transcription = redis_asr_client.get(f"asr:{task_id}")
|
|
if transcription:
|
|
return {
|
|
"status": "completed",
|
|
"transcription": transcription.decode('utf-8')
|
|
}
|
|
return {"status": status}
|
|
|
|
return {"status": "not_found"}
|
|
|
|
@v1_chat_app.get("/tts_audio/{task_id}")
|
|
async def get_tts_audio(task_id: str, api_key_info: dict = Depends(verify_api_key)):
|
|
task_status = redis_task_client.get(f"task_status:tts:{task_id}")
|
|
if task_status:
|
|
status = task_status.decode('utf-8')
|
|
if status == "completed":
|
|
# 从任务信息中获取使用的音色
|
|
task_info = redis_task_client.get(f"task_info:tts:{task_id}")
|
|
if task_info:
|
|
task_data = json.loads(task_info)
|
|
voice = task_data.get('voice', 'girl') # 默认使用 'girl'
|
|
# 'default' 和 'girl' 都使用 girl 的 Redis
|
|
redis_tts = voice_to_redis['girl'] if voice in ['default', 'girl'] else voice_to_redis[voice]
|
|
|
|
# 从对应音色的 Redis 获取音频文件路径
|
|
audio_info = redis_tts.get(f"tts:{task_data['text_hash']}")
|
|
if audio_info:
|
|
audio_path = json.loads(audio_info)['path']
|
|
if os.path.exists(audio_path):
|
|
file_name = os.path.basename(audio_path)
|
|
return FileResponse(audio_path, media_type="audio/wav", filename=file_name)
|
|
else:
|
|
raise HTTPException(status_code=404, detail="音频文件不存在")
|
|
elif status == "queued" or status == "processing":
|
|
raise HTTPException(status_code=202, detail="音频文件正在生成中")
|
|
else:
|
|
raise HTTPException(status_code=500, detail="任务处理失败")
|
|
|
|
raise HTTPException(status_code=404, detail="任务不存在")
|
|
@v1_chat_app.get("/getvoice")
|
|
async def get_available_voices(api_key_info: dict = Depends(verify_api_key)):
|
|
valid_voices = ["default", "girl", "woman", "man", "leijun", "dufu", "hejiong", "mahuateng", "lidan", "yuhua", "liuzhenyun", "dabing", "luoxiang", "xuzhiyuan"]
|
|
return {"available_voices": valid_voices}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8008)
|
|
|
|
|