- Added nil checks for response objects in the Execute method to prevent potential nil pointer dereferences. - Updated the streamChat method to include assistant identification details (ID, name, avatar) in the message structure, improving context for messages. - Introduced a new field in the Message struct to mark messages as new, enhancing message tracking capabilities. These changes improve the robustness and maintainability of the Neo API, paving the way for better message handling and assistant functionalities.
409 lines
11 KiB
Go
409 lines
11 KiB
Go
package message
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"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/openai"
|
|
)
|
|
|
|
// Message the message
|
|
type Message struct {
|
|
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:"is_new,omitempty"` // Mark as a new message from neo
|
|
Actions []Action `json:"actions,omitempty"` // Conversation Actions for frontend
|
|
Attachments []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:"-"`
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// Attachment represents a file attachment
|
|
type Attachment struct {
|
|
Name string `json:"name,omitempty"`
|
|
URL string `json:"url,omitempty"`
|
|
Description string `json:"description,omitempty"`
|
|
Type string `json:"type,omitempty"`
|
|
ContentType string `json:"content_type,omitempty"`
|
|
Bytes int64 `json:"bytes,omitempty"`
|
|
CreatedAt int64 `json:"created_at,omitempty"`
|
|
FileID string `json:"file_id,omitempty"`
|
|
ChatID string `json:"chat_id,omitempty"`
|
|
AssistantID string `json:"assistant_id,omitempty"`
|
|
}
|
|
|
|
// 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{}{}}
|
|
}
|
|
|
|
// NewString create a new message from string
|
|
func NewString(content 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
|
|
}
|
|
return &Message{Text: content}, nil
|
|
}
|
|
|
|
// NewOpenAI create a new message from OpenAI response
|
|
func NewOpenAI(data []byte) *Message {
|
|
if data == nil || len(data) == 0 {
|
|
return nil
|
|
}
|
|
|
|
msg := New()
|
|
text := string(data)
|
|
data = []byte(strings.TrimPrefix(text, "data: "))
|
|
|
|
switch {
|
|
|
|
case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`):
|
|
var toolCalls openai.ToolCalls
|
|
if err := jsoniter.Unmarshal(data, &toolCalls); err != nil {
|
|
msg.Text = err.Error() + "\n" + string(data)
|
|
return msg
|
|
}
|
|
|
|
msg.Type = "tool_calls"
|
|
if len(toolCalls.Choices) > 0 && len(toolCalls.Choices[0].Delta.ToolCalls) > 0 {
|
|
msg.Props["id"] = toolCalls.Choices[0].Delta.ToolCalls[0].ID
|
|
msg.Props["name"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
|
|
msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
|
|
}
|
|
|
|
case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"content":`):
|
|
var message openai.Message
|
|
if err := jsoniter.Unmarshal(data, &message); err != nil {
|
|
msg.Text = err.Error() + "\n" + string(data)
|
|
return msg
|
|
}
|
|
|
|
msg.Type = "text"
|
|
if len(message.Choices) > 0 {
|
|
msg.Text = message.Choices[0].Delta.Content
|
|
}
|
|
|
|
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
|
|
|
|
default:
|
|
str := strings.TrimPrefix(strings.Trim(string(data), "\""), "data: ")
|
|
msg.Type = "error"
|
|
msg.Text = str
|
|
}
|
|
|
|
return msg
|
|
}
|
|
|
|
// String returns the string representation
|
|
func (m *Message) String() string {
|
|
if m.Text != "" {
|
|
return m.Text
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// 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"
|
|
}
|
|
}
|
|
|
|
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 isNew, ok := msg["is_new"].(bool); ok {
|
|
m.IsNew = isNew
|
|
}
|
|
|
|
if assistantID, ok := msg["assistant_id"].(string); ok {
|
|
m.AssistantID = 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
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// Write writes the message to response writer
|
|
func (m *Message) Write(w gin.ResponseWriter) bool {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
message := "Write Response Exception: (if client close the connection, it's normal) \n %s\n\n"
|
|
color.Red(message, r)
|
|
}
|
|
}()
|
|
|
|
data, err := jsoniter.Marshal(m)
|
|
if err != nil {
|
|
log.Error("%s", err.Error())
|
|
return false
|
|
}
|
|
|
|
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", err.Error())
|
|
return false
|
|
}
|
|
w.Flush()
|
|
return true
|
|
}
|
|
|
|
// Append appends content to the byte slice
|
|
func (m *Message) Append(content []byte) []byte {
|
|
return append(content, []byte(m.Text)...)
|
|
}
|
|
|
|
// 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
|
|
}
|