Merge pull request #1371 from trheyi/main

Enhance Assistant stream functionality and remove history handling
This commit is contained in:
Max 2025-12-08 19:17:02 +08:00 committed by GitHub
commit 1659f59d72
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 6018 additions and 394 deletions

View file

@ -22,7 +22,14 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
log.Trace("[AGENT] Stream started: assistant=%s, contextID=%s", ast.ID, ctx.ID) log.Trace("[AGENT] Stream started: assistant=%s, contextID=%s", ast.ID, ctx.ID)
defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID) defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID)
// Validate user permissions
var err error var err error
err = ast.checkPermissions(ctx)
if err != nil {
return nil, err
}
// Start stream time
streamStartTime := time.Now() streamStartTime := time.Now()
// Set up interrupt handler if interrupt controller is available // Set up interrupt handler if interrupt controller is available
@ -65,6 +72,10 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
// Now ctx.Capabilities is set, so output adapters can use it // Now ctx.Capabilities is set, so output adapters can use it
ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime) ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime)
// Initialize chat, prepare kb collection (optional) etc.
// Use async version to not block the main flow
ast.InitializeConversationAsync(ctx, opts)
// Initialize agent trace node // Initialize agent trace node
agentNode := ast.initAgentTraceNode(ctx, inputMessages) agentNode := ast.initAgentTraceNode(ctx, inputMessages)

218
agent/assistant/chat.go Normal file
View file

@ -0,0 +1,218 @@
package assistant
import (
"fmt"
"strings"
"sync"
agentcontext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n"
"github.com/yaoapp/yao/kb"
kbapi "github.com/yaoapp/yao/kb/api"
"github.com/yaoapp/yao/trace/types"
)
// kbCollectionCreating tracks collections currently being created to avoid duplicate creation
var kbCollectionCreating sync.Map
// WithHistory merges the input messages with chat history and traces it
// This method can be overridden or extended to implement actual history loading
func (ast *Assistant) WithHistory(ctx *agentcontext.Context, input []agentcontext.Message, agentNode types.Node, options ...*agentcontext.Options) ([]agentcontext.Message, error) {
// TODO: Implement actual history loading logic here
// For now, just simulate a check and return the input messages as is
// Simulate error check (this is where actual history loading would happen)
// if some_condition {
// ast.traceAgentFail(agentNode, err)
// return nil, err
// }
fullMessages := input
// Log the chat history
ast.traceAgentHistory(ctx, agentNode, fullMessages)
return fullMessages, nil
}
// InitializeConversation prepares KB collection for the conversation (synchronous)
func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error {
var opts *agentcontext.Options
if len(options) > 0 && options[0] != nil {
opts = options[0]
} else {
opts = &agentcontext.Options{}
}
// SKIP: History (for internal calls like title/prompt etc.)
if opts.Skip != nil && opts.Skip.History {
return nil
}
// Check if authorized info is available
if ctx.Authorized == nil {
fmt.Printf(">>> Warning: no authorized info, skipping KB collection preparation\n")
return nil
}
// Prepare kb collection
err := ast.prepareKBCollection(ctx, opts)
if err != nil {
// Log but don't fail the chat
fmt.Printf(">>> Warning: failed to prepare KB collection: %v\n", err)
}
return nil
}
// InitializeConversationAsync prepares KB collection asynchronously
func (ast *Assistant) InitializeConversationAsync(ctx *agentcontext.Context, options ...*agentcontext.Options) {
go ast.InitializeConversation(ctx, options...)
}
// prepareKBCollection prepares kb collection (internal method)
func (ast *Assistant) prepareKBCollection(ctx *agentcontext.Context, opts *agentcontext.Options) error {
// Get global KB setting
kbSetting := GetGlobalKBSetting()
if kbSetting == nil || kbSetting.Chat == nil {
return nil // No KB configuration for chat, skip
}
// Check if KB API is initialized
if kb.API == nil {
return fmt.Errorf("KB API not initialized")
}
// Check if authorized info is available
if ctx.Authorized == nil {
return fmt.Errorf("authorized information not available")
}
chatKB := kbSetting.Chat
// Debug: log locale information
fmt.Printf(">>> prepareKBCollection: locale=%s\n", ctx.Locale)
// Get KB collection ID for this chat session
// Same team + user always produces the same ID (idempotent)
collectionID := GetChatKBID(ctx.Authorized.TeamID, ctx.Authorized.UserID)
// Check if this collection is currently being created by another goroutine
if _, isCreating := kbCollectionCreating.LoadOrStore(collectionID, true); isCreating {
fmt.Printf(">>> KB collection %s is already being created, skipping\n", collectionID)
return nil
}
// Ensure cleanup even if panic occurs
defer kbCollectionCreating.Delete(collectionID)
// Check if collection already exists
existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID)
if err != nil {
// If check fails, log and continue to create (let create handle conflicts)
fmt.Printf(">>> Warning: failed to check collection existence: %v, will attempt to create\n", err)
} else if existsResult != nil && existsResult.Exists {
// Collection exists, no need to create
fmt.Printf(">>> KB collection already exists: %s\n", collectionID)
return nil
}
// Create new collection for this chat session
createParams := &kbapi.CreateCollectionParams{
ID: collectionID,
EmbeddingProviderID: chatKB.EmbeddingProviderID,
EmbeddingOptionID: chatKB.EmbeddingOptionID,
Locale: chatKB.Locale,
Config: chatKB.Config,
Metadata: mergeChatMetadata(chatKB.Metadata, ctx),
AuthScope: ctx.Authorized.WithCreateScope(make(map[string]interface{})),
}
_, err = kb.API.CreateCollection(ctx.Context, createParams)
if err != nil {
return fmt.Errorf("failed to create KB collection: %w", err)
}
fmt.Printf(">>> Created KB collection: %s for team=%s, user=%s\n",
collectionID, ctx.Authorized.TeamID, ctx.Authorized.UserID)
_ = opts
return nil
}
// GetChatKBID returns the KB collection ID for a chat session
// Same team + user always returns the same ID (deterministic)
// Format: chat_{team}_{user} or chat_user_{user} if no team
func GetChatKBID(teamID, userID string) string {
// Sanitize IDs: replace invalid chars with underscores
cleanTeamID := sanitizeCollectionID(teamID)
cleanUserID := sanitizeCollectionID(userID)
if cleanTeamID != "" {
return fmt.Sprintf("chat_%s_%s", cleanTeamID, cleanUserID)
}
return fmt.Sprintf("chat_user_%s", cleanUserID)
}
// sanitizeCollectionID replaces invalid characters with underscores
// Collection IDs only allow: a-z, A-Z, 0-9, and underscore
func sanitizeCollectionID(id string) string {
if id == "" {
return ""
}
// Replace any character that is not alphanumeric or underscore with underscore
result := make([]byte, len(id))
for i := 0; i < len(id); i++ {
c := id[i]
if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' {
result[i] = c
} else {
result[i] = '_'
}
}
return string(result)
}
// mergeChatMetadata merges default metadata with chat context information
func mergeChatMetadata(defaultMetadata map[string]interface{}, ctx *agentcontext.Context) map[string]interface{} {
metadata := make(map[string]interface{})
// Copy default metadata
for k, v := range defaultMetadata {
metadata[k] = v
}
// Add chat-specific metadata (only for internal tracking, not displayed)
metadata["chat_id"] = ctx.ChatID
metadata["team_id"] = ctx.Authorized.TeamID
metadata["user_id"] = ctx.Authorized.UserID
// Get locale from context, default to zh-CN if not set
locale := ctx.Locale
if locale == "" {
locale = "zh-CN"
}
locale = strings.ToLower(locale)
// Use i18n for name and description (fixed, not showing user/team IDs)
if _, exists := metadata["name"]; !exists {
metadata["name"] = i18n.T(locale, "kb.chat.name")
}
if _, exists := metadata["description"]; !exists {
metadata["description"] = i18n.T(locale, "kb.chat.description")
}
fmt.Printf(">>> mergeChatMetadata: locale=%s, name=%v, description=%v\n", locale, metadata["name"], metadata["description"]) // Debug log
return metadata
}
func (ast *Assistant) saveChat(ctx *agentcontext.Context, input []agentcontext.Message, opts *agentcontext.Options) error {
_ = ctx
_ = input
_ = opts
return nil
}

View file

@ -0,0 +1,302 @@
package assistant_test
import (
"context"
"fmt"
"sync"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/yao/agent/assistant"
agentcontext "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/testutils"
"github.com/yaoapp/yao/kb"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
)
func TestGetChatKBID(t *testing.T) {
t.Run("WithTeamAndUser", func(t *testing.T) {
teamID := "5659-5504-2879"
userID := "4287-9400-2030-0504"
collectionID := assistant.GetChatKBID(teamID, userID)
// Should sanitize dashes to underscores
expected := "chat_5659_5504_2879_4287_9400_2030_0504"
assert.Equal(t, expected, collectionID)
t.Logf("✓ Collection ID with team: %s", collectionID)
})
t.Run("WithoutTeam", func(t *testing.T) {
teamID := ""
userID := "4287-9400-2030-0504"
collectionID := assistant.GetChatKBID(teamID, userID)
// Should use chat_user_ prefix
expected := "chat_user_4287_9400_2030_0504"
assert.Equal(t, expected, collectionID)
t.Logf("✓ Collection ID without team: %s", collectionID)
})
t.Run("Idempotent", func(t *testing.T) {
teamID := "test-team-123"
userID := "test-user-456"
id1 := assistant.GetChatKBID(teamID, userID)
id2 := assistant.GetChatKBID(teamID, userID)
id3 := assistant.GetChatKBID(teamID, userID)
// Same input should always produce same output
assert.Equal(t, id1, id2)
assert.Equal(t, id2, id3)
t.Logf("✓ Idempotent: %s", id1)
})
t.Run("SanitizeSpecialChars", func(t *testing.T) {
teamID := "team-with-dashes@123"
userID := "user.with.dots!"
collectionID := assistant.GetChatKBID(teamID, userID)
// Should only contain alphanumeric and underscores
assert.Regexp(t, "^[a-zA-Z0-9_]+$", collectionID)
t.Logf("✓ Sanitized ID: %s", collectionID)
})
t.Run("EmptyUserID", func(t *testing.T) {
teamID := "test-team"
userID := ""
collectionID := assistant.GetChatKBID(teamID, userID)
// Should handle empty user ID gracefully
expected := "chat_test_team_"
assert.Equal(t, expected, collectionID)
t.Logf("✓ Empty user ID handled: %s", collectionID)
})
}
func TestPrepareKBCollection(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Skip if KB not configured
kbSetting := assistant.GetGlobalKBSetting()
if kbSetting == nil || kbSetting.Chat == nil {
t.Skip("KB chat settings not configured in agent/kb.yml, skipping test")
}
// Get assistant
ast, err := assistant.Get("mohe")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("CreateNewCollection", func(t *testing.T) {
// Use unique IDs based on timestamp to avoid conflicts
timestamp := fmt.Sprintf("%d", time.Now().UnixNano())
teamID := fmt.Sprintf("test_team_%s", timestamp)
userID := fmt.Sprintf("test_user_%s", timestamp)
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_chat_prepare_001",
Authorized: &oauthtypes.AuthorizedInfo{
TeamID: teamID,
UserID: userID,
},
}
opts := &agentcontext.Options{}
// This should create a new KB collection
err := ast.InitializeConversation(ctx, opts)
// Should not return error
assert.NoError(t, err)
t.Logf("✓ KB collection prepared successfully")
// Clean up
collectionID := assistant.GetChatKBID(teamID, userID)
_, _ = kb.API.RemoveCollection(ctx.Context, collectionID)
})
t.Run("IdempotentCollectionCreation", func(t *testing.T) {
// Use unique IDs based on timestamp to avoid conflicts
timestamp := fmt.Sprintf("%d", time.Now().UnixNano())
teamID := fmt.Sprintf("idem_team_%s", timestamp)
userID := fmt.Sprintf("idem_user_%s", timestamp)
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_chat_idempotent",
Authorized: &oauthtypes.AuthorizedInfo{
TeamID: teamID,
UserID: userID,
},
}
opts := &agentcontext.Options{}
// First call - creates collection
err1 := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err1)
// Second call - should skip because collection exists
err2 := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err2)
// Third call - still no error
err3 := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err3)
t.Logf("✓ Idempotent collection preparation works correctly")
// Clean up after test
collectionID := assistant.GetChatKBID(teamID, userID)
_, _ = kb.API.RemoveCollection(ctx.Context, collectionID)
})
t.Run("HandleMissingAuthorizedInfo", func(t *testing.T) {
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_chat_no_auth",
Authorized: nil, // Missing authorized info
}
opts := &agentcontext.Options{}
// Should not error, just skip KB preparation
err := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err)
t.Logf("✓ Correctly skipped KB preparation when authorized info is missing")
})
t.Run("ConcurrentCreation", func(t *testing.T) {
// Use unique IDs based on timestamp to avoid conflicts
timestamp := fmt.Sprintf("%d", time.Now().UnixNano())
teamID := fmt.Sprintf("concurrent_team_%s", timestamp)
userID := fmt.Sprintf("concurrent_user_%s", timestamp)
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_chat_concurrent",
Authorized: &oauthtypes.AuthorizedInfo{
TeamID: teamID,
UserID: userID,
},
}
opts := &agentcontext.Options{}
// Launch 5 concurrent calls to create the same collection
var wg sync.WaitGroup
errors := make([]error, 5)
for i := 0; i < 5; i++ {
wg.Add(1)
go func(idx int) {
defer wg.Done()
errors[idx] = ast.InitializeConversation(ctx, opts)
}(i)
}
// Wait for all goroutines to complete
wg.Wait()
// All calls should succeed (no errors, or just warning logs)
// Note: Some goroutines may skip due to concurrent creation lock
for i, err := range errors {
assert.NoError(t, err, "Goroutine %d should not error", i)
}
// Wait a bit for async operations to complete
time.Sleep(200 * time.Millisecond)
// Verify collection was created (at least by one goroutine)
collectionID := assistant.GetChatKBID(teamID, userID)
existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID)
if err != nil || existsResult == nil || !existsResult.Exists {
// Collection might not have been created due to errors, that's okay for this test
// The main goal is to verify no panics or race conditions occurred
t.Logf("⚠ Collection not created (might have failed), but no panics occurred: %v", err)
} else {
t.Logf("✓ Concurrent creation handled correctly, collection: %s", collectionID)
}
// Clean up
_, _ = kb.API.RemoveCollection(ctx.Context, collectionID)
})
}
func TestInitializeConversation(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean(t)
// Skip if KB not configured
kbSetting := assistant.GetGlobalKBSetting()
if kbSetting == nil || kbSetting.Chat == nil {
t.Skip("KB chat settings not configured in agent/kb.yml, skipping test")
}
ast, err := assistant.Get("mohe")
require.NoError(t, err)
require.NotNil(t, ast)
t.Run("FullInitialization", func(t *testing.T) {
// Use unique IDs based on timestamp to avoid conflicts
timestamp := fmt.Sprintf("%d", time.Now().UnixNano())
teamID := fmt.Sprintf("init_team_%s", timestamp)
userID := fmt.Sprintf("init_user_%s", timestamp)
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_init_chat_001",
Authorized: &oauthtypes.AuthorizedInfo{
TeamID: teamID,
UserID: userID,
},
}
opts := &agentcontext.Options{}
// Should initialize conversation without error
err := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err)
t.Logf("✓ Conversation initialized successfully")
// Verify collection was created
collectionID := assistant.GetChatKBID(teamID, userID)
existsResult, err := kb.API.CollectionExists(ctx.Context, collectionID)
assert.NoError(t, err)
assert.NotNil(t, existsResult)
assert.True(t, existsResult.Exists, "KB collection should be created")
t.Logf("✓ KB collection created: %s", collectionID)
// Clean up
_, _ = kb.API.RemoveCollection(ctx.Context, collectionID)
})
t.Run("SkipHistoryFlag", func(t *testing.T) {
ctx := &agentcontext.Context{
Context: context.Background(),
ChatID: "test_skip_history",
Authorized: &oauthtypes.AuthorizedInfo{
TeamID: "skip_team",
UserID: "skip_user",
},
}
opts := &agentcontext.Options{
Skip: &agentcontext.Skip{
History: true,
},
}
// Should skip initialization when history flag is set
err := ast.InitializeConversation(ctx, opts)
assert.NoError(t, err)
t.Logf("✓ Correctly skipped with history flag")
})
}

View file

