Files
yingqing/core/backend/apps/ai/views.py
T
zycandClaude Opus 4.8 0a519ca892 feat: 任务中心只展示生图任务(模特上身图/平台套图/图片创作)
- AITaskViewSet 支持 ?task_type=a,b,c 过滤
- 前端任务中心请求 task_type=person_image,product_image,过滤掉
  entity_extraction/script_*/storyboard 等流水线内部任务,只留生图工作室产出

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-27 11:04:40 +08:00

101 lines
5.3 KiB
Python

from rest_framework import status
from rest_framework.response import Response
from rest_framework.views import APIView
from rest_framework.viewsets import ReadOnlyModelViewSet
from apps.assets.serializers import AssetSerializer
from apps.common.api import TeamScopedViewSetMixin, get_current_team
from apps.common.celery_health import require_worker
from .models import AITask, ModelConfig
from .serializers import AITaskSerializer, ModelConfigSerializer
from .services import enqueue_standalone_images
class GenerateImageView(APIView):
"""独立生图(不绑项目)· 图片创作/模特图/平台套图共用 —— **异步**。
POST /api/ai/generate-image/ 提交生成,秒级返回 RESERVED 任务列表(慢出图交给 worker)。
GET /api/ai/generate-image/?ids=… 轮询这些任务的状态;成功的任务带回成图 asset。
"""
def post(self, request):
require_worker() # 异步出图依赖 worker 兜底执行,没 worker 直接拒绝(否则任务永远 RESERVED)
prompt = str(request.data.get("prompt") or "").strip()
if not prompt:
return Response({"detail": "prompt 不能为空"}, status=status.HTTP_400_BAD_REQUEST)
mode = str(request.data.get("mode") or "image")
try:
count = int(request.data.get("count") or 1)
except (TypeError, ValueError):
count = 1
product_id = str(request.data.get("product_id") or "").strip() or None
reference_product = bool(request.data.get("reference_product"))
model_id = str(request.data.get("model_id") or "").strip() or None
model_entity_id = str(request.data.get("model_entity_id") or "").strip() or None
ratio = str(request.data.get("ratio") or "").strip() or None
image_model = str(request.data.get("image_model") or "").strip() or None
team = get_current_team(request.user)
try:
tasks = enqueue_standalone_images(team=team, user=request.user, prompt=prompt, mode=mode, count=count, product_id=product_id, reference_product=reference_product, model_id=model_id, model_entity_id=model_entity_id, ratio=ratio, image_model=image_model)
except ValueError as exc: # 无可用模型 / 余额不足等,立即反馈
return Response({"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST)
return Response(
{"tasks": [{"id": str(t.id), "status": t.status} for t in tasks]},
status=status.HTTP_202_ACCEPTED,
)
def get(self, request):
team = get_current_team(request.user)
ids = [s for s in str(request.query_params.get("ids") or "").split(",") if s]
if not ids:
return Response({"tasks": []})
tasks = AITask.objects.filter(team=team, id__in=ids).prefetch_related(
"generated_assets", "generated_assets__files"
)
data = [
{
"id": str(t.id),
"status": t.status,
"error_message": t.error_message,
"assets": AssetSerializer(
[a for a in t.generated_assets.all() if not a.is_deleted], many=True
).data,
}
for t in tasks
]
return Response({"tasks": data})
class AITaskViewSet(TeamScopedViewSetMixin, ReadOnlyModelViewSet):
# 序列化器不含 request_payload/response_payload(单条可达 3MB+ base64 图),defer 掉:
# 否则只为序列化 14 个小字段也会把几十 MB blob 从库里拉回(远程库实测 40 条要 30s+)。
# 默认按创建时间倒序:任务中心 = 历史流水,最新的(多为成功)排最前。
# 缺省排序时 MySQL 按主键(UUID)乱序返回,会把一批旧失败记录顶到首页,
# 前端只取首页 → 误判「全部失败」。order_by 保证稳定且新任务优先。
queryset = AITask.objects.select_related("team", "project", "model_config", "model_config__provider").defer("request_payload", "response_payload").order_by("-created_at")
serializer_class = AITaskSerializer
search_fields = ["idempotency_key", "provider_task_id", "project__name"]
ordering_fields = ["created_at", "updated_at", "completed_at"]
def get_queryset(self):
# 可选 ?task_type=a,b,c 过滤:生图工作室的任务中心只想看生图任务(模特上身图/平台套图/
# 图片创作 = person_image / product_image),不掺脚本/实体抽取/故事板等流水线内部任务。
queryset = super().get_queryset()
raw = self.request.query_params.get("task_type", "").strip()
if raw:
types = [t.strip() for t in raw.split(",") if t.strip()]
if types:
queryset = queryset.filter(task_type__in=types)
return queryset
class ModelConfigViewSet(ReadOnlyModelViewSet):
# 按创建序固定排序:最早创建的 active 模型排第一 = 前端选择器默认项,与 get_default_model 口径一致
# (否则 DB 默认序不稳定,可能默认选到 Gemini 等;用户要默认 = 豆包 2.0 Pro,它最早创建)
queryset = ModelConfig.objects.select_related("provider").filter(status=ModelConfig.Status.ACTIVE).order_by("created_at")
serializer_class = ModelConfigSerializer
search_fields = ["name", "display_name", "capability"]
ordering_fields = ["created_at", "display_name"]