opencode/yao
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[add] Neo save chat history (80%)

Browse source
This commit is contained in:
Max 2023-04-30 19:22:23 +08:00
parent 7ee7de96c9
commit 5f67f64f40
7 changed files with 375 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
neo/conversation/types.go Normal file
View file

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

169
neo/conversation/xun.go
View file

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

152
neo/conversation/xun_test.go Normal file
View file

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

18
neo/load.go
View file

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

39
neo/neo.go
View file

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

2
neo/neo_test.go
View file

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

14
neo/types.go
View file

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