@ -1,31 +0,0 @@
package assistant
import (
"github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/trace/types"
)
// WithHistory merges the input messages with chat history and traces it
// This method can be overridden or extended to implement actual history loading
func (ast *Assistant) WithHistory(
ctx *context.Context,
inputMessages []context.Message,
agentNode types.Node,
) ([]context.Message, error) {
// TODO: Implement actual history loading logic here
// For now, just simulate a check and return the input messages as is
// Simulate error check (this is where actual history loading would happen)
// if some_condition {
// ast.traceAgentFail(agentNode, err)
// return nil, err
// }
fullMessages := inputMessages
// Log the chat history
ast.traceAgentHistory(ctx, agentNode, fullMessages)
return fullMessages, nil
}

View file

@ -23,6 +23,11 @@ func (s *Script) Execute(ctx *context.Context, method string, args ...interface{
} }
defer scriptCtx.Close() defer scriptCtx.Close()
// Set authorized information if available
if ctx.Authorized != nil {
scriptCtx.WithAuthorized(ctx.Authorized.AuthorizedToMap())
}
// The first argument is the context // The first argument is the context
args = append([]interface{}{ctx}, args...) args = append([]interface{}{ctx}, args...)

View file

@ -23,9 +23,10 @@ var loaded = NewCache(200) // 200 is the default capacity
var storage store.Store = nil var storage store.Store = nil
var search interface{} = nil var search interface{} = nil
var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{} var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
var defaultConnector string = "" // default connector var defaultConnector string = "" // default connector
var globalUses *context.Uses = nil // global uses configuration from agent.yml var globalUses *context.Uses = nil // global uses configuration from agent.yml
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
var globalKBSetting *store.KBSetting = nil // global KB setting from agent/kb.yml
// LoadBuiltIn load the built-in assistants // LoadBuiltIn load the built-in assistants
func LoadBuiltIn() error { func LoadBuiltIn() error {
@ -161,6 +162,16 @@ func GetGlobalPrompts(ctx map[string]string) []store.Prompt {
return store.Prompts(globalPrompts).Parse(ctx) return store.Prompts(globalPrompts).Parse(ctx)
} }
// SetGlobalKBSetting set the global KB setting from agent/kb.yml
func SetGlobalKBSetting(kbSetting *store.KBSetting) {
globalKBSetting = kbSetting
}
// GetGlobalKBSetting returns the global KB setting
func GetGlobalKBSetting() *store.KBSetting {
return globalKBSetting
}
// SetCache set the cache // SetCache set the cache
func SetCache(capacity int) { func SetCache(capacity int) {
ClearCache() ClearCache()

View file

@ -0,0 +1,14 @@
package assistant
import (
"fmt"
"github.com/yaoapp/yao/agent/context"
)
func (ast *Assistant) checkPermissions(ctx *context.Context) error {
if ctx.Authorized == nil {
return fmt.Errorf("authorized information not found")
}
return nil
}

View file

@ -20,6 +20,11 @@ var scriptsMutex sync.Mutex
// Execute execute the script // Execute execute the script
func (s *Script) Execute(ctx context.Context, method string, args ...interface{}) (interface{}, error) { func (s *Script) Execute(ctx context.Context, method string, args ...interface{}) (interface{}, error) {
return s.ExecuteWithAuthorized(ctx, method, nil, args...)
}
// ExecuteWithAuthorized execute the script with authorized information
func (s *Script) ExecuteWithAuthorized(ctx context.Context, method string, authorized map[string]interface{}, args ...interface{}) (interface{}, error) {
if s == nil || s.Script == nil { if s == nil || s.Script == nil {
return nil, nil return nil, nil
} }
@ -30,6 +35,11 @@ func (s *Script) Execute(ctx context.Context, method string, args ...interface{}
} }
defer scriptCtx.Close() defer scriptCtx.Close()
// Set authorized information if available
if authorized != nil {
scriptCtx.WithAuthorized(authorized)
}
// Call the method with provided arguments as-is // Call the method with provided arguments as-is
result, err := scriptCtx.CallWith(ctx, method, args...) result, err := scriptCtx.CallWith(ctx, method, args...)
@ -359,8 +369,14 @@ func makeScriptHandler(script *Script) process.Handler {
// Get arguments from process // Get arguments from process
args := p.Args args := p.Args
// Execute the script // Convert authorized info to map if available
result, err := script.Execute(p.Context, method, args...) var authorized map[string]interface{}
if p.Authorized != nil {
authorized = p.Authorized.AuthorizedToMap()
}
// Execute the script with authorized information
result, err := script.ExecuteWithAuthorized(p.Context, method, authorized, args...)
if err != nil { if err != nil {
exception.New(err.Error(), 500).Throw() exception.New(err.Error(), 500).Throw()
} }

View file

@ -1,10 +1,12 @@
package assistant package assistant
import ( import (
"context"
"testing" "testing"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require" "github.com/stretchr/testify/require"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/test" "github.com/yaoapp/yao/test"
) )
@ -155,3 +157,151 @@ func TestGenerateScriptID(t *testing.T) {
// TestLoadScriptsThreadSafety tests concurrent script loading // TestLoadScriptsThreadSafety tests concurrent script loading
// Note: This test is commented out due to path format differences // Note: This test is commented out due to path format differences
// Thread safety is ensured by the scriptsMutex in LoadScripts function // Thread safety is ensured by the scriptsMutex in LoadScripts function
func TestExecuteWithAuthorized(t *testing.T) {
test.Prepare(t, config.Conf)
defer test.Clean()
t.Run("ExecuteWithAuthorizedInfo", func(t *testing.T) {
// Create a script that returns the authorized info from __yao_data
scriptSource := `
function GetAuth() {
if (typeof __yao_data !== 'undefined' && __yao_data.AUTHORIZED) {
return __yao_data.AUTHORIZED;
}
return null;
}
`
data := map[string]interface{}{
"scripts": map[string]interface{}{
"auth_test": scriptSource,
},
}
_, scripts, err := LoadScriptsFromData(data, "test.authorized")
require.NoError(t, err)
require.NotNil(t, scripts)
require.Contains(t, scripts, "auth_test")
script := scripts["auth_test"]
// Create authorized info
authorized := map[string]interface{}{
"user_id": "user123",
"team_id": "team456",
"scope": "read write",
"constraints": map[string]interface{}{
"team_only": true,
},
}
// Execute with authorized info
ctx := context.Background()
result, err := script.ExecuteWithAuthorized(ctx, "GetAuth", authorized)
require.NoError(t, err)
require.NotNil(t, result)
// Verify the authorized info was passed correctly
resultMap, ok := result.(map[string]interface{})
require.True(t, ok, "Result should be a map")
assert.Equal(t, "user123", resultMap["user_id"])
assert.Equal(t, "team456", resultMap["team_id"])
assert.Equal(t, "read write", resultMap["scope"])
constraints, ok := resultMap["constraints"].(map[string]interface{})
require.True(t, ok, "Constraints should be a map")
assert.Equal(t, true, constraints["team_only"])
t.Logf("✓ Authorized info passed correctly to script")
})
t.Run("ExecuteWithoutAuthorizedInfo", func(t *testing.T) {
// Create a script that checks for authorized info
scriptSource := `
function CheckAuth() {
if (typeof __yao_data !== 'undefined' && __yao_data.AUTHORIZED) {
return { hasAuth: true, data: __yao_data.AUTHORIZED };
}
return { hasAuth: false };
}
`
data := map[string]interface{}{
"scripts": map[string]interface{}{
"no_auth_test": scriptSource,
},
}
_, scripts, err := LoadScriptsFromData(data, "test.noauth")
require.NoError(t, err)
require.NotNil(t, scripts)
require.Contains(t, scripts, "no_auth_test")
script := scripts["no_auth_test"]
// Execute without authorized info
ctx := context.Background()
result, err := script.Execute(ctx, "CheckAuth")
require.NoError(t, err)
require.NotNil(t, result)
resultMap, ok := result.(map[string]interface{})
require.True(t, ok)
assert.Equal(t, false, resultMap["hasAuth"])
t.Logf("✓ Script executed correctly without authorized info")
})
t.Run("MakeScriptHandlerWithAuthorized", func(t *testing.T) {
// Create a script that returns authorized user_id
scriptSource := `
function GetUserID() {
if (typeof __yao_data !== 'undefined' && __yao_data.AUTHORIZED) {
return __yao_data.AUTHORIZED.user_id || null;
}
return null;
}
`
data := map[string]interface{}{
"scripts": map[string]interface{}{
"handler_test": scriptSource,
},
}
_, scripts, err := LoadScriptsFromData(data, "test.handler")
require.NoError(t, err)
require.NotNil(t, scripts)
require.Contains(t, scripts, "handler_test")
script := scripts["handler_test"]
// Create a process handler
handler := makeScriptHandler(script)
require.NotNil(t, handler)
// Create a mock process with authorized info
ctx := context.Background()
p := &process.Process{
Method: "GetUserID",
Args: []interface{}{},
Context: ctx,
Authorized: &process.AuthorizedInfo{
UserID: "user999",
TeamID: "team888",
Scope: "admin",
},
}
// Execute the handler
result := handler(p)
require.NotNil(t, result)
// Verify the result
assert.Equal(t, "user999", result)
t.Logf("✓ Process handler correctly passed authorized info")
})
}

View file

@ -94,6 +94,10 @@ func init() {
"mcp.list_samples.description": "List samples for '%s' from MCP client '%s'", "mcp.list_samples.description": "List samples for '%s' from MCP client '%s'",
"mcp.get_sample.label": "MCP: Get Sample", "mcp.get_sample.label": "MCP: Get Sample",
"mcp.get_sample.description": "Get sample #%d for '%s' from MCP client '%s'", "mcp.get_sample.description": "Get sample #%d for '%s' from MCP client '%s'",
// KB: Chat collection
"kb.chat.name": "Chat Knowledge Base",
"kb.chat.description": "Auto-created knowledge base collection for chat sessions",
}, },
} }
@ -156,6 +160,10 @@ func init() {
"common.status.completed": "已完成", "common.status.completed": "已完成",
"common.status.failed": "失败", "common.status.failed": "失败",
"common.status.retrying": "重试中", "common.status.retrying": "重试中",
// KB: Chat collection
"kb.chat.name": "聊天知识库",
"kb.chat.description": "自动为聊天会话创建的知识库集合",
}, },
} }
@ -246,6 +254,10 @@ func init() {
"mcp.list_samples.description": "从 MCP 客户端 '%s' 列出 '%s' 的示例", "mcp.list_samples.description": "从 MCP 客户端 '%s' 列出 '%s' 的示例",
"mcp.get_sample.label": "MCP: 获取示例", "mcp.get_sample.label": "MCP: 获取示例",
"mcp.get_sample.description": "从 MCP 客户端 '%s' 获取 '%s' 的第 %d 个示例", "mcp.get_sample.description": "从 MCP 客户端 '%s' 获取 '%s' 的第 %d 个示例",
// KB: Chat collection
"kb.chat.name": "聊天知识库",
"kb.chat.description": "自动为聊天会话创建的知识库集合",
}, },
} }
} }

View file

