[add] neo (50%)

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/...)

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) {

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)
 }

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)

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
+}

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
+}

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
+}

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
+}

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() {}

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)
+}

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)
+}

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)

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()