ai-agent/internal/ai/platform_test.go

package ai

import (
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"code.tczkiot.com/wlw/ai-agent/contract"
	"code.tczkiot.com/wlw/ai-agent/internal/pkg/enums"
)

type platformAITestProvider struct {
	config      *contract.PlatformAIConfig
	configCalls *int
	source      string
}

func (p platformAITestProvider) ModelSource(context.Context) (string, error) {
	return p.source, nil
}

func (p platformAITestProvider) Config(context.Context) (*contract.PlatformAIConfig, error) {
	if p.configCalls != nil {
		*p.configCalls = *p.configCalls + 1
	}
	return p.config, nil
}

func (p platformAITestProvider) Status(context.Context) (*contract.PlatformAIStatus, error) {
	return &contract.PlatformAIStatus{Enabled: true}, nil
}

func TestEmbeddingUsesPlatformAIProvider(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/embeddings" {
			t.Errorf("unexpected path: %s", r.URL.Path)
		}
		if got := r.Header.Get("Authorization"); got != "Bearer test-platform-key" {
			t.Errorf("unexpected authorization header: %q", got)
		}
		var body map[string]any
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Errorf("decode request: %v", err)
			return
		}
		if got := body["model"]; got != "qwen3.7-text-embedding" {
			t.Errorf("unexpected embedding model: %v", got)
		}
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{
			"object":"list",
			"data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],
			"model":"qwen3.7-text-embedding",
			"usage":{"prompt_tokens":2,"total_tokens":2}
		}`))
	}))
	defer server.Close()

	SetPlatformAIProvider(platformAITestProvider{
		source: contract.ModelSourcePlatform,
		config: &contract.PlatformAIConfig{
			APIKey:             "test-platform-key",
			BaseURL:            server.URL + "/v1",
			EmbeddingDimension: 3,
			EmbeddingModel:     "qwen3.7-text-embedding",
			HTTPClient:         server.Client(),
			MaxRetryCount:      0,
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	result, err := Embedding.GenerateEmbedding(context.Background(), "hello")
	if err != nil {
		t.Fatalf("generate platform embedding: %v", err)
	}
	if result.ModelName != "qwen3.7-text-embedding" || result.Dimension != 3 || result.TokensUsed != 2 {
		t.Fatalf("unexpected embedding result: %+v", result)
	}
}

func TestPlatformEmbeddingRequiresEnabledConfiguration(t *testing.T) {
	SetPlatformAIProvider(platformAITestProvider{
		source: contract.ModelSourcePlatform,
		config: &contract.PlatformAIConfig{
			APIKey:             "license-signed",
			BaseURL:            "https://example.com/v1",
			EmbeddingDimension: 0,
			EmbeddingModel:     "",
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	_, err := Embedding.GetModel(context.Background())
	if err == nil {
		t.Fatal("expected disabled platform embedding to be rejected")
	}
}

func TestResolveAIConfigConsumesIndependentPlatformTaskModels(t *testing.T) {
	SetPlatformAIProvider(platformAITestProvider{
		source: contract.ModelSourcePlatform,
		config: &contract.PlatformAIConfig{
			APIKey:             "platform-managed",
			BaseURL:            "https://platform.example/v1",
			ChatEnabled:        true,
			ChatModel:          "qwen-plus",
			VisionEnabled:      true,
			VisionModel:        "qwen3-vl-plus",
			EmbeddingEnabled:   true,
			EmbeddingModel:     "qwen3.7-text-embedding",
			EmbeddingDimension: 1024,
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	chat, err := ResolveAIConfig(context.Background(), enums.AIModelTypeLLM, 0)
	if err != nil {
		t.Fatalf("resolve platform chat: %v", err)
	}
	if !chat.ChatEnabled || chat.ModelName != "qwen-plus" || !chat.VisionEnabled || chat.VisionModel != "qwen3-vl-plus" {
		t.Fatalf("unexpected platform chat config: %+v", chat)
	}

	embedding, err := ResolveAIConfig(context.Background(), enums.AIModelTypeEmbedding, 0)
	if err != nil {
		t.Fatalf("resolve platform embedding: %v", err)
	}
	if !embedding.EmbeddingEnabled || embedding.ModelName != "qwen3.7-text-embedding" || embedding.Dimension != 1024 {
		t.Fatalf("unexpected platform embedding config: %+v", embedding)
	}
}

func TestResolveAIConfigRejectsExplicitlyDisabledPlatformChat(t *testing.T) {
	SetPlatformAIProvider(platformAITestProvider{
		source: contract.ModelSourcePlatform,
		config: &contract.PlatformAIConfig{
			APIKey:      "platform-managed",
			BaseURL:     "https://platform.example/v1",
			ChatEnabled: false,
			ChatModel:   "qwen-plus",
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	_, err := ResolveAIConfig(context.Background(), enums.AIModelTypeLLM, 0)
	if err == nil || !strings.Contains(err.Error(), "chat model is not enabled") {
		t.Fatalf("expected explicit disabled chat error, got %v", err)
	}
}

func TestResolveVisionAIConfigDoesNotDependOnPlatformChat(t *testing.T) {
	SetPlatformAIProvider(platformAITestProvider{
		source: contract.ModelSourcePlatform,
		config: &contract.PlatformAIConfig{
			APIKey:        "platform-managed",
			BaseURL:       "https://platform.example/v1",
			ChatEnabled:   false,
			ChatModel:     "qwen-plus",
			VisionEnabled: true,
			VisionModel:   "qwen3-vl-plus",
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	config, err := ResolveVisionAIConfig(context.Background(), 0)
	if err != nil {
		t.Fatalf("resolve independent platform vision: %v", err)
	}
	if config.ChatEnabled || !config.VisionEnabled || config.ModelName != "qwen3-vl-plus" {
		t.Fatalf("unexpected independent platform vision config: %+v", config)
	}
}

func TestPlatformRerankNeverFallsBackToCustomConfig(t *testing.T) {
	configCalls := 0
	SetPlatformAIProvider(platformAITestProvider{
		source:      contract.ModelSourcePlatform,
		configCalls: &configCalls,
		config: &contract.PlatformAIConfig{
			APIKey:    "license-signed",
			BaseURL:   "https://platform.example/v1",
			ModelName: "platform-default",
		},
	})
	t.Cleanup(func() { SetPlatformAIProvider(nil) })

	if _, err := ResolveAIConfig(context.Background(), enums.AIModelTypeRerank, 0); !errors.Is(err, ErrPlatformModelUnsupported) {
		t.Fatal("expected unsupported platform rerank to fail without reading a local config")
	}
	if configCalls != 0 {
		t.Fatalf("platform Config() calls = %d, want 0 for unsupported rerank", configCalls)
	}
}