Add assistant management functionality in conversation module

- Introduced SaveAssistant, DeleteAssistant, and GetAssistants methods across Mongo, Redis, and Weaviate implementations to manage assistant records.
- Enhanced the AssistantFilter and AssistantResponse types for improved pagination and filtering capabilities.
- Updated the conversation interface to include new assistant-related methods, ensuring consistent functionality across different storage backends.
- Added tests for assistant CRUD operations and pagination, ensuring robust functionality and reliability in managing assistants.
This commit is contained in:
Max 2024-12-29 10:03:38 +08:00
parent 8166503058
commit b17778afea
6 changed files with 465 additions and 36 deletions

View file

@ -1,5 +1,7 @@
package conversation
import "github.com/yaoapp/xun"
// Mongo conversation
type Mongo struct{}
@ -58,3 +60,29 @@ func (conv *Mongo) DeleteChat(sid string, cid string) error {
func (conv *Mongo) DeleteAllChats(sid string) error {
return nil
}
// SaveAssistant creates or updates an assistant
func (conv *Mongo) SaveAssistant(assistant map[string]interface{}) error {
return nil
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Mongo) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves assistants with pagination and tag filtering
func (conv *Mongo) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{
P: xun.P{
Items: []interface{}{},
Total: 0,
TotalPages: 0,
PageSize: filter.PageSize,
CurrentPage: filter.Page,
NextPage: 0,
PreviousPage: 0,
LastPage: 0,
},
}, nil
}

View file

@ -1,5 +1,7 @@
package conversation
import "github.com/yaoapp/xun"
// Redis conversation
type Redis struct{}
@ -58,3 +60,29 @@ func (conv *Redis) DeleteChat(sid string, cid string) error {
func (conv *Redis) DeleteAllChats(sid string) error {
return nil
}
// SaveAssistant creates or updates an assistant
func (conv *Redis) SaveAssistant(assistant map[string]interface{}) error {
return nil
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Redis) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves assistants with pagination and tag filtering
func (conv *Redis) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{
P: xun.P{
Items: []interface{}{},
Total: 0,
TotalPages: 0,
PageSize: filter.PageSize,
CurrentPage: filter.Page,
NextPage: 0,
PreviousPage: 0,
LastPage: 0,
},
}, nil
}

View file

