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

from datetime import datetime
from typing import TYPE_CHECKING, Literal

from core_domain import ChatCompletionResponseContract
from core_shared import JSONValue
from pydantic import BaseModel, Field, HttpUrl

if TYPE_CHECKING:
    from app.db.models import ModelDefinition

ModelStatus = Literal["active", "disabled"]


class ModelCreateRequest(BaseModel):
    code: str = Field(min_length=1, max_length=64)
    name: str = Field(min_length=1, max_length=128)
    provider_type: str = "openai_compatible"
    provider_base_url: HttpUrl | str
    provider_api_key: str | None = None
    model_name: str = Field(min_length=1, max_length=128)
    status: ModelStatus = "active"
    description: str | None = None
    capabilities_json: list[str] = Field(default_factory=lambda: ["chat"])
    context_window: int | None = Field(default=None, ge=1)
    max_output_tokens: int | None = Field(default=None, ge=1)
    default_temperature: float | None = Field(default=None, ge=0, le=2)
    timeout_seconds: float = Field(default=60.0, ge=1, le=300)
    metadata_json: dict[str, JSONValue] = Field(default_factory=dict)


class ModelUpdateRequest(BaseModel):
    code: str | None = Field(default=None, min_length=1, max_length=64)
    name: str | None = Field(default=None, min_length=1, max_length=128)
    provider_type: str | None = None
    provider_base_url: HttpUrl | str | None = None
    provider_api_key: str | None = None
    model_name: str | None = Field(default=None, min_length=1, max_length=128)
    status: ModelStatus | None = None
    description: str | None = None
    capabilities_json: list[str] | None = None
    context_window: int | None = Field(default=None, ge=1)
    max_output_tokens: int | None = Field(default=None, ge=1)
    default_temperature: float | None = Field(default=None, ge=0, le=2)
    timeout_seconds: float | None = Field(default=None, ge=1, le=300)
    metadata_json: dict[str, JSONValue] | None = None


class ModelStatusUpdateRequest(BaseModel):
    status: ModelStatus


class ModelResponse(BaseModel):
    id: str
    code: str
    name: str
    provider_type: str
    provider_base_url: str
    has_provider_api_key: bool
    model_name: str
    status: ModelStatus
    description: str | None = None
    capabilities_json: list[str] = Field(default_factory=list)
    context_window: int | None = None
    max_output_tokens: int | None = None
    default_temperature: float | None = None
    timeout_seconds: float
    metadata_json: dict[str, JSONValue] | None = None
    created_time: datetime
    updated_time: datetime

    @classmethod
    def from_entity(cls, entity: "ModelDefinition") -> "ModelResponse":
        return cls(
            id=entity.id,
            code=entity.code,
            name=entity.name,
            provider_type=entity.provider_type,
            provider_base_url=entity.provider_base_url,
            has_provider_api_key=bool(entity.provider_api_key),
            model_name=entity.model_name,
            status=entity.status,
            description=entity.description,
            capabilities_json=list(entity.capabilities_json or []),
            context_window=entity.context_window,
            max_output_tokens=entity.max_output_tokens,
            default_temperature=entity.default_temperature,
            timeout_seconds=entity.timeout_seconds,
            metadata_json=entity.metadata_json,
            created_time=entity.created_time,
            updated_time=entity.updated_time,
        )


class ModelTestRequest(BaseModel):
    prompt: str = Field(default="Reply with a short readiness check.", min_length=1)
    system_prompt: str | None = "You are a concise model connectivity checker."
    temperature: float | None = Field(default=None, ge=0, le=2)
    max_tokens: int | None = Field(default=128, ge=1)


class ModelTestResponse(BaseModel):
    model: ModelResponse
    response: ChatCompletionResponseContract