204 lines
7.7 KiB
Python
204 lines
7.7 KiB
Python
"""全能创作 · 会话与消息的写入服务(契约 §1/§2)。
|
|
|
|
只放「怎么把一条消息安全落库」这类底座能力;agent 循环、工具执行、SSE 在
|
|
后续的 creation_agent.py 里,别混进来。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
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,不另开一条,避免中间态刷屏。"""
|
|
message.kind = CreationMessage.Kind.ERROR
|
|
message.text = (error or "生成失败")[:500]
|
|
message.save(update_fields=["kind", "text", "updated_at"])
|
|
conversation = message.conversation
|
|
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
|
|
task = _message_task(message)
|
|
if task is None:
|
|
return False
|
|
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):
|
|
fail_generating_message(message, task.error_message or "生成失败")
|
|
return True
|
|
return False
|
|
|
|
|
|
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,所以对话流仍然是往下叠加;
|
|
这里改的只是同一次生成自己的中间态。
|
|
"""
|
|
message.kind = CreationMessage.Kind.RESULT
|
|
message.payload = {**(message.payload or {}), **meta, "assets": assets}
|
|
message.save(update_fields=["kind", "payload", "updated_at"])
|
|
conversation = message.conversation
|
|
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
|