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

from datetime import datetime
from typing import TYPE_CHECKING

from core_domain import (
    AgentDefinitionContract,
    AgentMemoryPolicyContract,
    AgentModelConfigContract,
    AgentRunContract,
    AgentRunStatus,
    AgentSkillRefContract,
    AgentStatus,
    AgentToolInvocationContract,
    AgentToolRefContract,
    AgentVersionContract,
    AgentVersionStatus,
)
from core_shared import JSONValue
from pydantic import BaseModel, ConfigDict, Field

if TYPE_CHECKING:
    from app.db.models import AgentDefinition, AgentRun, AgentToolInvocation, AgentVersion


class AgentCreateRequest(BaseModel):
    code: str
    name: str
    description: str | None = None
    agent_type: str = "assistant"
    owner_user_id: str | None = None
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)


class AgentStatusUpdateRequest(BaseModel):
    status: AgentStatus


class AgentResponse(AgentDefinitionContract):
    @classmethod
    def from_entity(cls, entity: "AgentDefinition") -> "AgentResponse":
        return cls.model_validate(entity, from_attributes=True)


class AgentVersionCreateRequest(BaseModel):
    model_config = ConfigDict(populate_by_name=True)
    agent_id: str
    status: AgentVersionStatus = "draft"
    role: str = "assistant"
    goal: str | None = None
    system_prompt: str
    model_config_data: AgentModelConfigContract = Field(
        default_factory=AgentModelConfigContract,
        alias="model_config")
    memory_policy: AgentMemoryPolicyContract = Field(default_factory=AgentMemoryPolicyContract)
    tool_refs: list[AgentToolRefContract] = Field(default_factory=list)
    skill_refs: list[AgentSkillRefContract] = Field(default_factory=list)


class AgentVersionResponse(AgentVersionContract):
    @classmethod
    def from_entity(cls, entity: "AgentVersion") -> "AgentVersionResponse":
        return cls.model_validate(entity, from_attributes=True)


class AgentRunCreateRequest(BaseModel):
    agent_id: str
    agent_version_id: str | None = None
    session_id: str | None = None
    input_text: str | None = None
    input_json: dict[str, JSONValue] | None = None


class AgentRunStatusUpdateRequest(BaseModel):
    status: AgentRunStatus
    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 AgentRunExecuteRequest(BaseModel):
    worker_key: str | None = None
    dry_run: bool = False


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


class AgentRunResponse(AgentRunContract):
    @classmethod
    def from_entity(cls, entity: "AgentRun") -> "AgentRunResponse":
        return cls.model_validate(entity, from_attributes=True)


class AgentToolInvocationResponse(AgentToolInvocationContract):
    @classmethod
    def from_entity(cls, entity: "AgentToolInvocation") -> "AgentToolInvocationResponse":
        return cls.model_validate(entity, from_attributes=True)


class AgentRunExecuteResponse(BaseModel):
    run: AgentRunResponse
    model: str | None = None
    dry_run: bool = False


class AgentWorkerExecuteNextResponse(BaseModel):
    run: AgentRunResponse
    model: str | None = None
    dry_run: bool = False
    released_lease_count: int = 0


class AgentHealthResponse(BaseModel):
    service: str
    status: str
    database: str
    checked_time: datetime