From 14f447c7822bc90f66f3be8608930e63b2fd01d7 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 16 Apr 2025 14:29:50 +0800 Subject: [PATCH] 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. --- neo/assistant/api.go | 24 +++++-- neo/message/contents.go | 153 ++++++++++++++++++++++++++++++++-------- neo/message/message.go | 121 ++----------------------------- neo/neo.go | 9 ++- 4 files changed, 152 insertions(+), 155 deletions(-) diff --git a/neo/assistant/api.go b/neo/assistant/api.go index 4dc384fd..91daa934 100644 --- a/neo/assistant/api.go +++ b/neo/assistant/api.go @@ -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 = "\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\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 diff --git a/neo/message/contents.go b/neo/message/contents.go index 6e0631af..e50bc222 100644 --- a/neo/message/contents.go +++ b/neo/message/contents.go @@ -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) +} diff --git a/neo/message/message.go b/neo/message/message.go index b57147a2..b93ed410 100644 --- a/neo/message/message.go +++ b/neo/message/message.go @@ -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 diff --git a/neo/neo.go b/neo/neo.go index 5bdfe438..5f241b8d 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -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 }