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

Refactor conversation history handling and enhance Assistant structure

Browse source 
- Updated SaveHistory method across Mongo, Redis, and Weaviate to include a context parameter, improving flexibility in message storage.
- Modified the Assistant struct in types.go to add new fields: Type, Avatar, and Flows, enhancing the representation of assistant attributes.
- Adjusted the Xun implementation to support the new context handling in SaveHistory, ensuring messages can now include contextual information.
- Updated tests to reflect changes in message structure and ensure proper functionality with the new context parameter.
This commit is contained in:
Max 2024-12-20 12:30:31 +08:00
parent 33b28b50f4
commit af52dd35a1
8 changed files with 140 additions and 153 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
neo/assistant/types.go
View file

@ -30,13 +30,16 @@ type QueryParam struct {
// Assistant the assistant
type Assistant struct {
ID string `json:"assistant_id"` // Assistant ID
Name string `json:"name,omitempty"` // Assistant Name
Connector string `json:"connector"` // AI Connector
Description string `json:"description,omitempty"` // Assistant Description
Option map[string]interface{} `json:"option,omitempty"` // AI Option
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
API API `json:"-" yaml:"-"` // Assistant API
ID string `json:"assistant_id"` // Assistant ID
Type string `json:"type,omitempty"` // Assistant Type, default is assistant
Name string `json:"name,omitempty"` // Assistant Name
Avatar string `json:"avatar,omitempty"` // Assistant Avatar
Connector string `json:"connector"` // AI Connector
Description string `json:"description,omitempty"` // Assistant Description
Option map[string]interface{} `json:"option,omitempty"` // AI Option
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
Flows []map[string]interface{} `json:"flows,omitempty"` // Assistant Flows
API API `json:"-" yaml:"-"` // Assistant API
}
// File the file

2
neo/conversation/mongo.go
View file

@ -30,7 +30,7 @@ func (conv *Mongo) GetHistory(sid string, cid string) ([]map[string]interface{},
}
// SaveHistory save the history
func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
func (conv *Mongo) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}

2
neo/conversation/redis.go
View file

@ -30,7 +30,7 @@ func (conv *Redis) GetHistory(sid string, cid string) ([]map[string]interface{},
}
// SaveHistory save the history
func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
func (conv *Redis) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}

4
neo/conversation/types.go
View file

@ -43,9 +43,7 @@ type Conversation interface {
GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error)
GetChat(sid string, cid string) (*ChatInfo, error)
GetHistory(sid string, cid string) ([]map[string]interface{}, error)
SaveHistory(sid string, messages []map[string]interface{}, cid string) error
GetRequest(sid string, rid string) ([]map[string]interface{}, error)
SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error
SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error
DeleteChat(sid string, cid string) error
DeleteAllChats(sid string) error
UpdateChatTitle(sid string, cid string, title string) error

2
neo/conversation/weaviate.go
View file

@ -30,7 +30,7 @@ func (conv *Weaviate) GetHistory(sid string, cid string) ([]map[string]interface
}
// SaveHistory save the history
func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
func (conv *Weaviate) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
return nil
}

217
neo/conversation/xun.go
View file

