Add chat deletion functionality across conversation implementations

- Implemented DeleteChat and DeleteAllChats methods in Mongo, Redis, Weaviate, and Xun conversation handlers to allow users to delete specific chats and all chats associated with their account. - Updated the Conversation interface in types.go to include the new deletion methods, ensuring consistent API across different storage backends. - Added unit tests for DeleteChat and DeleteAllChats methods in xun_test.go to verify functionality and ensure proper deletion of chat histories. - Enhanced error handling in deletion methods to provide robust feedback during chat removal operations.
YaoApp · Dec 17, 2024 · 7f9307d · 7f9307d
1 parent 9db67a2
commit 7f9307d
Show file tree

Hide file tree

Showing 6 changed files with 167 additions and 1 deletion.
diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go
@@ -48,3 +48,13 @@ func (conv *Mongo) SaveRequest(sid string, rid string, cid string, messages []ma
 func (conv *Mongo) GetChat(sid string, cid string) (*ChatInfo, error) {
 	return nil, nil
 }
+
+// DeleteChat deletes a specific chat and its history
+func (conv *Mongo) DeleteChat(sid string, cid string) error {
+	return nil
+}
+
+// DeleteAllChats deletes all chats and their histories for a user
+func (conv *Mongo) DeleteAllChats(sid string) error {
+	return nil
+}
diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go
@@ -48,3 +48,13 @@ func (conv *Redis) SaveRequest(sid string, rid string, cid string, messages []ma
 func (conv *Redis) GetChat(sid string, cid string) (*ChatInfo, error) {
 	return nil, nil
 }
+
+// DeleteChat deletes a specific chat and its history
+func (conv *Redis) DeleteChat(sid string, cid string) error {
+	return nil
+}
+
+// DeleteAllChats deletes all chats and their histories for a user
+func (conv *Redis) DeleteAllChats(sid string) error {
+	return nil
+}
diff --git a/neo/conversation/types.go b/neo/conversation/types.go
@@ -40,11 +40,13 @@ type ChatGroupResponse struct {
 
 // Conversation the store interface
 type Conversation interface {
-	UpdateChatTitle(sid string, cid string, title string) error
 	GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error)
 	GetChat(sid string, cid string) (*ChatInfo, error)
 	GetHistory(sid string, cid string) ([]map[string]interface{}, error)
 	SaveHistory(sid string, messages []map[string]interface{}, cid string) error
 	GetRequest(sid string, rid string) ([]map[string]interface{}, error)
 	SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error
+	DeleteChat(sid string, cid string) error
+	DeleteAllChats(sid string) error
+	UpdateChatTitle(sid string, cid string, title string) error
 }
diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go
@@ -48,3 +48,13 @@ func (conv *Weaviate) SaveRequest(sid string, rid string, cid string, messages [
 func (conv *Weaviate) GetChat(sid string, cid string) (*ChatInfo, error) {
 	return nil, nil
 }
+
+// DeleteChat deletes a specific chat and its history
+func (conv *Weaviate) DeleteChat(sid string, cid string) error {
+	return nil
+}
+
+// DeleteAllChats deletes all chats and their histories for a user
+func (conv *Weaviate) DeleteAllChats(sid string) error {
+	return nil
+}
diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go
@@ -559,6 +559,11 @@ func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) {
 		return nil, err
 	}
 
+	// Return nil if chat_id is nil (means no chat found)
+	if row.Get("chat_id") == nil {
+		return nil, nil
+	}
+
 	chat := map[string]interface{}{
 		"chat_id": row.Get("chat_id"),
 		"title":   row.Get("title"),
@@ -575,3 +580,50 @@ func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) {
 		History: history,
 	}, nil
 }
+
+// DeleteChat deletes a specific chat and its history
+func (conv *Xun) DeleteChat(sid string, cid string) error {
+	userID, err := conv.getUserID(sid)
+	if err != nil {
+		return err
+	}
+
+	// Delete history records first
+	_, err = conv.newQuery().
+		Where("sid", userID).
+		Where("cid", cid).
+		Delete()
+	if err != nil {
+		return err
+	}
+
+	// Then delete the chat
+	_, err = conv.newQueryChat().
+		Where("sid", userID).
+		Where("chat_id", cid).
+		Limit(1).
+		Delete()
+	return err
+}
+
+// DeleteAllChats deletes all chats and their histories for a user
+func (conv *Xun) DeleteAllChats(sid string) error {
+	userID, err := conv.getUserID(sid)
+	if err != nil {
+		return err
+	}
+
+	// Delete history records first
+	_, err = conv.newQuery().
+		Where("sid", userID).
+		Delete()
+	if err != nil {
+		return err
+	}
+
+	// Then delete all chats
+	_, err = conv.newQueryChat().
+		Where("sid", userID).
+		Delete()
+	return err
+}
diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go
@@ -321,3 +321,85 @@ func TestXunGetChats(t *testing.T) {
 
 	assert.Greater(t, len(groups.Groups), 0)
 }
+
+func TestXunDeleteChat(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+	defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
+	defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
+
+	conv, err := NewXun(Setting{
+		Connector: "default",
+		Table:     "__unit_test_conversation",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create a test chat
+	sid := "test_user"
+	cid := "test_chat"
+	messages := []map[string]interface{}{
+		{"role": "user", "content": "test message"},
+	}
+
+	// Save the chat and history
+	err = conv.SaveHistory(sid, messages, cid)
+	assert.Nil(t, err)
+
+	// Verify chat exists
+	chat, err := conv.GetChat(sid, cid)
+	assert.Nil(t, err)
+	assert.NotNil(t, chat)
+
+	// Delete the chat
+	err = conv.DeleteChat(sid, cid)
+	assert.Nil(t, err)
+
+	// Verify chat is deleted
+	chat, err = conv.GetChat(sid, cid)
+	assert.Nil(t, err)
+	assert.Equal(t, (*ChatInfo)(nil), chat)
+}
+
+func TestXunDeleteAllChats(t *testing.T) {
+	test.Prepare(t, config.Conf)
+	defer test.Clean()
+	defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
+	defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
+
+	conv, err := NewXun(Setting{
+		Connector: "default",
+		Table:     "__unit_test_conversation",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Create multiple test chats
+	sid := "test_user"
+	messages := []map[string]interface{}{
+		{"role": "user", "content": "test message"},
+	}
+
+	// Save multiple chats
+	for i := 0; i < 3; i++ {
+		cid := fmt.Sprintf("test_chat_%d", i)
+		err = conv.SaveHistory(sid, messages, cid)
+		assert.Nil(t, err)
+	}
+
+	// Verify chats exist
+	response, err := conv.GetChats(sid, ChatFilter{})
+	assert.Nil(t, err)
+	assert.Greater(t, response.Total, int64(0))
+
+	// Delete all chats
+	err = conv.DeleteAllChats(sid)
+	assert.Nil(t, err)
+
+	// Verify all chats are deleted
+	response, err = conv.GetChats(sid, ChatFilter{})
+	assert.Nil(t, err)
+	assert.Equal(t, int64(0), response.Total)
+}