Refactor store imports and update store initialization logic

- Updated import paths for store packages to use specific types for better clarity and maintainability.
- Refactored store initialization in the agent to utilize the new store types for Mongo, Redis, and Xun.
- Removed deprecated store-related files to streamline the codebase and improve overall organization.
- Ensured that all references to the store in the agent files are consistent with the new structure.
This commit is contained in:
Max 2025-11-07 18:27:43 +08:00
parent 52f74e5edd
commit 447ab4c15b
23 changed files with 1812 additions and 1770 deletions

View file

@ -13,7 +13,7 @@ import (
"github.com/yaoapp/yao/agent/assistant"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/helper"
"github.com/yaoapp/yao/openapi/oauth"
)

View file

@ -18,7 +18,7 @@ import (
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/message"
chatMessage "github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
)
// Get get the assistant by id

View file

@ -7,7 +7,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
sui "github.com/yaoapp/yao/sui/core"
)

View file

@ -14,7 +14,7 @@ import (
"github.com/yaoapp/gou/fs"
v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
agentvision "github.com/yaoapp/yao/agent/vision"
"github.com/yaoapp/yao/openai"
"github.com/yaoapp/yao/share"

View file

@ -4,7 +4,7 @@ import (
"fmt"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
)
// Tool represents a tool

View file

@ -8,7 +8,7 @@ import (
v8 "github.com/yaoapp/gou/runtime/v8"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
api "github.com/yaoapp/yao/openai"
)

View file

@ -8,7 +8,10 @@ import (
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store"
mongoStore "github.com/yaoapp/yao/agent/store/mongo"
redisStore "github.com/yaoapp/yao/agent/store/redis"
store "github.com/yaoapp/yao/agent/store/types"
xunStore "github.com/yaoapp/yao/agent/store/xun"
"github.com/yaoapp/yao/config"
)
@ -122,7 +125,7 @@ func initStore() error {
var err error
if Agent.StoreSetting.Connector == "default" || Agent.StoreSetting.Connector == "" {
Agent.Store, err = store.NewXun(Agent.StoreSetting)
Agent.Store, err = xunStore.NewXun(Agent.StoreSetting)
return err
}
@ -133,15 +136,15 @@ func initStore() error {
}
if conn.Is(connector.DATABASE) {
Agent.Store, err = store.NewXun(Agent.StoreSetting)
Agent.Store, err = xunStore.NewXun(Agent.StoreSetting)
return err
} else if conn.Is(connector.REDIS) {
Agent.Store = store.NewRedis()
Agent.Store = redisStore.NewRedis()
return nil
} else if conn.Is(connector.MONGO) {
Agent.Store = store.NewMongo()
Agent.Store = mongoStore.NewMongo()
return nil
}

View file

@ -9,7 +9,7 @@ import (
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
)
// GetAgent returns the Agent instance

View file

@ -1,26 +1,28 @@
package store
package mongo
import "github.com/yaoapp/yao/agent/store/types"
// Mongo represents a MongoDB-based conversation storage
type Mongo struct{}
// NewMongo create a new mongo store
func NewMongo() Store {
func NewMongo() types.Store {
return &Mongo{}
}
// GetChats retrieves a list of chats
func (m *Mongo) GetChats(sid string, filter ChatFilter, locale ...string) (*ChatGroupResponse, error) {
return &ChatGroupResponse{}, nil
func (m *Mongo) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
return &types.ChatGroupResponse{}, nil
}
// GetChat retrieves a single chat's information
func (m *Mongo) GetChat(sid string, cid string, locale ...string) (*ChatInfo, error) {
return &ChatInfo{}, nil
func (m *Mongo) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetChatWithFilter retrieves a single chat's information with filter options
func (m *Mongo) GetChatWithFilter(sid string, cid string, filter ChatFilter, locale ...string) (*ChatInfo, error) {
return &ChatInfo{}, nil
func (m *Mongo) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetHistory retrieves chat history
@ -29,7 +31,7 @@ func (m *Mongo) GetHistory(sid string, cid string, locale ...string) ([]map[stri
}
// GetHistoryWithFilter retrieves chat history with filter options
func (m *Mongo) GetHistoryWithFilter(sid string, cid string, filter ChatFilter, locale ...string) ([]map[string]interface{}, error) {
func (m *Mongo) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
@ -54,7 +56,7 @@ func (m *Mongo) UpdateChatTitle(sid string, cid string, title string) error {
}
// SaveAssistant saves assistant information
func (m *Mongo) SaveAssistant(assistant *AssistantModel) (string, error) {
func (m *Mongo) SaveAssistant(assistant *types.AssistantModel) (string, error) {
return assistant.ID, nil
}
@ -64,23 +66,23 @@ func (m *Mongo) DeleteAssistant(assistantID string) error {
}
// GetAssistants retrieves a list of assistants
func (m *Mongo) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error) {
return &AssistantList{}, nil
func (m *Mongo) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
return &types.AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (*AssistantModel, error) {
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions (not implemented)
func (m *Mongo) DeleteAssistants(filter AssistantFilter) (int64, error) {
func (m *Mongo) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
return 0, nil
}
// GetAssistantTags retrieves all unique tags from assistants
func (m *Mongo) GetAssistantTags(locale ...string) ([]Tag, error) {
return []Tag{}, nil
func (m *Mongo) GetAssistantTags(locale ...string) ([]types.Tag, error) {
return []types.Tag{}, nil
}
// Close closes the store and releases any resources

View file

@ -1,26 +1,28 @@
package store
import "github.com/yaoapp/yao/agent/store/types"
// Redis represents a Redis-based conversation storage
type Redis struct{}
// NewRedis create a new redis store
func NewRedis() Store {
func NewRedis() types.Store {
return &Redis{}
}
// GetChats retrieves a list of chats
func (r *Redis) GetChats(sid string, filter ChatFilter, locale ...string) (*ChatGroupResponse, error) {
return &ChatGroupResponse{}, nil
func (r *Redis) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
return &types.ChatGroupResponse{}, nil
}
// GetChat retrieves a single chat's information
func (r *Redis) GetChat(sid string, cid string, locale ...string) (*ChatInfo, error) {
return &ChatInfo{}, nil
func (r *Redis) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetChatWithFilter retrieves a single chat's information with filter options
func (r *Redis) GetChatWithFilter(sid string, cid string, filter ChatFilter, locale ...string) (*ChatInfo, error) {
return &ChatInfo{}, nil
func (r *Redis) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
return &types.ChatInfo{}, nil
}
// GetHistory retrieves chat history
@ -29,7 +31,7 @@ func (r *Redis) GetHistory(sid string, cid string, locale ...string) ([]map[stri
}
// GetHistoryWithFilter retrieves chat history with filter options
func (r *Redis) GetHistoryWithFilter(sid string, cid string, filter ChatFilter, locale ...string) ([]map[string]interface{}, error) {
func (r *Redis) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
return []map[string]interface{}{}, nil
}
@ -54,7 +56,7 @@ func (r *Redis) UpdateChatTitle(sid string, cid string, title string) error {
}
// SaveAssistant saves assistant information
func (r *Redis) SaveAssistant(assistant *AssistantModel) (string, error) {
func (r *Redis) SaveAssistant(assistant *types.AssistantModel) (string, error) {
return assistant.ID, nil
}
@ -64,23 +66,23 @@ func (r *Redis) DeleteAssistant(assistantID string) error {
}
// GetAssistants retrieves a list of assistants
func (r *Redis) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error) {
return &AssistantList{}, nil
func (r *Redis) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
return &types.AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
func (r *Redis) GetAssistant(assistantID string, locale ...string) (*AssistantModel, error) {
func (r *Redis) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}
// DeleteAssistants deletes assistants based on filter conditions (not implemented)
func (r *Redis) DeleteAssistants(filter AssistantFilter) (int64, error) {
func (r *Redis) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
return 0, nil
}
// GetAssistantTags retrieves all unique tags from assistants
func (r *Redis) GetAssistantTags(locale ...string) ([]Tag, error) {
return []Tag{}, nil
func (r *Redis) GetAssistantTags(locale ...string) ([]types.Tag, error) {
return []types.Tag{}, nil
}
// Close closes the store and releases any resources

View file

@ -1,4 +1,4 @@
package store
package types
// Store defines the conversation storage interface
// Provides basic operations required for conversation management

View file

@ -1,4 +1,4 @@
package store
package types
import (
"fmt"

View file

@ -1,4 +1,4 @@
package store
package types
import (
"testing"

View file

@ -1,4 +1,4 @@
package store
package types
import "github.com/yaoapp/yao/agent/i18n"

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,656 @@
package xun
import (
"fmt"
"math"
"strings"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store/types"
)
// SaveAssistant saves assistant information
func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error) {
if assistant == nil {
return "", fmt.Errorf("assistant cannot be nil")
}
// Validate required fields
if assistant.Name == "" {
return "", fmt.Errorf("field name is required")
}
if assistant.Type == "" {
return "", fmt.Errorf("field type is required")
}
if assistant.Connector == "" {
return "", fmt.Errorf("field connector is required")
}
// Generate assistant_id if not provided
if assistant.ID == "" {
var err error
assistant.ID, err = conv.GenerateAssistantID()
if err != nil {
return "", err
}
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistant.ID).
Exists()
if err != nil {
return "", err
}
// Convert model to map for database storage
data := make(map[string]interface{})
data["assistant_id"] = assistant.ID
data["type"] = assistant.Type
data["connector"] = assistant.Connector
data["built_in"] = assistant.BuiltIn
data["sort"] = assistant.Sort
data["readonly"] = assistant.Readonly
data["public"] = assistant.Public
data["mentionable"] = assistant.Mentionable
data["automated"] = assistant.Automated
data["created_at"] = assistant.CreatedAt
data["updated_at"] = assistant.UpdatedAt
// Handle nullable string fields from assistant.mod.yao
// Store as nil if empty string (this matches database nullable: true fields)
if assistant.Name != "" {
data["name"] = assistant.Name
} else {
data["name"] = nil
}
if assistant.Avatar != "" {
data["avatar"] = assistant.Avatar
} else {
data["avatar"] = nil
}
if assistant.Description != "" {
data["description"] = assistant.Description
} else {
data["description"] = nil
}
if assistant.Path != "" {
data["path"] = assistant.Path
} else {
data["path"] = nil
}
// Share field: nullable: false with default "private"
// Apply default if empty
if assistant.Share != "" {
data["share"] = assistant.Share
} else {
data["share"] = "private" // Apply default value
}
// Permission management fields - store as nil if empty
if assistant.YaoCreatedBy != "" {
data["__yao_created_by"] = assistant.YaoCreatedBy
} else {
data["__yao_created_by"] = nil
}
if assistant.YaoUpdatedBy != "" {
data["__yao_updated_by"] = assistant.YaoUpdatedBy
} else {
data["__yao_updated_by"] = nil
}
if assistant.YaoTeamID != "" {
data["__yao_team_id"] = assistant.YaoTeamID
} else {
data["__yao_team_id"] = nil
}
if assistant.YaoTenantID != "" {
data["__yao_tenant_id"] = assistant.YaoTenantID
} else {
data["__yao_tenant_id"] = nil
}
// Handle simple types
if assistant.Options != nil {
jsonStr, err := jsoniter.MarshalToString(assistant.Options)
if err != nil {
return "", fmt.Errorf("failed to marshal options: %w", err)
}
data["options"] = jsonStr
}
if assistant.Tags != nil {
jsonStr, err := jsoniter.MarshalToString(assistant.Tags)
if err != nil {
return "", fmt.Errorf("failed to marshal tags: %w", err)
}
data["tags"] = jsonStr
}
// Handle interface{} fields - they should already be in the correct format
jsonFields := map[string]interface{}{
"prompts": assistant.Prompts,
"kb": assistant.KB,
"mcp": assistant.MCP,
"workflow": assistant.Workflow,
"tools": assistant.Tools,
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
}
for field, value := range jsonFields {
if value != nil {
jsonStr, err := jsoniter.MarshalToString(value)
if err != nil {
return "", fmt.Errorf("failed to marshal %s: %w", field, err)
}
data[field] = jsonStr
}
}
// Update or insert
if exists {
_, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistant.ID).
Update(data)
if err != nil {
return "", err
}
return assistant.ID, nil
}
err = conv.query.New().
Table(conv.getAssistantTable()).
Insert(data)
if err != nil {
return "", err
}
return assistant.ID, nil
}
// DeleteAssistant deletes an assistant by assistant_id
func (conv *Xun) DeleteAssistant(assistantID string) error {
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Exists()
if err != nil {
return err
}
if !exists {
return fmt.Errorf("assistant %s not found", assistantID)
}
_, err = conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
Delete()
return err
}
// GetAssistants retrieves assistants with pagination and filtering
func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (*types.AssistantList, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
// Apply tag filter if provided
if len(filter.Tags) > 0 {
qb.Where(func(qb query.Query) {
for i, tag := range filter.Tags {
// For each tag, we need to match it as part of a JSON array
// This will match both single tag arrays ["tag1"] and multi-tag arrays ["tag1","tag2"]
pattern := fmt.Sprintf("%%\"%s\"%%", tag)
if i == 0 {
qb.Where("tags", "like", pattern)
} else {
qb.OrWhere("tags", "like", pattern)
}
}
})
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where(func(qb query.Query) {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
})
}
// Apply type filter if provided
if filter.Type != "" {
qb.Where("type", filter.Type)
}
// Apply connector filter if provided
if filter.Connector != "" {
qb.Where("connector", filter.Connector)
}
// Apply assistant_id filter if provided
if filter.AssistantID != "" {
qb.Where("assistant_id", filter.AssistantID)
}
// Apply assistantIDs filter if provided
if len(filter.AssistantIDs) > 0 {
qb.WhereIn("assistant_id", filter.AssistantIDs)
}
// Apply mentionable filter if provided
if filter.Mentionable != nil {
qb.Where("mentionable", *filter.Mentionable)
}
// Apply automated filter if provided
if filter.Automated != nil {
qb.Where("automated", *filter.Automated)
}
// Apply built_in filter if provided
if filter.BuiltIn != nil {
qb.Where("built_in", *filter.BuiltIn)
}
// Set defaults for pagination
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.Page <= 0 {
filter.Page = 1
}
// Get total count
total, err := qb.Clone().Count()
if err != nil {
return nil, err
}
// Calculate pagination
offset := (filter.Page - 1) * filter.PageSize
totalPages := int(math.Ceil(float64(total) / float64(filter.PageSize)))
nextPage := filter.Page + 1
if nextPage > totalPages {
nextPage = 0
}
prevPage := filter.Page - 1
if prevPage < 1 {
prevPage = 0
}
// Apply select fields if provided
if len(filter.Select) > 0 {
selectFields := make([]interface{}, len(filter.Select))
for i, field := range filter.Select {
selectFields[i] = field
}
qb.Select(selectFields...)
}
// Get paginated results
rows, err := qb.OrderBy("sort", "asc").
OrderBy("updated_at", "desc").
Offset(offset).
Limit(filter.PageSize).
Get()
if err != nil {
return nil, err
}
// Convert rows to types.AssistantModel slice
assistants := make([]*types.AssistantModel, 0, len(rows))
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales"}
for _, row := range rows {
data := row.ToMap()
if data == nil {
continue
}
// Parse JSON fields
conv.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel using existing helper function
model, err := types.ToAssistantModel(data)
if err != nil {
log.Error("Failed to convert row to types.AssistantModel: %s", err.Error())
continue
}
// Apply i18n translations if locale is provided
if len(locale) > 0 && model != nil {
lang := strings.ToLower(locale[0])
// Translate name if locales are available
if model.Locales != nil {
if localeData, ok := model.Locales[lang]; ok {
if messages, ok := localeData.Messages["name"]; ok {
if nameStr, ok := messages.(string); ok {
model.Name = nameStr
}
}
if messages, ok := localeData.Messages["description"]; ok {
if descStr, ok := messages.(string); ok {
model.Description = descStr
}
}
}
}
}
assistants = append(assistants, model)
}
return &types.AssistantList{
Data: assistants,
Page: filter.Page,
PageSize: filter.PageSize,
PageCount: totalPages,
Next: nextPage,
Prev: prevPage,
Total: int(total),
}, nil
}
// GetAssistant retrieves a single assistant by ID
func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
row, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
First()
if err != nil {
return nil, err
}
if row == nil {
return nil, fmt.Errorf("assistant %s not found", assistantID)
}
data := row.ToMap()
if len(data) == 0 {
return nil, fmt.Errorf("the assistant %s is empty", assistantID)
}
// Parse JSON fields
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales"}
conv.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
model := &types.AssistantModel{
ID: getString(data, "assistant_id"),
Type: getString(data, "type"),
Name: getString(data, "name"),
Avatar: getString(data, "avatar"),
Connector: getString(data, "connector"),
Path: getString(data, "path"),
BuiltIn: getBool(data, "built_in"),
Sort: getInt(data, "sort"),
Description: getString(data, "description"),
Readonly: getBool(data, "readonly"),
Public: getBool(data, "public"),
Share: getString(data, "share"),
Mentionable: getBool(data, "mentionable"),
Automated: getBool(data, "automated"),
CreatedAt: getInt64(data, "created_at"),
UpdatedAt: getInt64(data, "updated_at"),
YaoCreatedBy: getString(data, "__yao_created_by"),
YaoUpdatedBy: getString(data, "__yao_updated_by"),
YaoTeamID: getString(data, "__yao_team_id"),
YaoTenantID: getString(data, "__yao_tenant_id"),
}
// Handle Tags
if tags, ok := data["tags"].([]interface{}); ok {
model.Tags = make([]string, len(tags))
for i, tag := range tags {
if s, ok := tag.(string); ok {
model.Tags[i] = s
}
}
}
// Handle Options
if options, ok := data["options"].(map[string]interface{}); ok {
model.Options = options
}
// Handle typed fields with conversion
if prompts, has := data["prompts"]; has && prompts != nil {
// Try to unmarshal to []Prompt
raw, err := jsoniter.Marshal(prompts)
if err == nil {
var p []types.Prompt
if err := jsoniter.Unmarshal(raw, &p); err == nil {
model.Prompts = p
}
}
}
if kb, has := data["kb"]; has && kb != nil {
kbConverted, err := types.ToKnowledgeBase(kb)
if err == nil {
model.KB = kbConverted
}
}
if mcp, has := data["mcp"]; has && mcp != nil {
mcpConverted, err := types.ToMCPServers(mcp)
if err == nil {
model.MCP = mcpConverted
}
}
if workflow, has := data["workflow"]; has && workflow != nil {
wf, err := types.ToWorkflow(workflow)
if err == nil {
model.Workflow = wf
}
}
if tools, has := data["tools"]; has && tools != nil {
raw, err := jsoniter.Marshal(tools)
if err == nil {
var tc types.ToolCalls
if err := jsoniter.Unmarshal(raw, &tc); err == nil {
model.Tools = &tc
}
}
}
if placeholder, has := data["placeholder"]; has && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)
if err == nil {
var ph types.Placeholder
if err := jsoniter.Unmarshal(raw, &ph); err == nil {
model.Placeholder = &ph
}
}
}
if locales, has := data["locales"]; has && locales != nil {
raw, err := jsoniter.Marshal(locales)
if err == nil {
var loc i18n.Map
if err := jsoniter.Unmarshal(raw, &loc); err == nil {
model.Locales = loc
}
}
}
return model, nil
}
// DeleteAssistants deletes assistants based on filter conditions
func (conv *Xun) DeleteAssistants(filter types.AssistantFilter) (int64, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
// Apply tag filter if provided
if len(filter.Tags) > 0 {
qb.Where(func(qb query.Query) {
for i, tag := range filter.Tags {
pattern := fmt.Sprintf("%%\"%s\"%%", tag)
if i == 0 {
qb.Where("tags", "like", pattern)
} else {
qb.OrWhere("tags", "like", pattern)
}
}
})
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where(func(qb query.Query) {
qb.Where("name", "like", fmt.Sprintf("%%%s%%", filter.Keywords)).
OrWhere("description", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
})
}
// Apply connector filter if provided
if filter.Connector != "" {
qb.Where("connector", filter.Connector)
}
// Apply assistant_id filter if provided
if filter.AssistantID != "" {
qb.Where("assistant_id", filter.AssistantID)
}
// Apply assistantIDs filter if provided
if len(filter.AssistantIDs) > 0 {
qb.WhereIn("assistant_id", filter.AssistantIDs)
}
// Apply mentionable filter if provided
if filter.Mentionable != nil {
qb.Where("mentionable", *filter.Mentionable)
}
// Apply automated filter if provided
if filter.Automated != nil {
qb.Where("automated", *filter.Automated)
}
// Apply built_in filter if provided
if filter.BuiltIn != nil {
qb.Where("built_in", *filter.BuiltIn)
}
// Execute delete and return number of deleted records
return qb.Delete()
}
// GetAssistantTags retrieves all unique tags from assistants
func (conv *Xun) GetAssistantTags(locale ...string) ([]types.Tag, error) {
q := conv.newQuery().Table(conv.getAssistantTable())
rows, err := q.Select("tags").Where("type", "assistant").GroupBy("tags").Get()
if err != nil {
return nil, err
}
tagSet := map[string]bool{}
for _, row := range rows {
if tags, ok := row["tags"].(string); ok && tags != "" {
var tagList []string
if err := jsoniter.UnmarshalFromString(tags, &tagList); err == nil {
for _, tag := range tagList {
tagSet[tag] = true
}
}
}
}
lang := "en"
if len(locale) > 0 {
lang = locale[0]
}
// Convert map keys to slice
tags := make([]types.Tag, 0, len(tagSet))
for tag := range tagSet {
tags = append(tags, types.Tag{
Value: tag,
Label: i18n.TranslateGlobal(lang, tag).(string),
})
}
return tags, nil
}
// GetHistoryWithFilter get the history with filter options
func (conv *Xun) GetHistoryWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
qb := conv.newQuery().
Select("role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "uid", "silent", "created_at", "updated_at").
Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
// Apply silent filter if provided, otherwise exclude silent messages by default
if filter.Silent != nil {
if *filter.Silent {
// Include all messages (both silent and non-silent)
} else {
// Only include non-silent messages
qb.Where("silent", false)
}
} else {
// Default behavior: exclude silent messages
qb.Where("silent", false)
}
if conv.setting.TTL > 0 {
qb.Where("expired_at", ">", time.Now())
}
limit := 20
if conv.setting.MaxSize > 0 {
limit = conv.setting.MaxSize
}
if filter.PageSize > 0 {
limit = filter.PageSize
}
// Apply pagination if provided
if filter.Page > 0 {
offset := (filter.Page - 1) * limit
qb.Offset(offset)
}
rows, err := qb.Limit(limit).Get()
if err != nil {
return nil, err
}
res := []map[string]interface{}{}
for _, row := range rows {
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"assistant_id": row.Get("assistant_id"),
"assistant_name": row.Get("assistant_name"),
"assistant_avatar": row.Get("assistant_avatar"),
"mentions": row.Get("mentions"),
"uid": row.Get("uid"),
"silent": row.Get("silent"),
"created_at": row.Get("created_at"),
"updated_at": row.Get("updated_at"),
}
res = append([]map[string]interface{}{message}, res...)
}
return res, nil
}

427
agent/store/xun/chat.go Normal file
View file

@ -0,0 +1,427 @@
package xun
import (
"fmt"
"math"
"strings"
"time"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store/types"
)
// UpdateChatTitle update the chat title
func (conv *Xun) UpdateChatTitle(sid string, cid string, title string) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
_, err = conv.newQueryChat().
Where("sid", userID).
Where("chat_id", cid).
Update(map[string]interface{}{
"title": title,
"updated_at": time.Now(),
})
return err
}
// GetChat get the chat info and its history
func (conv *Xun) GetChat(sid string, cid string, locale ...string) (*types.ChatInfo, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
// Get chat info
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id").
Where("sid", userID).
Where("chat_id", cid)
row, err := qb.First()
if err != nil {
return nil, err
}
// Return nil if chat_id is nil (means no chat found)
if row.Get("chat_id") == nil {
return nil, nil
}
chat := map[string]interface{}{
"chat_id": row.Get("chat_id"),
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
}
// Get assistant details if assistant_id exists
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistant, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("name", "avatar").
Where("assistant_id", assistantID).
First()
if err != nil {
return nil, err
}
name := assistant.Get("name")
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(assistantID.(string), lang, name).(string)
}
if assistant != nil {
chat["assistant_name"] = name
chat["assistant_avatar"] = assistant.Get("avatar")
}
}
// Get chat history with default filter (silent=false)
history, err := conv.GetHistory(sid, cid, locale...)
if err != nil {
return nil, err
}
return &types.ChatInfo{
Chat: chat,
History: history,
}, nil
}
// GetChatWithFilter get the chat info and its history with filter options
func (conv *Xun) GetChatWithFilter(sid string, cid string, filter types.ChatFilter, locale ...string) (*types.ChatInfo, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
// Get chat info
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id").
Where("sid", userID).
Where("chat_id", cid)
row, err := qb.First()
if err != nil {
return nil, err
}
// Return nil if chat_id is nil (means no chat found)
if row.Get("chat_id") == nil {
return nil, nil
}
chat := map[string]interface{}{
"chat_id": row.Get("chat_id"),
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
}
// Get assistant details if assistant_id exists
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistant, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("name", "avatar").
Where("assistant_id", assistantID).
First()
if err != nil {
return nil, err
}
if assistant != nil {
chat["assistant_name"] = assistant.Get("name")
chat["assistant_avatar"] = assistant.Get("avatar")
}
}
// Get chat history with filter
history, err := conv.GetHistoryWithFilter(sid, cid, filter, locale...)
if err != nil {
return nil, err
}
return &types.ChatInfo{
Chat: chat,
History: history,
}, nil
}
// DeleteChat deletes a specific chat and its history
func (conv *Xun) DeleteChat(sid string, cid string) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Delete history records first
_, err = conv.newQuery().
Where("sid", userID).
Where("cid", cid).
Delete()
if err != nil {
return err
}
// Then delete the chat
_, err = conv.newQueryChat().
Where("sid", userID).
Where("chat_id", cid).
Limit(1).
Delete()
return err
}
// DeleteAllChats deletes all chats and their histories for a user
func (conv *Xun) DeleteAllChats(sid string) error {
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Delete history records first
_, err = conv.newQuery().
Where("sid", userID).
Delete()
if err != nil {
return err
}
// Then delete all chats
_, err = conv.newQueryChat().
Where("sid", userID).
Delete()
return err
}
// GetChats get the chat list with grouping by date
func (conv *Xun) GetChats(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
// Default behavior: exclude silent chats
if filter.Silent == nil {
silentFalse := false
filter.Silent = &silentFalse
}
return conv.getChatsWithFilter(sid, filter, locale...)
}
// getChatsWithFilter get the chats with filter options
func (conv *Xun) getChatsWithFilter(sid string, filter types.ChatFilter, locale ...string) (*types.ChatGroupResponse, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
// Set default values
if filter.Page <= 0 {
filter.Page = 1
}
if filter.PageSize <= 0 {
filter.PageSize = 20
}
if filter.Order == "" {
filter.Order = "desc"
}
// Get total count
qbCount := conv.newQueryChat().
Where("sid", userID)
// Apply silent filter if provided
if filter.Silent != nil {
if *filter.Silent {
// Include all chats (both silent and non-silent)
} else {
// Only include non-silent chats
qbCount.Where("silent", false)
}
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qbCount.Where("title", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
}
total, err := qbCount.Count()
if err != nil {
return nil, err
}
// Calculate last page
lastPage := int(math.Ceil(float64(total) / float64(filter.PageSize)))
if lastPage < 1 {
lastPage = 1
}
// Get chats with pagination
qb := conv.newQueryChat().
Select("chat_id", "title", "assistant_id", "silent", "created_at", "updated_at").
Where("sid", userID)
// Apply silent filter if provided
if filter.Silent != nil {
if *filter.Silent {
// Include all chats (both silent and non-silent)
} else {
// Only include non-silent chats
qb.Where("silent", false)
}
}
// Apply keyword filter if provided
if filter.Keywords != "" {
qb.Where("title", "like", fmt.Sprintf("%%%s%%", filter.Keywords))
}
// Apply pagination
offset := (filter.Page - 1) * filter.PageSize
qb.OrderBy("updated_at", filter.Order).
Offset(offset).
Limit(filter.PageSize)
rows, err := qb.Get()
if err != nil {
return nil, err
}
// Group chats by date
today := time.Now().Truncate(24 * time.Hour)
yesterday := today.AddDate(0, 0, -1)
thisWeekStart := today.AddDate(0, 0, -int(today.Weekday()))
lastWeekStart := thisWeekStart.AddDate(0, 0, -7)
lastWeekEnd := thisWeekStart.AddDate(0, 0, -1)
groups := map[string][]map[string]interface{}{
"Today": {},
"Yesterday": {},
"This Week": {},
"Last Week": {},
"Even Earlier": {},
}
// Collect assistant IDs to fetch their details
assistantIDs := []interface{}{}
for _, row := range rows {
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
assistantIDs = append(assistantIDs, assistantID)
}
}
// Fetch assistant details
assistantMap := map[string]map[string]interface{}{}
if len(assistantIDs) > 0 {
assistants, err := conv.query.New().
Table(conv.getAssistantTable()).
Select("assistant_id", "name", "avatar").
WhereIn("assistant_id", assistantIDs).
Get()
if err != nil {
return nil, err
}
for _, assistant := range assistants {
if id := assistant.Get("assistant_id"); id != nil {
name := assistant.Get("name")
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(id.(string), lang, name).(string)
}
assistantMap[fmt.Sprintf("%v", id)] = map[string]interface{}{
"name": name,
"avatar": assistant.Get("avatar"),
}
}
}
}
for _, row := range rows {
chatID := row.Get("chat_id")
if chatID == nil || chatID == "" {
continue
}
chat := map[string]interface{}{
"chat_id": chatID,
"title": row.Get("title"),
"assistant_id": row.Get("assistant_id"),
"silent": row.Get("silent"),
}
// Add assistant details if available
if assistantID := row.Get("assistant_id"); assistantID != nil && assistantID != "" {
if assistant, ok := assistantMap[fmt.Sprintf("%v", assistantID)]; ok {
name := assistant["name"]
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
name = i18n.Translate(assistantID.(string), lang, name).(string)
}
chat["assistant_name"] = name
chat["assistant_avatar"] = assistant["avatar"]
}
}
var dbDatetime = row.Get("updated_at")
if dbDatetime == nil {
dbDatetime = row.Get("created_at")
}
var createdAt time.Time
switch v := dbDatetime.(type) {
case time.Time:
createdAt = v
case string:
parsed, err := time.Parse("2006-01-02 15:04:05.999999-07:00", v)
if err != nil {
// Try alternative format
parsed, err = time.Parse(time.RFC3339, v)
if err != nil {
continue
}
}
createdAt = parsed
default:
continue
}
createdDate := createdAt.Truncate(24 * time.Hour)
switch {
case createdDate.Equal(today):
groups["Today"] = append(groups["Today"], chat)
case createdDate.Equal(yesterday):
groups["Yesterday"] = append(groups["Yesterday"], chat)
case createdDate.After(thisWeekStart) && createdDate.Before(today):
groups["This Week"] = append(groups["This Week"], chat)
case createdDate.After(lastWeekStart) && createdDate.Before(lastWeekEnd.AddDate(0, 0, 1)):
groups["Last Week"] = append(groups["Last Week"], chat)
default:
groups["Even Earlier"] = append(groups["Even Earlier"], chat)
}
}
// Convert to ordered slice and apply i18n
result := []types.ChatGroup{}
for _, label := range []string{"Today", "Yesterday", "This Week", "Last Week", "Even Earlier"} {
if len(groups[label]) > 0 {
translatedLabel := label
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
translatedLabel = i18n.TranslateGlobal(lang, label).(string)
}
result = append(result, types.ChatGroup{
Label: translatedLabel,
Chats: groups[label],
})
}
}
return &types.ChatGroupResponse{
Groups: result,
Page: filter.Page,
PageSize: filter.PageSize,
Total: total,
LastPage: lastPage,
}, nil
}

250
agent/store/xun/history.go Normal file
View file

@ -0,0 +1,250 @@
package xun
import (
"fmt"
"strings"
"time"
"github.com/google/uuid"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/i18n"
)
// GetHistory get the history
func (conv *Xun) GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error) {
userID, err := conv.getUserID(sid)
if err != nil {
return nil, err
}
qb := conv.newQuery().
Select("role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "uid", "silent", "created_at", "updated_at").
Where("sid", userID).
Where("cid", cid).
OrderBy("id", "desc")
// By default, exclude silent messages
qb.Where("silent", false)
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 {
assistantName := row.Get("assistant_name")
assistantID := row.Get("assistant_id")
if len(locale) > 0 && assistantID != nil {
lang := strings.ToLower(locale[0])
assistantName = i18n.Translate(assistantID.(string), lang, assistantName).(string)
}
message := map[string]interface{}{
"role": row.Get("role"),
"name": row.Get("name"),
"content": row.Get("content"),
"context": row.Get("context"),
"assistant_id": row.Get("assistant_id"),
"assistant_name": assistantName,
"assistant_avatar": row.Get("assistant_avatar"),
"mentions": row.Get("mentions"),
"uid": row.Get("uid"),
"silent": row.Get("silent"),
"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, context map[string]interface{}) error {
if cid == "" {
cid = uuid.New().String() // Generate a new UUID if cid is empty
}
userID, err := conv.getUserID(sid)
if err != nil {
return err
}
// Get assistant_id from context
var assistantID interface{} = nil
if context != nil {
if id, ok := context["assistant_id"].(string); ok && id != "" {
assistantID = id
}
}
// Get silent flag from context
var silent bool = false
var historyVisible bool = true
if context != nil {
if silentVal, ok := context["silent"]; ok {
switch v := silentVal.(type) {
case bool:
silent = v
case string:
silent = v == "true" || v == "1" || v == "yes"
case int:
silent = v != 0
case float64:
silent = v != 0
}
}
// Get history visible from context
if historyVisibleVal, ok := context["history_visible"]; ok {
switch v := historyVisibleVal.(type) {
case bool:
historyVisible = v
case string:
historyVisible = v == "true" || v == "1" || v == "yes"
case int:
historyVisible = v != 0
case float64:
historyVisible = v != 0
}
}
}
// First ensure chat record exists
exists, err := conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Exists()
if err != nil {
return err
}
if !exists {
// Create new chat record
err = conv.newQueryChat().
Insert(map[string]interface{}{
"chat_id": cid,
"sid": userID,
"assistant_id": assistantID,
"silent": silent || historyVisible == false,
"created_at": time.Now(),
})
if err != nil {
return err
}
} else {
// Update assistant_id and silent if needed
_, err = conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Update(map[string]interface{}{
"assistant_id": assistantID,
"silent": silent || historyVisible == false,
})
if err != nil {
return err
}
}
// Save message history
var expiredAt interface{} = nil
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 {
// Type assertion safety checks
role, ok := message["role"].(string)
if !ok {
return fmt.Errorf("invalid role type in message: %v", message["role"])
}
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
}
}
// Process mentions if present
var mentionsRaw interface{} = nil
if mentions, ok := message["mentions"].([]interface{}); ok && len(mentions) > 0 {
mentionsRaw, err = jsoniter.MarshalToString(mentions)
if err != nil {
return err
}
}
value := map[string]interface{}{
"role": role,
"name": "",
"content": content,
"sid": userID,
"cid": cid,
"uid": userID,
"context": contextRaw,
"mentions": mentionsRaw,
"assistant_id": nil,
"assistant_name": nil,
"assistant_avatar": nil,
"silent": silent,
"created_at": now,
"updated_at": nil,
"expired_at": expiredAt,
}
if name, ok := message["name"].(string); ok {
value["name"] = name
}
// Add assistant fields if present
if assistantID, ok := message["assistant_id"].(string); ok {
value["assistant_id"] = assistantID
}
if assistantName, ok := message["assistant_name"].(string); ok {
value["assistant_name"] = assistantName
}
if assistantAvatar, ok := message["assistant_avatar"].(string); ok {
value["assistant_avatar"] = assistantAvatar
}
values = append(values, value)
}
err = conv.newQuery().Insert(values)
if err != nil {
return err
}
// Update Chat updated_at
_, err = conv.newQueryChat().
Where("chat_id", cid).
Where("sid", userID).
Update(map[string]interface{}{"updated_at": now})
if err != nil {
return err
}
return nil
}

40
agent/store/xun/utils.go Normal file
View file

@ -0,0 +1,40 @@
package xun
// Helper functions for type conversion
func getString(data map[string]interface{}, key string) string {
if v, ok := data[key].(string); ok {
return v
}
return ""
}
func getBool(data map[string]interface{}, key string) bool {
if v, ok := data[key].(bool); ok {
return v
}
return false
}
func getInt(data map[string]interface{}, key string) int {
switch v := data[key].(type) {
case int:
return v
case int64:
return int(v)
case float64:
return int(v)
}
return 0
}
func getInt64(data map[string]interface{}, key string) int64 {
switch v := data[key].(type) {
case int64:
return v
case int:
return int64(v)
case float64:
return int64(v)
}
return 0
}

329
agent/store/xun/xun.go Normal file
View file

@ -0,0 +1,329 @@
package xun
import (
"fmt"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/xun/capsule"
"github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/xun/dbal/schema"
"github.com/yaoapp/yao/agent/store/types"
)
// Package store provides functionality for managing chat conversations and assistants.
// Xun implements the Store interface using a database backend.
// It provides functionality for:
// - Managing chat conversations and their message histories
// - Organizing chats with pagination and date-based grouping
// - Handling chat metadata like titles and creation dates
// - Managing AI assistants with their configurations and metadata
// - Supporting data expiration through TTL settings
type Xun struct {
query query.Query
schema schema.Schema
setting types.Setting
cleanTicker *time.Ticker
cleanStop chan bool
}
// Public interface methods:
//
// NewXun creates a new conversation instance with the given settings
// GetChats retrieves a paginated list of chats grouped by date
// GetChat retrieves a specific chat and its message history
// GetChatWithFilter retrieves a specific chat with filter options
// GetHistory retrieves the message history for a specific chat
// GetHistoryWithFilter retrieves the message history with filter options
// SaveHistory saves new messages to a chat's historys
// DeleteChat deletes a specific chat and its history
// DeleteAllChats deletes all chats and their histories for a user
// UpdateChatTitle updates the title of a specific chat
// SaveAssistant creates or updates an assistant
// DeleteAssistant deletes an assistant by assistant_id
// GetAssistants retrieves a paginated list of assistants with filtering
// GetAssistant retrieves a single assistant by assistant_id
// DeleteAssistants deletes assistants based on filter conditions
// GetAssistantTags retrieves all unique tags from assistants
// Close closes the store and releases any resources
// NewXun create a new xun store
func NewXun(setting types.Setting) (types.Store, error) {
conv := &Xun{setting: setting}
if setting.Connector == "default" || setting.Connector == "" {
conv.query = capsule.Global.Query()
conv.schema = capsule.Global.Schema()
} else {
conn, err := connector.Select(setting.Connector)
if err != nil {
return nil, fmt.Errorf("select store connector %s error: %s", setting.Connector, err.Error())
}
conv.query, err = conn.Query()
if err != nil {
return nil, fmt.Errorf("query store connector %s error: %s", setting.Connector, err.Error())
}
conv.schema, err = conn.Schema()
if err != nil {
return nil, err
}
}
err := conv.initialize()
if err != nil {
return nil, err
}
return conv, nil
}
// Rename the following functions to start with lowercase letters to make them private:
func (conv *Xun) newQuery() query.Query {
qb := conv.query.New()
qb.Table(conv.getHistoryTable())
return qb
}
func (conv *Xun) newQueryChat() query.Query {
qb := conv.query.New()
qb.Table(conv.getChatTable())
return qb
}
func (conv *Xun) clean() {
nums, err := conv.newQuery().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: %d", nums)
}
}
// startAutoClean starts the automatic cleanup routine
func (conv *Xun) startAutoClean() {
if conv.cleanTicker != nil {
conv.stopAutoClean() // Stop existing ticker if any
}
conv.cleanTicker = time.NewTicker(1 * time.Hour) // Clean every hour
conv.cleanStop = make(chan bool)
go func() {
for {
select {
case <-conv.cleanTicker.C:
conv.clean()
case <-conv.cleanStop:
return
}
}
}()
log.Trace("Started automatic cleanup")
}
// stopAutoClean stops the automatic cleanup routine
func (conv *Xun) stopAutoClean() {
if conv.cleanTicker != nil {
conv.cleanTicker.Stop()
conv.cleanTicker = nil
}
if conv.cleanStop != nil {
close(conv.cleanStop)
conv.cleanStop = nil
}
log.Trace("Stopped automatic cleanup")
}
// Close stops the automatic cleanup and closes resources
func (conv *Xun) Close() error {
conv.stopAutoClean()
return nil
}
// Rename Init to initialize to avoid conflicts
func (conv *Xun) initialize() error {
// Start automatic cleanup if TTL is enabled
if conv.setting.TTL > 0 {
conv.startAutoClean()
}
return nil
}
func (conv *Xun) initHistoryTable() error {
historyTable := conv.getHistoryTable()
has, err := conv.schema.HasTable(historyTable)
if err != nil {
return err
}
// Create the history table
if !has {
err = conv.schema.CreateTable(historyTable, func(table schema.Blueprint) {
table.ID("id")
table.String("sid", 255).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.String("assistant_id", 200).Null().Index()
table.String("assistant_name", 200).Null()
table.String("assistant_avatar", 200).Null()
table.JSON("mentions").Null()
table.Boolean("silent").SetDefault(false).Index()
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
table.TimestampTz("expired_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the conversation history table: %s", historyTable)
}
// Validate the table
tab, err := conv.schema.GetTable(historyTable)
if err != nil {
return err
}
fields := []string{"id", "sid", "cid", "uid", "role", "name", "content", "context", "assistant_id", "assistant_name", "assistant_avatar", "mentions", "silent", "created_at", "updated_at", "expired_at"}
for _, field := range fields {
if !tab.HasColumn(field) {
return fmt.Errorf("%s is required", field)
}
}
return nil
}
func (conv *Xun) initChatTable() error {
chatTable := conv.getChatTable()
has, err := conv.schema.HasTable(chatTable)
if err != nil {
return err
}
// Create the chat table
if !has {
err = conv.schema.CreateTable(chatTable, func(table schema.Blueprint) {
table.ID("id")
table.String("chat_id", 200).Unique().Index()
table.String("title", 200).Null()
table.String("assistant_id", 200).Null().Index()
table.String("sid", 255).Index()
table.Boolean("silent").SetDefault(false).Index()
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").Index()
table.TimestampTz("updated_at").Null().Index()
})
if err != nil {
return err
}
log.Trace("Create the chat table: %s", chatTable)
}
// Validate the table
tab, err := conv.schema.GetTable(chatTable)
if err != nil {
return err
}
fields := []string{"id", "chat_id", "title", "assistant_id", "sid", "silent", "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) {
// TODO: get the user id from the authentication system
return "guest", nil
}
func (conv *Xun) getHistoryTable() string {
m := model.Select("__yao.agent.history")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.history"
}
func (conv *Xun) getChatTable() string {
m := model.Select("__yao.agent.chat")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.chat"
}
func (conv *Xun) getAssistantTable() string {
m := model.Select("__yao.agent.assistant")
if m != nil && m.MetaData.Table.Name != "" {
return m.MetaData.Table.Name
}
return "__yao.agent.assistant"
}
// parseJSONFields parses JSON string fields into their corresponding Go types
func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
for _, field := range fields {
if val := data[field]; val != nil {
if strVal, ok := val.(string); ok && strVal != "" {
var parsed interface{}
if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil {
data[field] = parsed
}
}
}
}
}
// GenerateAssistantID generates a random-looking 6-digit ID
func (conv *Xun) GenerateAssistantID() (string, error) {
maxAttempts := 10 // Maximum number of attempts to generate a unique ID
for i := 0; i < maxAttempts; i++ {
// Generate a random number using timestamp and some bit operations
timestamp := time.Now().UnixNano()
random := (timestamp ^ (timestamp >> 12)) % 1000000
hash := fmt.Sprintf("%06d", random)
// Check if this ID already exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", hash).
Exists()
if err != nil {
return "", err
}
if !exists {
return hash, nil
}
// If ID exists, wait a bit and try again
time.Sleep(time.Millisecond)
}
return "", fmt.Errorf("failed to generate unique ID after %d attempts", maxAttempts)
}

View file

@ -1,10 +1,11 @@
package store
package xun
import (
"fmt"
"testing"
"time"
"github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
@ -23,7 +24,7 @@ func TestSaveAssistant(t *testing.T) {
defer test.Clean()
// Create a new xun store
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -32,7 +33,7 @@ func TestSaveAssistant(t *testing.T) {
defer store.Close()
t.Run("CreateNewAssistant", func(t *testing.T) {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Test Assistant",
Type: "assistant",
Connector: "openai",
@ -67,7 +68,7 @@ func TestSaveAssistant(t *testing.T) {
t.Run("UpdateExistingAssistant", func(t *testing.T) {
// Create initial assistant
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Update Test Assistant",
Type: "assistant",
Connector: "openai",
@ -118,7 +119,7 @@ func TestSaveAssistant(t *testing.T) {
}
// Test missing name
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Type: "assistant",
Connector: "openai",
}
@ -128,7 +129,7 @@ func TestSaveAssistant(t *testing.T) {
}
// Test missing type
assistant = &AssistantModel{
assistant = &types.AssistantModel{
Name: "Test",
Connector: "openai",
}
@ -138,7 +139,7 @@ func TestSaveAssistant(t *testing.T) {
}
// Test missing connector
assistant = &AssistantModel{
assistant = &types.AssistantModel{
Name: "Test",
Type: "assistant",
}
@ -149,12 +150,12 @@ func TestSaveAssistant(t *testing.T) {
})
t.Run("ComplexDataTypes", func(t *testing.T) {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Complex Assistant",
Type: "assistant",
Connector: "openai",
Share: "private",
Prompts: []Prompt{
Prompts: []types.Prompt{
{Role: "system", Content: "You are a helpful assistant"},
{Role: "user", Content: "Hello"},
},
@ -163,7 +164,7 @@ func TestSaveAssistant(t *testing.T) {
"max_tokens": 2000,
},
Tags: []string{"complex", "testing", "data"},
Placeholder: &Placeholder{
Placeholder: &types.Placeholder{
Title: "Type your message",
Description: "Enter your message here...",
Prompts: []string{"What can I help you with?"},
@ -200,7 +201,7 @@ func TestDeleteAssistant(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -210,7 +211,7 @@ func TestDeleteAssistant(t *testing.T) {
t.Run("DeleteExistingAssistant", func(t *testing.T) {
// Create assistant
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Delete Test",
Type: "assistant",
Connector: "openai",
@ -248,7 +249,7 @@ func TestGetAssistant(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -258,7 +259,7 @@ func TestGetAssistant(t *testing.T) {
t.Run("GetExistingAssistant", func(t *testing.T) {
// Create assistant
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Get Test",
Type: "assistant",
Connector: "openai",
@ -316,7 +317,7 @@ func TestGetAssistants(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -325,7 +326,7 @@ func TestGetAssistants(t *testing.T) {
defer store.Close()
// Clean up existing data before creating test assistants
deleted, err := store.DeleteAssistants(AssistantFilter{})
deleted, err := store.DeleteAssistants(types.AssistantFilter{})
if err != nil {
t.Logf("Warning: Failed to clean up existing assistants: %v", err)
} else if deleted > 0 {
@ -333,7 +334,7 @@ func TestGetAssistants(t *testing.T) {
}
// Create test assistants
assistants := []AssistantModel{
assistants := []types.AssistantModel{
{
Name: "Assistant 1",
Type: "assistant",
@ -379,7 +380,7 @@ func TestGetAssistants(t *testing.T) {
}
t.Run("GetAllAssistants", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Page: 1,
PageSize: 20,
})
@ -397,7 +398,7 @@ func TestGetAssistants(t *testing.T) {
})
t.Run("FilterByType", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Type: "assistant",
Page: 1,
PageSize: 20,
@ -414,7 +415,7 @@ func TestGetAssistants(t *testing.T) {
})
t.Run("FilterByConnector", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Connector: "openai",
Page: 1,
PageSize: 20,
@ -431,7 +432,7 @@ func TestGetAssistants(t *testing.T) {
})
t.Run("FilterByTags", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Tags: []string{"automation"},
Page: 1,
PageSize: 20,
@ -460,7 +461,7 @@ func TestGetAssistants(t *testing.T) {
})
t.Run("FilterByKeywords", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Keywords: "Second",
Page: 1,
PageSize: 20,
@ -485,7 +486,7 @@ func TestGetAssistants(t *testing.T) {
t.Run("FilterByMentionable", func(t *testing.T) {
mentionableTrue := true
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Mentionable: &mentionableTrue,
Page: 1,
PageSize: 20,
@ -507,7 +508,7 @@ func TestGetAssistants(t *testing.T) {
t.Run("FilterByAutomated", func(t *testing.T) {
automatedFalse := false
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Automated: &automatedFalse,
Page: 1,
PageSize: 20,
@ -525,7 +526,7 @@ func TestGetAssistants(t *testing.T) {
t.Run("Pagination", func(t *testing.T) {
// Test first page
response1, err := store.GetAssistants(AssistantFilter{
response1, err := store.GetAssistants(types.AssistantFilter{
Page: 1,
PageSize: 2,
})
@ -547,7 +548,7 @@ func TestGetAssistants(t *testing.T) {
// Test second page if there are enough records
if response1.Total > 2 {
response2, err := store.GetAssistants(AssistantFilter{
response2, err := store.GetAssistants(types.AssistantFilter{
Page: 2,
PageSize: 2,
})
@ -562,7 +563,7 @@ func TestGetAssistants(t *testing.T) {
})
t.Run("FieldSelection", func(t *testing.T) {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Select: []string{"assistant_id", "name", "type"},
Page: 1,
PageSize: 20,
@ -587,7 +588,7 @@ func TestGetAssistants(t *testing.T) {
t.Run("FilterByAssistantID", func(t *testing.T) {
if len(createdIDs) > 0 {
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
AssistantID: createdIDs[0],
Page: 1,
PageSize: 20,
@ -609,7 +610,7 @@ func TestGetAssistants(t *testing.T) {
t.Run("FilterByAssistantIDs", func(t *testing.T) {
if len(createdIDs) >= 2 {
filterIDs := []string{createdIDs[0], createdIDs[1]}
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
AssistantIDs: filterIDs,
Page: 1,
PageSize: 20,
@ -630,7 +631,7 @@ func TestDeleteAssistants(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -642,7 +643,7 @@ func TestDeleteAssistants(t *testing.T) {
// Create assistants with specific tag
tag := fmt.Sprintf("delete-test-%d", time.Now().UnixNano())
for i := 0; i < 3; i++ {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: fmt.Sprintf("Delete Test %d", i),
Type: "assistant",
Connector: "openai",
@ -656,7 +657,7 @@ func TestDeleteAssistants(t *testing.T) {
}
// Delete by tag
count, err := store.DeleteAssistants(AssistantFilter{
count, err := store.DeleteAssistants(types.AssistantFilter{
Tags: []string{tag},
})
if err != nil {
@ -672,7 +673,7 @@ func TestDeleteAssistants(t *testing.T) {
// Create assistants with specific connector
connector := fmt.Sprintf("test-connector-%d", time.Now().UnixNano())
for i := 0; i < 2; i++ {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: fmt.Sprintf("Connector Test %d", i),
Type: "assistant",
Connector: connector,
@ -685,7 +686,7 @@ func TestDeleteAssistants(t *testing.T) {
}
// Delete by connector
count, err := store.DeleteAssistants(AssistantFilter{
count, err := store.DeleteAssistants(types.AssistantFilter{
Connector: connector,
})
if err != nil {
@ -700,7 +701,7 @@ func TestDeleteAssistants(t *testing.T) {
t.Run("DeleteByKeywords", func(t *testing.T) {
// Create assistants with specific keyword
keyword := fmt.Sprintf("unique-keyword-%d", time.Now().UnixNano())
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: fmt.Sprintf("Assistant with %s", keyword),
Type: "assistant",
Connector: "openai",
@ -713,7 +714,7 @@ func TestDeleteAssistants(t *testing.T) {
}
// Delete by keyword
count, err := store.DeleteAssistants(AssistantFilter{
count, err := store.DeleteAssistants(types.AssistantFilter{
Keywords: keyword,
})
if err != nil {
@ -727,7 +728,7 @@ func TestDeleteAssistants(t *testing.T) {
t.Run("DeleteByAssistantID", func(t *testing.T) {
// Create an assistant
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Single Delete Test",
Type: "assistant",
Connector: "openai",
@ -739,7 +740,7 @@ func TestDeleteAssistants(t *testing.T) {
}
// Delete by ID
count, err := store.DeleteAssistants(AssistantFilter{
count, err := store.DeleteAssistants(types.AssistantFilter{
AssistantID: id,
})
if err != nil {
@ -757,7 +758,7 @@ func TestGetAssistantTags(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -768,7 +769,7 @@ func TestGetAssistantTags(t *testing.T) {
t.Run("GetUniqueTags", func(t *testing.T) {
// Create assistants with various tags
uniqueTag := fmt.Sprintf("tag-test-%d", time.Now().UnixNano())
assistants := []AssistantModel{
assistants := []types.AssistantModel{
{
Name: "Tags Test 1",
Type: "assistant",
@ -829,7 +830,7 @@ func TestGenerateAssistantID(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -868,7 +869,7 @@ func TestAssistantPermissionFields(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -877,7 +878,7 @@ func TestAssistantPermissionFields(t *testing.T) {
defer store.Close()
t.Run("SaveWithPermissionFields", func(t *testing.T) {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Permission Test Assistant",
Type: "assistant",
Connector: "openai",
@ -918,7 +919,7 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Run("UpdatePermissionFields", func(t *testing.T) {
// Create assistant
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Update Permission Test",
Type: "assistant",
Connector: "openai",
@ -961,7 +962,7 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Run("EmptyPermissionFields", func(t *testing.T) {
// Create assistant without permission fields
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "No Permission Fields",
Type: "assistant",
Connector: "openai",
@ -999,7 +1000,7 @@ func TestEmptyStringAsNull(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -1013,7 +1014,7 @@ func TestEmptyStringAsNull(t *testing.T) {
// - name (nullable: true, but required by validation)
// - avatar, description, path (nullable: true)
// - share (nullable: false, but empty should trigger default)
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Test Null Fields", // Required by validation
Type: "assistant",
Connector: "openai",
@ -1059,7 +1060,7 @@ func TestEmptyStringAsNull(t *testing.T) {
t.Run("NonEmptyStringsPreserved", func(t *testing.T) {
// Create assistant with non-empty values
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: "Test Non-Empty Fields",
Type: "assistant",
Connector: "openai",
@ -1102,7 +1103,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
store, err := NewXun(Setting{
store, err := NewXun(types.Setting{
Connector: "default",
})
if err != nil {
@ -1114,7 +1115,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
// Step 1: Create multiple assistants
assistantIDs := []string{}
for i := 0; i < 3; i++ {
assistant := &AssistantModel{
assistant := &types.AssistantModel{
Name: fmt.Sprintf("Workflow Assistant %d", i),
Type: "assistant",
Connector: "openai",
@ -1134,7 +1135,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Logf("Created %d assistants", len(assistantIDs))
// Step 2: Retrieve all assistants
response, err := store.GetAssistants(AssistantFilter{
response, err := store.GetAssistants(types.AssistantFilter{
Tags: []string{"workflow"},
Page: 1,
PageSize: 20,
@ -1204,7 +1205,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
}
// Step 6: Bulk delete remaining assistants
count, err := store.DeleteAssistants(AssistantFilter{
count, err := store.DeleteAssistants(types.AssistantFilter{
Tags: []string{"workflow"},
})
if err != nil {
@ -1214,7 +1215,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Logf("Bulk deleted %d assistants", count)
// Verify bulk deletion
finalResponse, err := store.GetAssistants(AssistantFilter{
finalResponse, err := store.GetAssistants(types.AssistantFilter{
Tags: []string{"workflow"},
Page: 1,
PageSize: 20,

View file

@ -1,4 +1,4 @@
package store
package xun
// import (
// "fmt"

View file

@ -3,7 +3,7 @@ package agent
import (
"github.com/gin-gonic/gin"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/store"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/vision"
)