feat: Enhance message handling with timestamps and token management
- Introduce BeginAt and EndAt fields in Message struct to track message timing. - Update NewText and NewType methods in Contents to accept Extra struct for better metadata handling. - Modify ScanTokens function to manage token IDs and beginning timestamps, improving token processing accuracy. - Refactor GenerateWithAI and streamChat methods to utilize new timestamp features, enhancing message context and flow.
This commit is contained in:
Max
parent
1fadfe0882
commit
14f447c782
4 changed files with 152 additions and 155 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue