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

Enhance Neo API assistant with improved message handling and hook support

Browse source 
- Updated HookCreate to support extended timeout and input message parsing
- Modified CreateResponse to include optional input messages
- Refactored assistant hooks to improve context and message handling
- Added support for more flexible streaming and done hook processing
- Adjusted timeout contexts and method signatures for better error management

This change improves the Neo API assistant's flexibility and robustness in handling complex message interactions and hook scenarios.
This commit is contained in:
Max 2025-01-29 16:05:36 +08:00
parent 4164f082e9
commit 63807e1944
4 changed files with 59 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

44
neo/assistant/api.go
View file

@ -46,12 +46,21 @@ func GetByConnector(connector string, name string) (*Assistant, error) {
// Execute implements the execute functionality
func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}) error {
contents := chatMessage.NewContents()
return ast.execute(c, ctx, input, options, contents)
}
// Execute implements the execute functionality
func (ast *Assistant) execute(c *gin.Context, ctx chatctx.Context, input string, options map[string]interface{}, contents *chatMessage.Contents) error {
messages, err := ast.withHistory(ctx, input)
if err != nil {
return err
}
contents := chatMessage.NewContents()
if contents == nil {
contents = chatMessage.NewContents()
}
options = ast.withOptions(options)
// Add RAG and Version support
@ -85,7 +94,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
// Handle next action
if res != nil && res.Next != nil {
return res.Next.Execute(c, ctx)
return res.Next.Execute(c, ctx, contents)
}
// Update options if provided
@ -103,7 +112,7 @@ func (ast *Assistant) Execute(c *gin.Context, ctx chatctx.Context, input string,
}
// Execute the next action
func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *chatMessage.Contents) error {
switch next.Action {
case "process":
@ -164,7 +173,7 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context) error {
if v, ok := next.Payload["options"].(map[string]interface{}); ok {
options = v
}
return assistant.Execute(c, ctx, input, options)
return assistant.execute(c, ctx, input, options, contents)
case "exit":
return nil
@ -264,11 +273,11 @@ func (ast *Assistant) streamChat(
value := msg.String()
if value != "" {
// Handle stream
res, err := ast.HookStream(c, ctx, messages, contents)
res, err := ast.HookStream(c, ctx, messages, msg, contents)
if err == nil && res != nil {
if res.Next != nil {
err = res.Next.Execute(c, ctx)
err = res.Next.Execute(c, ctx, contents)
if err != nil {
chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
}
@ -299,22 +308,17 @@ func (ast *Assistant) streamChat(
// msg.Write(c.Writer)
// }
fmt.Println("Done", contents.JSON())
res, hookErr := ast.HookDone(c, ctx, messages, contents)
if hookErr == nil && res != nil {
if res.Output != nil {
chatMessage.New().
Map(map[string]interface{}{
"text": res.Input,
"done": true,
}).
Write(c.Writer)
}
if res.Next != nil {
err := res.Next.Execute(c, ctx)
err := res.Next.Execute(c, ctx, contents)
if err != nil {
chatMessage.New().Error(err.Error()).Done().Write(c.Writer)
}
done <- true
return 0 // break
}
@ -331,6 +335,16 @@ func (ast *Assistant) streamChat(
Write(c.Writer)
}
// Output
if res.Output != nil {
chatMessage.New().
Map(map[string]interface{}{
"text": res.Input,
"done": true,
}).
Write(c.Writer)
}
done <- true
return 0 // break
}

23
neo/assistant/hooks.go
View file

@ -17,9 +17,7 @@ import (
// HookInit initialize the assistant
func (ast *Assistant) HookInit(c *gin.Context, context chatctx.Context, input []message.Message, options map[string]interface{}, contents *message.Contents) (*ResHookInit, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
defer cancel()
ctx := ast.createBackgroundContext()
v, err := ast.call(ctx, "Init", c, contents, context, input, options)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
@ -72,13 +70,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, contents *chatMessage.Contents) (*ResHookStream, error) {
func (ast *Assistant) HookStream(c *gin.Context, context chatctx.Context, input []message.Message, msg *message.Message, contents *chatMessage.Contents) (*ResHookStream, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
ctx, cancel := ast.createTimeoutContext(5 * time.Second)
defer cancel()
v, err := ast.call(ctx, "Stream", c, contents, context, input)
v, err := ast.call(ctx, "Stream", c, contents, context, input, msg, contents.JSON())
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
@ -140,7 +138,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
// Create timeout context
ctx := ast.createBackgroundContext()
v, err := ast.call(ctx, "Done", c, contents, context, input)
v, err := ast.call(ctx, "Done", c, contents, context, input, contents.Data)
if err != nil {
if err.Error() == HookErrorMethodNotFound {
return nil, nil
@ -148,10 +146,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
return nil, err
}
response := &ResHookDone{
Input: input,
Output: contents.Data,
}
response := &ResHookDone{Input: input, Output: contents.Data}
switch v := v.(type) {
case map[string]interface{}:
@ -198,7 +193,7 @@ func (ast *Assistant) HookDone(c *gin.Context, context chatctx.Context, input []
// HookFail Handle failure of assistant response
func (ast *Assistant) HookFail(c *gin.Context, context chatctx.Context, input []message.Message, err error, contents *chatMessage.Contents) (*ResHookFail, error) {
// Create timeout context
ctx, cancel := ast.createTimeoutContext(c)
ctx, cancel := ast.createTimeoutContext(5 * time.Second)
defer cancel()
v, callErr := ast.call(ctx, "Fail", c, contents, context, input, err.Error())
@ -240,8 +235,8 @@ 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(context.Background(), 5*time.Second)
func (ast *Assistant) createTimeoutContext(time time.Duration) (context.Context, context.CancelFunc) {
ctx, cancel := context.WithTimeout(context.Background(), time)
return ctx, cancel
}

20
neo/hooks.go
View file

@ -8,6 +8,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/process"
chatctx "github.com/yaoapp/yao/neo/context"
"github.com/yaoapp/yao/neo/message"
)
// HookCreate create the assistant
@ -25,7 +26,7 @@ func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{
}
// Create a context with 10 second timeout
timeoutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
defer cancel()
p, err := process.Of(neo.Create, ctx, messages, c.Writer)
@ -63,10 +64,23 @@ func (neo *DSL) HookCreate(ctx chatctx.Context, messages []map[string]interface{
chatID = ctx.ChatID
}
return CreateResponse{AssistantID: assistantID, ChatID: chatID}, nil
// Messages fixed input
input := []message.Message{}
if vv, has := v["input"]; has {
bytes, err := jsoniter.Marshal(vv)
if err != nil {
return CreateResponse{}, err
}
err = jsoniter.Unmarshal(bytes, &input)
if err != nil {
return CreateResponse{}, err
}
}
return CreateResponse{AssistantID: assistantID, ChatID: chatID, Input: input}, nil
}
return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID}, nil
return CreateResponse{AssistantID: assistantID, ChatID: ctx.ChatID, Input: nil}, nil
}
// HookPrepare executes the prepare hook before AI is called

6
neo/types.go
View file

@ -3,6 +3,7 @@ package neo
import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/neo/assistant"
"github.com/yaoapp/yao/neo/message"
"github.com/yaoapp/yao/neo/rag"
"github.com/yaoapp/yao/neo/store"
"github.com/yaoapp/yao/neo/vision"
@ -62,6 +63,7 @@ type FileUpload struct {
// CreateResponse the response of the create hook
type CreateResponse 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"`
Input []message.Message `json:"messages,omitempty"`
}