Merge pull request #1363 from trheyi/main

Enhance assistant
This commit is contained in:
Max 2025-12-02 17:15:46 +08:00 committed by GitHub
commit 9e0d0b2e24
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
46 changed files with 5290 additions and 1363 deletions

View file

@ -6,6 +6,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/assistant/handlers"
"github.com/yaoapp/yao/agent/context"
@ -303,7 +304,7 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// GetConnector get the connector object, capabilities, and error with priority: createResponse > ctx > ast
// Note: createResponse.Connector is already applied to ctx.Connector by applyContextAdjustments in create.go
// Returns: (connector, capabilities, error)
func (ast *Assistant) GetConnector(ctx *context.Context) (connector.Connector, *context.ModelCapabilities, error) {
func (ast *Assistant) GetConnector(ctx *context.Context) (connector.Connector, *openai.Capabilities, error) {
// Determine connector ID with priority
connectorID := ast.Connector
if ctx.Connector != "" {
@ -322,68 +323,56 @@ func (ast *Assistant) GetConnector(ctx *context.Context) (connector.Connector, *
}
// Get connector capabilities from settings
capabilities := ast.getConnectorCapabilities(connectorID)
capabilities := ast.getConnectorCapabilities(conn)
return conn, capabilities, nil
}
// getConnectorCapabilities get the capabilities of a connector from settings
func (ast *Assistant) getConnectorCapabilities(connectorID string) *context.ModelCapabilities {
// Initialize with default capabilities (all disabled)
falseVal := false
capabilities := &context.ModelCapabilities{
Vision: falseVal,
ToolCalls: &falseVal,
Audio: &falseVal,
Reasoning: &falseVal,
Streaming: &falseVal,
// Priority: 1. modelCapabilities mapping, 2. connector's Setting()["capabilities"]
func (ast *Assistant) getConnectorCapabilities(conn connector.Connector) *openai.Capabilities {
if conn == nil {
return &openai.Capabilities{
Vision: false,
ToolCalls: false,
Audio: false,
Reasoning: false,
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true,
}
}
// Get model capabilities from global configuration
modelCaps, exists := modelCapabilities[connectorID]
if !exists {
// Return default capabilities if model not found in configuration
// Get connector ID
connectorID := conn.ID()
// Priority 1: Check global modelCapabilities mapping
if modelCaps, exists := modelCapabilities[connectorID]; exists {
return &modelCaps
}
// Priority 2: Get capabilities from connector's Setting() method
// Modern connectors (post-upgrade) provide default capabilities via Setting()
settings := conn.Setting()
if caps, ok := settings["capabilities"]; ok {
if capabilities, ok := caps.(*openai.Capabilities); ok {
return capabilities
}
// Update capabilities based on model configuration
// Vision can be bool or string (VisionFormat)
if modelCaps.Vision != nil {
capabilities.Vision = modelCaps.Vision
}
// Handle both Tools (deprecated) and ToolCalls
if modelCaps.ToolCalls || modelCaps.Tools {
v := true
capabilities.ToolCalls = &v
// Fallback: Return minimal default capabilities
// This should rarely happen with upgraded connectors
return &openai.Capabilities{
Vision: false,
ToolCalls: false,
Audio: false,
Reasoning: false,
Streaming: false,
JSON: false,
Multimodal: false,
TemperatureAdjustable: true, // Default to true for non-reasoning models
}
if modelCaps.Audio {
v := true
capabilities.Audio = &v
}
if modelCaps.Reasoning {
v := true
capabilities.Reasoning = &v
}
if modelCaps.Streaming {
v := true
capabilities.Streaming = &v
}
if modelCaps.JSON {
v := true
capabilities.JSON = &v
}
if modelCaps.Multimodal {
v := true
capabilities.Multimodal = &v
}
return capabilities
}
// Info get the assistant information

View file

@ -4,7 +4,6 @@ import (
"fmt"
"path"
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/fs"
"github.com/yaoapp/yao/agent/i18n"
store "github.com/yaoapp/yao/agent/store/types"
@ -16,31 +15,6 @@ func Get(id string) (*Assistant, error) {
return LoadStore(id)
}
// GetByConnector get the assistant by connector
func GetByConnector(connector string, name string) (*Assistant, error) {
id := "connector:" + connector
assistant, exists := loaded.Get(id)
if exists {
return assistant, nil
}
data := map[string]interface{}{
"assistant_id": id,
"connector": connector,
"description": "Default assistant for " + connector,
"name": name,
"type": "assistant",
}
assistant, err := loadMap(data)
if err != nil {
return nil, err
}
loaded.Put(assistant)
return assistant, nil
}
// GetPlaceholder returns the placeholder of the assistant
func (ast *Assistant) GetPlaceholder(locale string) *store.Placeholder {
@ -97,15 +71,18 @@ func (ast *Assistant) Map() map[string]interface{} {
"share": ast.Share,
"avatar": ast.Avatar,
"connector": ast.Connector,
"connector_options": ast.ConnectorOptions,
"path": ast.Path,
"built_in": ast.BuiltIn,
"sort": ast.Sort,
"description": ast.Description,
"options": ast.Options,
"prompts": ast.Prompts,
"prompt_presets": ast.PromptPresets,
"disable_global_prompts": ast.DisableGlobalPrompts,
"source": ast.Source,
"kb": ast.KB,
"mcp": ast.MCP,
"tools": ast.Tools,
"workflow": ast.Workflow,
"tags": ast.Tags,
"mentionable": ast.Mentionable,
@ -175,6 +152,8 @@ func (ast *Assistant) Clone() *Assistant {
Share: ast.Share,
Mentionable: ast.Mentionable,
Automated: ast.Automated,
DisableGlobalPrompts: ast.DisableGlobalPrompts,
Source: ast.Source,
CreatedAt: ast.CreatedAt,
UpdatedAt: ast.UpdatedAt,
},
@ -247,17 +226,28 @@ func (ast *Assistant) Clone() *Assistant {
copy(clone.Prompts, ast.Prompts)
}
// Deep copy tools
if ast.Tools != nil {
clone.Tools = &store.ToolCalls{}
if ast.Tools.Tools != nil {
clone.Tools.Tools = make([]store.Tool, len(ast.Tools.Tools))
copy(clone.Tools.Tools, ast.Tools.Tools)
// Deep copy prompt presets
if ast.PromptPresets != nil {
clone.PromptPresets = make(map[string][]store.Prompt)
for k, v := range ast.PromptPresets {
prompts := make([]store.Prompt, len(v))
copy(prompts, v)
clone.PromptPresets[k] = prompts
}
}
if ast.Tools.Prompts != nil {
clone.Tools.Prompts = make([]store.Prompt, len(ast.Tools.Prompts))
copy(clone.Tools.Prompts, ast.Tools.Prompts)
// Deep copy connector options
if ast.ConnectorOptions != nil {
clone.ConnectorOptions = &store.ConnectorOptions{
Optional: ast.ConnectorOptions.Optional,
}
if ast.ConnectorOptions.Connectors != nil {
clone.ConnectorOptions.Connectors = make([]string, len(ast.ConnectorOptions.Connectors))
copy(clone.ConnectorOptions.Connectors, ast.ConnectorOptions.Connectors)
}
if ast.ConnectorOptions.Filters != nil {
clone.ConnectorOptions.Filters = make([]store.ModelCapability, len(ast.ConnectorOptions.Filters))
copy(clone.ConnectorOptions.Filters, ast.ConnectorOptions.Filters)
}
}
@ -328,29 +318,7 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
ast.Connector = v
}
if v, has := data["tools"]; has {
switch tools := v.(type) {
case []store.Tool:
ast.Tools = &store.ToolCalls{
Tools: tools,
Prompts: ast.Prompts,
}
case *store.ToolCalls:
ast.Tools = tools
default:
raw, err := jsoniter.Marshal(tools)
if err != nil {
return err
}
ast.Tools = &store.ToolCalls{}
err = jsoniter.Unmarshal(raw, &ast.Tools)
if err != nil {
return err
}
}
}
// Note: tools field is deprecated, now handled by MCP
if v, ok := data["type"].(string); ok {
ast.Type = v
@ -364,6 +332,9 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
if v, ok := data["automated"].(bool); ok {
ast.Automated = v
}
if v, ok := data["disable_global_prompts"].(bool); ok {
ast.DisableGlobalPrompts = v
}
if v, ok := data["readonly"].(bool); ok {
ast.Readonly = v
}
@ -379,6 +350,27 @@ func (ast *Assistant) Update(data map[string]interface{}) error {
if v, ok := data["options"].(map[string]interface{}); ok {
ast.Options = v
}
if v, ok := data["source"].(string); ok {
ast.Source = v
}
// ConnectorOptions
if v, has := data["connector_options"]; has {
connOpts, err := store.ToConnectorOptions(v)
if err != nil {
return err
}
ast.ConnectorOptions = connOpts
}
// PromptPresets
if v, has := data["prompt_presets"]; has {
presets, err := store.ToPromptPresets(v)
if err != nil {
return err
}
ast.PromptPresets = presets
}
// KB
if v, has := data["kb"]; has {

View file

@ -3,8 +3,10 @@ package assistant
import (
"fmt"
"github.com/spf13/cast"
"github.com/yaoapp/gou/json"
"github.com/yaoapp/yao/agent/context"
store "github.com/yaoapp/yao/agent/store/types"
)
// BuildRequest build the LLM request
@ -48,29 +50,217 @@ func (ast *Assistant) buildMessages(ctx *context.Context, messages []context.Mes
finalMessages = append([]context.Message{mcpSamplesMsg}, finalMessages...)
}
// ⚠️ Just for testing, will remove later
// If we have prompts, prepend them to the beginning
if len(ast.Prompts) > 0 {
promptMessages := make([]context.Message, 0, len(ast.Prompts))
for _, prompt := range ast.Prompts {
msg := context.Message{
Role: context.MessageRole(prompt.Role),
Content: prompt.Content,
}
// Add name if provided
if prompt.Name != "" {
name := prompt.Name
msg.Name = &name
}
promptMessages = append(promptMessages, msg)
}
// Prepend prompt messages to the beginning
// Build and prepend system prompts (global + assistant prompts)
promptMessages := ast.buildSystemPrompts(ctx, createResponse)
if len(promptMessages) > 0 {
finalMessages = append(promptMessages, finalMessages...)
}
return finalMessages, nil
}
// buildSystemPrompts builds system prompt messages from global prompts and assistant prompts
// Order: Global prompts (if not disabled) -> Assistant prompts (or preset)
// Variables are parsed with context information
//
// Priority for prompt preset selection:
// 1. createResponse.PromptPreset (highest)
// 2. ctx.Metadata["__prompt_preset"]
// 3. ast.Prompts (default)
//
// Priority for disable global prompts:
// 1. createResponse.DisableGlobalPrompts (highest)
// 2. ctx.Metadata["__disable_global_prompts"]
// 3. ast.DisableGlobalPrompts (default)
func (ast *Assistant) buildSystemPrompts(ctx *context.Context, createResponse *context.HookCreateResponse) []context.Message {
// Build context variables from ctx and ast
ctxVars := ast.buildContextVariables(ctx)
// Determine if global prompts should be disabled
disableGlobal := ast.shouldDisableGlobalPrompts(ctx, createResponse)
// Get assistant prompts (default or preset)
assistantPrompts := ast.getAssistantPrompts(ctx, createResponse)
var allPrompts []store.Prompt
// 1. Add global prompts (if not disabled)
if !disableGlobal && len(globalPrompts) > 0 {
// Parse global prompts with context variables
parsedGlobal := store.Prompts(globalPrompts).Parse(ctxVars)
allPrompts = append(allPrompts, parsedGlobal...)
}
// 2. Add assistant prompts (default or preset)
if len(assistantPrompts) > 0 {
// Parse assistant prompts with context variables
parsedAssistant := store.Prompts(assistantPrompts).Parse(ctxVars)
allPrompts = append(allPrompts, parsedAssistant...)
}
// Convert to context.Message slice
if len(allPrompts) == 0 {
return nil
}
messages := make([]context.Message, 0, len(allPrompts))
for _, prompt := range allPrompts {
msg := context.Message{
Role: context.MessageRole(prompt.Role),
Content: prompt.Content,
}
if prompt.Name != "" {
name := prompt.Name
msg.Name = &name
}
messages = append(messages, msg)
}
return messages
}
// shouldDisableGlobalPrompts determines if global prompts should be disabled
// Priority: createResponse > ctx.Metadata > ast.DisableGlobalPrompts
func (ast *Assistant) shouldDisableGlobalPrompts(ctx *context.Context, createResponse *context.HookCreateResponse) bool {
// Priority 1: Hook response (highest)
if createResponse != nil && createResponse.DisableGlobalPrompts != nil {
return *createResponse.DisableGlobalPrompts
}
// Priority 2: ctx.Metadata["__disable_global_prompts"]
if ctx != nil && ctx.Metadata != nil {
if disable, ok := ctx.Metadata["__disable_global_prompts"].(bool); ok {
return disable
}
}
// Priority 3: Assistant configuration (default)
return ast.DisableGlobalPrompts
}
// getAssistantPrompts returns the assistant prompts based on preset selection
// Priority: createResponse.PromptPreset > ctx.Metadata["__prompt_preset"] > ast.Prompts
func (ast *Assistant) getAssistantPrompts(ctx *context.Context, createResponse *context.HookCreateResponse) []store.Prompt {
// Get preset key
presetKey := ast.getPromptPresetKey(ctx, createResponse)
// If preset key is specified and exists, use it
if presetKey != "" && ast.PromptPresets != nil {
if presets, ok := ast.PromptPresets[presetKey]; ok && len(presets) > 0 {
return presets
}
}
// Fallback to default prompts
return ast.Prompts
}
// getPromptPresetKey returns the prompt preset key
// Priority: createResponse.PromptPreset > ctx.Metadata["__prompt_preset"]
func (ast *Assistant) getPromptPresetKey(ctx *context.Context, createResponse *context.HookCreateResponse) string {
// Priority 1: Hook response (highest)
if createResponse != nil && createResponse.PromptPreset != "" {
return createResponse.PromptPreset
}
// Priority 2: ctx.Metadata["__prompt_preset"]
if ctx != nil && ctx.Metadata != nil {
if preset, ok := ctx.Metadata["__prompt_preset"].(string); ok && preset != "" {
return preset
}
}
// No preset specified
return ""
}
// buildContextVariables extracts context variables from Context and Assistant for prompt parsing
func (ast *Assistant) buildContextVariables(ctx *context.Context) map[string]string {
vars := make(map[string]string)
// Get locale from ctx (default to empty)
locale := ""
if ctx != nil && ctx.Locale != "" {
locale = ctx.Locale
}
// Assistant info (with locale support)
if ast != nil {
if ast.ID != "" {
vars["ASSISTANT_ID"] = ast.ID
}
// Use localized name and description
name := ast.GetName(locale)
if name != "" {
vars["ASSISTANT_NAME"] = name
}
description := ast.GetDescription(locale)
if description != "" {
vars["ASSISTANT_DESCRIPTION"] = description
}
if ast.Type != "" {
vars["ASSISTANT_TYPE"] = ast.Type
}
}
if ctx == nil {
return vars
}
// Basic context info
if ctx.ChatID != "" {
vars["CHAT_ID"] = ctx.ChatID
}
if ctx.Locale != "" {
vars["LOCALE"] = ctx.Locale
}
if ctx.Theme != "" {
vars["THEME"] = ctx.Theme
}
if ctx.Route != "" {
vars["ROUTE"] = ctx.Route
}
if ctx.Referer != "" {
vars["REFERER"] = ctx.Referer
}
// Client info (only non-sensitive fields)
if ctx.Client.Type != "" {
vars["CLIENT_TYPE"] = ctx.Client.Type
}
// Authorized info (only internal IDs, no PII)
// Note: USER_SUBJECT and CLIENT_IP are excluded for privacy/GDPR compliance
if ctx.Authorized != nil {
if ctx.Authorized.UserID != "" {
vars["USER_ID"] = ctx.Authorized.UserID
}
if ctx.Authorized.TeamID != "" {
vars["TEAM_ID"] = ctx.Authorized.TeamID
}
if ctx.Authorized.TenantID != "" {
vars["TENANT_ID"] = ctx.Authorized.TenantID
}
}
// Metadata - custom variables from ctx.Metadata
// All metadata keys are exposed as $CTX.{KEY}
// Supports string, int, uint, float, bool types
if ctx.Metadata != nil {
for key, value := range ctx.Metadata {
if value == nil {
continue
}
strVal := cast.ToString(value)
if strVal != "" {
vars[key] = strVal
}
}
}
return vars
}
// buildCompletionOptions builds completion options from multiple sources
// Priority (lowest to highest, later overrides earlier): ast > ctx > createResponse
// The priority means: if createResponse has a value, use it; else use ctx; else use ast

View file

@ -0,0 +1,840 @@
package assistant_test
import (
stdContext "context"
"strings"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// containsString is a helper to check if a content (string or interface{}) contains a substring
func containsString(content interface{}, substr string) bool {
switch v := content.(type) {
case string:
return strings.Contains(v, substr)
default:
return false
}
}
// newPromptTestContext creates a context suitable for prompt testing with Create Hook
func newPromptTestContext(chatID, assistantID string) *context.Context {
return &context.Context{
Context: stdContext.Background(),
ChatID: chatID,
AssistantID: assistantID,
Locale: "en-us",
Theme: "light",
Client: context.Client{
Type: "web",
UserAgent: "TestAgent/1.0",
IP: "127.0.0.1",
},
Referer: context.RefererAPI,
Accept: context.AcceptWebCUI,
Metadata: make(map[string]interface{}),
Authorized: &types.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client-id",
UserID: "test-user-123",
TeamID: "test-team-456",
TenantID: "test-tenant-789",
SessionID: "test-session-id",
},
}
}
func TestBuildSystemPromptsIntegration(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
t.Run("AssistantWithLocale", func(t *testing.T) {
// Load an assistant with locales
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
ctx := &context.Context{
Locale: "zh-cn",
Authorized: &types.AuthorizedInfo{
UserID: "test-user-123",
TeamID: "test-team-456",
},
Metadata: map[string]interface{}{
"CUSTOM_VAR": "custom-value",
"INT_VAR": 42,
"BOOL_VAR": true,
},
}
// Build request to test the full flow
messages := []context.Message{
{Role: context.RoleUser, Content: "Hello"},
}
finalMessages, options, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
require.NotNil(t, options)
// Should have system prompts prepended
assert.Greater(t, len(finalMessages), 1)
// First messages should be system prompts
hasSystemPrompt := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
hasSystemPrompt = true
break
}
}
assert.True(t, hasSystemPrompt, "Should have system prompts")
})
t.Run("DisableGlobalPrompts", func(t *testing.T) {
// Load fullfields assistant which has disable_global_prompts: true
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
require.True(t, ast.DisableGlobalPrompts)
ctx := &context.Context{
Locale: "en-us",
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Hello"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Should still have assistant prompts
hasSystemPrompt := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
hasSystemPrompt = true
break
}
}
assert.True(t, hasSystemPrompt, "Should have assistant prompts even with global disabled")
})
t.Run("MetadataTypeConversion", func(t *testing.T) {
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{
Metadata: map[string]interface{}{
"STRING_VAL": "hello",
"INT_VAL": 123,
"INT64_VAL": int64(456),
"FLOAT_VAL": 3.14,
"BOOL_TRUE": true,
"BOOL_FALSE": false,
"UINT_VAL": uint(789),
"NIL_VAL": nil,
"EMPTY_VAL": "",
"ZERO_INT": 0,
"ZERO_FLOAT": 0.0,
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test metadata"},
}
// This should not panic
_, _, err = ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
})
t.Run("AuthorizedInfoPrivacy", func(t *testing.T) {
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{
Authorized: &types.AuthorizedInfo{
UserID: "user-123",
Subject: "user@example.com", // PII - should not be exposed
TeamID: "team-456",
TenantID: "tenant-789",
},
Client: context.Client{
Type: "web",
IP: "192.168.1.1", // Should not be exposed
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test privacy"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Check that sensitive info is not in any system prompts
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.NotContains(t, msg.Content, "user@example.com", "Subject should not be in prompts")
assert.NotContains(t, msg.Content, "192.168.1.1", "IP should not be in prompts")
}
}
})
t.Run("ContextVariablesInPrompts", func(t *testing.T) {
// Set up global prompts with variables
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "User ID: $CTX.USER_ID, Team: $CTX.TEAM_ID, Custom: $CTX.MY_VAR"},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{
Authorized: &types.AuthorizedInfo{
UserID: "user-abc",
TeamID: "team-xyz",
},
Metadata: map[string]interface{}{
"MY_VAR": "my-value",
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test variables"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Find the global prompt and verify variables are replaced
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && !found {
if assert.Contains(t, msg.Content, "User ID: user-abc") {
found = true
assert.Contains(t, msg.Content, "Team: team-xyz")
assert.Contains(t, msg.Content, "Custom: my-value")
}
}
}
assert.True(t, found, "Should find global prompt with replaced variables")
})
t.Run("SystemVariablesReplacement", func(t *testing.T) {
// Set up global prompts with $SYS.* variables
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "Time: $SYS.TIME, Date: $SYS.DATE, Datetime: $SYS.DATETIME, Weekday: $SYS.WEEKDAY"},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test system variables"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Find the global prompt and verify $SYS.* variables are replaced
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
// Should NOT contain $SYS. prefix (variables should be replaced)
if !assert.NotContains(t, msg.Content, "$SYS.TIME") {
continue
}
if !assert.NotContains(t, msg.Content, "$SYS.DATE") {
continue
}
if !assert.NotContains(t, msg.Content, "$SYS.DATETIME") {
continue
}
if !assert.NotContains(t, msg.Content, "$SYS.WEEKDAY") {
continue
}
// Should contain "Time:", "Date:", etc. with actual values
assert.Contains(t, msg.Content, "Time:")
assert.Contains(t, msg.Content, "Date:")
assert.Contains(t, msg.Content, "Datetime:")
assert.Contains(t, msg.Content, "Weekday:")
found = true
break
}
}
assert.True(t, found, "Should find global prompt with replaced $SYS.* variables")
})
t.Run("EnvVariablesReplacement", func(t *testing.T) {
// Set test environment variable
t.Setenv("TEST_PROMPT_VAR", "env-test-value")
// Set up global prompts with $ENV.* variables
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "Env Value: $ENV.TEST_PROMPT_VAR, Not Exist: $ENV.NOT_EXIST_VAR_XYZ"},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test env variables"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Find the global prompt and verify $ENV.* variables are replaced
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
// Should NOT contain $ENV. prefix for existing vars
if !assert.NotContains(t, msg.Content, "$ENV.TEST_PROMPT_VAR") {
continue
}
// Should contain the actual env value
assert.Contains(t, msg.Content, "Env Value: env-test-value")
// Non-existent env var should be replaced with empty string
assert.Contains(t, msg.Content, "Not Exist: ")
assert.NotContains(t, msg.Content, "$ENV.NOT_EXIST_VAR_XYZ")
found = true
break
}
}
assert.True(t, found, "Should find global prompt with replaced $ENV.* variables")
})
t.Run("AllVariableTypesReplacement", func(t *testing.T) {
// Set test environment variable
t.Setenv("TEST_APP_NAME", "MyTestApp")
// Set up global prompts with all variable types
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: `System Info:
- Time: $SYS.TIME
- Date: $SYS.DATE
- App: $ENV.TEST_APP_NAME
- User: $CTX.USER_ID
- Custom: $CTX.CUSTOM_KEY
- Assistant: $CTX.ASSISTANT_NAME`},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{
Authorized: &types.AuthorizedInfo{
UserID: "all-vars-user",
},
Metadata: map[string]interface{}{
"CUSTOM_KEY": "custom-value-123",
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test all variables"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Find the global prompt and verify ALL variable types are replaced
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && !found {
content := msg.Content
// Check $SYS.* replaced
if assert.NotContains(t, content, "$SYS.TIME") &&
assert.NotContains(t, content, "$SYS.DATE") {
// Check $ENV.* replaced
assert.NotContains(t, content, "$ENV.TEST_APP_NAME")
assert.Contains(t, content, "App: MyTestApp")
// Check $CTX.* replaced
assert.NotContains(t, content, "$CTX.USER_ID")
assert.Contains(t, content, "User: all-vars-user")
assert.NotContains(t, content, "$CTX.CUSTOM_KEY")
assert.Contains(t, content, "Custom: custom-value-123")
// Check assistant name from $CTX.ASSISTANT_NAME
assert.NotContains(t, content, "$CTX.ASSISTANT_NAME")
found = true
}
}
}
assert.True(t, found, "Should find global prompt with all variable types replaced")
})
t.Run("PromptPresetFromHook", func(t *testing.T) {
// Load fullfields assistant which has prompt_presets
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
require.NotNil(t, ast.PromptPresets)
require.Contains(t, ast.PromptPresets, "chat.friendly")
ctx := &context.Context{}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test preset from hook"},
}
// Hook returns prompt_preset
createResponse := &context.HookCreateResponse{
PromptPreset: "chat.friendly",
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should have system prompts from the preset
hasSystemPrompt := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
hasSystemPrompt = true
// Verify it's from the friendly preset (check content)
assert.Contains(t, msg.Content, "friendly", "Should use friendly preset prompts")
break
}
}
assert.True(t, hasSystemPrompt, "Should have system prompts from preset")
})
t.Run("PromptPresetFromMetadata", func(t *testing.T) {
// Load fullfields assistant which has prompt_presets
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
ctx := &context.Context{
Metadata: map[string]interface{}{
"__prompt_preset": "chat.professional",
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test preset from metadata"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Should have system prompts from the preset
hasSystemPrompt := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
hasSystemPrompt = true
// Verify it's from the professional preset
assert.Contains(t, msg.Content, "professional", "Should use professional preset prompts")
break
}
}
assert.True(t, hasSystemPrompt, "Should have system prompts from preset")
})
t.Run("PromptPresetHookOverridesMetadata", func(t *testing.T) {
// Load fullfields assistant
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
ctx := &context.Context{
Metadata: map[string]interface{}{
"__prompt_preset": "chat.professional", // Lower priority
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test hook overrides metadata"},
}
// Hook returns different preset (higher priority)
createResponse := &context.HookCreateResponse{
PromptPreset: "chat.friendly",
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should use hook's preset, not metadata's
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.Contains(t, msg.Content, "friendly", "Hook preset should override metadata preset")
break
}
}
})
t.Run("PromptPresetNotFound", func(t *testing.T) {
// Load fullfields assistant
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
ctx := &context.Context{
Metadata: map[string]interface{}{
"__prompt_preset": "non.existent.preset",
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test non-existent preset"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Should fallback to default prompts (not crash)
hasSystemPrompt := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
hasSystemPrompt = true
break
}
}
assert.True(t, hasSystemPrompt, "Should fallback to default prompts when preset not found")
})
t.Run("DisableGlobalPromptsFromHook", func(t *testing.T) {
// Set global prompts
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "GLOBAL_PROMPT_MARKER"},
})
defer assistant.SetGlobalPrompts(nil)
// Load an assistant that does NOT disable global prompts
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
require.False(t, ast.DisableGlobalPrompts)
ctx := &context.Context{}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test disable from hook"},
}
// Hook disables global prompts
disableTrue := true
createResponse := &context.HookCreateResponse{
DisableGlobalPrompts: &disableTrue,
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should NOT have global prompt
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.NotContains(t, msg.Content, "GLOBAL_PROMPT_MARKER", "Global prompts should be disabled by hook")
}
}
})
t.Run("DisableGlobalPromptsFromMetadata", func(t *testing.T) {
// Set global prompts
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "GLOBAL_PROMPT_MARKER_2"},
})
defer assistant.SetGlobalPrompts(nil)
// Load an assistant that does NOT disable global prompts
ast, err := assistant.Get("yaobots")
require.NoError(t, err)
ctx := &context.Context{
Metadata: map[string]interface{}{
"__disable_global_prompts": true,
},
}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test disable from metadata"},
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Should NOT have global prompt
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.NotContains(t, msg.Content, "GLOBAL_PROMPT_MARKER_2", "Global prompts should be disabled by metadata")
}
}
})
t.Run("EnableGlobalPromptsOverrideAssistant", func(t *testing.T) {
// Set global prompts
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "GLOBAL_ENABLED_MARKER"},
})
defer assistant.SetGlobalPrompts(nil)
// Load fullfields assistant which has disable_global_prompts: true
ast, err := assistant.Get("tests.fullfields")
require.NoError(t, err)
require.True(t, ast.DisableGlobalPrompts)
ctx := &context.Context{}
messages := []context.Message{
{Role: context.RoleUser, Content: "Test enable override"},
}
// Hook enables global prompts (overrides assistant's disable)
disableFalse := false
createResponse := &context.HookCreateResponse{
DisableGlobalPrompts: &disableFalse,
}
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should have global prompt (hook enabled it)
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && msg.Content == "GLOBAL_ENABLED_MARKER" {
found = true
break
}
}
assert.True(t, found, "Global prompts should be enabled by hook override")
})
}
// TestPromptPresetAssistant tests the tests.promptpreset assistant with Create Hook
func TestPromptPresetAssistant(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
t.Run("LoadPromptPresetAssistant", func(t *testing.T) {
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
require.NotNil(t, ast)
assert.Equal(t, "tests.promptpreset", ast.ID)
assert.Equal(t, "Prompt Preset Test", ast.Name)
assert.False(t, ast.DisableGlobalPrompts)
// Should have prompt presets loaded
require.NotNil(t, ast.PromptPresets)
assert.Contains(t, ast.PromptPresets, "mode.friendly")
assert.Contains(t, ast.PromptPresets, "mode.professional")
// Should have script
assert.NotNil(t, ast.Script)
})
t.Run("CreateHookSelectsFriendlyPreset", func(t *testing.T) {
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-friendly-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "use friendly mode please"},
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.friendly", createResponse.PromptPreset)
// Build request
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should have friendly preset marker in one of the system messages
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && containsString(msg.Content, "FRIENDLY_PRESET_MARKER") {
found = true
break
}
}
assert.True(t, found, "Should use friendly preset from Create Hook")
})
t.Run("CreateHookSelectsProfessionalPreset", func(t *testing.T) {
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-professional-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "use professional tone"},
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.professional", createResponse.PromptPreset)
// Build request
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should have professional preset marker in one of the system messages
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && containsString(msg.Content, "PROFESSIONAL_PRESET_MARKER") {
found = true
break
}
}
assert.True(t, found, "Should use professional preset from Create Hook")
})
t.Run("CreateHookDisablesGlobalPrompts", func(t *testing.T) {
// Set global prompts
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "GLOBAL_MARKER_FOR_DISABLE_TEST"},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-disable-global-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "disable global prompts"},
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, createResponse)
require.NotNil(t, createResponse.DisableGlobalPrompts)
assert.True(t, *createResponse.DisableGlobalPrompts)
// Build request
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should NOT have global prompt
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.NotContains(t, msg.Content, "GLOBAL_MARKER_FOR_DISABLE_TEST")
}
}
})
t.Run("CreateHookPresetAndDisableGlobal", func(t *testing.T) {
// Set global prompts
assistant.SetGlobalPrompts([]store.Prompt{
{Role: "system", Content: "GLOBAL_MARKER_COMBINED_TEST"},
})
defer assistant.SetGlobalPrompts(nil)
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-combined-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "friendly no global"},
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "mode.friendly", createResponse.PromptPreset)
require.NotNil(t, createResponse.DisableGlobalPrompts)
assert.True(t, *createResponse.DisableGlobalPrompts)
// Build request
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should have friendly preset but NOT global
hasFriendly := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem {
assert.NotContains(t, msg.Content, "GLOBAL_MARKER_COMBINED_TEST")
if containsString(msg.Content, "FRIENDLY_PRESET_MARKER") {
hasFriendly = true
}
}
}
assert.True(t, hasFriendly, "Should have friendly preset")
})
t.Run("CreateHookUnknownPresetFallback", func(t *testing.T) {
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-unknown-preset-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "unknown preset test"},
}
// Call Create hook
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, createResponse)
assert.Equal(t, "non.existent.preset", createResponse.PromptPreset)
// Build request - should not error, fallback to default
finalMessages, _, err := ast.BuildRequest(ctx, messages, createResponse)
require.NoError(t, err)
// Should fallback to default prompts
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && containsString(msg.Content, "DEFAULT_PROMPT_MARKER") {
found = true
break
}
}
assert.True(t, found, "Should fallback to default prompts when preset not found")
})
t.Run("CreateHookReturnsNull", func(t *testing.T) {
ast, err := assistant.Get("tests.promptpreset")
require.NoError(t, err)
ctx := newPromptTestContext("chat-null-test", "tests.promptpreset")
messages := []context.Message{
{Role: context.RoleUser, Content: "just a normal message"},
}
// Call Create hook - should return nil
createResponse, err := ast.Script.Create(ctx, messages)
require.NoError(t, err)
assert.Nil(t, createResponse)
// Build request with nil createResponse
finalMessages, _, err := ast.BuildRequest(ctx, messages, nil)
require.NoError(t, err)
// Should use default prompts
found := false
for _, msg := range finalMessages {
if msg.Role == context.RoleSystem && containsString(msg.Content, "DEFAULT_PROMPT_MARKER") {
found = true
break
}
}
assert.True(t, found, "Should use default prompts when hook returns null")
})
}

