Atlas/internal/provider/model_schema.go

package provider

import (
	"encoding/json"
	"sort"
	"strings"
)

const (
	ModelKindTextToImage  = "text_to_image"
	ModelKindImageToImage = "image_to_image"
)

type InputField struct {
	Name        string `json:"name"`
	Title       string `json:"title"`
	Description string `json:"description"`
	Type        string `json:"type"`
	Default     any    `json:"default,omitempty"`
	Required    bool   `json:"required"`
	Order       int    `json:"order"`
	Widget      string `json:"widget"` // text | textarea | number | integer | image | switch
}

type ModelDetail struct {
	ID          string       `json:"id"`
	ModelName   string       `json:"model_name,omitempty"`
	Version     string       `json:"version,omitempty"`
	Name        string       `json:"name"`
	Description string       `json:"description,omitempty"`
	Kind        string       `json:"kind"`
	InputSchema []InputField `json:"input_schema"`
}

// PredictionVersion returns the upstream prediction version identifier.
func (m ModelDetail) PredictionVersion() string {
	if m.ModelName != "" {
		return m.ModelName
	}
	if idx := strings.LastIndex(m.ID, "/"); idx >= 0 && idx < len(m.ID)-1 {
		return m.ID[idx+1:]
	}
	return m.ID
}

type rawProperty struct {
	Type        string      `json:"type"`
	Title       string      `json:"title"`
	Description string      `json:"description"`
	Default     any         `json:"default"`
	Order       int         `json:"x-order"`
}

func parseReplicateModelDetails(body []byte) []ModelDetail {
	var resp struct {
		Results []struct {
			Owner       string `json:"owner"`
			Name        string `json:"name"`
			Description string `json:"description"`
			LatestVersion struct {
				ID            string `json:"id"`
				OpenAPISchema struct {
					Components struct {
						Schemas struct {
							Input struct {
								Properties map[string]json.RawMessage `json:"properties"`
								Required   []string                   `json:"required"`
							} `json:"Input"`
						} `json:"schemas"`
					} `json:"components"`
				} `json:"openapi_schema"`
			} `json:"latest_version"`
		} `json:"results"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil
	}
	out := make([]ModelDetail, 0, len(resp.Results))
	for _, item := range resp.Results {
		if item.Owner == "" || item.Name == "" {
			continue
		}
		id := item.Owner + "/" + item.Name
		name := item.Name
		if item.Description != "" {
			name = item.Description
		}
		schema := parseInputSchema(item.LatestVersion.OpenAPISchema.Components.Schemas.Input.Properties, item.LatestVersion.OpenAPISchema.Components.Schemas.Input.Required)
		version := item.LatestVersion.ID
		if version == "" {
			version = item.Name
		}
		out = append(out, ModelDetail{
			ID:          id,
			ModelName:   item.Name,
			Version:     version,
			Name:        name,
			Description: item.Description,
			Kind:        inferModelKind(schema),
			InputSchema: schema,
		})
	}
	return out
}

func parseInputSchema(properties map[string]json.RawMessage, required []string) []InputField {
	if len(properties) == 0 {
		return defaultPromptSchema()
	}
	reqSet := map[string]struct{}{}
	for _, name := range required {
		reqSet[name] = struct{}{}
	}
	fields := make([]InputField, 0, len(properties))
	for name, raw := range properties {
		var prop rawProperty
		if err := json.Unmarshal(raw, &prop); err != nil {
			continue
		}
		title := prop.Title
		if title == "" {
			title = name
		}
		fieldType := prop.Type
		if fieldType == "" {
			fieldType = "string"
		}
		fields = append(fields, InputField{
			Name:        name,
			Title:       title,
			Description: prop.Description,
			Type:        fieldType,
			Default:     prop.Default,
			Required:    containsRequired(reqSet, name),
			Order:       prop.Order,
			Widget:      inferWidget(name, fieldType),
		})
	}
	sort.Slice(fields, func(i, j int) bool {
		if fields[i].Order == fields[j].Order {
			return fields[i].Name < fields[j].Name
		}
		return fields[i].Order < fields[j].Order
	})
	return fields
}

func inferModelKind(schema []InputField) string {
	for _, f := range schema {
		if f.Name == "image" && f.Required {
			return ModelKindImageToImage
		}
	}
	for _, f := range schema {
		if f.Name == "image" && f.Widget == "image" {
			return ModelKindImageToImage
		}
	}
	return ModelKindTextToImage
}

func inferWidget(name, fieldType string) string {
	switch fieldType {
	case "integer", "number":
		return fieldType
	case "boolean":
		return "switch"
	}
	if name == "image" || strings.HasSuffix(name, "_image") || strings.Contains(name, "image_url") {
		return "image"
	}
	if name == "prompt" || name == "negative_prompt" || strings.Contains(name, "prompt") {
		return "textarea"
	}
	return "text"
}

func defaultPromptSchema() []InputField {
	return []InputField{{
		Name:     "prompt",
		Title:    "正向提示词",
		Type:     "string",
		Required: true,
		Order:    0,
		Widget:   "textarea",
	}}
}

func manualModelDetail(modelID, name string) ModelDetail {
	if name == "" {
		name = modelID
	}
	modelName := modelID
	if idx := strings.LastIndex(modelID, "/"); idx >= 0 && idx < len(modelID)-1 {
		modelName = modelID[idx+1:]
	}
	schema := defaultPromptSchema()
	kind := ModelKindTextToImage
	if strings.Contains(strings.ToLower(modelID), "edit") || strings.Contains(strings.ToLower(name), "edit") {
		schema = append(schema, InputField{
			Name:     "image",
			Title:    "输入图 URL",
			Type:     "string",
			Required: true,
			Order:    2,
			Widget:   "image",
		})
		kind = ModelKindImageToImage
	}
	return ModelDetail{
		ID:          modelID,
		ModelName:   modelName,
		Version:     modelName,
		Name:        name,
		Kind:        kind,
		InputSchema: schema,
	}
}

func containsRequired(set map[string]struct{}, name string) bool {
	_, ok := set[name]
	return ok
}

func ValidateModelInput(schema []InputField, params map[string]any) error {
	if len(schema) == 0 {
		if p, _ := params["prompt"].(string); strings.TrimSpace(p) == "" {
			return errRequired("prompt")
		}
		return nil
	}
	for _, field := range schema {
		if !field.Required {
			continue
		}
		val, ok := params[field.Name]
		if !ok || isEmptyValue(val) {
			return errRequired(field.Name)
		}
	}
	return nil
}

func isEmptyValue(v any) bool {
	switch t := v.(type) {
	case string:
		return strings.TrimSpace(t) == ""
	case nil:
		return true
	default:
		return false
	}
}

type validationError string

func (e validationError) Error() string { return string(e) }

func errRequired(name string) error {
	return validationError(name + " required")
}

func DefaultParamValues(schema []InputField) map[string]any {
	out := map[string]any{}
	for _, field := range schema {
		if field.Default != nil {
			out[field.Name] = field.Default
		}
	}
	return out
}