[add] Neo save chat history (80%)

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

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 {
+func NewXun(setting Setting) (*Xun, error) {
-	return &Xun{}
+
+	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
 }

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

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",
+		ID:      "neo",
-		Prompts:             []aigc.Prompt{},
+		Prompts: []aigc.Prompt{},
-		Option:              map[string]interface{}{},
+		Option:  map[string]interface{}{},
-		Allows:              []string{},
+		Allows:  []string{},
-		ConversationSetting: ConversationSetting{Table: "yao_neo_conversation", MaxSize: 100, Connector: "default"},
+		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 {

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
+						text := message.Choices[0].Delta.Content
-						data, _ := jsoniter.Marshal(map[string]interface{}{"text": 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()
+		neo.Conversation, err = conversation.NewXun(neo.ConversationSetting)
-		return nil
+		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()
+		neo.Conversation, err = conversation.NewXun(neo.ConversationSetting)
-		return nil
+		return err
 
 	} else if conn.Is(connector.REDIS) {
 		neo.Conversation = conversation.NewRedis()

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

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)