Enhance streaming chat processing with token parsing and delta message handling

- Implemented token scanning for 'think' and 'tool' types in stream chat
- Added support for delta message streaming with improved message chunking
- Updated contents handling to remove last empty data and track token states
- Refined message type processing and streaming logic in assistant API
- Modified message appending and type handling to support new streaming approach

neo/assistant/api.go

@@ -282,6 +282,7 @@ func (ast *Assistant) streamChat(
 	contents *chatMessage.Contents) error {
 
 	errorRaw := ""
+	isFirst := true
 	err := ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
 		select {
 		case <-clientBreak:
@@ -312,10 +313,44 @@ func (ast *Assistant) streamChat(
 				return 0 // break
 			}
 
-			// Append content and send message
+			delta := msg.String()
-			msg.AppendTo(contents)
+
-			value := msg.String()
+			// Chunk the delta
-			if value != "" {
+			if delta != "" {
+
+				msg.AppendTo(contents) // Append content and send message
+
+				// Scan the tokens
+				breakpoint := false
+				contents.ScanTokens(func(token string, begin bool, text string, tails string) {
+					msg.Type = token
+					msg.Text = ""                                    // clear the text
+					msg.Props = map[string]interface{}{"text": text} // Update props
+
+					// End of the token clear the text
+					if begin {
+						return
+					}
+
+					// Ignore the end of the token
+					if !begin && tails == "" {
+						breakpoint = true
+						return
+					}
+
+					// New message with the tails
+					newMsg, err := chatMessage.NewString(tails)
+					if err != nil {
+						return
+					}
+					messages = append(messages, *newMsg)
+				})
+
+				// If the breakpoint is true, continue the stream
+				if breakpoint {
+					return 1 // continue
+				}
+
 				// Handle stream
 				res, err := ast.HookStream(c, ctx, messages, msg, contents)
 				if err == nil && res != nil {
@@ -335,26 +370,33 @@ func (ast *Assistant) streamChat(
 					}
 				}
 
-				chatMessage.New().
+				// Write the message to the client
-					Map(map[string]interface{}{
+				output := chatMessage.New().Map(map[string]interface{}{
-						"assistant_id":     ast.ID,
+					"text":  delta,
-						"assistant_name":   ast.Name,
+					"type":  msg.Type,
-						"assistant_avatar": ast.Avatar,
-						"text":             value,
 					"done":  msg.IsDone,
-					}).
+					"delta": true,
-					Write(c.Writer)
+				})
+
+				if isFirst {
+					output.Assistant(ast.ID, ast.Name, ast.Avatar)
+					isFirst = false
+				}
+				output.Write(c.Writer)
 			}
 
 			// Complete the stream
 			if msg.IsDone {
+
 				// if value == "" {
 				// 	msg.Write(c.Writer)
 				// }
 
+				// Remove the last empty data
+				contents.RemoveLastEmpty()
+
 				res, hookErr := ast.HookDone(c, ctx, messages, contents)
 				if hookErr == nil && res != nil {
-
 					if res.Next != nil {
 						err := res.Next.Execute(c, ctx, contents)
 						if err != nil {
@@ -365,13 +407,15 @@ func (ast *Assistant) streamChat(
 						return 0 // break
 					}
 
-				} else if value != "" {
+				} else if delta != "" {
 					chatMessage.New().
 						Map(map[string]interface{}{
 							"assistant_id":     ast.ID,
 							"assistant_name":   ast.Name,
 							"assistant_avatar": ast.Avatar,
-							"text":             value,
+							"text":             delta,
+							"type":             "text",
+							"delta":            true,
 							"done":             true,
 						}).
 						Write(c.Writer)
@@ -384,16 +428,15 @@ func (ast *Assistant) streamChat(
 					return 0 // break
 				}
 
-				// Output
+				msg := chatMessage.New().Done()
-				if res.Output != nil {
+				if res != nil && res.Output != nil {
-					chatMessage.New().
+					msg = chatMessage.New().
 						Map(map[string]interface{}{
 							"text": res.Input,
 							"done": true,
-						}).
+						})
-						Write(c.Writer)
 				}
-
+				msg.Write(c.Writer)
 				done <- true
 				return 0 // break
 			}
@@ -517,9 +560,9 @@ func (ast *Assistant) withPrompts(messages []chatMessage.Message) []chatMessage.
 				"role": "system",
 				"content": "## Tool Calls Response Rules:\n" +
 					"1. The response should be a valid JSON object:\n" +
-					"  1.1. e.g: <tool_calls>{\"arguments\":{\"assistant_id\":\"xxxx\"},\"function\":\"select_assistant\"}</tool_calls>\n" +
+					"  1.1. e.g: <tool>{\"function\":\"function_name\",\"arguments\":{\"arg1\":\"xxxx\"}}</tool>\n" +
 					"  1.2. strict the example format, do not add any additional information.\n" +
-					"  1.3. The JSON object should be wrapped by <tool_calls> and </tool_calls>.\n" +
+					"  1.3. The JSON object should be wrapped by <tool> and </tool>.\n" +
 					"2. The structure of the JSON object is { \"arguments\": {...}, function:\"function_name\"}\n" +
 					"3. The function_name should be the name of the function defined in tool_calls.\n" +
 					"4. The arguments should be the arguments of the function defined in tool_calls.\n",
@@ -606,8 +649,8 @@ func (ast *Assistant) requestMessages(ctx context.Context, messages []chatMessag
 
 	for index, message := range messages {
 
-		// Ignore the tool call message
+		// Ignore the tool, think, error
-		if message.Type == "tool_calls" {
+		if message.Type == "tool" || message.Type == "think" || message.Type == "error" {
 			continue
 		}
 

neo/message/contents.go

@@ -2,6 +2,7 @@ package message
 
 import (
 	"fmt"
+	"strings"
 
 	jsoniter "github.com/json-iterator/go"
 )
@@ -15,15 +16,21 @@ const (
 	ContentStatusError
 )
 
+var tokens = map[string][2]string{
+	"think": {"<think>", "</think>"},
+	"tool":  {"<tool>", "</tool>"},
+}
+
 // Contents the contents
 type Contents struct {
 	Current int    `json:"current"` // the current content index
 	Data    []Data `json:"data"`    // the data
+	token   string // the current token
 }
 
 // Data the data of the content
 type Data struct {
-	Type      string                 `json:"type"`                // text, function, error, ...
+	Type      string                 `json:"type"`                // text, function, error, think, tool
 	ID        string                 `json:"id"`                  // the id of the content
 	Function  string                 `json:"function"`            // the function name
 	Bytes     []byte                 `json:"bytes"`               // the content bytes
@@ -39,6 +46,59 @@ func NewContents() *Contents {
 	}
 }
 
+// ScanTokens scan the tokens
+func (c *Contents) ScanTokens(cb func(token string, begin bool, text string, tails string)) {
+
+	text := strings.TrimSpace(c.Text())
+
+	// check the end of the token
+	if c.token != "" {
+		token := tokens[c.token]
+
+		// Check the end of the token
+		if index := strings.Index(text, token[1]); index >= 0 {
+			tails := ""
+			if index > 0 {
+				tails = text[index+len(token[1]):]
+			}
+
+			text = strings.TrimLeft(text[len(token[0]):index], "\n")
+			c.Data[c.Current].Bytes = []byte(text)
+			c.UpdateType(c.token, map[string]interface{}{"text": text})
+			c.NewText([]byte(tails)) // Create new text with the tails
+			cb(c.token, false, text, tails)
+			c.token = "" // clear the token
+			return
+		}
+
+		// call the callback for the begin of the token
+		cb(c.token, true, text, "")
+		return
+	}
+
+	// scan the begin of the token
+	for name, token := range tokens {
+		if index := strings.Index(text, token[0]); index >= 0 {
+			c.token = name
+			text = strings.TrimSpace(text[index+len(token[0]):])
+			cb(name, true, text, "") // call the callback
+		}
+	}
+}
+
+// RemoveLastEmpty remove the last empty data
+func (c *Contents) RemoveLastEmpty() {
+	if c.Current == -1 {
+		return
+	}
+
+	// Remove the last empty data
+	if len(c.Data[c.Current].Bytes) == 0 && c.Data[c.Current].Type == "text" {
+		c.Data = c.Data[:c.Current]
+		c.Current--
+	}
+}
+
 // NewText create a new text data and append to the contents
 func (c *Contents) NewText(bytes []byte) *Contents {
 	c.Data = append(c.Data, Data{
@@ -49,10 +109,10 @@ func (c *Contents) NewText(bytes []byte) *Contents {
 	return c
 }
 
-// NewFunction create a new function data and append to the contents
+// NewTool create a new tool data and append to the contents
-func (c *Contents) NewFunction(function string, arguments []byte) *Contents {
+func (c *Contents) NewTool(function string, arguments []byte) *Contents {
 	c.Data = append(c.Data, Data{
-		Type:      "function",
+		Type:      "tool",
 		Function:  function,
 		Arguments: arguments,
 	})
@@ -82,10 +142,10 @@ func (c *Contents) UpdateType(typ string, props map[string]interface{}) *Content
 	return c
 }
 
-// SetFunctionID set the id of the current function content
+// SetToolID set the id of the current tool content
-func (c *Contents) SetFunctionID(id string) *Contents {
+func (c *Contents) SetToolID(id string) *Contents {
 	if c.Current == -1 {
-		c.NewFunction("", []byte{})
+		c.NewTool("", []byte{})
 	}
 	c.Data[c.Current].ID = id
 	return c
@@ -111,10 +171,10 @@ func (c *Contents) AppendText(bytes []byte) *Contents {
 	return c
 }
 
-// AppendFunction append the function to the current content
+// AppendTool append the tool to the current content
-func (c *Contents) AppendFunction(arguments []byte) *Contents {
+func (c *Contents) AppendTool(arguments []byte) *Contents {
 	if c.Current == -1 {
-		c.NewFunction("", arguments)
+		c.NewTool("", arguments)
 		return c
 	}
 	c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...)
@@ -145,6 +205,14 @@ func (c *Contents) Text() string {
 	return string(c.Data[c.Current].Bytes)
 }
 
+// CurrentType returns the type of the current content
+func (c *Contents) CurrentType() string {
+	if c.Current == -1 {
+		return ""
+	}
+	return c.Data[c.Current].Type
+}
+
 // Map returns the map representation
 func (data *Data) Map() (map[string]interface{}, error) {
 	v := map[string]interface{}{"type": data.Type}

neo/message/message.go

@@ -21,6 +21,7 @@ type Message struct {
 	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:"new,omitempty"`              // Mark as a new message from neo
+	IsDelta         bool                   `json:"delta,omitempty"`            // Mark as a delta 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 ...
@@ -199,7 +200,7 @@ func NewOpenAI(data []byte) *Message {
 			return msg
 		}
 
-		msg.Type = "tool_calls"
+		msg.Type = "tool_calls_native"
 		if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 {
 			msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID
 			msg.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
@@ -283,7 +284,7 @@ func (m *Message) String() string {
 	}
 
 	switch typ {
-	case "text":
+	case "text", "think", "tool":
 		return m.Text
 
 	case "error":
@@ -308,6 +309,12 @@ func (m *Message) SetText(text string) *Message {
 	return m
 }
 
+// SetProps set the props
+func (m *Message) SetProps(props map[string]interface{}) *Message {
+	m.Props = props
+	return m
+}
+
 // Error set the error
 func (m *Message) Error(message interface{}) *Message {
 	m.Type = "error"
@@ -347,7 +354,7 @@ func (m *Message) AppendTo(contents *Contents) *Message {
 	}
 
 	switch m.Type {
-	case "text":
+	case "text", "think", "tool":
 		if m.Text != "" {
 			if m.IsNew {
 				contents.NewText([]byte(m.Text))
@@ -358,22 +365,22 @@ func (m *Message) AppendTo(contents *Contents) *Message {
 		}
 		return m
 
-	case "tool_calls":
+	case "tool_calls_native":
 
 		// Set function name
 		new := false
-		if name, ok := m.Props["function"].(string); ok && name != "" {
+		if name, ok := m.Props["tool"].(string); ok && name != "" {
-			contents.NewFunction(name, []byte(m.Text))
+			contents.NewTool(name, []byte(m.Text))
 			new = true
 		}
 
 		// Set id
 		if id, ok := m.Props["id"].(string); ok && id != "" {
-			contents.SetFunctionID(id)
+			contents.SetToolID(id)
 		}
 
 		if !new {
-			contents.AppendFunction([]byte(m.Text))
+			contents.AppendTool([]byte(m.Text))
 		}
 		return m
 
@@ -460,6 +467,10 @@ func (m *Message) Map(msg map[string]interface{}) *Message {
 		m.IsNew = isNew
 	}
 
+	if isDelta, ok := msg["delta"].(bool); ok {
+		m.IsDelta = isDelta
+	}
+
 	if assistantID, ok := msg["assistant_id"].(string); ok {
 		m.AssistantID = assistantID