ai-video-fullstack/backend/test_schemas.py

"""Contracts for persisted text test cases and batch execution results."""

from __future__ import annotations

import json
import re
from datetime import datetime
from typing import Annotated, Any, Literal, Union

from pydantic import Field, field_validator, model_validator

from schemas import CamelModel


TestCaseInputMode = Literal[
    "fixed_script_text",
    "fixed_script_turn_voice",
    "fixed_script_continuous_voice",
    "user_sim_text",
    "user_sim_voice",
]
ContextRole = Literal["agent", "user", "tool_call", "tool_result"]
EvaluationKind = Literal["reply", "tool_call", "overall"]
EvaluationStatus = Literal["pass", "fail"]
CaseStatus = Literal["waiting", "running", "pass", "fail", "error", "skipped"]
RunStatus = Literal["queued", "running", "completed", "cancelled"]
ErrorStage = Literal["pipeline", "model", "tool", "evaluation"]
StopReason = Literal["manual", "assertion_failure", "execution_error"]


class ContextTurn(CamelModel):
    role: ContextRole
    content: str = Field(max_length=20_000)
    tool_name: str | None = Field(default=None, max_length=128)
    tool_call_id: str | None = Field(default=None, max_length=128)
    is_error: bool | None = None

    @model_validator(mode="after")
    def validate_content(self):
        self.content = self.content.strip()
        self.tool_call_id = (self.tool_call_id or "").strip() or None
        if not self.content:
            raise ValueError("上下文内容不能为空")
        if self.role in {"tool_call", "tool_result"}:
            self.tool_name = (self.tool_name or "").strip()
            if not self.tool_name:
                raise ValueError("工具上下文必须填写工具名称")
            try:
                payload = json.loads(self.content)
            except json.JSONDecodeError as exc:
                raise ValueError("工具上下文必须是有效 JSON") from exc
            if self.role == "tool_call" and not isinstance(payload, dict):
                raise ValueError("Tool Call 参数必须是 JSON 对象")
        return self


