201 lines
8.2 KiB
Go
201 lines
8.2 KiB
Go
package v1
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/gochat/gochat/internal/llm"
|
|
"github.com/gochat/gochat/internal/model"
|
|
"github.com/gochat/gochat/internal/repository"
|
|
"github.com/gochat/gochat/internal/service"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/sqlite"
|
|
"gorm.io/gorm"
|
|
)
|
|
|
|
type mockCaptainTaskHandlerLLM struct {
|
|
response *llm.ChatResponse
|
|
err error
|
|
streamChunks []llm.StreamChunk
|
|
streamErr error
|
|
}
|
|
|
|
func (m *mockCaptainTaskHandlerLLM) ChatCompletion(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) {
|
|
return m.response, m.err
|
|
}
|
|
func (m *mockCaptainTaskHandlerLLM) CreateEmbedding(_ context.Context, _ llm.EmbeddingRequest) (*llm.EmbeddingResponse, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
func (m *mockCaptainTaskHandlerLLM) ChatCompletionStream(_ context.Context, _ llm.ChatRequest, onChunk func(llm.StreamChunk) error) error {
|
|
if m.streamErr != nil {
|
|
return m.streamErr
|
|
}
|
|
for _, chunk := range m.streamChunks {
|
|
if err := onChunk(chunk); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func setupCaptainTaskHandlerTest(t *testing.T, provider llm.Provider) (*CaptainTaskHandler, *gorm.DB) {
|
|
t.Helper()
|
|
gin.SetMode(gin.TestMode)
|
|
dbName := fmt.Sprintf("file:%s?mode=memory&cache=private", t.Name())
|
|
db, err := gorm.Open(sqlite.Open(dbName), &gorm.Config{})
|
|
require.NoError(t, err)
|
|
require.NoError(t, db.AutoMigrate(
|
|
&model.Account{},
|
|
&model.Conversation{},
|
|
&model.Message{},
|
|
&model.CaptainAssistant{},
|
|
&model.CaptainAssistantResponse{},
|
|
&model.CaptainCustomTool{},
|
|
&model.CaptainDocument{},
|
|
&model.CopilotSuggestionMessage{},
|
|
))
|
|
assistantRepo := repository.NewCaptainAssistantRepo(db)
|
|
responseRepo := repository.NewCaptainAssistantResponseRepo(db)
|
|
customToolRepo := repository.NewCaptainCustomToolRepo(db)
|
|
conversationRepo := repository.NewConversationRepo(db)
|
|
messageRepo := repository.NewMessageRepo(db)
|
|
suggestionRepo := repository.NewCopilotSuggestionRepo(db)
|
|
svc := service.NewCaptainTaskService(assistantRepo, responseRepo, customToolRepo, conversationRepo, messageRepo, provider, nil, suggestionRepo)
|
|
return NewCaptainTaskHandler(svc), db
|
|
}
|
|
|
|
func TestCaptainTaskHandler_Summarize_ChatwootRawPayload(t *testing.T) {
|
|
provider := &mockCaptainTaskHandlerLLM{response: &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Role: "assistant", Content: "Short summary"}}}}}
|
|
handler, db := setupCaptainTaskHandlerTest(t, provider)
|
|
displayID := uint(123)
|
|
conv := &model.Conversation{AccountID: 1, DisplayID: &displayID, Status: "open", ChannelType: "web_widget", Channel: "web_widget"}
|
|
require.NoError(t, db.Create(conv).Error)
|
|
require.NoError(t, db.Create(&model.Message{ConversationID: conv.ID, AccountID: 1, SenderType: "contact", MessageType: "incoming", Content: "Need help"}).Error)
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Params = gin.Params{{Key: "account_id", Value: "1"}}
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/summarize", bytes.NewReader([]byte(`{"conversation_display_id":123}`)))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
handler.Summarize(c)
|
|
|
|
require.Equal(t, http.StatusOK, w.Code)
|
|
var resp map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
|
assert.Equal(t, "Short summary", resp["message"])
|
|
assert.NotContains(t, resp, "success")
|
|
}
|
|
|
|
func TestCaptainTaskHandler_ReplySuggestion_UsesAccountIDRouteParam(t *testing.T) {
|
|
provider := &mockCaptainTaskHandlerLLM{response: &llm.ChatResponse{Choices: []llm.ChatChoice{{Message: llm.ChatMessage{Role: "assistant", Content: "Happy to help"}}}}}
|
|
handler, db := setupCaptainTaskHandlerTest(t, provider)
|
|
displayID := uint(3)
|
|
conv := &model.Conversation{AccountID: 1, DisplayID: &displayID, Status: "open", ChannelType: "web_widget", Channel: "web_widget"}
|
|
require.NoError(t, db.Create(conv).Error)
|
|
require.NoError(t, db.Create(&model.Message{ConversationID: conv.ID, AccountID: 1, SenderType: "contact", MessageType: "incoming", Content: "Need help"}).Error)
|
|
|
|
router := gin.New()
|
|
router.POST("/api/v1/accounts/:account_id/captain/tasks/reply_suggestion", handler.ReplySuggestion)
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/reply_suggestion", bytes.NewReader([]byte(`{"conversation_display_id":3}`)))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
router.ServeHTTP(w, req)
|
|
|
|
require.Equal(t, http.StatusOK, w.Code)
|
|
var resp map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
|
assert.Equal(t, "Happy to help", resp["message"])
|
|
}
|
|
|
|
func TestCaptainTaskHandler_Rewrite_NoProviderRawDisabled(t *testing.T) {
|
|
handler, _ := setupCaptainTaskHandlerTest(t, nil)
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Params = gin.Params{{Key: "account_id", Value: "1"}}
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/rewrite", bytes.NewReader([]byte(`{"content":"hello","operation":"professional"}`)))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
handler.Rewrite(c)
|
|
|
|
require.Equal(t, http.StatusUnprocessableEntity, w.Code)
|
|
var resp map[string]interface{}
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp))
|
|
assert.Equal(t, "Captain is disabled", resp["error"])
|
|
assert.NotContains(t, resp, "success")
|
|
}
|
|
|
|
func TestCaptainTaskHandler_StreamRewrite_NoProviderDisabledSSE(t *testing.T) {
|
|
handler, _ := setupCaptainTaskHandlerTest(t, nil)
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Params = gin.Params{{Key: "account_id", Value: "1"}}
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/rewrite/stream", bytes.NewReader([]byte(`{"content":"hello","operation":"professional"}`)))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
handler.StreamRewrite(c)
|
|
|
|
body := w.Body.String()
|
|
assert.Equal(t, "text/event-stream", w.Header().Get("Content-Type"))
|
|
assert.Contains(t, body, "event: error")
|
|
assert.Contains(t, body, `"error": "Captain is disabled"`)
|
|
assert.Contains(t, body, `"status": 422`)
|
|
assert.Contains(t, body, "event: done")
|
|
}
|
|
|
|
func TestCaptainTaskHandler_StreamSummarize_ChatwootDisplayID(t *testing.T) {
|
|
provider := &mockCaptainTaskHandlerLLM{streamChunks: []llm.StreamChunk{
|
|
{Choices: []llm.StreamChoice{{Delta: llm.StreamDelta{Content: "Short"}}}},
|
|
{Choices: []llm.StreamChoice{{Delta: llm.StreamDelta{Content: " summary"}, FinishReason: "stop"}}},
|
|
}}
|
|
handler, db := setupCaptainTaskHandlerTest(t, provider)
|
|
displayID := uint(456)
|
|
conv := &model.Conversation{AccountID: 1, DisplayID: &displayID, Status: "open", ChannelType: "web_widget", Channel: "web_widget"}
|
|
require.NoError(t, db.Create(conv).Error)
|
|
require.NoError(t, db.Create(&model.Message{ConversationID: conv.ID, AccountID: 1, SenderType: "contact", MessageType: "incoming", Content: "Need help"}).Error)
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Params = gin.Params{{Key: "account_id", Value: "1"}}
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/summarize/stream", bytes.NewReader([]byte(`{"conversation_display_id":456}`)))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
handler.StreamSummarize(c)
|
|
|
|
body := w.Body.String()
|
|
assert.Contains(t, body, "event: message")
|
|
assert.Contains(t, body, `"content": "Short"`)
|
|
assert.Contains(t, body, `"content": " summary"`)
|
|
assert.Contains(t, body, "event: done")
|
|
assert.NotContains(t, body, "success")
|
|
}
|
|
|
|
func TestCaptainTaskHandler_StreamRewrite_InvalidOperationSSE(t *testing.T) {
|
|
handler, _ := setupCaptainTaskHandlerTest(t, &mockCaptainTaskHandlerLLM{})
|
|
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Params = gin.Params{{Key: "account_id", Value: "1"}}
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/accounts/1/captain/tasks/rewrite/stream", strings.NewReader(`{"content":"hello","operation":"pirate"}`))
|
|
c.Request.Header.Set("Content-Type", "application/json")
|
|
|
|
handler.StreamRewrite(c)
|
|
|
|
body := w.Body.String()
|
|
assert.Contains(t, body, "event: error")
|
|
assert.Contains(t, body, `"error": "Invalid operation: pirate"`)
|
|
assert.Contains(t, body, `"status": 422`)
|
|
assert.Contains(t, body, "event: done")
|
|
}
|