From b19a6b953a2ebceae1de7edf31cc86f92aed95bc Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 3 Feb 2025 15:46:25 +0800 Subject: [PATCH] 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 | 95 ++++++++++++++++++++++++++++++----------- neo/message/contents.go | 88 +++++++++++++++++++++++++++++++++----- neo/message/message.go | 27 ++++++++---- 3 files changed, 166 insertions(+), 44 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 35f488ee..b04ae7e0 100644 --- a/neo/assistant/api.go +++ b/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 - msg.AppendTo(contents) - value := msg.String() - if value != "" { + delta := msg.String() + + // Chunk the delta + 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(). - Map(map[string]interface{}{ - "assistant_id": ast.ID, - "assistant_name": ast.Name, - "assistant_avatar": ast.Avatar, - "text": value, - "done": msg.IsDone, - }). - Write(c.Writer) + // Write the message to the client + output := chatMessage.New().Map(map[string]interface{}{ + "text": delta, + "type": msg.Type, + "done": msg.IsDone, + "delta": true, + }) + + 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 - if res.Output != nil { - chatMessage.New(). + msg := chatMessage.New().Done() + if res != nil && res.Output != nil { + 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: {\"arguments\":{\"assistant_id\":\"xxxx\"},\"function\":\"select_assistant\"}\n" + + " 1.1. e.g: {\"function\":\"function_name\",\"arguments\":{\"arg1\":\"xxxx\"}}\n" + " 1.2. strict the example format, do not add any additional information.\n" + - " 1.3. The JSON object should be wrapped by and .\n" + + " 1.3. The JSON object should be wrapped by and .\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 - if message.Type == "tool_calls" { + // Ignore the tool, think, error + if message.Type == "tool" || message.Type == "think" || message.Type == "error" { continue } diff --git a/neo/message/contents.go b/neo/message/contents.go index 04b30229..eb0f8788 100644 --- a/neo/message/contents.go +++ b/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": {"", ""}, + "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 -func (c *Contents) NewFunction(function string, arguments []byte) *Contents { +// NewTool create a new tool data and append to the 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 -func (c *Contents) SetFunctionID(id string) *Contents { +// SetToolID set the id of the current tool content +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 -func (c *Contents) AppendFunction(arguments []byte) *Contents { +// AppendTool append the tool to the current content +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} diff --git a/neo/message/message.go b/neo/message/message.go index bbcbe1c1..ae74d89a 100644 --- a/neo/message/message.go +++ b/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 != "" { - contents.NewFunction(name, []byte(m.Text)) + if name, ok := m.Props["tool"].(string); ok && name != "" { + 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