From c8643bbe45eb0bea0e500c77ab4fe64f6ca00ad8 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 14 Dec 2024 14:53:01 +0800 Subject: [PATCH] Refactor HookCreate and assistant management in Neo API to improve response handling and streamline assistant creation. Update HookCreate to return structured CreateResponse with AssistantID and ChatID, enhancing error handling. Modify Load function to create the default assistant after querying the assistant list, ensuring proper initialization. Refactor newAssistant methods for better clarity and maintainability, and update types to include CreateResponse struct for improved response management. --- neo/assistant/types.go | 14 ++++----- neo/hooks.go | 39 +++++++++++++++++++++---- neo/load.go | 13 +++++---- neo/neo.go | 66 ++++++++++++++++++++++++++++++++++++------ neo/types.go | 6 ++++ 5 files changed, 110 insertions(+), 28 deletions(-) diff --git a/neo/assistant/types.go b/neo/assistant/types.go index bffb1ea7..ce52d3f5 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -27,11 +27,11 @@ type QueryParam struct { // Assistant the assistant type Assistant struct { - ID string `json:"assistant_id"` // Assistant ID - Name string `json:"name,omitempty"` // Assistant Name - Description string `json:"description"` // Assistant Description - Connector string `json:"connector"` // AI Connector - Option map[string]interface{} `json:"option"` // AI Option - Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts - API API `json:"-" yaml:"-"` // Assistant API + ID string `json:"assistant_id"` // Assistant ID + Name string `json:"name,omitempty"` // Assistant Name + Connector string `json:"connector"` // AI Connector + Description string `json:"description,omitempty"` // Assistant Description + Option map[string]interface{} `json:"option,omitempty"` // AI Option + Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts + API API `json:"-" yaml:"-"` // Assistant API } diff --git a/neo/hooks.go b/neo/hooks.go index 544a0ca4..c189c74f 100644 --- a/neo/hooks.go +++ b/neo/hooks.go @@ -11,9 +11,9 @@ import ( ) // HookCreate create the assistant -func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) error { +func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gin.Context) (CreateResponse, error) { if neo.Create == "" { - return nil + return CreateResponse{AssistantID: neo.Use, ChatID: ctx.ChatID}, nil } // Create a context with 10 second timeout @@ -22,21 +22,48 @@ func (neo *DSL) HookCreate(ctx Context, messages []map[string]interface{}, c *gi p, err := process.Of(neo.Create, ctx, messages, c.Writer) if err != nil { - return err + return CreateResponse{}, err } err = p.WithContext(timeoutCtx).Execute() if err != nil { - return err + return CreateResponse{}, err } defer p.Release() // Check if context was canceled if timeoutCtx.Err() != nil { - return timeoutCtx.Err() + return CreateResponse{}, timeoutCtx.Err() } - return nil + value := p.Value() + switch v := value.(type) { + case CreateResponse: + return v, nil + + case map[string]interface{}: + assistantID := "" + if id, ok := v["assistant_id"].(string); ok { + assistantID = id + } + + if assistantID == "" && neo.Use != "" { + assistantID = neo.Use + } + chatID := "" + if id, ok := v["chat_id"].(string); ok { + chatID = id + } + + if chatID == "" { + chatID = ctx.ChatID + } + + return CreateResponse{AssistantID: assistantID, ChatID: chatID}, nil + } + + // Default assistant + return CreateResponse{AssistantID: neo.Use, ChatID: ctx.ChatID}, nil } // HookAssistants query the assistant list from the assistant list hook diff --git a/neo/load.go b/neo/load.go index 9bbbdd8d..dec5fb2f 100644 --- a/neo/load.go +++ b/neo/load.go @@ -51,12 +51,6 @@ func Load(cfg config.Config) error { return err } - // Create Default Assistant - Neo.Assistant, err = Neo.createDefaultAssistant() - if err != nil { - return err - } - // Query Assistant List ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -73,6 +67,13 @@ func Load(cfg config.Config) error { if err != nil { return fmt.Errorf("Neo assistant list failed: %w", err) } + + // Create Default Assistant + Neo.Assistant, err = Neo.createDefaultAssistant() + if err != nil { + return err + } + return nil case <-ctx.Done(): return fmt.Errorf("Neo assistant list timeout: %w", ctx.Err()) diff --git a/neo/neo.go b/neo/neo.go index ee189ba9..54a863c0 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -30,13 +30,29 @@ func (neo *DSL) Answer(ctx Context, question string, c *gin.Context) error { return err } - err = neo.HookCreate(ctx, messages, c) + // Get the assistant_id, chat_id + res, err := neo.HookCreate(ctx, messages, c) if err != nil { msg := message.New().Error(err).Done() msg.Write(c.Writer) return err } + // Select Assistant + ast := neo.Assistant + if res.AssistantID != "" { + ast, err = neo.newAssistant(res.AssistantID) + if err != nil { + msg := message.New().Error(err).Done() + msg.Write(c.Writer) + return err + } + } + + // Chat with AI + + fmt.Println(ast) + // Get the assistant_id, chat_id time.Sleep(1 * time.Second) @@ -69,14 +85,38 @@ func (neo *DSL) updateAssistantList(list []assistant.Assistant) { } } -// createDefaultAssistant create a default assistant -func (neo *DSL) createDefaultAssistant() (assistant.API, error) { +// newAssistant create a new assistant +func (neo *DSL) newAssistant(id string) (assistant.API, error) { + // Try to find assistant in AssistantList first + if id != "" && neo.AssistantMaps != nil { + if ast, ok := neo.AssistantMaps[id]; ok { - // Moapi - if neo.Connector == "" || strings.HasPrefix(neo.Connector, "moapi") { + if ast.API != nil { + return ast.API, nil + } + api, err := neo.newAssistantByConfig(&ast) + if err != nil { + return nil, err + } + ast.API = api + return api, nil + } + } + return neo.newAssistantByConnector(id) +} + +// newAssistantByConfig create a new assistant from assistant configuration +func (neo *DSL) newAssistantByConfig(ast *assistant.Assistant) (assistant.API, error) { + return neo.newAssistantByConnector(ast.Connector) +} + +// newAssistantByConnector create a new assistant from connector id +func (neo *DSL) newAssistantByConnector(id string) (assistant.API, error) { + // Moapi connector + if id == "" || strings.HasPrefix(id, "moapi") { model := "gpt-3.5-turbo" - if strings.HasPrefix(neo.Connector, "moapi:") { - model = strings.TrimPrefix(neo.Connector, "moapi:") + if strings.HasPrefix(id, "moapi:") { + model = strings.TrimPrefix(id, "moapi:") } conn, err := connector.New(`moapi`, `__yao.moapi`, []byte(`{"model": "`+model+`"}`)) @@ -92,9 +132,9 @@ func (neo *DSL) createDefaultAssistant() (assistant.API, error) { } // Other connector - conn, err := connector.Select(neo.Connector) + conn, err := connector.Select(id) if err != nil { - return nil, fmt.Errorf("Neo assistant connector %s not support", neo.Connector) + return nil, fmt.Errorf("Neo assistant connector %s not support", id) } if conn.Is(connector.OPENAI) { @@ -113,6 +153,14 @@ func (neo *DSL) createDefaultAssistant() (assistant.API, error) { return api, nil } +// createDefaultAssistant create a default assistant +func (neo *DSL) createDefaultAssistant() (assistant.API, error) { + if neo.Use != "" { + return neo.newAssistant(neo.Use) + } + return neo.newAssistant(neo.Connector) +} + // // AnswerOld reply the message // func (neo *DSL) AnswerOld(ctx Context, question string, c *gin.Context) error { // // get the chat messages diff --git a/neo/types.go b/neo/types.go index e8fc85f3..c26f86d3 100644 --- a/neo/types.go +++ b/neo/types.go @@ -52,6 +52,12 @@ type Field struct { Bind string `json:"bind,omitempty"` } +// CreateResponse the response of the create hook +type CreateResponse struct { + AssistantID string `json:"assistant_id,omitempty"` + ChatID string `json:"chat_id,omitempty"` +} + // AI the AI interface type AI interface { ChatCompletions(messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception)