From 9db67a24564b043e0ed249a0a058f8a46d165b49 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 17 Dec 2024 10:18:37 +0800 Subject: [PATCH] Refactor chat retrieval in Neo API to support filtering and pagination - Updated GetChats method across conversation implementations to accept a ChatFilter struct, enabling keyword filtering, pagination, and ordering. - Enhanced handleChatList method in the Neo API to construct a filter from query parameters, improving the flexibility of chat retrieval. - Introduced new types (ChatFilter, ChatGroup, ChatGroupResponse) to facilitate structured responses and better organization of chat data. - Improved error handling and response structure for chat retrieval, ensuring robust feedback and clarity in API responses. --- neo/api.go | 26 ++++++-- neo/conversation/mongo.go | 10 ++- neo/conversation/redis.go | 10 ++- neo/conversation/types.go | 25 +++++++- neo/conversation/weaviate.go | 10 ++- neo/conversation/xun.go | 114 ++++++++++++++++++++++++++++++----- neo/conversation/xun_test.go | 75 +++++++++++++++++++++++ 7 files changed, 245 insertions(+), 25 deletions(-) diff --git a/neo/api.go b/neo/api.go index c830f1a0..6a214fbe 100644 --- a/neo/api.go +++ b/neo/api.go @@ -5,6 +5,7 @@ import ( "io" "net/url" "path/filepath" + "strconv" "strings" "github.com/gin-gonic/gin" @@ -12,6 +13,7 @@ import ( "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/process" "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/neo/message" ) @@ -123,17 +125,33 @@ func (neo *DSL) handleChatList(c *gin.Context) { return } - // Get keywords from query parameter - keywords := c.Query("keywords") + // Create filter from query parameters + filter := conversation.ChatFilter{ + Keywords: c.Query("keywords"), + Order: c.Query("order"), + } - list, err := neo.Conversation.GetChats(sid, keywords) + // Parse page and pagesize + if page := c.Query("page"); page != "" { + if n, err := strconv.Atoi(page); err == nil { + filter.Page = n + } + } + + if pageSize := c.Query("pagesize"); pageSize != "" { + if n, err := strconv.Atoi(pageSize); err == nil { + filter.PageSize = n + } + } + + response, err := neo.Conversation.GetChats(sid, filter) if err != nil { c.JSON(500, gin.H{"message": err.Error(), "code": 500}) c.Done() return } - c.JSON(200, map[string]interface{}{"data": list}) + c.JSON(200, map[string]interface{}{"data": response}) c.Done() } diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index f2ee634d..d129ffe8 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -14,8 +14,14 @@ func (conv *Mongo) UpdateChatTitle(sid string, cid string, title string) error { } // GetChats get the chat list -func (conv *Mongo) GetChats(sid string, keywords ...string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Mongo) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index 5515dd21..2f9639ee 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -14,8 +14,14 @@ func (conv *Redis) UpdateChatTitle(sid string, cid string, title string) error { } // GetChats get the chat list -func (conv *Redis) GetChats(sid string, keywords ...string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Redis) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 566157ec..9bf67e0c 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -15,10 +15,33 @@ type ChatInfo struct { History []map[string]interface{} `json:"history"` } +// ChatFilter represents the filter parameters for GetChats +type ChatFilter struct { + Keywords string `json:"keywords,omitempty"` + Page int `json:"page,omitempty"` // 页码,从1开始 + PageSize int `json:"pagesize,omitempty"` // 每页数量 + Order string `json:"order,omitempty"` // desc/asc +} + +// ChatGroup represents a group of chats by date +type ChatGroup struct { + Label string `json:"label"` + Chats []map[string]interface{} `json:"chats"` +} + +// ChatGroupResponse represents paginated chat groups +type ChatGroupResponse struct { + Groups []ChatGroup `json:"groups"` + Page int `json:"page"` // 当前页码 + PageSize int `json:"pagesize"` // 每页数量 + Total int64 `json:"total"` // 总记录数 + LastPage int `json:"last_page"` // 最后一页页码 +} + // Conversation the store interface type Conversation interface { UpdateChatTitle(sid string, cid string, title string) error - GetChats(sid string, keywords ...string) ([]map[string]interface{}, 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 diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 6770af4c..710f16c7 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -14,8 +14,14 @@ func (conv *Weaviate) UpdateChatTitle(sid string, cid string, title string) erro } // GetChats get the chat list -func (conv *Weaviate) GetChats(sid string, keywords ...string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil +func (conv *Weaviate) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { + return &ChatGroupResponse{ + Groups: []ChatGroup{}, + Page: filter.Page, + PageSize: filter.PageSize, + Total: 0, + LastPage: 1, + }, nil } // GetHistory get the history diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index c1dfed36..cbc436b4 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -2,6 +2,7 @@ package conversation import ( "fmt" + "math" "strings" "time" @@ -240,39 +241,124 @@ func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error { return err } -// GetChats get the chat list -func (conv *Xun) GetChats(sid string, keywords ...string) ([]map[string]interface{}, error) { +// GetChats get the chat list with grouping by date +func (conv *Xun) GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error) { userID, err := conv.getUserID(sid) if err != nil { return nil, err } + // Set defaults + if filter.PageSize <= 0 { + filter.PageSize = 100 + } + if filter.Page <= 0 { + filter.Page = 1 + } + if filter.Order == "" { + filter.Order = "desc" + } + + // Build base query qb := conv.newQueryChat(). - Select("chat_id", "title"). + Select("chat_id", "title", "created_at"). Where("sid", userID) - // Add title search if keywords provided - if len(keywords) > 0 && keywords[0] != "" { - keyword := strings.TrimSpace(keywords[0]) // Trim whitespace from keyword + // Add keyword filter + if filter.Keywords != "" { + keyword := strings.TrimSpace(filter.Keywords) if keyword != "" { qb.Where("title", "like", "%"+keyword+"%") } } - rows, err := qb.Get() + // Get total count + total, err := qb.Clone().Count() if err != nil { return nil, err } - res := []map[string]interface{}{} - for _, row := range rows { - res = append(res, map[string]interface{}{ - "chat_id": row.Get("chat_id"), - "title": row.Get("title"), - }) + // Calculate pagination + offset := (filter.Page - 1) * filter.PageSize + lastPage := int(math.Ceil(float64(total) / float64(filter.PageSize))) + + // Get paginated results + rows, err := qb.OrderBy("created_at", filter.Order). + Offset(offset). + Limit(filter.PageSize). + Get() + if err != nil { + return nil, err } - return res, nil + // Group chats by date + today := time.Now().Truncate(24 * time.Hour) + yesterday := today.AddDate(0, 0, -1) + thisWeekStart := today.AddDate(0, 0, -int(today.Weekday())) + lastWeekStart := thisWeekStart.AddDate(0, 0, -7) + + groups := map[string][]map[string]interface{}{ + "Today": {}, + "Yesterday": {}, + "This Week": {}, + "Last Week": {}, + "Even Earlier": {}, + } + + for _, row := range rows { + chat := map[string]interface{}{ + "chat_id": row.Get("chat_id"), + "title": row.Get("title"), + } + + createdAt, ok := row.Get("created_at").(time.Time) + if !ok { + // Try to parse string if it's not already time.Time + if timeStr, ok := row.Get("created_at").(string); ok { + var err error + createdAt, err = time.Parse(time.RFC3339, timeStr) + if err != nil { + continue + } + } else { + continue + } + } + + createdDate := createdAt.Truncate(24 * time.Hour) + + switch { + case createdDate.Equal(today): + groups["Today"] = append(groups["Today"], chat) + case createdDate.Equal(yesterday): + groups["Yesterday"] = append(groups["Yesterday"], chat) + case createdDate.After(thisWeekStart) || createdDate.Equal(thisWeekStart): + groups["This Week"] = append(groups["This Week"], chat) + case createdDate.After(lastWeekStart) || createdDate.Equal(lastWeekStart): + groups["Last Week"] = append(groups["Last Week"], chat) + default: + groups["Even Earlier"] = append(groups["Even Earlier"], chat) + } + } + + // Convert to ordered slice + result := []ChatGroup{} + for _, label := range []string{"Today", "Yesterday", "This Week", "Last Week", "Even Earlier"} { + if len(groups[label]) > 0 { + result = append(result, ChatGroup{ + Label: label, + Chats: groups[label], + }) + } + } + + return &ChatGroupResponse{ + Groups: result, + Page: filter.Page, + PageSize: filter.PageSize, + Total: total, + LastPage: lastPage, + }, nil } // GetHistory get the history diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index e3fc0698..28dea36d 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -1,7 +1,9 @@ package conversation import ( + "fmt" "testing" + "time" "github.com/stretchr/testify/assert" "github.com/yaoapp/gou/connector" @@ -246,3 +248,76 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) { } assert.Equal(t, 2, len(allData)) } + +func TestXunGetChats(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") + + // Drop both tables before test + err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + if err != nil { + t.Fatal(err) + } + + // Save some test chats + sid := "test_user" + messages := []map[string]interface{}{ + {"role": "user", "content": "test message"}, + } + + // Create chats with different dates + for i := 0; i < 5; i++ { + chatID := fmt.Sprintf("chat_%d", i) + // First create the chat with a title + err = conv.newQueryChat().Insert(map[string]interface{}{ + "chat_id": chatID, + "title": fmt.Sprintf("Test Chat %d", i), + "sid": sid, + "created_at": time.Now(), + }) + if err != nil { + t.Fatal(err) + } + + // Then save the history + err = conv.SaveHistory(sid, messages, chatID) + if err != nil { + t.Fatal(err) + } + } + + // Test getting chats with default filter + filter := ChatFilter{ + PageSize: 10, + Order: "desc", + } + groups, err := conv.GetChats(sid, filter) + if err != nil { + t.Fatal(err) + } + + assert.Greater(t, len(groups.Groups), 0) + + // Test with keywords + filter.Keywords = "test" + groups, err = conv.GetChats(sid, filter) + if err != nil { + t.Fatal(err) + } + + assert.Greater(t, len(groups.Groups), 0) +}