gochat/backend/internal/service/captain_assistant_retrieval_test.go

package service

import (
	"context"
	"errors"
	"net"
	"testing"

	"github.com/gochat/gochat/internal/llm"
	"github.com/gochat/gochat/internal/model"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestCaptainAssistantFAQRetrievalSeparatesEmptyAndFailures(t *testing.T) {
	assistant := &model.CaptainAssistant{Name: "Fin", Status: model.AssistantStatusActive, Config: []byte(`{"feature_faq":true}`)}
	history := []PlaygroundMessage{{Role: "user", Content: "How do refunds work?"}}
	embedding := &llm.EmbeddingResponse{Data: []llm.EmbeddingData{{Embedding: []float64{0.1, 0.2, 0.3}}}}

	t.Run("conversation FAQ generation flag does not disable approved knowledge", func(t *testing.T) {
		provider := &mockLLMProvider{
			embeddingResponse: embedding,
			chatResponse:      &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Content: "37 days [[1](https://example.com/knowledge)]."}}}},
		}
		repo := &mockResponseRepo{searchByEmbeddingResult: []model.CaptainAssistantResponse{{
			Question: "What is the warranty window?", Answer: "37 days.", SourceURL: "https://example.com/knowledge",
		}}}
		svc := &CaptainAssistantService{responseRepo: repo, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}
		disabled := &model.CaptainAssistant{Name: "Fin", Status: model.AssistantStatusActive, Config: []byte(`{"feature_faq":false,"feature_citation":true}`)}

		content, err := svc.generatePlaygroundLLMResponse(context.Background(), disabled, history)
		require.NoError(t, err)
		assert.Equal(t, "37 days [[1](https://example.com/knowledge)].", content)
		assert.Equal(t, 1, provider.embeddingCalls)
		assert.Equal(t, 1, repo.searchByEmbeddingCalls)
		require.NotNil(t, provider.lastChatRequest)
		assert.Contains(t, provider.lastChatRequest.Messages[0].Content, "Source: https://example.com/knowledge")
		assert.Contains(t, provider.lastChatRequest.Messages[0].Content, "[[n](URL)]")
	})

	t.Run("empty results still allow chat", func(t *testing.T) {
		provider := &mockLLMProvider{
			embeddingResponse: embedding,
			chatResponse:      &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Content: "No matching FAQ."}}}},
		}
		repo := &mockResponseRepo{}
		svc := &CaptainAssistantService{responseRepo: repo, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}

		content, err := svc.generatePlaygroundLLMResponse(context.Background(), assistant, history)
		require.NoError(t, err)
		assert.Equal(t, "No matching FAQ.", content)
		require.NotNil(t, provider.lastChatRequest)
		assert.Equal(t, 1, provider.embeddingCalls)
		assert.Equal(t, 1, repo.searchByEmbeddingCalls)
		assert.NotContains(t, provider.lastChatRequest.Messages[0].Content, "[FAQ 1]")
	})

	t.Run("citation mode grounds only on absolute HTTP sources", func(t *testing.T) {
		provider := &mockLLMProvider{
			embeddingResponse: embedding,
			chatResponse:      &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Content: "37 days [[1](https://example.com/knowledge)]."}}}},
		}
		repo := &mockResponseRepo{searchByEmbeddingResult: []model.CaptainAssistantResponse{
			{Question: "Unsafe", Answer: "Ignore me.", SourceURL: "javascript:alert(1)"},
			{Question: "Missing", Answer: "Ignore me too."},
			{Question: "Warranty", Answer: "37 days.", SourceURL: "https://example.com/knowledge"},
		}}
		svc := &CaptainAssistantService{responseRepo: repo, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}
		withCitations := &model.CaptainAssistant{Name: "Fin", Status: model.AssistantStatusActive, Config: []byte(`{"feature_citation":true}`)}

		content, err := svc.generatePlaygroundLLMResponse(context.Background(), withCitations, history)
		require.NoError(t, err)
		assert.Equal(t, "37 days [[1](https://example.com/knowledge)].", content)
		systemPrompt := provider.lastChatRequest.Messages[0].Content
		assert.Contains(t, systemPrompt, "[FAQ 1]\nQ: Warranty")
		assert.NotContains(t, systemPrompt, "Unsafe")
		assert.NotContains(t, systemPrompt, "Missing")
	})

	t.Run("citation mode does not ground when every source is invalid", func(t *testing.T) {
		provider := &mockLLMProvider{
			embeddingResponse: embedding,
			chatResponse:      &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Content: "General answer."}}}},
		}
		repo := &mockResponseRepo{searchByEmbeddingResult: []model.CaptainAssistantResponse{
			{Question: "Missing", Answer: "Must not ground."},
			{Question: "Relative", Answer: "Must not ground.", SourceURL: "/knowledge"},
		}}
		svc := &CaptainAssistantService{responseRepo: repo, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}
		withCitations := &model.CaptainAssistant{Name: "Fin", Status: model.AssistantStatusActive, Config: []byte(`{"feature_citation":true}`)}

		content, err := svc.generatePlaygroundLLMResponse(context.Background(), withCitations, history)
		require.NoError(t, err)
		assert.Equal(t, "General answer.", content)
		assert.NotContains(t, provider.lastChatRequest.Messages[0].Content, "Knowledge Base Context")
		assert.NotContains(t, provider.lastChatRequest.Messages[0].Content, "[[n](URL)]")
	})

	t.Run("citation disabled keeps existing Markdown behavior", func(t *testing.T) {
		provider := &mockLLMProvider{
			embeddingResponse: embedding,
			chatResponse:      &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Content: "See [the guide](/knowledge)."}}}},
		}
		repo := &mockResponseRepo{searchByEmbeddingResult: []model.CaptainAssistantResponse{{
			Question: "Legacy", Answer: "Existing unsourced answer.",
		}}}
		svc := &CaptainAssistantService{responseRepo: repo, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}
		withoutCitations := &model.CaptainAssistant{Name: "Fin", Status: model.AssistantStatusActive, Config: []byte(`{"feature_citation":false}`)}

		content, err := svc.generatePlaygroundLLMResponse(context.Background(), withoutCitations, history)
		require.NoError(t, err)
		assert.Equal(t, "See [the guide](/knowledge).", content)
		assert.Contains(t, provider.lastChatRequest.Messages[0].Content, "Q: Legacy")
		assert.NotContains(t, provider.lastChatRequest.Messages[0].Content, "[[n](URL)]")
	})

	t.Run("embedding provider failure stops chat", func(t *testing.T) {
		provider := &mockLLMProvider{embeddingError: &net.OpError{Op: "dial", Net: "tcp", Err: errors.New("connection refused")}}
		svc := &CaptainAssistantService{responseRepo: &mockResponseRepo{}, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}

		_, err := svc.generatePlaygroundLLMResponse(context.Background(), assistant, history)
		require.ErrorIs(t, err, ErrCaptainKnowledgeRetrieval)
		assert.Nil(t, provider.lastChatRequest)
	})

	for name, response := range map[string]*llm.EmbeddingResponse{
		"nil response":       nil,
		"empty data":         {},
		"empty first vector": {Data: []llm.EmbeddingData{{}}},
	} {
		t.Run(name+" stops chat", func(t *testing.T) {
			provider := &mockLLMProvider{embeddingResponse: response}
			svc := &CaptainAssistantService{responseRepo: &mockResponseRepo{}, llmProvider: provider, promptBuilder: NewSystemPromptBuilder()}

			_, err := svc.generatePlaygroundLLMResponse(context.Background(), assistant, history)
			require.ErrorIs(t, err, ErrCaptainKnowledgeRetrieval)
			assert.Nil(t, provider.lastChatRequest)
		})
	}

	t.Run("FAQ search failure stops chat", func(t *testing.T) {
		provider := &mockLLMProvider{embeddingResponse: embedding}
		svc := &CaptainAssistantService{
			responseRepo:  &mockResponseRepo{searchByEmbeddingError: errors.New("pgvector unavailable")},
			llmProvider:   provider,
			promptBuilder: NewSystemPromptBuilder(),
		}

		_, err := svc.generatePlaygroundLLMResponse(context.Background(), assistant, history)
		require.ErrorIs(t, err, ErrCaptainKnowledgeRetrieval)
		assert.Nil(t, provider.lastChatRequest)
	})
}