mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-09 12:54:47 +08:00
563d855780
### What problem does this PR solve? Implement OpenAI chat completions in GO POST /api/v1/openai/<chat_id>/chat/completions OpenAI chat cli: internal/development.md ### Type of change - [x] Refactoring
712 lines
20 KiB
Go
712 lines
20 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
|
|
"ragflow/internal/entity"
|
|
)
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Fake implementations
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type fakeSessionStore struct {
|
|
mu sync.Mutex
|
|
sessions map[string]*entity.ChatSession
|
|
dialogs map[string]*entity.Chat
|
|
dialogExists map[string]bool // key: tenantID|chatID
|
|
getByIDErr error
|
|
createErr error
|
|
updateByIDErr error
|
|
deleteByIDErr error
|
|
getDialogErr error
|
|
// record calls
|
|
createCalled []*entity.ChatSession
|
|
updateCalled []struct {
|
|
id string
|
|
updates map[string]interface{}
|
|
}
|
|
deleteByIDIDs []string
|
|
}
|
|
|
|
func newFakeSessionStore() *fakeSessionStore {
|
|
return &fakeSessionStore{
|
|
sessions: make(map[string]*entity.ChatSession),
|
|
dialogs: make(map[string]*entity.Chat),
|
|
dialogExists: make(map[string]bool),
|
|
}
|
|
}
|
|
|
|
func (f *fakeSessionStore) GetByID(id string) (*entity.ChatSession, error) {
|
|
if f.getByIDErr != nil {
|
|
return nil, f.getByIDErr
|
|
}
|
|
s, ok := f.sessions[id]
|
|
if !ok {
|
|
return nil, errors.New("record not found")
|
|
}
|
|
return s, nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) Create(conv *entity.ChatSession) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if f.createErr != nil {
|
|
return f.createErr
|
|
}
|
|
f.sessions[conv.ID] = conv
|
|
f.createCalled = append(f.createCalled, conv)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) UpdateByID(id string, updates map[string]interface{}) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if f.updateByIDErr != nil {
|
|
return f.updateByIDErr
|
|
}
|
|
f.updateCalled = append(f.updateCalled, struct {
|
|
id string
|
|
updates map[string]interface{}
|
|
}{id, updates})
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) DeleteByID(id string) error {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if f.deleteByIDErr != nil {
|
|
return f.deleteByIDErr
|
|
}
|
|
f.deleteByIDIDs = append(f.deleteByIDIDs, id)
|
|
delete(f.sessions, id)
|
|
return nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) ListByChatID(chatID string) ([]*entity.ChatSession, error) {
|
|
var result []*entity.ChatSession
|
|
for _, s := range f.sessions {
|
|
if s.DialogID == chatID {
|
|
result = append(result, s)
|
|
}
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) GetDialogByID(chatID string) (*entity.Chat, error) {
|
|
if f.getDialogErr != nil {
|
|
return nil, f.getDialogErr
|
|
}
|
|
d, ok := f.dialogs[chatID]
|
|
if !ok {
|
|
return nil, errors.New("dialog not found")
|
|
}
|
|
return d, nil
|
|
}
|
|
|
|
func (f *fakeSessionStore) CheckDialogExists(tenantID, chatID string) (bool, error) {
|
|
key := tenantID + "|" + chatID
|
|
return f.dialogExists[key], nil
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type fakeTenantStore struct {
|
|
tenantIDs []string
|
|
err error
|
|
}
|
|
|
|
func (f *fakeTenantStore) GetTenantIDsByUserID(userID string) ([]string, error) {
|
|
return f.tenantIDs, f.err
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
|
|
type fakePipeline struct {
|
|
resultChan <-chan AsyncChatResult
|
|
err error
|
|
}
|
|
|
|
func (f *fakePipeline) AsyncChat(ctx context.Context, chat *entity.Chat, messages []map[string]interface{}, stream bool, kwargs map[string]interface{}) (<-chan AsyncChatResult, error) {
|
|
return f.resultChan, f.err
|
|
}
|
|
|
|
func makeResultChan(results ...AsyncChatResult) <-chan AsyncChatResult {
|
|
ch := make(chan AsyncChatResult, len(results))
|
|
for _, r := range results {
|
|
ch <- r
|
|
}
|
|
close(ch)
|
|
return ch
|
|
}
|
|
|
|
// ===================================================================
|
|
// SetChatSession tests
|
|
// ===================================================================
|
|
|
|
func TestSetChatSession_CreateNew(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
dialog := &entity.Chat{ID: "dialog-1", PromptConfig: entity.JSONMap{"prologue": "Welcome!"}}
|
|
store.dialogs["dialog-1"] = dialog
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
resp, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
DialogID: "dialog-1",
|
|
IsNew: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.ID == "" {
|
|
t.Fatal("expected session ID to be generated")
|
|
}
|
|
if resp.DialogID != "dialog-1" {
|
|
t.Fatalf("expected dialog_id=dialog-1, got %s", resp.DialogID)
|
|
}
|
|
if len(store.createCalled) != 1 {
|
|
t.Fatalf("expected 1 Create call, got %d", len(store.createCalled))
|
|
}
|
|
|
|
// Verify prologue is in the message
|
|
var msgObj map[string]interface{}
|
|
if err := json.Unmarshal(store.createCalled[0].Message, &msgObj); err != nil {
|
|
t.Fatalf("failed to unmarshal message: %v", err)
|
|
}
|
|
msgs, _ := msgObj["messages"].([]interface{})
|
|
if len(msgs) != 1 {
|
|
t.Fatalf("expected 1 initial message, got %d", len(msgs))
|
|
}
|
|
firstMsg, _ := msgs[0].(map[string]interface{})
|
|
if firstMsg["role"] != "assistant" || firstMsg["content"] != "Welcome!" {
|
|
t.Fatalf("unexpected prologue message: %#v", firstMsg)
|
|
}
|
|
}
|
|
|
|
func TestSetChatSession_CreateNewDefaultPrologue(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.dialogs["dialog-1"] = &entity.Chat{ID: "dialog-1"}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
resp, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
DialogID: "dialog-1",
|
|
IsNew: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.ID == "" {
|
|
t.Fatal("expected session ID")
|
|
}
|
|
// Default prologue
|
|
var msgObj map[string]interface{}
|
|
json.Unmarshal(store.createCalled[0].Message, &msgObj)
|
|
msgs, _ := msgObj["messages"].([]interface{})
|
|
firstMsg, _ := msgs[0].(map[string]interface{})
|
|
if !strings.Contains(firstMsg["content"].(string), "Hi! I'm your assistant") {
|
|
t.Fatalf("expected default prologue, got %q", firstMsg["content"])
|
|
}
|
|
}
|
|
|
|
func TestSetChatSession_CreateNewDialogNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
DialogID: "nonexistent",
|
|
IsNew: true,
|
|
})
|
|
if err == nil || err.Error() != "Dialog not found" {
|
|
t.Fatalf("expected 'Dialog not found' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestSetChatSession_UpdateExisting(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1", Name: strPtr("old name"),
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
resp, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
SessionID: "session-1",
|
|
Name: "new name",
|
|
IsNew: false,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.ID != "session-1" {
|
|
t.Fatalf("expected session-1, got %s", resp.ID)
|
|
}
|
|
if len(store.updateCalled) != 1 {
|
|
t.Fatalf("expected UpdateByID call, got %d", len(store.updateCalled))
|
|
}
|
|
}
|
|
|
|
func TestSetChatSession_UpdateNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.updateByIDErr = errors.New("Chat session not found")
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
SessionID: "missing",
|
|
IsNew: false,
|
|
})
|
|
if err == nil || err.Error() != "Chat session not found" {
|
|
t.Fatalf("expected 'Chat session not found' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestSetChatSession_NameTruncation(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.dialogs["dialog-1"] = &entity.Chat{ID: "dialog-1"}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
longName := strings.Repeat("x", 300)
|
|
resp, err := svc.SetChatSession("user-1", &SetChatSessionRequest{
|
|
DialogID: "dialog-1",
|
|
Name: longName,
|
|
IsNew: true,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp.Name == nil || len(*resp.Name) > 255 {
|
|
t.Fatalf("expected name truncated to <=255, got len=%d", len(*resp.Name))
|
|
}
|
|
}
|
|
|
|
// ===================================================================
|
|
// RemoveChatSessions tests
|
|
// ===================================================================
|
|
|
|
func TestRemoveChatSessions_Success(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["conv-1"] = &entity.ChatSession{ID: "conv-1", DialogID: "dialog-1"}
|
|
store.sessions["conv-2"] = &entity.ChatSession{ID: "conv-2", DialogID: "dialog-1"}
|
|
store.dialogExists["user-1|dialog-1"] = true
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-1"}},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
err := svc.RemoveChatSessions("user-1", []string{"conv-1", "conv-2"})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if len(store.deleteByIDIDs) != 2 {
|
|
t.Fatalf("expected 2 deletes, got %d", len(store.deleteByIDIDs))
|
|
}
|
|
}
|
|
|
|
func TestRemoveChatSessions_SessionNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-1"}},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
err := svc.RemoveChatSessions("user-1", []string{"missing"})
|
|
if err == nil || !strings.Contains(err.Error(), "not found") {
|
|
t.Fatalf("expected 'not found' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRemoveChatSessions_NotOwner(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["conv-1"] = &entity.ChatSession{ID: "conv-1", DialogID: "dialog-1"}
|
|
// No tenant matches — dialogExists stays false for all combinations
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-other"}},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
err := svc.RemoveChatSessions("user-1", []string{"conv-1"})
|
|
if err == nil || !strings.Contains(err.Error(), "Only owner") {
|
|
t.Fatalf("expected 'Only owner' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
// ===================================================================
|
|
// ListChatSessions tests
|
|
// ===================================================================
|
|
|
|
func TestListChatSessions_Success(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["s1"] = &entity.ChatSession{ID: "s1", DialogID: "chat-1"}
|
|
store.sessions["s2"] = &entity.ChatSession{ID: "s2", DialogID: "chat-1"}
|
|
store.dialogExists["tenant-1|chat-1"] = true
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-1"}},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
resp, err := svc.ListChatSessions("user-1", "chat-1")
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if len(resp.Sessions) != 2 {
|
|
t.Fatalf("expected 2 sessions, got %d", len(resp.Sessions))
|
|
}
|
|
}
|
|
|
|
func TestListChatSessions_NotOwner(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{tenantIDs: []string{"tenant-other"}},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.ListChatSessions("user-1", "chat-1")
|
|
if err == nil || !strings.Contains(err.Error(), "only owner") {
|
|
t.Fatalf("expected 'only owner' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
// ===================================================================
|
|
// Completion tests
|
|
// ===================================================================
|
|
|
|
func TestCompletion_Success(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
session := &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
store.sessions["session-1"] = session
|
|
store.dialogs["dialog-1"] = &entity.Chat{
|
|
ID: "dialog-1", TenantID: "tenant-1", LLMID: "chat@factory",
|
|
LLMSetting: entity.JSONMap{},
|
|
}
|
|
|
|
pipeline := &fakePipeline{
|
|
resultChan: makeResultChan(
|
|
AsyncChatResult{Answer: "Hello", Reference: map[string]interface{}{"chunks": []interface{}{}}},
|
|
AsyncChatResult{Answer: " world", Final: true, Reference: map[string]interface{}{"chunks": []interface{}{}}},
|
|
),
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: pipeline,
|
|
}
|
|
|
|
result, err := svc.Completion("user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1")
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
ans, _ := result["answer"].(string)
|
|
if ans != "Hello world" {
|
|
t.Fatalf("expected answer 'Hello world', got %q", ans)
|
|
}
|
|
}
|
|
|
|
func TestCompletion_EmptyMessages(t *testing.T) {
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: &fakeSessionStore{},
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.Completion("user-1", "session-1", nil, "", nil, "msg-1")
|
|
if err == nil || err.Error() != "messages cannot be empty" {
|
|
t.Fatalf("expected 'messages cannot be empty', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletion_LastMessageNotFromUser(t *testing.T) {
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: &fakeSessionStore{},
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.Completion("user-1", "session-1", []map[string]interface{}{
|
|
{"role": "assistant", "content": "hello"},
|
|
}, "", nil, "msg-1")
|
|
if err == nil || !strings.Contains(err.Error(), "not from user") {
|
|
t.Fatalf("expected 'not from user' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletion_ConversationNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.Completion("user-1", "missing", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1")
|
|
if err == nil || err.Error() != "Conversation not found" {
|
|
t.Fatalf("expected 'Conversation not found', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletion_DialogNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
_, err := svc.Completion("user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1")
|
|
if err == nil || err.Error() != "Dialog not found" {
|
|
t.Fatalf("expected 'Dialog not found', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletion_PipelineError(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
store.dialogs["dialog-1"] = &entity.Chat{
|
|
ID: "dialog-1", TenantID: "tenant-1", LLMID: "chat@factory",
|
|
LLMSetting: entity.JSONMap{},
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{err: errors.New("model unavailable")},
|
|
}
|
|
|
|
_, err := svc.Completion("user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1")
|
|
if err == nil || err.Error() != "model unavailable" {
|
|
t.Fatalf("expected 'model unavailable' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
// ===================================================================
|
|
// CompletionStream tests
|
|
// ===================================================================
|
|
|
|
func readStreamChan(ch <-chan string, n int) []string {
|
|
var msgs []string
|
|
for i := 0; i < n; i++ {
|
|
select {
|
|
case msg, ok := <-ch:
|
|
if !ok {
|
|
return msgs
|
|
}
|
|
msgs = append(msgs, msg)
|
|
default:
|
|
return msgs
|
|
}
|
|
}
|
|
return msgs
|
|
}
|
|
|
|
func TestCompletionStream_Success(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
store.dialogs["dialog-1"] = &entity.Chat{
|
|
ID: "dialog-1", TenantID: "tenant-1", LLMID: "chat@factory",
|
|
LLMSetting: entity.JSONMap{},
|
|
}
|
|
|
|
pipeline := &fakePipeline{
|
|
resultChan: makeResultChan(
|
|
AsyncChatResult{Answer: "stream", Reference: map[string]interface{}{"chunks": []interface{}{}}},
|
|
AsyncChatResult{Answer: " answer", Reference: map[string]interface{}{"chunks": []interface{}{}}},
|
|
),
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: pipeline,
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1", streamChan)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
// Should receive data events and final signal
|
|
msgs := readStreamChan(streamChan, 5)
|
|
if len(msgs) < 3 {
|
|
t.Fatalf("expected at least 3 stream messages, got %d: %v", len(msgs), msgs)
|
|
}
|
|
// Check final signal
|
|
finalFound := false
|
|
for _, m := range msgs {
|
|
if strings.Contains(m, `"data":true`) {
|
|
finalFound = true
|
|
break
|
|
}
|
|
}
|
|
if !finalFound {
|
|
t.Fatal("expected final=true signal in stream")
|
|
}
|
|
}
|
|
|
|
func TestCompletionStream_EmptyMessages(t *testing.T) {
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: &fakeSessionStore{},
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "session-1", nil, "", nil, "msg-1", streamChan)
|
|
if err == nil || err.Error() != "messages cannot be empty" {
|
|
t.Fatalf("expected 'messages cannot be empty', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletionStream_LastMessageNotFromUser(t *testing.T) {
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: &fakeSessionStore{},
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "session-1", []map[string]interface{}{
|
|
{"role": "assistant", "content": "hello"},
|
|
}, "", nil, "msg-1", streamChan)
|
|
if err == nil || !strings.Contains(err.Error(), "not from user") {
|
|
t.Fatalf("expected 'not from user' error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletionStream_ConversationNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "missing", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1", streamChan)
|
|
if err == nil || err.Error() != "Conversation not found" {
|
|
t.Fatalf("expected 'Conversation not found', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletionStream_DialogNotFound(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{},
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1", streamChan)
|
|
if err == nil || err.Error() != "Dialog not found" {
|
|
t.Fatalf("expected 'Dialog not found', got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestCompletionStream_PipelineError(t *testing.T) {
|
|
store := newFakeSessionStore()
|
|
store.sessions["session-1"] = &entity.ChatSession{
|
|
ID: "session-1", DialogID: "dialog-1",
|
|
Message: json.RawMessage(`{"messages":[]}`),
|
|
Reference: json.RawMessage(`[]`),
|
|
}
|
|
store.dialogs["dialog-1"] = &entity.Chat{
|
|
ID: "dialog-1", TenantID: "tenant-1", LLMID: "chat@factory",
|
|
LLMSetting: entity.JSONMap{},
|
|
}
|
|
|
|
svc := &ChatSessionService{
|
|
chatSessionDAO: store,
|
|
userTenantDAO: &fakeTenantStore{},
|
|
pipeline: &fakePipeline{err: errors.New("model unavailable")},
|
|
}
|
|
|
|
streamChan := make(chan string, 10)
|
|
err := svc.CompletionStream(context.Background(), "user-1", "session-1", []map[string]interface{}{
|
|
{"role": "user", "content": "hi"},
|
|
}, "", nil, "msg-1", streamChan)
|
|
if err == nil || err.Error() != "model unavailable" {
|
|
t.Fatalf("expected 'model unavailable' error, got %v", err)
|
|
}
|
|
}
|