Enhance Neo API assistant interaction with improved message handling and streaming support
- Added message handling in the Execute method to support input messages from the response. - Updated the handleChatStream method to pass the context to the streamChat method, improving context management during chat streaming. - Introduced the HookStream method to handle streaming responses, allowing for custom output and next action handling based on the assistant's response. - Enhanced the ResHookStream struct to include silent output control and next action management, providing better flexibility in assistant interactions. These changes improve the robustness and maintainability of the Neo API, paving the way for enhanced assistant functionalities and message management.
This commit is contained in:
Max
parent
cb2cd0c317
commit
ca8993f4c8
3 changed files with 116 additions and 28 deletions
|
|
@ -81,6 +81,11 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
|
|||
options = res.Options
|
||||
}
|
||||
|
||||
// messages
|
||||
if res.Input != nil {
|
||||
messages = res.Input
|
||||
}
|
||||
|
||||
// Only proceed with chat stream if no specific next action was handled
|
||||
return ast.handleChatStream(c, ctx, messages, options)
|
||||
}
|
||||
|
|
@ -93,7 +98,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
|
|||
|
||||
// Chat with AI in background
|
||||
go func() {
|
||||
err := ast.streamChat(c, messages, options, clientBreak, done, &content)
|
||||
err := ast.streamChat(c, ctx, messages, options, clientBreak, done, &content)
|
||||
if err != nil {
|
||||
chatMessage.New().Error(err).Done().Write(c.Writer)
|
||||
}
|
||||
|
|
@ -113,7 +118,7 @@ func (ast *Assistant) handleChatStream(c *gin.Context, ctx chatctx.Context, mess
|
|||
}
|
||||
|
||||
// streamChat handles the streaming chat interaction
|
||||
func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, options map[string]interface{},
|
||||
func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.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 {
|
||||
|
|
@ -138,12 +143,35 @@ func (ast *Assistant) streamChat(c *gin.Context, messages []message.Message, opt
|
|||
*content = msg.Append(*content)
|
||||
value := msg.String()
|
||||
if value != "" {
|
||||
chatMessage.New().
|
||||
Map(map[string]interface{}{
|
||||
"text": value,
|
||||
"done": msg.IsDone,
|
||||
}).
|
||||
Write(c.Writer)
|
||||
|
||||
// Handle stream
|
||||
res, err := ast.HookStream(c, ctx, messages, value)
|
||||
if err != nil {
|
||||
return 0 // break
|
||||
}
|
||||
|
||||
// Custom output from hook
|
||||
if res.Output != "" {
|
||||
value = res.Output
|
||||
}
|
||||
|
||||
// Custom next action from hook
|
||||
if res.Next != nil {
|
||||
switch res.Next.Action {
|
||||
case "exit":
|
||||
done <- true
|
||||
return 0 // break
|
||||
}
|
||||
}
|
||||
|
||||
if !res.Silent {
|
||||
chatMessage.New().
|
||||
Map(map[string]interface{}{
|
||||
"text": value,
|
||||
"done": msg.IsDone,
|
||||
}).
|
||||
Write(c.Writer)
|
||||
}
|
||||
}
|
||||
|
||||
// Complete the stream
|
||||
|
|
|
|||
|
|
@ -10,26 +10,6 @@ import (
|
|||
"github.com/yaoapp/yao/neo/message"
|
||||
)
|
||||
|
||||
const (
|
||||
// HookErrorMethodNotFound is the error message for method not found
|
||||
HookErrorMethodNotFound = "method not found"
|
||||
)
|
||||
|
||||
// ResHookInit the response of the init hook
|
||||
type ResHookInit struct {
|
||||
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, input []message.Message, options map[string]interface{}) (*ResHookInit, error) {
|
||||
// Create timeout context
|
||||
|
|
@ -54,6 +34,16 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
|
|||
response.ChatID = res
|
||||
}
|
||||
|
||||
if res, ok := v["next"].(map[string]interface{}); ok {
|
||||
response.Next = &NextAction{}
|
||||
if name, ok := res["action"].(string); ok {
|
||||
response.Next.Action = name
|
||||
}
|
||||
if payload, ok := res["payload"].(map[string]interface{}); ok {
|
||||
response.Next.Payload = payload
|
||||
}
|
||||
}
|
||||
|
||||
case string:
|
||||
response.AssistantID = v
|
||||
response.ChatID = context.ChatID
|
||||
|
|
@ -66,6 +56,49 @@ func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []
|
|||
return response, nil
|
||||
}
|
||||
|
||||
// HookStream Handle streaming response from LLM
|
||||
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, c.Writer)
|
||||
if err != nil {
|
||||
if err.Error() == HookErrorMethodNotFound {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
response := &ResHookStream{}
|
||||
switch v := v.(type) {
|
||||
case map[string]interface{}:
|
||||
if res, ok := v["output"].(string); ok {
|
||||
response.Output = res
|
||||
}
|
||||
if res, ok := v["next"].(map[string]interface{}); ok {
|
||||
response.Next = &NextAction{}
|
||||
if name, ok := res["action"].(string); ok {
|
||||
response.Next.Action = name
|
||||
}
|
||||
if payload, ok := res["payload"].(map[string]interface{}); ok {
|
||||
response.Next.Payload = payload
|
||||
}
|
||||
}
|
||||
|
||||
// Custom silent from hook
|
||||
if res, ok := v["silent"].(bool); ok {
|
||||
response.Silent = res
|
||||
}
|
||||
|
||||
case string:
|
||||
response.Output = v
|
||||
}
|
||||
|
||||
return response, nil
|
||||
}
|
||||
|
||||
// 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)
|
||||
|
|
|
|||
|
|
@ -13,6 +13,11 @@ import (
|
|||
api "github.com/yaoapp/yao/openai"
|
||||
)
|
||||
|
||||
const (
|
||||
// HookErrorMethodNotFound is the error message for method not found
|
||||
HookErrorMethodNotFound = "method not found"
|
||||
)
|
||||
|
||||
// API the assistant API interface
|
||||
type API interface {
|
||||
Chat(ctx context.Context, messages []message.Message, option map[string]interface{}, cb func(data []byte) int) error
|
||||
|
|
@ -23,6 +28,28 @@ type API interface {
|
|||
HookInit(c *gin.Context, ctx chatctx.Context, input []message.Message, options map[string]interface{}) (*ResHookInit, error)
|
||||
}
|
||||
|
||||
// ResHookInit the response of the init hook
|
||||
type ResHookInit struct {
|
||||
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"`
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// NextAction the next action
|
||||
type NextAction struct {
|
||||
Action string `json:"action"`
|
||||
Payload map[string]interface{} `json:"payload,omitempty"`
|
||||
}
|
||||
|
||||
// RAG the RAG interface
|
||||
type RAG struct {
|
||||
Engine driver.Engine
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue