gochat/internal/service/rag_service_test.go

package service

import (
	"context"
	"fmt"
	"testing"

	"github.com/pgvector/pgvector-go"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

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

// ========== Mock Repos for RAG Service Tests ==========

// mockAssistantRepo mocks AssistantRepoIface for RAG service tests.
type mockAssistantRepo struct {
	getByIDResult *model.CaptainAssistant
	getByIDError  error
}

func (m *mockAssistantRepo) GetByID(ctx context.Context, id uint) (*model.CaptainAssistant, error) {
	return m.getByIDResult, m.getByIDError
}

// mockResponseRepo mocks ResponseRepoIface for RAG service tests.
type mockResponseRepo struct {
	searchByEmbeddingResult []model.CaptainAssistantResponse
	searchByEmbeddingError  error
	getByIDResult           *model.CaptainAssistantResponse
	getByIDError            error
	updateError             error
}

func (m *mockResponseRepo) SearchByEmbedding(ctx context.Context, assistantID uint, embedding pgvector.Vector, limit int) ([]model.CaptainAssistantResponse, error) {
	return m.searchByEmbeddingResult, m.searchByEmbeddingError
}

func (m *mockResponseRepo) GetByID(ctx context.Context, id uint) (*model.CaptainAssistantResponse, error) {
	return m.getByIDResult, m.getByIDError
}

func (m *mockResponseRepo) Update(ctx context.Context, resp *model.CaptainAssistantResponse) error {
	return m.updateError
}

// ========== RAG Service Tests ==========

func setupRAGServiceWithMocks(
	mockProvider *mockLLMProvider,
	mockAssistRepo *mockAssistantRepo,
	mockRespRepo *mockResponseRepo,
) *RAGService {
	return NewRAGService(mockRespRepo, mockAssistRepo, mockProvider)
}

// --- Query ---
func TestRAGService_Query(t *testing.T) {
	mockProvider := &mockLLMProvider{}
	assistant := &model.CaptainAssistant{
		AccountID: 1,
		Name:      "RAGAssistant",
		Status:    model.AssistantStatusActive,
		Config:    marshalAssistantConfig(model.DefaultAssistantConfig()),
	}
	assistant.ID = 1

	mockAssistRepo := &mockAssistantRepo{
		getByIDResult: assistant,
	}
	mockRespRepo := &mockResponseRepo{
		searchByEmbeddingResult: []model.CaptainAssistantResponse{
			{
				AccountID:   1,
				AssistantID: 1,
				Question:    "What is GoChat?",
				Answer:      "GoChat is an open-source customer engagement platform.",
				Status:      model.ResponseStatusApproved,
			},
		},
	}
	// Set the ID on the response after creation (can't set promoted Base.ID in literal)
	mockRespRepo.searchByEmbeddingResult[0].ID = 10

	// Setup mock LLM: embedding generation returns a valid vector
	mockProvider.embeddingResponse = &llm.EmbeddingResponse{
		Data: []llm.EmbeddingData{
			{Embedding: make([]float64, 1536)},
		},
	}
	// Setup mock LLM: chat completion returns a valid answer
	mockProvider.chatResponse = &llm.ChatResponse{
		Choices: []llm.ChatChoice{
			{Message: llm.ChatMessage{Role: "assistant", Content: "GoChat is an open-source platform for customer engagement."}},
		},
	}

	svc := setupRAGServiceWithMocks(mockProvider, mockAssistRepo, mockRespRepo)

	req := &RAGQueryRequest{
		AssistantID: 1,
		Question:    "What is GoChat?",
		TopK:        5,
	}

	result, err := svc.Query(context.Background(), 1, req)
	require.NoError(t, err)
	assert.NotEmpty(t, result.Answer)
	assert.Equal(t, uint(1), result.AssistantID)
	assert.Len(t, result.Sources, 1)
}

func TestRAGService_Query_AssistantNotFound(t *testing.T) {
	mockProvider := &mockLLMProvider{}
	mockAssistRepo := &mockAssistantRepo{
		getByIDError: fmt.Errorf("record not found"),
	}
	mockRespRepo := &mockResponseRepo{}

	svc := setupRAGServiceWithMocks(mockProvider, mockAssistRepo, mockRespRepo)

	req := &RAGQueryRequest{
		AssistantID: 9999,
		Question:    "What is GoChat?",
	}

	_, err := svc.Query(context.Background(), 1, req)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "assistant not found")
}

