Refactor Neo API assistant message handling and content structure

- Replaced the Content struct with a new Contents struct to improve message handling and organization.
- Updated methods across the Assistant and Message structs to utilize the new Contents structure, enhancing the clarity and efficiency of message processing.
- Improved error handling by ensuring proper message appending and content management during chat streaming.
- Removed the deprecated Content struct, streamlining the codebase and reducing complexity.

These changes enhance the maintainability and robustness of the Neo API assistant, paving the way for improved message handling and assistant functionalities.

neo/assistant/api.go

@@ -56,6 +56,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	// Run init hook
 	res, err := ast.HookInit(c, ctx, messages, options)
 	if err != nil {
+		chatMessage.New().
+			Assistant(ast.ID, ast.Name, ast.Avatar).
+			Error(err).
+			Done().
+			Write(c.Writer)
 		return err
 	}
 
@@ -63,6 +68,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	if res != nil && res.AssistantID != ctx.AssistantID {
 		newAst, err := Get(res.AssistantID)
 		if err != nil {
+			chatMessage.New().
+				Assistant(ast.ID, ast.Name, ast.Avatar).
+				Error(err).
+				Done().
+				Write(c.Writer)
 			return err
 		}
 		*ast = *newAst
@@ -163,16 +173,16 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
 func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}) error {
 	clientBreak := make(chan bool, 1)
 	done := make(chan bool, 1)
-	content := chatMessage.NewContent("text")
+	contents := chatMessage.NewContents()
 
 	// Chat with AI in background
 	go func() {
-		err := ast.streamChat(c, ctx, messages, options, clientBreak, done, content)
+		err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents)
 		if err != nil {
 			chatMessage.New().Error(err).Done().Write(c.Writer)
 		}
 
-		ast.saveChatHistory(ctx, messages, content)
+		ast.saveChatHistory(ctx, messages, contents)
 		done <- true
 	}()
 
