From d43b7b89277610a5f3c5a8ca631e3deab56155af Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 18 Jan 2025 10:19:17 +0800 Subject: [PATCH] Enhance Neo API assistant message storage and structure - Removed unused utility imports and streamlined the saveChatHistory method to improve message storage efficiency. - Updated the Message struct to include new fields for assistant identification and mentions, enhancing message context and traceability. - Modified the history table structure to accommodate new assistant-related fields and mentions, ensuring comprehensive data storage. - Improved the SaveHistory method to process and store mentions alongside messages, providing better context for user interactions. These changes enhance the robustness and maintainability of the Neo API, paving the way for improved message handling and assistant functionalities. --- neo/assistant/api.go | 33 +++++++++++++------- neo/message/message.go | 29 +++++++++++------ neo/store/xun.go | 70 ++++++++++++++++++++++++++++++------------ 3 files changed, 93 insertions(+), 39 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 6ce85c89..6447f82c 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -10,7 +10,6 @@ import ( jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/fs" "github.com/yaoapp/gou/process" - "github.com/yaoapp/kun/utils" chatctx "github.com/yaoapp/yao/neo/context" chatMessage "github.com/yaoapp/yao/neo/message" ) @@ -319,15 +318,29 @@ func (ast *Assistant) streamChat( // saveChatHistory saves the chat history if storage is available func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, content *chatMessage.Content) { if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 { - storage.SaveHistory( - ctx.Sid, - []map[string]interface{}{ - {"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid}, - {"role": "assistant", "content": content.String(), "name": ctx.Sid}, + userMessage := messages[len(messages)-1] + data := []map[string]interface{}{ + { + "role": "user", + "content": userMessage.Content(), + "name": ctx.Sid, }, - ctx.ChatID, - nil, - ) + { + "role": "assistant", + "content": content.String(), + "name": ctx.Sid, + "assistant_id": ast.ID, + "assistant_name": ast.Name, + "assistant_avatar": ast.Avatar, + }, + } + + // Add mentions + if userMessage.Mentions != nil { + data[0]["mentions"] = userMessage.Mentions + } + + storage.SaveHistory(ctx.Sid, data, ctx.ChatID, ctx.Map()) } } @@ -521,8 +534,6 @@ func (ast *Assistant) withAttachments(ctx context.Context, msg *chatMessage.Mess }, }) } - - utils.Dump(contents) return contents, nil } diff --git a/neo/message/message.go b/neo/message/message.go index e9ce0cea..6f559bc9 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -16,15 +16,26 @@ import ( // Message the message type Message struct { - Text string `json:"text,omitempty"` // text content - Type string `json:"type,omitempty"` // error, text, plan, table, form, page, file, video, audio, image, markdown, json ... - Props map[string]interface{} `json:"props,omitempty"` // props for the types - IsDone bool `json:"done,omitempty"` - Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend - Attachments []Attachment `json:"attachments,omitempty"` // File attachments - Role string `json:"role,omitempty"` // user, assistant, system ... - Name string `json:"name,omitempty"` // name for the message - Data map[string]interface{} `json:"-"` + Text string `json:"text,omitempty"` // text content + Type string `json:"type,omitempty"` // error, text, plan, table, form, page, file, video, audio, image, markdown, json ... + Props map[string]interface{} `json:"props,omitempty"` // props for the types + IsDone bool `json:"done,omitempty"` + Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend + Attachments []Attachment `json:"attachments,omitempty"` // File attachments + Role string `json:"role,omitempty"` // user, assistant, system ... + Name string `json:"name,omitempty"` // name for the message + AssistantID string `json:"assistant_id,omitempty"` // assistant_id (for assistant role = assistant ) + AssistantName string `json:"assistant_name,omitempty"` // assistant_name (for assistant role = assistant ) + AssistantAvatar string `json:"assistant_avatar,omitempty"` // assistant_avatar (for assistant role = assistant ) + Mentions []Mention `json:"menions,omitempty"` // Mentions for the message ( for user role = user ) + Data map[string]interface{} `json:"-"` +} + +// Mention represents a mention +type Mention struct { + ID string `json:"assistant_id"` // assistant_id + Name string `json:"name"` // name + Avatar string `json:"avatar,omitempty"` // avatar } // Attachment represents a file attachment diff --git a/neo/store/xun.go b/neo/store/xun.go index aadb4583..1bf8a94b 100644 --- a/neo/store/xun.go +++ b/neo/store/xun.go @@ -141,6 +141,10 @@ func (conv *Xun) initHistoryTable() error { table.String("name", 200).Null().Index() table.Text("content").Null() table.JSON("context").Null() + table.String("assistant_id", 200).Null().Index() + table.String("assistant_name", 200).Null() + table.String("assistant_avatar", 200).Null() + table.JSON("mentions").Null() table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() table.TimestampTz("updated_at").Null().Index() table.TimestampTz("expired_at").Null().Index() @@ -158,7 +162,7 @@ func (conv *Xun) initHistoryTable() error { return err } - fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"} + fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "created_at", "updated_at", "expired_at"} for _, field := range fields { if !tab.HasColumn(field) { return fmt.Errorf("%s is required", field) @@ -450,7 +454,7 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e } qb := conv.newQuery(). - Select("role", "name", "content", "context", "uid", "created_at", "updated_at"). + Select("role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "uid", "created_at", "updated_at"). Where("sid", userID). Where("cid", cid). OrderBy("id", "desc") @@ -472,13 +476,17 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e res := []map[string]interface{}{} for _, row := range rows { message := map[string]interface{}{ - "role": row.Get("role"), - "name": row.Get("name"), - "content": row.Get("content"), - "context": row.Get("context"), - "uid": row.Get("uid"), - "created_at": row.Get("created_at"), - "updated_at": row.Get("updated_at"), + "role": row.Get("role"), + "name": row.Get("name"), + "content": row.Get("content"), + "context": row.Get("context"), + "assistant_id": row.Get("assistant_id"), + "assistant_name": row.Get("assistant_name"), + "assistant_avatar": row.Get("assistant_avatar"), + "mentions": row.Get("mentions"), + "uid": row.Get("uid"), + "created_at": row.Get("created_at"), + "updated_at": row.Get("updated_at"), } res = append([]map[string]interface{}{message}, res...) } @@ -551,23 +559,47 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid } } + // Process mentions if present + var mentionsRaw interface{} = nil + if mentions, ok := message["mentions"].([]interface{}); ok && len(mentions) > 0 { + mentionsRaw, err = jsoniter.MarshalToString(mentions) + if err != nil { + return err + } + } + value := map[string]interface{}{ - "role": role, - "name": "", - "content": content, - "sid": userID, - "cid": cid, - "uid": userID, - "context": contextRaw, - "created_at": now, - "updated_at": nil, - "expired_at": expiredAt, + "role": role, + "name": "", + "content": content, + "sid": userID, + "cid": cid, + "uid": userID, + "context": contextRaw, + "mentions": mentionsRaw, + "assistant_id": nil, + "assistant_name": nil, + "assistant_avatar": nil, + "created_at": now, + "updated_at": nil, + "expired_at": expiredAt, } if name, ok := message["name"].(string); ok { value["name"] = name } + // Add assistant fields if present + if assistantID, ok := message["assistant_id"].(string); ok { + value["assistant_id"] = assistantID + } + if assistantName, ok := message["assistant_name"].(string); ok { + value["assistant_name"] = assistantName + } + if assistantAvatar, ok := message["assistant_avatar"].(string); ok { + value["assistant_avatar"] = assistantAvatar + } + values = append(values, value) }