easyai-ai-gateway/apps/api/internal/clients/parameter_correction.go

package clients

import (
	"container/list"
	"errors"
	"fmt"
	"math"
	"net/url"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"sync"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

const parameterCorrectionCacheCapacity = 512

var safeUpstreamCorrectionParameters = stringSet(
	"reasoning_effort", "reasoning.effort", "enable_thinking", "thinking_budget",
	"temperature", "top_p", "frequency_penalty", "presence_penalty", "seed", "stop", "n",
	"max_tokens", "max_completion_tokens", "max_output_tokens",
	"logprobs", "top_logprobs", "service_tier", "verbosity",
)

var (
	anchoredParameterPattern   = regexp.MustCompile(`(?i)^(?:unknown|unsupported|invalid) parameter\s*:?\s*[` + "`" + `'\"]?([a-z][a-z0-9_.]*)`)
	restrictedParameterPattern = regexp.MustCompile(`(?i)^the value of the ([a-z][a-z0-9_.]*) parameter is restricted to\s+([^\s,.;]+)`)
	fixedValuePattern          = regexp.MustCompile(`(?i)(?:restricted to|must be(?: set)?(?: to)?|only supports?(?: the)? value(?: of)?)\s*[` + "`" + `'\"]?([^\s,.;` + "`" + `'\"]+)`)
	betweenPattern             = regexp.MustCompile(`(?i)\bbetween\s+(-?[0-9]+(?:\.[0-9]+)?)\s+and\s+(-?[0-9]+(?:\.[0-9]+)?)\b`)
	maximumPattern             = regexp.MustCompile(`(?i)(?:less than or equal to|at most|maximum(?: value)?(?: is| of)?|<=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
	minimumPattern             = regexp.MustCompile(`(?i)(?:greater than or equal to|at least|minimum(?: value)?(?: is| of)?|>=)\s+(-?[0-9]+(?:\.[0-9]+)?)`)
)

type parameterCorrectionAction string

const (
	parameterCorrectionRemove parameterCorrectionAction = "remove"
	parameterCorrectionSet    parameterCorrectionAction = "set"
	parameterCorrectionClamp  parameterCorrectionAction = "clamp"
)

type parameterCorrectionRule struct {
	Param  string
	Action parameterCorrectionAction
	Value  any
	Min    *float64
	Max    *float64
}

type parameterCorrectionScope struct {
	Provider string
	BaseURL  string
	Protocol string
	Kind     string
	Model    string
}

func newParameterCorrectionScope(request Request, endpointKind string) parameterCorrectionScope {
	return parameterCorrectionScope{
		Provider: normalizedReasoningString(request.Candidate.Provider),
		BaseURL:  normalizedCorrectionBaseURL(request.Candidate.BaseURL),
		Protocol: strings.TrimSpace(request.UpstreamProtocol),
		Kind:     strings.TrimSpace(endpointKind),
		Model:    normalizedReasoningString(upstreamModelName(request.Candidate)),
	}
}

func normalizedCorrectionBaseURL(raw string) string {
	trimmed := strings.TrimRight(strings.TrimSpace(raw), "/")
	parsed, err := url.Parse(trimmed)
	if err != nil || parsed.Host == "" {
		return strings.ToLower(trimmed)
	}
	parsed.Scheme = strings.ToLower(parsed.Scheme)
	parsed.Host = strings.ToLower(parsed.Host)
	parsed.RawQuery = ""
	parsed.Fragment = ""
	return strings.TrimRight(parsed.String(), "/")
}

func (scope parameterCorrectionScope) prefix() string {
	return strings.Join([]string{scope.Provider, scope.BaseURL, scope.Protocol, scope.Kind, scope.Model}, "\x1f") + "\x1f"
}

func (scope parameterCorrectionScope) key(param string) string {
	return scope.prefix() + param
}

type parameterCorrectionCacheEntry struct {
	key  string
	rule parameterCorrectionRule
}

type ParameterCorrectionCache struct {
	mu       sync.Mutex
	capacity int
	entries  map[string]*list.Element
	lru      *list.List
}

func NewParameterCorrectionCache() *ParameterCorrectionCache {
	return &ParameterCorrectionCache{
		capacity: parameterCorrectionCacheCapacity,
		entries:  make(map[string]*list.Element),
		lru:      list.New(),
	}
}

func (cache *ParameterCorrectionCache) apply(scope parameterCorrectionScope, body map[string]any) []string {
	if cache == nil {
		return nil
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()
	prefix := scope.prefix()
	elements := make([]*list.Element, 0)
	for key, element := range cache.entries {
		if strings.HasPrefix(key, prefix) {
			elements = append(elements, element)
		}
	}
	sort.Slice(elements, func(i, j int) bool {
		return elements[i].Value.(parameterCorrectionCacheEntry).rule.Param < elements[j].Value.(parameterCorrectionCacheEntry).rule.Param
	})
	applied := make([]string, 0, len(elements))
	for _, element := range elements {
		rule := element.Value.(parameterCorrectionCacheEntry).rule
		if applyParameterCorrectionRule(body, rule) {
			applied = append(applied, rule.Param)
			cache.lru.MoveToFront(element)
		}
	}
	return applied
}

func (cache *ParameterCorrectionCache) commit(scope parameterCorrectionScope, rules []parameterCorrectionRule) {
	if cache == nil || len(rules) == 0 {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()
	for _, rule := range rules {
		key := scope.key(rule.Param)
		if existing := cache.entries[key]; existing != nil {
			existing.Value = parameterCorrectionCacheEntry{key: key, rule: rule}
			cache.lru.MoveToFront(existing)
			continue
		}
		element := cache.lru.PushFront(parameterCorrectionCacheEntry{key: key, rule: rule})
		cache.entries[key] = element
		for cache.lru.Len() > cache.capacity {
			oldest := cache.lru.Back()
			if oldest == nil {
				break
			}
			delete(cache.entries, oldest.Value.(parameterCorrectionCacheEntry).key)
			cache.lru.Remove(oldest)
		}
	}
}

func (cache *ParameterCorrectionCache) invalidate(scope parameterCorrectionScope, param string) {
	if cache == nil {
		return
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()
	key := scope.key(param)
	if element := cache.entries[key]; element != nil {
		delete(cache.entries, key)
		cache.lru.Remove(element)
	}
}

func (cache *ParameterCorrectionCache) size() int {
	if cache == nil {
		return 0
	}
	cache.mu.Lock()
	defer cache.mu.Unlock()
	return cache.lru.Len()
}

func deriveParameterCorrection(err error, body map[string]any, candidate store.RuntimeModelCandidate) (parameterCorrectionRule, bool) {
	var clientErr *ClientError
	if !errors.As(err, &clientErr) || (clientErr.StatusCode != 400 && clientErr.StatusCode != 422) {
		return parameterCorrectionRule{}, false
	}
	param, code, message := structuredParameterError(clientErr)
	if param == "" {
		param = anchoredParameterFromMessage(message)
	}
	param = normalizeCorrectionParam(param)
	if !isSafeCorrectionParam(param) {
		return parameterCorrectionRule{}, false
	}
	lowerMessage := strings.ToLower(message)
	lowerCode := strings.ToLower(code)
	if conflictRule, ok := deriveConflictCorrection(param, lowerMessage); ok {
		return conflictRule, true
	}
	if strings.Contains(lowerMessage, "unknown parameter") ||
		strings.Contains(lowerMessage, "unsupported parameter") ||
		strings.Contains(lowerMessage, "not supported") ||
		strings.Contains(lowerMessage, "does not support") ||
		strings.Contains(lowerCode, "unknown_parameter") ||
		strings.Contains(lowerCode, "unsupported_parameter") {
		return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
	}
	if param == "reasoning_effort" || param == "reasoning.effort" {
		if effort := nearestErrorSupportedReasoningEffort(message, normalizedReasoningString(valueAtCorrectionPath(body, param)), candidate); effort != "" {
			return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: effort}, true
		}
	}
	if match := fixedValuePattern.FindStringSubmatch(message); len(match) == 2 {
		if value, ok := parseCorrectionScalar(match[1]); ok {
			return parameterCorrectionRule{Param: param, Action: parameterCorrectionSet, Value: value}, true
		}
	}
	if min, max, ok := explicitNumericBounds(message); ok {
		if isOutputLimitCorrectionParam(param) {
			min = nil
		}
		return parameterCorrectionRule{Param: param, Action: parameterCorrectionClamp, Min: min, Max: max}, true
	}
	return parameterCorrectionRule{}, false
}

func requestIDFromParameterError(err error) string {
	var clientErr *ClientError
	if errors.As(err, &clientErr) {
		return clientErr.RequestID
	}
	return ""
}

func structuredParameterError(clientErr *ClientError) (string, string, string) {
	param := strings.TrimSpace(clientErr.Param)
	code := strings.TrimSpace(clientErr.Code)
	message := strings.TrimSpace(clientErr.Message)
	if clientErr.Wire == nil {
		return param, code, message
	}
	errorObject, _ := clientErr.Wire.Body["error"].(map[string]any)
	if errorObject == nil {
		return param, code, message
	}
	if value := strings.TrimSpace(stringFromAny(errorObject["param"])); value != "" {
		param = value
	}
	if value := strings.TrimSpace(stringFromAny(errorObject["code"])); value != "" {
		code = value
	}
	if value := strings.TrimSpace(stringFromAny(errorObject["message"])); value != "" {
		message = value
	}
	return param, code, message
}

func anchoredParameterFromMessage(message string) string {
	trimmed := strings.TrimSpace(message)
	for _, pattern := range []*regexp.Regexp{anchoredParameterPattern, restrictedParameterPattern} {
		if match := pattern.FindStringSubmatch(trimmed); len(match) >= 2 {
			return match[1]
		}
	}
	return ""
}

func normalizeCorrectionParam(param string) string {
	return strings.ToLower(strings.Trim(strings.TrimSpace(param), "`'\"[]"))
}

func isSafeCorrectionParam(param string) bool {
	_, ok := safeUpstreamCorrectionParameters[param]
	return ok
}

func deriveConflictCorrection(param string, message string) (parameterCorrectionRule, bool) {
	if !strings.Contains(message, "cannot be used together") &&
		!strings.Contains(message, "mutually exclusive") &&
		!strings.Contains(message, "not allowed when") {
		return parameterCorrectionRule{}, false
	}
	if strings.Contains(message, "thinking_budget") && strings.Contains(message, "reasoning_effort") {
		return parameterCorrectionRule{Param: "reasoning_effort", Action: parameterCorrectionRemove}, true
	}
	return parameterCorrectionRule{Param: param, Action: parameterCorrectionRemove}, true
}

func nearestErrorSupportedReasoningEffort(message string, requested string, candidate store.RuntimeModelCandidate) string {
	allowed := make(map[string]struct{})
	lowerMessage := strings.ToLower(message)
	for _, effort := range chatReasoningEffortOrder {
		if regexp.MustCompile(`\b` + regexp.QuoteMeta(effort) + `\b`).MatchString(lowerMessage) {
			allowed[effort] = struct{}{}
		}
	}
	if len(allowed) == 0 {
		if resolved, state := resolveCandidateReasoningEffort(requested, candidate); state == reasoningCapabilitySupported {
			return resolved
		}
		return ""
	}
	requestedIndex := reasoningEffortIndex(requested)
	best := ""
	bestDistance := len(chatReasoningEffortOrder) + 1
	for index, effort := range chatReasoningEffortOrder {
		if _, ok := allowed[effort]; !ok {
			continue
		}
		distance := int(math.Abs(float64(index - requestedIndex)))
		if distance < bestDistance {
			best = effort
			bestDistance = distance
		}
	}
	return best
}

func parseCorrectionScalar(raw string) (any, bool) {
	trimmed := strings.Trim(strings.TrimSpace(raw), "`'\"")
	switch strings.ToLower(trimmed) {
	case "true":
		return true, true
	case "false":
		return false, true
	}
	if number, err := strconv.ParseFloat(trimmed, 64); err == nil {
		if math.Trunc(number) == number {
			return int(number), true
		}
		return number, true
	}
	if isOpenAIReasoningEffort(strings.ToLower(trimmed)) {
		return strings.ToLower(trimmed), true
	}
	return nil, false
}

func explicitNumericBounds(message string) (*float64, *float64, bool) {
	if match := betweenPattern.FindStringSubmatch(message); len(match) == 3 {
		minimum, minErr := strconv.ParseFloat(match[1], 64)
		maximum, maxErr := strconv.ParseFloat(match[2], 64)
		if minErr == nil && maxErr == nil {
			return &minimum, &maximum, true
		}
	}
	var minimum *float64
	var maximum *float64
	if match := minimumPattern.FindStringSubmatch(message); len(match) == 2 {
		if value, err := strconv.ParseFloat(match[1], 64); err == nil {
			minimum = &value
		}
	}
	if match := maximumPattern.FindStringSubmatch(message); len(match) == 2 {
		if value, err := strconv.ParseFloat(match[1], 64); err == nil {
			maximum = &value
		}
	}
	return minimum, maximum, minimum != nil || maximum != nil
}

func isOutputLimitCorrectionParam(param string) bool {
	return param == "max_tokens" || param == "max_completion_tokens" || param == "max_output_tokens"
}

func applyParameterCorrectionRule(body map[string]any, rule parameterCorrectionRule) bool {
	current := valueAtCorrectionPath(body, rule.Param)
	switch rule.Action {
	case parameterCorrectionRemove:
		if current == nil {
			return false
		}
		deleteAtCorrectionPath(body, rule.Param)
		return true
	case parameterCorrectionSet:
		if current == nil {
			return false
		}
		if fmt.Sprint(current) == fmt.Sprint(rule.Value) {
			return false
		}
		setAtCorrectionPath(body, rule.Param, rule.Value)
		return true
	case parameterCorrectionClamp:
		number, ok := finiteFloatFromAny(current)
		if !ok {
			return false
		}
		corrected := number
		if rule.Min != nil && corrected < *rule.Min {
			corrected = *rule.Min
		}
		if rule.Max != nil && corrected > *rule.Max {
			corrected = *rule.Max
		}
		if corrected == number {
			return false
		}
		if math.Trunc(corrected) == corrected {
			setAtCorrectionPath(body, rule.Param, int(corrected))
		} else {
			setAtCorrectionPath(body, rule.Param, corrected)
		}
		return true
	default:
		return false
	}
}

func valueAtCorrectionPath(body map[string]any, param string) any {
	if param != "reasoning.effort" {
		return body[param]
	}
	reasoning, _ := body["reasoning"].(map[string]any)
	return reasoning["effort"]
}

func setAtCorrectionPath(body map[string]any, param string, value any) {
	if param != "reasoning.effort" {
		body[param] = value
		return
	}
	reasoning, _ := body["reasoning"].(map[string]any)
	if reasoning == nil {
		reasoning = map[string]any{}
		body["reasoning"] = reasoning
	}
	reasoning["effort"] = value
}

func deleteAtCorrectionPath(body map[string]any, param string) {
	if param != "reasoning.effort" {
		delete(body, param)
		return
	}
	reasoning, _ := body["reasoning"].(map[string]any)
	delete(reasoning, "effort")
	if len(reasoning) == 0 {
		delete(body, "reasoning")
	}
}