@@ -194,7 +204,7 @@ func (ast *Assistant) streamChat(
 	options map[string]interface{},
 	clientBreak chan bool,
 	done chan bool,
-	content *chatMessage.Content) error {
+	contents *chatMessage.Contents) error {
 
 	return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
 		select {
@@ -210,7 +220,7 @@ func (ast *Assistant) streamChat(
 			// Handle error
 			if msg.Type == "error" {
 				value := msg.String()
-				res, hookErr := ast.HookFail(c, ctx, messages, content.String(), fmt.Errorf("%s", value))
+				res, hookErr := ast.HookFail(c, ctx, messages, contents.JSON(), fmt.Errorf("%s", value))
 				if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
 					value = res.Output
 					if res.Error != "" {
@@ -221,26 +231,12 @@ func (ast *Assistant) streamChat(
 				return 0 // break
 			}
 
-			// Handle tool call
-			if msg.Type == "tool_calls" {
-				content.SetType("function") // Set type to function
-				// Set id
-				if id, ok := msg.Props["id"].(string); ok && id != "" {
-					content.SetID(id)
-				}
-
-				// Set name
-				if name, ok := msg.Props["name"].(string); ok && name != "" {
-					content.SetName(name)
-				}
-			}
-
 			// Append content and send message
+			msg.Append(contents)
 			value := msg.String()
-			content.Append(value)
 			if value != "" {
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, content.String(), content.Type == "function")
+				res, err := ast.HookStream(c, ctx, messages, contents.JSON())
 				if err == nil && res != nil {
 					if res.Output != "" {
 						value = res.Output
@@ -278,9 +274,7 @@ func (ast *Assistant) streamChat(
 					msg.Write(c.Writer)
 				}
 
-				// Call HookDone
+				res, hookErr := ast.HookDone(c, ctx, messages, contents.JSON())
-				content.SetStatus(chatMessage.ContentStatusDone)
-				res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function")
 				if hookErr == nil && res != nil {
 					if res.Output != "" {
 						chatMessage.New().
@@ -322,8 +316,8 @@ 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) {
+func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, contents *chatMessage.Contents) {
-	if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 {
+	if len(contents.Data) > 0 && ctx.Sid != "" && len(messages) > 0 {
 		userMessage := messages[len(messages)-1]
 		data := []map[string]interface{}{
 			{
@@ -333,7 +327,7 @@ func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessag
 			},
 			{
 				"role":             "assistant",
-				"content":          content.String(),
+				"content":          contents.JSON(),
 				"name":             ctx.Sid,
 				"assistant_id":     ast.ID,
 				"assistant_name":   ast.Name,

neo/assistant/hooks.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	jsoniter "github.com/json-iterator/go"
 	chatctx "github.com/yaoapp/yao/neo/context"
 	"github.com/yaoapp/yao/neo/message"
 )
@@ -34,6 +35,17 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
 			response.ChatID = res
 		}
 
+		// input
+		if input, has := v["input"]; has {
+			raw, _ := jsoniter.MarshalToString(input)
+			vv := []message.Message{}
+			err := jsoniter.UnmarshalFromString(raw, &vv)
+			if err != nil {
+				return nil, err
+			}
+			response.Input = vv
+		}
+
 		if res, ok := v["next"].(map[string]interface{}); ok {
 			response.Next = &NextAction{}
 			if name, ok := res["action"].(string); ok {
@@ -57,13 +69,13 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
 }
 
 // HookStream Handle streaming response from LLM
-func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookStream, error) {
+func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookStream, error) {
 
 	// Create timeout context
 	ctx, cancel := ast.createTimeoutContext(c)
 	defer cancel()
 
-	v, err := ast.call(ctx, "Stream", context, input, output, toolcall, c.Writer)
+	v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
 			return nil, nil
@@ -100,12 +112,12 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 }
 
 // HookDone Handle completion of assistant response
-func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookDone, error) {
+func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string) (*ResHookDone, error) {
 	// Create timeout context
 	ctx, cancel := ast.createTimeoutContext(c)
 	defer cancel()
 
-	v, err := ast.call(ctx, "Done", context, input, output, toolcall, c.Writer)
+	v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
 	if err != nil {
 		if err.Error() == HookErrorMethodNotFound {
 			return nil, nil
@@ -185,7 +197,7 @@ func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []
 
 // 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)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 	return ctx, cancel
 }
 

neo/assistant/load.go

@@ -28,6 +28,10 @@ var defaultConnector string = "" // default connector
 
 // LoadBuiltIn load the built-in assistants
 func LoadBuiltIn() error {
+
+	// Clear the cache
+	loaded.Clear()
+
 	root := `/assistants`
 	app, err := fs.Get("app")
 	if err != nil {

neo/message/content.go

@@ -1,100 +0,0 @@
-package message
-
-import (
-	jsoniter "github.com/json-iterator/go"
-)
-
-const (
-	// ContentStatusPending the content status pending
-	ContentStatusPending = iota
-	// ContentStatusDone the content status done
-	ContentStatusDone
-	// ContentStatusError the content status error
-	ContentStatusError
-)
-
-// Content the content
-type Content struct {
-	ID     string `json:"id"`
-	Name   string `json:"name"`
-	Bytes  []byte `json:"bytes"`
-	Type   string `json:"type"`   // text, function, error
-	Status uint8  `json:"status"` // 0: pending, 1: done
-}
-
-// NewContent create a new content
-func NewContent(typ string) *Content {
-	if typ == "" {
-		typ = "text"
-	}
-
-	return &Content{
-		Bytes:  []byte{},
-		Type:   typ,
-		Status: ContentStatusPending,
-	}
-}
-
-// String the content string
-func (c *Content) String() string {
-	if c.Type == "function" {
-
-		var arguments interface{} = string(c.Bytes)
-		if c.Status == ContentStatusDone {
-			var vv interface{} = nil
-			err := jsoniter.Unmarshal(c.Bytes, &vv)
-			if err != nil {
-				return ""
-			}
-			arguments = vv
-		}
-
-		data := map[string]interface{}{
-			"id":   c.ID,
-			"type": "function",
-			"text": c.Name,
-			"props": map[string]interface{}{
-				"id":        c.ID,
-				"name":      c.Name,
-				"arguments": arguments,
-			},
-			"function": map[string]interface{}{
-				"name":      c.Name,
-				"arguments": arguments,
-			},
-		}
-
-		raw, err := jsoniter.MarshalToString(data)
-		if err != nil {
-			return ""
-		}
-
-		return raw
-	}
-	return string(c.Bytes)
-}
-
-// SetID set the content id
-func (c *Content) SetID(id string) {
-	c.ID = id
-}
-
-// SetName set the content name
-func (c *Content) SetName(name string) {
-	c.Name = name
-}
-
-// SetType set the content type
-func (c *Content) SetType(typ string) {
-	c.Type = typ
-}
-
-// Append append the content
-func (c *Content) Append(data string) {
-	c.Bytes = append(c.Bytes, []byte(data)...)
-}
-
-// SetStatus set the content status
-func (c *Content) SetStatus(status uint8) {
-	c.Status = status
-}

neo/message/contents.go

@@ -0,0 +1,150 @@
+package message
+
+import (
+	jsoniter "github.com/json-iterator/go"
+)
+
+const (
+	// ContentStatusPending the content status pending
+	ContentStatusPending = iota
+	// ContentStatusDone the content status done
+	ContentStatusDone
+	// ContentStatusError the content status error
+	ContentStatusError
+)
+
+// Contents the contents
+type Contents struct {
+	Current int    `json:"current"` // the current content index
+	Data    []Data `json:"data"`    // the data
+}
+
+// Data the data of the content
+type Data struct {
+	Type      string `json:"type"`      // text, function, error, ...
+	ID        string `json:"id"`        // the id of the content
+	Function  string `json:"function"`  // the function name
+	Bytes     []byte `json:"bytes"`     // the content bytes
+	Arguments []byte `json:"arguments"` // the function arguments
+}
+
+// NewContents create a new contents
+func NewContents() *Contents {
+	return &Contents{
+		Current: -1,
+		Data:    []Data{},
+	}
+}
+
+// NewText create a new text data and append to the contents
+func (c *Contents) NewText(bytes []byte) *Contents {
+	c.Data = append(c.Data, Data{
+		Type:  "text",
+		Bytes: bytes,
+	})
+	c.Current++
+	return c
+}
+
+// NewFunction create a new function data and append to the contents
+func (c *Contents) NewFunction(function string, arguments []byte) *Contents {
+	c.Data = append(c.Data, Data{
+		Type:      "function",
+		Function:  function,
+		Arguments: arguments,
+	})
+	c.Current++
+	return c
+}
+
+// SetFunctionID set the id of the current function content
+func (c *Contents) SetFunctionID(id string) *Contents {
+	if c.Current == -1 {
+		c.NewFunction("", []byte{})
+	}
+	c.Data[c.Current].ID = id
+	return c
+}
+
+// NewError create a new error data and append to the contents
+func (c *Contents) NewError(err []byte) *Contents {
+	c.Data = append(c.Data, Data{
+		Type:  "error",
+		Bytes: err,
+	})
+	c.Current++
+	return c
+}
+
+// AppendText append the text to the current content
+func (c *Contents) AppendText(bytes []byte) *Contents {
+	if c.Current == -1 {
+		c.NewText(bytes)
+		return c
+	}
+	c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...)
+	return c
+}
+
+// AppendFunction append the function to the current content
+func (c *Contents) AppendFunction(arguments []byte) *Contents {
+	if c.Current == -1 {
+		c.NewFunction("", arguments)
+		return c
+	}
+	c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...)
+	return c
+}
+
+// AppendError append the error to the current content
+func (c *Contents) AppendError(err []byte) *Contents {
+	if c.Current == -1 {
+		c.NewError(err)
+		return c
+	}
+	c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, err...)
+	return c
+}
+
+// JSON returns the json representation
+func (c *Contents) JSON() string {
+	raw, _ := jsoniter.MarshalToString(c.Data)
+	return raw
+}
+
+// Text returns the text of the current content
+func (c *Contents) Text() string {
+	if c.Current == -1 {
+		return ""
+	}
+	return string(c.Data[c.Current].Bytes)
+}
+
+// MarshalJSON returns the json representation
+func (data *Data) MarshalJSON() ([]byte, error) {
+
+	v := map[string]interface{}{"type": data.Type}
+
+	if data.ID != "" {
+		v["id"] = data.ID
+	}
+
+	if data.Bytes != nil {
+		v["bytes"] = string(data.Bytes)
+	}
+
+	if data.Arguments != nil {
+		var vv interface{} = nil
+		err := jsoniter.Unmarshal(data.Arguments, &vv)
+		if err != nil {
+			return nil, err
+		}
+		v["arguments"] = vv
+	}
+
+	if data.Function != "" {
+		v["function"] = data.Function
+	}
+
+	return jsoniter.Marshal(v)
+}

neo/message/message.go

@@ -99,7 +99,7 @@ func NewOpenAI(data []byte) *Message {
 		msg.Type = "tool_calls"
 		if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 {
 			msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID
-			msg.Props["name"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
+			msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
 			msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
 		}
 
@@ -184,6 +184,32 @@ func (m *Message) SetContent(content string) *Message {
 	return m
 }
 
+// Append append the contents
+func (m *Message) Append(contents *Contents) *Message {
+
+	switch m.Type {
+	case "text":
+		if m.Text != "" {
+			contents.AppendText([]byte(m.Text))
+		}
+
+	case "tool_calls":
+
+		// Set function name
+		if name, ok := m.Props["function"].(string); ok && name != "" {
+			contents.NewFunction(name, []byte(m.Text))
+		}
+
+		// Set id
+		if id, ok := m.Props["id"].(string); ok && id != "" {
+			contents.SetFunctionID(id)
+		}
+
+		contents.AppendFunction([]byte(m.Text))
+	}
+	return m
+}
+
 // Content get the content
 func (m *Message) Content() string {
 	content := map[string]interface{}{"text": m.Text}
@@ -291,6 +317,14 @@ func (m *Message) Done() *Message {
 	return m
 }
 
+// Assistant set the assistant
+func (m *Message) Assistant(id string, name string, avatar string) *Message {
+	m.AssistantID = id
+	m.AssistantName = name
+	m.AssistantAvatar = avatar
+	return m
+}
+
 // Action add an action
 func (m *Message) Action(name string, t string, payload interface{}, next string) *Message {
 	if m.Data != nil {
@@ -339,11 +373,6 @@ func (m *Message) Write(w gin.ResponseWriter) bool {
 	return true
 }
 
-// Append appends content to the byte slice
-func (m *Message) Append(content []byte) []byte {
-	return append(content, []byte(m.Text)...)
-}
-
 // WriteError writes an error message to response writer
 func (m *Message) WriteError(w gin.ResponseWriter, message string) {
 	errMsg := strings.Trim(exception.New(message, 500).Message, "\"")

neo/neo.go

@@ -100,7 +100,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 	clientBreak := make(chan bool, 1)
 	done := make(chan bool, 1)
 	fail := make(chan error, 1)
-	content := []byte{}
+	contents := message.NewContents()
 
 	// Chat with AI in background
 	go func() {
@@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 				}
 
 				// Append content and send message
-				content = msg.Append(content)
+				msg.Append(contents)
 				if !silent {
 					value := msg.String()
 					if value != "" {
@@ -166,7 +166,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 	// Wait for completion or client disconnect
 	select {
 	case <-done:
-		return string(content), nil
+		return contents.Text(), nil
 	case err := <-fail:
 		return "", err
 	case <-c.Writer.CloseNotify():