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,