From 228e8cf482b31937a81d2a6339efd8e06f949e3d Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Sat, 18 Jan 2025 11:20:55 +0800
Subject: [PATCH] Enhance Neo API assistant message handling and response
 structure

- Added nil checks for response objects in the Execute method to prevent potential nil pointer dereferences.
- Updated the streamChat method to include assistant identification details (ID, name, avatar) in the message structure, improving context for messages.
- Introduced a new field in the Message struct to mark messages as new, enhancing message tracking capabilities.

These changes improve the robustness and maintainability of the Neo API, paving the way for better message handling and assistant functionalities.
---
 neo/assistant/api.go   | 28 +++++++++++++++++-----------
 neo/message/message.go | 26 ++++++++++++++++++++++----
 2 files changed, 39 insertions(+), 15 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 6447f82c..41f2b0d4 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -60,7 +60,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	}
 
 	// Switch to the new assistant if necessary
-	if res.AssistantID != ctx.AssistantID {
+	if res != nil && res.AssistantID != ctx.AssistantID {
 		newAst, err := Get(res.AssistantID)
 		if err != nil {
 			return err
@@ -69,17 +69,17 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
 	}
 
 	// Handle next action
-	if res.Next != nil {
+	if res != nil && res.Next != nil {
 		return res.Next.Execute(c, ctx)
 	}
 
 	// Update options if provided
-	if res.Options != nil {
+	if res != nil && res.Options != nil {
 		options = res.Options
 	}
 
 	// messages
-	if res.Input != nil {
+	if res != nil && res.Input != nil {
 		messages = res.Input
 	}
 
@@ -263,17 +263,20 @@ func (ast *Assistant) streamChat(
 
 				chatMessage.New().
 					Map(map[string]interface{}{
-						"text": value,
-						"done": msg.IsDone,
+						"assistant_id":     ast.ID,
+						"assistant_name":   ast.Name,
+						"assistant_avatar": ast.Avatar,
+						"text":             value,
+						"done":             msg.IsDone,
 					}).
 					Write(c.Writer)
 			}
 
 			// Complete the stream
 			if msg.IsDone {
-				// if value == "" {
-				// 	msg.Write(c.Writer)
-				// }
+				if value == "" {
+					msg.Write(c.Writer)
+				}
 
 				// Call HookDone
 				content.SetStatus(chatMessage.ContentStatusDone)
@@ -300,8 +303,11 @@ func (ast *Assistant) streamChat(
 				} else if value != "" {
 					chatMessage.New().
 						Map(map[string]interface{}{
-							"text": value,
-							"done": true,
+							"assistant_id":     ast.ID,
+							"assistant_name":   ast.Name,
+							"assistant_avatar": ast.Avatar,
+							"text":             value,
+							"done":             true,
 						}).
 						Write(c.Writer)
 				}
diff --git a/neo/message/message.go b/neo/message/message.go
index 6f559bc9..39551102 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -16,10 +16,11 @@ import (
 
 // Message the message
 type Message struct {
-	Text            string                 `json:"text,omitempty"`  // text content
-	Type            string                 `json:"type,omitempty"`  // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
-	Props           map[string]interface{} `json:"props,omitempty"` // props for the types
-	IsDone          bool                   `json:"done,omitempty"`
+	Text            string                 `json:"text,omitempty"`             // text content
+	Type            string                 `json:"type,omitempty"`             // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
+	Props           map[string]interface{} `json:"props,omitempty"`            // props for the types
+	IsDone          bool                   `json:"done,omitempty"`             // Mark as a done message from neo
+	IsNew           bool                   `json:"is_new,omitempty"`           // Mark as a new message from neo
 	Actions         []Action               `json:"actions,omitempty"`          // Conversation Actions for frontend
 	Attachments     []Attachment           `json:"attachments,omitempty"`      // File attachments
 	Role            string                 `json:"role,omitempty"`             // user, assistant, system ...
@@ -244,6 +245,23 @@ func (m *Message) Map(msg map[string]interface{}) *Message {
 	if done, ok := msg["done"].(bool); ok {
 		m.IsDone = done
 	}
+
+	if isNew, ok := msg["is_new"].(bool); ok {
+		m.IsNew = isNew
+	}
+
+	if assistantID, ok := msg["assistant_id"].(string); ok {
+		m.AssistantID = assistantID
+	}
+
+	if assistantName, ok := msg["assistant_name"].(string); ok {
+		m.AssistantName = assistantName
+	}
+
+	if assistantAvatar, ok := msg["assistant_avatar"].(string); ok {
+		m.AssistantAvatar = assistantAvatar
+	}
+
 	if actions, ok := msg["actions"].([]interface{}); ok {
 		for _, action := range actions {
 			if v, ok := action.(map[string]interface{}); ok {