133 lines
4.1 KiB
Python
133 lines
4.1 KiB
Python
import json
|
|
import logging
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.models.image_engine import ImageEngine
|
|
from app.models.video_engine import VideoEngine
|
|
|
|
logger = logging.getLogger("videogen")
|
|
|
|
|
|
async def resolve_video_engine(
|
|
db: AsyncSession,
|
|
engine_id: str,
|
|
callable_models: list[dict],
|
|
) -> VideoEngine:
|
|
"""根据 engine_id 解析视频引擎，并验证是否在 api-key 的可调用列表中。"""
|
|
# 验证授权
|
|
allowed = {m["engine_id"] for m in callable_models if m.get("engine_type") == "video"}
|
|
if engine_id not in allowed:
|
|
from fastapi import HTTPException, status
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"该 API Key 无权使用引擎 {engine_id}",
|
|
)
|
|
|
|
result = await db.execute(
|
|
select(VideoEngine).where(
|
|
VideoEngine.id == engine_id,
|
|
VideoEngine.is_active == True,
|
|
VideoEngine.deleted_at.is_(None),
|
|
).limit(1)
|
|
)
|
|
engine = result.scalar_one_or_none()
|
|
if not engine:
|
|
from fastapi import HTTPException, status
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"视频引擎 {engine_id} 不存在或未启用",
|
|
)
|
|
return engine
|
|
|
|
|
|
async def resolve_image_engine(
|
|
db: AsyncSession,
|
|
engine_id: str,
|
|
callable_models: list[dict],
|
|
) -> ImageEngine:
|
|
"""根据 engine_id 解析图片引擎，并验证是否在 api-key 的可调用列表中。"""
|
|
allowed = {m["engine_id"] for m in callable_models if m.get("engine_type") == "image"}
|
|
if engine_id not in allowed:
|
|
from fastapi import HTTPException, status
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"该 API Key 无权使用引擎 {engine_id}",
|
|
)
|
|
|
|
result = await db.execute(
|
|
select(ImageEngine).where(
|
|
ImageEngine.id == engine_id,
|
|
ImageEngine.is_active == True,
|
|
ImageEngine.deleted_at.is_(None),
|
|
).limit(1)
|
|
)
|
|
engine = result.scalar_one_or_none()
|
|
if not engine:
|
|
from fastapi import HTTPException, status
|
|
raise HTTPException(
|
|
status_code=status.HTTP_404_NOT_FOUND,
|
|
detail=f"图片引擎 {engine_id} 不存在或未启用",
|
|
)
|
|
return engine
|
|
|
|
|
|
def build_engine_snapshot(engine: VideoEngine | ImageEngine) -> dict:
|
|
"""构建引擎配置快照。"""
|
|
snapshot = {
|
|
"id": str(engine.id),
|
|
"name": str(engine.name),
|
|
"provider": str(getattr(engine, "provider", "")),
|
|
"api_base": str(engine.api_base),
|
|
"model_name": str(engine.model_name),
|
|
}
|
|
# 可选字段
|
|
for field in [
|
|
"supported_ratios", "supported_resolutions", "supported_durations",
|
|
"default_size", "multi_generation_enabled", "max_generation_count",
|
|
]:
|
|
val = getattr(engine, field, None)
|
|
if val is not None:
|
|
if isinstance(val, str):
|
|
try:
|
|
val = json.loads(val)
|
|
except (json.JSONDecodeError, TypeError):
|
|
pass
|
|
snapshot[field] = val
|
|
return snapshot
|
|
|
|
|
|
async def resolve_engine_by_model_name(
|
|
db: AsyncSession,
|
|
model_name: str,
|
|
callable_models: list[dict],
|
|
engine_type: str,
|
|
) -> tuple[str, VideoEngine | ImageEngine]:
|
|
"""根据模型名称查找对应的引擎。
|
|
|
|
Returns:
|
|
(engine_id, engine 对象)
|
|
"""
|
|
# 在 callable_models 中查找
|
|
target = None
|
|
for m in callable_models:
|
|
if m.get("model_name") == model_name and m.get("engine_type") == engine_type:
|
|
target = m
|
|
break
|
|
|
|
if not target:
|
|
from fastapi import HTTPException, status
|
|
raise HTTPException(
|
|
status_code=status.HTTP_403_FORBIDDEN,
|
|
detail=f"该 API Key 无权使用模型 {model_name}",
|
|
)
|
|
|
|
engine_id = target["engine_id"]
|
|
if engine_type == "video":
|
|
engine = await resolve_video_engine(db, engine_id, callable_models)
|
|
else:
|
|
engine = await resolve_image_engine(db, engine_id, callable_models)
|
|
|
|
return engine_id, engine
|