@ -86,6 +86,12 @@ func Load(cfg config.Config) error {
return err return err
} }
// Initialize KB Configuration
err = initKBConfig()
if err != nil {
return err
}
// Initialize Assistant // Initialize Assistant
err = initAssistant() err = initAssistant()
if err != nil { if err != nil {
@ -209,6 +215,10 @@ func initAssistant() error {
assistant.SetModelCapabilities(agentDSL.Models) assistant.SetModelCapabilities(agentDSL.Models)
} }
if agentDSL.KB != nil {
assistant.SetGlobalKBSetting(agentDSL.KB)
}
// Load Built-in Assistants // Load Built-in Assistants
err := assistant.LoadBuiltIn() err := assistant.LoadBuiltIn()
if err != nil { if err != nil {
@ -225,6 +235,29 @@ func initAssistant() error {
return nil return nil
} }
// initKBConfig initialize the knowledge base configuration from agent/kb.yml
func initKBConfig() error {
path := filepath.Join("agent", "kb.yml")
if exists, _ := application.App.Exists(path); !exists {
return nil // KB config is optional
}
// Read the KB configuration
bytes, err := application.App.Read(path)
if err != nil {
return err
}
var kbSetting store.KBSetting
err = application.Parse("kb.yml", bytes, &kbSetting)
if err != nil {
return err
}
agentDSL.KB = &kbSetting
return nil
}
// defaultAssistant get the default assistant // defaultAssistant get the default assistant
func defaultAssistant() (*assistant.Assistant, error) { func defaultAssistant() (*assistant.Assistant, error) {
if agentDSL.Uses == nil || agentDSL.Uses.Default == "" { if agentDSL.Uses == nil || agentDSL.Uses.Default == "" {

View file

@ -59,6 +59,41 @@ func TestLoad(t *testing.T) {
assert.NotNil(t, agent.Models) assert.NotNil(t, agent.Models)
assert.Greater(t, len(agent.Models), 0) assert.Greater(t, len(agent.Models), 0)
}) })
t.Run("LoadKBConfig", func(t *testing.T) {
// KB configuration should be loaded from agent/kb.yml
assert.NotNil(t, agent.KB)
assert.NotNil(t, agent.KB.Chat)
// Verify chat KB settings
assert.Equal(t, "__yao.openai", agent.KB.Chat.EmbeddingProviderID)
assert.Equal(t, "text-embedding-3-small", agent.KB.Chat.EmbeddingOptionID)
assert.Equal(t, "zh-CN", agent.KB.Chat.Locale)
// Verify config
assert.NotNil(t, agent.KB.Chat.Config)
assert.Equal(t, "hnsw", agent.KB.Chat.Config.IndexType.String())
assert.Equal(t, "cosine", agent.KB.Chat.Config.Distance.String())
// Verify metadata
assert.NotNil(t, agent.KB.Chat.Metadata)
assert.Equal(t, "chat_session", agent.KB.Chat.Metadata["category"])
assert.Equal(t, true, agent.KB.Chat.Metadata["auto_created"])
// Verify document defaults
assert.NotNil(t, agent.KB.Chat.DocumentDefaults)
assert.NotNil(t, agent.KB.Chat.DocumentDefaults.Chunking)
assert.Equal(t, "__yao.structured", agent.KB.Chat.DocumentDefaults.Chunking.ProviderID)
assert.Equal(t, "standard", agent.KB.Chat.DocumentDefaults.Chunking.OptionID)
assert.NotNil(t, agent.KB.Chat.DocumentDefaults.Extraction)
assert.Equal(t, "__yao.openai", agent.KB.Chat.DocumentDefaults.Extraction.ProviderID)
assert.Equal(t, "gpt-4o-mini", agent.KB.Chat.DocumentDefaults.Extraction.OptionID)
assert.NotNil(t, agent.KB.Chat.DocumentDefaults.Converter)
assert.Equal(t, "__yao.utf8", agent.KB.Chat.DocumentDefaults.Converter.ProviderID)
assert.Equal(t, "standard-text", agent.KB.Chat.DocumentDefaults.Converter.OptionID)
})
} }
func TestGetGlobalPrompts(t *testing.T) { func TestGetGlobalPrompts(t *testing.T) {

File diff suppressed because it is too large Load diff

View file

@ -4,6 +4,7 @@ import (
"encoding/json" "encoding/json"
"fmt" "fmt"
graphragtypes "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/xun/dbal/query" "github.com/yaoapp/xun/dbal/query"
"github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/context"
"github.com/yaoapp/yao/agent/i18n" "github.com/yaoapp/yao/agent/i18n"
@ -95,6 +96,34 @@ type Prompt struct {
Name string `json:"name,omitempty"` Name string `json:"name,omitempty"`
} }
// KBSetting Knowledge Base configuration for agent (from agent/kb.yml)
type KBSetting struct {
Chat *ChatKBSetting `json:"chat,omitempty" yaml:"chat,omitempty"` // Chat session KB settings
}
// ChatKBSetting represents KB settings for chat sessions
type ChatKBSetting struct {
EmbeddingProviderID string `json:"embedding_provider_id" yaml:"embedding_provider_id"` // Embedding provider ID
EmbeddingOptionID string `json:"embedding_option_id" yaml:"embedding_option_id"` // Embedding option ID
Locale string `json:"locale,omitempty" yaml:"locale,omitempty"` // Locale for content processing
Config *graphragtypes.CreateCollectionOptions `json:"config,omitempty" yaml:"config,omitempty"` // Vector index configuration
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"` // Collection metadata defaults
DocumentDefaults *DocumentDefaults `json:"document_defaults,omitempty" yaml:"document_defaults,omitempty"` // Document processing defaults
}
// DocumentDefaults represents default settings for document processing
type DocumentDefaults struct {
Chunking *ProviderOption `json:"chunking,omitempty" yaml:"chunking,omitempty"` // Chunking provider configuration
Extraction *ProviderOption `json:"extraction,omitempty" yaml:"extraction,omitempty"` // Extraction provider configuration
Converter *ProviderOption `json:"converter,omitempty" yaml:"converter,omitempty"` // Converter provider configuration
}
// ProviderOption represents a provider and option ID pair
type ProviderOption struct {
ProviderID string `json:"provider_id" yaml:"provider_id"` // Provider ID
OptionID string `json:"option_id" yaml:"option_id"` // Option ID within the provider
}
// KnowledgeBase the knowledge base configuration // KnowledgeBase the knowledge base configuration
type KnowledgeBase struct { type KnowledgeBase struct {
Collections []string `json:"collections,omitempty"` // Knowledge base collection IDs Collections []string `json:"collections,omitempty"` // Knowledge base collection IDs

View file

@ -5,6 +5,7 @@ import (
"github.com/yaoapp/yao/agent" "github.com/yaoapp/yao/agent"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/test" "github.com/yaoapp/yao/test"
) )
@ -16,8 +17,14 @@ import (
func Prepare(t *testing.T, opts ...interface{}) { func Prepare(t *testing.T, opts ...interface{}) {
test.Prepare(t, config.Conf, opts...) test.Prepare(t, config.Conf, opts...)
// Load KB (required for agent KB features)
_, err := kb.Load(config.Conf)
if err != nil {
t.Fatal(err)
}
// Load agent // Load agent
err := agent.Load(config.Conf) err = agent.Load(config.Conf)
if err != nil { if err != nil {
t.Fatal(err) t.Fatal(err)
} }

View file

@ -18,6 +18,7 @@ type DSL struct {
// Global External Settings - model capabilities, tools, etc. // Global External Settings - model capabilities, tools, etc.
// =============================== // ===============================
Models map[string]openai.Capabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration Models map[string]openai.Capabilities `json:"models,omitempty" yaml:"models,omitempty"` // The model capabilities configuration
KB *store.KBSetting `json:"kb,omitempty" yaml:"kb,omitempty"` // The knowledge base configuration loaded from agent/kb.yml
// Internal // Internal
// =============================== // ===============================

15
kb/api/api.go Normal file
View file

@ -0,0 +1,15 @@
package api
import (
"github.com/yaoapp/gou/graphrag/types"
kbtypes "github.com/yaoapp/yao/kb/types"
)
// NewAPI creates a new API instance with the provided KB dependencies
func NewAPI(graphRag types.GraphRag, config *kbtypes.Config, providers *kbtypes.ProviderConfig) API {
return &KBInstance{
GraphRag: graphRag,
Config: config,
Providers: providers,
}
}

667
kb/api/collection.go Normal file
View file

@ -0,0 +1,667 @@
package api
import (
"context"
"fmt"
graphragtypes "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/maps"
)
// CreateCollection creates a new collection with the provided parameters
func (instance *KBInstance) CreateCollection(ctx context.Context, params *CreateCollectionParams) (*CreateCollectionResult, error) {
// Basic validation (before provider settings)
if params.ID == "" {
return nil, fmt.Errorf("invalid parameters: id is required")
}
if params.EmbeddingProviderID == "" {
return nil, fmt.Errorf("invalid parameters: embedding_provider_id is required")
}
if params.EmbeddingOptionID == "" {
return nil, fmt.Errorf("invalid parameters: embedding_option_id is required")
}
// Get provider settings to resolve dimension and properties
providerSettings, err := instance.getProviderSettings(params.EmbeddingProviderID, params.EmbeddingOptionID, params.Locale)
if err != nil {
return nil, fmt.Errorf("failed to resolve provider settings: %w", err)
}
// Set dimension from provider settings
if params.Config != nil {
params.Config.Dimension = providerSettings.Dimension
}
// Validate full parameters after dimension is set
if err := validateCreateParams(params); err != nil {
return nil, fmt.Errorf("invalid parameters: %w", err)
}
// Prepare metadata
metadata := params.Metadata
if metadata == nil {
metadata = make(map[string]interface{})
}
// Add embedding information to metadata
metadata["__embedding_provider"] = params.EmbeddingProviderID
metadata["__embedding_option"] = params.EmbeddingOptionID
if providerSettings.Properties != nil {
metadata["__embedding_properties"] = providerSettings.Properties
}
if params.Locale != "" {
metadata["__locale"] = params.Locale
}
// Prepare database record
dbData := map[string]interface{}{
"collection_id": params.ID,
"name": metadata["name"],
"description": metadata["description"],
"status": "creating",
"embedding_provider_id": params.EmbeddingProviderID,
"embedding_option_id": params.EmbeddingOptionID,
"embedding_properties": providerSettings.Properties,
"locale": params.Locale,
}
// Add config options to database if provided
if params.Config != nil {
if params.Config.Distance != "" {
dbData["distance"] = params.Config.Distance
}
if params.Config.IndexType != "" {
dbData["index_type"] = params.Config.IndexType
}
if params.Config.M > 0 {
dbData["m"] = params.Config.M
}
if params.Config.EfConstruction > 0 {
dbData["ef_construction"] = params.Config.EfConstruction
}
if params.Config.EfSearch > 0 {
dbData["ef_search"] = params.Config.EfSearch
}
if params.Config.NumLists > 0 {
dbData["num_lists"] = params.Config.NumLists
}
if params.Config.NumProbes > 0 {
dbData["num_probes"] = params.Config.NumProbes
}
}
// Add share field from metadata if provided
if share, ok := metadata["share"].(string); ok {
if share == "private" || share == "team" {
dbData["share"] = share
}
}
// Merge auth scope fields
if params.AuthScope != nil {
for k, v := range params.AuthScope {
dbData[k] = v
}
}
// Create database record first
_, err = instance.Config.CreateCollection(maps.MapStrAny(dbData))
if err != nil {
return nil, fmt.Errorf("failed to save collection metadata: %w", err)
}
// Read back the database record to get auto-generated fields (created_at, updated_at)
dbRecord, err := instance.Config.FindCollection(params.ID, model.QueryParam{})
if err != nil {
// Rollback on error
rollbackErr := instance.Config.RemoveCollection(params.ID)
if rollbackErr != nil {
log.Error("Failed to rollback collection database record: %v", rollbackErr)
}
return nil, fmt.Errorf("failed to read created collection: %w", err)
}
// Add all database fields to metadata for GraphRag
// This ensures GraphRag metadata contains complete information for vector search filtering
// Timestamps
if createdAt, ok := dbRecord["created_at"]; ok {
metadata["created_at"] = createdAt
// If updated_at is not set, use created_at (for newly created records)
if updatedAt, ok := dbRecord["updated_at"]; ok && updatedAt != nil {
metadata["updated_at"] = updatedAt
} else {
metadata["updated_at"] = createdAt
}
}
// Auth scope fields (for permission-based vector search)
if createdBy, ok := dbRecord["__yao_created_by"]; ok && createdBy != nil {
metadata["__yao_created_by"] = createdBy
}
if teamID, ok := dbRecord["__yao_team_id"]; ok && teamID != nil {
metadata["__yao_team_id"] = teamID
}
if tenantID, ok := dbRecord["__yao_tenant_id"]; ok && tenantID != nil {
metadata["__yao_tenant_id"] = tenantID
}
// Collection ID (for consistency with OpenAPI created collections)
metadata["collection_id"] = params.ID
// Collection properties
if share, ok := dbRecord["share"]; ok && share != nil {
metadata["share"] = share
}
if preset, ok := dbRecord["preset"]; ok {
metadata["preset"] = preset
}
if public, ok := dbRecord["public"]; ok {
metadata["public"] = public
}
if sort, ok := dbRecord["sort"]; ok {
metadata["sort"] = sort
}
if status, ok := dbRecord["status"]; ok && status != nil {
metadata["status"] = status
}
if uid, ok := dbRecord["uid"]; ok {
metadata["uid"] = uid
}
if cover, ok := dbRecord["cover"]; ok {
metadata["cover"] = cover
}
if documentCount, ok := dbRecord["document_count"]; ok {
metadata["document_count"] = documentCount
}
collectionConfig := graphragtypes.CollectionConfig{
ID: params.ID,
Metadata: metadata,
Config: params.Config,
}
// Create collection in GraphRag
collectionID, err := instance.GraphRag.CreateCollection(ctx, collectionConfig)
if err != nil {
// Rollback: remove the database record
rollbackErr := instance.Config.RemoveCollection(params.ID)
if rollbackErr != nil {
log.Error("Failed to rollback collection database record: %v", rollbackErr)
}
return nil, fmt.Errorf("failed to create collection: %w", err)
}
// Update status to active after successful creation
updateErr := instance.updateCollectionWithSync(ctx, params.ID, maps.MapStrAny{"status": "active"})
if updateErr != nil {
log.Error("Failed to update collection status to active: %v", updateErr)
}
return &CreateCollectionResult{
CollectionID: collectionID,
Message: "Collection created successfully",
}, nil
}
// RemoveCollection removes an existing collection by ID
func (instance *KBInstance) RemoveCollection(ctx context.Context, collectionID string) (*RemoveCollectionResult, error) {
if collectionID == "" {
return nil, fmt.Errorf("collection ID is required")
}
// Try to remove from GraphRag (vector/graph stores)
// Don't fail if collection doesn't exist there - we still want to clean up database
removed := false
graphRagErr := error(nil)
removedFromGraphRag, err := instance.GraphRag.RemoveCollection(ctx, collectionID)
if err != nil {
// Log the error but continue to database cleanup
log.Warn("Failed to remove collection from GraphRag: %v (will continue with database cleanup)", err)
graphRagErr = err
} else {
removed = removedFromGraphRag
}
// Always attempt to clean up database, even if GraphRag removal failed
// This ensures we can recover from inconsistent states
documentsRemoved := 0
// Count documents in this collection
if count, err := instance.Config.DocumentCount(collectionID); err == nil {
documentsRemoved = count
}
// Remove all documents belonging to this collection
dbCleanupSuccess := true
if err := instance.Config.RemoveDocumentsByCollectionID(collectionID); err != nil {
log.Error("Failed to remove documents from collection %s: %v", collectionID, err)
dbCleanupSuccess = false
} else {
log.Info("Removed %d documents from collection %s", documentsRemoved, collectionID)
}
// Remove the collection itself from database
if err := instance.Config.RemoveCollection(collectionID); err != nil {
log.Error("Failed to remove collection from database: %v", err)
dbCleanupSuccess = false
} else {
log.Info("Successfully removed collection %s and %d documents from database", collectionID, documentsRemoved)
}
// Determine final result and error
// If both GraphRag and database cleanup failed, return error
if graphRagErr != nil && !dbCleanupSuccess {
return nil, fmt.Errorf("failed to remove collection: GraphRag error: %v", graphRagErr)
}
// If collection didn't exist in GraphRag but was cleaned from database, still consider it successful
if !removed && dbCleanupSuccess {
log.Info("Collection %s was not found in GraphRag but was cleaned from database", collectionID)
}
return &RemoveCollectionResult{
CollectionID: collectionID,
Removed: removed || dbCleanupSuccess, // Consider successful if either succeeded
DocumentsRemoved: documentsRemoved,
Message: "Collection removed successfully",
}, nil
}
// GetCollection retrieves a collection by ID
func (instance *KBInstance) GetCollection(ctx context.Context, collectionID string) (map[string]interface{}, error) {
if collectionID == "" {
return nil, fmt.Errorf("collection ID is required")
}
collection, err := instance.GraphRag.GetCollection(ctx, collectionID)
if err != nil {
// Check if it's a "not found" error
if err.Error() == fmt.Sprintf("collection with ID '%s' not found", collectionID) {
return nil, fmt.Errorf("collection not found")
}
return nil, fmt.Errorf("failed to get collection: %w", err)
}
// Convert CollectionInfo to map[string]interface{}
// Use a hybrid structure: flatten metadata to top level AND include metadata object
// This ensures backward compatibility with both access patterns:
// - collection.id / collection.collection_id (for ID)
// - collection.metadata.name (for nested access)
result := make(map[string]interface{})
result["id"] = collection.ID // Primary ID field for frontend
result["collection_id"] = collection.ID // Alias for backward compatibility
// Flatten metadata fields to top level for backward compatibility
if collection.Metadata != nil {
for k, v := range collection.Metadata {
result[k] = v
}
// Also include the metadata object itself
result["metadata"] = collection.Metadata
}
if collection.Config != nil {
result["config"] = collection.Config
}
return result, nil
}
// CollectionExists checks if a collection exists by ID
func (instance *KBInstance) CollectionExists(ctx context.Context, collectionID string) (*CollectionExistsResult, error) {
if collectionID == "" {
return nil, fmt.Errorf("collection ID is required")
}
exists, err := instance.GraphRag.CollectionExists(ctx, collectionID)
if err != nil {
return nil, fmt.Errorf("failed to check collection existence: %w", err)
}
return &CollectionExistsResult{
CollectionID: collectionID,
Exists: exists,
}, nil
}
// ListCollections lists collections with pagination and filtering
func (instance *KBInstance) ListCollections(ctx context.Context, filter *ListCollectionsFilter) (*ListCollectionsResult, error) {
page := filter.Page
if page <= 0 {
page = DefaultPage
}
pageSize := filter.PageSize
if pageSize <= 0 {
pageSize = DefaultPageSize
} else if pageSize > MaxPageSize {
pageSize = MaxPageSize
}
// Process select fields
selectFields := filter.Select
if len(selectFields) == 0 {
selectFields = DefaultCollectionFields
} else {
// Filter valid fields
validFields := []interface{}{}
for _, field := range selectFields {
if fieldStr, ok := field.(string); ok && AvailableCollectionFields[fieldStr] {
validFields = append(validFields, field)
}
}
if len(validFields) == 0 {
selectFields = DefaultCollectionFields
} else {
selectFields = validFields
}
}
// Build query parameters
param := model.QueryParam{Select: selectFields}
// Build wheres
var wheres []model.QueryWhere
// Add auth filters
if len(filter.AuthFilters) > 0 {
wheres = append(wheres, filter.AuthFilters...)
}
// Filter by keywords (search in name and description)
if filter.Keywords != "" {
wheres = append(wheres, model.QueryWhere{
Column: "name",
Value: "%" + filter.Keywords + "%",
OP: "like",
})
wheres = append(wheres, model.QueryWhere{
Column: "description",
Value: "%" + filter.Keywords + "%",
OP: "like",
Method: "orwhere",
})
}
// Filter by status
if len(filter.Status) > 0 {
statusValues := []interface{}{}
for _, status := range filter.Status {
if status != "" {
statusValues = append(statusValues, status)
}
}
if len(statusValues) > 0 {
if len(statusValues) == 1 {
wheres = append(wheres, model.QueryWhere{
Column: "status",
Value: statusValues[0],
})
} else {
wheres = append(wheres, model.QueryWhere{
Column: "status",
Value: statusValues,
OP: "in",
})
}
}
}
// Filter by system flag
if filter.System != nil {
wheres = append(wheres, model.QueryWhere{
Column: "system",
Value: *filter.System,
})
}
// Filter by embedding_provider_id
if filter.EmbeddingProviderID != "" {
wheres = append(wheres, model.QueryWhere{
Column: "embedding_provider_id",
Value: filter.EmbeddingProviderID,
})
}
param.Wheres = wheres
// Process sort orders
orders := filter.Sort
if len(orders) == 0 {
orders = DefaultSort
} else {
// Validate sort fields
validOrders := []model.QueryOrder{}
for _, order := range orders {
if ValidCollectionSortFields[order.Column] {
validOrders = append(validOrders, order)
}
}
if len(validOrders) == 0 {
orders = DefaultSort
} else {
orders = validOrders
}
}
param.Orders = orders
// Query collections
result, err := instance.Config.SearchCollections(param, page, pageSize)
if err != nil {
return nil, fmt.Errorf("failed to search collections: %w", err)
}
// Convert maps.MapStr result to ListCollectionsResult
listResult := &ListCollectionsResult{
Page: page,
PageSize: pageSize,
Data: make([]map[string]interface{}, 0), // Initialize as empty array, not nil
}
// Extract pagination data from result
if data, ok := result["data"].([]map[string]interface{}); ok {
listResult.Data = data
} else if data, ok := result["data"].([]interface{}); ok {
// Convert []interface{} to []map[string]interface{}
converted := make([]map[string]interface{}, 0, len(data))
for _, item := range data {
if mapItem, ok := item.(map[string]interface{}); ok {
converted = append(converted, mapItem)
}
}
listResult.Data = converted
} else if data, ok := result["data"].([]maps.MapStr); ok {
// Handle []maps.MapStr type (most likely from model.Paginate)
converted := make([]map[string]interface{}, 0, len(data))
for _, item := range data {
converted = append(converted, map[string]interface{}(item))
}
listResult.Data = converted
}
if next, ok := result["next"].(int); ok {
listResult.Next = next
}
if prev, ok := result["prev"].(int); ok {
listResult.Prev = prev
}
if total, ok := result["total"].(int); ok {
listResult.Total = total
}
if pagecnt, ok := result["pagecnt"].(int); ok {
listResult.PageCnt = pagecnt
}
return listResult, nil
}
// UpdateCollectionMetadata updates the metadata of an existing collection
func (instance *KBInstance) UpdateCollectionMetadata(ctx context.Context, collectionID string, params *UpdateMetadataParams) (*UpdateMetadataResult, error) {
if collectionID == "" {
return nil, fmt.Errorf("collection ID is required")
}
if len(params.Metadata) == 0 {
return nil, fmt.Errorf("metadata is required and cannot be empty")
}
err := instance.GraphRag.UpdateCollectionMetadata(ctx, collectionID, params.Metadata)
if err != nil {
return nil, fmt.Errorf("failed to update collection metadata: %w", err)
}
// Update collection metadata in database after successful GraphRag update
// Prepare update data from metadata
updateData := maps.MapStrAny{}
if name, ok := params.Metadata["name"]; ok {
updateData["name"] = name
}
if description, ok := params.Metadata["description"]; ok {
updateData["description"] = description
}
if status, ok := params.Metadata["status"]; ok {
updateData["status"] = status
}
// Merge auth scope fields
if params.AuthScope != nil {
for k, v := range params.AuthScope {
updateData[k] = v
}
}
if len(updateData) > 0 {
// Only update database, don't sync to GraphRag again
if err := instance.Config.UpdateCollection(collectionID, updateData); err != nil {
log.Error("Failed to update collection in database: %v", err)
}
}
return &UpdateMetadataResult{
CollectionID: collectionID,
Message: "Collection metadata updated successfully",
}, nil
}
// Helper methods
// validateCreateParams validates the create collection parameters
func validateCreateParams(params *CreateCollectionParams) error {
if params.ID == "" {
return fmt.Errorf("id is required")
}
if params.EmbeddingProviderID == "" {
return fmt.Errorf("embedding_provider_id is required")
}
if params.EmbeddingOptionID == "" {
return fmt.Errorf("embedding_option_id is required")
}
// Validate CreateCollectionOptions if provided
if params.Config != nil {
if err := params.Config.Validate(); err != nil && err.Error() != "collection name cannot be empty" {
return fmt.Errorf("invalid config: %w", err)
}
}
return nil
}
// ProviderSettings represents the resolved provider configuration
type ProviderSettings struct {
Dimension int `json:"dimension"`
Connector string `json:"connector"`
Properties map[string]interface{} `json:"properties"`
}
// getProviderSettings reads and resolves provider settings by provider ID and option value
func (instance *KBInstance) getProviderSettings(providerID, optionValue, locale string) (*ProviderSettings, error) {
// Default locale to "en" if empty
if locale == "" {
locale = DefaultLocale
}
// Get the specific provider from instance
provider, err := instance.Providers.GetProvider("embedding", providerID, locale)
if err != nil {
return nil, fmt.Errorf("failed to get provider %s: %v", providerID, err)
}
// Find the target option
targetOption, found := provider.GetOption(optionValue)
if !found {
return nil, fmt.Errorf("option not found: %s for provider %s", optionValue, providerID)
}
// Extract settings from option properties
settings := &ProviderSettings{
Properties: make(map[string]interface{}),
}
// Copy all properties
if targetOption.Properties != nil {
for key, value := range targetOption.Properties {
settings.Properties[key] = value
}
}
// Extract dimension
if dim, ok := targetOption.Properties["dimensions"]; ok {
if dimInt, ok := dim.(int); ok {
settings.Dimension = dimInt
} else if dimFloat, ok := dim.(float64); ok {
settings.Dimension = int(dimFloat)
}
}
// Extract connector
if connector, ok := targetOption.Properties["connector"]; ok {
if connStr, ok := connector.(string); ok {
settings.Connector = connStr
}
}
return settings, nil
}
// updateCollectionWithSync updates collection metadata in database and syncs to GraphRag
func (instance *KBInstance) updateCollectionWithSync(ctx context.Context, collectionID string, data maps.MapStrAny) error {
// Create a copy of data for GraphRag to avoid contamination from database operations
originalData := make(maps.MapStrAny)
for k, v := range data {
originalData[k] = v
}
// Update collection in database
if err := instance.Config.UpdateCollection(collectionID, data); err != nil {
return fmt.Errorf("failed to update collection in database: %w", err)
}
// Sync to GraphRag metadata
// Convert the original (unmodified) data to map[string]interface{}
metadata := make(map[string]interface{})
for k, v := range originalData {
metadata[k] = v
}
// Update GraphRag metadata
if err := instance.GraphRag.UpdateCollectionMetadata(ctx, collectionID, metadata); err != nil {
return fmt.Errorf("failed to sync collection metadata to GraphRag: %w", err)
}
return nil
}

732
kb/api/collection_test.go Normal file
View file

@ -0,0 +1,732 @@
package api_test
import (
"context"
"fmt"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
graphragtypes "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/kb/api"
"github.com/yaoapp/yao/test"
)
func TestMain(m *testing.M) {
// Setup test environment
test.Prepare(&testing.T{}, config.Conf)
defer test.Clean()
// Load knowledge base
_, err := kb.Load(config.Conf)
if err != nil {
panic("Failed to load knowledge base: " + err.Error())
}
// Run tests and exit with status code
os.Exit(m.Run())
}
func TestCreateCollection(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_create_%d", time.Now().UnixNano())
// Clean up after test
defer func() {
_, _ = kb.API.RemoveCollection(ctx, testCollectionID)
}()
t.Run("CreateCollectionSuccess", func(t *testing.T) {
params := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Test Collection",
"description": "Test Description",
"share": "team",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Locale: "en",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
M: 16,
EfConstruction: 200,
EfSearch: 64,
// Dimension will be set automatically by the API from provider settings
},
AuthScope: map[string]interface{}{
"__yao_created_by": "test_user",
"__yao_team_id": "test_team",
},
}
result, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
assert.NotNil(t, result)
if result != nil {
assert.Equal(t, testCollectionID, result.CollectionID)
assert.Contains(t, result.Message, "successfully")
t.Logf("Created collection: %s", result.CollectionID)
}
// ✅ Verify that auth scope fields are stored in GraphRag metadata
collection, err := kb.API.GetCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.NotNil(t, collection)
// Check metadata object
metadata, ok := collection["metadata"].(map[string]interface{})
assert.True(t, ok, "metadata should be a map")
// Verify auth scope fields in metadata (for permission-based vector search)
assert.Equal(t, "test_user", metadata["__yao_created_by"], "created_by should be in metadata")
assert.Equal(t, "test_team", metadata["__yao_team_id"], "team_id should be in metadata")
t.Logf("✅ Auth scope fields verified in metadata: created_by=%v, team_id=%v",
metadata["__yao_created_by"], metadata["__yao_team_id"])
// Verify they are also flattened at top level
assert.Equal(t, "test_user", collection["__yao_created_by"], "created_by should be at top level")
assert.Equal(t, "test_team", collection["__yao_team_id"], "team_id should be at top level")
// ✅ Verify other database fields in metadata
assert.Equal(t, "team", metadata["share"], "share should be in metadata")
assert.Equal(t, "active", metadata["status"], "status should be in metadata")
assert.NotNil(t, metadata["preset"], "preset should be in metadata")
assert.NotNil(t, metadata["public"], "public should be in metadata")
assert.NotNil(t, metadata["sort"], "sort should be in metadata")
t.Logf("✅ Database fields verified in metadata: share=%v, status=%v, preset=%v, public=%v",
metadata["share"], metadata["status"], metadata["preset"], metadata["public"])
})
t.Run("CreateCollectionMissingID", func(t *testing.T) {
params := &api.CreateCollectionParams{
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
result, err := kb.API.CreateCollection(ctx, params)
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "id is required")
})
t.Run("CreateCollectionMissingProvider", func(t *testing.T) {
params := &api.CreateCollectionParams{
ID: "test_missing_provider",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
result, err := kb.API.CreateCollection(ctx, params)
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "embedding_provider_id is required")
})
t.Run("CreateCollectionInvalidProvider", func(t *testing.T) {
params := &api.CreateCollectionParams{
ID: "test_invalid_provider",
EmbeddingProviderID: "invalid_provider",
EmbeddingOptionID: "invalid_option",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
result, err := kb.API.CreateCollection(ctx, params)
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "provider")
})
}
func TestGetCollection(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_get_%d", time.Now().UnixNano())
// Create a test collection first
params := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Test Get Collection",
"description": "Test Description",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Locale: "en",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
// Clean up after test
defer func() {
_, _ = kb.API.RemoveCollection(ctx, testCollectionID)
}()
t.Run("GetCollectionSuccess", func(t *testing.T) {
collection, err := kb.API.GetCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.NotNil(t, collection)
// Check that both id and collection_id are present
assert.Equal(t, testCollectionID, collection["id"])
assert.Equal(t, testCollectionID, collection["collection_id"])
// Check that metadata is present
assert.NotNil(t, collection["metadata"])
metadata, ok := collection["metadata"].(map[string]interface{})
assert.True(t, ok)
assert.Equal(t, "Test Get Collection", metadata["name"])
// Check that fields are also flattened at top level
assert.Equal(t, "Test Get Collection", collection["name"])
// Check that config is present
assert.NotNil(t, collection["config"])
// ✅ Check that timestamps are present in metadata (for frontend)
assert.NotNil(t, metadata["created_at"], "created_at should be present in metadata")
assert.NotNil(t, metadata["updated_at"], "updated_at should be present in metadata")
t.Logf("Timestamps in metadata: created_at=%v, updated_at=%v", metadata["created_at"], metadata["updated_at"])
// ✅ Check that timestamps are also flattened at top level
assert.NotNil(t, collection["created_at"], "created_at should be present at top level")
assert.NotNil(t, collection["updated_at"], "updated_at should be present at top level")
t.Logf("Timestamps at top level: created_at=%v, updated_at=%v", collection["created_at"], collection["updated_at"])
// Note: This test doesn't create collection with auth scope, so permission fields won't be present
// See TestCreateCollection/CreateCollectionSuccess for auth scope verification
t.Logf("Retrieved collection: %v", collection["id"])
})
t.Run("GetCollectionNotFound", func(t *testing.T) {
collection, err := kb.API.GetCollection(ctx, "nonexistent_collection")
assert.Error(t, err)
assert.Nil(t, collection)
assert.Contains(t, err.Error(), "not found")
})
t.Run("GetCollectionEmptyID", func(t *testing.T) {
collection, err := kb.API.GetCollection(ctx, "")
assert.Error(t, err)
assert.Nil(t, collection)
assert.Contains(t, err.Error(), "required")
})
}
func TestCollectionExists(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_exists_%d", time.Now().UnixNano())
// Create a test collection
params := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Test Exists Collection",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
// Clean up after test
defer func() {
_, _ = kb.API.RemoveCollection(ctx, testCollectionID)
}()
t.Run("CollectionExistsTrue", func(t *testing.T) {
result, err := kb.API.CollectionExists(ctx, testCollectionID)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.True(t, result.Exists)
assert.Equal(t, testCollectionID, result.CollectionID)
})
t.Run("CollectionExistsFalse", func(t *testing.T) {
result, err := kb.API.CollectionExists(ctx, "nonexistent_collection")
assert.NoError(t, err)
assert.NotNil(t, result)
assert.False(t, result.Exists)
})
t.Run("CollectionExistsEmptyID", func(t *testing.T) {
result, err := kb.API.CollectionExists(ctx, "")
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "required")
})
}
func TestRemoveCollection(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_remove_%d", time.Now().UnixNano())
// Create a test collection
params := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Test Remove Collection",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
t.Run("RemoveCollectionSuccess", func(t *testing.T) {
result, err := kb.API.RemoveCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.True(t, result.Removed)
assert.Equal(t, testCollectionID, result.CollectionID)
assert.Contains(t, result.Message, "successfully")
// Verify collection is removed
exists, err := kb.API.CollectionExists(ctx, testCollectionID)
assert.NoError(t, err)
assert.False(t, exists.Exists)
})
t.Run("RemoveCollectionNotFound", func(t *testing.T) {
// The new implementation is more tolerant - it attempts database cleanup
// even if the collection doesn't exist in GraphRag
// This is considered successful as long as database cleanup succeeds
result, err := kb.API.RemoveCollection(ctx, "nonexistent_collection")
// Should succeed (database cleanup succeeds even if collection doesn't exist)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.True(t, result.Removed)
t.Logf("✓ Handled non-existent collection gracefully (database cleanup succeeded)")
})
t.Run("RemoveCollectionEmptyID", func(t *testing.T) {
result, err := kb.API.RemoveCollection(ctx, "")
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "required")
})
t.Run("RemoveCollectionInconsistentState", func(t *testing.T) {
// Test removing a collection that exists in database but not in vector store
// This simulates an inconsistent state that can occur after failed operations
inconsistentCollectionID := fmt.Sprintf("test_inconsistent_%d", time.Now().UnixNano())
// Create a test collection first
params := &api.CreateCollectionParams{
ID: inconsistentCollectionID,
Metadata: map[string]interface{}{
"name": "Test Inconsistent Collection",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
// Now remove it normally first time
result, err := kb.API.RemoveCollection(ctx, inconsistentCollectionID)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.True(t, result.Removed)
// Verify it's gone
exists, err := kb.API.CollectionExists(ctx, inconsistentCollectionID)
assert.NoError(t, err)
assert.False(t, exists.Exists)
t.Logf("✓ Successfully removed collection in inconsistent state")
})
}
func TestListCollections(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
// Create multiple test collections
timestamp := time.Now().UnixNano()
testCollections := []string{
fmt.Sprintf("test_list_1_%d", timestamp),
fmt.Sprintf("test_list_2_%d", timestamp),
fmt.Sprintf("test_list_3_%d", timestamp),
}
for i, collectionID := range testCollections {
params := &api.CreateCollectionParams{
ID: collectionID,
Metadata: map[string]interface{}{
"name": "Test List Collection " + string(rune('A'+i)),
"description": "Description " + string(rune('A'+i)),
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
}
// Clean up after test
defer func() {
for _, collectionID := range testCollections {
_, _ = kb.API.RemoveCollection(ctx, collectionID)
}
}()
t.Run("ListCollectionsDefault", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotNil(t, result.Data)
assert.GreaterOrEqual(t, len(result.Data), 3) // At least our 3 test collections
assert.Equal(t, 1, result.Page)
assert.Equal(t, 20, result.PageSize)
t.Logf("Found %d collections", len(result.Data))
})
t.Run("ListCollectionsWithPagination", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 2,
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.LessOrEqual(t, len(result.Data), 2)
assert.Equal(t, 1, result.Page)
assert.Equal(t, 2, result.PageSize)
})
t.Run("ListCollectionsWithKeywords", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Keywords: "Test List Collection A",
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.GreaterOrEqual(t, len(result.Data), 1)
// Check that returned collections match the keyword
for _, item := range result.Data {
name, ok := item["name"].(string)
if ok {
assert.Contains(t, name, "Test List Collection")
}
}
})
t.Run("ListCollectionsWithStatus", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Status: []string{"active"},
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
// All collections should have status "active"
for _, item := range result.Data {
status, ok := item["status"].(string)
if ok {
assert.Equal(t, "active", status)
}
}
})
t.Run("ListCollectionsWithSort", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Sort: []model.QueryOrder{
{Column: "created_at", Option: "desc"},
},
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.GreaterOrEqual(t, len(result.Data), 3)
})
t.Run("ListCollectionsWithSelect", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Select: []interface{}{"id", "collection_id", "name", "status"},
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.GreaterOrEqual(t, len(result.Data), 3)
// Check that returned fields are limited
for _, item := range result.Data {
assert.NotNil(t, item["collection_id"])
assert.NotNil(t, item["name"])
}
})
t.Run("ListCollectionsEmptyResult", func(t *testing.T) {
filter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Keywords: "nonexistent_keyword_xyz123",
}
result, err := kb.API.ListCollections(ctx, filter)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.NotNil(t, result.Data)
assert.Equal(t, 0, len(result.Data))
})
}
func TestUpdateCollectionMetadata(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_update_%d", time.Now().UnixNano())
// Create a test collection
params := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Original Name",
"description": "Original Description",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
_, err := kb.API.CreateCollection(ctx, params)
assert.NoError(t, err)
// Clean up after test
defer func() {
_, _ = kb.API.RemoveCollection(ctx, testCollectionID)
}()
t.Run("UpdateMetadataSuccess", func(t *testing.T) {
updateParams := &api.UpdateMetadataParams{
Metadata: map[string]interface{}{
"name": "Updated Name",
"description": "Updated Description",
},
AuthScope: map[string]interface{}{
"__yao_updated_by": "test_user",
},
}
result, err := kb.API.UpdateCollectionMetadata(ctx, testCollectionID, updateParams)
assert.NoError(t, err)
assert.NotNil(t, result)
assert.Equal(t, testCollectionID, result.CollectionID)
assert.Contains(t, result.Message, "successfully")
// Verify the update
collection, err := kb.API.GetCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.Equal(t, "Updated Name", collection["name"])
assert.Equal(t, "Updated Description", collection["description"])
})
t.Run("UpdateMetadataEmptyID", func(t *testing.T) {
updateParams := &api.UpdateMetadataParams{
Metadata: map[string]interface{}{
"name": "Updated Name",
},
}
result, err := kb.API.UpdateCollectionMetadata(ctx, "", updateParams)
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "required")
})
t.Run("UpdateMetadataEmptyMetadata", func(t *testing.T) {
updateParams := &api.UpdateMetadataParams{
Metadata: map[string]interface{}{},
}
result, err := kb.API.UpdateCollectionMetadata(ctx, testCollectionID, updateParams)
assert.Error(t, err)
assert.Nil(t, result)
assert.Contains(t, err.Error(), "empty")
})
t.Run("UpdateMetadataNotFound", func(t *testing.T) {
updateParams := &api.UpdateMetadataParams{
Metadata: map[string]interface{}{
"name": "Updated Name",
},
}
result, err := kb.API.UpdateCollectionMetadata(ctx, "nonexistent_collection", updateParams)
assert.Error(t, err)
assert.Nil(t, result)
})
}
func TestCollectionIntegration(t *testing.T) {
if kb.API == nil {
t.Skip("KB API not initialized")
}
ctx := context.Background()
testCollectionID := fmt.Sprintf("test_integration_%d", time.Now().UnixNano())
t.Run("FullCollectionLifecycle", func(t *testing.T) {
// 1. Create Collection
createParams := &api.CreateCollectionParams{
ID: testCollectionID,
Metadata: map[string]interface{}{
"name": "Integration Test Collection",
"description": "Full lifecycle test",
"share": "team",
},
EmbeddingProviderID: "__yao.openai",
EmbeddingOptionID: "text-embedding-3-small",
Locale: "en",
Config: &graphragtypes.CreateCollectionOptions{
Distance: "cosine",
IndexType: "hnsw",
},
}
createResult, err := kb.API.CreateCollection(ctx, createParams)
assert.NoError(t, err)
assert.NotNil(t, createResult)
t.Logf("Created collection: %s", createResult.CollectionID)
// 2. Check Exists
existsResult, err := kb.API.CollectionExists(ctx, testCollectionID)
assert.NoError(t, err)
assert.True(t, existsResult.Exists)
t.Logf("Collection exists: %v", existsResult.Exists)
// 3. Get Collection
collection, err := kb.API.GetCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.Equal(t, testCollectionID, collection["id"])
assert.Equal(t, testCollectionID, collection["collection_id"])
assert.Equal(t, "Integration Test Collection", collection["name"])
t.Logf("Retrieved collection: %s", collection["name"])
// 4. Update Metadata
updateParams := &api.UpdateMetadataParams{
Metadata: map[string]interface{}{
"name": "Updated Integration Test",
"description": "Updated description",
},
}
updateResult, err := kb.API.UpdateCollectionMetadata(ctx, testCollectionID, updateParams)
assert.NoError(t, err)
assert.NotNil(t, updateResult)
t.Logf("Updated collection metadata")
// 5. Verify Update
updatedCollection, err := kb.API.GetCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.Equal(t, "Updated Integration Test", updatedCollection["name"])
t.Logf("Verified update: %s", updatedCollection["name"])
// 6. List Collections (should include our test collection)
listFilter := &api.ListCollectionsFilter{
Page: 1,
PageSize: 20,
Keywords: "Updated Integration Test",
}
listResult, err := kb.API.ListCollections(ctx, listFilter)
assert.NoError(t, err)
assert.GreaterOrEqual(t, len(listResult.Data), 1)
t.Logf("Found collection in list")
// 7. Remove Collection
removeResult, err := kb.API.RemoveCollection(ctx, testCollectionID)
assert.NoError(t, err)
assert.True(t, removeResult.Removed)
t.Logf("Removed collection: %s", removeResult.CollectionID)
// 8. Verify Removal
existsAfterRemove, err := kb.API.CollectionExists(ctx, testCollectionID)
assert.NoError(t, err)
assert.False(t, existsAfterRemove.Exists)
t.Logf("Verified removal: exists=%v", existsAfterRemove.Exists)
})
}

