diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 1c16d32d..d6285d3d 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -321,7 +321,8 @@ func (ast *Assistant) streamChat( msg.AppendTo(contents) // Append content and send message // Scan the tokens - contents.ScanTokens(func(token string, begin bool, text string, tails string) { + contents.ScanTokens(func(token string, id string, begin bool, text string, tails string) { + msg.ID = id msg.Type = token msg.Text = "" // clear the text msg.Props = map[string]interface{}{"text": text} // Update props @@ -332,7 +333,7 @@ func (ast *Assistant) streamChat( } // New message with the tails - newMsg, err := chatMessage.NewString(tails) + newMsg, err := chatMessage.NewString(tails, id) if err != nil { return } diff --git a/neo/message/contents.go b/neo/message/contents.go index fff3c847..10c30785 100644 --- a/neo/message/contents.go +++ b/neo/message/contents.go @@ -4,6 +4,7 @@ import ( "fmt" "strings" + "github.com/google/uuid" jsoniter "github.com/json-iterator/go" ) @@ -26,6 +27,7 @@ type Contents struct { Current int `json:"current"` // the current content index Data []Data `json:"data"` // the data token string // the current token + id string // the id of the contents } // Data the data of the content @@ -47,7 +49,7 @@ func NewContents() *Contents { } // ScanTokens scan the tokens -func (c *Contents) ScanTokens(cb func(token string, begin bool, text string, tails string)) { +func (c *Contents) ScanTokens(cb func(token string, id string, begin bool, text string, tails string)) { text := strings.TrimSpace(c.Text()) @@ -61,15 +63,15 @@ func (c *Contents) ScanTokens(cb func(token string, begin bool, text string, tai if index > 0 { tails = text[index+len(token[1]):] } - 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.UpdateType(c.token, map[string]interface{}{"text": text}, c.id) + c.NewText([]byte(tails), c.id) // Create new text with the tails + cb(c.token, c.id, false, text, tails) c.token = "" // clear the token return } // call the callback for the begin of the token - cb(c.token, true, text, "") + cb(c.token, c.id, true, text, "") return } @@ -77,7 +79,8 @@ func (c *Contents) ScanTokens(cb func(token string, begin bool, text string, tai for name, token := range tokens { if index := strings.Index(text, token[0]); index >= 0 { c.token = name - cb(name, true, text, "") // call the callback + c.id = uuid.New().String() + cb(name, c.id, true, text, "") // call the callback } } } @@ -96,43 +99,58 @@ func (c *Contents) RemoveLastEmpty() { } // NewText create a new text data and append to the contents -func (c *Contents) NewText(bytes []byte) *Contents { - c.Data = append(c.Data, Data{ - Type: "text", - Bytes: bytes, - }) +func (c *Contents) NewText(bytes []byte, id ...string) *Contents { + + data := Data{Type: "text", Bytes: bytes} + if len(id) > 0 && id[0] != "" { + data.ID = id[0] + } + c.Data = append(c.Data, data) c.Current++ return c } // 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{ +func (c *Contents) NewTool(function string, arguments []byte, id ...string) *Contents { + + data := Data{ Type: "tool", Function: function, Arguments: arguments, - }) + } + if len(id) > 0 && id[0] != "" { + data.ID = id[0] + } + c.Data = append(c.Data, data) c.Current++ return c } // NewType create a new type data and append to the contents -func (c *Contents) NewType(typ string, props map[string]interface{}) *Contents { - c.Data = append(c.Data, Data{ +func (c *Contents) NewType(typ string, props map[string]interface{}, id ...string) *Contents { + + data := Data{ Type: typ, Props: props, - }) + } + if len(id) > 0 && id[0] != "" { + data.ID = id[0] + } + c.Data = append(c.Data, data) c.Current++ return c } // UpdateType update the type of the current content -func (c *Contents) UpdateType(typ string, props map[string]interface{}) *Contents { +func (c *Contents) UpdateType(typ string, props map[string]interface{}, id ...string) *Contents { if c.Current == -1 { - c.NewType(typ, props) + c.NewType(typ, props, id...) return c } + if len(id) > 0 && id[0] != "" { + c.Data[c.Current].ID = id[0] + } c.Data[c.Current].Type = typ c.Data[c.Current].Props = props return c @@ -158,19 +176,23 @@ func (c *Contents) NewError(err []byte) *Contents { } // AppendText append the text to the current content -func (c *Contents) AppendText(bytes []byte) *Contents { +func (c *Contents) AppendText(bytes []byte, id ...string) *Contents { if c.Current == -1 { - c.NewText(bytes) + c.NewText(bytes, id...) return c } + + if len(id) > 0 && id[0] != "" { + c.Data[c.Current].ID = id[0] + } c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...) return c } // AppendTool append the tool to the current content -func (c *Contents) AppendTool(arguments []byte) *Contents { +func (c *Contents) AppendTool(arguments []byte, id ...string) *Contents { if c.Current == -1 { - c.NewTool("", arguments) + c.NewTool("", arguments, id...) return c } c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...) diff --git a/neo/message/message.go b/neo/message/message.go index ae74d89a..f6b4e169 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -16,6 +16,7 @@ import ( // Message the message type Message struct { + ID string `json:"id,omitempty"` // 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 @@ -135,7 +136,7 @@ func NewContent(content string) ([]Message, error) { } // NewString create a new message from string -func NewString(content string) (*Message, error) { +func NewString(content string, id ...string) (*Message, error) { if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") { var msg Message if err := jsoniter.UnmarshalFromString(content, &msg); err != nil { @@ -143,6 +144,9 @@ func NewString(content string) (*Message, error) { } return &msg, nil } + if len(id) > 0 { + return &Message{ID: id[0], Text: content}, nil + } return &Message{Text: content}, nil } @@ -357,10 +361,10 @@ func (m *Message) AppendTo(contents *Contents) *Message { case "text", "think", "tool": if m.Text != "" { if m.IsNew { - contents.NewText([]byte(m.Text)) + contents.NewText([]byte(m.Text), m.ID) return m } - contents.AppendText([]byte(m.Text)) + contents.AppendText([]byte(m.Text), m.ID) return m } return m