func TestRAGService_Query_EmbeddingGenerationFails(t *testing.T) {
	mockProvider := &mockLLMProvider{
		embeddingError: fmt.Errorf("embedding service unavailable"),
	}
	assistant := &model.CaptainAssistant{
		AccountID: 1,
		Name:      "RAGAssistant",
		Status:    model.AssistantStatusActive,
		Config:    marshalAssistantConfig(model.DefaultAssistantConfig()),
	}
	assistant.ID = 1

	mockAssistRepo := &mockAssistantRepo{
		getByIDResult: assistant,
	}
	mockRespRepo := &mockResponseRepo{}

	svc := setupRAGServiceWithMocks(mockProvider, mockAssistRepo, mockRespRepo)

	req := &RAGQueryRequest{
		AssistantID: 1,
		Question:    "What is GoChat?",
	}

	_, err := svc.Query(context.Background(), 1, req)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "embedding generation failed")
}

func TestRAGService_Query_NoRelevantFAQs(t *testing.T) {
	mockProvider := &mockLLMProvider{}
	assistant := &model.CaptainAssistant{
		AccountID: 1,
		Name:      "RAGAssistant",
		Status:    model.AssistantStatusActive,
		Config:    marshalAssistantConfig(model.DefaultAssistantConfig()),
	}
	assistant.ID = 1

	mockAssistRepo := &mockAssistantRepo{
		getByIDResult: assistant,
	}
	// No FAQ results found — triggers queryWithoutContext
	mockRespRepo := &mockResponseRepo{
		searchByEmbeddingResult: []model.CaptainAssistantResponse{},
	}

	// Setup mock LLM: embedding generation returns a valid vector
	mockProvider.embeddingResponse = &llm.EmbeddingResponse{
		Data: []llm.EmbeddingData{
			{Embedding: make([]float64, 1536)},
		},
	}
	// Setup mock LLM: fallback chat completion (no context)
	mockProvider.chatResponse = &llm.ChatResponse{
		Choices: []llm.ChatChoice{
			{Message: llm.ChatMessage{Role: "assistant", Content: "Based on general knowledge, GoChat may be a communication tool."}},
		},
	}

	svc := setupRAGServiceWithMocks(mockProvider, mockAssistRepo, mockRespRepo)

	req := &RAGQueryRequest{
		AssistantID: 1,
		Question:    "What is GoChat?",
		TopK:        5,
	}

	result, err := svc.Query(context.Background(), 1, req)
	require.NoError(t, err)
	assert.NotEmpty(t, result.Answer)
	assert.Equal(t, uint(1), result.AssistantID)
	assert.Len(t, result.Sources, 0)
	// Low confidence when no context
	assert.Equal(t, 0.3, result.Confidence)
}

func TestRAGService_Query_DraftAssistant(t *testing.T) {
	mockProvider := &mockLLMProvider{}
	assistant := &model.CaptainAssistant{
		AccountID: 1,
		Name:      "RAGAssistant",
		Status:    model.AssistantStatusDraft,
		Config:    marshalAssistantConfig(model.DefaultAssistantConfig()),
	}
	assistant.ID = 1

	mockAssistRepo := &mockAssistantRepo{
		getByIDResult: assistant,
	}
	mockRespRepo := &mockResponseRepo{}

	svc := setupRAGServiceWithMocks(mockProvider, mockAssistRepo, mockRespRepo)

	req := &RAGQueryRequest{
		AssistantID: 1,
		Question:    "What is GoChat?",
	}

	_, err := svc.Query(context.Background(), 1, req)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "assistant is not active")
}