Refactor Neo API assistant interaction and streamline message handling
- Updated the Answer method to utilize the new Execute method, simplifying the assistant interaction process. - Introduced a new Execute method in the Assistant struct to encapsulate the chat execution logic, enhancing clarity and maintainability. - Refactored the HookInit method to accept input messages and options, improving the initialization process for assistants. - Enhanced the ResHookInit struct to include next action handling and input messages, providing better control over assistant responses. - Removed obsolete chat handling methods, streamlining the codebase and improving overall structure. These changes improve the robustness and maintainability of the Neo API, paving the way for future enhancements in assistant functionalities and message management.
This commit is contained in:
Max
parent
ffb9ede647
commit
cb2cd0c317
4 changed files with 198 additions and 162 deletions
|
|
@ -6,7 +6,9 @@ import (
|
|||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/gou/fs"
|
||||
chatctx "github.com/yaoapp/yao/neo/context"
|
||||
"github.com/yaoapp/yao/neo/message"
|
||||
chatMessage "github.com/yaoapp/yao/neo/message"
|
||||
)
|
||||
|
|
@ -41,6 +43,184 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
|
|||
return assistant, nil
|
||||
}
|
||||
|
||||
// Execute implements the execute functionality
|
||||
func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error {
|
||||
messages, err := ast.withHistory(ctx, input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
options = ast.withOptions(options)
|
||||
|
||||
// Run init hook
|
||||
res, err := ast.HookInit(c, ctx, messages, options)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Switch to the new assistant if necessary
|
||||
if res.AssistantID != ctx.AssistantID {
|
||||
newAst, err := Get(res.AssistantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*ast = *newAst
|
||||
}
|
||||
|
||||
// Handle next action
|
||||
if res.Next != nil {
|
||||
switch res.Next.Action {
|
||||
case "exit":
|
||||
return nil
|
||||
// Add other actions here if needed
|
||||
}
|
||||
}
|
||||
|
||||
// Update options if provided
|
||||
if res.Options != nil {
|
||||
options = res.Options
|
||||
}
|
||||
|
||||
// Only proceed with chat stream if no specific next action was handled
|
||||
return ast.handleChatStream(c, ctx, messages, options)
|
||||
}
|
||||
|
||||
// handleChatStream manages the streaming chat interaction with the AI
|
||||
func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, messages []message.Message, options map[string]interface{}) error {
|
||||
clientBreak := make(chan bool, 1)
|
||||
done := make(chan bool, 1)
|
||||
content := []byte{}
|
||||
|
||||
// Chat with AI in background
|
||||
go func() {
|
||||
err := ast.streamChat(c, messages, options, clientBreak, done, &content)
|
||||
if err != nil {
|
||||
chatMessage.New().Error(err).Done().Write(c.Writer)
|
||||
}
|
||||
|
||||
ast.saveChatHistory(ctx, messages, content)
|
||||
done <- true
|
||||
}()
|
||||
|
||||
// Wait for completion or client disconnect
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-c.Writer.CloseNotify():
|
||||
clientBreak <- true
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// streamChat handles the streaming chat interaction
|
||||
func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, options map[string]interface{},
|
||||
clientBreak chan bool, done chan bool, content *[]byte) error {
|
||||
|
||||
return ast.Chat(c.Request.Context(), messages, options, func(data []byte) int {
|
||||
select {
|
||||
case <-clientBreak:
|
||||
return 0 // break
|
||||
|
||||
default:
|
||||
msg := chatMessage.NewOpenAI(data)
|
||||
if msg == nil {
|
||||
return 1 // continue
|
||||
}
|
||||
|
||||
// Handle error
|
||||
if msg.Type == "error" {
|
||||
value := msg.String()
|
||||
chatMessage.New().Error(value).Done().Write(c.Writer)
|
||||
return 0 // break
|
||||
}
|
||||
|
||||
// Append content and send message
|
||||
*content = msg.Append(*content)
|
||||
value := msg.String()
|
||||
if value != "" {
|
||||
chatMessage.New().
|
||||
Map(map[string]interface{}{
|
||||
"text": value,
|
||||
"done": msg.IsDone,
|
||||
}).
|
||||
Write(c.Writer)
|
||||
}
|
||||
|
||||
// Complete the stream
|
||||
if msg.IsDone {
|
||||
if value == "" {
|
||||
msg.Write(c.Writer)
|
||||
}
|
||||
done <- true
|
||||
return 0 // break
|
||||
}
|
||||
|
||||
return 1 // continue
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// saveChatHistory saves the chat history if storage is available
|
||||
func (ast *Assistant) saveChatHistory(ctx chatctx.Context, messages []message.Message, content []byte) {
|
||||
if len(content) > 0 && ctx.Sid != "" && len(messages) > 0 {
|
||||
storage.SaveHistory(
|
||||
ctx.Sid,
|
||||
[]map[string]interface{}{
|
||||
{"role": "user", "content": messages[len(messages)-1].Content(), "name": ctx.Sid},
|
||||
{"role": "assistant", "content": string(content), "name": ctx.Sid},
|
||||
},
|
||||
ctx.ChatID,
|
||||
nil,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (ast *Assistant) withOptions(options map[string]interface{}) map[string]interface{} {
|
||||
if options == nil {
|
||||
options = map[string]interface{}{}
|
||||
}
|
||||
|
||||
if ast.Options != nil {
|
||||
for key, value := range ast.Options {
|
||||
options[key] = value
|
||||
}
|
||||
}
|
||||
return options
|
||||
}
|
||||
|
||||
func (ast *Assistant) withPrompts(messages []message.Message) []message.Message {
|
||||
if ast.Prompts != nil {
|
||||
for _, prompt := range ast.Prompts {
|
||||
name := ast.Name
|
||||
if prompt.Name != "" {
|
||||
name = prompt.Name
|
||||
}
|
||||
messages = append(messages, *message.New().Map(map[string]interface{}{"role": prompt.Role, "content": prompt.Content, "name": name}))
|
||||
}
|
||||
}
|
||||
return messages
|
||||
}
|
||||
|
||||
func (ast *Assistant) withHistory(ctx chatctx.Context, input string) ([]message.Message, error) {
|
||||
messages := []message.Message{}
|
||||
messages = ast.withPrompts(messages)
|
||||
if storage != nil {
|
||||
history, err := storage.GetHistory(ctx.Sid, ctx.ChatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Add history messages
|
||||
for _, h := range history {
|
||||
messages = append(messages, *message.New().Map(h))
|
||||
}
|
||||
}
|
||||
|
||||
// Add user message
|
||||
messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": input, "name": ctx.Sid}))
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// Chat implements the chat functionality
|
||||
func (ast *Assistant) Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error {
|
||||
if ast.openai == nil {
|
||||
|
|
|
|||
|
|
@ -17,17 +17,26 @@ const (
|
|||
|
||||
// ResHookInit the response of the init hook
|
||||
type ResHookInit struct {
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
ChatID string `json:"chat_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
ChatID string `json:"chat_id,omitempty"`
|
||||
Next *NextAction `json:"next,omitempty"`
|
||||
Input []message.Message `json:"input,omitempty"`
|
||||
Options map[string]interface{} `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
// NextAction the next action
|
||||
type NextAction struct {
|
||||
Action string `json:"action"`
|
||||
Payload map[string]interface{} `json:"payload,omitempty"`
|
||||
}
|
||||
|
||||
// HookInit initialize the assistant
|
||||
func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, messages []message.Message) (*ResHookInit, error) {
|
||||
func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error) {
|
||||
// Create timeout context
|
||||
ctx, cancel := ast.createTimeoutContext(c)
|
||||
defer cancel()
|
||||
|
||||
v, err := ast.call(ctx, "Init", context, messages, c.Writer)
|
||||
v, err := ast.call(ctx, "Init", context, input, c.Writer)
|
||||
if err != nil {
|
||||
if err.Error() == HookErrorMethodNotFound {
|
||||
return nil, nil
|
||||
|
|
|
|||
|
|
@ -19,7 +19,8 @@ type API interface {
|
|||
Upload(ctx context.Context, file *multipart.FileHeader, reader io.Reader, option map[string]interface{}) (*File, error)
|
||||
Download(ctx context.Context, fileID string) (*FileResponse, error)
|
||||
ReadBase64(ctx context.Context, fileID string) (string, error)
|
||||
HookInit(c *gin.Context, ctx chatctx.Context, messages []message.Message) (*ResHookInit, error)
|
||||
Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error
|
||||
HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error)
|
||||
}
|
||||
|
||||
// RAG the RAG interface
|
||||
|
|
|
|||
160
neo/neo.go
160
neo/neo.go
|
|
@ -4,7 +4,6 @@ import (
|
|||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/kun/log"
|
||||
|
|
@ -13,44 +12,17 @@ import (
|
|||
"github.com/yaoapp/yao/neo/message"
|
||||
)
|
||||
|
||||
// Lock the assistant list
|
||||
var lock sync.Mutex = sync.Mutex{}
|
||||
|
||||
// Answer reply the message
|
||||
func (neo *DSL) Answer(ctx chatctx.Context, question string, c *gin.Context) error {
|
||||
messages, err := neo.withHistory(ctx, question)
|
||||
if err != nil {
|
||||
msg := message.New().Error(err).Done()
|
||||
msg.Write(c.Writer)
|
||||
return err
|
||||
}
|
||||
|
||||
var res *assistant.ResHookInit = nil
|
||||
var err error
|
||||
var ast assistant.API = neo.Assistant
|
||||
|
||||
if ctx.AssistantID != "" {
|
||||
ast, err = neo.Select(ctx.AssistantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Init the assistant
|
||||
res, err = ast.HookInit(c, ctx, messages)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Switch to the new assistant if necessary
|
||||
if res.AssistantID != ctx.AssistantID {
|
||||
ast, err = neo.Select(res.AssistantID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// Chat with AI
|
||||
return neo.chat(ast, ctx, messages, c)
|
||||
return ast.Execute(c, ctx, question, nil)
|
||||
}
|
||||
|
||||
// Select select an assistant
|
||||
|
|
@ -87,6 +59,7 @@ func (neo *DSL) GenerateChatTitle(ctx chatctx.Context, input string, c *gin.Cont
|
|||
2. The title should be a single sentence.
|
||||
3. The title should be in same language as the chat.
|
||||
4. The title should be no more than 50 characters.
|
||||
5. ANSWER ONLY THE TITLE CONTENT, FOR EXAMPLE: Chat with AI is a valid title, but "Chat with AI" is not a valid title.
|
||||
`
|
||||
isSilent := false
|
||||
if len(silent) > 0 {
|
||||
|
|
@ -273,130 +246,3 @@ func (neo *DSL) Download(ctx chatctx.Context, c *gin.Context) (*assistant.FileRe
|
|||
// Download file using the assistant
|
||||
return ast.Download(ctx.Context, fileID)
|
||||
}
|
||||
|
||||
// chat chat with AI
|
||||
func (neo *DSL) chat(ast assistant.API, ctx chatctx.Context, messages []message.Message, c *gin.Context) error {
|
||||
if ast == nil {
|
||||
msg := message.New().Error("assistant is not initialized").Done()
|
||||
msg.Write(c.Writer)
|
||||
return fmt.Errorf("assistant is not initialized")
|
||||
}
|
||||
|
||||
clientBreak := make(chan bool, 1)
|
||||
done := make(chan bool, 1)
|
||||
content := []byte{}
|
||||
|
||||
// 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
|
||||
|
||||
default:
|
||||
msg := message.NewOpenAI(data)
|
||||
if msg == nil {
|
||||
return 1 // continue
|
||||
}
|
||||
|
||||
// Handle error
|
||||
if msg.Type == "error" {
|
||||
value := msg.String()
|
||||
message.New().Error(value).Done().Write(c.Writer)
|
||||
return 0 // break
|
||||
}
|
||||
|
||||
// Append content and send message
|
||||
content = msg.Append(content)
|
||||
value := msg.String()
|
||||
if value != "" {
|
||||
message.New().
|
||||
Map(map[string]interface{}{
|
||||
"text": value,
|
||||
"done": msg.IsDone,
|
||||
}).
|
||||
Write(c.Writer)
|
||||
}
|
||||
|
||||
// Complete the stream
|
||||
if msg.IsDone {
|
||||
if value == "" {
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
func (neo *DSL) withHistory(ctx chatctx.Context, question string) ([]message.Message, error) {
|
||||
history, err := neo.Store.GetHistory(ctx.Sid, ctx.ChatID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Add history messages
|
||||
messages := []message.Message{}
|
||||
for _, h := range history {
|
||||
messages = append(messages, *message.New().Map(h))
|
||||
}
|
||||
|
||||
// Add user message
|
||||
messages = append(messages, *message.New().Map(map[string]interface{}{"role": "user", "content": question, "name": ctx.Sid}))
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
// saveHistory save the history
|
||||
func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages []message.Message) {
|
||||
if len(content) > 0 && sid != "" && len(messages) > 0 {
|
||||
err := neo.Store.SaveHistory(
|
||||
sid,
|
||||
[]map[string]interface{}{
|
||||
{"role": "user", "content": messages[len(messages)-1].Content(), "name": sid},
|
||||
{"role": "assistant", "content": string(content), "name": sid},
|
||||
},
|
||||
chatID,
|
||||
nil,
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
log.Error("Save history error: %s", err.Error())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sendMessage sends a message to the client
|
||||
func (neo *DSL) sendMessage(w gin.ResponseWriter, data interface{}) error {
|
||||
if msg, ok := data.(map[string]interface{}); ok {
|
||||
if !message.New().Map(msg).Write(w) {
|
||||
return fmt.Errorf("failed to write message to stream")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("invalid message data type")
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue