gochat/backend/internal/service/copilot_feature_model_integration_test.go

package service

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

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/datatypes"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/gochat/gochat/internal/repository"
)

func TestCopilotFeatureModelsReachProviderRequests(t *testing.T) {
	type capturedRequest struct {
		Model       string
		Temperature float64
		MaxTokens   int
	}
	var (
		mu       sync.Mutex
		captured []capturedRequest
	)
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var request llm.ChatRequest
		require.NoError(t, json.NewDecoder(r.Body).Decode(&request))
		mu.Lock()
		captured = append(captured, capturedRequest{
			Model:       request.Model,
			Temperature: request.Temperature,
			MaxTokens:   request.MaxTokens,
		})
		mu.Unlock()
		w.Header().Set("Content-Type", "application/json")
		_, _ = w.Write([]byte(`{"id":"chat-1","choices":[{"message":{"role":"assistant","content":"billing,vip"},"finish_reason":"stop"}]}`))
	}))
	defer server.Close()

	db, err := gorm.Open(sqlite.Open("file::memory:"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	require.NoError(t, db.AutoMigrate(
		&model.Account{},
		&model.CaptainAssistant{},
		&model.Conversation{},
		&model.Message{},
	))
	t.Cleanup(func() {
		sqlDB, dbErr := db.DB()
		require.NoError(t, dbErr)
		require.NoError(t, sqlDB.Close())
	})

	account := &model.Account{
		Name: "Feature Models",
		CaptainModels: datatypes.JSON([]byte(`{
			"editor":"editor-model",
			"copilot":"copilot-model",
			"assistant":"assistant-model",
			"label_suggestion":"label-model"
		}`)),
	}
	require.NoError(t, db.Create(account).Error)
	assistant := &model.CaptainAssistant{
		AccountID: account.ID,
		Name:      "Feature Assistant",
		Status:    model.AssistantStatusActive,
		Config:    json.RawMessage(`{"temperature":0.2}`),
	}
	require.NoError(t, db.Create(assistant).Error)
	conversation := &model.Conversation{AccountID: account.ID, Status: "open"}
	require.NoError(t, db.Create(conversation).Error)
	require.NoError(t, db.Create(&model.Message{
		AccountID:      account.ID,
		ConversationID: conversation.ID,
		MessageType:    string(model.MessageTypeIncoming),
		SenderType:     "contact",
		Content:        "I need billing help",
	}).Error)

	accountRepo := repository.NewAccountRepo(db)
	manager := llm.NewProviderManager()
	manager.SetAccountModelResolver(func(ctx context.Context, accountID uint, feature string) (string, error) {
		current, findErr := accountRepo.FindByID(ctx, accountID)
		if findErr != nil {
			return "", findErr
		}
		models := map[string]string{}
		require.NoError(t, json.Unmarshal(current.CaptainModels, &models))
		return models[feature], nil
	})
	require.NoError(t, manager.Configure(llm.RuntimeProviderConfig{
		ChatProvider:  "openai",
		ChatBaseURL:   server.URL,
		ChatAPIKey:    "test-key",
		ChatModel:     "platform-model",
		EmbeddingMode: llm.EmbeddingModeReuseChat,
		Temperature:   0.9,
		MaxTokens:     777,
	}))

	editorService := NewCaptainTaskService(nil, nil, nil, nil, nil, manager, nil)
	_, err = editorService.Rewrite(context.Background(), account.ID, &TaskRewriteRequest{Content: "draft", Operation: "improve"})
	require.NoError(t, err)

	copilotService := NewCopilotService(nil, nil, nil, manager)
	_, err = copilotService.SummarizeConversation(context.Background(), account.ID, "customer conversation")
	require.NoError(t, err)

	assistantService := NewCaptainAssistantService(repository.NewCaptainAssistantRepo(db), nil, nil, nil, manager)
	_, err = assistantService.GenerateResponse(context.Background(), assistant.ID, "help me")
	require.NoError(t, err)

	labelService := NewCaptainTaskExtendedService(
		repository.NewConversationRepo(db),
		repository.NewMessageRepo(db),
		nil,
		nil,
		manager,
	)
	_, err = labelService.LabelSuggestion(context.Background(), account.ID, &ChatwootLabelSuggestionRequest{
		ConversationDisplayID: conversation.ID,
	})
	require.NoError(t, err)

	mu.Lock()
	defer mu.Unlock()
	require.Len(t, captured, 4)
	assert.Equal(t, []string{"editor-model", "copilot-model", "assistant-model", "label-model"}, []string{
		captured[0].Model,
		captured[1].Model,
		captured[2].Model,
		captured[3].Model,
	})
	assert.Equal(t, 0.9, captured[0].Temperature)
	assert.Equal(t, 0.9, captured[1].Temperature)
	assert.Equal(t, 0.2, captured[2].Temperature)
	assert.Equal(t, 0.9, captured[3].Temperature)
	for _, request := range captured {
		assert.Equal(t, 777, request.MaxTokens)
	}
}