Files
ragflow/internal/service/chat_session_test.go
qinling0210 7862f69f39 Implement chat completions in go (#16491)
### Summary

POST   /api/v1/chat/completions
2026-07-01 15:52:52 +08:00

631 lines
17 KiB
Go

package service
import (
"context"
"encoding/json"
"errors"
"reflect"
"strings"
"sync"
"testing"
"gorm.io/gorm"
"ragflow/internal/common"
"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
getDialogErr error
// record calls
createCalled []*entity.ChatSession
updateCalled []struct {
id string
updates map[string]interface{}
}
}
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, gorm.ErrRecordNotFound
}
return s, nil
}
func (f *fakeSessionStore) GetBySessionIDAndChatID(sessionID, chatID string) (*entity.ChatSession, error) {
s, err := f.GetByID(sessionID)
if err != nil {
return nil, err
}
if s.DialogID != chatID {
return nil, gorm.ErrRecordNotFound
}
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
}
s, ok := f.sessions[id]
if !ok {
return gorm.ErrRecordNotFound
}
f.updateCalled = append(f.updateCalled, struct {
id string
updates map[string]interface{}
}{id, updates})
for k, v := range updates {
switch k {
case "name":
if str, ok := v.(string); ok {
s.Name = &str
}
case "message":
if raw, ok := v.([]byte); ok {
s.Message = append(json.RawMessage(nil), raw...)
}
case "reference":
if raw, ok := v.([]byte); ok {
s.Reference = append(json.RawMessage(nil), raw...)
}
}
}
return nil
}
func (f *fakeSessionStore) DeleteByID(id string) error {
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
}
// ===================================================================
// 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)
}
}
// ===================================================================
// GetSession / UpdateSession tests
// ===================================================================
func TestGetSession_Success(t *testing.T) {
store := newFakeSessionStore()
store.sessions["session-1"] = &entity.ChatSession{
ID: "session-1",
DialogID: "chat-1",
Name: strPtr("session"),
Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`),
Reference: json.RawMessage(`[
{"chunks":[{"chunk_id":"chunk-1","content_with_weight":"hello","doc_id":"doc-1","docnm_kwd":"Doc 1","kb_id":"kb-1"}]},
[]
]`),
UserID: strPtr("user-1"),
}
icon := "avatar.png"
store.dialogs["chat-1"] = &entity.Chat{ID: "chat-1", Icon: &icon}
store.dialogExists["user-1|chat-1"] = true
svc := &ChatSessionService{
chatSessionDAO: store,
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
resp, code, err := svc.GetSession("user-1", "chat-1", "session-1")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("unexpected code: %v", code)
}
if resp.ChatID != "chat-1" {
t.Fatalf("chat_id=%q", resp.ChatID)
}
if resp.Avatar == nil || *resp.Avatar != "avatar.png" {
t.Fatalf("avatar=%v", resp.Avatar)
}
if len(resp.Messages) != 1 || resp.Messages[0]["content"] != "hello" {
t.Fatalf("messages=%#v", resp.Messages)
}
if len(resp.Reference) != 2 {
t.Fatalf("reference len=%d", len(resp.Reference))
}
firstRef, ok := resp.Reference[0].(map[string]interface{})
if !ok {
t.Fatalf("reference[0] type=%T", resp.Reference[0])
}
chunks, ok := firstRef["chunks"].([]FormattedChunk)
if !ok {
t.Fatalf("chunks type=%T", firstRef["chunks"])
}
if len(chunks) != 1 || chunks[0].ID != "chunk-1" {
t.Fatalf("chunks=%#v", chunks)
}
if _, ok := resp.Reference[1].([]interface{}); !ok {
t.Fatalf("reference[1] changed unexpectedly: %T", resp.Reference[1])
}
}
func TestGetSession_NotOwner(t *testing.T) {
svc := &ChatSessionService{
chatSessionDAO: newFakeSessionStore(),
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
_, code, err := svc.GetSession("user-1", "chat-1", "session-1")
if err == nil || err.Error() != "No authorization." {
t.Fatalf("err=%v", err)
}
if code != common.CodeAuthenticationError {
t.Fatalf("code=%v", code)
}
}
func TestGetSession_WrongChat(t *testing.T) {
store := newFakeSessionStore()
store.sessions["session-1"] = &entity.ChatSession{ID: "session-1", DialogID: "chat-2"}
store.dialogExists["user-1|chat-1"] = true
svc := &ChatSessionService{
chatSessionDAO: store,
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
_, code, err := svc.GetSession("user-1", "chat-1", "session-1")
if err == nil || err.Error() != "Session does not belong to this chat!" {
t.Fatalf("err=%v", err)
}
if code != common.CodeDataError {
t.Fatalf("code=%v", code)
}
}
func TestUpdateSession_Success(t *testing.T) {
store := newFakeSessionStore()
store.sessions["session-1"] = &entity.ChatSession{
ID: "session-1",
DialogID: "chat-1",
Name: strPtr("old"),
Message: json.RawMessage(`[{"role":"assistant","content":"hello"}]`),
}
store.dialogExists["user-1|chat-1"] = true
svc := &ChatSessionService{
chatSessionDAO: store,
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
longName := " " + strings.Repeat("x", 260) + " "
resp, code, err := svc.UpdateSession("user-1", "chat-1", "session-1", map[string]interface{}{
"name": longName,
"user_id": "spoof",
"chat_id": "spoof-chat",
})
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("code=%v", code)
}
if resp.Name == nil || len(*resp.Name) != 255 {
t.Fatalf("name=%v", resp.Name)
}
if len(store.updateCalled) != 1 {
t.Fatalf("update calls=%d", len(store.updateCalled))
}
if _, ok := store.updateCalled[0].updates["user_id"]; ok {
t.Fatalf("unexpected user_id update: %#v", store.updateCalled[0].updates)
}
if _, ok := store.updateCalled[0].updates["chat_id"]; ok {
t.Fatalf("unexpected chat_id update: %#v", store.updateCalled[0].updates)
}
if !reflect.DeepEqual(resp.Messages, []map[string]interface{}{{"role": "assistant", "content": "hello"}}) {
t.Fatalf("messages=%#v", resp.Messages)
}
}
func TestUpdateSession_ValidationErrors(t *testing.T) {
store := newFakeSessionStore()
store.sessions["session-1"] = &entity.ChatSession{ID: "session-1", DialogID: "chat-1"}
store.dialogExists["user-1|chat-1"] = true
svc := &ChatSessionService{
chatSessionDAO: store,
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
cases := []struct {
name string
req map[string]interface{}
message string
code common.ErrorCode
}{
{name: "empty body", req: map[string]interface{}{}, message: "Request body cannot be empty", code: common.CodeArgumentError},
{name: "message", req: map[string]interface{}{"message": []interface{}{}}, message: "`messages` cannot be changed.", code: common.CodeDataError},
{name: "messages", req: map[string]interface{}{"messages": []interface{}{}}, message: "`messages` cannot be changed.", code: common.CodeDataError},
{name: "reference", req: map[string]interface{}{"reference": []interface{}{}}, message: "`reference` cannot be changed.", code: common.CodeDataError},
{name: "empty name", req: map[string]interface{}{"name": " "}, message: "`name` can not be empty.", code: common.CodeDataError},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
_, code, err := svc.UpdateSession("user-1", "chat-1", "session-1", tc.req)
if err == nil || err.Error() != tc.message {
t.Fatalf("err=%v", err)
}
if code != tc.code {
t.Fatalf("code=%v", code)
}
})
}
}
func TestUpdateSession_NotFound(t *testing.T) {
store := newFakeSessionStore()
store.dialogExists["user-1|chat-1"] = true
svc := &ChatSessionService{
chatSessionDAO: store,
userTenantDAO: &fakeTenantStore{},
pipeline: &fakePipeline{},
}
_, code, err := svc.UpdateSession("user-1", "chat-1", "missing", map[string]interface{}{"name": "renamed"})
if err == nil || err.Error() != "Session not found!" {
t.Fatalf("err=%v", err)
}
if code != common.CodeDataError {
t.Fatalf("code=%v", code)
}
}
func TestParseCollections_ReturnEmptySlicesForMissingOrNull(t *testing.T) {
messageInputs := []json.RawMessage{
nil,
json.RawMessage(`null`),
json.RawMessage(`{"messages":null}`),
}
for _, input := range messageInputs {
got := parseMessages(input)
if got == nil || len(got) != 0 {
t.Fatalf("parseMessages(%s)=%#v", string(input), got)
}
}
referenceInputs := []json.RawMessage{
nil,
json.RawMessage(`null`),
}
for _, input := range referenceInputs {
got := parseReferenceList(input)
if got == nil || len(got) != 0 {
t.Fatalf("parseReferenceList(%s)=%#v", string(input), got)
}
}
}
func TestParseCollections_ReturnNilForMalformedData(t *testing.T) {
messageInputs := []json.RawMessage{
json.RawMessage(`not-json`),
json.RawMessage(`{"unexpected":[]}`),
}
for _, input := range messageInputs {
if got := parseMessages(input); got != nil {
t.Fatalf("parseMessages(%s)=%#v, want nil", string(input), got)
}
}
referenceInputs := []json.RawMessage{
json.RawMessage(`not-json`),
json.RawMessage(`{"unexpected":[]}`),
}
for _, input := range referenceInputs {
if got := parseReferenceList(input); got != nil {
t.Fatalf("parseReferenceList(%s)=%#v, want nil", string(input), got)
}
}
}
// ===================================================================
// chunksFormat tests — verifies field normalization after the rewrite.
// ===================================================================
func TestChunksFormat_NormalizesRawFieldNames(t *testing.T) {
svc := &ChatSessionService{}
ref := map[string]interface{}{
"chunks": []map[string]interface{}{
{
"chunk_id": "c1",
"content_with_weight": "hello world",
"content_ltks": "hello world ltks",
"doc_id": "d1",
"docnm_kwd": "Document 1",
"kb_id": "kb1",
"image_id": "img1",
"img_id": "img2",
"positions": []int{0, 10},
"position_int": []int{1, 11},
"doc_type_kwd": "pdf",
"similarity": 0.95,
"vector_similarity": 0.9,
"term_similarity": 0.85,
"row_id": "r1",
"url": "http://example.com",
"document_metadata": map[string]interface{}{"author": "Alice"},
},
},
}
result := svc.chunksFormat(ref)
if len(result) != 1 {
t.Fatalf("expected 1 chunk, got %d", len(result))
}
c := result[0]
if c["id"] != "c1" {
t.Fatalf("id=%v", c["id"])
}
if c["content"] != "hello world" {
t.Fatalf("content=%v", c["content"])
}
if c["document_id"] != "d1" {
t.Fatalf("document_id=%v", c["document_id"])
}
if c["document_name"] != "Document 1" {
t.Fatalf("document_name=%v", c["document_name"])
}
if c["dataset_id"] != "kb1" {
t.Fatalf("dataset_id=%v", c["dataset_id"])
}
if c["image_id"] != "img1" {
t.Fatalf("image_id=%v", c["image_id"])
}
if c["doc_type"] != "pdf" {
t.Fatalf("doc_type=%v", c["doc_type"])
}
if c["similarity"] != 0.95 {
t.Fatalf("similarity=%v", c["similarity"])
}
if c["url"] != "http://example.com" {
t.Fatalf("url=%v", c["url"])
}
pos, ok := c["positions"].([]int)
if !ok || len(pos) != 2 || pos[0] != 0 {
t.Fatalf("positions=%v (%T)", c["positions"], c["positions"])
}
// Raw keys must be normalized away.
if _, exists := c["content_with_weight"]; exists {
t.Fatal("content_with_weight should not be present after normalization")
}
if _, exists := c["content_ltks"]; exists {
t.Fatal("content_ltks should not be present after normalization")
}
}
func TestChunksFormat_PreservesAlreadyNormalizedFields(t *testing.T) {
svc := &ChatSessionService{}
ref := map[string]interface{}{
"chunks": []map[string]interface{}{
{
"id": "c2",
"content": "already normalized",
"document_id": "d2",
"document_name": "Doc 2",
},
},
}
result := svc.chunksFormat(ref)
if len(result) != 1 {
t.Fatalf("expected 1 chunk, got %d", len(result))
}
c := result[0]
if c["id"] != "c2" {
t.Fatalf("id=%v", c["id"])
}
if c["content"] != "already normalized" {
t.Fatalf("content=%v", c["content"])
}
}
func TestChunksFormat_EmptyReference(t *testing.T) {
svc := &ChatSessionService{}
if n := len(svc.chunksFormat(nil)); n != 0 {
t.Fatalf("nil ref: expected 0, got %d", n)
}
if n := len(svc.chunksFormat(map[string]interface{}{})); n != 0 {
t.Fatalf("empty ref: expected 0, got %d", n)
}
if n := len(svc.chunksFormat(map[string]interface{}{"chunks": nil})); n != 0 {
t.Fatalf("nil chunks: expected 0, got %d", n)
}
if n := len(svc.chunksFormat(map[string]interface{}{"chunks": []map[string]interface{}{}})); n != 0 {
t.Fatalf("empty chunks: expected 0, got %d", n)
}
}
func TestChunksFormat_ChunksAsInterfaceSlice(t *testing.T) {
svc := &ChatSessionService{}
ref := map[string]interface{}{
"chunks": []interface{}{
map[string]interface{}{
"chunk_id": "c3",
"content_with_weight": "from interface slice",
},
},
}
result := svc.chunksFormat(ref)
if len(result) != 1 {
t.Fatalf("expected 1 chunk, got %d", len(result))
}
if result[0]["id"] != "c3" {
t.Fatalf("id=%v", result[0]["id"])
}
if result[0]["content"] != "from interface slice" {
t.Fatalf("content=%v", result[0]["content"])
}
}
func TestChunksFormat_IgnoresNonMapItems(t *testing.T) {
svc := &ChatSessionService{}
ref := map[string]interface{}{
"chunks": []interface{}{
"not a map",
map[string]interface{}{
"chunk_id": "c4",
"content_with_weight": "valid chunk",
},
},
}
result := svc.chunksFormat(ref)
if len(result) != 1 {
t.Fatalf("expected 1 chunk (non-maps skipped), got %d", len(result))
}
if result[0]["id"] != "c4" {
t.Fatalf("id=%v", result[0]["id"])
}
}
func TestChunksFormat_UnsupportedTypeReturnsEmpty(t *testing.T) {
svc := &ChatSessionService{}
ref := map[string]interface{}{"chunks": "not a slice"}
result := svc.chunksFormat(ref)
if len(result) != 0 {
t.Fatalf("expected empty for string type, got %d", len(result))
}
}