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

Merge pull request #928 from trheyi/main

Browse source 
feat: Refactor token management and enhance message scanning
This commit is contained in:
Max 2025-04-16 16:23:34 +08:00 • committed by GitHub
commit ac22599865
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 141 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
neo/assistant/api.go
View file

@ -15,6 +15,7 @@ import (
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/kun/log"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
chatMessage "github.com/yaoapp/yao/neo/message"
)
@ -431,6 +432,7 @@ func (ast *Assistant) streamChat(
end.Retry = ctx.Retry
end.Silent = ctx.Silent
end.EndAt = time.Now().UnixNano()
end.BeginAt = beginAt
end.Callback(cb).Write(c.Writer)
end.AppendTo(contents)
@ -439,7 +441,11 @@ func (ast *Assistant) streamChat(
// Clear the token and make a new line
contents.NewText([]byte{}, chatMessage.Extra{ID: currentMessageID})
contents.ClearToken()
// Clear the token
contents.ClearToken(tokenID)
beginAt = 0
tokenID = ""
}
// for native tool_calls response, keep the first tool_calls_native message
@ -481,33 +487,36 @@ func (ast *Assistant) streamChat(
msg.AppendTo(contents) // Append content
// Scan the tokens
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()
contents.ScanTokens(currentMessageID, tokenID, beginAt, func(params message.ScanCallbackParams) {
currentMessageID = params.MessageID
msg.ID = params.MessageID
msg.Type = params.Token
msg.Text = "" // clear the text
msg.Props = map[string]interface{}{"text": params.Text} // Update props
msg.BeginAt = params.BeginAt
msg.EndAt = params.EndAt
// End of the token clear the text
if beginAt != 0 {
tokenID = tid
msg.BeginAt = beginAt
if params.Begin {
tokenID = params.TokenID
beginAt = params.BeginAt
return
}
if params.End {
tokenID = ""
beginAt = 0
return
}
// New message with the tails
if tails != "" {
newMsg, err := chatMessage.NewString(tails, id)
if params.Tails != "" {
newMsg, err := chatMessage.NewString(params.Tails, params.MessageID)
if err != nil {
return
}
messages = append(messages, *newMsg)
}
// Reset the begin at and token id
beginAt = 0
tokenID = ""
})
// Handle stream
@ -557,6 +566,12 @@ func (ast *Assistant) streamChat(
output.Assistant(ast.ID, ast.Name, ast.Avatar)
isFirst = false
}
if msg.Type == "think" || msg.Type == "tool" {
output.BeginAt = msg.BeginAt
output.EndAt = msg.EndAt
}
output.Callback(cb).Write(c.Writer)
}

103
neo/message/contents.go
View file

@ -26,10 +26,12 @@ var tokens = map[string][2]string{
// 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
id string // the id of the contents
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
stack [][]string // the token stack
mapping map[string]string // the mapping of the token stack
}
// Data the data of the content
@ -49,6 +51,19 @@ type Extra struct {
End int64 `json:"end,omitempty"` // the end time
}
// ScanCallbackParams the params of the scan callback
type ScanCallbackParams struct {
Token string
MessageID string
TokenID string
BeginAt int64
EndAt int64
Begin bool
End bool
Text string
Tails string
}
// NewContents create a new contents
func NewContents() *Contents {
return &Contents{
@ -58,14 +73,15 @@ func NewContents() *Contents {
}
// ScanTokens scan the tokens
func (c *Contents) ScanTokens(messageID string, tokenID string, beginAt int64, cb func(token string, messageID string, tokenID string, beginAt int64, text string, tails string)) {
func (c *Contents) ScanTokens(messageID string, tokenID string, beginAt int64, cb func(params ScanCallbackParams)) {
text := strings.TrimSpace(c.Text())
// check the end of the token
if c.token != "" {
token := tokens[c.token]
token := c.GetToken(c.token)
tokenType := c.GetTokenType(c.token)
// Check the end of the token
if index := strings.Index(text, token[1]); index >= 0 {
tails := ""
@ -78,42 +94,101 @@ func (c *Contents) ScanTokens(messageID string, tokenID string, beginAt int64, c
End: time.Now().UnixNano(),
}
c.UpdateType(c.token, map[string]interface{}{"text": text}, extra)
c.UpdateType(tokenType, 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
cb(ScanCallbackParams{Token: tokenType, MessageID: c.id, TokenID: tokenID, BeginAt: beginAt, Begin: false, End: true, Text: text, Tails: tails, EndAt: extra.End})
c.ClearToken(c.token) // clear the token
return
}
// call the callback for the scanning of the token
cb(c.token, c.id, tokenID, beginAt, text, "")
cb(ScanCallbackParams{Token: tokenType, MessageID: c.id, TokenID: tokenID, BeginAt: beginAt, Begin: false, End: false, Text: text, Tails: "", EndAt: 0})
return
}
// scan the begin of the token
begin := false
for name, token := range tokens {
if index := strings.Index(text, token[0]); index >= 0 {
c.token = name
c.id = messageID
if c.id == "" {
c.id = GenerateNumericID("M")
}
tokenType := name
if tokenID != "" {
tokenType = c.GetTokenType(tokenID)
}
// First time scanning the token, generate the token ID and begin time
if tokenID == "" {
if tokenID == "" || tokenType != name {
tokenID = GenerateNumericID("T")
beginAt = time.Now().UnixNano()
begin = true
c.token = tokenID
c.AppendToken(tokenID, name)
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
cb(ScanCallbackParams{Token: name, MessageID: c.id, TokenID: tokenID, BeginAt: beginAt, Begin: begin, End: false, Text: text, Tails: "", EndAt: 0}) // call the callback
}
}
}
// ClearToken clear the token
func (c *Contents) ClearToken() {
func (c *Contents) ClearToken(id string) {
c.token = ""
next := 0
if c.stack == nil {
c.stack = [][]string{}
}
if c.mapping == nil {
c.mapping = map[string]string{}
}
for i, node := range c.stack {
if node[0] == id {
next = i + 1
delete(c.mapping, id)
break
}
}
// Remove the token from the stack, and set the next token
if next > 0 && next < len(c.stack) {
c.stack = c.stack[next:]
c.token = c.stack[len(c.stack)-1][0]
}
}
// AppendToken append the token to the stack
func (c *Contents) AppendToken(id string, name string) {
if c.stack == nil {
c.stack = [][]string{}
}
if c.mapping == nil {
c.mapping = map[string]string{}
}
c.stack = append(c.stack, []string{id, name})
c.mapping[id] = name
c.token = id
}
// GetTokenType get the token type from the stack
func (c *Contents) GetTokenType(id string) string {
return c.mapping[id]
}
// GetToken get the token from the stack
func (c *Contents) GetToken(name string) [2]string {
typ, ok := c.mapping[name]
if !ok {
return [2]string{}
}
return tokens[typ]
}
// RemoveLastEmpty remove the last empty data

32
neo/neo.go
View file

@ -172,32 +172,40 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
// Clear the token and make a new line
contents.NewText([]byte{}, message.Extra{ID: currentMessageID})
contents.ClearToken()
contents.ClearToken(currentMessageID)
}
// Append content and send message
msg.AppendTo(contents)
// Scan the tokens
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
contents.ScanTokens(currentMessageID, tokenID, beginAt, func(params message.ScanCallbackParams) {
currentMessageID = params.MessageID
msg.ID = params.MessageID
msg.Type = params.Token
msg.Text = "" // clear the text
msg.Props = map[string]interface{}{"text": params.Text} // Update props
// End of the token clear the text
if beginAt != 0 {
if params.Begin {
msg.BeginAt = beginAt
return
}
// New message with the tails
newMsg, err := message.NewString(tails, id)
if err != nil {
// End of the token clear the text
if params.End {
msg.EndAt = params.EndAt
return
}
msgList = append(msgList, *newMsg)
// New message with the tails
if params.Tails != "" {
newMsg, err := message.NewString(params.Tails, params.MessageID)
if err != nil {
return
}
msgList = append(msgList, *newMsg)
}
})
if !silent {