model-gateway/common/util/prompt.go

package util

import (
	"encoding/json"
	"fmt"
	"model-gateway/model/entity"
	"model-gateway/service/gateway"
	"path/filepath"
	"strings"

	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/util/gconv"
)

// ======================== 表单校验与构建 ========================

// ValidateAndParseForm 校验表单并转为嵌套 map
func ValidateAndParseForm(forms []entity.Form) (map[string]any, error) {
	result := gjson.New("{}")
	for _, form := range forms {
		if form.Key == "" {
			continue
		}
		if form.Required && (form.Value == nil || gconv.String(form.Value) == "") {
			return nil, fmt.Errorf("字段 %s 为必填", form.Label)
		}
		if form.Value == nil {
			continue
		}

		val, err := validateAndConvert(form)
		if err != nil {
			return nil, err
		}
		_ = result.Set(form.Key, val)
	}
	return result.Map(), nil
}

// validateAndConvert 验证表单字段并转为标准格式
func validateAndConvert(form entity.Form) (any, error) {
	val := form.Value
	fc := form.FieldConstraint

	switch form.Type {
	case "string":
		s := gconv.String(val)
		if fc.MaxLength > 0 && len(s) > fc.MaxLength {
			return nil, fmt.Errorf("字段 %s 超过最大长度 %d", form.Label, fc.MaxLength)
		}
		if fc.MinLength > 0 && len(s) < fc.MinLength {
			return nil, fmt.Errorf("字段 %s 不足最小长度 %d", form.Label, fc.MinLength)
		}
		return s, nil

	case "number":
		f := gconv.Float64(val)
		if fc.Min != nil && f < gconv.Float64(fc.Min) {
			return nil, fmt.Errorf("字段 %s 不能小于 %v", form.Label, fc.Min)
		}
		if fc.Max != nil && f > gconv.Float64(fc.Max) {
			return nil, fmt.Errorf("字段 %s 不能大于 %v", form.Label, fc.Max)
		}
		switch fc.NumberType {
		case "float", "positiveFloat", "negativeFloat":
			return f, nil
		default:
			return int(f), nil
		}

	case "select", "radio":
		v, ok := val.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("字段 %s 格式错误", form.Label)
		}
		return v, nil

	case "upload":
		var urls []string
		switch v := val.(type) {
		case []any:
			for _, u := range v {
				urls = append(urls, gconv.String(u))
			}
		case []string:
			urls = v
		case string:
			if v != "" {
				urls = []string{v}
			}
		}
		if fc.MaxCount > 0 && len(urls) > fc.MaxCount {
			return nil, fmt.Errorf("字段 %s 上传数量超过上限 %d", form.Label, fc.MaxCount)
		}
		if fc.Accept != "" && len(urls) > 0 {
			allowed := strings.Split(fc.Accept, ",")
			for _, fileUrl := range urls {
				ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(fileUrl), "."))
				if !containsExt(allowed, ext) {
					return nil, fmt.Errorf("字段 %s 不支持的文件格式: %s", form.Label, ext)
				}
			}
		}
		return urls, nil

	default:
		return gconv.String(val), nil
	}
}

func containsExt(allowed []string, ext string) bool {
	for _, a := range allowed {
		if strings.TrimSpace(a) == ext {
			return true
		}
	}
	return false
}

// BuildTemplateFromForm 从 Form 构建模板
func BuildTemplateFromForm(forms []entity.Form) map[string]any {
	jsonObj := gjson.New("{}")
	for _, form := range forms {
		if form.Key == "" {
			continue
		}
		if form.DefaultValue != nil {
			_ = jsonObj.Set(form.Key, form.DefaultValue)
		} else {
			switch form.Type {
			case "number":
				_ = jsonObj.Set(form.Key, 0)
			case "boolean":
				_ = jsonObj.Set(form.Key, false)
			default:
				_ = jsonObj.Set(form.Key, "")
			}
		}
	}
	return jsonObj.Map()
}

// ======================== 提示词拼接 ========================

// GetFormByRole 根据 role 获取单个 Form
func GetFormByRole(forms []entity.Form, role string) (entity.Form, bool) {
	for _, f := range forms {
		if f.Role == role {
			return f, true
		}
	}
	return entity.Form{}, false
}

// GetRoleContentPath 从 Form 数组中提取指定 role 的 Key 路径
func GetRoleContentPath(forms []entity.Form, role string) string {
	for _, form := range forms {
		if form.Role == role {
			return form.Key
		}
	}
	return ""
}

// MergePrompt 将多个提示词拼接到指定路径的 content 中
func MergePrompt(messages map[string]any, contentPath string, prompts ...string) map[string]any {
	var parts []string
	for _, p := range prompts {
		if p != "" {
			parts = append(parts, p)
		}
	}
	if len(parts) == 0 {
		return messages
	}

	appendContent := strings.Join(parts, "\n")
	msgJson := gjson.New(messages)

	existing := msgJson.Get(contentPath).String()
	if existing != "" {
		appendContent = existing + "\n" + appendContent
	}
	_ = msgJson.Set(contentPath, appendContent)

	return msgJson.Map()
}

// ExtractUserContent 将 messages 转为自然语言文本
func ExtractUserContent(messages map[string]any) string {
	return formatMapToText(messages, "")
}

