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

package httpapi

import (
	"context"
	"slices"
	"testing"

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

func TestPlatformModelResponseExposesTextGenerateContextAndThinkingFromNestedCapabilities(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "gemini-3-pro-preview",
		ModelType: store.StringList{"text_generate"},
		Capabilities: map[string]any{
			"text_generate": map[string]any{
				"max_context_tokens":   1000000,
				"supportThinking":      true,
				"thinkingEffortLevels": []any{"minimal", "low", "medium", "high"},
			},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	if textGenerate["max_context_tokens"] != 1000000 {
		t.Fatalf("expected text_generate.max_context_tokens 1000000, got %#v", textGenerate["max_context_tokens"])
	}
	if textGenerate["supportThinking"] != true {
		t.Fatalf("expected text_generate.supportThinking true, got %#v", textGenerate["supportThinking"])
	}
	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"minimal", "low", "medium", "high"})
	if _, ok := response.Capabilities["contextWindow"]; ok {
		t.Fatalf("expected contextWindow root alias to be omitted, got %#v", response.Capabilities["contextWindow"])
	}
	if _, ok := response.Capabilities["thinkingEffortLevels"]; ok {
		t.Fatalf("expected thinkingEffortLevels root alias to be omitted, got %#v", response.Capabilities["thinkingEffortLevels"])
	}
}

func TestPlatformModelResponseCopiesRootTextCapabilityFieldsToTextGenerate(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "legacy-root-capability-model",
		ModelType: store.StringList{"text_generate"},
		Capabilities: map[string]any{
			"maxContextTokens":     128000,
			"supportThinking":      true,
			"thinkingEffortLevels": []any{"low", "medium"},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	if textGenerate["max_context_tokens"] != 128000 {
		t.Fatalf("expected text_generate.max_context_tokens 128000, got %#v", textGenerate["max_context_tokens"])
	}
	if textGenerate["supportThinking"] != true {
		t.Fatalf("expected text_generate.supportThinking true, got %#v", textGenerate["supportThinking"])
	}
	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"low", "medium"})
}

func TestPlatformModelResponseExposesEmptyThinkingEffortLevelsWhenOnlyThinkingSwitchIsConfigured(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "thinking-switch-model",
		ModelType: store.StringList{"text_generate"},
		Capabilities: map[string]any{
			"text_generate": map[string]any{
				"max_context_tokens":        262144,
				"max_thinking_tokens":       32768,
				"supportThinking":           true,
				"supportThinkingModeSwitch": true,
				"supportStructuredOutput":   true,
			},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	if textGenerate["supportThinking"] != true {
		t.Fatalf("expected text_generate.supportThinking true, got %#v", textGenerate["supportThinking"])
	}
	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{})
}

func TestPlatformModelResponseInheritsMissingTextGenerateThinkingLevelsFromBaseModel(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "glm-4-7-251222",
		ModelType: store.StringList{"text_generate"},
		BaseCapabilities: map[string]any{
			"originalTypes": []any{"text_generate"},
			"text_generate": map[string]any{
				"max_context_tokens":   204800,
				"supportThinking":      true,
				"thinkingEffortLevels": []any{"none", "minimal", "low", "medium", "high"},
			},
		},
		Capabilities: map[string]any{
			"originalTypes": []any{"text_generate"},
			"text_generate": map[string]any{
				"max_context_tokens":        204800,
				"max_thinking_tokens":       131072,
				"supportThinking":           true,
				"supportThinkingModeSwitch": true,
			},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"none", "minimal", "low", "medium", "high"})
	if textGenerate["max_thinking_tokens"] != 131072 {
		t.Fatalf("expected platform text_generate.max_thinking_tokens to be preserved, got %#v", textGenerate["max_thinking_tokens"])
	}
}

func TestPlatformModelResponseUsesOriginalTypesWhenModelTypeIsMissing(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "catalog-snapshot-model",
		Capabilities: map[string]any{
			"originalTypes": []any{"text_generate"},
			"text_generate": map[string]any{
				"max_context_tokens":   262144,
				"supportThinking":      true,
				"thinkingEffortLevels": []any{"high", "max"},
			},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	if textGenerate["max_context_tokens"] != 262144 {
		t.Fatalf("expected text_generate.max_context_tokens 262144, got %#v", textGenerate["max_context_tokens"])
	}
	if textGenerate["supportThinking"] != true {
		t.Fatalf("expected text_generate.supportThinking true, got %#v", textGenerate["supportThinking"])
	}
	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"high", "max"})
}

func TestPlatformModelResponsePreservesTextGenerateFieldsOverFallbacks(t *testing.T) {
	model := store.PlatformModel{
		ModelName: "reasoning-model-with-tools",
		ModelType: store.StringList{"text_generate", "tools_call"},
		Capabilities: map[string]any{
			"maxContextTokens": 999999,
			"supportThinking":  false,
			"text_generate": map[string]any{
				"max_context_tokens":   1000000,
				"supportThinking":      true,
				"thinkingEffortLevels": []any{"minimal", "low", "medium"},
			},
			"tools_call": map[string]any{
				"thinkingEffortLevels": []any{"medium", "high"},
			},
		},
	}

	response := (&Server{}).platformModelResponse(context.Background(), model)
	textGenerate := textGenerateCapabilities(t, response)

	if textGenerate["max_context_tokens"] != 1000000 {
		t.Fatalf("expected text_generate.max_context_tokens to stay 1000000, got %#v", textGenerate["max_context_tokens"])
	}
	if textGenerate["supportThinking"] != true {
		t.Fatalf("expected text_generate.supportThinking to stay true, got %#v", textGenerate["supportThinking"])
	}
	assertStringListValue(t, textGenerate["thinkingEffortLevels"], []string{"minimal", "low", "medium"})
}

func textGenerateCapabilities(t *testing.T, model store.PlatformModel) map[string]any {
	t.Helper()
	capabilities, ok := model.Capabilities["text_generate"].(map[string]any)
	if !ok {
		t.Fatalf("expected capabilities.text_generate object, got %#v", model.Capabilities["text_generate"])
	}
	return capabilities
}

func assertStringListValue(t *testing.T, got any, want []string) {
	t.Helper()
	var items []string
	switch value := got.(type) {
	case []string:
		items = value
	case []any:
		items = make([]string, 0, len(value))
		for _, item := range value {
			text, ok := item.(string)
			if !ok {
				t.Fatalf("expected string list %v, got non-string item %#v in %#v", want, item, got)
			}
			items = append(items, text)
		}
	default:
		t.Fatalf("expected string list %v, got %#v", want, got)
	}
	if !slices.Equal(items, want) {
		t.Fatalf("expected string list %v, got %v", want, items)
	}
}