feat: Refactor token management and enhance message scanning
- Update ScanTokens function to utilize ScanCallbackParams for improved clarity and structure in token processing. - Modify ClearToken method to accept an ID parameter, enhancing token stack management. - Refactor GenerateWithAI and streamChat methods to leverage new token handling features, improving message flow and context. - Introduce BeginAt and EndAt fields in message processing for better timing accuracy.
This commit is contained in:
Max
parent
77b991ed20
commit
712c8ade62
3 changed files with 141 additions and 43 deletions
|
|
@ -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)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
32
neo/neo.go
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue