diff --git a/neo/neo.go b/neo/neo.go index 6c8dbc28..0b610c01 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -130,6 +130,29 @@ func (neo *DSL) API(router *gin.Engine, path string) error { } switch cmd { + + case "ModelList": + c.JSON(200, gin.H{"data": neo.Models, "code": 200}) + c.Done() + + case "SelectModel": + model, ok := payload["model"].(string) + if !ok { + c.JSON(400, gin.H{"message": "model is required", "code": 400}) + c.Done() + return + } + + err := neo.Select(model) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + return + } + + c.JSON(200, gin.H{"message": "success", "code": 200}) + c.Done() + case "ExitCommandMode": err := command.Exit(sid) if err != nil { @@ -524,11 +547,17 @@ func (neo *DSL) getGuardHandlers() ([]gin.HandlerFunc, error) { // NewAI create a new AI func (neo *DSL) newAI() error { - if neo.Connector == "" { - ai, err := openai.NewMoapi("gpt-3.5-turbo") + if neo.Connector == "" || strings.HasPrefix(neo.Connector, "moapi") { + model := "gpt-3.5-turbo" + if neo.Connector != "" { + model = strings.TrimPrefix(neo.Connector, "moapi:") + } + + ai, err := openai.NewMoapi(model) if err != nil { return err } + neo.AI = ai return nil } @@ -550,6 +579,16 @@ func (neo *DSL) newAI() error { return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) } +// Select select the model +func (neo *DSL) Select(model string) error { + ai, err := openai.NewMoapi(model) + if err != nil { + return err + } + neo.AI = ai + return nil +} + // newConversation create a new conversation func (neo *DSL) newConversation() error { diff --git a/neo/types.go b/neo/types.go index 37129e96..f6846caa 100644 --- a/neo/types.go +++ b/neo/types.go @@ -22,6 +22,7 @@ type DSL struct { Prompts []aigc.Prompt `json:"prompts,omitempty"` Allows []string `json:"allows,omitempty"` Command Command `json:"command,omitempty"` + Models []string `json:"models,omitempty"` AI aigc.AI `json:"-" yaml:"-"` Conversation conversation.Conversation `json:"-" yaml:"-"` GuardHandlers []gin.HandlerFunc `json:"-" yaml:"-"`