Merge pull request #1286 from trheyi/main
Add UpdateAssistant method for modifying assistant fields
This commit is contained in:
commit
a21906d684
9 changed files with 2395 additions and 23 deletions
|
|
@ -60,6 +60,11 @@ func (m *Mongo) SaveAssistant(assistant *types.AssistantModel) (string, error) {
|
||||||
return assistant.ID, nil
|
return assistant.ID, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// UpdateAssistant updates specific fields of an assistant
|
||||||
|
func (m *Mongo) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// DeleteAssistant deletes an assistant
|
// DeleteAssistant deletes an assistant
|
||||||
func (m *Mongo) DeleteAssistant(assistantID string) error {
|
func (m *Mongo) DeleteAssistant(assistantID string) error {
|
||||||
return nil
|
return nil
|
||||||
|
|
|
||||||
|
|
@ -60,6 +60,11 @@ func (r *Redis) SaveAssistant(assistant *types.AssistantModel) (string, error) {
|
||||||
return assistant.ID, nil
|
return assistant.ID, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// UpdateAssistant updates specific fields of an assistant
|
||||||
|
func (r *Redis) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
// DeleteAssistant deletes an assistant
|
// DeleteAssistant deletes an assistant
|
||||||
func (r *Redis) DeleteAssistant(assistantID string) error {
|
func (r *Redis) DeleteAssistant(assistantID string) error {
|
||||||
return nil
|
return nil
|
||||||
|
|
|
||||||
|
|
@ -66,6 +66,12 @@ type Store interface {
|
||||||
// Returns: Assistant ID and potential error
|
// Returns: Assistant ID and potential error
|
||||||
SaveAssistant(assistant *AssistantModel) (string, error)
|
SaveAssistant(assistant *AssistantModel) (string, error)
|
||||||
|
|
||||||
|
// UpdateAssistant updates assistant fields
|
||||||
|
// assistantID: Assistant ID
|
||||||
|
// updates: Map of fields to update
|
||||||
|
// Returns: Potential error
|
||||||
|
UpdateAssistant(assistantID string, updates map[string]interface{}) error
|
||||||
|
|
||||||
// DeleteAssistant deletes an assistant
|
// DeleteAssistant deletes an assistant
|
||||||
// assistantID: Assistant ID
|
// assistantID: Assistant ID
|
||||||
// Returns: Potential error
|
// Returns: Potential error
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ package xun
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
"math"
|
"math"
|
||||||
|
"time"
|
||||||
|
|
||||||
jsoniter "github.com/json-iterator/go"
|
jsoniter "github.com/json-iterator/go"
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
|
|
@ -172,6 +173,89 @@ func (conv *Xun) SaveAssistant(assistant *types.AssistantModel) (string, error)
|
||||||
return assistant.ID, nil
|
return assistant.ID, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// UpdateAssistant updates specific fields of an assistant
|
||||||
|
func (conv *Xun) UpdateAssistant(assistantID string, updates map[string]interface{}) error {
|
||||||
|
if assistantID == "" {
|
||||||
|
return fmt.Errorf("assistant_id is required")
|
||||||
|
}
|
||||||
|
if len(updates) == 0 {
|
||||||
|
return fmt.Errorf("no fields to update")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if assistant exists
|
||||||
|
exists, err := conv.query.New().
|
||||||
|
Table(conv.getAssistantTable()).
|
||||||
|
Where("assistant_id", assistantID).
|
||||||
|
Exists()
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if !exists {
|
||||||
|
return fmt.Errorf("assistant %s not found", assistantID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Prepare update data
|
||||||
|
data := make(map[string]interface{})
|
||||||
|
|
||||||
|
// List of fields that need JSON marshaling
|
||||||
|
jsonFields := []string{"options", "tags", "prompts", "kb", "mcp", "workflow", "tools", "placeholder", "locales"}
|
||||||
|
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"}
|
||||||
|
nullableFieldSet := make(map[string]bool)
|
||||||
|
for _, field := range nullableStringFields {
|
||||||
|
nullableFieldSet[field] = true
|
||||||
|
}
|
||||||
|
|
||||||
|
// Process each update field
|
||||||
|
for key, value := range updates {
|
||||||
|
// Skip system fields that shouldn't be updated directly
|
||||||
|
if key == "assistant_id" || key == "created_at" {
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
|
||||||
|
// Handle JSON fields
|
||||||
|
if jsonFieldSet[key] {
|
||||||
|
if value != nil {
|
||||||
|
jsonStr, err := jsoniter.MarshalToString(value)
|
||||||
|
if err != nil {
|
||||||
|
return fmt.Errorf("failed to marshal %s: %w", key, err)
|
||||||
|
}
|
||||||
|
data[key] = jsonStr
|
||||||
|
} else {
|
||||||
|
data[key] = nil
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
// Handle regular fields
|
||||||
|
// Convert empty strings to nil for nullable fields
|
||||||
|
if strVal, ok := value.(string); ok && strVal == "" && nullableFieldSet[key] {
|
||||||
|
data[key] = nil
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
data[key] = value
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Always update updated_at timestamp
|
||||||
|
data["updated_at"] = types.ToMySQLTime(time.Now().UnixNano())
|
||||||
|
|
||||||
|
if len(data) == 0 {
|
||||||
|
return fmt.Errorf("no valid fields to update")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Perform update
|
||||||
|
_, err = conv.query.New().
|
||||||
|
Table(conv.getAssistantTable()).
|
||||||
|
Where("assistant_id", assistantID).
|
||||||
|
Update(data)
|
||||||
|
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
// DeleteAssistant deletes an assistant by assistant_id
|
// DeleteAssistant deletes an assistant by assistant_id
|
||||||
func (conv *Xun) DeleteAssistant(assistantID string) error {
|
func (conv *Xun) DeleteAssistant(assistantID string) error {
|
||||||
// Check if assistant exists
|
// Check if assistant exists
|
||||||
|
|
|
||||||
|
|
@ -2,6 +2,7 @@ package xun
|
||||||
|
|
||||||
import (
|
import (
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
|
@ -1745,6 +1746,461 @@ func TestGetAssistantsWithQueryFilter(t *testing.T) {
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// TestUpdateAssistant tests the UpdateAssistant method for incremental updates
|
||||||
|
func TestUpdateAssistant(t *testing.T) {
|
||||||
|
test.Prepare(t, config.Conf)
|
||||||
|
defer test.Clean()
|
||||||
|
|
||||||
|
store, err := NewXun(types.Setting{
|
||||||
|
Connector: "default",
|
||||||
|
})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create store: %v", err)
|
||||||
|
}
|
||||||
|
defer store.Close()
|
||||||
|
|
||||||
|
t.Run("UpdateSingleField", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Original Name",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Description: "Original description",
|
||||||
|
Tags: []string{"original"},
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update only description
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"description": "Updated description",
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify update
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.Description != "Updated description" {
|
||||||
|
t.Errorf("Expected description 'Updated description', got '%s'", retrieved.Description)
|
||||||
|
}
|
||||||
|
// Other fields should remain unchanged
|
||||||
|
if retrieved.Name != "Original Name" {
|
||||||
|
t.Errorf("Expected name 'Original Name', got '%s'", retrieved.Name)
|
||||||
|
}
|
||||||
|
if len(retrieved.Tags) != 1 || retrieved.Tags[0] != "original" {
|
||||||
|
t.Errorf("Expected tags [original], got %v", retrieved.Tags)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateMultipleFields", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Test Assistant",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Description: "Test description",
|
||||||
|
Sort: 100,
|
||||||
|
Mentionable: false,
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update multiple fields
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"name": "Updated Name",
|
||||||
|
"description": "Updated description",
|
||||||
|
"sort": 200,
|
||||||
|
"mentionable": true,
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify all updates
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.Name != "Updated Name" {
|
||||||
|
t.Errorf("Expected name 'Updated Name', got '%s'", retrieved.Name)
|
||||||
|
}
|
||||||
|
if retrieved.Description != "Updated description" {
|
||||||
|
t.Errorf("Expected description 'Updated description', got '%s'", retrieved.Description)
|
||||||
|
}
|
||||||
|
if retrieved.Sort != 200 {
|
||||||
|
t.Errorf("Expected sort 200, got %d", retrieved.Sort)
|
||||||
|
}
|
||||||
|
if !retrieved.Mentionable {
|
||||||
|
t.Error("Expected mentionable to be true")
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateJSONFields", func(t *testing.T) {
|
||||||
|
// Create assistant with complex fields
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "JSON Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Tags: []string{"tag1", "tag2"},
|
||||||
|
Options: map[string]interface{}{"temperature": 0.7},
|
||||||
|
Prompts: []types.Prompt{
|
||||||
|
{Role: "system", Content: "Original system prompt"},
|
||||||
|
},
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update JSON fields
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"tags": []string{"updated", "new-tags"},
|
||||||
|
"options": map[string]interface{}{
|
||||||
|
"temperature": 0.9,
|
||||||
|
"max_tokens": 2000,
|
||||||
|
},
|
||||||
|
"prompts": []types.Prompt{
|
||||||
|
{Role: "system", Content: "Updated system prompt"},
|
||||||
|
{Role: "user", Content: "New user prompt"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update JSON fields: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify updates
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if len(retrieved.Tags) != 2 || retrieved.Tags[0] != "updated" {
|
||||||
|
t.Errorf("Expected tags [updated, new-tags], got %v", retrieved.Tags)
|
||||||
|
}
|
||||||
|
if temp, ok := retrieved.Options["temperature"].(float64); !ok || temp != 0.9 {
|
||||||
|
t.Errorf("Expected temperature 0.9, got %v", retrieved.Options["temperature"])
|
||||||
|
}
|
||||||
|
if len(retrieved.Prompts) != 2 {
|
||||||
|
t.Errorf("Expected 2 prompts, got %d", len(retrieved.Prompts))
|
||||||
|
}
|
||||||
|
if retrieved.Prompts[0].Content != "Updated system prompt" {
|
||||||
|
t.Errorf("Expected updated system prompt, got '%s'", retrieved.Prompts[0].Content)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateKBAndMCP", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "KB MCP Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update KB and MCP
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"kb": map[string]interface{}{
|
||||||
|
"collections": []string{"collection1", "collection2"},
|
||||||
|
},
|
||||||
|
"mcp": map[string]interface{}{
|
||||||
|
"servers": []string{"server1", "server2"},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update KB and MCP: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify updates
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.KB == nil || len(retrieved.KB.Collections) != 2 {
|
||||||
|
t.Errorf("Expected 2 KB collections, got %v", retrieved.KB)
|
||||||
|
}
|
||||||
|
if retrieved.MCP == nil || len(retrieved.MCP.Servers) != 2 {
|
||||||
|
t.Errorf("Expected 2 MCP servers, got %v", retrieved.MCP)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdatePermissionFields", func(t *testing.T) {
|
||||||
|
// Create assistant with permission fields
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Permission Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Share: "private",
|
||||||
|
YaoCreatedBy: "user-1",
|
||||||
|
YaoTeamID: "team-1",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update permission fields
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"__yao_updated_by": "user-2",
|
||||||
|
"__yao_tenant_id": "tenant-1",
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update permission fields: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify updates
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.YaoUpdatedBy != "user-2" {
|
||||||
|
t.Errorf("Expected YaoUpdatedBy 'user-2', got '%s'", retrieved.YaoUpdatedBy)
|
||||||
|
}
|
||||||
|
if retrieved.YaoTenantID != "tenant-1" {
|
||||||
|
t.Errorf("Expected YaoTenantID 'tenant-1', got '%s'", retrieved.YaoTenantID)
|
||||||
|
}
|
||||||
|
// Created by should remain unchanged
|
||||||
|
if retrieved.YaoCreatedBy != "user-1" {
|
||||||
|
t.Errorf("Expected YaoCreatedBy 'user-1', got '%s'", retrieved.YaoCreatedBy)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateWithEmptyStrings", func(t *testing.T) {
|
||||||
|
// Create assistant with values
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Empty String Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Avatar: "https://example.com/avatar.png",
|
||||||
|
Description: "Some description",
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update with empty strings (should become NULL)
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"avatar": "",
|
||||||
|
"description": "",
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update with empty strings: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify empty strings are stored as NULL
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.Avatar != "" {
|
||||||
|
t.Errorf("Expected empty avatar, got '%s'", retrieved.Avatar)
|
||||||
|
}
|
||||||
|
if retrieved.Description != "" {
|
||||||
|
t.Errorf("Expected empty description, got '%s'", retrieved.Description)
|
||||||
|
}
|
||||||
|
// Name should remain unchanged
|
||||||
|
if retrieved.Name != "Empty String Test" {
|
||||||
|
t.Errorf("Expected name 'Empty String Test', got '%s'", retrieved.Name)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateNonExistentAssistant", func(t *testing.T) {
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"name": "Updated Name",
|
||||||
|
}
|
||||||
|
|
||||||
|
err := store.UpdateAssistant("nonexistent-id", updates)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error when updating non-existent assistant")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "not found") {
|
||||||
|
t.Errorf("Expected 'not found' error, got: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateWithEmptyID", func(t *testing.T) {
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"name": "Updated Name",
|
||||||
|
}
|
||||||
|
|
||||||
|
err := store.UpdateAssistant("", updates)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error when updating with empty ID")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "required") {
|
||||||
|
t.Errorf("Expected 'required' error, got: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateWithEmptyUpdates", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Empty Updates Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to update with empty map
|
||||||
|
updates := map[string]interface{}{}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err == nil {
|
||||||
|
t.Error("Expected error when updating with no fields")
|
||||||
|
}
|
||||||
|
if !strings.Contains(err.Error(), "no fields to update") {
|
||||||
|
t.Errorf("Expected 'no fields to update' error, got: %v", err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateTimestampAutomaticallySet", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "Timestamp Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get original updated_at
|
||||||
|
original, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Wait a bit to ensure timestamp difference
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
|
||||||
|
// Update assistant
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"description": "Updated to test timestamp",
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get updated assistant
|
||||||
|
updated, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve updated assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify description was updated (main test objective)
|
||||||
|
if updated.Description != "Updated to test timestamp" {
|
||||||
|
t.Errorf("Expected description 'Updated to test timestamp', got '%s'", updated.Description)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Only check timestamp if both are set (some stores may not return timestamps)
|
||||||
|
if original.UpdatedAt > 0 && updated.UpdatedAt > 0 {
|
||||||
|
if updated.UpdatedAt <= original.UpdatedAt {
|
||||||
|
t.Errorf("Expected updated_at to increase, original=%d, updated=%d", original.UpdatedAt, updated.UpdatedAt)
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
t.Logf("Skipping timestamp comparison (original=%d, updated=%d)", original.UpdatedAt, updated.UpdatedAt)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("UpdateSkipsSystemFields", func(t *testing.T) {
|
||||||
|
// Create assistant
|
||||||
|
assistant := &types.AssistantModel{
|
||||||
|
Name: "System Fields Test",
|
||||||
|
Type: "assistant",
|
||||||
|
Connector: "openai",
|
||||||
|
Share: "private",
|
||||||
|
}
|
||||||
|
|
||||||
|
id, err := store.SaveAssistant(assistant)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to create assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get original
|
||||||
|
original, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Try to update system fields (should be ignored)
|
||||||
|
updates := map[string]interface{}{
|
||||||
|
"assistant_id": "new-id-123", // Should be ignored
|
||||||
|
"created_at": int64(123456789), // Should be ignored
|
||||||
|
"name": "Valid Update", // Should be applied
|
||||||
|
}
|
||||||
|
|
||||||
|
err = store.UpdateAssistant(id, updates)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to update assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify system fields unchanged, but name updated
|
||||||
|
retrieved, err := store.GetAssistant(id)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatalf("Failed to retrieve assistant: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if retrieved.ID != id {
|
||||||
|
t.Errorf("Expected ID to remain %s, got %s", id, retrieved.ID)
|
||||||
|
}
|
||||||
|
if retrieved.CreatedAt != original.CreatedAt {
|
||||||
|
t.Errorf("Expected created_at to remain unchanged")
|
||||||
|
}
|
||||||
|
if retrieved.Name != "Valid Update" {
|
||||||
|
t.Errorf("Expected name 'Valid Update', got '%s'", retrieved.Name)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
// TestAssistantCompleteWorkflow tests a complete workflow
|
// TestAssistantCompleteWorkflow tests a complete workflow
|
||||||
func TestAssistantCompleteWorkflow(t *testing.T) {
|
func TestAssistantCompleteWorkflow(t *testing.T) {
|
||||||
test.Prepare(t, config.Conf)
|
test.Prepare(t, config.Conf)
|
||||||
|
|
|
||||||
|
|
@ -18,9 +18,10 @@ func Attach(group *gin.RouterGroup, oauth types.OAuth) {
|
||||||
|
|
||||||
// Assistant CRUD - Standard REST endpoints
|
// Assistant CRUD - Standard REST endpoints
|
||||||
group.GET("/assistants", ListAssistants) // GET /assistants - List assistants
|
group.GET("/assistants", ListAssistants) // GET /assistants - List assistants
|
||||||
group.POST("/assistants", n.HandleAssistantSave) // POST /assistants - Create/Update assistant
|
group.POST("/assistants", CreateAssistant) // POST /assistants - Create assistant
|
||||||
group.GET("/assistants/tags", ListAssistantTags) // GET /assistants/tags - Get all assistant tags with permission filtering
|
group.GET("/assistants/tags", ListAssistantTags) // GET /assistants/tags - Get all assistant tags with permission filtering
|
||||||
group.GET("/assistants/:id", GetAssistant) // GET /assistants/:id - Get assistant details with permission verification
|
group.GET("/assistants/:id", GetAssistant) // GET /assistants/:id - Get assistant details with permission verification
|
||||||
|
group.PUT("/assistants/:id", UpdateAssistant) // PUT /assistants/:id - Update assistant
|
||||||
group.DELETE("/assistants/:id", n.HandleAssistantDelete) // DELETE /assistants/:id - Delete assistant
|
group.DELETE("/assistants/:id", n.HandleAssistantDelete) // DELETE /assistants/:id - Delete assistant
|
||||||
|
|
||||||
// Assistant Actions
|
// Assistant Actions
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,15 @@
|
||||||
package agent
|
package agent
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"strconv"
|
"strconv"
|
||||||
"strings"
|
"strings"
|
||||||
|
|
||||||
"github.com/gin-gonic/gin"
|
"github.com/gin-gonic/gin"
|
||||||
"github.com/yaoapp/kun/log"
|
"github.com/yaoapp/kun/log"
|
||||||
"github.com/yaoapp/yao/agent"
|
"github.com/yaoapp/yao/agent"
|
||||||
|
"github.com/yaoapp/yao/agent/assistant"
|
||||||
agenttypes "github.com/yaoapp/yao/agent/store/types"
|
agenttypes "github.com/yaoapp/yao/agent/store/types"
|
||||||
"github.com/yaoapp/yao/openapi/oauth/authorized"
|
"github.com/yaoapp/yao/openapi/oauth/authorized"
|
||||||
"github.com/yaoapp/yao/openapi/oauth/types"
|
"github.com/yaoapp/yao/openapi/oauth/types"
|
||||||
|
|
@ -213,8 +216,8 @@ func GetAssistant(c *gin.Context) {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// Check permission
|
// Check read permission
|
||||||
hasPermission, err := checkAssistantPermission(authInfo, assistant)
|
hasPermission, err := checkAssistantPermission(authInfo, assistantID, true)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
log.Error("Failed to check permission for assistant %s: %v", assistantID, err)
|
log.Error("Failed to check permission for assistant %s: %v", assistantID, err)
|
||||||
errorResp := &response.ErrorResponse{
|
errorResp := &response.ErrorResponse{
|
||||||
|
|
@ -313,9 +316,200 @@ func ListAssistantTags(c *gin.Context) {
|
||||||
response.RespondWithSuccess(c, response.StatusOK, tags)
|
response.RespondWithSuccess(c, response.StatusOK, tags)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// CreateAssistant creates a new assistant
|
||||||
|
func CreateAssistant(c *gin.Context) {
|
||||||
|
// Get authorized information
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
|
||||||
|
// Get Agent instance from global variable
|
||||||
|
agentInstance := agent.GetAgent()
|
||||||
|
if agentInstance == nil || agentInstance.Store == nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrServerError.Code,
|
||||||
|
ErrorDescription: "Agent store not initialized",
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse request body
|
||||||
|
var assistantData map[string]interface{}
|
||||||
|
if err := c.ShouldBindJSON(&assistantData); err != nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrInvalidRequest.Code,
|
||||||
|
ErrorDescription: "Invalid request body: " + err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert to AssistantModel
|
||||||
|
model, err := agenttypes.ToAssistantModel(assistantData)
|
||||||
|
if err != nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrInvalidRequest.Code,
|
||||||
|
ErrorDescription: "Invalid assistant data: " + err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Attach create scope to the assistant data
|
||||||
|
if authInfo != nil {
|
||||||
|
scope := authInfo.AccessScope()
|
||||||
|
model.YaoCreatedBy = scope.CreatedBy
|
||||||
|
model.YaoUpdatedBy = scope.UpdatedBy
|
||||||
|
model.YaoTeamID = scope.TeamID
|
||||||
|
model.YaoTenantID = scope.TenantID
|
||||||
|
}
|
||||||
|
|
||||||
|
// Save assistant using Store
|
||||||
|
id, err := agentInstance.Store.SaveAssistant(model)
|
||||||
|
if err != nil {
|
||||||
|
log.Error("Failed to create assistant: %v", err)
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrServerError.Code,
|
||||||
|
ErrorDescription: "Failed to create assistant: " + err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update the assistant map with the returned ID
|
||||||
|
assistantData["assistant_id"] = id
|
||||||
|
|
||||||
|
// Clear cache and reload assistant to make it effective
|
||||||
|
cache := assistant.GetCache()
|
||||||
|
if cache != nil {
|
||||||
|
cache.Remove(id)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reload the assistant to ensure it's available in cache with updated data
|
||||||
|
_, err = assistant.Get(id)
|
||||||
|
if err != nil {
|
||||||
|
// Just log the error, don't fail the request
|
||||||
|
log.Error("Error reloading assistant %s: %v", id, err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Return success response
|
||||||
|
response.RespondWithSuccess(c, response.StatusOK, assistantData)
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateAssistant updates an existing assistant
|
||||||
|
func UpdateAssistant(c *gin.Context) {
|
||||||
|
// Get authorized information
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
|
||||||
|
// Get Agent instance from global variable
|
||||||
|
agentInstance := agent.GetAgent()
|
||||||
|
if agentInstance == nil || agentInstance.Store == nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrServerError.Code,
|
||||||
|
ErrorDescription: "Agent store not initialized",
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get assistant ID from URL parameter
|
||||||
|
assistantID := c.Param("id")
|
||||||
|
if assistantID == "" {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrInvalidRequest.Code,
|
||||||
|
ErrorDescription: "assistant_id is required",
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check update permission
|
||||||
|
hasPermission, err := checkAssistantPermission(authInfo, assistantID, false)
|
||||||
|
if err != nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrServerError.Code,
|
||||||
|
ErrorDescription: err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusForbidden, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// 403 Forbidden
|
||||||
|
if !hasPermission {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrAccessDenied.Code,
|
||||||
|
ErrorDescription: "Forbidden: No permission to update this assistant",
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusForbidden, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Parse request body with update data
|
||||||
|
var updateData map[string]interface{}
|
||||||
|
if err := c.ShouldBindJSON(&updateData); err != nil {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrInvalidRequest.Code,
|
||||||
|
ErrorDescription: "Invalid request body: " + err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusBadRequest, errorResp)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Add update metadata
|
||||||
|
if authInfo != nil {
|
||||||
|
scope := authInfo.AccessScope()
|
||||||
|
updateData["__yao_updated_by"] = scope.UpdatedBy
|
||||||
|
}
|
||||||
|
|
||||||
|
// Update assistant using Store
|
||||||
|
err = agentInstance.Store.UpdateAssistant(assistantID, updateData)
|
||||||
|
if err != nil {
|
||||||
|
log.Error("Failed to update assistant %s: %v", assistantID, err)
|
||||||
|
// Check if it's a "not found" error
|
||||||
|
if strings.Contains(err.Error(), "not found") {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrInvalidRequest.Code,
|
||||||
|
ErrorDescription: "Assistant not found: " + assistantID,
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusNotFound, errorResp)
|
||||||
|
} else {
|
||||||
|
errorResp := &response.ErrorResponse{
|
||||||
|
Code: response.ErrServerError.Code,
|
||||||
|
ErrorDescription: "Failed to update assistant: " + err.Error(),
|
||||||
|
}
|
||||||
|
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clear cache and reload assistant to make it effective
|
||||||
|
cache := assistant.GetCache()
|
||||||
|
if cache != nil {
|
||||||
|
cache.Remove(assistantID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Reload the assistant to ensure it's available in cache with updated data
|
||||||
|
updatedAssistant, err := assistant.Get(assistantID)
|
||||||
|
if err != nil {
|
||||||
|
// Just log the error, don't fail the request
|
||||||
|
log.Error("Error reloading assistant %s: %v", assistantID, err)
|
||||||
|
// Return simple success response
|
||||||
|
response.RespondWithSuccess(c, response.StatusOK, map[string]interface{}{"assistant_id": assistantID})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert updated assistant to map for response
|
||||||
|
responseData, _ := json.Marshal(updatedAssistant)
|
||||||
|
var responseMap map[string]interface{}
|
||||||
|
json.Unmarshal(responseData, &responseMap)
|
||||||
|
|
||||||
|
// Return success response with updated assistant data
|
||||||
|
response.RespondWithSuccess(c, response.StatusOK, responseMap)
|
||||||
|
}
|
||||||
|
|
||||||
// checkAssistantPermission checks if the user has permission to access the assistant
|
// checkAssistantPermission checks if the user has permission to access the assistant
|
||||||
// Similar logic to checkFilePermission in openapi/file/file.go
|
// Similar logic to checkCollectionPermission in openapi/kb/collection.go
|
||||||
func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistant *agenttypes.AssistantModel) (bool, error) {
|
// readable: true for read permission, false for write permission
|
||||||
|
func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistantID string, readable ...bool) (bool, error) {
|
||||||
// No auth info, allow access
|
// No auth info, allow access
|
||||||
if authInfo == nil {
|
if authInfo == nil {
|
||||||
return true, nil
|
return true, nil
|
||||||
|
|
@ -326,39 +520,53 @@ func checkAssistantPermission(authInfo *types.AuthorizedInfo, assistant *agentty
|
||||||
return true, nil
|
return true, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// If assistant is public, allow access
|
// Get Agent instance
|
||||||
|
agentInstance := agent.GetAgent()
|
||||||
|
if agentInstance == nil || agentInstance.Store == nil {
|
||||||
|
return false, fmt.Errorf("agent store not initialized")
|
||||||
|
}
|
||||||
|
|
||||||
|
// Get assistant from store
|
||||||
|
assistant, err := agentInstance.Store.GetAssistant(assistantID)
|
||||||
|
if err != nil {
|
||||||
|
return false, fmt.Errorf("assistant not found: %s", assistantID)
|
||||||
|
}
|
||||||
|
|
||||||
|
// If readable mode, check if the assistant is accessible for reading
|
||||||
|
if len(readable) > 0 && readable[0] {
|
||||||
|
// If assistant is public, allow read access
|
||||||
if assistant.Public {
|
if assistant.Public {
|
||||||
return true, nil
|
return true, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// Team only permission validation for read
|
||||||
|
if assistant.Share == "team" && authInfo.Constraints.TeamOnly {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if user is the creator - always allow creator to access their own assistant
|
||||||
|
if assistant.YaoCreatedBy != "" && assistant.YaoCreatedBy == authInfo.UserID {
|
||||||
|
return true, nil
|
||||||
|
}
|
||||||
|
|
||||||
// Combined Team and Owner permission validation
|
// Combined Team and Owner permission validation
|
||||||
if authInfo.Constraints.TeamOnly && authInfo.Constraints.OwnerOnly {
|
if authInfo.Constraints.TeamOnly && authInfo.Constraints.OwnerOnly {
|
||||||
if assistant.YaoCreatedBy == authInfo.UserID && assistant.YaoTeamID == authInfo.TeamID {
|
if assistant.YaoTeamID != "" && assistant.YaoTeamID == authInfo.TeamID {
|
||||||
return true, nil
|
return true, nil
|
||||||
}
|
}
|
||||||
return false, nil
|
return false, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// Team only permission validation
|
// Team only permission validation
|
||||||
if authInfo.Constraints.TeamOnly {
|
if authInfo.Constraints.TeamOnly && assistant.YaoTeamID != "" && assistant.YaoTeamID == authInfo.TeamID {
|
||||||
// Check if assistant belongs to the same team
|
|
||||||
if assistant.YaoTeamID == authInfo.TeamID {
|
|
||||||
// Allow if user created it or if it's shared with team
|
|
||||||
if assistant.YaoCreatedBy == authInfo.UserID || assistant.Share == "team" {
|
|
||||||
return true, nil
|
return true, nil
|
||||||
}
|
}
|
||||||
}
|
|
||||||
return false, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// Owner only permission validation
|
// Owner only permission validation (already handled above by creator check)
|
||||||
if authInfo.Constraints.OwnerOnly {
|
if authInfo.Constraints.OwnerOnly {
|
||||||
// Check if user created the assistant and team_id is empty (not a team resource)
|
|
||||||
if assistant.YaoCreatedBy == authInfo.UserID && assistant.YaoTeamID == "" {
|
|
||||||
return true, nil
|
|
||||||
}
|
|
||||||
return false, nil
|
return false, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
return false, nil
|
return false, fmt.Errorf("no permission to access assistant: %s", assistantID)
|
||||||
}
|
}
|
||||||
|
|
|
||||||
561
openapi/tests/agent/assistant_create_test.go
Normal file
561
openapi/tests/agent/assistant_create_test.go
Normal file
|
|
@ -0,0 +1,561 @@
|
||||||
|
package openapi_test
|
||||||
|
|
||||||
|
import (
|
||||||
|
"bytes"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"net/http"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/yaoapp/yao/openapi"
|
||||||
|
"github.com/yaoapp/yao/openapi/tests/testutils"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TestCreateAssistant tests the create assistant endpoint
|
||||||
|
func TestCreateAssistant(t *testing.T) {
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
|
||||||
|
// Get base URL from server config
|
||||||
|
baseURL := ""
|
||||||
|
if openapi.Server != nil && openapi.Server.Config != nil {
|
||||||
|
baseURL = openapi.Server.Config.BaseURL
|
||||||
|
}
|
||||||
|
|
||||||
|
// Register test client and get token
|
||||||
|
client := testutils.RegisterTestClient(t, "Agent Create Test Client", []string{"https://localhost/callback"})
|
||||||
|
defer testutils.CleanupTestClient(t, client.ClientID)
|
||||||
|
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
|
||||||
|
|
||||||
|
t.Run("CreateAssistantSuccess", func(t *testing.T) {
|
||||||
|
// Create a new assistant
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
"description": "A test assistant created by automated tests",
|
||||||
|
"tags": []string{"test", "automation"},
|
||||||
|
"public": false,
|
||||||
|
"share": "private",
|
||||||
|
"mentionable": true,
|
||||||
|
"automated": false,
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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()
|
||||||
|
|
||||||
|
// Expect successful response
|
||||||
|
assert.Equal(t, http.StatusOK, resp.StatusCode, "Should successfully create assistant")
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
// Verify response contains assistant_id
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID, "Response should have assistant_id")
|
||||||
|
assert.NotEmpty(t, assistantID, "Assistant ID should not be empty")
|
||||||
|
|
||||||
|
t.Logf("Successfully created assistant with ID: %s", assistantID)
|
||||||
|
|
||||||
|
// Clean up: delete the created assistant
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
t.Logf("Cleaned up test assistant: %s", assistantID)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantWithMinimalFields", func(t *testing.T) {
|
||||||
|
// Create assistant with only required fields
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Minimal Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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 create assistant with minimal fields")
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID, "Response should have assistant_id")
|
||||||
|
t.Logf("Created minimal assistant with ID: %s", assistantID)
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantWithAllFields", func(t *testing.T) {
|
||||||
|
// Create assistant with all possible fields
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Complete Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
"description": "A complete test assistant with all fields",
|
||||||
|
"avatar": "https://example.com/avatar.png",
|
||||||
|
"tags": []string{"test", "complete", "all-fields"},
|
||||||
|
"public": false,
|
||||||
|
"share": "private",
|
||||||
|
"mentionable": true,
|
||||||
|
"automated": false,
|
||||||
|
"readonly": false,
|
||||||
|
"built_in": false,
|
||||||
|
"sort": 100,
|
||||||
|
"placeholder": map[string]interface{}{
|
||||||
|
"en-us": "Ask me anything...",
|
||||||
|
"zh-cn": "有什么可以帮您的...",
|
||||||
|
},
|
||||||
|
"prompts": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"role": "system",
|
||||||
|
"content": "You are a helpful assistant.",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"options": map[string]interface{}{
|
||||||
|
"temperature": 0.7,
|
||||||
|
"max_tokens": 2000,
|
||||||
|
},
|
||||||
|
"workflow": map[string]interface{}{
|
||||||
|
"steps": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"name": "step1",
|
||||||
|
"action": "process",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"tools": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"name": "search",
|
||||||
|
"description": "Search the web",
|
||||||
|
"parameters": map[string]interface{}{
|
||||||
|
"query": "string",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"kb": map[string]interface{}{
|
||||||
|
"collections": []string{"collection1", "collection2"},
|
||||||
|
"enabled": true,
|
||||||
|
},
|
||||||
|
"mcp": map[string]interface{}{
|
||||||
|
"servers": []map[string]interface{}{
|
||||||
|
{
|
||||||
|
"name": "server1",
|
||||||
|
"url": "http://localhost:3000",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
},
|
||||||
|
"locales": map[string]interface{}{
|
||||||
|
"zh-cn": map[string]interface{}{
|
||||||
|
"name": "完整测试助手",
|
||||||
|
"description": "包含所有字段的完整测试助手",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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 create assistant with all fields")
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID, "Response should have assistant_id")
|
||||||
|
t.Logf("Created complete assistant with ID: %s", assistantID)
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantMissingRequiredFields", func(t *testing.T) {
|
||||||
|
// Test with missing required fields
|
||||||
|
testCases := []struct {
|
||||||
|
name string
|
||||||
|
data map[string]interface{}
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "MissingName",
|
||||||
|
data: map[string]interface{}{
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "MissingType",
|
||||||
|
data: map[string]interface{}{
|
||||||
|
"name": "Test Assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "MissingConnector",
|
||||||
|
data: map[string]interface{}{
|
||||||
|
"name": "Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tc := range testCases {
|
||||||
|
t.Run(tc.name, func(t *testing.T) {
|
||||||
|
jsonData, err := json.Marshal(tc.data)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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()
|
||||||
|
|
||||||
|
// Should return 400 Bad Request or 500 Internal Server Error
|
||||||
|
assert.True(t, resp.StatusCode == http.StatusBadRequest || resp.StatusCode == http.StatusInternalServerError,
|
||||||
|
"Should return error for missing required fields")
|
||||||
|
|
||||||
|
var errorResponse map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&errorResponse)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
t.Logf("Correctly rejected request with %s (status: %d)", tc.name, resp.StatusCode)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantInvalidJSON", func(t *testing.T) {
|
||||||
|
// Test with invalid JSON
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", bytes.NewBufferString("{invalid json}"))
|
||||||
|
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.StatusBadRequest, resp.StatusCode, "Should return 400 for invalid JSON")
|
||||||
|
|
||||||
|
var errorResponse map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&errorResponse)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.Contains(t, errorResponse, "error", "Error response should have 'error' field")
|
||||||
|
|
||||||
|
t.Logf("Correctly rejected invalid JSON")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantUnauthorized", func(t *testing.T) {
|
||||||
|
// Test without authentication
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", bytes.NewBuffer(jsonData))
|
||||||
|
assert.NoError(t, err)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
// No Authorization header
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.NotNil(t, resp)
|
||||||
|
defer resp.Body.Close()
|
||||||
|
|
||||||
|
assert.Equal(t, http.StatusUnauthorized, resp.StatusCode, "Should return 401 without authentication")
|
||||||
|
t.Logf("Correctly rejected unauthorized request")
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantWithLocales", func(t *testing.T) {
|
||||||
|
// Create assistant with localized content
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Multilingual Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
"locales": map[string]interface{}{
|
||||||
|
"zh-cn": map[string]interface{}{
|
||||||
|
"name": "多语言测试助手",
|
||||||
|
"description": "这是一个多语言测试助手",
|
||||||
|
},
|
||||||
|
"ja-jp": map[string]interface{}{
|
||||||
|
"name": "多言語テストアシスタント",
|
||||||
|
"description": "これは多言語テストアシスタントです",
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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 create assistant with locales")
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID, "Response should have assistant_id")
|
||||||
|
t.Logf("Created multilingual assistant with ID: %s", assistantID)
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantWithTeamScope", func(t *testing.T) {
|
||||||
|
// Create assistant - should automatically attach team scope from auth
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Team Scoped Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
"share": "team",
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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 create team-scoped assistant")
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID, "Response should have assistant_id")
|
||||||
|
t.Logf("Created team-scoped assistant with ID: %s", assistantID)
|
||||||
|
|
||||||
|
// Verify the assistant was created with proper scope
|
||||||
|
// Get the assistant to check if __yao_created_by and __yao_team_id were set
|
||||||
|
getReq, err := http.NewRequest("GET", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
getReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
|
||||||
|
getResp, err := http.DefaultClient.Do(getReq)
|
||||||
|
if err == nil {
|
||||||
|
defer getResp.Body.Close()
|
||||||
|
if getResp.StatusCode == http.StatusOK {
|
||||||
|
var assistant map[string]interface{}
|
||||||
|
json.NewDecoder(getResp.Body).Decode(&assistant)
|
||||||
|
t.Logf("Assistant scope fields: __yao_created_by=%v, __yao_team_id=%v",
|
||||||
|
assistant["__yao_created_by"], assistant["__yao_team_id"])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
t.Run("CreateAssistantVerifyCacheReload", func(t *testing.T) {
|
||||||
|
// Create assistant and verify it's immediately available
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Cache Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, err := json.Marshal(assistantData)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
req, err := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", 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)
|
||||||
|
|
||||||
|
var response map[string]interface{}
|
||||||
|
err = json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
|
||||||
|
assistantID, hasID := response["assistant_id"].(string)
|
||||||
|
assert.True(t, hasID)
|
||||||
|
|
||||||
|
// Immediately try to get the assistant - should be available (cache reloaded)
|
||||||
|
getReq, err := http.NewRequest("GET", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
getReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
|
||||||
|
getResp, err := http.DefaultClient.Do(getReq)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
assert.NotNil(t, getResp)
|
||||||
|
defer getResp.Body.Close()
|
||||||
|
|
||||||
|
// Should be immediately available thanks to cache reload
|
||||||
|
if getResp.StatusCode == http.StatusOK {
|
||||||
|
t.Logf("Assistant immediately available after creation (cache reloaded successfully)")
|
||||||
|
} else {
|
||||||
|
t.Logf("Assistant not immediately available (status: %d) - cache reload may have failed", getResp.StatusCode)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Clean up
|
||||||
|
deleteReq, err := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+assistantID, nil)
|
||||||
|
assert.NoError(t, err)
|
||||||
|
deleteReq.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
deleteResp, err := http.DefaultClient.Do(deleteReq)
|
||||||
|
if err == nil {
|
||||||
|
defer deleteResp.Body.Close()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// BenchmarkCreateAssistant benchmarks the create assistant endpoint
|
||||||
|
func BenchmarkCreateAssistant(b *testing.B) {
|
||||||
|
// Convert testing.B to testing.T for Prepare/Clean
|
||||||
|
t := &testing.T{}
|
||||||
|
serverURL := testutils.Prepare(t)
|
||||||
|
defer testutils.Clean()
|
||||||
|
|
||||||
|
baseURL := ""
|
||||||
|
if openapi.Server != nil && openapi.Server.Config != nil {
|
||||||
|
baseURL = openapi.Server.Config.BaseURL
|
||||||
|
}
|
||||||
|
|
||||||
|
client := testutils.RegisterTestClient(t, "Agent Create Benchmark Client", []string{"https://localhost/callback"})
|
||||||
|
defer testutils.CleanupTestClient(t, client.ClientID)
|
||||||
|
tokenInfo := testutils.ObtainAccessToken(t, serverURL, client.ClientID, client.ClientSecret, "https://localhost/callback", "openid profile")
|
||||||
|
|
||||||
|
assistantData := map[string]interface{}{
|
||||||
|
"name": "Benchmark Test Assistant",
|
||||||
|
"type": "assistant",
|
||||||
|
"connector": "openai",
|
||||||
|
}
|
||||||
|
|
||||||
|
jsonData, _ := json.Marshal(assistantData)
|
||||||
|
|
||||||
|
// Track created assistants for cleanup
|
||||||
|
createdIDs := make([]string, 0, b.N)
|
||||||
|
|
||||||
|
// Reset timer after setup
|
||||||
|
b.ResetTimer()
|
||||||
|
|
||||||
|
// Run benchmark
|
||||||
|
for i := 0; i < b.N; i++ {
|
||||||
|
assistantData["name"] = fmt.Sprintf("Benchmark Test Assistant %d", i)
|
||||||
|
jsonData, _ = json.Marshal(assistantData)
|
||||||
|
|
||||||
|
req, _ := http.NewRequest("POST", serverURL+baseURL+"/agent/assistants", bytes.NewBuffer(jsonData))
|
||||||
|
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
req.Header.Set("Content-Type", "application/json")
|
||||||
|
|
||||||
|
resp, err := http.DefaultClient.Do(req)
|
||||||
|
if err != nil {
|
||||||
|
b.Fatalf("Request failed: %v", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if resp.StatusCode == http.StatusOK {
|
||||||
|
var response map[string]interface{}
|
||||||
|
json.NewDecoder(resp.Body).Decode(&response)
|
||||||
|
if id, ok := response["assistant_id"].(string); ok {
|
||||||
|
createdIDs = append(createdIDs, id)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
resp.Body.Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Cleanup created assistants
|
||||||
|
b.StopTimer()
|
||||||
|
for _, id := range createdIDs {
|
||||||
|
req, _ := http.NewRequest("DELETE", serverURL+baseURL+"/agent/assistants/"+id, nil)
|
||||||
|
req.Header.Set("Authorization", "Bearer "+tokenInfo.AccessToken)
|
||||||
|
resp, _ := http.DefaultClient.Do(req)
|
||||||
|
if resp != nil {
|
||||||
|
resp.Body.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
1046
openapi/tests/agent/assistant_update_test.go
Normal file
1046
openapi/tests/agent/assistant_update_test.go
Normal file
File diff suppressed because it is too large
Load diff
Loading…
Add table
Reference in a new issue