57
kb/api/consts.go Normal file
View file

@ -0,0 +1,57 @@
package api
import "github.com/yaoapp/gou/model"
// Collection field definitions
var (
// AvailableCollectionFields defines all available fields for security filtering
AvailableCollectionFields = map[string]bool{
"id": true, "collection_id": true, "name": true, "description": true,
"status": true, "preset": true, "public": true, "share": true, "sort": true, "cover": true,
"document_count": true, "embedding_provider_id": true, "embedding_option_id": true,
"embedding_properties": true, "locale": true, "dimension": true,
"distance_metric": true, "hnsw_m": true, "ef_construction": true,
"ef_search": true, "num_lists": true, "num_probes": true,
"created_at": true, "updated_at": true,
}
// DefaultCollectionFields defines the default compact field list
DefaultCollectionFields = []interface{}{
"id", "collection_id", "name", "description", "status", "preset", "public", "share",
"sort", "cover", "document_count", "embedding_provider_id", "embedding_option_id",
"locale", "dimension", "distance_metric", "created_at", "updated_at",
}
// ValidCollectionSortFields defines valid fields for sorting
ValidCollectionSortFields = map[string]bool{
"created_at": true,
"updated_at": true,
"name": true,
"sort": true,
"document_count": true,
"status": true,
}
)
// Default pagination settings
const (
DefaultPage = 1
DefaultPageSize = 20
MaxPageSize = 100
)
// Default sort settings
const (
DefaultSortField = "created_at"
DefaultSortOrder = "desc"
)
// DefaultSort defines the default sort order for collection queries
var DefaultSort = []model.QueryOrder{
{Column: DefaultSortField, Option: DefaultSortOrder},
}
// Default locale
const (
DefaultLocale = "en"
)

