From c98a402e541b746e66a2c99613cd8a1367a5da7a Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 14 Dec 2024 17:45:40 +0800 Subject: [PATCH] Refactor Neo API chat handling to improve AI interaction and error management. Introduce background processing for chat with AI, enhancing responsiveness and client handling. Update assistant initialization logic to streamline the creation of OpenAI assistants and ensure proper error handling. Modify message struct to support structured error responses, improving clarity in communication. Additionally, clean up unused code and enhance overall maintainability of the assistant management system. --- neo/assistant/base/base.go | 15 +- neo/assistant/base/chat.go | 15 +- neo/assistant/openai/chat.go | 15 +- neo/assistant/openai/openai.go | 15 +- neo/assistant/types.go | 2 +- neo/message/json.go | 29 +-- neo/message/types.go | 2 +- neo/neo.go | 416 ++++++++++----------------------- neo/types.go | 13 -- openai/openai.go | 41 +++- 10 files changed, 222 insertions(+), 341 deletions(-) diff --git a/neo/assistant/base/base.go b/neo/assistant/base/base.go index df363f3a..f6d1376a 100644 --- a/neo/assistant/base/base.go +++ b/neo/assistant/base/base.go @@ -5,6 +5,7 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/yao/neo/assistant" + "github.com/yaoapp/yao/openai" ) // Base the base assistant @@ -12,14 +13,22 @@ type Base struct { ID string `json:"assistant_id"` Prompts []assistant.Prompt `json:"prompts,omitempty"` Connector connector.Connector `json:"-" yaml:"-"` + openai *openai.OpenAI } // New create a new base assistant func New(connector connector.Connector, prompts []assistant.Prompt, id ...string) (*Base, error) { - if len(id) > 0 { - return &Base{Connector: connector, ID: id[0], Prompts: prompts}, nil + + setting := connector.Setting() + api, err := openai.NewOpenAI(setting) + if err != nil { + return nil, err } - return &Base{Connector: connector, Prompts: prompts}, nil + + if len(id) > 0 { + return &Base{Connector: connector, ID: id[0], Prompts: prompts, openai: api}, nil + } + return &Base{Connector: connector, Prompts: prompts, openai: api}, nil } // List list all assistants diff --git a/neo/assistant/base/chat.go b/neo/assistant/base/chat.go index 6024ca15..7df5c58a 100644 --- a/neo/assistant/base/chat.go +++ b/neo/assistant/base/chat.go @@ -2,9 +2,20 @@ package base import ( "context" + "fmt" ) // Chat the chat -func (ast *Base) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, error) { - return nil, nil +func (ast *Base) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { + + if ast.openai == nil { + return fmt.Errorf("api is not initialized") + } + + _, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb) + if ext != nil { + return fmt.Errorf("openai chat completions with error: %s", ext.Message) + } + + return nil } diff --git a/neo/assistant/openai/chat.go b/neo/assistant/openai/chat.go index b1f77d19..765c6f3a 100644 --- a/neo/assistant/openai/chat.go +++ b/neo/assistant/openai/chat.go @@ -2,6 +2,7 @@ package openai import ( "context" + "fmt" ) // Chat the chat struct @@ -14,6 +15,16 @@ type Chat struct { func (ast *OpenAI) NewChat() {} // Chat the chat -func (ast *OpenAI) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, error) { - return nil, nil +func (ast *OpenAI) Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error { + + if ast.openai == nil { + return fmt.Errorf("openai is not initialized") + } + + _, ext := ast.openai.ChatCompletionsWith(ctx, messages, option, cb) + if ext != nil { + return fmt.Errorf("openai chat completions with error: %s", ext.Message) + } + + return nil } diff --git a/neo/assistant/openai/openai.go b/neo/assistant/openai/openai.go index fb87ee6a..ee10d12c 100644 --- a/neo/assistant/openai/openai.go +++ b/neo/assistant/openai/openai.go @@ -5,20 +5,29 @@ import ( "github.com/yaoapp/gou/connector" "github.com/yaoapp/yao/neo/assistant" + api "github.com/yaoapp/yao/openai" ) // OpenAI the openai assistant type OpenAI struct { ID string `json:"assistant_id"` // the assistant id Connector connector.Connector `json:"-" yaml:"-"` + openai *api.OpenAI } // New create a new openai assistant func New(connector connector.Connector, id ...string) (*OpenAI, error) { - if len(id) > 0 { - return &OpenAI{ID: id[0], Connector: connector}, nil + + setting := connector.Setting() + openai, err := api.NewOpenAI(setting) + if err != nil { + return nil, err } - return &OpenAI{Connector: connector}, nil + + if len(id) > 0 { + return &OpenAI{ID: id[0], Connector: connector, openai: openai}, nil + } + return &OpenAI{Connector: connector, openai: openai}, nil } // Current set the current assistant diff --git a/neo/assistant/types.go b/neo/assistant/types.go index ce52d3f5..33566c05 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -6,7 +6,7 @@ import ( // API the assistant API interface type API interface { - Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, error) + Chat(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) error List(ctx context.Context, param QueryParam) ([]Assistant, error) } diff --git a/neo/message/json.go b/neo/message/json.go index 767858af..1ebda8a1 100644 --- a/neo/message/json.go +++ b/neo/message/json.go @@ -1,6 +1,7 @@ package message import ( + "fmt" "strings" "github.com/fatih/color" @@ -54,7 +55,13 @@ func NewOpenAI(data []byte) *JSON { break default: - msg.Error = text + + str := string(data) + // Remove "data: " and " + str = strings.TrimPrefix(str, "data: ") + str = strings.Trim(str, "\"") + msg.Type = "error" + msg.Text = str } return &JSON{msg} @@ -83,10 +90,13 @@ func (json *JSON) Text(text string) *JSON { // Error set the error func (json *JSON) Error(message interface{}) *JSON { + json.Message.Type = "error" if err, ok := message.(error); ok { - json.Message.Error = err.Error() + json.Message.Text = err.Error() } else if msg, ok := message.(string); ok { - json.Message.Error = msg + json.Message.Text = msg + } else { + json.Message.Text = fmt.Sprintf("%v", message) } return json } @@ -101,12 +111,8 @@ func (json *JSON) Map(msg map[string]interface{}) *JSON { json.Message.Text = text } - if err, ok := msg["error"].(string); ok { - json.Message.Error = err - } - - if err, ok := msg["error"].(error); ok { - json.Message.Error = err.Error() + if typ, ok := msg["type"].(string); ok { + json.Message.Text = typ } if done, ok := msg["done"].(bool); ok { @@ -222,11 +228,6 @@ func (json *JSON) Write(w gin.ResponseWriter) bool { } }() - if json.Message != nil && json.Message.Error != "" { - json.writeError(w, json.Message.Error) - return false - } - data, err := jsoniter.Marshal(json.Message) if err != nil { log.Error("%s", err.Error()) diff --git a/neo/message/types.go b/neo/message/types.go index 468339df..ed54cf31 100644 --- a/neo/message/types.go +++ b/neo/message/types.go @@ -3,7 +3,7 @@ package message // Message the message type Message struct { Text string `json:"text,omitempty"` - Error string `json:"error,omitempty"` + Type string `json:"type,omitempty"` Done bool `json:"done,omitempty"` Confirm bool `json:"confirm,omitempty"` Command *Command `json:"command,omitempty"` diff --git a/neo/neo.go b/neo/neo.go index 54a863c0..e1021dc0 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -4,18 +4,16 @@ import ( "fmt" "strings" "sync" - "time" - "github.com/fatih/color" "github.com/gin-gonic/gin" "github.com/yaoapp/gou/connector" - "github.com/yaoapp/gou/process" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/assistant/base" "github.com/yaoapp/yao/neo/assistant/openai" "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/neo/message" + "github.com/yaoapp/yao/share" ) // Lock the assistant list @@ -50,26 +48,86 @@ func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { } // Chat with AI + return neo.chat(ast, ctx, messages, c) +} - fmt.Println(ast) +// chat chat with AI +func (neo *DSL) chat(ast assistant.API, ctx Context, messages []map[string]interface{}, c *gin.Context) error { - // Get the assistant_id, chat_id - time.Sleep(1 * time.Second) + if ast == nil { + msg := message.New().Error("assistant is not initialized").Done() + msg.Write(c.Writer) + return fmt.Errorf("assistant is not initialized") + } - // Send a text message to the client - msg := message.New().Map(map[string]interface{}{ - "text": "Hello, world!", - "done": true, - }) - msg.Write(c.Writer) + clientBreak := make(chan bool, 1) + done := make(chan bool, 1) + content := []byte{} - // Select Assistant + // Chat with AI in background + go func() { + err := ast.Chat(c.Request.Context(), messages, neo.Option, func(data []byte) int { + select { + case <-clientBreak: + return 0 // break - // Prepare Messages + default: + msg := message.NewOpenAI(data) + if msg == nil { + return 1 // continue + } - // Call AI + // Handle error + if msg.Type == "error" { + message.New().Error(msg.Message.Text).Done().Write(c.Writer) + return 0 // break + } - return nil + // Append content and send message + content = msg.Append(content) + if msg.Message != nil && msg.Message.Text != "" { + message.New(). + Map(map[string]interface{}{ + "text": msg.Message.Text, + "done": msg.Message.Done, + }). + Write(c.Writer) + } + + // Complete the stream + if msg.Message != nil && msg.Message.Done { + if msg.Message.Text == "" { + msg.Write(c.Writer) + } + done <- true + return 0 // break + } + + return 1 // continue + } + }) + + if err != nil { + log.Error("Chat error: %s", err.Error()) + message.New().Error(err).Done().Write(c.Writer) + } + + // Save chat history + if len(content) > 0 { + neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages) + } + + done <- true + }() + + // Wait for completion or client disconnect + select { + case <-done: + return nil + case <-c.Writer.CloseNotify(): + clientBreak <- true + return nil + } } // updateAssistantList update the assistant list @@ -114,21 +172,7 @@ func (neo *DSL) newAssistantByConfig(ast *assistant.Assistant) (assistant.API, e func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { // Moapi connector if id == "" || strings.HasPrefix(id, "moapi") { - model := "gpt-3.5-turbo" - if strings.HasPrefix(id, "moapi:") { - model = strings.TrimPrefix(id, "moapi:") - } - - conn, err := connector.New(`moapi`, `__yao.moapi`, []byte(`{"model": "`+model+`"}`)) - if err != nil { - return nil, fmt.Errorf("Create moapi assistant error: %s", err.Error()) - } - - api, err := openai.New(conn, neo.Use) - if err != nil { - return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) - } - return api, nil + return neo.newMoapiAssistant(id) } // Other connector @@ -153,6 +197,42 @@ func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { return api, nil } +// newMoapiAssistant creates a new moapi assistant +func (neo *DSL) newMoapiAssistant(id string) (assistant.API, error) { + model := "gpt-3.5-turbo" + if strings.HasPrefix(id, "moapi:") { + model = strings.TrimPrefix(id, "moapi:") + } + + // Get the moapi setting + url := share.MoapiHosts[0] + if share.App.Moapi.Mirrors != nil { + url = share.App.Moapi.Mirrors[0] + } + key := share.App.Moapi.Secret + organization := share.App.Moapi.Organization + + if !strings.HasPrefix(url, "http") { + url = "https://" + url + } + + // Check the moapi secret + if key == "" { + return nil, fmt.Errorf("The moapi secret is empty") + } + + conn, err := connector.New(`moapi`, `__yao.moapi`, []byte(`{"name":"Moapi", "options":{"model": "`+model+`", "key": "`+key+`", "organization": "`+organization+`", "host": "`+url+`"}}`)) + if err != nil { + return nil, fmt.Errorf("Create moapi assistant error: %s", err.Error()) + } + + api, err := openai.New(conn, neo.Use) + if err != nil { + return nil, fmt.Errorf("Create openai assistant error: %s", err.Error()) + } + return api, nil +} + // createDefaultAssistant create a default assistant func (neo *DSL) createDefaultAssistant() (assistant.API, error) { if neo.Use != "" { @@ -161,144 +241,6 @@ func (neo *DSL) createDefaultAssistant() (assistant.API, error) { return neo.newAssistant(neo.Connector) } -// // AnswerOld reply the message -// func (neo *DSL) AnswerOld(ctx Context, question string, c *gin.Context) error { -// // get the chat messages -// messages, err := neo.chatMessages(ctx, question) -// if err != nil { -// return err -// } - -// clientBreak := make(chan bool, 1) -// done := make(chan bool, 1) -// content := []byte{} - -// // Execute the command or chat with AI in the background -// go func() { - -// // chat with AI -// c.Header("Content-Type", "text/event-stream;charset=utf-8") -// c.Header("Cache-Control", "no-cache") -// c.Header("Connection", "keep-alive") - -// _, ex := neo.AI.ChatCompletionsWith(ctx, messages, neo.Option, func(data []byte) int { - -// select { -// case <-clientBreak: -// return 0 // break -// default: - -// msg := message.NewOpenAI(data) -// if msg == nil { -// return 1 // continue success -// } - -// if msg.Error != "" { -// neo.send(ctx, msg, messages, content, c) -// return 0 // break -// } - -// content = msg.Append(content) -// err := neo.send(ctx, msg, messages, content, c) -// if err != nil { -// c.Status(500) -// return 0 // break -// } - -// // Complete the stream -// if msg.IsDone() { -// done <- true -// return 0 // break -// } - -// return 1 // continue success -// } -// }) - -// // Throw the error -// if ex != nil { -// log.Error("Neo chat error: %s", ex.Message) -// c.Status(200) -// done <- true -// return -// } - -// // save the history -// neo.saveHistory(ctx.Sid, ctx.ChatID, content, messages) -// c.Status(200) - -// // Complete the stream -// done <- true - -// }() - -// select { -// case <-done: -// return nil -// case <-c.Writer.CloseNotify(): -// clientBreak <- true -// return nil -// } - -// } - -// Send send the message to the stream -func (neo *DSL) send(ctx Context, msg *message.JSON, messages []map[string]interface{}, content []byte, c *gin.Context) error { - - w := c.Writer - - if msg.Message != nil && msg.Message.Error != "" { - msg.Write(w) - return nil - } - - // Directly write the message - if neo.Write == "" { - ok := msg.Write(c.Writer) - if !ok { - return fmt.Errorf("Stream write error") - } - return nil - } - - // Execute the custom write hook get the response - args := []interface{}{ctx, messages, msg, string(content), w} - p, err := process.Of(neo.Write, args...) - if err != nil { - msg.Write(w) - color.Red("Neo custom write error: %s", err.Error()) - return fmt.Errorf("Stream write error: %s", err.Error()) - } - - err = p.WithSID(ctx.Sid).Execute() - if err != nil { - log.Error("Neo custom write error: %s", err.Error()) - msg.Write(w) - return nil - } - defer p.Release() - - res := p.Value() - if res == nil { - color.Red("Neo custom write return null") - return fmt.Errorf("Neo custom write return null") - } - - // Send the custom write response to the stream - if messages, ok := res.([]interface{}); ok { - for _, new := range messages { - if v, ok := new.(map[string]interface{}); ok { - newMsg := message.New().Map(v) - newMsg.Write(w) - } - } - return nil - } - - color.Red("Neo custom write should return an array of response") - return fmt.Errorf("Neo should return an array of response") -} - // prompts get the prompts func (neo *DSL) prompts() []map[string]interface{} { prompts := []map[string]interface{}{} @@ -313,55 +255,6 @@ func (neo *DSL) prompts() []map[string]interface{} { return prompts } -// prepare the messages -func (neo *DSL) prepare(ctx Context, messages []map[string]interface{}) []map[string]interface{} { - if neo.Prepare == "" { - return []map[string]interface{}{} - } - - prompts := []map[string]interface{}{} - p, err := process.Of(neo.Prepare, ctx, messages) - if err != nil { - color.Red("Neo prepare error: %s", err.Error()) - return prompts - } - - err = p.WithSID(ctx.Sid).Execute() - if err != nil { - color.Red("Neo prepare execute error: %s", err.Error()) - return prompts - } - defer p.Release() - - data := p.Value() - items, ok := data.([]interface{}) - if !ok { - color.Red("Neo prepare response is not array") - return prompts - } - - for i, item := range items { - v, ok := item.(map[string]interface{}) - if !ok { - color.Red("Neo prepare response [%d] is not map", i) - continue - } - - if _, ok := v["role"]; !ok { - color.Red(`Neo prepare response [%d]["role"] required`, i) - continue - } - - if _, ok := v["content"]; !ok { - color.Red(`Neo prepare response [%d]["content"] required`, i) - continue - } - prompts = append(prompts, v) - } - - return prompts -} - // chatMessages get the chat messages func (neo *DSL) chatMessages(ctx Context, content string) ([]map[string]interface{}, error) { @@ -395,51 +288,6 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages } } -// // NewAI create a new AI -// func (neo *DSL) newAI() error { - -// if neo.Connector == "" || strings.HasPrefix(neo.Connector, "moapi") { -// model := "gpt-3.5-turbo" -// if strings.HasPrefix(neo.Connector, "moapi:") { -// model = strings.TrimPrefix(neo.Connector, "moapi:") -// } - -// ai, err := openai.NewMoapi(model) -// if err != nil { -// return err -// } - -// neo.AI = ai -// return nil -// } - -// conn, err := connector.Select(neo.Connector) -// if err != nil { -// return err -// } - -// if conn.Is(connector.OPENAI) { -// ai, err := openai.New(neo.Connector) -// if err != nil { -// return err -// } -// neo.AI = ai -// return nil -// } - -// return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) -// } - -// // Select select the model -// func (neo *DSL) Select(model string) error { -// ai, err := openai.NewMoapi(model) -// if err != nil { -// return err -// } -// neo.AI = ai -// return nil -// } - // createConversation create a new conversation func (neo *DSL) createConversation() error { @@ -475,37 +323,11 @@ func (neo *DSL) createConversation() error { return fmt.Errorf("%s conversation connector %s not support", neo.ID, neo.ConversationSetting.Connector) } -// // NewAI create a new AI -// func (neo *DSL) newAI() error { - -// if neo.Connector == "" || strings.HasPrefix(neo.Connector, "moapi") { -// model := "gpt-3.5-turbo" -// if strings.HasPrefix(neo.Connector, "moapi:") { -// model = strings.TrimPrefix(neo.Connector, "moapi:") -// } - -// ai, err := openai.NewMoapi(model) -// if err != nil { -// return err -// } - -// neo.AI = ai -// return nil -// } - -// conn, err := connector.Select(neo.Connector) -// if err != nil { -// return err -// } - -// if conn.Is(connector.OPENAI) { -// ai, err := openai.New(neo.Connector) -// if err != nil { -// return err -// } -// neo.AI = ai -// return nil -// } - -// return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) -// } +// sendMessage sends a message to the client +func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error { + msg := message.New().Map(data.(map[string]interface{})) + if !msg.Write(w) { + return fmt.Errorf("failed to write message to stream") + } + return nil +} diff --git a/neo/types.go b/neo/types.go index c26f86d3..c0f9694a 100644 --- a/neo/types.go +++ b/neo/types.go @@ -4,7 +4,6 @@ import ( "context" "github.com/gin-gonic/gin" - "github.com/yaoapp/kun/exception" "github.com/yaoapp/yao/neo/assistant" "github.com/yaoapp/yao/neo/conversation" ) @@ -57,15 +56,3 @@ type CreateResponse struct { AssistantID string `json:"assistant_id,omitempty"` ChatID string `json:"chat_id,omitempty"` } - -// AI the AI interface -type AI interface { - ChatCompletions(messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) - ChatCompletionsWith(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) - GetContent(response interface{}) (string, *exception.Exception) - Embeddings(input interface{}, user string) (interface{}, *exception.Exception) - Tiktoken(input string) (int, error) - MaxToken() int -} - -// Prompt a prompt diff --git a/openai/openai.go b/openai/openai.go index fa0eec68..825a09e4 100644 --- a/openai/openai.go +++ b/openai/openai.go @@ -54,12 +54,43 @@ func New(id string) (*OpenAI, error) { } setting := c.Setting() + return NewOpenAI(setting) +} + +// NewOpenAI create a new OpenAI instance by setting +func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) { + + key := "" + if v, ok := setting["key"].(string); ok { + key = v + } + + model := "gpt-3.5-turbo" + if v, ok := setting["model"].(string); ok { + model = v + } + + host := "https://api.openai.com" + if v, ok := setting["host"].(string); ok { + host = v + } + + organization := "" + if v, ok := setting["organization"].(string); ok { + organization = v + } + + maxToken := 2048 + if v, ok := setting["max_token"].(int); ok { + maxToken = v + } + return &OpenAI{ - key: setting["key"].(string), - model: setting["model"].(string), - host: setting["host"].(string), - organization: "", - maxToken: 2048, + key: key, + model: model, + host: host, + organization: organization, + maxToken: maxToken, }, nil }