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

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

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, ModelProviderDefinition

ModelStatus = Literal["active", "disabled"]
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 ModelCreateRequest(BaseModel):
    code: str | None = Field(default=None, min_length=1, max_length=64)
    name: str = Field(min_length=1, max_length=128)
    provider_id: str | None = None
    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_id: str | None = None
    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_id: str | None = None
    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_id=entity.provider_id,
            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 ModelDto(BaseModel):
    id: str
    name: str
    providerId: str | None = None
    providerType: str
    providerBaseUrl: str
    hasProviderApiKey: bool
    modelName: str
    description: str | None = None
    capabilities: list[str] = Field(default_factory=list)
    contextWindow: int | None = None
    maxOutputTokens: int | None = None
    defaultTemperature: float | None = None
    timeoutSeconds: float
    metadata: dict[str, JSONValue] | None = None
    createdTime: datetime
    updatedTime: datetime

    @classmethod
    def from_entity(cls, entity: "ModelDefinition") -> "ModelDto":
        return cls(
            id=entity.id,
            name=entity.name,
            providerId=entity.provider_id,
            providerType=entity.provider_type,
            providerBaseUrl=entity.provider_base_url,
            hasProviderApiKey=bool(entity.provider_api_key),
            modelName=entity.model_name,
            description=entity.description,
            capabilities=list(entity.capabilities_json or []),
            contextWindow=entity.context_window,
            maxOutputTokens=entity.max_output_tokens,
            defaultTemperature=entity.default_temperature,
            timeoutSeconds=entity.timeout_seconds,
            metadata=entity.metadata_json,
            createdTime=entity.created_time,
            updatedTime=entity.updated_time)


class ModelCreateRequestDto(BaseModel):
    name: str = Field(min_length=1, max_length=128)
    providerId: str | None = None
    providerType: str = "openai_compatible"
    providerBaseUrl: HttpUrl | str | None = None
    providerApiKey: str | None = None
    modelName: str = Field(min_length=1, max_length=128)
    description: str | None = None
    capabilities: list[str] = Field(default_factory=lambda: ["chat"])
    contextWindow: int | None = Field(default=None, ge=1)
    maxOutputTokens: int | None = Field(default=None, ge=1)
    defaultTemperature: float | None = Field(default=None, ge=0, le=2)
    timeoutSeconds: float = Field(default=60.0, ge=1, le=300)
    metadata: dict[str, JSONValue] = Field(default_factory=dict)


class ModelUpdateRequestDto(BaseModel):
    modelId: str
    name: str | None = Field(default=None, min_length=1, max_length=128)
    providerId: str | None = None
    providerType: str | None = None
    providerBaseUrl: HttpUrl | str | None = None
    providerApiKey: str | None = None
    modelName: str | None = Field(default=None, min_length=1, max_length=128)
    description: str | None = None
    capabilities: list[str] | None = None
    contextWindow: int | None = Field(default=None, ge=1)
    maxOutputTokens: int | None = Field(default=None, ge=1)
    defaultTemperature: float | None = Field(default=None, ge=0, le=2)
    timeoutSeconds: float | None = Field(default=None, ge=1, le=300)
    metadata: dict[str, JSONValue] | None = None


class ModelDeleteRequestDto(BaseModel):
    modelId: str


class DeleteData(BaseModel):
    deleted: bool
    modelId: str | None = None
    providerId: str | None = None


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


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


class ModelTestData(BaseModel):
    model: ModelDto
    response: ChatCompletionResponseContract


class ModelItemDto(BaseModel):
    modelId: str
    displayName: str
    modelType: str
    ownedBy: str | None = None
    contextWindow: int | None = None


class ModelProviderDto(BaseModel):
    id: str
    name: str
    providerType: str
    baseUrl: str
    apiKeyRef: str
    models: list[ModelItemDto]
    defaultModel: str | None = None
    extraConfig: dict[str, JSONValue] = Field(default_factory=dict)
    createdTime: datetime
    updatedTime: datetime

    @classmethod
    def from_entity(cls, entity: "ModelProviderDefinition") -> "ModelProviderDto":
        return cls(
            id=entity.id,
            name=entity.name,
            providerType=entity.provider_type,
            baseUrl=entity.base_url,
            apiKeyRef=_mask_api_key(entity.api_key),
            models=[
                ModelItemDto(**_to_camel_model_item(item))
                for item in entity.models_json or []
            ],
            defaultModel=entity.default_model,
            extraConfig=entity.extra_config_json or {},
            createdTime=entity.created_time,
            updatedTime=entity.updated_time)


class ModelProviderCreateRequestDto(BaseModel):
    name: str = Field(min_length=1, max_length=128)
    providerType: str = "openai_compatible"
    baseUrl: HttpUrl | str
    apiKey: str | None = None
    models: list[ModelItemDto] = Field(default_factory=list)
    defaultModel: str | None = None
    extraConfig: dict[str, JSONValue] = Field(default_factory=dict)


class ModelProviderUpdateRequestDto(BaseModel):
    providerId: str
    name: str | None = Field(default=None, min_length=1, max_length=128)
    baseUrl: HttpUrl | str | None = None
    apiKey: str | None = None
    models: list[ModelItemDto] | None = None
    defaultModel: str | None = None
    extraConfig: dict[str, JSONValue] | None = None


class ModelProviderDeleteRequestDto(BaseModel):
    providerId: str


class ModelProviderTestRequestDto(BaseModel):
    providerId: str


class ModelProviderTestData(BaseModel):
    success: bool
    message: str
    latencyMs: int | None = None
    modelList: list[str] = Field(default_factory=list)


class DiscoverModelsRequestDto(BaseModel):
    providerId: str | None = None
    providerType: str | None = None
    baseUrl: HttpUrl | str | None = None
    apiKey: str | None = None


class DiscoverModelsData(BaseModel):
    providerType: str
    models: list[ModelItemDto]


def _mask_api_key(api_key: str | None) -> str:
    if not api_key:
        return ""
    return f"{api_key[:3]}***masked"


def _to_snake_model_item(item: ModelItemDto) -> dict[str, JSONValue]:
    data: dict[str, JSONValue] = {
        "model_id": item.modelId,
        "display_name": item.displayName,
        "model_type": item.modelType,
    }
    if item.ownedBy is not None:
        data["owned_by"] = item.ownedBy
    if item.contextWindow is not None:
        data["context_window"] = item.contextWindow
    return data


def _to_camel_model_item(item: dict[str, JSONValue]) -> dict[str, JSONValue]:
    return {
        "modelId": str(item.get("model_id") or item.get("modelId") or ""),
        "displayName": str(item.get("display_name") or item.get("displayName") or ""),
        "modelType": str(item.get("model_type") or item.get("modelType") or "chat"),
        "ownedBy": (
            item.get("owned_by")
            if isinstance(item.get("owned_by"), str)
            else item.get("ownedBy")
        ),
        "contextWindow": (
            item.get("context_window")
            if isinstance(item.get("context_window"), int)
            else item.get("contextWindow")
        ),
    }