From 6203b5ccf230b98094d9b88c077a5423a87f3cb9 Mon Sep 17 00:00:00 2001
From: Max <shadow.iqka@qq.com>
Date: Wed, 16 Apr 2025 16:52:32 +0800
Subject: [PATCH] feat: Enhance message structure with ToolID for improved
 context

- Add ToolID field to Message struct to associate messages with specific tools.
- Update streamChat method to include ToolID in message handling for better tracking.
- Modify message properties to incorporate ToolID, enhancing clarity in message processing.
---
 neo/assistant/api.go   | 7 +++++--
 neo/message/message.go | 1 +
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/neo/assistant/api.go b/neo/assistant/api.go
index 2e338662..9f1dd423 100644
--- a/neo/assistant/api.go
+++ b/neo/assistant/api.go
@@ -433,6 +433,7 @@ func (ast *Assistant) streamChat(
 				end.Silent = ctx.Silent
 				end.End = time.Now().UnixNano()
 				end.Begin = beganAt
+				end.ToolID = tokenID
 
 				end.Callback(cb).Write(c.Writer)
 				end.AppendTo(contents)
@@ -491,10 +492,11 @@ func (ast *Assistant) streamChat(
 					currentMessageID = params.MessageID
 					msg.ID = params.MessageID
 					msg.Type = params.Token
-					msg.Text = ""                                           // clear the text
-					msg.Props = map[string]interface{}{"text": params.Text} // Update props
+					msg.Text = ""                                                                 // clear the text
+					msg.Props = map[string]interface{}{"text": params.Text, "id": params.TokenID} // Update props
 					msg.Begin = params.BeganAt
 					msg.End = params.EndAt
+					msg.ToolID = params.TokenID
 
 					// End of the token clear the text
 					if params.Begin {
@@ -570,6 +572,7 @@ func (ast *Assistant) streamChat(
 				if msg.Type == "think" || msg.Type == "tool" {
 					output.Begin = msg.Begin
 					output.End = msg.End
+					output.ToolID = msg.ToolID
 				}
 
 				output.Callback(cb).Write(c.Writer)
diff --git a/neo/message/message.go b/neo/message/message.go
index e42d7844..897aeb5b 100644
--- a/neo/message/message.go
+++ b/neo/message/message.go
@@ -21,6 +21,7 @@ var locker = sync.Mutex{}
 // Message the message
 type Message struct {
 	ID              string                 `json:"id,omitempty"`               // id for the message
+	ToolID          string                 `json:"tool_id,omitempty"`          // tool_id for the message
 	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