func formatMapToText(v any, prefix string) string {
	var b strings.Builder
	switch val := v.(type) {
	case map[string]any:
		for k, vv := range val {
			switch vv.(type) {
			case map[string]any, []any:
				b.WriteString(formatMapToText(vv, prefix))
			default:
				b.WriteString(fmt.Sprintf("%s：%v，", k, vv))
			}
		}
	case []any:
		for _, elem := range val {
			b.WriteString(formatMapToText(elem, ""))
		}
	case string:
		b.WriteString(fmt.Sprintf("%s，", val))
	}
	return b.String()
}

// InjectHistory 拼接历史提示词
func InjectHistory(messages map[string]any, history []gateway.SessionHistoryItem, mergeOrder []string) map[string]any {
	msgs, _ := messages["messages"].([]any)

	// 按 role 分组
	grouped := make(map[string][]any)
	for _, m := range msgs {
		if msg, ok := m.(map[string]any); ok {
			role := gconv.String(msg["role"])
			grouped[role] = append(grouped[role], msg)
		}
	}

	result := make([]any, 0, len(msgs)+len(history))
	for _, part := range mergeOrder {
		switch part {
		case "history":
			for _, h := range history {
				result = append(result, map[string]any{
					"role":    h.Role,
					"content": h.Content,
				})
			}
		default:
			if msgs, ok := grouped[part]; ok {
				result = append(result, msgs...)
			}
		}
	}

	messages["messages"] = result
	return messages
}

// SplitByAttachment 根据附件数量拆分多轮，未超出返回空
func SplitByAttachment(messages map[string]any, forms []entity.Form) []map[string]any {
	// 1) 获取约束
	maxVideo := getMaxCount(forms, "video")
	maxImage := getMaxCount(forms, "image")
	maxAudio := getMaxCount(forms, "audio")

	// 2) 提取 system 消息和 user content
	msgs := gjson.New(messages).Get("messages").Array()
	var systemMsg map[string]any
	var userContent []any

	for _, m := range msgs {
		msg := m.(map[string]any)
		switch msg["role"] {
		case "system":
			systemMsg = msg
		case "user":
			if c, ok := msg["content"].([]any); ok {
				userContent = c
			}
		}
	}
	if systemMsg == nil || len(userContent) == 0 {
		return nil
	}

	// 3) 分离 text 和各类附件
	var texts []any
	var videos, images, audios []any

	for _, c := range userContent {
		item := c.(map[string]any)
		switch item["type"] {
		case "video_url":
			videos = append(videos, item)
		case "image_url":
			images = append(images, item)
		case "input_audio":
			audios = append(audios, item)
		default:
			texts = append(texts, item)
		}
	}

	// 4) 判断是否需要拆分
	if len(videos) <= maxVideo && len(images) <= maxImage && len(audios) <= maxAudio {
		return nil
	}

	// 5) 按约束分片
	videoChunks := chunk(videos, maxVideo)
	imageChunks := chunk(images, maxImage)
	audioChunks := chunk(audios, maxAudio)

	maxRounds := len(videoChunks)
	if len(imageChunks) > maxRounds {
		maxRounds = len(imageChunks)
	}
	if len(audioChunks) > maxRounds {
		maxRounds = len(audioChunks)
	}

	// 6) 构建每轮 messages
	var rounds []map[string]any
	for i := 0; i < maxRounds; i++ {
		var newContent []any
		newContent = append(newContent, texts...)
		if i < len(videoChunks) {
			newContent = append(newContent, videoChunks[i]...)
		} else {
			newContent = append(newContent, videos...) // 未超出，每轮都带
		}
		if i < len(imageChunks) {
			newContent = append(newContent, imageChunks[i]...)
		} else {
			newContent = append(newContent, images...)
		}
		if i < len(audioChunks) {
			newContent = append(newContent, audioChunks[i]...)
		} else {
			newContent = append(newContent, audios...)
		}

		rounds = append(rounds, map[string]any{
			"model":      gjson.New(messages).Get("model").Val(),
			"max_tokens": gjson.New(messages).Get("max_tokens").Val(),
			"messages": []any{
				systemMsg,
				map[string]any{"role": "user", "content": newContent},
			},
		})
	}

	return rounds
}

// getMaxCount 从表单获取指定 role 的最大数量
func getMaxCount(forms []entity.Form, role string) int {
	for _, f := range forms {
		if f.Role == role {
			return f.FieldConstraint.MaxCount
		}
	}
	return 0
}

// chunk 将数组按 size 分片
func chunk(items []any, size int) [][]any {
	if size <= 0 {
		return nil
	}
	var chunks [][]any
	for i := 0; i < len(items); i += size {
		end := i + size
		if end > len(items) {
			end = len(items)
		}
		chunks = append(chunks, items[i:end])
	}
	return chunks
}

// ======================== 请求模板 ========================

// BuildRequestBody 替换模板占位符构建请求体
func BuildRequestBody(template map[string]any, modelName, systemPrompt, userContent string) map[string]any {
	escapedSystem := escapeJSON(systemPrompt)
	escapedUser := escapeJSON(userContent)

	str := gjson.New(template).MustToJsonString()
	str = strings.ReplaceAll(str, `"{{model}}"`, `"`+modelName+`"`)
	str = strings.ReplaceAll(str, `"{{system}}"`, escapedSystem)
	str = strings.ReplaceAll(str, `"{{user}}"`, escapedUser)

	return gjson.New(str).Map()
}

func escapeJSON(s string) string {
	b, _ := json.Marshal(s)
	return string(b)
}