gochat/backend/internal/llm/provider_manager_test.go

package llm

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestProviderManagerRequiresPageConfiguration(t *testing.T) {
	manager := NewProviderManager()
	_, err := manager.ChatCompletion(context.Background(), ChatRequest{})
	require.ErrorIs(t, err, ErrProviderNotConfigured)
}

func TestProviderManagerUsesAccountFeatureModelAndGenerationSettings(t *testing.T) {
	var request ChatRequest
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"chat-1","choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()

	manager := NewProviderManager()
	manager.SetAccountModelResolver(func(_ context.Context, accountID uint, feature string) (string, error) {
		assert.Equal(t, uint(42), accountID)
		assert.Equal(t, "editor", feature)
		return "account-editor-model", nil
	})
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "platform-model",
		EmbeddingMode: EmbeddingModeReuseChat,
		Temperature:   0.25,
		MaxTokens:     777,
	}))

	ctx := WithAccountFeature(context.Background(), 42, "editor")
	_, err := manager.ChatCompletion(ctx, ChatRequest{Model: "hard-coded-model"})
	require.NoError(t, err)
	assert.Equal(t, "account-editor-model", request.Model)
	assert.Equal(t, 0.25, request.Temperature)
	assert.Equal(t, 777, request.MaxTokens)
}

func TestProviderManagerResolvesAccountFeatureModel(t *testing.T) {
	manager := NewProviderManager()
	manager.SetAccountModelResolver(func(_ context.Context, accountID uint, feature string) (string, error) {
		assert.Equal(t, uint(42), accountID)
		assert.Equal(t, "assistant", feature)
		return "account-assistant-model", nil
	})
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai",
		ChatAPIKey:    "test-key",
		ChatModel:     "platform-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	}))

	model, err := manager.ResolveChatModel(WithAccountFeature(context.Background(), 42, "assistant"))
	require.NoError(t, err)
	assert.Equal(t, "account-assistant-model", model)
}

func TestProviderManagerCompatibleKeepsConfiguredModelAndAPIError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var request ChatRequest
		require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
		assert.Equal(t, "provider-model", request.Model)
		w.Header().Set("Content-Type", "application/json")
		if request.Messages[0].Content == "fail" {
			w.WriteHeader(http.StatusBadRequest)
			_, _ = w.Write([]byte(`{"error":{"message":"rejected","type":"invalid_request_error","code":"model_not_found"}}`))
			return
		}
		_, _ = w.Write([]byte(`{"id":"chat-1","choices":[{"message":{"role":"assistant","content":"ok"}}]}`))
	}))
	defer server.Close()

	manager := NewProviderManager()
	manager.SetAccountModelResolver(func(context.Context, uint, string) (string, error) {
		return "stale-openai-model", nil
	})
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai_compatible",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "provider-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	}))

	ctx := WithAccountFeature(context.Background(), 42, "assistant")
	_, err := manager.ChatCompletion(ctx, ChatRequest{Messages: []ChatMessage{{Role: "user", Content: "ok"}}})
	require.NoError(t, err)

	_, err = manager.ChatCompletion(ctx, ChatRequest{Messages: []ChatMessage{{Role: "user", Content: "fail"}}})
	var apiErr *APIError
	require.ErrorAs(t, err, &apiErr)
	assert.Equal(t, http.StatusBadRequest, apiErr.StatusCode)
	assert.Equal(t, "invalid_request_error", apiErr.Type)
	assert.Equal(t, "model_not_found", apiErr.Code)
}

func TestProviderManagerPreservesExplicitAssistantTemperature(t *testing.T) {
	var request ChatRequest
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"chat-1","choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()

	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai_compatible",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "platform-model",
		EmbeddingMode: EmbeddingModeReuseChat,
		Temperature:   0.9,
		MaxTokens:     777,
	}))

	ctx := WithTemperatureOverride(context.Background(), 0.2)
	_, err := manager.ChatCompletion(ctx, ChatRequest{})
	require.NoError(t, err)
	assert.Equal(t, 0.2, request.Temperature)
	assert.Equal(t, 777, request.MaxTokens)
}

