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