agent-platform/services/runtime-service/app/schemas/run.py

from datetime import datetime
from typing import TYPE_CHECKING

from pydantic import BaseModel, Field

from core_domain import (
    InitialNodeContract,
    NodeExecutionRequestContract,
    NodeRunContract,
    NodeRunStatusUpdateContract,
    RunBootstrapContract,
    RunCreateContract,
    RunExecutionRequestContract,
    WorkflowRunStatusUpdateContract,
    WorkflowRunContract,
)
from core_shared import JSONValue

if TYPE_CHECKING:
    from app.db.models import ExecutionLog, NodeArtifact, NodeRun, TraceSpan, WorkflowRun


class InitialNodeCreateRequest(InitialNodeContract):
    pass


class RunCreateRequest(RunCreateContract):
    initial_node: InitialNodeCreateRequest | None = None


class WorkflowRunResponse(WorkflowRunContract):

    @classmethod
    def from_entity(cls, entity: "WorkflowRun") -> "WorkflowRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class NodeRunResponse(NodeRunContract):

    @classmethod
    def from_entity(cls, entity: "NodeRun") -> "NodeRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class RunBootstrapResponse(RunBootstrapContract):
    run: WorkflowRunResponse
    initial_node: NodeRunResponse | None = None


class WorkflowRunStatusUpdateRequest(WorkflowRunStatusUpdateContract):
    pass


class NodeRunStatusUpdateRequest(NodeRunStatusUpdateContract):
    pass


class NodeRunExecuteRequest(NodeExecutionRequestContract):
    pass


class NodeRunExecuteResponse(BaseModel):
    run: WorkflowRunResponse
    node_run: NodeRunResponse
    executor_name: str


class RunExecuteRequest(RunExecutionRequestContract):
    pass


class RunExecuteResponse(BaseModel):
    run: WorkflowRunResponse
    node_runs: list[NodeRunResponse]
    executor_names: list[str]


class WorkerExecuteNextRequest(BaseModel):
    worker_key: str
    lease_seconds: int | None = Field(default=None, gt=0)


class WorkerExecuteNextResponse(BaseModel):
    run: WorkflowRunResponse
    node_run: NodeRunResponse
    executor_name: str
    released_lease_count: int = 0


class ExecutionLogResponse(BaseModel):
    id: str
    tenant_id: str
    run_id: str
    node_run_id: str | None = None
    event_type: str
    level: str
    message: str
    detail_json: dict[str, JSONValue] | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "ExecutionLog") -> "ExecutionLogResponse":
        return cls.model_validate(entity, from_attributes=True)


class NodeArtifactResponse(BaseModel):
    id: str
    tenant_id: str
    run_id: str
    node_run_id: str
    node_id: str
    artifact_type: str
    name: str
    mime_type: str | None = None
    content_text: str | None = None
    content_json: dict[str, JSONValue] | None = None
    storage_uri: str | None = None
    size_bytes: int | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "NodeArtifact") -> "NodeArtifactResponse":
        return cls.model_validate(entity, from_attributes=True)


class TraceSpanResponse(BaseModel):
    id: str
    tenant_id: str
    run_id: str
    node_run_id: str | None = None
    parent_span_id: str | None = None
    span_type: str
    name: str
    status: str
    started_time: datetime
    ended_time: datetime | None = None
    duration_ms: int | None = None
    attributes_json: dict[str, JSONValue] | None = None
    error_code: str | None = None
    error_message: str | None = None
    created_time: datetime

    @classmethod
    def from_entity(cls, entity: "TraceSpan") -> "TraceSpanResponse":
        return cls.model_validate(entity, from_attributes=True)