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

@@ -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

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
+	Type  string                 `json:"type"`            // text, function, error, think, tool
-	ID    string                 `json:"id"`    // the id of the content
+	ID    string                 `json:"id"`              // the id of the content
-	Bytes []byte                 `json:"bytes"` // the content bytes
+	Bytes []byte                 `json:"bytes"`           // the content bytes
-	Props map[string]interface{} `json:"props"` // the props
+	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
+			extra := Extra{
-			cb(c.token, c.id, false, text, tails)
+				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
+		// call the callback for the scanning of the token
-		cb(c.token, c.id, true, text, "")
+		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] != "" {
+	if len(extra) > 0 {
-		c.Data[c.Current].ID = id[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] != "" {
+	if len(extra) > 0 {
-		c.Data[c.Current].ID = id[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)
+}

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

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
 					}