Add GetStorage function and enhance loadSource logic
- Introduced GetStorage function for testing purposes, allowing retrieval of the current storage instance. - Enhanced loadSource function to load scripts from the source field if present, improving assistant initialization. - Updated comments for clarity on TypeScript handling in loadSource, ensuring better understanding of script loading mechanics.
This commit is contained in:
Max
parent
ee44432e69
commit
4b747a5bfa
3 changed files with 958 additions and 4 deletions
|
|
@ -131,6 +131,11 @@ 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]gouOpenAI.Capabilities) {
|
||||
modelCapabilities = capabilities
|
||||
|
|
@ -214,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 {
|
||||
|
|
|
|||
937
agent/assistant/load_store_test.go
Normal file
937
agent/assistant/load_store_test.go
Normal 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)
|
||||
})
|
||||
}
|
||||
|
|
@ -2,6 +2,7 @@ package assistant
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
v8 "github.com/yaoapp/gou/runtime/v8"
|
||||
|
|
@ -9,17 +10,19 @@ import (
|
|||
)
|
||||
|
||||
// loadSource loads hook script from source code string
|
||||
// The source field stores TypeScript code directly
|
||||
// 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
|
||||
}
|
||||
|
||||
// Generate a virtual file path for the script
|
||||
file := fmt.Sprintf("assistants/%s/source.ts", assistantID)
|
||||
// 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.MakeScript([]byte(source), file, 5*time.Second, true)
|
||||
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)
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue