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

from __future__ import annotations

from langchain_core.messages import AIMessage

from production_build_agents.contracts.models import (
    CandidateArtifact,
    ExecutorCandidate,
    Finding,
    ArtifactBindingClaim,
)
from tests.support.fake_models import ToolAwareFakeChatModel


DEFAULT_SOURCE_PATHS = [
    "$.帖子类型",
    "$.核心制作点[0]",
]


def build_executor_candidate(
    *,
    run_id: str = "test-run",
    plan_id: str = "GlobalDataPlan",
    task_id: str = "Task1",
    plan_version: int = 1,
    deliverable_type: str = "structured_data",
    artifact_uri: str | None = None,
    expectation_id: str = "Requirement1-Expectation1",
    evidence_tool_call_ids: list[str] | None = None,
    source_paths: list[str] | None = None,
) -> ExecutorCandidate:
    selected_source_paths = (
        source_paths
        if source_paths is not None
        else DEFAULT_SOURCE_PATHS
    )
    artifacts = []
    findings = [
        Finding(
            statement="主视觉应保持制作表给出的统一设定",
            source_paths=selected_source_paths,
        )
    ]
    payload = {"global_constraints": ["统一主视觉"]}
    if deliverable_type == "image":
        payload = {}
        artifacts = [
            CandidateArtifact(
                artifact_type="image",
                uri=artifact_uri or "https://example.test/reference.png",
                description="全局视觉参考图",
            )
        ]
        findings = []
    elif deliverable_type == "video":
        payload = {}
        artifacts = [
            CandidateArtifact(
                artifact_type="video",
                uri=artifact_uri or "https://example.test/video.mp4",
                description="任务视频",
            )
        ]
        findings = []
    elif deliverable_type == "research_result":
        payload = {}
        artifacts = []
        findings = [
            Finding(
                statement="检索得到一条可用参考资料",
                source_urls=[artifact_uri or "https://example.test/source"],
            )
        ]
    elif deliverable_type == "reference_collection":
        payload = {}
        artifacts = [
            CandidateArtifact(
                artifact_type="reference",
                uri=artifact_uri or "https://example.test/reference",
                description="参考资料",
            )
        ]
        findings = []
    return ExecutorCandidate(
        schema_version="0.3",
        run_id=run_id,
        plan_id=plan_id,
        task_id=task_id,
        plan_version=plan_version,
        deliverable_type=deliverable_type,
        payload=payload,
        artifacts=artifacts,
        artifact_binding_claims=[
            ArtifactBindingClaim(
                expectation_id=expectation_id,
                artifact_uri=(
                    artifacts[0].uri
                    if artifacts
                    else "delivery_artifact.json"
                ),
                evidence_tool_call_ids=evidence_tool_call_ids or [],
            )
        ],
        findings=findings,
        unresolved=[],
        summary="已完成当前 Task",
    )


def build_executor_model(
    *candidates: ExecutorCandidate,
) -> ToolAwareFakeChatModel:
    return ToolAwareFakeChatModel(
        responses=[
            AIMessage(content=candidate.model_dump_json())
            for candidate in candidates
        ]
    )