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

Refactor assistant handling and storage integration

Browse source 
- Updated assistant data handling to utilize structured models, improving code clarity and maintainability.
- Enhanced the assistant save and load functions to work with the new AssistantModel structure.
- Refactored mention handling to streamline the conversion process from raw data to structured mentions.
- Improved error handling in the assistant save process, ensuring better feedback for invalid data.
- Updated API responses to use the new gin.H format for consistency across the application.
This commit is contained in:
Max 2025-11-07 11:20:59 +08:00
parent 296a37f031
commit 1d5f27cd98
22 changed files with 3340 additions and 792 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
agent/api.go
View file

@ -537,12 +537,12 @@ func (agent *DSL) handleMentions(c *gin.Context) {
// Convert assistants to mentions
mentions := []Mention{}
for _, item := range response.Data {
for _, assistant := range response.Data {
mention := Mention{
ID: item["assistant_id"].(string),
Name: item["name"].(string),
Type: item["type"].(string),
Avatar: item["avatar"].(string),
ID: assistant.ID,
Name: assistant.Name,
Type: assistant.Type,
Avatar: assistant.Avatar,
}
mentions = append(mentions, mention)
}
@ -899,7 +899,7 @@ func (agent *DSL) HandleAssistantDetail(c *gin.Context) {
return
}
c.JSON(200, map[string]interface{}{"data": response.Data[0]})
c.JSON(200, gin.H{"data": response.Data[0]})
c.Done()
}
@ -912,7 +912,15 @@ func (agent *DSL) HandleAssistantSave(c *gin.Context) {
return
}
id, err := agent.Store.SaveAssistant(assistantData)
// Convert to AssistantModel
model, err := store.ToAssistantModel(assistantData)
if err != nil {
c.JSON(400, gin.H{"message": fmt.Sprintf("invalid assistant data: %s", err.Error()), "code": 400})
c.Done()
return
}
id, err := agent.Store.SaveAssistant(model)
if err != nil {
c.JSON(500, gin.H{"message": err.Error(), "code": 500})
c.Done()
@ -927,11 +935,11 @@ func (agent *DSL) HandleAssistantSave(c *gin.Context) {
// Remove the assistant from cache to ensure fresh data on next load
cache := assistant.GetCache()
if cache != nil {
cache.Remove(id.(string))
cache.Remove(id)
}
// Reload the assistant to ensure it's available in cache with updated data
_, err = assistant.Get(id.(string))
_, err = assistant.Get(id)
if err != nil {
// Just log the error, don't fail the request
fmt.Printf("Error reloading assistant %s: %v\n", id, err)

17
agent/assistant/api.go
View file

@ -18,6 +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"
)
// Get get the assistant by id
@ -259,7 +260,7 @@ func (next *NextAction) Execute(c *gin.Context, ctx chatctx.Context, contents *c
}
// GetPlaceholder returns the placeholder of the assistant
func (ast *Assistant) GetPlaceholder(locale string) *Placeholder {
func (ast *Assistant) GetPlaceholder(locale string) *store.Placeholder {
prompts := []string{}
if ast.Placeholder.Prompts != nil {
@ -267,7 +268,7 @@ func (ast *Assistant) GetPlaceholder(locale string) *Placeholder {
}
title := i18n.Translate(ast.ID, locale, ast.Placeholder.Title).(string)
description := i18n.Translate(ast.ID, locale, ast.Placeholder.Description).(string)
return &Placeholder{
return &store.Placeholder{
Title: title,
Description: description,
Prompts: prompts,
@ -795,10 +796,18 @@ func (ast *Assistant) withPrompts(messages []chatMessage.Message) []chatMessage.
if ast.Tools != nil && ast.Tools.Tools != nil && len(ast.Tools.Tools) > 0 {
settings, has := connectorSettings[ast.Connector]
if !has || !settings.Tools {
raw, _ := jsoniter.MarshalToString(ast.Tools.Tools)
// Convert store tools to runtime tools if not already done
if ast.runtimeTools == nil {
runtimeTools, err := ToRuntimeTools(ast.Tools.Tools)
if err == nil {
ast.runtimeTools = runtimeTools
}
}
raw, _ := jsoniter.MarshalToString(ast.runtimeTools)
examples := []string{}
for _, tool := range ast.Tools.Tools {
for _, tool := range ast.runtimeTools {
example := tool.Example()
examples = append(examples, example)
}

173
agent/assistant/assistant.go
View file

@ -6,6 +6,8 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/store"
sui "github.com/yaoapp/yao/sui/core"
)
@ -15,7 +17,7 @@ func (ast *Assistant) Save() error {
return fmt.Errorf("storage is not set")
}
_, err := storage.SaveAssistant(ast.Map())
_, err := storage.SaveAssistant(&ast.AssistantModel)
if err != nil {
return err
}
@ -35,6 +37,8 @@ func (ast *Assistant) Map() map[string]interface{} {
"type": ast.Type,
"name": ast.Name,
"readonly": ast.Readonly,
"public": ast.Public,
"share": ast.Share,
"avatar": ast.Avatar,
"connector": ast.Connector,
"path": ast.Path,
@ -43,14 +47,17 @@ func (ast *Assistant) Map() map[string]interface{} {
"description": ast.Description,
"options": ast.Options,
"prompts": ast.Prompts,
"kb": ast.KB,
"mcp": ast.MCP,
"tools": ast.Tools,
"workflow": ast.Workflow,
"tags": ast.Tags,
"mentionable": ast.Mentionable,
"automated": ast.Automated,
"placeholder": ast.Placeholder,
"locales": ast.Locales,
"created_at": timeToMySQLFormat(ast.CreatedAt),
"updated_at": timeToMySQLFormat(ast.UpdatedAt),
"created_at": store.ToMySQLTime(ast.CreatedAt),
"updated_at": store.ToMySQLTime(ast.UpdatedAt),
}
}
@ -97,20 +104,27 @@ func (ast *Assistant) Clone() *Assistant {
}
clone := &Assistant{
ID: ast.ID,
Type: ast.Type,
Name: ast.Name,
Avatar: ast.Avatar,
Connector: ast.Connector,
Path: ast.Path,
BuiltIn: ast.BuiltIn,
Sort: ast.Sort,
Description: ast.Description,
Readonly: ast.Readonly,
Mentionable: ast.Mentionable,
Automated: ast.Automated,
Script: ast.Script,
openai: ast.openai,
AssistantModel: store.AssistantModel{
ID: ast.ID,
Type: ast.Type,
Name: ast.Name,
Avatar: ast.Avatar,
Connector: ast.Connector,
Path: ast.Path,
BuiltIn: ast.BuiltIn,
Sort: ast.Sort,
Description: ast.Description,
Readonly: ast.Readonly,
Public: ast.Public,
Share: ast.Share,
Mentionable: ast.Mentionable,
Automated: ast.Automated,
CreatedAt: ast.CreatedAt,
UpdatedAt: ast.UpdatedAt,
},
Search: ast.Search,
Script: ast.Script,
openai: ast.openai,
}
// Deep copy tags
@ -119,6 +133,36 @@ func (ast *Assistant) Clone() *Assistant {
copy(clone.Tags, ast.Tags)
}
// Deep copy KB
if ast.KB != nil {
clone.KB = &store.KnowledgeBase{}
if ast.KB.Collections != nil {
clone.KB.Collections = make([]string, len(ast.KB.Collections))
copy(clone.KB.Collections, ast.KB.Collections)
}
if ast.KB.Options != nil {
clone.KB.Options = make(map[string]interface{})
for k, v := range ast.KB.Options {
clone.KB.Options[k] = v
}
}
}
// Deep copy MCP
if ast.MCP != nil {
clone.MCP = &store.MCPServers{}
if ast.MCP.Servers != nil {
clone.MCP.Servers = make([]string, len(ast.MCP.Servers))
copy(clone.MCP.Servers, ast.MCP.Servers)
}
if ast.MCP.Options != nil {
clone.MCP.Options = make(map[string]interface{})
for k, v := range ast.MCP.Options {
clone.MCP.Options[k] = v
}
}
}
// Deep copy options
if ast.Options != nil {
clone.Options = make(map[string]interface{})
@ -129,29 +173,66 @@ func (ast *Assistant) Clone() *Assistant {
// Deep copy prompts
if ast.Prompts != nil {
clone.Prompts = make([]Prompt, len(ast.Prompts))
clone.Prompts = make([]store.Prompt, len(ast.Prompts))
copy(clone.Prompts, ast.Prompts)
}
// Deep copy tools
if ast.Tools != nil {
clone.Tools = &ToolCalls{}
clone.Tools = &store.ToolCalls{}
if ast.Tools.Tools != nil {
clone.Tools.Tools = make([]Tool, len(ast.Tools.Tools))
clone.Tools.Tools = make([]store.Tool, len(ast.Tools.Tools))
copy(clone.Tools.Tools, ast.Tools.Tools)
}
if ast.Tools.Prompts != nil {
clone.Tools.Prompts = make([]Prompt, len(ast.Tools.Prompts))
clone.Tools.Prompts = make([]store.Prompt, len(ast.Tools.Prompts))
copy(clone.Tools.Prompts, ast.Tools.Prompts)
}
}
// Deep copy workflow
if ast.Workflow != nil {
clone.Workflow = make(map[string]interface{})
for k, v := range ast.Workflow {
clone.Workflow[k] = v
clone.Workflow = &store.Workflow{}
if ast.Workflow.Workflows != nil {
clone.Workflow.Workflows = make([]string, len(ast.Workflow.Workflows))
copy(clone.Workflow.Workflows, ast.Workflow.Workflows)
}
if ast.Workflow.Options != nil {
clone.Workflow.Options = make(map[string]interface{})
for k, v := range ast.Workflow.Options {
clone.Workflow.Options[k] = v
}
}
}
// Deep copy placeholder
if ast.Placeholder != nil {
clone.Placeholder = &store.Placeholder{
Title: ast.Placeholder.Title,
Description: ast.Placeholder.Description,
}
if ast.Placeholder.Prompts != nil {
clone.Placeholder.Prompts = make([]string, len(ast.Placeholder.Prompts))
copy(clone.Placeholder.Prompts, ast.Placeholder.Prompts)
}
}
// Deep copy locales
if ast.Locales != nil {
clone.Locales = make(i18n.Map)
for k, v := range ast.Locales {
// Deep copy messages
messages := make(map[string]any)
if v.Messages != nil {
for mk, mv := range v.Messages {
messages[mk] = mv
}
}
clone.Locales[k] = i18n.I18n{
Locale: v.Locale,
Messages: messages,
}
}
}
@ -179,13 +260,13 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
if v, has := data["tools"]; has {
switch tools := v.(type) {
case []Tool:
ast.Tools = &ToolCalls{
case []store.Tool:
ast.Tools = &store.ToolCalls{
Tools: tools,
Prompts: ast.Prompts,
}
case *ToolCalls:
case *store.ToolCalls:
ast.Tools = tools
default:
@ -193,7 +274,7 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
if err != nil {
return err
}
ast.Tools = &ToolCalls{}
ast.Tools = &store.ToolCalls{}
err = jsoniter.Unmarshal(raw, &ast.Tools)
if err != nil {
return err
@ -213,6 +294,15 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
if v, ok := data["automated"].(bool); ok {
ast.Automated = v
}
if v, ok := data["readonly"].(bool); ok {
ast.Readonly = v
}
if v, ok := data["public"].(bool); ok {
ast.Public = v
}
if v, ok := data["share"].(string); ok {
ast.Share = v
}
if v, ok := data["tags"].([]string); ok {
ast.Tags = v
}
@ -220,5 +310,32 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
ast.Options = v
}
// KB
if v, has := data["kb"]; has {
kb, err := store.ToKnowledgeBase(v)
if err != nil {
return err
}
ast.KB = kb
}
// MCP
if v, has := data["mcp"]; has {
mcp, err := store.ToMCPServers(v)
if err != nil {
return err
}
ast.MCP = mcp
}
// Workflow
if v, has := data["workflow"]; has {
workflow, err := store.ToWorkflow(v)
if err != nil {
return err
}
ast.Workflow = workflow
}
return ast.Validate()
}

79
agent/assistant/load.go
View file

@ -55,8 +55,7 @@ func LoadBuiltIn() error {
// Get all existing built-in assistants
for _, assistant := range res.Data {
assistantID := assistant["assistant_id"].(string)
deletedBuiltIn[assistantID] = true
deletedBuiltIn[assistant.ID] = true
}
}
@ -182,14 +181,14 @@ func LoadStore(id string) (*Assistant, error) {
return nil, fmt.Errorf("storage is not set")
}
data, err := storage.GetAssistant(id)
storeModel, err := storage.GetAssistant(id)
if err != nil {
return nil, err
}
// Load from path
if data["path"] != nil {
assistant, err = LoadPath(data["path"].(string))
if storeModel.Path != "" {
assistant, err = LoadPath(storeModel.Path)
if err != nil {
return nil, err
}
@ -197,8 +196,11 @@ func LoadStore(id string) (*Assistant, error) {
return assistant, nil
}
// Load from store
assistant, err = loadMap(data)
// Create assistant from store model
assistant = &Assistant{AssistantModel: *storeModel}
// Initialize the assistant
err = assistant.initialize()
if err != nil {
return nil, err
}
@ -345,7 +347,7 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
if err != nil {
return nil, err
}
assistant.Placeholder = &Placeholder{}
assistant.Placeholder = &store.Placeholder{}
err = jsoniter.Unmarshal(placeholder, assistant.Placeholder)
if err != nil {
return nil, err
@ -357,13 +359,13 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return nil, err
}
assistant.Placeholder = &Placeholder{}
assistant.Placeholder = &store.Placeholder{}
err = jsoniter.Unmarshal(raw, assistant.Placeholder)
if err != nil {
return nil, err
}
case *Placeholder:
case *store.Placeholder:
assistant.Placeholder = vv
case nil:
@ -386,6 +388,16 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Readonly = v
}
// Public
if v, ok := data["public"].(bool); ok {
assistant.Public = v
}
// Share
if v, ok := data["share"].(string); ok {
assistant.Share = v
}
// built_in
if v, ok := data["built_in"].(bool); ok {
assistant.BuiltIn = v
@ -470,11 +482,11 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
if prompts, has := data["prompts"]; has {
switch v := prompts.(type) {
case []Prompt:
case []store.Prompt:
assistant.Prompts = v
case string:
var prompts []Prompt
var prompts []store.Prompt
err := yaml.Unmarshal([]byte(v), &prompts)
if err != nil {
return nil, err
@ -487,7 +499,7 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return nil, err
}
var prompts []Prompt
var prompts []store.Prompt
err = jsoniter.Unmarshal(raw, &prompts)
if err != nil {
return nil, err
@ -499,13 +511,13 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
// tools
if tools, has := data["tools"]; has {
switch vv := tools.(type) {
case []Tool:
assistant.Tools = &ToolCalls{
case []store.Tool:
assistant.Tools = &store.ToolCalls{
Tools: vv,
Prompts: assistant.Prompts,
}
case ToolCalls:
case store.ToolCalls:
assistant.Tools = &vv
default:
@ -514,7 +526,7 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return nil, fmt.Errorf("tools format error %s", err.Error())
}
var tools ToolCalls
var tools store.ToolCalls
err = jsoniter.Unmarshal(raw, &tools)
if err != nil {
return nil, fmt.Errorf("tools format error %s", err.Error())
@ -523,6 +535,33 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
}
// kb
if kb, has := data["kb"]; has {
knowledgeBase, err := store.ToKnowledgeBase(kb)
if err != nil {
return nil, err
}
assistant.KB = knowledgeBase
}
// mcp
if mcp, has := data["mcp"]; has {
mcpServers, err := store.ToMCPServers(mcp)
if err != nil {
return nil, err
}
assistant.MCP = mcpServers
}
// workflow
if workflow, has := data["workflow"]; has {
wf, err := store.ToWorkflow(workflow)
if err != nil {
return nil, err
}
assistant.Workflow = wf
}
// script
if data["script"] != nil {
switch v := data["script"].(type) {
@ -668,7 +707,7 @@ func (ast *Assistant) initialize() error {
return nil
}
func loadTools(file string) (*ToolCalls, int64, error) {
func loadTools(file string) (*store.ToolCalls, int64, error) {
app, err := fs.Get("app")
if err != nil {
@ -686,10 +725,10 @@ func loadTools(file string) (*ToolCalls, int64, error) {
}
if len(content) == 0 {
return &ToolCalls{Tools: []Tool{}, Prompts: []Prompt{}}, ts.UnixNano(), nil
return &store.ToolCalls{Tools: []store.Tool{}, Prompts: []store.Prompt{}}, ts.UnixNano(), nil
}
var tools ToolCalls
var tools store.ToolCalls
err = application.Parse(file, content, &tools)
if err != nil {
return nil, 0, err

59
agent/assistant/tool.go
View file

@ -4,6 +4,7 @@ import (
"fmt"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/yao/agent/store"
)
// Tool represents a tool
@ -95,3 +96,61 @@ func generateExampleValue(name string, prop SchemaProperty) interface{} {
return fmt.Sprintf("<%s>", name)
}
}
// ToRuntimeTool converts store.Tool to assistant.Tool (OpenAI format)
func ToRuntimeTool(storeTool store.Tool) (Tool, error) {
var tool Tool
// Marshal and unmarshal to convert between formats
raw, err := jsoniter.Marshal(storeTool)
if err != nil {
return tool, fmt.Errorf("failed to marshal store tool: %w", err)
}
// Try to unmarshal as OpenAI format first
err = jsoniter.Unmarshal(raw, &tool)
if err == nil && tool.Function.Name != "" {
return tool, nil
}
// If it's a simple format, convert it
tool.Type = "function"
if storeTool.Type != "" {
tool.Type = storeTool.Type
}
tool.Function.Name = storeTool.Name
tool.Function.Description = storeTool.Description
// Convert parameters
if storeTool.Parameters != nil {
raw, err := jsoniter.Marshal(storeTool.Parameters)
if err != nil {
return tool, fmt.Errorf("failed to marshal parameters: %w", err)
}
var params Parameter
err = jsoniter.Unmarshal(raw, &params)
if err != nil {
return tool, fmt.Errorf("failed to unmarshal parameters: %w", err)
}
tool.Function.Parameters = params
}
return tool, nil
}
// ToRuntimeTools converts []store.Tool to []assistant.Tool
func ToRuntimeTools(storeTools []store.Tool) ([]Tool, error) {
if storeTools == nil {
return nil, nil
}
tools := make([]Tool, 0, len(storeTools))
for _, storeTool := range storeTools {
tool, err := ToRuntimeTool(storeTool)
if err != nil {
return nil, err
}
tools = append(tools, tool)
}
return tools, nil
}

54
agent/assistant/types.go
View file

@ -7,8 +7,8 @@ import (
"github.com/gin-gonic/gin"
v8 "github.com/yaoapp/gou/runtime/v8"
chatctx "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/agent/message"
"github.com/yaoapp/yao/agent/store"
api "github.com/yaoapp/yao/openai"
)
@ -24,7 +24,7 @@ type API interface {
// Download(ctx context.Context, fileID string) (*FileResponse, error)
// ReadBase64(ctx context.Context, fileID string) (string, error)
GetPlaceholder(locale string) *Placeholder
GetPlaceholder(locale string) *store.Placeholder
Execute(c *gin.Context, ctx chatctx.Context, input interface{}, options map[string]interface{}, callback ...interface{}) (interface{}, error)
Call(c *gin.Context, payload APIPayload) (interface{}, error)
}
@ -98,43 +98,18 @@ type QueryParam struct {
// Assistant the assistant
type Assistant struct {
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
Path string `json:"path,omitempty"` // Assistant Path
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description
Tags []string `json:"tags,omitempty"` // Assistant Tags
Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly
Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
Tools *ToolCalls `json:"tools,omitempty"` // Assistant Tools
Workflow map[string]interface{} `json:"workflow,omitempty"` // Assistant Workflow
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Search *SearchOption `json:"search,omitempty" yaml:"search,omitempty"` // Whether this assistant supports search
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script
store.AssistantModel
Search *SearchOption `json:"search,omitempty" yaml:"search,omitempty"` // Whether this assistant supports search
Script *v8.Script `json:"-" yaml:"-"` // Assistant Script
// Internal
// ===============================
openai *api.OpenAI // OpenAI API
search bool // Whether this assistant supports search
vision bool // Whether this assistant supports vision
toolCalls bool // Whether this assistant supports tool_calls
initHook bool // Whether this assistant has an init hook
}
// ToolCalls the tool calls
type ToolCalls struct {
Tools []Tool `json:"tools,omitempty"`
Prompts []Prompt `json:"prompts,omitempty"`
openai *api.OpenAI // OpenAI API
search bool // Whether this assistant supports search
vision bool // Whether this assistant supports vision
toolCalls bool // Whether this assistant supports tool_calls
initHook bool // Whether this assistant has an init hook
runtimeTools []Tool // Converted tools for business logic (OpenAI format)
}
// ConnectorSetting the connector setting
@ -143,13 +118,6 @@ type ConnectorSetting struct {
Tools bool `json:"tools,omitempty" yaml:"tools,omitempty"`
}
// Placeholder the assistant placeholder
type Placeholder struct {
Title string `json:"title,omitempty"`
Description string `json:"description,omitempty"`
Prompts []string `json:"prompts,omitempty"`
}
// VisionCapableModels list of LLM models that support vision capabilities
var VisionCapableModels = map[string]bool{
// OpenAI Models

11
agent/assistant/utils.go
View file

@ -43,17 +43,6 @@ func getTimestamp(v interface{}) (int64, error) {
return 0, fmt.Errorf("invalid timestamp type %T", v)
}
func stringToTimestamp(v string) (int64, error) {
return strconv.ParseInt(v, 10, 64)
}
func timeToMySQLFormat(ts int64) string {
if ts == 0 {
return "0000-00-00 00:00:00"
}
return time.Unix(ts/1e9, ts%1e9).Format("2006-01-02 15:04:05")
}
// stringHash returns the sha256 hash of the string
func stringHash(v string) string {
h := sha256.New()

8
agent/load.go
View file

@ -63,8 +63,8 @@ func Load(cfg config.Config) error {
return err
}
// Initialize Connectors
err = initConnectors()
// Initialize Connector settings
err = initConnectorSettings()
if err != nil {
return err
}
@ -95,7 +95,7 @@ func initGlobalI18n() error {
}
// initConnectors initialize the connectors
func initConnectors() error {
func initConnectorSettings() error {
path := filepath.Join("agent", "connectors.yml")
if exists, _ := application.App.Exists(path); !exists {
return nil
@ -129,7 +129,7 @@ func initStore() error {
// other connector
conn, err := connector.Select(Agent.StoreSetting.Connector)
if err != nil {
return err
return fmt.Errorf("load connectors error: %s", err.Error())
}
if conn.Is(connector.DATABASE) {

16
agent/process.go
View file

@ -79,7 +79,13 @@ func processAssistantCreate(process *process.Process) interface{} {
exception.New("Agent store is not initialized", 500).Throw()
}
id, err := agent.Store.SaveAssistant(data)
// Convert to AssistantModel
model, err := store.ToAssistantModel(data)
if err != nil {
exception.New("Invalid assistant data: %s", 400, err.Error()).Throw()
}
id, err := agent.Store.SaveAssistant(model)
if err != nil {
exception.New("Failed to create assistant: %s", 500, err.Error()).Throw()
}
@ -97,7 +103,13 @@ func processAssistantSave(process *process.Process) interface{} {
exception.New("Agent store is not initialized", 500).Throw()
}
id, err := agent.Store.SaveAssistant(data)
// Convert to AssistantModel
model, err := store.ToAssistantModel(data)
if err != nil {
exception.New("Invalid assistant data: %s", 400, err.Error()).Throw()
}
id, err := agent.Store.SaveAssistant(model)
if err != nil {
exception.New("Failed to save assistant: %s", 500, err.Error()).Throw()
}

97
agent/store/Interfaces.go Normal file
View file

@ -0,0 +1,97 @@
package store
// Store defines the conversation storage interface
// Provides basic operations required for conversation management
type Store interface {
// GetChats retrieves a list of chats
// sid: Session ID
// filter: Filter conditions
// Returns: Grouped chat list and potential error
GetChats(sid string, filter ChatFilter, locale ...string) (*ChatGroupResponse, error)
// GetChat retrieves a single chat's information
// sid: Session ID
// cid: Chat ID
// Returns: Chat information and potential error
GetChat(sid string, cid string, locale ...string) (*ChatInfo, error)
// GetChatWithFilter retrieves a single chat's information with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: Chat information and potential error
GetChatWithFilter(sid string, cid string, filter ChatFilter, locale ...string) (*ChatInfo, error)
// GetHistory retrieves chat history
// sid: Session ID
// cid: Chat ID
// Returns: History record list and potential error
GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error)
// GetHistoryWithFilter retrieves chat history with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: History record list and potential error
GetHistoryWithFilter(sid string, cid string, filter ChatFilter, locale ...string) ([]map[string]interface{}, error)
// SaveHistory saves chat history
// sid: Session ID
// messages: Message list
// cid: Chat ID
// context: Context information
// Returns: Potential error
SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error
// DeleteChat deletes a single chat
// sid: Session ID
// cid: Chat ID
// Returns: Potential error
DeleteChat(sid string, cid string) error
// DeleteAllChats deletes all chats
// sid: Session ID
// Returns: Potential error
DeleteAllChats(sid string) error
// UpdateChatTitle updates chat title
// sid: Session ID
// cid: Chat ID
// title: New title
// Returns: Potential error
UpdateChatTitle(sid string, cid string, title string) error
// SaveAssistant saves assistant information
// assistant: Assistant information
// Returns: Assistant ID and potential error
SaveAssistant(assistant *AssistantModel) (string, error)
// DeleteAssistant deletes an assistant
// assistantID: Assistant ID
// Returns: Potential error
DeleteAssistant(assistantID string) error
// GetAssistants retrieves a paginated list of assistants with filtering
// filter: Filter conditions for querying assistants
// locale: Optional locale for i18n translations
// Returns: Paginated assistant list and potential error
GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error)
// GetAssistantTags retrieves all unique tags from assistants
// Returns: List of tags and potential error
GetAssistantTags(locale ...string) ([]Tag, error)
// GetAssistant retrieves a single assistant by ID
// assistantID: Assistant ID
// Returns: Assistant information and potential error
GetAssistant(assistantID string, locale ...string) (*AssistantModel, error)
// DeleteAssistants deletes assistants based on filter conditions
// filter: Filter conditions
// Returns: Number of deleted records and potential error
DeleteAssistants(filter AssistantFilter) (int64, error)
// Close closes the store and releases any resources
// Returns: Potential error
Close() error
}

362
agent/store/convert.go Normal file
View file

@ -0,0 +1,362 @@
package store
import (
"fmt"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/spf13/cast"
"github.com/yaoapp/yao/agent/i18n"
)
// ToKnowledgeBase converts various types to KnowledgeBase
func ToKnowledgeBase(v interface{}) (*KnowledgeBase, error) {
if v == nil {
return nil, nil
}
switch kb := v.(type) {
case *KnowledgeBase:
return kb, nil
case KnowledgeBase:
return &kb, nil
case []string:
return &KnowledgeBase{Collections: kb}, nil
case []interface{}:
var collections []string
for _, item := range kb {
collections = append(collections, cast.ToString(item))
}
return &KnowledgeBase{Collections: collections}, nil
default:
raw, err := jsoniter.Marshal(kb)
if err != nil {
return nil, fmt.Errorf("kb format error: %s", err.Error())
}
var knowledgeBase KnowledgeBase
err = jsoniter.Unmarshal(raw, &knowledgeBase)
if err != nil {
return nil, fmt.Errorf("kb format error: %s", err.Error())
}
return &knowledgeBase, nil
}
}
// ToMCPServers converts various types to MCPServers
func ToMCPServers(v interface{}) (*MCPServers, error) {
if v == nil {
return nil, nil
}
switch mcp := v.(type) {
case *MCPServers:
return mcp, nil
case MCPServers:
return &mcp, nil
case []string:
return &MCPServers{Servers: mcp}, nil
case []interface{}:
var servers []string
for _, item := range mcp {
servers = append(servers, cast.ToString(item))
}
return &MCPServers{Servers: servers}, nil
default:
raw, err := jsoniter.Marshal(mcp)
if err != nil {
return nil, fmt.Errorf("mcp format error: %s", err.Error())
}
var mcpServers MCPServers
err = jsoniter.Unmarshal(raw, &mcpServers)
if err != nil {
return nil, fmt.Errorf("mcp format error: %s", err.Error())
}
return &mcpServers, nil
}
}
// ToWorkflow converts various types to Workflow
func ToWorkflow(v interface{}) (*Workflow, error) {
if v == nil {
return nil, nil
}
switch workflow := v.(type) {
case *Workflow:
return workflow, nil
case Workflow:
return &workflow, nil
case []string:
return &Workflow{Workflows: workflow}, nil
case []interface{}:
var workflows []string
for _, item := range workflow {
workflows = append(workflows, cast.ToString(item))
}
return &Workflow{Workflows: workflows}, nil
default:
raw, err := jsoniter.Marshal(workflow)
if err != nil {
return nil, fmt.Errorf("workflow format error: %s", err.Error())
}
var wf Workflow
err = jsoniter.Unmarshal(raw, &wf)
if err != nil {
return nil, fmt.Errorf("workflow format error: %s", err.Error())
}
return &wf, nil
}
}
// ToMySQLTime converts various types to MySQL datetime format
func ToMySQLTime(v interface{}) string {
switch val := v.(type) {
case int64:
if val == 0 {
return "0000-00-00 00:00:00"
}
return time.Unix(val/1e9, val%1e9).Format("2006-01-02 15:04:05")
case int:
if val == 0 {
return "0000-00-00 00:00:00"
}
return time.Unix(int64(val)/1e9, int64(val)%1e9).Format("2006-01-02 15:04:05")
case string:
// If already in MySQL format, return as-is
if _, err := time.Parse("2006-01-02 15:04:05", val); err == nil {
return val
}
// Try RFC3339 format
if ts, err := time.Parse(time.RFC3339, val); err == nil {
return ts.Format("2006-01-02 15:04:05")
}
// Try parsing as Unix timestamp
if ts, err := cast.ToInt64E(val); err == nil {
if ts == 0 {
return "0000-00-00 00:00:00"
}
return time.Unix(ts/1e9, ts%1e9).Format("2006-01-02 15:04:05")
}
return val
case time.Time:
if val.IsZero() {
return "0000-00-00 00:00:00"
}
return val.Format("2006-01-02 15:04:05")
case nil:
return "0000-00-00 00:00:00"
default:
return "0000-00-00 00:00:00"
}
}
// ToAssistantModel converts various types to AssistantModel
func ToAssistantModel(v interface{}) (*AssistantModel, error) {
if v == nil {
return nil, nil
}
// If already an AssistantModel, return it
switch model := v.(type) {
case *AssistantModel:
return model, nil
case AssistantModel:
return &model, nil
}
// Convert to map first if needed
var data map[string]interface{}
switch v := v.(type) {
case map[string]interface{}:
data = v
default:
// Try to marshal and unmarshal
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, fmt.Errorf("failed to marshal to AssistantModel: %w", err)
}
err = jsoniter.Unmarshal(raw, &data)
if err != nil {
return nil, fmt.Errorf("failed to unmarshal to map: %w", err)
}
}
model := &AssistantModel{}
// Basic string fields
if id, ok := data["assistant_id"].(string); ok {
model.ID = id
}
if typ, ok := data["type"].(string); ok {
model.Type = typ
}
if name, ok := data["name"].(string); ok {
model.Name = name
}
if avatar, ok := data["avatar"].(string); ok {
model.Avatar = avatar
}
if connector, ok := data["connector"].(string); ok {
model.Connector = connector
}
if path, ok := data["path"].(string); ok {
model.Path = path
}
if description, ok := data["description"].(string); ok {
model.Description = description
}
if share, ok := data["share"].(string); ok {
model.Share = share
}
// Boolean fields (handle both bool and int types from database)
model.BuiltIn = getBoolValue(data, "built_in")
model.Readonly = getBoolValue(data, "readonly")
model.Public = getBoolValue(data, "public")
model.Mentionable = getBoolValue(data, "mentionable")
model.Automated = getBoolValue(data, "automated")
// Integer fields
if sort, ok := data["sort"].(int); ok {
model.Sort = sort
} else if sort, ok := data["sort"].(float64); ok {
model.Sort = int(sort)
}
if createdAt, ok := data["created_at"].(int64); ok {
model.CreatedAt = createdAt
} else if createdAt, ok := data["created_at"].(float64); ok {
model.CreatedAt = int64(createdAt)
}
if updatedAt, ok := data["updated_at"].(int64); ok {
model.UpdatedAt = updatedAt
} else if updatedAt, ok := data["updated_at"].(float64); ok {
model.UpdatedAt = int64(updatedAt)
}
// Tags (string array)
if tags, ok := data["tags"]; ok && tags != nil {
raw, err := jsoniter.Marshal(tags)
if err == nil {
var t []string
if err := jsoniter.Unmarshal(raw, &t); err == nil {
model.Tags = t
}
}
}
// Options (map)
if options, ok := data["options"].(map[string]interface{}); ok {
model.Options = options
}
// Prompts
if prompts, ok := data["prompts"]; ok && prompts != nil {
raw, err := jsoniter.Marshal(prompts)
if err == nil {
var p []Prompt
if err := jsoniter.Unmarshal(raw, &p); err == nil {
model.Prompts = p
}
}
}
// KB
if kb, ok := data["kb"]; ok && kb != nil {
kbConverted, err := ToKnowledgeBase(kb)
if err == nil {
model.KB = kbConverted
}
}
// MCP
if mcp, ok := data["mcp"]; ok && mcp != nil {
mcpConverted, err := ToMCPServers(mcp)
if err == nil {
model.MCP = mcpConverted
}
}
// Workflow
if workflow, ok := data["workflow"]; ok && workflow != nil {
wf, err := ToWorkflow(workflow)
if err == nil {
model.Workflow = wf
}
}
// Tools
if tools, ok := data["tools"]; ok && tools != nil {
raw, err := jsoniter.Marshal(tools)
if err == nil {
var tc ToolCalls
if err := jsoniter.Unmarshal(raw, &tc); err == nil {
model.Tools = &tc
}
}
}
// Placeholder
if placeholder, ok := data["placeholder"]; ok && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)
if err == nil {
var ph Placeholder
if err := jsoniter.Unmarshal(raw, &ph); err == nil {
model.Placeholder = &ph
}
}
}
// Locales
if locales, ok := data["locales"]; ok && 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
}
// getBoolValue extracts a boolean value from a map, handling both bool and numeric types
func getBoolValue(data map[string]interface{}, key string) bool {
if v, ok := data[key]; ok && v != nil {
switch val := v.(type) {
case bool:
return val
case int:
return val != 0
case int64:
return val != 0
case float64:
return val != 0
case string:
return val == "true" || val == "1"
}
}
return false
}

869
agent/store/convert_test.go Normal file
View file

@ -0,0 +1,869 @@
package store
import (
"testing"
"time"
)
// TestToKnowledgeBase tests the ToKnowledgeBase conversion function
func TestToKnowledgeBase(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToKnowledgeBase(nil)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != nil {
t.Errorf("Expected nil result, got: %v", result)
}
})
t.Run("KnowledgeBasePointer", func(t *testing.T) {
kb := &KnowledgeBase{Collections: []string{"col1", "col2"}}
result, err := ToKnowledgeBase(kb)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != kb {
t.Errorf("Expected same pointer")
}
})
t.Run("KnowledgeBaseValue", func(t *testing.T) {
kb := KnowledgeBase{Collections: []string{"col1", "col2"}}
result, err := ToKnowledgeBase(kb)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Collections) != 2 {
t.Errorf("Expected 2 collections, got %d", len(result.Collections))
}
})
t.Run("StringSlice", func(t *testing.T) {
collections := []string{"col1", "col2", "col3"}
result, err := ToKnowledgeBase(collections)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Collections) != 3 {
t.Errorf("Expected 3 collections, got %d", len(result.Collections))
}
if result.Collections[0] != "col1" {
t.Errorf("Expected 'col1', got '%s'", result.Collections[0])
}
})
t.Run("InterfaceSlice", func(t *testing.T) {
collections := []interface{}{"col1", "col2", 123}
result, err := ToKnowledgeBase(collections)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Collections) != 3 {
t.Errorf("Expected 3 collections, got %d", len(result.Collections))
}
if result.Collections[2] != "123" {
t.Errorf("Expected '123', got '%s'", result.Collections[2])
}
})
t.Run("MapInput", func(t *testing.T) {
data := map[string]interface{}{
"collections": []string{"col1", "col2"},
}
result, err := ToKnowledgeBase(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Collections) != 2 {
t.Errorf("Expected 2 collections, got %d", len(result.Collections))
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToKnowledgeBase(invalidData)
if err == nil {
t.Error("Expected error for invalid input")
}
})
t.Run("InvalidJSONUnmarshal", func(t *testing.T) {
// Test with data that marshals but can't unmarshal to KnowledgeBase
data := map[string]interface{}{
"invalid_field": "should cause unmarshal to fail gracefully",
}
result, err := ToKnowledgeBase(data)
// Should not error, just return empty KnowledgeBase
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
})
}
// TestToMCPServers tests the ToMCPServers conversion function
func TestToMCPServers(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToMCPServers(nil)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != nil {
t.Errorf("Expected nil result, got: %v", result)
}
})
t.Run("MCPServersPointer", func(t *testing.T) {
mcp := &MCPServers{Servers: []string{"server1", "server2"}}
result, err := ToMCPServers(mcp)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != mcp {
t.Errorf("Expected same pointer")
}
})
t.Run("MCPServersValue", func(t *testing.T) {
mcp := MCPServers{Servers: []string{"server1", "server2"}}
result, err := ToMCPServers(mcp)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Servers) != 2 {
t.Errorf("Expected 2 servers, got %d", len(result.Servers))
}
})
t.Run("StringSlice", func(t *testing.T) {
servers := []string{"server1", "server2", "server3"}
result, err := ToMCPServers(servers)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Servers) != 3 {
t.Errorf("Expected 3 servers, got %d", len(result.Servers))
}
if result.Servers[0] != "server1" {
t.Errorf("Expected 'server1', got '%s'", result.Servers[0])
}
})
t.Run("InterfaceSlice", func(t *testing.T) {
servers := []interface{}{"server1", "server2", 456}
result, err := ToMCPServers(servers)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Servers) != 3 {
t.Errorf("Expected 3 servers, got %d", len(result.Servers))
}
if result.Servers[2] != "456" {
t.Errorf("Expected '456', got '%s'", result.Servers[2])
}
})
t.Run("MapInput", func(t *testing.T) {
data := map[string]interface{}{
"servers": []string{"server1", "server2"},
}
result, err := ToMCPServers(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Servers) != 2 {
t.Errorf("Expected 2 servers, got %d", len(result.Servers))
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToMCPServers(invalidData)
if err == nil {
t.Error("Expected error for invalid input")
}
})
t.Run("InvalidJSONUnmarshal", func(t *testing.T) {
// Test with data that marshals but can't unmarshal to MCPServers
data := map[string]interface{}{
"invalid_field": "should cause unmarshal to fail gracefully",
}
result, err := ToMCPServers(data)
// Should not error, just return empty MCPServers
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
})
}
// TestToWorkflow tests the ToWorkflow conversion function
func TestToWorkflow(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToWorkflow(nil)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != nil {
t.Errorf("Expected nil result, got: %v", result)
}
})
t.Run("WorkflowPointer", func(t *testing.T) {
wf := &Workflow{Workflows: []string{"wf1", "wf2"}}
result, err := ToWorkflow(wf)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != wf {
t.Errorf("Expected same pointer")
}
})
t.Run("WorkflowValue", func(t *testing.T) {
wf := Workflow{Workflows: []string{"wf1", "wf2"}}
result, err := ToWorkflow(wf)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Workflows) != 2 {
t.Errorf("Expected 2 workflows, got %d", len(result.Workflows))
}
})
t.Run("StringSlice", func(t *testing.T) {
workflows := []string{"wf1", "wf2", "wf3"}
result, err := ToWorkflow(workflows)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Workflows) != 3 {
t.Errorf("Expected 3 workflows, got %d", len(result.Workflows))
}
if result.Workflows[0] != "wf1" {
t.Errorf("Expected 'wf1', got '%s'", result.Workflows[0])
}
})
t.Run("InterfaceSlice", func(t *testing.T) {
workflows := []interface{}{"wf1", "wf2", 789}
result, err := ToWorkflow(workflows)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Workflows) != 3 {
t.Errorf("Expected 3 workflows, got %d", len(result.Workflows))
}
if result.Workflows[2] != "789" {
t.Errorf("Expected '789', got '%s'", result.Workflows[2])
}
})
t.Run("MapInput", func(t *testing.T) {
data := map[string]interface{}{
"workflows": []string{"wf1", "wf2"},
}
result, err := ToWorkflow(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Workflows) != 2 {
t.Errorf("Expected 2 workflows, got %d", len(result.Workflows))
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToWorkflow(invalidData)
if err == nil {
t.Error("Expected error for invalid input")
}
})
t.Run("InvalidJSONUnmarshal", func(t *testing.T) {
// Test with data that marshals but can't unmarshal to Workflow
data := map[string]interface{}{
"invalid_field": "should cause unmarshal to fail gracefully",
}
result, err := ToWorkflow(data)
// Should not error, just return empty Workflow
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
})
}
// TestToMySQLTime tests the ToMySQLTime conversion function
func TestToMySQLTime(t *testing.T) {
t.Run("Int64Zero", func(t *testing.T) {
result := ToMySQLTime(int64(0))
if result != "0000-00-00 00:00:00" {
t.Errorf("Expected '0000-00-00 00:00:00', got '%s'", result)
}
})
t.Run("Int64Timestamp", func(t *testing.T) {
// Unix timestamp in nanoseconds: 1609459200000000000 = 2021-01-01 00:00:00 UTC
timestamp := int64(1609459200000000000)
result := ToMySQLTime(timestamp)
// Should be in format "2021-01-01 00:00:00" or similar depending on timezone
if len(result) != 19 {
t.Errorf("Expected 19 character timestamp, got %d: '%s'", len(result), result)
}
})
t.Run("IntZero", func(t *testing.T) {
result := ToMySQLTime(int(0))
if result != "0000-00-00 00:00:00" {
t.Errorf("Expected '0000-00-00 00:00:00', got '%s'", result)
}
})
t.Run("IntTimestamp", func(t *testing.T) {
timestamp := int(1609459200000000000)
result := ToMySQLTime(timestamp)
if len(result) != 19 {
t.Errorf("Expected 19 character timestamp, got %d: '%s'", len(result), result)
}
})
t.Run("StringMySQLFormat", func(t *testing.T) {
mysqlTime := "2021-01-01 12:30:45"
result := ToMySQLTime(mysqlTime)
if result != mysqlTime {
t.Errorf("Expected '%s', got '%s'", mysqlTime, result)
}
})
t.Run("StringRFC3339", func(t *testing.T) {
rfc3339Time := "2021-01-01T12:30:45Z"
result := ToMySQLTime(rfc3339Time)
expected := "2021-01-01 12:30:45"
if result != expected {
t.Errorf("Expected '%s', got '%s'", expected, result)
}
})
t.Run("StringUnixTimestamp", func(t *testing.T) {
// Unix timestamp in seconds as string
result := ToMySQLTime("1609459200000000000")
if len(result) != 19 {
t.Errorf("Expected 19 character timestamp, got %d: '%s'", len(result), result)
}
})
t.Run("StringInvalidFormat", func(t *testing.T) {
invalidTime := "not-a-valid-time"
result := ToMySQLTime(invalidTime)
// Should return the original string when it can't be parsed
if result != invalidTime {
t.Errorf("Expected '%s', got '%s'", invalidTime, result)
}
})
t.Run("TimeZero", func(t *testing.T) {
zeroTime := time.Time{}
result := ToMySQLTime(zeroTime)
if result != "0000-00-00 00:00:00" {
t.Errorf("Expected '0000-00-00 00:00:00', got '%s'", result)
}
})
t.Run("TimeNormal", func(t *testing.T) {
normalTime := time.Date(2021, 1, 1, 12, 30, 45, 0, time.UTC)
result := ToMySQLTime(normalTime)
expected := "2021-01-01 12:30:45"
if result != expected {
t.Errorf("Expected '%s', got '%s'", expected, result)
}
})
t.Run("NilInput", func(t *testing.T) {
result := ToMySQLTime(nil)
if result != "0000-00-00 00:00:00" {
t.Errorf("Expected '0000-00-00 00:00:00', got '%s'", result)
}
})
t.Run("UnknownType", func(t *testing.T) {
// Test with unsupported type
result := ToMySQLTime(struct{}{})
if result != "0000-00-00 00:00:00" {
t.Errorf("Expected '0000-00-00 00:00:00', got '%s'", result)
}
})
}
// TestToAssistantModel tests the ToAssistantModel conversion function
func TestToAssistantModel(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToAssistantModel(nil)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != nil {
t.Errorf("Expected nil result, got: %v", result)
}
})
t.Run("AssistantModelPointer", func(t *testing.T) {
model := &AssistantModel{
ID: "test-id",
Name: "Test Assistant",
}
result, err := ToAssistantModel(model)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != model {
t.Errorf("Expected same pointer")
}
})
t.Run("AssistantModelValue", func(t *testing.T) {
model := AssistantModel{
ID: "test-id",
Name: "Test Assistant",
}
result, err := ToAssistantModel(model)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ID != "test-id" {
t.Errorf("Expected 'test-id', got '%s'", result.ID)
}
})
t.Run("MapWithAllFields", func(t *testing.T) {
data := map[string]interface{}{
"assistant_id": "test-id",
"type": "assistant",
"name": "Test Assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"path": "/path/to/assistant",
"description": "Test description",
"share": "team",
"built_in": true,
"readonly": false,
"public": true,
"mentionable": true,
"automated": false,
"sort": 100,
"created_at": int64(1609459200),
"updated_at": int64(1609459300),
"tags": []string{"tag1", "tag2"},
"options": map[string]interface{}{
"temperature": 0.7,
},
"prompts": []map[string]interface{}{
{"role": "system", "content": "You are helpful"},
},
"kb": map[string]interface{}{
"collections": []string{"col1"},
},
"mcp": map[string]interface{}{
"servers": []string{"server1"},
},
"workflow": map[string]interface{}{
"workflows": []string{"wf1"},
},
"tools": map[string]interface{}{
"calls": []string{"tool1"},
},
"placeholder": map[string]interface{}{
"title": "Enter message",
},
"locales": map[string]interface{}{
"en": map[string]interface{}{
"name": "English Name",
},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
// Verify all fields
if result.ID != "test-id" {
t.Errorf("Expected ID 'test-id', got '%s'", result.ID)
}
if result.Type != "assistant" {
t.Errorf("Expected Type 'assistant', got '%s'", result.Type)
}
if result.Name != "Test Assistant" {
t.Errorf("Expected Name 'Test Assistant', got '%s'", result.Name)
}
if result.Avatar != "https://example.com/avatar.png" {
t.Errorf("Expected Avatar URL, got '%s'", result.Avatar)
}
if result.Connector != "openai" {
t.Errorf("Expected Connector 'openai', got '%s'", result.Connector)
}
if result.Path != "/path/to/assistant" {
t.Errorf("Expected Path, got '%s'", result.Path)
}
if result.Description != "Test description" {
t.Errorf("Expected Description, got '%s'", result.Description)
}
if result.Share != "team" {
t.Errorf("Expected Share 'team', got '%s'", result.Share)
}
if !result.BuiltIn {
t.Error("Expected BuiltIn to be true")
}
if result.Readonly {
t.Error("Expected Readonly to be false")
}
if !result.Public {
t.Error("Expected Public to be true")
}
if !result.Mentionable {
t.Error("Expected Mentionable to be true")
}
if result.Automated {
t.Error("Expected Automated to be false")
}
if result.Sort != 100 {
t.Errorf("Expected Sort 100, got %d", result.Sort)
}
if result.CreatedAt != 1609459200 {
t.Errorf("Expected CreatedAt 1609459200, got %d", result.CreatedAt)
}
if result.UpdatedAt != 1609459300 {
t.Errorf("Expected UpdatedAt 1609459300, got %d", result.UpdatedAt)
}
if len(result.Tags) != 2 {
t.Errorf("Expected 2 tags, got %d", len(result.Tags))
}
if result.Options == nil {
t.Error("Expected Options to be set")
}
if len(result.Prompts) != 1 {
t.Errorf("Expected 1 prompt, got %d", len(result.Prompts))
}
if result.KB == nil {
t.Error("Expected KB to be set")
}
if result.MCP == nil {
t.Error("Expected MCP to be set")
}
if result.Workflow == nil {
t.Error("Expected Workflow to be set")
}
if result.Tools == nil {
t.Error("Expected Tools to be set")
}
if result.Placeholder == nil {
t.Error("Expected Placeholder to be set")
}
if result.Locales == nil {
t.Error("Expected Locales to be set")
}
})
t.Run("MapWithFloatNumbers", func(t *testing.T) {
data := map[string]interface{}{
"sort": float64(150),
"created_at": float64(1609459200),
"updated_at": float64(1609459300),
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.Sort != 150 {
t.Errorf("Expected Sort 150, got %d", result.Sort)
}
if result.CreatedAt != 1609459200 {
t.Errorf("Expected CreatedAt 1609459200, got %d", result.CreatedAt)
}
if result.UpdatedAt != 1609459300 {
t.Errorf("Expected UpdatedAt 1609459300, got %d", result.UpdatedAt)
}
})
t.Run("MapWithNilFields", func(t *testing.T) {
data := map[string]interface{}{
"assistant_id": "test-id",
"tags": nil,
"options": nil,
"prompts": nil,
"kb": nil,
"mcp": nil,
"workflow": nil,
"tools": nil,
"placeholder": nil,
"locales": nil,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ID != "test-id" {
t.Errorf("Expected ID 'test-id', got '%s'", result.ID)
}
// All nil fields should remain nil
if result.Tags != nil {
t.Error("Expected Tags to be nil")
}
if result.Options != nil {
t.Error("Expected Options to be nil")
}
})
t.Run("StructInput", func(t *testing.T) {
type CustomStruct struct {
AssistantID string `json:"assistant_id"`
Name string `json:"name"`
Type string `json:"type"`
}
input := CustomStruct{
AssistantID: "custom-id",
Name: "Custom Assistant",
Type: "bot",
}
result, err := ToAssistantModel(input)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ID != "custom-id" {
t.Errorf("Expected ID 'custom-id', got '%s'", result.ID)
}
if result.Name != "Custom Assistant" {
t.Errorf("Expected Name 'Custom Assistant', got '%s'", result.Name)
}
if result.Type != "bot" {
t.Errorf("Expected Type 'bot', got '%s'", result.Type)
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToAssistantModel(invalidData)
if err == nil {
t.Error("Expected error for invalid input")
}
})
t.Run("EmptyMap", func(t *testing.T) {
data := map[string]interface{}{}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
// All fields should have default values
if result.ID != "" {
t.Errorf("Expected empty ID, got '%s'", result.ID)
}
})
}
// TestToAssistantModelComplexTypes tests complex type conversions in ToAssistantModel
func TestToAssistantModelComplexTypes(t *testing.T) {
t.Run("CompleteLocales", func(t *testing.T) {
data := map[string]interface{}{
"locales": map[string]interface{}{
"en": map[string]interface{}{
"locale": "en",
"messages": map[string]interface{}{
"name": "English Name",
"description": "English Description",
},
},
"zh": map[string]interface{}{
"locale": "zh",
"messages": map[string]interface{}{
"name": "中文名称",
"description": "中文描述",
},
},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.Locales == nil {
t.Fatal("Expected Locales to be set")
}
if len(result.Locales) != 2 {
t.Errorf("Expected 2 locales, got %d", len(result.Locales))
}
})
t.Run("ComplexPrompts", func(t *testing.T) {
data := map[string]interface{}{
"prompts": []interface{}{
map[string]interface{}{
"role": "system",
"content": "You are a helpful assistant",
},
map[string]interface{}{
"role": "user",
"content": "Hello",
},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result.Prompts) != 2 {
t.Errorf("Expected 2 prompts, got %d", len(result.Prompts))
}
})
}
// TestGetBoolValue tests the getBoolValue helper function
func TestGetBoolValue(t *testing.T) {
t.Run("BoolTrue", func(t *testing.T) {
data := map[string]interface{}{"key": true}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true")
}
})
t.Run("BoolFalse", func(t *testing.T) {
data := map[string]interface{}{"key": false}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false")
}
})
t.Run("IntNonZero", func(t *testing.T) {
data := map[string]interface{}{"key": 1}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true for non-zero int")
}
})
t.Run("IntZero", func(t *testing.T) {
data := map[string]interface{}{"key": 0}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for zero int")
}
})
t.Run("Int64NonZero", func(t *testing.T) {
data := map[string]interface{}{"key": int64(1)}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true for non-zero int64")
}
})
t.Run("Int64Zero", func(t *testing.T) {
data := map[string]interface{}{"key": int64(0)}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for zero int64")
}
})
t.Run("Float64NonZero", func(t *testing.T) {
data := map[string]interface{}{"key": float64(1.5)}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true for non-zero float64")
}
})
t.Run("Float64Zero", func(t *testing.T) {
data := map[string]interface{}{"key": float64(0)}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for zero float64")
}
})
t.Run("StringTrue", func(t *testing.T) {
data := map[string]interface{}{"key": "true"}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true for string 'true'")
}
})
t.Run("StringOne", func(t *testing.T) {
data := map[string]interface{}{"key": "1"}
result := getBoolValue(data, "key")
if !result {
t.Error("Expected true for string '1'")
}
})
t.Run("StringFalse", func(t *testing.T) {
data := map[string]interface{}{"key": "false"}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for string 'false'")
}
})
t.Run("StringOther", func(t *testing.T) {
data := map[string]interface{}{"key": "other"}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for other string values")
}
})
t.Run("NilValue", func(t *testing.T) {
data := map[string]interface{}{"key": nil}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for nil value")
}
})
t.Run("MissingKey", func(t *testing.T) {
data := map[string]interface{}{}
result := getBoolValue(data, "missing")
if result {
t.Error("Expected false for missing key")
}
})
t.Run("UnsupportedType", func(t *testing.T) {
data := map[string]interface{}{"key": struct{}{}}
result := getBoolValue(data, "key")
if result {
t.Error("Expected false for unsupported type")
}
})
}

60
agent/store/mongo.go
View file

@ -54,8 +54,8 @@ func (m *Mongo) UpdateChatTitle(sid string, cid string, title string) error {
}
// SaveAssistant saves assistant information
func (m *Mongo) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return assistant["assistant_id"], nil
func (m *Mongo) SaveAssistant(assistant *AssistantModel) (string, error) {
return assistant.ID, nil
}
// DeleteAssistant deletes an assistant
@ -64,12 +64,12 @@ func (m *Mongo) DeleteAssistant(assistantID string) error {
}
// GetAssistants retrieves a list of assistants
func (m *Mongo) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantResponse, error) {
return &AssistantResponse{}, nil
func (m *Mongo) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error) {
return &AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (map[string]interface{}, error) {
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (*AssistantModel, error) {
return nil, nil
}
@ -83,56 +83,6 @@ func (m *Mongo) GetAssistantTags(locale ...string) ([]Tag, error) {
return []Tag{}, nil
}
// SaveAttachment saves attachment information
func (m *Mongo) SaveAttachment(attachment map[string]interface{}) (interface{}, error) {
return attachment["file_id"], nil
}
// DeleteAttachment deletes an attachment
func (m *Mongo) DeleteAttachment(fileID string) error {
return nil
}
// GetAttachments retrieves a list of attachments
func (m *Mongo) GetAttachments(filter AttachmentFilter, locale ...string) (*AttachmentResponse, error) {
return &AttachmentResponse{}, nil
}
// GetAttachment retrieves a single attachment by file ID
func (m *Mongo) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) {
return nil, nil
}
// DeleteAttachments deletes attachments based on filter conditions
func (m *Mongo) DeleteAttachments(filter AttachmentFilter) (int64, error) {
return 0, nil
}
// SaveKnowledge saves knowledge collection information
func (m *Mongo) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) {
return knowledge["collection_id"], nil
}
// DeleteKnowledge deletes a knowledge collection
func (m *Mongo) DeleteKnowledge(collectionID string) error {
return nil
}
// GetKnowledges retrieves a list of knowledge collections
func (m *Mongo) GetKnowledges(filter KnowledgeFilter, locale ...string) (*KnowledgeResponse, error) {
return &KnowledgeResponse{}, nil
}
// GetKnowledge retrieves a single knowledge collection by ID
func (m *Mongo) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) {
return nil, nil
}
// DeleteKnowledges deletes knowledge collections based on filter conditions
func (m *Mongo) DeleteKnowledges(filter KnowledgeFilter) (int64, error) {
return 0, nil
}
// Close closes the store and releases any resources
func (m *Mongo) Close() error {
return nil

60
agent/store/redis.go
View file

@ -54,8 +54,8 @@ func (r *Redis) UpdateChatTitle(sid string, cid string, title string) error {
}
// SaveAssistant saves assistant information
func (r *Redis) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
return assistant["assistant_id"], nil
func (r *Redis) SaveAssistant(assistant *AssistantModel) (string, error) {
return assistant.ID, nil
}
// DeleteAssistant deletes an assistant
@ -64,12 +64,12 @@ func (r *Redis) DeleteAssistant(assistantID string) error {
}
// GetAssistants retrieves a list of assistants
func (r *Redis) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantResponse, error) {
return &AssistantResponse{}, nil
func (r *Redis) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error) {
return &AssistantList{}, nil
}
// GetAssistant retrieves a single assistant by ID
func (r *Redis) GetAssistant(assistantID string, locale ...string) (map[string]interface{}, error) {
func (r *Redis) GetAssistant(assistantID string, locale ...string) (*AssistantModel, error) {
return nil, nil
}
@ -83,56 +83,6 @@ func (r *Redis) GetAssistantTags(locale ...string) ([]Tag, error) {
return []Tag{}, nil
}
// SaveAttachment saves attachment information
func (r *Redis) SaveAttachment(attachment map[string]interface{}) (interface{}, error) {
return attachment["file_id"], nil
}
// DeleteAttachment deletes an attachment
func (r *Redis) DeleteAttachment(fileID string) error {
return nil
}
// GetAttachments retrieves a list of attachments
func (r *Redis) GetAttachments(filter AttachmentFilter, locale ...string) (*AttachmentResponse, error) {
return &AttachmentResponse{}, nil
}
// GetAttachment retrieves a single attachment by file ID
func (r *Redis) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) {
return nil, nil
}
// DeleteAttachments deletes attachments based on filter conditions
func (r *Redis) DeleteAttachments(filter AttachmentFilter) (int64, error) {
return 0, nil
}
// SaveKnowledge saves knowledge collection information
func (r *Redis) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) {
return knowledge["collection_id"], nil
}
// DeleteKnowledge deletes a knowledge collection
func (r *Redis) DeleteKnowledge(collectionID string) error {
return nil
}
// GetKnowledges retrieves a list of knowledge collections
func (r *Redis) GetKnowledges(filter KnowledgeFilter, locale ...string) (*KnowledgeResponse, error) {
return &KnowledgeResponse{}, nil
}
// GetKnowledge retrieves a single knowledge collection by ID
func (r *Redis) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) {
return nil, nil
}
// DeleteKnowledges deletes knowledge collections based on filter conditions
func (r *Redis) DeleteKnowledges(filter KnowledgeFilter) (int64, error) {
return 0, nil
}
// Close closes the store and releases any resources
func (r *Redis) Close() error {
return nil

234
agent/store/types.go
View file

@ -1,5 +1,7 @@
package store
import "github.com/yaoapp/yao/agent/i18n"
// Setting represents the conversation configuration structure
// Used to configure basic conversation parameters including connector, user field, table name, etc.
type Setting struct {
@ -60,16 +62,16 @@ type AssistantFilter struct {
Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty
}
// AssistantResponse represents the assistant response structure
// Used for returning paginated assistant lists
type AssistantResponse struct {
Data []map[string]interface{} `json:"data"` // The paginated data
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Number of items per page
PageCnt int `json:"pagecnt"` // Total number of pages
Next int `json:"next"` // Next page number
Prev int `json:"prev"` // Previous page number
Total int64 `json:"total"` // Total number of items
// AssistantList represents the paginated assistant list response structure
// Used for returning paginated assistant lists with metadata
type AssistantList struct {
Data []*AssistantModel `json:"data"` // List of assistants
Page int `json:"page"` // Current page number (1-based)
PageSize int `json:"pagesize"` // Number of items per page
PageCount int `json:"pagecount"` // Total number of pages
Next int `json:"next"` // Next page number (0 if no next page)
Prev int `json:"prev"` // Previous page number (0 if no previous page)
Total int `json:"total"` // Total number of items across all pages
}
// Tag represents a tag
@ -78,153 +80,83 @@ type Tag struct {
Label string `json:"label"`
}
// Store defines the conversation storage interface
// Provides basic operations required for conversation management
type Store interface {
// GetChats retrieves a list of chats
// sid: Session ID
// filter: Filter conditions
// Returns: Grouped chat list and potential error
GetChats(sid string, filter ChatFilter, locale ...string) (*ChatGroupResponse, error)
// GetChat retrieves a single chat's information
// sid: Session ID
// cid: Chat ID
// Returns: Chat information and potential error
GetChat(sid string, cid string, locale ...string) (*ChatInfo, error)
// GetChatWithFilter retrieves a single chat's information with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: Chat information and potential error
GetChatWithFilter(sid string, cid string, filter ChatFilter, locale ...string) (*ChatInfo, error)
// GetHistory retrieves chat history
// sid: Session ID
// cid: Chat ID
// Returns: History record list and potential error
GetHistory(sid string, cid string, locale ...string) ([]map[string]interface{}, error)
// GetHistoryWithFilter retrieves chat history with filter options
// sid: Session ID
// cid: Chat ID
// filter: Filter conditions
// Returns: History record list and potential error
GetHistoryWithFilter(sid string, cid string, filter ChatFilter, locale ...string) ([]map[string]interface{}, error)
// SaveHistory saves chat history
// sid: Session ID
// messages: Message list
// cid: Chat ID
// context: Context information
// Returns: Potential error
SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error
// DeleteChat deletes a single chat
// sid: Session ID
// cid: Chat ID
// Returns: Potential error
DeleteChat(sid string, cid string) error
// DeleteAllChats deletes all chats
// sid: Session ID
// Returns: Potential error
DeleteAllChats(sid string) error
// UpdateChatTitle updates chat title
// sid: Session ID
// cid: Chat ID
// title: New title
// Returns: Potential error
UpdateChatTitle(sid string, cid string, title string) error
// SaveAssistant saves assistant information
// assistant: Assistant information
// Returns: Potential error
SaveAssistant(assistant map[string]interface{}) (interface{}, error)
// DeleteAssistant deletes an assistant
// assistantID: Assistant ID
// Returns: Potential error
DeleteAssistant(assistantID string) error
// GetAssistants retrieves a list of assistants
// filter: Filter conditions
// Returns: Paginated assistant list and potential error
GetAssistants(filter AssistantFilter, locale ...string) (*AssistantResponse, error)
// GetAssistantTags retrieves all unique tags from assistants
// Returns: List of tags and potential error
GetAssistantTags(locale ...string) ([]Tag, error)
// GetAssistant retrieves a single assistant by ID
// assistantID: Assistant ID
// Returns: Assistant information and potential error
GetAssistant(assistantID string, locale ...string) (map[string]interface{}, error)
// DeleteAssistants deletes assistants based on filter conditions
// filter: Filter conditions
// Returns: Number of deleted records and potential error
DeleteAssistants(filter AssistantFilter) (int64, error)
// Close closes the store and releases any resources
// Returns: Potential error
Close() error
// Prompt a prompt
type Prompt struct {
Role string `json:"role"`
Content string `json:"content"`
Name string `json:"name,omitempty"`
}
// AttachmentFilter represents the attachment filter structure
// Used for filtering and pagination when retrieving attachment lists
type AttachmentFilter struct {
UID string `json:"uid,omitempty"` // Filter by user ID
Guest *bool `json:"guest,omitempty"` // Filter by guest status
Manager string `json:"manager,omitempty"` // Filter by upload manager
ContentType string `json:"content_type,omitempty"` // Filter by content type
Name string `json:"name,omitempty"` // Filter by filename
Public *bool `json:"public,omitempty"` // Filter by public status
Gzip *bool `json:"gzip,omitempty"` // Filter by gzip compression
CollectionID string `json:"collection_id,omitempty"` // Filter by knowledge collection ID
Status string `json:"status,omitempty"` // Filter by processing status (uploading, uploaded, indexing, indexed, upload_failed, index_failed)
Keywords string `json:"keywords,omitempty"` // Search in filename
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Items per page
Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty
// KnowledgeBase the knowledge base configuration
type KnowledgeBase struct {
Collections []string `json:"collections,omitempty"` // Knowledge base collection IDs
Options map[string]interface{} `json:"options,omitempty"` // Additional options for knowledge base
}
// AttachmentResponse represents the attachment response structure
// Used for returning paginated attachment lists
type AttachmentResponse struct {
Data []map[string]interface{} `json:"data"` // The paginated data
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Number of items per page
PageCnt int `json:"pagecnt"` // Total number of pages
Next int `json:"next"` // Next page number
Prev int `json:"prev"` // Previous page number
Total int64 `json:"total"` // Total number of items
// MCPServers the MCP servers configuration
type MCPServers struct {
Servers []string `json:"servers,omitempty"` // MCP server IDs
Options map[string]interface{} `json:"options,omitempty"` // Additional options for MCP servers
}
// KnowledgeFilter represents the knowledge filter structure
// Used for filtering and pagination when retrieving knowledge lists
type KnowledgeFilter struct {
UID string `json:"uid,omitempty"` // Filter by user ID
Name string `json:"name,omitempty"` // Filter by collection name
Keywords string `json:"keywords,omitempty"` // Search in name and description
Public *bool `json:"public,omitempty"` // Filter by public status
Readonly *bool `json:"readonly,omitempty"` // Filter by readonly status
System *bool `json:"system,omitempty"` // Filter by system status
Page int `json:"page,omitempty"` // Page number, starting from 1
PageSize int `json:"pagesize,omitempty"` // Items per page
Select []string `json:"select,omitempty"` // Fields to return, returns all fields if empty
// Workflow the workflow configuration
type Workflow struct {
Workflows []string `json:"workflows,omitempty"` // Workflow IDs
Options map[string]interface{} `json:"options,omitempty"` // Additional workflow options
}
// KnowledgeResponse represents the knowledge response structure
// Used for returning paginated knowledge lists
type KnowledgeResponse struct {
Data []map[string]interface{} `json:"data"` // The paginated data
Page int `json:"page"` // Current page number
PageSize int `json:"pagesize"` // Number of items per page
PageCnt int `json:"pagecnt"` // Total number of pages
Next int `json:"next"` // Next page number
Prev int `json:"prev"` // Previous page number
Total int64 `json:"total"` // Total number of items
// Tool represents a tool configuration for storage
type Tool struct {
Type string `json:"type,omitempty"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]interface{} `json:"parameters,omitempty"`
}
// ToolCalls the tool calls
type ToolCalls struct {
Tools []Tool `json:"tools,omitempty"`
Prompts []Prompt `json:"prompts,omitempty"`
}
// Placeholder the assistant placeholder
type Placeholder struct {
Title string `json:"title,omitempty"`
Description string `json:"description,omitempty"`
Prompts []string `json:"prompts,omitempty"`
}
// AssistantModel the assistant database model
type AssistantModel struct {
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
Path string `json:"path,omitempty"` // Assistant Path
BuiltIn bool `json:"built_in,omitempty"` // Whether this is a built-in assistant
Sort int `json:"sort,omitempty"` // Assistant Sort
Description string `json:"description,omitempty"` // Assistant Description
Tags []string `json:"tags,omitempty"` // Assistant Tags
Readonly bool `json:"readonly,omitempty"` // Whether this assistant is readonly
Public bool `json:"public,omitempty"` // Whether this assistant is shared across all teams in the platform
Share string `json:"share,omitempty"` // Assistant sharing scope (private/team)
Mentionable bool `json:"mentionable,omitempty"` // Whether this assistant is mentionable
Automated bool `json:"automated,omitempty"` // Whether this assistant is automated
Options map[string]interface{} `json:"options,omitempty"` // AI Options
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts
KB *KnowledgeBase `json:"kb,omitempty"` // Knowledge base configuration
MCP *MCPServers `json:"mcp,omitempty"` // MCP servers configuration
Tools *ToolCalls `json:"tools,omitempty"` // Assistant Tools
Workflow *Workflow `json:"workflow,omitempty"` // Workflow configuration
Placeholder *Placeholder `json:"placeholder,omitempty"` // Assistant Placeholder
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
CreatedAt int64 `json:"created_at"` // Creation timestamp
UpdatedAt int64 `json:"updated_at"` // Last update timestamp
// Permission management fields (not exposed in JSON API responses)
YaoCreatedBy string `json:"-"` // User who created the assistant (not exposed in JSON)
YaoUpdatedBy string `json:"-"` // User who last updated the assistant (not exposed in JSON)
YaoTeamID string `json:"-"` // Team ID for team-based access control (not exposed in JSON)
YaoTenantID string `json:"-"` // Tenant ID for multi-tenancy support (not exposed in JSON)
}

480
agent/store/xun.go
View file

@ -9,6 +9,7 @@ import (
"github.com/google/uuid"
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"
@ -16,16 +17,14 @@ import (
"github.com/yaoapp/yao/agent/i18n"
)
// Package conversation provides functionality for managing chat conversations and assistants.
// Package store provides functionality for managing chat conversations and assistants.
// Xun implements the Conversation interface using a database backend.
// 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
// - Managing file attachments with metadata and access control
// - Managing knowledge collections for AI assistants
// - Supporting data expiration through TTL settings
type Xun struct {
query query.Query
@ -38,41 +37,38 @@ type Xun struct {
// Public interface methods:
//
// NewXun creates a new conversation instance with the given settings
// UpdateChatTitle updates the title of a specific chat
// 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 history
// 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
// SaveAttachment creates or updates an attachment
// DeleteAttachment deletes an attachment by file_id
// GetAttachments retrieves a paginated list of attachments with filtering
// GetAttachment retrieves a single attachment by file_id
// SaveKnowledge creates or updates a knowledge collection
// DeleteKnowledge deletes a knowledge collection by collection_id
// GetKnowledges retrieves a paginated list of knowledge collections with filtering
// GetKnowledge retrieves a single knowledge collection by collection_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 Setting) (Store, error) {
conv := &Xun{setting: setting}
if setting.Connector == "default" {
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, err
return nil, fmt.Errorf("select store connector %s error: %s", setting.Connector, err.Error())
}
conv.query, err = conn.Query()
if err != nil {
return nil, err
return nil, fmt.Errorf("query store connector %s error: %s", setting.Connector, err.Error())
}
conv.schema, err = conn.Schema()
@ -162,21 +158,6 @@ func (conv *Xun) Close() error {
// Rename Init to initialize to avoid conflicts
func (conv *Xun) initialize() error {
// Initialize chat table
if err := conv.initChatTable(); err != nil {
return err
}
// Initialize history table
if err := conv.initHistoryTable(); err != nil {
return err
}
// Initialize assistant table
if err := conv.initAssistantTable(); err != nil {
return err
}
// Start automatic cleanup if TTL is enabled
if conv.setting.TTL > 0 {
conv.startAutoClean()
@ -277,78 +258,32 @@ 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() // assistant name
table.String("avatar", 200).Null() // assistant avatar
table.String("connector", 200).NotNull() // assistant connector
table.String("description", 600).Null().Index() // assistant description
table.String("path", 200).Null() // assistant storage path
table.Integer("sort").SetDefault(9999).Index() // assistant sort order
table.Boolean("built_in").SetDefault(false).Index() // whether this is a built-in assistant
table.JSON("placeholder").Null() // assistant placeholder
table.JSON("options").Null() // assistant options
table.JSON("prompts").Null() // assistant prompts
table.JSON("workflow").Null() // assistant workflow
table.JSON("knowledge").Null() // assistant knowledge
table.JSON("tools").Null() // assistant tools
table.JSON("tags").Null() // assistant tags
table.Boolean("readonly").SetDefault(false).Index() // assistant readonly
table.JSON("permissions").Null() // assistant permissions
table.JSON("locales").Null() // assistant i18n
table.Boolean("automated").SetDefault(true).Index() // assistant autoable
table.Boolean("mentionable").SetDefault(true).Index() // Whether this assistant can appear in @ mention list
table.TimestampTz("created_at").SetDefaultRaw("CURRENT_TIMESTAMP").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", "path", "sort", "built_in", "placeholder", "options", "prompts", "workflow", "knowledge", "tools", "tags", "readonly", "permissions", "locales", "automated", "mentionable", "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"
}
@ -1022,64 +957,108 @@ func (conv *Xun) parseJSONFields(data map[string]interface{}, fields []string) {
}
// SaveAssistant saves assistant information
func (conv *Xun) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
func (conv *Xun) SaveAssistant(assistant *AssistantModel) (string, error) {
if assistant == nil {
return "", fmt.Errorf("assistant cannot be nil")
}
// Validate required fields
requiredFields := []string{"name", "type", "connector"}
for _, field := range requiredFields {
if _, ok := assistant[field]; !ok {
return nil, fmt.Errorf("field %s is required", field)
}
if assistant[field] == nil || assistant[field] == "" {
return nil, fmt.Errorf("field %s cannot be empty", field)
}
if assistant.Name == "" {
return "", fmt.Errorf("field name is required")
}
// Create a copy of the assistant map to avoid modifying the original
assistantCopy := make(map[string]interface{})
for k, v := range assistant {
assistantCopy[k] = v
if assistant.Type == "" {
return "", fmt.Errorf("field type is required")
}
// Process JSON fields
jsonFields := []string{"tags", "options", "prompts", "workflow", "knowledge", "tools", "permissions", "placeholder", "locales"}
for _, field := range jsonFields {
if val, ok := assistantCopy[field]; ok && val != nil {
// If it's a string, try to parse it first
if strVal, ok := val.(string); ok && strVal != "" {
var parsed interface{}
if err := jsoniter.UnmarshalFromString(strVal, &parsed); err == nil {
assistantCopy[field] = parsed
}
}
}
if assistant.Connector == "" {
return "", fmt.Errorf("field connector is required")
}
// Generate assistant_id if not provided
if _, ok := assistantCopy["assistant_id"]; !ok {
if assistant.ID == "" {
var err error
assistantCopy["assistant_id"], err = conv.GenerateAssistantID()
assistant.ID, err = conv.GenerateAssistantID()
if err != nil {
return nil, err
return "", err
}
}
// Check if assistant exists
exists, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantCopy["assistant_id"]).
Where("assistant_id", assistant.ID).
Exists()
if err != nil {
return nil, err
return "", err
}
// Convert JSON fields to strings for storage
for _, field := range jsonFields {
if val, ok := assistantCopy[field]; ok && val != nil {
jsonStr, err := jsoniter.MarshalToString(val)
// Convert model to map for database storage
data := make(map[string]interface{})
data["assistant_id"] = assistant.ID
data["type"] = assistant.Type
data["name"] = assistant.Name
data["avatar"] = assistant.Avatar
data["connector"] = assistant.Connector
data["path"] = assistant.Path
data["built_in"] = assistant.BuiltIn
data["sort"] = assistant.Sort
data["description"] = assistant.Description
data["readonly"] = assistant.Readonly
data["public"] = assistant.Public
data["share"] = assistant.Share
data["mentionable"] = assistant.Mentionable
data["automated"] = assistant.Automated
data["created_at"] = assistant.CreatedAt
data["updated_at"] = assistant.UpdatedAt
// Permission management fields
if assistant.YaoCreatedBy != "" {
data["__yao_created_by"] = assistant.YaoCreatedBy
}
if assistant.YaoUpdatedBy != "" {
data["__yao_updated_by"] = assistant.YaoUpdatedBy
}
if assistant.YaoTeamID != "" {
data["__yao_team_id"] = assistant.YaoTeamID
}
if assistant.YaoTenantID != "" {
data["__yao_tenant_id"] = assistant.YaoTenantID
}
// 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 nil, fmt.Errorf("failed to marshal %s to JSON: %v", field, err)
return "", fmt.Errorf("failed to marshal %s: %w", field, err)
}
assistantCopy[field] = jsonStr
data[field] = jsonStr
}
}
@ -1087,21 +1066,21 @@ func (conv *Xun) SaveAssistant(assistant map[string]interface{}) (interface{}, e
if exists {
_, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantCopy["assistant_id"]).
Update(assistantCopy)
Where("assistant_id", assistant.ID).
Update(data)
if err != nil {
return nil, err
return "", err
}
return assistantCopy["assistant_id"], nil
return assistant.ID, nil
}
err = conv.query.New().
Table(conv.getAssistantTable()).
Insert(assistantCopy)
Insert(data)
if err != nil {
return nil, err
return "", err
}
return assistantCopy["assistant_id"], nil
return assistant.ID, nil
}
// DeleteAssistant deletes an assistant by assistant_id
@ -1127,7 +1106,7 @@ func (conv *Xun) DeleteAssistant(assistantID string) error {
}
// GetAssistants retrieves assistants with pagination and filtering
func (conv *Xun) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantResponse, error) {
func (conv *Xun) GetAssistants(filter AssistantFilter, locale ...string) (*AssistantList, error) {
qb := conv.query.New().
Table(conv.getAssistantTable())
@ -1235,53 +1214,62 @@ func (conv *Xun) GetAssistants(filter AssistantFilter, locale ...string) (*Assis
return nil, err
}
// Convert rows to map slice and parse JSON fields
data := make([]map[string]interface{}, len(rows))
jsonFields := []string{"tags", "options", "prompts", "workflow", "knowledge", "tools", "permissions", "placeholder"}
for i, row := range rows {
data[i] = row
// Only parse JSON fields if they are selected or no select filter is provided
if filter.Select == nil || len(filter.Select) == 0 {
conv.parseJSONFields(data[i], jsonFields)
} else {
// Parse only selected JSON fields
selectedJSONFields := []string{}
for _, field := range jsonFields {
for _, selected := range filter.Select {
if selected == field {
selectedJSONFields = append(selectedJSONFields, field)
break
// Convert rows to AssistantModel slice
assistants := make([]*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 AssistantModel using existing helper function
model, err := ToAssistantModel(data)
if err != nil {
log.Error("Failed to convert row to 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
}
}
}
}
if len(selectedJSONFields) > 0 {
conv.parseJSONFields(data[i], selectedJSONFields)
}
}
assistants = append(assistants, model)
}
// Translate Data
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
for i, row := range data {
assistantID := row["assistant_id"].(string)
data[i] = i18n.Translate(assistantID, lang, row).(map[string]interface{})
}
}
return &AssistantResponse{
Data: data,
Page: filter.Page,
PageSize: filter.PageSize,
PageCnt: totalPages,
Next: nextPage,
Prev: prevPage,
Total: total,
return &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) (map[string]interface{}, error) {
func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*AssistantModel, error) {
row, err := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
@ -1300,13 +1288,151 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (map[string]
}
// Parse JSON fields
jsonFields := []string{"tags", "options", "prompts", "workflow", "knowledge", "tools", "permissions", "placeholder"}
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales"}
conv.parseJSONFields(data, jsonFields)
if len(locale) > 0 {
lang := strings.ToLower(locale[0])
return i18n.Translate(assistantID, lang, data).(map[string]interface{}), nil
// Convert map to AssistantModel
model := &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"),
}
return data, nil
// 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 []Prompt
if err := jsoniter.Unmarshal(raw, &p); err == nil {
model.Prompts = p
}
}
}
if kb, has := data["kb"]; has && kb != nil {
kbConverted, err := ToKnowledgeBase(kb)
if err == nil {
model.KB = kbConverted
}
}
if mcp, has := data["mcp"]; has && mcp != nil {
mcpConverted, err := ToMCPServers(mcp)
if err == nil {
model.MCP = mcpConverted
}
}
if workflow, has := data["workflow"]; has && workflow != nil {
wf, err := 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 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 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
}
// 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
}
// DeleteAssistants deletes assistants based on filter conditions

1107
agent/store/xun_assistant_test.go Normal file
View file

File diff suppressed because it is too large Load diff

288
data/bindata.go
View file

File diff suppressed because it is too large Load diff

10
engine/load.go
View file

@ -297,11 +297,11 @@ func Load(cfg config.Config, options LoadOption) (warnings []Warning, err error)
warnings = append(warnings, Warning{Widget: "Knowledge Base", Error: err})
}
// Load Neo
// Load Agent
err = agent.Load(cfg)
if err != nil {
// printErr(cfg.Mode, "Neo", err)
warnings = append(warnings, Warning{Widget: "Neo", Error: err})
// printErr(cfg.Mode, "Agent", err)
warnings = append(warnings, Warning{Widget: "Agent", Error: err})
}
for name, hook := range LoadHooks {
@ -527,10 +527,10 @@ func Reload(cfg config.Config, options LoadOption) (err error) {
}
// Load Neo
// Load Agent
err = agent.Load(cfg)
if err != nil {
printErr(cfg.Mode, "Neo", err)
printErr(cfg.Mode, "Agent", err)
}
// Load OpenAPI

46
yao/models/agent/assistant.mod.yao
View file

@ -6,10 +6,7 @@
"builtin": true,
"readonly": true,
"sort": 9999,
"table": {
"name": "agent_assistant",
"comment": "Agent assistant table"
},
"table": { "name": "agent_assistant", "comment": "Agent assistant table" },
"columns": [
{
"name": "id",
@ -56,7 +53,7 @@
"name": "connector",
"type": "string",
"label": "Connector",
"comment": "Assistant connector",
"comment": "Assistant default connector, if not set, use the global default connector",
"length": 200,
"nullable": false
},
@ -122,10 +119,17 @@
"nullable": true
},
{
"name": "knowledge",
"name": "kb",
"type": "json",
"label": "Knowledge",
"comment": "Assistant knowledge",
"label": "Knowledge Base",
"comment": "Assistant knowledge base collections",
"nullable": true
},
{
"name": "mcp",
"type": "json",
"label": "MCP Servers",
"comment": "MCP servers available for the assistant to use",
"nullable": true
},
{
@ -150,13 +154,29 @@
"default": false,
"index": true
},
{
"name": "permissions",
"type": "json",
"label": "Permissions",
"comment": "Assistant permissions",
"nullable": true
"name": "public",
"type": "boolean",
"label": "Public Assistant",
"comment": "Whether this assistant is shared across all teams in the platform",
"default": false,
"nullable": false
},
{
"name": "share",
"type": "enum",
"label": "Share",
"comment": "Assistant sharing scope",
"option": [
"private", // Only visible to the owner
"team" // Visible to all team members
],
"default": "private",
"nullable": false,
"index": true
},
{
"name": "locales",
"type": "json",

20
yao/models/agent/chat.mod.yao
View file

@ -6,10 +6,7 @@
"builtin": true,
"readonly": true,
"sort": 9999,
"table": {
"name": "agent_chat",
"comment": "Agent chat table"
},
"table": { "name": "agent_chat", "comment": "Agent chat table" },
"columns": [
{
"name": "id",
@ -44,15 +41,6 @@
"nullable": true,
"index": true
},
{
"name": "sid",
"type": "string",
"label": "Session ID",
"comment": "Session identifier",
"length": 255,
"nullable": false,
"index": true
},
{
"name": "silent",
"type": "boolean",
@ -77,12 +65,6 @@
}
},
"indexes": [
{
"name": "idx_agent_chat_session_assistant",
"columns": ["sid", "assistant_id"],
"type": "index",
"comment": "Index for session and assistant queries"
},
{
"name": "idx_agent_chat_silent",
"columns": ["silent", "created_at"],

56
yao/models/agent/history.mod.yao
View file

@ -6,10 +6,7 @@
"builtin": true,
"readonly": true,
"sort": 9999,
"table": {
"name": "agent_history",
"comment": "Agent chat history table"
},
"table": { "name": "agent_history", "comment": "Agent chat history table" },
"columns": [
{
"name": "id",
@ -18,16 +15,7 @@
"comment": "Unique record identifier"
},
{
"name": "sid",
"type": "string",
"label": "Session ID",
"comment": "Session identifier",
"length": 255,
"nullable": false,
"index": true
},
{
"name": "cid",
"name": "chat_id",
"type": "string",
"label": "Chat ID",
"comment": "Chat identifier",
@ -35,24 +23,6 @@
"nullable": true,
"index": true
},
{
"name": "uid",
"type": "string",
"label": "User ID",
"comment": "User identifier",
"length": 255,
"nullable": true,
"index": true
},
{
"name": "role",
"type": "string",
"label": "Role",
"comment": "Message role (user/assistant/system)",
"length": 200,
"nullable": true,
"index": true
},
{
"name": "name",
"type": "string",
@ -128,7 +98,7 @@
"chat": {
"type": "hasOne",
"model": "__yao.agent.chat",
"key": "cid",
"key": "chat_id",
"foreign": "chat_id"
},
"assistant": {
@ -140,29 +110,11 @@
"user": {
"type": "hasOne",
"model": "__yao.user",
"key": "uid",
"key": "user_id",
"foreign": "user_id"
}
},
"indexes": [
{
"name": "idx_agent_history_session_chat",
"columns": ["sid", "cid"],
"type": "index",
"comment": "Index for session and chat queries"
},
{
"name": "idx_agent_history_user_role",
"columns": ["uid", "role"],
"type": "index",
"comment": "Index for user and role queries"
},
{
"name": "idx_agent_history_assistant",
"columns": ["assistant_id", "created_at"],
"type": "index",
"comment": "Index for assistant history queries"
},
{
"name": "idx_agent_history_expired",
"columns": ["expired_at", "silent"],