34
kb/api/interfaces.go Normal file
View file

@ -0,0 +1,34 @@
package api
import (
"context"
"github.com/yaoapp/gou/graphrag/types"
kbtypes "github.com/yaoapp/yao/kb/types"
)
// API defines the unified interface for all KB operations
type API interface {
// Collection operations
CreateCollection(ctx context.Context, params *CreateCollectionParams) (*CreateCollectionResult, error)
RemoveCollection(ctx context.Context, collectionID string) (*RemoveCollectionResult, error)
GetCollection(ctx context.Context, collectionID string) (map[string]interface{}, error)
CollectionExists(ctx context.Context, collectionID string) (*CollectionExistsResult, error)
ListCollections(ctx context.Context, filter *ListCollectionsFilter) (*ListCollectionsResult, error)
UpdateCollectionMetadata(ctx context.Context, collectionID string, params *UpdateMetadataParams) (*UpdateMetadataResult, error)
// Document operations (future)
// AddDocument(ctx context.Context, params *AddDocumentParams) (*AddDocumentResult, error)
// RemoveDocument(ctx context.Context, documentID string) (*RemoveDocumentResult, error)
// ...
// Segment operations (future)
// ...
}
// KBInstance holds the KB instance dependencies required by the API
type KBInstance struct {
GraphRag types.GraphRag // GraphRag instance for vector/graph operations
Config *kbtypes.Config // KB configuration
Providers *kbtypes.ProviderConfig // Provider configurations
}

73
kb/api/types.go Normal file
View file

@ -0,0 +1,73 @@
package api
import (
"github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/gou/model"
)
// CreateCollectionParams represents the parameters for creating a collection
type CreateCollectionParams struct {
ID string `json:"id" yaml:"id"`
Metadata map[string]interface{} `json:"metadata,omitempty" yaml:"metadata,omitempty"`
EmbeddingProviderID string `json:"embedding_provider_id" yaml:"embedding_provider_id"`
EmbeddingOptionID string `json:"embedding_option_id" yaml:"embedding_option_id"`
Locale string `json:"locale,omitempty" yaml:"locale,omitempty"`
Config *types.CreateCollectionOptions `json:"config,omitempty" yaml:"config,omitempty"`
AuthScope map[string]interface{} `json:"-" yaml:"-"` // Internal: authentication scope fields
}
// CreateCollectionResult represents the result of creating a collection
type CreateCollectionResult struct {
CollectionID string `json:"collection_id" yaml:"collection_id"`
Message string `json:"message" yaml:"message"`
}
// RemoveCollectionResult represents the result of removing a collection
type RemoveCollectionResult struct {
CollectionID string `json:"collection_id" yaml:"collection_id"`
Removed bool `json:"removed" yaml:"removed"`
DocumentsRemoved int `json:"documents_removed" yaml:"documents_removed"`
Message string `json:"message" yaml:"message"`
}
// CollectionExistsResult represents the result of checking if a collection exists
type CollectionExistsResult struct {
CollectionID string `json:"collection_id" yaml:"collection_id"`
Exists bool `json:"exists" yaml:"exists"`
}
// ListCollectionsFilter represents the filter options for listing collections
type ListCollectionsFilter struct {
Page int `json:"page" yaml:"page"`
PageSize int `json:"pagesize" yaml:"pagesize"`
Keywords string `json:"keywords,omitempty" yaml:"keywords,omitempty"`
Status []string `json:"status,omitempty" yaml:"status,omitempty"`
System *bool `json:"system,omitempty" yaml:"system,omitempty"`
EmbeddingProviderID string `json:"embedding_provider_id,omitempty" yaml:"embedding_provider_id,omitempty"`
Select []interface{} `json:"select,omitempty" yaml:"select,omitempty"`
Sort []model.QueryOrder `json:"sort,omitempty" yaml:"sort,omitempty"`
AuthFilters []model.QueryWhere `json:"-" yaml:"-"` // Internal: authentication filters
}
// ListCollectionsResult represents the result of listing collections
type ListCollectionsResult struct {
Data []map[string]interface{} `json:"data" yaml:"data"`
Next int `json:"next" yaml:"next"`
Prev int `json:"prev" yaml:"prev"`
Page int `json:"page" yaml:"page"`
PageSize int `json:"pagesize" yaml:"pagesize"`
Total int `json:"total" yaml:"total"`
PageCnt int `json:"pagecnt" yaml:"pagecnt"`
}
// UpdateMetadataParams represents the parameters for updating collection metadata
type UpdateMetadataParams struct {
Metadata map[string]interface{} `json:"metadata" yaml:"metadata"`
AuthScope map[string]interface{} `json:"-" yaml:"-"` // Internal: authentication scope fields for update
}
// UpdateMetadataResult represents the result of updating collection metadata
type UpdateMetadataResult struct {
CollectionID string `json:"collection_id" yaml:"collection_id"`
Message string `json:"message" yaml:"message"`
}

View file

