From ffb9ede647391c94e040a0748fa11b8c9ff9b26f Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 13 Jan 2025 14:47:17 +0800 Subject: [PATCH] Refactor Neo API message handling and enhance assistant interaction - Updated the Answer method to utilize the new withHistory function, improving message retrieval and context management. - Refactored the chat method to accept a slice of message.Message instead of a map, enhancing type safety and clarity. - Introduced a new withHistory function to streamline the process of retrieving chat history and user messages, improving maintainability. - Enhanced the HookInit method to accept a gin.Context, allowing for better context handling during assistant initialization. - Updated the assistant API interface to reflect changes in message handling, ensuring consistency across the codebase. These changes improve the robustness and maintainability of the Neo API, paving the way for future enhancements in assistant functionalities and message management. --- neo/assistant/api.go | 22 +++++++-------- neo/assistant/hooks.go | 49 +++++++++++++++++++++++++------- neo/assistant/types.go | 5 ++-- neo/message/message.go | 64 ++++++++++++++++++++++++++++++++++++++++++ neo/neo.go | 31 ++++++++++---------- 5 files changed, 133 insertions(+), 38 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 8481ed9d..685e1bf2 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -7,6 +7,7 @@ import ( "strings" "github.com/yaoapp/gou/fs" + "github.com/yaoapp/yao/neo/message" chatMessage "github.com/yaoapp/yao/neo/message" ) @@ -41,7 +42,7 @@ func GetByConnector(connector string, name string) (*Assistant, error) { } // Chat implements the chat functionality -func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { +func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error { if ast.openai == nil { return fmt.Errorf("openai is not initialized") } @@ -59,13 +60,12 @@ func (ast *Assistant) Chat(ctx context.Context, messages []map[string]interface{ return nil } -func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string]interface{}) ([]map[string]interface{}, error) { +func (ast *Assistant) requestMessages(ctx context.Context, messages []message.Message) ([]map[string]interface{}, error) { newMessages := []map[string]interface{}{} - // With Prompts if ast.Prompts != nil { for _, prompt := range ast.Prompts { - message := map[string]interface{}{ + msg := map[string]interface{}{ "role": prompt.Role, "content": prompt.Content, } @@ -75,20 +75,20 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string name = prompt.Name } - message["name"] = name - newMessages = append(newMessages, message) + msg["name"] = name + newMessages = append(newMessages, msg) } } length := len(messages) for index, message := range messages { - role, ok := message["role"].(string) - if !ok { + role := message.Role + if role == "" { return nil, fmt.Errorf("role must be string") } - content, ok := message["content"].(string) - if !ok { + content := message.Text + if content == "" { return nil, fmt.Errorf("content must be string") } @@ -97,7 +97,7 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []map[string "content": content, } - if name, ok := message["name"].(string); ok { + if name := message.Name; name != "" { newMessage["name"] = name } diff --git a/neo/assistant/hooks.go b/neo/assistant/hooks.go index b9c8fb7b..57f1329b 100644 --- a/neo/assistant/hooks.go +++ b/neo/assistant/hooks.go @@ -1,8 +1,11 @@ package assistant import ( + "context" "fmt" + "time" + "github.com/gin-gonic/gin" chatctx "github.com/yaoapp/yao/neo/context" "github.com/yaoapp/yao/neo/message" ) @@ -19,8 +22,12 @@ type ResHookInit struct { } // HookInit initialize the assistant -func (ast *Assistant) HookInit(context chatctx.Context, messages []message.Message) (*ResHookInit, error) { - v, err := ast.call("Init", context, messages) +func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, messages []message.Message) (*ResHookInit, error) { + // Create timeout context + ctx, cancel := ast.createTimeoutContext(c) + defer cancel() + + v, err := ast.call(ctx, "Init", context, messages, c.Writer) if err != nil { if err.Error() == HookErrorMethodNotFound { return nil, nil @@ -50,25 +57,47 @@ func (ast *Assistant) HookInit(context chatctx.Context, messages []message.Messa return response, nil } -// Call the script method -func (ast *Assistant) call(method string, context chatctx.Context, args ...any) (interface{}, error) { +// createTimeoutContext creates a timeout context with 5 seconds timeout +func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) { + ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second) + return ctx, cancel +} +// Call the script method +func (ast *Assistant) call(ctx context.Context, method string, context chatctx.Context, args ...any) (interface{}, error) { if ast.Script == nil { return nil, nil } - ctx, err := ast.Script.NewContext(context.Sid, nil) + scriptCtx, err := ast.Script.NewContext(context.Sid, nil) if err != nil { return nil, err } - defer ctx.Close() + defer scriptCtx.Close() // Check if the method exists - if !ctx.Global().Has(method) { + if !scriptCtx.Global().Has(method) { return nil, fmt.Errorf(HookErrorMethodNotFound) } - // Call the method - args = append([]interface{}{context.Map()}, args...) - return ctx.Call(method, args...) + // Create done channel for handling cancellation + done := make(chan struct{}) + var result interface{} + var callErr error + + go func() { + defer close(done) + // Call the method + args = append([]interface{}{context.Map()}, args...) + result, callErr = scriptCtx.Call(method, args...) + }() + + // Wait for either context cancellation or method completion + select { + case <-ctx.Done(): + scriptCtx.Close() // Force close the script context + return nil, ctx.Err() + case <-done: + return result, callErr + } } diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 71dddf19..54bc47f8 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -5,6 +5,7 @@ import ( "io" "mime/multipart" + "github.com/gin-gonic/gin" "github.com/yaoapp/gou/rag/driver" v8 "github.com/yaoapp/gou/runtime/v8" chatctx "github.com/yaoapp/yao/neo/context" @@ -14,11 +15,11 @@ import ( // API the assistant API interface type API interface { - Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error + Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error) Download(ctx context.Context, fileID string) (*FileResponse, error) ReadBase64(ctx context.Context, fileID string) (string, error) - HookInit(ctx chatctx.Context, messages []message.Message) (*ResHookInit, error) + HookInit(c *gin.Context, ctx chatctx.Context, messages []message.Message) (*ResHookInit, error) } // RAG the RAG interface diff --git a/neo/message/message.go b/neo/message/message.go index 48c3a770..cad14749 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -22,6 +22,8 @@ type Message struct { 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:"-"` } @@ -134,12 +136,74 @@ func (m *Message) Error(message interface{}) *Message { return m } +// SetContent set the content +func (m *Message) SetContent(content string) *Message { + if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") { + var msg Message + if err := jsoniter.UnmarshalFromString(content, &msg); err != nil { + m.Text = err.Error() + "\n" + content + return m + } + *m = msg + } else { + m.Text = content + m.Type = "text" + } + return m +} + +// Content get the content +func (m *Message) Content() string { + content := map[string]interface{}{"text": m.Text} + if m.Attachments != nil { + content["attachments"] = m.Attachments + } + + if m.Type != "" { + content["type"] = m.Type + } + contentRaw, _ := jsoniter.MarshalToString(content) + return contentRaw +} + +// ToMap convert to map +func (m *Message) ToMap() map[string]interface{} { + return map[string]interface{}{ + "content": m.Content(), + "role": m.Role, + "name": m.Name, + } +} + // Map set from map func (m *Message) Map(msg map[string]interface{}) *Message { if msg == nil { return m } + // Content {"text": "xxxx", "attachments": ... } + if content, ok := msg["content"].(string); ok { + if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") { + var msg Message + if err := jsoniter.UnmarshalFromString(content, &msg); err != nil { + m.Text = err.Error() + "\n" + content + return m + } + *m = msg + } else { + m.Text = content + m.Type = "text" + } + } + + if role, ok := msg["role"].(string); ok { + m.Role = role + } + + if name, ok := msg["name"].(string); ok { + m.Name = name + } + if text, ok := msg["text"].(string); ok { m.Text = text } diff --git a/neo/neo.go b/neo/neo.go index 884f4a68..83bd91b2 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -18,7 +18,7 @@ var lock sync.Mutex = sync.Mutex{} // Answer reply the message func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error { - messages, err := neo.chatMessages(ctx, question) + messages, err := neo.withHistory(ctx, question) if err != nil { msg := message.New().Error(err).Done() msg.Write(c.Writer) @@ -36,7 +36,7 @@ func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) err } // Init the assistant - res, err = ast.HookInit(ctx, []message.Message{{Text: question}}) + res, err = ast.HookInit(c, ctx, messages) if err != nil { return err } @@ -131,7 +131,11 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st // Chat with AI in background go func() { - err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int { + msgList := make([]message.Message, len(messages)) + for i, msg := range messages { + msgList[i] = *message.New().Map(msg) + } + err := ast.Chat(c.Request.Context(), msgList, neo.Option, func(data []byte) int { select { case <-clientBreak: return 0 // break @@ -271,7 +275,7 @@ func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileRe } // chat chat with AI -func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []map[string]interface{}, c *gin.Context) error { +func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []message.Message, c *gin.Context) error { if ast == nil { msg := message.New().Error("assistant is not initialized").Done() msg.Write(c.Writer) @@ -350,33 +354,30 @@ func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []map[stri } } -// chatMessages get the chat messages -func (neo *DSL) chatMessages(ctx chatctx.Context, content ...string) ([]map[string]interface{}, error) { - +func (neo *DSL) withHistory(ctx chatctx.Context, question string) ([]message.Message, error) { history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID) if err != nil { return nil, err } - messages := []map[string]interface{}{} - messages = append(messages, history...) - if len(content) == 0 { - return messages, nil + // Add history messages + messages := []message.Message{} + for _, h := range history { + messages = append(messages, *message.New().Map(h)) } // Add user message - messages = append(messages, map[string]interface{}{"role": "user", "content": content[0], "name": ctx.Sid}) + messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": question, "name": ctx.Sid})) return messages, nil } // saveHistory save the history -func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []map[string]interface{}) { - +func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []message.Message) { if len(content) > 0 && sid != "" && len(messages) > 0 { err := neo.Store.SaveHistory( sid, []map[string]interface{}{ - {"role": "user", "content": messages[len(messages)-1]["content"], "name": sid}, + {"role": "user", "content": messages[len(messages)-1].Content(), "name": sid}, {"role": "assistant", "content": string(content), "name": sid}, }, chatID,