@ -7,6 +7,7 @@ import (
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/session"
"github.com/yaoapp/kun/log"
@ -23,13 +24,16 @@ type Xun struct {
}
type row struct {
Role string `json:"role"`
Name string `json:"name"` // User name
Content string `json:"content"`
Sid string `json:"sid"`
Rid string `json:"rid"`
Cid string `json:"cid"` // Chat ID from chat history
ExpiredAt interface{} `json:"expired_at"`
Role string `json:"role"` // Message role
Name string `json:"name"` // User name
Content string `json:"content"` // Message content
Sid string `json:"sid"` // Session ID
Cid string `json:"cid"` // Chat ID from chat history
UID string `json:"uid"` // User ID
Context map[string]interface{} `json:"context"` // Message context
CreatedAt time.Time `json:"created_at"` // Created time
UpdatedAt *time.Time `json:"updated_at"` // Updated time
ExpiredAt interface{} `json:"expired_at"` // Expired time
}
// Public interface methods and constructor remain exported:
@ -111,6 +115,11 @@ func (conv *Xun) initialize() error {
return err
}
// Initialize assistant table
if err := conv.initAssistantTable(); err != nil {
return err
}
return nil
}
@ -126,11 +135,12 @@ func (conv *Xun) initHistoryTable() error {
err = conv.schema.CreateTable(historyTable, func(table schema.Blueprint) {
table.ID("id")
table.String("sid", 255).Index()
table.String("rid", 255).Null().Index()
table.String("cid", 200).Null().Index()
table.String("uid", 255).Null().Index()
table.String("role", 200).Null().Index()
table.String("name", 200).Null().Index()
table.Text("content").Null()
table.JSON("context").Null()
table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index()
table.TimestampTz("updated_at").Null().Index()
table.TimestampTz("expired_at").Null().Index()
@ -148,7 +158,7 @@ func (conv *Xun) initHistoryTable() error {
return err
}
fields := []string{"id", "sid", "rid", "cid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
@ -198,6 +208,52 @@ func (conv *Xun) initChatTable() error {
return nil
}
func (conv *Xun) initAssistantTable() error {
assistantTable := conv.getAssistantTable()
has, err := conv.schema.HasTable(assistantTable)
if err != nil {
return err
}
// Create the assistant table
if !has {
err = conv.schema.CreateTable(assistantTable, func(table schema.Blueprint) {
table.ID("id")
table.String("assistant_id", 200).Unique().Index()
table.String("type", 200).SetDefault("assistant").Index() // default is assistant
table.String("name", 200).Null()
table.String("avatar", 200).Null()
table.String("connector", 200).NotNull()
table.Text("description").Null()
table.JSON("option").Null()
table.JSON("prompts").Null()
table.JSON("flows").Null()
table.TimestampTz("created_at").SetDefaultRaw("NOW()").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the assistant table: %s", assistantTable)
}
// Validate the table
tab, err := conv.schema.GetTable(assistantTable)
if err != nil {
return err
}
fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "option", "prompts", "flows", "created_at", "updated_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
}
}
return nil
}
func (conv *Xun) getUserID(sid string) (string, error) {
field := "user_id"
if conv.setting.UserField != "" {
@ -217,13 +273,17 @@ func (conv *Xun) getUserID(sid string) (string, error) {
}
func (conv *Xun) getHistoryTable() string {
return conv.setting.Table
return conv.setting.Table + "_history"
}
func (conv *Xun) getChatTable() string {
return conv.setting.Table + "_chat"
}
func (conv *Xun) getAssistantTable() string {
return conv.setting.Table + "_assistant"
}
// UpdateChatTitle update the chat title
func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error {
userID, err := conv.getUserID(sid)
@ -380,7 +440,7 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e
}
qb := conv.newQuery().
Select("role", "name", "content").
Select("role", "name", "content", "context", "uid", "created_at", "updated_at").
Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
@ -401,18 +461,23 @@ func (conv *Xun) GetHistory(sid string, cid string) ([]map[string]interface{}, e
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...)
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"uid": row.Get("uid"),
"created_at": row.Get("created_at"),
"updated_at": row.Get("updated_at"),
}
res = append([]map[string]interface{}{message}, res...)
}
return res, nil
}
// SaveHistory save the history
func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string) error {
func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
if cid == "" {
cid = uuid.New().String() // Generate a new UUID if cid is empty
@ -450,24 +515,49 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid
// Save message history
defer conv.clean()
var expiredAt interface{} = nil
values := []row{}
values := []map[string]interface{}{}
if conv.setting.TTL > 0 {
expiredAt = time.Now().Add(time.Duration(conv.setting.TTL) * time.Second)
}
now := time.Now()
for _, message := range messages {
value := row{
Role: message["role"].(string),
Name: "",
Content: message["content"].(string),
Sid: userID,
Cid: cid,
ExpiredAt: expiredAt,
// Type assertion safety checks
role, ok := message["role"].(string)
if !ok {
return fmt.Errorf("invalid role type in message: %v", message["role"])
}
if message["name"] != nil {
value.Name = message["name"].(string)
content, ok := message["content"].(string)
if !ok {
return fmt.Errorf("invalid content type in message: %v", message["content"])
}
var contextRaw interface{} = nil
if context != nil {
contextRaw, err = jsoniter.MarshalToString(context)
if err != nil {
return err
}
}
value := map[string]interface{}{
"role": role,
"name": "",
"content": content,
"sid": userID,
"cid": cid,
"uid": userID,
"context": contextRaw,
"created_at": now,
"updated_at": nil,
"expired_at": expiredAt,
}
if name, ok := message["name"].(string); ok {
value["name"] = name
}
values = append(values, value)
}
@ -479,79 +569,6 @@ func (conv *Xun) SaveHistory(sid string, messages []map[string]interface{}, cid
return nil
}
// GetRequest get the request history
func (conv *Xun) GetRequest(sid string, rid string) ([]map[string]interface{}, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
qb := conv.newQuery().
Select("role", "name", "content", "sid").
Where("rid", rid).
Where("sid", userID).
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
}
// SaveRequest save the request history
func (conv *Xun) SaveRequest(sid string, rid string, cid string, messages []map[string]interface{}) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
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: userID,
Cid: cid,
Rid: rid,
ExpiredAt: expiredAt,
}
if message["name"] != nil {
value.Name = message["name"].(string)
}
values = append(values, value)
}
return conv.newQuery().Insert(values)
}
// GetChat get the chat info and its history
func (conv *Xun) GetChat(sid string, cid string) (*ChatInfo, error) {
userID, err := conv.getUserID(sid)

48
neo/conversation/xun_test.go
View file

@ -45,7 +45,7 @@ func TestNewXunDefault(t *testing.T) {
t.Fatal(err)
}
fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
assert.Equal(t, true, tab.HasColumn(field))
}
@ -103,7 +103,7 @@ func TestNewXunConnector(t *testing.T) {
t.Fatal(err)
}
fields := []string{"id", "sid", "cid", "rid", "role", "name", "content", "created_at", "updated_at", "expired_at"}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
assert.Equal(t, true, tab.HasColumn(field))
}
@ -143,7 +143,7 @@ func TestXunSaveAndGetHistory(t *testing.T) {
err = conv.SaveHistory("123456", []map[string]interface{}{
{"role": "user", "name": "user1", "content": "hello"},
{"role": "assistant", "name": "user1", "content": "Hello there, how"},
}, cid)
}, cid, nil)
assert.Nil(t, err)
// get the history
@ -154,38 +154,6 @@ func TestXunSaveAndGetHistory(t *testing.T) {
assert.Equal(t, 2, len(data))
}
func TestXunSaveAndGetRequest(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.SaveRequest("123456", "912836", "test.command", []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.GetRequest("123456", "912836")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, 2, len(data))
}
func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
@ -209,7 +177,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
{"role": "user", "name": "user1", "content": "hello"},
{"role": "assistant", "name": "assistant1", "content": "Hi! How can I help you?"},
}
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
// get the history for specific cid
@ -224,7 +192,7 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
moreMessages := []map[string]interface{}{
{"role": "user", "name": "user1", "content": "another message"},
}
err = conv.SaveHistory(sid, moreMessages, anotherCID)
err = conv.SaveHistory(sid, moreMessages, anotherCID, nil)
assert.Nil(t, err)
// get history for the first cid - should still be 2 messages
@ -294,7 +262,7 @@ func TestXunGetChats(t *testing.T) {
}
// Then save the history
err = conv.SaveHistory(sid, messages, chatID)
err = conv.SaveHistory(sid, messages, chatID, nil)
if err != nil {
t.Fatal(err)
}
@ -344,7 +312,7 @@ func TestXunDeleteChat(t *testing.T) {
}
// Save the chat and history
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
// Verify chat exists
@ -385,7 +353,7 @@ func TestXunDeleteAllChats(t *testing.T) {
// Save multiple chats
for i := 0; i < 3; i++ {
cid := fmt.Sprintf("test_chat_%d", i)
err = conv.SaveHistory(sid, messages, cid)
err = conv.SaveHistory(sid, messages, cid, nil)
assert.Nil(t, err)
}

1
neo/neo.go
View file

@ -498,6 +498,7 @@ func (neo *DSL) saveHistory(sid string, chatID string, content []byte, messages
{"role": "assistant", "content": string(content), "name": sid},
},
chatID,
nil,
)
if err != nil {