class ToolParamAssertion(CamelModel):
    name: str = Field(min_length=1, max_length=128)
    match_mode: Literal["exact", "regex", "llm"]
    value: str = Field(min_length=1, max_length=4_000)

    @field_validator("name", "value")
    @classmethod
    def strip_required_text(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("参数名称和值不能为空")
        return value

    @model_validator(mode="after")
    def validate_regex(self):
        if self.match_mode == "regex":
            try:
                re.compile(self.value)
            except re.error as exc:
                raise ValueError(f"正则表达式无效: {exc}") from exc
        return self


class ToolSimulationResponse(CamelModel):
    outcome: Literal["success", "error"] = "success"
    body: str = Field(default='{"status":"ok"}', max_length=100_000)
    delay_ms: int = Field(default=0, ge=0, le=60_000)

    @model_validator(mode="after")
    def validate_json_body(self):
        if not self.body.strip():
            raise ValueError("Mock 工具返回值不能为空")
        try:
            json.loads(self.body)
        except json.JSONDecodeError as exc:
            raise ValueError("Mock 工具返回值必须是有效 JSON") from exc
        return self


class AnyToolSimulationScope(CamelModel):
    phase: Literal["any"] = "any"


class StartupToolSimulationScope(CamelModel):
    phase: Literal["startup"] = "startup"


class TurnToolSimulationScope(CamelModel):
    phase: Literal["turn"] = "turn"
    turn_id: str = Field(min_length=1, max_length=128)


ToolSimulationScope = Annotated[
    Union[
        AnyToolSimulationScope,
        StartupToolSimulationScope,
        TurnToolSimulationScope,
    ],
    Field(discriminator="phase"),
]


class ToolSimulationStep(ToolSimulationResponse):
    id: str = Field(min_length=1, max_length=128)
    scope: ToolSimulationScope = Field(default_factory=AnyToolSimulationScope)


class ToolSimulationBase(CamelModel):
    tool_id: str = Field(min_length=1, max_length=40)
    function_name: str = Field(min_length=1, max_length=128)

    @field_validator("tool_id", "function_name")
    @classmethod
    def strip_tool_identity(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("必须选择有效的工具")
        return value


class FixedToolSimulation(ToolSimulationBase):
    mode: Literal["fixed"] = "fixed"
    response: ToolSimulationResponse = Field(default_factory=ToolSimulationResponse)


class SequenceToolSimulation(ToolSimulationBase):
    mode: Literal["sequence"] = "sequence"
    steps: list[ToolSimulationStep] = Field(min_length=1, max_length=10_000)


ToolSimulation = Annotated[
    Union[FixedToolSimulation, SequenceToolSimulation],
    Field(discriminator="mode"),
]


class ReplyExpectedBehavior(CamelModel):
    id: str = Field(min_length=1, max_length=128)
    type: Literal["reply"] = "reply"
    assertion_type: Literal["keyword", "llm"]
    keywords: list[str] = Field(default_factory=list, max_length=50)
    keyword_match_mode: Literal["any", "all"] = "any"
    negate_keywords: bool = False
    llm_criteria: str = Field(default="", max_length=4_000)

    @model_validator(mode="after")
    def validate_assertion(self):
        self.keywords = [item.strip() for item in self.keywords if item.strip()]
        self.llm_criteria = self.llm_criteria.strip()
        if self.assertion_type == "keyword" and not self.keywords:
            raise ValueError("关键词断言至少需要一个关键词")
        if self.assertion_type == "llm" and not self.llm_criteria:
            raise ValueError("LLM 判断要求不能为空")
        return self


class ToolCallExpectedBehavior(CamelModel):
    id: str = Field(min_length=1, max_length=128)
    type: Literal["tool_call"] = "tool_call"
    tool_id: str = Field(min_length=1, max_length=40)
    function_name: str = Field(min_length=1, max_length=128)
    expectation: Literal["called", "not_called"]
    min_calls: int = Field(default=1, ge=0, le=100)
    max_calls: int | None = Field(default=None, ge=0, le=100)
    param_assertions: list[ToolParamAssertion] = Field(default_factory=list, max_length=50)

    @model_validator(mode="after")
    def validate_call_expectation(self):
        self.tool_id = self.tool_id.strip()
        self.function_name = self.function_name.strip()
        if not self.tool_id or not self.function_name:
            raise ValueError("必须选择有效的工具")
        if self.expectation == "called" and self.min_calls < 1:
            raise ValueError("应调用工具时最少调用次数必须大于 0")
        if self.max_calls is not None and self.max_calls < self.min_calls:
            raise ValueError("最多调用次数不能小于最少调用次数")
        return self


ExpectedBehavior = Annotated[
    Union[ReplyExpectedBehavior, ToolCallExpectedBehavior],
    Field(discriminator="type"),
]


class TestImageRef(CamelModel):
    asset_id: str = Field(min_length=1, max_length=40)
    filename: str = Field(default="image.jpg", max_length=255)
    width: int = Field(ge=1)
    height: int = Field(ge=1)
    size_bytes: int = Field(default=0, ge=0)

    @field_validator("asset_id", "filename")
    @classmethod
    def strip_image_text(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("图片资源信息不完整")
        return value


class FixedInputTurn(CamelModel):
    id: str = Field(min_length=1, max_length=128)
    user_input: str = Field(default="", max_length=20_000)
    image: TestImageRef | None = None
    behaviors: list[ExpectedBehavior] = Field(default_factory=list, max_length=100)

    @model_validator(mode="after")
    def validate_input(self):
        self.user_input = self.user_input.strip()
        if not self.user_input and self.image is None:
            raise ValueError("用户输入和图片至少填写一项")
        return self


class OverallCriterion(CamelModel):
    id: str = Field(min_length=1, max_length=128)
    type: Literal["llm"] = "llm"
    name: str = Field(min_length=1, max_length=80)
    criteria: str = Field(min_length=1, max_length=2_000)

    @field_validator("name", "criteria")
    @classmethod
    def strip_criterion(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("整体评估标准不能为空")
        return value


class TestCaseDefinition(CamelModel):
    context_turns: list[ContextTurn] = Field(default_factory=list, max_length=100)
    turns: list[FixedInputTurn] = Field(min_length=1, max_length=100)
    tool_simulations: list[ToolSimulation] = Field(default_factory=list, max_length=100)
    overall_criteria: list[OverallCriterion] = Field(default_factory=list, max_length=50)

    @model_validator(mode="before")
    @classmethod
    def migrate_legacy_tool_mocks(cls, raw):
        """Normalize the former per-assertion mockResponse contract on read."""

        if not isinstance(raw, dict):
            return raw
        if "tool_simulations" in raw or "toolSimulations" in raw:
            return raw

        groups: dict[tuple[str, str], list[dict[str, Any]]] = {}
        for turn_index, turn in enumerate(raw.get("turns") or []):
            if not isinstance(turn, dict):
                continue
            turn_id = str(turn.get("id") or "").strip()
            for behavior_index, behavior in enumerate(turn.get("behaviors") or []):
                if not isinstance(behavior, dict):
                    continue
                if behavior.get("type") != "tool_call":
                    continue
                if behavior.get("expectation") != "called":
                    continue
                mock = behavior.get("mock_response") or behavior.get("mockResponse")
                if not isinstance(mock, dict):
                    continue
                tool_id = str(
                    behavior.get("tool_id") or behavior.get("toolId") or ""
                ).strip()
                function_name = str(
                    behavior.get("function_name")
                    or behavior.get("functionName")
                    or ""
                ).strip()
                if not tool_id or not function_name:
                    continue
                response = {
                    "outcome": str(mock.get("outcome") or "success"),
                    "body": str(mock.get("body") or '{"status":"ok"}'),
                    "delayMs": int(
                        mock.get("delay_ms")
                        if mock.get("delay_ms") is not None
                        else mock.get("delayMs") or 0
                    ),
                }
                groups.setdefault((tool_id, function_name), []).append(
                    {
                        "response": response,
                        "turnId": turn_id,
                        "turnIndex": turn_index,
                        "behaviorIndex": behavior_index,
                        "minCalls": max(
                            1,
                            int(
                                behavior.get("min_calls")
                                if behavior.get("min_calls") is not None
                                else behavior.get("minCalls") or 1
                            ),
                        ),
                    }
                )

        migrated: list[dict[str, Any]] = []
        for (tool_id, function_name), entries in groups.items():
            response_signatures = {
                json.dumps(entry["response"], ensure_ascii=False, sort_keys=True)
                for entry in entries
            }
            if len(response_signatures) == 1:
                migrated.append(
                    {
                        "toolId": tool_id,
                        "functionName": function_name,
                        "mode": "fixed",
                        "response": entries[0]["response"],
                    }
                )
                continue

            steps: list[dict[str, Any]] = []
            for entry in entries:
                for repeat_index in range(entry["minCalls"]):
                    steps.append(
                        {
                            "id": (
                                f"legacy_step_{entry['turnIndex'] + 1}_"
                                f"{entry['behaviorIndex'] + 1}_{repeat_index + 1}"
                            ),
                            "scope": {
                                "phase": "turn",
                                "turnId": entry["turnId"],
                            },
                            **entry["response"],
                        }
                    )
            migrated.append(
                {
                    "toolId": tool_id,
                    "functionName": function_name,
                    "mode": "sequence",
                    "steps": steps,
                }
            )

        normalized = dict(raw)
        normalized["toolSimulations"] = migrated
        return normalized

    @model_validator(mode="after")
    def validate_has_expectation(self):
        if not any(turn.behaviors for turn in self.turns) and not self.overall_criteria:
            raise ValueError("至少配置一项有效的预期行为或整体评估标准")
        return self

    @model_validator(mode="after")
    def validate_tool_simulations(self):
        tool_ids: set[str] = set()
        step_ids: set[str] = set()
        turn_ids = {turn.id for turn in self.turns}
        simulations_by_tool: dict[str, ToolSimulation] = {}

        for simulation in self.tool_simulations:
            if simulation.tool_id in tool_ids:
                raise ValueError(
                    f"工具 {simulation.function_name} 只能配置一条工具模拟"
                )
            tool_ids.add(simulation.tool_id)
            simulations_by_tool[simulation.tool_id] = simulation
            if simulation.mode != "sequence":
                continue
            for step in simulation.steps:
                if step.id in step_ids:
                    raise ValueError(f"工具模拟步骤 ID 重复: {step.id}")
                step_ids.add(step.id)
                if step.scope.phase == "turn" and step.scope.turn_id not in turn_ids:
                    raise ValueError(
                        f"工具模拟步骤引用了不存在的轮次: {step.scope.turn_id}"
                    )

        for turn in self.turns:
            for behavior in turn.behaviors:
                if behavior.type != "tool_call" or behavior.expectation != "called":
                    continue
                simulation = simulations_by_tool.get(behavior.tool_id)
                if simulation is None:
                    raise ValueError(
                        f"工具 {behavior.function_name} 配置了应调用断言，但缺少工具模拟"
                    )
                if simulation.function_name != behavior.function_name:
                    raise ValueError(
                        f"工具 {behavior.function_name} 的断言与模拟不匹配"
                    )
                if simulation.mode == "fixed":
                    continue
                covers_turn = any(
                    step.scope.phase == "any"
                    or (
                        step.scope.phase == "turn"
                        and step.scope.turn_id == turn.id
                    )
                    for step in simulation.steps
                )
                if not covers_turn:
                    raise ValueError(
                        f"工具 {behavior.function_name} 缺少覆盖本轮的模拟步骤"
                    )
        return self

    @model_validator(mode="after")
    def validate_tool_context_pairs(self):
        pending: list[tuple[str | None, str]] = []
        seen_ids: set[str] = set()
        for index, turn in enumerate(self.context_turns):
            if turn.role == "tool_call":
                call_id = (turn.tool_call_id or "").strip() or None
                if call_id and call_id in seen_ids:
                    raise ValueError(f"上下文第 {index + 1} 条 Tool Call ID 重复")
                if call_id:
                    seen_ids.add(call_id)
                pending.append((call_id, turn.tool_name or ""))
                continue
            if turn.role != "tool_result":
                continue

            result_id = (turn.tool_call_id or "").strip() or None
            matched_index = -1
            for pending_index in range(len(pending) - 1, -1, -1):
                call_id, tool_name = pending[pending_index]
                if result_id:
                    matches = call_id == result_id
                else:
                    matches = tool_name == turn.tool_name
                if matches:
                    matched_index = pending_index
                    break
            if matched_index < 0:
                raise ValueError(
                    f"上下文第 {index + 1} 条 Tool Result 没有匹配的 Tool Call"
                )
            _call_id, tool_name = pending.pop(matched_index)
            if tool_name != turn.tool_name:
                raise ValueError(
                    f"上下文第 {index + 1} 条 Tool Result 的工具名称不匹配"
                )

        if pending:
            names = "、".join(tool_name for _call_id, tool_name in pending)
            raise ValueError(f"上下文 Tool Call 缺少对应的 Tool Result：{names}")
        return self


class TestSuiteCreate(CamelModel):
    name: str = Field(min_length=1, max_length=128)
    description: str = Field(default="", max_length=2_048)

    @field_validator("name")
    @classmethod
    def strip_name(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("测试集名称不能为空")
        return value


class TestSuiteUpdate(TestSuiteCreate):
    pass


class TestSuiteOut(CamelModel):
    id: str
    name: str
    description: str
    case_count: int = 0
    passed_count: int = 0
    run_count: int = 0
    updated_at: datetime


class TestCaseWrite(CamelModel):
    name: str = Field(min_length=1, max_length=128)
    description: str = Field(default="", max_length=2_048)
    input_mode: TestCaseInputMode = "fixed_script_text"
    context_turns: list[ContextTurn] = Field(default_factory=list, max_length=100)
    turns: list[FixedInputTurn] = Field(min_length=1, max_length=100)
    tool_simulations: list[ToolSimulation] = Field(default_factory=list, max_length=100)
    overall_criteria: list[OverallCriterion] = Field(default_factory=list, max_length=50)

    @field_validator("name")
    @classmethod
    def strip_name(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("测试用例名称不能为空")
        return value

    @model_validator(mode="after")
    def validate_mvp_definition(self):
        if self.input_mode != "fixed_script_text":
            raise ValueError("第一版只支持固定脚本 · 多模态")
        TestCaseDefinition(
            context_turns=self.context_turns,
            turns=self.turns,
            tool_simulations=self.tool_simulations,
            overall_criteria=self.overall_criteria,
        )
        return self

    def definition(self) -> TestCaseDefinition:
        return TestCaseDefinition(
            context_turns=self.context_turns,
            turns=self.turns,
            tool_simulations=self.tool_simulations,
            overall_criteria=self.overall_criteria,
        )


class TestCaseOut(TestCaseWrite):
    id: str
    suite_id: str
    last_result: Literal["pass", "fail", "not_run"]
    sort_order: int
    updated_at: datetime


class TestCaseOrderIn(CamelModel):
    case_ids: list[str] = Field(min_length=1, max_length=1_000)


class TestCaseBulkDeleteIn(CamelModel):
    case_ids: list[str] = Field(min_length=1, max_length=1_000)


class BatchRunConfig(CamelModel):
    concurrency: int = Field(default=3, ge=1, le=20)
    timeout_secs: int = Field(default=60, ge=1, le=600)
    failure_strategy: Literal["continue", "stop_on_fail"] = "continue"
    error_retry_count: int = Field(default=0, ge=0, le=3)
    error_strategy: Literal["continue", "stop_on_error"] = "continue"


class BatchRunCreate(CamelModel):
    assistant_id: str = Field(min_length=1, max_length=40)
    evaluator_model_resource_id: str = Field(min_length=1, max_length=40)
    case_ids: list[str] = Field(min_length=1, max_length=1_000)
    dynamic_variables: dict[str, str | int | float | bool] = Field(
        default_factory=dict,
        max_length=50,
    )
    config: BatchRunConfig = Field(default_factory=BatchRunConfig)
    title: str | None = Field(default=None, max_length=256)

    @field_validator("case_ids")
    @classmethod
    def unique_case_ids(cls, values: list[str]) -> list[str]:
        normalized = list(dict.fromkeys(value.strip() for value in values if value.strip()))
        if not normalized:
            raise ValueError("至少选择一个测试用例")
        return normalized


class BatchExecutionError(CamelModel):
    code: str
    message: str
    stage: ErrorStage
    retryable: bool


class BatchEvaluationResult(CamelModel):
    id: str
    label: str
    kind: EvaluationKind
    status: EvaluationStatus
    expected: str
    actual: str
    reason: str = ""


class BatchToolCallRecord(CamelModel):
    id: str
    function_name: str
    arguments_json: str
    result_json: str
    outcome: Literal["success", "error"]
    duration_ms: int = 0


class BatchTurnResult(CamelModel):
    id: str
    index: int
    user_input: str
    image: TestImageRef | None = None
    assistant_reply: str
    tool_calls: list[BatchToolCallRecord] = Field(default_factory=list)
    runtime_records: list[dict[str, Any]] = Field(default_factory=list)
    evaluations: list[BatchEvaluationResult] = Field(default_factory=list)


class BatchStartupResult(CamelModel):
    completed: bool = False
    greeting: str = ""
    assistant_outputs: list[str] = Field(default_factory=list)
    auto_confirmed_messages: list[dict[str, Any]] = Field(default_factory=list)
    runtime_records: list[dict[str, Any]] = Field(default_factory=list)


class BatchRunCaseOut(CamelModel):
    id: str
    name: str
    status: CaseStatus
    turns: list[BatchTurnResult] = Field(default_factory=list)
    startup: BatchStartupResult = Field(default_factory=BatchStartupResult)
    overall_criteria: list[BatchEvaluationResult] = Field(default_factory=list)
    attempt_count: int
    max_attempts: int
    execution_error: BatchExecutionError | None = None


class BatchRunSnapshotOut(CamelModel):
    id: str
    status: RunStatus
    title: str
    assistant_name: str
    config: dict[str, Any]
    cases: list[BatchRunCaseOut]
    started_at: datetime
    finished_at: datetime | None
    stop_reason: StopReason | None


class PromptImprovementChange(CamelModel):
    before: str = Field(max_length=100_000)
    after: str = Field(min_length=1, max_length=100_000)
    reason: str = Field(min_length=1, max_length=4_000)
    evidence: list[str] = Field(default_factory=list, max_length=8)


class PromptImprovementManualRecommendation(CamelModel):
    category: Literal["tool", "test_case", "other"]
    title: str = Field(min_length=1, max_length=200)
    detail: str = Field(min_length=1, max_length=4_000)


class PromptImprovementOut(CamelModel):
    diagnosis: str = Field(min_length=1, max_length=4_000)
    prompt_change: PromptImprovementChange | None = None
    manual_recommendations: list[PromptImprovementManualRecommendation] = Field(
        default_factory=list,
        max_length=20,
    )
    evaluator_name: str = Field(min_length=1, max_length=128)