@ -1,5 +1,7 @@
package conversation
import "github.com/yaoapp/xun"
// Setting the conversation config
type Setting struct {
Connector string `json:"connector,omitempty"`
@ -38,6 +40,18 @@ type ChatGroupResponse struct {
LastPage int `json:"last_page"` // 最后一页页码
}
// AssistantFilter represents the filter parameters for GetAssistants
type AssistantFilter struct {
Tags []string `json:"tags,omitempty"`
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Items per page
}
// AssistantResponse represents paginated assistant results
type AssistantResponse struct {
xun.P
}
// Conversation the store interface
type Conversation interface {
GetChats(sid string, filter ChatFilter) (*ChatGroupResponse, error)
@ -47,4 +61,7 @@ type Conversation interface {
DeleteChat(sid string, cid string) error
DeleteAllChats(sid string) error
UpdateChatTitle(sid string, cid string, title string) error
SaveAssistant(assistant map[string]interface{}) error
DeleteAssistant(assistantID string) error
GetAssistants(filter AssistantFilter) (*AssistantResponse, error)
}

View file

@ -1,5 +1,7 @@
package conversation
import "github.com/yaoapp/xun"
// Weaviate Database conversation
type Weaviate struct{}
@ -58,3 +60,29 @@ func (conv *Weaviate) DeleteChat(sid string, cid string) error {
func (conv *Weaviate) DeleteAllChats(sid string) error {
return nil
}
// SaveAssistant creates or updates an assistant
func (conv *Weaviate) SaveAssistant(assistant map[string]interface{}) error {
return nil
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Weaviate) DeleteAssistant(assistantID string) error {
return nil
}
// GetAssistants retrieves assistants with pagination and tag filtering
func (conv *Weaviate) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
return &AssistantResponse{
P: xun.P{
Items: []interface{}{},
Total: 0,
TotalPages: 0,
PageSize: filter.PageSize,
CurrentPage: filter.Page,
NextPage: 0,
PreviousPage: 0,
LastPage: 0,
},
}, nil
}

View file

@ -217,6 +217,7 @@ func (conv *Xun) initAssistantTable() error {
table.JSON("flows").Null() // assistant flows
table.JSON("files").Null() // assistant files
table.JSON("functions").Null() // assistant functions
table.JSON("tags").Null() // assistant tags
table.Boolean("readonly").SetDefault(false).Index() // assistant readonly
table.JSON("permissions").Null() // assistant permissions
table.Boolean("automated").SetDefault(true).Index() // assistant autoable
@ -237,7 +238,7 @@ func (conv *Xun) initAssistantTable() error {
return err
}
fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "mentionable", "created_at", "updated_at"}
fields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "mentionable", "created_at", "updated_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
@ -648,3 +649,83 @@ func (conv *Xun) DeleteAllChats(sid string) error {
Delete()
return err
}
// SaveAssistant creates or updates an assistant
func (conv *Xun) SaveAssistant(assistant map[string]interface{}) error {
assistantID, ok := assistant["assistant_id"].(string)
if !ok || assistantID == "" {
assistantID = uuid.New().String()
assistant["assistant_id"] = assistantID
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Exists()
if err != nil {
return err
}
now := time.Now()
assistant["updated_at"] = now
if exists {
// Update existing assistant
_, err = conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Update(assistant)
} else {
// Create new assistant
assistant["created_at"] = now
err = conv.query.New().
Table(conv.getAssistantTable()).
Insert(assistant)
}
return err
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Xun) DeleteAssistant(assistantID string) error {
_, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Delete()
return err
}
// GetAssistants retrieves assistants with pagination and tag filtering
func (conv *Xun) GetAssistants(filter AssistantFilter) (*AssistantResponse, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
// Apply tag filter if provided
if filter.Tags != nil && len(filter.Tags) > 0 {
for i, tag := range filter.Tags {
if i == 0 {
qb.Where("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag))
} else {
qb.OrWhere("tags", "like", fmt.Sprintf("%%\"%s\"%%", tag))
}
}
}
// Set defaults for pagination
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.Page <= 0 {
filter.Page = 1
}
// Get paginated results
paginator, err := qb.OrderBy("created_at", "desc").
Paginate(filter.PageSize, filter.Page)
if err != nil {
return nil, err
}
return &AssistantResponse{P: paginator}, nil
}

View file

@ -5,8 +5,10 @@ import (
"testing"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/stretchr/testify/assert"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/xun"
"github.com/yaoapp/xun/capsule"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
@ -15,9 +17,21 @@ import (
func TestNewXunDefault(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
@ -32,15 +46,29 @@ func TestNewXunDefault(t *testing.T) {
return
}
has, err := capsule.Schema().HasTable("__unit_test_conversation")
// Check history table
has, err := capsule.Schema().HasTable("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the table
tab, err := conv.schema.GetTable(conv.setting.Table)
// Check chat table
has, err = capsule.Schema().HasTable("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// Check assistant table
has, err = capsule.Schema().HasTable("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the history table
tab, err := conv.schema.GetTable(conv.getHistoryTable())
if err != nil {
t.Fatal(err)
}
@ -50,17 +78,27 @@ func TestNewXunDefault(t *testing.T) {
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")
// validate the chat table
tab, err = conv.schema.GetTable(conv.getChatTable())
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"}
for _, field := range chatFields {
assert.Equal(t, true, tab.HasColumn(field))
}
// validate the assistant table
tab, err = conv.schema.GetTable(conv.getAssistantTable())
if err != nil {
t.Fatal(err)
}
assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"}
for _, field := range assistantFields {
assert.Equal(t, true, tab.HasColumn(field))
}
}
func TestNewXunConnector(t *testing.T) {
@ -77,9 +115,14 @@ func TestNewXunConnector(t *testing.T) {
t.Fatal(err)
}
defer sch.DropTableIfExists("__unit_test_conversation")
defer sch.DropTableIfExists("__unit_test_conversation_history")
defer sch.DropTableIfExists("__unit_test_conversation_chat")
defer sch.DropTableIfExists("__unit_test_conversation_assistant")
sch.DropTableIfExists("__unit_test_conversation_history")
sch.DropTableIfExists("__unit_test_conversation_chat")
sch.DropTableIfExists("__unit_test_conversation_assistant")
sch.DropTableIfExists("__unit_test_conversation")
conv, err := NewXun(Setting{
Connector: "mysql",
Table: "__unit_test_conversation",
@ -90,15 +133,29 @@ func TestNewXunConnector(t *testing.T) {
return
}
has, err := sch.HasTable("__unit_test_conversation")
// Check history table
has, err := sch.HasTable("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the table
tab, err := conv.schema.GetTable(conv.setting.Table)
// Check chat table
has, err = sch.HasTable("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// Check assistant table
has, err = sch.HasTable("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
// validate the history table
tab, err := conv.schema.GetTable(conv.getHistoryTable())
if err != nil {
t.Fatal(err)
}
@ -108,26 +165,41 @@ func TestNewXunConnector(t *testing.T) {
assert.Equal(t, true, tab.HasColumn(field))
}
conv, err = NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
has, err = sch.HasTable("__unit_test_conversation")
// validate the chat table
tab, err = conv.schema.GetTable(conv.getChatTable())
if err != nil {
t.Fatal(err)
}
assert.Equal(t, true, has)
chatFields := []string{"id", "chat_id", "title", "sid", "created_at", "updated_at"}
for _, field := range chatFields {
assert.Equal(t, true, tab.HasColumn(field))
}
// validate the assistant table
tab, err = conv.schema.GetTable(conv.getAssistantTable())
if err != nil {
t.Fatal(err)
}
assistantFields := []string{"id", "assistant_id", "type", "name", "avatar", "connector", "description", "options", "prompts", "flows", "files", "functions", "tags", "readonly", "permissions", "automated", "mentionable", "created_at", "updated_at"}
for _, field := range assistantFields {
assert.Equal(t, true, tab.HasColumn(field))
}
}
func TestXunSaveAndGetHistory(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
@ -157,9 +229,15 @@ func TestXunSaveAndGetHistory(t *testing.T) {
func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
err = capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
if err != nil {
t.Fatal(err)
}
@ -220,11 +298,11 @@ func TestXunSaveAndGetHistoryWithCID(t *testing.T) {
func TestXunGetChats(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
// Drop both tables before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation")
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
if err != nil {
t.Fatal(err)
}
@ -293,7 +371,7 @@ func TestXunGetChats(t *testing.T) {
func TestXunDeleteChat(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
conv, err := NewXun(Setting{
@ -333,7 +411,7 @@ func TestXunDeleteChat(t *testing.T) {
func TestXunDeleteAllChats(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_chat")
conv, err := NewXun(Setting{
@ -371,3 +449,172 @@ func TestXunDeleteAllChats(t *testing.T) {
assert.Nil(t, err)
assert.Equal(t, int64(0), response.Total)
}
func TestXunAssistantCRUD(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
// Drop assistant table before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
if err != nil {
t.Fatal(err)
}
// Test creating a new assistant
tagsJSON, err := jsoniter.MarshalToString([]string{"tag1", "tag2", "tag3"})
if err != nil {
t.Fatal(err)
}
optionsJSON, err := jsoniter.MarshalToString(map[string]interface{}{
"model": "gpt-4",
})
if err != nil {
t.Fatal(err)
}
assistant := map[string]interface{}{
"name": "Test Assistant",
"type": "assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"description": "Test Description",
"tags": tagsJSON,
"options": optionsJSON,
}
// Test SaveAssistant (Create)
err = conv.SaveAssistant(assistant)
assert.Nil(t, err)
assistantID := assistant["assistant_id"].(string)
assert.NotEmpty(t, assistantID)
// Test GetAssistants with no filter
resp, err := conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.P.Items))
// Test GetAssistants with tag filter (single tag)
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag1"},
})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.P.Items))
// Test GetAssistants with tag filter (multiple tags)
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag1", "tag4"},
})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.P.Items))
// Test GetAssistants with non-existent tag
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"nonexistent"},
})
assert.Nil(t, err)
assert.Equal(t, 0, len(resp.P.Items))
// Test SaveAssistant (Update)
assistant["name"] = "Updated Assistant"
err = conv.SaveAssistant(assistant)
assert.Nil(t, err)
resp, err = conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
assert.Equal(t, 1, len(resp.P.Items))
item := resp.P.Items[0].(xun.R)
assert.Equal(t, "Updated Assistant", item["name"])
// Test DeleteAssistant
err = conv.DeleteAssistant(assistantID)
assert.Nil(t, err)
resp, err = conv.GetAssistants(AssistantFilter{})
assert.Nil(t, err)
assert.Equal(t, 0, len(resp.P.Items))
}
func TestXunAssistantPagination(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_history")
defer capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
// Drop assistant table before test
err := capsule.Schema().DropTableIfExists("__unit_test_conversation_assistant")
if err != nil {
t.Fatal(err)
}
conv, err := NewXun(Setting{
Connector: "default",
Table: "__unit_test_conversation",
})
if err != nil {
t.Fatal(err)
}
// Create multiple assistants for pagination testing
for i := 0; i < 25; i++ {
tagsJSON, err := jsoniter.MarshalToString([]string{fmt.Sprintf("tag%d", i%5)})
if err != nil {
t.Fatal(err)
}
assistant := map[string]interface{}{
"name": fmt.Sprintf("Assistant %d", i),
"type": "assistant",
"connector": "openai",
"description": fmt.Sprintf("Description %d", i),
"tags": tagsJSON,
}
err = conv.SaveAssistant(assistant)
assert.Nil(t, err)
}
// Test first page
resp, err := conv.GetAssistants(AssistantFilter{
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 10, len(resp.P.Items))
assert.Equal(t, 25, resp.P.Total)
assert.Equal(t, 3, resp.P.LastPage)
// Test second page
resp, err = conv.GetAssistants(AssistantFilter{
Page: 2,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 10, len(resp.P.Items))
// Test last page
resp, err = conv.GetAssistants(AssistantFilter{
Page: 3,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 5, len(resp.P.Items))
// Test filtering with tags
resp, err = conv.GetAssistants(AssistantFilter{
Tags: []string{"tag0"},
Page: 1,
PageSize: 10,
})
assert.Nil(t, err)
assert.Equal(t, 5, len(resp.P.Items))
}