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 {
	return s.platformModelResponseWithRuleSets(model, s.responsePricingRuleSetConfigs(ctx, []store.PlatformModel{model}))
}

func (s *Server) platformModelResponseWithRuleSets(model store.PlatformModel, ruleSetConfigs map[string]map[string]any) store.PlatformModel {
	model.Capabilities = store.EffectivePlatformModelCapabilities(model.BaseCapabilities, model.Capabilities)
	model.Capabilities = enrichResponseCapabilities(model)
	model = withEffectiveResponseBillingConfig(model, ruleSetConfigs)
	return store.FilterPlatformModelBillingConfig(model)
}

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

func (s *Server) responsePricingRuleSetConfigs(ctx context.Context, models []store.PlatformModel) map[string]map[string]any {
	configs := map[string]map[string]any{}
	if s.store == nil {
		return configs
	}
	ids := map[string]bool{}
	for _, model := range models {
		for _, id := range []string{firstNonEmpty(model.BasePricingRuleSetID, model.PlatformPricingRuleSetID), model.PricingRuleSetID} {
			if id != "" {
				ids[id] = true
			}
		}
	}
	for id := range ids {
		if config, err := s.store.PricingRuleSetBillingConfig(ctx, id); err == nil && len(config) > 0 {
			configs[id] = config
		}
	}
	return configs
}

func withEffectiveResponseBillingConfig(model store.PlatformModel, ruleSetConfigs map[string]map[string]any) store.PlatformModel {
	inheritedRuleSetConfig := ruleSetConfigs[firstNonEmpty(model.BasePricingRuleSetID, model.PlatformPricingRuleSetID)]
	modelRuleSetConfig := ruleSetConfigs[model.PricingRuleSetID]
	model.BillingConfig = store.ResolveEffectiveBillingConfig(store.EffectiveBillingConfigInput{
		BaseConfig:                model.BaseBillingConfig,
		LegacyPlatformModelConfig: model.BillingConfig,
		InheritedRuleSetConfig:    inheritedRuleSetConfig,
		ModelRuleSetConfig:        modelRuleSetConfig,
		Override:                  model.BillingConfigOverride,
	})
	return model
}

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
	}
}