opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Enhance streaming chat processing with token parsing and delta message handling

Browse source 
- 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
This commit is contained in:
Max 2025-02-03 15:46:25 +08:00
parent 941593b6ef
commit b19a6b953a
3 changed files with 166 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

95
neo/assistant/api.go
View file

@ -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: <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
if message.Type == "tool_calls" {
// Ignore the tool, think, error
if message.Type == "tool" || message.Type == "think" || message.Type == "error" {
continue
}

88
neo/message/contents.go
View file

@ -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
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}

27
neo/message/message.go
View file

@ -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