core/server/providers/capabilities_test.go

package providers

import "testing"

func TestLookupKnownModel_ExactMatch(t *testing.T) {
	caps, ok := LookupKnownModel("gpt-4o")
	if !ok {
		t.Fatal("expected gpt-4o to be found")
	}
	if caps.MaxOutputTokens != 16384 {
		t.Errorf("gpt-4o max output: got %d, want 16384", caps.MaxOutputTokens)
	}
	if !caps.ToolCalling {
		t.Error("gpt-4o should support tool calling")
	}
	if !caps.Vision {
		t.Error("gpt-4o should support vision")
	}
}

func TestLookupKnownModel_PrefixMatch(t *testing.T) {
	caps, ok := LookupKnownModel("claude-sonnet-4-20250514")
	if !ok {
		t.Fatal("expected claude-sonnet-4-20250514 to match prefix claude-sonnet-4")
	}
	if caps.MaxOutputTokens != 16000 {
		t.Errorf("claude-sonnet-4 max output: got %d, want 16000", caps.MaxOutputTokens)
	}
}

func TestLookupKnownModel_ProviderPrefix(t *testing.T) {
	caps, ok := LookupKnownModel("anthropic/claude-sonnet-4-20250514")
	if !ok {
		t.Fatal("expected provider-prefixed model to match")
	}
	if caps.MaxOutputTokens != 16000 {
		t.Errorf("got %d, want 16000", caps.MaxOutputTokens)
	}
}

func TestLookupKnownModel_NotFound(t *testing.T) {
	_, ok := LookupKnownModel("totally-unknown-model-xyz")
	if ok {
		t.Error("expected unknown model to not be found")
	}
}

func TestInferCapabilities_ToolCalling(t *testing.T) {
	tests := []struct {
		model string
		want  bool
	}{
		{"llama-3.1-70b-instruct", true},
		{"mistral-7b", true},
		{"qwen-2.5-72b", true},
		{"phi-3-mini", true},
		{"random-smallmodel", false},
	}
	for _, tt := range tests {
		caps := InferCapabilities(tt.model)
		if caps.ToolCalling != tt.want {
			t.Errorf("InferCapabilities(%q).ToolCalling = %v, want %v", tt.model, caps.ToolCalling, tt.want)
		}
	}
}

func TestInferCapabilities_Vision(t *testing.T) {
	caps := InferCapabilities("llava-v1.6-34b")
	if !caps.Vision {
		t.Error("llava should infer vision capability")
	}
}

func TestInferCapabilities_Reasoning(t *testing.T) {
	caps := InferCapabilities("deepseek-r1-distill-llama-70b")
	if !caps.Reasoning {
		t.Error("deepseek-r1 should infer reasoning capability")
	}
}

func TestResolveMaxOutput_FromCaps(t *testing.T) {
	// Explicit value in caps takes priority
	caps := ModelCapabilities{MaxOutputTokens: 32000}
	got := ResolveMaxOutput("whatever-model", caps)
	if got != 32000 {
		t.Errorf("got %d, want 32000", got)
	}
}

func TestResolveMaxOutput_FromKnownTable(t *testing.T) {
	// No value in caps, falls to known table
	caps := ModelCapabilities{}
	got := ResolveMaxOutput("claude-opus-4-20250514", caps)
	if got != 32000 {
		t.Errorf("got %d, want 32000", got)
	}
}

func TestResolveMaxOutput_FromContext(t *testing.T) {
	// Unknown model but has context window
	caps := ModelCapabilities{MaxContext: 32768}
	got := ResolveMaxOutput("unknown-model-abc", caps)
	// 32768 / 8 = 4096
	if got != 4096 {
		t.Errorf("got %d, want 4096 (derived from context/8)", got)
	}
}

func TestResolveMaxOutput_ContextClamp(t *testing.T) {
	// Very large context → clamp derived output to 16384
	caps := ModelCapabilities{MaxContext: 1048576}
	got := ResolveMaxOutput("unknown-large-model", caps)
	if got != 16384 {
		t.Errorf("got %d, want 16384 (clamped)", got)
	}

	// Very small context → clamp derived output to 2048
	caps = ModelCapabilities{MaxContext: 4096}
	got = ResolveMaxOutput("unknown-tiny-model", caps)
	if got != 2048 {
		t.Errorf("got %d, want 2048 (clamped)", got)
	}
}

func TestResolveMaxOutput_LastResort(t *testing.T) {
	// Totally unknown, no context
	caps := ModelCapabilities{}
	got := ResolveMaxOutput("mystery-model", caps)
	if got != 4096 {
		t.Errorf("got %d, want 4096 (last resort)", got)
	}
}

func TestMergeCapabilities(t *testing.T) {
	// Provider reports tool_calling and max_output — these should be authoritative.
	// Known table for claude-sonnet-4 has vision, thinking, etc — should fill gaps.
	providerCaps := ModelCapabilities{
		ToolCalling:     true,
		MaxOutputTokens: 16384,
	}
	merged := MergeCapabilities(providerCaps, "claude-sonnet-4-20250514")

	if !merged.ToolCalling {
		t.Error("tool_calling should be preserved from provider")
	}
	if merged.MaxOutputTokens != 16384 {
		t.Errorf("max_output should be 16384 from provider, got %d", merged.MaxOutputTokens)
	}
	// Vision should be filled from known table (claude-sonnet-4 has it)
	if !merged.Vision {
		t.Error("vision should be filled from known table")
	}
	// MaxContext should be filled from known table
	if merged.MaxContext == 0 {
		t.Error("max_context should be filled from known table")
	}
}

func TestMergeCapabilities_ProviderFalseWins(t *testing.T) {
	// Provider explicitly reports vision:false — should NOT be overridden by known table
	providerCaps := ModelCapabilities{
		ToolCalling:     true,
		Vision:          false, // provider says no vision
		MaxOutputTokens: 8192,
		MaxContext:       65536,
	}
	// Even though gpt-4o has vision in known table, provider caps are authoritative
	// Because HasProviderData() is true, the caller passes these as authoritative
	merged := MergeCapabilities(providerCaps, "some-unknown-model")

	if merged.Vision {
		t.Error("vision should remain false — provider is authoritative")
	}
}

func TestMergeCapabilities_EmptyProvider(t *testing.T) {
	// Empty provider caps — should fall through entirely to known table
	merged := MergeCapabilities(ModelCapabilities{}, "gpt-4o")

	if !merged.ToolCalling {
		t.Error("should get tool_calling from known table when provider is empty")
	}
	if !merged.Vision {
		t.Error("should get vision from known table when provider is empty")
	}
}

func TestHasProviderData(t *testing.T) {
	empty := ModelCapabilities{}
	if empty.HasProviderData() {
		t.Error("empty caps should not have provider data")
	}

	withTool := ModelCapabilities{ToolCalling: true}
	if !withTool.HasProviderData() {
		t.Error("caps with tool_calling should have provider data")
	}

	withContext := ModelCapabilities{MaxContext: 128000}
	if !withContext.HasProviderData() {
		t.Error("caps with max_context should have provider data")
	}
}