diff --git a/agent/assistant/agent.go b/agent/assistant/agent.go index eca73f0a..4d655f63 100644 --- a/agent/assistant/agent.go +++ b/agent/assistant/agent.go @@ -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) defer log.Trace("[AGENT] Stream ended: assistant=%s, contextID=%s", ast.ID, ctx.ID) + // Validate user permissions var err error + err = ast.checkPermissions(ctx) + if err != nil { + return nil, err + } + + // Start stream time streamStartTime := time.Now() // 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 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 agentNode := ast.initAgentTraceNode(ctx, inputMessages) diff --git a/agent/assistant/chat.go b/agent/assistant/chat.go new file mode 100644 index 00000000..2091507b --- /dev/null +++ b/agent/assistant/chat.go @@ -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 +} diff --git a/agent/assistant/chat_test.go b/agent/assistant/chat_test.go new file mode 100644 index 00000000..7fb02da3 --- /dev/null +++ b/agent/assistant/chat_test.go @@ -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") + }) +} diff --git a/agent/assistant/history.go b/agent/assistant/history.go deleted file mode 100644 index bbecb73c..00000000 --- a/agent/assistant/history.go +++ /dev/null @@ -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 -} diff --git a/agent/assistant/hook/script.go b/agent/assistant/hook/script.go index 0f59a246..363dbe53 100644 --- a/agent/assistant/hook/script.go +++ b/agent/assistant/hook/script.go @@ -23,6 +23,11 @@ func (s *Script) Execute(ctx *context.Context, method string, args ...interface{ } defer scriptCtx.Close() + // Set authorized information if available + if ctx.Authorized != nil { + scriptCtx.WithAuthorized(ctx.Authorized.AuthorizedToMap()) + } + // The first argument is the context args = append([]interface{}{ctx}, args...) diff --git a/agent/assistant/load.go b/agent/assistant/load.go index 28381f66..2ac9974d 100644 --- a/agent/assistant/load.go +++ b/agent/assistant/load.go @@ -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() diff --git a/agent/assistant/permission.go b/agent/assistant/permission.go new file mode 100644 index 00000000..ea59815a --- /dev/null +++ b/agent/assistant/permission.go @@ -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 +} diff --git a/agent/assistant/scripts.go b/agent/assistant/scripts.go index e04c7709..574bd90c 100644 --- a/agent/assistant/scripts.go +++ b/agent/assistant/scripts.go @@ -20,6 +20,11 @@ var scriptsMutex sync.Mutex // Execute execute the script 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 { return nil, nil } @@ -30,6 +35,11 @@ func (s *Script) Execute(ctx context.Context, method string, args ...interface{} } defer scriptCtx.Close() + // Set authorized information if available + if authorized != nil { + scriptCtx.WithAuthorized(authorized) + } + // Call the method with provided arguments as-is result, err := scriptCtx.CallWith(ctx, method, args...) @@ -359,8 +369,14 @@ func makeScriptHandler(script *Script) process.Handler { // Get arguments from process args := p.Args - // Execute the script - result, err := script.Execute(p.Context, method, args...) + // Convert authorized info to map if available + 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 { exception.New(err.Error(), 500).Throw() } diff --git a/agent/assistant/scripts_test.go b/agent/assistant/scripts_test.go index bb0aea1f..11a1bb0f 100644 --- a/agent/assistant/scripts_test.go +++ b/agent/assistant/scripts_test.go @@ -1,10 +1,12 @@ package assistant import ( + "context" "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "github.com/yaoapp/gou/process" "github.com/yaoapp/yao/config" "github.com/yaoapp/yao/test" ) @@ -155,3 +157,151 @@ func TestGenerateScriptID(t *testing.T) { // TestLoadScriptsThreadSafety tests concurrent script loading // Note: This test is commented out due to path format differences // 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") + }) +} diff --git a/agent/i18n/builtin.go b/agent/i18n/builtin.go index 380f4f2c..b912d798 100644 --- a/agent/i18n/builtin.go +++ b/agent/i18n/builtin.go @@ -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": "自动为聊天会话创建的知识库集合", }, } } diff --git a/agent/load.go b/agent/load.go index 44fa1494..63738a52 100644 --- a/agent/load.go +++ b/agent/load.go @@ -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 == "" { diff --git a/agent/load_test.go b/agent/load_test.go index 602cc094..7f238951 100644 --- a/agent/load_test.go +++ b/agent/load_test.go @@ -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) { diff --git a/agent/store/CHAT_STORAGE_DESIGN.md b/agent/store/CHAT_STORAGE_DESIGN.md new file mode 100644 index 00000000..d4269260 --- /dev/null +++ b/agent/store/CHAT_STORAGE_DESIGN.md @@ -0,0 +1,1061 @@ +# Chat Storage Design + +This document describes the design for storing chat conversations, messages, and execution steps in the YAO Agent system. + +## Table of Contents + +- [Overview](#overview) +- [Architecture](#architecture) +- [Data Models](#data-models) +- [Write Strategy](#write-strategy) +- [API Interface](#api-interface) +- [Usage Examples](#usage-examples) +- [Related Documents](#related-documents) + +## Overview + +The chat storage system is designed to: + +1. **Store user-visible messages** - All messages sent via `ctx.Send()`, including text, images, loading states, etc. +2. **Support resume/retry** - Track execution steps to enable recovery from interruptions or failures +3. **Efficient writes** - Batch message writes at request end + +### Design Goals + +| Goal | Solution | +| ------------------------ | ------------------------------------------------ | +| Complete chat history | Store final content of all `ctx.Send()` messages | +| Resume from interruption | Track step status and input/output | +| Retry failed operations | Store step input for re-execution | +| Minimize database writes | Batch writes at request end | + +### Non-Goals + +- **Tracing/debugging** - Handled by separate [Trace module](../../trace/README.md) +- **Streaming replay** - Not needed, history shows final content only +- **Request tracking/billing** - Handled by [OpenAPI Request module](../../openapi/request/REQUEST_DESIGN.md) + +### Relationship with OpenAPI Request + +The Agent storage focuses on **chat content and execution state**, while request tracking (billing, rate limiting, auditing) is handled globally by the OpenAPI layer: + +| Concern | Module | Table | +| ---------------- | ----------------- | ----------------- | +| Request tracking | `openapi/request` | `openapi_request` | +| Billing (tokens) | `openapi/request` | `openapi_request` | +| Rate limiting | `openapi/request` | - | +| Chat sessions | `agent/store` | `agent_chat` | +| Chat messages | `agent/store` | `agent_message` | +| Resume/Retry | `agent/store` | `agent_resume` | + +The `request_id` from OpenAPI middleware is passed to Agent and stored in messages/steps for correlation. + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Chat Storage │ +├─────────────────────────────────────────────────────────────┤ +│ │ +│ ┌─────────────────┐ │ +│ │ Chat │ Metadata: title, assistant, user │ +│ └────────┬────────┘ │ +│ │ │ +│ │ 1:N │ +│ ▼ │ +│ ┌─────────────────┐ │ +│ │ Message │ User-visible: type, props, role │ +│ └────────┬────────┘ │ +│ │ │ +│ │ N:N (via request_id) │ +│ ▼ │ +│ ┌─────────────────┐ │ +│ │ Resume │ Recovery: type, status, input/output │ +│ │ (only on fail) │ Only saved when interrupted/failed │ +│ └─────────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────┘ +``` + +## Data Models + +### 1. Chat Table + +Stores chat metadata and session information. + +**Table Name:** `agent_chat` + +| Column | Type | Nullable | Index | Description | +| ----------------- | ----------- | -------- | ------ | -------------------------------- | +| `id` | ID | No | PK | Auto-increment primary key | +| `chat_id` | string(64) | No | Unique | Unique chat identifier | +| `title` | string(500) | Yes | - | Chat title | +| `assistant_id` | string(200) | No | Yes | Associated assistant ID | +| `mode` | string(50) | No | - | Chat mode (default: "chat") | +| `status` | enum | No | Yes | Status: `active`, `archived` | +| `preset` | boolean | No | - | Whether this is a preset chat | +| `public` | boolean | No | - | Whether shared across all teams | +| `share` | enum | No | Yes | Sharing scope: `private`, `team` | +| `sort` | integer | No | - | Sort order for display | +| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Creation timestamp | +| `updated_at` | timestamp | No | - | Last update timestamp | + +**Model Options:** + +```json +{ + "option": { + "soft_deletes": true, + "permission": true, + "timestamps": true + } +} +``` + +**Note:** `permission: true` enables Yao's built-in permission management, which automatically adds the following fields: + +| Field | Type | Description | +| ------------------ | ----------- | ------------------------------ | +| `__yao_created_by` | string(200) | User ID who created the record | +| `__yao_updated_by` | string(200) | User ID who last updated | +| `__yao_team_id` | string(200) | Team ID for team-level access | +| `__yao_tenant_id` | string(200) | Tenant ID for multi-tenancy | + +These fields are automatically managed by the framework and used for access control filtering. + +**Indexes:** + +| Name | Columns | Type | +| -------------------- | ----------------- | ----- | +| `idx_chat_assistant` | `assistant_id` | index | +| `idx_chat_status` | `status` | index | +| `idx_chat_share` | `share` | index | +| `idx_chat_last_msg` | `last_message_at` | index | + +### 2. Message Table + +Stores user-visible messages (both user input and assistant responses). + +**Table Name:** `agent_message` + +| Column | Type | Nullable | Index | Description | +| -------------- | ----------- | -------- | ------ | ----------------------------------------- | +| `id` | ID | No | PK | Auto-increment primary key | +| `message_id` | string(64) | No | Unique | Unique message identifier | +| `chat_id` | string(64) | No | Yes | Parent chat ID | +| `request_id` | string(64) | Yes | Yes | Request ID for grouping | +| `role` | enum | No | Yes | Role: `user`, `assistant` | +| `type` | string(50) | No | - | Message type (text, image, loading, etc.) | +| `props` | json | No | - | Message properties (content, url, etc.) | +| `block_id` | string(64) | Yes | Yes | Block grouping ID | +| `thread_id` | string(64) | Yes | Yes | Thread grouping ID | +| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) | +| `sequence` | integer | No | Yes | Message order within chat | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Creation timestamp | +| `updated_at` | timestamp | No | - | Last update timestamp | + +**Indexes:** + +| Name | Columns | Type | +| ------------------- | --------------------- | ----- | +| `idx_msg_chat_seq` | `chat_id`, `sequence` | index | +| `idx_msg_request` | `request_id` | index | +| `idx_msg_block` | `block_id` | index | +| `idx_msg_assistant` | `assistant_id` | index | + +**Message Types (Built-in):** + +All built-in types defined in `agent/output/BUILTIN_TYPES.md` are stored. See that document for complete Props structures. + +| Type | Description | Props Example | Stored? | +| ------------ | -------------------------------- | ------------------------------------------------------------------------------------------- | ----------- | +| `user_input` | User input (frontend display) | `{"content": "Hello", "role": "user", "name": "John"}` | ✅ Yes | +| `text` | Text/Markdown content | `{"content": "Hello **world**!"}` | ✅ Yes | +| `thinking` | Reasoning process (o1, DeepSeek) | `{"content": "Let me analyze..."}` | ✅ Yes | +| `loading` | Loading/processing indicator | `{"message": "Searching knowledge base..."}` | ✅ Yes | +| `tool_call` | LLM tool/function call | `{"id": "call_abc123", "name": "get_weather", "arguments": "{\"location\":\"SF\"}"}` | ✅ Yes | +| `error` | Error message | `{"message": "Connection timeout", "code": "TIMEOUT", "details": "..."}` | ✅ Yes | +| `image` | Image content | `{"url": "...", "alt": "...", "width": 200, "height": 200, "detail": "auto"}` | ✅ Yes | +| `audio` | Audio content | `{"url": "...", "format": "mp3", "duration": 120.5, "transcript": "...", "controls": true}` | ✅ Yes | +| `video` | Video content | `{"url": "...", "format": "mp4", "thumbnail": "...", "width": 640, "height": 360}` | ✅ Yes | +| `action` | System action (CUI only) | `{"name": "open_panel", "payload": {"panel_id": "user_profile"}}` | ✅ Yes | +| `event` | Lifecycle event (CUI only) | `{"event": "stream_start", "message": "...", "data": {...}}` | ⚠️ Optional | + +**Note on `event` type:** Lifecycle events (`stream_start`, `stream_end`, etc.) are typically transient and may not need persistent storage. Consider storing only significant events or skipping entirely based on use case. + +**Tool Call Storage:** + +Tool calls from LLM responses are stored as `tool_call` type messages. The raw tool call data is preserved in `props`: + +```json +{ + "message_id": "msg_001", + "chat_id": "chat_123", + "role": "assistant", + "type": "tool_call", + "props": { + "id": "call_abc123", + "name": "get_weather", + "arguments": "{\"location\": \"San Francisco\", \"unit\": \"celsius\"}" + }, + "block_id": "B1", + "sequence": 5 +} +``` + +**Tool Result Storage:** + +Tool execution results can be stored as `text` type with metadata indicating it's a tool result: + +```json +{ + "message_id": "msg_002", + "chat_id": "chat_123", + "role": "assistant", + "type": "text", + "props": { + "content": "The weather in San Francisco is 18°C and sunny." + }, + "metadata": { + "tool_call_id": "call_abc123", + "tool_name": "get_weather", + "is_tool_result": true + }, + "block_id": "B1", + "sequence": 6 +} +``` + +**Custom Types:** + +Any type not in the built-in list is considered a custom type and stored with its original structure: + +```json +{ + "type": "chart", + "props": { + "chartType": "bar", + "data": [...], + "options": {...} + } +} +``` + +**Multimodal User Input:** + +User input with multimodal content (text + images + files) is stored as `user_input` type: + +```json +{ + "message_id": "msg_000", + "chat_id": "chat_123", + "role": "user", + "type": "user_input", + "props": { + "content": [ + { "type": "text", "text": "What's in this image?" }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/photo.jpg", + "detail": "high" + } + } + ], + "role": "user", + "name": "John" + }, + "sequence": 1 +} +``` + +### 3. Resume Table + +Stores execution state for resume/retry functionality. **Only written when request is interrupted or failed.** + +**Table Name:** `agent_resume` + +| Column | Type | Nullable | Index | Description | +| ----------------- | ----------- | -------- | ------ | -------------------------------- | +| `id` | ID | No | PK | Auto-increment primary key | +| `resume_id` | string(64) | No | Unique | Unique resume record identifier | +| `chat_id` | string(64) | No | Yes | Parent chat ID | +| `request_id` | string(64) | No | Yes | Request ID | +| `assistant_id` | string(200) | No | Yes | Assistant executing this step | +| `stack_id` | string(64) | No | Yes | Stack node ID for this execution | +| `stack_parent_id` | string(64) | Yes | Yes | Parent stack ID (for A2A calls) | +| `stack_depth` | integer | No | - | Call depth (0=root, 1+=nested) | +| `type` | enum | No | Yes | Step type | +| `status` | enum | No | Yes | Status: `interrupted`, `failed` | +| `input` | json | Yes | - | Step input data | +| `output` | json | Yes | - | Step output data (partial) | +| `space_snapshot` | json | Yes | - | Space data snapshot for recovery | +| `error` | text | Yes | - | Error message if failed | +| `sequence` | integer | No | Yes | Step order within request | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Creation timestamp | +| `updated_at` | timestamp | No | - | Last update timestamp | + +**Space Snapshot:** + +The `space_snapshot` field stores the shared data space (`ctx.Space`) at each step for recovery purposes. + +```typescript +// Example: In Next hook, set data to Space before delegate +ctx.space.Set("choose_prompt", "query"); +return { + delegate: { agent_id: "expense", messages: payload.messages }, +}; +``` + +If interrupted during delegate, the `space_snapshot` allows restoring `ctx.Space` state: + +```json +{ + "choose_prompt": "query", + "user_preferences": { "currency": "USD" } +} +``` + +**Resume Step Types:** + +| Type | Description | Input | Output | +| ------------- | --------------------- | ---------------------- | ------------------------------------- | +| `input` | User input received | `{messages: [...]}` | - | +| `hook_create` | Create hook execution | `{messages: [...]}` | `{messages: [...], ...}` | +| `llm` | LLM completion call | `{messages: [...]}` | `{content: "...", tool_calls: [...]}` | +| `tool` | Tool/MCP execution | `{server, tool, args}` | `{result: ...}` | +| `hook_next` | Next hook execution | `{completion, tools}` | `{data: ...}` | +| `delegate` | A2A delegation | `{agent_id, messages}` | `{response: ...}` | + +**Resume Status (only two values - table only stores failed/interrupted):** + +| Status | Description | Action | +| ------------- | ----------------- | -------- | +| `failed` | Failed with error | Retry | +| `interrupted` | User interrupted | Continue | + +**Indexes:** + +| Name | Columns | Type | +| ---------------------- | ------------------------ | ----- | +| `idx_resume_chat` | `chat_id` | index | +| `idx_resume_request` | `request_id`, `sequence` | index | +| `idx_resume_status` | `status` | index | +| `idx_resume_stack` | `stack_id` | index | +| `idx_resume_parent` | `stack_parent_id` | index | +| `idx_resume_assistant` | `assistant_id` | index | + +## Write Strategy + +### Two-Write Strategy + +All data is buffered in memory during execution and written to database only **twice**: + +1. **Write 1 (Entry)**: When `Stream()` starts - save user input message +2. **Write 2 (Exit)**: When `Stream()` exits - batch save messages (and steps only on error/interrupt) + +**Note**: Request tracking (status, tokens, duration) is handled by [OpenAPI Request Middleware](../../openapi/request/REQUEST_DESIGN.md). + +``` +Stream() Entry + │ + ├── 【Write 1】Save user input + │ - User message (role=user) + │ + ├── Execution (all in memory) + │ - ctx.Send() → messageBuffer + │ - ctx.Append() → update messageBuffer + │ - ctx.Replace() → update messageBuffer + │ - Each step → stepBuffer + │ + └── 【Write 2】Save final state (via defer) + │ + ├── Always: + │ - Batch write all assistant messages + │ - Update token usage in openapi_request (via request_id) + │ + └── Only on error/interrupt: + - Batch write all steps (for resume/retry) +``` + +### Write Points + +| Event | Message Table | Step Table | Token Usage | +| ---------------- | -------------------- | ----------------------------------- | ----------- | +| Stream entry | Write 1 (user input) | - | - | +| During execution | Buffer in memory | Buffer in memory | - | +| **Completed** | **Batch write all** | **❌ Skip (no need to resume)** | ✅ Update | +| On interrupt | Batch write buffered | ✅ Batch write (status=interrupted) | ✅ Update | +| On error | Batch write buffered | ✅ Batch write (status=failed) | ✅ Update | + +**Why skip Steps on success?** + +- Steps are only needed for resume/retry operations +- If completed successfully, there's nothing to resume +- Reduces database writes and keeps Step table clean + +### Why Two Writes? + +| Scenario | What Happens | Data Safe? | +| ------------------ | ----------------------------------- | ---------- | +| Normal completion | `defer` triggers → Write 2 executes | ✅ | +| User clicks stop | `defer` triggers → Write 2 executes | ✅ | +| LLM timeout | `defer` triggers → Write 2 executes | ✅ | +| Tool failure | `defer` triggers → Write 2 executes | ✅ | +| Network disconnect | `defer` triggers → Write 2 executes | ✅ | +| Process crash | Service is down, user must retry | N/A | + +**Note**: Process crash is a catastrophic failure handled at infrastructure level, not application level. + +### Write Count Comparison + +For a typical request: user input → hook_create → llm → tool → llm → hook_next → 5 messages + +| Strategy | Database Writes | Notes | +| ---------------------- | --------------- | ------------------ | +| Write per operation | 1 + 5 + 5 = 11 | One write per step | +| **Two-write strategy** | **2** | Entry + Exit only | + +### Implementation + +````go +func (ast *Assistant) Stream(ctx, inputMessages, options) { + // ========== Write 1: Entry ========== + userMsg := createUserMessage(ctx, inputMessages) + chatStore.SaveMessages(ctx.ChatID, []*Message{userMsg}) + + // ========== Memory Buffers ========== + messageBuffer := NewMessageBuffer() + stepBuffer := NewStepBuffer() + + // Track current step for error handling + var currentStep *Step + + defer func() { + // ========== Write 2: Exit (always executes) ========== + // Determine final status for incomplete steps + finalStatus := "completed" + if ctx.IsInterrupted() { + finalStatus = "interrupted" + } + if r := recover(); r != nil { + finalStatus = "failed" + } + + // Update status of any incomplete step + if currentStep != nil && currentStep.Status == "running" { + currentStep.Status = finalStatus + } + + // Batch write all buffered data + chatStore.SaveMessages(ctx.ChatID, messageBuffer.GetAll()) + chatStore.SaveSteps(stepBuffer.GetAll()) + + // Update token usage in OpenAPI request record + if ctx.RequestID != "" && completionResponse != nil { + request.UpdateTokenUsage( + ctx.RequestID, + completionResponse.Usage.PromptTokens, + completionResponse.Usage.CompletionTokens, + ) + } + }() + + // ========== Execution (all in memory) ========== + // Note: request_id = ctx.RequestID (from OpenAPI middleware) + + // hook_create + currentStep = stepBuffer.Add(createStep(ctx, "hook_create", "running", input, nil)) + createResponse := ast.HookScript.Create(...) + currentStep.Output = createResponse + currentStep.Status = "completed" + + // llm + currentStep = stepBuffer.Add(createStep(ctx, "llm", "running", messages, nil)) + completionResponse := ast.executeLLMStream(...) + currentStep.Output = completionResponse + currentStep.Status = "completed" + + // tool (if any) + for _, toolCall := range completionResponse.ToolCalls { + currentStep = stepBuffer.Add(createStep(ctx, "tool", "running", toolCall, nil)) + result := executeToolCall(toolCall) + currentStep.Output = result + currentStep.Status = "completed" + } + + // hook_next + currentStep = stepBuffer.Add(createStep(ctx, "hook_next", "running", payload, nil)) + nextResponse := ast.HookScript.Next(...) + currentStep.Output = nextResponse + currentStep.Status = "completed" + currentStep = nil // All done + + // Messages are automatically buffered via ctx.Send() +} + +// createResumeRecord creates a resume record with context information +// Only called when request fails or is interrupted +func createResumeRecord(ctx *Context, stepType, status string, input, output interface{}, err error) *Resume { + // Capture Space snapshot for recovery + var spaceSnapshot map[string]interface{} + if ctx.Space != nil { + spaceSnapshot = ctx.Space.Snapshot() // Get all key-value pairs + } + + errorMsg := "" + if err != nil { + errorMsg = err.Error() + } + + return &Resume{ + ResumeID: generateID(), + ChatID: ctx.ChatID, // ChatID + RequestID: ctx.RequestID, // From OpenAPI middleware + AssistantID: ctx.AssistantID, + StackID: ctx.Stack.ID, + StackParentID: ctx.Stack.ParentID, + StackDepth: ctx.Stack.Depth, + Type: stepType, + Status: status, // "failed" or "interrupted" + Input: input, + Output: output, + SpaceSnapshot: spaceSnapshot, // Shared space data for recovery + Error: errorMsg, + Sequence: nextSequence(), + } +} + +## API Interface + +### ChatStore Interface + +```go +// ChatStore defines the chat storage interface +type ChatStore interface { + // Chat Management + CreateChat(chat *Chat) error + GetChat(chatID string) (*Chat, error) + UpdateChat(chatID string, updates map[string]interface{}) error + DeleteChat(chatID string) error + ListChats(filter ChatFilter) (*ChatList, error) + + // Message Management + SaveMessages(chatID string, messages []*Message) error + GetMessages(chatID string, filter MessageFilter) ([]*Message, error) + UpdateMessage(messageID string, updates map[string]interface{}) error + DeleteMessages(chatID string, messageIDs []string) error + + // Resume Management (only called on failure/interrupt) + SaveResume(records []*Resume) error + GetResume(chatID string) ([]*Resume, error) + GetLastResume(chatID string) (*Resume, error) + GetResumeByStackID(stackID string) ([]*Resume, error) + GetStackPath(stackID string) ([]string, error) // Returns [root_stack_id, ..., current_stack_id] + DeleteResume(chatID string) error // Clean up after successful resume +} + +// SpaceStore defines the interface for Space snapshot operations +// Note: Space itself uses plan.Space interface, this is for persistence +type SpaceStore interface { + // Snapshot returns all key-value pairs in the space + Snapshot() map[string]interface{} + + // Restore sets multiple key-value pairs from a snapshot + Restore(data map[string]interface{}) error +} +```` + +### Data Structures + +```go +// Chat represents a chat session +type Chat struct { + ChatID string `json:"chat_id"` + Title string `json:"title,omitempty"` + AssistantID string `json:"assistant_id"` + Mode string `json:"mode"` + Status string `json:"status"` + Preset bool `json:"preset"` + Public bool `json:"public"` + Share string `json:"share"` // "private" or "team" + Sort int `json:"sort"` + LastMessageAt *time.Time `json:"last_message_at,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// Message represents a chat message +type Message struct { + MessageID string `json:"message_id"` + ChatID string `json:"chat_id"` + RequestID string `json:"request_id,omitempty"` + Role string `json:"role"` + Type string `json:"type"` + Props map[string]interface{} `json:"props"` + BlockID string `json:"block_id,omitempty"` + ThreadID string `json:"thread_id,omitempty"` + AssistantID string `json:"assistant_id,omitempty"` + Sequence int `json:"sequence"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// Resume represents an execution state for recovery (only stored on failure/interrupt) +type Resume struct { + ResumeID string `json:"resume_id"` + ChatID string `json:"chat_id"` + RequestID string `json:"request_id"` + AssistantID string `json:"assistant_id"` + StackID string `json:"stack_id"` + StackParentID string `json:"stack_parent_id,omitempty"` + StackDepth int `json:"stack_depth"` + Type string `json:"type"` + Status string `json:"status"` // "failed" or "interrupted" + Input map[string]interface{} `json:"input,omitempty"` + Output map[string]interface{} `json:"output,omitempty"` + SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"` // Shared space data for recovery + Error string `json:"error,omitempty"` + Sequence int `json:"sequence"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} +``` + +### Filter Structures + +```go +// ChatFilter for listing chats +type ChatFilter struct { + UserID string `json:"user_id,omitempty"` + TeamID string `json:"team_id,omitempty"` + AssistantID string `json:"assistant_id,omitempty"` + Status string `json:"status,omitempty"` + Keywords string `json:"keywords,omitempty"` + Page int `json:"page,omitempty"` + PageSize int `json:"pagesize,omitempty"` +} + +// MessageFilter for listing messages +type MessageFilter struct { + RequestID string `json:"request_id,omitempty"` + Role string `json:"role,omitempty"` + BlockID string `json:"block_id,omitempty"` + Limit int `json:"limit,omitempty"` + Offset int `json:"offset,omitempty"` +} + +// ChatList paginated response +type ChatList struct { + Data []*Chat `json:"data"` + Page int `json:"page"` + PageSize int `json:"pagesize"` + PageCount int `json:"pagecount"` + Total int `json:"total"` +} +``` + +## Usage Examples + +### 1. Complete Message Storage Example + +A typical conversation with various message types stored in `agent_message`: + +``` +User: "What's the weather in SF? Also show me a chart." + +Timeline: +1. User sends multimodal input +2. Hook shows loading state +3. LLM thinks and calls tool +4. Tool returns result +5. LLM generates text response +6. Hook sends image chart +``` + +**Stored Messages:** + +```json +[ + // 1. User input (role=user, type=user_input) + { + "message_id": "msg_001", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "user", + "type": "user_input", + "props": { + "content": "What's the weather in SF? Also show me a chart.", + "role": "user" + }, + "sequence": 1 + }, + + // 2. Loading state from Create hook (role=assistant, type=loading) + { + "message_id": "msg_002", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "loading", + "props": { + "message": "Searching knowledge base..." + }, + "block_id": "B1", + "assistant_id": "weather_assistant", + "sequence": 2 + }, + + // 3. LLM thinking process (role=assistant, type=thinking) + { + "message_id": "msg_003", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "thinking", + "props": { + "content": "User wants weather info for San Francisco. I should use the get_weather tool..." + }, + "block_id": "B2", + "assistant_id": "weather_assistant", + "sequence": 3 + }, + + // 4. LLM tool call (role=assistant, type=tool_call) + { + "message_id": "msg_004", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "tool_call", + "props": { + "id": "call_weather_001", + "name": "get_weather", + "arguments": "{\"location\": \"San Francisco\", \"unit\": \"celsius\"}" + }, + "block_id": "B2", + "assistant_id": "weather_assistant", + "sequence": 4 + }, + + // 5. Tool result (role=assistant, type=text, with tool metadata) + { + "message_id": "msg_005", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "text", + "props": { + "content": "Weather data retrieved: 18°C, sunny, humidity 65%" + }, + "block_id": "B2", + "metadata": { + "tool_call_id": "call_weather_001", + "tool_name": "get_weather", + "is_tool_result": true + }, + "assistant_id": "weather_assistant", + "sequence": 5 + }, + + // 6. LLM text response (role=assistant, type=text) + { + "message_id": "msg_006", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "text", + "props": { + "content": "The weather in San Francisco is currently **18°C** and sunny with 65% humidity. Perfect weather for outdoor activities!" + }, + "block_id": "B2", + "assistant_id": "weather_assistant", + "sequence": 6 + }, + + // 7. Chart image from Next hook (role=assistant, type=image) + { + "message_id": "msg_007", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "image", + "props": { + "url": "https://charts.example.com/weather_sf.png", + "alt": "San Francisco 7-day weather forecast", + "width": 800, + "height": 400 + }, + "block_id": "B3", + "assistant_id": "weather_assistant", + "sequence": 7 + } +] +``` + +**Streaming IDs (from `STREAMING.md`):** + +During streaming, messages include additional fields for real-time delivery: + +| Field | Purpose | Stored? | +| ------------ | ------------------------------ | ------- | +| `chunk_id` | Deduplication, ordering, debug | ❌ No | +| `message_id` | Delta merge target | ✅ Yes | +| `block_id` | UI block/section grouping | ✅ Yes | +| `thread_id` | Concurrent stream distinction | ✅ Yes | +| `delta` | Whether this is a delta chunk | ❌ No | +| `delta_path` | Path for delta merge | ❌ No | + +**Note:** `chunk_id`, `delta`, and `delta_path` are transient streaming control fields and are NOT stored. Only the final merged content is persisted. + +### 2. Error Message Storage + +When errors occur, they are stored as `error` type: + +```json +{ + "message_id": "msg_err_001", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "error", + "props": { + "message": "Failed to connect to weather service", + "code": "SERVICE_UNAVAILABLE", + "details": "Connection timeout after 30 seconds" + }, + "block_id": "B2", + "assistant_id": "weather_assistant", + "sequence": 5 +} +``` + +### 3. Action Message Storage (CUI clients) + +System actions are stored but only processed by CUI clients: + +```json +{ + "message_id": "msg_action_001", + "chat_id": "chat_123", + "request_id": "req_abc", + "role": "assistant", + "type": "action", + "props": { + "name": "open_panel", + "payload": { + "panel_id": "weather_details", + "location": "San Francisco" + } + }, + "block_id": "B2", + "assistant_id": "weather_assistant", + "sequence": 6 +} +``` + +### 4. Audio/Video Message Storage + +Multimedia content storage: + +```json +// Audio message +{ + "message_id": "msg_audio_001", + "chat_id": "chat_123", + "role": "assistant", + "type": "audio", + "props": { + "url": "https://storage.example.com/audio/response.mp3", + "format": "mp3", + "duration": 45.5, + "transcript": "Here's the weather forecast for today...", + "controls": true + }, + "sequence": 7 +} + +// Video message +{ + "message_id": "msg_video_001", + "chat_id": "chat_123", + "role": "assistant", + "type": "video", + "props": { + "url": "https://storage.example.com/video/weather_report.mp4", + "format": "mp4", + "thumbnail": "https://storage.example.com/video/weather_report_thumb.jpg", + "duration": 120.0, + "width": 1280, + "height": 720, + "controls": true + }, + "sequence": 8 +} +``` + +### 5. Load Chat History + +```go +// Get chat list +chats, _ := chatStore.ListChats(ChatFilter{ + UserID: "user123", + Status: "active", + Page: 1, + PageSize: 20, +}) + +// Get messages for a chat +messages, _ := chatStore.GetMessages("chat_123", MessageFilter{ + Limit: 100, +}) + +// Return to frontend +return map[string]interface{}{ + "chat": chat, + "messages": messages, +} +``` + +### 6. Resume from Interruption + +```go +func (ast *Assistant) Resume(ctx *Context) error { + // 1. Find last resume record + record, _ := chatStore.GetLastResume(ctx.ChatID) + if record == nil { + return nil // Nothing to resume + } + + // 2. Restore Space data from snapshot + if record.SpaceSnapshot != nil && ctx.Space != nil { + for key, value := range record.SpaceSnapshot { + ctx.Space.Set(key, value) + } + } + + // 3. Check if this is an A2A nested call + if record.StackDepth > 0 { + // Need to rebuild the call stack + return ast.ResumeNestedCall(ctx, record) + } + + // 4. Resume based on step type + var err error + switch record.Type { + case "llm": + // Re-execute LLM call with saved input + messages := record.Input["messages"].([]Message) + err = ast.executeLLMStream(ctx, messages, ...) + + case "tool": + // Retry tool call + err = ast.retryToolCall(ctx, record) + + case "hook_next": + // Re-execute hook + err = ast.executeHookNext(ctx, record.Input) + + case "delegate": + // Resume delegated agent call + agentID := record.Input["agent_id"].(string) + messages := record.Input["messages"].([]Message) + err = ast.delegateToAgent(ctx, agentID, messages) + } + + // 5. Clean up resume records on success + if err == nil { + chatStore.DeleteResume(ctx.ChatID) + } + + return err +} +``` + +### 7. Resume A2A Nested Calls + +For agent-to-agent (A2A) recursive calls, the stack information is essential for proper recovery. + +```go +func (ast *Assistant) ResumeNestedCall(ctx *Context, step *Step) error { + // 1. Rebuild the call stack from root to interrupted point + stackPath, _ := chatStore.GetStackPath(step.StackID) + // stackPath: [root_stack_id, parent_stack_id, ..., current_stack_id] + + // 2. Get all steps for each stack level + for _, stackID := range stackPath { + steps, _ := chatStore.GetStepsByStackID(stackID) + // Restore context for each level + } + + // 3. Resume from the interrupted assistant + targetAssistant := assistant.Select(step.AssistantID) + return targetAssistant.Stream(ctx, step.Input["messages"], ...) +} +``` + +### 8. Handle Interruption + +Interruption is handled automatically by the `defer` block in the two-write strategy. When `ctx.IsInterrupted()` returns true, the status is set to `interrupted` and all buffered data is saved. + +```go +// Inside the defer block (see Write Strategy - Implementation) +if ctx.IsInterrupted() { + status = "interrupted" +} +// Then batch write all buffered messages and steps +``` + +## A2A (Agent-to-Agent) Call Example + +When Assistant A delegates to Assistant B, the step records look like: + +``` +Request: User asks "analyze this data and visualize it" + +Step Records: +┌─────┬─────────────┬─────────────┬──────────┬───────┬───────┬─────────────┬─────────────────────────────┐ +│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │ space_snapshot │ +├─────┼─────────────┼─────────────┼──────────┼───────┼───────┼─────────────┼─────────────────────────────┤ +│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │ {} │ +│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │ {} │ +│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ {"choose_prompt": "query"} │ ← Space data set before delegate +│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │ {"choose_prompt": "query"} │ +│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ {"choose_prompt": "query"} │ ← interrupted here +└─────┴─────────────┴─────────────┴──────────┴───────┴───────┴─────────────┴─────────────────────────────┘ + +Resume Flow: +1. Find step with status="interrupted" → step 5 +2. Restore Space from space_snapshot: {"choose_prompt": "query"} +3. Check stack_depth=1 → nested call +4. Get stack path: [stk_001, stk_002] +5. Resume visualizer assistant with step 5's input +6. When visualizer completes, update step 3 (delegate) to completed +``` + +**Space Snapshot Use Case (from expense assistant):** + +```typescript +// In Next hook, before delegating to another agent +ctx.space.Set("choose_prompt", "query"); +return { + delegate: { agent_id: "expense", messages: payload.messages }, +}; + +// If interrupted during delegate, Resume will: +// 1. Restore space_snapshot → ctx.space now has "choose_prompt": "query" +// 2. The delegated agent's Create hook can read: ctx.space.GetDel("choose_prompt") +``` + +## Related Documents + +- [OpenAPI Request Design](../../openapi/request/REQUEST_DESIGN.md) - Global request tracking, billing, rate limiting +- [Trace Module](../../trace/README.md) - Detailed execution tracing for debugging +- [Agent Context](../context/README.md) - Context and message handling diff --git a/agent/store/types/types.go b/agent/store/types/types.go index 85f4dbf8..78d0344e 100644 --- a/agent/store/types/types.go +++ b/agent/store/types/types.go @@ -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 diff --git a/agent/testutils/testutils.go b/agent/testutils/testutils.go index a6fb5875..5c79c7a7 100644 --- a/agent/testutils/testutils.go +++ b/agent/testutils/testutils.go @@ -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) } diff --git a/agent/types/types.go b/agent/types/types.go index 99f009c7..c83ff390 100644 --- a/agent/types/types.go +++ b/agent/types/types.go @@ -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 // =============================== diff --git a/kb/api/api.go b/kb/api/api.go new file mode 100644 index 00000000..b02ade9d --- /dev/null +++ b/kb/api/api.go @@ -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, + } +} diff --git a/kb/api/collection.go b/kb/api/collection.go new file mode 100644 index 00000000..051ea758 --- /dev/null +++ b/kb/api/collection.go @@ -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 +} diff --git a/kb/api/collection_test.go b/kb/api/collection_test.go new file mode 100644 index 00000000..17332cee --- /dev/null +++ b/kb/api/collection_test.go @@ -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) + }) +} diff --git a/kb/api/consts.go b/kb/api/consts.go new file mode 100644 index 00000000..1520aa94 --- /dev/null +++ b/kb/api/consts.go @@ -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" +) diff --git a/kb/api/interfaces.go b/kb/api/interfaces.go new file mode 100644 index 00000000..1436341c --- /dev/null +++ b/kb/api/interfaces.go @@ -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 +} diff --git a/kb/api/types.go b/kb/api/types.go new file mode 100644 index 00000000..8a9b6af0 --- /dev/null +++ b/kb/api/types.go @@ -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"` +} diff --git a/kb/kb.go b/kb/kb.go index f4c75fbc..83acab88 100644 --- a/kb/kb.go +++ b/kb/kb.go @@ -10,6 +10,7 @@ import ( "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/kun/log" "github.com/yaoapp/yao/config" + "github.com/yaoapp/yao/kb/api" // Register the built-in providers _ "github.com/yaoapp/yao/kb/providers" @@ -21,6 +22,9 @@ import ( // Instance is the GraphRag instance var Instance types.GraphRag = nil +// API is the Knowledge Base API instance +var API api.API = nil + // KnowledgeBase is the Knowledge Base instance type KnowledgeBase struct { Config *kbtypes.Config // Knowledge Base configuration @@ -86,6 +90,10 @@ func Load(appConfig config.Config) (*KnowledgeBase, error) { // Set the instance to the global variable Instance = instance + + // Create and set the API instance + API = api.NewAPI(graphRag, &config, providers) + return instance, nil } diff --git a/openapi/kb/collection.go b/openapi/kb/collection.go index c140e934..c0fc0db3 100644 --- a/openapi/kb/collection.go +++ b/openapi/kb/collection.go @@ -9,9 +9,9 @@ import ( "github.com/gin-gonic/gin" "github.com/yaoapp/gou/graphrag/types" "github.com/yaoapp/gou/model" - "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" "github.com/yaoapp/yao/openapi/response" @@ -20,37 +20,6 @@ import ( // 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 type ProviderSettings struct { Dimension int `json:"dimension"` @@ -61,6 +30,16 @@ type ProviderSettings struct { // CreateCollection creates a new collection 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 req, collectionData, err := PrepareCreateCollection(c) if err != nil { @@ -74,12 +53,56 @@ func CreateCollection(c *gin.Context) { // Attach create scope to the collection data authInfo := authorized.GetInfo(c) + var authScope map[string]interface{} if authInfo != nil { 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 - if kb.Instance == nil { + // Build API params + 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{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", @@ -88,68 +111,6 @@ func CreateCollection(c *gin.Context) { 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) // Get collection ID from URL parameter @@ -163,16 +124,6 @@ func RemoveCollection(c *gin.Context) { 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 hasPermission, err := checkCollectionPermission(authInfo, collectionID) if err != nil { @@ -194,60 +145,38 @@ func RemoveCollection(c *gin.Context) { return } - // Call the actual RemoveCollection method - removed, err := kb.Instance.RemoveCollection(c.Request.Context(), collectionID) + // Call API to remove collection + result, err := kb.API.RemoveCollection(c.Request.Context(), collectionID) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to remove collection: " + err.Error(), + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) 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{ - "message": "Collection removed successfully", - "collection_id": collectionID, - "removed": removed, - "documents_removed": documentsRemoved, + "message": result.Message, + "collection_id": result.CollectionID, + "removed": result.Removed, + "documents_removed": result.DocumentsRemoved, } response.RespondWithSuccess(c, response.StatusOK, successData) } // CollectionExists checks if a collection exists 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 collectionID := c.Param("collectionID") if collectionID == "" { @@ -259,8 +188,28 @@ func CollectionExists(c *gin.Context) { return } - // Check if kb.Instance is available - if kb.Instance == nil { + // Call API to check collection existence + 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{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", @@ -269,26 +218,6 @@ func CollectionExists(c *gin.Context) { 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") if collectionID == "" { errorResp := &response.ErrorResponse{ @@ -299,21 +228,11 @@ func GetCollection(c *gin.Context) { 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 - } - - // Use the dedicated GetCollection method - collection, err := kb.Instance.GetCollection(c.Request.Context(), collectionID) + // Call API to get collection + collection, err := kb.API.GetCollection(c.Request.Context(), collectionID) if err != nil { // 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{ Code: response.ErrInvalidRequest.Code, ErrorDescription: "Collection not found", @@ -324,7 +243,7 @@ func GetCollection(c *gin.Context) { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to get collection: " + err.Error(), + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) return @@ -336,11 +255,8 @@ func GetCollection(c *gin.Context) { // ListCollections lists collections with pagination func ListCollections(c *gin.Context) { - // Get authorized information - authInfo := authorized.GetInfo(c) - - // Check if kb.Instance is available - if kb.Instance == nil { + // Check if kb.API is available + if kb.API == nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", @@ -349,6 +265,9 @@ func ListCollections(c *gin.Context) { return } + // Get authorized information + authInfo := authorized.GetInfo(c) + // Parse pagination parameters page := 1 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 var selectFields []interface{} if selectParam := strings.TrimSpace(c.Query("select")); selectParam != "" { requestedFields := strings.Split(selectParam, ",") for _, field := range requestedFields { field = strings.TrimSpace(field) - if field != "" && availableCollectionFields[field] { + if field != "" && kbapi.AvailableCollectionFields[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 - param := model.QueryParam{Select: selectFields} - - // Add filters - var wheres []model.QueryWhere - - // Apply permission-based filtering - wheres = append(wheres, 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", - }) + // Build filter for API + filter := &kbapi.ListCollectionsFilter{ + Page: page, + PageSize: pagesize, + Keywords: strings.TrimSpace(c.Query("keywords")), + EmbeddingProviderID: strings.TrimSpace(c.Query("embedding_provider_id")), + Select: selectFields, + Sort: orders, + AuthFilters: AuthFilter(c, authInfo), } - // Filter by status (support multiple values separated by comma) + // Parse status parameter if statusParam := strings.TrimSpace(c.Query("status")); statusParam != "" { statusList := strings.Split(statusParam, ",") - var statusValues []interface{} for _, status := range statusList { status = strings.TrimSpace(status) if status != "" { - statusValues = append(statusValues, 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.Status = append(filter.Status, status) } } } - // Filter by system flag + // Parse system parameter if systemParam := strings.TrimSpace(c.Query("system")); systemParam != "" { switch systemParam { case "true", "1": - wheres = append(wheres, model.QueryWhere{ - Column: "system", - Value: true, - }) + systemVal := true + filter.System = &systemVal case "false", "0": - wheres = append(wheres, model.QueryWhere{ - Column: "system", - Value: false, - }) + systemVal := false + filter.System = &systemVal } } - // Filter by embedding_provider_id - if providerID := strings.TrimSpace(c.Query("embedding_provider_id")); providerID != "" { - 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) + // Call API to list collections + result, err := kb.API.ListCollections(c.Request.Context(), filter) if err != nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to search collections: " + err.Error(), + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) 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 @@ -576,8 +419,8 @@ func UpdateCollectionMetadata(c *gin.Context) { return } - // Check if kb.Instance is available - if kb.Instance == nil { + // Check if kb.API is available + if kb.API == nil { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, ErrorDescription: "Knowledge base not initialized", @@ -608,45 +451,31 @@ func UpdateCollectionMetadata(c *gin.Context) { return } - // Call the actual UpdateCollectionMetadata method - err = kb.Instance.UpdateCollectionMetadata(c.Request.Context(), collectionID, req.Metadata) + // Build API params + 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 { errorResp := &response.ErrorResponse{ Code: response.ErrServerError.Code, - ErrorDescription: "Failed to update collection metadata: " + err.Error(), + ErrorDescription: err.Error(), } response.RespondWithError(c, response.StatusInternalServerError, errorResp) 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{ - "message": "Collection metadata updated successfully", - "collection_id": collectionID, + "message": result.Message, + "collection_id": result.CollectionID, } response.RespondWithSuccess(c, response.StatusOK, successData) } diff --git a/openapi/kb/collection_process.go b/openapi/kb/collection_process.go new file mode 100644 index 00000000..42d36816 --- /dev/null +++ b/openapi/kb/collection_process.go @@ -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, ¶ms); 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, ¶ms) + 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, ¶ms); 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, ¶ms) + 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 +} diff --git a/openapi/kb/kb.go b/openapi/kb/kb.go index a60b6679..7f1cb5ca 100644 --- a/openapi/kb/kb.go +++ b/openapi/kb/kb.go @@ -11,6 +11,15 @@ import ( func init() { // Register kb process handlers 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.addtext": ProcessAddText, "documents.addurl": ProcessAddURL, diff --git a/openapi/oauth/authorized/utils.go b/openapi/oauth/authorized/utils.go index 799ba6e0..3ac1c350 100644 --- a/openapi/oauth/authorized/utils.go +++ b/openapi/oauth/authorized/utils.go @@ -8,9 +8,37 @@ import ( // ProcessAuthInfo extracts authorized information from the process func ProcessAuthInfo(p *process.Process) *types.AuthorizedInfo { - // TODO: Implement this function - // Get authorized information from the process context - info := &types.AuthorizedInfo{} + if p == nil { + return nil + } + + // 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 } diff --git a/openapi/oauth/authorized/utils_test.go b/openapi/oauth/authorized/utils_test.go new file mode 100644 index 00000000..d47f3e54 --- /dev/null +++ b/openapi/oauth/authorized/utils_test.go @@ -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) + }) +} diff --git a/openapi/oauth/types/authorized_test.go b/openapi/oauth/types/authorized_test.go index cd164827..0454b565 100644 --- a/openapi/oauth/types/authorized_test.go +++ b/openapi/oauth/types/authorized_test.go @@ -462,3 +462,152 @@ func TestCopyScopesIntegration(t *testing.T) { assert.Equal(t, "tenant789", updateResult["__yao_tenant_id"]) 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) + } + } + }) + } +} diff --git a/openapi/oauth/types/types.go b/openapi/oauth/types/types.go index e205c723..8bd573a3 100644 --- a/openapi/oauth/types/types.go +++ b/openapi/oauth/types/types.go @@ -627,6 +627,64 @@ type AuthorizedInfo struct { 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 type JWTClaims struct { jwt.StandardClaims diff --git a/openapi/request/REQUEST_DESIGN.md b/openapi/request/REQUEST_DESIGN.md new file mode 100644 index 00000000..cbd20246 --- /dev/null +++ b/openapi/request/REQUEST_DESIGN.md @@ -0,0 +1,1172 @@ +# OpenAPI Request Design + +This document describes the design for global request tracking, billing, rate limiting, and auditing in the YAO OpenAPI layer. + +## Table of Contents + +- [Overview](#overview) +- [Architecture](#architecture) +- [Storage Strategy](#storage-strategy) +- [Data Model](#data-model) +- [Middleware Design](#middleware-design) +- [Rate Limiting](#rate-limiting) +- [Billing Integration](#billing-integration) +- [API Interface](#api-interface) +- [Integration with Services](#integration-with-services) + +## Overview + +The Request module provides a unified layer for: + +1. **Request Tracking** - Record all API requests with unique IDs +2. **Billing** - Track token usage and API calls for billing +3. **Rate Limiting** - Enforce request limits per user/team +4. **Auditing** - Provide audit trail for compliance + +### Design Goals + +| Goal | Solution | +| ------------------- | ------------------------------------------------ | +| Unified tracking | Single middleware for all API endpoints | +| Accurate billing | Token usage updated by services after completion | +| Flexible rate limit | Configurable limits per user/team/endpoint | +| Low overhead | KV for real-time, SQL for archive | + +### Scope + +| In Scope | Out of Scope | +| ------------------------ | ------------------------------ | +| All `/api/*` endpoints | Static file serving | +| Token usage tracking | Detailed request/response logs | +| Rate limiting | Request body storage | +| Request duration metrics | Response caching | + +## Architecture + +``` +┌─────────────────────────────────────────────────────────────┐ +│ HTTP Request │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ OAuth Guard │ +│ - Token validation │ +│ - Set AuthorizedInfo in context │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Request Middleware │ +│ - Generate request_id │ +│ - KV: Rate limit check │ +│ - KV: Quota check │ +│ - KV: Request status tracking │ +│ - Async: Archive to SQL │ +└─────────────────────────────────────────────────────────────┘ + │ + ▼ +┌─────────────────────────────────────────────────────────────┐ +│ Service Handlers │ +│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │ +│ │ Agent │ │ KB │ │ LLM │ │ File │ ... │ +│ └─────────┘ └─────────┘ └─────────┘ └─────────┘ │ +│ │ │ +│ └── Update token usage via request_id │ +└─────────────────────────────────────────────────────────────┘ +``` + +## Storage Strategy + +### Two-Layer Storage + +| Layer | Storage | Purpose | TTL | +| -------------- | -------- | ---------------------------- | --------- | +| **Real-time** | KV/Redis | Rate limiting, quota, status | 1h - 7d | +| **Persistent** | SQL | Archive, billing, audit | Permanent | + +### Why Hybrid? + +| Scenario | KV (Redis) | SQL | +| ----------------- | ----------------- | ---------------- | +| Rate limit check | ⚡ < 1ms | ❌ Too slow | +| Quota check | ⚡ < 1ms | ❌ Too slow | +| Request status | ⚡ Fast update | ❌ Too slow | +| Billing report | ❌ No aggregation | ✅ SUM/GROUP BY | +| Audit query | ❌ No persistence | ✅ Full history | +| Complex filtering | ❌ Key-only | ✅ WHERE clauses | + +### KV Keys Design + +``` +# Rate Limiting (TTL: 60s) +ratelimit:user:{user_id}:{service} → count +ratelimit:team:{team_id}:{service} → count +ratelimit:ip:{ip} → count + +# Request Status (TTL: 1h) +request:{request_id} → {status, service, created_at, ...} + +# Token Usage - Daily (TTL: 7d) +tokens:user:{user_id}:{YYYY-MM-DD} → {input, output, total} +tokens:team:{team_id}:{YYYY-MM-DD} → {input, output, total} + +# Quota (TTL: 24h for daily, 30d for monthly) +quota:user:{user_id}:daily → remaining_tokens +quota:team:{team_id}:monthly → remaining_tokens +``` + +### Data Flow + +``` +Request arrives + │ + ├── 1. KV: Rate limit check + │ INCR ratelimit:user:{id}:{service} + │ if > limit → 429 Too Many Requests + │ + ├── 2. KV: Quota check + │ GET quota:user:{id}:daily + │ if <= 0 → 429 Quota Exceeded + │ + ├── 3. KV: Record request status + │ SET request:{id} {status: "running", ...} EX 3600 + │ + ├── 4. Execute request... + │ + ├── 5. KV: Update tokens + │ HINCRBY tokens:user:{id}:{date} input {n} + │ HINCRBY tokens:user:{id}:{date} output {n} + │ DECRBY quota:user:{id}:daily {total} + │ + ├── 6. KV: Update request status + │ SET request:{id} {status: "completed", duration_ms: ...} + │ + └── 7. Async: Archive to SQL + INSERT INTO openapi_request ... + +``` + +## Data Model + +### KV Data Structures + +#### Rate Limit Counter + +```go +// Key: ratelimit:{type}:{id}:{service} +// Value: integer count +// TTL: 60 seconds (sliding window) + +type RateLimitKey struct { + Type string // "user", "team", "ip" + ID string // user_id, team_id, or IP + Service string // "agent", "kb", "llm", etc. +} + +func (k RateLimitKey) String() string { + return fmt.Sprintf("ratelimit:%s:%s:%s", k.Type, k.ID, k.Service) +} +``` + +#### Request Status + +```go +// Key: request:{request_id} +// Value: JSON object +// TTL: 1 hour + +type RequestStatus struct { + RequestID string `json:"request_id"` + UserID string `json:"user_id"` + TeamID string `json:"team_id,omitempty"` + Service string `json:"service"` + ResourceID string `json:"resource_id,omitempty"` + Status string `json:"status"` // running, completed, failed + CreatedAt time.Time `json:"created_at"` + CompletedAt time.Time `json:"completed_at,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` + Error string `json:"error,omitempty"` +} +``` + +#### Token Usage (Daily) + +```go +// Key: tokens:{type}:{id}:{date} +// Value: Hash {input, output, total} +// TTL: 7 days + +type TokenUsage struct { + Input int64 `json:"input"` + Output int64 `json:"output"` + Total int64 `json:"total"` +} +``` + +#### Quota + +```go +// Key: quota:{type}:{id}:{period} +// Value: remaining tokens (integer) +// TTL: 24h (daily) or 30d (monthly) + +type QuotaKey struct { + Type string // "user", "team" + ID string + Period string // "daily", "monthly" +} +``` + +### SQL Table (Archive) + +**Table Name:** `openapi_request` + +**Purpose:** Long-term storage for billing reports, audit logs, and analytics. + +| Column | Type | Nullable | Index | Description | +| --------------- | ----------- | -------- | ------ | ------------------------------------------------ | +| `id` | ID | No | PK | Auto-increment primary key | +| `request_id` | string(64) | No | Unique | Unique request identifier | +| `user_id` | string(200) | No | Yes | User ID from auth | +| `team_id` | string(200) | Yes | Yes | Team ID from auth | +| `session_id` | string(200) | Yes | Yes | Session ID | +| `endpoint` | string(200) | No | Yes | API endpoint path | +| `method` | string(10) | No | - | HTTP method (GET, POST, etc.) | +| `service` | string(50) | No | Yes | Service type: `agent`, `kb`, `llm`, `file`, etc. | +| `resource_id` | string(200) | Yes | Yes | Resource ID (assistant_id, collection_id, etc.) | +| `status` | enum | No | Yes | `pending`, `running`, `completed`, `failed` | +| `status_code` | integer | Yes | - | HTTP response status code | +| `referer` | string(50) | Yes | - | Request source (api, jssdk, agent, etc.) | +| `client_type` | string(50) | Yes | - | Client type (web, ios, android, etc.) | +| `client_ip` | string(50) | Yes | Yes | Client IP address | +| `input_tokens` | integer | Yes | - | Input token count (LLM calls) | +| `output_tokens` | integer | Yes | - | Output token count (LLM calls) | +| `total_tokens` | integer | Yes | Yes | Total token count | +| `duration_ms` | integer | Yes | Yes | Request duration in milliseconds | +| `error` | text | Yes | - | Error message if failed | +| `metadata` | json | Yes | - | Additional metadata | +| `created_at` | timestamp | No | Yes | Request start time | +| `completed_at` | timestamp | Yes | Yes | Request completion time | + +**Indexes:** + +| Name | Columns | Type | Purpose | +| ------------------ | --------------------------------------- | ----- | ------------------------ | +| `idx_req_user` | `user_id`, `created_at` | index | User request history | +| `idx_req_team` | `team_id`, `created_at` | index | Team request history | +| `idx_req_endpoint` | `endpoint`, `created_at` | index | Endpoint analytics | +| `idx_req_service` | `service`, `created_at` | index | Service analytics | +| `idx_req_status` | `status` | index | Find incomplete requests | +| `idx_req_billing` | `team_id`, `created_at`, `total_tokens` | index | Billing queries | +| `idx_req_ip` | `client_ip`, `created_at` | index | IP-based rate limiting | + +### Service Types + +| Service | Description | Resource ID Example | +| ------- | -------------------- | ------------------- | +| `agent` | Chat/Agent API | `assistant_id` | +| `kb` | Knowledge Base API | `collection_id` | +| `llm` | Direct LLM API | `connector_id` | +| `file` | File upload/download | `file_id` | +| `user` | User management | `user_id` | +| `team` | Team management | `team_id` | +| `mcp` | MCP server calls | `server_id` | + +### Status Values + +| Status | Description | Set By | +| ----------- | ------------------------------ | ---------- | +| `pending` | Request received, not started | Middleware | +| `running` | Request being processed | Middleware | +| `completed` | Request completed successfully | Middleware | +| `failed` | Request failed with error | Middleware | + +## Middleware Design + +### Modular Middleware Architecture + +Each middleware is independent and can be composed based on business needs. + +``` +┌─────────────────────────────────────────────────────────────┐ +│ Available Middlewares │ +├─────────────────────────────────────────────────────────────┤ +│ │ +│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ +│ │ RequestID │ │ RateLimit │ │ Quota │ │ +│ │ (Basic) │ │ (Protect) │ │ (Billing) │ │ +│ └─────────────┘ └─────────────┘ └─────────────┘ │ +│ │ +│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ +│ │ Metrics │ │ Archive │ │ Billing │ │ +│ │ (Monitor) │ │ (Audit) │ │ (Charge) │ │ +│ └─────────────┘ └─────────────┘ └─────────────┘ │ +│ │ +└─────────────────────────────────────────────────────────────┘ +``` + +### Middleware List + +| Middleware | File | Purpose | Dependencies | +| ----------- | --------------- | -------------------------------- | ------------------ | +| `RequestID` | `request_id.go` | Generate and track request ID | None | +| `RateLimit` | `ratelimit.go` | Request frequency limiting | KV, RequestID | +| `Quota` | `quota.go` | Token quota enforcement | KV, RequestID | +| `Metrics` | `metrics.go` | Request duration, status metrics | RequestID | +| `Archive` | `archive.go` | Persist request to SQL | SQL, RequestID | +| `Billing` | `billing.go` | Token usage tracking & charging | KV, SQL, RequestID | + +### Usage Examples + +#### Example 1: Full Protection (Agent API) + +```go +// Agent API needs all protections +agent := api.Group("/chat") +agent.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting + request.Quota(kv, config), // Token quota + request.Metrics(), // Duration tracking + request.Archive(sql), // Audit logging + request.Billing(kv, sql), // Token billing +) +agent.POST("/completions", handler.ChatCompletions) +``` + +#### Example 2: Light Protection (File API) + +```go +// File API only needs basic tracking +file := api.Group("/file") +file.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting + request.Metrics(), // Duration tracking +) +file.POST("/upload", handler.Upload) +``` + +#### Example 3: Internal API (No Billing) + +```go +// Internal API skips billing +internal := api.Group("/internal") +internal.Use( + request.RequestID(), // Generate request_id + request.Metrics(), // Duration tracking + request.Archive(sql), // Audit logging only +) +internal.GET("/health", handler.Health) +``` + +#### Example 4: Public API (Rate Limit Only) + +```go +// Public endpoints only need rate limiting +public := api.Group("/public") +public.Use( + request.RequestID(), // Generate request_id + request.RateLimit(kv, config), // Rate limiting by IP +) +public.GET("/models", handler.ListModels) +``` + +--- + +### Middleware Implementations + +#### 1. RequestID Middleware (Base) + +```go +// request_id.go +package request + +// RequestID generates and sets request ID +func RequestID() gin.HandlerFunc { + return func(c *gin.Context) { + requestID := generateRequestID() + c.Set("request_id", requestID) + c.Header("X-Request-ID", requestID) + + // Also set start time for metrics + c.Set("request_start_time", time.Now()) + + // Detect and set service info + service := detectService(c.FullPath()) + c.Set("request_service", service) + c.Set("request_resource_id", extractResourceID(c, service)) + + c.Next() + } +} + +func generateRequestID() string { + return fmt.Sprintf("req_%s", nanoid.New()) +} + +func detectService(endpoint string) string { + switch { + case strings.HasPrefix(endpoint, "/api/chat"): + return ServiceAgent + case strings.HasPrefix(endpoint, "/api/agent"): + return ServiceAgent + case strings.HasPrefix(endpoint, "/api/kb"): + return ServiceKB + case strings.HasPrefix(endpoint, "/api/llm"): + return ServiceLLM + case strings.HasPrefix(endpoint, "/api/file"): + return ServiceFile + case strings.HasPrefix(endpoint, "/api/user"): + return ServiceUser + case strings.HasPrefix(endpoint, "/api/team"): + return ServiceTeam + case strings.HasPrefix(endpoint, "/api/mcp"): + return ServiceMCP + default: + return ServiceOther + } +} +``` + +#### 2. RateLimit Middleware + +```go +// ratelimit.go +package request + +// RateLimit enforces request frequency limits +func RateLimit(kv KVStore, config *RateLimitConfig) gin.HandlerFunc { + return func(c *gin.Context) { + if config == nil || !config.Enabled { + c.Next() + return + } + + authInfo := authorized.GetInfo(c) + service := c.GetString("request_service") + + // Check user rate limit + userKey := fmt.Sprintf("ratelimit:user:%s:%s", authInfo.UserID, service) + userCount, _ := kv.Incr(userKey, 60*time.Second) + if userCount > int64(config.GetUserLimit(service)) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": fmt.Sprintf("User rate limit exceeded: %d requests per minute", config.GetUserLimit(service)), + "retry_after": 60, + }) + return + } + + // Check team rate limit + if authInfo.TeamID != "" { + teamKey := fmt.Sprintf("ratelimit:team:%s:%s", authInfo.TeamID, service) + teamCount, _ := kv.Incr(teamKey, 60*time.Second) + if teamCount > int64(config.GetTeamLimit(service)) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": "Team rate limit exceeded", + "retry_after": 60, + }) + return + } + } + + // Check IP rate limit + ipKey := fmt.Sprintf("ratelimit:ip:%s", c.ClientIP()) + ipCount, _ := kv.Incr(ipKey, 60*time.Second) + if ipCount > int64(config.GetIPLimit()) { + c.AbortWithStatusJSON(429, gin.H{ + "error": "rate_limit_exceeded", + "message": "IP rate limit exceeded", + "retry_after": 60, + }) + return + } + + c.Next() + } +} +``` + +#### 3. Quota Middleware + +```go +// quota.go +package request + +// Quota enforces token quota limits +func Quota(kv KVStore, config *QuotaConfig) gin.HandlerFunc { + return func(c *gin.Context) { + if config == nil || !config.Enabled { + c.Next() + return + } + + authInfo := authorized.GetInfo(c) + + // Check user daily quota + userQuotaKey := fmt.Sprintf("quota:user:%s:daily", authInfo.UserID) + remaining, exists := kv.Get(userQuotaKey) + + if !exists { + // Initialize quota for the day + limit := config.GetUserDailyLimit(authInfo.UserID) + kv.Set(userQuotaKey, limit, 24*time.Hour) + remaining = limit + } + + if remaining <= 0 { + c.AbortWithStatusJSON(429, gin.H{ + "error": "quota_exceeded", + "message": "Daily token quota exceeded", + "reset_at": getNextDayStart(), + }) + return + } + + // Check team monthly quota + if authInfo.TeamID != "" { + teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", authInfo.TeamID) + teamRemaining, exists := kv.Get(teamQuotaKey) + + if !exists { + limit := config.GetTeamMonthlyLimit(authInfo.TeamID) + kv.Set(teamQuotaKey, limit, 30*24*time.Hour) + teamRemaining = limit + } + + if teamRemaining <= 0 { + c.AbortWithStatusJSON(429, gin.H{ + "error": "quota_exceeded", + "message": "Team monthly token quota exceeded", + "reset_at": getNextMonthStart(), + }) + return + } + } + + c.Next() + } +} +``` + +#### 4. Metrics Middleware + +```go +// metrics.go +package request + +// Metrics tracks request duration and status +func Metrics() gin.HandlerFunc { + return func(c *gin.Context) { + startTime := c.GetTime("request_start_time") + if startTime.IsZero() { + startTime = time.Now() + } + + c.Next() + + // Calculate duration + duration := time.Since(startTime) + c.Set("request_duration_ms", duration.Milliseconds()) + + // Determine status + status := "completed" + if c.Writer.Status() >= 400 { + status = "failed" + } + c.Set("request_status", status) + + // TODO: Export to Prometheus/metrics system + // metrics.RequestDuration.WithLabelValues(service, status).Observe(duration.Seconds()) + // metrics.RequestTotal.WithLabelValues(service, status).Inc() + } +} +``` + +#### 5. Archive Middleware + +```go +// archive.go +package request + +// Archive persists request to SQL for audit +func Archive(sql SQLStore) gin.HandlerFunc { + return func(c *gin.Context) { + c.Next() + + // Get request info from context + requestID := c.GetString("request_id") + if requestID == "" { + return + } + + authInfo := authorized.GetInfo(c) + startTime := c.GetTime("request_start_time") + durationMs := c.GetInt64("request_duration_ms") + status := c.GetString("request_status") + if status == "" { + status = "completed" + } + + completedAt := time.Now() + + // Async archive to SQL + go func() { + sql.Archive(&Request{ + RequestID: requestID, + UserID: authInfo.UserID, + TeamID: authInfo.TeamID, + SessionID: authInfo.SessionID, + Endpoint: c.FullPath(), + Method: c.Request.Method, + Service: c.GetString("request_service"), + ResourceID: c.GetString("request_resource_id"), + Status: status, + StatusCode: c.Writer.Status(), + Referer: c.GetHeader("X-Yao-Referer"), + ClientType: getClientType(c.GetHeader("User-Agent")), + ClientIP: c.ClientIP(), + DurationMs: durationMs, + Error: c.GetString("request_error"), + CreatedAt: startTime, + CompletedAt: &completedAt, + }) + }() + } +} +``` + +#### 6. Billing Middleware + +```go +// billing.go +package request + +// Billing tracks token usage (called by services after completion) +func Billing(kv KVStore, sql SQLStore) gin.HandlerFunc { + return func(c *gin.Context) { + c.Next() + + // Token usage is updated by services via UpdateTokenUsage() + // This middleware just ensures the billing context is available + c.Set("billing_kv", kv) + c.Set("billing_sql", sql) + } +} + +// UpdateTokenUsage is called by services after completion +func UpdateTokenUsage(c *gin.Context, input, output int) error { + kv, ok := c.Get("billing_kv") + if !ok { + return nil // Billing not enabled + } + + sql, _ := c.Get("billing_sql") + requestID := c.GetString("request_id") + authInfo := authorized.GetInfo(c) + + return updateTokenUsageInternal( + kv.(KVStore), + sql.(SQLStore), + requestID, + authInfo.UserID, + authInfo.TeamID, + input, + output, + ) +} +``` + +## Rate Limiting + +### Configuration + +```yaml +# openapi.yml +rate_limit: + enabled: true + + # Default limits (requests per minute) + default: + per_user: 60 + per_team: 300 + per_ip: 100 + + # Service-specific limits + services: + agent: + per_user: 30 + per_team: 150 + llm: + per_user: 20 + per_team: 100 + kb: + per_user: 60 + per_team: 300 + + # Token limits (per day) + tokens: + per_user: 100000 + per_team: 1000000 + +# quota configuration +quota: + enabled: true + + # Default quotas + default: + user_daily: 100000 # tokens per day + team_monthly: 10000000 # tokens per month + + + # Can be overridden per user/team in database +``` + +### Rate Limit Check (KV-based) + +```go +func checkRateLimit(kv KVStore, authInfo *types.AuthorizedInfo, service, clientIP string) error { + config := GetRateLimitConfig() + if !config.Enabled { + return nil + } + + // 1. Check per-user limit (INCR with TTL) + userKey := fmt.Sprintf("ratelimit:user:%s:%s", authInfo.UserID, service) + userCount, _ := kv.Incr(userKey, 60*time.Second) // TTL 60s + if userCount > int64(config.GetUserLimit(service)) { + return fmt.Errorf("user rate limit exceeded: %d requests per minute", config.GetUserLimit(service)) + } + + // 2. Check per-team limit + if authInfo.TeamID != "" { + teamKey := fmt.Sprintf("ratelimit:team:%s:%s", authInfo.TeamID, service) + teamCount, _ := kv.Incr(teamKey, 60*time.Second) + if teamCount > int64(config.GetTeamLimit(service)) { + return fmt.Errorf("team rate limit exceeded") + } + } + + // 3. Check per-IP limit + ipKey := fmt.Sprintf("ratelimit:ip:%s", clientIP) + ipCount, _ := kv.Incr(ipKey, 60*time.Second) + if ipCount > int64(config.GetIPLimit()) { + return fmt.Errorf("IP rate limit exceeded") + } + + return nil +} +``` + +### Quota Check (KV-based) + +```go +func checkQuota(kv KVStore, authInfo *types.AuthorizedInfo) error { + config := GetQuotaConfig() + if !config.Enabled { + return nil + } + + // Check user daily quota + userQuotaKey := fmt.Sprintf("quota:user:%s:daily", authInfo.UserID) + remaining, exists := kv.Get(userQuotaKey) + + if !exists { + // Initialize quota for the day + limit := config.GetUserDailyLimit(authInfo.UserID) + kv.Set(userQuotaKey, limit, 24*time.Hour) + remaining = limit + } + + if remaining <= 0 { + return fmt.Errorf("daily token quota exceeded") + } + + // Check team monthly quota if applicable + if authInfo.TeamID != "" { + teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", authInfo.TeamID) + teamRemaining, exists := kv.Get(teamQuotaKey) + + if !exists { + limit := config.GetTeamMonthlyLimit(authInfo.TeamID) + kv.Set(teamQuotaKey, limit, 30*24*time.Hour) + teamRemaining = limit + } + + if teamRemaining <= 0 { + return fmt.Errorf("team monthly token quota exceeded") + } + } + + return nil +} +``` + +## Billing Integration + +### Token Usage Update + +Services update token usage after completion. This updates both KV (real-time) and SQL (archive). + +```go +// Called by Agent/LLM services after completion +func UpdateTokenUsage(kv KVStore, sql SQLStore, requestID string, userID, teamID string, input, output int) error { + total := input + output + date := time.Now().Format("2006-01-02") + + // 1. KV: Update daily token usage + userTokenKey := fmt.Sprintf("tokens:user:%s:%s", userID, date) + kv.HIncrBy(userTokenKey, "input", int64(input)) + kv.HIncrBy(userTokenKey, "output", int64(output)) + kv.HIncrBy(userTokenKey, "total", int64(total)) + kv.Expire(userTokenKey, 7*24*time.Hour) // Keep for 7 days + + if teamID != "" { + teamTokenKey := fmt.Sprintf("tokens:team:%s:%s", teamID, date) + kv.HIncrBy(teamTokenKey, "input", int64(input)) + kv.HIncrBy(teamTokenKey, "output", int64(output)) + kv.HIncrBy(teamTokenKey, "total", int64(total)) + kv.Expire(teamTokenKey, 7*24*time.Hour) + } + + // 2. KV: Deduct from quota + userQuotaKey := fmt.Sprintf("quota:user:%s:daily", userID) + kv.DecrBy(userQuotaKey, int64(total)) + + if teamID != "" { + teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", teamID) + kv.DecrBy(teamQuotaKey, int64(total)) + } + + // 3. SQL: Update request record (async) + go sql.UpdateTokens(requestID, input, output) + + return nil +} +``` + +### Billing Queries + +```sql +-- Daily token usage by team +SELECT + DATE(created_at) as date, + team_id, + service, + SUM(total_tokens) as tokens, + COUNT(*) as requests +FROM openapi_request +WHERE team_id = ? + AND created_at >= ? AND created_at < ? + AND status = 'completed' +GROUP BY DATE(created_at), team_id, service + +-- Monthly billing summary +SELECT + team_id, + service, + SUM(total_tokens) as total_tokens, + SUM(input_tokens) as input_tokens, + SUM(output_tokens) as output_tokens, + COUNT(*) as request_count, + AVG(duration_ms) as avg_duration +FROM openapi_request +WHERE created_at >= ? AND created_at < ? + AND status = 'completed' +GROUP BY team_id, service + +-- User quota check +SELECT SUM(total_tokens) as used +FROM openapi_request +WHERE user_id = ? + AND created_at >= CURDATE() + AND status = 'completed' +``` + +## API Interface + +### KV Store Interface + +```go +// KVStore defines the KV storage interface for real-time operations +type KVStore interface { + // Basic operations + Get(key string) (int64, bool) + Set(key string, value int64, ttl time.Duration) error + Incr(key string, ttl time.Duration) (int64, error) + DecrBy(key string, delta int64) (int64, error) + Expire(key string, ttl time.Duration) error + Del(key string) error + + // Hash operations (for token usage) + HGet(key, field string) (int64, error) + HSet(key, field string, value int64) error + HIncrBy(key, field string, delta int64) (int64, error) + HGetAll(key string) (map[string]int64, error) + + // Request status (JSON) + SetRequestStatus(requestID string, status *RequestStatus, ttl time.Duration) error + GetRequestStatus(requestID string) (*RequestStatus, error) +} +``` + +### SQL Store Interface + +```go +// SQLStore defines the SQL storage interface for archiving and analytics +type SQLStore interface { + // Archive stores a completed request + Archive(req *Request) error + + // UpdateTokens updates token usage for a request + UpdateTokens(requestID string, input, output int) error + + // Get retrieves a request by ID + Get(requestID string) (*Request, error) + + // List lists requests with filters + List(filter *RequestFilter) (*RequestList, error) + + // GetUsage gets usage statistics + GetUsage(filter *UsageFilter) (*UsageStats, error) +} +``` + +### Data Structures + +```go +// Request represents an API request record +type Request struct { + RequestID string `json:"request_id"` + UserID string `json:"user_id"` + TeamID string `json:"team_id,omitempty"` + SessionID string `json:"session_id,omitempty"` + Endpoint string `json:"endpoint"` + Method string `json:"method"` + Service string `json:"service"` + ResourceID string `json:"resource_id,omitempty"` + Status Status `json:"status"` + StatusCode int `json:"status_code,omitempty"` + Referer string `json:"referer,omitempty"` + ClientType string `json:"client_type,omitempty"` + ClientIP string `json:"client_ip,omitempty"` + InputTokens int `json:"input_tokens,omitempty"` + OutputTokens int `json:"output_tokens,omitempty"` + TotalTokens int `json:"total_tokens,omitempty"` + DurationMs int64 `json:"duration_ms,omitempty"` + Error string `json:"error,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + CreatedAt time.Time `json:"created_at"` + CompletedAt *time.Time `json:"completed_at,omitempty"` +} + +// CompletionInfo contains info for completing a request +type CompletionInfo struct { + StatusCode int + DurationMs int64 + Error string +} + +// RequestFilter for listing requests +type RequestFilter struct { + UserID string `json:"user_id,omitempty"` + TeamID string `json:"team_id,omitempty"` + Service string `json:"service,omitempty"` + Status Status `json:"status,omitempty"` + StartTime time.Time `json:"start_time,omitempty"` + EndTime time.Time `json:"end_time,omitempty"` + Page int `json:"page,omitempty"` + PageSize int `json:"pagesize,omitempty"` +} + +// UsageFilter for usage statistics +type UsageFilter struct { + UserID string `json:"user_id,omitempty"` + TeamID string `json:"team_id,omitempty"` + Service string `json:"service,omitempty"` + StartTime time.Time `json:"start_time"` + EndTime time.Time `json:"end_time"` + GroupBy string `json:"group_by,omitempty"` // day, week, month +} + +// UsageStats contains usage statistics +type UsageStats struct { + TotalRequests int64 `json:"total_requests"` + TotalTokens int64 `json:"total_tokens"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + AvgDurationMs float64 `json:"avg_duration_ms"` + ByService map[string]int64 `json:"by_service,omitempty"` + ByDay []DailyUsage `json:"by_day,omitempty"` +} + +type DailyUsage struct { + Date string `json:"date"` + Requests int64 `json:"requests"` + Tokens int64 `json:"tokens"` +} +``` + +## Integration with Services + +### Route Registration Example + +```go +// openapi/openapi.go +func (s *OpenAPI) RegisterRoutes(r *gin.Engine) { + api := r.Group("/api") + + // 1. OAuth Guard (authentication) - for all routes + api.Use(oauth.Guard) + + // 2. Register different route groups with different middleware combinations + s.registerAgentRoutes(api) + s.registerKBRoutes(api) + s.registerLLMRoutes(api) + s.registerFileRoutes(api) + s.registerPublicRoutes(api) +} + +func (s *OpenAPI) registerAgentRoutes(api *gin.RouterGroup) { + // Agent API: Full protection + billing + agent := api.Group("/chat") + agent.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Quota(s.kv, s.quotaConfig), + request.Metrics(), + request.Archive(s.sql), + request.Billing(s.kv, s.sql), + ) + agent.POST("/completions", s.handler.ChatCompletions) +} + +func (s *OpenAPI) registerKBRoutes(api *gin.RouterGroup) { + // KB API: Rate limit + archive (no token billing) + kb := api.Group("/kb") + kb.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Metrics(), + request.Archive(s.sql), + ) + kb.POST("/search", s.handler.KBSearch) + kb.POST("/upload", s.handler.KBUpload) +} + +func (s *OpenAPI) registerFileRoutes(api *gin.RouterGroup) { + // File API: Light protection + file := api.Group("/file") + file.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), + request.Metrics(), + ) + file.POST("/upload", s.handler.FileUpload) + file.GET("/download/:id", s.handler.FileDownload) +} + +func (s *OpenAPI) registerPublicRoutes(api *gin.RouterGroup) { + // Public API: Rate limit only (no auth required) + public := api.Group("/public") + public.Use( + request.RequestID(), + request.RateLimit(s.kv, s.rateLimitConfig), // IP-based only + ) + public.GET("/models", s.handler.ListModels) + public.GET("/health", s.handler.Health) +} +``` + +### Agent Service Integration + +```go +// agent/context/openapi.go +func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest, *Context, *Options, error) { + // Get request ID from middleware + requestID := c.GetString("request_id") + + // Create context with request ID + ctx := New(c.Request.Context(), authInfo, chatID) + ctx.RequestID = requestID // Use global request_id + ctx.GinContext = c // Keep gin context for billing + + // ... +} + +// agent/assistant/agent.go +func (ast *Assistant) Stream(ctx, inputMessages, options) { + defer func() { + // Update token usage via billing middleware + if ctx.GinContext != nil && completionResponse != nil && completionResponse.Usage != nil { + request.UpdateTokenUsage( + ctx.GinContext, + completionResponse.Usage.PromptTokens, + completionResponse.Usage.CompletionTokens, + ) + } + }() + + // ... +} +``` + +### LLM Service Integration + +```go +// llm/api/completion.go +func (api *API) Completion(c *gin.Context) { + // ... execute LLM call ... + + // Update token usage + if response.Usage != nil { + request.UpdateTokenUsage(c, response.Usage.PromptTokens, response.Usage.CompletionTokens) + } +} +``` + +## Summary + +### Middleware Components + +| Middleware | File | Purpose | Storage | +| ----------- | --------------- | ------------------------ | -------- | +| `RequestID` | `request_id.go` | Generate request ID | - | +| `RateLimit` | `ratelimit.go` | Frequency limiting | KV | +| `Quota` | `quota.go` | Token quota enforcement | KV | +| `Metrics` | `metrics.go` | Duration/status tracking | - | +| `Archive` | `archive.go` | Persist to SQL | SQL | +| `Billing` | `billing.go` | Token usage tracking | KV + SQL | + +### Storage Components + +| Component | File | Purpose | +| --------- | ---------- | ---------------------------- | +| KV Store | `kv.go` | Real-time: rate limit, quota | +| SQL Store | `sql.go` | Archive: billing, audit | +| Types | `types.go` | Data structures | + +### Middleware Combinations by Use Case + +| Use Case | RequestID | RateLimit | Quota | Metrics | Archive | Billing | +| ------------ | --------- | --------- | ----- | ------- | ------- | ------- | +| Agent API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| LLM API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| KB API | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | +| File API | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | +| Public API | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | +| Internal API | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | + +### Key Points + +1. **Modular design**: Each middleware is independent and composable +2. **Business-driven composition**: Routes choose which middleware to use +3. **Two-layer storage**: KV for real-time, SQL for archive +4. **KV operations are synchronous**: Rate limit and quota checks must be fast +5. **SQL writes are async**: Archive happens in background goroutine +6. **Services update tokens via gin context**: `request.UpdateTokenUsage(c, input, output)` +7. **KV data has TTL**: Auto-expires to prevent memory bloat +8. **SQL data is permanent**: For billing and compliance diff --git a/openapi/tests/kb/collection_process_test.go b/openapi/tests/kb/collection_process_test.go new file mode 100644 index 00000000..56f3f497 --- /dev/null +++ b/openapi/tests/kb/collection_process_test.go @@ -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") + }) +}