agent-platform/services/team-service/app/schemas/team.py

from datetime import datetime
from typing import TYPE_CHECKING, Generic, TypeVar

from core_domain import (
    TeamDefinitionContract,
    TeamMemberContract,
    TeamRunContract,
    TeamRunStatus,
    TeamStatus,
    TeamConfigContract,
)
from core_shared import JSONValue
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from app.db.models import TeamDefinition, TeamRun, TeamConfig

T = TypeVar("T")


class ApiErrorResponse(BaseModel):
    errorType: str
    message: str
    details: dict[str, JSONValue] = Field(default_factory=dict)


class ApiResponse(BaseModel, Generic[T]):
    success: bool = True
    data: T | None = None
    error: ApiErrorResponse | None = None
    requestId: str
    serverTime: datetime


class PageRequest(BaseModel):
    page: int = Field(default=1, ge=1)
    pageSize: int = Field(default=20, ge=1, le=200)
    keyword: str | None = None

    @property
    def offset(self) -> int:
        return (self.page - 1) * self.pageSize


class PageResult(BaseModel, Generic[T]):
    items: list[T]
    total: int
    page: int
    pageSize: int
    hasMore: bool

    @classmethod
    def from_items(
        cls,
        *,
        items: list[T],
        total: int,
        page: int,
        page_size: int) -> "PageResult[T]":
        return cls(
            items=items,
            total=total,
            page=page,
            pageSize=page_size,
            hasMore=page * page_size < total)


class TeamCreateRequest(BaseModel):
    code: str
    name: str
    description: str | None = None
    team_type: str = "collaborative"
    owner_user_id: str | None = None
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)


class TeamStatusUpdateRequest(BaseModel):
    status: TeamStatus


class TeamResponse(TeamDefinitionContract):
    @classmethod
    def from_entity(cls, entity: "TeamDefinition") -> "TeamResponse":
        return cls.model_validate(entity, from_attributes=True)


class TeamConfigCreateRequest(BaseModel):
    team_id: str
    coordination_mode: str = "supervisor"
    objective: str | None = None
    member_refs: list[TeamMemberContract] = Field(default_factory=list)
    policy_json: dict[str, JSONValue] = Field(default_factory=dict)


class TeamConfigResponse(TeamConfigContract):
    @classmethod
    def from_entity(cls, entity: "TeamConfig") -> "TeamConfigResponse":
        return cls.model_validate(entity, from_attributes=True)


class TeamRunCreateRequest(BaseModel):
    team_id: str
    team_config_id: str | None = None
    session_id: str | None = None
    input_text: str | None = None
    input_json: dict[str, JSONValue] | None = None
    enqueue: bool = True


class TeamRunStatusUpdateRequest(BaseModel):
    status: TeamRunStatus
    worker_key: str | None = None
    output_text: str | None = None
    output_json: dict[str, JSONValue] | None = None
    error_code: str | None = None
    error_message: str | None = None


class TeamRunExecuteRequest(BaseModel):
    worker_key: str | None = None
    dry_run: bool = True


class TeamRunResponse(TeamRunContract):
    @classmethod
    def from_entity(cls, entity: "TeamRun") -> "TeamRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class TeamRunExecuteResponse(BaseModel):
    run: TeamRunResponse
    member_run_count: int = 0
    dry_run: bool = True


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


class TeamWorkerExecuteNextResponse(BaseModel):
    run: TeamRunResponse
    member_run_count: int = 0
    dry_run: bool = True
    released_lease_count: int = 0


class TeamDto(BaseModel):
    id: str
    name: str
    description: str | None = None
    teamType: str
    status: TeamStatus
    ownerUserId: str | None = None
    metadata: dict[str, JSONValue] = Field(default_factory=dict)
    createdTime: datetime

    @classmethod
    def from_entity(cls, entity: "TeamDefinition") -> "TeamDto":
        return cls(
            id=entity.id,
            name=entity.name,
            description=entity.description,
            teamType=entity.team_type,
            status=entity.status,
            ownerUserId=entity.owner_user_id,
            metadata=entity.metadata_json or {},
            createdTime=entity.created_time)


class TeamListRequestDto(PageRequest):
    status: TeamStatus | None = None


class TeamCreateRequestDto(BaseModel):
    name: str
    description: str | None = None
    teamType: str = "collaborative"
    ownerUserId: str | None = None
    metadata: dict[str, JSONValue] = Field(default_factory=dict)


