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

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["上上签", "中上签", "中下签", "下下签"]
    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


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


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

    divination_derived: DerivedDivinationData | None = None


WorkerAgentOutput = WorkerAgentOutputLite | WorkerAgentOutputRich


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