From 60ef6c1e5c94684902fca3e0be38d286ba760b14 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 29 Apr 2023 16:48:37 +0800 Subject: [PATCH] [add] neo (50%) --- Makefile | 2 +- aigc/aigc_test.go | 2 +- aigc/load_test.go | 2 +- aigc/types.go | 7 +- neo/conversation/mongo.go | 19 +++ neo/conversation/redis.go | 19 +++ neo/conversation/weaviate.go | 19 +++ neo/conversation/xun.go | 19 +++ neo/load.go | 49 ++++++++ neo/neo.go | 237 +++++++++++++++++++++++++++++++++++ neo/types.go | 36 ++++++ openai/openai.go | 43 ++++++- openai/openai_test.go | 60 +++++++++ 13 files changed, 506 insertions(+), 8 deletions(-) create mode 100644 neo/conversation/mongo.go create mode 100644 neo/conversation/redis.go create mode 100644 neo/conversation/weaviate.go create mode 100644 neo/conversation/xun.go create mode 100644 neo/load.go create mode 100644 neo/neo.go create mode 100644 neo/types.go diff --git a/Makefile b/Makefile index a32e35ab..dd43e0aa 100644 --- a/Makefile +++ b/Makefile @@ -9,7 +9,7 @@ COMMIT := $(shell git log | head -n 1 | awk '{print substr($$2, 0, 12)}') NOW := $(shell date +"%FT%T%z") # ROOT_DIR := $(shell dirname $(realpath $(firstword $(MAKEFILE_LIST)))) -TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|tests|openai|aigc|share*') +TESTFOLDER := $(shell $(GO) list ./... | grep -vE 'examples|tests|openai|aigc|neo|share*') TESTTAGS ?= "" # TESTWIDGETS := $(shell $(GO) list ./widgets/...) diff --git a/aigc/aigc_test.go b/aigc/aigc_test.go index 84524ad5..4d988784 100644 --- a/aigc/aigc_test.go +++ b/aigc/aigc_test.go @@ -22,7 +22,7 @@ func TestCall(t *testing.T) { if ex != nil { t.Fatal(ex.Message) } - assert.Equal(t, "Hello", content) + assert.Contains(t, content, "Hello") } func TestCallWithProcess(t *testing.T) { diff --git a/aigc/load_test.go b/aigc/load_test.go index 016212dc..f683ff95 100644 --- a/aigc/load_test.go +++ b/aigc/load_test.go @@ -24,5 +24,5 @@ func check(t *testing.T) { assert.True(t, ids["translate"]) assert.True(t, ids["draw"]) - assert.Equal(t, 2, len(Autopilots)) + assert.GreaterOrEqual(t, len(Autopilots), 2) } diff --git a/aigc/types.go b/aigc/types.go index f9934356..98d16970 100644 --- a/aigc/types.go +++ b/aigc/types.go @@ -1,6 +1,10 @@ package aigc -import "github.com/yaoapp/kun/exception" +import ( + "context" + + "github.com/yaoapp/kun/exception" +) // DSL the connector DSL type DSL struct { @@ -29,6 +33,7 @@ type Optional struct { // AI the AI interface type AI interface { ChatCompletions(messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) + ChatCompletionsWith(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) GetContent(response interface{}) (string, *exception.Exception) Embeddings(input interface{}, user string) (interface{}, *exception.Exception) Tiktoken(input string) (int, error) diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go new file mode 100644 index 00000000..415dc3a5 --- /dev/null +++ b/neo/conversation/mongo.go @@ -0,0 +1,19 @@ +package conversation + +// Mongo conversation +type Mongo struct{} + +// NewMongo create a new conversation +func NewMongo() *Mongo { + return &Mongo{} +} + +// GetHistory get the history +func (conv *Mongo) GetHistory(sid string) ([]map[string]interface{}, error) { + return []map[string]interface{}{}, nil +} + +// SaveHistory save the history +func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}) error { + return nil +} diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go new file mode 100644 index 00000000..6481c21f --- /dev/null +++ b/neo/conversation/redis.go @@ -0,0 +1,19 @@ +package conversation + +// Redis conversation +type Redis struct{} + +// NewRedis create a new conversation +func NewRedis() *Redis { + return &Redis{} +} + +// GetHistory get the history +func (conv *Redis) GetHistory(sid string) ([]map[string]interface{}, error) { + return []map[string]interface{}{}, nil +} + +// SaveHistory save the history +func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}) error { + return nil +} diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go new file mode 100644 index 00000000..88542d54 --- /dev/null +++ b/neo/conversation/weaviate.go @@ -0,0 +1,19 @@ +package conversation + +// Weaviate Database conversation +type Weaviate struct{} + +// NewWeaviate create a new conversation +func NewWeaviate() *Weaviate { + return &Weaviate{} +} + +// GetHistory get the history +func (conv *Weaviate) GetHistory(sid string) ([]map[string]interface{}, error) { + return []map[string]interface{}{}, nil +} + +// SaveHistory save the history +func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}) error { + return nil +} diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go new file mode 100644 index 00000000..0f0baf30 --- /dev/null +++ b/neo/conversation/xun.go @@ -0,0 +1,19 @@ +package conversation + +// Xun Database conversation +type Xun struct{} + +// NewXun create a new conversation +func NewXun() *Xun { + return &Xun{} +} + +// GetHistory get the history +func (conv *Xun) GetHistory(sid string) ([]map[string]interface{}, error) { + return []map[string]interface{}{}, nil +} + +// SaveHistory save the history +func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}) error { + return nil +} diff --git a/neo/load.go b/neo/load.go new file mode 100644 index 00000000..45c3542c --- /dev/null +++ b/neo/load.go @@ -0,0 +1,49 @@ +package neo + +import ( + "path/filepath" + + "github.com/yaoapp/gou/application" + "github.com/yaoapp/yao/aigc" + "github.com/yaoapp/yao/config" +) + +var neo *Neo + +// Load load AIGC +func Load(cfg config.Config) error { + + setting := Neo{ + ID: "neo", + Prompts: []aigc.Prompt{}, + Option: map[string]interface{}{}, + Allows: []string{}, + ConversationSetting: ConversationSetting{Table: "yao_neo_conversation", MaxSize: 100, Connector: "default"}, + } + + bytes, err := application.App.Read(filepath.Join("neo", "neo.yml")) + if err != nil { + return err + } + + err = application.Parse("neo.yml", bytes, &neo) + if err != nil { + return err + } + + *neo = setting + err = neo.newAI() + if err != nil { + return err + } + + err = neo.newConversation() + if err != nil { + return err + } + + return nil +} + +// LoadCommands load the commands +func (neo *Neo) LoadCommands() {} diff --git a/neo/neo.go b/neo/neo.go new file mode 100644 index 00000000..13da4bdd --- /dev/null +++ b/neo/neo.go @@ -0,0 +1,237 @@ +package neo + +import ( + "context" + "fmt" + "io" + "net/http" + "net/url" + + "github.com/gin-gonic/gin" + "github.com/yaoapp/gou/api" + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/gou/process" + "github.com/yaoapp/yao/helper" + "github.com/yaoapp/yao/neo/conversation" + "github.com/yaoapp/yao/openai" +) + +// API is a method on the Neo type +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}) + } + + // set the guard + err := neo.setGuard(router) + if err != nil { + return err + } + + // Cross-Domain + neo.crossDomain(router, path, allows...) + + // api router + router.GET(path, func(c *gin.Context) { + + sid := c.GetString("__sid") + content := c.GetString("content") + if content == "" { + c.JSON(400, gin.H{"message": "content is required", "code": 400}) + return + } + + messages := append([]map[string]interface{}{}, prompts...) + history, err := neo.Conversation.GetHistory(sid) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + } + + messages = append(messages, history...) + messages = append(messages, map[string]interface{}{"role": "user", "content": content, "user": sid}) + + // reply the content + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + err = neo.Answer(ctx, c, messages) + if err != nil { + c.JSON(500, gin.H{"message": err.Error(), "code": 500}) + c.Done() + } + + }) + + return nil +} + +// Answer the message +func (neo *Neo) Answer(ctx context.Context, c *gin.Context, messages []map[string]interface{}) error { + + chanStream := make(chan []byte, 1) + chanError := make(chan error, 1) + + go func() { + defer func() { + close(chanStream) + close(chanError) + }() + + _, ex := neo.AI.ChatCompletions(messages, neo.Option, func(data []byte) int { + chanStream <- data + return 1 + }) + + if ex != nil { + chanError <- fmt.Errorf("AI chat error: %s", ex.Message) + } + }() + + c.Header("Content-Type", "text/event-stream;charset=utf-8") + ok := c.Stream(func(w io.Writer) bool { + select { + case err := <-chanError: + if err != nil { + c.JSON(http.StatusInternalServerError, err.Error()) + } + return false + + case msg := <-chanStream: + msg = append(msg, []byte("\n")...) + w.Write(msg) + return true + + case <-ctx.Done(): + return false + } + }) + + if !ok { + c.Status(500) + return nil + } + + c.Status(200) + return nil +} + +func (neo *Neo) crossDomain(router *gin.Engine, path string, allows ...string) { + + if len(allows) == 0 { + return + } + + allowsMap := map[string]bool{} + for _, allow := range allows { + allowsMap[allow] = true + } + + router.Use(func(c *gin.Context) { + referer := c.Request.Referer() + if referer != "" { + + if !api.IsAllowed(c, allowsMap) { + c.AbortWithStatus(403) + return + } + + url, _ := url.Parse(referer) + referer = fmt.Sprintf("%s://%s", url.Scheme, url.Host) + c.Writer.Header().Set("Access-Control-Allow-Origin", referer) + c.Writer.Header().Set("Access-Control-Allow-Credentials", "true") + c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With") + c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT") + c.AbortWithStatus(204) + } + }) + + router.OPTIONS(path, func(c *gin.Context) { c.Status(200) }) +} + +func (neo *Neo) setGuard(router *gin.Engine) error { + + if neo.Guard == "" { + router.Use(func(c *gin.Context) { + token := c.Query("token") + if token == "" { + c.JSON(403, gin.H{"message": "token is required", "code": 403}) + c.Abort() + return + } + + user := helper.JwtValidate(token) + c.Set("__sid", user.SID) + c.Next() + }) + return nil + } + + // validate the custom guard + _, err := process.Of(neo.Guard) + if err != nil { + return err + } + + // custom guard + router.Use(api.ProcessGuard(neo.Guard)) + return nil +} + +// NewAI create a new AI +func (neo *Neo) newAI() error { + + if neo.Connector == "" { + return fmt.Errorf("%s connector is required", neo.ID) + } + + conn, err := connector.Select(neo.Connector) + if err != nil { + return err + } + + if conn.Is(connector.OPENAI) { + ai, err := openai.New(neo.Connector) + if err != nil { + return err + } + neo.AI = ai + } + + return fmt.Errorf("%s connector %s not support, should be a openai", neo.ID, neo.Connector) +} + +// newConversation create a new conversation +func (neo *Neo) newConversation() error { + + if neo.ConversationSetting.Connector == "default" || neo.ConversationSetting.Connector == "" { + neo.Conversation = conversation.NewXun() + return nil + } + + conn, err := connector.Select(neo.ConversationSetting.Connector) + if err != nil { + return err + } + + if conn.Is(connector.DATABASE) { + neo.Conversation = conversation.NewXun() + return nil + + } else if conn.Is(connector.REDIS) { + neo.Conversation = conversation.NewRedis() + return nil + + } else if conn.Is(connector.MONGO) { + neo.Conversation = conversation.NewMongo() + return nil + + } else if conn.Is(connector.WEAVIATE) { + neo.Conversation = conversation.NewWeaviate() + return nil + } + + return fmt.Errorf("%s conversation connector %s not support", neo.ID, neo.ConversationSetting.Connector) +} diff --git a/neo/types.go b/neo/types.go new file mode 100644 index 00000000..b3f837ef --- /dev/null +++ b/neo/types.go @@ -0,0 +1,36 @@ +package neo + +import "github.com/yaoapp/yao/aigc" + +// Neo AI assistant +type Neo struct { + ID string `json:"-"` + 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"` + Allows []string `json:"allows,omitempty"` + AI aigc.AI `json:"-"` + Conversation Conversation `json:"-"` + Command Command `json:"-"` +} + +// ConversationSetting the conversation config +type ConversationSetting struct { + Connector string `json:"connector,omitempty"` + Table string `json:"table,omitempty"` + MaxSize int `json:"max_size,omitempty"` +} + +// Conversation the store interface +type Conversation interface { + GetHistory(sid string) ([]map[string]interface{}, error) + SaveHistory(sid string, messages []map[string]interface{}) error +} + +// Command the command interface +type Command interface { + Match(messages []map[string]interface{}) (bool, error) +} diff --git a/openai/openai.go b/openai/openai.go index c6a3b860..a1bb7edf 100644 --- a/openai/openai.go +++ b/openai/openai.go @@ -1,6 +1,7 @@ package openai import ( + "context" "encoding/base64" "fmt" @@ -58,7 +59,24 @@ func (openai OpenAI) Completions(prompt interface{}, option map[string]interface if cb != nil { option["stream"] = true - return nil, openai.stream("/v1/completions", option, cb) + return nil, openai.stream(context.Background(), "/v1/completions", option, cb) + } + + option["stream"] = false + return openai.post("/v1/completions", option) +} + +// CompletionsWith Creates a completion for the provided prompt and parameters. +// https://platform.openai.com/docs/api-reference/completions/create +func (openai OpenAI) CompletionsWith(ctx context.Context, prompt interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) { + if option == nil { + option = map[string]interface{}{} + } + option["prompt"] = prompt + + if cb != nil { + option["stream"] = true + return nil, openai.stream(ctx, "/v1/completions", option, cb) } option["stream"] = false @@ -75,7 +93,24 @@ func (openai OpenAI) ChatCompletions(messages []map[string]interface{}, option m if cb != nil { option["stream"] = true - return nil, openai.stream("/v1/chat/completions", option, cb) + return nil, openai.stream(context.Background(), "/v1/chat/completions", option, cb) + } + + option["stream"] = false + return openai.post("/v1/chat/completions", option) +} + +// ChatCompletionsWith Creates a model response for the given chat conversation. +// https://platform.openai.com/docs/api-reference/chat/create +func (openai OpenAI) ChatCompletionsWith(ctx context.Context, messages []map[string]interface{}, option map[string]interface{}, cb func(data []byte) int) (interface{}, *exception.Exception) { + if option == nil { + option = map[string]interface{}{} + } + option["messages"] = messages + + if cb != nil { + option["stream"] = true + return nil, openai.stream(ctx, "/v1/chat/completions", option, cb) } option["stream"] = false @@ -304,7 +339,7 @@ func (openai OpenAI) postFileWithoutModel(path string, files map[string][]byte, } // stream post request -func (openai OpenAI) stream(path string, payload map[string]interface{}, cb func(data []byte) int) *exception.Exception { +func (openai OpenAI) stream(ctx context.Context, path string, payload map[string]interface{}, cb func(data []byte) int) *exception.Exception { url := fmt.Sprintf("%s%s", openai.host, path) key := fmt.Sprintf("Bearer %s", openai.key) payload["model"] = openai.model @@ -314,7 +349,7 @@ func (openai OpenAI) stream(path string, payload map[string]interface{}, cb func "Content-Type": {"application/json; charset=utf-8"}, "Authorization": {key}, }). - Stream("POST", payload, cb) + Stream(ctx, "POST", payload, cb) if err != nil { return exception.New(err.Error(), 500) diff --git a/openai/openai_test.go b/openai/openai_test.go index 0189c96d..7437ab0c 100644 --- a/openai/openai_test.go +++ b/openai/openai_test.go @@ -1,8 +1,10 @@ package openai import ( + "context" "encoding/base64" "testing" + "time" "github.com/google/uuid" "github.com/stretchr/testify/assert" @@ -52,6 +54,35 @@ func TestCompletions(t *testing.T) { assert.NotEmpty(t, res) } +func TestCompletionsWith(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + openai := prepare(t, "text-davinci-003") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + time.Sleep(200 * time.Millisecond) + cancel() + }() + + res := []byte{} + _, err := openai.CompletionsWith(ctx, "Write an article about internet ", nil, func(data []byte) int { + res = append(res, data...) + if len(data) == 0 { + res = append(res, []byte("\n")...) + } + + if string(data) == "data: [DONE]" { + return 0 + } + + return 1 + }) + + assert.Contains(t, err.Message, "context canceled") +} + func TestChatCompletions(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean() @@ -92,6 +123,35 @@ func TestChatCompletions(t *testing.T) { assert.NotEmpty(t, res) } +func TestChatCompletionsWith(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + openai := prepare(t, "gpt-3_5-turbo") + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { + time.Sleep(200 * time.Millisecond) + cancel() + }() + + res := []byte{} + _, err := openai.ChatCompletionsWith(ctx, []map[string]interface{}{{"role": "user", "content": "Write an article about internet"}}, nil, func(data []byte) int { + res = append(res, data...) + if len(data) == 0 { + res = append(res, []byte("\n")...) + } + + if string(data) == "data: [DONE]" { + return 0 + } + + return 1 + }) + + assert.Contains(t, err.Message, "context canceled") +} + func TestEdits(t *testing.T) { test.Prepare(t, config.Conf) defer test.Clean()