"""全能创作 · 会话与消息的写入服务(契约 §1/§2)。 只放「怎么把一条消息安全落库」这类底座能力;agent 循环、工具执行、SSE 在 后续的 creation_agent.py 里,别混进来。 """ from __future__ import annotations from datetime import timedelta import uuid from django.core.cache import cache from django.db import transaction from django.db.models import Max from django.utils import timezone from .models import CreationConversation, CreationMessage @transaction.atomic def append_message( conversation: CreationConversation, *, role: str, kind: str = CreationMessage.Kind.TEXT, text: str = "", payload: dict | None = None, refs: list | None = None, task=None, ) -> CreationMessage: """往会话尾部追加一条消息,并刷新 last_active_at。 seq 在事务里 select_for_update 锁住会话行再取 max+1 —— SSE 流式期间可能有 并发写(用户抢答 / 轮询回填),不锁会撞 uniq_creation_message_seq。 """ locked = CreationConversation.objects.select_for_update().get(pk=conversation.pk) next_seq = (locked.messages.aggregate(m=Max("seq"))["m"] or 0) + 1 message = CreationMessage.objects.create( conversation=locked, role=role, kind=kind, text=text, payload=payload or {}, refs=refs or [], task=task, seq=next_seq, ) locked.last_active_at = timezone.now() locked.save(update_fields=["last_active_at", "updated_at"]) conversation.last_active_at = locked.last_active_at return message def _assets_from_task(task) -> list[dict]: """任务落库的资产 → 结果卡要用的 {id,url,cover,type}。 URL 必须走长期直链:预签名链 1 小时过期,写进消息 payload 第二天就打不开。 """ from apps.ai.services import asset_stable_url from apps.assets.models import Asset items = [] assets = Asset.objects.filter( origin_task=task, is_deleted=False, purged_at__isnull=True, ).prefetch_related("files") for asset in assets: url, cover = asset_stable_url(asset) if not url: continue items.append({ "id": str(asset.id), "url": url, "cover": cover or url, "type": "video" if asset.asset_type == Asset.Type.VIDEO else "image", }) return items def _meta_from_task(task, message: CreationMessage) -> dict: payload = message.payload or {} req = task.request_payload or {} model = "" if task.model_config_id: model = task.model_config.display_name or task.model_config.name return { "model": model or payload.get("model") or req.get("model") or "", "ratio": req.get("ratio") or payload.get("ratio") or "", "prompt": payload.get("prompt") or req.get("prompt") or "", } def _message_task(message: CreationMessage): """GENERATING 消息挂的任务:优先 FK,payload.task_id 兜底(旧数据 / 序列化往返)。""" if message.task_id: return message.task task_id = (message.payload or {}).get("task_id") if not task_id: return None from .models import AITask return AITask.objects.select_related("model_config").filter(id=task_id).first() @transaction.atomic def fail_generating_message(message: CreationMessage, error: str) -> CreationMessage: """生成失败:GENERATING 原地改成 ERROR,不另开一条,避免中间态刷屏。""" is_person_reference = (message.payload or {}).get("kind") == "person_reference" message.kind = CreationMessage.Kind.ERROR message.text = (error or "生成失败")[:500] message.save(update_fields=["kind", "text", "updated_at"]) conversation = message.conversation if is_person_reference: memory = dict(conversation.memory or {}) memory["person_source_pending"] = False memory.pop("person_source_ready", None) conversation.memory = memory conversation.status = CreationConversation.Status.RUNNING conversation.agent_status = CreationConversation.AgentStatus.AWAITING_USER conversation.last_active_at = timezone.now() conversation.save(update_fields=[ "memory", "status", "agent_status", "last_active_at", "updated_at", ]) append_message( conversation, role="assistant", text="这次人物参考没生成成功。可以重新选择人物来源。", payload={ "reply_hint": "选择一种方式继续…", "reply_options": [{"label": "重新选择人物", "text": "重新选择人物来源"}], }, ) return message conversation.status = CreationConversation.Status.FAILED conversation.last_active_at = timezone.now() conversation.save(update_fields=["status", "last_active_at", "updated_at"]) return message def sync_generating_message(message: CreationMessage) -> bool: """看挂着的 AITask 是否已经终态,是就把 GENERATING 改成 RESULT / ERROR。 出图/出片在 worker 里跑,agent 只提交。前端轮询 GET 会话时靠这个回填; worker 结束时也会调一次,不用干等到下一次轮询。 返回是否改了这条消息。 """ from .models import AITask if message.kind != CreationMessage.Kind.GENERATING: return False payload = message.payload or {} if payload.get("kind") == "video_segments": return _sync_segmented_video_message(message) task = _message_task(message) if task is None: return False if payload.get("kind") == "video_merge" and task.status in (AITask.Status.SUBMITTED, AITask.Status.POLLING): # Celery 不在线的本地环境同样可完成用户已确认的合并;此前从不在普通出片完成时调用。 run_segmented_video_merge(str(task.id)) task.refresh_from_db() if task.status == AITask.Status.SUCCEEDED: assets = _assets_from_task(task) if not assets: return False # 状态已成功但资产还没落(极端竞态),下轮再试 finish_generating_message(message, assets=assets, meta=_meta_from_task(task, message)) return True if task.status in (AITask.Status.FAILED, AITask.Status.CANCELLED): from .generation_errors import public_error_for_task public_error = public_error_for_task(task) fail_generating_message( message, public_error.fallback_message if public_error else (task.error_message or "生成失败"), ) return True return False def _sync_segmented_video_message(message: CreationMessage) -> bool: """同步全能创作的多段出片。 每段仍是普通 FREE_VIDEO 任务,故可沿用既有计费、轮询和资产落库;这里仅把它们聚合成 一张结果卡。任一片段先完成就先写回 GENERATING 卡供用户预览;所有分段完成后才转成 可合并的结果卡,绝不在此处触发 ffmpeg。 """ from .free_video import finalize_free_video from .models import AITask payload = message.payload or {} ids = [str(value) for value in payload.get("task_ids") or [] if value] if not ids: return False by_id = { str(task.id): task for task in AITask.objects.filter(id__in=ids).select_related("model_config") } tasks = [by_id.get(task_id) for task_id in ids] if any(task is None for task in tasks): fail_generating_message(message, "分段视频任务不完整,请重新生成。") return True # 本地没有 worker 时,用户的会话轮询本身即可推进每个片段;已在 worker 中的任务则幂等返回。 refreshed = [] for task in tasks: assert task is not None if task.status in (AITask.Status.SUBMITTED, AITask.Status.POLLING): try: task = finalize_free_video(task=task) except Exception: # noqa: BLE001 - 单段网络抖动不能使整组直接失败 task.refresh_from_db() else: task.refresh_from_db() refreshed.append(task) failed = next( (task for task in refreshed if task.status in (AITask.Status.FAILED, AITask.Status.CANCELLED)), None, ) if failed is not None: from .generation_errors import public_error_for_task number = next((index + 1 for index, task in enumerate(refreshed) if task.id == failed.id), 1) public_error = public_error_for_task(failed, operation="video_generate") detail = public_error.fallback_message if public_error else (failed.error_message or "请重试") fail_generating_message(message, f"第 {number} 段生成失败:{detail}") return True assets: list[dict] = [] completed_segments = 0 for index, task in enumerate(refreshed, start=1): if task.status != AITask.Status.SUCCEEDED: continue task_assets = _assets_from_task(task) # 上游状态先成功、资产稍后才落库时,保留 GENERATING,下一轮再展示这段。 if not task_assets: continue completed_segments += 1 for asset in task_assets: assets.append({**asset, "label": f"第 {index} 段", "segment_index": index}) # 先完成的片段必须立刻回到前端,不能等另一段慢任务一起完成才出现。 # 保持同一条 GENERATING 消息,避免把一个 60 秒视频拆成多条对话消息。 if not all(task.status == AITask.Status.SUCCEEDED for task in refreshed): progress = { **payload, "assets": assets, "completed_segment_count": completed_segments, } if progress != payload: message.payload = progress message.save(update_fields=["payload", "updated_at"]) return True return False # 任务都成功但有片段的资产还在落库,继续保持生成中,避免最终结果缺片。 if completed_segments != len(refreshed): progress = { **payload, "assets": assets, "completed_segment_count": completed_segments, } if progress != payload: message.payload = progress message.save(update_fields=["payload", "updated_at"]) return True return False first = refreshed[0] finish_generating_message( message, assets=assets, meta={ **_meta_from_task(first, message), "kind": "video_segments", "segment_count": len(refreshed), "total_duration": payload.get("total_duration") or "", "needs_merge": True, "segments": payload.get("segments") or [], }, ) return True def start_segmented_video_merge(*, conversation: CreationConversation, message: CreationMessage, user): """用户明确点击后才建立合并任务;这之前绝不下载片段或调用 ffmpeg。""" from .models import AITask payload = dict(message.payload or {}) if message.kind != CreationMessage.Kind.RESULT or not payload.get("needs_merge"): raise ValueError("这条结果不需要合并") state = str(payload.get("merge_state") or "") if state in {"queued", "processing", "completed"}: raise ValueError("这组片段已经在合并中或已合并") task_ids = [str(value) for value in payload.get("task_ids") or [] if value] source_tasks = list(AITask.objects.filter(id__in=task_ids).select_related("model_config")) if len(source_tasks) != len(task_ids): raise ValueError("分段视频不完整,无法合并") model_config = source_tasks[0].model_config merge_task = AITask.objects.create( team=conversation.team, created_by=user, project=None, task_type=AITask.Type.EXPORT, status=AITask.Status.SUBMITTED, model_config=model_config, idempotency_key=f"omni-video-merge:{conversation.id}:{message.id}:{uuid.uuid4()}", request_payload={ "feature": "omni_video_merge", "source_message_id": str(message.id), "source_task_ids": task_ids, "prompt": payload.get("prompt") or "", }, ) payload["merge_state"] = "queued" payload["merge_task_id"] = str(merge_task.id) message.payload = payload message.save(update_fields=["payload", "updated_at"]) generating = append_message( conversation, role="assistant", kind=CreationMessage.Kind.GENERATING, payload={"task_id": str(merge_task.id), "kind": "video_merge", "prompt": payload.get("prompt") or ""}, task=merge_task, ) return merge_task, generating def run_segmented_video_merge(task_id: str) -> None: """后台执行用户已确认的合并。任务创建不等于执行;仅此函数调用 ffmpeg。""" from .free_video import _store_free_video_media from .models import AITask from .services import asset_stable_url from .video_replace import _concat_shot_media from apps.assets.models import Asset task = AITask.objects.select_related("team", "created_by").filter(id=task_id).first() if task is None or task.task_type != AITask.Type.EXPORT: return source_message_id = str((task.request_payload or {}).get("source_message_id") or "") source_message = CreationMessage.objects.filter(id=source_message_id).first() if source_message is None: return try: with transaction.atomic(): locked = AITask.objects.select_for_update().get(id=task.id) if locked.status not in (AITask.Status.SUBMITTED, AITask.Status.POLLING): return locked.status = AITask.Status.POSTPROCESSING locked.save(update_fields=["status", "updated_at"]) urls: list[str] = [] for source_id in (locked.request_payload or {}).get("source_task_ids") or []: source_asset = ( Asset.objects.filter(origin_task_id=source_id, is_deleted=False, purged_at__isnull=True) .prefetch_related("files") .first() ) url, _cover = asset_stable_url(source_asset) if not url: raise ValueError("有片段尚未准备好,暂时不能合并") urls.append(url) merged_bytes = _concat_shot_media(urls) _store_free_video_media(task=locked, video_bytes=merged_bytes) with transaction.atomic(): locked = AITask.objects.select_for_update().get(id=task.id) if locked.status != AITask.Status.POSTPROCESSING: return locked.status = AITask.Status.SUCCEEDED locked.completed_at = timezone.now() locked.save(update_fields=["status", "completed_at", "updated_at"]) source = CreationMessage.objects.select_for_update().get(id=source_message.id) source_payload = dict(source.payload or {}) source_payload["merge_state"] = "completed" source.payload = source_payload source.save(update_fields=["payload", "updated_at"]) except Exception as exc: # noqa: BLE001 - 合并失败不影响已完成的分段视频 import logging logging.getLogger(__name__).exception("omni video merge failed for %s", task_id) with transaction.atomic(): locked = AITask.objects.select_for_update().get(id=task.id) if locked.status == AITask.Status.POSTPROCESSING: locked.status = AITask.Status.FAILED locked.error_code = "MergeError" locked.error_message = str(exc)[:500] locked.completed_at = timezone.now() locked.save(update_fields=["status", "error_code", "error_message", "completed_at", "updated_at"]) source = CreationMessage.objects.select_for_update().filter(id=source_message.id).first() if source is not None: source_payload = dict(source.payload or {}) source_payload["merge_state"] = "failed" source.payload = source_payload source.save(update_fields=["payload", "updated_at"]) finally: # 复用既有「生成中 → 结果/错误」转换,生成的合成片会成为独立结果卡。 fresh = AITask.objects.filter(id=task_id).first() if fresh is not None: sync_generating_for_task(fresh) def sync_generating_messages(conversation: CreationConversation) -> int: """把一条会话里所有已结束的 GENERATING 回填。返回改了几条。""" pending = list( conversation.messages.filter(kind=CreationMessage.Kind.GENERATING) .select_related("task", "task__model_config") ) return sum(1 for message in pending if sync_generating_message(message)) def sync_generating_for_task(task) -> int: """worker / poll 终态后:只扫挂在这个任务上的 GENERATING。失败不能向外抛。""" try: pending = list( CreationMessage.objects.filter( task=task, kind=CreationMessage.Kind.GENERATING, ).select_related("conversation", "task", "task__model_config") ) return sum(1 for message in pending if sync_generating_message(message)) except Exception: # noqa: BLE001 — 回填失败不能把已经成功的出图任务打成失败 import logging logging.getLogger(__name__).warning( "omni create: sync generating for task %s failed", getattr(task, "id", "?"), exc_info=True, ) return 0 @transaction.atomic def finish_generating_message(message: CreationMessage, *, assets: list[dict], meta: dict) -> CreationMessage: """把「生成中」原地改成「结果」(契约 §2:不新增消息,避免中间态刷屏)。 重生成是**新开一条** GENERATING → RESULT,所以对话流仍然是往下叠加; 这里改的只是同一次生成自己的中间态。 """ original_payload = dict(message.payload or {}) is_person_reference = original_payload.get("kind") == "person_reference" message.kind = CreationMessage.Kind.RESULT message.payload = {**original_payload, **meta, "assets": assets} message.save(update_fields=["kind", "payload", "updated_at"]) conversation = message.conversation if is_person_reference: from apps.assets.models import Asset, Model asset_id = str((assets[0] if assets else {}).get("id") or "") asset = Asset.objects.filter( id=asset_id, team=conversation.team, is_deleted=False, purged_at__isnull=True, ).first() if asset is None: return fail_generating_message(message, "人物参考已生成,但未找到可锁定的图片资产") model = Model.objects.filter( team=conversation.team, portrait_asset=asset, is_deleted=False, purged_at__isnull=True, ).first() is_pet = "宠物" in str(conversation.preset or "") if model is None: model = Model.objects.create( team=conversation.team, created_by=conversation.created_by, name="平台生成宠物角色" if is_pet else "平台生成出镜人物", source=Model.Source.AI, portrait_asset=asset, description="由全能创作生成并锁定的宠物角色。" if is_pet else "由全能创作生成并锁定的视频出镜人物。", metadata={"feature": "omni_create", "conversation_id": str(conversation.id)}, ) pin_refs(conversation, [{ "type": "model", "id": str(model.id), "name": model.name, "cover": (assets[0] if assets else {}).get("cover") or (assets[0] if assets else {}).get("url") or "", }]) memory = dict(conversation.memory or {}) memory["person_source"] = "platform_generate" memory["person_source_pending"] = False memory["person_source_ready"] = True memory["person_model_id"] = str(model.id) conversation.memory = memory conversation.status = CreationConversation.Status.RUNNING conversation.agent_status = CreationConversation.AgentStatus.AWAITING_USER conversation.last_active_at = timezone.now() conversation.save(update_fields=[ "memory", "status", "agent_status", "last_active_at", "updated_at", ]) append_message( conversation, role="assistant", text=( "宠物角色已生成。你看这只宠物合适吗?是否使用这个宠物角色继续创作?" if is_pet else "人物参考已生成。你看这位角色合适吗?是否使用这个角色继续创作?" ), payload={ "reply_hint": ( "回复「使用这个角色」继续,或说明想要的宠物品种与外观…" if is_pet else "回复「使用这个角色」继续,或说明想调整的地方…" ), "reply_options": [ {"label": "使用这个角色", "text": "使用这个角色继续创作"}, {"label": "重新生成一个", "text": "重新生成一个宠物角色" if is_pet else "重新生成一个角色"}, {"label": "上传其他宠物" if is_pet else "上传其他人物", "text": "我上传宠物参考图" if is_pet else "我上传人物参考图"}, ], }, ) return message conversation.status = CreationConversation.Status.COMPLETED conversation.last_active_at = timezone.now() conversation.save(update_fields=["status", "last_active_at", "updated_at"]) return message def pin_refs(conversation: CreationConversation, refs: list[dict]) -> list[dict]: """把本轮引用的实体并进「实体锁定」。按 (type, id) 去重,保留首次出现的顺序 —— 每轮无条件带上它们,这是多次生成之间锁脸/锁商品的唯一手段(契约 §5)。 """ merged = list(conversation.pinned_refs or []) seen = {(item.get("type"), str(item.get("id"))) for item in merged} changed = False for ref in refs or []: ref_type, ref_id = ref.get("type"), ref.get("id") if not ref_type or not ref_id: continue # 缺 type/id 的 ref 解析不出实体,直接丢 key = (ref_type, str(ref_id)) if key in seen: continue merged.append(ref) seen.add(key) changed = True if changed: conversation.pinned_refs = merged conversation.save(update_fields=["pinned_refs", "updated_at"]) return merged # ---------------------------------------------------------------- Agent 在途锁(一账号同时只整理一份方案) AGENT_LOCK_TTL_SECONDS = 20 * 60 AGENT_STALE_MINUTES = 20 AGENT_BUSY_DETAIL = "当前已有对话正在整理方案,请等待完成后再试" def agent_lock_key(team_id) -> str: return f"omni:agent:{team_id}" def release_agent_lock(team_id, conversation_id=None) -> None: """释放团队级 agent 锁。传 conversation_id 时只在锁仍指向该会话时删,避免误伤。""" key = agent_lock_key(team_id) if conversation_id is None: cache.delete(key) return if str(cache.get(key) or "") == str(conversation_id): cache.delete(key) def acquire_agent_lock(team_id, conversation_id) -> bool: """cache.add 原子占坑。成功返回 True;已有占用返回 False。""" return bool(cache.add(agent_lock_key(team_id), str(conversation_id), timeout=AGENT_LOCK_TTL_SECONDS)) def cleanup_stale_agent_planning(team) -> int: """把超时仍卡在 planning 的会话清回 idle,并尝试释放对应 Redis 锁。""" from django.db.models import Q cutoff = timezone.now() - timedelta(minutes=AGENT_STALE_MINUTES) stale = list( CreationConversation.objects.filter( team=team, agent_status=CreationConversation.AgentStatus.PLANNING, is_deleted=False, ).filter( # started_at 为空也当过期(异常写入) Q(agent_started_at__lt=cutoff) | Q(agent_started_at__isnull=True) )[:20] ) cleared = 0 now = timezone.now() for conv in stale: updated = CreationConversation.objects.filter( pk=conv.pk, agent_status=CreationConversation.AgentStatus.PLANNING, ).update( agent_status=CreationConversation.AgentStatus.IDLE, agent_started_at=None, updated_at=now, ) if updated: release_agent_lock(team.id, conv.id) cleared += 1 return cleared def find_planning_conversation(team, *, exclude_id=None): """返回团队里仍在 planning 的会话(已先清过期)。""" cleanup_stale_agent_planning(team) qs = CreationConversation.objects.filter( team=team, agent_status=CreationConversation.AgentStatus.PLANNING, is_deleted=False, ) if exclude_id is not None: qs = qs.exclude(pk=exclude_id) return qs.order_by("-agent_started_at", "-last_active_at").first() def team_agent_busy(team, *, exclude_id=None) -> bool: """DB planning 或 Redis 锁任一占用 → 忙。exclude_id 用于同一会话续跑自检。""" if find_planning_conversation(team, exclude_id=exclude_id) is not None: return True holder = cache.get(agent_lock_key(team.id)) if not holder: return False if exclude_id is not None and str(holder) == str(exclude_id): return False # 锁还在但持有会话已不在 planning → 孤儿锁,清掉放行 still = CreationConversation.objects.filter( id=holder, team=team, agent_status=CreationConversation.AgentStatus.PLANNING, is_deleted=False, ).exists() if not still: release_agent_lock(team.id, holder) return False return True def begin_agent_planning(conversation: CreationConversation) -> bool: """占 Redis 锁并把会话标成 planning。失败(冲突)返回 False,调用方回 409。""" team_id = conversation.team_id if team_agent_busy(conversation.team, exclude_id=conversation.id): return False if not acquire_agent_lock(team_id, conversation.id): # 锁被别人占了;若锁指向自己(重入),允许续写状态 holder = cache.get(agent_lock_key(team_id)) if str(holder or "") != str(conversation.id): return False now = timezone.now() CreationConversation.objects.filter(pk=conversation.pk).update( agent_status=CreationConversation.AgentStatus.PLANNING, agent_started_at=now, last_active_at=now, updated_at=now, ) conversation.agent_status = CreationConversation.AgentStatus.PLANNING conversation.agent_started_at = now conversation.last_active_at = now clear_agent_cancel(conversation.id) return True def finish_agent_planning( conversation: CreationConversation, *, awaiting_user: bool = False, ) -> None: """Celery finally:写回 agent_status 并释放锁。""" status = ( CreationConversation.AgentStatus.AWAITING_USER if awaiting_user else CreationConversation.AgentStatus.IDLE ) now = timezone.now() CreationConversation.objects.filter(pk=conversation.pk).update( agent_status=status, agent_started_at=None, last_active_at=now, updated_at=now, ) conversation.agent_status = status conversation.agent_started_at = None release_agent_lock(conversation.team_id, conversation.id) # 正常结束也清掉取消标记,避免下一轮误判 clear_agent_cancel(conversation.id) def agent_cancel_key(conversation_id) -> str: return f"omni:agent:cancel:{conversation_id}" def clear_agent_cancel(conversation_id) -> None: cache.delete(agent_cancel_key(conversation_id)) def is_agent_cancel_requested(conversation_id) -> bool: """Celery tool round 之间读:用户点了终止则 True。""" return bool(cache.get(agent_cancel_key(conversation_id))) def request_agent_cancel(conversation: CreationConversation) -> bool: """用户终止整理方案:打 Redis 取消标,立刻 idle + 释放团队锁。 仅当当前仍是 planning 时成功。Celery 在下一轮工具间隙看到标后停跑, 已落库消息保留。返回 True=已受理, False=当时不是 planning(调用方回 409)。 """ # 先打标,再收态 —— worker 在长 LLM 调用回来后也能看见 cache.set(agent_cancel_key(conversation.id), "1", timeout=AGENT_LOCK_TTL_SECONDS) updated = CreationConversation.objects.filter( pk=conversation.pk, agent_status=CreationConversation.AgentStatus.PLANNING, ).update( agent_status=CreationConversation.AgentStatus.IDLE, agent_started_at=None, last_active_at=timezone.now(), updated_at=timezone.now(), ) if not updated: # 已经不是 planning(竞态完成/别人清了);仍留标一会儿无害,清掉避免脏状态 clear_agent_cancel(conversation.id) conversation.refresh_from_db(fields=["agent_status", "agent_started_at"]) return False conversation.agent_status = CreationConversation.AgentStatus.IDLE conversation.agent_started_at = None release_agent_lock(conversation.team_id, conversation.id) return True