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

Merge pull request #839 from trheyi/main

Browse source 
Enhance Neo API assistant with improved message handling and hook sup…
This commit is contained in:
Max 2025-01-29 16:14:37 +08:00 • committed by GitHub
commit 4911aa5a5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
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"`
}