video-image-production-build/tests/support/planner_fixtures.py

from __future__ import annotations

import hashlib
import tempfile
from pathlib import Path

from langchain_core.messages import AIMessage

from production_build_agents.contracts.models import (
    ArtifactExpectation,
    GlobalDataPlan,
    GlobalDataRequirement,
    PlannedTask,
    TaskPackage,
)
from production_build_agents.agents.planner.agent import materialize_task_package
from tests.support.fake_models import ToolAwareFakeChatModel

GLOBAL_DATA_AUDIT_PATHS = [
    "$.帖子类型",
    "$.核心制作点[0]",
]
_PLAN_FIXTURE_DIR = tempfile.TemporaryDirectory(
    prefix="production-build-plan-fixtures-"
)
_PLAN_FIXTURE_ROOT = Path(_PLAN_FIXTURE_DIR.name)


def verification_capabilities(artifact_type: str) -> list[str]:
    return {
        "structured_data": ["document_content"],
        "document": ["document_content"],
        "research_result": ["external_source", "document_content"],
        "reference": ["external_source"],
        "image": ["visual_content"],
        "video": ["visual_content"],
        "audio": ["source_identity", "technical_integrity"],
    }[artifact_type]


def build_planned_task(
    *,
    task_id: str = "Task1",
    objective: str = "整理全片需要统一遵循的视觉约束",
    reason: str = "后续素材需要共享同一组明确约束",
    depends_on: list[str] | None = None,
    expectation_ids: list[str] | None = None,
    skill_id: str = "structured-analysis",
    source_paths: list[str] | None = None,
    deliverable_type: str = "structured_data",
    priority: int = 10,
    max_attempts: int = 2,
) -> PlannedTask:
    return PlannedTask(
        task_id=task_id,
        objective=objective,
        reason=reason,
        depends_on=depends_on or [],
        expectation_ids=expectation_ids or ["Requirement1-Expectation1"],
        replaces_artifact_ids=[],
        skill_id=skill_id,
        source_paths=source_paths or GLOBAL_DATA_AUDIT_PATHS,
        deliverable_type=deliverable_type,
        priority=priority,
        max_attempts=max_attempts,
    )


def build_global_data_plan(
    *,
    tasks: list[PlannedTask] | None = None,
    plan_version: int = 1,
) -> GlobalDataPlan:
    selected_tasks = tasks or [build_planned_task()]
    requirement_source_paths = list(
        dict.fromkeys(
            source_path
            for task in selected_tasks
            for source_path in task.source_paths
        )
    )
    artifact_types: list[str] = []
    task_artifact_types: list[str] = []
    for task in selected_tasks:
        artifact_type = {
            "structured_data": "structured_data",
            "document": "document",
            "research_result": "research_result",
            "image": "image",
            "video": "video",
            "reference_collection": (
                "image"
                if task.skill_id == "reference-inspection"
                else "reference"
            ),
        }[task.deliverable_type]
        task_artifact_types.append(artifact_type)
        if artifact_type not in artifact_types:
            artifact_types.append(artifact_type)
    selected_tasks = [
        task.model_copy(
            update={
                "expectation_ids": [
                    "Requirement1-Expectation"
                    f"{artifact_types.index(artifact_type) + 1}"
                ]
            }
        )
        for task, artifact_type in zip(
            selected_tasks,
            task_artifact_types,
        )
    ]
    return GlobalDataPlan(
        plan_id="GlobalDataPlan",
        plan_version=plan_version,
        goal="准备正式视频制作需要的全局资料",
        stage_requirements=[
            GlobalDataRequirement(
                requirement_id="Requirement1",
                description="准备正式生产需要复用的全局资料",
                importance="critical",
                source_paths=requirement_source_paths,
                artifact_expectations=[
                    ArtifactExpectation(
                        expectation_id=(
                            f"Requirement1-Expectation{index}"
                        ),
                        artifact_type=artifact_type,
                        minimum_count=1,
                        usage_scope="全片正式生产",
                        verification_capabilities=verification_capabilities(
                            artifact_type
                        ),
                    )
                    for index, artifact_type in enumerate(
                        artifact_types,
                        start=1,
                    )
                ],
            )
        ],
        tasks=selected_tasks,
        revision_summary="先完成最必要的全局资料，再进入正式制作。",
    )


def build_task_package(
    production_brief_path: Path,
    *,
    run_id: str = "test-run",
    planned_task: PlannedTask | None = None,
    plan_version: int = 1,
    dependency_artifacts: dict[str, str] | None = None,
    expectation_artifact_type: str | None = None,
    plan_path: Path | None = None,
) -> TaskPackage:
    task = planned_task or build_planned_task()
    plan_tasks = [task]
    for index, dependency in enumerate(task.depends_on, start=1):
        plan_tasks.insert(
            index - 1,
            build_planned_task(
                task_id=dependency,
                objective=f"准备 {dependency} 的前置资料",
            ),
        )
    plan = build_global_data_plan(tasks=plan_tasks, plan_version=plan_version)
    if expectation_artifact_type is not None:
        requirement = plan.stage_requirements[0]
        expectation = requirement.artifact_expectations[0].model_copy(
            update={
                "artifact_type": expectation_artifact_type,
                "verification_capabilities": verification_capabilities(
                    expectation_artifact_type
                ),
            }
        )
        plan = plan.model_copy(
            update={
                "stage_requirements": [
                    requirement.model_copy(
                        update={"artifact_expectations": [expectation]}
                    )
                ]
            }
        )
    plan_json = plan.model_dump_json(indent=2)
    if plan_path is None:
        digest = hashlib.sha256(plan_json.encode("utf-8")).hexdigest()[:16]
        plan_path = _PLAN_FIXTURE_ROOT / (
            f"{plan.plan_id}.v{plan.plan_version}.{digest}.json"
        )
    plan_path.parent.mkdir(parents=True, exist_ok=True)
    plan_path.write_text(plan_json, encoding="utf-8")
    selected_task = next(
        item for item in plan.tasks if item.task_id == task.task_id
    )
    return materialize_task_package(
        plan,
        selected_task,
        run_id=run_id,
        production_brief_path=production_brief_path,
        plan_path=plan_path,
        dependency_artifacts=dependency_artifacts,
    )


def build_planner_model(
    *plans: GlobalDataPlan,
    read_paths: list[str] | None = None,
) -> ToolAwareFakeChatModel:
    selected_paths = (
        read_paths
        if read_paths is not None
        else GLOBAL_DATA_AUDIT_PATHS
    )
    responses: list[AIMessage] = []
    for index, plan in enumerate(plans, start=1):
        responses.extend(
            [
                AIMessage(
                    content="",
                    tool_calls=[
                        {
                            "name": "read_production_brief",
                            "args": {"source_paths": selected_paths},
                            "id": f"planner-read-production-brief-{index}",
                            "type": "tool_call",
                        }
                    ],
                ),
                AIMessage(content=plan.model_dump_json()),
            ]
        )
    return ToolAwareFakeChatModel(
        responses=responses,
    )