150 lines
5.9 KiB
Python
150 lines
5.9 KiB
Python
"""
|
||
app/services/task_service.py — 任务创建 service
|
||
校验有无 key → 建 GenerationTask → 只推 task_id 入队，绝不传 key（基石B）。
|
||
"""
|
||
|
||
import logging
|
||
|
||
from sqlalchemy.orm import Session
|
||
|
||
from app.core.response import raise_business
|
||
from app.middleware.workspace_guard import CurrentUser
|
||
from app.models.task import GenerationTask
|
||
from app.models.workspace import UserApiKey
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _check_user_has_key(db: Session, user_id: int, workspace_id: int) -> None:
|
||
"""校验用户在此 workspace 是否有可用 API Key（openai/apiports均可），没有则引导去配置。"""
|
||
key = (
|
||
db.query(UserApiKey)
|
||
.filter(
|
||
UserApiKey.user_id == user_id,
|
||
UserApiKey.workspace_id == workspace_id,
|
||
UserApiKey.provider.in_(["openai", "apiports"]), # G6坑修复：接受主备通道名
|
||
)
|
||
.first()
|
||
)
|
||
if not key:
|
||
raise_business("尚未配置 API Key，请先在设置中录入")
|
||
|
||
|
||
def _check_concurrency_limit(db: Session, user_id: int, workspace_id: int) -> None:
|
||
"""
|
||
校验该用户未完成任务数未超并发上限（红线=5,可配置)。
|
||
只算 pending/generating(真占 worker),挑选/审核态不计。超限引导稍后再试。
|
||
"""
|
||
from app.core.config import get_settings
|
||
from app.constants.enums import TaskStatus
|
||
|
||
limit = get_settings().MAX_CONCURRENT_TASKS_PER_USER
|
||
running = (
|
||
db.query(GenerationTask)
|
||
.filter(
|
||
GenerationTask.operator_id == user_id,
|
||
GenerationTask.workspace_id == workspace_id,
|
||
GenerationTask.status.in_([TaskStatus.PENDING.value, TaskStatus.GENERATING.value]),
|
||
)
|
||
.count()
|
||
)
|
||
if running >= limit:
|
||
raise_business(f"您有 {running} 个任务正在生成，已达并发上限 {limit} 个，请等待完成后再发起")
|
||
|
||
|
||
def create_generation_task(
|
||
db: Session,
|
||
current_user: CurrentUser,
|
||
body, # CreateTaskRequest
|
||
) -> GenerationTask:
|
||
"""
|
||
建 GenerationTask 并推 Celery 队列。
|
||
只传 task_id，绝不传 key（基石B）。
|
||
"""
|
||
if body.track == "ai":
|
||
# 轨A：先检查有没有 key
|
||
_check_user_has_key(db, current_user.user_id, current_user.workspace_id)
|
||
# 并发上限：只算正在消耗生成资源的任务(pending/generating)，
|
||
# 已生成完等挑选/审核的不占 worker。红线=每用户5个(可配置)。
|
||
_check_concurrency_limit(db, current_user.user_id, current_user.workspace_id)
|
||
|
||
from app.models.product import Product
|
||
|
||
product = db.query(Product).filter(
|
||
Product.id == body.product_id,
|
||
Product.workspace_id == current_user.workspace_id,
|
||
).first()
|
||
if not product:
|
||
raise_business("产品不存在")
|
||
|
||
# 禁降级铁律：本次产品入镜(need_product_image=True)时，产品必须已上传参考图，
|
||
# 否则拒绝建任务（不允许降级纯文生图，防产品包装跑偏/过抽检失败）。
|
||
need_img = getattr(body, "need_product_image", True)
|
||
if need_img:
|
||
if not (product.image_path or "").strip():
|
||
raise_business("该产品未上传参考图，无法生成产品入镜内容；请先到产品库上传产品图，或关闭「产品入镜」开关")
|
||
|
||
# 第2环：关联标杆笔记ID存库（JSON list）。pipeline 据此读 features_json 注入文案 prompt。
|
||
import json
|
||
_bids = getattr(body, "benchmark_ids", None) or []
|
||
if _bids:
|
||
from app.models.product import BenchmarkNote
|
||
|
||
ids = [int(i) for i in _bids]
|
||
rows = (
|
||
db.query(BenchmarkNote)
|
||
.filter(
|
||
BenchmarkNote.id.in_(ids),
|
||
BenchmarkNote.workspace_id == current_user.workspace_id,
|
||
BenchmarkNote.product_id == body.product_id,
|
||
)
|
||
.all()
|
||
)
|
||
by_id = {int(b.id): b for b in rows}
|
||
missing = [i for i in ids if i not in by_id]
|
||
if missing:
|
||
raise_business(f"标杆笔记不存在或不属于当前产品：{missing[:3]}")
|
||
not_ready = [
|
||
i for i in ids
|
||
if by_id[i].analyze_status != "done" or by_id[i].analysis_source == "fallback"
|
||
]
|
||
if not_ready:
|
||
raise_business("所选标杆尚未分析完成，请先到系统管理里点击「分析标杆」")
|
||
benchmark_ids_json = json.dumps([int(i) for i in _bids]) if _bids else None
|
||
|
||
task = GenerationTask(
|
||
workspace_id=current_user.workspace_id,
|
||
product_id=body.product_id,
|
||
operator_id=current_user.user_id,
|
||
theme=body.theme,
|
||
text_count=body.text_count,
|
||
image_count=body.image_count,
|
||
track=body.track,
|
||
need_product_image=need_img,
|
||
benchmark_ids=benchmark_ids_json,
|
||
status="pending",
|
||
)
|
||
db.add(task)
|
||
db.commit()
|
||
db.refresh(task)
|
||
logger.info("GenerationTask created: id=%s ws=%s", task.id, current_user.workspace_id)
|
||
|
||
if body.track == "ai":
|
||
enqueue_generation(task.id)
|
||
|
||
return task
|
||
|
||
|
||
def enqueue_generation(task_id: int, regen_strategy: str | None = None,
|
||
regen_role: str | None = None, custom_prompt: str | None = None,
|
||
reuse_text: bool = False) -> None:
|
||
"""只推 task_id（+可选重生参数）入队，绝不推 key（基石B）。
|
||
regen_strategy/regen_role/custom_prompt 仅 R2 单张/单套重生时传，常规生成留 None。
|
||
reuse_text=True 仅导入轨「去生图」时传：首次出图但复用库内导入文案，跳过文案生成。"""
|
||
from app.workers.tasks import run_generation_pipeline
|
||
run_generation_pipeline.delay(task_id, regen_strategy=regen_strategy,
|
||
regen_role=regen_role, custom_prompt=custom_prompt,
|
||
reuse_text=reuse_text)
|
||
logger.info("Enqueued task_id=%s regen=%s/%s reuse_text=%s",
|
||
task_id, regen_strategy, regen_role, reuse_text)
|