diff --git a/neo/conversation/types.go b/neo/conversation/types.go new file mode 100644 index 00000000..a6a5ca70 --- /dev/null +++ b/neo/conversation/types.go @@ -0,0 +1,9 @@ +package conversation + +// Setting the conversation config +type Setting struct { + Connector string `json:"connector,omitempty"` + Table string `json:"table,omitempty"` + MaxSize int `json:"max_size,omitempty" yaml:"max_size,omitempty"` + TTL int `json:"ttl,omitempty" yaml:"ttl,omitempty"` +} diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 0f0baf30..cc15bce3 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -1,19 +1,180 @@ package conversation +import ( + "fmt" + "time" + + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/kun/log" + "github.com/yaoapp/xun/capsule" + "github.com/yaoapp/xun/dbal/query" + "github.com/yaoapp/xun/dbal/schema" +) + // Xun Database conversation -type Xun struct{} +type Xun struct { + query query.Query + schema schema.Schema + setting Setting +} + +type row struct { + Role string `json:"role"` + Name string `json:"name"` + Content string `json:"content"` + Sid string `json:"sid"` + ExpiredAt interface{} `json:"expired_at"` +} // NewXun create a new conversation -func NewXun() *Xun { - return &Xun{} +func NewXun(setting Setting) (*Xun, error) { + + conv := &Xun{setting: setting} + if setting.Connector == "default" { + conv.query = capsule.Global.Query() + conv.schema = capsule.Global.Schema() + + } else { + + conn, err := connector.Select(setting.Connector) + if err != nil { + return nil, err + } + + conv.query, err = conn.Query() + if err != nil { + return nil, err + } + + conv.schema, err = conn.Schema() + if err != nil { + return nil, err + } + } + + err := conv.Init() + if err != nil { + return nil, err + } + + return conv, nil } // GetHistory get the history func (conv *Xun) GetHistory(sid string) ([]map[string]interface{}, error) { - return []map[string]interface{}{}, nil + + qb := conv.query.Table(conv.setting.Table). + Select("role", "name", "content"). + Where("sid", sid). + OrderBy("id", "desc") + + if conv.setting.TTL > 0 { + qb.Where("expired_at", ">", time.Now()) + } + + limit := 20 + if conv.setting.MaxSize > 0 { + limit = conv.setting.MaxSize + } + + rows, err := qb.Limit(limit).Get() + if err != nil { + return nil, err + } + + res := []map[string]interface{}{} + for _, row := range rows { + res = append([]map[string]interface{}{{ + "role": row.Get("role"), + "name": row.Get("name"), + "content": row.Get("content"), + }}, res...) + } + + return res, nil } // SaveHistory save the history func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}) error { + + defer conv.clean() + var expiredAt interface{} = nil + values := []row{} + if conv.setting.TTL > 0 { + expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second) + } + + for _, message := range messages { + value := row{ + Role: message["role"].(string), + Name: "", + Content: message["content"].(string), + Sid: sid, + ExpiredAt: expiredAt, + } + + if message["name"] != nil { + value.Name = message["name"].(string) + } + values = append(values, value) + } + + return conv.query.Table(conv.setting.Table).Insert(values) +} + +func (conv *Xun) clean() { + nums, err := conv.query.Table(conv.setting.Table).Where("expired_at", "<=", time.Now()).Delete() + if err != nil { + log.Error("Clean the conversation table error: %s", err.Error()) + return + } + + if nums > 0 { + log.Trace("Clean the conversation table: %s %d", conv.setting.Table, nums) + } +} + +// Init init the conversation +func (conv *Xun) Init() error { + + has, err := conv.schema.HasTable(conv.setting.Table) + if err != nil { + return err + } + + // create the table + if !has { + err = conv.schema.CreateTable(conv.setting.Table, func(table schema.Blueprint) { + + table.ID("id") // The ID field + table.String("sid", 255).Index() + table.String("role", 200).Null().Index() + table.String("name", 200).Null().Index() + table.Text("content").Null() + + table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index() + table.TimestampTz("updated_at").Null().Index() + table.TimestampTz("expired_at").Null().Index() + }) + + if err != nil { + return err + } + log.Trace("Create the conversation table: %s", conv.setting.Table) + } + + // validate the table + tab, err := conv.schema.GetTable(conv.setting.Table) + if err != nil { + return err + } + + fields := []string{"id", "sid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + for _, field := range fields { + if !tab.HasColumn(field) { + return fmt.Errorf("%s is required", field) + } + } + return nil } diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go new file mode 100644 index 00000000..1ed02bf1 --- /dev/null +++ b/neo/conversation/xun_test.go @@ -0,0 +1,152 @@ +package conversation + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/yaoapp/gou/connector" + "github.com/yaoapp/xun/capsule" + "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/test" +) + +func TestNewXunDefault(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + + err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + + if err != nil { + t.Error(err) + return + } + + has, err := capsule.Schema().HasTable("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, true, has) + + // validate the table + tab, err := conv.schema.GetTable(conv.setting.Table) + if err != nil { + t.Fatal(err) + } + + fields := []string{"id", "sid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + for _, field := range fields { + assert.Equal(t, true, tab.HasColumn(field)) + } + + conv, err = NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + + has, err = capsule.Schema().HasTable("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, true, has) +} + +func TestNewXunConnector(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + conn, err := connector.Select("mysql") + if err != nil { + t.Fatal(err) + } + + sch, err := conn.Schema() + if err != nil { + t.Fatal(err) + } + + defer sch.DropTableIfExists("__unit_test_conversation") + + sch.DropTableIfExists("__unit_test_conversation") + conv, err := NewXun(Setting{ + Connector: "mysql", + Table: "__unit_test_conversation", + }) + + if err != nil { + t.Error(err) + return + } + + has, err := sch.HasTable("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, true, has) + + // validate the table + tab, err := conv.schema.GetTable(conv.setting.Table) + if err != nil { + t.Fatal(err) + } + + fields := []string{"id", "sid", "role", "name", "content", "created_at", "updated_at", "expired_at"} + for _, field := range fields { + assert.Equal(t, true, tab.HasColumn(field)) + } + + conv, err = NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + }) + + has, err = sch.HasTable("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + assert.Equal(t, true, has) +} + +func TestXunSaveAndGetHistory(t *testing.T) { + + test.Prepare(t, config.Conf) + defer test.Clean() + defer capsule.Schema().DropTableIfExists("__unit_test_conversation") + + err := capsule.Schema().DropTableIfExists("__unit_test_conversation") + if err != nil { + t.Fatal(err) + } + + conv, err := NewXun(Setting{ + Connector: "default", + Table: "__unit_test_conversation", + TTL: 3600, + }) + + // save the history + err = conv.SaveHistory("123456", []map[string]interface{}{ + {"role": "user", "name": "user1", "content": "hello"}, + {"role": "assistant", "name": "user1", "content": "Hello there, how"}, + }) + assert.Nil(t, err) + + // get the history + data, err := conv.GetHistory("123456") + if err != nil { + t.Fatal(err) + } + assert.Equal(t, 2, len(data)) +} diff --git a/neo/load.go b/neo/load.go index b81138f4..7f9679d5 100644 --- a/neo/load.go +++ b/neo/load.go @@ -6,6 +6,7 @@ import ( "github.com/yaoapp/gou/application" "github.com/yaoapp/yao/aigc" "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/neo/conversation" ) // Neo the neo AI assistant @@ -15,11 +16,14 @@ var Neo *DSL func Load(cfg config.Config) error { setting := DSL{ - ID: "neo", - Prompts: []aigc.Prompt{}, - Option: map[string]interface{}{}, - Allows: []string{}, - ConversationSetting: ConversationSetting{Table: "yao_neo_conversation", MaxSize: 100, Connector: "default"}, + ID: "neo", + Prompts: []aigc.Prompt{}, + Option: map[string]interface{}{}, + Allows: []string{}, + ConversationSetting: conversation.Setting{ + Table: "yao_neo_conversation", + Connector: "default", + }, } bytes, err := application.App.Read(filepath.Join("neo", "neo.yml")) @@ -32,6 +36,10 @@ func Load(cfg config.Config) error { return err } + if setting.ConversationSetting.MaxSize == 0 { + setting.ConversationSetting.MaxSize = 100 + } + Neo = &setting err = Neo.newAI() if err != nil { diff --git a/neo/neo.go b/neo/neo.go index 2f705fcd..70a390d6 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -14,6 +14,7 @@ import ( "github.com/yaoapp/gou/api" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/process" + "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/helper" "github.com/yaoapp/yao/neo/conversation" "github.com/yaoapp/yao/openai" @@ -63,6 +64,7 @@ func (neo *DSL) API(router *gin.Engine, path string) error { messages = append(messages, history...) messages = append(messages, map[string]interface{}{"role": "user", "content": content, "name": sid}) + // utils.Dump(messages) // reply the content ctx, cancel := context.WithCancel(context.Background()) @@ -101,6 +103,25 @@ func (neo *DSL) Answer(ctx context.Context, c *gin.Context, messages []map[strin } }() + // save the history + content := []byte{} + defer func() { + sid := c.GetString("__sid") + if len(content) > 0 && sid != "" && len(messages) > 0 { + err := neo.Conversation.SaveHistory( + sid, + []map[string]interface{}{ + {"role": "user", "content": messages[len(messages)-1]["content"], "name": sid}, + {"role": "assistant", "content": string(content), "name": sid}, + }, + ) + + if err != nil { + log.Error("Save history error: %s", err.Error()) + } + } + }() + c.Header("Content-Type", "text/event-stream;charset=utf-8") ok := c.Stream(func(w io.Writer) bool { select { @@ -125,8 +146,9 @@ func (neo *DSL) Answer(ctx context.Context, c *gin.Context, messages []map[strin } if len(message.Choices) > 0 { - content := message.Choices[0].Delta.Content - data, _ := jsoniter.Marshal(map[string]interface{}{"text": content}) + text := message.Choices[0].Delta.Content + content = append(content, []byte(text)...) + data, _ := jsoniter.Marshal(map[string]interface{}{"text": text}) w.Write([]byte(fmt.Sprintf("data: %s\n\n", data))) return true } @@ -136,9 +158,6 @@ func (neo *DSL) Answer(ctx context.Context, c *gin.Context, messages []map[strin return true } } - - // msg = append(msg, []byte("\n")...) - // w.Write(msg) return true } }) @@ -244,19 +263,21 @@ func (neo *DSL) newAI() error { // newConversation create a new conversation func (neo *DSL) newConversation() error { + var err error if neo.ConversationSetting.Connector == "default" || neo.ConversationSetting.Connector == "" { - neo.Conversation = conversation.NewXun() - return nil + neo.Conversation, err = conversation.NewXun(neo.ConversationSetting) + return err } + // other connector conn, err := connector.Select(neo.ConversationSetting.Connector) if err != nil { return err } if conn.Is(connector.DATABASE) { - neo.Conversation = conversation.NewXun() - return nil + neo.Conversation, err = conversation.NewXun(neo.ConversationSetting) + return err } else if conn.Is(connector.REDIS) { neo.Conversation = conversation.NewRedis() diff --git a/neo/neo_test.go b/neo/neo_test.go index 258304f0..3daffa65 100644 --- a/neo/neo_test.go +++ b/neo/neo_test.go @@ -49,7 +49,7 @@ func TestAPI(t *testing.T) { return 1 }) - assert.Contains(t, string(res), "[DONE]") + assert.Contains(t, string(res), `{"done":true}`) } func TestAPIAuth(t *testing.T) { diff --git a/neo/types.go b/neo/types.go index d8138d0e..f67e6a6a 100644 --- a/neo/types.go +++ b/neo/types.go @@ -1,6 +1,9 @@ package neo -import "github.com/yaoapp/yao/aigc" +import ( + "github.com/yaoapp/yao/aigc" + "github.com/yaoapp/yao/neo/conversation" +) // DSL AI assistant type DSL struct { @@ -8,7 +11,7 @@ type DSL struct { Name string `json:"name,omitempty"` Guard string `json:"guard,omitempty"` Connector string `json:"connector"` - ConversationSetting ConversationSetting `json:"conversation"` + ConversationSetting conversation.Setting `json:"conversation" yaml:"conversation"` Option map[string]interface{} `json:"option"` Prompts []aigc.Prompt `json:"prompts,omitempty"` Allows []string `json:"allows,omitempty"` @@ -17,13 +20,6 @@ type DSL struct { Command Command `json:"-" yaml:"-"` } -// 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)