@ -10,6 +10,7 @@ import (
"github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/kun/log" "github.com/yaoapp/kun/log"
"github.com/yaoapp/yao/config" "github.com/yaoapp/yao/config"
"github.com/yaoapp/yao/kb/api"
// Register the built-in providers // Register the built-in providers
_ "github.com/yaoapp/yao/kb/providers" _ "github.com/yaoapp/yao/kb/providers"
@ -21,6 +22,9 @@ import (
// Instance is the GraphRag instance // Instance is the GraphRag instance
var Instance types.GraphRag = nil var Instance types.GraphRag = nil
// API is the Knowledge Base API instance
var API api.API = nil
// KnowledgeBase is the Knowledge Base instance // KnowledgeBase is the Knowledge Base instance
type KnowledgeBase struct { type KnowledgeBase struct {
Config *kbtypes.Config // Knowledge Base configuration Config *kbtypes.Config // Knowledge Base configuration
@ -86,6 +90,10 @@ func Load(appConfig config.Config) (*KnowledgeBase, error) {
// Set the instance to the global variable // Set the instance to the global variable
Instance = instance Instance = instance
// Create and set the API instance
API = api.NewAPI(graphRag, &config, providers)
return instance, nil return instance, nil
} }

View file

@ -9,9 +9,9 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/gou/graphrag/types"
"github.com/yaoapp/gou/model" "github.com/yaoapp/gou/model"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/maps" "github.com/yaoapp/kun/maps"
"github.com/yaoapp/yao/kb" "github.com/yaoapp/yao/kb"
kbapi "github.com/yaoapp/yao/kb/api"
"github.com/yaoapp/yao/openapi/oauth/authorized" "github.com/yaoapp/yao/openapi/oauth/authorized"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types" oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
"github.com/yaoapp/yao/openapi/response" "github.com/yaoapp/yao/openapi/response"
@ -20,37 +20,6 @@ import (
// Collection Management Handlers // Collection Management Handlers
// Collection field definitions
var (
// availableCollectionFields defines all available fields for security filtering
availableCollectionFields = map[string]bool{
"id": true, "collection_id": true, "name": true, "description": true,
"status": true, "preset": true, "public": true, "share": true, "sort": true, "cover": true,
"document_count": true, "embedding_provider_id": true, "embedding_option_id": true,
"embedding_properties": true, "locale": true, "dimension": true,
"distance_metric": true, "hnsw_m": true, "ef_construction": true,
"ef_search": true, "num_lists": true, "num_probes": true,
"created_at": true, "updated_at": true,
}
// defaultCollectionFields defines the default compact field list
defaultCollectionFields = []interface{}{
"id", "collection_id", "name", "description", "status", "preset", "public", "share",
"sort", "cover", "document_count", "embedding_provider_id", "embedding_option_id",
"locale", "dimension", "distance_metric", "created_at", "updated_at",
}
// validCollectionSortFields defines valid fields for sorting
validCollectionSortFields = map[string]bool{
"created_at": true,
"updated_at": true,
"name": true,
"sort": true,
"document_count": true,
"status": true,
}
)
// ProviderSettings represents the resolved provider configuration // ProviderSettings represents the resolved provider configuration
type ProviderSettings struct { type ProviderSettings struct {
Dimension int `json:"dimension"` Dimension int `json:"dimension"`
@ -61,6 +30,16 @@ type ProviderSettings struct {
// CreateCollection creates a new collection // CreateCollection creates a new collection
func CreateCollection(c *gin.Context) { func CreateCollection(c *gin.Context) {
// Check if kb.API is available
if kb.API == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Prepare request and database data // Prepare request and database data
req, collectionData, err := PrepareCreateCollection(c) req, collectionData, err := PrepareCreateCollection(c)
if err != nil { if err != nil {
@ -74,12 +53,56 @@ func CreateCollection(c *gin.Context) {
// Attach create scope to the collection data // Attach create scope to the collection data
authInfo := authorized.GetInfo(c) authInfo := authorized.GetInfo(c)
var authScope map[string]interface{}
if authInfo != nil { if authInfo != nil {
collectionData = authInfo.WithCreateScope(collectionData) collectionData = authInfo.WithCreateScope(collectionData)
// Extract auth scope fields
authScope = make(map[string]interface{})
if createdBy, ok := collectionData["__yao_created_by"]; ok {
authScope["__yao_created_by"] = createdBy
}
if updatedBy, ok := collectionData["__yao_updated_by"]; ok {
authScope["__yao_updated_by"] = updatedBy
}
if teamID, ok := collectionData["__yao_team_id"]; ok {
authScope["__yao_team_id"] = teamID
}
} }
// Check if kb.Instance is available // Build API params
if kb.Instance == nil { params := &kbapi.CreateCollectionParams{
ID: req.ID,
Metadata: req.Metadata,
EmbeddingProviderID: req.Config.EmbeddingProviderID,
EmbeddingOptionID: req.Config.EmbeddingOptionID,
Locale: req.Config.Locale,
Config: req.Config.CreateCollectionOptions,
AuthScope: authScope,
}
// Call API to create collection
result, err := kb.API.CreateCollection(c.Request.Context(), params)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
successData := gin.H{
"message": result.Message,
"collection_id": result.CollectionID,
}
response.RespondWithSuccess(c, response.StatusCreated, successData)
}
// RemoveCollection removes an existing collection
func RemoveCollection(c *gin.Context) {
// Check if kb.API is available
if kb.API == nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized", ErrorDescription: "Knowledge base not initialized",
@ -88,68 +111,6 @@ func CreateCollection(c *gin.Context) {
return return
} }
// Get KB config
config, err := kb.GetConfig()
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to get KB config: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// First create database record
_, err = config.CreateCollection(maps.MapStrAny(collectionData))
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to save collection metadata: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Create CollectionConfig for GraphRag
collectionConfig := types.CollectionConfig{
ID: req.ID,
Metadata: req.Metadata,
Config: req.Config.CreateCollectionOptions,
}
// Call the actual CreateCollection method
collectionID, err := kb.Instance.CreateCollection(c.Request.Context(), collectionConfig)
if err != nil {
// Rollback: remove the database record
rollbackErr := config.RemoveCollection(req.ID)
if rollbackErr != nil {
log.Error("Failed to rollback collection database record: %v", rollbackErr)
}
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to create collection: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Update status to active after successful creation and sync to GraphRag
updateErr := UpdateCollectionWithSync(req.ID, maps.MapStrAny{"status": "active"}, config)
if updateErr != nil {
log.Error("Failed to update collection status to active: %v", updateErr)
}
successData := gin.H{
"message": "Collection created successfully",
"collection_id": collectionID,
}
response.RespondWithSuccess(c, response.StatusCreated, successData)
}
// RemoveCollection removes an existing collection
func RemoveCollection(c *gin.Context) {
authInfo := authorized.GetInfo(c) authInfo := authorized.GetInfo(c)
// Get collection ID from URL parameter // Get collection ID from URL parameter
@ -163,16 +124,6 @@ func RemoveCollection(c *gin.Context) {
return return
} }
// Check if kb.Instance is available
if kb.Instance == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Check remove permission // Check remove permission
hasPermission, err := checkCollectionPermission(authInfo, collectionID) hasPermission, err := checkCollectionPermission(authInfo, collectionID)
if err != nil { if err != nil {
@ -194,60 +145,38 @@ func RemoveCollection(c *gin.Context) {
return return
} }
// Call the actual RemoveCollection method // Call API to remove collection
removed, err := kb.Instance.RemoveCollection(c.Request.Context(), collectionID) result, err := kb.API.RemoveCollection(c.Request.Context(), collectionID)
if err != nil { if err != nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Failed to remove collection: " + err.Error(), ErrorDescription: err.Error(),
} }
response.RespondWithError(c, response.StatusInternalServerError, errorResp) response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return return
} }
if !removed {
errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Collection not found or could not be removed",
}
response.RespondWithError(c, response.StatusNotFound, errorResp)
return
}
// Remove collection and all its documents from database after successful GraphRag removal
documentsRemoved := 0
if config, err := kb.GetConfig(); err == nil {
// First, count documents in this collection (for reporting)
if count, err := config.DocumentCount(collectionID); err == nil {
documentsRemoved = count
}
// Remove all documents belonging to this collection
if err := config.RemoveDocumentsByCollectionID(collectionID); err != nil {
log.Error("Failed to remove documents from collection %s: %v", collectionID, err)
} else {
log.Info("Removed %d documents from collection %s", documentsRemoved, collectionID)
}
// Then remove the collection itself
if err := config.RemoveCollection(collectionID); err != nil {
log.Error("Failed to remove collection from database: %v", err)
} else {
log.Info("Successfully removed collection %s and %d documents", collectionID, documentsRemoved)
}
}
successData := gin.H{ successData := gin.H{
"message": "Collection removed successfully", "message": result.Message,
"collection_id": collectionID, "collection_id": result.CollectionID,
"removed": removed, "removed": result.Removed,
"documents_removed": documentsRemoved, "documents_removed": result.DocumentsRemoved,
} }
response.RespondWithSuccess(c, response.StatusOK, successData) response.RespondWithSuccess(c, response.StatusOK, successData)
} }
// CollectionExists checks if a collection exists // CollectionExists checks if a collection exists
func CollectionExists(c *gin.Context) { func CollectionExists(c *gin.Context) {
// Check if kb.API is available
if kb.API == nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Get collection ID from URL parameter // Get collection ID from URL parameter
collectionID := c.Param("collectionID") collectionID := c.Param("collectionID")
if collectionID == "" { if collectionID == "" {
@ -259,8 +188,28 @@ func CollectionExists(c *gin.Context) {
return return
} }
// Check if kb.Instance is available // Call API to check collection existence
if kb.Instance == nil { result, err := kb.API.CollectionExists(c.Request.Context(), collectionID)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
successData := gin.H{
"collection_id": result.CollectionID,
"exists": result.Exists,
}
response.RespondWithSuccess(c, response.StatusOK, successData)
}
// GetCollection retrieves a collection by ID
func GetCollection(c *gin.Context) {
// Check if kb.API is available
if kb.API == nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized", ErrorDescription: "Knowledge base not initialized",
@ -269,26 +218,6 @@ func CollectionExists(c *gin.Context) {
return return
} }
// Call the actual CollectionExists method
exists, err := kb.Instance.CollectionExists(c.Request.Context(), collectionID)
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to check collection existence: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
successData := gin.H{
"collection_id": collectionID,
"exists": exists,
}
response.RespondWithSuccess(c, response.StatusOK, successData)
}
// GetCollection retrieves a collection by ID
func GetCollection(c *gin.Context) {
collectionID := c.Param("collectionID") collectionID := c.Param("collectionID")
if collectionID == "" { if collectionID == "" {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
@ -299,21 +228,11 @@ func GetCollection(c *gin.Context) {
return return
} }
// Check if kb.Instance is available // Call API to get collection
if kb.Instance == nil { collection, err := kb.API.GetCollection(c.Request.Context(), collectionID)
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized",
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Use the dedicated GetCollection method
collection, err := kb.Instance.GetCollection(c.Request.Context(), collectionID)
if err != nil { if err != nil {
// Check if it's a "not found" error // Check if it's a "not found" error
if err.Error() == fmt.Sprintf("collection with ID '%s' not found", collectionID) { if err.Error() == "collection not found" || err.Error() == fmt.Sprintf("collection with ID '%s' not found", collectionID) {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrInvalidRequest.Code, Code: response.ErrInvalidRequest.Code,
ErrorDescription: "Collection not found", ErrorDescription: "Collection not found",
@ -324,7 +243,7 @@ func GetCollection(c *gin.Context) {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Failed to get collection: " + err.Error(), ErrorDescription: err.Error(),
} }
response.RespondWithError(c, response.StatusInternalServerError, errorResp) response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return return
@ -336,11 +255,8 @@ func GetCollection(c *gin.Context) {
// ListCollections lists collections with pagination // ListCollections lists collections with pagination
func ListCollections(c *gin.Context) { func ListCollections(c *gin.Context) {
// Get authorized information // Check if kb.API is available
authInfo := authorized.GetInfo(c) if kb.API == nil {
// Check if kb.Instance is available
if kb.Instance == nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized", ErrorDescription: "Knowledge base not initialized",
@ -349,6 +265,9 @@ func ListCollections(c *gin.Context) {
return return
} }
// Get authorized information
authInfo := authorized.GetInfo(c)
// Parse pagination parameters // Parse pagination parameters
page := 1 page := 1
if pageStr := c.Query("page"); pageStr != "" { if pageStr := c.Query("page"); pageStr != "" {
@ -364,180 +283,104 @@ func ListCollections(c *gin.Context) {
} }
} }
// Get KB config
config, err := kb.GetConfig()
if err != nil {
errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code,
ErrorDescription: "Failed to get KB config: " + err.Error(),
}
response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return
}
// Parse select parameter // Parse select parameter
var selectFields []interface{} var selectFields []interface{}
if selectParam := strings.TrimSpace(c.Query("select")); selectParam != "" { if selectParam := strings.TrimSpace(c.Query("select")); selectParam != "" {
requestedFields := strings.Split(selectParam, ",") requestedFields := strings.Split(selectParam, ",")
for _, field := range requestedFields { for _, field := range requestedFields {
field = strings.TrimSpace(field) field = strings.TrimSpace(field)
if field != "" && availableCollectionFields[field] { if field != "" && kbapi.AvailableCollectionFields[field] {
selectFields = append(selectFields, field) selectFields = append(selectFields, field)
} }
} }
// If no valid fields found, use default }
if len(selectFields) == 0 {
selectFields = defaultCollectionFields // Parse sort parameter
var orders []model.QueryOrder
if sortParam := strings.TrimSpace(c.Query("sort")); sortParam != "" {
sortItems := strings.Split(sortParam, ",")
for _, sortItem := range sortItems {
sortItem = strings.TrimSpace(sortItem)
if sortItem == "" {
continue
}
sortParts := strings.Fields(sortItem)
if len(sortParts) == 0 {
continue
}
sortField := sortParts[0]
sortOrder := "desc"
if len(sortParts) >= 2 {
sortOrder = strings.ToLower(sortParts[1])
}
// Validate sort field and order
if kbapi.ValidCollectionSortFields[sortField] && (sortOrder == "asc" || sortOrder == "desc") {
orders = append(orders, model.QueryOrder{
Column: sortField,
Option: sortOrder,
})
}
} }
} else {
selectFields = defaultCollectionFields
} }
// Build query parameters // Build filter for API
param := model.QueryParam{Select: selectFields} filter := &kbapi.ListCollectionsFilter{
Page: page,
// Add filters PageSize: pagesize,
var wheres []model.QueryWhere Keywords: strings.TrimSpace(c.Query("keywords")),
EmbeddingProviderID: strings.TrimSpace(c.Query("embedding_provider_id")),
// Apply permission-based filtering Select: selectFields,
wheres = append(wheres, AuthFilter(c, authInfo)...) Sort: orders,
AuthFilters: AuthFilter(c, authInfo),
// Filter by keywords (search in name and description)
if keywords := strings.TrimSpace(c.Query("keywords")); keywords != "" {
wheres = append(wheres, model.QueryWhere{
Column: "name",
Value: "%" + keywords + "%",
OP: "like",
})
wheres = append(wheres, model.QueryWhere{
Column: "description",
Value: "%" + keywords + "%",
OP: "like",
Wheres: []model.QueryWhere{},
Method: "orwhere",
})
} }
// Filter by status (support multiple values separated by comma) // Parse status parameter
if statusParam := strings.TrimSpace(c.Query("status")); statusParam != "" { if statusParam := strings.TrimSpace(c.Query("status")); statusParam != "" {
statusList := strings.Split(statusParam, ",") statusList := strings.Split(statusParam, ",")
var statusValues []interface{}
for _, status := range statusList { for _, status := range statusList {
status = strings.TrimSpace(status) status = strings.TrimSpace(status)
if status != "" { if status != "" {
statusValues = append(statusValues, status) filter.Status = append(filter.Status, status)
}
}
if len(statusValues) > 0 {
if len(statusValues) == 1 {
// Single status
wheres = append(wheres, model.QueryWhere{
Column: "status",
Value: statusValues[0],
})
} else {
// Multiple status - use IN clause
wheres = append(wheres, model.QueryWhere{
Column: "status",
Value: statusValues,
OP: "in",
})
} }
} }
} }
// Filter by system flag // Parse system parameter
if systemParam := strings.TrimSpace(c.Query("system")); systemParam != "" { if systemParam := strings.TrimSpace(c.Query("system")); systemParam != "" {
switch systemParam { switch systemParam {
case "true", "1": case "true", "1":
wheres = append(wheres, model.QueryWhere{ systemVal := true
Column: "system", filter.System = &systemVal
Value: true,
})
case "false", "0": case "false", "0":
wheres = append(wheres, model.QueryWhere{ systemVal := false
Column: "system", filter.System = &systemVal
Value: false,
})
} }
} }
// Filter by embedding_provider_id // Call API to list collections
if providerID := strings.TrimSpace(c.Query("embedding_provider_id")); providerID != "" { result, err := kb.API.ListCollections(c.Request.Context(), filter)
wheres = append(wheres, model.QueryWhere{
Column: "embedding_provider_id",
Value: providerID,
})
}
param.Wheres = wheres
// Add ordering
sortParam := strings.TrimSpace(c.Query("sort"))
if sortParam == "" {
sortParam = "created_at desc" // Default sort
}
// Parse sort parameter (format: "field1 direction1,field2 direction2")
var orders []model.QueryOrder
sortItems := strings.Split(sortParam, ",")
for _, sortItem := range sortItems {
sortItem = strings.TrimSpace(sortItem)
if sortItem == "" {
continue
}
// Parse each sort item (format: "field direction")
sortParts := strings.Fields(sortItem)
sortField := "created_at" // Default field
sortOrder := "desc" // Default order
if len(sortParts) >= 1 {
sortField = sortParts[0]
}
if len(sortParts) >= 2 {
sortOrder = strings.ToLower(sortParts[1])
}
// Validate sort field
if !validCollectionSortFields[sortField] {
continue // Skip invalid fields
}
// Validate sort order
if sortOrder != "asc" && sortOrder != "desc" {
sortOrder = "desc" // Default order
}
orders = append(orders, model.QueryOrder{
Column: sortField,
Option: sortOrder,
})
}
// If no valid orders found, use default
if len(orders) == 0 {
orders = []model.QueryOrder{
{Column: "created_at", Option: "desc"},
}
}
param.Orders = orders
// Query collections using KB config
result, err := config.SearchCollections(param, page, pagesize)
if err != nil { if err != nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Failed to search collections: " + err.Error(), ErrorDescription: err.Error(),
} }
response.RespondWithError(c, response.StatusInternalServerError, errorResp) response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return return
} }
c.JSON(http.StatusOK, result) // Return the result directly to maintain backward compatibility
c.JSON(http.StatusOK, gin.H{
"data": result.Data,
"next": result.Next,
"prev": result.Prev,
"page": result.Page,
"pagesize": result.PageSize,
"total": result.Total,
"pagecnt": result.PageCnt,
})
} }
// UpdateCollectionMetadata updates the metadata of an existing collection // UpdateCollectionMetadata updates the metadata of an existing collection
@ -576,8 +419,8 @@ func UpdateCollectionMetadata(c *gin.Context) {
return return
} }
// Check if kb.Instance is available // Check if kb.API is available
if kb.Instance == nil { if kb.API == nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Knowledge base not initialized", ErrorDescription: "Knowledge base not initialized",
@ -608,45 +451,31 @@ func UpdateCollectionMetadata(c *gin.Context) {
return return
} }
// Call the actual UpdateCollectionMetadata method // Build API params
err = kb.Instance.UpdateCollectionMetadata(c.Request.Context(), collectionID, req.Metadata) var authScope map[string]interface{}
if authInfo != nil {
authScope = authInfo.WithUpdateScope(maps.MapStrAny{})
}
params := &kbapi.UpdateMetadataParams{
Metadata: req.Metadata,
AuthScope: authScope,
}
// Call API to update collection metadata
result, err := kb.API.UpdateCollectionMetadata(c.Request.Context(), collectionID, params)
if err != nil { if err != nil {
errorResp := &response.ErrorResponse{ errorResp := &response.ErrorResponse{
Code: response.ErrServerError.Code, Code: response.ErrServerError.Code,
ErrorDescription: "Failed to update collection metadata: " + err.Error(), ErrorDescription: err.Error(),
} }
response.RespondWithError(c, response.StatusInternalServerError, errorResp) response.RespondWithError(c, response.StatusInternalServerError, errorResp)
return return
} }
// Update collection metadata in database after successful GraphRag update
// Note: Only update database here, don't sync to GraphRag again (already done above)
if config, err := kb.GetConfig(); err == nil {
// Prepare update data from metadata
updateData := maps.MapStrAny{}
if name, ok := req.Metadata["name"]; ok {
updateData["name"] = name
}
if description, ok := req.Metadata["description"]; ok {
updateData["description"] = description
}
if status, ok := req.Metadata["status"]; ok {
updateData["status"] = status
}
// Update __yao_updated_by
updateData = authInfo.WithUpdateScope(updateData)
if len(updateData) > 0 {
// Only update database, don't sync to GraphRag again to avoid duplicate updates
if err := config.UpdateCollection(collectionID, updateData); err != nil {
log.Error("Failed to update collection in database: %v", err)
}
}
}
successData := gin.H{ successData := gin.H{
"message": "Collection metadata updated successfully", "message": result.Message,
"collection_id": collectionID, "collection_id": result.CollectionID,
} }
response.RespondWithSuccess(c, response.StatusOK, successData) response.RespondWithSuccess(c, response.StatusOK, successData)
} }

View file

@ -0,0 +1,373 @@
package kb
import (
"encoding/json"
"github.com/yaoapp/gou/model"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/exception"
"github.com/yaoapp/kun/log"
"github.com/yaoapp/kun/maps"
"github.com/yaoapp/yao/kb"
kbapi "github.com/yaoapp/yao/kb/api"
"github.com/yaoapp/yao/openapi/oauth/authorized"
oauthtypes "github.com/yaoapp/yao/openapi/oauth/types"
)
// ProcessCreateCollection creates a new collection via Yao process
// Process: kb.collection.Create
//
// Args[0]: params (map) - Collection creation parameters
//
// {
// "id": "collection_id",
// "metadata": {
// "name": "Collection Name",
// "description": "Description"
// },
// "embedding_provider_id": "__yao.openai",
// "embedding_option_id": "text-embedding-3-small",
// "locale": "en",
// "config": {
// "distance": "cosine",
// "index_type": "hnsw",
// "m": 16,
// "ef_construction": 200,
// "ef_search": 64
// }
// }
//
// Returns: map with collection_id and message
func ProcessCreateCollection(process *process.Process) interface{} {
process.ValidateArgNums(1)
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
// Get authorized info from process
authInfo := authorized.ProcessAuthInfo(process)
// Parse parameters using JSON for type safety
paramsJSON, err := json.Marshal(process.Args[0])
if err != nil {
exception.New("Failed to encode parameters: "+err.Error(), 400).Throw()
}
var params kbapi.CreateCollectionParams
if err := json.Unmarshal(paramsJSON, &params); err != nil {
exception.New("Failed to decode parameters: "+err.Error(), 400).Throw()
}
// Apply auth scope from authorized info
if authInfo != nil {
authScope := authInfo.WithCreateScope(maps.MapStrAny{})
params.AuthScope = authScope
}
// Call API
result, err := kb.API.CreateCollection(process.Context, &params)
if err != nil {
log.Error("Failed to create collection: %v", err)
exception.New(err.Error(), 500).Throw()
}
return maps.MapStrAny{
"collection_id": result.CollectionID,
"message": result.Message,
}
}
// ProcessRemoveCollection removes a collection via Yao process
// Process: kb.collection.Remove
//
// Args[0]: collection_id (string) - Collection ID to remove
//
// Returns: map with collection_id, removed status, documents_removed count, and message
func ProcessRemoveCollection(process *process.Process) interface{} {
process.ValidateArgNums(1)
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
// Get authorized info from process
authInfo := authorized.ProcessAuthInfo(process)
collectionID := process.ArgsString(0)
if collectionID == "" {
exception.New("Collection ID is required", 400).Throw()
}
// Check remove permission
hasPermission, err := checkCollectionPermission(authInfo, collectionID)
if err != nil {
exception.New(err.Error(), 403).Throw()
}
if !hasPermission {
exception.New("Forbidden: No permission to remove collection", 403).Throw()
}
result, err := kb.API.RemoveCollection(process.Context, collectionID)
if err != nil {
log.Error("Failed to remove collection: %v", err)
exception.New(err.Error(), 500).Throw()
}
return maps.MapStrAny{
"collection_id": result.CollectionID,
"removed": result.Removed,
"documents_removed": result.DocumentsRemoved,
"message": result.Message,
}
}
// ProcessGetCollection retrieves a collection by ID via Yao process
// Process: kb.collection.Get
//
// Args[0]: collection_id (string) - Collection ID to retrieve
//
// Returns: map containing collection details
func ProcessGetCollection(process *process.Process) interface{} {
process.ValidateArgNums(1)
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
collectionID := process.ArgsString(0)
if collectionID == "" {
exception.New("Collection ID is required", 400).Throw()
}
collection, err := kb.API.GetCollection(process.Context, collectionID)
if err != nil {
log.Error("Failed to get collection: %v", err)
exception.New(err.Error(), 500).Throw()
}
return collection
}
// ProcessCollectionExists checks if a collection exists via Yao process
// Process: kb.collection.Exists
//
// Args[0]: collection_id (string) - Collection ID to check
//
// Returns: map with collection_id and exists status
func ProcessCollectionExists(process *process.Process) interface{} {
process.ValidateArgNums(1)
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
collectionID := process.ArgsString(0)
if collectionID == "" {
exception.New("Collection ID is required", 400).Throw()
}
result, err := kb.API.CollectionExists(process.Context, collectionID)
if err != nil {
log.Error("Failed to check collection existence: %v", err)
exception.New(err.Error(), 500).Throw()
}
return maps.MapStrAny{
"collection_id": result.CollectionID,
"exists": result.Exists,
}
}
// ProcessListCollections lists collections with pagination via Yao process
// Process: kb.collection.List
//
// Args[0]: filter (map) - Optional filter parameters
//
// {
// "page": 1,
// "pagesize": 20,
// "keywords": "search term",
// "status": ["active"],
// "embedding_provider_id": "__yao.openai",
// "system": false,
// "select": ["id", "name", "status"],
// "sort": [{"column": "created_at", "option": "desc"}]
// }
//
// Returns: map with data array and pagination info
func ProcessListCollections(process *process.Process) interface{} {
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
// Get authorized info from process
authInfo := authorized.ProcessAuthInfo(process)
// Default filter
filter := &kbapi.ListCollectionsFilter{
Page: kbapi.DefaultPage,
PageSize: kbapi.DefaultPageSize,
}
// Parse filter parameters using JSON (optional)
if process.NumOfArgs() > 0 {
filterJSON, err := json.Marshal(process.Args[0])
if err != nil {
exception.New("Failed to encode filter: "+err.Error(), 400).Throw()
}
if err := json.Unmarshal(filterJSON, filter); err != nil {
exception.New("Failed to decode filter: "+err.Error(), 400).Throw()
}
}
// Apply auth filters from authorized info
if authInfo != nil {
filter.AuthFilters = processAuthFilter(authInfo)
}
result, err := kb.API.ListCollections(process.Context, filter)
if err != nil {
log.Error("Failed to list collections: %v", err)
exception.New(err.Error(), 500).Throw()
}
return maps.MapStrAny{
"data": result.Data,
"next": result.Next,
"prev": result.Prev,
"page": result.Page,
"pagesize": result.PageSize,
"total": result.Total,
"pagecnt": result.PageCnt,
}
}
// ProcessUpdateCollectionMetadata updates collection metadata via Yao process
// Process: kb.collection.UpdateMetadata
//
// Args[0]: collection_id (string) - Collection ID
// Args[1]: params (map) - Update parameters
//
// {
// "metadata": {
// "name": "New Name",
// "description": "New Description"
// }
// }
//
// Returns: map with collection_id and message
func ProcessUpdateCollectionMetadata(process *process.Process) interface{} {
process.ValidateArgNums(2)
if kb.API == nil {
exception.New("Knowledge base not initialized", 500).Throw()
}
// Get authorized info from process
authInfo := authorized.ProcessAuthInfo(process)
collectionID := process.ArgsString(0)
if collectionID == "" {
exception.New("Collection ID is required", 400).Throw()
}
// Check update permission
hasPermission, err := checkCollectionPermission(authInfo, collectionID)
if err != nil {
exception.New(err.Error(), 403).Throw()
}
if !hasPermission {
exception.New("Forbidden: No permission to update collection", 403).Throw()
}
// Parse parameters using JSON
paramsJSON, err := json.Marshal(process.Args[1])
if err != nil {
exception.New("Failed to encode parameters: "+err.Error(), 400).Throw()
}
var params kbapi.UpdateMetadataParams
if err := json.Unmarshal(paramsJSON, &params); err != nil {
exception.New("Failed to decode parameters: "+err.Error(), 400).Throw()
}
if len(params.Metadata) == 0 {
exception.New("Metadata is required and cannot be empty", 400).Throw()
}
// Apply auth scope from authorized info
if authInfo != nil {
authScope := authInfo.WithUpdateScope(maps.MapStrAny{})
params.AuthScope = authScope
}
result, err := kb.API.UpdateCollectionMetadata(process.Context, collectionID, &params)
if err != nil {
log.Error("Failed to update collection metadata: %v", err)
exception.New(err.Error(), 500).Throw()
}
return maps.MapStrAny{
"collection_id": result.CollectionID,
"message": result.Message,
}
}
// Helper functions for Process handlers
// processAuthFilter applies permission-based filtering to query wheres for process handlers
// This function builds where clauses based on the user's authorization constraints
func processAuthFilter(authInfo *oauthtypes.AuthorizedInfo) []model.QueryWhere {
if authInfo == nil {
return []model.QueryWhere{}
}
var wheres []model.QueryWhere
scope := authInfo.AccessScope()
// Team only - User can access:
// 1. Public records (public = true)
// 2. Records in their team where:
// - They created the record (__yao_created_by matches)
// - OR the record is shared with team (share = "team")
if authInfo.Constraints.TeamOnly && authInfo.TeamID != "" && authInfo.UserID != "" {
wheres = append(wheres, model.QueryWhere{
Wheres: []model.QueryWhere{
{Column: "public", Value: true, Method: "orwhere"},
{Wheres: []model.QueryWhere{
{Column: "__yao_team_id", Value: scope.TeamID},
{Wheres: []model.QueryWhere{
{Column: "__yao_created_by", Value: scope.CreatedBy},
{Column: "share", Value: "team", Method: "orwhere"},
}},
}, Method: "orwhere"},
},
})
return wheres
}
// Owner only - User can access:
// 1. Public records (public = true)
// 2. Records they created where:
// - __yao_team_id is null (not team records)
// - __yao_created_by matches their user ID
if authInfo.Constraints.OwnerOnly && authInfo.UserID != "" {
wheres = append(wheres, model.QueryWhere{
Wheres: []model.QueryWhere{
{Column: "public", Value: true, Method: "orwhere"},
{Wheres: []model.QueryWhere{
{Column: "__yao_team_id", OP: "null"},
{Column: "__yao_created_by", Value: scope.CreatedBy},
}, Method: "orwhere"},
},
})
return wheres
}
return wheres
}

View file

@ -11,6 +11,15 @@ import (
func init() { func init() {
// Register kb process handlers // Register kb process handlers
process.RegisterGroup("kb", map[string]process.Handler{ process.RegisterGroup("kb", map[string]process.Handler{
// Collection processes
"collection.create": ProcessCreateCollection,
"collection.remove": ProcessRemoveCollection,
"collection.get": ProcessGetCollection,
"collection.exists": ProcessCollectionExists,
"collection.list": ProcessListCollections,
"collection.updatemetadata": ProcessUpdateCollectionMetadata,
// Document processes
"documents.addfile": ProcessAddFile, "documents.addfile": ProcessAddFile,
"documents.addtext": ProcessAddText, "documents.addtext": ProcessAddText,
"documents.addurl": ProcessAddURL, "documents.addurl": ProcessAddURL,

View file

@ -8,9 +8,37 @@ import (
// ProcessAuthInfo extracts authorized information from the process // ProcessAuthInfo extracts authorized information from the process
func ProcessAuthInfo(p *process.Process) *types.AuthorizedInfo { func ProcessAuthInfo(p *process.Process) *types.AuthorizedInfo {
// TODO: Implement this function if p == nil {
// Get authorized information from the process context return nil
info := &types.AuthorizedInfo{} }
// Get authorized info from process
processAuth := p.GetAuthorized()
if processAuth == nil {
return nil
}
// Convert process.AuthorizedInfo to types.AuthorizedInfo
info := &types.AuthorizedInfo{
Subject: processAuth.Subject,
ClientID: processAuth.ClientID,
UserID: processAuth.UserID,
Scope: processAuth.Scope,
TeamID: processAuth.TeamID,
TenantID: processAuth.TenantID,
SessionID: processAuth.SessionID,
RememberMe: processAuth.RememberMe,
}
// Convert constraints
info.Constraints = types.DataConstraints{
OwnerOnly: processAuth.Constraints.OwnerOnly,
CreatorOnly: processAuth.Constraints.CreatorOnly,
EditorOnly: processAuth.Constraints.EditorOnly,
TeamOnly: processAuth.Constraints.TeamOnly,
Extra: processAuth.Constraints.Extra,
}
return info return info
} }

View file

@ -0,0 +1,88 @@
package authorized
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/process"
)
func TestProcessAuthInfo(t *testing.T) {
t.Run("WithNilProcess", func(t *testing.T) {
result := ProcessAuthInfo(nil)
assert.Nil(t, result)
})
t.Run("WithProcessNoAuth", func(t *testing.T) {
p := &process.Process{}
result := ProcessAuthInfo(p)
// GetAuthorized returns empty struct instead of nil, so ProcessAuthInfo will return an empty AuthorizedInfo
require.NotNil(t, result)
assert.Empty(t, result.UserID)
assert.Empty(t, result.TeamID)
assert.Empty(t, result.Subject)
})
t.Run("WithProcessWithAuth", func(t *testing.T) {
p := &process.Process{
Authorized: &process.AuthorizedInfo{
Subject: "user123",
ClientID: "client456",
UserID: "u789",
Scope: "read write",
TeamID: "t123",
TenantID: "tenant456",
SessionID: "session789",
RememberMe: true,
Constraints: process.DataConstraints{
OwnerOnly: true,
CreatorOnly: false,
EditorOnly: false,
TeamOnly: true,
Extra: map[string]interface{}{
"department": "engineering",
},
},
},
}
result := ProcessAuthInfo(p)
require.NotNil(t, result)
assert.Equal(t, "user123", result.Subject)
assert.Equal(t, "client456", result.ClientID)
assert.Equal(t, "u789", result.UserID)
assert.Equal(t, "read write", result.Scope)
assert.Equal(t, "t123", result.TeamID)
assert.Equal(t, "tenant456", result.TenantID)
assert.Equal(t, "session789", result.SessionID)
assert.True(t, result.RememberMe)
assert.True(t, result.Constraints.OwnerOnly)
assert.False(t, result.Constraints.CreatorOnly)
assert.False(t, result.Constraints.EditorOnly)
assert.True(t, result.Constraints.TeamOnly)
assert.Equal(t, "engineering", result.Constraints.Extra["department"])
})
t.Run("WithPartialData", func(t *testing.T) {
p := &process.Process{
Authorized: &process.AuthorizedInfo{
UserID: "u123",
TeamID: "t456",
Constraints: process.DataConstraints{
TeamOnly: true,
},
},
}
result := ProcessAuthInfo(p)
require.NotNil(t, result)
assert.Equal(t, "u123", result.UserID)
assert.Equal(t, "t456", result.TeamID)
assert.True(t, result.Constraints.TeamOnly)
assert.False(t, result.Constraints.OwnerOnly)
})
}

View file

@ -462,3 +462,152 @@ func TestCopyScopesIntegration(t *testing.T) {
assert.Equal(t, "tenant789", updateResult["__yao_tenant_id"]) assert.Equal(t, "tenant789", updateResult["__yao_tenant_id"])
assert.Nil(t, updateResult["__yao_created_by"]) // Should not be copied assert.Nil(t, updateResult["__yao_created_by"]) // Should not be copied
} }
func TestAuthorizedToMap(t *testing.T) {
tests := []struct {
name string
auth *AuthorizedInfo
expected map[string]interface{}
}{
{
name: "Full AuthorizedInfo",
auth: &AuthorizedInfo{
Subject: "user123",
ClientID: "client456",
Scope: "read write",
SessionID: "session789",
UserID: "user123",
TeamID: "team456",
TenantID: "tenant789",
RememberMe: true,
Constraints: DataConstraints{
OwnerOnly: true,
CreatorOnly: false,
EditorOnly: false,
TeamOnly: true,
Extra: map[string]interface{}{
"department": "engineering",
},
},
},
expected: map[string]interface{}{
"sub": "user123",
"client_id": "client456",
"scope": "read write",
"session_id": "session789",
"user_id": "user123",
"team_id": "team456",
"tenant_id": "tenant789",
"remember_me": true,
"constraints": map[string]interface{}{
"owner_only": true,
"team_only": true,
"extra": map[string]interface{}{
"department": "engineering",
},
},
},
},
{
name: "Partial AuthorizedInfo",
auth: &AuthorizedInfo{
UserID: "user123",
TeamID: "team456",
},
expected: map[string]interface{}{
"user_id": "user123",
"team_id": "team456",
},
},
{
name: "AuthorizedInfo with only constraints",
auth: &AuthorizedInfo{
UserID: "user123",
Constraints: DataConstraints{
TeamOnly: true,
Extra: map[string]interface{}{
"region": "us-west",
},
},
},
expected: map[string]interface{}{
"user_id": "user123",
"constraints": map[string]interface{}{
"team_only": true,
"extra": map[string]interface{}{
"region": "us-west",
},
},
},
},
{
name: "Nil AuthorizedInfo",
auth: nil,
expected: nil,
},
{
name: "Empty AuthorizedInfo",
auth: &AuthorizedInfo{},
expected: map[string]interface{}{},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
result := tt.auth.AuthorizedToMap()
if tt.expected == nil {
assert.Nil(t, result)
return
}
assert.NotNil(t, result)
// Check all expected keys
for key, expectedValue := range tt.expected {
actualValue, ok := result[key]
if !ok {
t.Errorf("Key %s not found in result", key)
continue
}
// Special handling for nested maps (constraints)
if key == "constraints" {
expectedConstraints, _ := expectedValue.(map[string]interface{})
actualConstraints, ok := actualValue.(map[string]interface{})
assert.True(t, ok, "constraints should be map[string]interface{}")
for cKey, cExpectedValue := range expectedConstraints {
cActualValue, ok := actualConstraints[cKey]
assert.True(t, ok, "Constraint key %s should exist", cKey)
// Special handling for nested extra map
if cKey == "extra" {
expectedExtra, _ := cExpectedValue.(map[string]interface{})
actualExtra, ok := cActualValue.(map[string]interface{})
assert.True(t, ok, "extra should be map[string]interface{}")
for eKey, eExpectedValue := range expectedExtra {
eActualValue, ok := actualExtra[eKey]
assert.True(t, ok, "Extra key %s should exist", eKey)
assert.Equal(t, eExpectedValue, eActualValue)
}
} else {
assert.Equal(t, cExpectedValue, cActualValue)
}
}
} else {
assert.Equal(t, expectedValue, actualValue)
}
}
// Check no unexpected keys (except for empty maps)
if len(tt.expected) > 0 {
for key := range result {
_, ok := tt.expected[key]
assert.True(t, ok, "Unexpected key %s in result", key)
}
}
})
}
}

View file

@ -627,6 +627,64 @@ type AuthorizedInfo struct {
Constraints DataConstraints `json:"constraints,omitempty"` Constraints DataConstraints `json:"constraints,omitempty"`
} }
// AuthorizedToMap converts AuthorizedInfo to map[string]interface{}
// This is useful for passing authorized information to runtime bridges (e.g., V8)
func (auth *AuthorizedInfo) AuthorizedToMap() map[string]interface{} {
if auth == nil {
return nil
}
result := make(map[string]interface{})
if auth.Subject != "" {
result["sub"] = auth.Subject
}
if auth.ClientID != "" {
result["client_id"] = auth.ClientID
}
if auth.Scope != "" {
result["scope"] = auth.Scope
}
if auth.SessionID != "" {
result["session_id"] = auth.SessionID
}
if auth.UserID != "" {
result["user_id"] = auth.UserID
}
if auth.TeamID != "" {
result["team_id"] = auth.TeamID
}
if auth.TenantID != "" {
result["tenant_id"] = auth.TenantID
}
if auth.RememberMe {
result["remember_me"] = auth.RememberMe
}
// Add constraints if any are set
if auth.Constraints.OwnerOnly || auth.Constraints.CreatorOnly || auth.Constraints.EditorOnly || auth.Constraints.TeamOnly || len(auth.Constraints.Extra) > 0 {
constraints := make(map[string]interface{})
if auth.Constraints.OwnerOnly {
constraints["owner_only"] = true
}
if auth.Constraints.CreatorOnly {
constraints["creator_only"] = true
}
if auth.Constraints.EditorOnly {
constraints["editor_only"] = true
}
if auth.Constraints.TeamOnly {
constraints["team_only"] = true
}
if len(auth.Constraints.Extra) > 0 {
constraints["extra"] = auth.Constraints.Extra
}
result["constraints"] = constraints
}
return result
}
// JWTClaims represents JWT-specific claims structure // JWTClaims represents JWT-specific claims structure
type JWTClaims struct { type JWTClaims struct {
jwt.StandardClaims jwt.StandardClaims

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,458 @@
package openapi_test
import (
"context"
"fmt"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/yaoapp/gou/process"
"github.com/yaoapp/kun/maps"
"github.com/yaoapp/yao/kb"
"github.com/yaoapp/yao/openapi/tests/testutils"
)
// TestProcessCreateCollection tests the kb.collection.Create process
func TestProcessCreateCollection(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean()
// Ensure KB is initialized
if kb.API == nil {
t.Skip("Knowledge base not initialized - skipping test")
}
testCollectionID := fmt.Sprintf("test_process_create_%d", time.Now().UnixNano())
t.Run("CreateCollectionWithAuth", func(t *testing.T) {
// Register collection for cleanup
testutils.RegisterTestCollection(testCollectionID)
// Create process with authorized info
p := process.New("kb.collection.Create").
WithContext(context.Background()).
WithAuthorized(&process.AuthorizedInfo{
UserID: "test_user_123",
TeamID: "test_team_456",
Subject: "user@example.com",
ClientID: "test_client",
Scope: "openid profile",
Constraints: process.DataConstraints{
TeamOnly: true,
},
})
// Prepare parameters - use map for Process API
params := map[string]interface{}{
"id": testCollectionID,
"metadata": map[string]interface{}{
"name": "Process Test Collection",
"description": "Created via Process API with auth",
},
"embedding_provider_id": "__yao.openai",
"embedding_option_id": "text-embedding-3-small",
"locale": "en",
"config": map[string]interface{}{
"index_type": "hnsw",
"distance": "cosine",
},
}
p.Args = []interface{}{params}
// Execute process
result, err := p.Exec()
require.NoError(t, err)
require.NotNil(t, result)
// Verify result
resultMap, ok := result.(maps.MapStrAny)
require.True(t, ok, "Result should be a maps.MapStrAny")
assert.Equal(t, testCollectionID, resultMap["collection_id"])
assert.Contains(t, resultMap, "message")
t.Logf("✓ Successfully created collection via process: %s", testCollectionID)
// Verify auth scope was applied by checking the collection
collection, err := kb.API.GetCollection(context.Background(), testCollectionID)
require.NoError(t, err)
assert.NotNil(t, collection)
// Check if auth fields were set
if createdBy, ok := collection["__yao_created_by"]; ok {
assert.Equal(t, "test_user_123", createdBy)
t.Logf("✓ Auth scope applied: __yao_created_by = %v", createdBy)
}
if teamID, ok := collection["__yao_team_id"]; ok {
assert.Equal(t, "test_team_456", teamID)
t.Logf("✓ Auth scope applied: __yao_team_id = %v", teamID)
}
})
t.Run("CreateCollectionWithoutAuth", func(t *testing.T) {
testCollectionID2 := fmt.Sprintf("test_process_create_noauth_%d", time.Now().UnixNano())
testutils.RegisterTestCollection(testCollectionID2)
// Create process without authorized info
p := process.New("kb.collection.Create").
WithContext(context.Background())
params := map[string]interface{}{
"id": testCollectionID2,
"metadata": map[string]interface{}{
"name": "Process Test Collection No Auth",
"description": "Created via Process API without auth",
},
"embedding_provider_id": "__yao.openai",
"embedding_option_id": "text-embedding-3-small",
"locale": "en",
"config": map[string]interface{}{
"index_type": "hnsw",
"distance": "cosine",
},
}
p.Args = []interface{}{params}
// Execute process
result, err := p.Exec()
require.NoError(t, err)
require.NotNil(t, result)
resultMap, ok := result.(maps.MapStrAny)
require.True(t, ok)
assert.Equal(t, testCollectionID2, resultMap["collection_id"])
t.Logf("✓ Successfully created collection without auth: %s", testCollectionID2)
})
t.Run("CreateCollectionInvalidParams", func(t *testing.T) {
// Create process with invalid parameters
p := process.New("kb.collection.Create").
WithContext(context.Background())
// Missing required fields
params := map[string]interface{}{
"metadata": map[string]interface{}{
"name": "Invalid Collection",
},
// Missing id, embedding_provider_id, etc.
}
p.Args = []interface{}{params}
// Execute should throw exception or return error
defer func() {
if r := recover(); r != nil {
t.Logf("✓ Correctly rejected invalid parameters via panic: %v", r)
return
}
}()
result, err := p.Exec()
if err != nil {
t.Logf("✓ Correctly rejected invalid parameters via error: %v", err)
return
}
if result == nil {
t.Log("✓ Correctly rejected invalid parameters (nil result)")
return
}
t.Error("Should have thrown exception or returned error for invalid parameters")
})
}
// TestProcessListCollections tests the kb.collection.List process
func TestProcessListCollections(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean()
if kb.API == nil {
t.Skip("Knowledge base not initialized - skipping test")
}
t.Run("ListCollectionsWithAuth", func(t *testing.T) {
// Test listing with auth filters
p := process.New("kb.collection.List").
WithContext(context.Background()).
WithAuthorized(&process.AuthorizedInfo{
UserID: "test_user_789",
TeamID: "test_team_789",
Constraints: process.DataConstraints{
TeamOnly: true,
},
})
filter := map[string]interface{}{
"page": 1,
"pagesize": 20,
}
p.Args = []interface{}{filter}
// Execute process
result, err := p.Exec()
require.NoError(t, err)
require.NotNil(t, result)
resultMap, ok := result.(maps.MapStrAny)
require.True(t, ok, "Result should be a map")
assert.Contains(t, resultMap, "data")
assert.Contains(t, resultMap, "page")
assert.Contains(t, resultMap, "pagesize")
assert.Contains(t, resultMap, "total")
t.Logf("✓ Retrieved collections with auth filters")
})
t.Run("ListCollectionsNoFilter", func(t *testing.T) {
// Test listing without filter (should use defaults)
p := process.New("kb.collection.List").
WithContext(context.Background())
// No arguments - should use default filter
p.Args = []interface{}{}
result, err := p.Exec()
require.NoError(t, err)
require.NotNil(t, result)
resultMap, ok := result.(maps.MapStrAny)
require.True(t, ok)
assert.Contains(t, resultMap, "data")
assert.Equal(t, 1, resultMap["page"])
assert.Equal(t, 20, resultMap["pagesize"])
t.Logf("✓ Retrieved collections with default filter")
})
}
// TestProcessGetCollection tests the kb.collection.Get process
func TestProcessGetCollection(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean()
if kb.API == nil {
t.Skip("Knowledge base not initialized - skipping test")
}
testCollectionID := fmt.Sprintf("test_process_get_%d", time.Now().UnixNano())
testutils.RegisterTestCollection(testCollectionID)
t.Run("GetCollectionNotFound", func(t *testing.T) {
p := process.New("kb.collection.Get").
WithContext(context.Background())
p.Args = []interface{}{"nonexistent_collection_id"}
// Execute should throw exception or return error
defer func() {
if r := recover(); r != nil {
t.Logf("✓ Correctly rejected nonexistent collection via panic: %v", r)
return
}
}()
result, err := p.Exec()
if err != nil {
t.Logf("✓ Correctly rejected nonexistent collection via error: %v", err)
return
}
if result == nil {
t.Log("✓ Correctly rejected nonexistent collection (nil result)")
return
}
t.Error("Should have thrown exception or returned error for nonexistent collection")
})
}
// TestProcessCollectionExists tests the kb.collection.Exists process
func TestProcessCollectionExists(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean()
if kb.API == nil {
t.Skip("Knowledge base not initialized - skipping test")
}
testCollectionID := fmt.Sprintf("test_process_exists_%d", time.Now().UnixNano())
testutils.RegisterTestCollection(testCollectionID)
t.Run("CollectionExistsBeforeCreation", func(t *testing.T) {
p := process.New("kb.collection.Exists").
WithContext(context.Background())
p.Args = []interface{}{testCollectionID}
result, err := p.Exec()
require.NoError(t, err)
require.NotNil(t, result)
resultMap, ok := result.(maps.MapStrAny)
require.True(t, ok)
assert.Equal(t, testCollectionID, resultMap["collection_id"])
assert.Equal(t, false, resultMap["exists"])
t.Logf("✓ Correctly reported collection does not exist")
})
}
// TestProcessCollectionIntegration tests the full collection lifecycle via Process API
func TestProcessCollectionIntegration(t *testing.T) {
testutils.Prepare(t)
defer testutils.Clean()
if kb.API == nil {
t.Skip("Knowledge base not initialized - skipping test")
}
testCollectionID := fmt.Sprintf("test_process_integration_%d", time.Now().UnixNano())
testutils.RegisterTestCollection(testCollectionID)
ctx := context.Background()
authInfo := &process.AuthorizedInfo{
UserID: "integration_user",
TeamID: "integration_team",
Subject: "integration@example.com",
ClientID: "integration_client",
Scope: "openid profile",
Constraints: process.DataConstraints{
TeamOnly: true,
},
}
t.Run("FullLifecycleViaProcess", func(t *testing.T) {
// Step 1: Check collection doesn't exist
p1 := process.New("kb.collection.Exists").WithContext(ctx)
p1.Args = []interface{}{testCollectionID}
result1, err := p1.Exec()
require.NoError(t, err)
existsResult := result1.(maps.MapStrAny)
assert.Equal(t, false, existsResult["exists"])
t.Logf("✓ Step 1: Confirmed collection doesn't exist")
// Step 2: Create collection
p2 := process.New("kb.collection.Create").WithContext(ctx).WithAuthorized(authInfo)
p2.Args = []interface{}{
map[string]interface{}{
"id": testCollectionID,
"metadata": map[string]interface{}{
"name": "Integration Test Collection",
"description": "Full lifecycle test",
},
"embedding_provider_id": "__yao.openai",
"embedding_option_id": "text-embedding-3-small",
"locale": "en",
"config": map[string]interface{}{
"index_type": "hnsw",
"distance": "cosine",
},
},
}
result2, err := p2.Exec()
require.NoError(t, err)
createResult := result2.(maps.MapStrAny)
assert.Equal(t, testCollectionID, createResult["collection_id"])
t.Logf("✓ Step 2: Created collection")
// Step 3: Verify collection exists
p3 := process.New("kb.collection.Exists").WithContext(ctx)
p3.Args = []interface{}{testCollectionID}
result3, err := p3.Exec()
require.NoError(t, err)
existsResult2 := result3.(maps.MapStrAny)
assert.Equal(t, true, existsResult2["exists"])
t.Logf("✓ Step 3: Confirmed collection exists")
// Step 4: Get collection
p4 := process.New("kb.collection.Get").WithContext(ctx)
p4.Args = []interface{}{testCollectionID}
result4, err := p4.Exec()
require.NoError(t, err)
// GetCollection returns map[string]interface{}
var getResult map[string]interface{}
if mapStrAny, ok := result4.(maps.MapStrAny); ok {
getResult = mapStrAny
} else if m, ok := result4.(map[string]interface{}); ok {
getResult = m
} else {
t.Fatalf("Unexpected result type: %T", result4)
}
assert.Equal(t, "Integration Test Collection", getResult["name"])
t.Logf("✓ Step 4: Retrieved collection details")
// Step 5: Update metadata
p5 := process.New("kb.collection.UpdateMetadata").WithContext(ctx).WithAuthorized(authInfo)
p5.Args = []interface{}{
testCollectionID,
map[string]interface{}{
"metadata": map[string]interface{}{
"name": "Updated Integration Collection",
"description": "Updated via process",
},
},
}
result5, err := p5.Exec()
require.NoError(t, err)
updateResult := result5.(maps.MapStrAny)
assert.Equal(t, testCollectionID, updateResult["collection_id"])
t.Logf("✓ Step 5: Updated collection metadata")
// Step 6: Verify update
p6 := process.New("kb.collection.Get").WithContext(ctx)
p6.Args = []interface{}{testCollectionID}
result6, err := p6.Exec()
require.NoError(t, err)
// GetCollection returns map[string]interface{}
var getResult2 map[string]interface{}
if mapStrAny, ok := result6.(maps.MapStrAny); ok {
getResult2 = mapStrAny
} else if m, ok := result6.(map[string]interface{}); ok {
getResult2 = m
} else {
t.Fatalf("Unexpected result type: %T", result6)
}
assert.Equal(t, "Updated Integration Collection", getResult2["name"])
t.Logf("✓ Step 6: Verified metadata update")
// Step 7: List collections (should include ours)
p7 := process.New("kb.collection.List").WithContext(ctx).WithAuthorized(authInfo)
p7.Args = []interface{}{
map[string]interface{}{
"page": 1,
"pagesize": 100,
},
}
result7, err := p7.Exec()
require.NoError(t, err)
listResult := result7.(maps.MapStrAny)
assert.Contains(t, listResult, "data")
t.Logf("✓ Step 7: Listed collections")
// Step 8: Remove collection
p8 := process.New("kb.collection.Remove").WithContext(ctx).WithAuthorized(authInfo)
p8.Args = []interface{}{testCollectionID}
result8, err := p8.Exec()
require.NoError(t, err)
removeResult := result8.(maps.MapStrAny)
assert.Equal(t, true, removeResult["removed"])
t.Logf("✓ Step 8: Removed collection")
// Step 9: Verify collection no longer exists
p9 := process.New("kb.collection.Exists").WithContext(ctx)
p9.Args = []interface{}{testCollectionID}
result9, err := p9.Exec()
require.NoError(t, err)
existsResult3 := result9.(maps.MapStrAny)
assert.Equal(t, false, existsResult3["exists"])
t.Logf("✓ Step 9: Confirmed collection no longer exists")
t.Logf("✅ Full lifecycle completed successfully via Process API")
})
}