opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Refactor Neo API assistant message handling and content structure

Browse source 
- Replaced the Content struct with a new Contents struct to improve message handling and organization.
- Updated methods across the Assistant and Message structs to utilize the new Contents structure, enhancing the clarity and efficiency of message processing.
- Improved error handling by ensuring proper message appending and content management during chat streaming.
- Removed the deprecated Content struct, streamlining the codebase and reducing complexity.

These changes enhance the maintainability and robustness of the Neo API assistant, paving the way for improved message handling and assistant functionalities.
This commit is contained in:
Max 2025-01-19 12:23:23 +08:00
parent 5b30ef5fde
commit 5d5fa4c270
7 changed files with 230 additions and 141 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
neo/assistant/api.go
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,26 +231,12 @@ 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.Append(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.JSON())
if err == nil && res != nil {
if res.Output != "" {
value = res.Output
@ -278,9 +274,7 @@ func (ast *Assistant) streamChat(
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.JSON())
if hookErr == nil && res != nil {
if res.Output != "" {
chatMessage.New().
@ -322,8 +316,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 +327,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,

22
neo/assistant/hooks.go
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 string) (*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
@ -100,12 +112,12 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
}
// 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 string) (*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
@ -185,7 +197,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
}

4
neo/assistant/load.go
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 {

100
neo/message/content.go
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
}

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

@ -0,0 +1,150 @@
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)
}
// 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["bytes"] = 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)
}

41
neo/message/message.go
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
}
// Append append the contents
func (m *Message) Append(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, "\"")

6
neo/neo.go
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.Append(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():