class TeamDetailRequestDto(BaseModel):
    teamId: str


class TeamUpdateRequestDto(BaseModel):
    teamId: str
    name: str | None = None
    description: str | None = None
    teamType: str | None = None
    status: TeamStatus | None = None
    ownerUserId: str | None = None
    metadata: dict[str, JSONValue] | None = None


class TeamDeleteRequestDto(BaseModel):
    teamId: str


class TeamConfigDto(BaseModel):
    id: str
    teamId: str
    coordinationMode: str
    objective: str | None = None
    memberRefs: list[dict[str, JSONValue]]
    policy: dict[str, JSONValue]
    createdTime: datetime

    @classmethod
    def from_entity(cls, entity: "TeamConfig") -> "TeamConfigDto":
        return cls(
            id=entity.id,
            teamId=entity.team_id,
            coordinationMode=entity.coordination_mode,
            objective=entity.objective,
            memberRefs=entity.member_refs_json,
            policy=entity.policy_json,
            createdTime=entity.created_time)


class TeamConfigListRequestDto(PageRequest):
    teamId: str | None = None


class TeamConfigCreateRequestDto(BaseModel):
    teamId: str
    coordinationMode: str = "supervisor"
    objective: str | None = None
    memberRefs: list[dict[str, JSONValue]] = Field(default_factory=list)
    policy: dict[str, JSONValue] = Field(default_factory=dict)


class TeamConfigDetailRequestDto(BaseModel):
    configId: str


class TeamConfigUpdateRequestDto(BaseModel):
    configId: str
    coordinationMode: str | None = None
    objective: str | None = None
    memberRefs: list[dict[str, JSONValue]] | None = None
    policy: dict[str, JSONValue] | None = None


class TeamConfigDeleteRequestDto(BaseModel):
    configId: str


class TeamRunDto(BaseModel):
    id: str
    teamId: str
    teamConfigId: str
    sessionId: str | None = None
    inputText: str | None = None
    inputJson: dict[str, JSONValue] | None = None
    outputText: str | None = None
    outputJson: dict[str, JSONValue] | None = None
    status: TeamRunStatus
    workerKey: str | None = None
    queuedTime: datetime | None = None
    leaseExpireTime: datetime | None = None
    startedTime: datetime | None = None
    finishedTime: datetime | None = None
    errorCode: str | None = None
    errorMessage: str | None = None
    createdTime: datetime

    @classmethod
    def from_entity(cls, entity: "TeamRun") -> "TeamRunDto":
        return cls(
            id=entity.id,
            teamId=entity.team_id,
            teamConfigId=entity.team_config_id,
            sessionId=entity.session_id,
            inputText=entity.input_text,
            inputJson=entity.input_json,
            outputText=entity.output_text,
            outputJson=entity.output_json,
            status=entity.status,
            workerKey=entity.worker_key,
            queuedTime=entity.queued_time,
            leaseExpireTime=entity.lease_expire_time,
            startedTime=entity.started_time,
            finishedTime=entity.finished_time,
            errorCode=entity.error_code,
            errorMessage=entity.error_message,
            createdTime=entity.created_time)


class TeamRunListRequestDto(PageRequest):
    teamId: str | None = None
    sessionId: str | None = None
    status: TeamRunStatus | None = None


class TeamRunCreateRequestDto(BaseModel):
    teamId: str
    teamConfigId: str | None = None
    sessionId: str | None = None
    inputText: str | None = None
    inputJson: dict[str, JSONValue] | None = None
    enqueue: bool = True


class TeamRunDetailRequestDto(BaseModel):
    teamRunId: str


class TeamRunStatusUpdateRequestDto(BaseModel):
    teamRunId: str
    status: TeamRunStatus
    workerKey: str | None = None
    outputText: str | None = None
    outputJson: dict[str, JSONValue] | None = None
    errorCode: str | None = None
    errorMessage: str | None = None


class TeamRunExecuteRequestDto(BaseModel):
    teamRunId: str
    workerKey: str | None = None
    dryRun: bool = True


class TeamRunExecuteData(BaseModel):
    run: TeamRunDto
    memberRunCount: int = 0
    dryRun: bool = True


class TeamRunDeleteRequestDto(BaseModel):
    teamRunId: str


class DeleteData(BaseModel):
    deleted: bool
    teamId: str | None = None
    configId: str | None = None
    teamRunId: str | None = None