eryao/backend/src/schemas/agent/runtime_models.py

from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from schemas.agent.ui_hints import UiHintsPayload
from schemas.domain.divination import DerivedDivinationData


class RunStatus(str, Enum):
    SUCCESS = "success"
    FAILED = "failed"


class ToolStatus(str, Enum):
    SUCCESS = "success"
    FAILURE = "failure"
    PARTIAL = "partial"


class ErrorInfo(BaseModel):
    model_config = ConfigDict(extra="forbid")

    code: str
    message: str
    retryable: bool = False
    details: dict[str, Any] | None = None


class ToolAgentOutput(BaseModel):
    model_config = ConfigDict(extra="forbid")

    tool_name: str
    tool_call_id: str
    tool_call_args: dict[str, Any] | None = None
    status: ToolStatus
    result: str
    error: ErrorInfo | None = None


class WorkerAgentOutputLite(BaseModel):
    model_config = ConfigDict(extra="forbid")

    status: RunStatus = RunStatus.SUCCESS
    sign_level: Literal["上上签", "中上签", "中下签", "下下签"]
    summary: str = Field(min_length=1, max_length=300)
    conclusion: list[str] = Field(min_length=1, max_length=6)
    focus_points: list[str] = Field(default_factory=list, max_length=6)
    advice: list[str] = Field(min_length=1, max_length=6)
    keywords: list[str] = Field(min_length=3, max_length=8)
    answer: str = Field(min_length=1, max_length=4000)
    error: ErrorInfo | None = None

    # Backward-compatible shadow fields for legacy consumers.
    key_points: list[str] = Field(default_factory=list, max_length=6)
    result_type: str = Field(default="structured_payload")
    suggested_actions: list[str] = Field(default_factory=list, max_length=6)
    divination_derived: DerivedDivinationData | None = None

    @model_validator(mode="after")
    def sync_compatibility_fields(self) -> WorkerAgentOutputLite:
        if not self.key_points and self.focus_points:
            self.key_points = list(self.focus_points)
        if not self.suggested_actions and self.advice:
            self.suggested_actions = list(self.advice)
        return self


class WorkerAgentOutputRich(WorkerAgentOutputLite):
    ui_hints: UiHintsPayload | None = None


class AgentOutput(WorkerAgentOutputRich):
    model_config = ConfigDict(extra="forbid")


WorkerAgentOutput = WorkerAgentOutputLite | WorkerAgentOutputRich


def resolve_worker_output_model() -> type[WorkerAgentOutputLite]:
    return WorkerAgentOutputLite