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

from typing import TYPE_CHECKING

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

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


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 TeamVersionCreateRequest(BaseModel):
    team_id: str
    status: TeamVersionStatus = "draft"
    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 TeamVersionResponse(TeamVersionContract):
    @classmethod
    def from_entity(cls, entity: "TeamVersion") -> "TeamVersionResponse":
        return cls.model_validate(entity, from_attributes=True)


class TeamRunCreateRequest(BaseModel):
    team_id: str
    team_version_id: str | None = None
    session_id: str | None = None
    input_text: str | None = None
    input_json: dict[str, JSONValue] | None = None


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