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

@@ -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
+			msg["name"] = name
-			newMessages = append(newMessages, message)
+			newMessages = append(newMessages, msg)
 		}
 	}
 
 	length := len(messages)
 	for index, message := range messages {
-		role, ok := message["role"].(string)
+		role := message.Role
-		if !ok {
+		if role == "" {
 			return nil, fmt.Errorf("role must be string")
 		}
 
-		content, ok := message["content"].(string)
+		content := message.Text
-		if !ok {
+		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
 		}
 

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) {
+func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, messages []message.Message) (*ResHookInit, error) {
-	v, err := ast.call("Init", context, messages)
+	// 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
+// createTimeoutContext creates a timeout context with 5 seconds timeout
-func (ast *Assistant) call(method string, context chatctx.Context, args ...any) (interface{}, error) {
+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
+	// Create done channel for handling cancellation
-	args = append([]interface{}{context.Map()}, args...)
+	done := make(chan struct{})
-	return ctx.Call(method, args...)
+	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
+	}
 }

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

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
 	}

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) withHistory(ctx chatctx.Context, question string) ([]message.Message, error) {
-func (neo *DSL) chatMessages(ctx chatctx.Context, content ...string) ([]map[string]interface{}, error) {
-
 	history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID)
 	if err != nil {
 		return nil, err
 	}
 
-	messages := []map[string]interface{}{}
+	// Add history messages
-	messages = append(messages, history...)
+	messages := []message.Message{}
-	if len(content) == 0 {
+	for _, h := range history {
-		return messages, nil
+		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,