From af52dd35a1a1b4d8a6d730f49760d82f6742bf75 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 20 Dec 2024 12:30:31 +0800 Subject: [PATCH] Refactor conversation history handling and enhance Assistant structure - 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. --- neo/assistant/types.go | 17 +-- neo/conversation/mongo.go | 2 +- neo/conversation/redis.go | 2 +- neo/conversation/types.go | 4 +- neo/conversation/weaviate.go | 2 +- neo/conversation/xun.go | 217 +++++++++++++++++++---------------- neo/conversation/xun_test.go | 48 ++------ neo/neo.go | 1 + 8 files changed, 140 insertions(+), 153 deletions(-) diff --git a/neo/assistant/types.go b/neo/assistant/types.go index 03930582..35ad239d 100644 --- a/neo/assistant/types.go +++ b/neo/assistant/types.go @@ -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 diff --git a/neo/conversation/mongo.go b/neo/conversation/mongo.go index 98c938ca..a30461fe 100644 --- a/neo/conversation/mongo.go +++ b/neo/conversation/mongo.go @@ -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 } diff --git a/neo/conversation/redis.go b/neo/conversation/redis.go index afb1aa74..7da3f4d2 100644 --- a/neo/conversation/redis.go +++ b/neo/conversation/redis.go @@ -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 } diff --git a/neo/conversation/types.go b/neo/conversation/types.go index 205a99d6..63375e90 100644 --- a/neo/conversation/types.go +++ b/neo/conversation/types.go @@ -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 diff --git a/neo/conversation/weaviate.go b/neo/conversation/weaviate.go index 7b251851..0131251e 100644 --- a/neo/conversation/weaviate.go +++ b/neo/conversation/weaviate.go @@ -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 } diff --git a/neo/conversation/xun.go b/neo/conversation/xun.go index 61a25b55..8ab43dc7 100644 --- a/neo/conversation/xun.go +++ b/neo/conversation/xun.go @@ -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) diff --git a/neo/conversation/xun_test.go b/neo/conversation/xun_test.go index d42f32e2..34037b5c 100644 --- a/neo/conversation/xun_test.go +++ b/neo/conversation/xun_test.go @@ -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) } diff --git a/neo/neo.go b/neo/neo.go index a83ec685..0583c1d9 100644 --- a/neo/neo.go +++ b/neo/neo.go @@ -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 {