geo/server/internal/tenant/app/kol_variable_renderer.go

package app

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"regexp"
	"strings"
)

type KolVariableType string

const (
	KolVariableTypeInput    KolVariableType = "input"
	KolVariableTypeTextarea KolVariableType = "textarea"
	KolVariableTypeSelect   KolVariableType = "select"
	KolVariableTypeNumber   KolVariableType = "number"
	KolVariableTypeCheckbox KolVariableType = "checkbox"
)

type KolVariableDefinition struct {
	ID          string          `json:"id"`
	Key         string          `json:"key"`
	Label       string          `json:"label"`
	Type        KolVariableType `json:"type"`
	Required    bool            `json:"required"`
	Placeholder *string         `json:"placeholder,omitempty"`
	Options     []string        `json:"options,omitempty"`
}

type KolSchemaJSON struct {
	Variables []KolVariableDefinition `json:"variables"`
}

var kolPlaceholderRE = regexp.MustCompile(`\{\{\s*([^{}]+?)\s*\}\}`)

func RenderPrompt(content string, schema KolSchemaJSON, values map[string]any) (string, error) {
	schema = normalizeKolSchema(schema)

	byToken := make(map[string]KolVariableDefinition, len(schema.Variables)*2)
	for _, variable := range schema.Variables {
		id := strings.TrimSpace(variable.ID)
		key := strings.TrimSpace(variable.Key)
		if id != "" {
			byToken[id] = variable
		}
		if key != "" {
			byToken[key] = variable
		}
	}

	missingSet := make(map[string]struct{})
	missing := make([]string, 0)
	addMissing := func(id string) {
		if id == "" {
			return
		}
		if _, exists := missingSet[id]; exists {
			return
		}
		missingSet[id] = struct{}{}
		missing = append(missing, id)
	}

	rendered := kolPlaceholderRE.ReplaceAllStringFunc(content, func(match string) string {
		submatches := kolPlaceholderRE.FindStringSubmatch(match)
		if len(submatches) < 2 {
			return match
		}

		token := strings.TrimSpace(submatches[1])
		variable, exists := byToken[token]
		if !exists {
			addMissing(token)
			return match
		}

		value, hasValue := lookupKolVariableValue(values, variable, token)
		if hasValue {
			return fmt.Sprint(value)
		}
		if variable.Required {
			addMissing(kolVariableDisplayName(variable))
		}
		return ""
	})

	if len(missing) > 0 {
		return "", fmt.Errorf("missing variables: %s", strings.Join(missing, ", "))
	}

	return rendered, nil
}

func HashPrompt(rendered string) string {
	sum := sha256.Sum256([]byte(rendered))
	return hex.EncodeToString(sum[:])
}

func ValidateSchema(schema KolSchemaJSON) error {
	schema = normalizeKolSchema(schema)

	seenIDs := make(map[string]struct{}, len(schema.Variables))
	seenKeys := make(map[string]struct{}, len(schema.Variables))
	for _, variable := range schema.Variables {
		id := strings.TrimSpace(variable.ID)
		if id == "" || !strings.HasPrefix(id, "var_") {
			return fmt.Errorf("variable id must be non-empty and start with var_")
		}
		if _, exists := seenIDs[id]; exists {
			return fmt.Errorf("duplicate variable id: %s", id)
		}
		seenIDs[id] = struct{}{}

		key := strings.TrimSpace(variable.Key)
		if key == "" {
			return fmt.Errorf("variable key must be non-empty")
		}
		if _, exists := seenKeys[key]; exists {
			return fmt.Errorf("duplicate variable key: %s", key)
		}
		seenKeys[key] = struct{}{}

		switch KolVariableType(strings.TrimSpace(string(variable.Type))) {
		case KolVariableTypeInput, KolVariableTypeTextarea, KolVariableTypeSelect, KolVariableTypeNumber, KolVariableTypeCheckbox:
		default:
			return fmt.Errorf("unsupported variable type: %s", variable.Type)
		}
	}

	return nil
}

func JSONEncodeSchema(schema KolSchemaJSON) ([]byte, error) {
	return json.Marshal(normalizeKolSchema(schema))
}

func normalizeKolSchema(schema KolSchemaJSON) KolSchemaJSON {
	if schema.Variables == nil {
		schema.Variables = []KolVariableDefinition{}
	}
	return schema
}

func lookupKolVariableValue(values map[string]any, variable KolVariableDefinition, token string) (any, bool) {
	candidates := []string{
		strings.TrimSpace(token),
		strings.TrimSpace(variable.Key),
		strings.TrimSpace(variable.ID),
	}

	seen := make(map[string]struct{}, len(candidates))
	for _, candidate := range candidates {
		if candidate == "" {
			continue
		}
		if _, exists := seen[candidate]; exists {
			continue
		}
		seen[candidate] = struct{}{}

		if value, ok := values[candidate]; ok {
			return value, true
		}
	}

	return nil, false
}

func kolVariableDisplayName(variable KolVariableDefinition) string {
	if label := strings.TrimSpace(variable.Label); label != "" {
		return label
	}
	if key := strings.TrimSpace(variable.Key); key != "" {
		return key
	}
	return strings.TrimSpace(variable.ID)
}