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)