Implement Knowledge Base configuration loading and enhance Assistant initialization
- Added functionality to load Knowledge Base (KB) configuration from `agent/kb.yml`, allowing dynamic settings for chat sessions. - Introduced `initKBConfig` function to read and parse KB settings, integrating them into the Assistant's initialization process. - Enhanced the Assistant's conversation initialization to prepare KB collections asynchronously, improving performance during chat interactions. - Updated tests to verify the correct loading and application of KB settings, ensuring robust integration with the Assistant's functionality. - Refactored metadata handling for KB collections to include additional fields for improved context management during chat sessions.
This commit is contained in:
Max
parent
5f803ef4d4
commit
31b627aba0
13 changed files with 814 additions and 79 deletions
|
|
@ -73,11 +73,8 @@ func (ast *Assistant) Stream(ctx *context.Context, inputMessages []context.Messa
|
|||
ast.sendAgentStreamStart(ctx, streamHandler, streamStartTime)
|
||||
|
||||
// Initialize chat, prepare kb collection (optional) etc.
|
||||
err = ast.initializeConversation(ctx, inputMessages, opts)
|
||||
if err != nil {
|
||||
ast.sendStreamEndOnError(ctx, streamHandler, streamStartTime, err)
|
||||
return nil, err
|
||||
}
|
||||
// Use async version to not block the main flow
|
||||
ast.InitializeConversationAsync(ctx, opts)
|
||||
|
||||
// Initialize agent trace node
|
||||
agentNode := ast.initAgentTraceNode(ctx, inputMessages)
|
||||
|
|
|
|||
|
|
@ -2,14 +2,22 @@ package assistant
|
|||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
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 *context.Context, input []context.Message, agentNode types.Node, options ...*context.Options) ([]context.Message, error) {
|
||||
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
|
||||
|
|
@ -28,14 +36,14 @@ func (ast *Assistant) WithHistory(ctx *context.Context, input []context.Message,
|
|||
return fullMessages, nil
|
||||
}
|
||||
|
||||
// initializeConversation initialize the conversation
|
||||
func (ast *Assistant) initializeConversation(ctx *context.Context, input []context.Message, options ...*context.Options) error {
|
||||
// InitializeConversation prepares KB collection for the conversation (synchronous)
|
||||
func (ast *Assistant) InitializeConversation(ctx *agentcontext.Context, options ...*agentcontext.Options) error {
|
||||
|
||||
var opts *context.Options
|
||||
var opts *agentcontext.Options
|
||||
if len(options) > 0 && options[0] != nil {
|
||||
opts = options[0]
|
||||
} else {
|
||||
opts = &context.Options{}
|
||||
opts = &agentcontext.Options{}
|
||||
}
|
||||
|
||||
// SKIP: History (for internal calls like title/prompt etc.)
|
||||
|
|
@ -43,35 +51,166 @@ func (ast *Assistant) initializeConversation(ctx *context.Context, input []conte
|
|||
return nil
|
||||
}
|
||||
|
||||
chatid := ctx.ChatID
|
||||
teamid := ctx.Authorized.TeamID
|
||||
userid := ctx.Authorized.UserID
|
||||
fmt.Printf(">>> initializeChat: chatid=%s, teamid=%s, userid=%s\n", chatid, teamid, userid)
|
||||
|
||||
// Prepare kb collection (optional)
|
||||
err := ast.prepareKBCollection(ctx, input, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
// Check if authorized info is available
|
||||
if ctx.Authorized == nil {
|
||||
fmt.Printf(">>> Warning: no authorized info, skipping KB collection preparation\n")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Save chat
|
||||
err = ast.saveChat(ctx, input, opts)
|
||||
// Prepare kb collection
|
||||
err := ast.prepareKBCollection(ctx, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
// Log but don't fail the chat
|
||||
fmt.Printf(">>> Warning: failed to prepare KB collection: %v\n", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Prepare kb collection (optional)
|
||||
func (ast *Assistant) prepareKBCollection(ctx *context.Context, input []context.Message, opts *context.Options) error {
|
||||
_ = ctx
|
||||
// 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
|
||||
_ = input
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ast *Assistant) saveChat(ctx *context.Context, input []context.Message, opts *context.Options) error {
|
||||
// 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
|
||||
|
|
|
|||
302
agent/assistant/chat_test.go
Normal file
302
agent/assistant/chat_test.go
Normal 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")
|
||||
})
|
||||
}
|
||||
|
|
@ -23,9 +23,10 @@ var loaded = NewCache(200) // 200 is the default capacity
|
|||
var storage store.Store = nil
|
||||
var search interface{} = nil
|
||||
var modelCapabilities map[string]gouOpenAI.Capabilities = map[string]gouOpenAI.Capabilities{}
|
||||
var defaultConnector string = "" // default connector
|
||||
var globalUses *context.Uses = nil // global uses configuration from agent.yml
|
||||
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
|
||||
var defaultConnector string = "" // default connector
|
||||
var globalUses *context.Uses = nil // global uses configuration from agent.yml
|
||||
var globalPrompts []store.Prompt = nil // global prompts from agent/prompts.yml
|
||||
var globalKBSetting *store.KBSetting = nil // global KB setting from agent/kb.yml
|
||||
|
||||
// LoadBuiltIn load the built-in assistants
|
||||
func LoadBuiltIn() error {
|
||||
|
|
@ -161,6 +162,16 @@ func GetGlobalPrompts(ctx map[string]string) []store.Prompt {
|
|||
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
|
||||
func SetCache(capacity int) {
|
||||
ClearCache()
|
||||
|
|
|
|||
|
|
@ -94,6 +94,10 @@ func init() {
|
|||
"mcp.list_samples.description": "List samples for '%s' from MCP client '%s'",
|
||||
"mcp.get_sample.label": "MCP: Get Sample",
|
||||
"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.failed": "失败",
|
||||
"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.get_sample.label": "MCP: 获取示例",
|
||||
"mcp.get_sample.description": "从 MCP 客户端 '%s' 获取 '%s' 的第 %d 个示例",
|
||||
|
||||
// KB: Chat collection
|
||||
"kb.chat.name": "聊天知识库",
|
||||
"kb.chat.description": "自动为聊天会话创建的知识库集合",
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -86,6 +86,12 @@ func Load(cfg config.Config) error {
|
|||
return err
|
||||
}
|
||||
|
||||
// Initialize KB Configuration
|
||||
err = initKBConfig()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// Initialize Assistant
|
||||
err = initAssistant()
|
||||
if err != nil {
|
||||
|
|
@ -209,6 +215,10 @@ func initAssistant() error {
|
|||
assistant.SetModelCapabilities(agentDSL.Models)
|
||||
}
|
||||
|
||||
if agentDSL.KB != nil {
|
||||
assistant.SetGlobalKBSetting(agentDSL.KB)
|
||||
}
|
||||
|
||||
// Load Built-in Assistants
|
||||
err := assistant.LoadBuiltIn()
|
||||
if err != nil {
|
||||
|
|
@ -225,6 +235,29 @@ func initAssistant() error {
|
|||
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
|
||||
func defaultAssistant() (*assistant.Assistant, error) {
|
||||
if agentDSL.Uses == nil || agentDSL.Uses.Default == "" {
|
||||
|
|
|
|||
|
|
@ -59,6 +59,41 @@ func TestLoad(t *testing.T) {
|
|||
assert.NotNil(t, agent.Models)
|
||||
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) {
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
graphragtypes "github.com/yaoapp/gou/graphrag/types"
|
||||
"github.com/yaoapp/xun/dbal/query"
|
||||
"github.com/yaoapp/yao/agent/context"
|
||||
"github.com/yaoapp/yao/agent/i18n"
|
||||
|
|
@ -95,6 +96,34 @@ type Prompt struct {
|
|||
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
|
||||
type KnowledgeBase struct {
|
||||
Collections []string `json:"collections,omitempty"` // Knowledge base collection IDs
|
||||
|
|
|
|||
|
|
@ -5,6 +5,7 @@ import (
|
|||
|
||||
"github.com/yaoapp/yao/agent"
|
||||
"github.com/yaoapp/yao/config"
|
||||
"github.com/yaoapp/yao/kb"
|
||||
"github.com/yaoapp/yao/test"
|
||||
)
|
||||
|
||||
|
|
@ -16,8 +17,14 @@ import (
|
|||
func Prepare(t *testing.T, opts ...interface{}) {
|
||||
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
|
||||
err := agent.Load(config.Conf)
|
||||
err = agent.Load(config.Conf)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ type DSL struct {
|
|||
// Global External Settings - model capabilities, tools, etc.
|
||||
// ===============================
|
||||
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
|
||||
// ===============================
|
||||
|
|
|
|||
|
|
@ -113,6 +113,71 @@ func (instance *KBInstance) CreateCollection(ctx context.Context, params *Create
|
|||
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,
|
||||
|
|
@ -149,16 +214,22 @@ func (instance *KBInstance) RemoveCollection(ctx context.Context, collectionID s
|
|||
return nil, fmt.Errorf("collection ID is required")
|
||||
}
|
||||
|
||||
removed, err := instance.GraphRag.RemoveCollection(ctx, collectionID)
|
||||
// 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 {
|
||||
return nil, fmt.Errorf("failed to remove collection: %w", err)
|
||||
// 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
|
||||
}
|
||||
|
||||
if !removed {
|
||||
return nil, fmt.Errorf("collection not found or could not be removed")
|
||||
}
|
||||
|
||||
// Remove collection and documents from database after successful GraphRag removal
|
||||
// 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
|
||||
|
|
@ -167,22 +238,36 @@ func (instance *KBInstance) RemoveCollection(ctx context.Context, collectionID s
|
|||
}
|
||||
|
||||
// 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
|
||||
// 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", collectionID, documentsRemoved)
|
||||
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,
|
||||
Removed: removed || dbCleanupSuccess, // Consider successful if either succeeded
|
||||
DocumentsRemoved: documentsRemoved,
|
||||
Message: "Collection removed successfully",
|
||||
}, nil
|
||||
|
|
|
|||
|
|
@ -77,6 +77,34 @@ func TestCreateCollection(t *testing.T) {
|
|||
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) {
|
||||
|
|
@ -182,6 +210,19 @@ func TestGetCollection(t *testing.T) {
|
|||
// 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"])
|
||||
})
|
||||
|
||||
|
|
@ -293,9 +334,16 @@ func TestRemoveCollection(t *testing.T) {
|
|||
})
|
||||
|
||||
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")
|
||||
assert.Error(t, err)
|
||||
assert.Nil(t, result)
|
||||
|
||||
// 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) {
|
||||
|
|
@ -304,6 +352,42 @@ func TestRemoveCollection(t *testing.T) {
|
|||
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) {
|
||||
|
|
|
|||
|
|
@ -7,67 +7,67 @@ import (
|
|||
|
||||
// CreateCollectionParams represents the parameters for creating a collection
|
||||
type CreateCollectionParams struct {
|
||||
ID string `json:"id"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
EmbeddingProviderID string `json:"embedding_provider_id"`
|
||||
EmbeddingOptionID string `json:"embedding_option_id"`
|
||||
Locale string `json:"locale,omitempty"`
|
||||
Config *types.CreateCollectionOptions `json:"config,omitempty"`
|
||||
AuthScope map[string]interface{} `json:"-"` // Internal: authentication scope fields
|
||||
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"`
|
||||
Message string `json:"message"`
|
||||
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"`
|
||||
Removed bool `json:"removed"`
|
||||
DocumentsRemoved int `json:"documents_removed"`
|
||||
Message string `json:"message"`
|
||||
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"`
|
||||
Exists bool `json:"exists"`
|
||||
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"`
|
||||
PageSize int `json:"pagesize"`
|
||||
Keywords string `json:"keywords,omitempty"`
|
||||
Status []string `json:"status,omitempty"`
|
||||
System *bool `json:"system,omitempty"`
|
||||
EmbeddingProviderID string `json:"embedding_provider_id,omitempty"`
|
||||
Select []interface{} `json:"select,omitempty"`
|
||||
Sort []model.QueryOrder `json:"sort,omitempty"`
|
||||
AuthFilters []model.QueryWhere `json:"-"` // Internal: authentication filters
|
||||
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"`
|
||||
Next int `json:"next"`
|
||||
Prev int `json:"prev"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pagesize"`
|
||||
Total int `json:"total"`
|
||||
PageCnt int `json:"pagecnt"`
|
||||
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"`
|
||||
AuthScope map[string]interface{} `json:"-"` // Internal: authentication scope fields for update
|
||||
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"`
|
||||
Message string `json:"message"`
|
||||
CollectionID string `json:"collection_id" yaml:"collection_id"`
|
||||
Message string `json:"message" yaml:"message"`
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue