GopherGate/internal/router/heuristic_test.go

package router

import (
	"testing"

	"gophergate/internal/db"
)

func TestRouteHeuristic_ConditionRules(t *testing.T) {
	targets := []string{
		"deepseek-v4-flash", // index 0
		"gemini-3-flash",    // index 1
		"grok-build-0.1",    // index 2
		"kimi-k2.6",         // index 3
		"mimo-v2.5-pro",     // index 4
		"grok-4.3",          // index 5
		"deepseek-v4-pro",   // index 6
	}

	rulesJSON := `[
		{
			"rule_id": "fast_flow_extraction",
			"conditions": {
				"any_of_tags": ["fast-flow", "classification"],
				"max_input_tokens_lt": 8000,
				"requires_reasoning": false
			},
			"primary_model": "deepseek-v4-flash",
			"fallback_model": "grok-build-0.1"
		},
		{
			"rule_id": "multimodal_long_context",
			"conditions": {
				"any_of_tags": ["standard-pro", "long-doc"],
				"has_multimodal_input": true
			},
			"primary_model": "gemini-3-flash",
			"fallback_model": "mimo-v2.5-pro"
		},
		{
			"rule_id": "regional_fallback_general",
			"conditions": {
				"is_default_fallback": true
			},
			"primary_model": "kimi-k2.6"
		}
	]`

	group := db.ModelGroup{
		ID:             "dustins_stack",
		Strategy:       "heuristic",
		HeuristicRules: &rulesJSON,
	}

	// 1. Test Match Fast Flow (condition success)
	ctx1 := &RouteContext{
		UserMessage:        "classify this JSON",
		InputTokens:        500,
		HasMultimodalInput: false,
		RequiresReasoning:  false,
		Tags:               []string{"fast-flow", "classification"},
	}
	dec1, err := routeHeuristic(group, targets, ctx1)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec1.SelectedModel != "deepseek-v4-flash" {
		t.Fatalf("expected deepseek-v4-flash, got %s", dec1.SelectedModel)
	}

	// 2. Test Multimodal Long Context (condition success)
	ctx2 := &RouteContext{
		UserMessage:        "explain this video",
		InputTokens:        15000,
		HasMultimodalInput: true,
		RequiresReasoning:  false,
		Tags:               []string{"standard-pro", "video-analysis"},
	}
	dec2, err := routeHeuristic(group, targets, ctx2)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec2.SelectedModel != "gemini-3-flash" {
		t.Fatalf("expected gemini-3-flash, got %s", dec2.SelectedModel)
	}

	// 3. Test Fallback general rule
	ctx3 := &RouteContext{
		UserMessage:        "hello there",
		InputTokens:        100,
		HasMultimodalInput: false,
		RequiresReasoning:  false,
		Tags:               []string{"general"},
	}
	dec3, err := routeHeuristic(group, targets, ctx3)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec3.SelectedModel != "kimi-k2.6" {
		t.Fatalf("expected kimi-k2.6, got %s", dec3.SelectedModel)
	}
}

func TestRouteHeuristic_LegacyRules(t *testing.T) {
	targets := []string{"gpt-4o-mini", "deepseek-v4-pro", "kimi-k2.6"}

	// Legacy pattern-based rule with regex
	rulesJSON := `[
		{"pattern": "\\b(agent|agents|tool use)\\b", "target": 1},
		{"pattern": "summarize", "target": 2}
	]`

	group := db.ModelGroup{
		ID:             "heavy-logic",
		Strategy:       "heuristic",
		HeuristicRules: &rulesJSON,
	}

	// 1. Test regex match
	ctx1 := &RouteContext{
		UserMessage: "We need an agent to do tool use",
	}
	dec1, err := routeHeuristic(group, targets, ctx1)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec1.SelectedModel != "deepseek-v4-pro" {
		t.Fatalf("expected deepseek-v4-pro, got %s", dec1.SelectedModel)
	}

	// 2. Test literal match
	ctx2 := &RouteContext{
		UserMessage: "Please summarize this text",
	}
	dec2, err := routeHeuristic(group, targets, ctx2)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec2.SelectedModel != "kimi-k2.6" {
		t.Fatalf("expected kimi-k2.6, got %s", dec2.SelectedModel)
	}
}

func TestRouteHeuristic_SoftwareDevelopmentCondition(t *testing.T) {
	targets := []string{
		"gemini-3-flash",
		"kimi-k2.7-code",
		"deepseek-v4-pro",
		"mimo-v2.5-pro",
	}

	rulesJSON := `[
		{
			"rule_id": "agentic_code_and_tools",
			"conditions": {
				"requires_tool_calling": true,
				"is_software_development": true
			},
			"primary_model": "kimi-k2.7-code",
			"fallback_model": "deepseek-v4-pro"
		},
		{
			"rule_id": "agentic_general_tasks",
			"conditions": {
				"requires_tool_calling": true
			},
			"primary_model": "mimo-v2.5-pro"
		}
	]`

	group := db.ModelGroup{
		ID:             "dustins_stack",
		Strategy:       "heuristic",
		HeuristicRules: &rulesJSON,
	}

	// 1. Tool calling + Software dev should route to kimi-k2.7-code
	ctx1 := &RouteContext{
		UserMessage:           "Write a python script to parse logs",
		RequiresToolCalling:   true,
		IsSoftwareDevelopment: true,
	}
	dec1, err := routeHeuristic(group, targets, ctx1)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec1.SelectedModel != "kimi-k2.7-code" {
		t.Fatalf("expected kimi-k2.7-code, got %s", dec1.SelectedModel)
	}

	// 2. Tool calling but NOT Software dev should route to mimo-v2.5-pro
	ctx2 := &RouteContext{
		UserMessage:           "Search the web for weather in New York",
		RequiresToolCalling:   true,
		IsSoftwareDevelopment: false,
	}
	dec2, err := routeHeuristic(group, targets, ctx2)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
	if dec2.SelectedModel != "mimo-v2.5-pro" {
		t.Fatalf("expected mimo-v2.5-pro, got %s (reason: %s)", dec2.SelectedModel, dec2.Reason)
	}
}