from django.db.models import Count from django.utils import timezone from rest_framework import status from rest_framework.decorators import action from rest_framework.response import Response from rest_framework.views import APIView from rest_framework.viewsets import ModelViewSet, 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, ImageConversation, ModelConfig from .serializers import AITaskSerializer, ImageConversationSerializer, 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 # 平台套图:前端传规范化平台 id(taobao/douyin/…),用于后端注入平台版式块(优化版) platform_id = str(request.data.get("platform_id") or "").strip() or None conversation_id = str(request.data.get("conversation_id") or "").strip() or None # 用户在图片创作里上传的参考图(已先传成 Asset),按 id 列表带入 → 生成时作多图参考(image_edit) raw_refs = request.data.get("reference_image_ids") or [] if isinstance(raw_refs, str): raw_refs = [s for s in raw_refs.split(",") if s.strip()] reference_image_ids = [str(r).strip() for r in raw_refs if str(r).strip()] team = get_current_team(request.user) # 对话归属:传了 id 就用现成对话(限本团队);没传则自动开一条新对话,标题取 prompt 前 24 字。 conversation = None if conversation_id: conversation = ImageConversation.objects.filter(team=team, id=conversation_id, is_deleted=False).first() if conversation is None: conversation = ImageConversation.objects.create( team=team, created_by=request.user, mode=mode if mode in dict(ImageConversation.Mode.choices) else ImageConversation.Mode.IMAGE, product_id=product_id, title=(prompt[:24] or "默认创作"), ) 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, conversation=conversation, reference_image_ids=reference_image_ids, platform_id=platform_id) except ValueError as exc: # 无可用模型 / 余额不足等,立即反馈 return Response({"detail": str(exc)}, status=status.HTTP_400_BAD_REQUEST) # 本次提交即刷新对话活跃时间,左栏「最近」据此置顶 ImageConversation.objects.filter(id=conversation.id).update(last_active_at=timezone.now()) return Response( {"conversation_id": str(conversation.id), "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() # 从 request_payload(已 defer)里只抽 batch_id / mode 两个 JSON 标量供前端「按批分组 + 标签」用: # KeyTextTransform 在 SQL 层 JSON_EXTRACT,不会把几 MB 的 payload 整列拉回(避开 payload 性能坑)。 from django.db.models.fields.json import KeyTextTransform queryset = queryset.annotate( rp_batch_id=KeyTextTransform("batch_id", "request_payload"), rp_mode=KeyTextTransform("mode", "request_payload"), rp_product_id=KeyTextTransform("product_id", "request_payload"), ) 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 # YYX#row22:只有「工作台图片生成」(mode∈model/cover/image)才计入未读; # 脚本/实体抽取/故事板等流水线内部任务不进这块未读统计。 _GEN_MODES = ("model", "cover", "image") def _unread_base(self): """本团队、属于图片生成、且未读(read_at is null)的任务集合(已 annotate rp_mode/rp_product_id)。 过滤必须走 request_payload__mode 路径查找而非 rp_mode 注解:KeyTextTransform 注解上的 exact/in 查找在 MySQL 会把裸字符串塞进 JSON_EXTRACT 当文档解析 → 3141 全查询炸 (sqlite 编译路径不同,单测测不出来);路径查找 RHS 会被正确 JSON 编码。""" return self.get_queryset().filter(request_payload__mode__in=self._GEN_MODES, read_at__isnull=True) @action(detail=False, methods=["get"], url_path="unread") def unread(self, request): """未读生成任务汇总:总数(导航胶囊) + 按商品分组(商品预览角标)。""" rows = self._unread_base().values_list("rp_product_id", flat=True) total = 0 by_product: dict[str, int] = {} for pid in rows: total += 1 if pid: by_product[str(pid)] = by_product.get(str(pid), 0) + 1 return Response({"total": total, "by_product": by_product}) @action(detail=False, methods=["get"], url_path="workbench") def workbench(self, request): """R100:工作台(模特上身图/平台套图/图片创作)的生成记录改从后端持久数据拉取 —— 原先只存前端 localStorage(1 小时过期、换浏览器/清缓存即空),任务中心有记录而工作台丢。 与任务中心同源 = 同一批 AITask;按 request_payload.mode(白名单)+ product 过滤, 附带每个任务的存活成图资产,前端按 batch_id 归批还原批次流。 性能:仍只用 KeyTextTransform 在 SQL 层抽 payload 标量、defer 两个巨型 payload 列,不整列拖出。""" from django.db.models.fields.json import KeyTextTransform mode = str(request.query_params.get("mode") or "").strip() if mode not in self._GEN_MODES: return Response({"detail": "mode 仅支持 model / cover / image"}, status=status.HTTP_400_BAD_REQUEST) try: limit = min(500, max(1, int(request.query_params.get("limit") or 200))) except (TypeError, ValueError): limit = 200 qs = ( AITask.objects.filter( team=self.get_team(), project__isnull=True, # 工作台独立生图不挂项目;排除流水线内部任务 task_type__in=[AITask.Type.PERSON_IMAGE, AITask.Type.PRODUCT_IMAGE], ) .defer("request_payload", "response_payload") .annotate( rp_mode=KeyTextTransform("mode", "request_payload"), rp_batch_id=KeyTextTransform("batch_id", "request_payload"), rp_product_id=KeyTextTransform("product_id", "request_payload"), rp_prompt=KeyTextTransform("prompt", "request_payload"), rp_ratio=KeyTextTransform("ratio", "request_payload"), rp_platform_id=KeyTextTransform("platform_id", "request_payload"), rp_model_id=KeyTextTransform("model_id", "request_payload"), rp_model_entity_id=KeyTextTransform("model_entity_id", "request_payload"), ) .filter(request_payload__mode=mode) # 路径查找,不能用 rp_mode 注解比较(MySQL 3141,见 _unread_base) ) product_id = str(request.query_params.get("product") or "").strip() if product_id: qs = qs.filter(request_payload__product_id=product_id) tasks = list(qs.order_by("-created_at").prefetch_related("generated_assets", "generated_assets__files")[:limit]) tasks.reverse() # 旧 → 新,与对话流/工作台批次流的时间序一致 data = [ { "id": str(t.id), "status": t.status, "error_message": t.error_message, "prompt": t.rp_prompt or "", "batch_id": t.rp_batch_id or "", "ratio": t.rp_ratio or "", "product_id": t.rp_product_id or "", "model_id": t.rp_model_id or "", "model_entity_id": t.rp_model_entity_id or "", "platform_id": t.rp_platform_id or "", "created_at": t.created_at, # 软删的图不再出现在工作台记录里(R109:删除资产库图片 → 任务记录联动) "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}) @action(detail=False, methods=["post"], url_path="mark-read") def mark_read(self, request): """标记已读 → 清零未读胶囊。 body 可选:product_id(只清该商品) / batch_id(只清该批) / ids(指定任务); 都不传 = 把当前团队所有未读生成任务标记已读(进任务中心时调用)。""" qs = self._unread_base() product_id = str(request.data.get("product_id") or "").strip() batch_id = str(request.data.get("batch_id") or "").strip() ids = request.data.get("ids") or [] if isinstance(ids, str): ids = [s for s in ids.split(",") if s.strip()] if product_id: qs = qs.filter(request_payload__product_id=product_id) if batch_id: qs = qs.filter(request_payload__batch_id=batch_id) if ids: qs = qs.filter(id__in=[str(i).strip() for i in ids if str(i).strip()]) updated = qs.update(read_at=timezone.now()) return Response({"updated": updated}) class ImageConversationViewSet(TeamScopedViewSetMixin, ModelViewSet): """图片创作工作室的「对话」CRUD。 list 按 ?mode= 过滤、排除软删、按 last_active_at 倒序(左栏「最近」); create 开新对话;partial_update 重命名;destroy 软删(不连带删图)。 detail action `tasks` 返回该对话下的生图任务 + 成图 asset,供切换对话时回填批次流。 """ serializer_class = ImageConversationSerializer queryset = ImageConversation.objects.filter(is_deleted=False).select_related("product").order_by("-last_active_at") def get_queryset(self): queryset = super().get_queryset().annotate(_task_count=Count("tasks")) mode = self.request.query_params.get("mode", "").strip() if mode: queryset = queryset.filter(mode=mode) return queryset def perform_destroy(self, instance): # 软删:对话从列表消失,但其 AITask.conversation 置空由 DB on_delete=SET_NULL 不触发(我们没真删), # 成图始终留在资产库。仅打标记。 instance.is_deleted = True instance.save(update_fields=["is_deleted", "updated_at"]) @action(detail=True, methods=["get"]) def tasks(self, request, pk=None): from apps.assets.models import Asset from apps.assets.serializers import _asset_preview conversation = self.get_object() tasks = ( AITask.objects.filter(conversation=conversation) .prefetch_related("generated_assets", "generated_assets__files") .order_by("created_at") ) # 参考图 id → {name,url}:跨任务可能重复,缓存一次解析,供切换/刷新后批次头回显「参考了哪些图」 ref_cache: dict[str, dict] = {} def resolve_refs(ids): out = [] for rid in ids or []: rid = str(rid) if rid not in ref_cache: a = Asset.objects.filter(id=rid).prefetch_related("files").first() ref_cache[rid] = {"name": a.name, "url": _asset_preview(a)} if a else None if ref_cache[rid]: out.append(ref_cache[rid]) return out data = [ { "id": str(t.id), "status": t.status, "error_message": t.error_message, "prompt": (t.request_payload or {}).get("prompt", ""), "batch_id": (t.request_payload or {}).get("batch_id", ""), "ratio": (t.request_payload or {}).get("ratio") or "", "reference_images": resolve_refs((t.request_payload or {}).get("reference_image_ids")), "created_at": t.created_at, "assets": AssetSerializer( [a for a in t.generated_assets.all() if not a.is_deleted], many=True ).data, } for t in tasks ] return Response({"conversation_id": str(conversation.id), "tasks": data}) 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"]