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

Merge pull request #926 from trheyi/main

Browse source 
feat: Enhance message handling with timestamps and token management
This commit is contained in:
Max 2025-04-16 14:30:38 +08:00 • committed by GitHub
commit 17686a77c6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 152 additions and 155 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
neo/assistant/api.go
View file

@ -370,6 +370,8 @@ func (ast *Assistant) streamChat(
toolsCount := 0
currentMessageID := ""
tokenID := ""
beginAt := int64(0)
var retry error = nil
var result interface{} = nil // To save the result
var content string = "" // To save the content
@ -414,6 +416,7 @@ func (ast *Assistant) streamChat(
// for api reasoning_content response
if msg.Type == "think" {
if isFirstThink {
msg.BeginAt = time.Now().UnixNano()
msg.Text = "<think>\n" + msg.Text // add the think begin tag
isFirstThink = false
isThinking = true
@ -427,14 +430,15 @@ func (ast *Assistant) streamChat(
end.ID = currentMessageID
end.Retry = ctx.Retry
end.Silent = ctx.Silent
end.EndAt = time.Now().UnixNano()
end.Callback(cb).Write(c.Writer)
end.AppendTo(contents)
contents.UpdateType("think", map[string]interface{}{"text": contents.Text()}, currentMessageID)
contents.UpdateType("think", map[string]interface{}{"text": contents.Text()}, chatMessage.Extra{ID: currentMessageID, End: time.Now().UnixNano()})
isThinking = false
// Clear the token and make a new line
contents.NewText([]byte{}, currentMessageID)
contents.NewText([]byte{}, chatMessage.Extra{ID: currentMessageID})
contents.ClearToken()
}
@ -460,11 +464,12 @@ func (ast *Assistant) streamChat(
}
toolsCount++
msg.BeginAt = time.Now().UnixNano()
}
if msg.IsEndTool {
msg.Text = msg.Text + "\n</tool>\n" // add the tool_calls close tag
msg.EndAt = time.Now().UnixNano()
}
}
@ -473,18 +478,21 @@ func (ast *Assistant) streamChat(
// Chunk the delta
if delta != "" {
msg.AppendTo(contents) // Append content and send message
msg.AppendTo(contents) // Append content
// Scan the tokens
contents.ScanTokens(currentMessageID, func(token string, id string, begin bool, text string, tails string) {
contents.ScanTokens(currentMessageID, tokenID, beginAt, func(token string, id string, tid string, beginAt int64, text string, tails string) {
currentMessageID = id
msg.ID = id
msg.Type = token
msg.Text = "" // clear the text
msg.Props = map[string]interface{}{"text": text} // Update props
msg.EndAt = time.Now().Unix()
// End of the token clear the text
if begin {
if beginAt != 0 {
tokenID = tid
msg.BeginAt = beginAt
return
}
@ -496,6 +504,10 @@ func (ast *Assistant) streamChat(
}
messages = append(messages, *newMsg)
}
// Reset the begin at and token id
beginAt = 0
tokenID = ""
})
// Handle stream

153
neo/message/contents.go
View file

@ -1,7 +1,10 @@
package message
import (
"fmt"
"math/rand"
"strings"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
@ -31,10 +34,19 @@ type Contents struct {
// Data the data of the content
type Data struct {
Type string `json:"type"` // text, function, error, think, tool
ID string `json:"id"` // the id of the content
Bytes []byte `json:"bytes"` // the content bytes
Props map[string]interface{} `json:"props"` // the props
Type string `json:"type"` // text, function, error, think, tool
ID string `json:"id"` // the id of the content
Bytes []byte `json:"bytes"` // the content bytes
Props map[string]interface{} `json:"props"` // the props
Begin int64 `json:"begin,omitempty"` // the begin time
End int64 `json:"end,omitempty"` // the end time
}
// Extra the extra of the content
type Extra struct {
ID string `json:"id,omitempty"` // the id of the content
Begin int64 `json:"begin,omitempty"` // the begin time
End int64 `json:"end,omitempty"` // the end time
}
// NewContents create a new contents
@ -46,7 +58,7 @@ func NewContents() *Contents {
}
// ScanTokens scan the tokens
func (c *Contents) ScanTokens(currentID string, cb func(token string, id string, begin bool, text string, tails string)) {
func (c *Contents) ScanTokens(messageID string, tokenID string, beginAt int64, cb func(token string, messageID string, tokenID string, beginAt int64, text string, tails string)) {
text := strings.TrimSpace(c.Text())
@ -60,15 +72,21 @@ func (c *Contents) ScanTokens(currentID string, cb func(token string, id string,
if index > 0 {
tails = text[index+len(token[1]):]
}
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)
extra := Extra{
ID: c.id,
End: time.Now().UnixNano(),
}
c.UpdateType(c.token, map[string]interface{}{"text": text}, extra)
c.NewText([]byte(tails), extra) // Create new text with the tails
cb(c.token, c.id, tokenID, beginAt, text, tails)
c.ClearToken() // clear the token
return
}
// call the callback for the begin of the token
cb(c.token, c.id, true, text, "")
// call the callback for the scanning of the token
cb(c.token, c.id, tokenID, beginAt, text, "")
return
}
@ -76,11 +94,19 @@ func (c *Contents) ScanTokens(currentID string, cb func(token string, id string,
for name, token := range tokens {
if index := strings.Index(text, token[0]); index >= 0 {
c.token = name
c.id = currentID
c.id = messageID
if c.id == "" {
c.id = uuid.New().String()
c.id = GenerateNumericID("M")
}
cb(name, c.id, true, text, "") // call the callback
// First time scanning the token, generate the token ID and begin time
if tokenID == "" {
tokenID = GenerateNumericID("T")
beginAt = time.Now().UnixNano()
c.UpdateType(name, map[string]interface{}{"text": text, "id": tokenID}, Extra{ID: c.id, Begin: beginAt, End: beginAt})
}
cb(name, c.id, tokenID, beginAt, text, "") // call the callback
}
}
}
@ -104,44 +130,79 @@ func (c *Contents) RemoveLastEmpty() {
}
// NewText create a new text data and append to the contents
func (c *Contents) NewText(bytes []byte, id ...string) *Contents {
func (c *Contents) NewText(bytes []byte, extra ...Extra) *Contents {
data := Data{Type: "text", Bytes: bytes}
if len(id) > 0 && id[0] != "" {
data.ID = id[0]
if len(extra) > 0 {
if extra[0].Begin != 0 {
data.Begin = extra[0].Begin
}
if extra[0].End != 0 {
data.End = extra[0].End
}
if extra[0].ID != "" {
data.ID = extra[0].ID
}
}
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{}, id ...string) *Contents {
func (c *Contents) NewType(typ string, props map[string]interface{}, extra ...Extra) *Contents {
data := Data{
Type: typ,
Props: props,
}
if len(id) > 0 && id[0] != "" {
data.ID = id[0]
if len(extra) > 0 {
if extra[0].Begin != 0 {
data.Begin = extra[0].Begin
}
if extra[0].End != 0 {
data.End = extra[0].End
}
if extra[0].ID != "" {
data.ID = extra[0].ID
}
}
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{}, id ...string) *Contents {
func (c *Contents) UpdateType(typ string, props map[string]interface{}, extra ...Extra) *Contents {
if c.Current == -1 {
c.NewType(typ, props, id...)
c.NewType(typ, props, extra...)
return c
}
if len(id) > 0 && id[0] != "" {
c.Data[c.Current].ID = id[0]
if len(extra) > 0 {
if extra[0].Begin != 0 {
c.Data[c.Current].Begin = extra[0].Begin
}
if extra[0].End != 0 {
c.Data[c.Current].End = extra[0].End
}
if extra[0].ID != "" {
c.Data[c.Current].ID = extra[0].ID
}
}
c.Data[c.Current].Type = typ
c.Data[c.Current].Props = props
if props != nil {
if c.Data[c.Current].Props == nil {
c.Data[c.Current].Props = map[string]interface{}{}
}
for k, v := range props {
c.Data[c.Current].Props[k] = v
}
}
return c
}
@ -156,14 +217,22 @@ func (c *Contents) NewError(err []byte) *Contents {
}
// AppendText append the text to the current content
func (c *Contents) AppendText(bytes []byte, id ...string) *Contents {
func (c *Contents) AppendText(bytes []byte, extra ...Extra) *Contents {
if c.Current == -1 {
c.NewText(bytes, id...)
c.NewText(bytes, extra...)
return c
}
if len(id) > 0 && id[0] != "" {
c.Data[c.Current].ID = id[0]
if len(extra) > 0 {
if extra[0].ID != "" {
c.Data[c.Current].ID = extra[0].ID
}
if extra[0].Begin != 0 {
c.Data[c.Current].Begin = extra[0].Begin
}
if extra[0].End != 0 {
c.Data[c.Current].End = extra[0].End
}
}
c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...)
return c
@ -237,5 +306,31 @@ func (data *Data) MarshalJSON() ([]byte, error) {
v["props"] = data.Props
}
// Add the begin and end time
if data.Begin != 0 {
v["begin"] = data.Begin
}
if data.End != 0 {
v["end"] = data.End
}
return jsoniter.Marshal(v)
}
// GenerateNumericID generates a 10-digit number using UUID as seed
func GenerateNumericID(prefix string) string {
// Generate UUID and use it as seed
id := uuid.New()
seed := int64(id[0])<<56 | int64(id[1])<<48 | int64(id[2])<<40 | int64(id[3])<<32 |
int64(id[4])<<24 | int64(id[5])<<16 | int64(id[6])<<8 | int64(id[7])
// Create a new random source using the seed
source := rand.NewSource(seed)
r := rand.New(source)
// Generate a number between 1000000000 and 9999999999 (10 digits)
num := r.Int63n(9000000000) + 1000000000
return fmt.Sprintf("%s%d", prefix, num)
}

121
neo/message/message.go
View file

@ -44,6 +44,8 @@ type Message struct {
IsBeginTool bool `json:"-"` // is new tool for the message for native tool_calls
IsEndTool bool `json:"-"` // is end tool for the message for native tool_calls
Result any `json:"result,omitempty"` // result for the message
BeginAt int64 `json:"begin_at,omitempty"` // begin at for the message // timestamp
EndAt int64 `json:"end_at,omitempty"` // end at for the message // timestamp
}
// Mention represents a mention
@ -338,121 +340,6 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
}
return msg
// switch {
// case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`) && !strings.Contains(text, `"tool_calls":null`):
// var toolCalls openai.ToolCalls
// if err := jsoniter.Unmarshal(data, &toolCalls); err != nil {
// color.Red("JSON parse error: %s", err.Error())
// color.White(string(data))
// msg.Text = "JSON parse error\n" + string(data)
// msg.Type = "error"
// msg.IsDone = true
// return msg
// }
// msg.Type = "tool_calls_native"
// if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 {
// id := toolCalls.Choices[0].Delta.ToolCalls[0].ID
// function := toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
// arguments := toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
// text := arguments
// if id != "" {
// text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
// }
// msg.Text = text
// }
// case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`):
// var message openai.MessageWithReasoningContent
// if err := jsoniter.Unmarshal(data, &message); err != nil {
// color.Red("JSON parse error: %s", err.Error())
// color.White(string(data))
// msg.Text = "JSON parse error\n" + string(data)
// msg.Type = "error"
// msg.IsDone = true
// return msg
// }
// msg.Type = "text"
// if len(message.Choices) > 0 {
// if reasoningContent, ok := message.Choices[0].Delta["reasoning_content"].(string); ok {
// msg.Text = reasoningContent
// msg.Type = "think"
// return msg
// }
// if content, ok := message.Choices[0].Delta["content"].(string); ok && content != "" {
// msg.Text = content
// msg.Type = "text"
// return msg
// }
// if isThinking {
// msg.Type = "think"
// msg.Text = ""
// return msg
// }
// msg.Text = ""
// return msg
// }
// case strings.Index(text, `{"code":`) == 0 || strings.Index(text, `"statusCode":`) > 0:
// var errorMessage openai.Error
// if err := jsoniter.UnmarshalFromString(text, &errorMessage); err != nil {
// color.Red("JSON parse error: %s", err.Error())
// color.White(string(data))
// msg.Text = "JSON parse error\n" + string(data)
// msg.Type = "error"
// msg.IsDone = true
// return msg
// }
// msg.Type = "error"
// msg.Text = errorMessage.Message
// msg.IsDone = true
// break
// case strings.Contains(text, `{"error":{`):
// var errorMessage openai.ErrorMessage
// if err := jsoniter.Unmarshal(data, &errorMessage); err != nil {
// color.Red("JSON parse error: %s", err.Error())
// color.White(string(data))
// msg.Text = "JSON parse error\n" + string(data)
// msg.Type = "error"
// msg.IsDone = true
// return msg
// }
// msg.Type = "error"
// msg.Text = errorMessage.Error.Message
// msg.IsDone = true
// break
// case strings.Contains(text, `"usage":`) && !strings.Contains(text, `"chat.completion.chunk`):
// msg.IsDone = true
// break
// case strings.Contains(text, `[DONE]`):
// msg.IsDone = true
// case strings.Contains(text, `"finish_reason":"stop"`):
// msg.IsDone = true
// case strings.Contains(text, `"finish_reason":"tool_calls"`):
// msg.IsDone = true
// // Not a data message
// case !strings.Contains(text, `data: `):
// msg.Pending = true
// msg.Text = text
// default:
// str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ")
// msg.Type = "error"
// msg.Text = str
// }
}
// String returns the string representation
@ -536,10 +423,10 @@ func (m *Message) AppendTo(contents *Contents) *Message {
case "text", "think", "tool", "tool_calls_native":
if m.Text != "" {
if m.IsNew {
contents.NewText([]byte(m.Text), m.ID)
contents.NewText([]byte(m.Text), Extra{ID: m.ID, Begin: m.BeginAt, End: m.EndAt})
return m
}
contents.AppendText([]byte(m.Text), m.ID)
contents.AppendText([]byte(m.Text), Extra{ID: m.ID, Begin: m.BeginAt, End: m.EndAt})
return m
}
return m

9
neo/neo.go
View file

@ -128,6 +128,8 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
isFirstThink := true
isThinking := false
currentMessageID := ""
tokenID := ""
beginAt := int64(0)
err := ast.Chat(c.Request.Context(), msgList, neo.Option, func(data []byte) int {
select {
case <-clientBreak:
@ -169,7 +171,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
isThinking = false
// Clear the token and make a new line
contents.NewText([]byte{}, currentMessageID)
contents.NewText([]byte{}, message.Extra{ID: currentMessageID})
contents.ClearToken()
}
@ -177,7 +179,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
msg.AppendTo(contents)
// Scan the tokens
contents.ScanTokens(currentMessageID, func(token string, id string, begin bool, text string, tails string) {
contents.ScanTokens(currentMessageID, tokenID, beginAt, func(token string, id string, tid string, beginAt int64, text string, tails string) {
currentMessageID = id
msg.ID = id
msg.Type = token
@ -185,7 +187,8 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
msg.Props = map[string]interface{}{"text": text} // Update props
// End of the token clear the text
if begin {
if beginAt != 0 {
msg.BeginAt = beginAt
return
}