Merge pull request #823 from trheyi/main

Refactor Neo API assistant message handling and content structure
This commit is contained in:
Max 2025-01-19 15:25:49 +08:00 committed by GitHub
commit 8ee8429997
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 313 additions and 157 deletions

View file

@ -56,6 +56,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
// Run init hook
res, err := ast.HookInit(c, ctx, messages, options)
if err != nil {
chatMessage.New().
Assistant(ast.ID, ast.Name, ast.Avatar).
Error(err).
Done().
Write(c.Writer)
return err
}
@ -63,6 +68,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
if res != nil && res.AssistantID != ctx.AssistantID {
newAst, err := Get(res.AssistantID)
if err != nil {
chatMessage.New().
Assistant(ast.ID, ast.Name, ast.Avatar).
Error(err).
Done().
Write(c.Writer)
return err
}
*ast = *newAst
@ -163,16 +173,16 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []chatMessage.Message, options map[string]interface{}) error {
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
content := chatMessage.NewContent("text")
contents := chatMessage.NewContents()
// Chat with AI in background
go func() {
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, content)
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, contents)
if err != nil {
chatMessage.New().Error(err).Done().Write(c.Writer)
}
ast.saveChatHistory(ctx, messages, content)
ast.saveChatHistory(ctx, messages, contents)
done <- true
}()
@ -194,7 +204,7 @@ func (ast *Assistant) streamChat(
options map[string]interface{},
clientBreak chan bool,
done chan bool,
content *chatMessage.Content) error {
contents *chatMessage.Contents) error {
return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
select {
@ -210,7 +220,7 @@ func (ast *Assistant) streamChat(
// Handle error
if msg.Type == "error" {
value := msg.String()
res, hookErr := ast.HookFail(c, ctx, messages, content.String(), fmt.Errorf("%s", value))
res, hookErr := ast.HookFail(c, ctx, messages, contents.JSON(), fmt.Errorf("%s", value))
if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
value = res.Output
if res.Error != "" {
@ -221,30 +231,13 @@ func (ast *Assistant) streamChat(
return 0 // break
}
// Handle tool call
if msg.Type == "tool_calls" {
content.SetType("function") // Set type to function
// Set id
if id, ok := msg.Props["id"].(string); ok && id != "" {
content.SetID(id)
}
// Set name
if name, ok := msg.Props["name"].(string); ok && name != "" {
content.SetName(name)
}
}
// Append content and send message
msg.AppendTo(contents)
value := msg.String()
content.Append(value)
if value != "" {
// Handle stream
res, err := ast.HookStream(c, ctx, messages, content.String(), content.Type == "function")
res, err := ast.HookStream(c, ctx, messages, contents.Data)
if err == nil && res != nil {
if res.Output != "" {
value = res.Output
}
if res.Next != nil {
err = res.Next.Execute(c, ctx)
@ -274,18 +267,16 @@ func (ast *Assistant) streamChat(
// Complete the stream
if msg.IsDone {
if value == "" {
msg.Write(c.Writer)
}
// if value == "" {
// msg.Write(c.Writer)
// }
// Call HookDone
content.SetStatus(chatMessage.ContentStatusDone)
res, hookErr := ast.HookDone(c, ctx, messages, content.String(), content.Type == "function")
res, hookErr := ast.HookDone(c, ctx, messages, contents.Data)
if hookErr == nil && res != nil {
if res.Output != "" {
if res.Output != nil {
chatMessage.New().
Map(map[string]interface{}{
"text": res.Output,
"text": res.Input,
"done": true,
}).
Write(c.Writer)
@ -322,8 +313,8 @@ func (ast *Assistant) streamChat(
}
// saveChatHistory saves the chat history if storage is available
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, content *chatMessage.Content) {
if len(content.Bytes) > 0 && ctx.Sid != "" && len(messages) > 0 {
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessage.Message, contents *chatMessage.Contents) {
if len(contents.Data) > 0 && ctx.Sid != "" && len(messages) > 0 {
userMessage := messages[len(messages)-1]
data := []map[string]interface{}{
{
@ -333,7 +324,7 @@ func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []chatMessag
},
{
"role": "assistant",
"content": content.String(),
"content": contents.JSON(),
"name": ctx.Sid,
"assistant_id": ast.ID,
"assistant_name": ast.Name,

View file

@ -6,6 +6,7 @@ import (
"time"
"github.com/gin-gonic/gin"
jsoniter "github.com/json-iterator/go"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
@ -34,6 +35,17 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
response.ChatID = res
}
// input
if input, has := v["input"]; has {
raw, _ := jsoniter.MarshalToString(input)
vv := []message.Message{}
err := jsoniter.UnmarshalFromString(raw, &vv)
if err != nil {
return nil, err
}
response.Input = vv
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
@ -57,13 +69,13 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
}
// HookStream Handle streaming response from LLM
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookStream, error) {
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, output []message.Data) (*ResHookStream, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Stream", context, input, output, toolcall, c.Writer)
v, err := ast.call(ctx, "Stream", context, input, output, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
@ -75,8 +87,24 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(res, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
if res, ok := v["output"].([]interface{}); ok {
vv := []message.Data{}
raw, _ := jsoniter.MarshalToString(res)
err := jsoniter.UnmarshalFromString(raw, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
@ -93,19 +121,24 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
}
case string:
response.Output = v
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(v, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
return response, nil
}
// HookDone Handle completion of assistant response
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output string, toolcall bool) (*ResHookDone, error) {
func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []message.Message, output []message.Data) (*ResHookDone, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, err := ast.call(ctx, "Done", context, input, output, toolcall, c.Writer)
v, err := ast.call(ctx, "Done", context, input, output, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
@ -121,8 +154,24 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(res, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
if res, ok := v["output"].([]interface{}); ok {
vv := []message.Data{}
raw, _ := jsoniter.MarshalToString(res)
err := jsoniter.UnmarshalFromString(raw, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
if res, ok := v["next"].(map[string]interface{}); ok {
response.Next = &NextAction{}
if name, ok := res["action"].(string); ok {
@ -133,7 +182,12 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
}
}
case string:
response.Output = v
vv := []message.Data{}
err := jsoniter.UnmarshalFromString(v, &vv)
if err != nil {
return nil, err
}
response.Output = vv
}
return response, nil
@ -185,7 +239,7 @@ func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []
// createTimeoutContext creates a timeout context with 5 seconds timeout
func (ast *Assistant) createTimeoutContext(c *gin.Context) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithTimeout(c.Request.Context(), 5*time.Second)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
return ctx, cancel
}

View file

@ -28,6 +28,10 @@ var defaultConnector string = "" // default connector
// LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error {
// Clear the cache
loaded.Clear()
root := `/assistants`
app, err := fs.Get("app")
if err != nil {

View file

@ -39,16 +39,16 @@ type ResHookInit struct {
// ResHookStream the response of the stream hook
type ResHookStream struct {
Silent bool `json:"silent,omitempty"` // Whether to suppress the output
Next *NextAction `json:"next,omitempty"` // The next action
Output string `json:"output,omitempty"` // The output
Silent bool `json:"silent,omitempty"` // Whether to suppress the output
Next *NextAction `json:"next,omitempty"` // The next action
Output []message.Data `json:"output,omitempty"` // The output
}
// ResHookDone the response of the done hook
type ResHookDone struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output string `json:"output,omitempty"`
Output []message.Data `json:"output,omitempty"`
}
// ResHookFail the response of the fail hook

View file

@ -1,100 +0,0 @@
package message
import (
jsoniter "github.com/json-iterator/go"
)
const (
// ContentStatusPending the content status pending
ContentStatusPending = iota
// ContentStatusDone the content status done
ContentStatusDone
// ContentStatusError the content status error
ContentStatusError
)
// Content the content
type Content struct {
ID string `json:"id"`
Name string `json:"name"`
Bytes []byte `json:"bytes"`
Type string `json:"type"` // text, function, error
Status uint8 `json:"status"` // 0: pending, 1: done
}
// NewContent create a new content
func NewContent(typ string) *Content {
if typ == "" {
typ = "text"
}
return &Content{
Bytes: []byte{},
Type: typ,
Status: ContentStatusPending,
}
}
// String the content string
func (c *Content) String() string {
if c.Type == "function" {
var arguments interface{} = string(c.Bytes)
if c.Status == ContentStatusDone {
var vv interface{} = nil
err := jsoniter.Unmarshal(c.Bytes, &vv)
if err != nil {
return ""
}
arguments = vv
}
data := map[string]interface{}{
"id": c.ID,
"type": "function",
"text": c.Name,
"props": map[string]interface{}{
"id": c.ID,
"name": c.Name,
"arguments": arguments,
},
"function": map[string]interface{}{
"name": c.Name,
"arguments": arguments,
},
}
raw, err := jsoniter.MarshalToString(data)
if err != nil {
return ""
}
return raw
}
return string(c.Bytes)
}
// SetID set the content id
func (c *Content) SetID(id string) {
c.ID = id
}
// SetName set the content name
func (c *Content) SetName(name string) {
c.Name = name
}
// SetType set the content type
func (c *Content) SetType(typ string) {
c.Type = typ
}
// Append append the content
func (c *Content) Append(data string) {
c.Bytes = append(c.Bytes, []byte(data)...)
}
// SetStatus set the content status
func (c *Content) SetStatus(status uint8) {
c.Status = status
}

178
neo/message/contents.go Normal file
View file

@ -0,0 +1,178 @@
package message
import (
jsoniter "github.com/json-iterator/go"
)
const (
// ContentStatusPending the content status pending
ContentStatusPending = iota
// ContentStatusDone the content status done
ContentStatusDone
// ContentStatusError the content status error
ContentStatusError
)
// Contents the contents
type Contents struct {
Current int `json:"current"` // the current content index
Data []Data `json:"data"` // the data
}
// Data the data of the content
type Data struct {
Type string `json:"type"` // text, function, error, ...
ID string `json:"id"` // the id of the content
Function string `json:"function"` // the function name
Bytes []byte `json:"bytes"` // the content bytes
Arguments []byte `json:"arguments"` // the function arguments
}
// NewContents create a new contents
func NewContents() *Contents {
return &Contents{
Current: -1,
Data: []Data{},
}
}
// NewText create a new text data and append to the contents
func (c *Contents) NewText(bytes []byte) *Contents {
c.Data = append(c.Data, Data{
Type: "text",
Bytes: bytes,
})
c.Current++
return c
}
// NewFunction create a new function data and append to the contents
func (c *Contents) NewFunction(function string, arguments []byte) *Contents {
c.Data = append(c.Data, Data{
Type: "function",
Function: function,
Arguments: arguments,
})
c.Current++
return c
}
// SetFunctionID set the id of the current function content
func (c *Contents) SetFunctionID(id string) *Contents {
if c.Current == -1 {
c.NewFunction("", []byte{})
}
c.Data[c.Current].ID = id
return c
}
// NewError create a new error data and append to the contents
func (c *Contents) NewError(err []byte) *Contents {
c.Data = append(c.Data, Data{
Type: "error",
Bytes: err,
})
c.Current++
return c
}
// AppendText append the text to the current content
func (c *Contents) AppendText(bytes []byte) *Contents {
if c.Current == -1 {
c.NewText(bytes)
return c
}
c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, bytes...)
return c
}
// AppendFunction append the function to the current content
func (c *Contents) AppendFunction(arguments []byte) *Contents {
if c.Current == -1 {
c.NewFunction("", arguments)
return c
}
c.Data[c.Current].Arguments = append(c.Data[c.Current].Arguments, arguments...)
return c
}
// AppendError append the error to the current content
func (c *Contents) AppendError(err []byte) *Contents {
if c.Current == -1 {
c.NewError(err)
return c
}
c.Data[c.Current].Bytes = append(c.Data[c.Current].Bytes, err...)
return c
}
// JSON returns the json representation
func (c *Contents) JSON() string {
raw, _ := jsoniter.MarshalToString(c.Data)
return raw
}
// Text returns the text of the current content
func (c *Contents) Text() string {
if c.Current == -1 {
return ""
}
return string(c.Data[c.Current].Bytes)
}
// Map returns the map representation
func (data *Data) Map() (map[string]interface{}, error) {
v := map[string]interface{}{"type": data.Type}
if data.ID != "" {
v["id"] = data.ID
}
if data.Bytes != nil {
v["text"] = string(data.Bytes)
}
if data.Arguments != nil {
var vv interface{} = nil
err := jsoniter.Unmarshal(data.Arguments, &vv)
if err != nil {
return nil, err
}
v["arguments"] = vv
}
if data.Function != "" {
v["function"] = data.Function
}
return v, nil
}
// MarshalJSON returns the json representation
func (data *Data) MarshalJSON() ([]byte, error) {
v := map[string]interface{}{"type": data.Type}
if data.ID != "" {
v["id"] = data.ID
}
if data.Bytes != nil {
v["text"] = string(data.Bytes)
}
if data.Arguments != nil {
var vv interface{} = nil
err := jsoniter.Unmarshal(data.Arguments, &vv)
if err != nil {
return nil, err
}
v["arguments"] = vv
}
if data.Function != "" {
v["function"] = data.Function
}
return jsoniter.Marshal(v)
}

View file

@ -99,7 +99,7 @@ func NewOpenAI(data []byte) *Message {
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.Props["function"] = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Name
msg.Text = toolCalls.Choices[0].Delta.ToolCalls[0].Function.Arguments
}
@ -184,6 +184,32 @@ func (m *Message) SetContent(content string) *Message {
return m
}
// AppendTo append the contents
func (m *Message) AppendTo(contents *Contents) *Message {
switch m.Type {
case "text":
if m.Text != "" {
contents.AppendText([]byte(m.Text))
}
case "tool_calls":
// Set function name
if name, ok := m.Props["function"].(string); ok && name != "" {
contents.NewFunction(name, []byte(m.Text))
}
// Set id
if id, ok := m.Props["id"].(string); ok && id != "" {
contents.SetFunctionID(id)
}
contents.AppendFunction([]byte(m.Text))
}
return m
}
// Content get the content
func (m *Message) Content() string {
content := map[string]interface{}{"text": m.Text}
@ -291,6 +317,14 @@ func (m *Message) Done() *Message {
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 {
@ -339,11 +373,6 @@ func (m *Message) Write(w gin.ResponseWriter) bool {
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, "\"")

View file

@ -100,7 +100,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
clientBreak := make(chan bool, 1)
done := make(chan bool, 1)
fail := make(chan error, 1)
content := []byte{}
contents := message.NewContents()
// Chat with AI in background
go func() {
@ -126,7 +126,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
}
// Append content and send message
content = msg.Append(content)
msg.AppendTo(contents)
if !silent {
value := msg.String()
if value != "" {
@ -166,7 +166,7 @@ func (neo *DSL) GenerateWithAI(ctx chatctx.Context, input string, messageType st
// Wait for completion or client disconnect
select {
case <-done:
return string(content), nil
return contents.Text(), nil
case err := <-fail:
return "", err
case <-c.Writer.CloseNotify():