Merge pull request #823 from trheyi/main

Refactor Neo API assistant message handling and content structure

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,30 +231,13 @@ 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.AppendTo(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.Data)
 				if err == nil && res != nil {
-					if res.Output != "" {
-						value = res.Output
-					}
 
 					if res.Next != nil {
 						err = res.Next.Execute(c, ctx)
@@ -274,18 +267,16 @@ func (ast *Assistant) streamChat(
 
 			// Complete the stream
 			if msg.IsDone {
-				if value == "" {
+				// if value == "" {
-					msg.Write(c.Writer)
+				// 	msg.Write(c.Writer)
-				}
+				// }
 
-				// Call HookDone
+				res, hookErr := ast.HookDone(c, ctx, messages, contents.Data)
-				content.SetStatus(chatMessage.ContentStatusDone)
-				res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function")
 				if hookErr == nil && res != nil {
-					if res.Output != "" {
+					if res.Output != nil {
 						chatMessage.New().
 							Map(map[string]interface{}{
-								"text": res.Output,
+								"text": res.Input,
 								"done": true,
 							}).
 							Write(c.Writer)
@@ -322,8 +313,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 +324,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 []message.Data) (*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
@@ -75,8 +87,24 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 	switch v := v.(type) {
 	case map[string]interface{}:
 		if res, ok := v["output"].(string); ok {
-			response.Output = res
+			vv := []message.Data{}
+			err := jsoniter.UnmarshalFromString(res, &vv)
+			if err != nil {
+				return nil, err
+			}
+			response.Output = vv
 		}
+
+		if res, ok := v["output"].([]interface{}); ok {
+			vv := []message.Data{}
+			raw, _ := jsoniter.MarshalToString(res)
+			err := jsoniter.UnmarshalFromString(raw, &vv)
+			if err != nil {
+				return nil, err
+			}
+			response.Output = vv
+		}
+
 		if res, ok := v["next"].(map[string]interface{}); ok {
 			response.Next = &NextAction{}
 			if name, ok := res["action"].(string); ok {
@@ -93,19 +121,24 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
 		}
 
 	case string:
-		response.Output = v
+		vv := []message.Data{}
+		err := jsoniter.UnmarshalFromString(v, &vv)
+		if err != nil {
+			return nil, err
+		}
+		response.Output = vv
 	}
 
 	return response, nil
 }
 
 // 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 []message.Data) (*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
@@ -121,8 +154,24 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 	switch v := v.(type) {
 	case map[string]interface{}:
 		if res, ok := v["output"].(string); ok {
-			response.Output = res
+			vv := []message.Data{}
+			err := jsoniter.UnmarshalFromString(res, &vv)
+			if err != nil {
+				return nil, err
+			}
+			response.Output = vv
 		}
+
+		if res, ok := v["output"].([]interface{}); ok {
+			vv := []message.Data{}
+			raw, _ := jsoniter.MarshalToString(res)
+			err := jsoniter.UnmarshalFromString(raw, &vv)
+			if err != nil {
+				return nil, err
+			}
+			response.Output = vv
+		}
+
 		if res, ok := v["next"].(map[string]interface{}); ok {
 			response.Next = &NextAction{}
 			if name, ok := res["action"].(string); ok {
@@ -133,7 +182,12 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
 			}
 		}
 	case string:
-		response.Output = v
+		vv := []message.Data{}
+		err := jsoniter.UnmarshalFromString(v, &vv)
+		if err != nil {
+			return nil, err
+		}
+		response.Output = vv
 	}
 
 	return response, nil
@@ -185,7 +239,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/assistant/types.go

@@ -39,16 +39,16 @@ type ResHookInit struct {
 
 // ResHookStream the response of the stream hook
 type ResHookStream struct {
-	Silent bool        `json:"silent,omitempty"` // Whether to suppress the output
+	Silent bool           `json:"silent,omitempty"` // Whether to suppress the output
-	Next   *NextAction `json:"next,omitempty"`   // The next action
+	Next   *NextAction    `json:"next,omitempty"`   // The next action
-	Output string      `json:"output,omitempty"` // The output
+	Output []message.Data `json:"output,omitempty"` // The output
 }
 
 // ResHookDone the response of the done hook
 type ResHookDone struct {
 	Next   *NextAction       `json:"next,omitempty"`
 	Input  []message.Message `json:"input,omitempty"`
-	Output string            `json:"output,omitempty"`
+	Output []message.Data    `json:"output,omitempty"`
 }
 
 // ResHookFail the response of the fail hook

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,178 @@
+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)
+}
+
+// Map returns the map representation
+func (data *Data) Map() (map[string]interface{}, error) {
+	v := map[string]interface{}{"type": data.Type}
+
+	if data.ID != "" {
+		v["id"] = data.ID
+	}
+
+	if data.Bytes != nil {
+		v["text"] = 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 v, nil
+}
+
+// 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["text"] = 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
 }
 
+// AppendTo append the contents
+func (m *Message) AppendTo(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.AppendTo(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():