feat(captain): align streaming fallbacks
This commit is contained in:
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -21,8 +22,10 @@ import (
|
||||
)
|
||||
|
||||
type mockCaptainTaskHandlerLLM struct {
|
||||
response *llm.ChatResponse
|
||||
err error
|
||||
response *llm.ChatResponse
|
||||
err error
|
||||
streamChunks []llm.StreamChunk
|
||||
streamErr error
|
||||
}
|
||||
|
||||
func (m *mockCaptainTaskHandlerLLM) ChatCompletion(_ context.Context, _ llm.ChatRequest) (*llm.ChatResponse, error) {
|
||||
@@ -31,7 +34,16 @@ func (m *mockCaptainTaskHandlerLLM) ChatCompletion(_ context.Context, _ llm.Chat
|
||||
func (m *mockCaptainTaskHandlerLLM) CreateEmbedding(_ context.Context, _ llm.EmbeddingRequest) (*llm.EmbeddingResponse, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (m *mockCaptainTaskHandlerLLM) ChatCompletionStream(_ context.Context, _ llm.ChatRequest, _ func(llm.StreamChunk) error) error {
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
@@ -101,3 +113,67 @@ func TestCaptainTaskHandler_Rewrite_NoProviderRawDisabled(t *testing.T) {
|
||||
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: "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: "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: "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")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user