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

Enhance Neo API assistant response handling with new hook methods

Browse source 
- Introduced HookDone and HookFail methods to manage completion and failure scenarios in assistant responses, improving error handling and output customization.
- Updated streamChat method to utilize these hooks, allowing for more flexible response management based on the assistant's output and error states.
- Enhanced ResHookDone and ResHookFail structs to include next action handling, input messages, and error information, providing better control over 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 2025-01-13 17:32:51 +08:00
parent ca8993f4c8
commit e440f1ff81
3 changed files with 146 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

70
neo/assistant/api.go
View file

@ -135,6 +135,13 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
// Handle error
if msg.Type == "error" {
value := msg.String()
res, hookErr := ast.HookFail(c, ctx, messages, string(*content), fmt.Errorf("%s", value))
if hookErr == nil && res != nil && (res.Output != "" || res.Error != "") {
value = res.Output
if res.Error != "" {
value = res.Error
}
}
chatMessage.New().Error(value).Done().Write(c.Writer)
return 0 // break
}
@ -143,35 +150,27 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
*content = msg.Append(*content)
value := msg.String()
if value != "" {
// 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":
res, err := ast.HookStream(c, ctx, messages, string(*content))
if err == nil && res != nil {
if res.Output != "" {
value = res.Output
}
if res.Next != nil && res.Next.Action == "exit" {
done <- true
return 0 // break
}
if res.Silent {
return 1 // continue
}
}
if !res.Silent {
chatMessage.New().
Map(map[string]interface{}{
"text": value,
"done": msg.IsDone,
}).
Write(c.Writer)
}
chatMessage.New().
Map(map[string]interface{}{
"text": value,
"done": msg.IsDone,
}).
Write(c.Writer)
}
// Complete the stream
@ -179,6 +178,31 @@ func (ast *Assistant) streamChat(c *gin.Context, ctx chatctx.Context, messages [
if value == "" {
msg.Write(c.Writer)
}
// Call HookDone
res, hookErr := ast.HookDone(c, ctx, messages, string(*content))
if hookErr == nil && res != nil {
if res.Output != "" {
chatMessage.New().
Map(map[string]interface{}{
"text": res.Output,
"done": true,
}).
Write(c.Writer)
}
if res.Next != nil && res.Next.Action == "exit" {
done <- true
return 0 // break
}
} else if value != "" {
chatMessage.New().
Map(map[string]interface{}{
"text": value,
"done": true,
}).
Write(c.Writer)
}
done <- true
return 0 // break
}

84
neo/assistant/hooks.go
View file

@ -99,6 +99,90 @@ func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input
return response, nil
}
// HookDone Handle completion of assistant response
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, c.Writer)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, err
}
response := &ResHookDone{
Input: input,
Output: output,
}
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
}
}
case string:
response.Output = v
}
return response, nil
}
// HookFail Handle failure of assistant response
func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, output string, err error) (*ResHookFail, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
v, callErr := ast.call(ctx, "Fail", context, input, output, err.Error(), c.Writer)
if callErr != nil {
if callErr.Error() == HookErrorMethodNotFound {
return nil, nil
}
return nil, callErr
}
response := &ResHookFail{
Input: input,
Output: output,
Error: err.Error(),
}
switch v := v.(type) {
case map[string]interface{}:
if res, ok := v["output"].(string); ok {
response.Output = res
}
if res, ok := v["error"].(string); ok {
response.Error = 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.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)

15
neo/assistant/types.go
View file

@ -44,6 +44,21 @@ type ResHookStream struct {
Output string `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"`
}
// ResHookFail the response of the fail hook
type ResHookFail struct {
Next *NextAction `json:"next,omitempty"`
Input []message.Message `json:"input,omitempty"`
Output string `json:"output,omitempty"`
Error string `json:"error,omitempty"`
}
// NextAction the next action
type NextAction struct {
Action string `json:"action"`