View file

@ -4,13 +4,13 @@ import (
"fmt"
"os"
"path/filepath"
"regexp"
"strings"
"time"
jsoniter "github.com/json-iterator/go"
"github.com/spf13/cast"
"github.com/yaoapp/gou/application"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/fs"
v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/yao/agent/assistant/hook"
@ -26,9 +26,10 @@ import (
var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil
var search interface{} = nil
var modelCapabilities map[string]ModelCapabilities = map[string]ModelCapabilities{}
var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
// LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error {
@ -130,8 +131,13 @@ func SetStorage(s store.Store) {
storage = s
}
// GetStorage returns the storage (for testing purposes)
func GetStorage() store.Store {
return storage
}
// SetModelCapabilities set the model capabilities configuration
func SetModelCapabilities(capabilities map[string]ModelCapabilities) {
func SetModelCapabilities(capabilities map[string]gouOpenAI.Capabilities) {
modelCapabilities = capabilities
}
@ -145,6 +151,20 @@ func SetGlobalUses(uses *context.Uses) {
globalUses = uses
}
// SetGlobalPrompts set the global prompts from agent/prompts.yml
func SetGlobalPrompts(prompts []store.Prompt) {
globalPrompts = prompts
}
// GetGlobalPrompts returns the global prompts with variables parsed
// ctx: context variables for parsing $CTX.* variables
func GetGlobalPrompts(ctx map[string]string) []store.Prompt {
if len(globalPrompts) == 0 {
return nil
}
return store.Prompts(globalPrompts).Parse(ctx)
}
// SetCache set the cache
func SetCache(capacity int) {
ClearCache()
@ -180,7 +200,8 @@ func LoadStore(id string) (*Assistant, error) {
return nil, fmt.Errorf("storage is not set")
}
storeModel, err := storage.GetAssistant(id)
// Request all fields when loading assistant from store
storeModel, err := storage.GetAssistant(id, store.AssistantFullFields)
if err != nil {
return nil, err
}
@ -198,6 +219,15 @@ func LoadStore(id string) (*Assistant, error) {
// Create assistant from store model
assistant = &Assistant{AssistantModel: *storeModel}
// Load script from source field if present
if assistant.Source != "" {
script, err := loadSource(assistant.Source, assistant.ID)
if err != nil {
return nil, err
}
assistant.Script = script
}
// Initialize the assistant
err = assistant.initialize()
if err != nil {
@ -266,10 +296,10 @@ func LoadPath(path string) (*Assistant, error) {
updatedAt := int64(0)
// prompts
// prompts (default prompts from prompts.yml)
promptsfile := filepath.Join(path, "prompts.yml")
if has, _ := app.Exists(promptsfile); has {
prompts, ts, err := loadPrompts(promptsfile, path)
prompts, ts, err := store.LoadPrompts(promptsfile, path)
if err != nil {
return nil, err
}
@ -278,6 +308,19 @@ func LoadPath(path string) (*Assistant, error) {
updatedAt = ts
}
// prompt_presets (from prompts directory, key is filename without extension)
promptsDir := filepath.Join(path, "prompts")
if has, _ := app.Exists(promptsDir); has {
presets, ts, err := store.LoadPromptPresets(promptsDir, path)
if err != nil {
return nil, err
}
if len(presets) > 0 {
data["prompt_presets"] = presets
updatedAt = max(updatedAt, ts)
}
}
// load script
scriptfile := filepath.Join(path, "src", "index.ts")
if has, _ := app.Exists(scriptfile); has {
@ -382,6 +425,11 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Automated = v
}
// DisableGlobalPrompts
if v, ok := data["disable_global_prompts"].(bool); ok {
assistant.DisableGlobalPrompts = v
}
// Readonly
if v, ok := data["readonly"].(bool); ok {
assistant.Readonly = v
@ -417,6 +465,15 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Connector = connector
}
// connector_options
if connOpts, has := data["connector_options"]; has {
opts, err := store.ToConnectorOptions(connOpts)
if err != nil {
return nil, err
}
assistant.ConnectorOptions = opts
}
// tags
if v, has := data["tags"]; has {
switch vv := v.(type) {
@ -508,33 +565,25 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
}
}
// tools
if tools, has := data["tools"]; has {
switch vv := tools.(type) {
case []store.Tool:
assistant.Tools = &store.ToolCalls{
Tools: vv,
Prompts: assistant.Prompts,
}
case store.ToolCalls:
assistant.Tools = &vv
default:
raw, err := jsoniter.Marshal(tools)
// prompt_presets
if presets, has := data["prompt_presets"]; has {
promptPresets, err := store.ToPromptPresets(presets)
if err != nil {
return nil, fmt.Errorf("tools format error %s", err.Error())
return nil, err
}
assistant.PromptPresets = promptPresets
}
var tools store.ToolCalls
err = jsoniter.Unmarshal(raw, &tools)
if err != nil {
return nil, fmt.Errorf("tools format error %s", err.Error())
}
assistant.Tools = &tools
}
// source (hook script code) - store the source code
if source, ok := data["source"].(string); ok {
assistant.Source = source
}
// tools - deprecated, now handled by MCP
// if tools, has := data["tools"]; has {
// ... removed ...
// }
// kb
if kb, has := data["kb"]; has {
knowledgeBase, err := store.ToKnowledgeBase(kb)
@ -562,7 +611,29 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
assistant.Workflow = wf
}
// script
// uses (wrapper configurations for vision, audio, etc.)
if uses, has := data["uses"]; has {
switch v := uses.(type) {
case *context.Uses:
assistant.Uses = v
case context.Uses:
assistant.Uses = &v
default:
raw, err := jsoniter.Marshal(v)
if err != nil {
return nil, err
}
var usesConfig context.Uses
err = jsoniter.Unmarshal(raw, &usesConfig)
if err != nil {
return nil, err
}
assistant.Uses = &usesConfig
}
}
// script loading priority: script field > source field
// If script field exists, use it; otherwise try source field
if data["script"] != nil {
switch v := data["script"].(type) {
case string:
@ -577,6 +648,13 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
case *v8.Script:
assistant.Script = &hook.Script{Script: v}
}
} else if assistant.Source != "" {
// Load from source field if script is not provided
script, err := loadSource(assistant.Source, assistant.ID)
if err != nil {
return nil, err
}
assistant.Script = script
}
// created_at
@ -606,43 +684,6 @@ func loadMap(data map[string]interface{}) (*Assistant, error) {
return assistant, nil
}
func loadPrompts(file string, root string) (string, int64, error) {
app, err := fs.Get("app")
if err != nil {
return "", 0, err
}
ts, err := app.ModTime(file)
if err != nil {
return "", 0, err
}
prompts, err := app.ReadFile(file)
if err != nil {
return "", 0, err
}
re := regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`)
prompts = re.ReplaceAllFunc(prompts, func(s []byte) []byte {
asset := re.FindStringSubmatch(string(s))[1]
assetFile := filepath.Join(root, "assets", asset)
assetContent, err := app.ReadFile(assetFile)
if err != nil {
return []byte("")
}
// Add proper YAML formatting for content
lines := strings.Split(string(assetContent), "\n")
formattedContent := "|\n"
for _, line := range lines {
formattedContent += " " + line + "\n"
}
return []byte(formattedContent)
})
return string(prompts), ts.UnixNano(), nil
}
func loadScript(file string, root string) (*hook.Script, int64, error) {
app, err := fs.Get("app")

View file

@ -0,0 +1,937 @@
package assistant_test
import (
stdContext "context"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/openapi/oauth/types"
)
// TestLoadStoreWithSource tests loading assistant from database with Source field
func TestLoadStoreWithSource(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Create assistant with Source
assistantID := "test.store-with-source"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Assistant With Source",
Type: "assistant",
Connector: "gpt-4o",
Description: "Test assistant loaded from store with source code",
Prompts: []store.Prompt{
{Role: "system", Content: "You are a helpful assistant."},
},
Options: map[string]interface{}{
"temperature": 0.7,
},
Tags: []string{"Test", "Source"},
// Simple Create hook that returns null
Source: `
// @ts-nocheck
function Create(ctx, messages) {
return null;
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
// Save to database
err := ast.Save()
require.NoError(t, err)
// Cleanup after test
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
// Clear cache to ensure fresh load from database
assistant.GetCache().Clear()
// Load from store
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify basic fields
assert.Equal(t, assistantID, loaded.ID)
assert.Equal(t, "Test Assistant With Source", loaded.Name)
assert.Equal(t, "assistant", loaded.Type)
assert.Equal(t, "Test assistant loaded from store with source code", loaded.Description)
// Verify prompts
require.NotNil(t, loaded.Prompts)
assert.Len(t, loaded.Prompts, 1)
assert.Equal(t, "system", loaded.Prompts[0].Role)
assert.Equal(t, "You are a helpful assistant.", loaded.Prompts[0].Content)
// Verify options
assert.NotNil(t, loaded.Options)
assert.Equal(t, 0.7, loaded.Options["temperature"])
// Verify tags
assert.NotNil(t, loaded.Tags)
assert.Contains(t, loaded.Tags, "Test")
assert.Contains(t, loaded.Tags, "Source")
// Verify script was compiled from source
assert.NotNil(t, loaded.Script, "Script should be compiled from Source field")
// Verify source is stored
assert.NotEmpty(t, loaded.Source)
}
// TestLoadStoreWithoutSource tests loading assistant from database without Source field
func TestLoadStoreWithoutSource(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Create assistant without Source
assistantID := "test.store-without-source"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Assistant Without Source",
Type: "assistant",
Connector: "gpt-4o",
Description: "Test assistant loaded from store without source code",
Prompts: []store.Prompt{
{Role: "system", Content: "You are a helpful assistant without hooks."},
},
Options: map[string]interface{}{
"temperature": 0.5,
"max_tokens": 1000,
},
Tags: []string{"Test", "NoSource"},
CreatedAt: now,
UpdatedAt: now,
},
}
// Save to database
err := ast.Save()
require.NoError(t, err)
// Cleanup after test
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
// Clear cache to ensure fresh load from database
assistant.GetCache().Clear()
// Load from store
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify basic fields
assert.Equal(t, assistantID, loaded.ID)
assert.Equal(t, "Test Assistant Without Source", loaded.Name)
assert.Equal(t, "assistant", loaded.Type)
assert.Equal(t, "Test assistant loaded from store without source code", loaded.Description)
// Verify prompts
require.NotNil(t, loaded.Prompts)
assert.Len(t, loaded.Prompts, 1)
assert.Equal(t, "system", loaded.Prompts[0].Role)
// Verify options
assert.NotNil(t, loaded.Options)
assert.Equal(t, 0.5, loaded.Options["temperature"])
assert.Equal(t, float64(1000), loaded.Options["max_tokens"])
// Verify tags
assert.NotNil(t, loaded.Tags)
assert.Contains(t, loaded.Tags, "Test")
assert.Contains(t, loaded.Tags, "NoSource")
// Verify script is nil (no source)
assert.Nil(t, loaded.Script, "Script should be nil when no Source field")
assert.Empty(t, loaded.Source)
}
// newStoreTestContext creates a Context for testing with commonly used fields pre-populated.
func newStoreTestContext(chatID, assistantID string) *context.Context {
return &context.Context{
Context: stdContext.Background(),
ChatID: chatID,
AssistantID: assistantID,
Connector: "",
Locale: "en-us",
Theme: "light",
Client: context.Client{
Type: "web",
UserAgent: "TestAgent/1.0",
IP: "127.0.0.1",
},
Referer: context.RefererAPI,
Accept: context.AcceptWebCUI,
Route: "",
Metadata: make(map[string]interface{}),
Authorized: &types.AuthorizedInfo{
Subject: "test-user",
ClientID: "test-client-id",
Scope: "openid profile email",
SessionID: "test-session-id",
UserID: "test-user-123",
TeamID: "test-team-456",
TenantID: "test-tenant-789",
},
}
}
// TestLoadStoreWithSourceExecuteHook tests that Source-based script is properly compiled and can execute
func TestLoadStoreWithSourceExecuteHook(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Create assistant with a working Create hook
assistantID := "test.store-source-hook"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Source Hook",
Type: "assistant",
Connector: "gpt-4o",
Prompts: []store.Prompt{
{Role: "system", Content: "Default prompt"},
},
// Create hook that modifies temperature and adds metadata
Source: `
// @ts-nocheck
function Create(ctx: any, messages: any[]): any {
return {
temperature: 0.9,
metadata: {
hook_executed: true,
chat_id: ctx.chat_id
}
};
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
// Save to database
err := ast.Save()
require.NoError(t, err)
// Cleanup after test
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
// Clear cache
assistant.GetCache().Clear()
// Load from store
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
require.NotNil(t, loaded.Script, "Script should be compiled from Source")
// Verify the script object exists and is usable
assert.NotNil(t, loaded.Script.Script)
// Execute the Create hook
ctx := newStoreTestContext("test-chat-id", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err, "Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
// Verify temperature was set
require.NotNil(t, res.Temperature, "Temperature should be set")
assert.Equal(t, 0.9, *res.Temperature, "Temperature should be 0.9")
// Verify metadata was set
require.NotNil(t, res.Metadata, "Metadata should be set")
assert.Equal(t, true, res.Metadata["hook_executed"], "hook_executed should be true")
assert.Equal(t, "test-chat-id", res.Metadata["chat_id"], "chat_id should match context")
}
// TestLoadStoreWithPromptPresets tests loading assistant with prompt presets from database
func TestLoadStoreWithPromptPresets(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-with-presets"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test With Presets",
Type: "assistant",
Connector: "gpt-4o",
Prompts: []store.Prompt{
{Role: "system", Content: "Default prompt"},
},
PromptPresets: map[string][]store.Prompt{
"friendly": {
{Role: "system", Content: "You are a friendly assistant."},
},
"professional": {
{Role: "system", Content: "You are a professional assistant."},
},
"mode.casual": {
{Role: "system", Content: "You are a casual assistant."},
},
},
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify prompt presets
require.NotNil(t, loaded.PromptPresets)
assert.Len(t, loaded.PromptPresets, 3)
friendlyPreset, ok := loaded.PromptPresets["friendly"]
assert.True(t, ok)
assert.Len(t, friendlyPreset, 1)
assert.Equal(t, "You are a friendly assistant.", friendlyPreset[0].Content)
professionalPreset, ok := loaded.PromptPresets["professional"]
assert.True(t, ok)
assert.Len(t, professionalPreset, 1)
assert.Equal(t, "You are a professional assistant.", professionalPreset[0].Content)
casualPreset, ok := loaded.PromptPresets["mode.casual"]
assert.True(t, ok)
assert.Len(t, casualPreset, 1)
assert.Equal(t, "You are a casual assistant.", casualPreset[0].Content)
}
// TestLoadStoreWithDisableGlobalPrompts tests loading assistant with disable_global_prompts flag
func TestLoadStoreWithDisableGlobalPrompts(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-disable-global"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Disable Global Prompts",
Type: "assistant",
Connector: "gpt-4o",
DisableGlobalPrompts: true,
Prompts: []store.Prompt{
{Role: "system", Content: "Only this prompt should be used."},
},
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
assert.True(t, loaded.DisableGlobalPrompts)
}
// TestLoadStoreCaching tests that loaded assistants are cached
func TestLoadStoreCaching(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-caching"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Caching",
Type: "assistant",
Connector: "gpt-4o",
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
// First load
ast1, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, ast1)
// Second load - should be from cache
ast2, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, ast2)
// Should be the same instance (from cache)
assert.Same(t, ast1, ast2)
}
// TestLoadStoreNotFound tests loading non-existent assistant
func TestLoadStoreNotFound(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistant.GetCache().Clear()
_, err := assistant.Get("non-existent-assistant-id-12345")
assert.Error(t, err)
}
// TestLoadStoreWithAllFields tests loading assistant with comprehensive fields
func TestLoadStoreWithAllFields(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-all-fields"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test All Fields",
Type: "assistant",
Avatar: "/api/icons/test.png",
Connector: "gpt-4o",
Description: "Test assistant with all fields",
Tags: []string{"Test", "AllFields", "Complete"},
Readonly: true,
Public: true,
Share: "team",
Mentionable: true,
Automated: false,
Sort: 100,
Options: map[string]interface{}{
"temperature": 0.8,
"max_tokens": 2000,
},
Prompts: []store.Prompt{
{Role: "system", Content: "You are a test assistant."},
{Role: "system", Content: "Follow all instructions carefully."},
},
PromptPresets: map[string][]store.Prompt{
"default": {
{Role: "system", Content: "Default mode prompt."},
},
},
DisableGlobalPrompts: true,
Placeholder: &store.Placeholder{
Title: "Test Placeholder",
Description: "This is a test placeholder",
Prompts: []string{"Test prompt 1", "Test prompt 2"},
},
Source: `
// @ts-nocheck
function Create(ctx: any, messages: any[]): any {
return {
temperature: 0.5,
metadata: {
assistant_name: "Test All Fields",
executed: true
}
};
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
// Verify all fields
assert.Equal(t, assistantID, loaded.ID)
assert.Equal(t, "Test All Fields", loaded.Name)
assert.Equal(t, "assistant", loaded.Type)
assert.Equal(t, "/api/icons/test.png", loaded.Avatar)
assert.Equal(t, "Test assistant with all fields", loaded.Description)
// Boolean fields
assert.True(t, loaded.Readonly)
assert.True(t, loaded.Public)
assert.Equal(t, "team", loaded.Share)
assert.True(t, loaded.Mentionable)
assert.False(t, loaded.Automated)
assert.True(t, loaded.DisableGlobalPrompts)
assert.Equal(t, 100, loaded.Sort)
// Tags
assert.Len(t, loaded.Tags, 3)
assert.Contains(t, loaded.Tags, "Test")
assert.Contains(t, loaded.Tags, "AllFields")
assert.Contains(t, loaded.Tags, "Complete")
// Options
assert.Equal(t, 0.8, loaded.Options["temperature"])
assert.Equal(t, float64(2000), loaded.Options["max_tokens"])
// Prompts
assert.Len(t, loaded.Prompts, 2)
// Prompt presets
assert.NotNil(t, loaded.PromptPresets)
assert.Contains(t, loaded.PromptPresets, "default")
// Placeholder
assert.NotNil(t, loaded.Placeholder)
assert.Equal(t, "Test Placeholder", loaded.Placeholder.Title)
assert.Equal(t, "This is a test placeholder", loaded.Placeholder.Description)
assert.Len(t, loaded.Placeholder.Prompts, 2)
// Script from source
assert.NotNil(t, loaded.Script)
assert.NotEmpty(t, loaded.Source)
// Execute the Create hook to verify it works
ctx := newStoreTestContext("test-chat-all-fields", assistantID)
messages := []context.Message{{Role: "user", Content: "Test message"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err, "Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
// Verify hook returned expected values
require.NotNil(t, res.Temperature, "Temperature should be set")
assert.Equal(t, 0.5, *res.Temperature, "Temperature should be 0.5")
require.NotNil(t, res.Metadata, "Metadata should be set")
assert.Equal(t, "Test All Fields", res.Metadata["assistant_name"], "assistant_name should match")
assert.Equal(t, true, res.Metadata["executed"], "executed should be true")
}
// TestLoadStoreHookWithTypeScript tests that TypeScript features work in Source field
func TestLoadStoreHookWithTypeScript(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-typescript-hook"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test TypeScript Hook",
Type: "assistant",
Connector: "gpt-4o",
Prompts: []store.Prompt{
{Role: "system", Content: "Default prompt"},
},
// TypeScript code with type annotations and interfaces
Source: `
// TypeScript interfaces
interface CreateContext {
chat_id: string;
assistant_id: string;
locale: string;
authorized?: {
user_id: string;
team_id: string;
};
}
interface Message {
role: string;
content: string | object;
}
interface CreateResponse {
temperature?: number;
messages?: Message[];
metadata?: Record<string, any>;
}
// Create hook with full TypeScript syntax
function Create(ctx: CreateContext, messages: Message[]): CreateResponse | null {
// Type-safe access to context
const chatId: string = ctx.chat_id || "unknown";
const locale: string = ctx.locale || "en-us";
const userId: string = ctx.authorized?.user_id || "anonymous";
// Process messages
const userMessages: Message[] = messages.filter((m: Message) => m.role === "user");
const messageCount: number = userMessages.length;
// Return typed response
return {
temperature: 0.7,
messages: [
{
role: "system",
content: "TypeScript hook executed successfully"
}
],
metadata: {
chat_id: chatId,
locale: locale,
user_id: userId,
message_count: messageCount,
typescript_features: true
}
};
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
require.NotNil(t, loaded.Script, "Script should be compiled from TypeScript Source")
// Execute the Create hook
ctx := newStoreTestContext("ts-test-chat", assistantID)
messages := []context.Message{
{Role: "user", Content: "Hello"},
{Role: "assistant", Content: "Hi there"},
{Role: "user", Content: "How are you?"},
}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err, "TypeScript Create hook should execute without error")
require.NotNil(t, res, "Create hook should return a response")
// Verify temperature
require.NotNil(t, res.Temperature)
assert.Equal(t, 0.7, *res.Temperature)
// Verify messages
require.Len(t, res.Messages, 1)
assert.Equal(t, context.RoleSystem, res.Messages[0].Role)
assert.Equal(t, "TypeScript hook executed successfully", res.Messages[0].Content)
// Verify metadata
require.NotNil(t, res.Metadata)
assert.Equal(t, "ts-test-chat", res.Metadata["chat_id"])
assert.Equal(t, "en-us", res.Metadata["locale"])
assert.Equal(t, "test-user-123", res.Metadata["user_id"])
assert.Equal(t, float64(2), res.Metadata["message_count"]) // 2 user messages
assert.Equal(t, true, res.Metadata["typescript_features"])
}
// TestLoadStoreHookReturnNull tests that hook returning null works correctly
func TestLoadStoreHookReturnNull(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-hook-null"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Hook Return Null",
Type: "assistant",
Connector: "gpt-4o",
Source: `
function Create(ctx: any, messages: any[]): any {
// Return null to indicate no modifications
return null;
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
require.NotNil(t, loaded.Script)
ctx := newStoreTestContext("null-test-chat", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err, "Hook returning null should not error")
assert.Nil(t, res, "Hook returning null should return nil response")
}
// TestLoadStoreHookWithPromptPreset tests that hook can return prompt_preset
func TestLoadStoreHookWithPromptPreset(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-hook-preset"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Hook Prompt Preset",
Type: "assistant",
Connector: "gpt-4o",
Prompts: []store.Prompt{
{Role: "system", Content: "Default prompt"},
},
PromptPresets: map[string][]store.Prompt{
"friendly": {
{Role: "system", Content: "You are a friendly assistant."},
},
"professional": {
{Role: "system", Content: "You are a professional assistant."},
},
},
Source: `
function Create(ctx: any, messages: any[]): any {
// Check first message to determine preset
const firstMsg = messages[0];
if (firstMsg && typeof firstMsg.content === "string") {
if (firstMsg.content.includes("friendly")) {
return { prompt_preset: "friendly" };
}
if (firstMsg.content.includes("professional")) {
return { prompt_preset: "professional" };
}
}
return null;
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
require.NotNil(t, loaded.Script)
// Test friendly preset selection
t.Run("SelectFriendlyPreset", func(t *testing.T) {
ctx := newStoreTestContext("preset-test-1", assistantID)
messages := []context.Message{{Role: "user", Content: "Be friendly please"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, res)
assert.Equal(t, "friendly", res.PromptPreset)
})
// Test professional preset selection
t.Run("SelectProfessionalPreset", func(t *testing.T) {
ctx := newStoreTestContext("preset-test-2", assistantID)
messages := []context.Message{{Role: "user", Content: "Be professional"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, res)
assert.Equal(t, "professional", res.PromptPreset)
})
// Test no preset (returns null)
t.Run("NoPreset", func(t *testing.T) {
ctx := newStoreTestContext("preset-test-3", assistantID)
messages := []context.Message{{Role: "user", Content: "Hello"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err)
assert.Nil(t, res)
})
}
// TestLoadStoreHookDisableGlobalPrompts tests that hook can disable global prompts
func TestLoadStoreHookDisableGlobalPrompts(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
assistantID := "test.store-hook-disable-global"
now := time.Now().UnixNano()
ast := &assistant.Assistant{
AssistantModel: store.AssistantModel{
ID: assistantID,
Name: "Test Hook Disable Global",
Type: "assistant",
Connector: "gpt-4o",
Source: `
function Create(ctx: any, messages: any[]): any {
const firstMsg = messages[0];
if (firstMsg && typeof firstMsg.content === "string") {
if (firstMsg.content.includes("disable_global")) {
return { disable_global_prompts: true };
}
if (firstMsg.content.includes("enable_global")) {
return { disable_global_prompts: false };
}
}
return null;
}
`,
CreatedAt: now,
UpdatedAt: now,
},
}
err := ast.Save()
require.NoError(t, err)
defer func() {
storage := assistant.GetStorage()
if storage != nil {
storage.DeleteAssistant(assistantID)
}
assistant.GetCache().Clear()
}()
assistant.GetCache().Clear()
loaded, err := assistant.Get(assistantID)
require.NoError(t, err)
require.NotNil(t, loaded)
require.NotNil(t, loaded.Script)
// Test disable global prompts
t.Run("DisableGlobalPrompts", func(t *testing.T) {
ctx := newStoreTestContext("disable-test-1", assistantID)
messages := []context.Message{{Role: "user", Content: "disable_global prompts"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, res)
require.NotNil(t, res.DisableGlobalPrompts)
assert.True(t, *res.DisableGlobalPrompts)
})
// Test enable global prompts
t.Run("EnableGlobalPrompts", func(t *testing.T) {
ctx := newStoreTestContext("disable-test-2", assistantID)
messages := []context.Message{{Role: "user", Content: "enable_global prompts"}}
res, err := loaded.Script.Create(ctx, messages)
require.NoError(t, err)
require.NotNil(t, res)
require.NotNil(t, res.DisableGlobalPrompts)
assert.False(t, *res.DisableGlobalPrompts)
})
}

View file

@ -1,435 +1,502 @@
package assistant
// func prepare(t *testing.T) {
// test.Prepare(t, config.Conf)
// }
import (
"testing"
// func TestLoad_LoadPath(t *testing.T) {
// prepare(t)
// defer test.Clean()
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
store "github.com/yaoapp/yao/agent/store/types"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
// assistant, err := LoadPath("/assistants/modi")
// if err != nil {
// t.Fatal(err)
// }
func prepare(t *testing.T) {
test.Prepare(t, config.Conf)
}
// // Validate basic properties
// assert.NotNil(t, assistant)
// assert.Equal(t, "modi", assistant.ID)
// assert.Equal(t, "Modi", assistant.Name)
// assert.Equal(t, "https://api.dicebear.com/7.x/bottts/svg?seed=Modi", assistant.Avatar)
// assert.Equal(t, "deepseek", assistant.Connector)
// assert.NotNil(t, assistant.Prompts)
// assert.NotNil(t, assistant.Script)
// TestLoadPath tests loading assistant from path
func TestLoadPath(t *testing.T) {
prepare(t)
defer test.Clean()
// // Test non-existent assistant
// _, err = LoadPath("/assistants/non-existent")
// assert.Error(t, err)
// }
t.Run("LoadFullFieldsAssistant", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// func TestLoad_LoadStore(t *testing.T) {
// prepare(t)
// defer test.Clean()
// Basic fields
assert.Equal(t, "tests.fullfields", assistant.ID)
assert.Equal(t, "Full Fields Test Assistant", assistant.Name)
assert.Equal(t, "assistant", assistant.Type)
assert.Equal(t, "/api/__yao/app/icons/app.png", assistant.Avatar)
assert.Equal(t, "gpt-4o", assistant.Connector)
assert.Equal(t, "/assistants/tests/fullfields", assistant.Path)
assert.Equal(t, "Test assistant with all available fields for unit testing", assistant.Description)
// // Test with nil storage
// _, err := LoadStore("test-id")
// assert.Error(t, err)
// assert.Contains(t, err.Error(), "storage is not set")
// Boolean fields
assert.True(t, assistant.Public)
assert.True(t, assistant.Readonly)
assert.True(t, assistant.Mentionable)
assert.False(t, assistant.Automated)
assert.True(t, assistant.DisableGlobalPrompts)
// // Setup mock storage
// mockStore := &mockStore{
// data: map[string]map[string]interface{}{
// "test-id": {
// "assistant_id": "test-id",
// "name": "Test Assistant",
// "avatar": "test-avatar",
// "connector": "gpt-3_5-turbo",
// },
// },
// }
// SetStorage(mockStore)
// defer SetStorage(nil)
// Share field
assert.Equal(t, "team", assistant.Share)
// // Test loading from store
// assistant, err := LoadStore("test-id")
// assert.NoError(t, err)
// assert.NotNil(t, assistant)
// assert.Equal(t, "test-id", assistant.ID)
// assert.Equal(t, "Test Assistant", assistant.Name)
// assert.Equal(t, "test-avatar", assistant.Avatar)
// assert.Equal(t, "gpt-3_5-turbo", assistant.Connector)
// Sort field
assert.Equal(t, 100, assistant.Sort)
// // Test cache functionality
// assistant2, err := LoadStore("test-id")
// assert.NoError(t, err)
// assert.Equal(t, assistant, assistant2) // Should be the same instance from cache
// Tags
assert.NotNil(t, assistant.Tags)
assert.Contains(t, assistant.Tags, "Test")
assert.Contains(t, assistant.Tags, "Development")
assert.Contains(t, assistant.Tags, "FullFields")
// // Test non-existent assistant
// _, err = LoadStore("non-existent")
// assert.Error(t, err)
// }
// Options
assert.NotNil(t, assistant.Options)
assert.Equal(t, 0.7, assistant.Options["temperature"])
assert.Equal(t, float64(2000), assistant.Options["max_tokens"])
// func TestLoad_Cache(t *testing.T) {
// prepare(t)
// defer test.Clean()
// Prompts (default prompts from prompts.yml)
assert.NotNil(t, assistant.Prompts)
assert.GreaterOrEqual(t, len(assistant.Prompts), 1)
assert.Equal(t, "system", assistant.Prompts[0].Role)
// // Clear any existing cache first
// ClearCache()
// Script (from src/index.ts)
assert.NotNil(t, assistant.Script)
})
// // Test cache operations
// SetCache(2) // Set small cache size for testing
// assert.Equal(t, 2, loaded.capacity, "Cache capacity should be 2")
t.Run("LoadConnectorOptions", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Create test assistants
// assistant1 := &Assistant{ID: "id1", Name: "Assistant 1"}
// assistant2 := &Assistant{ID: "id2", Name: "Assistant 2"}
// assistant3 := &Assistant{ID: "id3", Name: "Assistant 3"}
// ConnectorOptions
assert.NotNil(t, assistant.ConnectorOptions)
assert.True(t, assistant.ConnectorOptions.Optional)
assert.NotNil(t, assistant.ConnectorOptions.Connectors)
assert.Contains(t, assistant.ConnectorOptions.Connectors, "gpt-4o")
assert.Contains(t, assistant.ConnectorOptions.Connectors, "gpt-4o-mini")
assert.Contains(t, assistant.ConnectorOptions.Connectors, "deepseek")
assert.NotNil(t, assistant.ConnectorOptions.Filters)
assert.Len(t, assistant.ConnectorOptions.Filters, 2)
})
// // Test Put and Get
// loaded.Put(assistant1)
// assert.Equal(t, 1, loaded.Len(), "Cache should have 1 item")
t.Run("LoadPromptPresets", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// loaded.Put(assistant2)
// assert.Equal(t, 2, loaded.Len(), "Cache should have 2 items")
// PromptPresets (from prompts directory)
assert.NotNil(t, assistant.PromptPresets)
// // Test cache hit
// cached, exists := loaded.Get("id1")
// assert.True(t, exists)
// assert.Equal(t, assistant1, cached)
// Top-level presets: chat.yml -> "chat", task.yml -> "task"
chatPreset, hasChat := assistant.PromptPresets["chat"]
assert.True(t, hasChat, "Should have 'chat' preset")
assert.NotEmpty(t, chatPreset)
// // Test cache eviction (LRU)
// // At this point: assistant1 is most recently used (due to Get), then assistant2
// loaded.Put(assistant3) // This should evict assistant2 since it's least recently used
// assert.Equal(t, 2, loaded.Len(), "Cache should still have 2 items")
// _, exists = loaded.Get("id2")
// assert.False(t, exists, "assistant2 should have been evicted (least recently used)")
// _, exists = loaded.Get("id1")
// assert.True(t, exists, "assistant1 should still be in cache (was accessed recently)")
// _, exists = loaded.Get("id3")
// assert.True(t, exists, "assistant3 should be in cache (most recently added)")
taskPreset, hasTask := assistant.PromptPresets["task"]
assert.True(t, hasTask, "Should have 'task' preset")
assert.NotEmpty(t, taskPreset)
// // Test clear cache
// ClearCache()
// assert.Nil(t, loaded)
// Nested presets: chat/friendly.yml -> "chat.friendly"
friendlyPreset, hasFriendly := assistant.PromptPresets["chat.friendly"]
assert.True(t, hasFriendly, "Should have 'chat.friendly' preset")
assert.NotEmpty(t, friendlyPreset)
// // Test setting new cache capacity
// SetCache(100)
// assert.NotNil(t, loaded)
// }
professionalPreset, hasProfessional := assistant.PromptPresets["chat.professional"]
assert.True(t, hasProfessional, "Should have 'chat.professional' preset")
assert.NotEmpty(t, professionalPreset)
// func TestLoad_Validate(t *testing.T) {
// tests := []struct {
// name string
// ast *Assistant
// wantErr bool
// }{
// {
// name: "valid assistant",
// ast: &Assistant{
// ID: "test-id",
// Name: "Test Assistant",
// Connector: "test-connector",
// },
// wantErr: false,
// },
// {
// name: "missing id",
// ast: &Assistant{
// Name: "Test Assistant",
// Connector: "test-connector",
// },
// wantErr: true,
// },
// {
// name: "missing name",
// ast: &Assistant{
// ID: "test-id",
// Connector: "test-connector",
// },
// wantErr: true,
// },
// {
// name: "missing connector",
// ast: &Assistant{
// ID: "test-id",
// Name: "Test Assistant",
// },
// wantErr: true,
// },
// }
// task/analysis.yml -> "task.analysis"
analysisPreset, hasAnalysis := assistant.PromptPresets["task.analysis"]
assert.True(t, hasAnalysis, "Should have 'task.analysis' preset")
assert.NotEmpty(t, analysisPreset)
})
// for _, tt := range tests {
// t.Run(tt.name, func(t *testing.T) {
// err := tt.ast.Validate()
// if (err != nil) != tt.wantErr {
// t.Errorf("Assistant.Validate() error = %v, wantErr %v", err, tt.wantErr)
// }
// })
// }
// }
t.Run("LoadKnowledgeBase", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// func TestLoad_Clone(t *testing.T) {
// // Create a test assistant with all fields populated
// original := &Assistant{
// ID: "test-id",
// Type: "test-type",
// Name: "Test Assistant",
// Avatar: "test-avatar",
// Connector: "test-connector",
// Path: "test-path",
// BuiltIn: true,
// Sort: 1,
// Description: "test description",
// Tags: []string{"tag1", "tag2"},
// Readonly: true,
// Mentionable: true,
// Automated: true,
// Options: map[string]interface{}{"key": "value"},
// Prompts: []Prompt{{Role: "system", Content: "test"}},
// Workflow: map[string]interface{}{"step": "test"},
// }
// KB
assert.NotNil(t, assistant.KB)
assert.NotNil(t, assistant.KB.Collections)
assert.Contains(t, assistant.KB.Collections, "test-collection")
assert.NotNil(t, assistant.KB.Options)
assert.Equal(t, float64(5), assistant.KB.Options["top_k"])
})
// // Clone the assistant
// clone := original.Clone()
t.Run("LoadMCPServers", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Verify all fields are correctly cloned
// assert.Equal(t, original.ID, clone.ID)
// assert.Equal(t, original.Type, clone.Type)
// assert.Equal(t, original.Name, clone.Name)
// assert.Equal(t, original.Avatar, clone.Avatar)
// assert.Equal(t, original.Connector, clone.Connector)
// assert.Equal(t, original.Path, clone.Path)
// assert.Equal(t, original.BuiltIn, clone.BuiltIn)
// assert.Equal(t, original.Sort, clone.Sort)
// assert.Equal(t, original.Description, clone.Description)
// assert.Equal(t, original.Tags, clone.Tags)
// assert.Equal(t, original.Readonly, clone.Readonly)
// assert.Equal(t, original.Mentionable, clone.Mentionable)
// assert.Equal(t, original.Automated, clone.Automated)
// assert.Equal(t, original.Options, clone.Options)
// assert.Equal(t, original.Prompts, clone.Prompts)
// assert.Equal(t, original.Workflow, clone.Workflow)
// MCP
assert.NotNil(t, assistant.MCP)
assert.NotNil(t, assistant.MCP.Servers)
assert.Len(t, assistant.MCP.Servers, 1)
assert.Equal(t, "echo", assistant.MCP.Servers[0].ServerID)
assert.Contains(t, assistant.MCP.Servers[0].Tools, "ping")
assert.Contains(t, assistant.MCP.Servers[0].Tools, "echo")
})
// // Verify deep copy by modifying original
// original.Tags[0] = "modified"
// original.Options["key"] = "modified"
// original.Workflow["step"] = "modified"
// assert.NotEqual(t, original.Tags[0], clone.Tags[0])
// assert.NotEqual(t, original.Options["key"], clone.Options["key"])
// assert.NotEqual(t, original.Workflow["step"], clone.Workflow["step"])
t.Run("LoadWorkflow", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Test nil case
// var nilAssistant *Assistant
// assert.Nil(t, nilAssistant.Clone())
// }
// Workflow
assert.NotNil(t, assistant.Workflow)
assert.NotNil(t, assistant.Workflow.Workflows)
assert.Contains(t, assistant.Workflow.Workflows, "test-workflow")
assert.NotNil(t, assistant.Workflow.Options)
assert.Equal(t, float64(10), assistant.Workflow.Options["max_steps"])
})
// func TestLoad_Update(t *testing.T) {
// // Create a test assistant
// ast := &Assistant{
// ID: "test-id",
// Name: "Original Name",
// Connector: "original-connector",
// }
t.Run("LoadPlaceholder", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Test updating various fields
// updates := map[string]interface{}{
// "name": "Updated Name",
// "avatar": "updated-avatar",
// "description": "Updated description",
// "connector": "updated-connector",
// "type": "updated-type",
// "sort": 2,
// "mentionable": true,
// "automated": true,
// "tags": []string{"new-tag"},
// "options": map[string]interface{}{"new": "value"},
// }
// Placeholder
assert.NotNil(t, assistant.Placeholder)
assert.Equal(t, "Full Fields Test", assistant.Placeholder.Title)
assert.Equal(t, "Test assistant with complete field coverage", assistant.Placeholder.Description)
assert.NotNil(t, assistant.Placeholder.Prompts)
assert.Len(t, assistant.Placeholder.Prompts, 3)
})
// err := ast.Update(updates)
// assert.NoError(t, err)
t.Run("LoadLocales", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Verify updates
// assert.Equal(t, "Updated Name", ast.Name)
// assert.Equal(t, "updated-avatar", ast.Avatar)
// assert.Equal(t, "Updated description", ast.Description)
// assert.Equal(t, "updated-connector", ast.Connector)
// assert.Equal(t, "updated-type", ast.Type)
// assert.Equal(t, 2, ast.Sort)
// assert.True(t, ast.Mentionable)
// assert.True(t, ast.Automated)
// assert.Equal(t, []string{"new-tag"}, ast.Tags)
// assert.Equal(t, map[string]interface{}{"new": "value"}, ast.Options)
// Locales
assert.NotNil(t, assistant.Locales)
// // Test nil assistant
// var nilAssistant *Assistant
// err = nilAssistant.Update(updates)
// assert.Error(t, err)
enLocale, hasEn := assistant.Locales["en-us"]
assert.True(t, hasEn, "Should have en-us locale")
assert.NotNil(t, enLocale)
// // Test invalid update that would make the assistant invalid
// invalidUpdates := map[string]interface{}{
// "name": "",
// }
// err = ast.Update(invalidUpdates)
// assert.Error(t, err)
// }
zhLocale, hasZh := assistant.Locales["zh-cn"]
assert.True(t, hasZh, "Should have zh-cn locale")
assert.NotNil(t, zhLocale)
})
// func TestLoadBuiltIn(t *testing.T) {
// prepare(t)
// defer test.Clean()
t.Run("LoadNonExistentAssistant", func(t *testing.T) {
_, err := LoadPath("/assistants/non-existent")
assert.Error(t, err)
})
}
// // Clear any existing cache and storage
// ClearCache()
// SetStorage(nil)
// TestLoadPathMCPTest tests loading the MCP test assistant
func TestLoadPathMCPTest(t *testing.T) {
prepare(t)
defer test.Clean()
// // Create a mock store to verify built-in assistants are saved
// mockStore := &mockStore{
// data: make(map[string]map[string]interface{}),
// }
// SetStorage(mockStore)
// SetCache(100)
assistant, err := LoadPath("/assistants/tests/mcptest")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Test loading built-in assistants
// err := LoadBuiltIn()
// assert.NoError(t, err)
assert.Equal(t, "tests.mcptest", assistant.ID)
assert.Equal(t, "MCP Test Assistant", assistant.Name)
assert.Equal(t, "gpt-4o", assistant.Connector)
// // Verify Modi assistant was loaded
// assistant, exists := loaded.Get("modi")
// assert.True(t, exists, "Modi assistant should be loaded in cache")
// if exists {
// assert.Equal(t, "modi", assistant.ID)
// assert.Equal(t, "Modi", assistant.Name)
// assert.Equal(t, "deepseek", assistant.Connector)
// assert.True(t, assistant.BuiltIn)
// assert.True(t, assistant.Readonly)
// assert.NotNil(t, assistant.Prompts)
// assert.NotNil(t, assistant.Script)
// }
// MCP configuration
assert.NotNil(t, assistant.MCP)
assert.Len(t, assistant.MCP.Servers, 1)
assert.Equal(t, "echo", assistant.MCP.Servers[0].ServerID)
// }
// Locales
assert.NotNil(t, assistant.Locales)
assert.Contains(t, assistant.Locales, "en-us")
assert.Contains(t, assistant.Locales, "zh-cn")
}
// // mockStore implements store.Store interface for testing
// type mockStore struct {
// data map[string]map[string]interface{}
// }
// TestLoadPathBuildRequest tests loading the build request test assistant
func TestLoadPathBuildRequest(t *testing.T) {
prepare(t)
defer test.Clean()
// func (m *mockStore) GetAssistant(id string, locale ...string) (map[string]interface{}, error) {
// if data, ok := m.data[id]; ok {
// return data, nil
// }
// return nil, fmt.Errorf("assistant not found: %s", id)
// }
assistant, err := LoadPath("/assistants/tests/buildrequest")
require.NoError(t, err)
require.NotNil(t, assistant)
// // Add other required interface methods with empty implementations
// func (m *mockStore) GetThread(id string) (map[string]interface{}, error) { return nil, nil }
// func (m *mockStore) GetMessage(id string) (map[string]interface{}, error) { return nil, nil }
// func (m *mockStore) GetFile(id string) (map[string]interface{}, error) { return nil, nil }
// func (m *mockStore) CreateAssistant(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateThread(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateMessage(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) CreateFile(data map[string]interface{}) (map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) UpdateAssistant(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateThread(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateMessage(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) UpdateFile(id string, data map[string]interface{}) error { return nil }
// func (m *mockStore) DeleteAssistant(id string) error { return nil }
// func (m *mockStore) DeleteThread(id string) error { return nil }
// func (m *mockStore) DeleteMessage(id string) error { return nil }
// func (m *mockStore) DeleteFile(id string) error { return nil }
// func (m *mockStore) ListAssistants(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListThreads(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListMessages(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) ListFiles(query map[string]interface{}) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) DeleteAllChats(id string) error { return nil }
// func (m *mockStore) DeleteChat(id string, chatID string) error { return nil }
// func (m *mockStore) GetAssistants(filter store.AssistantFilter, locale ...string) (*store.AssistantResponse, error) {
// return nil, nil
// }
// func (m *mockStore) GetChat(id string, chatID string, locale ...string) (*store.ChatInfo, error) {
// return nil, nil
// }
// func (m *mockStore) GetChatWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) (*store.ChatInfo, error) {
// return nil, nil
// }
// func (m *mockStore) GetChats(id string, filter store.ChatFilter, locale ...string) (*store.ChatGroupResponse, error) {
// return nil, nil
// }
// func (m *mockStore) GetHistory(id string, chatID string, locale ...string) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) GetHistoryWithFilter(id string, chatID string, filter store.ChatFilter, locale ...string) ([]map[string]interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) SaveAssistant(assistant map[string]interface{}) (interface{}, error) {
// return nil, nil
// }
// func (m *mockStore) SaveHistory(sid string, messages []map[string]interface{}, cid string, context map[string]interface{}) error {
// return nil
// }
// func (m *mockStore) UpdateChatTitle(sid string, cid string, title string) error { return nil }
// func (m *mockStore) DeleteAssistants(filter store.AssistantFilter) (int64, error) { return 0, nil }
// func (m *mockStore) GetAssistantTags(locale ...string) ([]store.Tag, error) {
// return []store.Tag{}, nil
// }
assert.Equal(t, "tests.buildrequest", assistant.ID)
assert.Equal(t, "Build Request Test", assistant.Name)
// // Attachment related methods
// func (m *mockStore) SaveAttachment(attachment map[string]interface{}) (interface{}, error) {
// return attachment["file_id"], nil
// }
// Script should be loaded
assert.NotNil(t, assistant.Script)
// func (m *mockStore) DeleteAttachment(fileID string) error {
// return nil
// }
// Options
assert.NotNil(t, assistant.Options)
assert.Equal(t, 0.5, assistant.Options["temperature"])
}
// func (m *mockStore) GetAttachments(filter store.AttachmentFilter, locale ...string) (*store.AttachmentResponse, error) {
// return &store.AttachmentResponse{}, nil
// }
// TestCache tests the assistant cache functionality
func TestCache(t *testing.T) {
// Clear any existing cache
ClearCache()
// func (m *mockStore) GetAttachment(fileID string, locale ...string) (map[string]interface{}, error) {
// return nil, nil
// }
// Set small cache for testing
SetCache(3)
assert.NotNil(t, loaded)
// func (m *mockStore) DeleteAttachments(filter store.AttachmentFilter) (int64, error) {
// return 0, nil
// }
// Create test assistants
ast1 := &Assistant{AssistantModel: store.AssistantModel{ID: "id1", Name: "Assistant 1"}}
ast2 := &Assistant{AssistantModel: store.AssistantModel{ID: "id2", Name: "Assistant 2"}}
ast3 := &Assistant{AssistantModel: store.AssistantModel{ID: "id3", Name: "Assistant 3"}}
ast4 := &Assistant{AssistantModel: store.AssistantModel{ID: "id4", Name: "Assistant 4"}}
// // Knowledge related methods
// func (m *mockStore) SaveKnowledge(knowledge map[string]interface{}) (interface{}, error) {
// return knowledge["collection_id"], nil
// }
t.Run("PutAndGet", func(t *testing.T) {
loaded.Put(ast1)
assert.Equal(t, 1, loaded.Len())
// func (m *mockStore) DeleteKnowledge(collectionID string) error {
// return nil
// }
cached, exists := loaded.Get("id1")
assert.True(t, exists)
assert.Equal(t, ast1, cached)
})
// func (m *mockStore) GetKnowledges(filter store.KnowledgeFilter, locale ...string) (*store.KnowledgeResponse, error) {
// return &store.KnowledgeResponse{}, nil
// }
t.Run("CacheEviction", func(t *testing.T) {
loaded.Put(ast2)
loaded.Put(ast3)
assert.Equal(t, 3, loaded.Len())
// func (m *mockStore) GetKnowledge(collectionID string, locale ...string) (map[string]interface{}, error) {
// return nil, nil
// }
// Access ast1 to make it recently used
loaded.Get("id1")
// func (m *mockStore) DeleteKnowledges(filter store.KnowledgeFilter) (int64, error) {
// return 0, nil
// }
// Add ast4, should evict ast2 (least recently used)
loaded.Put(ast4)
assert.Equal(t, 3, loaded.Len())
// // Close closes the store and releases any resources
// func (m *mockStore) Close() error {
// return nil
// }
_, exists := loaded.Get("id2")
assert.False(t, exists, "ast2 should be evicted")
_, exists = loaded.Get("id1")
assert.True(t, exists, "ast1 should still exist")
_, exists = loaded.Get("id4")
assert.True(t, exists, "ast4 should exist")
})
t.Run("ClearCache", func(t *testing.T) {
ClearCache()
assert.Nil(t, loaded)
})
t.Run("SetCacheAfterClear", func(t *testing.T) {
SetCache(100)
assert.NotNil(t, loaded)
})
}
// TestClone tests the assistant Clone method
func TestClone(t *testing.T) {
prepare(t)
defer test.Clean()
t.Run("CloneFullFieldsAssistant", func(t *testing.T) {
original, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
clone := original.Clone()
require.NotNil(t, clone)
// Basic fields should be equal
assert.Equal(t, original.ID, clone.ID)
assert.Equal(t, original.Name, clone.Name)
assert.Equal(t, original.Type, clone.Type)
assert.Equal(t, original.Connector, clone.Connector)
assert.Equal(t, original.Description, clone.Description)
// Verify deep copy - modifying original should not affect clone
if len(original.Tags) > 0 {
originalTag := original.Tags[0]
original.Tags[0] = "modified"
assert.NotEqual(t, original.Tags[0], clone.Tags[0])
original.Tags[0] = originalTag // restore
}
if original.Options != nil {
original.Options["test_key"] = "test_value"
_, exists := clone.Options["test_key"]
assert.False(t, exists, "Clone should not have modified key")
delete(original.Options, "test_key") // cleanup
}
})
t.Run("CloneNil", func(t *testing.T) {
var nilAssistant *Assistant
assert.Nil(t, nilAssistant.Clone())
})
}
// TestUpdate tests the assistant Update method
func TestUpdate(t *testing.T) {
prepare(t)
defer test.Clean()
t.Run("UpdateBasicFields", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
updates := map[string]interface{}{
"name": "Updated Name",
"description": "Updated description",
"tags": []string{"updated", "tags"},
}
err = assistant.Update(updates)
require.NoError(t, err)
assert.Equal(t, "Updated Name", assistant.Name)
assert.Equal(t, "Updated description", assistant.Description)
assert.Equal(t, []string{"updated", "tags"}, assistant.Tags)
})
t.Run("UpdateConnectorOptions", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
updates := map[string]interface{}{
"connector_options": map[string]interface{}{
"optional": false,
"connectors": []string{"new-connector"},
},
}
err = assistant.Update(updates)
require.NoError(t, err)
assert.NotNil(t, assistant.ConnectorOptions)
assert.False(t, assistant.ConnectorOptions.Optional)
assert.Contains(t, assistant.ConnectorOptions.Connectors, "new-connector")
})
t.Run("UpdatePromptPresets", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
updates := map[string]interface{}{
"prompt_presets": map[string]interface{}{
"custom": []map[string]interface{}{
{"role": "system", "content": "Custom preset"},
},
},
}
err = assistant.Update(updates)
require.NoError(t, err)
assert.NotNil(t, assistant.PromptPresets)
customPreset, exists := assistant.PromptPresets["custom"]
assert.True(t, exists)
assert.Len(t, customPreset, 1)
})
t.Run("UpdateSource", func(t *testing.T) {
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
updates := map[string]interface{}{
"source": "function Create(ctx, messages) { return { messages: messages }; }",
}
err = assistant.Update(updates)
require.NoError(t, err)
assert.Equal(t, "function Create(ctx, messages) { return { messages: messages }; }", assistant.Source)
})
t.Run("UpdateNilAssistant", func(t *testing.T) {
var nilAssistant *Assistant
err := nilAssistant.Update(map[string]interface{}{"name": "test"})
assert.Error(t, err)
})
}
// TestMap tests the assistant Map method
func TestMap(t *testing.T) {
prepare(t)
defer test.Clean()
assistant, err := LoadPath("/assistants/tests/fullfields")
require.NoError(t, err)
m := assistant.Map()
require.NotNil(t, m)
// Check all fields are present
assert.Equal(t, assistant.ID, m["assistant_id"])
assert.Equal(t, assistant.Name, m["name"])
assert.Equal(t, assistant.Type, m["type"])
assert.Equal(t, assistant.Connector, m["connector"])
assert.Equal(t, assistant.Description, m["description"])
assert.Equal(t, assistant.Path, m["path"])
assert.Equal(t, assistant.Tags, m["tags"])
assert.Equal(t, assistant.Options, m["options"])
assert.Equal(t, assistant.Prompts, m["prompts"])
assert.Equal(t, assistant.KB, m["kb"])
assert.Equal(t, assistant.MCP, m["mcp"])
assert.Equal(t, assistant.Workflow, m["workflow"])
assert.Equal(t, assistant.Placeholder, m["placeholder"])
assert.Equal(t, assistant.Locales, m["locales"])
// New fields
assert.Equal(t, assistant.ConnectorOptions, m["connector_options"])
assert.Equal(t, assistant.PromptPresets, m["prompt_presets"])
assert.Equal(t, assistant.Source, m["source"])
}
// TestValidate tests the assistant Validate method
func TestValidate(t *testing.T) {
tests := []struct {
name string
ast *Assistant
wantErr bool
}{
{
name: "ValidAssistant",
ast: &Assistant{
AssistantModel: store.AssistantModel{
ID: "test-id",
Name: "Test Assistant",
Connector: "gpt-4o",
},
},
wantErr: false,
},
{
name: "MissingID",
ast: &Assistant{
AssistantModel: store.AssistantModel{
Name: "Test Assistant",
Connector: "gpt-4o",
},
},
wantErr: true,
},
{
name: "MissingName",
ast: &Assistant{
AssistantModel: store.AssistantModel{
ID: "test-id",
Connector: "gpt-4o",
},
},
wantErr: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := tt.ast.Validate()
if tt.wantErr {
assert.Error(t, err)
} else {
assert.NoError(t, err)
}
})
}
}

39
agent/assistant/source.go Normal file
View file

@ -0,0 +1,39 @@
package assistant
import (
"fmt"
"strings"
"time"
v8 "github.com/yaoapp/gou/runtime/v8"
"github.com/yaoapp/yao/agent/assistant/hook"
)
// loadSource loads hook script from source code string
// The source field stores TypeScript code directly (but without imports)
// Priority: script field > source field (if script exists, source is ignored)
// Note: Uses MakeScriptInMemory which supports TypeScript syntax without file resolution.
func loadSource(source string, assistantID string) (*hook.Script, error) {
if source == "" {
return nil, nil
}
// Use virtual .ts path for TypeScript support
// MakeScriptInMemory handles TypeScript transform without file system access
virtualFile := fmt.Sprintf("assistants/%s/source.ts", strings.ReplaceAll(assistantID, ".", "/"))
script, err := v8.MakeScriptInMemory([]byte(source), virtualFile, 5*time.Second, true)
if err != nil {
return nil, fmt.Errorf("failed to compile source script: %w", err)
}
return &hook.Script{Script: script}, nil
}
// TODO: Future enhancement - support multiple files merged with special comment delimiter
// Format: // file: index.ts
// This would allow splitting large scripts into multiple logical files while storing as single source
// func loadSourceMultiFile(source string, assistantID string) (*hook.Script, error) {
// // Parse source by "// file: xxx.ts" delimiter
// // Merge and compile
// }

View file

@ -3,6 +3,7 @@ package assistant
import (
"fmt"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -48,7 +49,7 @@ func (ast *Assistant) traceCreateHook(agentNode types.Node, createResponse *cont
}
// traceConnectorCapabilities logs the connector capabilities to the agent trace node
func (ast *Assistant) traceConnectorCapabilities(agentNode types.Node, capabilities *context.ModelCapabilities) {
func (ast *Assistant) traceConnectorCapabilities(agentNode types.Node, capabilities *openai.Capabilities) {
if agentNode == nil {
return
}

View file

@ -40,19 +40,6 @@ type Assistant struct {
// toolCalls bool // Whether this assistant supports tool_calls
}
// ModelCapabilities defines the capabilities of a language model
// This configuration is loaded from agent/models.yml
type ModelCapabilities struct {
Vision interface{} `json:"vision,omitempty" yaml:"vision,omitempty"` // Supports vision/image input: bool or VisionFormat string ("openai", "claude"/"base64", "default")
Tools bool `json:"tools,omitempty" yaml:"tools,omitempty"` // Supports tool/function calling (deprecated, use ToolCalls)
ToolCalls bool `json:"tool_calls,omitempty" yaml:"tool_calls,omitempty"` // Supports tool/function calling
Audio bool `json:"audio,omitempty" yaml:"audio,omitempty"` // Supports audio input/output
Reasoning bool `json:"reasoning,omitempty" yaml:"reasoning,omitempty"` // Supports reasoning/thinking mode (o1, DeepSeek R1)
Streaming bool `json:"streaming,omitempty" yaml:"streaming,omitempty"` // Supports streaming responses
JSON bool `json:"json,omitempty" yaml:"json,omitempty"` // Supports JSON mode
Multimodal bool `json:"multimodal,omitempty" yaml:"multimodal,omitempty"` // Supports multimodal input
}
// VisionCapableModels list of LLM models that support vision capabilities
var VisionCapableModels = map[string]bool{
// OpenAI Models

View file

@ -3,6 +3,7 @@ package context
import (
"time"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/output"
"github.com/yaoapp/yao/agent/output/message"
)
@ -290,18 +291,10 @@ func (ctx *Context) getOutput() (*output.Output, error) {
Accept: string(ctx.Accept),
}
// Convert ModelCapabilities to message.ModelCapabilities
// Set ModelCapabilities (now using openai.Capabilities directly)
if ctx.Capabilities != nil {
options.Capabilities = &message.ModelCapabilities{
Vision: ctx.Capabilities.Vision,
ToolCalls: ctx.Capabilities.ToolCalls,
Audio: ctx.Capabilities.Audio,
Reasoning: ctx.Capabilities.Reasoning,
Streaming: ctx.Capabilities.Streaming,
JSON: ctx.Capabilities.JSON,
Multimodal: ctx.Capabilities.Multimodal,
TemperatureAdjustable: ctx.Capabilities.TemperatureAdjustable,
}
caps := openai.Capabilities(*ctx.Capabilities)
options.Capabilities = &caps
}
var err error

View file

@ -5,6 +5,7 @@ import (
"sync"
"time"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/gou/store"
"github.com/yaoapp/yao/agent/output"
@ -238,7 +239,7 @@ type Context struct {
Skip *Skip `json:"skip,omitempty"` // Skip configuration (history, trace, etc.), nil means don't skip anything
// Model capabilities (set by assistant, used by output adapters)
Capabilities *ModelCapabilities `json:"-"` // Model capabilities for the current connector
Capabilities *openai.Capabilities `json:"-"` // Model capabilities for the current connector
// Interrupt control (all interrupt-related logic is encapsulated in InterruptController)
Interrupt *InterruptController `json:"-"` // Interrupt controller for handling user interrupts during streaming
@ -325,6 +326,10 @@ type HookCreateResponse struct {
// MCP configuration - allow hook to add/override MCP servers for this request
MCPServers []MCPServerConfig `json:"mcp_servers,omitempty"`
// Prompt configuration
PromptPreset string `json:"prompt_preset,omitempty"` // Select prompt preset (e.g., "chat.friendly", "task.analysis")
DisableGlobalPrompts *bool `json:"disable_global_prompts,omitempty"` // Temporarily disable global prompts for this request
// Context adjustments - allow hook to modify context fields
AssistantID string `json:"assistant_id,omitempty"` // Override assistant ID
Connector string `json:"connector,omitempty"` // Override connector

View file

@ -1,6 +1,9 @@
package context
import "github.com/yaoapp/yao/agent/output/message"
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/output/message"
)
// Uses represents the wrapper configurations for assistant
// Used to specify which assistant or MCP server to use for vision, audio, search, and fetch operations
@ -28,17 +31,13 @@ const (
VisionFormatDefault VisionFormat = "default"
)
// ModelCapabilities defines the capabilities of a language model
// Used by LLM to select appropriate provider and validate requests
type ModelCapabilities message.ModelCapabilities
// GetVisionSupport returns whether vision is supported and the format
func (m *ModelCapabilities) GetVisionSupport() (bool, VisionFormat) {
if m == nil || m.Vision == nil {
func GetVisionSupport(cap *openai.Capabilities) (bool, VisionFormat) {
if cap == nil || cap.Vision == nil {
return false, VisionFormatNone
}
switch v := m.Vision.(type) {
switch v := cap.Vision.(type) {
case bool:
// Legacy bool format
return v, VisionFormatDefault
@ -65,7 +64,7 @@ func (m *ModelCapabilities) GetVisionSupport() (bool, VisionFormat) {
type CompletionOptions struct {
// Model capabilities (used by LLM to select appropriate provider)
// nil means capabilities are not specified/checked
Capabilities *ModelCapabilities `json:"capabilities,omitempty"`
Capabilities *openai.Capabilities `json:"capabilities,omitempty"`
// User-specified tools for vision, audio, search, and fetch processing
Uses *Uses `json:"uses,omitempty"`

View file

@ -1,6 +1,7 @@
package adapters
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/context"
)
@ -28,7 +29,7 @@ type ReasoningAdapter struct {
// NewReasoningAdapter creates a new reasoning adapter
// If cap.TemperatureAdjustable is provided, it overrides the default behavior
func NewReasoningAdapter(format ReasoningFormat, cap *context.ModelCapabilities) *ReasoningAdapter {
func NewReasoningAdapter(format ReasoningFormat, cap *openai.Capabilities) *ReasoningAdapter {
supportsEffort := false
supportsTemperature := true
@ -49,8 +50,8 @@ func NewReasoningAdapter(format ReasoningFormat, cap *context.ModelCapabilities)
}
// Override with explicit capability if provided
if cap != nil && cap.TemperatureAdjustable != nil {
supportsTemperature = *cap.TemperatureAdjustable
if cap != nil {
supportsTemperature = cap.TemperatureAdjustable
}
return &ReasoningAdapter{

View file

@ -4,6 +4,7 @@ import (
"fmt"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/context"
)
@ -11,11 +12,11 @@ import (
// Provides common functionality for all LLM providers
type Provider struct {
Connector connector.Connector
Capabilities *context.ModelCapabilities
Capabilities *openai.Capabilities
}
// NewProvider create a new base provider
func NewProvider(conn connector.Connector, capabilities *context.ModelCapabilities) *Provider {
func NewProvider(conn connector.Connector, capabilities *openai.Capabilities) *Provider {
return &Provider{
Connector: conn,
Capabilities: capabilities,
@ -74,33 +75,33 @@ func (p *Provider) SupportsVision() bool {
if p.Capabilities == nil {
return false
}
supported, _ := p.Capabilities.GetVisionSupport()
supported, _ := context.GetVisionSupport(p.Capabilities)
return supported
}
// SupportsAudio check if this provider supports audio
func (p *Provider) SupportsAudio() bool {
return p.Capabilities != nil && p.Capabilities.Audio != nil && *p.Capabilities.Audio
return p.Capabilities != nil && p.Capabilities.Audio
}
// SupportsTools check if this provider supports tool calls
func (p *Provider) SupportsTools() bool {
return p.Capabilities != nil && p.Capabilities.ToolCalls != nil && *p.Capabilities.ToolCalls
return p.Capabilities != nil && p.Capabilities.ToolCalls
}
// SupportsStreaming check if this provider supports streaming
func (p *Provider) SupportsStreaming() bool {
return p.Capabilities != nil && p.Capabilities.Streaming != nil && *p.Capabilities.Streaming
return p.Capabilities != nil && p.Capabilities.Streaming
}
// SupportsJSON check if this provider supports JSON mode
func (p *Provider) SupportsJSON() bool {
return p.Capabilities != nil && p.Capabilities.JSON != nil && *p.Capabilities.JSON
return p.Capabilities != nil && p.Capabilities.JSON
}
// SupportsReasoning check if this provider supports reasoning mode
func (p *Provider) SupportsReasoning() bool {
return p.Capabilities != nil && p.Capabilities.Reasoning != nil && *p.Capabilities.Reasoning
return p.Capabilities != nil && p.Capabilities.Reasoning
}
// GetConnectorSetting gets a setting value from the connector

View file

@ -5,6 +5,7 @@ import (
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -60,15 +61,13 @@ func TestClaudeSonnet4StreamBasic(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &falseVal, // Claude Sonnet 4 (non-thinking) doesn't expose reasoning
ToolCalls: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: false, // Claude Sonnet 4 (non-thinking) doesn't expose reasoning
ToolCalls: true,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -140,15 +139,13 @@ func TestClaudeSonnet4PostBasic(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &falseVal,
Reasoning: &falseVal,
ToolCalls: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: false,
Reasoning: false,
ToolCalls: true,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -206,15 +203,13 @@ func TestClaudeSonnet4WithToolCalls(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &falseVal,
Reasoning: &falseVal,
ToolCalls: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: false,
Reasoning: false,
ToolCalls: true,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -244,8 +239,8 @@ func TestClaudeSonnet4WithToolCalls(t *testing.T) {
options.Tools = []map[string]interface{}{simpleTool}
options.ToolChoice = "auto"
// Set lower max_tokens for faster response
maxTokens := 50
// Set enough tokens for tool call response
maxTokens := 150
options.MaxTokens = &maxTokens
llmInstance, err := llm.New(conn, options)
@ -256,7 +251,7 @@ func TestClaudeSonnet4WithToolCalls(t *testing.T) {
messages := []context.Message{
{
Role: context.RoleUser,
Content: "Call get_info with query='A' and count=1",
Content: "Please use the get_info function to retrieve information. Pass 'A' as the query parameter and 1 as the count parameter.",
},
}
@ -299,15 +294,13 @@ func TestClaudeSonnet4Vision(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &falseVal,
Reasoning: &falseVal,
ToolCalls: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: false,
Reasoning: false,
ToolCalls: true,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -376,15 +369,13 @@ func TestClaudeSonnet4ThinkingStream(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &trueVal, // Claude Thinking mode exposes reasoning
ToolCalls: &falseVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: true, // Claude Thinking mode exposes reasoning
ToolCalls: false,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -456,15 +447,13 @@ func TestClaudeSonnet4ThinkingPost(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &falseVal,
Reasoning: &trueVal,
ToolCalls: &falseVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: false,
Reasoning: true,
ToolCalls: false,
Vision: "claude", // Claude requires base64 format
Multimodal: true,
},
}
@ -550,14 +539,12 @@ func TestClaudeTemperatureHandling(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &falseVal,
Reasoning: &tt.reasoning,
ToolCalls: &trueVal,
Vision: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: false,
Reasoning: tt.reasoning,
ToolCalls: true,
Vision: true,
},
}

View file

@ -7,6 +7,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -28,16 +29,14 @@ func TestDeepSeekR1StreamBasic(t *testing.T) {
}
// Create LLM instance with capabilities
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &trueVal, // DeepSeek R1 supports reasoning
ToolCalls: &falseVal, // R1 doesn't support native tool calls
Vision: &falseVal,
Audio: &falseVal,
Multimodal: &falseVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: true, // DeepSeek R1 supports reasoning
ToolCalls: false, // R1 doesn't support native tool calls
Vision: false,
Audio: false,
Multimodal: false,
},
}
@ -207,15 +206,13 @@ func TestDeepSeekR1PostBasic(t *testing.T) {
}
// Create LLM instance
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
ToolCalls: &falseVal,
Vision: &falseVal,
Audio: &falseVal,
Multimodal: &falseVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
ToolCalls: false,
Vision: false,
Audio: false,
Multimodal: false,
},
}
@ -295,16 +292,14 @@ func TestDeepSeekR1LogicPuzzle(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &trueVal,
ToolCalls: &falseVal,
Vision: &falseVal,
Audio: &falseVal,
Multimodal: &falseVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: true,
ToolCalls: false,
Vision: false,
Audio: false,
Multimodal: false,
},
}

View file

@ -5,6 +5,7 @@ import (
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -24,16 +25,14 @@ func TestDeepSeekV3StreamBasic(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &falseVal, // V3 doesn't support reasoning
ToolCalls: &trueVal, // V3 supports tool calls
Vision: &falseVal,
Audio: &falseVal,
Multimodal: &falseVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: false, // V3 doesn't support reasoning
ToolCalls: true, // V3 supports tool calls
Vision: false,
Audio: false,
Multimodal: false,
},
}
@ -135,15 +134,13 @@ func TestDeepSeekV3PostBasic(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal,
ToolCalls: &trueVal,
Vision: &falseVal,
Audio: &falseVal,
Multimodal: &falseVal,
Capabilities: &openai.Capabilities{
Reasoning: false,
ToolCalls: true,
Vision: false,
Audio: false,
Multimodal: false,
},
}
@ -226,12 +223,10 @@ func TestDeepSeekV3WithToolCalls(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false,
ToolCalls: true,
},
}
@ -318,13 +313,11 @@ func TestDeepSeekV3NoReasoningEffort(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
effort := "high"
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal, // V3 doesn't support reasoning
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false, // V3 doesn't support reasoning
ToolCalls: true,
},
ReasoningEffort: &effort, // Should be ignored by adapter
}

View file

@ -5,6 +5,7 @@ import (
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -24,14 +25,13 @@ func TestGPT5StreamBasic(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
Reasoning: &trueVal, // GPT-5 supports reasoning
ToolCalls: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
Reasoning: true, // GPT-5 supports reasoning
ToolCalls: true,
Vision: true,
Multimodal: true,
},
}
@ -105,11 +105,10 @@ func TestGPT5ReasoningEffort(t *testing.T) {
for _, effort := range effortLevels {
t.Run("effort_"+effort, func(t *testing.T) {
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
ToolCalls: true,
},
ReasoningEffort: &effort,
}
@ -172,11 +171,10 @@ func TestGPT5PostWithToolCalls(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
ToolCalls: true,
},
}
@ -259,12 +257,11 @@ func TestGPT5Vision(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
Vision: &trueVal,
Multimodal: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
Vision: true,
Multimodal: true,
},
}
@ -331,13 +328,11 @@ func TestGPT5ReasoningEffortWithGPT4o(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
effort := "high"
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal, // GPT-4o doesn't support reasoning
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false, // GPT-4o doesn't support reasoning
ToolCalls: true,
},
ReasoningEffort: &effort, // Should be ignored by adapter
}

View file

@ -8,6 +8,7 @@ import (
jsoniter "github.com/json-iterator/go"
"github.com/yaoapp/gou/connector"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/http"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/agent/context"
@ -142,7 +143,7 @@ func buildAPIURL(host, endpoint string) string {
}
// New create a new OpenAI provider with capability adapters
func New(conn connector.Connector, capabilities *context.ModelCapabilities) *Provider {
func New(conn connector.Connector, capabilities *gouOpenAI.Capabilities) *Provider {
return &Provider{
Provider: base.NewProvider(conn, capabilities),
adapters: buildAdapters(capabilities),
@ -150,7 +151,7 @@ func New(conn connector.Connector, capabilities *context.ModelCapabilities) *Pro
}
// buildAdapters builds capability adapters based on model capabilities
func buildAdapters(cap *context.ModelCapabilities) []adapters.CapabilityAdapter {
func buildAdapters(cap *gouOpenAI.Capabilities) []adapters.CapabilityAdapter {
if cap == nil {
return []adapters.CapabilityAdapter{}
}
@ -158,12 +159,10 @@ func buildAdapters(cap *context.ModelCapabilities) []adapters.CapabilityAdapter
result := make([]adapters.CapabilityAdapter, 0)
// Tool call adapter
if cap.ToolCalls != nil {
result = append(result, adapters.NewToolCallAdapter(*cap.ToolCalls))
}
result = append(result, adapters.NewToolCallAdapter(cap.ToolCalls))
// Vision adapter
visionSupport, visionFormat := cap.GetVisionSupport()
visionSupport, visionFormat := context.GetVisionSupport(cap)
if visionSupport {
result = append(result, adapters.NewVisionAdapter(true, visionFormat))
} else if cap.Vision != nil {
@ -172,14 +171,11 @@ func buildAdapters(cap *context.ModelCapabilities) []adapters.CapabilityAdapter
}
// Audio adapter
if cap.Audio != nil {
result = append(result, adapters.NewAudioAdapter(*cap.Audio))
}
result = append(result, adapters.NewAudioAdapter(cap.Audio))
// Reasoning adapter (always add to handle reasoning_effort and temperature parameters)
// Even if the model doesn't support reasoning, we need the adapter to strip reasoning_effort
if cap.Reasoning != nil {
if *cap.Reasoning {
if cap.Reasoning {
// Detect reasoning format based on capabilities
format := detectReasoningFormat(cap)
result = append(result, adapters.NewReasoningAdapter(format, cap))
@ -187,16 +183,15 @@ func buildAdapters(cap *context.ModelCapabilities) []adapters.CapabilityAdapter
// Model doesn't support reasoning, use None format to strip reasoning parameters
result = append(result, adapters.NewReasoningAdapter(adapters.ReasoningFormatNone, cap))
}
}
return result
}
// detectReasoningFormat detects the reasoning format based on capabilities
func detectReasoningFormat(cap *context.ModelCapabilities) adapters.ReasoningFormat {
func detectReasoningFormat(cap *gouOpenAI.Capabilities) adapters.ReasoningFormat {
// TODO: Implement better detection logic
// For now, default to OpenAI o1 format if reasoning is supported
if cap.Reasoning != nil && *cap.Reasoning {
if cap.Reasoning {
return adapters.ReasoningFormatOpenAI
}
return adapters.ReasoningFormatNone

View file

@ -8,6 +8,7 @@ import (
"time"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -29,11 +30,10 @@ func TestOpenAIStreamBasic(t *testing.T) {
}
// Create LLM instance with capabilities
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -117,10 +117,9 @@ func TestOpenAIPostBasic(t *testing.T) {
}
// Create LLM instance
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
}
@ -192,11 +191,10 @@ func TestOpenAIStreamWithToolCalls(t *testing.T) {
}
// Create LLM instance with tool call capabilities
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -307,10 +305,9 @@ func TestOpenAIPostWithToolCalls(t *testing.T) {
}
// Create LLM instance with tool call capabilities
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
}
@ -424,11 +421,10 @@ func TestOpenAIStreamWithInvalidToolCall(t *testing.T) {
}
// Create LLM instance
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -523,11 +519,10 @@ func TestOpenAIStreamRetry(t *testing.T) {
}
// Create LLM instance
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal, // Need this to select OpenAI provider
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true, // Need this to select OpenAI provider
},
}
@ -582,11 +577,10 @@ func TestOpenAIStreamChunkTypes(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -649,11 +643,10 @@ func TestOpenAIStreamErrorCallback(t *testing.T) {
t.Fatalf("Failed to create test connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -711,11 +704,10 @@ func TestOpenAIToolCallValidationRetry(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
Tools: []map[string]interface{}{
{
@ -815,11 +807,10 @@ func TestOpenAIJSONMode(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSON,
@ -902,10 +893,9 @@ func TestOpenAIJSONModePost(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSON,
@ -973,8 +963,6 @@ func TestOpenAIJSONSchema(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
// Define a strict JSON schema
// Note: For OpenAI strict mode, 'required' must include ALL properties
schema := map[string]interface{}{
@ -1009,9 +997,9 @@ func TestOpenAIJSONSchema(t *testing.T) {
}
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSONSchema,
@ -1019,7 +1007,7 @@ func TestOpenAIJSONSchema(t *testing.T) {
Name: "user_info",
Description: "User information schema",
Schema: schema,
Strict: &trueVal,
Strict: func() *bool { v := true; return &v }(),
},
},
}
@ -1121,8 +1109,6 @@ func TestOpenAIJSONSchemaPost(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
// Simple schema for testing
// Note: For OpenAI strict mode, 'required' must include ALL properties
schema := map[string]interface{}{
@ -1144,8 +1130,8 @@ func TestOpenAIJSONSchemaPost(t *testing.T) {
}
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
ToolCalls: true,
},
ResponseFormat: &context.ResponseFormat{
Type: context.ResponseFormatJSONSchema,
@ -1153,7 +1139,7 @@ func TestOpenAIJSONSchemaPost(t *testing.T) {
Name: "api_response",
Description: "API response format",
Schema: schema,
Strict: &trueVal,
Strict: func() *bool { v := true; return &v }(),
},
},
}
@ -1263,11 +1249,10 @@ func TestOpenAIStreamLifecycleEvents(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -1369,11 +1354,10 @@ func TestOpenAIStreamContextCancellation(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true,
},
}
@ -1443,13 +1427,12 @@ func TestOpenAIStreamWithTemperature(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
temperature := 0.7 // Moderate temperature
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Streaming: &trueVal,
ToolCalls: &trueVal, // Need this to select OpenAI provider
Capabilities: &openai.Capabilities{
Streaming: true,
ToolCalls: true, // Need this to select OpenAI provider
},
Temperature: &temperature,
}

View file

@ -5,6 +5,7 @@ import (
"testing"
"github.com/yaoapp/gou/connector"
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/gou/plan"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/llm"
@ -23,11 +24,10 @@ func TestTemperatureGPT5AutoReset(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
invalidTemp := 0.7 // GPT-5 doesn't support this
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
},
Temperature: &invalidTemp, // Should be reset to 1.0
}
@ -73,11 +73,10 @@ func TestTemperatureDeepSeekR1AutoReset(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
invalidTemp := 0.5 // DeepSeek R1 doesn't support this
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
},
Temperature: &invalidTemp, // Should be reset to 1.0
}
@ -126,13 +125,11 @@ func TestTemperatureGPT4oPreserved(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
customTemp := 0.3 // GPT-4o should preserve this
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal, // Not a reasoning model
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false, // Not a reasoning model
ToolCalls: true,
},
Temperature: &customTemp, // Should be preserved
}
@ -178,13 +175,11 @@ func TestTemperatureDeepSeekV3Preserved(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
customTemp := 0.8 // DeepSeek V3 should preserve this
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal, // Not a reasoning model
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false, // Not a reasoning model
ToolCalls: true,
},
Temperature: &customTemp, // Should be preserved
}
@ -230,11 +225,10 @@ func TestTemperatureGPT5Default(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
defaultTemp := 1.0 // GPT-5's valid temperature
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: true,
},
Temperature: &defaultTemp, // Should work fine
}
@ -293,16 +287,14 @@ func TestTemperatureNoTemperatureProvided(t *testing.T) {
t.Fatalf("Failed to select connector: %v", err)
}
trueVal := true
falseVal := false
options := &context.CompletionOptions{
Capabilities: &context.ModelCapabilities{
Reasoning: &falseVal,
ToolCalls: &trueVal,
Capabilities: &openai.Capabilities{
Reasoning: false,
ToolCalls: true,
},
}
if tc.reasoning {
options.Capabilities.Reasoning = &trueVal
options.Capabilities.Reasoning = true
}
// Temperature not set - should use API default

View file

@ -6,6 +6,7 @@ import (
"github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/connector"
gouOpenAI "github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
@ -79,6 +80,12 @@ func Load(cfg config.Config) error {
return err
}
// Initialize Global Prompts
err = initGlobalPrompts()
if err != nil {
return err
}
// Initialize Assistant
err = initAssistant()
if err != nil {
@ -103,6 +110,25 @@ func initGlobalI18n() error {
return nil
}
// initGlobalPrompts initialize the global prompts from agent/prompts.yml
func initGlobalPrompts() error {
prompts, _, err := store.LoadGlobalPrompts()
if err != nil {
return err
}
agentDSL.GlobalPrompts = prompts
return nil
}
// GetGlobalPrompts returns the global prompts
// ctx: context variables for parsing $CTX.* variables
func GetGlobalPrompts(ctx map[string]string) []store.Prompt {
if agentDSL == nil || len(agentDSL.GlobalPrompts) == 0 {
return nil
}
return store.Prompts(agentDSL.GlobalPrompts).Parse(ctx)
}
// initModelCapabilities initialize the model capabilities configuration
func initModelCapabilities() error {
path := filepath.Join("agent", "models.yml")
@ -116,7 +142,7 @@ func initModelCapabilities() error {
return err
}
var models map[string]assistant.ModelCapabilities = map[string]assistant.ModelCapabilities{}
var models map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
err = application.Parse("models.yml", bytes, &models)
if err != nil {
return err
@ -174,6 +200,11 @@ func initAssistant() error {
assistant.SetGlobalUses(globalUses)
}
// Set global prompts
if len(agentDSL.GlobalPrompts) > 0 {
assistant.SetGlobalPrompts(agentDSL.GlobalPrompts)
}
if agentDSL.Models != nil {
assistant.SetModelCapabilities(agentDSL.Models)
}

View file

@ -1,16 +1,208 @@
package agent
// func TestLoad(t *testing.T) {
// test.Prepare(t, config.Conf)
// defer test.Clean()
import (
"strings"
"testing"
"time"
// err := Load(config.Conf)
// if err != nil {
// t.Fatal(err)
// }
// check(t)
// }
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test"
)
// func check(t *testing.T) {
// assert.NotNil(t, Agent)
// }
func prepare(t *testing.T) {
test.Prepare(t, config.Conf)
err := Load(config.Conf)
require.NoError(t, err)
}
func TestLoad(t *testing.T) {
prepare(t)
defer test.Clean()
agent := GetAgent()
require.NotNil(t, agent)
t.Run("LoadAgentSettings", func(t *testing.T) {
// Cache setting
assert.NotEmpty(t, agent.Cache)
// Store setting
assert.NotNil(t, agent.Store)
assert.Greater(t, agent.StoreSetting.MaxSize, 0)
// Uses setting
assert.NotNil(t, agent.Uses)
assert.NotEmpty(t, agent.Uses.Default)
})
t.Run("LoadDefaultAssistant", func(t *testing.T) {
assert.NotNil(t, agent.Assistant)
})
t.Run("LoadGlobalPrompts", func(t *testing.T) {
// Global prompts should be loaded from agent/prompts.yml
assert.NotNil(t, agent.GlobalPrompts)
assert.Greater(t, len(agent.GlobalPrompts), 0)
// First prompt should be system role
assert.Equal(t, "system", agent.GlobalPrompts[0].Role)
// Content should contain system context info (with variables not yet parsed)
assert.Contains(t, agent.GlobalPrompts[0].Content, "$SYS.")
})
t.Run("LoadModelCapabilities", func(t *testing.T) {
// Model capabilities should be loaded from agent/models.yml
assert.NotNil(t, agent.Models)
assert.Greater(t, len(agent.Models), 0)
})
}
func TestGetGlobalPrompts(t *testing.T) {
prepare(t)
defer test.Clean()
t.Run("ParseWithoutContext", func(t *testing.T) {
prompts := GetGlobalPrompts(nil)
require.NotNil(t, prompts)
require.Greater(t, len(prompts), 0)
// $SYS.* variables should be replaced
content := prompts[0].Content
assert.NotContains(t, content, "$SYS.DATETIME")
assert.NotContains(t, content, "$SYS.TIMEZONE")
assert.NotContains(t, content, "$SYS.WEEKDAY")
// Should contain actual time values
now := time.Now()
assert.Contains(t, content, now.Format("2006-01-02"))
})
t.Run("ParseWithContext", func(t *testing.T) {
ctx := map[string]string{
"USER_ID": "test-user-123",
"LOCALE": "zh-CN",
}
prompts := GetGlobalPrompts(ctx)
require.NotNil(t, prompts)
require.Greater(t, len(prompts), 0)
// $SYS.* variables should be replaced
content := prompts[0].Content
assert.NotContains(t, content, "$SYS.DATETIME")
})
t.Run("ParseSystemTimeVariables", func(t *testing.T) {
prompts := GetGlobalPrompts(nil)
require.NotNil(t, prompts)
content := prompts[0].Content
now := time.Now()
// Should contain current date
assert.Contains(t, content, now.Format("2006-01-02"))
// Should contain timezone
assert.Contains(t, content, now.Location().String())
// Should contain weekday
assert.Contains(t, content, now.Weekday().String())
})
}
func TestGetGlobalPromptsWithDisableFlag(t *testing.T) {
prepare(t)
defer test.Clean()
agent := GetAgent()
require.NotNil(t, agent)
t.Run("GlobalPromptsExist", func(t *testing.T) {
// Verify global prompts are loaded
assert.NotNil(t, agent.GlobalPrompts)
assert.Greater(t, len(agent.GlobalPrompts), 0)
})
t.Run("AssistantCanDisableGlobalPrompts", func(t *testing.T) {
// The fullfields test assistant has disable_global_prompts: true
// This test verifies the flag is properly loaded
// The actual merging logic is in the assistant module
prompts := GetGlobalPrompts(nil)
assert.NotNil(t, prompts)
// Global prompts should still be available
// The assistant decides whether to use them based on DisableGlobalPrompts flag
})
}
func TestGlobalPromptsContent(t *testing.T) {
prepare(t)
defer test.Clean()
agent := GetAgent()
require.NotNil(t, agent)
require.NotNil(t, agent.GlobalPrompts)
require.Greater(t, len(agent.GlobalPrompts), 0)
t.Run("SystemContextPrompt", func(t *testing.T) {
// Find system prompt
var systemPrompt string
for _, p := range agent.GlobalPrompts {
if p.Role == "system" {
systemPrompt = p.Content
break
}
}
assert.NotEmpty(t, systemPrompt)
assert.Contains(t, systemPrompt, "System Context")
})
t.Run("VariablesInRawPrompts", func(t *testing.T) {
// Raw prompts should contain unparsed variables
content := agent.GlobalPrompts[0].Content
assert.True(t,
strings.Contains(content, "$SYS.") ||
strings.Contains(content, "$ENV.") ||
strings.Contains(content, "$CTX."),
"Raw prompts should contain variable placeholders")
})
}
func TestAssistantGlobalPrompts(t *testing.T) {
prepare(t)
defer test.Clean()
t.Run("AssistantModuleReceivesGlobalPrompts", func(t *testing.T) {
// Verify assistant module has global prompts
prompts := assistant.GetGlobalPrompts(nil)
require.NotNil(t, prompts)
require.Greater(t, len(prompts), 0)
// Should be parsed (no $SYS.* variables)
content := prompts[0].Content
assert.NotContains(t, content, "$SYS.DATETIME")
})
t.Run("AssistantModuleParsesWithContext", func(t *testing.T) {
ctx := map[string]string{
"USER_ID": "assistant-test-user",
"LOCALE": "en-US",
}
prompts := assistant.GetGlobalPrompts(ctx)
require.NotNil(t, prompts)
// $SYS.* should be replaced
content := prompts[0].Content
assert.NotContains(t, content, "$SYS.")
// Should contain current time info
now := time.Now()
assert.Contains(t, content, now.Format("2006-01-02"))
})
}

View file

@ -35,7 +35,7 @@ type AdapterConfig struct {
Locale string
}
// ModelCapabilities is a simplified version of context.ModelCapabilities
// ModelCapabilities is a simplified version of openai.Capabilities
// We use a local type to avoid circular dependencies
type ModelCapabilities struct {
Reasoning *bool // Supports reasoning/thinking mode (o1, DeepSeek R1)

View file

@ -22,9 +22,10 @@ type Writer struct {
func NewWriter(options message.Options) (*Writer, error) {
// Get model capabilities from context (set by assistant)
var capabilities *ModelCapabilities
if options.Capabilities != nil && options.Capabilities.Reasoning != nil {
if options.Capabilities != nil && options.Capabilities.Reasoning {
v := true
capabilities = &ModelCapabilities{
Reasoning: options.Capabilities.Reasoning,
Reasoning: &v,
}
}

View file

@ -3,6 +3,7 @@ package message
import (
"net/http"
"github.com/yaoapp/gou/connector/openai"
traceTypes "github.com/yaoapp/yao/trace/types"
)
@ -12,23 +13,10 @@ type Options struct {
Accept string
Writer http.ResponseWriter
Trace traceTypes.Manager
Capabilities *ModelCapabilities
Capabilities *openai.Capabilities
Locale string
}
// ModelCapabilities defines the capabilities of a language model
// Used by LLM to select appropriate provider and validate requests
type ModelCapabilities struct {
Vision interface{} `json:"vision,omitempty"` // Supports vision/image input: bool or VisionFormat string ("openai", "claude"/"base64", "default")
ToolCalls *bool `json:"tool_calls,omitempty"` // Supports tool/function calling
Audio *bool `json:"audio,omitempty"` // Supports audio input/output
Reasoning *bool `json:"reasoning,omitempty"` // Supports reasoning/thinking mode (o1, DeepSeek R1)
Streaming *bool `json:"streaming,omitempty"` // Supports streaming responses
JSON *bool `json:"json,omitempty"` // Supports JSON mode
Multimodal *bool `json:"multimodal,omitempty"` // Supports multimodal input (text + images + audio)
TemperatureAdjustable *bool `json:"temperature_adjustable,omitempty"` // Supports temperature adjustment (reasoning models typically don't)
}
// Message represents a universal message structure (DSL)
// All messages are expressed through Type + Props, without predefining specific types
type Message struct {

View file

@ -76,7 +76,8 @@ func (m *Mongo) GetAssistants(filter types.AssistantFilter, locale ...string) (*
}
// GetAssistant retrieves a single assistant by ID
func (m *Mongo) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (m *Mongo) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}

View file

@ -76,7 +76,8 @@ func (r *Redis) GetAssistants(filter types.AssistantFilter, locale ...string) (*
}
// GetAssistant retrieves a single assistant by ID
func (r *Redis) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
// fields: Optional list of fields to retrieve. If empty, a default set of fields will be returned.
func (r *Redis) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
return nil, nil
}

View file

@ -217,6 +217,9 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
if path, ok := data["path"].(string); ok {
model.Path = path
}
if source, ok := data["source"].(string); ok {
model.Source = source
}
if description, ok := data["description"].(string); ok {
model.Description = description
}
@ -277,6 +280,31 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// PromptPresets
if promptPresets, ok := data["prompt_presets"]; ok && promptPresets != nil {
raw, err := jsoniter.Marshal(promptPresets)
if err == nil {
var pp map[string][]Prompt
if err := jsoniter.Unmarshal(raw, &pp); err == nil {
model.PromptPresets = pp
}
}
}
// DisableGlobalPrompts
model.DisableGlobalPrompts = getBoolValue(data, "disable_global_prompts")
// ConnectorOptions
if connectorOptions, ok := data["connector_options"]; ok && connectorOptions != nil {
raw, err := jsoniter.Marshal(connectorOptions)
if err == nil {
var co ConnectorOptions
if err := jsoniter.Unmarshal(raw, &co); err == nil {
model.ConnectorOptions = &co
}
}
}
// KB
if kb, ok := data["kb"]; ok && kb != nil {
kbConverted, err := ToKnowledgeBase(kb)
@ -301,17 +329,6 @@ func ToAssistantModel(v interface{}) (*AssistantModel, error) {
}
}
// 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)
@ -448,3 +465,56 @@ func ParseModelID(modelID string) string {
}
return parts[len(parts)-1]
}
// ToConnectorOptions converts various types to ConnectorOptions
func ToConnectorOptions(v interface{}) (*ConnectorOptions, error) {
if v == nil {
return nil, nil
}
switch opts := v.(type) {
case *ConnectorOptions:
return opts, nil
case ConnectorOptions:
return &opts, nil
default:
raw, err := jsoniter.Marshal(opts)
if err != nil {
return nil, fmt.Errorf("connector_options format error: %s", err.Error())
}
var connOpts ConnectorOptions
err = jsoniter.Unmarshal(raw, &connOpts)
if err != nil {
return nil, fmt.Errorf("connector_options format error: %s", err.Error())
}
return &connOpts, nil
}
}
// ToPromptPresets converts various types to map[string][]Prompt
func ToPromptPresets(v interface{}) (map[string][]Prompt, error) {
if v == nil {
return nil, nil
}
switch presets := v.(type) {
case map[string][]Prompt:
return presets, nil
default:
raw, err := jsoniter.Marshal(presets)
if err != nil {
return nil, fmt.Errorf("prompt_presets format error: %s", err.Error())
}
var result map[string][]Prompt
err = jsoniter.Unmarshal(raw, &result)
if err != nil {
return nil, fmt.Errorf("prompt_presets format error: %s", err.Error())
}
return result, nil
}
}

View file

@ -428,6 +428,11 @@ func TestToAssistantModel(t *testing.T) {
"name": "Test Assistant",
"avatar": "https://example.com/avatar.png",
"connector": "openai",
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai", "anthropic"},
"filters": []string{"vision", "tool_calls"},
},
"path": "/path/to/assistant",
"description": "Test description",
"share": "team",
@ -446,6 +451,16 @@ func TestToAssistantModel(t *testing.T) {
"prompts": []map[string]interface{}{
{"role": "system", "content": "You are helpful"},
},
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "You are a chat assistant"},
},
"task": []map[string]interface{}{
{"role": "system", "content": "You are a task assistant"},
},
},
"disable_global_prompts": true,
"source": "function hook() { return 'test'; }",
"kb": map[string]interface{}{
"collections": []string{"col1"},
},
@ -455,9 +470,6 @@ func TestToAssistantModel(t *testing.T) {
"workflow": map[string]interface{}{
"workflows": []string{"wf1"},
},
"tools": map[string]interface{}{
"calls": []string{"tool1"},
},
"placeholder": map[string]interface{}{
"title": "Enter message",
},
@ -489,9 +501,25 @@ func TestToAssistantModel(t *testing.T) {
if result.Connector != "openai" {
t.Errorf("Expected Connector 'openai', got '%s'", result.Connector)
}
if result.ConnectorOptions == nil {
t.Error("Expected ConnectorOptions to be set")
} else {
if !result.ConnectorOptions.Optional {
t.Error("Expected ConnectorOptions.Optional to be true")
}
if len(result.ConnectorOptions.Connectors) != 2 {
t.Errorf("Expected 2 connectors in options, got %d", len(result.ConnectorOptions.Connectors))
}
if len(result.ConnectorOptions.Filters) != 2 {
t.Errorf("Expected 2 filters, got %d", len(result.ConnectorOptions.Filters))
}
}
if result.Path != "/path/to/assistant" {
t.Errorf("Expected Path, got '%s'", result.Path)
}
if result.Source != "function hook() { return 'test'; }" {
t.Errorf("Expected Source, got '%s'", result.Source)
}
if result.Description != "Test description" {
t.Errorf("Expected Description, got '%s'", result.Description)
}
@ -531,6 +559,26 @@ func TestToAssistantModel(t *testing.T) {
if len(result.Prompts) != 1 {
t.Errorf("Expected 1 prompt, got %d", len(result.Prompts))
}
if result.PromptPresets == nil {
t.Error("Expected PromptPresets to be set")
} else {
if len(result.PromptPresets) != 2 {
t.Errorf("Expected 2 prompt presets, got %d", len(result.PromptPresets))
}
if chatPrompts, ok := result.PromptPresets["chat"]; !ok {
t.Error("Expected 'chat' prompt preset")
} else if len(chatPrompts) != 1 {
t.Errorf("Expected 1 chat prompt, got %d", len(chatPrompts))
}
if taskPrompts, ok := result.PromptPresets["task"]; !ok {
t.Error("Expected 'task' prompt preset")
} else if len(taskPrompts) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(taskPrompts))
}
}
if !result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be true")
}
if result.KB == nil {
t.Error("Expected KB to be set")
}
@ -540,9 +588,6 @@ func TestToAssistantModel(t *testing.T) {
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")
}
@ -583,7 +628,6 @@ func TestToAssistantModel(t *testing.T) {
"kb": nil,
"mcp": nil,
"workflow": nil,
"tools": nil,
"placeholder": nil,
"locales": nil,
}
@ -659,6 +703,221 @@ func TestToAssistantModel(t *testing.T) {
})
}
// TestToAssistantModelNewFields tests the newly added fields
func TestToAssistantModelNewFields(t *testing.T) {
t.Run("ConnectorOptions", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai", "anthropic", "azure"},
"filters": []string{"vision", "tool_calls", "audio"},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions == nil {
t.Fatal("Expected ConnectorOptions to be set")
}
if !result.ConnectorOptions.Optional {
t.Error("Expected Optional to be true")
}
if len(result.ConnectorOptions.Connectors) != 3 {
t.Errorf("Expected 3 connectors, got %d", len(result.ConnectorOptions.Connectors))
}
if len(result.ConnectorOptions.Filters) != 3 {
t.Errorf("Expected 3 filters, got %d", len(result.ConnectorOptions.Filters))
}
})
t.Run("PromptPresets", func(t *testing.T) {
data := map[string]interface{}{
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "You are a helpful chat assistant"},
{"role": "user", "content": "Example question"},
},
"task": []map[string]interface{}{
{"role": "system", "content": "You are a task completion assistant"},
},
"analyze": []map[string]interface{}{
{"role": "system", "content": "You are a data analysis assistant"},
},
},
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.PromptPresets == nil {
t.Fatal("Expected PromptPresets to be set")
}
if len(result.PromptPresets) != 3 {
t.Errorf("Expected 3 prompt preset modes, got %d", len(result.PromptPresets))
}
if chatPrompts, ok := result.PromptPresets["chat"]; !ok {
t.Error("Expected 'chat' mode in prompt presets")
} else if len(chatPrompts) != 2 {
t.Errorf("Expected 2 prompts in chat mode, got %d", len(chatPrompts))
}
if taskPrompts, ok := result.PromptPresets["task"]; !ok {
t.Error("Expected 'task' mode in prompt presets")
} else if len(taskPrompts) != 1 {
t.Errorf("Expected 1 prompt in task mode, got %d", len(taskPrompts))
}
if analyzePrompts, ok := result.PromptPresets["analyze"]; !ok {
t.Error("Expected 'analyze' mode in prompt presets")
} else if len(analyzePrompts) != 1 {
t.Errorf("Expected 1 prompt in analyze mode, got %d", len(analyzePrompts))
}
})
t.Run("Source", func(t *testing.T) {
hookScript := `
function beforeChat(context) {
console.log('Hook called');
return context;
}
`
data := map[string]interface{}{
"source": hookScript,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.Source != hookScript {
t.Errorf("Expected Source to match, got '%s'", result.Source)
}
})
t.Run("AllNewFields", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": map[string]interface{}{
"optional": true,
"connectors": []string{"openai"},
"filters": []string{"vision"},
},
"prompt_presets": map[string]interface{}{
"chat": []map[string]interface{}{
{"role": "system", "content": "Chat mode"},
},
},
"disable_global_prompts": true,
"source": "function test() {}",
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions == nil {
t.Error("Expected ConnectorOptions to be set")
}
if result.PromptPresets == nil {
t.Error("Expected PromptPresets to be set")
}
if !result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be true")
}
if result.Source == "" {
t.Error("Expected Source to be set")
}
})
t.Run("NilNewFields", func(t *testing.T) {
data := map[string]interface{}{
"connector_options": nil,
"prompt_presets": nil,
"disable_global_prompts": nil,
"source": nil,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.ConnectorOptions != nil {
t.Error("Expected ConnectorOptions to be nil")
}
if result.PromptPresets != nil {
t.Error("Expected PromptPresets to be nil")
}
if result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be false")
}
if result.Source != "" {
t.Error("Expected Source to be empty")
}
})
t.Run("DisableGlobalPrompts", func(t *testing.T) {
// Test with true
data := map[string]interface{}{
"disable_global_prompts": true,
}
result, err := ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be true")
}
// Test with false
data = map[string]interface{}{
"disable_global_prompts": false,
}
result, err = ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be false")
}
// Test with int 1
data = map[string]interface{}{
"disable_global_prompts": 1,
}
result, err = ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be true for int 1")
}
// Test with string "true"
data = map[string]interface{}{
"disable_global_prompts": "true",
}
result, err = ToAssistantModel(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.DisableGlobalPrompts {
t.Error("Expected DisableGlobalPrompts to be true for string 'true'")
}
})
}
// TestToAssistantModelComplexTypes tests complex type conversions in ToAssistantModel
func TestToAssistantModelComplexTypes(t *testing.T) {
t.Run("CompleteLocales", func(t *testing.T) {
@ -951,6 +1210,273 @@ func TestModelID(t *testing.T) {
})
}
// TestToConnectorOptions tests the ToConnectorOptions conversion function
func TestToConnectorOptions(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToConnectorOptions(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("ConnectorOptionsPointer", func(t *testing.T) {
opts := &ConnectorOptions{
Optional: true,
Connectors: []string{"openai", "anthropic"},
Filters: []ModelCapability{CapVision, CapToolCalls},
}
result, err := ToConnectorOptions(opts)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result != opts {
t.Errorf("Expected same pointer")
}
})
t.Run("ConnectorOptionsValue", func(t *testing.T) {
opts := ConnectorOptions{
Optional: true,
Connectors: []string{"openai", "anthropic"},
Filters: []ModelCapability{CapVision, CapToolCalls},
}
result, err := ToConnectorOptions(opts)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.Optional {
t.Error("Expected Optional to be true")
}
if len(result.Connectors) != 2 {
t.Errorf("Expected 2 connectors, got %d", len(result.Connectors))
}
if len(result.Filters) != 2 {
t.Errorf("Expected 2 filters, got %d", len(result.Filters))
}
})
t.Run("MapInput", func(t *testing.T) {
data := map[string]interface{}{
"optional": true,
"connectors": []string{"openai", "anthropic", "azure"},
"filters": []string{"vision", "tool_calls", "audio"},
}
result, err := ToConnectorOptions(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.Optional {
t.Error("Expected Optional to be true")
}
if len(result.Connectors) != 3 {
t.Errorf("Expected 3 connectors, got %d", len(result.Connectors))
}
if len(result.Filters) != 3 {
t.Errorf("Expected 3 filters, got %d", len(result.Filters))
}
})
t.Run("MapInputOptionalOnly", func(t *testing.T) {
data := map[string]interface{}{
"optional": true,
}
result, err := ToConnectorOptions(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if !result.Optional {
t.Error("Expected Optional to be true")
}
if result.Connectors != nil {
t.Error("Expected Connectors to be nil")
}
if result.Filters != nil {
t.Error("Expected Filters to be nil")
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToConnectorOptions(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 ConnectorOptions
data := map[string]interface{}{
"invalid_field": "should cause unmarshal to fail gracefully",
}
result, err := ToConnectorOptions(data)
// Should not error, just return empty ConnectorOptions
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
})
}
// TestToPromptPresets tests the ToPromptPresets conversion function
func TestToPromptPresets(t *testing.T) {
t.Run("NilInput", func(t *testing.T) {
result, err := ToPromptPresets(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("MapStringPromptSlice", func(t *testing.T) {
presets := map[string][]Prompt{
"chat": {
{Role: "system", Content: "You are a chat assistant"},
},
"task": {
{Role: "system", Content: "You are a task assistant"},
},
}
result, err := ToPromptPresets(presets)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 2 {
t.Errorf("Expected 2 presets, got %d", len(result))
}
if len(result["chat"]) != 1 {
t.Errorf("Expected 1 chat prompt, got %d", len(result["chat"]))
}
if len(result["task"]) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(result["task"]))
}
})
t.Run("MapInput", func(t *testing.T) {
data := map[string]interface{}{
"chat": []interface{}{
map[string]interface{}{"role": "system", "content": "Chat mode system prompt"},
map[string]interface{}{"role": "user", "content": "Example user message"},
},
"task": []interface{}{
map[string]interface{}{"role": "system", "content": "Task mode system prompt"},
},
"analyze": []interface{}{
map[string]interface{}{"role": "system", "content": "Analyze mode system prompt"},
},
}
result, err := ToPromptPresets(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 3 {
t.Errorf("Expected 3 presets, got %d", len(result))
}
if len(result["chat"]) != 2 {
t.Errorf("Expected 2 chat prompts, got %d", len(result["chat"]))
}
if len(result["task"]) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(result["task"]))
}
if len(result["analyze"]) != 1 {
t.Errorf("Expected 1 analyze prompt, got %d", len(result["analyze"]))
}
})
t.Run("EmptyMap", func(t *testing.T) {
data := map[string]interface{}{}
result, err := ToPromptPresets(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if result == nil {
t.Error("Expected non-nil result")
}
if len(result) != 0 {
t.Errorf("Expected empty map, got %d entries", len(result))
}
})
t.Run("SinglePreset", func(t *testing.T) {
data := map[string]interface{}{
"default": []interface{}{
map[string]interface{}{"role": "system", "content": "Default prompt"},
},
}
result, err := ToPromptPresets(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result) != 1 {
t.Errorf("Expected 1 preset, got %d", len(result))
}
if _, ok := result["default"]; !ok {
t.Error("Expected 'default' key in result")
}
})
t.Run("InvalidInput", func(t *testing.T) {
// Test with data that can't be marshaled
invalidData := make(chan int)
_, err := ToPromptPresets(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 map[string][]Prompt
// This is a string that can be marshaled but won't unmarshal to the expected type
data := "not a map"
_, err := ToPromptPresets(data)
if err == nil {
t.Error("Expected error for invalid JSON unmarshal")
}
})
t.Run("PromptWithAllFields", func(t *testing.T) {
data := map[string]interface{}{
"advanced": []interface{}{
map[string]interface{}{
"role": "system",
"content": "Advanced system prompt",
"name": "system-prompt",
},
map[string]interface{}{
"role": "user",
"content": "User example",
"name": "user-example",
},
map[string]interface{}{
"role": "assistant",
"content": "Assistant response",
"name": "assistant-response",
},
},
}
result, err := ToPromptPresets(data)
if err != nil {
t.Errorf("Expected no error, got: %v", err)
}
if len(result["advanced"]) != 3 {
t.Errorf("Expected 3 prompts in advanced, got %d", len(result["advanced"]))
}
if result["advanced"][0].Role != "system" {
t.Errorf("Expected role 'system', got '%s'", result["advanced"][0].Role)
}
if result["advanced"][0].Content != "Advanced system prompt" {
t.Errorf("Expected content 'Advanced system prompt', got '%s'", result["advanced"][0].Content)
}
})
}
// TestParseModelID tests the ParseModelID function
func TestParseModelID(t *testing.T) {
t.Run("ValidModelID", func(t *testing.T) {

View file

@ -10,6 +10,7 @@ var AssistantAllowedFields = map[string]bool{
"name": true,
"avatar": true,
"connector": true,
"connector_options": true,
"description": true,
"path": true,
"sort": true,
@ -17,15 +18,18 @@ var AssistantAllowedFields = map[string]bool{
"placeholder": true,
"options": true,
"prompts": true,
"prompt_presets": true,
"disable_global_prompts": true,
"workflow": true,
"kb": true,
"mcp": true,
"tools": true,
"source": true,
"tags": true,
"readonly": true,
"public": true,
"share": true,
"locales": true,
"uses": true,
"automated": true,
"mentionable": true,
"created_at": true,
@ -37,6 +41,7 @@ var AssistantAllowedFields = map[string]bool{
}
// AssistantDefaultFields defines the default fields to select for assistants when no specific fields are requested
// These are lightweight fields suitable for list views and basic information display
var AssistantDefaultFields = []string{
"assistant_id",
"type",
@ -44,6 +49,7 @@ var AssistantDefaultFields = []string{
"avatar",
"connector",
"description",
"tags", // Tags for categorization (lightweight)
"sort",
"built_in",
"readonly",
@ -51,8 +57,52 @@ var AssistantDefaultFields = []string{
"share",
"automated",
"mentionable",
"kb", // Knowledge base configuration (lightweight)
"mcp", // MCP servers configuration (lightweight)
"created_at",
"updated_at",
"__yao_created_by", // Permission: creator user ID
"__yao_updated_by", // Permission: updater user ID
"__yao_team_id", // Permission: team ID
"__yao_tenant_id", // Permission: tenant ID
}
// AssistantFullFields defines all available fields including complex/large fields
// Use this when you need complete assistant data for backend processing
var AssistantFullFields = []string{
"assistant_id",
"type",
"name",
"avatar",
"connector",
"connector_options",
"description",
"path",
"sort",
"built_in",
"placeholder",
"options",
"prompts",
"prompt_presets",
"disable_global_prompts",
"workflow",
"kb",
"mcp",
"source",
"tags",
"readonly",
"public",
"share",
"locales",
"uses",
"automated",
"mentionable",
"created_at",
"updated_at",
"__yao_created_by",
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
// ValidateAssistantFields validates and filters assistant select fields against the whitelist

View file

@ -118,12 +118,16 @@ func TestAssistantAllowedFields(t *testing.T) {
complexFields := []string{
"options",
"prompts",
"prompt_presets",
"disable_global_prompts",
"workflow",
"kb",
"mcp",
"tools",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
for _, field := range complexFields {
if !AssistantAllowedFields[field] {
@ -139,6 +143,12 @@ func TestAssistantDefaultFields(t *testing.T) {
"assistant_id",
"name",
"type",
"kb", // Knowledge base is essential for assistant functionality
"mcp", // MCP servers are essential for assistant functionality
"__yao_created_by", // Permission fields are essential for access control
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
defaultFieldsMap := make(map[string]bool)
@ -155,15 +165,17 @@ func TestAssistantDefaultFields(t *testing.T) {
t.Run("DoesNotContainSensitiveFields", func(t *testing.T) {
// Default fields should not include complex/large fields by default
// Note: kb, mcp, and tags are lightweight and included in defaults
sensitiveFields := []string{
"options",
"prompts",
"prompt_presets",
"workflow",
"kb",
"mcp",
"tools",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
defaultFieldsMap := make(map[string]bool)
@ -178,3 +190,82 @@ func TestAssistantDefaultFields(t *testing.T) {
}
})
}
func TestAssistantFullFields(t *testing.T) {
t.Run("ContainsAllAllowedFields", func(t *testing.T) {
// Full fields should contain all fields from allowed fields
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for field := range AssistantAllowedFields {
if field == "id" {
// "id" is an alias for "assistant_id", skip
continue
}
if !fullFieldsMap[field] {
t.Errorf("Allowed field %s is missing from full fields", field)
}
}
})
t.Run("AllFieldsAreAllowed", func(t *testing.T) {
// All fields in full list should be in allowed fields
for _, field := range AssistantFullFields {
if !AssistantAllowedFields[field] {
t.Errorf("Full field %s is not in allowed fields", field)
}
}
})
t.Run("ContainsComplexFields", func(t *testing.T) {
// Full fields should include all complex/large fields
complexFields := []string{
"options",
"prompts",
"prompt_presets",
"disable_global_prompts",
"workflow",
"kb",
"mcp",
"placeholder",
"locales",
"uses",
"connector_options",
"source",
}
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for _, field := range complexFields {
if !fullFieldsMap[field] {
t.Errorf("Complex field %s is missing from full fields", field)
}
}
})
t.Run("ContainsPermissionFields", func(t *testing.T) {
// Full fields should include permission fields
permissionFields := []string{
"__yao_created_by",
"__yao_updated_by",
"__yao_team_id",
"__yao_tenant_id",
}
fullFieldsMap := make(map[string]bool)
for _, field := range AssistantFullFields {
fullFieldsMap[field] = true
}
for _, field := range permissionFields {
if !fullFieldsMap[field] {
t.Errorf("Permission field %s is missing from full fields", field)
}
}
})
}

284
agent/store/types/prompt.go Normal file
View file

@ -0,0 +1,284 @@
package types
import (
"os"
"path/filepath"
"regexp"
"strings"
"time"
"github.com/yaoapp/gou/application"
"github.com/yaoapp/gou/fs"
"gopkg.in/yaml.v3"
)
// Prompts is a slice of Prompt with helper methods
type Prompts []Prompt
// SystemVariables defines the available system variables
// These are computed at parse time
var SystemVariables = map[string]func() string{
"TIME": func() string { return time.Now().Format("15:04:05") },
"DATE": func() string { return time.Now().Format("2006-01-02") },
"DATETIME": func() string { return time.Now().Format("2006-01-02 15:04:05") },
"TIMEZONE": func() string { return time.Now().Location().String() },
"WEEKDAY": func() string { return time.Now().Weekday().String() },
"YEAR": func() string { return time.Now().Format("2006") },
"MONTH": func() string { return time.Now().Format("01") },
"DAY": func() string { return time.Now().Format("02") },
"HOUR": func() string { return time.Now().Format("15") },
"MINUTE": func() string { return time.Now().Format("04") },
"SECOND": func() string { return time.Now().Format("05") },
"UNIX": func() string { return time.Now().Format("1136239445") }, // Unix timestamp
}
// Regular expressions for variable replacement
var (
reSysVar = regexp.MustCompile(`\$SYS\.([A-Z_]+)`)
reEnvVar = regexp.MustCompile(`\$ENV\.([A-Za-z_][A-Za-z0-9_]*)`)
reCtxVar = regexp.MustCompile(`\$CTX\.([A-Za-z_][A-Za-z0-9_]*)`)
reAssetRef = regexp.MustCompile(`@assets/([^\s]+\.(md|yml|yaml|json|txt))`)
)
// LoadPrompts loads prompts from a YAML file
// Handles @assets/* replacement at load time
// file: prompt file path relative to app root (e.g., "assistants/test/prompts.yml")
// root: resource root directory for assets (e.g., "assistants/test")
// Returns: prompts slice, modification timestamp, error
func LoadPrompts(file string, root string) ([]Prompt, int64, error) {
app, err := fs.Get("app")
if err != nil {
return nil, 0, err
}
ts, err := app.ModTime(file)
if err != nil {
return nil, 0, err
}
content, err := app.ReadFile(file)
if err != nil {
return nil, 0, err
}
// Replace @assets/xxx references with file content
content = replaceAssets(content, root, app)
// Parse prompts
var prompts []Prompt
err = yaml.Unmarshal(content, &prompts)
if err != nil {
return nil, 0, err
}
return prompts, ts.UnixNano(), nil
}
// LoadPromptsRaw loads raw prompt content from a YAML file
// Handles @assets/* replacement at load time
// Returns raw YAML string for further processing
func LoadPromptsRaw(file string, root string) (string, int64, error) {
app, err := fs.Get("app")
if err != nil {
return "", 0, err
}
ts, err := app.ModTime(file)
if err != nil {
return "", 0, err
}
content, err := app.ReadFile(file)
if err != nil {
return "", 0, err
}
// Replace @assets/xxx references with file content
content = replaceAssets(content, root, app)
return string(content), ts.UnixNano(), nil
}
// LoadGlobalPrompts loads global prompts from agent/prompts.yml
// Returns: prompts slice, modification timestamp, error
func LoadGlobalPrompts() ([]Prompt, int64, error) {
file := filepath.Join("agent", "prompts.yml")
// Check if file exists
exists, _ := application.App.Exists(file)
if !exists {
return nil, 0, nil
}
return LoadPrompts(file, "agent")
}
// LoadPromptPresets loads prompt presets from a directory
// Supports multi-level directories, key is path with "/" replaced by "."
// e.g., prompts/chat/friendly.yml -> "chat.friendly"
func LoadPromptPresets(dir string, root string) (map[string][]Prompt, int64, error) {
app, err := fs.Get("app")
if err != nil {
return nil, 0, err
}
// Check if directory exists
exists, _ := app.Exists(dir)
if !exists {
return nil, 0, nil
}
// Read directory recursively
files, err := app.ReadDir(dir, true)
if err != nil {
return nil, 0, err
}
presets := make(map[string][]Prompt)
var latestTs int64
for _, file := range files {
// Only process .yml/.yaml files
if !strings.HasSuffix(file, ".yml") && !strings.HasSuffix(file, ".yaml") {
continue
}
ts, err := app.ModTime(file)
if err != nil {
return nil, 0, err
}
if ts.UnixNano() > latestTs {
latestTs = ts.UnixNano()
}
// Read file content
content, err := app.ReadFile(file)
if err != nil {
return nil, 0, err
}
// Replace @assets/xxx references with file content
content = replaceAssets(content, root, app)
// Parse prompts
var prompts []Prompt
err = yaml.Unmarshal(content, &prompts)
if err != nil {
return nil, 0, err
}
// Build key: get relative path from dir, remove extension and replace "/" with "."
relPath := strings.TrimPrefix(file, dir+"/")
key := strings.TrimSuffix(relPath, filepath.Ext(relPath))
key = strings.ReplaceAll(key, "/", ".")
presets[key] = prompts
}
return presets, latestTs, nil
}
// replaceAssets replaces @assets/xxx references with file content
func replaceAssets(content []byte, root string, app fs.FileSystem) []byte {
return reAssetRef.ReplaceAllFunc(content, func(s []byte) []byte {
matches := reAssetRef.FindStringSubmatch(string(s))
if len(matches) < 2 {
return s
}
asset := matches[1]
assetFile := filepath.Join(root, "assets", asset)
assetContent, err := app.ReadFile(assetFile)
if err != nil {
return []byte("")
}
// Add proper YAML formatting for content (multiline string)
lines := strings.Split(string(assetContent), "\n")
formattedContent := "|\n"
for _, line := range lines {
formattedContent += " " + line + "\n"
}
return []byte(formattedContent)
})
}
// Parse parses a single prompt, replacing variables
// ctx: context variables map, key corresponds to $CTX.{key}
// Returns a new Prompt with variables replaced
func (p Prompt) Parse(ctx map[string]string) Prompt {
result := Prompt{
Role: p.Role,
Content: parseVariables(p.Content, ctx),
Name: p.Name,
}
return result
}
// Parse parses all prompts in the slice, replacing variables
// ctx: context variables map, key corresponds to $CTX.{key}
// Returns a new Prompts slice with variables replaced
func (ps Prompts) Parse(ctx map[string]string) Prompts {
result := make(Prompts, len(ps))
for i, p := range ps {
result[i] = p.Parse(ctx)
}
return result
}
// parseVariables replaces all variable types in content
func parseVariables(content string, ctx map[string]string) string {
// Replace $SYS.* variables
content = reSysVar.ReplaceAllStringFunc(content, func(s string) string {
matches := reSysVar.FindStringSubmatch(s)
if len(matches) < 2 {
return s
}
varName := matches[1]
if fn, ok := SystemVariables[varName]; ok {
return fn()
}
return s // Keep original if not found
})
// Replace $ENV.* variables
content = reEnvVar.ReplaceAllStringFunc(content, func(s string) string {
matches := reEnvVar.FindStringSubmatch(s)
if len(matches) < 2 {
return s
}
varName := matches[1]
return os.Getenv(varName)
})
// Replace $CTX.* variables
if ctx != nil {
content = reCtxVar.ReplaceAllStringFunc(content, func(s string) string {
matches := reCtxVar.FindStringSubmatch(s)
if len(matches) < 2 {
return s
}
varName := matches[1]
if val, ok := ctx[varName]; ok {
return val
}
return "" // Empty string if not found in ctx
})
}
return content
}
// Merge merges two prompt slices
// globalPrompts are prepended to assistantPrompts
func Merge(globalPrompts, assistantPrompts []Prompt) []Prompt {
if len(globalPrompts) == 0 {
return assistantPrompts
}
if len(assistantPrompts) == 0 {
return globalPrompts
}
result := make([]Prompt, 0, len(globalPrompts)+len(assistantPrompts))
result = append(result, globalPrompts...)
result = append(result, assistantPrompts...)
return result
}

View file

@ -0,0 +1,432 @@
package types
import (
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
)
func TestPromptParse(t *testing.T) {
tests := []struct {
name string
prompt Prompt
ctx map[string]string
validate func(t *testing.T, result Prompt)
}{
{
name: "ParseSysTimeVariables",
prompt: Prompt{
Role: "system",
Content: "Current time: $SYS.TIME, Date: $SYS.DATE",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
assert.Equal(t, "system", result.Role)
// Check that variables are replaced (not exact match due to time)
assert.NotContains(t, result.Content, "$SYS.TIME")
assert.NotContains(t, result.Content, "$SYS.DATE")
assert.Contains(t, result.Content, "Current time:")
assert.Contains(t, result.Content, "Date:")
},
},
{
name: "ParseSysDatetimeVariable",
prompt: Prompt{
Role: "system",
Content: "Now: $SYS.DATETIME, Timezone: $SYS.TIMEZONE",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
assert.NotContains(t, result.Content, "$SYS.DATETIME")
assert.NotContains(t, result.Content, "$SYS.TIMEZONE")
},
},
{
name: "ParseSysWeekdayVariable",
prompt: Prompt{
Role: "system",
Content: "Today is $SYS.WEEKDAY",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
weekday := time.Now().Weekday().String()
assert.Contains(t, result.Content, weekday)
},
},
{
name: "ParseSysYearMonthDay",
prompt: Prompt{
Role: "system",
Content: "Year: $SYS.YEAR, Month: $SYS.MONTH, Day: $SYS.DAY",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
now := time.Now()
assert.Contains(t, result.Content, now.Format("2006"))
assert.Contains(t, result.Content, now.Format("01"))
assert.Contains(t, result.Content, now.Format("02"))
},
},
{
name: "ParseSysHourMinuteSecond",
prompt: Prompt{
Role: "system",
Content: "Hour: $SYS.HOUR, Minute: $SYS.MINUTE, Second: $SYS.SECOND",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
assert.NotContains(t, result.Content, "$SYS.HOUR")
assert.NotContains(t, result.Content, "$SYS.MINUTE")
assert.NotContains(t, result.Content, "$SYS.SECOND")
},
},
{
name: "ParseEnvVariable",
prompt: Prompt{
Role: "system",
Content: "App: $ENV.TEST_APP_NAME",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
assert.Contains(t, result.Content, "App: TestApp")
},
},
{
name: "ParseEnvVariableNotFound",
prompt: Prompt{
Role: "system",
Content: "Value: $ENV.NOT_EXIST_VAR_12345",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
// Should be replaced with empty string
assert.Equal(t, "Value: ", result.Content)
},
},
{
name: "ParseCtxVariables",
prompt: Prompt{
Role: "system",
Content: "User: $CTX.USER_ID, Locale: $CTX.LOCALE",
},
ctx: map[string]string{
"USER_ID": "user-123",
"LOCALE": "zh-CN",
},
validate: func(t *testing.T, result Prompt) {
assert.Equal(t, "User: user-123, Locale: zh-CN", result.Content)
},
},
{
name: "ParseCtxVariableNotFound",
prompt: Prompt{
Role: "system",
Content: "Value: $CTX.NOT_EXIST",
},
ctx: map[string]string{
"OTHER": "value",
},
validate: func(t *testing.T, result Prompt) {
// Should be replaced with empty string
assert.Equal(t, "Value: ", result.Content)
},
},
{
name: "ParseCtxWithNilMap",
prompt: Prompt{
Role: "system",
Content: "Value: $CTX.SOMETHING",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
// Should keep original when ctx is nil
assert.Equal(t, "Value: $CTX.SOMETHING", result.Content)
},
},
{
name: "ParseMixedVariables",
prompt: Prompt{
Role: "system",
Content: "Time: $SYS.TIME, App: $ENV.TEST_APP_NAME, User: $CTX.USER_ID",
},
ctx: map[string]string{
"USER_ID": "user-456",
},
validate: func(t *testing.T, result Prompt) {
assert.NotContains(t, result.Content, "$SYS.TIME")
assert.Contains(t, result.Content, "App: TestApp")
assert.Contains(t, result.Content, "User: user-456")
},
},
{
name: "ParseUnknownSysVariable",
prompt: Prompt{
Role: "system",
Content: "Value: $SYS.UNKNOWN_VAR",
},
ctx: nil,
validate: func(t *testing.T, result Prompt) {
// Should keep original if not found
assert.Equal(t, "Value: $SYS.UNKNOWN_VAR", result.Content)
},
},
{
name: "ParsePreservesRoleAndName",
prompt: Prompt{
Role: "user",
Content: "Hello $CTX.NAME",
Name: "test_user",
},
ctx: map[string]string{
"NAME": "World",
},
validate: func(t *testing.T, result Prompt) {
assert.Equal(t, "user", result.Role)
assert.Equal(t, "Hello World", result.Content)
assert.Equal(t, "test_user", result.Name)
},
},
{
name: "ParseCustomCtxVariables",
prompt: Prompt{
Role: "system",
Content: "Custom: $CTX.MY_CUSTOM_VAR, Another: $CTX.ANOTHER_VAR",
},
ctx: map[string]string{
"MY_CUSTOM_VAR": "custom-value",
"ANOTHER_VAR": "another-value",
},
validate: func(t *testing.T, result Prompt) {
assert.Equal(t, "Custom: custom-value, Another: another-value", result.Content)
},
},
{
name: "ParseMultilineContent",
prompt: Prompt{
Role: "system",
Content: `# System Context
Current Time: $SYS.TIME
User: $CTX.USER_ID
App: $ENV.TEST_APP_NAME`,
},
ctx: map[string]string{
"USER_ID": "user-789",
},
validate: func(t *testing.T, result Prompt) {
assert.Contains(t, result.Content, "# System Context")
assert.NotContains(t, result.Content, "$SYS.TIME")
assert.Contains(t, result.Content, "User: user-789")
assert.Contains(t, result.Content, "App: TestApp")
},
},
}
// Set test environment variable
os.Setenv("TEST_APP_NAME", "TestApp")
defer os.Unsetenv("TEST_APP_NAME")
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.prompt.Parse(tt.ctx)
tt.validate(t, result)
})
}
}
func TestPromptsParse(t *testing.T) {
os.Setenv("TEST_APP_NAME", "TestApp")
defer os.Unsetenv("TEST_APP_NAME")
prompts := Prompts{
{Role: "system", Content: "Time: $SYS.TIME"},
{Role: "system", Content: "User: $CTX.USER_ID"},
{Role: "user", Content: "App: $ENV.TEST_APP_NAME"},
}
ctx := map[string]string{
"USER_ID": "user-123",
}
result := prompts.Parse(ctx)
assert.Len(t, result, 3)
assert.NotContains(t, result[0].Content, "$SYS.TIME")
assert.Equal(t, "User: user-123", result[1].Content)
assert.Equal(t, "App: TestApp", result[2].Content)
}
func TestMergePrompts(t *testing.T) {
tests := []struct {
name string
globalPrompts []Prompt
assistantPrompts []Prompt
expectedLen int
validate func(t *testing.T, result []Prompt)
}{
{
name: "MergeBothNonEmpty",
globalPrompts: []Prompt{
{Role: "system", Content: "Global prompt 1"},
{Role: "system", Content: "Global prompt 2"},
},
assistantPrompts: []Prompt{
{Role: "system", Content: "Assistant prompt 1"},
},
expectedLen: 3,
validate: func(t *testing.T, result []Prompt) {
assert.Equal(t, "Global prompt 1", result[0].Content)
assert.Equal(t, "Global prompt 2", result[1].Content)
assert.Equal(t, "Assistant prompt 1", result[2].Content)
},
},
{
name: "MergeGlobalEmpty",
globalPrompts: []Prompt{},
assistantPrompts: []Prompt{
{Role: "system", Content: "Assistant prompt"},
},
expectedLen: 1,
validate: func(t *testing.T, result []Prompt) {
assert.Equal(t, "Assistant prompt", result[0].Content)
},
},
{
name: "MergeAssistantEmpty",
globalPrompts: []Prompt{
{Role: "system", Content: "Global prompt"},
},
assistantPrompts: []Prompt{},
expectedLen: 1,
validate: func(t *testing.T, result []Prompt) {
assert.Equal(t, "Global prompt", result[0].Content)
},
},
{
name: "MergeBothEmpty",
globalPrompts: []Prompt{},
assistantPrompts: []Prompt{},
expectedLen: 0,
validate: func(t *testing.T, result []Prompt) {
assert.Empty(t, result)
},
},
{
name: "MergeGlobalNil",
globalPrompts: nil,
assistantPrompts: []Prompt{
{Role: "system", Content: "Assistant prompt"},
},
expectedLen: 1,
validate: func(t *testing.T, result []Prompt) {
assert.Equal(t, "Assistant prompt", result[0].Content)
},
},
{
name: "MergeAssistantNil",
globalPrompts: []Prompt{
{Role: "system", Content: "Global prompt"},
},
assistantPrompts: nil,
expectedLen: 1,
validate: func(t *testing.T, result []Prompt) {
assert.Equal(t, "Global prompt", result[0].Content)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := Merge(tt.globalPrompts, tt.assistantPrompts)
assert.Len(t, result, tt.expectedLen)
if tt.validate != nil {
tt.validate(t, result)
}
})
}
}
func TestSystemVariables(t *testing.T) {
// Test that all system variables are defined and return non-empty values
expectedVars := []string{
"TIME", "DATE", "DATETIME", "TIMEZONE", "WEEKDAY",
"YEAR", "MONTH", "DAY", "HOUR", "MINUTE", "SECOND", "UNIX",
}
for _, varName := range expectedVars {
t.Run(varName, func(t *testing.T) {
fn, ok := SystemVariables[varName]
assert.True(t, ok, "SystemVariables should contain %s", varName)
value := fn()
assert.NotEmpty(t, value, "SystemVariables[%s]() should return non-empty value", varName)
})
}
}
func TestParseVariablesEdgeCases(t *testing.T) {
os.Setenv("TEST_VAR", "test-value")
defer os.Unsetenv("TEST_VAR")
tests := []struct {
name string
content string
ctx map[string]string
expected string
}{
{
name: "EmptyContent",
content: "",
ctx: nil,
expected: "",
},
{
name: "NoVariables",
content: "Hello, World!",
ctx: nil,
expected: "Hello, World!",
},
{
name: "PartialVariableSyntax",
content: "Value: $SYS Value: $ENV Value: $CTX",
ctx: nil,
expected: "Value: $SYS Value: $ENV Value: $CTX",
},
{
name: "VariableInMiddleOfWord",
content: "prefix$SYS.TIMEsuffix",
ctx: nil,
expected: "prefix$SYS.TIMEsuffix", // Should not match - variable must be followed by valid char
},
{
name: "MultipleOccurrences",
content: "$CTX.VAR and $CTX.VAR again",
ctx: map[string]string{"VAR": "value"},
expected: "value and value again",
},
{
name: "SpecialCharactersInValue",
content: "User: $CTX.USER",
ctx: map[string]string{"USER": "user@example.com"},
expected: "User: user@example.com",
},
{
name: "UnicodeInValue",
content: "Name: $CTX.NAME",
ctx: map[string]string{"NAME": "用户名"},
expected: "Name: 用户名",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := parseVariables(tt.content, tt.ctx)
if tt.name == "VariableInMiddleOfWord" {
// This case depends on regex behavior - just check it doesn't crash
assert.NotEmpty(t, result)
} else {
assert.Equal(t, tt.expected, result)
}
})
}
}

View file

@ -91,8 +91,10 @@ type Store interface {
// GetAssistant retrieves a single assistant by ID
// assistantID: Assistant ID
// fields: List of fields to select, empty/nil means default fields (AssistantDefaultFields)
// locale: Optional locale for i18n translations
// Returns: Assistant information and potential error
GetAssistant(assistantID string, locale ...string) (*AssistantModel, error)
GetAssistant(assistantID string, fields []string, locale ...string) (*AssistantModel, error)
// DeleteAssistants deletes assistants based on filter conditions
// filter: Filter conditions

View file

@ -221,13 +221,37 @@ type Placeholder struct {
Prompts []string `json:"prompts,omitempty"`
}
// ModelCapability defines the available model capability filters
type ModelCapability string
// Model capability constants for filtering connectors
const (
CapVision ModelCapability = "vision"
CapAudio ModelCapability = "audio"
CapToolCalls ModelCapability = "tool_calls"
CapReasoning ModelCapability = "reasoning"
CapStreaming ModelCapability = "streaming"
CapJSON ModelCapability = "json"
CapMultimodal ModelCapability = "multimodal"
CapTemperatureAdjustable ModelCapability = "temperature_adjustable"
)
// ConnectorOptions the connector selection options
// Allows defining optional connector selection with filtering capabilities
type ConnectorOptions struct {
Optional bool `json:"optional,omitempty"` // Whether connector is optional for user selection
Connectors []string `json:"connectors,omitempty"` // List of available connectors, empty means all connectors are available
Filters []ModelCapability `json:"filters,omitempty"` // Filter by model capabilities, conditions can be stacked
}
// 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
Connector string `json:"connector"` // AI Connector (default connector)
ConnectorOptions *ConnectorOptions `json:"connector_options,omitempty"` // Connector selection options for user to choose from
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
@ -239,12 +263,14 @@ type AssistantModel struct {
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
Prompts []Prompt `json:"prompts,omitempty"` // AI Prompts (default prompts)
PromptPresets map[string][]Prompt `json:"prompt_presets,omitempty"` // Prompt presets organized by mode (e.g., "chat", "task", etc.)
DisableGlobalPrompts bool `json:"disable_global_prompts,omitempty"` // Whether to disable global prompts, default is false
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
Source string `json:"source,omitempty"` // Hook script source code
Locales i18n.Map `json:"locales,omitempty"` // Assistant Locales
Uses *context.Uses `json:"uses,omitempty"` // Assistant-specific wrapper configurations for vision, audio, etc. If not set, use global settings
CreatedAt int64 `json:"created_at"` // Creation timestamp

View file

@ -59,6 +59,7 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
data["public"] = assistant.Public
data["mentionable"] = assistant.Mentionable
data["automated"] = assistant.Automated
data["disable_global_prompts"] = assistant.DisableGlobalPrompts
// Set timestamps
now := time.Now().UnixNano()
@ -102,6 +103,11 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
} else {
data["path"] = nil
}
if assistant.Source != "" {
data["source"] = assistant.Source
} else {
data["source"] = nil
}
// Share field: nullable: false with default "private"
// Apply default if empty
@ -153,10 +159,11 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
// Handle interface{} fields - they should already be in the correct format
jsonFields := map[string]interface{}{
"prompts": assistant.Prompts,
"prompt_presets": assistant.PromptPresets,
"connector_options": assistant.ConnectorOptions,
"kb": assistant.KB,
"mcp": assistant.MCP,
"workflow": assistant.Workflow,
"tools": assistant.Tools,
"placeholder": assistant.Placeholder,
"locales": assistant.Locales,
"uses": assistant.Uses,
@ -218,14 +225,14 @@ func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interfac
data := make(map[string]interface{})
// List of fields that need JSON marshaling
jsonFields := []string{"options", "tags", "prompts", "kb", "mcp", "workflow", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"options", "tags", "prompts", "prompt_presets", "connector_options", "kb", "mcp", "workflow", "placeholder", "locales", "uses"}
jsonFieldSet := make(map[string]bool)
for _, field := range jsonFields {
jsonFieldSet[field] = true
}
// List of nullable string fields
nullableStringFields := []string{"name", "avatar", "description", "path", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableStringFields := []string{"name", "avatar", "description", "path", "source", "__yao_created_by", "__yao_updated_by", "__yao_team_id", "__yao_tenant_id"}
nullableFieldSet := make(map[string]bool)
for _, field := range nullableStringFields {
nullableFieldSet[field] = true
@ -418,7 +425,7 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
// Convert rows to types.AssistantModel slice
assistants := make([]*types.AssistantModel, 0, len(rows))
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
for _, row := range rows {
data := row.ToMap()
@ -456,11 +463,27 @@ func (conv *Xun) GetAssistants(filter types.AssistantFilter, locale ...string) (
}
// GetAssistant retrieves a single assistant by ID
func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.AssistantModel, error) {
row, err := conv.query.New().
func (conv *Xun) GetAssistant(assistantID string, fields []string, locale ...string) (*types.AssistantModel, error) {
qb := conv.query.New().
Table(conv.getAssistantTable()).
Where("assistant_id", assistantID).
First()
Where("assistant_id", assistantID)
// Apply select fields with security validation
// If no fields specified, use default fields
fieldsToSelect := fields
if len(fieldsToSelect) == 0 {
fieldsToSelect = types.AssistantDefaultFields
}
// ValidateAssistantFields will validate fields against whitelist
sanitized := types.ValidateAssistantFields(fieldsToSelect)
selectFields := make([]interface{}, len(sanitized))
for i, field := range sanitized {
selectFields[i] = field
}
qb.Select(selectFields...)
row, err := qb.First()
if err != nil {
return nil, err
}
@ -475,7 +498,7 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
// Parse JSON fields
jsonFields := []string{"tags", "options", "prompts", "workflow", "kb", "mcp", "tools", "placeholder", "locales", "uses"}
jsonFields := []string{"tags", "options", "prompts", "prompt_presets", "connector_options", "workflow", "kb", "mcp", "placeholder", "locales", "uses"}
conv.parseJSONFields(data, jsonFields)
// Convert map to types.AssistantModel
@ -486,6 +509,7 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
Avatar: getString(data, "avatar"),
Connector: getString(data, "connector"),
Path: getString(data, "path"),
Source: getString(data, "source"),
BuiltIn: getBool(data, "built_in"),
Sort: getInt(data, "sort"),
Description: getString(data, "description"),
@ -494,6 +518,7 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
Share: getString(data, "share"),
Mentionable: getBool(data, "mentionable"),
Automated: getBool(data, "automated"),
DisableGlobalPrompts: getBool(data, "disable_global_prompts"),
CreatedAt: getInt64(data, "created_at"),
UpdatedAt: getInt64(data, "updated_at"),
YaoCreatedBy: getString(data, "__yao_created_by"),
@ -529,6 +554,26 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
}
if promptPresets, has := data["prompt_presets"]; has && promptPresets != nil {
raw, err := jsoniter.Marshal(promptPresets)
if err == nil {
var pp map[string][]types.Prompt
if err := jsoniter.Unmarshal(raw, &pp); err == nil {
model.PromptPresets = pp
}
}
}
if connectorOptions, has := data["connector_options"]; has && connectorOptions != nil {
raw, err := jsoniter.Marshal(connectorOptions)
if err == nil {
var co types.ConnectorOptions
if err := jsoniter.Unmarshal(raw, &co); err == nil {
model.ConnectorOptions = &co
}
}
}
if kb, has := data["kb"]; has && kb != nil {
kbConverted, err := types.ToKnowledgeBase(kb)
if err == nil {
@ -550,16 +595,6 @@ func (conv *Xun) GetAssistant(assistantID string, locale ...string) (*types.Assi
}
}
if tools, has := data["tools"]; has && tools != nil {
raw, err := jsoniter.Marshal(tools)
if err == nil {
var tc types.ToolCalls
if err := jsoniter.Unmarshal(raw, &tc); err == nil {
model.Tools = &tc
}
}
}
if placeholder, has := data["placeholder"]; has && placeholder != nil {
raw, err := jsoniter.Marshal(placeholder)
if err == nil {

View file

@ -103,8 +103,8 @@ func TestSaveAssistant(t *testing.T) {
t.Errorf("Expected ID %s, got %s", id, updatedID)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - request all fields to see the update
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve updated assistant: %v", err)
}
@ -183,8 +183,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save complex assistant: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - request all fields for complex data
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve complex assistant: %v", err)
}
@ -237,8 +237,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with MCP: %v", err)
}
// Retrieve and verify MCP configuration
retrieved, err := store.GetAssistant(id)
// Retrieve and verify MCP configuration - mcp is in default fields
retrieved, err := store.GetAssistant(id, []string{})
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -316,8 +316,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant with MCP: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - mcp is in default fields
retrieved, err := store.GetAssistant(id, []string{})
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -353,8 +353,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with uses: %v", err)
}
// Retrieve and verify uses configuration
retrieved, err := store.GetAssistant(id)
// Retrieve and verify uses configuration - uses is NOT in default fields, need to request all
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -396,8 +396,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant without uses: %v", err)
}
// Retrieve and verify uses is nil
retrieved, err := store.GetAssistant(id)
// Retrieve and verify uses is nil - request all fields to check uses
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -425,8 +425,8 @@ func TestSaveAssistant(t *testing.T) {
t.Fatalf("Failed to save assistant with partial uses: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - request all fields for uses
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -451,6 +451,194 @@ func TestSaveAssistant(t *testing.T) {
t.Errorf("Expected fetch to be empty, got '%s'", retrieved.Uses.Fetch)
}
})
t.Run("ConnectorOptions", func(t *testing.T) {
// Test assistant with connector options
assistant := &types.AssistantModel{
Name: "Connector Options Test",
Type: "assistant",
Connector: "openai",
Share: "private",
ConnectorOptions: &types.ConnectorOptions{
Optional: true,
Connectors: []string{"openai", "anthropic"},
Filters: []types.ModelCapability{types.CapVision, types.CapToolCalls},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with connector options: %v", err)
}
// Retrieve and verify - connector_options is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.ConnectorOptions == nil {
t.Fatal("Expected connector options to be set")
}
if !retrieved.ConnectorOptions.Optional {
t.Error("Expected optional to be true")
}
if len(retrieved.ConnectorOptions.Connectors) != 2 {
t.Errorf("Expected 2 connectors, got %d", len(retrieved.ConnectorOptions.Connectors))
}
if len(retrieved.ConnectorOptions.Filters) != 2 {
t.Errorf("Expected 2 filters, got %d", len(retrieved.ConnectorOptions.Filters))
}
if retrieved.ConnectorOptions.Filters[0] != types.CapVision {
t.Errorf("Expected first filter to be vision, got '%s'", retrieved.ConnectorOptions.Filters[0])
}
t.Logf("Successfully saved and retrieved connector options for assistant %s", id)
})
t.Run("PromptPresets", func(t *testing.T) {
// Test assistant with prompt presets
assistant := &types.AssistantModel{
Name: "Prompt Presets Test",
Type: "assistant",
Connector: "openai",
Share: "private",
PromptPresets: map[string][]types.Prompt{
"chat": {
{Role: "system", Content: "You are a friendly chatbot"},
{Role: "user", Content: "Hello!"},
},
"task": {
{Role: "system", Content: "You are a task executor"},
},
},
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with prompt presets: %v", err)
}
// Retrieve and verify - prompt_presets is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.PromptPresets == nil {
t.Fatal("Expected prompt presets to be set")
}
if len(retrieved.PromptPresets) != 2 {
t.Errorf("Expected 2 preset groups, got %d", len(retrieved.PromptPresets))
}
chatPrompts, ok := retrieved.PromptPresets["chat"]
if !ok {
t.Fatal("Expected 'chat' preset to exist")
}
if len(chatPrompts) != 2 {
t.Errorf("Expected 2 chat prompts, got %d", len(chatPrompts))
}
if chatPrompts[0].Role != "system" {
t.Errorf("Expected system role, got '%s'", chatPrompts[0].Role)
}
taskPrompts, ok := retrieved.PromptPresets["task"]
if !ok {
t.Fatal("Expected 'task' preset to exist")
}
if len(taskPrompts) != 1 {
t.Errorf("Expected 1 task prompt, got %d", len(taskPrompts))
}
t.Logf("Successfully saved and retrieved prompt presets for assistant %s", id)
})
t.Run("SourceField", func(t *testing.T) {
// Test assistant with source code
sourceCode := `function onMessage(msg) {
console.log("Received:", msg);
return { status: "ok" };
}`
assistant := &types.AssistantModel{
Name: "Source Field Test",
Type: "assistant",
Connector: "openai",
Share: "private",
Source: sourceCode,
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with source: %v", err)
}
// Retrieve and verify - source is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.Source != sourceCode {
t.Errorf("Expected source code to match, got '%s'", retrieved.Source)
}
t.Logf("Successfully saved and retrieved source code for assistant %s", id)
})
t.Run("AllNewFieldsTogether", func(t *testing.T) {
// Test assistant with all new fields together
assistant := &types.AssistantModel{
Name: "All New Fields Test",
Type: "assistant",
Connector: "openai",
Share: "private",
ConnectorOptions: &types.ConnectorOptions{
Optional: false,
Connectors: []string{"openai"},
Filters: []types.ModelCapability{types.CapVision},
},
PromptPresets: map[string][]types.Prompt{
"default": {
{Role: "system", Content: "Default system prompt"},
},
},
Source: "// Hook code here",
}
id, err := store.SaveAssistant(assistant)
if err != nil {
t.Fatalf("Failed to save assistant with all new fields: %v", err)
}
// Retrieve and verify all new fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
if retrieved.ConnectorOptions == nil {
t.Error("Expected connector options to be set")
}
if retrieved.PromptPresets == nil {
t.Error("Expected prompt presets to be set")
}
if retrieved.Source == "" {
t.Error("Expected source to be set")
}
t.Logf("Successfully saved and retrieved all new fields for assistant %s", id)
})
}
// TestDeleteAssistant tests deleting a single assistant
@ -487,7 +675,7 @@ func TestDeleteAssistant(t *testing.T) {
}
// Verify deletion
_, err = store.GetAssistant(id)
_, err = store.GetAssistant(id, nil)
if err == nil {
t.Error("Expected error when getting deleted assistant")
}
@ -534,8 +722,8 @@ func TestGetAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Retrieve it
retrieved, err := store.GetAssistant(id)
// Retrieve it with default fields (tags are now in default fields)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -562,7 +750,7 @@ func TestGetAssistant(t *testing.T) {
})
t.Run("GetNonExistentAssistant", func(t *testing.T) {
_, err := store.GetAssistant("nonexistent-id")
_, err := store.GetAssistant("nonexistent-id", nil)
if err == nil {
t.Error("Expected error when getting non-existent assistant")
}
@ -1313,8 +1501,8 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Fatalf("Failed to save assistant with permission fields: %v", err)
}
// Retrieve and verify
retrieved, err := store.GetAssistant(id)
// Retrieve and verify - default fields include permission fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1361,8 +1549,8 @@ func TestAssistantPermissionFields(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - default fields include permission fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get updated assistant: %v", err)
}
@ -1393,7 +1581,7 @@ func TestAssistantPermissionFields(t *testing.T) {
}
// Retrieve and verify fields are empty
retrieved, err := store.GetAssistant(id)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1448,7 +1636,7 @@ func TestEmptyStringAsNull(t *testing.T) {
}
// Retrieve and verify empty strings are returned (not stored as empty strings)
retrieved, err := store.GetAssistant(id)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1493,8 +1681,8 @@ func TestEmptyStringAsNull(t *testing.T) {
t.Fatalf("Failed to save assistant: %v", err)
}
// Retrieve and verify values are preserved
retrieved, err := store.GetAssistant(id)
// Retrieve and verify values are preserved - path is sensitive, need full fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant: %v", err)
}
@ -1576,8 +1764,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
},
}
// Test English locale
retrievedEN, err := store.GetAssistant(id, "en")
// Test English locale - request all fields for placeholder
retrievedEN, err := store.GetAssistant(id, types.AssistantFullFields, "en")
if err != nil {
t.Fatalf("Failed to get assistant with EN locale: %v", err)
}
@ -1604,8 +1792,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
t.Errorf("Expected first prompt 'How can I help you?', got '%s'", retrievedEN.Placeholder.Prompts[0])
}
// Test Chinese locale
retrievedZH, err := store.GetAssistant(id, "zh-cn")
// Test Chinese locale - request all fields for placeholder
retrievedZH, err := store.GetAssistant(id, types.AssistantFullFields, "zh-cn")
if err != nil {
t.Fatalf("Failed to get assistant with ZH locale: %v", err)
}
@ -1623,8 +1811,8 @@ func TestGetAssistantWithLocale(t *testing.T) {
t.Errorf("Expected placeholder title '与我聊天', got '%s'", retrievedZH.Placeholder.Title)
}
// Test without locale (should return original {{...}} values)
retrievedNoLocale, err := store.GetAssistant(id)
// Test without locale (should return original {{...}} values) - request all fields for placeholder
retrievedNoLocale, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant without locale: %v", err)
}
@ -2039,8 +2227,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
retrieved, err := store.GetAssistant(id)
// Verify update - need full fields to see tags
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2087,8 +2275,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify all updates
retrieved, err := store.GetAssistant(id)
// Verify all updates - use default fields (includes name, description, sort, mentionable)
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2144,8 +2332,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update JSON fields: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - need full fields for tags, options, prompts
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2193,8 +2381,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update KB and MCP: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - KB and MCP are in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2247,8 +2435,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update MCP: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - MCP is in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2325,8 +2513,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update uses: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - uses is NOT in default fields
retrieved, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2361,8 +2549,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update uses again: %v", err)
}
// Verify second update
retrieved2, err := store.GetAssistant(id)
// Verify second update - uses is NOT in default fields
retrieved2, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2384,8 +2572,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to set uses to nil: %v", err)
}
// Verify uses is nil
retrieved3, err := store.GetAssistant(id)
// Verify uses is nil - uses is NOT in default fields
retrieved3, err := store.GetAssistant(id, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2422,8 +2610,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update permission fields: %v", err)
}
// Verify updates
retrieved, err := store.GetAssistant(id)
// Verify updates - permission fields are in default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2467,8 +2655,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update with empty strings: %v", err)
}
// Verify empty strings are stored as NULL
retrieved, err := store.GetAssistant(id)
// Verify empty strings are stored as NULL - default fields include avatar, description
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2553,8 +2741,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Get original updated_at
original, err := store.GetAssistant(id)
// Get original updated_at - default fields include updated_at
original, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2572,8 +2760,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Get updated assistant
updated, err := store.GetAssistant(id)
// Get updated assistant - default fields include description, updated_at
updated, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve updated assistant: %v", err)
}
@ -2607,8 +2795,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to create assistant: %v", err)
}
// Get original
original, err := store.GetAssistant(id)
// Get original - default fields
original, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2625,8 +2813,8 @@ func TestUpdateAssistant(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify system fields unchanged, but name updated
retrieved, err := store.GetAssistant(id)
// Verify system fields unchanged, but name updated - default fields
retrieved, err := store.GetAssistant(id, nil)
if err != nil {
t.Fatalf("Failed to retrieve assistant: %v", err)
}
@ -2693,9 +2881,9 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Errorf("Expected at least 3 assistants, got %d", len(response.Data))
}
// Step 3: Update one assistant
// Step 3: Update one assistant - need full fields for tags
updatedID := assistantIDs[1]
updatedAssistant, err := store.GetAssistant(updatedID)
updatedAssistant, err := store.GetAssistant(updatedID, types.AssistantFullFields)
if err != nil {
t.Fatalf("Failed to get assistant for update: %v", err)
}
@ -2708,8 +2896,8 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
t.Fatalf("Failed to update assistant: %v", err)
}
// Verify update
verifyAssistant, err := store.GetAssistant(updatedID)
// Verify update - default fields include description
verifyAssistant, err := store.GetAssistant(updatedID, nil)
if err != nil {
t.Fatalf("Failed to verify update: %v", err)
}
@ -2725,7 +2913,7 @@ func TestAssistantCompleteWorkflow(t *testing.T) {
}
// Verify deletion
_, err = store.GetAssistant(assistantIDs[0])
_, err = store.GetAssistant(assistantIDs[0], nil)
if err == nil {
t.Error("Expected error when getting deleted assistant")
}

View file

@ -1,6 +1,7 @@
package types
import (
"github.com/yaoapp/gou/connector/openai"
"github.com/yaoapp/yao/agent/assistant"
store "github.com/yaoapp/yao/agent/store/types"
)
@ -16,13 +17,14 @@ type DSL struct {
// Global External Settings - model capabilities, tools, etc.
// ===============================
Models map[string]assistant.ModelCapabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
Models map[string]openai.Capabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
// Internal
// ===============================
// ID string `json:"-" yaml:"-"` // The id of the instance
Assistant assistant.API `json:"-" yaml:"-"` // The default assistant
Store store.Store `json:"-" yaml:"-"` // The store of the assistant
GlobalPrompts []store.Prompt `json:"-" yaml:"-"` // Global prompts loaded from agent/prompts.yml
}
// Uses the default assistant settings

File diff suppressed because one or more lines are too long

View file

@ -192,6 +192,17 @@ func GetAssistant(c *gin.Context) {
return
}
// Parse select fields (optional - if not provided, returns default fields)
// Query parameter: ?select=field1,field2,field3
var fields []string
if selectParam := c.Query("select"); selectParam != "" {
fields = strings.Split(selectParam, ",")
// Trim whitespace from each field
for i, field := range fields {
fields[i] = strings.TrimSpace(field)
}
}
// Parse locale (optional - if not provided, returns raw data without i18n translation)
// This is useful for form editing scenarios where you need the original values
var assistant *agenttypes.AssistantModel
@ -200,10 +211,10 @@ func GetAssistant(c *gin.Context) {
if loc := c.Query("locale"); loc != "" {
// If locale is specified, get assistant with translation
locale := strings.ToLower(strings.TrimSpace(loc))
assistant, err = agentInstance.Store.GetAssistant(assistantID, locale)
assistant, err = agentInstance.Store.GetAssistant(assistantID, fields, locale)
} else {
// If no locale specified, get raw data without translation
assistant, err = agentInstance.Store.GetAssistant(assistantID)
assistant, err = agentInstance.Store.GetAssistant(assistantID, fields)
}
if err != nil {
log.Error("Failed to get assistant %s: %v", assistantID, err)
@ -521,8 +532,8 @@ func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistantID string
return false, fmt.Errorf("agent store not initialized")
}
// Get assistant from store
assistant, err := agentInstance.Store.GetAssistant(assistantID)
// Get assistant from store - only need default fields for permission check
assistant, err := agentInstance.Store.GetAssistant(assistantID, nil)
if err != nil {
return false, fmt.Errorf("assistant not found: %s", assistantID)
}

View file

@ -126,7 +126,7 @@ func AuthQueryFilter(c *gin.Context, authInfo *types.AuthorizedInfo) func(query.
}
// FilterBuiltInFields filters sensitive fields for built-in assistants in a list
// For built-in assistants, code-level fields (prompts, workflow, tools, kb, mcp, options) should be cleared
// For built-in assistants, code-level fields (prompts, prompt_presets, workflow, kb, mcp, options, source) should be cleared
func FilterBuiltInFields(assistants []*agenttypes.AssistantModel) {
if assistants == nil {
return
@ -138,7 +138,7 @@ func FilterBuiltInFields(assistants []*agenttypes.AssistantModel) {
}
// FilterBuiltInAssistant filters sensitive fields for a single built-in assistant
// For built-in assistants, code-level fields (prompts, workflow, tools, kb, mcp, options) should be cleared
// For built-in assistants, code-level fields (prompts, prompt_presets, workflow, kb, mcp, options, source) should be cleared
// This function can be used for both single assistant and list of assistants
func FilterBuiltInAssistant(assistant *agenttypes.AssistantModel) {
if assistant == nil {
@ -148,10 +148,11 @@ func FilterBuiltInAssistant(assistant *agenttypes.AssistantModel) {
if assistant.BuiltIn {
// Clear code-level sensitive fields for built-in assistants
assistant.Prompts = nil
assistant.PromptPresets = nil
assistant.Workflow = nil
assistant.Tools = nil
assistant.KB = nil
assistant.MCP = nil
assistant.Options = nil
assistant.Source = ""
}
}

View file

@ -130,6 +130,17 @@ func GetModelDetails(c *gin.Context) {
return
}
// For model API, we only need minimal fields: assistant_id, name, connector, created_at, and permission fields
modelFields := []string{
"assistant_id",
"name",
"connector",
"created_at",
"built_in",
"__yao_team_id",
"__yao_created_by",
}
// Parse locale (optional - for assistant name translation)
// Priority: 1. Query parameter "locale", 2. Header "Accept-Language", 3. Metadata
locale := context.GetLocale(c, nil)
@ -138,9 +149,9 @@ func GetModelDetails(c *gin.Context) {
var err error
if locale != "" {
assistant, err = agentInstance.Store.GetAssistant(assistantID, locale)
assistant, err = agentInstance.Store.GetAssistant(assistantID, modelFields, locale)
} else {
assistant, err = agentInstance.Store.GetAssistant(assistantID)
assistant, err = agentInstance.Store.GetAssistant(assistantID, modelFields)
}
if err != nil {

View file

@ -620,62 +620,7 @@ func TestUpdateAssistant(t *testing.T) {
t.Logf("Successfully updated assistant mcp settings: %s", assistantID)
})
t.Run("UpdateAssistantTools", func(t *testing.T) {
// Create a test assistant
assistantID := createTestAssistant("Tools Update Test")
defer func() {
deleteReq, _ := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
resp, _ := http.DefaultClient.Do(deleteReq)
if resp != nil {
resp.Body.Close()
}
}()
// Update tools
updateData := map[string]interface{}{
"tools": []map[string]interface{}{
{
"name": "web_search",
"description": "Search the web for information",
"parameters": map[string]interface{}{
"query": map[string]interface{}{
"type": "string",
"description": "Search query",
"required": true,
},
},
},
{
"name": "calculator",
"description": "Perform calculations",
"parameters": map[string]interface{}{
"expression": map[string]interface{}{
"type": "string",
"description": "Mathematical expression",
"required": true,
},
},
},
},
}
jsonData, err := json.Marshal(updateData)
assert.NoError(t, err)
req, err := http.NewRequest("PUT", serverURL+baseURL+"/agent/assistants/"+assistantID, bytes.NewBuffer(jsonData))
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
req.Header.Set("Content-Type", "application/json")
resp, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.NotNil(t, resp)
defer resp.Body.Close()
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully update tools")
t.Logf("Successfully updated assistant tools: %s", assistantID)
})
// Note: UpdateAssistantTools test removed - tools field is deprecated and replaced by MCP
t.Run("UpdateAssistantWorkflow", func(t *testing.T) {
// Create a test assistant
@ -776,12 +721,7 @@ func TestUpdateAssistant(t *testing.T) {
},
},
},
"tools": []map[string]interface{}{
{
"name": "updated_tool",
"description": "Updated tool description",
},
},
// Note: tools field removed - now handled by MCP
"kb": map[string]interface{}{
"collections": []string{"updated-collection"},
"enabled": true,

View file

@ -57,6 +57,13 @@
"length": 200,
"nullable": false
},
{
"name": "connector_options",
"type": "json",
"label": "Connector Options",
"comment": "Connector selection options: optional flag, available connectors list, and capability filters",
"nullable": true
},
{
"name": "description",
"type": "string",
@ -108,9 +115,24 @@
"name": "prompts",
"type": "json",
"label": "Prompts",
"comment": "Assistant prompts",
"comment": "Assistant default prompts",
"nullable": true
},
{
"name": "prompt_presets",
"type": "json",
"label": "Prompt Presets",
"comment": "Prompt presets organized by mode (e.g., chat, task, etc.)",
"nullable": true
},
{
"name": "disable_global_prompts",
"type": "boolean",
"label": "Disable Global Prompts",
"comment": "Whether to disable global prompts for this assistant",
"default": false,
"index": true
},
{
"name": "workflow",
"type": "json",
@ -133,10 +155,10 @@
"nullable": true
},
{
"name": "tools",
"type": "json",
"label": "Tools",
"comment": "Assistant tools",
"name": "source",
"type": "text",
"label": "Source",
"comment": "Hook script source code",
"nullable": true
},
{