Refactor Neo API assistant message handling and data structures

- Updated message appending methods to use AppendTo for improved clarity and consistency.
- Changed output types in HookStream and HookDone methods from string to []message.Data, enhancing data handling capabilities.
- Introduced a new Map method in the Data struct for better representation of message data.
- Streamlined JSON marshaling in the Data struct to ensure accurate data serialization.

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

neo/assistant/api.go

@@ -232,15 +232,12 @@ func (ast *Assistant) streamChat(
 			}
 
 			// Append content and send message
-			msg.Append(contents)
+			msg.AppendTo(contents)
 			value := msg.String()
 			if value != "" {
 				// Handle stream
-				res, err := ast.HookStream(c, ctx, messages, contents.JSON())
+				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)
@@ -270,16 +267,16 @@ func (ast *Assistant) streamChat(
 
 			// Complete the stream
 			if msg.IsDone {
-				if value == "" {
-					msg.Write(c.Writer)
-				}
+				// if value == "" {
+				// 	msg.Write(c.Writer)
+				// }
 
-				res, hookErr := ast.HookDone(c, ctx, messages, contents.JSON())
+				res, hookErr := ast.HookDone(c, ctx, messages, contents.Data)
 				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)

neo/assistant/hooks.go

@@ -69,7 +69,7 @@ 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) (*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)
@@ -87,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 {
@@ -105,14 +121,19 @@ 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) (*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()
@@ -133,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 {
@@ -145,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

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
-	Next   *NextAction `json:"next,omitempty"`   // The next action
-	Output string      `json:"output,omitempty"` // The output
+	Silent bool           `json:"silent,omitempty"` // Whether to suppress the output
+	Next   *NextAction    `json:"next,omitempty"`   // The next action
+	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/contents.go

@@ -120,6 +120,34 @@ func (c *Contents) Text() string {
 	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) {
 
@@ -130,7 +158,7 @@ func (data *Data) MarshalJSON() ([]byte, error) {
 	}
 
 	if data.Bytes != nil {
-		v["bytes"] = string(data.Bytes)
+		v["text"] = string(data.Bytes)
 	}
 
 	if data.Arguments != nil {

neo/message/message.go

@@ -184,8 +184,8 @@ func (m *Message) SetContent(content string) *Message {
 	return m
 }
 
-// Append append the contents
-func (m *Message) Append(contents *Contents) *Message {
+// AppendTo append the contents
+func (m *Message) AppendTo(contents *Contents) *Message {
 
 	switch m.Type {
 	case "text":

neo/neo.go

@@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
 				}
 
 				// Append content and send message
-				msg.Append(contents)
+				msg.AppendTo(contents)
 				if !silent {
 					value := msg.String()
 					if value != "" {