Maas-Model-platform/backend/app/schemas/local_model.py

"""
本地模型数据传输对象定义
"""
from datetime import datetime
from typing import List, Optional
from pydantic import BaseModel, ConfigDict, Field


class LocalModelCreate(BaseModel):
    """创建本地模型请求"""
    name: str = Field(..., min_length=1, max_length=100)
    supplier: str = Field("Custom", min_length=1, max_length=100)
    base_url: str = Field(..., min_length=1, max_length=500)
    api_key: Optional[str] = Field(None, max_length=500)
    visibility: Optional[str] = Field('user')
    categories: Optional[List[int]] = Field([0], description="模型分类数组: 0=LLM,1=多模态,2=TTS,3=STT,4=图像生成,5=视频生成,6=图像编辑,7=Embedding,8=Rerank")


class LocalModelUpdate(BaseModel):
    """更新本地模型请求"""
    name: Optional[str] = Field(None, min_length=1, max_length=100)
    supplier: Optional[str] = Field(None, min_length=1, max_length=100)
    base_url: Optional[str] = Field(None, min_length=1, max_length=500)
    api_key: Optional[str] = Field(None, max_length=500)
    visibility: Optional[str] = Field(None)
    categories: Optional[List[int]] = Field(None, description="模型分类数组")


class LocalModelResponse(BaseModel):
    """本地模型响应"""
    id: int
    name: str
    supplier: str = Field("Custom")
    base_url: str = Field("")
    has_api_key: bool = Field(False)
    visibility: str = Field('user')
    categories: List[int] = Field([0], description="模型分类数组")
    category: int = Field(0, description="主分类（兼容字段，取 categories[0]）")
    created_at: datetime
    updated_at: datetime

    model_config = ConfigDict(from_attributes=True)

    @classmethod
    def from_model(cls, model) -> "LocalModelResponse":
        cats = model.categories or [0]
        return cls(
            id=model.id,
            name=model.name,
            supplier=model.supplier or "Custom",
            base_url=model.base_url or "",
            has_api_key=bool(model.local_api_key),
            visibility=model.visibility or 'user',
            categories=cats,
            category=cats[0] if cats else 0,
            created_at=model.created_at,
            updated_at=model.updated_at
        )


class ConnectionTestRequest(BaseModel):
    """连接测试请求"""
    base_url: str = Field(..., min_length=1, max_length=500)
    api_key: Optional[str] = Field(None, max_length=500)
    model_name: Optional[str] = Field(None, max_length=100)
    categories: Optional[List[int]] = Field(None)


class ConnectionTestResponse(BaseModel):
    """连接测试响应"""
    success: bool
    message: str