yao/neo/message/message.go
Max 2f13dbbffb Enhance attachment model and add local uploader support
- Introduced a new `Groups` field in the `Attachment` struct to support categorization of attachments.
- Integrated `types.MetaInfo` into `ManagerOption` for improved metadata handling.
- Added a local uploader for attachments, enabling file uploads from a specified local path.
- Updated asset management functions to reflect changes in file metadata and structure.
2025-07-25 18:15:51 +08:00

695 lines
19 KiB
Go

package message
import (
"fmt"
"os"
"strings"
"sync"
"github.com/fatih/color"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/helper"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/maps"
"github.com/yaoapp/yao/attachment"
"github.com/yaoapp/yao/openai"
)
var locker = sync.Mutex{}
// Message the message
type Message struct {
ID string `json:"id,omitempty"` // id for the message
ToolID string `json:"tool_id,omitempty"` // tool_id for the message
Text string `json:"text,omitempty"` // text content
Type string `json:"type,omitempty"` // error, text, plan, table, form, page, file, video, audio, image, markdown, json ...
Props map[string]interface{} `json:"props,omitempty"` // props for the types
IsDone bool `json:"done,omitempty"` // Mark as a done message from neo
IsNew bool `json:"new,omitempty"` // Mark as a new message from neo
IsDelta bool `json:"delta,omitempty"` // Mark as a delta message from neo
Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend
Attachments []attachment.Attachment `json:"attachments,omitempty"` // File attachments
Role string `json:"role,omitempty"` // user, assistant, system ...
Name string `json:"name,omitempty"` // name for the message
AssistantID string `json:"assistant_id,omitempty"` // assistant_id (for assistant role = assistant )
AssistantName string `json:"assistant_name,omitempty"` // assistant_name (for assistant role = assistant )
AssistantAvatar string `json:"assistant_avatar,omitempty"` // assistant_avatar (for assistant role = assistant )
Mentions []Mention `json:"menions,omitempty"` // Mentions for the message ( for user role = user )
Data map[string]interface{} `json:"-"` // data for the message
Pending bool `json:"-"` // pending for the message
Hidden bool `json:"hidden,omitempty"` // hidden for the message (not show in the UI and history)
Retry bool `json:"retry,omitempty"` // retry for the message
Silent bool `json:"silent,omitempty"` // silent for the message (not show in the UI and history)
IsTool bool `json:"-"` // is tool for the message for native tool_calls
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
Begin int64 `json:"begin,omitempty"` // begin at for the message // timestamp
End int64 `json:"end,omitempty"` // end at for the message // timestamp
}
// Mention represents a mention
type Mention struct {
ID string `json:"assistant_id"` // assistant_id
Name string `json:"name"` // name
Avatar string `json:"avatar,omitempty"` // avatar
}
// Action the action
type Action struct {
Name string `json:"name,omitempty"`
Type string `json:"type"`
Payload interface{} `json:"payload,omitempty"`
}
// New create a new message
func New() *Message {
return &Message{Actions: []Action{}, Props: map[string]interface{}{}}
}
// NewHistory create a new message from history
func NewHistory(history map[string]interface{}) ([]Message, error) {
if history == nil {
return []Message{}, nil
}
var copy map[string]interface{} = map[string]interface{}{}
for key, value := range history {
if key != "content" {
copy[key] = value
}
}
globalMessage := New().Map(copy)
messages := []Message{}
if content, ok := history["content"].(string); ok {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message = *globalMessage
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
messages = append(messages, msg)
} else if strings.HasPrefix(content, "[") && strings.HasSuffix(content, "]") {
var msgs []Message
if err := jsoniter.UnmarshalFromString(content, &msgs); err != nil {
return nil, err
}
for _, msg := range msgs {
msg.AssistantID = globalMessage.AssistantID
msg.AssistantName = globalMessage.AssistantName
msg.AssistantAvatar = globalMessage.AssistantAvatar
msg.Role = globalMessage.Role
msg.Name = globalMessage.Name
msg.Mentions = globalMessage.Mentions
messages = append(messages, msg)
}
} else {
messages = append(messages, Message{Text: content})
}
}
return messages, nil
}
// NewContent create a new message from content
func NewContent(content string) ([]Message, error) {
messages := []Message{}
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
messages = append(messages, msg)
} else if strings.HasPrefix(content, "[") && strings.HasSuffix(content, "]") {
var msgs []Message
if err := jsoniter.UnmarshalFromString(content, &msgs); err != nil {
return nil, err
}
for _, msg := range msgs {
messages = append(messages, msg)
}
} else {
messages = append(messages, Message{Text: content})
}
return messages, nil
}
// NewString create a new message from string
func NewString(content string, id ...string) (*Message, error) {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
return nil, err
}
return &msg, nil
}
if len(id) > 0 {
return &Message{ID: id[0], Text: content}, nil
}
return &Message{Text: content}, nil
}
// NewStringError create a new message from string error
func NewStringError(content string) (*Message, error) {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg = New()
var errorMessage openai.ErrorMessage
if err := jsoniter.UnmarshalFromString(content, &errorMessage); err != nil {
msg.Text = err.Error() + "\n" + content
return msg, nil
}
msg.Type = "error"
msg.Text = errorMessage.Error.Message
return msg, nil
}
return &Message{Text: content}, nil
}
// NewMap create a new message from map
func NewMap(content map[string]interface{}) (*Message, error) {
return New().Map(content), nil
}
// NewAny create a new message from any content
func NewAny(content interface{}) (*Message, error) {
switch v := content.(type) {
case string:
return NewString(v)
case map[string]interface{}:
return NewMap(v)
}
return nil, fmt.Errorf("unknown content type: %T", content)
}
// NewOpenAI create a new message from OpenAI response
func NewOpenAI(data []byte, isThinking bool) *Message {
// For debug environment, print the response data
if os.Getenv("YAO_AGENT_PRINT_RESPONSE_DATA") == "true" {
log.Trace("[Response Data] %s", string(data))
}
if data == nil || len(data) == 0 {
return nil
}
msg := New()
text := string(data)
data = []byte(strings.TrimPrefix(text, "data: "))
switch {
case strings.Contains(text, `"object":"chat.completion.chunk"`): // Delta content
var chunk openai.ChatCompletionChunk
err := jsoniter.Unmarshal(data, &chunk)
if 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
}
// Empty content, then it is a pending message
if len(chunk.Choices) == 0 {
msg.Pending = true
return msg
}
// Tool calls
if len(chunk.Choices[0].Delta.ToolCalls) > 0 || chunk.Choices[0].FinishReason == "tool_calls" {
msg.Type = "tool_calls_native"
text := ""
if len(chunk.Choices[0].Delta.ToolCalls) > 0 {
id := chunk.Choices[0].Delta.ToolCalls[0].ID
function := chunk.Choices[0].Delta.ToolCalls[0].Function.Name
arguments := chunk.Choices[0].Delta.ToolCalls[0].Function.Arguments
text = arguments
if id != "" {
msg.IsBeginTool = true
msg.IsNew = true // mark as a new message
text = fmt.Sprintf(`{"id": "%s", "function": "%s", "arguments": %s`, id, function, arguments)
}
}
if chunk.Choices[0].FinishReason == "tool_calls" {
msg.IsEndTool = true
}
msg.Text = text
return msg
}
// Text content
if chunk.Choices[0].Delta.Content != "" {
msg.Type = "text"
msg.Text = chunk.Choices[0].Delta.Content
msg.IsDone = chunk.Choices[0].FinishReason == "stop" // is done when the content is finished
return msg
}
// Done messages
if chunk.Choices[0].FinishReason == "stop" || chunk.Choices[0].FinishReason == "tool_calls" {
msg.IsDone = true
return msg
}
// Reasoning content
if chunk.Choices[0].Delta.ReasoningContent != "" {
msg.Type = "think"
msg.Text = chunk.Choices[0].Delta.ReasoningContent
return msg
}
// Content is empty and is thinking, then it is a thinking message pending
if isThinking {
msg.Type = "think"
msg.Text = ""
return msg
}
msg.Text = ""
return msg
case strings.Contains(text, `"usage":`): // usage content
msg.IsDone = true
break
case strings.Contains(text, `[DONE]`):
msg.IsDone = true
return msg
case len(data) > 2 && data[0] == '{' && data[len(data)-1] == '}': // JSON content (error)
var error openai.Error
var errorMessage openai.ErrorMessage
if strings.Contains(string(data), `"error":`) {
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
}
error = errorMessage.Error
} else {
err := jsoniter.Unmarshal(data, &error)
if 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
}
}
message := error.Message
if message == "" {
message = "Unknown error occurred\n" + string(data)
}
msg.Type = "error"
msg.Text = message
msg.IsDone = true
return msg
case !strings.Contains(text, `data: `): // unknown message or uncompleted message
msg.Pending = true
msg.Text = text
return msg
default: // unknown message
str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ")
msg.Type = "error"
msg.Text = str
return msg
}
return msg
}
// String returns the string representation
func (m *Message) String() string {
typ := m.Type
if typ == "" {
typ = "text"
}
switch typ {
case "text", "think", "tool", "tool_calls_native":
return m.Text
case "error":
return m.Text
default:
raw, _ := jsoniter.MarshalToString(map[string]interface{}{"type": m.Type, "props": m.Props})
return raw
}
}
// SetText set the text
func (m *Message) SetText(text string) *Message {
m.Text = text
if m.Data != nil {
if replaced := helper.Bind(text, m.Data); replaced != nil {
if replacedText, ok := replaced.(string); ok {
m.Text = replacedText
}
}
}
return m
}
// SetProps set the props
func (m *Message) SetProps(props map[string]interface{}) *Message {
m.Props = props
return m
}
// Error set the error
func (m *Message) Error(message interface{}) *Message {
m.Type = "error"
switch v := message.(type) {
case error:
m.Text = v.Error()
case string:
m.Text = v
default:
m.Text = fmt.Sprintf("%v", message)
}
return m
}
// SetResult set the result
func (m *Message) SetResult(result any) *Message {
m.Result = result
m.Type = "result" // set the type to result
return m
}
// SetContent set the content
func (m *Message) SetContent(content string) *Message {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
return m
}
// AppendTo append the contents
func (m *Message) AppendTo(contents *Contents) *Message {
// Set type
if m.Type == "" {
m.Type = "text"
}
switch m.Type {
case "text", "think", "tool", "tool_calls_native":
if m.Text != "" {
if m.IsNew {
contents.NewText([]byte(m.Text), Extra{ID: m.ID, Begin: m.Begin, End: m.End})
return m
}
contents.AppendText([]byte(m.Text), Extra{ID: m.ID, Begin: m.Begin, End: m.End})
return m
}
return m
case "loading", "error", "action", "progress", "plan", "result": // Ignore progress, loading, plan and error messages
return m
default:
if m.IsNew {
contents.NewType(m.Type, m.Props)
return m
}
contents.UpdateType(m.Type, m.Props)
return m
}
}
// Content get the content
func (m *Message) Content() string {
content := map[string]interface{}{"text": m.Text}
if m.Attachments != nil {
content["attachments"] = m.Attachments
}
if m.Type != "" {
content["type"] = m.Type
}
contentRaw, _ := jsoniter.MarshalToString(content)
return contentRaw
}
// ToMap convert to map
func (m *Message) ToMap() map[string]interface{} {
return map[string]interface{}{
"content": m.Content(),
"role": m.Role,
"name": m.Name,
}
}
// Map set from map
func (m *Message) Map(msg map[string]interface{}) *Message {
if msg == nil {
return m
}
// Content {"text": "xxxx", "attachments": ... }
if content, ok := msg["content"].(string); ok {
if strings.HasPrefix(content, "{") && strings.HasSuffix(content, "}") {
var msg Message
if err := jsoniter.UnmarshalFromString(content, &msg); err != nil {
m.Text = err.Error() + "\n" + content
return m
}
*m = msg
} else {
m.Text = content
m.Type = "text"
}
}
// attachments
if attachments, has := msg["attachments"]; has {
raw, _ := jsoniter.Marshal(attachments)
m.Attachments = []attachment.Attachment{}
if err := jsoniter.Unmarshal(raw, &m.Attachments); err != nil {
color.Red("JSON parse error: %s", err.Error())
color.White(string(raw))
}
}
if role, ok := msg["role"].(string); ok {
m.Role = role
}
if name, ok := msg["name"].(string); ok {
m.Name = name
}
if text, ok := msg["text"].(string); ok {
m.Text = text
}
if typ, ok := msg["type"].(string); ok {
m.Type = typ
}
if done, ok := msg["done"].(bool); ok {
m.IsDone = done
}
if props, ok := msg["props"].(map[string]interface{}); ok {
m.Props = props
}
if isNew, ok := msg["new"].(bool); ok {
m.IsNew = isNew
}
if isDelta, ok := msg["delta"].(bool); ok {
m.IsDelta = isDelta
}
if assistantID, ok := msg["assistant_id"].(string); ok {
m.AssistantID = assistantID
// Set name
if m.Role == "assistant" {
m.Name = m.AssistantID
}
}
if assistantName, ok := msg["assistant_name"].(string); ok {
m.AssistantName = assistantName
}
if assistantAvatar, ok := msg["assistant_avatar"].(string); ok {
m.AssistantAvatar = assistantAvatar
}
if actions, ok := msg["actions"].([]interface{}); ok {
for _, action := range actions {
if v, ok := action.(map[string]interface{}); ok {
action := Action{}
if name, ok := v["name"].(string); ok {
action.Name = name
}
if t, ok := v["type"].(string); ok {
action.Type = t
}
if payload, ok := v["payload"].(map[string]interface{}); ok {
action.Payload = payload
}
m.Actions = append(m.Actions, action)
}
}
}
if data, ok := msg["data"].(map[string]interface{}); ok {
m.Data = data
}
return m
}
// Done set the done flag
func (m *Message) Done() *Message {
m.IsDone = true
return m
}
// Assistant set the assistant
func (m *Message) Assistant(id string, name string, avatar string) *Message {
m.AssistantID = id
m.AssistantName = name
m.AssistantAvatar = avatar
return m
}
// Action add an action
func (m *Message) Action(name string, t string, payload interface{}, next string) *Message {
if m.Data != nil {
payload = helper.Bind(payload, m.Data)
}
m.Actions = append(m.Actions, Action{
Name: name,
Type: t,
Payload: payload,
})
return m
}
// Bind replace with data
func (m *Message) Bind(data map[string]interface{}) *Message {
if data == nil {
return m
}
m.Data = maps.Of(data).Dot()
return m
}
// Callback callback the message
func (m *Message) Callback(fn interface{}) *Message {
if fn != nil {
switch v := fn.(type) {
case func(msg *Message):
if v == nil {
break
}
v(m)
break
case func():
if v == nil {
break
}
v()
break
default:
fmt.Println("no match callback")
break
}
}
return m
}
// WriteError writes an error message to response writer
func (m *Message) WriteError(w gin.ResponseWriter, message string) {
errMsg := strings.Trim(exception.New(message, 500).Message, "\"")
data := []byte(fmt.Sprintf(`{"text":"%s","type":"error"`, errMsg))
if m.IsDone {
data = []byte(fmt.Sprintf(`{"text":"%s","type":"error","done":true`, errMsg))
}
data = append([]byte("data: "), data...)
data = append(data, []byte("}\n\n")...)
if _, err := w.Write(data); err != nil {
color.Red("Write JSON Message Error: %s", message)
}
w.Flush()
}
// MarshalJSON implements json.Marshaler interface
func (m *Message) MarshalJSON() ([]byte, error) {
type Alias Message
return jsoniter.Marshal(&struct {
*Alias
}{
Alias: (*Alias)(m),
})
}
// UnmarshalJSON implements json.Unmarshaler interface
func (m *Message) UnmarshalJSON(data []byte) error {
type Alias Message
aux := &struct {
*Alias
}{
Alias: (*Alias)(m),
}
if err := jsoniter.Unmarshal(data, &aux); err != nil {
return err
}
return nil
}
// MarshalJSON implements json.Marshaler interface
func (a *Action) MarshalJSON() ([]byte, error) {
type Alias Action
return jsoniter.Marshal(&struct {
*Alias
}{
Alias: (*Alias)(a),
})
}
// UnmarshalJSON implements json.Unmarshaler interface
func (a *Action) UnmarshalJSON(data []byte) error {
type Alias Action
aux := &struct {
*Alias
}{
Alias: (*Alias)(a),
}
if err := jsoniter.Unmarshal(data, &aux); err != nil {
return err
}
return nil
}
// Write writes the message to response writer using the message queue
func (m *Message) Write(w gin.ResponseWriter) bool {
return WriteMessageAsync(m, w)
}