From 1bafeabdbac104423cae645151dd8c8a2696eb1f Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 29 Apr 2023 19:16:13 +0800 Subject: [PATCH] [add] Neo tests ( 60% ) --- aigc/aigc.go | 4 +- aigc/types.go | 6 +-- neo/load.go | 4 +- neo/load_test.go | 24 +++++++++ neo/neo.go | 21 +++++--- neo/neo_test.go | 128 +++++++++++++++++++++++++++++++++++++++++++++++ neo/types.go | 10 ++-- 7 files changed, 178 insertions(+), 19 deletions(-) create mode 100644 neo/load_test.go create mode 100644 neo/neo_test.go diff --git a/aigc/aigc.go b/aigc/aigc.go index bc998acb..69b69958 100644 --- a/aigc/aigc.go +++ b/aigc/aigc.go @@ -30,8 +30,8 @@ func (ai *DSL) Call(content string, user string, option map[string]interface{}) messages := []map[string]interface{}{} for _, prompt := range ai.Prompts { message := map[string]interface{}{"role": prompt.Role, "content": prompt.Content} - if prompt.User != "" { - message["user"] = prompt.User + if prompt.Name != "" { + message["name"] = prompt.Name } messages = append(messages, message) } diff --git a/aigc/types.go b/aigc/types.go index 98d16970..08f40eea 100644 --- a/aigc/types.go +++ b/aigc/types.go @@ -8,20 +8,20 @@ import ( // DSL the connector DSL type DSL struct { - ID string `json:"-"` + ID string `json:"-" yaml:"-"` Name string `json:"name,omitempty"` Connector string `json:"connector"` Process string `json:"process,omitempty"` Prompts []Prompt `json:"prompts"` Optional Optional `json:"optional,omitempty"` - AI AI `json:"-"` + AI AI `json:"-" yaml:"-"` } // Prompt a prompt type Prompt struct { Role string `json:"role"` Content string `json:"content"` - User string `json:"user,omitempty"` + Name string `json:"name,omitempty"` } // Optional optional diff --git a/neo/load.go b/neo/load.go index 45c3542c..e6ea1d3c 100644 --- a/neo/load.go +++ b/neo/load.go @@ -26,12 +26,12 @@ func Load(cfg config.Config) error { return err } - err = application.Parse("neo.yml", bytes, &neo) + err = application.Parse("neo.yml", bytes, &setting) if err != nil { return err } - *neo = setting + neo = &setting err = neo.newAI() if err != nil { return err diff --git a/neo/load_test.go b/neo/load_test.go new file mode 100644 index 00000000..b7216317 --- /dev/null +++ b/neo/load_test.go @@ -0,0 +1,24 @@ +package neo + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +func TestLoad(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + err := Load(config.Conf) + if err != nil { + t.Fatal(err) + } + check(t) +} + +func check(t *testing.T) { + assert.NotNil(t, neo) +} diff --git a/neo/neo.go b/neo/neo.go index 13da4bdd..fa27bf01 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -8,6 +8,7 @@ import ( "net/url" "github.com/gin-gonic/gin" + "github.com/google/uuid" "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/process" @@ -21,7 +22,11 @@ func (neo *Neo) API(router *gin.Engine, path string, allows ...string) error { prompts := []map[string]interface{}{} for _, prompt := range neo.Prompts { - prompts = append(prompts, map[string]interface{}{"role": prompt.Role, "content": prompt.Content, "user": prompt.User}) + message := map[string]interface{}{"role": prompt.Role, "content": prompt.Content} + if prompt.Name != "" { + message["name"] = prompt.Name + } + prompts = append(prompts, message) } // set the guard @@ -37,7 +42,11 @@ func (neo *Neo) API(router *gin.Engine, path string, allows ...string) error { router.GET(path, func(c *gin.Context) { sid := c.GetString("__sid") - content := c.GetString("content") + if sid == "" { + sid = uuid.New().String() + } + + content := c.Query("content") if content == "" { c.JSON(400, gin.H{"message": "content is required", "code": 400}) return @@ -51,7 +60,7 @@ func (neo *Neo) API(router *gin.Engine, path string, allows ...string) error { } messages = append(messages, history...) - messages = append(messages, map[string]interface{}{"role": "user", "content": content, "user": sid}) + messages = append(messages, map[string]interface{}{"role": "user", "content": content, "name": sid}) // reply the content ctx, cancel := context.WithCancel(context.Background()) @@ -80,7 +89,7 @@ func (neo *Neo) Answer(ctx context.Context, c *gin.Context, messages []map[strin close(chanError) }() - _, ex := neo.AI.ChatCompletions(messages, neo.Option, func(data []byte) int { + _, ex := neo.AI.ChatCompletionsWith(ctx, messages, neo.Option, func(data []byte) int { chanStream <- data return 1 }) @@ -103,9 +112,6 @@ func (neo *Neo) Answer(ctx context.Context, c *gin.Context, messages []map[strin msg = append(msg, []byte("\n")...) w.Write(msg) return true - - case <-ctx.Done(): - return false } }) @@ -198,6 +204,7 @@ func (neo *Neo) newAI() error { return err } neo.AI = ai + return nil } return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) diff --git a/neo/neo_test.go b/neo/neo_test.go new file mode 100644 index 00000000..1b85ed27 --- /dev/null +++ b/neo/neo_test.go @@ -0,0 +1,128 @@ +package neo + +import ( + "context" + "fmt" + "net" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + httpTest "github.com/yaoapp/gou/http" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/test" + _ "github.com/yaoapp/yao/utils" +) + +func TestAPI(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + // test router + router := testRouter(t) + err := neo.API(router, "/neo/chat") + if err != nil { + t.Fatal(err) + } + + // test server + host, shutdown := testServer(t, router) + defer shutdown() + + // test request + url := fmt.Sprintf("%s/neo/chat?content=hello&token=%s", host, testToken(t)) + res := []byte{} + req := httpTest.New(url). + WithHeader(http.Header{"Content-Type": []string{"application/json"}}) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + // send request + req.Stream(ctx, "GET", nil, func(data []byte) int { + res = append(res, data...) + return 1 + }) + + assert.Contains(t, string(res), "[DONE]") +} + +func TestAPIAuth(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + router := testRouter(t) + err := neo.API(router, "/neo/chat") + if err != nil { + t.Fatal(err) + } + + response := httptest.NewRecorder() + req, _ := http.NewRequest("GET", "/neo/chat?content=hello", nil) + assert.Panics(t, func() { + router.ServeHTTP(response, req) + }) +} + +func testServer(t *testing.T, router *gin.Engine) (string, func()) { + + // Listen + l, err := net.Listen("tcp4", ":0") + if err != nil { + t.Fatal(err) + } + + srv := &http.Server{Addr: ":0", Handler: router} + + // start serve + go func() { + if err := srv.Serve(l); err != nil && err != http.ErrServerClosed { + fmt.Println("[TestServer] Error:", err) + return + } + }() + + addr := strings.Split(l.Addr().String(), ":") + if len(addr) != 2 { + t.Fatal("invalid address") + } + + host := fmt.Sprintf("http://127.0.0.1:%s", addr[1]) + time.Sleep(50 * time.Millisecond) + + shutdown := func() { + srv.Close() + l.Close() + } + return host, shutdown +} + +func testRouter(t *testing.T) *gin.Engine { + + err := Load(config.Conf) + if err != nil { + t.Fatal(err) + } + + router := gin.New() + gin.SetMode(gin.ReleaseMode) + return router +} + +func testToken(t *testing.T) string { + token := helper.JwtMake(1, + map[string]interface{}{ + "id": 1, + "name": "Test", + }, + map[string]interface{}{ + "exp": 3600, + "sid": "123456", + }) + return token.Token +} diff --git a/neo/types.go b/neo/types.go index b3f837ef..85ef066c 100644 --- a/neo/types.go +++ b/neo/types.go @@ -4,17 +4,17 @@ import "github.com/yaoapp/yao/aigc" // Neo AI assistant type Neo struct { - ID string `json:"-"` + ID string `json:"-" yaml:"-"` Name string `json:"name,omitempty"` Guard string `json:"guard,omitempty"` Connector string `json:"connector"` ConversationSetting ConversationSetting `json:"conversation"` Option map[string]interface{} `json:"option"` - Prompts []aigc.Prompt `json:"prompts"` + Prompts []aigc.Prompt `json:"prompts,omitempty"` Allows []string `json:"allows,omitempty"` - AI aigc.AI `json:"-"` - Conversation Conversation `json:"-"` - Command Command `json:"-"` + AI aigc.AI `json:"-" yaml:"-"` + Conversation Conversation `json:"-" yaml:"-"` + Command Command `json:"-" yaml:"-"` } // ConversationSetting the conversation config