func TestProviderManagerUsesSeparateEmbeddingProvider(t *testing.T) {
	var chatCalls atomic.Int32
	chatServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		chatCalls.Add(1)
		http.NotFound(w, r)
	}))
	defer chatServer.Close()

	var embeddingRequest EmbeddingRequest
	var authorization string
	embeddingServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		authorization = r.Header.Get("Authorization")
		require.NoError(t, json.NewDecoder(r.Body).Decode(&embeddingRequest))
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2,0.3]}],"model":"embed-model"}`))
	}))
	defer embeddingServer.Close()

	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:        "openai_compatible",
		ChatBaseURL:         chatServer.URL,
		ChatAPIKey:          "chat-key",
		ChatModel:           "chat-model",
		EmbeddingMode:       EmbeddingModeSeparate,
		EmbeddingProvider:   "openai_compatible",
		EmbeddingBaseURL:    embeddingServer.URL,
		EmbeddingAPIKey:     "embedding-key",
		EmbeddingModel:      "embed-model",
		EmbeddingDimensions: 3,
	}))

	_, err := manager.CreateEmbedding(context.Background(), EmbeddingRequest{Model: "ignored", Input: []string{"hello"}})
	require.NoError(t, err)
	assert.Equal(t, int32(0), chatCalls.Load())
	assert.Equal(t, "Bearer embedding-key", authorization)
	assert.Equal(t, "embed-model", embeddingRequest.Model)
	assert.Equal(t, 3, embeddingRequest.Dimensions)
}

func TestProviderManagerUsesAnthropicChatWithSeparateCompatibleEmbedding(t *testing.T) {
	var (
		anthropicRequest anthropicRequest
		embeddingRequest EmbeddingRequest
		anthropicKey     string
		embeddingAuth    string
	)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case "/v1/messages":
			anthropicKey = r.Header.Get("x-api-key")
			require.NoError(t, json.NewDecoder(r.Body).Decode(&anthropicRequest))
			_, _ = w.Write([]byte(`{"id":"msg-1","type":"message","role":"assistant","content":[{"type":"text","text":"anthropic ok"}],"model":"claude-test","stop_reason":"end_turn","usage":{"input_tokens":2,"output_tokens":3}}`))
		case "/embeddings":
			embeddingAuth = r.Header.Get("Authorization")
			require.NoError(t, json.NewDecoder(r.Body).Decode(&embeddingRequest))
			_, _ = w.Write([]byte(`{"data":[{"index":0,"embedding":[0.1,0.2,0.3]}],"model":"embed-test"}`))
		default:
			http.NotFound(w, r)
		}
	}))
	defer server.Close()

	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:        "anthropic",
		ChatBaseURL:         server.URL,
		ChatAPIKey:          "anthropic-key",
		ChatModel:           "claude-test",
		EmbeddingMode:       EmbeddingModeSeparate,
		EmbeddingProvider:   "openai_compatible",
		EmbeddingBaseURL:    server.URL,
		EmbeddingAPIKey:     "embedding-key",
		EmbeddingModel:      "embed-test",
		EmbeddingDimensions: 3,
		Temperature:         0.4,
		MaxTokens:           321,
	}))

	chatResponse, err := manager.ChatCompletion(context.Background(), ChatRequest{
		Messages: []ChatMessage{
			{Role: "system", Content: "system prompt"},
			{Role: "user", Content: "hello"},
		},
	})
	require.NoError(t, err)
	require.Len(t, chatResponse.Choices, 1)
	assert.Equal(t, "anthropic ok", chatResponse.Choices[0].Message.Content)
	assert.Equal(t, "anthropic-key", anthropicKey)
	assert.Equal(t, "claude-test", anthropicRequest.Model)
	assert.Equal(t, "system prompt", anthropicRequest.System)
	assert.Equal(t, 0.4, anthropicRequest.Temperature)
	assert.Equal(t, 321, anthropicRequest.MaxTokens)

	embeddingResponse, err := manager.CreateEmbedding(context.Background(), EmbeddingRequest{Input: []string{"hello"}})
	require.NoError(t, err)
	require.Len(t, embeddingResponse.Data, 1)
	assert.Equal(t, "Bearer embedding-key", embeddingAuth)
	assert.Equal(t, "embed-test", embeddingRequest.Model)
	assert.Equal(t, 3, embeddingRequest.Dimensions)
}

func TestProviderManagerExplicitZeroRetriesDoesNotRetry(t *testing.T) {
	var calls atomic.Int32
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		calls.Add(1)
		http.Error(w, "temporary failure", http.StatusInternalServerError)
	}))
	defer server.Close()

	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai_compatible",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "test-model",
		EmbeddingMode: EmbeddingModeReuseChat,
		MaxRetries:    0,
	}))

	_, err := manager.ChatCompletion(context.Background(), ChatRequest{})
	require.Error(t, err)
	assert.Equal(t, int32(1), calls.Load())
}

func TestProviderManagerFailedConfigurationKeepsWorkingSnapshot(t *testing.T) {
	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai",
		ChatAPIKey:    "working-key",
		ChatModel:     "working-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	}))

	err := manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "anthropic",
		ChatAPIKey:    "bad-key",
		ChatModel:     "bad-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	})
	require.Error(t, err)
	snapshot, configured := manager.Snapshot()
	require.True(t, configured)
	assert.Equal(t, "working-model", snapshot.ChatModel)
}

func TestProviderManagerClearRemovesActiveProvider(t *testing.T) {
	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai",
		ChatAPIKey:    "test-key",
		ChatModel:     "test-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	}))
	manager.Clear()
	_, configured := manager.Snapshot()
	assert.False(t, configured)
	_, err := manager.ChatCompletion(context.Background(), ChatRequest{})
	require.ErrorIs(t, err, ErrProviderNotConfigured)
}

func TestProviderManagerUsesConfiguredDefaultModel(t *testing.T) {
	var requestedModel string
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var payload ChatRequest
		require.NoError(t, json.NewDecoder(r.Body).Decode(&payload))
		requestedModel = payload.Model
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"chat-1","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()

	manager := NewProviderManager()
	require.NoError(t, manager.Configure(RuntimeProviderConfig{
		ChatProvider:  "openai_compatible",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "page-model",
		EmbeddingMode: EmbeddingModeReuseChat,
	}))

	resp, err := manager.ChatCompletion(context.Background(), ChatRequest{
		Messages: []ChatMessage{{Role: "user", Content: "hello"}},
	})
	require.NoError(t, err)
	assert.Equal(t, "page-model", requestedModel)
	require.Len(t, resp.Choices, 1)
	assert.Equal(t, "ok", resp.Choices[0].Message.Content)
}