easyai-ai-gateway/apps/api/internal/httpapi/model_response.go

package httpapi

import (
	"context"

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

func (s *Server) platformModelResponse(ctx context.Context, model store.PlatformModel) store.PlatformModel {
	model.Capabilities = store.EffectivePlatformModelCapabilities(model.BaseCapabilities, model.Capabilities)
	model.Capabilities = enrichResponseCapabilities(model)
	model = s.withEffectiveResponseBillingConfig(ctx, model)
	return store.FilterPlatformModelBillingConfig(model)
}

func (s *Server) platformModelResponses(ctx context.Context, models []store.PlatformModel) []store.PlatformModel {
	items := make([]store.PlatformModel, len(models))
	for i, model := range models {
		items[i] = s.platformModelResponse(ctx, model)
	}
	return items
}

func (s *Server) withEffectiveResponseBillingConfig(ctx context.Context, model store.PlatformModel) store.PlatformModel {
	config := model.BillingConfig
	if model.PricingRuleSetID != "" {
		if ruleSetConfig, err := s.store.PricingRuleSetBillingConfig(ctx, model.PricingRuleSetID); err == nil && len(ruleSetConfig) > 0 {
			config = ruleSetConfig
		}
	}
	if len(model.BillingConfigOverride) > 0 {
		config = mergeResponseBillingConfig(config, model.BillingConfigOverride)
	}
	model.BillingConfig = config
	return model
}

func mergeResponseBillingConfig(base map[string]any, override map[string]any) map[string]any {
	if len(base) == 0 && len(override) == 0 {
		return nil
	}
	out := make(map[string]any, len(base)+len(override))
	for key, value := range base {
		out[key] = value
	}
	for key, value := range override {
		out[key] = value
	}
	return out
}

func enrichResponseCapabilities(model store.PlatformModel) map[string]any {
	if len(model.Capabilities) == 0 {
		return model.Capabilities
	}
	textGenerate, ok := enrichedTextGenerateCapabilities(model)
	if !ok {
		return model.Capabilities
	}

	out := make(map[string]any, len(model.Capabilities)+1)
	for key, value := range model.Capabilities {
		out[key] = value
	}
	out["text_generate"] = textGenerate
	return out
}

func enrichedTextGenerateCapabilities(model store.PlatformModel) (map[string]any, bool) {
	textGenerate := nestedCapabilities(model.Capabilities, "text_generate")
	if textGenerate == nil && !declaresModelType(model, "text_generate") {
		return nil, false
	}

	patch := map[string]any{}
	if _, ok := textGenerate["max_context_tokens"]; !ok {
		if value, ok := textGenerateContextTokens(model, textGenerate); ok {
			patch["max_context_tokens"] = value
		}
	}
	if _, ok := textGenerate["supportThinking"]; !ok {
		if value, ok := textGenerateSupportThinking(model, textGenerate); ok {
			patch["supportThinking"] = value
		}
	}
	if _, ok := textGenerate["thinkingEffortLevels"]; !ok {
		if value, ok := textGenerateThinkingEffortLevels(model, textGenerate); ok {
			patch["thinkingEffortLevels"] = value
		}
	}
	if len(patch) == 0 {
		return nil, false
	}

	out := make(map[string]any, len(textGenerate)+len(patch))
	for key, value := range textGenerate {
		out[key] = value
	}
	for key, value := range patch {
		out[key] = value
	}
	return out, true
}

func textGenerateContextTokens(model store.PlatformModel, textGenerate map[string]any) (any, bool) {
	if value, ok := capabilityValue(textGenerate, "maxContextTokens"); ok {
		return value, true
	}
	return capabilityValue(model.Capabilities, "max_context_tokens", "maxContextTokens", "contextWindow")
}

func textGenerateSupportThinking(model store.PlatformModel, textGenerate map[string]any) (any, bool) {
	if value, ok := capabilityValue(model.Capabilities, "supportThinking"); ok {
		return value, true
	}
	if _, ok := capabilityValue(textGenerate, "thinkingEffortLevels"); ok {
		return true, true
	}
	if _, ok := capabilityValue(model.Capabilities, "thinkingEffortLevels"); ok {
		return true, true
	}
	return capabilityValue(model.Capabilities, "reasoning")
}

func textGenerateThinkingEffortLevels(model store.PlatformModel, textGenerate map[string]any) (any, bool) {
	if value, ok := capabilityValue(model.Capabilities, "thinkingEffortLevels"); ok {
		return value, true
	}
	if hasTextGenerateThinkingCapability(model, textGenerate) {
		return []string{}, true
	}
	return nil, false
}

func hasTextGenerateThinkingCapability(model store.PlatformModel, textGenerate map[string]any) bool {
	if boolCapabilityValue(textGenerate, "supportThinking", "supportThinkingModeSwitch") {
		return true
	}
	if boolCapabilityValue(model.Capabilities, "supportThinking", "supportThinkingModeSwitch", "reasoning") {
		return true
	}
	_, ok := capabilityValue(textGenerate, "max_thinking_tokens", "maxThinkingTokens")
	return ok
}

func declaresModelType(model store.PlatformModel, modelType string) bool {
	if containsString(model.ModelType, modelType) {
		return true
	}
	if originalTypes, ok := stringListValue(model.Capabilities["originalTypes"]); ok {
		return containsString(originalTypes, modelType)
	}
	return false
}

func containsString(items []string, want string) bool {
	for _, item := range items {
		if item == want {
			return true
		}
	}
	return false
}

func nestedCapabilities(capabilities map[string]any, key string) map[string]any {
	if nested, ok := capabilities[key].(map[string]any); ok {
		return nested
	}
	return nil
}

func capabilityValue(capabilities map[string]any, keys ...string) (any, bool) {
	if len(capabilities) == 0 {
		return nil, false
	}
	for _, key := range keys {
		if value, ok := capabilities[key]; ok {
			return value, true
		}
	}
	return nil, false
}

func boolCapabilityValue(capabilities map[string]any, keys ...string) bool {
	if len(capabilities) == 0 {
		return false
	}
	for _, key := range keys {
		flag, ok := capabilities[key].(bool)
		if ok && flag {
			return true
		}
	}
	return false
}

func stringListValue(value any) ([]string, bool) {
	switch items := value.(type) {
	case []string:
		return items, len(items) > 0
	case []any:
		out := make([]string, 0, len(items))
		for _, item := range items {
			text, ok := item.(string)
			if !ok || text == "" {
				continue
			}
			out = append(out, text)
		}
		return out, len(out) > 0
	default:
		return nil, false
	}
}