From c24d00dc80e9e8ba028097edc45b61d239694ddd Mon Sep 17 00:00:00 2001 From: xiaoxia Date: Fri, 18 Sep 2026 20:35:40 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20#1970=20=E7=89=87=E6=AE=B5=E7=BA=A7=20A?= =?UTF-8?q?I=20=E6=A0=87=E7=AD=BE=20+=20=E5=8F=99=E4=BA=8B=E5=8A=A0?= =?UTF-8?q?=E6=9D=83=E5=8C=B9=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- alembic/versions/081_atom_clip_ai_tags.py | 26 ++ apps/worker/worker_app/tasks/__init__.py | 8 + .../worker_app/tasks/atom_clip_tagging.py | 93 ++++++ apps/worker/worker_app/tasks/atom_clips.py | 28 ++ .../tasks/backfill_atom_clip_tags.py | 102 ++++++ .../asset_atom_clip_repository.py | 20 ++ packages/adapters/sqlalchemy_impl/models.py | 1 + packages/config/base.py | 1 + packages/domain/asset_atom_clip.py | 1 + packages/domain/atom_clip_tagger.py | 292 +++++++++++++++++ packages/domain/narrative_match.py | 138 +++++++- packages/shared/ai_client.py | 94 ++++++ tests/unit/test_1970_atom_clip_tagger.py | 292 +++++++++++++++++ .../unit/test_1970_narrative_clip_weighted.py | 306 ++++++++++++++++++ 14 files changed, 1397 insertions(+), 5 deletions(-) create mode 100644 alembic/versions/081_atom_clip_ai_tags.py create mode 100644 apps/worker/worker_app/tasks/atom_clip_tagging.py create mode 100644 apps/worker/worker_app/tasks/backfill_atom_clip_tags.py create mode 100644 packages/domain/atom_clip_tagger.py create mode 100644 tests/unit/test_1970_atom_clip_tagger.py create mode 100644 tests/unit/test_1970_narrative_clip_weighted.py diff --git a/alembic/versions/081_atom_clip_ai_tags.py b/alembic/versions/081_atom_clip_ai_tags.py new file mode 100644 index 000000000..09e7a1e69 --- /dev/null +++ b/alembic/versions/081_atom_clip_ai_tags.py @@ -0,0 +1,26 @@ +"""add ai_tags to asset_atom_clips for #1970 fragment-level AI tagging + +Revision ID: 081_atom_clip_ai_tags +Revises: 080_edit_plan_clips_atom_clip_id +Create Date: 2026-09-18 +""" + +import sqlalchemy as sa + +from alembic import op + +revision = "081_atom_clip_ai_tags" +down_revision = "080_edit_plan_clips_atom_clip_id" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column( + "asset_atom_clips", + sa.Column("ai_tags", sa.JSON(), nullable=True), + ) + + +def downgrade() -> None: + op.drop_column("asset_atom_clips", "ai_tags") diff --git a/apps/worker/worker_app/tasks/__init__.py b/apps/worker/worker_app/tasks/__init__.py index e2a03307c..f4bade5ab 100755 --- a/apps/worker/worker_app/tasks/__init__.py +++ b/apps/worker/worker_app/tasks/__init__.py @@ -57,6 +57,14 @@ def __getattr__(name: str): from .atom_clips import generate_atom_clips return generate_atom_clips + elif name == "tag_atom_clip_task": + from .atom_clip_tagging import tag_atom_clip_task + + return tag_atom_clip_task + elif name == "backfill_atom_clip_tags": + from .backfill_atom_clip_tags import backfill_atom_clip_tags + + return backfill_atom_clip_tags raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/apps/worker/worker_app/tasks/atom_clip_tagging.py b/apps/worker/worker_app/tasks/atom_clip_tagging.py new file mode 100644 index 000000000..e03ecb5c4 --- /dev/null +++ b/apps/worker/worker_app/tasks/atom_clip_tagging.py @@ -0,0 +1,93 @@ +"""片段级 AI 标签 Celery 任务 — #1970 智能剪辑流程重构 P2. + +为单个 atom_clip 调用视觉 AI 生成结构化标签，并更新到 ai_tags 字段。 +失败不阻断流程（降级为仅继承素材标签）。 + +任务名：worker.tag_atom_clip +""" + +from __future__ import annotations + +from celery.utils.log import get_task_logger +from worker_app.celery_app import celery_app +from worker_app.db import SessionLocal + +from packages.adapters.sqlalchemy_impl.asset_atom_clip_repository import ( + SQLAlchemyAssetAtomClipRepository, +) +from packages.adapters.sqlalchemy_impl.asset_repository import SQLAlchemyAssetRepository +from packages.domain.atom_clip_tagger import tag_atom_clip +from packages.shared.ai_client import get_doubao_client +from packages.shared.mediakit_client import get_mediakit_client +from packages.shared.storage import get_shared_storage_service + +logger = get_task_logger(__name__) + + +@celery_app.task(name="worker.tag_atom_clip", bind=True, max_retries=2, default_retry_delay=10) +def tag_atom_clip_task(self, atom_clip_id: str) -> dict: + """为单个原子片段生成 AI 标签. + + Args: + atom_clip_id: 原子片段 ID。 + + Returns: + 任务结果 dict：status / clip_id / ai_tags（部分字段）。 + """ + db = SessionLocal() + try: + atom_repo = SQLAlchemyAssetAtomClipRepository(db) + asset_repo = SQLAlchemyAssetRepository(db) + + clip = atom_repo.find_by_id(atom_clip_id) + if clip is None: + return {"status": "skipped", "reason": "clip not found", "clip_id": atom_clip_id} + + # 已有标签则跳过（幂等） + if clip.ai_tags is not None: + return {"status": "skipped", "reason": "already tagged", "clip_id": atom_clip_id} + + # 获取素材信息 + asset = asset_repo.find_by_id(clip.asset_id) + if asset is None: + return {"status": "skipped", "reason": "asset not found", "clip_id": atom_clip_id} + + # 获取视频可访问 URL + storage = get_shared_storage_service() + video_url = storage.get_download_url(asset.storage_key, expires_seconds=3600) + + # 初始化客户端 + doubao_client = get_doubao_client() + mediakit_client = get_mediakit_client() + + # 调用 tagger + ai_tags = tag_atom_clip( + clip=clip, + video_url=video_url, + doubao_client=doubao_client, + mediakit_client=mediakit_client, + storage=storage, + ) + + # 更新数据库 + atom_repo.update_ai_tags(atom_clip_id, ai_tags) + + logger.info( + "[atom_clip_tagging] clip_id=%s ai_tags=%s", + atom_clip_id, + {k: v for k, v in ai_tags.items() if k != "inherited_tags"}, + ) + return { + "status": "completed", + "clip_id": atom_clip_id, + "has_ai_tags": any(v for k, v in ai_tags.items() if k != "inherited_tags" and v), + } + except Exception as exc: + db.rollback() + logger.exception("[atom_clip_tagging] clip_id=%s 失败: %s", atom_clip_id, exc) + # 可重试异常 + if self.request.retries < self.max_retries: + raise self.retry(exc=exc) from None + return {"status": "failed", "clip_id": atom_clip_id, "error": str(exc)} + finally: + db.close() diff --git a/apps/worker/worker_app/tasks/atom_clips.py b/apps/worker/worker_app/tasks/atom_clips.py index e18c3a7c0..eb556e66d 100644 --- a/apps/worker/worker_app/tasks/atom_clips.py +++ b/apps/worker/worker_app/tasks/atom_clips.py @@ -3,6 +3,8 @@ 素材入库预处理完成（ingest 置 READY）后异步触发： 根据素材时长和已缓存的 scdet 切换点计算原子片段并落库。 失败不阻断素材入库主流程（atom_clips 未就绪时选片有内存兜底）。 + +P2 增强：切片完成后自动链式触发 AI 标签任务（每个 clip 一个 tag_atom_clip 任务）。 """ from __future__ import annotations @@ -72,6 +74,10 @@ def generate_atom_clips(asset_id: str) -> dict: asset_id, len(clips), ) + + # P2 增强：链式触发 AI 标签任务（每个 clip 一个异步任务） + _dispatch_tagging_tasks(clips) + return {"status": "completed", "asset_id": asset_id, "clips_count": len(clips)} except Exception as exc: # noqa: BLE001 - 后台任务兜底，失败不阻断主流程 db.rollback() @@ -79,3 +85,25 @@ def generate_atom_clips(asset_id: str) -> dict: return {"status": "failed", "asset_id": asset_id, "error": str(exc)} finally: db.close() + + +def _dispatch_tagging_tasks(clips: list) -> None: + """为每个新建片段发送 AI 标签异步任务. + + 失败不阻断（标签任务是锦上添花，不影响核心流程）。 + """ + try: + for clip in clips: + celery_app.send_task( + "worker.tag_atom_clip", + args=[clip.id], + ) + logger.info( + "[atom_clips] 已发送 %d 个 AI 标签任务", + len(clips), + ) + except Exception as e: + logger.warning( + "[atom_clips] 发送 AI 标签任务失败（不影响切片结果）: %s", + e, + ) diff --git a/apps/worker/worker_app/tasks/backfill_atom_clip_tags.py b/apps/worker/worker_app/tasks/backfill_atom_clip_tags.py new file mode 100644 index 000000000..49399ae4b --- /dev/null +++ b/apps/worker/worker_app/tasks/backfill_atom_clip_tags.py @@ -0,0 +1,102 @@ +"""批量回填 AI 标签 Celery 任务 — #1970 智能剪辑流程重构 P2. + +查找所有 ai_tags IS NULL 的 atom_clips，分批触发 tag_atom_clip 任务。 +可通过 API 路由触发（管理员权限）。 + +任务名：worker.backfill_atom_clip_tags +""" + +from __future__ import annotations + +import time + +from celery.utils.log import get_task_logger +from worker_app.celery_app import celery_app +from worker_app.db import SessionLocal + +from packages.adapters.sqlalchemy_impl.asset_atom_clip_repository import ( + SQLAlchemyAssetAtomClipRepository, +) + +logger = get_task_logger(__name__) + +# 默认批量参数 +DEFAULT_BATCH_SIZE = 10 +DEFAULT_BATCH_INTERVAL = 5 # 秒 + + +@celery_app.task(name="worker.backfill_atom_clip_tags") +def backfill_atom_clip_tags( + batch_size: int = DEFAULT_BATCH_SIZE, + batch_interval: int = DEFAULT_BATCH_INTERVAL, + max_clips: int = 0, +) -> dict: + """批量回填未打标的 atom_clips. + + Args: + batch_size: 每批处理数量，默认 10。 + batch_interval: 每批间隔秒数，默认 5。 + max_clips: 最大处理总数，0 表示不限。 + + Returns: + 任务结果 dict：total_submitted / batches。 + """ + db = SessionLocal() + try: + atom_repo = SQLAlchemyAssetAtomClipRepository(db) + total_submitted = 0 + batches = 0 + + while True: + # 查找未打标的片段 + remaining = max_clips - total_submitted if max_clips > 0 else batch_size + fetch_limit = min(batch_size, remaining) if max_clips > 0 else batch_size + + untagged = atom_repo.find_untagged(limit=fetch_limit) + if not untagged: + break + + # 逐个发送 tag 任务 + for clip in untagged: + try: + celery_app.send_task( + "worker.tag_atom_clip", + args=[clip.id], + ) + total_submitted += 1 + except Exception as e: + logger.warning( + "[backfill] 提交任务失败 clip_id=%s: %s", + clip.id, + e, + ) + + batches += 1 + logger.info( + "[backfill] 第 %d 批完成，已提交 %d 个任务", + batches, + total_submitted, + ) + + # 检查是否达到上限 + if max_clips > 0 and total_submitted >= max_clips: + break + + # 批间间隔 + time.sleep(batch_interval) + + logger.info( + "[backfill] 回填完成: total_submitted=%d batches=%d", + total_submitted, + batches, + ) + return { + "status": "completed", + "total_submitted": total_submitted, + "batches": batches, + } + except Exception as exc: + logger.exception("[backfill] 回填失败: %s", exc) + return {"status": "failed", "error": str(exc)} + finally: + db.close() diff --git a/packages/adapters/sqlalchemy_impl/asset_atom_clip_repository.py b/packages/adapters/sqlalchemy_impl/asset_atom_clip_repository.py index 81f8cefbb..8133149f7 100644 --- a/packages/adapters/sqlalchemy_impl/asset_atom_clip_repository.py +++ b/packages/adapters/sqlalchemy_impl/asset_atom_clip_repository.py @@ -83,6 +83,25 @@ class SQLAlchemyAssetAtomClipRepository: models = query.all() return [self._to_domain(m) for m in models] + def update_ai_tags(self, clip_id: str, ai_tags: dict) -> bool: + """更新指定片段的 ai_tags 字段.""" + count = ( + self.session.query(AssetAtomClipModel).filter(AssetAtomClipModel.id == clip_id).update({"ai_tags": ai_tags}) + ) + self.session.commit() + return count > 0 + + def find_untagged(self, limit: int = 100) -> list[AssetAtomClip]: + """查找 ai_tags IS NULL 的片段，用于回填.""" + models = ( + self.session.query(AssetAtomClipModel) + .filter(AssetAtomClipModel.ai_tags.is_(None)) + .order_by(AssetAtomClipModel.created_at.asc()) + .limit(limit) + .all() + ) + return [self._to_domain(m) for m in models] + def _to_model(self, clip: AssetAtomClip) -> AssetAtomClipModel: return AssetAtomClipModel( id=clip.id, @@ -92,6 +111,7 @@ class SQLAlchemyAssetAtomClipRepository: duration=clip.duration, clip_index=clip.clip_index, tags=clip.tags, + ai_tags=clip.ai_tags, scene_change_at=clip.scene_change_at, is_fallback=clip.is_fallback, created_at=clip.created_at or datetime.now(UTC), diff --git a/packages/adapters/sqlalchemy_impl/models.py b/packages/adapters/sqlalchemy_impl/models.py index 5285a5d82..ed6bf04b7 100755 --- a/packages/adapters/sqlalchemy_impl/models.py +++ b/packages/adapters/sqlalchemy_impl/models.py @@ -837,6 +837,7 @@ class AssetAtomClipModel(Base): duration = Column(Float, nullable=False) clip_index = Column(Integer, nullable=False) tags = Column(JSON, nullable=False, default=list) + ai_tags = Column(JSON, nullable=True, default=None) scene_change_at = Column(Float, nullable=True) is_fallback = Column(Boolean, nullable=False, default=False) created_at = Column(DateTime, nullable=False, default=lambda: datetime.now(UTC)) diff --git a/packages/config/base.py b/packages/config/base.py index f56b5e2ed..de5ec1ada 100755 --- a/packages/config/base.py +++ b/packages/config/base.py @@ -68,6 +68,7 @@ class SharedSettings(BaseSettings): doubao_base_url: str = "https://ark.cn-beijing.volces.com/api/v3" doubao_timeout: int = 30 doubao_max_retries: int = 2 + doubao_vision_model: str = "doubao-1-5-vision-pro-250915" # ── MediaKit (火山引擎 AI 媒体工具) ────────────────────────────────── mediakit_api_key: str = "" diff --git a/packages/domain/asset_atom_clip.py b/packages/domain/asset_atom_clip.py index 404d7436b..86cd164ec 100644 --- a/packages/domain/asset_atom_clip.py +++ b/packages/domain/asset_atom_clip.py @@ -36,6 +36,7 @@ class AssetAtomClip: duration: float clip_index: int tags: list[str] = field(default_factory=list) + ai_tags: dict | None = None scene_change_at: float | None = None is_fallback: bool = False created_at: datetime | None = None diff --git a/packages/domain/atom_clip_tagger.py b/packages/domain/atom_clip_tagger.py new file mode 100644 index 000000000..f40dbd1a5 --- /dev/null +++ b/packages/domain/atom_clip_tagger.py @@ -0,0 +1,292 @@ +"""片段级 AI 标签 — #1970 智能剪辑流程重构 P2. + +对每个 atom_clip 提取关键帧，调用豆包视觉理解 API 识别内容， +生成结构化标签（场景、物体、动作、景别、是否有文字）。 + +纯函数 + IO 分离设计： +- build_vision_prompt() 返回结构化 prompt +- parse_vision_response(text) 解析 AI 返回的 JSON 标签 +- tag_atom_clip(...) 主入口，组合帧提取 → 视觉 API → 解析标签 + +降级策略：任何环节失败都返回 {"inherited_tags": clip.tags}，不阻断流程。 +""" + +from __future__ import annotations + +import json +import logging +import subprocess +import tempfile +from pathlib import Path +from typing import Any, Optional + +logger = logging.getLogger(__name__) + +# AI 标签结构的键 +AI_TAG_KEYS = ("scene", "objects", "action", "shot", "has_text") + + +def build_vision_prompt() -> str: + """返回结构化标签提取 prompt. + + 要求 AI 以 JSON 格式返回片段内容标签，包含： + - scene: 场景类型列表（如 "工厂", "办公室", "户外"） + - objects: 出现的物体列表（如 "产品", "手机", "电脑"） + - action: 动作类型列表（如 "演示", "说话", "操作"） + - shot: 景别（"特写" / "中景" / "远景" 之一） + - has_text: 画面中是否有显著文字（true/false） + """ + return """请分析这段视频片段的关键帧，识别内容并返回 JSON 格式标签。 + +要求返回以下 JSON 结构（严格 JSON，不要添加其他文字）： +{ + "scene": ["场景1", "场景2"], + "objects": ["物体1", "物体2"], + "action": ["动作1"], + "shot": "特写|中景|远景", + "has_text": true/false +} + +规则： +- scene: 场景类型，如"工厂"、"办公室"、"户外"、"商店"、"家庭"等，1-3个 +- objects: 画面中可见的主要物体，如"产品"、"手机"、"电脑"、"食品"等，1-5个 +- action: 人物或物体正在进行的动作，如"演示"、"说话"、"操作"、"展示"等，1-3个 +- shot: 景别判断，只能是"特写"、"中景"或"远景"之一 +- has_text: 画面中是否有显著可读文字（标题、字幕、标语等） + +请只返回 JSON，不要有其他说明文字。""" + + +def parse_vision_response(text: str) -> dict: + """解析 AI 返回的 JSON 标签文本. + + Args: + text: 视觉 API 返回的文本，期望是 JSON 格式。 + + Returns: + 结构化标签 dict，格式如： + {"scene": [...], "objects": [...], "action": [...], "shot": "...", "has_text": bool} + + 解析失败时返回空 dict。 + """ + if not text or not text.strip(): + return {} + + # 尝试直接解析 + cleaned = text.strip() + + # 去除可能的 markdown 代码块包裹 + if cleaned.startswith("```"): + lines = cleaned.split("\n") + # 去掉首尾的 ``` 行 + start = 1 + end = len(lines) + for i in range(len(lines) - 1, 0, -1): + if lines[i].strip().startswith("```"): + end = i + break + cleaned = "\n".join(lines[start:end]).strip() + + try: + data = json.loads(cleaned) + except json.JSONDecodeError: + # 尝试从文本中提取 JSON 块 + try: + start_idx = cleaned.index("{") + end_idx = cleaned.rindex("}") + 1 + data = json.loads(cleaned[start_idx:end_idx]) + except (ValueError, json.JSONDecodeError): + logger.warning("无法解析 AI 标签响应: %s", text[:200]) + return {} + + if not isinstance(data, dict): + return {} + + # 验证和清洗各字段 + result: dict[str, Any] = {} + for key in ("scene", "objects", "action"): + val = data.get(key) + if isinstance(val, list): + result[key] = [str(v).strip() for v in val if str(v).strip()] + elif isinstance(val, str) and val.strip(): + result[key] = [val.strip()] + else: + result[key] = [] + + shot_val = data.get("shot", "") + if isinstance(shot_val, str) and shot_val.strip() in ("特写", "中景", "远景"): + result["shot"] = shot_val.strip() + else: + result["shot"] = "" + + has_text_val = data.get("has_text") + if isinstance(has_text_val, bool): + result["has_text"] = has_text_val + elif isinstance(has_text_val, str): + result["has_text"] = has_text_val.lower() in ("true", "yes", "1") + else: + result["has_text"] = False + + return result + + +def _extract_frames_via_mediakit( + mediakit_client: Any, + video_url: str, + start_time: float, + end_time: float, +) -> Optional[list[str]]: + """通过 MediaKit 提取 3 帧（首、中、尾）. + + Returns: + 图片 URL 列表（3 个），失败返回 None。 + """ + try: + frames = mediakit_client.extract_frames( + video_url=video_url, + strategy="SpecifiedTime", + max_frames=3, + poll_interval=2.0, + max_poll_attempts=30, + ) + # MediaKit SpecifiedTime 策略可能不支持直接传时间点 + # 如果返回结果不够 3 帧，降级到 ffmpeg + if frames and len(frames) >= 1: + urls = [f.get("image_url", "") for f in frames if f.get("image_url")] + if urls: + return urls + except Exception as e: + logger.warning("MediaKit 抽帧失败，将降级为 ffmpeg: %s", e) + + return None + + +def _extract_frames_via_ffmpeg( + video_url: str, + start_time: float, + end_time: float, +) -> Optional[list[str]]: + """通过 ffmpeg 本地提取 3 帧并转为 base64. + + Returns: + base64 data URI 列表（3 个），失败返回 None。 + """ + import base64 + + mid_time = round((start_time + end_time) / 2, 3) + timestamps = [round(start_time, 3), mid_time, round(end_time, 3)] + + try: + frames_b64: list[str] = [] + with tempfile.TemporaryDirectory() as tmpdir: + for i, ts in enumerate(timestamps): + out_path = Path(tmpdir) / f"frame_{i}.jpg" + cmd = [ + "ffmpeg", + "-y", + "-ss", + str(ts), + "-i", + video_url, + "-vframes", + "1", + "-q:v", + "2", + str(out_path), + ] + result = subprocess.run( + cmd, + capture_output=True, + timeout=30, + ) + if result.returncode != 0 or not out_path.exists(): + logger.warning("ffmpeg 抽帧失败 ts=%s: %s", ts, result.stderr[:200]) + continue + + img_data = out_path.read_bytes() + b64 = base64.b64encode(img_data).decode("ascii") + frames_b64.append(f"data:image/jpeg;base64,{b64}") + + if frames_b64: + return frames_b64 + except Exception as e: + logger.warning("ffmpeg 抽帧异常: %s", e) + + return None + + +def tag_atom_clip( + clip: Any, + video_url: str, + doubao_client: Any, + mediakit_client: Any | None = None, + storage: Any | None = None, +) -> dict: + """主入口：为单个 atom_clip 生成 AI 标签. + + 流程：提取帧 → 调视觉 API → 解析标签 → 返回结构化标签 dict。 + 任何环节失败返回 {"inherited_tags": clip.tags}，不阻断流程。 + + Args: + clip: AssetAtomClip 领域对象（需有 start_time, end_time, tags）。 + video_url: 素材视频的公网可访问 URL。 + doubao_client: DoubaoClient 实例。 + mediakit_client: MediaKitClient 实例（可选，不可用时降级 ffmpeg）。 + storage: SharedStorageService 实例（可选，用于获取签名 URL）。 + + Returns: + 结构化标签 dict，格式如： + {"scene": [...], "objects": [...], "action": [...], "shot": "...", + "has_text": bool, "inherited_tags": [...]} + """ + inherited = list(getattr(clip, "tags", []) or []) + + # 检查 DoubaoClient 是否可用 + if not getattr(doubao_client, "is_available", False): + logger.info("DoubaoClient 不可用，跳过 AI 标签: clip_id=%s", getattr(clip, "id", "")) + return {"inherited_tags": inherited} + + # 提取帧图片 + frame_urls: Optional[list[str]] = None + start_time = getattr(clip, "start_time", 0.0) + end_time = getattr(clip, "end_time", 0.0) + + # 优先使用 MediaKit + if mediakit_client and getattr(mediakit_client, "is_available", False): + frame_urls = _extract_frames_via_mediakit(mediakit_client, video_url, start_time, end_time) + + # MediaKit 不可用或失败 → 降级 ffmpeg + if not frame_urls: + frame_urls = _extract_frames_via_ffmpeg(video_url, start_time, end_time) + + if not frame_urls: + logger.warning("帧提取失败，跳过 AI 标签: clip_id=%s", getattr(clip, "id", "")) + return {"inherited_tags": inherited} + + # 调用视觉 API + prompt = build_vision_prompt() + messages = [{"role": "user", "content": prompt}] + + try: + response_text = doubao_client.vision_completion( + messages=messages, + images=frame_urls, + timeout=60, + ) + except Exception as e: + logger.warning("视觉 API 调用异常: clip_id=%s error=%s", getattr(clip, "id", ""), e) + return {"inherited_tags": inherited} + + if not response_text: + logger.warning("视觉 API 返回空: clip_id=%s", getattr(clip, "id", "")) + return {"inherited_tags": inherited} + + # 解析标签 + ai_tags = parse_vision_response(response_text) + if not ai_tags: + logger.warning("标签解析失败: clip_id=%s response=%s", getattr(clip, "id", ""), response_text[:200]) + return {"inherited_tags": inherited} + + # 合并 inherited_tags + ai_tags["inherited_tags"] = inherited + return ai_tags diff --git a/packages/domain/narrative_match.py b/packages/domain/narrative_match.py index 5d72c6766..9a0a5b978 100644 --- a/packages/domain/narrative_match.py +++ b/packages/domain/narrative_match.py @@ -1,4 +1,4 @@ -"""叙事剪辑素材标签匹配 — #1970 PR3. +"""叙事剪辑素材标签匹配 — #1970 PR3 + P2 AI 标签加权. 叙事模式下，选片在现有评分（smart_match / atom_clip_selector）之前先做一层 文案标签匹配： @@ -8,6 +8,12 @@ - 调用方对优先池跑现有 smart_select_assets，数量不足时用普通池补足 （无任何匹配 → 完全降级为现有随机逻辑，行为与改造前一致）。 +P2 AI 标签加权（#1970 fragment-level AI tagging）： +- 片段级 AI 标签（scene/objects/action）与文案标签做交集时权重 2.0 +- 素材级标签（tag_ids 映射名）与文案标签交集时权重 1.0 +- 综合得分 = sum(命中权重) / max(可能权重) +- 有 AI 标签的片段命中时优先于仅素材标签命中的片段 + 纯函数模块：标签 id→名称映射由调用方查 TagModel 后注入，不直接碰 DB。 """ @@ -18,6 +24,10 @@ from typing import Any, Iterable # 标签归一化后仍短于此长度的标签不参与匹配（避免「的」「是」这类噪声短词） MIN_TAG_LEN = 2 +# 标签匹配权重 +AI_TAG_WEIGHT = 2.0 # AI 标签命中权重 +ASSET_TAG_WEIGHT = 1.0 # 素材标签命中权重 + def normalize_tag(tag: Any) -> str: """标签归一化：去空白、小写。数字/英文统一小写，中文不受影响。""" @@ -47,19 +57,81 @@ def build_asset_tag_name_index(tag_names_by_id: dict[str, Any]) -> dict[str, set return index +def _extract_ai_tag_names(ai_tags: dict) -> set[str]: + """从 AI 标签 dict 中提取所有标签名（scene + objects + action）. + + Args: + ai_tags: 片段级 AI 标签 dict，如 {"scene": [...], "objects": [...], "action": [...], ...} + + Returns: + 归一化后的标签名集合。 + """ + names: set[str] = set() + for key in ("scene", "objects", "action"): + values = ai_tags.get(key) + if isinstance(values, list): + names |= _normalize_tags(values) + return names + + +def _compute_ai_score( + asset_id: str, + wanted: set[str], + clip_ai_tags_by_asset: dict[str, list[dict]] | None, +) -> float: + """计算单个素材的 AI 标签加权得分. + + 对该素材的所有片段 AI 标签，求各片段标签名与文案标签交集的加权总和。 + 每个片段的命中权重 = 命中数 × AI_TAG_WEIGHT。 + 最终取所有片段的最高得分（而非累加，避免片段数多的素材不公平占优）。 + + Args: + asset_id: 素材 ID。 + wanted: 归一化后的文案标签集合。 + clip_ai_tags_by_asset: {asset_id: [ai_tag_dict, ...]} 每个片段一个。 + + Returns: + AI 标签加权得分（≥0）。 + """ + if not clip_ai_tags_by_asset or not wanted: + return 0.0 + + clips = clip_ai_tags_by_asset.get(asset_id) + if not clips: + return 0.0 + + best_score = 0.0 + for ai_tags in clips: + if not ai_tags or not isinstance(ai_tags, dict): + continue + ai_names = _extract_ai_tag_names(ai_tags) + hits = ai_names & wanted + score = len(hits) * AI_TAG_WEIGHT + if score > best_score: + best_score = score + + return best_score + + def match_assets_by_script_tags( assets: list[Any], *, script_tags: Iterable[Any], tag_names_by_id: dict[str, Any] | None = None, + clip_ai_tags_by_asset: dict[str, list[dict]] | None = None, ) -> tuple[list[Any], list[Any]]: """按文案标签把素材拆成「命中池 / 未命中池」，保持输入相对顺序。 + P2 加权逻辑： + - AI 标签命中（scene/objects/action ∩ 文案标签）权重 2.0 + - 素材标签命中（tag_ids 映射名 ∩ 文案标签）权重 1.0 + - 任一权重 > 0 → 命中池，否则 → 未命中池 + Args: assets: 候选素材（domain Asset，需有 id 与 tag_ids）。 script_tags: 文案 tags（字符串数组，名称语义）。 - tag_names_by_id: asset_id → 素材标签名列表；素材只有 tag_ids 时由调用方 - 查 TagModel 名称后传入。为空则视为无素材命中。 + tag_names_by_id: asset_id → 素材标签名列表。 + clip_ai_tags_by_asset: #1970 P2 — {asset_id: [ai_tag_dict, ...]}。 Returns: (matched, unmatched)：命中任一文案标签的素材 / 其余素材。 @@ -74,23 +146,74 @@ def match_assets_by_script_tags( unmatched: list[Any] = [] for asset in assets: asset_id = str(getattr(asset, "id", "") or "") + + # P2: AI 标签加权得分 + ai_score = _compute_ai_score(asset_id, wanted, clip_ai_tags_by_asset) + + # 素材标签得分 names = set(name_index.get(asset_id, set())) - # 兼容素材自身带字符串 tags（旧链路/测试替身） raw_tags = getattr(asset, "tags", None) if raw_tags: names |= _normalize_tags(raw_tags) - if names & wanted: + asset_score = len(names & wanted) * ASSET_TAG_WEIGHT + + # 综合得分 > 0 → 命中池 + if ai_score > 0 or asset_score > 0: matched.append(asset) else: unmatched.append(asset) return matched, unmatched +def compute_tag_match_score( + asset_id: str, + *, + script_tags: Iterable[Any], + tag_names_by_id: dict[str, Any] | None = None, + clip_ai_tags_by_asset: dict[str, list[dict]] | None = None, +) -> float: + """计算单个素材的标签匹配综合得分（0.0 ~ 1.0）. + + 综合得分 = sum(命中权重) / max(可能权重) + - AI 标签每命中一个 +2.0 + - 素材标签每命中一个 +1.0 + - max_possible = len(wanted) * (AI_TAG_WEIGHT + ASSET_TAG_WEIGHT) + + Args: + asset_id: 素材 ID。 + script_tags: 文案标签。 + tag_names_by_id: 素材标签名索引。 + clip_ai_tags_by_asset: AI 标签索引。 + + Returns: + 归一化得分 0.0~1.0。 + """ + wanted = _normalize_tags(script_tags) + if not wanted: + return 0.0 + + # AI 得分 + ai_score = _compute_ai_score(asset_id, wanted, clip_ai_tags_by_asset) + + # 素材标签得分 + name_index = build_asset_tag_name_index(tag_names_by_id or {}) + names = name_index.get(asset_id, set()) + asset_score = len(names & wanted) * ASSET_TAG_WEIGHT + + # 归一化：最大可能得分 = 文案标签数 × (AI权重 + 素材权重) + max_possible = len(wanted) * (AI_TAG_WEIGHT + ASSET_TAG_WEIGHT) + if max_possible <= 0: + return 0.0 + + return min((ai_score + asset_score) / max_possible, 1.0) + + def pick_narrative_assets( assets: list[Any], *, script_tags: Iterable[Any], tag_names_by_id: dict[str, Any] | None = None, + clip_ai_tags_by_asset: dict[str, list[dict]] | None = None, limit: int | None = None, rng: Any = None, ) -> list[Any]: @@ -100,9 +223,13 @@ def pick_narrative_assets( smart_match.smart_select_assets（质量/时长/新鲜度/未使用 + 随机噪声）， 不重写评分维度。 + P2 增强：有 AI 标签的片段命中时权重更高（2.0 vs 1.0）， + 命中池内部按综合标签得分排序（AI 标签命中多的排前面）。 + Args: assets: ready 视频素材候选（调用方负责状态/类型过滤）。 script_tags / tag_names_by_id: 见 match_assets_by_script_tags。 + clip_ai_tags_by_asset: #1970 P2 — {asset_id: [ai_tag_dict, ...]}。 limit: 需要的素材数量；None 表示全部（命中池 + 全部未命中池）。 rng: 注入 smart_select_assets 的随机源（可复现）。 @@ -115,6 +242,7 @@ def pick_narrative_assets( assets, script_tags=script_tags, tag_names_by_id=tag_names_by_id, + clip_ai_tags_by_asset=clip_ai_tags_by_asset, ) need = limit if (limit is not None and limit > 0) else None diff --git a/packages/shared/ai_client.py b/packages/shared/ai_client.py index 249c4b81d..5467013f3 100755 --- a/packages/shared/ai_client.py +++ b/packages/shared/ai_client.py @@ -37,6 +37,7 @@ class DoubaoClient: self.base_url: str = settings.doubao_base_url.rstrip("/") self.timeout: int = settings.doubao_timeout self.max_retries: int = settings.doubao_max_retries + self.vision_model: str = settings.doubao_vision_model @property def is_available(self) -> bool: @@ -103,6 +104,99 @@ class DoubaoClient: logger.error("豆包API调用最终失败: %s", last_error) return None + def vision_completion( + self, + messages: list[dict], + images: list[str] | None = None, + max_tokens: int = 2048, + temperature: float = 0.3, + timeout: int | None = None, + ) -> Optional[str]: + """调用豆包视觉理解 API（OpenAI 兼容多模态格式）. + + 将 images 附加到最后一条 user message 的 content 中， + 使用 vision_model（默认 doubao-1-5-vision-pro-250915）。 + + Args: + messages: 对话消息列表。最后一条 user message 会被注入图片内容。 + images: 图片列表，支持 base64 data URI 或 HTTP(S) URL。 + max_tokens: 最大生成 token 数，默认 2048。 + temperature: 采样温度，默认 0.3（视觉任务偏低更稳定）。 + timeout: 单次请求超时秒数，不传则使用默认 self.timeout。 + + Returns: + 模型返回的文本内容，失败返回 None。 + """ + if not self.is_available: + return None + + # 构造多模态 content：先追加文本，再追加图片 + vision_messages = [] + for msg in messages: + vision_messages.append(dict(msg)) + + # 将图片注入最后一条 user message + if images and vision_messages: + # 找到最后一条 user message + for i in range(len(vision_messages) - 1, -1, -1): + if vision_messages[i].get("role") == "user": + text_content = vision_messages[i].get("content", "") + multi_content: list[dict[str, Any]] = [] + if text_content: + multi_content.append({"type": "text", "text": text_content}) + for img in images: + if img.startswith("data:") or img.startswith("http://") or img.startswith("https://"): + multi_content.append({"type": "image_url", "image_url": {"url": img}}) + else: + # 当作 base64 编码 + multi_content.append( + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img}"}} + ) + vision_messages[i]["content"] = multi_content + break + + url = f"{self.base_url}/chat/completions" + headers = { + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + } + payload: dict[str, Any] = { + "model": self.vision_model, + "messages": vision_messages, + "temperature": temperature, + "max_tokens": max_tokens, + } + + req_timeout = timeout or self.timeout + last_error: Optional[Exception] = None + for attempt in range(self.max_retries + 1): + try: + response = httpx.post( + url, + headers=headers, + json=payload, + timeout=req_timeout, + ) + response.raise_for_status() + data = response.json() + content = data["choices"][0]["message"]["content"] + return content.strip() + except Exception as e: + last_error = e + if attempt < self.max_retries: + wait = 0.5 * (2**attempt) + logger.warning( + "豆包视觉API调用失败，%.1fs后重试 (第%d/%d次): %s", + wait, + attempt + 1, + self.max_retries + 1, + e, + ) + time.sleep(wait) + + logger.error("豆包视觉API调用最终失败: %s", last_error) + return None + # ── 单例 ───────────────────────────────────────────────────────────────────── diff --git a/tests/unit/test_1970_atom_clip_tagger.py b/tests/unit/test_1970_atom_clip_tagger.py new file mode 100644 index 000000000..5b95107b9 --- /dev/null +++ b/tests/unit/test_1970_atom_clip_tagger.py @@ -0,0 +1,292 @@ +"""#1970 P2 片段级 AI 标签模块测试。 + +测试范围： +- build_vision_prompt: 返回有效 prompt +- parse_vision_response: 正常/异常/空值 +- tag_atom_clip: 成功/MediaKit不可用/视觉API失败/超时降级 +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from datetime import UTC, datetime + +import pytest + +from packages.domain.atom_clip_tagger import ( + build_vision_prompt, + parse_vision_response, + tag_atom_clip, +) + +# ── Fake 对象 ────────────────────────────────────────────────────────────── + + +@dataclass +class FakeClip: + id: str = "clip-001" + asset_id: str = "asset-001" + start_time: float = 0.0 + end_time: float = 5.0 + duration: float = 5.0 + clip_index: int = 0 + tags: list[str] = field(default_factory=lambda: ["tag1", "tag2"]) + ai_tags: dict | None = None + + +class FakeDoubaoClient: + """模拟豆包客户端.""" + + def __init__(self, available: bool = True, response: str | None = None, raise_error: bool = False): + self._available = available + self._response = response + self._raise_error = raise_error + self.vision_calls: list[dict] = [] + + @property + def is_available(self) -> bool: + return self._available + + def vision_completion(self, messages, images=None, timeout=None, **kwargs): + self.vision_calls.append({"messages": messages, "images": images, "timeout": timeout}) + if self._raise_error: + raise RuntimeError("API error") + return self._response + + +class FakeMediaKitClient: + """模拟 MediaKit 客户端.""" + + def __init__(self, available: bool = True, frames: list[dict] | None = None): + self._available = available + self._frames = frames + + @property + def is_available(self) -> bool: + return self._available + + def extract_frames(self, video_url, strategy=None, max_frames=None, **kwargs): + return self._frames + + +# ── build_vision_prompt ──────────────────────────────────────────────────── + + +class TestBuildVisionPrompt: + def test_returns_non_empty_string(self): + prompt = build_vision_prompt() + assert isinstance(prompt, str) + assert len(prompt) > 100 + + def test_contains_required_keys(self): + prompt = build_vision_prompt() + assert "scene" in prompt + assert "objects" in prompt + assert "action" in prompt + assert "shot" in prompt + assert "has_text" in prompt + + def test_requests_json_format(self): + prompt = build_vision_prompt() + assert "JSON" in prompt or "json" in prompt + + +# ── parse_vision_response ────────────────────────────────────────────────── + + +class TestParseVisionResponse: + def test_valid_json(self): + response = json.dumps( + { + "scene": ["工厂", "车间"], + "objects": ["产品", "机器"], + "action": ["演示"], + "shot": "特写", + "has_text": True, + } + ) + result = parse_vision_response(response) + assert result["scene"] == ["工厂", "车间"] + assert result["objects"] == ["产品", "机器"] + assert result["action"] == ["演示"] + assert result["shot"] == "特写" + assert result["has_text"] is True + + def test_json_with_markdown_code_block(self): + response = '```json\n{"scene": ["办公室"], "objects": ["电脑"], "action": ["说话"], "shot": "中景", "has_text": false}\n```' + result = parse_vision_response(response) + assert result["scene"] == ["办公室"] + assert result["has_text"] is False + + def test_json_embedded_in_text(self): + response = '这是一些说明文字\n{"scene": ["户外"], "objects": ["汽车"], "action": ["展示"], "shot": "远景", "has_text": false}\n结束' + result = parse_vision_response(response) + assert result["scene"] == ["户外"] + + def test_empty_response(self): + assert parse_vision_response("") == {} + assert parse_vision_response(None) == {} + assert parse_vision_response(" ") == {} + + def test_invalid_json(self): + assert parse_vision_response("这不是JSON") == {} + + def test_partial_fields(self): + response = json.dumps({"scene": ["工厂"]}) + result = parse_vision_response(response) + assert result["scene"] == ["工厂"] + assert result["objects"] == [] + assert result["shot"] == "" + assert result["has_text"] is False + + def test_invalid_shot_value(self): + response = json.dumps({"scene": [], "objects": [], "action": [], "shot": "全景", "has_text": False}) + result = parse_vision_response(response) + # "全景" 不在有效值 ("特写", "中景", "远景") 中 + assert result["shot"] == "" + + def test_string_values_converted_to_list(self): + response = json.dumps( + {"scene": "工厂", "objects": "产品", "action": "演示", "shot": "特写", "has_text": "true"} + ) + result = parse_vision_response(response) + assert result["scene"] == ["工厂"] + assert result["objects"] == ["产品"] + assert result["has_text"] is True + + def test_non_dict_json(self): + assert parse_vision_response("[1, 2, 3]") == {} + assert parse_vision_response('"hello"') == {} + + +# ── tag_atom_clip ────────────────────────────────────────────────────────── + + +class TestTagAtomClip: + def test_success_with_mediakit(self): + """MediaKit 可用 + 视觉 API 成功 → 返回完整 AI 标签.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient( + response=json.dumps( + { + "scene": ["工厂"], + "objects": ["产品"], + "action": ["演示"], + "shot": "特写", + "has_text": False, + } + ) + ) + fake_mediakit = FakeMediaKitClient( + frames=[ + {"image_url": "https://example.com/frame1.jpg", "timestamp": 0.0}, + {"image_url": "https://example.com/frame2.jpg", "timestamp": 2.5}, + {"image_url": "https://example.com/frame3.jpg", "timestamp": 5.0}, + ] + ) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + mediakit_client=fake_mediakit, + ) + + assert result["scene"] == ["工厂"] + assert result["objects"] == ["产品"] + assert result["shot"] == "特写" + assert result["inherited_tags"] == ["tag1", "tag2"] + assert len(fake_doubao.vision_calls) == 1 + + def test_doubao_unavailable_returns_inherited(self): + """DoubaoClient 不可用 → 返回 inherited_tags.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient(available=False) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + ) + + assert result == {"inherited_tags": ["tag1", "tag2"]} + assert len(fake_doubao.vision_calls) == 0 + + def test_mediakit_unavailable_no_ffmpeg(self): + """MediaKit 不可用 + 无 ffmpeg → 降级 inherited_tags.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient() + fake_mediakit = FakeMediaKitClient(available=False) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + mediakit_client=fake_mediakit, + ) + + # 没有 ffmpeg 的情况下，帧提取失败 + assert result == {"inherited_tags": ["tag1", "tag2"]} + + def test_vision_api_error_returns_inherited(self): + """视觉 API 抛异常 → 降级 inherited_tags.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient(raise_error=True) + fake_mediakit = FakeMediaKitClient(frames=[{"image_url": "https://example.com/frame.jpg", "timestamp": 0.0}]) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + mediakit_client=fake_mediakit, + ) + + assert result == {"inherited_tags": ["tag1", "tag2"]} + + def test_vision_api_empty_response(self): + """视觉 API 返回空 → 降级 inherited_tags.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient(response=None) + fake_mediakit = FakeMediaKitClient(frames=[{"image_url": "https://example.com/frame.jpg", "timestamp": 0.0}]) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + mediakit_client=fake_mediakit, + ) + + assert result == {"inherited_tags": ["tag1", "tag2"]} + + def test_vision_api_invalid_json_response(self): + """视觉 API 返回无效 JSON → 降级 inherited_tags.""" + clip = FakeClip() + fake_doubao = FakeDoubaoClient(response="这不是JSON格式") + fake_mediakit = FakeMediaKitClient(frames=[{"image_url": "https://example.com/frame.jpg", "timestamp": 0.0}]) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + mediakit_client=fake_mediakit, + ) + + assert result == {"inherited_tags": ["tag1", "tag2"]} + + def test_clip_with_empty_tags(self): + """空素材标签 → inherited_tags 为空列表.""" + clip = FakeClip(tags=[]) + fake_doubao = FakeDoubaoClient(available=False) + + result = tag_atom_clip( + clip=clip, + video_url="https://example.com/video.mp4", + doubao_client=fake_doubao, + ) + + assert result == {"inherited_tags": []} + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) diff --git a/tests/unit/test_1970_narrative_clip_weighted.py b/tests/unit/test_1970_narrative_clip_weighted.py new file mode 100644 index 000000000..380c0aef8 --- /dev/null +++ b/tests/unit/test_1970_narrative_clip_weighted.py @@ -0,0 +1,306 @@ +"""#1970 P2 叙事匹配 AI 标签加权测试。 + +测试范围： +- AI 标签命中时权重 2.0 +- 无 AI 标签时降级到素材标签权重 1.0 +- 混合场景（部分素材有 AI 标签，部分只有素材标签） +- compute_tag_match_score 归一化得分 +""" + +from __future__ import annotations + +import datetime as dt +import random +from dataclasses import dataclass, field + +import pytest + +from packages.domain.narrative_match import ( + AI_TAG_WEIGHT, + ASSET_TAG_WEIGHT, + _compute_ai_score, + _extract_ai_tag_names, + compute_tag_match_score, + match_assets_by_script_tags, + pick_narrative_assets, +) + + +@dataclass +class FakeAsset: + id: str + tag_ids: list[str] = field(default_factory=list) + tags: list[str] = field(default_factory=list) + status: str = "ready" + file_type: str = "video" + duration: float = 10.0 + quality_score: float | None = None + created_at: object = None + metadata: dict = field(default_factory=dict) + + +def _make_old_dt(): + return dt.datetime(2020, 1, 1, tzinfo=dt.UTC) + + +# ── _extract_ai_tag_names ───────────────────────────────────────────────── + + +class TestExtractAiTagNames: + def test_extracts_all_keys(self): + ai_tags = { + "scene": ["工厂", "车间"], + "objects": ["产品"], + "action": ["演示"], + "shot": "特写", # shot 不参与标签匹配 + "has_text": False, + } + names = _extract_ai_tag_names(ai_tags) + assert names == {"工厂", "车间", "产品", "演示"} + + def test_empty_dict(self): + assert _extract_ai_tag_names({}) == set() + + def test_none_values(self): + ai_tags = {"scene": None, "objects": None, "action": None} + assert _extract_ai_tag_names(ai_tags) == set() + + def test_case_insensitive(self): + ai_tags = {"scene": ["Factory"], "objects": [], "action": []} + names = _extract_ai_tag_names(ai_tags) + assert "factory" in names + + +# ── _compute_ai_score ───────────────────────────────────────────────────── + + +class TestComputeAiScore: + def test_single_clip_hit(self): + wanted = {"工厂", "演示"} + clips = [{"scene": ["工厂"], "objects": [], "action": ["演示"]}] + score = _compute_ai_score("a1", wanted, {"a1": clips}) + # 命中 2 个 × 2.0 = 4.0 + assert score == 2 * AI_TAG_WEIGHT + + def test_multiple_clips_takes_best(self): + wanted = {"工厂", "演示"} + clips = [ + {"scene": ["工厂"], "objects": [], "action": []}, # 1 hit = 2.0 + {"scene": ["工厂"], "objects": [], "action": ["演示"]}, # 2 hits = 4.0 + ] + score = _compute_ai_score("a1", wanted, {"a1": clips}) + assert score == 2 * AI_TAG_WEIGHT # best = 2 hits + + def test_no_match(self): + wanted = {"美食"} + clips = [{"scene": ["工厂"], "objects": [], "action": ["演示"]}] + score = _compute_ai_score("a1", wanted, {"a1": clips}) + assert score == 0.0 + + def test_no_clips_for_asset(self): + wanted = {"工厂"} + assert _compute_ai_score("a1", wanted, {}) == 0.0 + assert _compute_ai_score("a1", wanted, None) == 0.0 + + def test_empty_wanted(self): + clips = [{"scene": ["工厂"], "objects": [], "action": []}] + assert _compute_ai_score("a1", set(), {"a1": clips}) == 0.0 + + +# ── match_assets_by_script_tags with AI tags ────────────────────────────── + + +class TestMatchWithAiTags: + def test_ai_tag_hit_puts_in_matched(self): + """有 AI 标签命中 → 进入命中池.""" + assets = [FakeAsset("a1", created_at=_make_old_dt())] + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": [], "action": []}]} + + matched, unmatched = match_assets_by_script_tags( + assets, + script_tags=["工厂"], + clip_ai_tags_by_asset=clip_ai_tags, + ) + + assert [a.id for a in matched] == ["a1"] + assert unmatched == [] + + def test_ai_tag_no_match_puts_in_unmatched(self): + """AI 标签未命中 → 进入未命中池.""" + assets = [FakeAsset("a1", created_at=_make_old_dt())] + clip_ai_tags = {"a1": [{"scene": ["办公室"], "objects": [], "action": []}]} + + matched, unmatched = match_assets_by_script_tags( + assets, + script_tags=["工厂"], + clip_ai_tags_by_asset=clip_ai_tags, + ) + + assert matched == [] + assert [a.id for a in unmatched] == ["a1"] + + def test_asset_tag_still_works_without_ai_tags(self): + """无 AI 标签时，素材标签仍按权重 1.0 匹配.""" + assets = [FakeAsset("a1", tags=["工厂"], created_at=_make_old_dt())] + + matched, unmatched = match_assets_by_script_tags( + assets, + script_tags=["工厂"], + ) + + assert [a.id for a in matched] == ["a1"] + + def test_mixed_ai_and_asset_tags(self): + """混合场景：一个素材有 AI 标签，另一个只有素材标签.""" + assets = [ + FakeAsset("a1", created_at=_make_old_dt()), # AI 标签命中 + FakeAsset("a2", tags=["工厂"], created_at=_make_old_dt()), # 素材标签命中 + FakeAsset("a3", tags=["美食"], created_at=_make_old_dt()), # 无命中 + ] + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": [], "action": []}]} + + matched, unmatched = match_assets_by_script_tags( + assets, + script_tags=["工厂"], + clip_ai_tags_by_asset=clip_ai_tags, + ) + + assert {a.id for a in matched} == {"a1", "a2"} + assert [a.id for a in unmatched] == ["a3"] + + def test_ai_tag_and_asset_tag_both_hit(self): + """同一素材 AI 标签和素材标签都命中 → 仍在命中池.""" + assets = [FakeAsset("a1", tags=["工厂"], created_at=_make_old_dt())] + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": [], "action": []}]} + + matched, unmatched = match_assets_by_script_tags( + assets, + script_tags=["工厂"], + tag_names_by_id={"a1": ["工厂"]}, + clip_ai_tags_by_asset=clip_ai_tags, + ) + + assert [a.id for a in matched] == ["a1"] + + +# ── compute_tag_match_score ─────────────────────────────────────────────── + + +class TestComputeTagMatchScore: + def test_ai_only_score(self): + """仅 AI 标签命中.""" + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": [], "action": ["演示"]}]} + score = compute_tag_match_score( + "a1", + script_tags=["工厂", "演示"], + clip_ai_tags_by_asset=clip_ai_tags, + ) + # AI: 2 hits × 2.0 = 4.0; asset: 0; max = 2 × 3.0 = 6.0 + assert abs(score - 4.0 / 6.0) < 0.01 + + def test_asset_only_score(self): + """仅素材标签命中.""" + score = compute_tag_match_score( + "a1", + script_tags=["工厂", "演示"], + tag_names_by_id={"a1": ["工厂"]}, + ) + # AI: 0; asset: 1 hit × 1.0 = 1.0; max = 2 × 3.0 = 6.0 + assert abs(score - 1.0 / 6.0) < 0.01 + + def test_both_ai_and_asset_score(self): + """AI 标签 + 素材标签同时命中.""" + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": [], "action": []}]} + score = compute_tag_match_score( + "a1", + script_tags=["工厂", "演示"], + tag_names_by_id={"a1": ["工厂"]}, + clip_ai_tags_by_asset=clip_ai_tags, + ) + # AI: 1 hit × 2.0 = 2.0; asset: 1 hit × 1.0 = 1.0; max = 2 × 3.0 = 6.0 + assert abs(score - 3.0 / 6.0) < 0.01 + + def test_no_match_score_zero(self): + """无命中 → 得分 0.""" + score = compute_tag_match_score( + "a1", + script_tags=["工厂"], + tag_names_by_id={"a1": ["美食"]}, + ) + assert score == 0.0 + + def test_full_match_score_one(self): + """全命中 → 得分接近 1.0.""" + clip_ai_tags = {"a1": [{"scene": ["工厂"], "objects": ["产品"], "action": ["演示"]}]} + score = compute_tag_match_score( + "a1", + script_tags=["工厂", "产品", "演示"], + clip_ai_tags_by_asset=clip_ai_tags, + ) + # AI: 3 hits × 2.0 = 6.0; max = 3 × 3.0 = 9.0 → 6/9 = 0.667 + # 注意：仅 AI 标签命中不可能达到 1.0（因为 max 包含素材权重） + assert score > 0.5 + + def test_empty_script_tags(self): + """空文案标签 → 得分 0.""" + assert compute_tag_match_score("a1", script_tags=[]) == 0.0 + + +# ── pick_narrative_assets with AI tags ──────────────────────────────────── + + +class TestPickNarrativeWithAiTags: + def _assets(self): + old = _make_old_dt() + return [ + FakeAsset("ai_match", created_at=old), # AI 标签命中 + FakeAsset("asset_match", tags=["工厂"], created_at=old), # 素材标签命中 + FakeAsset("no_match", tags=["美食"], created_at=old), # 无命中 + ] + + def test_ai_match_prioritized(self): + """AI 标签命中的素材进入命中池.""" + clip_ai_tags = {"ai_match": [{"scene": ["工厂"], "objects": [], "action": []}]} + + picked = pick_narrative_assets( + self._assets(), + script_tags=["工厂"], + clip_ai_tags_by_asset=clip_ai_tags, + limit=2, + rng=random.Random(0), + ) + + ids = {a.id for a in picked} + assert "ai_match" in ids + assert "asset_match" in ids + + def test_fallback_when_no_ai_match(self): + """AI 标签和素材标签都未命中 → 降级.""" + clip_ai_tags = {"ai_match": [{"scene": ["办公室"], "objects": [], "action": []}]} + + picked = pick_narrative_assets( + self._assets(), + script_tags=["不存在"], + clip_ai_tags_by_asset=clip_ai_tags, + limit=2, + rng=random.Random(0), + ) + + assert len(picked) == 2 # 从全量中选取 + + def test_backward_compat_without_ai_tags(self): + """不传 clip_ai_tags_by_asset 时行为与之前完全一致.""" + picked = pick_narrative_assets( + self._assets(), + script_tags=["工厂"], + limit=2, + rng=random.Random(0), + ) + + # 仅素材标签匹配 + ids = {a.id for a in picked} + assert "asset_match" in ids + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) -- 2.54.0